"""Session-level abort & kill switches. Three pieces: 1. ``AbortController``: a thread-safe flag set by a signal handler (SIGINT / SIGTERM) or by any code path that detects a fatal condition (cost cap, iteration cap, user request). 2. ``FileRollback``: snapshots original file bytes before each write/edit so that an abort can offer to revert reversible state. 3. ``kill_in_flight``: a helper that sends SIGINT to any python process running inside the sandbox container, stopping a hanging ``docker exec`` without tearing down the container itself. The controller is checked: - at the top of the inner agent loop (between LLM turns), - between tool calls inside ``handle_tool_calls``, - before each LLM request. """ from __future__ import annotations import hashlib import shutil import signal import subprocess import threading from pathlib import Path from typing import Any # --------------------------------------------------------------------------- # 6.3 AbortController # --------------------------------------------------------------------------- class AbortController: """Thread-safe abort flag for the agent session. A single instance is created per session. Signal handlers (SIGINT, SIGTERM) call ``trigger()``; the agent loop polls ``triggered`` between turns and between tool calls. Once triggered: - the inner loop stops dispatching new tools, - any in-flight tool is killed (``kill_in_flight``), - a ``session_abort`` audit event is emitted, - reversible file changes are offered for rollback. """ def __init__(self) -> None: self._triggered = False self._reason: str | None = None self._lock = threading.Lock() self._registered_signals: list[int] = [] @property def triggered(self) -> bool: with self._lock: return self._triggered @property def reason(self) -> str | None: with self._lock: return self._reason def trigger(self, reason: str = "abort requested") -> None: """Set the abort flag. Safe to call from a signal handler.""" with self._lock: if not self._triggered: self._triggered = True self._reason = reason def reset(self) -> None: """Clear the flag (used by tests).""" with self._lock: self._triggered = False self._reason = None def install_signal_handlers(self) -> None: """Register SIGINT / SIGTERM handlers that call ``trigger``. Only call this from the main thread (signal.signal requirement). The previous handlers are not saved, we deliberately replace them because the abort controller is the final arbiter. """ def _handler(signum, frame): name = signal.Signals(signum).name self.trigger(f"received {name}") for sig in (signal.SIGINT, signal.SIGTERM): try: signal.signal(sig, _handler) self._registered_signals.append(sig) except (ValueError, OSError): # Not in main thread, or signal not supported on this # platform, skip silently. pass def remove_signal_handlers(self) -> None: """Restore default signal handling.""" for sig in self._registered_signals: try: signal.signal(sig, signal.SIG_DFL) except (ValueError, OSError): pass self._registered_signals.clear() # --------------------------------------------------------------------------- # 6.3 kill_in_flight # --------------------------------------------------------------------------- def kill_in_flight(container: str) -> None: """Send SIGINT to any python process inside the sandbox container. This stops a hanging ``docker exec`` (e.g. a long ``run_bash``) without tearing down the container itself, so the ``finally`` block can still clean up. """ try: subprocess.run( ["docker", "exec", container, "pkill", "-INT", "python"], capture_output=True, timeout=10, ) except Exception: # Best-effort: if the container is already gone or pkill isn't # available, the exec subprocess's own timeout will # eventually clean up. pass # --------------------------------------------------------------------------- # 6.3 FileRollback # --------------------------------------------------------------------------- class FileRollback: """Snapshot original file bytes before each write/edit. On abort, ``offer_rollback`` walks the snapshot list and restores each file to its pre-edit state, prompting the user for confirmation. Only ``write_file`` and ``edit_file`` are reversible; ``run_bash`` side effects (e.g. ``git commit``) are not, the audit log is the only record for those. """ def __init__(self) -> None: # path -> backup path (in a temp dir) self._snapshots: list[tuple[str, Path]] = [] self._backup_dir: Path | None = None def _ensure_backup_dir(self) -> Path: if self._backup_dir is None: self._backup_dir = Path( __file__).resolve().parent / ".rollback_backups" self._backup_dir.mkdir(parents=True, exist_ok=True) return self._backup_dir def snapshot(self, path: str) -> None: """Save a copy of *path* if it exists, for later rollback.""" p = Path(path) if not p.exists() or not p.is_file(): return try: backup = self._ensure_backup_dir() / ( hashlib.sha256(str(p.resolve()).encode()).hexdigest()[:16] + "_" + p.name ) shutil.copy2(p, backup) self._snapshots.append((str(p.resolve()), backup)) except OSError: # If we can't snapshot, we just can't roll back, don't # block the tool call. pass @property def snapshot_count(self) -> int: return len(self._snapshots) def offer_rollback(self) -> int: """Prompt the user to revert all snapshotted files. Returns the number of files actually restored. """ if not self._snapshots: print(" [rollback] No reversible file changes to roll back.") return 0 print(f"\n [rollback] {len(self._snapshots)} file(s) were modified " "during this session.") try: answer = input(" Revert all changes? [y/n]: ").strip().lower() except EOFError: answer = "n" if answer not in ("y", "yes"): print(" [rollback] Keeping changes.") return 0 restored = 0 for original_path, backup_path in self._snapshots: try: shutil.copy2(backup_path, original_path) restored += 1 except OSError as e: print(f" [rollback] Could not restore {original_path}: {e}") print(f" [rollback] Restored {restored} file(s).") return restored def cleanup(self) -> None: """Remove the backup directory.""" if self._backup_dir and self._backup_dir.exists(): shutil.rmtree(self._backup_dir, ignore_errors=True) __all__ = [ "AbortController", "kill_in_flight", "FileRollback", ]