diff --git a/contributing/samples/gepa/experiment.py b/contributing/samples/gepa/experiment.py index 2f5d03a772..f68b349d9c 100644 --- a/contributing/samples/gepa/experiment.py +++ b/contributing/samples/gepa/experiment.py @@ -43,7 +43,6 @@ from tau_bench.types import EnvRunResult from tau_bench.types import RunConfig import tau_bench_agent as tau_bench_agent_lib - import utils diff --git a/contributing/samples/gepa/run_experiment.py b/contributing/samples/gepa/run_experiment.py index cfd850b3a3..1bc4ee58c8 100644 --- a/contributing/samples/gepa/run_experiment.py +++ b/contributing/samples/gepa/run_experiment.py @@ -25,7 +25,6 @@ from absl import flags import experiment from google.genai import types - import utils _OUTPUT_DIR = flags.DEFINE_string( diff --git a/src/google/adk/code_executors/__init__.py b/src/google/adk/code_executors/__init__.py index edeaf5d272..b3ff33f960 100644 --- a/src/google/adk/code_executors/__init__.py +++ b/src/google/adk/code_executors/__init__.py @@ -18,6 +18,7 @@ from .base_code_executor import BaseCodeExecutor from .built_in_code_executor import BuiltInCodeExecutor +from .isolated_code_executor import IsolatedCodeExecutor from .code_executor_context import CodeExecutorContext from .unsafe_local_code_executor import UnsafeLocalCodeExecutor diff --git a/src/google/adk/code_executors/unsafe_local_code_executor.py b/src/google/adk/code_executors/unsafe_local_code_executor.py index 6dd2ae9d8c..9d111b6d32 100644 --- a/src/google/adk/code_executors/unsafe_local_code_executor.py +++ b/src/google/adk/code_executors/unsafe_local_code_executor.py @@ -17,7 +17,8 @@ from contextlib import redirect_stdout import io import logging -import re +import subprocess +import sys from typing import Any from pydantic import Field @@ -38,7 +39,10 @@ def _prepare_globals(code: str, globals_: dict[str, Any]) -> None: class UnsafeLocalCodeExecutor(BaseCodeExecutor): - """A code executor that unsafely execute code in the current local context.""" + """A code executor that unsafely execute code in the current local context. + + This executor can be configured to run code in an isolated process. + """ # Overrides the BaseCodeExecutor attribute: this executor cannot be stateful. stateful: bool = Field(default=False, frozen=True, exclude=True) @@ -47,7 +51,9 @@ class UnsafeLocalCodeExecutor(BaseCodeExecutor): # optimize_data_file. optimize_data_file: bool = Field(default=False, frozen=True, exclude=True) - def __init__(self, **data): + use_separate_process: bool = False + + def __init__(self, use_separate_process: bool = False, **data): """Initializes the UnsafeLocalCodeExecutor.""" if 'stateful' in data and data['stateful']: raise ValueError('Cannot set `stateful=True` in UnsafeLocalCodeExecutor.') @@ -55,7 +61,8 @@ def __init__(self, **data): raise ValueError( 'Cannot set `optimize_data_file=True` in UnsafeLocalCodeExecutor.' ) - super().__init__(**data) + super().__init__(use_separate_process=use_separate_process, **data) + self.use_separate_process = use_separate_process @override def execute_code( @@ -63,23 +70,38 @@ def execute_code( invocation_context: InvocationContext, code_execution_input: CodeExecutionInput, ) -> CodeExecutionResult: + if self.use_separate_process: + logger.debug( + 'Executing code in isolated process:\n```\n%s\n```', + code_execution_input.code, + ) + process_result = subprocess.run( + [sys.executable, '-c', code_execution_input.code], + capture_output=True, + text=True, + ) + return CodeExecutionResult( + stdout=process_result.stdout, + stderr=process_result.stderr, + output_files=[], + ) + logger.debug('Executing code:\n```\n%s\n```', code_execution_input.code) # Execute the code. - output = '' - error = '' + stdout_capture = io.StringIO() + stderr_capture = io.StringIO() try: globals_ = {} _prepare_globals(code_execution_input.code, globals_) - stdout = io.StringIO() - with redirect_stdout(stdout): + with redirect_stdout(stdout_capture): exec(code_execution_input.code, globals_) - output = stdout.getvalue() except Exception as e: - error = str(e) + import traceback + stderr_capture.write(traceback.format_exc()) # Collect the final result. return CodeExecutionResult( - stdout=output, - stderr=error, + stdout=stdout_capture.getvalue(), + stderr=stderr_capture.getvalue(), 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 == "" diff --git a/tests/unittests/code_executors/test_unsafe_local_code_executor.py b/tests/unittests/code_executors/test_unsafe_local_code_executor.py index eeb10b34fa..a5ba040a7a 100644 --- a/tests/unittests/code_executors/test_unsafe_local_code_executor.py +++ b/tests/unittests/code_executors/test_unsafe_local_code_executor.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - +import os from unittest.mock import MagicMock from google.adk.agents.base_agent import BaseAgent @@ -101,3 +101,83 @@ def test_execute_code_empty(self, mock_invocation_context: InvocationContext): result = executor.execute_code(mock_invocation_context, code_input) assert result.stdout == "" assert result.stderr == "" + + def test_execute_code_isolated_simple_print( + self, mock_invocation_context: InvocationContext + ): + """Test execution with use_separate_process=True.""" + executor = UnsafeLocalCodeExecutor(use_separate_process=True) + code_input = CodeExecutionInput(code='print("hello isolated world")') + result = executor.execute_code(mock_invocation_context, code_input) + + assert isinstance(result, CodeExecutionResult) + assert result.stdout == "hello isolated world\n" + assert result.stderr == "" + assert result.output_files == [] + + def test_execute_code_isolated_with_error( + self, mock_invocation_context: InvocationContext + ): + """Test error handling with use_separate_process=True.""" + executor = UnsafeLocalCodeExecutor(use_separate_process=True) + code_input = CodeExecutionInput(code='raise ValueError("Isolated error")') + result = executor.execute_code(mock_invocation_context, code_input) + + assert isinstance(result, CodeExecutionResult) + assert result.stdout == "" + assert "Isolated error" in result.stderr + assert result.output_files == [] + + def test_execute_code_isolated_with_import( + self, mock_invocation_context: InvocationContext + ): + """Test imports with use_separate_process=True.""" + executor = UnsafeLocalCodeExecutor(use_separate_process=True) + 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_if_main( + self, mock_invocation_context: InvocationContext + ): + """Test code with `if __name__ == '__main__'`.""" + executor = UnsafeLocalCodeExecutor() + code = 'if __name__ == "__main__":\n print("executed as main")' + code_input = CodeExecutionInput(code=code) + result = executor.execute_code(mock_invocation_context, code_input) + + assert result.stdout == "executed as main\n" + assert result.stderr == "" + + def test_execute_code_isolated_memory_isolation_variable( + self, mock_invocation_context: InvocationContext + ): + """Test that variables do not persist between isolated executions.""" + executor = UnsafeLocalCodeExecutor(use_separate_process=True) + + # First execution defines a variable. + code_input1 = CodeExecutionInput(code="x = 100") + executor.execute_code(mock_invocation_context, code_input1) + + # Second execution tries to access the variable. + code_input2 = CodeExecutionInput(code="print(x)") + result = executor.execute_code(mock_invocation_context, code_input2) + + assert "name 'x' is not defined" in result.stderr + + def test_execute_code_isolated_memory_isolation_function( + self, mock_invocation_context: InvocationContext + ): + """Test that functions do not persist between isolated executions.""" + executor = UnsafeLocalCodeExecutor(use_separate_process=True) + + code_input1 = CodeExecutionInput(code="def my_func(): return 'isolated'") + executor.execute_code(mock_invocation_context, code_input1) + + code_input2 = CodeExecutionInput(code="print(my_func())") + result = executor.execute_code(mock_invocation_context, code_input2) + + assert "name 'my_func' is not defined" in result.stderr