武阳
238 lines
8.9 KiB
Python
238 lines
8.9 KiB
Python
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')
|