Add TTL for sandbox creation in AgentEngineSandboxCodeExecutor

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",
     ),
 )

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

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,
+    )
+