From 8f87082ae4aa0e4c116245b3c8a328e9a55ea6a0 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Fri, 21 Aug 2026 12:10:44 -0700 Subject: [PATCH] feat: Add Shell Sandbox environment support 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: 968624286 --- agentplatform/_genai/_agent_engines_utils.py | 20 +++ agentplatform/_genai/sandboxes.py | 76 ++++++++- agentplatform/_genai/types/__init__.py | 6 + agentplatform/_genai/types/common.py | 23 +++ .../unit/agentplatform/genai/test_sandbox.py | 160 ++++++++++++++++++ 5 files changed, 281 insertions(+), 4 deletions(-) diff --git a/agentplatform/_genai/_agent_engines_utils.py b/agentplatform/_genai/_agent_engines_utils.py index 41f9d180ec..87ff7395d0 100644 --- a/agentplatform/_genai/_agent_engines_utils.py +++ b/agentplatform/_genai/_agent_engines_utils.py @@ -234,6 +234,26 @@ logger = logging.getLogger("agentplatform_genai.agentengines") +def has_field(obj: Union[BaseModel, JsonDict], field_name: str) -> bool: + """Returns whether `obj` has `field_name` set to a non-None value. + + Supports both pydantic models (or any attribute-bearing object) and dicts. + + Args: + obj: The object to inspect. May be a pydantic model, a dict, or None. + field_name: The name of the field to check for. + + Returns: + True if `obj` is non-empty and `field_name` is set to a non-None value, + False otherwise. + """ + if not obj: + return False + if isinstance(obj, dict): + return obj.get(field_name) is not None + return getattr(obj, field_name, None) is not None + + @typing.runtime_checkable class Queryable(Protocol): """Protocol for Agent Engines that can be queried.""" diff --git a/agentplatform/_genai/sandboxes.py b/agentplatform/_genai/sandboxes.py index e19f93a2bd..9dfd08dcbd 100644 --- a/agentplatform/_genai/sandboxes.py +++ b/agentplatform/_genai/sandboxes.py @@ -706,15 +706,83 @@ def create( Returns: AgentEngineSandboxOperation: The operation for creating the sandbox. """ + if config is None: + config = types.CreateAgentEngineSandboxConfig() + elif isinstance(config, dict): + config = types.CreateAgentEngineSandboxConfig.model_validate(config) + + # A sandbox environment must be provided inline via `spec` (with an + # environment set), or by referencing an existing template or snapshot in + # `config`. + spec_has_environment = any( + _agent_engines_utils.has_field(spec, field_name) + for field_name in ( + "code_execution_environment", + "computer_use_environment", + "shell_environment", + ) + ) + if ( + not spec_has_environment + and not config.sandbox_environment_template + and not config.sandbox_environment_snapshot + ): + raise ValueError( + "A sandbox environment must be provided via `spec`, " + "`config.sandbox_environment_template`, or " + "`config.sandbox_environment_snapshot`." + ) + + if spec: + # Environments that can auto-provision a default sandbox + # environment template when the caller does not supply one. Ordered by + # precedence: the first matching environment is used. + environments = ( + ( + "shell_environment", + types.DefaultContainerCategory.DEFAULT_CONTAINER_CATEGORY_SHELL_SANDBOX, + "shell-sandbox-template", + ), + ( + "computer_use_environment", + types.DefaultContainerCategory.DEFAULT_CONTAINER_CATEGORY_COMPUTER_USE, + "computer-use-template", + ), + ) + + for field_name, category, display_name in environments: + if not _agent_engines_utils.has_field(spec, field_name): + continue + + if ( + not config.sandbox_environment_template + and not config.sandbox_environment_snapshot + ): + default_container_environment = ( + types.SandboxEnvironmentTemplateDefaultContainerEnvironment( + default_container_category=category, + ) + ) + template_operation = self.templates.create( + name=name, + display_name=display_name, + config=types.CreateSandboxEnvironmentTemplateConfig( + default_container_environment=default_container_environment, + ), + poll_interval_seconds=poll_interval_seconds, + ) + if not template_operation.response: + raise ValueError(f"Error creating {display_name}.") + config.sandbox_environment_template = ( + template_operation.response.name + ) + break + operation = self._create( name=name, spec=spec, config=config, ) - if config is None: - config = types.CreateAgentEngineSandboxConfig() - elif isinstance(config, dict): - config = types.CreateAgentEngineSandboxConfig.model_validate(config) if config.wait_for_completion: if not operation.done: operation = _agent_engines_utils._await_operation( diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index 24a8359701..e574831efd 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -1933,6 +1933,9 @@ from .common import SandboxEnvironmentSpecComputerUseEnvironmentOrDict from .common import SandboxEnvironmentSpecDict from .common import SandboxEnvironmentSpecOrDict +from .common import SandboxEnvironmentSpecShellEnvironment +from .common import SandboxEnvironmentSpecShellEnvironmentDict +from .common import SandboxEnvironmentSpecShellEnvironmentOrDict from .common import SandboxEnvironmentTemplate from .common import SandboxEnvironmentTemplateCustomContainerEnvironment from .common import SandboxEnvironmentTemplateCustomContainerEnvironmentDict @@ -3446,6 +3449,9 @@ "SandboxEnvironmentSpecComputerUseEnvironment", "SandboxEnvironmentSpecComputerUseEnvironmentDict", "SandboxEnvironmentSpecComputerUseEnvironmentOrDict", + "SandboxEnvironmentSpecShellEnvironment", + "SandboxEnvironmentSpecShellEnvironmentDict", + "SandboxEnvironmentSpecShellEnvironmentOrDict", "SandboxEnvironmentSpec", "SandboxEnvironmentSpecDict", "SandboxEnvironmentSpecOrDict", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index d849d5c68b..57fe7bc623 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -17161,6 +17161,23 @@ class SandboxEnvironmentSpecComputerUseEnvironmentDict(TypedDict, total=False): ] +class SandboxEnvironmentSpecShellEnvironment(_common.BaseModel): + """The shell environment with customized settings.""" + + pass + + +class SandboxEnvironmentSpecShellEnvironmentDict(TypedDict, total=False): + """The shell environment with customized settings.""" + + pass + + +SandboxEnvironmentSpecShellEnvironmentOrDict = Union[ + SandboxEnvironmentSpecShellEnvironment, SandboxEnvironmentSpecShellEnvironmentDict +] + + class SandboxEnvironmentSpec(_common.BaseModel): """The specification of a sandbox environment.""" @@ -17170,6 +17187,9 @@ class SandboxEnvironmentSpec(_common.BaseModel): computer_use_environment: Optional[SandboxEnvironmentSpecComputerUseEnvironment] = ( Field(default=None, description="""Optional. The computer use environment.""") ) + shell_environment: Optional[SandboxEnvironmentSpecShellEnvironment] = Field( + default=None, description="""Optional. The shell environment.""" + ) class SandboxEnvironmentSpecDict(TypedDict, total=False): @@ -17183,6 +17203,9 @@ class SandboxEnvironmentSpecDict(TypedDict, total=False): computer_use_environment: Optional[SandboxEnvironmentSpecComputerUseEnvironmentDict] """Optional. The computer use environment.""" + shell_environment: Optional[SandboxEnvironmentSpecShellEnvironmentDict] + """Optional. The shell environment.""" + SandboxEnvironmentSpecOrDict = Union[SandboxEnvironmentSpec, SandboxEnvironmentSpecDict] diff --git a/tests/unit/agentplatform/genai/test_sandbox.py b/tests/unit/agentplatform/genai/test_sandbox.py index 9672e473a9..c162ca23b8 100644 --- a/tests/unit/agentplatform/genai/test_sandbox.py +++ b/tests/unit/agentplatform/genai/test_sandbox.py @@ -22,7 +22,9 @@ from google.auth import credentials as auth_credentials import agentplatform from google.cloud import aiplatform +from agentplatform._genai import sandbox_templates from agentplatform._genai import sandboxes +from agentplatform._genai import types as agentplatform_types from google.cloud.aiplatform import initializer from vertexai._genai import ( sandboxes as vertexai_sandboxes, @@ -44,6 +46,11 @@ _TEST_SANDBOX_RESOURCE_NAME = ( f"{_TEST_AGENT_ENGINE_RESOURCE_NAME}/sandboxes/{_TEST_SANDBOX_ID}" ) +_TEST_SANDBOX_TEMPLATE_ID = "template-123" +_TEST_SANDBOX_TEMPLATE_RESOURCE_NAME = ( + f"{_TEST_AGENT_ENGINE_RESOURCE_NAME}" + f"/sandboxEnvironmentTemplates/{_TEST_SANDBOX_TEMPLATE_ID}" +) _TEST_AGENT_ENGINE_ENV_KEY = "GOOGLE_CLOUD_AGENT_ENGINE_ENV" _TEST_AGENT_ENGINE_ENV_VALUE = "test_env_value" _TEST_SERVICE_ACCOUNT_EMAIL = "test-sa@test-project.iam.gserviceaccount.com" @@ -136,6 +143,159 @@ def test_generate_browser_ws_headers( == "v1.stream, test_token, test_routing_token, 9222" ) + @mock.patch.object(sandboxes.Sandboxes, "_create") + def test_create_with_shell_environment_and_existing_template(self, mock_create): + mock_operation = mock.Mock() + mock_create.return_value = mock_operation + + result = self.client.agent_engines.sandboxes.create( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec={"shell_environment": {}}, + config={ + "sandbox_environment_template": _TEST_SANDBOX_TEMPLATE_RESOURCE_NAME, + "wait_for_completion": False, + }, + ) + + assert result is mock_operation + mock_create.assert_called_once() + _, kwargs = mock_create.call_args + assert kwargs["name"] == _TEST_AGENT_ENGINE_RESOURCE_NAME + assert ( + kwargs["config"].sandbox_environment_template + == _TEST_SANDBOX_TEMPLATE_RESOURCE_NAME + ) + + @mock.patch.object(sandboxes.Sandboxes, "_create") + @mock.patch.object(sandbox_templates.SandboxTemplates, "create") + def test_create_with_shell_environment_creates_template_when_absent( + self, mock_template_create, mock_create + ): + mock_template_operation = mock.Mock() + mock_template_operation.response.name = _TEST_SANDBOX_TEMPLATE_RESOURCE_NAME + mock_template_create.return_value = mock_template_operation + mock_create.return_value = mock.Mock() + + self.client.agent_engines.sandboxes.create( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec={"shell_environment": {}}, + config={"wait_for_completion": False}, + ) + + mock_template_create.assert_called_once() + _, template_kwargs = mock_template_create.call_args + template_config = template_kwargs["config"] + assert ( + template_config.default_container_environment.default_container_category + == agentplatform_types.DefaultContainerCategory.DEFAULT_CONTAINER_CATEGORY_SHELL_SANDBOX + ) + mock_create.assert_called_once() + _, create_kwargs = mock_create.call_args + assert ( + create_kwargs["config"].sandbox_environment_template + == _TEST_SANDBOX_TEMPLATE_RESOURCE_NAME + ) + + @mock.patch.object(sandboxes.Sandboxes, "_create") + @mock.patch.object(sandbox_templates.SandboxTemplates, "create") + def test_create_with_typed_shell_environment_creates_template_when_absent( + self, mock_template_create, mock_create + ): + mock_template_operation = mock.Mock() + mock_template_operation.response.name = _TEST_SANDBOX_TEMPLATE_RESOURCE_NAME + mock_template_create.return_value = mock_template_operation + mock_create.return_value = mock.Mock() + + shell_environment = agentplatform_types.SandboxEnvironmentSpecShellEnvironment() + self.client.agent_engines.sandboxes.create( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec=agentplatform_types.SandboxEnvironmentSpec( + shell_environment=shell_environment, + ), + config={"wait_for_completion": False}, + ) + + mock_template_create.assert_called_once() + _, template_kwargs = mock_template_create.call_args + template_config = template_kwargs["config"] + assert ( + template_config.default_container_environment.default_container_category + == agentplatform_types.DefaultContainerCategory.DEFAULT_CONTAINER_CATEGORY_SHELL_SANDBOX + ) + mock_create.assert_called_once() + _, create_kwargs = mock_create.call_args + assert ( + create_kwargs["config"].sandbox_environment_template + == _TEST_SANDBOX_TEMPLATE_RESOURCE_NAME + ) + + @mock.patch.object(sandboxes.Sandboxes, "_create") + @mock.patch.object(sandbox_templates.SandboxTemplates, "create") + def test_create_with_computer_use_environment_creates_template_when_absent( + self, mock_template_create, mock_create + ): + mock_template_operation = mock.Mock() + mock_template_operation.response.name = _TEST_SANDBOX_TEMPLATE_RESOURCE_NAME + mock_template_create.return_value = mock_template_operation + mock_create.return_value = mock.Mock() + + self.client.agent_engines.sandboxes.create( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec={"computer_use_environment": {}}, + config={"wait_for_completion": False}, + ) + + mock_template_create.assert_called_once() + _, template_kwargs = mock_template_create.call_args + template_config = template_kwargs["config"] + assert ( + template_config.default_container_environment.default_container_category + == agentplatform_types.DefaultContainerCategory.DEFAULT_CONTAINER_CATEGORY_COMPUTER_USE + ) + mock_create.assert_called_once() + _, create_kwargs = mock_create.call_args + assert ( + create_kwargs["config"].sandbox_environment_template + == _TEST_SANDBOX_TEMPLATE_RESOURCE_NAME + ) + + @mock.patch.object(sandboxes.Sandboxes, "_create") + @mock.patch.object(sandbox_templates.SandboxTemplates, "create") + def test_create_with_snapshot_does_not_create_template( + self, mock_template_create, mock_create + ): + mock_operation = mock.Mock() + mock_create.return_value = mock_operation + + result = self.client.agent_engines.sandboxes.create( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec={"computer_use_environment": {}}, + config={ + "sandbox_environment_snapshot": "projects/p/locations/l/agentEngines/ae/sandboxEnvironmentSnapshots/s1", + "wait_for_completion": False, + }, + ) + + assert result is mock_operation + mock_template_create.assert_not_called() + mock_create.assert_called_once() + _, create_kwargs = mock_create.call_args + assert ( + create_kwargs["config"].sandbox_environment_snapshot + == "projects/p/locations/l/agentEngines/ae/sandboxEnvironmentSnapshots/s1" + ) + + @mock.patch.object(sandboxes.Sandboxes, "_create") + def test_create_without_spec_template_or_snapshot_raises(self, mock_create): + for spec in (None, {}, agentplatform_types.SandboxEnvironmentSpec()): + with pytest.raises(ValueError, match="must be provided"): + self.client.agent_engines.sandboxes.create( + name=_TEST_AGENT_ENGINE_RESOURCE_NAME, + spec=spec, + ) + + mock_create.assert_not_called() + _MODULES = pytest.mark.parametrize( "module",