feat: execute-type param addition in GkeCodeExecutor

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/4111 from SHRUTI6991:execute-type/param-addition b1ec403e0927767d17c11cb9e894f6ccb4f08dd2
PiperOrigin-RevId: 877476098

pyproject.toml

@@ -155,6 +155,7 @@ extensions = [
   "crewai[tools];python_version>='3.11' and python_version<'3.12'",       # For CrewaiTool; chromadb/pypika fail on 3.12+
   "docker>=7.0.0",                              # For ContainerCodeExecutor
   "kubernetes>=29.0.0",                         # For GkeCodeExecutor
+  "k8s-agent-sandbox>=0.1.1.post2",             # For GkeCodeExecutor sandbox mode
   "langgraph>=0.2.60, <0.4.8",                  # For LangGraphAgent
   "litellm>=1.75.5, <2.0.0",                    # For LiteLlm class. Currently has OpenAI limitations. TODO: once LiteLlm fix it
   "llama-index-readers-file>=0.4.0",            # For retrieval using LlamaIndex.

src/google/adk/code_executors/gke_code_executor.py

@@ -19,12 +19,24 @@ import uuid
 
 import kubernetes as k8s
 from kubernetes.watch import Watch
+from pydantic import field_validator
+from typing_extensions import Literal
+from typing_extensions import override
+from typing_extensions import TYPE_CHECKING
 
 from ..agents.invocation_context import InvocationContext
 from .base_code_executor import BaseCodeExecutor
 from .code_execution_utils import CodeExecutionInput
 from .code_execution_utils import CodeExecutionResult
 
+try:
+  from agentic_sandbox import SandboxClient
+except ImportError:
+  SandboxClient = None
+
+if TYPE_CHECKING:
+  from agentic_sandbox import SandboxClient
+
 # Expose these for tests to monkeypatch.
 client = k8s.client
 config = k8s.config
@@ -36,9 +48,19 @@ logger = logging.getLogger("google_adk." + __name__)
 class GkeCodeExecutor(BaseCodeExecutor):
   """Executes Python code in a secure gVisor-sandboxed Pod on GKE.
 
-  This executor securely runs code by dynamically creating a Kubernetes Job for
-  each execution request. The user's code is mounted via a ConfigMap, and the
-  Pod is hardened with a strict security context and resource limits.
+  This executor supports two modes of execution: 'job' and 'sandbox'.
+
+  Job Mode (default):
+  Securely runs code by dynamically creating a Kubernetes Job for each execution
+  request. The user's code is mounted via a ConfigMap, and the Pod is hardened
+  with a strict security context and resource limits.
+
+  Sandbox Mode:
+  Executes code using the Agent Sandbox Client. This mode requires additional
+  infrastructure to be deployed in the cluster, specifically:
+  - Agent-sandbox controller
+  - Sandbox templates (e.g., python-sandbox-template)
+  - Sandbox router and gateway
 
   Key Features:
   - Sandboxed execution using the gVisor runtime.
@@ -70,6 +92,7 @@ class GkeCodeExecutor(BaseCodeExecutor):
   namespace: str = "default"
   image: str = "python:3.11-slim"
   timeout_seconds: int = 300
+  executor_type: Literal["job", "sandbox"] = "job"
   cpu_requested: str = "200m"
   mem_requested: str = "256Mi"
   # The maximum CPU the container can use, in "millicores". 1000m is 1 full CPU core.
@@ -79,6 +102,10 @@ class GkeCodeExecutor(BaseCodeExecutor):
   kubeconfig_path: str | None = None
   kubeconfig_context: str | None = None
 
+  # Sandbox constants
+  sandbox_gateway_name: str | None = None
+  sandbox_template: str | None = "python-sandbox-template"
+
   _batch_v1: k8s.client.BatchV1Api
   _core_v1: k8s.client.CoreV1Api
 
@@ -136,10 +163,46 @@ class GkeCodeExecutor(BaseCodeExecutor):
     self._batch_v1 = client.BatchV1Api()
     self._core_v1 = client.CoreV1Api()
 
-  def execute_code(
-      self,
-      invocation_context: InvocationContext,
-      code_execution_input: CodeExecutionInput,
+  @field_validator("executor_type")
+  @classmethod
+  def _check_sandbox_dependency(cls, v: str) -> str:
+    if v == "sandbox" and SandboxClient is None:
+      raise ImportError(
+          "k8s-agent-sandbox not found. To use Agent Sandbox, please install"
+          " google-adk with the extensions extra: pip install"
+          " google-adk[extensions]"
+      )
+    return v
+
+  def _execute_in_sandbox(self, code: str) -> CodeExecutionResult:
+    """Executes code using Agent Sandbox Client."""
+    try:
+      with SandboxClient(
+          template_name=self.sandbox_template,
+          gateway_name=self.sandbox_gateway_name,
+          namespace=self.namespace,
+      ) as sandbox:
+        # Execute the code as a python script
+        sandbox.write("script.py", code)
+        result = sandbox.run("python3 script.py")
+
+        return CodeExecutionResult(stdout=result.stdout, stderr=result.stderr)
+    except RuntimeError as e:
+      logger.error(
+          "SandboxClient failed to initialize or find gateway", exc_info=True
+      )
+      raise RuntimeError(f"Sandbox infrastructure error: {e}") from e
+    except TimeoutError as e:
+      logger.error("Sandbox timed out", exc_info=True)
+      # Returning a result instead of raising allows the Agent to process
+      # the error gracefully.
+      return CodeExecutionResult(stderr=f"Sandbox timed out: {e}")
+    except Exception as e:
+      logger.error("Sandbox execution failed: %s", e, exc_info=True)
+      raise
+
+  def _execute_as_job(
+      self, code: str, invocation_context: InvocationContext
   ) -> CodeExecutionResult:
     """Orchestrates the secure execution of a code snippet on GKE."""
     job_name = f"adk-exec-{uuid.uuid4().hex[:10]}"
@@ -150,7 +213,7 @@ class GkeCodeExecutor(BaseCodeExecutor):
       # 1. Create a ConfigMap to mount LLM-generated code into the Pod.
       # 2. Create a Job that runs the code from the ConfigMap.
       # 3. Set the Job as the ConfigMap's owner for automatic cleanup.
-      self._create_code_configmap(configmap_name, code_execution_input.code)
+      self._create_code_configmap(configmap_name, code)
       job_manifest = self._create_job_manifest(
           job_name, configmap_name, invocation_context
       )
@@ -162,7 +225,6 @@ class GkeCodeExecutor(BaseCodeExecutor):
       logger.info(
           f"Submitted Job '{job_name}' to namespace '{self.namespace}'."
       )
-      logger.debug("Executing code:\n```\n%s\n```", code_execution_input.code)
       return self._watch_job_completion(job_name)
 
     except ApiException as e:
@@ -186,6 +248,20 @@ class GkeCodeExecutor(BaseCodeExecutor):
           stderr=f"An unexpected executor error occurred: {e}"
       )
 
+  @override
+  def execute_code(
+      self,
+      invocation_context: InvocationContext,
+      code_execution_input: CodeExecutionInput,
+  ) -> CodeExecutionResult:
+    """Overrides the base method to route execution based on executor_type."""
+    code = code_execution_input.code
+    if self.executor_type == "sandbox":
+      return self._execute_in_sandbox(code)
+    else:
+      # Fallback to existing GKE Job logic
+      return self._execute_as_job(code, invocation_context)
+
   def _create_job_manifest(
       self,
       job_name: str,

tests/unittests/code_executors/test_gke_code_executor.py

@@ -71,19 +71,74 @@ class TestGkeCodeExecutor:
     assert executor.timeout_seconds == 300
     assert executor.cpu_requested == "200m"
     assert executor.mem_limit == "512Mi"
+    assert executor.executor_type == "job"
 
-  def test_init_with_overrides(self):
+  @patch("google.adk.code_executors.gke_code_executor.SandboxClient")
+  def test_init_with_overrides(self, mock_sandbox_client):
     """Tests that class attributes can be overridden at instantiation."""
     executor = GkeCodeExecutor(
         namespace="test-ns",
         image="custom-python:latest",
         timeout_seconds=60,
         cpu_limit="1000m",
+        executor_type="sandbox",
     )
     assert executor.namespace == "test-ns"
     assert executor.image == "custom-python:latest"
     assert executor.timeout_seconds == 60
     assert executor.cpu_limit == "1000m"
+    assert executor.executor_type == "sandbox"
+    assert executor.sandbox_template == "python-sandbox-template"
+
+  def test_init_backward_compatibility(self):
+    """Tests that the executor can be initialized with positional arguments."""
+    executor = GkeCodeExecutor(
+        "/path/to/kubeconfig",
+        "test-context",
+        namespace="test-ns",
+        image="test-image",
+        timeout_seconds=100,
+        executor_type="job",
+        cpu_requested="100m",
+        mem_requested="128Mi",
+        cpu_limit="200m",
+        mem_limit="256Mi",
+    )
+    assert executor.namespace == "test-ns"
+    assert executor.image == "test-image"
+    assert executor.timeout_seconds == 100
+    assert executor.executor_type == "job"
+    assert executor.cpu_requested == "100m"
+    assert executor.mem_requested == "128Mi"
+    assert executor.cpu_limit == "200m"
+    assert executor.mem_limit == "256Mi"
+    assert executor.kubeconfig_path == "/path/to/kubeconfig"
+    assert executor.kubeconfig_context == "test-context"
+
+  def test_init_partial_positional_args(self):
+    """Tests initialization with partial positional arguments."""
+    executor = GkeCodeExecutor("/path/to/kubeconfig")
+    assert executor.kubeconfig_path == "/path/to/kubeconfig"
+    assert executor.kubeconfig_context is None
+
+  def test_init_mixed_args(self):
+    """Tests initialization with mixed positional and keyword arguments."""
+    executor = GkeCodeExecutor(
+        "/path/to/kubeconfig",
+        kubeconfig_context="test-context",
+        namespace="test-ns",
+    )
+    assert executor.kubeconfig_path == "/path/to/kubeconfig"
+
+  def test_init_sandbox_missing_dependency(self):
+    """Tests that init raises ImportError if k8s-agent-sandbox is missing."""
+    with patch(
+        "google.adk.code_executors.gke_code_executor.SandboxClient", None
+    ):
+      with pytest.raises(ImportError, match="k8s-agent-sandbox not found"):
+        GkeCodeExecutor(executor_type="sandbox")
+
+        GkeCodeExecutor(executor_type="sandbox")
 
   @patch("google.adk.code_executors.gke_code_executor.Watch")
   def test_execute_code_success(
@@ -225,3 +280,170 @@ class TestGkeCodeExecutor:
     assert sec_context.allow_privilege_escalation is False
     assert sec_context.read_only_root_filesystem is True
     assert sec_context.capabilities.drop == ["ALL"]
+
+  @patch("google.adk.code_executors.gke_code_executor.SandboxClient")
+  def test_execute_code_forks_to_sandbox(
+      self,
+      mock_sandbox_client,
+      mock_invocation_context,
+      mock_k8s_clients,
+  ):
+    """Tests execute_code with executor_type='sandbox'.
+
+    Verifies that execute_code uses SandboxClient when executor_type is set to
+    'sandbox'.
+    """
+    # Setup Sandbox mock
+    mock_sandbox_instance = (
+        mock_sandbox_client.return_value.__enter__.return_value
+    )
+    mock_run_result = MagicMock()
+    mock_run_result.stdout = "sandbox stdout"
+    mock_run_result.stderr = None
+    mock_sandbox_instance.run.return_value = mock_run_result
+
+    # Instantiate with sandbox type
+    executor = GkeCodeExecutor(executor_type="sandbox")
+    code_input = CodeExecutionInput(code='print("sandbox")')
+
+    # Execute
+    result = executor.execute_code(mock_invocation_context, code_input)
+
+    # Assertions
+    assert result.stdout == "sandbox stdout"
+
+    # Verify SandboxClient was used
+    mock_sandbox_client.assert_called_once()
+    mock_sandbox_instance.run.assert_called_once()
+
+    # Verify Job path was NOT taken
+    mock_k8s_clients["batch_v1"].create_namespaced_job.assert_not_called()
+
+  @patch("google.adk.code_executors.gke_code_executor.SandboxClient")
+  def test_execute_code_sandbox_connection_error(
+      self,
+      mock_sandbox_client,
+      mock_invocation_context,
+  ):
+    """Tests handling of exceptions from SandboxClient."""
+    # Setup Sandbox mock to raise exception
+    mock_sandbox_client.return_value.__enter__.side_effect = Exception(
+        "Connection failed"
+    )
+
+    # Instantiate with sandbox type
+    executor = GkeCodeExecutor(executor_type="sandbox")
+    code_input = CodeExecutionInput(code='print("sandbox")')
+
+    # Execute & Assert
+    with pytest.raises(Exception, match="Connection failed"):
+      executor.execute_code(mock_invocation_context, code_input)
+
+  @patch("google.adk.code_executors.gke_code_executor.SandboxClient")
+  def test_execute_code_sandbox_runtime_error(
+      self,
+      mock_sandbox_client,
+      mock_invocation_context,
+  ):
+    """Tests handling of RuntimeError from SandboxClient."""
+    mock_sandbox_client.return_value.__enter__.side_effect = RuntimeError(
+        "Gateway not found"
+    )
+
+    executor = GkeCodeExecutor(executor_type="sandbox")
+    code_input = CodeExecutionInput(code='print("sandbox")')
+
+    with pytest.raises(
+        RuntimeError, match="Sandbox infrastructure error: Gateway not found"
+    ):
+      executor.execute_code(mock_invocation_context, code_input)
+
+  @patch("google.adk.code_executors.gke_code_executor.SandboxClient")
+  def test_execute_code_sandbox_timeout_error(
+      self,
+      mock_sandbox_client,
+      mock_invocation_context,
+  ):
+    """Tests handling of TimeoutError from SandboxClient."""
+    mock_sandbox_client.return_value.__enter__.side_effect = TimeoutError(
+        "Execution timed out"
+    )
+
+    executor = GkeCodeExecutor(executor_type="sandbox")
+    code_input = CodeExecutionInput(code='print("sandbox")')
+
+    result = executor.execute_code(mock_invocation_context, code_input)
+
+    assert result.stdout == ""
+    assert "Sandbox timed out: Execution timed out" in result.stderr
+
+  @patch("google.adk.code_executors.gke_code_executor.SandboxClient")
+  @patch("google.adk.code_executors.gke_code_executor.Watch")
+  def test_execute_code_forks_to_job(
+      self,
+      mock_watch,
+      mock_sandbox_client,
+      mock_invocation_context,
+      mock_k8s_clients,
+  ):
+    """Tests that execute_code uses K8s Job when executor_type='job'."""
+    # Setup K8s Job mocks (success path)
+    mock_job = MagicMock()
+    mock_job.status.succeeded = True
+    mock_watch.return_value.stream.return_value = [{"object": mock_job}]
+
+    mock_pod = MagicMock()
+    mock_pod.metadata.name = "pod-1"
+    mock_k8s_clients["core_v1"].list_namespaced_pod.return_value.items = [
+        mock_pod
+    ]
+    mock_k8s_clients["core_v1"].read_namespaced_pod_log.return_value = (
+        "job stdout"
+    )
+
+    # Instantiate with job type
+    executor = GkeCodeExecutor(executor_type="job")
+    code_input = CodeExecutionInput(code='print("job")')
+
+    # Execute
+    result = executor.execute_code(mock_invocation_context, code_input)
+
+    # Assertions
+    assert result.stdout == "job stdout"
+
+    # Verify Job path WAS taken
+    mock_k8s_clients["batch_v1"].create_namespaced_job.assert_called_once()
+
+    # Verify SandboxClient was NOT used
+    mock_sandbox_client.assert_not_called()
+
+  @patch("google.adk.code_executors.gke_code_executor.SandboxClient")
+  def test_execute_in_sandbox_returns_stderr(
+      self,
+      mock_sandbox_client,
+      mock_invocation_context,
+  ):
+    """Tests that stderr from the sandbox run is propagated to the result."""
+    # Setup Sandbox mock
+    mock_sandbox_instance = (
+        mock_sandbox_client.return_value.__enter__.return_value
+    )
+    mock_run_result = MagicMock()
+    mock_run_result.stdout = ""
+    mock_run_result.stderr = "oops\n"
+    mock_sandbox_instance.run.return_value = mock_run_result
+
+    # Instantiate with sandbox type
+    executor = GkeCodeExecutor(executor_type="sandbox")
+    code_input = CodeExecutionInput(
+        code="import sys; print('oops', file=sys.stderr)"
+    )
+
+    # Execute
+    result = executor.execute_code(mock_invocation_context, code_input)
+
+    # Assertions
+    assert result.stdout == ""
+    assert result.stderr == "oops\n"
+    mock_sandbox_instance.write.assert_called_with("script.py", code_input.code)
+    mock_sandbox_instance.run.assert_called_with("python3 script.py")