diff --git a/hermes_cli/gateway_client.py b/hermes_cli/gateway_client.py index cbaf355d36..fa7cc3f7e2 100644 --- a/hermes_cli/gateway_client.py +++ b/hermes_cli/gateway_client.py @@ -3,7 +3,8 @@ from __future__ import annotations import asyncio from collections import deque -from contextlib import asynccontextmanager, suppress +from contextlib import asynccontextmanager, contextmanager, suppress +import contextvars import json import os from pathlib import Path @@ -266,13 +267,31 @@ def _session_ticket(home: Path, endpoint, *, purpose="interactive", scope=None, return grant["ticket"] +# A restart's after-turn wait admits clients so a human can answer a parked prompt (#139041). A machine +# poller held in it (a Kanban worker reading its card's receipt) keeps the turn the wait is waiting +# on alive, so the restart, and ``hermes update`` behind it, is held until its cap. +_REFUSE_RESTART_WAIT: contextvars.ContextVar[bool] = contextvars.ContextVar("refuse_restart_wait", default=False) + + +@contextmanager +def outside_restart_wait(): + """``connect_gateway`` calls inside never attach to an owner in its restart's after-turn wait.""" + token = _REFUSE_RESTART_WAIT.set(True) + try: + yield + finally: + _REFUSE_RESTART_WAIT.reset(token) + + @asynccontextmanager async def connect_gateway(endpoint=None, subject=None, *, max_size=8 * 1024 * 1024): """Attach to this home's owner. ``endpoint``: an already-discovered ready owner to dial as-is (maintenance that must never start a gateway); default ensures one, starting it if needed. ``subject``: mint the ticket for this principal (the caller's authenticated dashboard user) rather than the socket identity; only a peer the owner already trusts may name one. - ``max_size``: the largest frame accepted; a ``session.resume`` answer carries the whole history.""" + ``max_size``: the largest frame accepted; a ``session.resume`` answer carries the whole history. + Inside ``outside_restart_wait()`` an owner in a restart's after-turn wait is refused as + ``draining`` (``restart_wait``) instead of attached.""" from websockets.asyncio.client import connect from hermes_constants import get_hermes_home from hermes_cli.gateway_runtime import attachable, ensure_gateway_runtime, gateway_profile_home @@ -296,7 +315,7 @@ async def connect_gateway(endpoint=None, subject=None, *, max_size=8 * 1024 * 10 # A client's own start: the daemon it may launch ends itself once idle (no client, adapter, # cron/kanban work or admission), and the next client starts a fresh one. result = await asyncio.to_thread(ensure_gateway_runtime, home, idle_exit=True) - if not attachable(result): + if not attachable(result) or (_REFUSE_RESTART_WAIT.get() and result.state != "ready"): detail = getattr(result, "detail", None) if result.reason_code == "runtime_exited" and detail: # A multi-line startup report (redacted traceback tail + where the full log is). diff --git a/hermes_cli/kanban_worker_client.py b/hermes_cli/kanban_worker_client.py index d82dfe0982..399c1c085a 100644 --- a/hermes_cli/kanban_worker_client.py +++ b/hermes_cli/kanban_worker_client.py @@ -58,7 +58,7 @@ class _Claim: async def run(params, board_db): from agent.retry_utils import jittered_backoff - from hermes_cli.gateway_client import connect_gateway, GatewayClientError + from hermes_cli.gateway_client import connect_gateway, GatewayClientError, outside_restart_wait from websockets.exceptions import WebSocketException if os.environ.get('HERMES_TUI_GATEWAY_URL'): raise GatewayClientError('Kanban requires the assigned local profile owner') @@ -71,15 +71,19 @@ async def run(params, board_db): while True: seen = claim.answered try: - async with connect_gateway() as client: - # kanban.run subscribes this socket to the card's session, but only receipts matter - # here: discard the turn's events, or they fill the client's bounded queue and its - # reader stops, with the next receipt behind it. - drain = asyncio.create_task(_discard(client.events)) - try: - receipt = await _await_receipt(client, params, board_db, claim) - finally: - drain.cancel() + # Never into a restart's after-turn wait: it cut this socket so the turn's submitter + # waits for the NEXT owner, and a worker polling its receipt inside the wait keeps the + # turn it waits on alive, holding `hermes update` until restart_after_turn_timeout. + with outside_restart_wait(): + async with connect_gateway() as client: + # kanban.run subscribes this socket to the card's session, but only receipts + # matter here: discard the turn's events, or they fill the client's bounded + # queue and its reader stops, with the next receipt behind it. + drain = asyncio.create_task(_discard(client.events)) + try: + receipt = await _await_receipt(client, params, board_db, claim) + finally: + drain.cancel() # A restart also surfaces as a refused or reset dial, a failed handshake, or an answer # that never came. except (GatewayClientError, WebSocketException, OSError, TimeoutError) as exc: diff --git a/tests/hermes_cli/test_kanban_worker_restart_wait.py b/tests/hermes_cli/test_kanban_worker_restart_wait.py new file mode 100644 index 0000000000..a38116b42f --- /dev/null +++ b/tests/hermes_cli/test_kanban_worker_restart_wait.py @@ -0,0 +1,109 @@ +"""A Kanban worker never re-enters an owner's restart after-turn wait. + +The wait cuts every socket opened before it so their clients go to the NEXT owner. A worker that +redials into the wait and keeps polling ``prompt.receipt`` stays alive, so ``_watch_submitter`` +never ends the card's turn and the restart (``hermes update``) is held until +``agent.restart_after_turn_timeout``. The wait still admits interactive clients (#139041). +""" + +import asyncio +from contextlib import asynccontextmanager +from types import SimpleNamespace + +import pytest + +from hermes_cli import gateway_client, gateway_runtime, kanban_worker_client +from gateway.drain_control import RESTART_WAIT_REASON + +GUE = gateway_client.GatewayUnavailableError + + +@pytest.fixture(autouse=True) +def _env(monkeypatch): + monkeypatch.delenv("HERMES_TUI_GATEWAY_URL", raising=False) + + +def _restart_wait_verdict(*_, **__): + return gateway_runtime.GatewayDiscovery( + "draining", gateway_runtime.GatewayEndpoint("p", "i", 1, 1, "http://127.0.0.1:1", "none", frozenset()), + reason_code=RESTART_WAIT_REASON) + + +@pytest.mark.asyncio +async def test_a_machine_client_is_refused_by_a_restart_wait_and_mints_no_ticket(monkeypatch): + monkeypatch.setattr(gateway_runtime, "ensure_gateway_runtime", _restart_wait_verdict) + minted = [] + + def ticket(*args, **kwargs): + minted.append(args) + raise RuntimeError("ticket minted") + monkeypatch.setattr(gateway_client, "_session_ticket", ticket) + + with pytest.raises(GUE) as raised, gateway_client.outside_restart_wait(): + async with gateway_client.connect_gateway(): + pass + assert (raised.value.state, raised.value.reason_code) == ("draining", RESTART_WAIT_REASON) + assert minted == [] + + # An interactive client still attaches to answer a parked prompt. + with pytest.raises(RuntimeError, match="ticket minted"): + async with gateway_client.connect_gateway(): + pass + assert len(minted) == 1 + + +class _RestartingOwner: + """First connection: the claim is admitted and pending, then the restart wait cuts the socket. + Every later dial lands in the wait, where the turn stays pending (its LLM call is held).""" + + def __init__(self): + self.dials = 0 + self.receipts_in_wait = 0 + + def connect(self): + owner = self + + @asynccontextmanager + async def _connect(*_, **__): + owner.dials += 1 + in_wait = owner.dials > 1 + # The verdict the real connect_gateway gives a restart-wait owner (tested above). + refuse = getattr(gateway_client, "_REFUSE_RESTART_WAIT", None) + if in_wait and refuse is not None and refuse.get(): + raise GUE(f"Gateway draining: {RESTART_WAIT_REASON}", "draining", RESTART_WAIT_REASON) + + class _Client: + events = asyncio.Queue() + + async def rpc(self, method, **params): + if method == "kanban.run": + return {"session_id": "s1", "receipt": {"status": "pending", "admission_id": "a1"}} + if not in_wait: + raise gateway_client.GatewayClientError("Gateway disconnected; turn outcome is unknown") + owner.receipts_in_wait += 1 + if owner.receipts_in_wait > 50: + raise AssertionError("the worker keeps polling its card inside the restart wait") + return {"status": "pending"} + + yield _Client() + + return _connect + + +@pytest.mark.asyncio +async def test_a_worker_cut_by_a_restart_wait_waits_outside_it(monkeypatch, tmp_path): + clock = [0.0] + monkeypatch.setattr(kanban_worker_client, "time", SimpleNamespace(monotonic=lambda: clock[0])) + + async def _sleep(seconds): + clock[0] += seconds + monkeypatch.setattr(kanban_worker_client.asyncio, "sleep", _sleep) + owner = _RestartingOwner() + monkeypatch.setattr(gateway_client, "connect_gateway", owner.connect()) + + # The worker waits for the next owner without reading anything from the restarting one, and an + # owner that never comes back still fails it at OWNER_WAIT_S (which ends the turn it held). + with pytest.raises(gateway_client.GatewayClientError, match="owner silent"): + await kanban_worker_client.run({"task_id": "t1"}, str(tmp_path / "kanban.db")) + assert owner.receipts_in_wait == 0 + assert owner.dials > 2