diff --git a/src/google/adk/code_executors/isolated_code_executor.py b/src/google/adk/code_executors/isolated_code_executor.py new file mode 100644 index 0000000000..7b74786c43 --- /dev/null +++ b/src/google/adk/code_executors/isolated_code_executor.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +import subprocess +import sys + +from pydantic import Field +from typing_extensions import override + +from ..agents.invocation_context import InvocationContext +from .base_code_executor import BaseCodeExecutor +from .code_execution_utils import CodeExecutionInput +from .code_execution_utils import CodeExecutionResult + + +class IsolatedCodeExecutor(BaseCodeExecutor): + """A code executor that executes code in an isolated process. + + This provides memory isolation from the main application, but it is not a + full security sandbox. The executed code runs with the same permissions as the + main application and can access the filesystem, network, etc. + """ + + # Overrides the BaseCodeExecutor attribute: this executor cannot be stateful. + stateful: bool = Field(default=False, frozen=True, exclude=True) + + # Overrides the BaseCodeExecutor attribute: this executor cannot + # optimize_data_file. + optimize_data_file: bool = Field(default=False, frozen=True, exclude=True) + + def __init__(self, **data): + """Initializes the IsolatedCodeExecutor.""" + if 'stateful' in data and data['stateful']: + raise ValueError('Cannot set `stateful=True` in IsolatedCodeExecutor.') + if 'optimize_data_file' in data and data['optimize_data_file']: + raise ValueError( + 'Cannot set `optimize_data_file=True` in IsolatedCodeExecutor.' + ) + super().__init__(**data) + + @override + def execute_code( + self, + invocation_context: InvocationContext, + code_execution_input: CodeExecutionInput, + ) -> CodeExecutionResult: + # Executes code by spawning a new python interpreter process. + code = code_execution_input.code + process_result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True + ) + + + return CodeExecutionResult( + stdout=process_result.stdout, + stderr=process_result.stderr, + output_files=[], + ) diff --git a/tests/unittests/code_executors/test_isolated_code_executor.py b/tests/unittests/code_executors/test_isolated_code_executor.py new file mode 100644 index 0000000000..70a1553b8a --- /dev/null +++ b/tests/unittests/code_executors/test_isolated_code_executor.py @@ -0,0 +1,112 @@ +import os +from unittest.mock import MagicMock + +from google.adk.agents.base_agent import BaseAgent +from google.adk.agents.invocation_context import InvocationContext +from google.adk.code_executors.code_execution_utils import CodeExecutionInput +from google.adk.code_executors.code_execution_utils import CodeExecutionResult +from google.adk.code_executors.isolated_code_executor import IsolatedCodeExecutor +from google.adk.sessions.base_session_service import BaseSessionService +from google.adk.sessions.session import Session +import pytest + + +@pytest.fixture +def mock_invocation_context() -> InvocationContext: + """Provides a mock InvocationContext.""" + mock_agent = MagicMock(spec=BaseAgent) + mock_session = MagicMock(spec=Session) + mock_session_service = MagicMock(spec=BaseSessionService) + return InvocationContext( + invocation_id="test_invocation", + agent=mock_agent, + session=mock_session, + session_service=mock_session_service, + ) + + +class TestIsolatedCodeExecutor: + + def test_init_default(self): + executor = IsolatedCodeExecutor() + assert not executor.stateful + assert not executor.optimize_data_file + + def test_init_stateful_raises_error(self): + with pytest.raises( + ValueError, + match="Cannot set `stateful=True` in IsolatedCodeExecutor.", + ): + IsolatedCodeExecutor(stateful=True) + + def test_init_optimize_data_file_raises_error(self): + with pytest.raises( + ValueError, + match=( + "Cannot set `optimize_data_file=True` in IsolatedCodeExecutor." + ), + ): + IsolatedCodeExecutor(optimize_data_file=True) + + def test_execute_code_simple_print( + self, mock_invocation_context: InvocationContext + ): + executor = IsolatedCodeExecutor() + code_input = CodeExecutionInput(code='print("hello world")') + result = executor.execute_code(mock_invocation_context, code_input) + + assert isinstance(result, CodeExecutionResult) + assert result.stdout == "hello world\n" + assert result.stderr == "" + assert result.output_files == [] + + def test_execute_code_with_error( + self, mock_invocation_context: InvocationContext + ): + executor = IsolatedCodeExecutor() + code_input = CodeExecutionInput(code='raise ValueError("Test error")') + result = executor.execute_code(mock_invocation_context, code_input) + + assert isinstance(result, CodeExecutionResult) + assert result.stdout == "" + assert "Test error" in result.stderr + assert result.output_files == [] + + def test_execute_code_variable_assignment( + self, mock_invocation_context: InvocationContext + ): + executor = IsolatedCodeExecutor() + code_input = CodeExecutionInput(code="x = 10\nprint(x * 2)") + result = executor.execute_code(mock_invocation_context, code_input) + + assert result.stdout == "20\n" + assert result.stderr == "" + + def test_execute_code_empty(self, mock_invocation_context: InvocationContext): + executor = IsolatedCodeExecutor() + code_input = CodeExecutionInput(code="") + result = executor.execute_code(mock_invocation_context, code_input) + assert result.stdout == "" + assert result.stderr == "" + + def test_execute_code_with_import( + self, mock_invocation_context: InvocationContext + ): + executor = IsolatedCodeExecutor() + code = "import os; print(os.linesep)" + code_input = CodeExecutionInput(code=code) + result = executor.execute_code(mock_invocation_context, code_input) + + assert result.stdout.strip() == os.linesep.strip() + assert result.stderr == "" + + def test_execute_code_multiline_output( + self, mock_invocation_context: InvocationContext + ): + executor = IsolatedCodeExecutor() + code = 'print("line 1")\nprint("line 2")' + code_input = CodeExecutionInput(code=code) + result = executor.execute_code(mock_invocation_context, code_input) + + assert result.stdout == "line 1\nline 2\n" + assert result.stderr == ""