From ee8d956413473d1bbbb025a470ad882c1487d8b8 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Tue, 24 Feb 2026 11:15:35 -0800 Subject: [PATCH] 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. PiperOrigin-RevId: 874705260 --- .../agent_engine_code_execution/README | 4 +- .../agent_engine_code_execution/agent.py | 7 +- .../agent_engine_sandbox_code_executor.py | 54 ++----- ...test_agent_engine_sandbox_code_executor.py | 133 ------------------ 4 files changed, 19 insertions(+), 179 deletions(-) diff --git a/contributing/samples/agent_engine_code_execution/README b/contributing/samples/agent_engine_code_execution/README index b0443ae2..8d5a4442 100644 --- a/contributing/samples/agent_engine_code_execution/README +++ b/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://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. +* 1. Follow https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/code-execution/overview to create a code execution sandbox environment. -* 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." +* 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. ## Sample prompt diff --git a/contributing/samples/agent_engine_code_execution/agent.py b/contributing/samples/agent_engine_code_execution/agent.py index a32e4ca4..d85989eb 100644 --- a/contributing/samples/agent_engine_code_execution/agent.py +++ b/contributing/samples/agent_engine_code_execution/agent.py @@ -85,10 +85,11 @@ When plotting trends, you should make sure to sort and order the data by the x-a """, code_executor=AgentEngineSandboxCodeExecutor( - # 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. + # Replace with your sandbox resource name if you already have one. + sandbox_resource_name="SANDBOX_RESOURCE_NAME", # "projects/vertex-agent-loadtest/locations/us-central1/reasoningEngines/6842889780301135872/sandboxEnvironments/6545148628569161728", - sandbox_resource_name=None, - # Replace with agent engine resource name used for creating sandbox environment. + # Replace with agent engine resource name used for creating sandbox if + # sandbox_resource_name is not set. agent_engine_resource_name="AGENT_ENGINE_RESOURCE_NAME", ), ) diff --git a/src/google/adk/code_executors/agent_engine_sandbox_code_executor.py b/src/google/adk/code_executors/agent_engine_sandbox_code_executor.py index 9348dbc4..69d1778a 100644 --- a/src/google/adk/code_executors/agent_engine_sandbox_code_executor.py +++ b/src/google/adk/code_executors/agent_engine_sandbox_code_executor.py @@ -21,7 +21,6 @@ import re 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 @@ -39,15 +38,10 @@ 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, @@ -73,19 +67,30 @@ class AgentEngineSandboxCodeExecutor(BaseCodeExecutor): 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 ) ) - self.agent_engine_resource_name = agent_engine_resource_name + # @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 else: raise ValueError( 'Either sandbox_resource_name or agent_engine_resource_name must be' @@ -98,39 +103,6 @@ class AgentEngineSandboxCodeExecutor(BaseCodeExecutor): 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, diff --git a/tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py b/tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py index 604685fe..6022527f 100644 --- a/tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py +++ b/tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py @@ -19,7 +19,6 @@ 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 @@ -28,10 +27,6 @@ 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 @@ -123,131 +118,3 @@ 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")'}, - )