From cbd14ebf99bbff22ed28273f095f1fc05793bed2 Mon Sep 17 00:00:00 2001 From: Emily Feng Date: Fri, 22 May 2026 17:09:49 +0000 Subject: [PATCH] feat: Add support for creating sandboxes from templates and snapshots This change allows AgentEngineSandboxComputer to create new sandboxes using either a specified sandbox template or a sandbox snapshot. The environment variables VMAAS_SANDBOX_TEMPLATE_NAME and VMAAS_SANDBOX_SNAPSHOT_NAME are introduced to configure this behavior. Co-authored-by: Emily Feng Change-Id: Iebdd980a16966ba765cacbce6d63d0d5b691650a --- .../sandbox_computer_use/agent.py | 7 +++ .../integrations/sandbox_computer_use/main.py | 3 + .../integrations/vmaas/sandbox_computer.py | 60 ++++++++++++++----- 3 files changed, 54 insertions(+), 16 deletions(-) diff --git a/contributing/samples/integrations/sandbox_computer_use/agent.py b/contributing/samples/integrations/sandbox_computer_use/agent.py index b95dbad8..f559dc16 100644 --- a/contributing/samples/integrations/sandbox_computer_use/agent.py +++ b/contributing/samples/integrations/sandbox_computer_use/agent.py @@ -24,6 +24,8 @@ Prerequisites: - GOOGLE_CLOUD_PROJECT: Your GCP project ID - VMAAS_SERVICE_ACCOUNT: Your service account email - VMAAS_SANDBOX_NAME: (Optional) Existing sandbox resource name for BYOS mode + - VMAAS_SANDBOX_TEMPLATE_NAME: (Optional) Sandbox template name to create a new sandbox (mutually exclusive with VMAAS_SANDBOX_NAME) + - VMAAS_SANDBOX_SNAPSHOT_NAME: (Optional) Sandbox snapshot name to create a new sandbox (mutually exclusive with VMAAS_SANDBOX_NAME) Usage: # Run via ADK web UI @@ -53,12 +55,17 @@ SERVICE_ACCOUNT = os.environ.get("VMAAS_SERVICE_ACCOUNT") SANDBOX_NAME = os.environ.get("SANDBOX_NAME") or os.environ.get( "VMAAS_SANDBOX_NAME" ) +SANDBOX_TEMPLATE_NAME = os.environ.get("VMAAS_SANDBOX_TEMPLATE_NAME") +SANDBOX_SNAPSHOT_NAME = os.environ.get("VMAAS_SANDBOX_SNAPSHOT_NAME") + # Create the sandbox computer sandbox_computer = AgentEngineSandboxComputer( project_id=PROJECT_ID, service_account_email=SERVICE_ACCOUNT, sandbox_name=SANDBOX_NAME, + sandbox_template_name=SANDBOX_TEMPLATE_NAME, + sandbox_snapshot_name=SANDBOX_SNAPSHOT_NAME, search_engine_url="https://www.google.com", ) diff --git a/contributing/samples/integrations/sandbox_computer_use/main.py b/contributing/samples/integrations/sandbox_computer_use/main.py index e6c60bc2..19fbd8a5 100644 --- a/contributing/samples/integrations/sandbox_computer_use/main.py +++ b/contributing/samples/integrations/sandbox_computer_use/main.py @@ -22,6 +22,9 @@ Prerequisites: - GOOGLE_CLOUD_PROJECT: Your GCP project ID - VMAAS_SERVICE_ACCOUNT: Your service account email with roles/iam.serviceAccountTokenCreator permission + - VMAAS_SANDBOX_NAME: (Optional) Existing sandbox resource name for BYOS mode + - VMAAS_SANDBOX_TEMPLATE_NAME: (Optional) Sandbox template name to create a new sandbox (mutually exclusive with VMAAS_SANDBOX_NAME) + - VMAAS_SANDBOX_SNAPSHOT_NAME: (Optional) Sandbox snapshot name to create a new sandbox (mutually exclusive with VMAAS_SANDBOX_NAME) Usage: cd contributing/samples diff --git a/src/google/adk/integrations/vmaas/sandbox_computer.py b/src/google/adk/integrations/vmaas/sandbox_computer.py index 1dd38e0a..9c1c9d68 100644 --- a/src/google/adk/integrations/vmaas/sandbox_computer.py +++ b/src/google/adk/integrations/vmaas/sandbox_computer.py @@ -90,6 +90,8 @@ class AgentEngineSandboxComputer(BaseComputer): location: str = "us-central1", service_account_email: str | None = None, sandbox_name: str | None = None, + sandbox_template_name: str | None = None, + sandbox_snapshot_name: str | None = None, sandbox_ttl_seconds: int = 3600, search_engine_url: str = "https://www.google.com", vertexai_client: "vertexai.Client | None" = None, @@ -97,25 +99,36 @@ class AgentEngineSandboxComputer(BaseComputer): """Initialize the sandbox computer. Args: - project_id: GCP project ID. If None, uses Application Default - Credentials project. + project_id: GCP project ID. If None, uses Application Default Credentials + project. location: Vertex AI location (default: us-central1). - service_account_email: Service account email for token generation. - Must have roles/iam.serviceAccountTokenCreator permission. - If None, attempts to use ADC service account. - sandbox_name: Existing sandbox resource name (BYOS mode). If provided, - the agent engine name is extracted from it. If None, creates new - agent engine and sandbox on demand. - Format: projects/{project}/locations/{location}/reasoningEngines/{id}/sandboxEnvironments/{id} + service_account_email: Service account email for token generation. Must + have roles/iam.serviceAccountTokenCreator permission. If None, attempts + to use ADC service account. + sandbox_name: Existing sandbox resource name (BYOS mode). If provided, the + agent engine name is extracted from it. If None, creates new agent + engine and sandbox on demand. + Format: + projects/{project}/locations/{location}/reasoningEngines/{id}/sandboxEnvironments/{id} + sandbox_template_name: Sandbox template resource name to use for creating + new sandboxes. Templates allow faster creation and custom environments. + Format: + projects/{project}/locations/{location}/sandboxEnvironmentTemplates/{id} + sandbox_snapshot_name: Sandbox snapshot resource name to use for restoring + sandbox state, enabling faster startup. + Format: + projects/{project}/locations/{location}/reasoningEngines/{id}/sandboxEnvironmentSnapshots/{id} sandbox_ttl_seconds: TTL for auto-created sandboxes (default: 1 hour). search_engine_url: URL to navigate to for search() method. - vertexai_client: Optional Vertex AI client instance. If None, creates - one lazily using project_id and location. + vertexai_client: Optional Vertex AI client instance. If None, creates one + lazily using project_id and location. """ self._project_id = project_id self._location = location self._service_account_email = service_account_email self._sandbox_name = sandbox_name + self._sandbox_template_name = sandbox_template_name + self._sandbox_snapshot_name = sandbox_snapshot_name self._sandbox_ttl_seconds = sandbox_ttl_seconds self._search_engine_url = search_engine_url self._screen_size = (1280, 720) @@ -210,13 +223,29 @@ class AgentEngineSandboxComputer(BaseComputer): from vertexai import types + config = { + "display_name": "adk_computer_use_sandbox", + } + spec = None + if self._sandbox_template_name: + config["sandbox_environment_template"] = self._sandbox_template_name + logger.info( + "Creating sandbox from template: %s", self._sandbox_template_name + ) + elif self._sandbox_snapshot_name: + config["sandbox_environment_snapshot"] = self._sandbox_snapshot_name + logger.info( + "Creating sandbox from snapshot: %s", self._sandbox_snapshot_name + ) + else: + spec = {"computer_use_environment": {}} + logger.info("Creating sandbox with computer use environment spec") + operation = await asyncio.to_thread( client.agent_engines.sandboxes.create, - spec={"computer_use_environment": {}}, + spec=spec, name=agent_engine_name, - config=types.CreateAgentEngineSandboxConfig( - display_name="adk_computer_use_sandbox" - ), + config=config, ) sandbox_name = operation.response.name @@ -249,7 +278,6 @@ class AgentEngineSandboxComputer(BaseComputer): token = await asyncio.to_thread( client.agent_engines.sandboxes.generate_access_token, service_account_email=self._service_account_email, - sandbox_id=sandbox_name, timeout=_DEFAULT_TOKEN_TIMEOUT, )