From 27548e392f8a8503609a6abce0fb8081e8fb24f5 Mon Sep 17 00:00:00 2001 From: George Weale Date: Thu, 30 Jul 2026 15:40:44 -0700 Subject: [PATCH] fix: kill runaway code on timeout in container and local executors Co-authored-by: George Weale PiperOrigin-RevId: 956773467 --- .../code_executors/container_code_executor.py | 102 +++++++- .../unsafe_local_code_executor.py | 64 ++++- .../test_container_code_executor.py | 246 ++++++++++++++++++ .../test_unsafe_local_code_executor.py | 114 ++++++++ 4 files changed, 523 insertions(+), 3 deletions(-) diff --git a/src/google/adk/code_executors/container_code_executor.py b/src/google/adk/code_executors/container_code_executor.py index 671c33ac..4d69c57e 100644 --- a/src/google/adk/code_executors/container_code_executor.py +++ b/src/google/adk/code_executors/container_code_executor.py @@ -33,6 +33,74 @@ from .code_execution_utils import CodeExecutionResult logger = logging.getLogger('google_adk.' + __name__) DEFAULT_IMAGE_TAG = 'adk-code-executor:latest' +# Reported by the supervisor below when it kills a run that hit the bound. +# Follows the convention of coreutils `timeout`; unlike 128 + SIGALRM it is not +# something the executed code produces by letting an alarm of its own fire. +_TIMEOUT_EXIT_CODE = 124 + +# Runs the code under a supervisor that enforces a hard wall-clock bound inside +# the container. The code runs in a forked child in its own process group; the +# supervisor waits for it and, when the bound expires, SIGKILLs the whole group +# and exits with `_TIMEOUT_EXIT_CODE`. The group is also swept once the code +# finishes normally, so nothing it started is left running in the shared +# container (a leftover process would also hold the exec's output open). +# +# Two properties matter for code that may be hostile: the deadline lives in a +# process the code never runs in, so it cannot be disarmed from inside (an +# alarm armed in the executing process could be cancelled with one call), and +# SIGKILL to the group reaches what the code spawned, not just its top frame. +# The code is passed as an argument rather than inlined so that tracebacks keep +# the original line numbers, and argv is restored to what `python3 -c` would +# have given so that code parsing arguments still works. +# +# Not covered: code that deliberately leaves the group (`os.setsid()`) or +# double-forks away survives the bound until the container is torn down. +_TIMEOUT_WRAPPER = """\ +import os, signal, sys + +_timeout = int(sys.argv[1]) +_source = sys.argv[2] +del sys.argv[1:] + +_pid = os.fork() +if _pid == 0: + try: + os.setpgid(0, 0) + except OSError: + pass + exec(compile(_source, '', 'exec'), {'__name__': '__main__'}) +else: + + def _sweep_group(): + try: + os.killpg(_pid, signal.SIGKILL) + except OSError: + pass + + def _expire(_signum, _frame): + _sweep_group() + try: + os.kill(_pid, signal.SIGKILL) + except OSError: + pass + os._exit(%d) + + try: + os.setpgid(_pid, _pid) + except OSError: + pass + signal.signal(signal.SIGALRM, _expire) + signal.alarm(_timeout) + _status = os.waitpid(_pid, 0)[1] + signal.alarm(0) + _sweep_group() + os._exit( + 128 + os.WTERMSIG(_status) + if os.WIFSIGNALED(_status) + else os.WEXITSTATUS(_status) + ) +""" % _TIMEOUT_EXIT_CODE + class ContainerCodeExecutor(BaseCodeExecutor): """A code executor that uses a custom container to execute code. @@ -87,6 +155,23 @@ class ContainerCodeExecutor(BaseCodeExecutor): code must make network requests and you trust it. """ + # Overrides the BaseCodeExecutor attribute: unlike the base default of None, + # the timeout here is always finite and must be positive (0 would mean no + # bound at all). + timeout_seconds: int = Field(default=300, gt=0) + """The wall-clock timeout in seconds for a single code execution. + + Every execution shares one long-lived container, so an unbounded run (e.g. a + loop emitted by the model) would keep burning that container's CPU for every + later caller. Defaults to 300, matching ``GkeCodeExecutor``. A computation + that legitimately runs longer than the timeout is killed, so raise it rather + than removing it; ``None`` is rejected, unlike on the base class. + + When the timeout expires the executed code is killed along with the process + group it runs in, so what it spawned goes with it. Code that deliberately + detaches from that group keeps running until the container is torn down. + """ + # Overrides the BaseCodeExecutor attribute: this executor cannot be stateful. stateful: bool = Field(default=False, frozen=True, exclude=True) @@ -151,7 +236,13 @@ class ContainerCodeExecutor(BaseCodeExecutor): output = '' error = '' exec_result = self._container.exec_run( - ['python3', '-c', code_execution_input.code], + [ + 'python3', + '-c', + _TIMEOUT_WRAPPER, + str(self.timeout_seconds), + code_execution_input.code, + ], demux=True, ) logger.debug('Executed code:\n```\n%s\n```', code_execution_input.code) @@ -165,6 +256,15 @@ class ContainerCodeExecutor(BaseCodeExecutor): ): error = exec_result.output[1].decode('utf-8') + if exec_result.exit_code == _TIMEOUT_EXIT_CODE: + # Appended rather than assigned: whatever the code managed to write to + # stderr before the alarm fired is still the useful diagnostic, but on + # its own it would hide the fact that the run was cut short. + timed_out = ( + f'Code execution timed out after {self.timeout_seconds} seconds.' + ) + error = f'{error}\n{timed_out}' if error else timed_out + # Collect the final result. return CodeExecutionResult( stdout=output, 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 64752fff..851b63a5 100644 --- a/src/google/adk/code_executors/unsafe_local_code_executor.py +++ b/src/google/adk/code_executors/unsafe_local_code_executor.py @@ -18,8 +18,10 @@ from contextlib import redirect_stdout import io import logging import multiprocessing +import os import queue import re +import signal import traceback from typing import Any @@ -33,11 +35,23 @@ from .code_execution_utils import CodeExecutionResult logger = logging.getLogger('google_adk.' + __name__) +# How long to wait for a timed-out execution to exit after SIGTERM before +# escalating to SIGKILL, so that the timeout itself cannot block forever. +_TERMINATE_GRACE_SECONDS = 5 + def _execute_in_process( code: str, globals_: dict[str, Any], result_queue: multiprocessing.Queue ) -> None: """Executes code in a separate process and puts result in queue.""" + # Detach into a new session/process group before running anything, so that a + # timed-out execution can be killed together with everything it spawned. + if hasattr(os, 'setsid'): + try: + os.setsid() + except OSError: + logger.debug('Could not detach the execution process group.') + stdout = io.StringIO() error = None try: @@ -48,6 +62,53 @@ def _execute_in_process( result_queue.put((stdout.getvalue(), error)) +def _execution_group( + process: multiprocessing.process.BaseProcess, +) -> int | None: + """Returns the group the execution detached into, or None if it has not.""" + if process.pid is None or not hasattr(os, 'killpg'): + return None + try: + group = os.getpgid(process.pid) + # Only report the group once the execution has detached into its own; + # otherwise the group is still ours and signalling it would take down the + # agent along with the code it is running. + return group if group != os.getpgid(0) else None + except OSError: + return None + + +def _signal_group(group: int, sig: int) -> None: + """Signals every process left in a group, tolerating an empty one.""" + try: + os.killpg(group, sig) + except OSError: + logger.debug('Could not signal the execution process group.') + + +def _kill_execution(process: multiprocessing.process.BaseProcess) -> None: + """Kills a timed-out execution along with any process it spawned.""" + # Resolved up front: once the execution process has been reaped its group can + # no longer be looked up through it, and the group is what holds whatever the + # code spawned. + group = _execution_group(process) + + # SIGTERM first, so the code and its children get the same grace period the + # execution process itself gets before anything is killed outright. + if group is not None: + _signal_group(group, signal.SIGTERM) + process.terminate() + process.join(_TERMINATE_GRACE_SECONDS) + + # Escalate unconditionally: the execution process exiting says nothing about + # a child of it that is ignoring SIGTERM. + if group is not None: + _signal_group(group, signal.SIGKILL) + if process.is_alive(): + process.kill() + process.join() + + 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): @@ -102,8 +163,7 @@ class UnsafeLocalCodeExecutor(BaseCodeExecutor): if err: error = err except queue.Empty: - process.terminate() - process.join() + _kill_execution(process) error = f'Code execution timed out after {self.timeout_seconds} seconds.' # Collect the final result. diff --git a/tests/unittests/code_executors/test_container_code_executor.py b/tests/unittests/code_executors/test_container_code_executor.py index 5574135b..09bc1243 100644 --- a/tests/unittests/code_executors/test_container_code_executor.py +++ b/tests/unittests/code_executors/test_container_code_executor.py @@ -14,9 +14,19 @@ """Tests for the ContainerCodeExecutor container hardening defaults.""" +import os +import signal +import subprocess +import sys +import textwrap +import time from unittest import mock +from google.adk.code_executors import container_code_executor +from google.adk.code_executors.code_execution_utils import CodeExecutionInput from google.adk.code_executors.container_code_executor import ContainerCodeExecutor +import pydantic +import pytest def _mock_docker_client(): @@ -56,3 +66,239 @@ def test_container_network_can_be_explicitly_enabled(mock_docker): _, kwargs = client.containers.run.call_args assert not kwargs['network_disabled'] + + +def _executed_command(container) -> list[str]: + """Returns the command of the last `exec_run` call on the container.""" + args, _ = container.exec_run.call_args + return args[0] + + +@mock.patch('google.adk.code_executors.container_code_executor.docker') +def test_execute_code_bounds_execution_by_default(mock_docker): + """Code runs under a finite timeout even when the caller sets none.""" + client = _mock_docker_client() + mock_docker.from_env.return_value = client + executor = ContainerCodeExecutor(image='test-image') + container = client.containers.run.return_value + container.exec_run.return_value = mock.MagicMock( + exit_code=0, output=(b'', b'') + ) + + executor.execute_code(mock.MagicMock(), CodeExecutionInput(code='x = 1')) + + # The container is shared by every invocation, so an unbounded run would pin + # it for all later callers. + assert executor.timeout_seconds == 300 + assert _executed_command(container) == [ + 'python3', + '-c', + container_code_executor._TIMEOUT_WRAPPER, + '300', + 'x = 1', + ] + + +@mock.patch('google.adk.code_executors.container_code_executor.docker') +def test_execute_code_passes_configured_timeout(mock_docker): + """The inherited `timeout_seconds` bounds the in-container execution.""" + client = _mock_docker_client() + mock_docker.from_env.return_value = client + executor = ContainerCodeExecutor(image='test-image', timeout_seconds=7) + container = client.containers.run.return_value + container.exec_run.return_value = mock.MagicMock( + exit_code=0, output=(b'', b'') + ) + + executor.execute_code( + mock.MagicMock(), CodeExecutionInput(code='while True: pass') + ) + + assert _executed_command(container) == [ + 'python3', + '-c', + container_code_executor._TIMEOUT_WRAPPER, + '7', + 'while True: pass', + ] + + +@mock.patch('google.adk.code_executors.container_code_executor.docker') +@pytest.mark.parametrize('timeout', [0, -1, None]) +def test_non_positive_timeout_is_rejected(mock_docker, timeout): + """A timeout of 0 or None would mean no bound, so it is refused up front.""" + mock_docker.from_env.return_value = _mock_docker_client() + + with pytest.raises(pydantic.ValidationError): + ContainerCodeExecutor(image='test-image', timeout_seconds=timeout) + + +@mock.patch('google.adk.code_executors.container_code_executor.docker') +def test_execute_code_reports_timeout(mock_docker): + """A run the supervisor cut short is reported as a timeout.""" + client = _mock_docker_client() + mock_docker.from_env.return_value = client + executor = ContainerCodeExecutor(image='test-image', timeout_seconds=7) + container = client.containers.run.return_value + container.exec_run.return_value = mock.MagicMock( + exit_code=container_code_executor._TIMEOUT_EXIT_CODE, output=(b'', b'') + ) + + result = executor.execute_code( + mock.MagicMock(), CodeExecutionInput(code='while True: pass') + ) + + assert 'timed out after 7 seconds' in result.stderr + + +@mock.patch('google.adk.code_executors.container_code_executor.docker') +def test_execute_code_reports_timeout_alongside_stderr(mock_docker): + """Output written before the alarm fired does not hide the timeout.""" + client = _mock_docker_client() + mock_docker.from_env.return_value = client + executor = ContainerCodeExecutor(image='test-image', timeout_seconds=7) + container = client.containers.run.return_value + container.exec_run.return_value = mock.MagicMock( + exit_code=container_code_executor._TIMEOUT_EXIT_CODE, + output=(b'', b'a warning from the code'), + ) + + result = executor.execute_code( + mock.MagicMock(), CodeExecutionInput(code='while True: pass') + ) + + assert 'a warning from the code' in result.stderr + assert 'timed out after 7 seconds' in result.stderr + + +# The wrapper below is a string of Python that only ever runs inside the +# container, so the tests run it directly on this host instead: no docker, no +# daemon, and only snippets written here. +_POSIX_ONLY = pytest.mark.skipif( + not hasattr(os, 'fork') or not hasattr(os, 'killpg'), + reason='The in-container bound is enforced with POSIX process groups.', +) + + +def _run_wrapper(timeout: int, code: str) -> subprocess.CompletedProcess: + """Runs the wrapper exactly as the container executor asks the container to.""" + return subprocess.run( + [ + sys.executable, + '-c', + container_code_executor._TIMEOUT_WRAPPER, + str(timeout), + code, + ], + capture_output=True, + text=True, + timeout=30, + check=False, + ) + + +def _is_alive(pid: int) -> bool: + """Returns whether `pid` is a live (non-zombie) process.""" + try: + with open(f'/proc/{pid}/stat', encoding='utf-8') as stat_file: + state = stat_file.read().rsplit(')', 1)[1].split()[0] + except OSError: + return False + return state != 'Z' + + +@_POSIX_ONLY +def test_wrapper_kills_a_run_that_hits_the_bound(): + """A loop that never returns is killed, not merely waited on.""" + started = time.monotonic() + + completed = _run_wrapper(1, 'while True: pass') + + assert completed.returncode == container_code_executor._TIMEOUT_EXIT_CODE + assert time.monotonic() - started < 15 + + +@_POSIX_ONLY +def test_wrapper_bound_cannot_be_disarmed_by_the_executed_code(): + """The deadline is held by a process the executed code never runs in.""" + completed = _run_wrapper( + 1, + 'import signal, time\n' + 'signal.alarm(0)\n' + 'signal.signal(signal.SIGALRM, signal.SIG_IGN)\n' + 'time.sleep(25)\n' + 'print("outlived the bound")\n', + ) + + assert completed.returncode == container_code_executor._TIMEOUT_EXIT_CODE + assert 'outlived the bound' not in completed.stdout + + +@_POSIX_ONLY +@pytest.mark.skipif( + not os.path.isdir('/proc'), reason='Liveness is checked through /proc.' +) +def test_wrapper_kills_what_the_code_spawned(tmp_path): + """Whatever the run started dies with it, so the container is not pinned.""" + pid_file = tmp_path / 'spawned.pid' + code = textwrap.dedent(f""" + import os + import time + + spawned = os.fork() + if spawned == 0: + time.sleep(60) + os._exit(0) + with open({str(pid_file)!r}, 'w') as f: + f.write(str(spawned)) + time.sleep(60) + """) + + completed = _run_wrapper(1, code) + + assert completed.returncode == container_code_executor._TIMEOUT_EXIT_CODE + assert pid_file.exists(), 'the code never got as far as spawning a process' + spawned_pid = int(pid_file.read_text()) + try: + deadline = time.monotonic() + 10 + while time.monotonic() < deadline and _is_alive(spawned_pid): + time.sleep(0.05) + assert not _is_alive(spawned_pid) + finally: + try: + os.kill(spawned_pid, signal.SIGKILL) + except OSError: + pass + + +@_POSIX_ONLY +def test_wrapper_leaves_argv_as_a_plain_python_c_run_would(): + """The timeout and the source do not leak into the code's own arguments.""" + completed = _run_wrapper( + 5, + 'import argparse, sys\n' + 'argparse.ArgumentParser().parse_args()\n' + 'print(sys.argv)\n', + ) + + assert completed.returncode == 0, completed.stderr + assert completed.stdout.strip() == "['-c']" + + +@_POSIX_ONLY +def test_wrapper_passes_through_output_and_exit_code(): + """A run that finishes on its own is reported exactly as it ended.""" + completed = _run_wrapper(5, 'import sys\nprint("hello")\nsys.exit(3)\n') + + assert completed.returncode == 3 + assert completed.stdout.strip() == 'hello' + + +@_POSIX_ONLY +def test_wrapper_reports_an_uncaught_error_against_the_original_line(): + """Wrapping the code does not shift the line numbers in its traceback.""" + completed = _run_wrapper(5, 'x = 1\nraise ValueError("boom")\n') + + assert completed.returncode == 1 + assert '"", line 2' in completed.stderr + assert 'boom' in completed.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 fa22e1bb..49b54d89 100644 --- a/tests/unittests/code_executors/test_unsafe_local_code_executor.py +++ b/tests/unittests/code_executors/test_unsafe_local_code_executor.py @@ -12,11 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +import multiprocessing +import os +import signal import textwrap +import time 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 import unsafe_local_code_executor 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.unsafe_local_code_executor import UnsafeLocalCodeExecutor @@ -25,6 +30,25 @@ from google.adk.sessions.session import Session import pytest +def _written_pid(pid_file) -> int | None: + """Returns the pid the executed code recorded, or None if it has not yet.""" + try: + recorded = pid_file.read_text().strip() + except OSError: + return None + return int(recorded) if recorded else None + + +def _is_alive(pid: int) -> bool: + """Returns whether `pid` is a live (non-zombie) process.""" + try: + with open(f"/proc/{pid}/stat", encoding="utf-8") as stat_file: + state = stat_file.read().rsplit(")", 1)[1].split()[0] + except OSError: + return False + return state != "Z" + + @pytest.fixture def mock_invocation_context() -> InvocationContext: """Provides a mock InvocationContext.""" @@ -131,3 +155,93 @@ class TestUnsafeLocalCodeExecutor: assert result.stdout == "" assert "Code execution timed out after 1 seconds." in result.stderr + + def test_kill_execution_signals_group_before_killing_it(self, monkeypatch): + """The group gets SIGTERM and its grace period before SIGKILL.""" + signalled = [] + monkeypatch.setattr( + unsafe_local_code_executor.os, + "killpg", + lambda group, sig: signalled.append((group, sig)), + ) + monkeypatch.setattr( + unsafe_local_code_executor, + "_execution_group", + lambda process: 4321, + ) + process = MagicMock() + process.is_alive.return_value = False + process.terminate.side_effect = lambda: signalled.append(("child", "term")) + + unsafe_local_code_executor._kill_execution(process) + + assert signalled == [ + (4321, signal.SIGTERM), + ("child", "term"), + (4321, signal.SIGKILL), + ] + process.join.assert_any_call( + unsafe_local_code_executor._TERMINATE_GRACE_SECONDS + ) + + @pytest.mark.skipif( + not hasattr(os, "killpg") + or not hasattr(os, "fork") + or not os.path.isdir("/proc"), + reason="Process-group teardown is checked on POSIX with /proc only.", + ) + def test_kill_execution_kills_what_the_code_spawned(self, tmp_path): + """Killing a live execution takes the processes it spawned with it.""" + pid_file = tmp_path / "spawned.pid" + # Forked rather than spawned through `sys.executable`, so the descendant + # exists within milliseconds and the test never waits on interpreter + # start-up. + code = textwrap.dedent(f""" + import os + import time + + spawned = os.fork() + if spawned == 0: + time.sleep(60) + os._exit(0) + with open({str(pid_file)!r}, 'w') as f: + f.write(str(spawned)) + time.sleep(60) + """) + ctx = multiprocessing.get_context("spawn") + result_queue = ctx.Queue() + process = ctx.Process( + target=unsafe_local_code_executor._execute_in_process, + args=(code, {}, result_queue), + daemon=True, + ) + process.start() + spawned_pid = None + try: + # Waiting for the pid to be written rather than for a fixed duration: + # the only thing that has to have happened is the fork. The file exists + # from the moment it is opened, so its content is what is polled for. + deadline = time.time() + 30 + while time.time() < deadline and not _written_pid(pid_file): + time.sleep(0.05) + spawned_pid = _written_pid(pid_file) + if spawned_pid is None: + pytest.skip("this environment could not start the execution process") + + unsafe_local_code_executor._kill_execution(process) + + assert not process.is_alive() + deadline = time.time() + 10 + while time.time() < deadline and _is_alive(spawned_pid): + time.sleep(0.05) + assert not _is_alive(spawned_pid) + finally: + if spawned_pid is not None: + try: + os.kill(spawned_pid, signal.SIGKILL) + except OSError: + pass + if process.is_alive(): + process.kill() + process.join() + result_queue.close()