from __future__ import annotations import os import re import signal import subprocess import sys import tempfile import time from pathlib import Path from unittest import TestCase from src.agent_tools import build_tool_context, default_tool_registry, execute_tool from src.agent_types import AgentPermissions, AgentRuntimeConfig class PythonExecToolTests(TestCase): def test_python_exec_runs_inline_code_with_shell_permission(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: config = AgentRuntimeConfig( cwd=Path(tmp_dir), permissions=AgentPermissions(allow_shell_commands=True), ) context = build_tool_context(config) result = execute_tool( default_tool_registry(), 'python_exec', {'code': 'print("hello from python")'}, context, ) self.assertTrue(result.ok) self.assertIn('exit_code=0', result.content) self.assertIn('hello from python', result.content) self.assertEqual(result.metadata.get('action'), 'python_exec') self.assertEqual(result.metadata.get('mode'), 'code') def test_python_exec_exposes_session_scratchpad(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: root = Path(tmp_dir) scratchpad = (root / 'session' / 'scratchpad').resolve() scratchpad.mkdir(parents=True) config = AgentRuntimeConfig( cwd=root, permissions=AgentPermissions(allow_shell_commands=True), ) context = build_tool_context(config, scratchpad_directory=scratchpad) result = execute_tool( default_tool_registry(), 'python_exec', {'code': 'import os\nprint(os.environ["PYTHON_EXEC_SCRATCHPAD"])'}, context, ) self.assertTrue(result.ok) self.assertIn(str(scratchpad), result.content) self.assertEqual(result.metadata.get('scratchpad_directory'), str(scratchpad)) def test_python_exec_runs_from_session_scratchpad(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: root = Path(tmp_dir) scratchpad = (root / 'session' / 'scratchpad').resolve() scratchpad.mkdir(parents=True) config = AgentRuntimeConfig( cwd=root, permissions=AgentPermissions(allow_shell_commands=True), ) context = build_tool_context(config, scratchpad_directory=scratchpad) result = execute_tool( default_tool_registry(), 'python_exec', {'code': 'import os\nprint(os.getcwd())'}, context, ) self.assertTrue(result.ok) self.assertIn(str(scratchpad), result.content) def test_python_exec_prefers_user_python_env(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: root = Path(tmp_dir) user_env = (root / 'account' / 'python' / '.venv').resolve() bin_dir = user_env / 'bin' bin_dir.mkdir(parents=True) python_bin = bin_dir / 'python' try: python_bin.symlink_to(sys.executable) except OSError: python_bin.write_text( '#!/bin/sh\nexec "$PYTHON_EXEC_TEST_REAL_PYTHON" "$@"\n', encoding='utf-8', ) python_bin.chmod(0o755) config = AgentRuntimeConfig( cwd=root, permissions=AgentPermissions(allow_shell_commands=True), python_env_dir=user_env, ) context = build_tool_context( config, extra_env={'PYTHON_EXEC_TEST_REAL_PYTHON': sys.executable}, ) result = execute_tool( default_tool_registry(), 'python_exec', {'code': 'import os, sys\nprint(sys.executable)\nprint(os.environ["PYTHON_EXEC_VENV"])'}, context, ) self.assertTrue(result.ok) self.assertIn(str(user_env), result.content) self.assertEqual(result.metadata.get('python_env_dir'), str(user_env)) def test_python_package_show_uses_user_python_env(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: root = Path(tmp_dir) user_env = (root / 'account' / 'python' / '.venv').resolve() bin_dir = user_env / 'bin' bin_dir.mkdir(parents=True) python_bin = bin_dir / 'python' try: python_bin.symlink_to(sys.executable) except OSError: python_bin.write_text( '#!/bin/sh\nexec "$PYTHON_EXEC_TEST_REAL_PYTHON" "$@"\n', encoding='utf-8', ) python_bin.chmod(0o755) config = AgentRuntimeConfig( cwd=root, permissions=AgentPermissions(allow_shell_commands=True), python_env_dir=user_env, ) context = build_tool_context( config, extra_env={'PYTHON_EXEC_TEST_REAL_PYTHON': sys.executable}, ) result = execute_tool( default_tool_registry(), 'python_package', {'action': 'show'}, context, ) self.assertTrue(result.ok) self.assertIn(str(user_env), result.content) self.assertEqual(result.metadata.get('python_env_dir'), str(user_env)) def test_python_exec_is_blocked_without_shell_permission(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: config = AgentRuntimeConfig(cwd=Path(tmp_dir)) context = build_tool_context(config) result = execute_tool( default_tool_registry(), 'python_exec', {'code': 'print("blocked")'}, context, ) self.assertFalse(result.ok) self.assertIn('--allow-shell', result.content) self.assertEqual(result.metadata.get('error_kind'), 'permission_denied') def test_python_exec_runs_workspace_script_with_args(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: root = Path(tmp_dir) script = root / 'tool_script.py' script.write_text( 'import sys\nprint("|".join(sys.argv[1:]))\n', encoding='utf-8', ) config = AgentRuntimeConfig( cwd=root, permissions=AgentPermissions(allow_shell_commands=True), ) context = build_tool_context(config) result = execute_tool( default_tool_registry(), 'python_exec', {'script_path': 'tool_script.py', 'args': ['a', 'b']}, context, ) self.assertTrue(result.ok) self.assertIn('a|b', result.content) self.assertEqual(result.metadata.get('mode'), 'script') def test_python_exec_timeout_kills_spawned_children(self) -> None: if os.name != 'posix': self.skipTest('process group cleanup test is POSIX-only') with tempfile.TemporaryDirectory() as tmp_dir: config = AgentRuntimeConfig( cwd=Path(tmp_dir), permissions=AgentPermissions(allow_shell_commands=True), ) context = build_tool_context(config) child_code = ( 'import subprocess, sys, time\n' 'child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(60)"])\n' 'print(child.pid, flush=True)\n' 'time.sleep(60)\n' ) result = execute_tool( default_tool_registry(), 'python_exec', { 'code': child_code, 'timeout_seconds': 1, 'max_output_chars': 2000, }, context, ) self.assertTrue(result.ok) self.assertTrue(result.metadata.get('timed_out')) match = re.search(r'\[stdout\]\s*(\d+)', result.content) self.assertIsNotNone(match, result.content) child_pid = int(match.group(1)) time.sleep(0.3) if self._pid_is_running(child_pid): try: os.kill(child_pid, signal.SIGKILL) except ProcessLookupError: pass self.fail(f'spawned child process is still running: {child_pid}') @staticmethod def _pid_is_running(pid: int) -> bool: status = subprocess.run( ['ps', '-p', str(pid), '-o', 'stat='], capture_output=True, text=True, check=False, ) state = status.stdout.strip() return bool(state) and not state.startswith('Z')