fix: Update agent_engine_sandbox_code_executor in ADK

1. For prototyping and testing purposes, sandbox name can be provided, and it will be used for all requests across the lifecycle of an agent
2. If no sandbox name is provided, agent engine name will be provided, and we will automatically create one sandbox per session, and the sandbox has TTL set for a year.
If the sandbox stored in the session hits the TTL, it will not be in "STATE_RUNNING" so a new sandbox will be created.

Co-authored-by: Lusha Wang <lusha@google.com>
PiperOrigin-RevId: 874415933

Author: lushasummer

contributing/samples/agent_engine_code_execution/README

@@ -7,9 +7,9 @@ This sample data science agent uses Agent Engine Code Execution Sandbox to execu
 
 ## How to use
 
-* 1. Follow https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/code-execution/overview to create a code execution sandbox environment.
+* 1. Follow https://docs.cloud.google.com/agent-builder/agent-engine/code-execution/quickstart#create-an-agent-engine-instance to create an agent engine instance. Replace the AGENT_ENGINE_RESOURCE_NAME with the one you just created. A new sandbox environment under this agent engine instance will be created for each session with TTL of 1 year. But sandbox can only main its state for up to 14 days. This is the recommended usage for production environments.
 
-* 2. Replace the SANDBOX_RESOURCE_NAME with the one you just created. If you dont want to create a new sandbox environment directly, the Agent Engine Code Execution Sandbox will create one for you by default using the AGENT_ENGINE_RESOURCE_NAME you specified, however, please ensure to clean up sandboxes after use; otherwise, it will consume quotas.
+* 2. For testing or protyping purposes, create a sandbox environment by following this guide: https://docs.cloud.google.com/agent-builder/agent-engine/code-execution/quickstart#create_a_sandbox. Replace the SANDBOX_RESOURCE_NAME with the one you just created. This will be used as the default sandbox environment for all the code executions throughout the lifetime of the agent. As the sandbox is re-used across sessions, all sessions will share the same Python environment and variable values."
 
 
 ## Sample prompt

contributing/samples/agent_engine_code_execution/agent.py

@@ -85,11 +85,10 @@ def base_system_instruction():
 
 """,
     code_executor=AgentEngineSandboxCodeExecutor(
-        # Replace with your sandbox resource name if you already have one.
-        sandbox_resource_name="SANDBOX_RESOURCE_NAME",
+        # Replace with your sandbox resource name if you already have one. Only use it for testing or prototyping purposes, because this will use the same sandbox for all requests.
         # "projects/vertex-agent-loadtest/locations/us-central1/reasoningEngines/6842889780301135872/sandboxEnvironments/6545148628569161728",
-        # Replace with agent engine resource name used for creating sandbox if
-        # sandbox_resource_name is not set.
+        sandbox_resource_name=None,
+        # Replace with agent engine resource name used for creating sandbox environment.
         agent_engine_resource_name="AGENT_ENGINE_RESOURCE_NAME",
     ),
 )

src/google/adk/code_executors/agent_engine_sandbox_code_executor.py

@@ -21,6 +21,7 @@
 from typing import Optional
 
 from typing_extensions import override
+from vertexai import types
 
 from ..agents.invocation_context import InvocationContext
 from .base_code_executor import BaseCodeExecutor
@@ -38,10 +39,15 @@ class AgentEngineSandboxCodeExecutor(BaseCodeExecutor):
     sandbox_resource_name: If set, load the existing resource name of the code
       interpreter extension instead of creating a new one. Format:
       projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789
+    agent_engine_resource_name: The resource name of the agent engine to use
+      to create the code execution sandbox. Format:
+      projects/123/locations/us-central1/reasoningEngines/456
   """
 
   sandbox_resource_name: str = None
 
+  agent_engine_resource_name: str = None
+
   def __init__(
       self,
       sandbox_resource_name: Optional[str] = None,
@@ -67,30 +73,19 @@ def __init__(
     agent_engine_resource_name_pattern = r'^projects/([a-zA-Z0-9-_]+)/locations/([a-zA-Z0-9-_]+)/reasoningEngines/(\d+)$'
 
     if sandbox_resource_name is not None:
-      self.sandbox_resource_name = sandbox_resource_name
       self._project_id, self._location = (
           self._get_project_id_and_location_from_resource_name(
               sandbox_resource_name, sandbox_resource_name_pattern
           )
       )
+      self.sandbox_resource_name = sandbox_resource_name
     elif agent_engine_resource_name is not None:
-      from vertexai import types
-
       self._project_id, self._location = (
           self._get_project_id_and_location_from_resource_name(
               agent_engine_resource_name, agent_engine_resource_name_pattern
           )
       )
-      # @TODO - Add TTL for sandbox creation after it is available
-      # in SDK.
-      operation = self._get_api_client().agent_engines.sandboxes.create(
-          spec={'code_execution_environment': {}},
-          name=agent_engine_resource_name,
-          config=types.CreateAgentEngineSandboxConfig(
-              display_name='default_sandbox'
-          ),
-      )
-      self.sandbox_resource_name = operation.response.name
+      self.agent_engine_resource_name = agent_engine_resource_name
     else:
       raise ValueError(
           'Either sandbox_resource_name or agent_engine_resource_name must be'
@@ -103,6 +98,39 @@ def execute_code(
       invocation_context: InvocationContext,
       code_execution_input: CodeExecutionInput,
   ) -> CodeExecutionResult:
+    if self.sandbox_resource_name is None:
+      sandbox_name = invocation_context.session.state.get('sandbox_name', None)
+      create_new_sandbox = False
+      if sandbox_name is None:
+        create_new_sandbox = True
+      else:
+        # Check if the sandbox is still running OR already expired due to ttl.
+        sandbox = self._get_api_client().agent_engines.sandboxes.get(
+            name=sandbox_name
+        )
+        if not sandbox or sandbox.state != 'STATE_RUNNING':
+          create_new_sandbox = True
+
+      if create_new_sandbox:
+        operation = self._get_api_client().agent_engines.sandboxes.create(
+            spec={'code_execution_environment': {}},
+            name=self.agent_engine_resource_name,
+            config=types.CreateAgentEngineSandboxConfig(
+                # VertexAiSessionService has a default TTL of 1 year, so we set
+                # the sandbox TTL to 1 year as well. For the current code
+                # execution sandbox, if it hasn't been used for 14 days, the
+                # state will be lost.
+                display_name='default_sandbox',
+                ttl='31536000s',
+            ),
+        )
+        self.sandbox_resource_name = operation.response.name
+        invocation_context.session.state['sandbox_name'] = (
+            self.sandbox_resource_name
+        )
+      else:
+        self.sandbox_resource_name = sandbox_name
+
     # Execute the code.
     input_data = {
         'code': code_execution_input.code,

tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py

@@ -19,6 +19,7 @@
 from google.adk.agents.invocation_context import InvocationContext
 from google.adk.code_executors.agent_engine_sandbox_code_executor import AgentEngineSandboxCodeExecutor
 from google.adk.code_executors.code_execution_utils import CodeExecutionInput
+from google.adk.sessions.session import Session
 import pytest
 
 
@@ -27,6 +28,10 @@ def mock_invocation_context() -> InvocationContext:
   """Fixture for a mock InvocationContext."""
   mock = MagicMock(spec=InvocationContext)
   mock.invocation_id = "test-invocation-123"
+  session = MagicMock(spec=Session)
+  mock.session = session
+  session.state = []
+
   return mock
 
 
@@ -118,3 +123,131 @@ def test_execute_code_success(
         name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789",
         input_data={"code": 'print("hello world")'},
     )
+
+  @patch("vertexai.Client")
+  def test_execute_code_recreates_sandbox_when_get_returns_none(
+      self,
+      mock_vertexai_client,
+      mock_invocation_context,
+  ):
+    # Setup Mocks
+    mock_api_client = MagicMock()
+    mock_vertexai_client.return_value = mock_api_client
+
+    # Existing sandbox name stored in session, but get() will return None
+    existing_sandbox_name = "projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/old"
+    mock_invocation_context.session.state = {
+        "sandbox_name": existing_sandbox_name
+    }
+
+    # Mock get to return None (simulating missing/expired sandbox)
+    mock_api_client.agent_engines.sandboxes.get.return_value = None
+
+    # Mock create operation to return a new sandbox resource name
+    operation_mock = MagicMock()
+    created_sandbox_name = "projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789"
+    operation_mock.response.name = created_sandbox_name
+    mock_api_client.agent_engines.sandboxes.create.return_value = operation_mock
+
+    # Mock execute_code response
+    mock_response = MagicMock()
+    mock_json_output = MagicMock()
+    mock_json_output.mime_type = "application/json"
+    mock_json_output.data = json.dumps(
+        {"stdout": "recreated sandbox run", "stderr": ""}
+    ).encode("utf-8")
+    mock_json_output.metadata = None
+    mock_response.outputs = [mock_json_output]
+    mock_api_client.agent_engines.sandboxes.execute_code.return_value = (
+        mock_response
+    )
+
+    # Execute using agent_engine_resource_name so a sandbox can be created
+    executor = AgentEngineSandboxCodeExecutor(
+        agent_engine_resource_name=(
+            "projects/123/locations/us-central1/reasoningEngines/456"
+        )
+    )
+    code_input = CodeExecutionInput(code='print("hello world")')
+    result = executor.execute_code(mock_invocation_context, code_input)
+
+    # Assert get was called for the existing sandbox
+    mock_api_client.agent_engines.sandboxes.get.assert_called_once_with(
+        name=existing_sandbox_name
+    )
+
+    # Assert create was called and session updated with new sandbox
+    mock_api_client.agent_engines.sandboxes.create.assert_called_once()
+    assert executor.sandbox_resource_name == created_sandbox_name
+    assert (
+        mock_invocation_context.session.state["sandbox_name"]
+        == created_sandbox_name
+    )
+
+    # Assert execute_code used the created sandbox name
+    mock_api_client.agent_engines.sandboxes.execute_code.assert_called_once_with(
+        name=created_sandbox_name,
+        input_data={"code": 'print("hello world")'},
+    )
+
+  @patch("vertexai.Client")
+  def test_execute_code_creates_sandbox_if_missing(
+      self,
+      mock_vertexai_client,
+      mock_invocation_context,
+  ):
+    # Setup Mocks
+    mock_api_client = MagicMock()
+    mock_vertexai_client.return_value = mock_api_client
+
+    # Mock create operation to return a sandbox resource name
+    operation_mock = MagicMock()
+    created_sandbox_name = "projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789"
+    operation_mock.response.name = created_sandbox_name
+    mock_api_client.agent_engines.sandboxes.create.return_value = operation_mock
+
+    # Mock execute_code response
+    mock_response = MagicMock()
+    mock_json_output = MagicMock()
+    mock_json_output.mime_type = "application/json"
+    mock_json_output.data = json.dumps(
+        {"stdout": "created sandbox run", "stderr": ""}
+    ).encode("utf-8")
+    mock_json_output.metadata = None
+    mock_response.outputs = [mock_json_output]
+    mock_api_client.agent_engines.sandboxes.execute_code.return_value = (
+        mock_response
+    )
+
+    # Ensure session.state behaves like a dict for storing sandbox_name
+    mock_invocation_context.session.state = {}
+
+    # Execute using agent_engine_resource_name so a sandbox will be created
+    executor = AgentEngineSandboxCodeExecutor(
+        agent_engine_resource_name=(
+            "projects/123/locations/us-central1/reasoningEngines/456"
+        ),
+        sandbox_resource_name=None,
+    )
+    code_input = CodeExecutionInput(code='print("hello world")')
+    result = executor.execute_code(mock_invocation_context, code_input)
+
+    # Assert sandbox creation was called and session state updated
+    mock_api_client.agent_engines.sandboxes.create.assert_called_once()
+    create_call_kwargs = (
+        mock_api_client.agent_engines.sandboxes.create.call_args.kwargs
+    )
+    assert create_call_kwargs["name"] == (
+        "projects/123/locations/us-central1/reasoningEngines/456"
+    )
+    assert executor.sandbox_resource_name == created_sandbox_name
+    assert (
+        mock_invocation_context.session.state["sandbox_name"]
+        == created_sandbox_name
+    )
+
+    # Assert execute_code used the created sandbox name
+    mock_api_client.agent_engines.sandboxes.execute_code.assert_called_once_with(
+        name=created_sandbox_name,
+        input_data={"code": 'print("hello world")'},
+    )