fix: Derive sandbox reasoning engine from template or snapshot name

PiperOrigin-RevId: 940679014

Author: google-genai-bot

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

@@ -133,13 +133,26 @@ def __init__(
     self._search_engine_url = search_engine_url
     self._screen_size = (1280, 720)
 
-    # Extract agent engine name from sandbox_name if provided
+    # Determine the agent engine to use. The backend requires that a sandbox
+    # is created under the same reasoning engine that owns the
+    # template/snapshot, so we derive the engine from whichever resource name
+    # is provided rather than creating a new (mismatched) engine. This lets
+    # users supply only a template or snapshot name without also pre-creating
+    # a sandbox (BYOS). All sandbox resource names embed the engine:
+    #   projects/.../reasoningEngines/{engine}/{sandboxEnvironments|
+    #   sandboxEnvironmentTemplates|sandboxEnvironmentSnapshots}/...
     self._agent_engine_name = None
-    if sandbox_name:
-      # Format: projects/.../reasoningEngines/.../sandboxEnvironments/...
-      parts = sandbox_name.split("/sandboxEnvironments/")
-      if len(parts) == 2:
-        self._agent_engine_name = parts[0]
+    for resource_name in (
+        sandbox_name,
+        sandbox_template_name,
+        sandbox_snapshot_name,
+    ):
+      if not resource_name:
+        continue
+      engine_name = resource_name.split("/sandboxEnvironment")[0]
+      if engine_name != resource_name and "/reasoningEngines/" in engine_name:
+        self._agent_engine_name = engine_name
+        break
 
     # Vertex client (lazy-initialized if not provided)
     self._client = vertexai_client

tests/unittests/integrations/vmaas/test_sandbox_computer.py

@@ -67,6 +67,77 @@ def test_init_with_byos(self):
     self.assertEqual(computer._agent_engine_name, agent_engine_name)
     self.assertEqual(computer._sandbox_name, sandbox_name)
 
+  def test_init_with_template_derives_agent_engine(self):
+    """Test agent engine is derived from sandbox_template_name."""
+    agent_engine_name = (
+        "projects/test/locations/us-central1/reasoningEngines/123"
+    )
+    template_name = f"{agent_engine_name}/sandboxEnvironmentTemplates/789"
+
+    computer = AgentEngineSandboxComputer(
+        project_id=self.project_id,
+        sandbox_template_name=template_name,
+    )
+
+    # Agent engine name should be derived from the template name so the
+    # sandbox is created under the template's reasoning engine (no BYOS).
+    self.assertEqual(computer._agent_engine_name, agent_engine_name)
+    self.assertEqual(computer._sandbox_template_name, template_name)
+
+  def test_init_with_snapshot_derives_agent_engine(self):
+    """Test agent engine is derived from sandbox_snapshot_name."""
+    agent_engine_name = (
+        "projects/test/locations/us-central1/reasoningEngines/123"
+    )
+    snapshot_name = f"{agent_engine_name}/sandboxEnvironmentSnapshots/789"
+
+    computer = AgentEngineSandboxComputer(
+        project_id=self.project_id,
+        sandbox_snapshot_name=snapshot_name,
+    )
+
+    self.assertEqual(computer._agent_engine_name, agent_engine_name)
+    self.assertEqual(computer._sandbox_snapshot_name, snapshot_name)
+
+  def test_init_sandbox_name_takes_precedence_over_template(self):
+    """Test sandbox_name wins when both sandbox_name and template are set."""
+    sandbox_engine = (
+        "projects/test/locations/us-central1/reasoningEngines/sandbox"
+    )
+    sandbox_name = f"{sandbox_engine}/sandboxEnvironments/456"
+    template_name = (
+        "projects/test/locations/us-central1/reasoningEngines/template"
+        "/sandboxEnvironmentTemplates/789"
+    )
+
+    computer = AgentEngineSandboxComputer(
+        project_id=self.project_id,
+        sandbox_name=sandbox_name,
+        sandbox_template_name=template_name,
+    )
+
+    self.assertEqual(computer._agent_engine_name, sandbox_engine)
+
+  def test_init_without_sandbox_source_has_no_agent_engine(self):
+    """Test agent engine is None when no sandbox source is provided."""
+    computer = AgentEngineSandboxComputer(project_id=self.project_id)
+    self.assertIsNone(computer._agent_engine_name)
+
+  async def test_ensure_agent_engine_with_template_name(self):
+    """Test _ensure_agent_engine reuses the template's reasoning engine."""
+    agent_engine_name = (
+        "projects/test/locations/us-central1/reasoningEngines/123"
+    )
+    template_name = f"{agent_engine_name}/sandboxEnvironmentTemplates/789"
+    computer = AgentEngineSandboxComputer(sandbox_template_name=template_name)
+    computer._session_state = {}
+
+    result = await computer._ensure_agent_engine()
+
+    self.assertEqual(result, agent_engine_name)
+    # Should not have created or stored a new engine.
+    self.assertNotIn(_STATE_KEY_AGENT_ENGINE_NAME, computer._session_state)
+
   def test_init_with_vertexai_client(self):
     """Test initialization with provided vertexai client."""
     mock_client = MagicMock()