diff --git a/agentplatform/_genai/sandboxes.py b/agentplatform/_genai/sandboxes.py index 406ce30164..bce278816b 100644 --- a/agentplatform/_genai/sandboxes.py +++ b/agentplatform/_genai/sandboxes.py @@ -26,7 +26,6 @@ from urllib.parse import urlencode from google import auth as google_auth -from google import genai from google.auth.transport import requests as google_auth_requests from google.genai import _api_module from google.genai import _common @@ -34,6 +33,7 @@ from google.genai._common import get_value_by_path as getv from google.genai._common import set_value_by_path as setv from google.genai.pagers import Pager +import requests from . import _runtimes_utils from . import types @@ -1408,15 +1408,27 @@ def send_command( headers["Authorization"] = f"Bearer {access_token}" headers["X-Sandbox-Routing-Token"] = routing_token headers["X-Sandbox-Port"] = port - endpoint = endpoint + path if path.startswith("/") else endpoint + "/" + path - http_options = genai_types.HttpOptions(headers=headers, base_url=endpoint) - http_client = genai.Client(vertexai=True, http_options=http_options) - # Full path is constructed in this function. The passed in path into request - # function will not be used. - response = http_client._api_client.request(http_method, path, request_dict) + url = endpoint + path if path.startswith("/") else endpoint + "/" + path + # The sandbox data plane authenticates via the Authorization, + # X-Sandbox-Routing-Token and X-Sandbox-Port headers set above, not via + # ADC. Routing this through genai.Client(vertexai=True) builds an + # authorized session with no credentials, whose request() then fails with + # "'NoneType' object has no attribute 'before_request'". Issue the request + # directly instead, mirroring generate_access_token above which also talks + # to a non-aiplatform endpoint with plain requests. + body = None + if request_dict: + body = json.dumps(request_dict) + headers.setdefault("Content-Type", "application/json") + response = requests.request( + http_method, + url, + headers=headers, + data=body, + ) return genai_types.HttpResponse( - headers=response.headers, - body=response.body, + headers=dict(response.headers), + body=response.text, ) def generate_browser_ws_headers( diff --git a/tests/unit/agentplatform/genai/test_sandbox.py b/tests/unit/agentplatform/genai/test_sandbox.py index ca8f5c73a1..c38ba9d549 100644 --- a/tests/unit/agentplatform/genai/test_sandbox.py +++ b/tests/unit/agentplatform/genai/test_sandbox.py @@ -29,11 +29,8 @@ from vertexai._genai import ( sandboxes as vertexai_sandboxes, ) -from google.genai import client -from google.genai import types as genai_types import pytest - _TEST_CREDENTIALS = mock.Mock(spec=auth_credentials.AnonymousCredentials()) _TEST_LOCATION = "us-central1" _TEST_PROJECT = "test-project" @@ -68,6 +65,7 @@ def google_auth_mock(): @pytest.mark.usefixtures("google_auth_mock") class TestSandbox: + def setup_method(self): importlib.reload(initializer) importlib.reload(aiplatform) @@ -82,18 +80,18 @@ def setup_method(self): def teardown_method(self): initializer.global_pool.shutdown(wait=True) - @mock.patch.object(client.Client, "_get_api_client") - def test_send_command(self, mock_get_api_client): + @mock.patch.object(sandboxes.requests, "request") + def test_send_command(self, mock_request): mock_sandbox = mock.Mock() mock_sandbox.connection_info.load_balancer_ip = None mock_sandbox.connection_info.load_balancer_hostname = ( "test-us-central1.example.vertexai.goog" ) mock_sandbox.connection_info.routing_token = "test_routing_token" - mock_http_client = mock_get_api_client.return_value - mock_http_client.request.return_value = genai_types.HttpResponse( - body=b"{}", headers={} - ) + mock_response = mock.Mock() + mock_response.text = "{}" + mock_response.headers = {} + mock_request.return_value = mock_response self.client.sandboxes.send_command( http_method="GET", @@ -102,21 +100,46 @@ def test_send_command(self, mock_get_api_client): path="test/path", ) - call_args = mock_get_api_client.call_args + # The sandbox data plane is called with requests, not the genai client. + call_args = mock_request.call_args assert call_args is not None - _, kwargs = call_args - http_options = kwargs["http_options"] - assert http_options.base_url == ( - "https://test-us-central1.example.vertexai.goog/test/path" + args, kwargs = call_args + assert args[0] == "GET" + assert args[1] == "https://test-us-central1.example.vertexai.goog/test/path" + assert kwargs["headers"]["Authorization"] == "Bearer test_token" + assert kwargs["headers"]["X-Sandbox-Routing-Token"] == "test_routing_token" + # A GET with no request body must not send a JSON payload. + assert kwargs["data"] is None + + @mock.patch.object(sandboxes.requests, "request") + def test_send_command_post_serializes_body(self, mock_request): + mock_sandbox = mock.Mock() + mock_sandbox.connection_info.load_balancer_ip = None + mock_sandbox.connection_info.load_balancer_hostname = ( + "test-us-central1.example.vertexai.goog" ) - assert http_options.headers["Authorization"] == "Bearer test_token" + mock_sandbox.connection_info.routing_token = "test_routing_token" + mock_response = mock.Mock() + mock_response.text = "{}" + mock_response.headers = {} + mock_request.return_value = mock_response - mock_http_client.request.assert_called_with("GET", "test/path", {}) + self.client.sandboxes.send_command( + http_method="POST", + access_token="test_token", + sandbox_environment=mock_sandbox, + path="cdp", + request_dict={"command": "Page.navigate"}, + ) + + _, kwargs = mock_request.call_args + assert kwargs["data"] == json.dumps({"command": "Page.navigate"}) + assert kwargs["headers"]["Content-Type"] == "application/json" @mock.patch.object(sandboxes.Sandboxes, "generate_access_token") - @mock.patch.object(client.Client, "_get_api_client") + @mock.patch.object(sandboxes.requests, "request") def test_generate_browser_ws_headers( - self, mock_get_api_client, mock_generate_access_token + self, mock_request, mock_generate_access_token ): mock_generate_access_token.return_value = "test_token" @@ -126,10 +149,10 @@ def test_generate_browser_ws_headers( "test-us-central1.example.vertexai.goog" ) mock_sandbox.connection_info.routing_token = "test_routing_token" - mock_http_client = mock_get_api_client.return_value - mock_http_client.request.return_value = genai_types.HttpResponse( - body=b'{"endpoint": "test/endpoint"}', headers={} - ) + mock_response = mock.Mock() + mock_response.text = '{"endpoint": "test/endpoint"}' + mock_response.headers = {} + mock_request.return_value = mock_response ws_url, headers = self.client.sandboxes.generate_browser_ws_headers( sandbox_environment=mock_sandbox, service_account_email=_TEST_SERVICE_ACCOUNT_EMAIL, @@ -150,7 +173,7 @@ def test_create_with_shell_environment_and_existing_template(self, mock_create): name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec={"shell_environment": {}}, config={ - "sandbox_environment_template": _TEST_SANDBOX_TEMPLATE_RESOURCE_NAME, + "sandbox_environment_template": (_TEST_SANDBOX_TEMPLATE_RESOURCE_NAME), "wait_for_completion": False, }, ) @@ -269,7 +292,9 @@ def test_create_with_snapshot_does_not_create_template( name=_TEST_AGENT_ENGINE_RESOURCE_NAME, spec={"computer_use_environment": {}}, config={ - "sandbox_environment_snapshot": "projects/p/locations/l/agentEngines/ae/sandboxEnvironmentSnapshots/s1", + "sandbox_environment_snapshot": ( + "projects/p/locations/l/agentEngines/ae/sandboxEnvironmentSnapshots/s1" + ), "wait_for_completion": False, }, ) diff --git a/vertexai/_genai/sandboxes.py b/vertexai/_genai/sandboxes.py index 0209d23078..df582e74df 100644 --- a/vertexai/_genai/sandboxes.py +++ b/vertexai/_genai/sandboxes.py @@ -25,7 +25,6 @@ from urllib.parse import urlencode from google import auth as google_auth -from google import genai from google.auth.transport import requests as google_auth_requests from google.genai import _api_module from google.genai import _common @@ -33,6 +32,7 @@ from google.genai._common import get_value_by_path as getv from google.genai._common import set_value_by_path as setv from google.genai.pagers import Pager +import requests from . import _agent_engines_utils from . import types @@ -983,15 +983,27 @@ def send_command( headers["Authorization"] = f"Bearer {access_token}" headers["X-Sandbox-Routing-Token"] = routing_token headers["X-Sandbox-Port"] = port - endpoint = endpoint + path if path.startswith("/") else endpoint + "/" + path - http_options = genai_types.HttpOptions(headers=headers, base_url=endpoint) - http_client = genai.Client(vertexai=True, http_options=http_options) - # Full path is constructed in this function. The passed in path into request - # function will not be used. - response = http_client._api_client.request(http_method, path, request_dict) + url = endpoint + path if path.startswith("/") else endpoint + "/" + path + # The sandbox data plane authenticates via the Authorization, + # X-Sandbox-Routing-Token and X-Sandbox-Port headers set above, not via + # ADC. Routing this through genai.Client(vertexai=True) builds an + # authorized session with no credentials, whose request() then fails with + # "'NoneType' object has no attribute 'before_request'". Issue the request + # directly instead, mirroring generate_access_token above which also talks + # to a non-aiplatform endpoint with plain requests. + body = None + if request_dict: + body = json.dumps(request_dict) + headers.setdefault("Content-Type", "application/json") + response = requests.request( + http_method, + url, + headers=headers, + data=body, + ) return genai_types.HttpResponse( - headers=response.headers, - body=response.body, + headers=dict(response.headers), + body=response.text, ) def generate_browser_ws_headers(