From 1cbd5910479f3be130c2c4427aefc312e9d34743 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 6 Sep 2026 01:34:15 +0100 Subject: [PATCH 1/2] fix: drain kanban workers before gateway restart --- gateway/kanban_watchers.py | 81 +++++++++- gateway/run_shutdown.py | 25 +++- hermes_cli/kanban_db.py | 2 + hermes_cli/kanban_db_dispatch.py | 99 +++++++++++++ .../test_gateway_restart_kanban_workers.py | 140 ++++++++++++++++++ 5 files changed, 345 insertions(+), 2 deletions(-) create mode 100644 tests/hermes_cli/test_gateway_restart_kanban_workers.py diff --git a/gateway/kanban_watchers.py b/gateway/kanban_watchers.py index c637086db4c08..f782b2185643f 100644 --- a/gateway/kanban_watchers.py +++ b/gateway/kanban_watchers.py @@ -230,6 +230,84 @@ def _kanban_dispatcher_boot(self) -> Optional[tuple]: "on config control alone.", _lock_path) return _load_config, _kb, kanban_cfg + def _kanban_running_workers(self, *, include_wedged: bool = True) -> list[dict]: + """Read live Kanban workers across every non-archived board. + + Workers are child processes with durable PIDs in board databases, not + gateway turns. Enumerating every board prevents a profile-scoped + gateway restart from leaving workers on another board to systemd's + KillMode escalation. + """ + try: + from hermes_cli import kanban_db as _kb + from hermes_cli import kanban_db_connect as _kbc + from hermes_cli import kanban_db_dispatch as _kbd + boards = _kb.list_boards(include_archived=False) + except Exception: + logger.exception("kanban restart drain: cannot enumerate boards") + return [] + + workers: list[dict] = [] + for board_meta in boards: + slug = board_meta.get("slug") or _kb.DEFAULT_BOARD + conn = None + try: + conn = _kbc.connect(board=slug) + workers.extend(_kbd.list_running_workers(conn, include_wedged=include_wedged)) + except Exception: + logger.debug("kanban restart drain: worker scan failed on board %s", slug, exc_info=True) + finally: + if conn is not None: + try: + conn.close() + except Exception: + pass + return workers + + def _active_kanban_worker_count(self) -> int: + """Count live workers, including wedged workers for force cleanup.""" + return len(self._kanban_running_workers(include_wedged=True)) + + def _wedged_kanban_worker_count(self) -> int: + """Count live workers excluded from the graceful restart wait.""" + return sum(1 for worker in self._kanban_running_workers(include_wedged=True) if worker.get("wedged")) + + def _interrupt_kanban_workers_for_restart(self) -> int: + """Durably block and SIGTERM every worker still live at the deadline.""" + try: + from hermes_cli import kanban_db as _kb + from hermes_cli import kanban_db_connect as _kbc + from hermes_cli import kanban_db_dispatch as _kbd + boards = _kb.list_boards(include_archived=False) + except Exception: + logger.exception("kanban restart drain: cannot enumerate boards for force cleanup") + return 0 + + interrupted = 0 + for board_meta in boards: + slug = board_meta.get("slug") or _kb.DEFAULT_BOARD + conn = None + try: + conn = _kbc.connect(board=slug) + prepared = _kbd.prepare_workers_for_gateway_restart( + conn, reason="gateway restart" + ) + interrupted += len(prepared) + if prepared: + logger.warning( + "kanban restart drain: blocked and SIGTERM'd %d worker(s) on board %s", + len(prepared), slug, + ) + except Exception: + logger.exception("kanban restart drain: force cleanup failed on board %s", slug) + finally: + if conn is not None: + try: + conn.close() + except Exception: + pass + return interrupted + async def _kanban_dispatcher_watcher(self) -> None: """Embedded kanban dispatcher — one tick every `dispatch_interval_seconds`. @@ -271,8 +349,9 @@ async def _kanban_dispatcher_watcher(self) -> None: try: # Emergency stop (`hermes pause`): no auto-decompose or # dispatch while paused; running workers finish naturally. - if not _kanban_dispatch_allowed(): + if self._draining or not _kanban_dispatch_allowed(): bad_ticks = 0 + ready_pending = False else: # Re-read the auto-decompose toggle live so disabling it # takes effect on the next tick, not on restart. diff --git a/gateway/run_shutdown.py b/gateway/run_shutdown.py index 8f0281eabec12..51dd7b2f13d8c 100644 --- a/gateway/run_shutdown.py +++ b/gateway/run_shutdown.py @@ -131,11 +131,17 @@ def elapsed(self) -> float: # Active-work accounting def _active_work_count(self) -> int: """All agent work the gateway must expose and drain as one total.""" + count_kanban = getattr(self, "_active_kanban_worker_count", None) + try: + active_kanban = max(0, int(count_kanban())) if callable(count_kanban) else 0 + except Exception: + active_kanban = 0 return ( self._running_agent_count() + self._active_cron_job_count() + self._active_api_run_count() + self._active_deferred_agent_worker_count() + + active_kanban ) @staticmethod @@ -1330,7 +1336,12 @@ def _idle_seconds(agent: Any) -> Optional[float]: def _awaitable_work_count(self) -> int: """Active work minus wedged turns — what the restart wait waits on.""" - return max(0, self._active_work_count() - self._wedged_agent_count()) + count_wedged_kanban = getattr(self, "_wedged_kanban_worker_count", None) + try: + wedged_kanban = max(0, int(count_wedged_kanban())) if callable(count_wedged_kanban) else 0 + except Exception: + wedged_kanban = 0 + return max(0, self._active_work_count() - self._wedged_agent_count() - wedged_kanban) async def _await_active_work_before_restart(self) -> bool: """Wait for in-flight work before ``stop()`` so the requesting turn isn't force-interrupted. @@ -1347,6 +1358,9 @@ async def _await_active_work_before_restart(self) -> bool: "past the inactivity timeout; skipping the after-turn wait " "and proceeding to stop()/drain which will interrupt them", active, ) + interrupt_kanban = getattr(self, "_interrupt_kanban_workers_for_restart", None) + if callable(interrupt_kanban): + interrupt_kanban() return False timeout = float(getattr(self, "_restart_after_turn_timeout", 0.0) or 0.0) if timeout <= 0: @@ -1354,6 +1368,9 @@ async def _await_active_work_before_restart(self) -> bool: "Restart requested with %d active work unit(s); " "restart_after_turn_timeout=0 — entering stop()/drain immediately", active, ) + interrupt_kanban = getattr(self, "_interrupt_kanban_workers_for_restart", None) + if callable(interrupt_kanban): + interrupt_kanban() return False logger.info( "Restart requested with %d active work unit(s); " @@ -1372,6 +1389,9 @@ async def _await_active_work_before_restart(self) -> bool: "still active; proceeding to stop()/drain which may " "interrupt remaining work (#77184)", timeout, self._active_work_count(), ) + interrupt_kanban = getattr(self, "_interrupt_kanban_workers_for_restart", None) + if callable(interrupt_kanban): + interrupt_kanban() return False if (now - last_status_at) >= 30.0: logger.info( @@ -1387,6 +1407,9 @@ async def _await_active_work_before_restart(self) -> bool: "Restart deferred wait: %d wedged work unit(s) remain; " "proceeding to stop()/drain which will interrupt them", self._active_work_count(), ) + interrupt_kanban = getattr(self, "_interrupt_kanban_workers_for_restart", None) + if callable(interrupt_kanban): + interrupt_kanban() return False logger.info("Restart deferred wait complete — active work drained; proceeding to stop()") return True diff --git a/hermes_cli/kanban_db.py b/hermes_cli/kanban_db.py index 47fdffc353b1f..4891a094ab2d5 100644 --- a/hermes_cli/kanban_db.py +++ b/hermes_cli/kanban_db.py @@ -4191,6 +4191,8 @@ def latest_summaries(conn: sqlite3.Connection, task_ids: Iterable[str]) -> dict[ _clear_failure_counter, _defer_reclaim_for_live_worker, _pid_alive, + list_running_workers, + prepare_workers_for_gateway_restart, _terminate_reclaimed_worker, _worker_survived_termination, _worker_terminal_timeout_env, diff --git a/hermes_cli/kanban_db_dispatch.py b/hermes_cli/kanban_db_dispatch.py index 743e91dbb5421..d3d3dfd1c28ff 100644 --- a/hermes_cli/kanban_db_dispatch.py +++ b/hermes_cli/kanban_db_dispatch.py @@ -506,6 +506,105 @@ def enforce_max_runtime(conn: sqlite3.Connection, *, signal_fn=None) -> list[str _STALE_HEARTBEAT_GAP_SECONDS = 3600 +def list_running_workers( + conn: sqlite3.Connection, + *, + include_wedged: bool = False, + now: Optional[int] = None, +) -> list[dict[str, Any]]: + """Return live host-local workers for a gateway lifecycle drain. + + Kanban workers are child processes of the dispatcher, not gateway turns; + their durable PID/heartbeat rows are the source of truth for restart + handoff. Workers with no heartbeat for the same one-hour threshold used + by stale detection are wedged and excluded from the graceful wait unless + ``include_wedged`` is requested for the force path. + """ + current = int(time.time()) if now is None else int(now) + rows = conn.execute( + "SELECT t.id, t.worker_pid, t.claim_lock, t.current_run_id, " + " COALESCE(r.started_at, t.started_at) AS active_started_at, " + " COALESCE(r.last_heartbeat_at, t.last_heartbeat_at) AS last_heartbeat_at " + "FROM tasks t " + "LEFT JOIN task_runs r ON r.id = t.current_run_id " + "WHERE t.status = 'running' AND t.worker_pid IS NOT NULL" + ).fetchall() + workers: list[dict[str, Any]] = [] + for row in rows: + claim_lock = row["claim_lock"] or "" + if not claim_lock.startswith(_kb._host_prefix()): + continue + pid = int(row["worker_pid"]) + if not _pid_alive(pid): + continue + last_activity = row["last_heartbeat_at"] or row["active_started_at"] + wedged = last_activity is None or current - int(last_activity) >= _STALE_HEARTBEAT_GAP_SECONDS + if wedged and not include_wedged: + continue + workers.append({ + "task_id": row["id"], + "worker_pid": pid, + "claim_lock": claim_lock, + "run_id": int(row["current_run_id"]) if row["current_run_id"] is not None else None, + "started_at": row["active_started_at"], + "last_heartbeat_at": row["last_heartbeat_at"], + "wedged": wedged, + }) + return workers + + +def prepare_workers_for_gateway_restart( + conn: sqlite3.Connection, + *, + reason: str = "gateway restart", + signal_fn=None, +) -> list[dict[str, Any]]: + """Block and SIGTERM live workers before a gateway restart force path. + + The durable block commits before SIGTERM. A worker later killed by the + supervisor is therefore an intentional restart interruption, not a crash + or retry-consuming dispatcher failure. The run-id CAS leaves a task + alone if it completed between the read and the transition. + """ + import signal + + prepared: list[dict[str, Any]] = [] + for worker in list_running_workers(conn, include_wedged=True): + run_id = worker["run_id"] + if run_id is None: + continue + if not _kb.block_task( + conn, + worker["task_id"], + reason=reason, + kind="transient", + expected_run_id=run_id, + ): + continue + payload = { + "worker_pid": worker["worker_pid"], + "run_id": run_id, + "reason": reason, + "wedged": worker["wedged"], + "signal": "SIGTERM", + } + with _kb.write_txn(conn): + _kb._append_event( + conn, worker["task_id"], "gateway_restart_worker", payload, run_id=run_id, + ) + kill = signal_fn if signal_fn is not None else os.kill + signaled = False + try: + kill(worker["worker_pid"], signal.SIGTERM) + signaled = True + except (ProcessLookupError, OSError): + # The durable block won the race; an already-exited child is not a + # second failure and does not need a signal. + pass + prepared.append({**payload, "task_id": worker["task_id"], "signaled": signaled}) + return prepared + + def detect_stale_running( conn: sqlite3.Connection, *, diff --git a/tests/hermes_cli/test_gateway_restart_kanban_workers.py b/tests/hermes_cli/test_gateway_restart_kanban_workers.py new file mode 100644 index 0000000000000..b923d99e9eb89 --- /dev/null +++ b/tests/hermes_cli/test_gateway_restart_kanban_workers.py @@ -0,0 +1,140 @@ +"""Regression coverage for gateway restart handoff of Kanban workers (#77184).""" + +import asyncio +import signal +from pathlib import Path + +import pytest + +from gateway.run_shutdown import GatewayShutdownMixin +from hermes_cli import kanban_db as kb +from hermes_cli import kanban_db_connect as kbc +from hermes_cli import kanban_db_dispatch as kbd + + +@pytest.fixture +def conn(tmp_path: Path): + connection = kbc.connect(db_path=tmp_path / "kanban.db") + try: + yield connection + finally: + connection.close() + + +def _running_task(conn, monkeypatch, *, heartbeat: int | None = 1, pid: int = 4242) -> str: + task_id = kb.create_task(conn, title="restart handoff", assignee="worker") + host = kb._claimer_id().split(":", 1)[0] + assert kb.claim_task(conn, task_id, claimer=f"{host}:dispatcher") is not None + kbd._set_worker_pid(conn, task_id, pid) + conn.execute( + "UPDATE tasks SET started_at = ?, last_heartbeat_at = ? WHERE id = ?", + (1, heartbeat, task_id), + ) + conn.execute( + "UPDATE task_runs SET started_at = ?, last_heartbeat_at = ? WHERE task_id = ?", + (1, heartbeat, task_id), + ) + conn.commit() + monkeypatch.setattr(kbd, "_pid_alive", lambda _pid: True) + return task_id + + +def test_running_worker_scan_excludes_wedged_from_graceful_wait(conn, monkeypatch): + fresh = _running_task(conn, monkeypatch, heartbeat=4990, pid=4242) + wedged = _running_task(conn, monkeypatch, heartbeat=1, pid=4243) + + graceful = kbd.list_running_workers(conn, now=5000) + assert [worker["task_id"] for worker in graceful] == [fresh] + + force = kbd.list_running_workers(conn, include_wedged=True, now=5000) + assert {worker["task_id"] for worker in force} == {fresh, wedged} + assert next(worker for worker in force if worker["task_id"] == wedged)["wedged"] is True + + +def test_restart_force_path_blocks_before_sigterm_without_retry_failure(conn, monkeypatch): + task_id = _running_task(conn, monkeypatch, heartbeat=1, pid=4242) + sent: list[tuple[int, int]] = [] + + prepared = kbd.prepare_workers_for_gateway_restart( + conn, + signal_fn=lambda pid, sig: sent.append((pid, sig)), + ) + + assert len(prepared) == 1 + assert prepared[0]["task_id"] == task_id + assert prepared[0]["reason"] == "gateway restart" + assert prepared[0]["signaled"] is True + row = conn.execute( + "SELECT status, worker_pid, consecutive_failures FROM tasks WHERE id = ?", + (task_id,), + ).fetchone() + assert dict(row) == { + "status": "blocked", + "worker_pid": None, + "consecutive_failures": 0, + } + assert sent == [(4242, signal.SIGTERM)] + event = conn.execute( + "SELECT kind, payload FROM task_events WHERE task_id = ? ORDER BY id DESC LIMIT 1", + (task_id,), + ).fetchone() + assert event["kind"] == "gateway_restart_worker" + assert "gateway restart" in event["payload"] + + +def test_restart_wait_excludes_wedged_kanban_workers_and_interrupts_them(): + interrupted: list[str] = [] + + class Runner(GatewayShutdownMixin): + _restart_after_turn_timeout = 1800 + + def _running_agent_count(self): + return 0 + + def _active_cron_job_count(self): + return 0 + + def _active_api_run_count(self): + return 0 + + def _active_deferred_agent_worker_count(self): + return 0 + + def _active_kanban_worker_count(self): + return 1 + + def _wedged_agent_count(self): + return 0 + + def _wedged_kanban_worker_count(self): + return 1 + + def _interrupt_kanban_workers_for_restart(self): + interrupted.append("gateway restart") + + def _scale_to_zero_status(self, *args, **kwargs): + pass + + result = asyncio.run(Runner()._await_active_work_before_restart()) + assert result is False + assert interrupted == ["gateway restart"] + + +def test_active_work_count_includes_kanban_workers(): + class Runner(GatewayShutdownMixin): + def _running_agent_count(self): + return 1 + + def _active_cron_job_count(self): + return 2 + + def _active_api_run_count(self): + return 3 + + def _active_deferred_agent_worker_count(self): + return 4 + + def _active_kanban_worker_count(self): + return 5 + + assert Runner()._active_work_count() == 15 From 8c3a6546356a5a05dbb52ce6ea5f7a44445da7ec Mon Sep 17 00:00:00 2001 From: devops Date: Sun, 6 Sep 2026 03:06:53 +0100 Subject: [PATCH 2/2] fix(gateway): fail closed on kanban worker scan errors --- gateway/kanban_watchers.py | 81 +++++---- gateway/run_shutdown.py | 158 +++++++++++++----- hermes_cli/kanban_db.py | 1 + hermes_cli/kanban_db_dispatch.py | 4 + .../test_restart_kanban_scan_fail_closed.py | 123 ++++++++++++++ 5 files changed, 293 insertions(+), 74 deletions(-) create mode 100644 tests/gateway/test_restart_kanban_scan_fail_closed.py diff --git a/gateway/kanban_watchers.py b/gateway/kanban_watchers.py index f782b2185643f..3469359cb8052 100644 --- a/gateway/kanban_watchers.py +++ b/gateway/kanban_watchers.py @@ -230,38 +230,49 @@ def _kanban_dispatcher_boot(self) -> Optional[tuple]: "on config control alone.", _lock_path) return _load_config, _kb, kanban_cfg + def _kanban_board_metadata(self): + """Return every non-archived board or fail closed.""" + from hermes_cli import kanban_db as _kb + from hermes_cli import kanban_db_connect as _kbc + from hermes_cli import kanban_db_dispatch as _kbd + + try: + boards = _kb.list_boards(include_archived=False) + except Exception as exc: + raise _kbd.KanbanWorkerScanError( + "kanban restart drain: board enumeration failed" + ) from exc + if not isinstance(boards, list): + raise _kbd.KanbanWorkerScanError( + "kanban restart drain: board enumeration returned a non-list" + ) + return _kb, _kbc, _kbd, boards + def _kanban_running_workers(self, *, include_wedged: bool = True) -> list[dict]: """Read live Kanban workers across every non-archived board. Workers are child processes with durable PIDs in board databases, not - gateway turns. Enumerating every board prevents a profile-scoped - gateway restart from leaving workers on another board to systemd's - KillMode escalation. + gateway turns. Any enumeration or per-board DB failure is fatal to this + scan: proceeding with a partial list could still leave a worker for + systemd's KillMode to SIGKILL. """ - try: - from hermes_cli import kanban_db as _kb - from hermes_cli import kanban_db_connect as _kbc - from hermes_cli import kanban_db_dispatch as _kbd - boards = _kb.list_boards(include_archived=False) - except Exception: - logger.exception("kanban restart drain: cannot enumerate boards") - return [] - + _kb, _kbc, _kbd, boards = self._kanban_board_metadata() workers: list[dict] = [] for board_meta in boards: slug = board_meta.get("slug") or _kb.DEFAULT_BOARD conn = None try: conn = _kbc.connect(board=slug) - workers.extend(_kbd.list_running_workers(conn, include_wedged=include_wedged)) - except Exception: - logger.debug("kanban restart drain: worker scan failed on board %s", slug, exc_info=True) + workers.extend( + _kbd.list_running_workers(conn, include_wedged=include_wedged) + ) + except Exception as exc: + raise _kbd.KanbanWorkerScanError( + f"kanban restart drain: worker scan failed on board {slug}" + ) from exc finally: if conn is not None: - try: - conn.close() - except Exception: - pass + conn.close() return workers def _active_kanban_worker_count(self) -> int: @@ -270,18 +281,18 @@ def _active_kanban_worker_count(self) -> int: def _wedged_kanban_worker_count(self) -> int: """Count live workers excluded from the graceful restart wait.""" - return sum(1 for worker in self._kanban_running_workers(include_wedged=True) if worker.get("wedged")) + return sum( + 1 + for worker in self._kanban_running_workers(include_wedged=True) + if worker.get("wedged") + ) def _interrupt_kanban_workers_for_restart(self) -> int: """Durably block and SIGTERM every worker still live at the deadline.""" - try: - from hermes_cli import kanban_db as _kb - from hermes_cli import kanban_db_connect as _kbc - from hermes_cli import kanban_db_dispatch as _kbd - boards = _kb.list_boards(include_archived=False) - except Exception: - logger.exception("kanban restart drain: cannot enumerate boards for force cleanup") - return 0 + # Preflight every board before changing any task. If any board cannot be + # inspected, propagate the scan error and keep the gateway draining. + self._kanban_running_workers(include_wedged=True) + _kb, _kbc, _kbd, boards = self._kanban_board_metadata() interrupted = 0 for board_meta in boards: @@ -296,16 +307,16 @@ def _interrupt_kanban_workers_for_restart(self) -> int: if prepared: logger.warning( "kanban restart drain: blocked and SIGTERM'd %d worker(s) on board %s", - len(prepared), slug, + len(prepared), + slug, ) - except Exception: - logger.exception("kanban restart drain: force cleanup failed on board %s", slug) + except Exception as exc: + raise _kbd.KanbanWorkerScanError( + f"kanban restart drain: force-interrupt failed on board {slug}" + ) from exc finally: if conn is not None: - try: - conn.close() - except Exception: - pass + conn.close() return interrupted async def _kanban_dispatcher_watcher(self) -> None: diff --git a/gateway/run_shutdown.py b/gateway/run_shutdown.py index 51dd7b2f13d8c..7b6e04a5ca6ec 100644 --- a/gateway/run_shutdown.py +++ b/gateway/run_shutdown.py @@ -132,10 +132,7 @@ def elapsed(self) -> float: def _active_work_count(self) -> int: """All agent work the gateway must expose and drain as one total.""" count_kanban = getattr(self, "_active_kanban_worker_count", None) - try: - active_kanban = max(0, int(count_kanban())) if callable(count_kanban) else 0 - except Exception: - active_kanban = 0 + active_kanban = max(0, int(count_kanban())) if callable(count_kanban) else 0 return ( self._running_agent_count() + self._active_cron_job_count() @@ -1335,82 +1332,154 @@ def _idle_seconds(agent: Any) -> Optional[float]: ) def _awaitable_work_count(self) -> int: - """Active work minus wedged turns — what the restart wait waits on.""" + """Active work minus wedged turns — what the restart wait waits on. + + Kanban visibility is safety-critical: an unreadable board must not be + converted into zero workers and allow ``stop()`` to run. + """ count_wedged_kanban = getattr(self, "_wedged_kanban_worker_count", None) - try: - wedged_kanban = max(0, int(count_wedged_kanban())) if callable(count_wedged_kanban) else 0 - except Exception: - wedged_kanban = 0 + wedged_kanban = ( + max(0, int(count_wedged_kanban())) + if callable(count_wedged_kanban) + else 0 + ) return max(0, self._active_work_count() - self._wedged_agent_count() - wedged_kanban) - async def _await_active_work_before_restart(self) -> bool: + async def _await_active_work_before_restart(self) -> Optional[bool]: """Wait for in-flight work before ``stop()`` so the requesting turn isn't force-interrupted. - Wedged turns are excluded (restart is their remedy). True when drained to zero, False when the - cap elapsed or only wedged work remains (caller proceeds to ``stop()``). + ``True`` means all work drained safely; ``False`` means the bounded wait + expired or only known wedged work remains and force cleanup is complete. + ``None`` means Kanban accounting was incomplete or force cleanup failed; + the caller must keep the gateway alive and draining rather than calling + ``stop()`` while systemd could SIGKILL an unaccounted worker. """ - active = self._active_work_count() + def _interrupt_kanban_or_abort() -> bool: + interrupt_kanban = getattr(self, "_interrupt_kanban_workers_for_restart", None) + if not callable(interrupt_kanban): + return True + try: + interrupt_kanban() + except Exception: + logger.exception( + "Restart drain cannot account for all Kanban boards; " + "leaving gateway draining without calling stop()" + ) + return False + return True + + try: + active = self._active_work_count() + except Exception: + logger.exception( + "Restart drain cannot account for all Kanban boards; " + "leaving gateway draining without calling stop()" + ) + return None if active <= 0: return True - if self._awaitable_work_count() <= 0: + + try: + awaitable = self._awaitable_work_count() + except Exception: + logger.exception( + "Restart drain lost Kanban worker visibility; " + "leaving gateway draining without calling stop()" + ) + return None + if awaitable <= 0: logger.warning( "Restart requested with %d active work unit(s), all wedged " "past the inactivity timeout; skipping the after-turn wait " - "and proceeding to stop()/drain which will interrupt them", active, + "and proceeding to stop()/drain which will interrupt them", + active, ) - interrupt_kanban = getattr(self, "_interrupt_kanban_workers_for_restart", None) - if callable(interrupt_kanban): - interrupt_kanban() - return False + return False if _interrupt_kanban_or_abort() else None + timeout = float(getattr(self, "_restart_after_turn_timeout", 0.0) or 0.0) if timeout <= 0: logger.info( "Restart requested with %d active work unit(s); " - "restart_after_turn_timeout=0 — entering stop()/drain immediately", active, + "restart_after_turn_timeout=0 — entering stop()/drain immediately", + active, ) - interrupt_kanban = getattr(self, "_interrupt_kanban_workers_for_restart", None) - if callable(interrupt_kanban): - interrupt_kanban() - return False + return False if _interrupt_kanban_or_abort() else None + logger.info( "Restart requested with %d active work unit(s); " "deferring stop() until they finish (cap=%.0fs) so in-flight " - "turns are not amputated (#77184)", active, timeout, + "turns are not amputated (#77184)", + active, + timeout, ) self._scale_to_zero_status("draining", "restart wait: status mark failed") loop = asyncio.get_running_loop() deadline = loop.time() + timeout last_status_at = 0.0 - while self._awaitable_work_count() > 0: + while True: + try: + awaitable = self._awaitable_work_count() + except Exception: + logger.exception( + "Restart drain lost Kanban worker visibility; " + "leaving gateway draining without calling stop()" + ) + return None + if awaitable <= 0: + break now = loop.time() if now >= deadline: + try: + active_now = self._active_work_count() + except Exception: + logger.exception( + "Restart drain cannot account for all Kanban boards; " + "leaving gateway draining without calling stop()" + ) + return None logger.warning( "Restart after-turn wait timed out after %.0fs with %d " "still active; proceeding to stop()/drain which may " - "interrupt remaining work (#77184)", timeout, self._active_work_count(), + "interrupt remaining work (#77184)", + timeout, + active_now, ) - interrupt_kanban = getattr(self, "_interrupt_kanban_workers_for_restart", None) - if callable(interrupt_kanban): - interrupt_kanban() - return False + return False if _interrupt_kanban_or_abort() else None if (now - last_status_at) >= 30.0: + try: + status_awaitable = self._awaitable_work_count() + except Exception: + logger.exception( + "Restart drain lost Kanban worker visibility; " + "leaving gateway draining without calling stop()" + ) + return None logger.info( "Restart deferred: waiting on %d active work unit(s) " "(%d wedged and excluded; %.0fs remaining before force drain)", - self._awaitable_work_count(), self._wedged_agent_count(), deadline - now, + status_awaitable, + self._wedged_agent_count(), + deadline - now, ) self._scale_to_zero_status("draining", "restart wait: status mark failed") last_status_at = now await asyncio.sleep(0.1) - if self._active_work_count() > 0: + + try: + active = self._active_work_count() + except Exception: + logger.exception( + "Restart drain cannot account for all Kanban boards; " + "leaving gateway draining without calling stop()" + ) + return None + if active > 0: logger.warning( "Restart deferred wait: %d wedged work unit(s) remain; " - "proceeding to stop()/drain which will interrupt them", self._active_work_count(), + "proceeding to stop()/drain which will interrupt them", + active, ) - interrupt_kanban = getattr(self, "_interrupt_kanban_workers_for_restart", None) - if callable(interrupt_kanban): - interrupt_kanban() - return False + return False if _interrupt_kanban_or_abort() else None logger.info("Restart deferred wait complete — active work drained; proceeding to stop()") return True @@ -1425,7 +1494,18 @@ def request_restart(self, *, detached: bool = False, via_service: bool = False) self._draining = True async def _run_restart() -> None: - await self._await_active_work_before_restart() + drain_result = await self._await_active_work_before_restart() + if drain_result is None: + # A partial board scan is unsafe: keeping the gateway alive and + # draining is preferable to letting systemd KillMode amputate + # an unaccounted worker. A later operator restart request can + # retry after the board/runtime fault is repaired. + self._restart_task_started = False + logger.error( + "Gateway restart deferred: Kanban worker accounting failed; " + "gateway remains alive and draining" + ) + return # Detached helper only AFTER the after-turn wait, or its drain_timeout+5 deadline fires mid-turn. if detached: with _log_suppressed(logging.ERROR, "Failed to launch detached gateway restart helper: %s"): diff --git a/hermes_cli/kanban_db.py b/hermes_cli/kanban_db.py index 4891a094ab2d5..086353a61a908 100644 --- a/hermes_cli/kanban_db.py +++ b/hermes_cli/kanban_db.py @@ -4188,6 +4188,7 @@ def latest_summaries(conn: sqlite3.Connection, task_ids: Iterable[str]) -> dict[ DEFAULT_FAILURE_LIMIT, DEFAULT_RATE_LIMIT_COOLDOWN_SECONDS, DispatchResult, + KanbanWorkerScanError, _clear_failure_counter, _defer_reclaim_for_live_worker, _pid_alive, diff --git a/hermes_cli/kanban_db_dispatch.py b/hermes_cli/kanban_db_dispatch.py index d3d3dfd1c28ff..0bbc4d92e4920 100644 --- a/hermes_cli/kanban_db_dispatch.py +++ b/hermes_cli/kanban_db_dispatch.py @@ -506,6 +506,10 @@ def enforce_max_runtime(conn: sqlite3.Connection, *, signal_fn=None) -> list[str _STALE_HEARTBEAT_GAP_SECONDS = 3600 +class KanbanWorkerScanError(RuntimeError): + """A gateway restart cannot prove that every board was inspected.""" + + def list_running_workers( conn: sqlite3.Connection, *, diff --git a/tests/gateway/test_restart_kanban_scan_fail_closed.py b/tests/gateway/test_restart_kanban_scan_fail_closed.py new file mode 100644 index 0000000000000..114cd69472728 --- /dev/null +++ b/tests/gateway/test_restart_kanban_scan_fail_closed.py @@ -0,0 +1,123 @@ +"""Fail-closed restart accounting for embedded Kanban workers.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import pytest + +from gateway.kanban_watchers import GatewayKanbanWatchersMixin +from gateway.run_shutdown import GatewayShutdownMixin +from hermes_cli import kanban_db as kb +from hermes_cli import kanban_db_connect as kbc +from hermes_cli import kanban_db_dispatch as kbd + + +class _Watcher(GatewayKanbanWatchersMixin): + pass + + +class _Conn: + def close(self): + pass + + +def test_worker_scan_fails_closed_on_board_enumeration_error(monkeypatch): + def fail_enumeration(*, include_archived=False): + raise OSError("boards root unavailable") + + monkeypatch.setattr(kb, "list_boards", fail_enumeration) + + with pytest.raises(kbd.KanbanWorkerScanError, match="board enumeration failed"): + _Watcher()._kanban_running_workers() + + +def test_worker_scan_fails_closed_on_later_board_error(monkeypatch): + monkeypatch.setattr( + kb, + "list_boards", + lambda **_: [{"slug": "jarvis-os"}, {"slug": "sycode-trading"}], + ) + monkeypatch.setattr(kbc, "connect", lambda *, board: _Conn()) + + calls = [] + + def scan_board(_conn, *, include_wedged=False): + calls.append(True) + if len(calls) == 2: + raise OSError("second board unavailable") + return [] + + monkeypatch.setattr(kbd, "list_running_workers", scan_board) + + with pytest.raises(kbd.KanbanWorkerScanError, match="board sycode-trading"): + _Watcher()._kanban_running_workers() + + +def test_force_cleanup_propagates_scan_failure_without_returning_zero(monkeypatch): + monkeypatch.setattr(kb, "list_boards", lambda **_: [{"slug": "jarvis-os"}]) + monkeypatch.setattr(kbc, "connect", lambda *, board: _Conn()) + monkeypatch.setattr( + kbd, + "list_running_workers", + lambda _conn, **_: (_ for _ in ()).throw(OSError("scan unavailable")), + ) + + with pytest.raises(kbd.KanbanWorkerScanError): + _Watcher()._interrupt_kanban_workers_for_restart() + + +def test_active_work_count_propagates_kanban_scan_error(): + class Runner(GatewayShutdownMixin): + def _running_agent_count(self): + return 0 + + def _active_cron_job_count(self): + return 0 + + def _active_api_run_count(self): + return 0 + + def _active_deferred_agent_worker_count(self): + return 0 + + def _active_kanban_worker_count(self): + raise kbd.KanbanWorkerScanError("board scan failed") + + with pytest.raises(kbd.KanbanWorkerScanError): + Runner()._active_work_count() + + +@pytest.mark.asyncio +async def test_restart_stays_draining_when_kanban_scan_fails(): + class Runner(GatewayShutdownMixin): + _restart_after_turn_timeout = 1800 + + def _running_agent_count(self): + return 0 + + def _active_cron_job_count(self): + return 0 + + def _active_api_run_count(self): + return 0 + + def _active_deferred_agent_worker_count(self): + return 0 + + def _active_kanban_worker_count(self): + raise kbd.KanbanWorkerScanError("board scan failed") + + def _scale_to_zero_status(self, *args, **kwargs): + pass + + runner = Runner() + runner._restart_task_started = False + runner.stop = AsyncMock() + + assert runner.request_restart(detached=False, via_service=True) is True + await runner._restart_task + + runner.stop.assert_not_awaited() + assert runner._draining is True + assert runner._restart_task_started is False