diff --git a/contributing/samples/agent_engine_code_execution/agent.py b/contributing/samples/agent_engine_code_execution/agent.py index d85989eb2d..a627fa9834 100644 --- a/contributing/samples/agent_engine_code_execution/agent.py +++ b/contributing/samples/agent_engine_code_execution/agent.py @@ -91,5 +91,8 @@ def base_system_instruction(): # 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", + # Optional: Set a TTL for the sandbox to automatically clean up resources. + # Format: duration string like "3600s" for 1 hour. + # sandbox_ttl="3600s", ), ) 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 f601d0455a..8de79d9f3f 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 @@ -38,14 +38,20 @@ 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 + sandbox_ttl: The time-to-live for the sandbox. The expiration time is + computed as: now + TTL. Format should be a duration string like "3600s" + for 1 hour. Only used when creating a new sandbox with + agent_engine_resource_name. """ sandbox_resource_name: str = None + sandbox_ttl: Optional[str] = None def __init__( self, sandbox_resource_name: Optional[str] = None, agent_engine_resource_name: Optional[str] = None, + sandbox_ttl: Optional[str] = None, **data, ): """Initializes the AgentEngineSandboxCodeExecutor. @@ -60,6 +66,9 @@ def __init__( projects/123/locations/us-central1/reasoningEngines/456, when both sandbox_resource_name and agent_engine_resource_name are set, agent_engine_resource_name will be ignored. + sandbox_ttl: The time-to-live for the sandbox. The expiration time is + computed as: now + TTL. Format should be a duration string like "3600s" + for 1 hour. Only used when creating a new sandbox. **data: Additional keyword arguments to be passed to the base class. """ super().__init__(**data) @@ -81,13 +90,13 @@ def __init__( agent_engine_resource_name, agent_engine_resource_name_pattern ) ) - # @TODO - Add TTL for sandbox creation after it is available - # in SDK. + self.sandbox_ttl = sandbox_ttl 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' + display_name='default_sandbox', + ttl=sandbox_ttl, ), ) self.sandbox_resource_name = operation.response.name 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 c948060184..2180143188 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 @@ -118,3 +118,69 @@ def test_execute_code_success( name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789", input_data={"code": 'print("hello world")'}, ) + + @patch("vertexai.types.CreateAgentEngineSandboxConfig") + @patch("vertexai.Client") + def test_init_with_agent_engine_resource_name_and_ttl( + self, + mock_vertexai_client, + mock_sandbox_config, + ): + """Tests sandbox creation with agent_engine_resource_name and TTL.""" + # Setup Mocks + mock_api_client = MagicMock() + mock_vertexai_client.return_value = mock_api_client + mock_operation = MagicMock() + mock_operation.response.name = ( + "projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" + ) + mock_api_client.agent_engines.sandboxes.create.return_value = mock_operation + + # Execute + executor = AgentEngineSandboxCodeExecutor( + agent_engine_resource_name="projects/123/locations/us-central1/reasoningEngines/456", + sandbox_ttl="3600s", + ) + + # Assert + assert executor.sandbox_resource_name == ( + "projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" + ) + assert executor.sandbox_ttl == "3600s" + mock_sandbox_config.assert_called_once_with( + display_name="default_sandbox", + ttl="3600s", + ) + + @patch("vertexai.types.CreateAgentEngineSandboxConfig") + @patch("vertexai.Client") + def test_init_with_agent_engine_resource_name_without_ttl( + self, + mock_vertexai_client, + mock_sandbox_config, + ): + """Tests sandbox creation with agent_engine_resource_name and no TTL.""" + # Setup Mocks + mock_api_client = MagicMock() + mock_vertexai_client.return_value = mock_api_client + mock_operation = MagicMock() + mock_operation.response.name = ( + "projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" + ) + mock_api_client.agent_engines.sandboxes.create.return_value = mock_operation + + # Execute + executor = AgentEngineSandboxCodeExecutor( + agent_engine_resource_name="projects/123/locations/us-central1/reasoningEngines/456", + ) + + # Assert + assert executor.sandbox_resource_name == ( + "projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" + ) + assert executor.sandbox_ttl is None + mock_sandbox_config.assert_called_once_with( + display_name="default_sandbox", + ttl=None, + ) +