diff --git a/agent/interrupt_scope.py b/agent/interrupt_scope.py new file mode 100644 index 0000000000..e08a66a823 --- /dev/null +++ b/agent/interrupt_scope.py @@ -0,0 +1,65 @@ +"""Host-owned cancellation for agents created deep inside synchronous work. + +A host that runs a blocking command on a worker thread (Hermes Console) never sees +the ``AIAgent`` a CLI subcommand forks inside it, so it cannot call ``interrupt()`` +when the user cancels. The host binds an :class:`InterruptScope` around the work; +every ``run_conversation()`` under that scope registers its agent, and +``scope.cancel()`` hard-interrupts them from any thread. Agents registering after +the cancel are interrupted immediately, so a cancel never loses the race with a +turn that has not started yet (#106179). +""" + +from __future__ import annotations + +import threading +from contextlib import contextmanager, nullcontext +from contextvars import ContextVar +from typing import Any, Iterator, Optional + +from agent.interrupt_compat import request_hard_interrupt + +_ACTIVE_SCOPE: ContextVar[Optional["InterruptScope"]] = ContextVar("hermes_interrupt_scope", default=None) + + +class InterruptScope: + def __init__(self) -> None: + self._lock = threading.Lock() + self._agents: list[Any] = [] + self.reason: Optional[str] = None + + def cancel(self, reason: str) -> None: + """Latch ``reason`` and hard-interrupt every agent running under this scope.""" + with self._lock: + self.reason = reason + agents = list(self._agents) + for agent in agents: + request_hard_interrupt(agent, reason, tool_reason="host cancelled the command") + + @contextmanager + def track(self, agent: Any) -> Iterator[None]: + with self._lock: + self._agents.append(agent) + reason = self.reason + if reason is not None: + request_hard_interrupt(agent, reason, tool_reason="host cancelled the command") + try: + yield + finally: + with self._lock: + self._agents.remove(agent) + + +@contextmanager +def bind_interrupt_scope(scope: Optional[InterruptScope]) -> Iterator[None]: + """Make ``scope`` the owner of every agent turn started in this context.""" + token = _ACTIVE_SCOPE.set(scope) + try: + yield + finally: + _ACTIVE_SCOPE.reset(token) + + +def track_in_interrupt_scope(agent: Any): + """Register ``agent`` with the bound scope for the duration of its turn (no-op without one).""" + scope = _ACTIVE_SCOPE.get() + return nullcontext() if scope is None else scope.track(agent) diff --git a/agent/turn_facade.py b/agent/turn_facade.py index 5f147c62cf..128ecd1eca 100644 --- a/agent/turn_facade.py +++ b/agent/turn_facade.py @@ -47,6 +47,7 @@ class TurnFacadeMixin: from agent.prompt_cache_scope import declared_conversation_scope_safe from agent.review_idle_queue import QUEUE as _review_queue from agent.subagent_lifecycle import bind_subagent_parent + from agent.interrupt_scope import track_in_interrupt_scope from agent.turn_facade_lease import admit_durable_turn_lease from hermes_cli.observability.relay_shared_metrics import finish_task_run, start_task_run @@ -115,7 +116,8 @@ class TurnFacadeMixin: ) # Keep the ContextVar scope local (agent tokens may be observed from another thread). - with bind_subagent_parent(self), scoped_runtime_main({}): + # A host that owns this thread (Hermes Console) may cancel the turn cross-thread. + with bind_subagent_parent(self), scoped_runtime_main({}), track_in_interrupt_scope(self): try: if lease is not None: lease.start() diff --git a/hermes_cli/web_routers/chat_ws.py b/hermes_cli/web_routers/chat_ws.py index 6d4dc9f84e..f7d9c217d6 100644 --- a/hermes_cli/web_routers/chat_ws.py +++ b/hermes_cli/web_routers/chat_ws.py @@ -6,7 +6,7 @@ reached through the late-binding seam (cycle-safe). """ import asyncio -import functools +import contextlib import json import logging import re @@ -15,6 +15,7 @@ from typing import Any, Dict, Optional from fastapi import APIRouter, FastAPI, HTTPException, WebSocket, WebSocketDisconnect +from agent.interrupt_scope import InterruptScope, bind_interrupt_scope from hermes_cli.pty_session import RegistryFull from hermes_cli.web_deps import LateState, late from hermes_cli.web_server_chat import ( @@ -154,16 +155,33 @@ async def _close_unless_sidecar_allowed(ws: WebSocket) -> bool: _CONSOLE_PROMPT = "hermes> " _CONSOLE_COMMAND_TIMEOUT_SECONDS = 60.0 +# Cancel/timeout interrupt the worker cooperatively; this bounds how long the prompt waits for it to exit. +_CONSOLE_UNWIND_TIMEOUT_SECONDS = 10.0 _CONSOLE_OUTPUT_LIMIT = 50000 -def _execute_console_line(engine: Any, line: str, *, confirmed: bool, profile: Optional[str]) -> Any: +def _execute_console_line( + engine: Any, line: str, *, confirmed: bool, profile: Optional[str], scope: Optional[InterruptScope] = None, +) -> Any: # _profile_scope swaps process-global skill module paths; keep it inside # the worker thread and never hold it across awaits. - with _profile_scope(profile): + with _profile_scope(profile), bind_interrupt_scope(scope): return engine.execute(line, confirmed=confirmed) +async def _unwind_console_worker(worker: Any, scope: InterruptScope, reason: str) -> None: + """Stop the command's worker after cancel/timeout: asyncio can only drop the waiter, so interrupt + any agent the command forked (closing its provider request) and wait for the thread to exit.""" + scope.cancel(f"Console command {reason}") + if worker.cancel(): # still queued: never ran + return + exited = asyncio.wrap_future(worker) + done, _ = await asyncio.wait({exited}, timeout=_CONSOLE_UNWIND_TIMEOUT_SECONDS) + if not done: + exited.cancel() + _log.warning("console worker still running %ss after %s", _CONSOLE_UNWIND_TIMEOUT_SECONDS, reason) + + class _ConsoleSender: """Serialises frames onto one console socket and owns the prompt suffix.""" @@ -281,18 +299,17 @@ async def console_ws(ws: WebSocket) -> None: async def run_command(line: str, *, confirmed: bool, command_id: int) -> None: nonlocal active_task, pending_confirmation, command_generation + scope = InterruptScope() + worker = _get_console_executor().submit( + _execute_console_line, engine, line, confirmed=confirmed, profile=profile, scope=scope, + ) try: - loop = asyncio.get_running_loop() - result = await asyncio.wait_for( - loop.run_in_executor( - _get_console_executor(), - functools.partial(_execute_console_line, engine, line, confirmed=confirmed, profile=profile), - ), - timeout=_CONSOLE_COMMAND_TIMEOUT_SECONDS, - ) + result = await asyncio.wait_for(asyncio.wrap_future(worker), timeout=_CONSOLE_COMMAND_TIMEOUT_SECONDS) except asyncio.CancelledError: + await _unwind_console_worker(worker, scope, "cancelled") raise except asyncio.TimeoutError: + await _unwind_console_worker(worker, scope, "timed out") if command_id == command_generation: pending_confirmation = None await out.error_then_complete( @@ -343,8 +360,11 @@ async def console_ws(ws: WebSocket) -> None: if frame_type == "cancel": if active_task and not active_task.done(): command_generation += 1 - active_task.cancel() - active_task = None + task, active_task = active_task, None + task.cancel() + # Report cancelled only once the worker (and any provider request it owns) is gone. + with contextlib.suppress(asyncio.CancelledError): + await task pending_confirmation = None await out.prompt(type="complete", status="cancelled") elif pending_confirmation: diff --git a/tests/hermes_cli/test_web_server_console_ws.py b/tests/hermes_cli/test_web_server_console_ws.py index 7b569d1b36..abeff215d4 100644 --- a/tests/hermes_cli/test_web_server_console_ws.py +++ b/tests/hermes_cli/test_web_server_console_ws.py @@ -86,3 +86,130 @@ def test_console_ws_cancel_returns_to_prompt(console_client, monkeypatch): complete = _recv_until(conn, "complete", status="cancelled") assert complete["prompt"] == "hermes> " + + +@pytest.fixture +def blocking_provider(): + """Loopback OpenAI-compatible server whose chat completion blocks until the peer closes + the socket (like a llama.cpp generation) or the test releases it.""" + import json + import select + import socket + import threading + from http.server import BaseHTTPRequestHandler, HTTPServer + + state = {"started": threading.Event(), "peer_closed": threading.Event(), "release": threading.Event()} + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *args): + pass + + def do_POST(self): + self.rfile.read(int(self.headers.get("Content-Length") or 0)) + if not self.path.endswith("/chat/completions"): # capability probes + self.send_response(404) + self.send_header("Content-Length", "0") + self.end_headers() + return + state["started"].set() + deadline = time.monotonic() + 30 + while not state["release"].is_set() and time.monotonic() < deadline: + readable, _, _ = select.select([self.connection], [], [], 0.05) + if readable and self.connection.recv(1, socket.MSG_PEEK) == b"": + state["peer_closed"].set() + return + body = json.dumps({"id": "x", "object": "chat.completion", "model": "test-model", "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "done"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + srv = HTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=srv.serve_forever, daemon=True).start() + state["base_url"] = f"http://127.0.0.1:{srv.server_port}/v1" + try: + yield state + finally: + state["release"].set() + srv.shutdown() + + +@pytest.mark.parametrize("stop", ["cancel", "timeout"]) +def test_console_cancel_stops_forked_agent_request_before_reporting(console_client, monkeypatch, blocking_provider, stop): + """#106179: cancelling (or timing out) a console command whose worker forked an AIAgent must interrupt + that agent — closing its in-flight provider request — and wait for the worker to exit BEFORE the + prompt reports cancelled/timeout. asyncio can only drop the waiter; the thread keeps decoding otherwise.""" + import threading + + from agent import curator + from hermes_cli.web_routers import chat_ws + + monkeypatch.setattr( + curator, "_resolve_review_provider", + lambda: ({"api_key": "test-key", "base_url": blocking_provider["base_url"]}, "test-model", "openai-compat", {}), + ) + worker_exited = threading.Event() + real_execute = chat_ws._execute_console_line + + def observed_execute(*args, **kwargs): + try: + return real_execute(*args, **kwargs) + finally: + worker_exited.set() + + monkeypatch.setattr(chat_ws, "_execute_console_line", observed_execute) + if stop == "timeout": + monkeypatch.setattr(chat_ws, "_CONSOLE_COMMAND_TIMEOUT_SECONDS", 2.0) + line = "curator run --consolidate --dry-run" + + with console_client.websocket_connect(_url()) as conn: + assert conn.receive_json()["type"] == "ready" + conn.send_json({"type": "input", "line": line}) + _recv_until(conn, "complete", status="confirm_required") + assert worker_exited.wait(10) # the confirm probe's worker, not the one under test + worker_exited.clear() + conn.send_json({"type": "confirm", "command": line}) + assert blocking_provider["started"].wait(60), "forked agent never reached the provider" + if stop == "cancel": + conn.send_json({"type": "cancel"}) + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + frame = conn.receive_json() + if frame.get("type") == "complete" and frame.get("status") in {"cancelled", "timeout"}: + break + else: + raise AssertionError("no cancelled/timeout frame") + observed = (frame["status"], blocking_provider["peer_closed"].is_set(), worker_exited.is_set()) + blocking_provider["release"].set() # a leaked worker (the bug) must not wedge socket teardown + worker_exited.wait(30) + assert observed == ("cancelled" if stop == "cancel" else "timeout", True, True), ( + "(status, provider request closed, worker exited) at the terminal frame") + + +def test_interrupt_scope_cancels_agents_that_start_after_the_cancel(): + """A turn that begins after the host cancelled must be interrupted on entry, else a cancel racing + agent construction leaves a live request behind.""" + from agent.interrupt_scope import InterruptScope, bind_interrupt_scope, track_in_interrupt_scope + + class Agent: + def __init__(self): + self.stops = [] + + def hard_interrupt(self, message=None, *, tool_reason=None): + self.stops.append(message) + + scope = InterruptScope() + early, late_agent, unscoped = Agent(), Agent(), Agent() + with bind_interrupt_scope(scope): + with track_in_interrupt_scope(early): + scope.cancel("Console command cancelled") + with track_in_interrupt_scope(late_agent): + pass + with track_in_interrupt_scope(unscoped): # no scope bound: nothing to register with + scope.cancel("Console command cancelled") + assert early.stops == ["Console command cancelled"] + assert late_agent.stops == ["Console command cancelled"] + assert unscoped.stops == []