Skip to content

Commit fc7f592

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
feat: Add execute_bash to the Agent Engine sandbox SDK.
Adds automatic default template provisioning for shell_environment and computer_use_environment to sandboxes.create in the GenAI / Agent Platform SDK. PiperOrigin-RevId: 971764450
1 parent a31d2ba commit fc7f592

2 files changed

Lines changed: 202 additions & 0 deletions

File tree

‎agentplatform/_genai/sandboxes.py‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1168,6 +1168,85 @@ def execute_code(
11681168

11691169
return response
11701170

1171+
def execute_bash(
1172+
self,
1173+
*,
1174+
name: str,
1175+
command: str,
1176+
cwd: Optional[str] = None,
1177+
timeout: Optional[int] = None,
1178+
port: str = "8080",
1179+
config: Optional[types.ExecuteCodeRuntimeSandboxConfigOrDict] = None,
1180+
) -> dict[str, Any]:
1181+
"""Runs a bash command in a shell sandbox.
1182+
1183+
The sandbox must have been created with a shell environment, either by
1184+
passing `spec={"shell_environment": {}}` to `create()` or by referencing a
1185+
template whose `default_container_category` is
1186+
`DEFAULT_CONTAINER_CATEGORY_SHELL_SANDBOX`.
1187+
1188+
Unlike `send_command`, this authenticates with the caller's own credentials,
1189+
so it needs no service account and no signed JWT.
1190+
1191+
Args:
1192+
name (str):
1193+
Required. The name of the agent engine sandbox to run the command
1194+
in, of the form
1195+
projects/*/locations/*/reasoningEngines/*/sandboxEnvironments/*.
1196+
command (str):
1197+
Required. The bash command to run.
1198+
cwd (str):
1199+
Optional. The working directory to run the command in. Defaults to
1200+
the sandbox's workspace directory.
1201+
timeout (int):
1202+
Optional. Seconds to allow the command to run before the sandbox
1203+
kills it. Defaults to the sandbox's own limit.
1204+
port (str):
1205+
Optional. The port the shell server listens on. Defaults to "8080".
1206+
config (ExecuteCodeRuntimeSandboxConfigOrDict):
1207+
Optional. The configuration for the request.
1208+
1209+
Returns:
1210+
dict[str, Any]: The result of the command, with keys `stdout`,
1211+
`stderr`, `returncode` and `duration_ms`. A command that exits
1212+
non-zero is reported here rather than raised.
1213+
1214+
Raises:
1215+
ValueError: If the sandbox returns no output.
1216+
"""
1217+
request: dict[str, Any] = {"command": command}
1218+
if cwd is not None:
1219+
request["cwd"] = cwd
1220+
if timeout is not None:
1221+
request["timeout"] = timeout
1222+
1223+
# The sandbox proxy addresses the container by URI and port, and forwards
1224+
# the JSON chunk as the POST body.
1225+
response = self._execute_code(
1226+
name=name,
1227+
inputs=[
1228+
types.Chunk(
1229+
mime_type="application/x.sandbox-request-uri",
1230+
data=b"/exec",
1231+
),
1232+
types.Chunk(
1233+
mime_type="application/x.sandbox-request-port",
1234+
data=port.encode("utf-8"),
1235+
),
1236+
types.Chunk(
1237+
mime_type="application/json",
1238+
data=json.dumps(request).encode("utf-8"),
1239+
),
1240+
],
1241+
config=config,
1242+
)
1243+
1244+
for output in response.outputs or []:
1245+
if output.data:
1246+
return json.loads(output.data.decode("utf-8"))
1247+
1248+
raise ValueError(f"The sandbox {name} returned no output for the command.")
1249+
11711250
def get(
11721251
self,
11731252
*,

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

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -294,6 +294,129 @@ def test_create_without_spec_template_or_snapshot_raises(self, mock_create):
294294

295295
mock_create.assert_not_called()
296296

297+
@staticmethod
298+
def _exec_response(payload):
299+
"""Builds the response the sandbox proxy returns for POST /exec."""
300+
return agentplatform_types.ExecuteSandboxEnvironmentResponse(
301+
outputs=[
302+
agentplatform_types.Chunk(
303+
mime_type="application/json",
304+
data=json.dumps(payload).encode("utf-8"),
305+
)
306+
]
307+
)
308+
309+
@staticmethod
310+
def _sent_chunks(mock_execute_code):
311+
"""Returns the uri, port and JSON body sent to the sandbox."""
312+
_, kwargs = mock_execute_code.call_args
313+
chunks = {chunk.mime_type: chunk.data for chunk in kwargs["inputs"]}
314+
return (
315+
chunks["application/x.sandbox-request-uri"],
316+
chunks["application/x.sandbox-request-port"],
317+
json.loads(chunks["application/json"]),
318+
)
319+
320+
@mock.patch.object(sandboxes.Sandboxes, "_execute_code")
321+
def test_execute_bash_runs_command_and_parses_result(self, mock_execute_code):
322+
"""A command is sent to /exec and its JSON result is returned."""
323+
expected = {
324+
"stdout": "hello\n",
325+
"stderr": "",
326+
"returncode": 0,
327+
"duration_ms": 5,
328+
}
329+
mock_execute_code.return_value = self._exec_response(expected)
330+
331+
result = self.client.sandboxes.execute_bash(
332+
name=_TEST_SANDBOX_RESOURCE_NAME,
333+
command="echo hello",
334+
)
335+
336+
assert result == expected
337+
_, kwargs = mock_execute_code.call_args
338+
assert kwargs["name"] == _TEST_SANDBOX_RESOURCE_NAME
339+
uri, port, body = self._sent_chunks(mock_execute_code)
340+
assert uri == b"/exec"
341+
assert port == b"8080"
342+
# cwd and timeout are left to the sandbox's own defaults when unset.
343+
assert body == {"command": "echo hello"}
344+
345+
@mock.patch.object(sandboxes.Sandboxes, "_execute_code")
346+
def test_execute_bash_forwards_cwd_and_timeout(self, mock_execute_code):
347+
"""cwd and timeout reach the container when the caller sets them."""
348+
mock_execute_code.return_value = self._exec_response({"stdout": "/tmp\n"})
349+
350+
self.client.sandboxes.execute_bash(
351+
name=_TEST_SANDBOX_RESOURCE_NAME,
352+
command="pwd",
353+
cwd="/tmp",
354+
timeout=30,
355+
)
356+
357+
_, _, body = self._sent_chunks(mock_execute_code)
358+
assert body == {"command": "pwd", "cwd": "/tmp", "timeout": 30}
359+
360+
@mock.patch.object(sandboxes.Sandboxes, "_execute_code")
361+
def test_execute_bash_uses_custom_port(self, mock_execute_code):
362+
"""A caller-supplied port overrides the shell server default."""
363+
mock_execute_code.return_value = self._exec_response({"stdout": ""})
364+
365+
self.client.sandboxes.execute_bash(
366+
name=_TEST_SANDBOX_RESOURCE_NAME,
367+
command="true",
368+
port="9222",
369+
)
370+
371+
_, port, _ = self._sent_chunks(mock_execute_code)
372+
assert port == b"9222"
373+
374+
@mock.patch.object(sandboxes.Sandboxes, "_execute_code")
375+
def test_execute_bash_reports_failure_without_raising(self, mock_execute_code):
376+
"""A non-zero exit is returned on the result, not raised."""
377+
expected = {
378+
"stdout": "",
379+
"stderr": "ls: cannot access '/nope': No such file or directory\n",
380+
"returncode": 2,
381+
"duration_ms": 4,
382+
}
383+
mock_execute_code.return_value = self._exec_response(expected)
384+
385+
result = self.client.sandboxes.execute_bash(
386+
name=_TEST_SANDBOX_RESOURCE_NAME,
387+
command="ls /nope",
388+
)
389+
390+
assert result == expected
391+
392+
@mock.patch.object(sandboxes.Sandboxes, "_execute_code")
393+
def test_execute_bash_forwards_config(self, mock_execute_code):
394+
"""The request config is passed through untouched."""
395+
mock_execute_code.return_value = self._exec_response({"stdout": ""})
396+
config = agentplatform_types.ExecuteCodeRuntimeSandboxConfig()
397+
398+
self.client.sandboxes.execute_bash(
399+
name=_TEST_SANDBOX_RESOURCE_NAME,
400+
command="true",
401+
config=config,
402+
)
403+
404+
_, kwargs = mock_execute_code.call_args
405+
assert kwargs["config"] is config
406+
407+
@mock.patch.object(sandboxes.Sandboxes, "_execute_code")
408+
def test_execute_bash_without_output_raises(self, mock_execute_code):
409+
"""An empty response is an error, not a silent success."""
410+
mock_execute_code.return_value = (
411+
agentplatform_types.ExecuteSandboxEnvironmentResponse(outputs=[])
412+
)
413+
414+
with pytest.raises(ValueError, match="returned no output"):
415+
self.client.sandboxes.execute_bash(
416+
name=_TEST_SANDBOX_RESOURCE_NAME,
417+
command="echo hello",
418+
)
419+
297420

298421
@pytest.mark.usefixtures("google_auth_mock")
299422
class TestSandboxTemplates:

0 commit comments

Comments
 (0)