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 <emilyfeng@google.com>
Change-Id: Iebdd980a16966ba765cacbce6d63d0d5b691650a

Author: Emily Feng

contributing/samples/integrations/sandbox_computer_use/agent.py

@@ -24,6 +24,8 @@
      - 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 @@
 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",
 )
 

contributing/samples/integrations/sandbox_computer_use/main.py

@@ -22,6 +22,9 @@
      - 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

src/google/adk/integrations/vmaas/sandbox_computer.py

@@ -90,32 +90,45 @@ def __init__(
       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,
   ):
     """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 @@ async def _get_sandbox(self) -> tuple[str, Any]:
 
     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 @@ async def _get_access_token(self, sandbox_name: str) -> str:
     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,
     )