diff --git a/src/tesslate_agent/agent/tools/shell_ops/execute.py b/src/tesslate_agent/agent/tools/shell_ops/execute.py index 97881b1..a4c89c5 100644 --- a/src/tesslate_agent/agent/tools/shell_ops/execute.py +++ b/src/tesslate_agent/agent/tools/shell_ops/execute.py @@ -13,6 +13,7 @@ from __future__ import annotations import logging +import re from typing import Any from tesslate_agent.agent.tools.output_formatter import ( @@ -28,6 +29,20 @@ _BYTES_PER_TOKEN = 4 +# Security Guardrail to prevent destructive shell commands +DANGEROUS_PATTERNS = [ + r"rm\s+-rf\s+/", # Root deletion + r"chmod\s+.*777", # Dangerous permission changes + r":\(\){ :\|:& };:", # Fork bombs + r"mv\s+.*\s+/dev/null", # Deleting data by moving to null + r"> /dev/sda", # Overwriting disk directly +] + +def is_command_safe(command: str) -> tuple[bool, str | None]: + for pattern in DANGEROUS_PATTERNS: + if re.search(pattern, command): + return False, f"Potentially dangerous command pattern detected: {pattern}" + return True, None @tool_retry async def shell_exec_executor( @@ -66,6 +81,15 @@ async def shell_exec_executor( if not command.endswith("\n"): command += "\n" + is_safe, error_msg = is_command_safe(command.strip()) + if not is_safe: + logger.warning("[SECURITY-BLOCK] session=%s command=%s", session_id, command.strip()) + return error_output( + message=error_msg, + suggestion="Refine the command to target specific directories or avoid destructive flags.", + details={"session_id": session_id, "tier": "local", "security_risk": "high"} + ) + try: PTY_SESSIONS.write(session_id, command) except KeyError: diff --git a/src/tesslate_agent/cli/runner.py b/src/tesslate_agent/cli/runner.py index 2e79d15..046a9e6 100644 --- a/src/tesslate_agent/cli/runner.py +++ b/src/tesslate_agent/cli/runner.py @@ -20,6 +20,7 @@ import os import sys import tempfile +import signal from collections.abc import Callable from pathlib import Path from typing import Any @@ -181,27 +182,67 @@ async def run_agent( timeout_seconds = max(1.0, timeout_ms / 1000.0) exit_code = EXIT_SUCCESS + # Fix: Explicit signal handling for graceful shutdown in containers + loop = asyncio.get_running_loop() + stop_event = asyncio.Event() + + def handle_exit_signal(): + logger.warning("Received exit signal, stopping agent gracefully...") + stop_event.set() + # Trigger an exception in the running _drive_agent task + for task in asyncio.all_tasks(loop): + if task.get_coro().__name__ == "_drive_agent": + task.cancel() + if sys.platform != "win32": + for sig in (signal.SIGINT, signal.SIGTERM): + loop.add_signal_handler(sig, handle_exit_signal) + try: - await asyncio.wait_for( - _drive_agent(agent, task, context, bridge, event_printer), + # Combined wait_for and cancellation event + agent_task = asyncio.create_task( + _drive_agent(agent, task, context, bridge, event_printer) + ) + + # Wait for either completion, timeout, or signal + done, pending = await asyncio.wait( + [agent_task], timeout=timeout_seconds, + return_when=asyncio.FIRST_COMPLETED ) - if bridge.has_error: - exit_code = EXIT_AGENT_ERROR - except asyncio.TimeoutError: - logger.error("agent run timed out after %.1fs", timeout_seconds) - bridge.mark_errored(f"timeout after {timeout_seconds:.1f}s") - exit_code = EXIT_AGENT_ERROR - except KeyboardInterrupt: + + if agent_task in done: + await agent_task + else: + # Handle Timeout or Signal + agent_task.cancel() + if stop_event.is_set(): + logger.error("agent run interrupted by signal") + bridge.mark_errored("interrupted by system signal (SIGINT/SIGTERM)") + exit_code = EXIT_AGENT_ERROR + else: + logger.error("agent run timed out after %.1fs", timeout_seconds) + bridge.mark_errored(f"timeout after {timeout_seconds:.1f}s") + exit_code = EXIT_AGENT_ERROR + + except (asyncio.CancelledError, asyncio.TimeoutError): + # Handled by signal logic above + pass + + except KeyboardInterrupt: # This handles Ctrl+C on Windows logger.error("agent run interrupted by user") bridge.mark_errored("interrupted by user") exit_code = EXIT_AGENT_ERROR + except Exception as exc: logger.exception("agent run failed with unexpected exception") bridge.mark_errored(f"unexpected exception: {exc}") exit_code = EXIT_AGENT_ERROR finally: - _write_trajectory(resolved_output, bridge.finalize()) + # Remove signal handlers to clean up the loop + if sys.platform != "win32": + for sig in (signal.SIGINT, signal.SIGTERM): + loop.remove_signal_handler(sig) + _write_trajectory(resolved_output, bridge.finalize()) return exit_code diff --git a/tests/agent/tools/shell_ops/test_execute_security.py b/tests/agent/tools/shell_ops/test_execute_security.py new file mode 100644 index 0000000..dc07cd9 --- /dev/null +++ b/tests/agent/tools/shell_ops/test_execute_security.py @@ -0,0 +1,15 @@ +import pytest +from tesslate_agent.agent.tools.shell_ops.execute import is_command_safe + +@pytest.mark.parametrize("command, expected_safe", [ + ("ls -la", True), + ("echo 'hello'", True), + ("rm -rf /", False), + ("chmod 777 sensitive_file.key", False), + (":(){ :|:& };:", False), +]) +def test_is_command_safe(command, expected_safe): + safe, msg = is_command_safe(command) + assert safe == expected_safe + if not safe: + assert "dangerous" in msg.lower() \ No newline at end of file diff --git a/tests/cli/test_runner_signals.py b/tests/cli/test_runner_signals.py new file mode 100644 index 0000000..4c491a8 --- /dev/null +++ b/tests/cli/test_runner_signals.py @@ -0,0 +1,32 @@ +import sys +from unittest.mock import MagicMock, patch + +# --- WINDOWS COMPATIBILITY FIX --- +# The agent imports `fcntl` for memory file locking, which does not exist on Windows. +# We mock it globally in the test environment so the test suite can run cross-platform. +if sys.platform == "win32": + sys.modules['fcntl'] = MagicMock() + +import pytest +import asyncio +from pathlib import Path +from tesslate_agent.cli.runner import run_agent + +@pytest.mark.anyio +async def test_graceful_shutdown_on_signal(): + # 1. Patch the final write function to prove the finally block executes + with patch('tesslate_agent.cli.runner._write_trajectory') as mock_write: + + # 2. Simulate the user pressing Ctrl+C (KeyboardInterrupt) during execution + with patch('tesslate_agent.cli.runner._drive_agent', side_effect=KeyboardInterrupt): + + # Start the task + runner_task = asyncio.create_task( + run_agent("test task", "openai/gpt-4o", Path("."), Path("output.json")) + ) + + await runner_task + + # 3. ASSERT: Verify the trajectory was written despite the interruption + assert mock_write.called + print("\n✅ PASS: Trajectory finalized successfully after simulated Ctrl+C.") \ No newline at end of file