diff --git a/cli.py b/cli.py index 379ca15df1a..956349e351e 100644 --- a/cli.py +++ b/cli.py @@ -15174,6 +15174,9 @@ class HermesCLI(CLIAgentSetupMixin, CLICommandsMixin): _drain_sk = get_current_session_key(default="") for _evt, _synth in process_registry.drain_notifications(session_key=_drain_sk): self._pending_input.put(_synth) + if _evt.get("type") == "async_delegation": + from tools.async_delegation import mark_completion_delivered + mark_completion_delivered(str(_evt.get("delegation_id") or "")) except Exception: pass continue @@ -15336,6 +15339,9 @@ class HermesCLI(CLIAgentSetupMixin, CLICommandsMixin): from tools.process_registry import process_registry for _evt, _synth in process_registry.drain_notifications(): self._pending_input.put(_synth) + if _evt.get("type") == "async_delegation": + from tools.async_delegation import mark_completion_delivered + mark_completion_delivered(str(_evt.get("delegation_id") or "")) except Exception: pass # Non-fatal — don't break the main loop diff --git a/gateway/run.py b/gateway/run.py index 313d9a14825..8dc9403fa04 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -15432,6 +15432,8 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew continue try: await self._inject_watch_notification(synth_text, evt) + from tools.async_delegation import mark_completion_delivered + mark_completion_delivered(str(evt.get("delegation_id") or "")) except Exception as e: logger.error("Async delegation injection error: %s", e) except Exception as e: diff --git a/tests/tools/test_async_delegation.py b/tests/tools/test_async_delegation.py index 7714d3c8c08..bf2187e740c 100644 --- a/tests/tools/test_async_delegation.py +++ b/tests/tools/test_async_delegation.py @@ -5,7 +5,11 @@ onto the shared process_registry.completion_queue, the rich re-injection block formatting, capacity rejection, and crash handling. """ +import json +import os import queue +import subprocess +import sys import threading import time @@ -223,6 +227,84 @@ def test_completed_records_pruned_to_cap(): assert len(ad.list_async_delegations()) <= ad._MAX_RETAINED_COMPLETED +def test_completion_is_persisted_and_delivery_can_be_acknowledged(tmp_path, monkeypatch): + """A finished child remains pending on disk until its queue consumer acks it.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + dispatched = ad.dispatch_async_delegation( + goal="durable", context="ctx", toolsets=["terminal"], role="leaf", + model="m", session_key="owner", parent_session_id="parent", + runner=lambda: {"status": "completed", "summary": "survived"}, + ) + assert _drain_one() is not None + + restored = queue.Queue() + assert ad.restore_undelivered_completions(restored) == 1 + row = ad.get_durable_delegation(dispatched["delegation_id"]) + assert row["origin_session"] == "owner" + assert row["state"] == "completed" + assert row["result"]["summary"] == "survived" + assert row["delivery_state"] == "pending" + assert row["delivery_attempts"] >= 2 + + assert ad.mark_completion_delivered(dispatched["delegation_id"]) + assert ad.restore_undelivered_completions(queue.Queue()) == 0 + assert ad.get_durable_delegation(dispatched["delegation_id"])["delivery_state"] == "delivered" + + +def test_real_process_restart_restores_owned_completion_once(tmp_path): + """Real-import E2E: a fresh interpreter restores a prior process's result.""" + repo = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) + env = {**os.environ, "HERMES_HOME": str(tmp_path), "PYTHONPATH": repo} + producer = r''' +import time +from tools import async_delegation as ad +r = ad.dispatch_async_delegation( + goal="restart", context=None, toolsets=None, role="leaf", model="m", + session_key="owner-session", parent_session_id="durable-parent", + runner=lambda: {"status": "completed", "summary": "after restart"}, +) +deadline = time.time() + 5 +while ad.active_count() and time.time() < deadline: + time.sleep(.01) +print(r["delegation_id"]) +''' + first = subprocess.run( + [sys.executable, "-c", producer], cwd=repo, env=env, + text=True, capture_output=True, timeout=15, check=True, + ) + delegation_id = first.stdout.strip().splitlines()[-1] + + consumer = r''' +import json +from tools.process_registry import process_registry +evt = process_registry.completion_queue.get_nowait() +print(json.dumps(evt, sort_keys=True)) +''' + second = subprocess.run( + [sys.executable, "-c", consumer], cwd=repo, env=env, + text=True, capture_output=True, timeout=15, check=True, + ) + evt = json.loads(second.stdout.strip().splitlines()[-1]) + assert evt["delegation_id"] == delegation_id + assert evt["session_key"] == "owner-session" + assert evt["parent_session_id"] == "durable-parent" + assert evt["summary"] == "after restart" + + acker = f''' +from tools import async_delegation as ad +assert ad.mark_completion_delivered({delegation_id!r}) +''' + subprocess.run( + [sys.executable, "-c", acker], cwd=repo, env=env, + text=True, capture_output=True, timeout=15, check=True, + ) + probe = subprocess.run( + [sys.executable, "-c", "from tools.process_registry import process_registry; print(process_registry.completion_queue.qsize())"], + cwd=repo, env=env, text=True, capture_output=True, timeout=15, check=True, + ) + assert probe.stdout.strip().splitlines()[-1] == "0" + + # --------------------------------------------------------------------------- # Integration: delegate_task(background=True) routing # --------------------------------------------------------------------------- diff --git a/tools/async_delegation.py b/tools/async_delegation.py index 9b4b0e9646a..011b398c78c 100644 --- a/tools/async_delegation.py +++ b/tools/async_delegation.py @@ -36,13 +36,16 @@ logic stays in one place. from __future__ import annotations +import json import logging +import sqlite3 import threading import time import uuid from concurrent.futures import ThreadPoolExecutor from typing import Any, Callable, Dict, List, Optional +from hermes_constants import get_hermes_home from tools.daemon_pool import DaemonThreadPoolExecutor from tools.thread_context import propagate_context_to_thread @@ -72,6 +75,133 @@ _records: Dict[str, Dict[str, Any]] = {} _DEFAULT_MAX_ASYNC_CHILDREN = 3 # How many completed records to retain for status queries before pruning. _MAX_RETAINED_COMPLETED = 50 +_DURABLE_RETENTION_SECONDS = 7 * 24 * 60 * 60 +_DB_LOCK = threading.Lock() + + +def _db_path(): + return get_hermes_home() / "delegations.db" + + +def _connect() -> sqlite3.Connection: + path = _db_path() + path.parent.mkdir(parents=True, exist_ok=True) + conn = sqlite3.connect(path, timeout=10) + conn.execute("PRAGMA journal_mode=WAL") + conn.execute( + """CREATE TABLE IF NOT EXISTS async_delegations ( + delegation_id TEXT PRIMARY KEY, + origin_session TEXT NOT NULL, + origin_ui_session_id TEXT NOT NULL DEFAULT '', + parent_session_id TEXT, + state TEXT NOT NULL, + dispatched_at REAL NOT NULL, + completed_at REAL, + updated_at REAL NOT NULL, + event_json TEXT, + result_json TEXT, + delivery_state TEXT NOT NULL DEFAULT 'pending', + delivery_attempts INTEGER NOT NULL DEFAULT 0, + delivered_at REAL + )""" + ) + return conn + + +def _persist_dispatch(record: Dict[str, Any]) -> None: + now = time.time() + with _DB_LOCK, _connect() as conn: + conn.execute( + """INSERT OR REPLACE INTO async_delegations + (delegation_id, origin_session, origin_ui_session_id, + parent_session_id, state, dispatched_at, updated_at, + delivery_state, delivery_attempts) + VALUES (?, ?, ?, ?, 'running', ?, ?, 'pending', 0)""", + (record["delegation_id"], record.get("session_key", ""), + record.get("origin_ui_session_id", ""), record.get("parent_session_id"), + record["dispatched_at"], now), + ) + cutoff = now - _DURABLE_RETENTION_SECONDS + conn.execute( + """DELETE FROM async_delegations WHERE delegation_id IN ( + SELECT delegation_id FROM async_delegations + WHERE (delivery_state = 'delivered' AND updated_at < ?) + OR (completed_at IS NOT NULL AND completed_at < ?) + ORDER BY updated_at ASC + )""", (cutoff, cutoff), + ) + conn.execute( + """DELETE FROM async_delegations WHERE delegation_id IN ( + SELECT delegation_id FROM async_delegations + WHERE state != 'running' ORDER BY updated_at DESC LIMIT -1 OFFSET ? + )""", (_MAX_RETAINED_COMPLETED,), + ) + + +def _persist_completion(event: Dict[str, Any], result: Dict[str, Any]) -> None: + now = time.time() + with _DB_LOCK, _connect() as conn: + conn.execute( + """UPDATE async_delegations SET state=?, completed_at=?, updated_at=?, + event_json=?, result_json=?, delivery_state='pending' + WHERE delegation_id=?""", + (event.get("status", "completed"), event.get("completed_at", now), now, + json.dumps(event), json.dumps(result), event["delegation_id"]), + ) + + +def _note_delivery_attempt(delegation_id: str) -> None: + with _DB_LOCK, _connect() as conn: + conn.execute( + "UPDATE async_delegations SET delivery_attempts=delivery_attempts+1, updated_at=? WHERE delegation_id=?", + (time.time(), delegation_id), + ) + + +def restore_undelivered_completions(target_queue) -> int: + """Enqueue durable pending completions as fresh turns after process start.""" + with _DB_LOCK, _connect() as conn: + rows = conn.execute( + """SELECT delegation_id, event_json FROM async_delegations + WHERE state != 'running' AND delivery_state='pending' AND event_json IS NOT NULL + ORDER BY completed_at, delegation_id""" + ).fetchall() + for delegation_id, payload in rows: + target_queue.put(json.loads(payload)) + conn.execute( + "UPDATE async_delegations SET delivery_attempts=delivery_attempts+1, updated_at=? WHERE delegation_id=?", + (time.time(), delegation_id), + ) + return len(rows) + + +def mark_completion_delivered(delegation_id: str) -> bool: + """Atomically acknowledge successful injection of a durable completion.""" + now = time.time() + with _DB_LOCK, _connect() as conn: + cur = conn.execute( + """UPDATE async_delegations SET delivery_state='delivered', delivered_at=?, updated_at=? + WHERE delegation_id=? AND delivery_state!='delivered'""", + (now, now, delegation_id), + ) + return cur.rowcount == 1 + + +def get_durable_delegation(delegation_id: str) -> Optional[Dict[str, Any]]: + with _DB_LOCK, _connect() as conn: + row = conn.execute( + """SELECT origin_session, state, dispatched_at, completed_at, + result_json, delivery_state, delivery_attempts + FROM async_delegations WHERE delegation_id=?""", (delegation_id,), + ).fetchone() + if row is None: + return None + return { + "delegation_id": delegation_id, "origin_session": row[0], "state": row[1], + "dispatched_at": row[2], "completed_at": row[3], + "result": json.loads(row[4]) if row[4] else None, + "delivery_state": row[5], "delivery_attempts": row[6], + } def _get_executor(max_workers: int) -> ThreadPoolExecutor: @@ -96,7 +226,7 @@ def _get_executor(max_workers: int) -> ThreadPoolExecutor: def active_count() -> int: """Number of async delegations currently running.""" with _records_lock: - return sum(1 for r in _records.values() if r.get("status") == "running") + return sum(1 for r in _records.values() if r.get("status") in {"running", "finalizing"}) def _new_delegation_id() -> str: @@ -206,6 +336,7 @@ def dispatch_async_delegation( } _records[delegation_id] = record + _persist_dispatch(record) executor = _get_executor(max_async_children) def _worker() -> None: @@ -252,14 +383,20 @@ def _finalize(delegation_id: str, result: Dict[str, Any], status: str) -> None: record = _records.get(delegation_id) if record is None: return - record["status"] = status + # Stay active until durable persistence and queue publication finish; + # otherwise process shutdown can kill this daemon worker in the narrow + # gap after status flips but before SQLite is committed. + record["status"] = "finalizing" record["completed_at"] = time.time() record["interrupt_fn"] = None # drop the closure; child is done - # Snapshot fields needed for the event while holding the lock. event_record = dict(record) - _prune_completed_locked() _push_completion_event(event_record, result, status) + with _records_lock: + record = _records.get(delegation_id) + if record is not None: + record["status"] = status + _prune_completed_locked() def _push_completion_event( @@ -309,8 +446,10 @@ def _push_completion_event( "completed_at": completed_at, "exit_reason": result.get("exit_reason"), } + _persist_completion(evt, result) try: process_registry.completion_queue.put(evt) + _note_delivery_attempt(str(record.get("delegation_id") or "")) except Exception as exc: # pragma: no cover logger.error( "Async delegation %s: failed to enqueue completion event; " @@ -393,6 +532,7 @@ def dispatch_async_delegation_batch( } _records[delegation_id] = record + _persist_dispatch(record) executor = _get_executor(max_async_children) def _worker() -> None: @@ -446,11 +586,10 @@ def _finalize_batch( record = _records.get(delegation_id) if record is None: return - record["status"] = status + record["status"] = "finalizing" record["completed_at"] = time.time() record["interrupt_fn"] = None event_record = dict(record) - _prune_completed_locked() try: from tools.process_registry import process_registry @@ -486,14 +625,22 @@ def _finalize_batch( "dispatched_at": dispatched_at, "completed_at": completed_at, } + _persist_completion(evt, combined) try: process_registry.completion_queue.put(evt) + _note_delivery_attempt(delegation_id) except Exception as exc: # pragma: no cover logger.error( "Async delegation batch %s: failed to enqueue completion event; " "result lost: %s", delegation_id, exc, ) + finally: + with _records_lock: + record = _records.get(delegation_id) + if record is not None: + record["status"] = status + _prune_completed_locked() def list_async_delegations() -> List[Dict[str, Any]]: diff --git a/tools/process_registry.py b/tools/process_registry.py index 0faa80ba6e2..04e3ec0ed2b 100644 --- a/tools/process_registry.py +++ b/tools/process_registry.py @@ -171,6 +171,13 @@ class ProcessRegistry: # gateway drain this after each agent turn to auto-trigger new turns. import queue as _queue_mod self.completion_queue: _queue_mod.Queue = _queue_mod.Queue() + # Rehydrate durable delegation completions only at registry startup. + # Consumers still inject them as fresh turns through this existing rail. + try: + from tools.async_delegation import restore_undelivered_completions + restore_undelivered_completions(self.completion_queue) + except Exception as exc: + logger.warning("Could not restore async delegation completions: %s", exc) # Track sessions whose completion was already consumed by the agent # via wait/log. Drain loops AND gateway/tui watchers skip notifications diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 8e825a22669..0771b488f69 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -8749,6 +8749,9 @@ def _notification_poller_loop( try: _emit("message.start", sid) _run_prompt_submit(rid, sid, session, text) + if evt.get("type") == "async_delegation": + from tools.async_delegation import mark_completion_delivered + mark_completion_delivered(str(evt.get("delegation_id") or "")) except Exception as exc: print( f"[tui_gateway] notification poller dispatch failed: " @@ -8801,6 +8804,9 @@ def _notification_poller_loop( try: _emit("message.start", sid) _run_prompt_submit(rid, sid, session, text) + if evt.get("type") == "async_delegation": + from tools.async_delegation import mark_completion_delivered + mark_completion_delivered(str(evt.get("delegation_id") or "")) except Exception as exc: print( f"[tui_gateway] notification poller dispatch failed: " @@ -9354,6 +9360,9 @@ def _run_prompt_submit(rid, sid: str, session: dict, text: Any) -> None: try: _emit("message.start", sid) _run_prompt_submit(rid, sid, session, synth) + if _evt.get("type") == "async_delegation": + from tools.async_delegation import mark_completion_delivered + mark_completion_delivered(str(_evt.get("delegation_id") or "")) except Exception as _n_exc: print( f"[tui_gateway] completion notification dispatch failed: "