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

tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py

@@ -19,6 +19,7 @@ from unittest.mock import patch
 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 @@ class TestAgentEngineSandboxCodeExecutor:
         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")'},
+    )