diff --git a/pyproject.toml b/pyproject.toml index 46064446..a3df7d63 100644 --- a/pyproject.toml +++ b/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. diff --git a/src/google/adk/code_executors/gke_code_executor.py b/src/google/adk/code_executors/gke_code_executor.py index 1dc46878..b44aa193 100644 --- a/src/google/adk/code_executors/gke_code_executor.py +++ b/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, diff --git a/tests/unittests/code_executors/test_gke_code_executor.py b/tests/unittests/code_executors/test_gke_code_executor.py index 3d62fd8d..300780ca 100644 --- a/tests/unittests/code_executors/test_gke_code_executor.py +++ b/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")