feat: Enhance AgentEngineSandboxCodeExecutor sample to automatically provision an Agent Engine if neither agent_engine_resource_name nor sandbox_resource_name is provided

The AgentEngineSandboxCodeExecutor now has three initialization modes:
1.  Create both an Agent Engine and sandbox if neither resource name is provided.
2.  Creating a new sandbox within a provided agent_engine_resource_name.
3.  Using a provided sandbox_resource_name.

PiperOrigin-RevId: 884088248

Author: google-genai-bot

contributing/samples/agent_engine_code_execution/agent.py

@@ -89,6 +89,6 @@ def base_system_instruction():
         # "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.
-        agent_engine_resource_name="AGENT_ENGINE_RESOURCE_NAME",
+        agent_engine_resource_name=None,
     ),
 )

src/google/adk/code_executors/agent_engine_sandbox_code_executor.py

@@ -17,7 +17,9 @@
 import json
 import logging
 import mimetypes
+import os
 import re
+import threading
 from typing import Optional
 
 from typing_extensions import override
@@ -46,6 +48,7 @@ class AgentEngineSandboxCodeExecutor(BaseCodeExecutor):
   sandbox_resource_name: str = None
 
   agent_engine_resource_name: str = None
+  _agent_engine_creation_lock: Optional[threading.Lock] = None
 
   def __init__(
       self,
@@ -61,42 +64,69 @@ def __init__(
         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:
+        to create the code execution sandbox. If not set, a new Agent Engine
+        will be created automatically. Format:
         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.
       **data: Additional keyword arguments to be passed to the base class.
     """
     super().__init__(**data)
+    self._agent_engine_creation_lock = threading.Lock()
     sandbox_resource_name_pattern = r'^projects/([a-zA-Z0-9-_]+)/locations/([a-zA-Z0-9-_]+)/reasoningEngines/(\d+)/sandboxEnvironments/(\d+)$'
     agent_engine_resource_name_pattern = r'^projects/([a-zA-Z0-9-_]+)/locations/([a-zA-Z0-9-_]+)/reasoningEngines/(\d+)$'
 
+    # Case 1: sandbox_resource_name is provided.
     if sandbox_resource_name is not None:
       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:
+
+    # Case 2: Agent Engine resource name is not provided.
+    elif agent_engine_resource_name is None:
+      # The Agent Engine will be auto-created lazily within execute_code().
+      self._project_id = os.environ.get('GOOGLE_CLOUD_PROJECT')
+      self._location = os.environ.get('GOOGLE_CLOUD_LOCATION', 'us-central1')
+      self.agent_engine_resource_name = None
+
+    # Case 3: Use the provided agent_engine_resource_name.
+    else:
       self._project_id, self._location = (
           self._get_project_id_and_location_from_resource_name(
-              agent_engine_resource_name, agent_engine_resource_name_pattern
+              agent_engine_resource_name,
+              agent_engine_resource_name_pattern,
           )
       )
       self.agent_engine_resource_name = agent_engine_resource_name
-    else:
-      raise ValueError(
-          'Either sandbox_resource_name or agent_engine_resource_name must be'
-          ' set.'
-      )
 
   @override
   def execute_code(
       self,
       invocation_context: InvocationContext,
       code_execution_input: CodeExecutionInput,
   ) -> CodeExecutionResult:
+    if (
+        self.sandbox_resource_name is None
+        and self.agent_engine_resource_name is None
+    ):
+      with self._agent_engine_creation_lock:
+        if self.agent_engine_resource_name is None:
+          logger.info(
+              'No Agent Engine resource name provided. Creating a new one...'
+          )
+          try:
+            # Create a default Agent Engine.
+            created_engine = self._get_api_client().agent_engines.create()
+            self.agent_engine_resource_name = created_engine.api_resource.name
+            logger.info(
+                'Created Agent Engine: %s', self.agent_engine_resource_name
+            )
+          except Exception as e:
+            logger.error('Failed to auto-create Agent Engine: %s', e)
+            raise
     # default to the sandbox resource name if set.
     sandbox_name = self.sandbox_resource_name
     if self.sandbox_resource_name is None:

tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import json
+import os
 from unittest.mock import MagicMock
 from unittest.mock import patch
 
@@ -249,3 +250,101 @@ def test_execute_code_creates_sandbox_if_missing(
         name=created_sandbox_name,
         input_data={"code": 'print("hello world")'},
     )
+
+  def test_init_with_agent_engine_resource_name(self):
+    """Tests init when only agent_engine_resource_name is provided."""
+    agent_engine_name = (
+        "projects/123/locations/us-central1/reasoningEngines/456"
+    )
+
+    executor = AgentEngineSandboxCodeExecutor(
+        agent_engine_resource_name=agent_engine_name
+    )
+
+    # Verify the engine name is set, and sandbox remains None.
+    assert executor.agent_engine_resource_name == agent_engine_name
+    assert executor.sandbox_resource_name is None
+    assert executor._project_id == "123"
+    assert executor._location == "us-central1"
+
+  @patch("vertexai.Client")
+  @patch.dict(
+      os.environ,
+      {
+          "GOOGLE_CLOUD_PROJECT": "test-project-456",
+          "GOOGLE_CLOUD_LOCATION": "us-central1",
+      },
+  )
+  def test_execute_code_with_auto_create_agent_engine(
+      self, mock_vertexai_client, mock_invocation_context
+  ):
+    """Tests that Agent Engine is created lazily in execute_code."""
+    # Setup Mocks
+    mock_api_client = MagicMock()
+    mock_vertexai_client.return_value = mock_api_client
+
+    # Mock Engine Creation
+    mock_created_engine = MagicMock()
+    mock_created_engine.api_resource.name = "projects/test-project-456/locations/us-central1/reasoningEngines/auto-created-ae-1"
+    mock_api_client.agent_engines.create.return_value = mock_created_engine
+
+    # Mock create operation to return a sandbox resource name
+    operation_mock = MagicMock()
+    created_sandbox_name = "projects/test-project-456/locations/us-central1/reasoningEngines/auto-created-ae-1/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
+    )
+
+    # Execute
+    executor = AgentEngineSandboxCodeExecutor()
+    code_input = CodeExecutionInput(code='print("hello world")')
+    executor.execute_code(mock_invocation_context, code_input)
+
+    # Assert
+    mock_api_client.agent_engines.create.assert_called_once()
+    assert (
+        executor.agent_engine_resource_name
+        == "projects/test-project-456/locations/us-central1/reasoningEngines/auto-created-ae-1"
+    )
+    assert executor.sandbox_resource_name is None
+    mock_api_client.agent_engines.sandboxes.create.assert_called_once()
+    assert (
+        mock_invocation_context.session.state["sandbox_name"]
+        == created_sandbox_name
+    )
+
+  @patch("vertexai.Client")
+  @patch.dict(
+      os.environ,
+      {
+          "GOOGLE_CLOUD_PROJECT": "test-project-456",
+          "GOOGLE_CLOUD_LOCATION": "us-central1",
+      },
+  )
+  def test_execute_code_auto_create_agent_engine_fails(
+      self, mock_vertexai_client, mock_invocation_context
+  ):
+    """Tests error handling when auto-creating Agent Engine fails."""
+    mock_api_client = MagicMock()
+    mock_vertexai_client.return_value = mock_api_client
+    mock_api_client.agent_engines.create.side_effect = Exception(
+        "Failed to auto-create Agent Engine"
+    )
+
+    executor = AgentEngineSandboxCodeExecutor()
+    code_input = CodeExecutionInput(code='print("hello world")')
+
+    with pytest.raises(Exception, match="Failed to auto-create Agent Engine"):
+      executor.execute_code(mock_invocation_context, code_input)