From 505bd42f65a8730793c7592ffa61d3f5b80f04cf Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Tue, 1 Sep 2026 11:00:33 -0700 Subject: [PATCH] feat: Add execute_bash to the Agent Engine sandbox SDK. Adds a helper function to simply bash execution in shell sandbox PiperOrigin-RevId: 974592800 --- agentplatform/_genai/sandboxes.py | 79 +++++++++++ .../unit/agentplatform/genai/test_sandbox.py | 123 ++++++++++++++++++ 2 files changed, 202 insertions(+) diff --git a/agentplatform/_genai/sandboxes.py b/agentplatform/_genai/sandboxes.py index e92a46a05b..406ce30164 100644 --- a/agentplatform/_genai/sandboxes.py +++ b/agentplatform/_genai/sandboxes.py @@ -1168,6 +1168,85 @@ def execute_code( return response + def execute_bash( + self, + *, + name: str, + command: str, + cwd: Optional[str] = None, + timeout: Optional[int] = None, + port: str = "8080", + config: Optional[types.ExecuteCodeRuntimeSandboxConfigOrDict] = None, + ) -> dict[str, Any]: + """Runs a bash command in a shell sandbox. + + The sandbox must have been created with a shell environment, either by + passing `spec={"shell_environment": {}}` to `create()` or by referencing a + template whose `default_container_category` is + `DEFAULT_CONTAINER_CATEGORY_SHELL_SANDBOX`. + + Unlike `send_command`, this authenticates with the caller's own credentials, + so it needs no service account and no signed JWT. + + Args: + name (str): + Required. The name of the agent engine sandbox to run the command + in, of the form + projects/*/locations/*/reasoningEngines/*/sandboxEnvironments/*. + command (str): + Required. The bash command to run. + cwd (str): + Optional. The working directory to run the command in. Defaults to + the sandbox's workspace directory. + timeout (int): + Optional. Seconds to allow the command to run before the sandbox + kills it. Defaults to the sandbox's own limit. + port (str): + Optional. The port the shell server listens on. Defaults to "8080". + config (ExecuteCodeRuntimeSandboxConfigOrDict): + Optional. The configuration for the request. + + Returns: + dict[str, Any]: The result of the command, with keys `stdout`, + `stderr`, `returncode` and `duration_ms`. A command that exits + non-zero is reported here rather than raised. + + Raises: + ValueError: If the sandbox returns no output. + """ + request: dict[str, Any] = {"command": command} + if cwd is not None: + request["cwd"] = cwd + if timeout is not None: + request["timeout"] = timeout + + # The sandbox proxy addresses the container by URI and port, and forwards + # the JSON chunk as the POST body. + response = self._execute_code( + name=name, + inputs=[ + types.Chunk( + mime_type="application/x.sandbox-request-uri", + data=b"/exec", + ), + types.Chunk( + mime_type="application/x.sandbox-request-port", + data=port.encode("utf-8"), + ), + types.Chunk( + mime_type="application/json", + data=json.dumps(request).encode("utf-8"), + ), + ], + config=config, + ) + + for output in response.outputs or []: + if output.data: + return json.loads(output.data.decode("utf-8")) + + raise ValueError(f"The sandbox {name} returned no output for the command.") + def get( self, *, diff --git a/tests/unit/agentplatform/genai/test_sandbox.py b/tests/unit/agentplatform/genai/test_sandbox.py index d128ff5311..ca8f5c73a1 100644 --- a/tests/unit/agentplatform/genai/test_sandbox.py +++ b/tests/unit/agentplatform/genai/test_sandbox.py @@ -294,6 +294,129 @@ def test_create_without_spec_template_or_snapshot_raises(self, mock_create): mock_create.assert_not_called() + @staticmethod + def _exec_response(payload): + """Builds the response the sandbox proxy returns for POST /exec.""" + return agentplatform_types.ExecuteSandboxEnvironmentResponse( + outputs=[ + agentplatform_types.Chunk( + mime_type="application/json", + data=json.dumps(payload).encode("utf-8"), + ) + ] + ) + + @staticmethod + def _sent_chunks(mock_execute_code): + """Returns the uri, port and JSON body sent to the sandbox.""" + _, kwargs = mock_execute_code.call_args + chunks = {chunk.mime_type: chunk.data for chunk in kwargs["inputs"]} + return ( + chunks["application/x.sandbox-request-uri"], + chunks["application/x.sandbox-request-port"], + json.loads(chunks["application/json"]), + ) + + @mock.patch.object(sandboxes.Sandboxes, "_execute_code") + def test_execute_bash_runs_command_and_parses_result(self, mock_execute_code): + """A command is sent to /exec and its JSON result is returned.""" + expected = { + "stdout": "hello\n", + "stderr": "", + "returncode": 0, + "duration_ms": 5, + } + mock_execute_code.return_value = self._exec_response(expected) + + result = self.client.sandboxes.execute_bash( + name=_TEST_SANDBOX_RESOURCE_NAME, + command="echo hello", + ) + + assert result == expected + _, kwargs = mock_execute_code.call_args + assert kwargs["name"] == _TEST_SANDBOX_RESOURCE_NAME + uri, port, body = self._sent_chunks(mock_execute_code) + assert uri == b"/exec" + assert port == b"8080" + # cwd and timeout are left to the sandbox's own defaults when unset. + assert body == {"command": "echo hello"} + + @mock.patch.object(sandboxes.Sandboxes, "_execute_code") + def test_execute_bash_forwards_cwd_and_timeout(self, mock_execute_code): + """cwd and timeout reach the container when the caller sets them.""" + mock_execute_code.return_value = self._exec_response({"stdout": "/tmp\n"}) + + self.client.sandboxes.execute_bash( + name=_TEST_SANDBOX_RESOURCE_NAME, + command="pwd", + cwd="/tmp", + timeout=30, + ) + + _, _, body = self._sent_chunks(mock_execute_code) + assert body == {"command": "pwd", "cwd": "/tmp", "timeout": 30} + + @mock.patch.object(sandboxes.Sandboxes, "_execute_code") + def test_execute_bash_uses_custom_port(self, mock_execute_code): + """A caller-supplied port overrides the shell server default.""" + mock_execute_code.return_value = self._exec_response({"stdout": ""}) + + self.client.sandboxes.execute_bash( + name=_TEST_SANDBOX_RESOURCE_NAME, + command="true", + port="9222", + ) + + _, port, _ = self._sent_chunks(mock_execute_code) + assert port == b"9222" + + @mock.patch.object(sandboxes.Sandboxes, "_execute_code") + def test_execute_bash_reports_failure_without_raising(self, mock_execute_code): + """A non-zero exit is returned on the result, not raised.""" + expected = { + "stdout": "", + "stderr": "ls: cannot access '/nope': No such file or directory\n", + "returncode": 2, + "duration_ms": 4, + } + mock_execute_code.return_value = self._exec_response(expected) + + result = self.client.sandboxes.execute_bash( + name=_TEST_SANDBOX_RESOURCE_NAME, + command="ls /nope", + ) + + assert result == expected + + @mock.patch.object(sandboxes.Sandboxes, "_execute_code") + def test_execute_bash_forwards_config(self, mock_execute_code): + """The request config is passed through untouched.""" + mock_execute_code.return_value = self._exec_response({"stdout": ""}) + config = agentplatform_types.ExecuteCodeRuntimeSandboxConfig() + + self.client.sandboxes.execute_bash( + name=_TEST_SANDBOX_RESOURCE_NAME, + command="true", + config=config, + ) + + _, kwargs = mock_execute_code.call_args + assert kwargs["config"] is config + + @mock.patch.object(sandboxes.Sandboxes, "_execute_code") + def test_execute_bash_without_output_raises(self, mock_execute_code): + """An empty response is an error, not a silent success.""" + mock_execute_code.return_value = ( + agentplatform_types.ExecuteSandboxEnvironmentResponse(outputs=[]) + ) + + with pytest.raises(ValueError, match="returned no output"): + self.client.sandboxes.execute_bash( + name=_TEST_SANDBOX_RESOURCE_NAME, + command="echo hello", + ) + @pytest.mark.usefixtures("google_auth_mock") class TestSandboxTemplates: