Skip to content
66 changes: 66 additions & 0 deletions src/google/adk/code_executors/isolated_code_executor.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
from __future__ import annotations

from contextlib import redirect_stdout
import io
import re
from typing import Any

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

import sys
import subprocess

#Don't think this is needed anymore but keeping it around just in case.

# def _prepare_globals(code: str, globals_: dict[str, Any]) -> None:
# """Prepare globals for code execution, injecting __name__ if needed."""
# if re.search(r"if\s+__name__\s*==\s*['\"]__main__['\"]", code):
# globals_['__name__'] = '__main__'

class IsolatedCodeExecutor(BaseCodeExecutor):
"""A code executor that safely executes code in an isolated environment through
the current local context."""

# 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
)
Comment on lines +48 to +52

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Consider adding a timeout to the subprocess.run call to prevent indefinite hanging in case the executed code enters an infinite loop or takes too long to execute. This will improve the robustness of the executor.

Also, it might be useful to capture and log the return code of the subprocess for debugging purposes.

    process_result = subprocess.run(
    [sys.executable, "-c", code],
    capture_output=True,
    text=True, # Enables decoding of stdout and stderr as text
    timeout=30 # Add a timeout to prevent indefinite hanging
    )

    if process_result.returncode != 0:
      print(f"Code execution failed with return code: {process_result.returncode}")


# Collect the final result.
return CodeExecutionResult(
stdout=process_result.stdout,
stderr=process_result.stderr,
output_files=[],
Comment on lines +55 to +58

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

It's important to handle potential exceptions that might occur during the code execution within the subprocess. For example, the code might raise an exception that isn't properly propagated back to the main process. Consider adding a try-except block around the subprocess.run call to catch and handle such exceptions, providing more informative error messages in the CodeExecutionResult.

Suggested change
return CodeExecutionResult(
stdout=process_result.stdout,
stderr=process_result.stderr,
output_files=[],
try:
process_result = subprocess.run(
[sys.executable, "-c", code],
capture_output=True,
text=True
)
except subprocess.TimeoutExpired as e:
return CodeExecutionResult(
stdout="",
stderr=f"Code execution timed out: {e}",
output_files=[],
)
except Exception as e:
return CodeExecutionResult(
stdout="",
stderr=f"Code execution failed: {e}",
output_files=[],
)

)