execute-type param addition in GkeCodeExecutor

pyproject.toml

@@ -27,6 +27,7 @@ dependencies = [
   # go/keep-sorted start
   "PyYAML>=6.0.2, <7.0.0",                                   # For APIHubToolset.
   "aiosqlite>=0.21.0",                                       # For SQLite database
+  "agentic_sandbox @ git+https://github.com/kubernetes-sigs/agent-sandbox.git@dbac66ecba5497ac493ca5e4ab5e0fcb1c945134#subdirectory=clients/python/agentic-sandbox-client", # For Agent Sandboxed Code Execution
   "anyio>=4.9.0, <5.0.0",                                    # For MCP Session Manager
   "authlib>=1.6.6, <2.0.0",                                  # For RestAPI Tool
   "click>=8.1.8, <9.0.0",                                    # For CLI tools

src/google/adk/code_executors/gke_code_executor.py

@@ -17,6 +17,7 @@
 import logging
 import uuid
 
+from agentic_sandbox import SandboxClient
 import kubernetes as k8s
 from kubernetes.watch import Watch
 
@@ -36,9 +37,19 @@
 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 +81,8 @@ class GkeCodeExecutor(BaseCodeExecutor):
   namespace: str = "default"
   image: str = "python:3.11-slim"
   timeout_seconds: int = 300
+  executor_type: str = "job"  # "job" or "sandbox"
+  sandbox_gateway_name: str | None = None
   cpu_requested: str = "200m"
   mem_requested: str = "256Mi"
   # The maximum CPU the container can use, in "millicores". 1000m is 1 full CPU core.
@@ -79,11 +92,17 @@ class GkeCodeExecutor(BaseCodeExecutor):
   kubeconfig_path: str | None = None
   kubeconfig_context: str | None = None
 
+  # Sandbox constants
+  sandbox_template: str | None = None
+
   _batch_v1: k8s.client.BatchV1Api
   _core_v1: k8s.client.CoreV1Api
 
   def __init__(
       self,
+      executor_type: str = "job",
+      sandbox_gateway_name: str | None = None,
+      sandbox_template: str = "python-sandbox-template",
       kubeconfig_path: str | None = None,
       kubeconfig_context: str | None = None,
       **data,
@@ -96,9 +115,17 @@ def __init__(
     3. Automatically via the default local kubeconfig file (~/.kube/config).
     """
     super().__init__(**data)
+    self.executor_type = executor_type
+    self.sandbox_gateway_name = sandbox_gateway_name
+    self.sandbox_template = sandbox_template
     self.kubeconfig_path = kubeconfig_path
     self.kubeconfig_context = kubeconfig_context
 
+    if executor_type not in ["job", "sandbox"]:
+      raise ValueError(
+          f"Invalid executor_type: '{executor_type}'. Must be 'job' or"
+          " 'sandbox'."
+      )
     if self.kubeconfig_path:
       try:
         logger.info(f"Using explicit kubeconfig from '{self.kubeconfig_path}'.")
@@ -136,10 +163,28 @@ def __init__(
     self._batch_v1 = client.BatchV1Api()
     self._core_v1 = client.CoreV1Api()
 
-  def execute_code(
-      self,
-      invocation_context: InvocationContext,
-      code_execution_input: CodeExecutionInput,
+  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
+        logger.debug("Executing code in sandbox:\n```\n%s\n```", code)
+        sandbox.write("script.py", code)
+        result = sandbox.run("python3 script.py")
+
+        return CodeExecutionResult(stdout=result.stdout, stderr=result.stderr)
+    except Exception as e:
+      logger.error("Sandbox execution failed", exc_info=True)
+      return CodeExecutionResult(
+          stderr=f"Sandbox execution failed: {str(e)}",
+      )
+
+  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 +195,7 @@ def execute_code(
       # 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 +207,7 @@ def execute_code(
       logger.info(
           f"Submitted Job '{job_name}' to namespace '{self.namespace}'."
       )
-      logger.debug("Executing code:\n```\n%s\n```", code_execution_input.code)
+      logger.debug("Executing code:\n```\n%s\n```", code)
       return self._watch_job_completion(job_name)
 
     except ApiException as e:
@@ -186,6 +231,19 @@ def execute_code(
           stderr=f"An unexpected executor error occurred: {e}"
       )
 
+  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,6 +71,7 @@ def test_init_defaults(self):
     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):
     """Tests that class attributes can be overridden at instantiation."""
@@ -79,11 +80,19 @@ def test_init_with_overrides(self):
         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_invalid_executor_type(self):
+    """Tests that init raises ValueError for invalid executor_type."""
+    with pytest.raises(ValueError, match="Invalid executor_type"):
+      GkeCodeExecutor(executor_type="invalid_type")
 
   @patch("google.adk.code_executors.gke_code_executor.Watch")
   def test_execute_code_success(
@@ -225,3 +234,131 @@ def test_create_job_manifest_structure(self, mock_invocation_context):
     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 that execute_code uses SandboxClient when executor_type='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_exception(
+      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
+    result = executor.execute_code(mock_invocation_context, code_input)
+
+    # Assertions
+    assert result.stdout == ""
+    assert "Sandbox execution failed: Connection failed" 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")