feat: Add Shell Sandbox environment support to the Agent Engine sandbox SDK

agentplatform/_genai/sandboxes.py

@@ -706,26 +706,90 @@ def create(
         Returns:
             AgentEngineSandboxOperation: The operation for creating the sandbox.
         """
-        if spec:
-            computer_use = False
+        if config is None:
+            config = types.CreateAgentEngineSandboxConfig()
+        elif isinstance(config, dict):
+            config = types.CreateAgentEngineSandboxConfig.model_validate(config)
+
+        def _spec_has(field_name: str) -> bool:
+            if spec is None:
+                return False
             if isinstance(spec, dict):
-                computer_use = spec.get("computer_use_environment") is not None
-            elif hasattr(spec, "computer_use_environment"):
-                computer_use = True
+                return spec.get(field_name) is not None
+            return getattr(spec, field_name, None) is not None
+
+        # 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(
+            _spec_has(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 _spec_has(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
 
-            if computer_use:
-                logging.warning(
-                    "The computer_use_environment feature in the sandboxes module is experimental and may change in future versions."
-                )
         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(

agentplatform/_genai/types/__init__.py

@@ -1781,6 +1781,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
@@ -3213,6 +3216,9 @@
     "SandboxEnvironmentSpecComputerUseEnvironment",
     "SandboxEnvironmentSpecComputerUseEnvironmentDict",
     "SandboxEnvironmentSpecComputerUseEnvironmentOrDict",
+    "SandboxEnvironmentSpecShellEnvironment",
+    "SandboxEnvironmentSpecShellEnvironmentDict",
+    "SandboxEnvironmentSpecShellEnvironmentOrDict",
     "SandboxEnvironmentSpec",
     "SandboxEnvironmentSpecDict",
     "SandboxEnvironmentSpecOrDict",

agentplatform/_genai/types/common.py

@@ -381,6 +381,10 @@ class DefaultContainerCategory(_common.CaseInSensitiveEnum):
     """The default value. This value is unused."""
     DEFAULT_CONTAINER_CATEGORY_COMPUTER_USE = "DEFAULT_CONTAINER_CATEGORY_COMPUTER_USE"
     """The default container image for Computer Use."""
+    DEFAULT_CONTAINER_CATEGORY_SHELL_SANDBOX = (
+        "DEFAULT_CONTAINER_CATEGORY_SHELL_SANDBOX"
+    )
+    """The default container image for Shell Sandbox."""
 
 
 class PostSnapshotAction(_common.CaseInSensitiveEnum):
@@ -16419,6 +16423,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."""
 
@@ -16428,6 +16449,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):
@@ -16441,6 +16465,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]
 

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",