diff --git a/agentplatform/_genai/sandboxes.py b/agentplatform/_genai/sandboxes.py index d25746864e..7d419c98e1 100644 --- a/agentplatform/_genai/sandboxes.py +++ b/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( diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index d317ab56e3..3c94e3532f 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/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", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index ec8498e4b5..43138a0a30 100644 --- a/agentplatform/_genai/types/common.py +++ b/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] 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",