Skip to content

Commit 2b28caf

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
fix: Fix Sandboxes.send_command failing to reach the sandbox data plane.
`Sandboxes.send_command` built an inner `genai.Client(vertexai=True, ...)` and called `_api_client.request(...)` to talk to the sandbox data plane. That genai client's authorized session has no credentials object, so every call failed with `AttributeError: 'NoneType' object has no attribute 'before_request'`, making all Computer Use / sandbox HTTP calls (GET /, /tabs, POST /cdp, etc.) unusable. The sandbox data plane does not use ADC; it authenticates via the `Authorization: Bearer <jwt>`, `X-Sandbox-Routing-Token` and `X-Sandbox-Port` headers that `send_command` already sets. Routing through a genai client was therefore both unnecessary and broken. Issue the request directly with `requests` instead, mirroring `generate_access_token`, which already talks to a non-aiplatform endpoint the same way. A JSON `Content-Type` is set only when a request body is present, and a caller-supplied `Content-Type` is preserved. `generate_browser_ws_headers`, which calls `send_command` internally, continues to parse the JSON body correctly. PiperOrigin-RevId: 979899791
1 parent f1506ea commit 2b28caf

3 files changed

Lines changed: 91 additions & 42 deletions

File tree

‎agentplatform/_genai/sandboxes.py‎

Lines changed: 21 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -26,14 +26,14 @@
2626
from urllib.parse import urlencode
2727

2828
from google import auth as google_auth
29-
from google import genai
3029
from google.auth.transport import requests as google_auth_requests
3130
from google.genai import _api_module
3231
from google.genai import _common
3332
from google.genai import types as genai_types
3433
from google.genai._common import get_value_by_path as getv
3534
from google.genai._common import set_value_by_path as setv
3635
from google.genai.pagers import Pager
36+
import requests
3737

3838
from . import _runtimes_utils
3939
from . import types
@@ -1408,15 +1408,27 @@ def send_command(
14081408
headers["Authorization"] = f"Bearer {access_token}"
14091409
headers["X-Sandbox-Routing-Token"] = routing_token
14101410
headers["X-Sandbox-Port"] = port
1411-
endpoint = endpoint + path if path.startswith("/") else endpoint + "/" + path
1412-
http_options = genai_types.HttpOptions(headers=headers, base_url=endpoint)
1413-
http_client = genai.Client(vertexai=True, http_options=http_options)
1414-
# Full path is constructed in this function. The passed in path into request
1415-
# function will not be used.
1416-
response = http_client._api_client.request(http_method, path, request_dict)
1411+
url = endpoint + path if path.startswith("/") else endpoint + "/" + path
1412+
# The sandbox data plane authenticates via the Authorization,
1413+
# X-Sandbox-Routing-Token and X-Sandbox-Port headers set above, not via
1414+
# ADC. Routing this through genai.Client(vertexai=True) builds an
1415+
# authorized session with no credentials, whose request() then fails with
1416+
# "'NoneType' object has no attribute 'before_request'". Issue the request
1417+
# directly instead, mirroring generate_access_token above which also talks
1418+
# to a non-aiplatform endpoint with plain requests.
1419+
body = None
1420+
if request_dict:
1421+
body = json.dumps(request_dict)
1422+
headers.setdefault("Content-Type", "application/json")
1423+
response = requests.request(
1424+
http_method,
1425+
url,
1426+
headers=headers,
1427+
data=body,
1428+
)
14171429
return genai_types.HttpResponse(
1418-
headers=response.headers,
1419-
body=response.body,
1430+
headers=dict(response.headers),
1431+
body=response.text,
14201432
)
14211433

14221434
def generate_browser_ws_headers(

‎tests/unit/agentplatform/genai/test_sandbox.py‎

Lines changed: 49 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -29,11 +29,8 @@
2929
from vertexai._genai import (
3030
sandboxes as vertexai_sandboxes,
3131
)
32-
from google.genai import client
33-
from google.genai import types as genai_types
3432
import pytest
3533

36-
3734
_TEST_CREDENTIALS = mock.Mock(spec=auth_credentials.AnonymousCredentials())
3835
_TEST_LOCATION = "us-central1"
3936
_TEST_PROJECT = "test-project"
@@ -68,6 +65,7 @@ def google_auth_mock():
6865

6966
@pytest.mark.usefixtures("google_auth_mock")
7067
class TestSandbox:
68+
7169
def setup_method(self):
7270
importlib.reload(initializer)
7371
importlib.reload(aiplatform)
@@ -82,18 +80,18 @@ def setup_method(self):
8280
def teardown_method(self):
8381
initializer.global_pool.shutdown(wait=True)
8482

85-
@mock.patch.object(client.Client, "_get_api_client")
86-
def test_send_command(self, mock_get_api_client):
83+
@mock.patch.object(sandboxes.requests, "request")
84+
def test_send_command(self, mock_request):
8785
mock_sandbox = mock.Mock()
8886
mock_sandbox.connection_info.load_balancer_ip = None
8987
mock_sandbox.connection_info.load_balancer_hostname = (
9088
"test-us-central1.example.vertexai.goog"
9189
)
9290
mock_sandbox.connection_info.routing_token = "test_routing_token"
93-
mock_http_client = mock_get_api_client.return_value
94-
mock_http_client.request.return_value = genai_types.HttpResponse(
95-
body=b"{}", headers={}
96-
)
91+
mock_response = mock.Mock()
92+
mock_response.text = "{}"
93+
mock_response.headers = {}
94+
mock_request.return_value = mock_response
9795

9896
self.client.sandboxes.send_command(
9997
http_method="GET",
@@ -102,21 +100,46 @@ def test_send_command(self, mock_get_api_client):
102100
path="test/path",
103101
)
104102

105-
call_args = mock_get_api_client.call_args
103+
# The sandbox data plane is called with requests, not the genai client.
104+
call_args = mock_request.call_args
106105
assert call_args is not None
107-
_, kwargs = call_args
108-
http_options = kwargs["http_options"]
109-
assert http_options.base_url == (
110-
"https://test-us-central1.example.vertexai.goog/test/path"
106+
args, kwargs = call_args
107+
assert args[0] == "GET"
108+
assert args[1] == "https://test-us-central1.example.vertexai.goog/test/path"
109+
assert kwargs["headers"]["Authorization"] == "Bearer test_token"
110+
assert kwargs["headers"]["X-Sandbox-Routing-Token"] == "test_routing_token"
111+
# A GET with no request body must not send a JSON payload.
112+
assert kwargs["data"] is None
113+
114+
@mock.patch.object(sandboxes.requests, "request")
115+
def test_send_command_post_serializes_body(self, mock_request):
116+
mock_sandbox = mock.Mock()
117+
mock_sandbox.connection_info.load_balancer_ip = None
118+
mock_sandbox.connection_info.load_balancer_hostname = (
119+
"test-us-central1.example.vertexai.goog"
111120
)
112-
assert http_options.headers["Authorization"] == "Bearer test_token"
121+
mock_sandbox.connection_info.routing_token = "test_routing_token"
122+
mock_response = mock.Mock()
123+
mock_response.text = "{}"
124+
mock_response.headers = {}
125+
mock_request.return_value = mock_response
113126

114-
mock_http_client.request.assert_called_with("GET", "test/path", {})
127+
self.client.sandboxes.send_command(
128+
http_method="POST",
129+
access_token="test_token",
130+
sandbox_environment=mock_sandbox,
131+
path="cdp",
132+
request_dict={"command": "Page.navigate"},
133+
)
134+
135+
_, kwargs = mock_request.call_args
136+
assert kwargs["data"] == json.dumps({"command": "Page.navigate"})
137+
assert kwargs["headers"]["Content-Type"] == "application/json"
115138

116139
@mock.patch.object(sandboxes.Sandboxes, "generate_access_token")
117-
@mock.patch.object(client.Client, "_get_api_client")
140+
@mock.patch.object(sandboxes.requests, "request")
118141
def test_generate_browser_ws_headers(
119-
self, mock_get_api_client, mock_generate_access_token
142+
self, mock_request, mock_generate_access_token
120143
):
121144
mock_generate_access_token.return_value = "test_token"
122145

@@ -126,10 +149,10 @@ def test_generate_browser_ws_headers(
126149
"test-us-central1.example.vertexai.goog"
127150
)
128151
mock_sandbox.connection_info.routing_token = "test_routing_token"
129-
mock_http_client = mock_get_api_client.return_value
130-
mock_http_client.request.return_value = genai_types.HttpResponse(
131-
body=b'{"endpoint": "test/endpoint"}', headers={}
132-
)
152+
mock_response = mock.Mock()
153+
mock_response.text = '{"endpoint": "test/endpoint"}'
154+
mock_response.headers = {}
155+
mock_request.return_value = mock_response
133156
ws_url, headers = self.client.sandboxes.generate_browser_ws_headers(
134157
sandbox_environment=mock_sandbox,
135158
service_account_email=_TEST_SERVICE_ACCOUNT_EMAIL,
@@ -150,7 +173,7 @@ def test_create_with_shell_environment_and_existing_template(self, mock_create):
150173
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
151174
spec={"shell_environment": {}},
152175
config={
153-
"sandbox_environment_template": _TEST_SANDBOX_TEMPLATE_RESOURCE_NAME,
176+
"sandbox_environment_template": (_TEST_SANDBOX_TEMPLATE_RESOURCE_NAME),
154177
"wait_for_completion": False,
155178
},
156179
)
@@ -269,7 +292,9 @@ def test_create_with_snapshot_does_not_create_template(
269292
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
270293
spec={"computer_use_environment": {}},
271294
config={
272-
"sandbox_environment_snapshot": "projects/p/locations/l/agentEngines/ae/sandboxEnvironmentSnapshots/s1",
295+
"sandbox_environment_snapshot": (
296+
"projects/p/locations/l/agentEngines/ae/sandboxEnvironmentSnapshots/s1"
297+
),
273298
"wait_for_completion": False,
274299
},
275300
)

‎vertexai/_genai/sandboxes.py‎

Lines changed: 21 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -25,14 +25,14 @@
2525
from urllib.parse import urlencode
2626

2727
from google import auth as google_auth
28-
from google import genai
2928
from google.auth.transport import requests as google_auth_requests
3029
from google.genai import _api_module
3130
from google.genai import _common
3231
from google.genai import types as genai_types
3332
from google.genai._common import get_value_by_path as getv
3433
from google.genai._common import set_value_by_path as setv
3534
from google.genai.pagers import Pager
35+
import requests
3636

3737
from . import _agent_engines_utils
3838
from . import types
@@ -983,15 +983,27 @@ def send_command(
983983
headers["Authorization"] = f"Bearer {access_token}"
984984
headers["X-Sandbox-Routing-Token"] = routing_token
985985
headers["X-Sandbox-Port"] = port
986-
endpoint = endpoint + path if path.startswith("/") else endpoint + "/" + path
987-
http_options = genai_types.HttpOptions(headers=headers, base_url=endpoint)
988-
http_client = genai.Client(vertexai=True, http_options=http_options)
989-
# Full path is constructed in this function. The passed in path into request
990-
# function will not be used.
991-
response = http_client._api_client.request(http_method, path, request_dict)
986+
url = endpoint + path if path.startswith("/") else endpoint + "/" + path
987+
# The sandbox data plane authenticates via the Authorization,
988+
# X-Sandbox-Routing-Token and X-Sandbox-Port headers set above, not via
989+
# ADC. Routing this through genai.Client(vertexai=True) builds an
990+
# authorized session with no credentials, whose request() then fails with
991+
# "'NoneType' object has no attribute 'before_request'". Issue the request
992+
# directly instead, mirroring generate_access_token above which also talks
993+
# to a non-aiplatform endpoint with plain requests.
994+
body = None
995+
if request_dict:
996+
body = json.dumps(request_dict)
997+
headers.setdefault("Content-Type", "application/json")
998+
response = requests.request(
999+
http_method,
1000+
url,
1001+
headers=headers,
1002+
data=body,
1003+
)
9921004
return genai_types.HttpResponse(
993-
headers=response.headers,
994-
body=response.body,
1005+
headers=dict(response.headers),
1006+
body=response.text,
9951007
)
9961008

9971009
def generate_browser_ws_headers(

0 commit comments

Comments
 (0)