From 47d853fdf26415fc36de10267512dca6fd2c6f91 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Tue, 14 Jul 2026 13:13:38 -0700 Subject: [PATCH] fix(delegation): fail closed on restored completions + stamp CLI dispatch identity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three-layer companion to the salvaged CLI drain-ownership fix (#64240): 1. restore_undelivered_completions stamps restored=True (in-memory only) on every durable completion re-enqueued at process start. 2. drain_notifications' legacy unfiltered branch re-queues restored events instead of consuming them — a fresh process can no longer adopt a dead session's delegation results (#64484). Same-process keyless events keep the legacy behavior. 3. delegate_tool's async dispatch now falls back to the parent agent's durable session_id when the approval-context key resolves empty (the CLI case), so the CLI's new positive-ownership drain can actually claim its own completions instead of failing closed on ''. --- .../test_restored_delegation_ownership.py | 149 ++++++++++++++++++ tools/async_delegation.py | 17 +- tools/delegate_tool.py | 14 ++ tools/process_registry.py | 13 ++ 4 files changed, 191 insertions(+), 2 deletions(-) create mode 100644 tests/tools/test_restored_delegation_ownership.py diff --git a/tests/tools/test_restored_delegation_ownership.py b/tests/tools/test_restored_delegation_ownership.py new file mode 100644 index 000000000000..1002decf3cd8 --- /dev/null +++ b/tests/tools/test_restored_delegation_ownership.py @@ -0,0 +1,149 @@ +"""Regression coverage for #64484 — durable-restored delegation completions +must never be adopted by a session that cannot positively prove ownership. + +Layers under test: +1. ``restore_undelivered_completions`` stamps every restored event with + ``restored=True`` (in-memory only). +2. ``ProcessRegistry.drain_notifications`` with NO filter (legacy + consume-everything CLI path) re-queues restored events instead of + consuming them. +3. Same-process (non-restored) keyless events keep the legacy behavior. +4. An owner with a matching session_key still receives its restored event. +""" + +import json +import queue + +from tools.process_registry import ProcessRegistry + + +def _make_registry(): + reg = ProcessRegistry.__new__(ProcessRegistry) + import threading + + reg._running = {} + reg._finished = {} + reg._lock = threading.Lock() + reg.completion_queue = queue.Queue() + reg._completion_consumed = set() + reg._poll_observed = set() + return reg + + +def _delegation_event(session_key="", restored=False, delegation_id="d1"): + evt = { + "type": "async_delegation", + "delegation_id": delegation_id, + "session_key": session_key, + "origin_ui_session_id": "", + "goal": "secret goal", + "status": "success", + "summary": "SECRET RESULT", + "api_calls": 3, + "duration_seconds": 1.5, + "dispatched_at": 1.0, + "completed_at": 2.0, + } + if restored: + evt["restored"] = True + return evt + + +def test_restore_stamps_restored_flag(tmp_path, monkeypatch): + """Every durable completion re-enqueued at startup carries restored=True.""" + import tools.async_delegation as ad + + monkeypatch.setattr(ad, "_db_path", lambda: tmp_path / "async_delegations.db") + record = { + "delegation_id": "d-old", + "goal": "old goal", + "context": None, + "toolsets": None, + "role": "leaf", + "model": "m", + "session_key": "OLD_SESSION_A", + "origin_ui_session_id": "", + "parent_session_id": "OLD_SESSION_A", + "status": "running", + "dispatched_at": 1.0, + "completed_at": None, + "interrupt_fn": None, + } + ad._persist_dispatch(record) + evt = _delegation_event(session_key="OLD_SESSION_A", delegation_id="d-old") + ad._persist_completion(evt, {"summary": "SECRET RESULT"}) + + q = queue.Queue() + restored = ad.restore_undelivered_completions(q) + assert restored == 1 + got = q.get_nowait() + assert got["restored"] is True + assert got["session_key"] == "OLD_SESSION_A" + + # The stamp is in-memory only — the durable payload is unchanged. + with ad._connect() as conn: + row = conn.execute( + "SELECT event_json FROM async_delegations WHERE delegation_id='d-old'" + ).fetchone() + assert "restored" not in json.loads(row[0]) + + +def test_unfiltered_drain_never_consumes_restored_events(): + """The legacy consume-everything branch must fail closed on restored events.""" + reg = _make_registry() + reg.completion_queue.put(_delegation_event(session_key="DEAD_SESSION", restored=True)) + + results = reg.drain_notifications() # no filter — legacy CLI post-turn shape + + assert results == [] + # Still queued for its real owner. + assert reg.completion_queue.qsize() == 1 + assert reg.completion_queue.get_nowait()["session_key"] == "DEAD_SESSION" + + +def test_unfiltered_drain_keeps_legacy_behavior_for_same_process_events(): + """Non-restored keyless events (created by this process) are still consumed.""" + reg = _make_registry() + reg.completion_queue.put(_delegation_event(session_key="")) + + results = reg.drain_notifications() + + assert len(results) == 1 + assert results[0][0]["delegation_id"] == "d1" + assert reg.completion_queue.empty() + + +def test_owner_session_key_drain_consumes_restored_event(): + """The owning session (key match) still receives its restored completion.""" + reg = _make_registry() + reg.completion_queue.put(_delegation_event(session_key="OWNER", restored=True)) + + results = reg.drain_notifications(session_key="OWNER") + + assert len(results) == 1 + assert results[0][0]["session_key"] == "OWNER" + assert reg.completion_queue.empty() + + +def test_foreign_session_key_drain_requeues_restored_event(): + """A different session's keyed drain must not claim the restored event.""" + reg = _make_registry() + reg.completion_queue.put(_delegation_event(session_key="OWNER", restored=True)) + + results = reg.drain_notifications(session_key="SOMEONE_ELSE") + + assert results == [] + assert reg.completion_queue.qsize() == 1 + + +def test_owns_event_callback_beats_restored_flag(): + """A positive-proof ownership callback consumes restored events it owns.""" + reg = _make_registry() + reg.completion_queue.put(_delegation_event(session_key="OWNER", restored=True)) + + results = reg.drain_notifications( + owns_event=lambda e: e.get("session_key") == "OWNER" + ) + + assert len(results) == 1 + assert reg.completion_queue.empty() diff --git a/tools/async_delegation.py b/tools/async_delegation.py index d2f2dc23b91b..c743decf644e 100644 --- a/tools/async_delegation.py +++ b/tools/async_delegation.py @@ -264,7 +264,17 @@ def recover_abandoned_delegations() -> int: def restore_undelivered_completions(target_queue) -> int: - """Enqueue durable pending completions as fresh turns after process start.""" + """Enqueue durable pending completions as fresh turns after process start. + + Every restored event is stamped ``restored=True`` (in-memory only — the + stamp is added after the durable payload is deserialized and is never + persisted). Restored events originate from a *previous* process, so no + consumer in THIS process implicitly owns them: drain paths that run + without an ownership filter (the legacy single-session behavior) must + leave them queued for a consumer that can positively prove ownership, + otherwise a brand-new session adopts a dead session's delegation + results seconds after boot (#64484). + """ recover_abandoned_delegations() with _DB_LOCK, _connect() as conn: rows = conn.execute( @@ -273,7 +283,10 @@ def restore_undelivered_completions(target_queue) -> int: ORDER BY completed_at, delegation_id""" ).fetchall() for _delegation_id, payload in rows: - target_queue.put(json.loads(payload)) + evt = json.loads(payload) + if isinstance(evt, dict): + evt["restored"] = True + target_queue.put(evt) return len(rows) diff --git a/tools/delegate_tool.py b/tools/delegate_tool.py index 24c2cd3325d5..c8f6fa5293a9 100644 --- a/tools/delegate_tool.py +++ b/tools/delegate_tool.py @@ -2844,6 +2844,20 @@ def delegate_task( _session_key = _agent_session_id except Exception: _origin_ui_session_id = "" + if not _session_key: + # CLI (single-process) path: the approval contextvar is only bound + # during gateway/TUI turns and HERMES_SESSION_KEY is not in the CLI + # environment, so the key resolves empty here. Since #64240 the CLI + # drains completions through a positive-ownership filter keyed on + # the durable AIAgent.session_id — an empty session_key would fail + # closed and the CLI could never claim its own completions, while + # a restored foreign event with an empty key could leak into any + # unfiltered consumer (#64484). Stamp the parent's durable session + # id instead; compression rotations are handled on the drain side + # via resolve_resume_session_id lineage resolution. + _agent_session_id = str(getattr(parent_agent, "session_id", "") or "") + if _agent_session_id: + _session_key = _agent_session_id _parent_session_id = getattr(parent_agent, "session_id", None) _child_agents = [c for (_, _, c) in children] diff --git a/tools/process_registry.py b/tools/process_registry.py index f7e6e8471b56..62dd8494e4ae 100644 --- a/tools/process_registry.py +++ b/tools/process_registry.py @@ -1211,6 +1211,19 @@ class ProcessRegistry: if evt_session_key != session_key: requeue.append(evt) continue + elif evt.get("restored"): + # Legacy unfiltered drain (no ownership callback, no + # session key). That behavior was safe when the in-memory + # queue could only hold events created by this very + # process — but durable restore (#63494) re-enqueues + # completions from PREVIOUS processes at startup, so an + # unfiltered consumer here would adopt a dead, unrelated + # session's conversation payload (#64484). Fail closed: + # leave restored events queued (still 'pending' on disk) + # for a consumer that can positively prove ownership, + # e.g. the owning session's --resume. + requeue.append(evt) + continue text = format_process_notification(evt) if text: results.append((evt, text))