fix(tui_gateway): hold foreign owners through the post-turn drain; keep owners across agent rebuilds

The post-turn safety-net drain took every event off the shared queue but held only the draining
session's owners. Idle session A's completion, sitting in session B's hands while B checked
ownership, left A in neither the queue nor a hold, so a reaper could close A and the requeued notice
was then dropped as unowned. The drain now dequeues through drained_for(), which holds every
dequeued event's owner until the foreign events are requeued and its own are routed.

_rebuild_session_agent (Bot Chat capability refresh) installed a replacement agent without the old
one's raw spawning ids, so a process started before a compression rotation stopped pinning its
session and fell outside the new agent's close kill set. The replacement now inherits them.

tests/tui_gateway/test_detached_process_notify_keepalive.py

@@ -8,10 +8,12 @@ with exit -15, no completion was ever delivered, and the promised follow-up sile
 
 from __future__ import annotations
 
+import contextlib
 import importlib
 import threading
 import time
 from pathlib import Path
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -292,3 +294,101 @@ def test_disabled_watch_does_not_pin_session(server, registry, detached):
     proc._watch_disabled = True
 
     assert server._session_is_reapable(sid, session)
+
+
+# ── the post-turn safety-net drain and agent rebuilds ──────────────────────
+# Session B finishing a turn drains the shared completion queue; an event owned by idle, detached session A passes
+# through B's hands before it is requeued. If A were reapable in that window, the reaper would close A and the
+# requeued notice would be dropped as unowned.
+
+
+def _post_turn_followups(server, namespace_overrides):
+    from tui_gateway import prompt_turn
+    from tui_gateway.method_ctx import rebind
+
+    namespace = dict(vars(server))
+    namespace.update(_drain_queued_prompt=lambda *_: False,
+                     _session_profile_runtime_scope=lambda *_: contextlib.nullcontext(), **namespace_overrides)
+    return rebind(prompt_turn._run_post_turn_followups, namespace)
+
+
+def test_foreign_notice_in_another_sessions_post_turn_drain_keeps_its_owner(server, registry, detached, monkeypatch):
+    sid, session, popped = detached
+    proc = _spawn(registry, session["session_key"], notify_on_complete=True)
+    proc.mark_exited(0)
+    registry._move_to_finished(proc)
+    other_sid, other = "other-sid", dict(session, session_key="other-owner", history_lock=threading.Lock())
+    server._sessions[other_sid] = other
+    monkeypatch.setattr(registry, "restore_completions", lambda: 0)
+    real_owns, observed = server._session_owns_notification_event, []
+
+    def owns(draining_sid, draining_session, event):
+        # B is deciding ownership of A's notice; the queue is empty, so only the in-flight hold covers A.
+        if draining_sid == other_sid and event["owner_task_id"] == session["session_key"]:
+            observed.append((registry.completion_queue.empty(), server._session_is_reapable(sid, session)))
+            server._schedule_ws_orphan_reap(sid)
+            server._pending_ws_reaps[sid].callback()
+        return real_owns(draining_sid, draining_session, event)
+
+    routed = []
+    _post_turn_followups(server, {"_session_owns_notification_event": owns,
+                                  "_notif_handle_ready": lambda *a, **k: routed.append(a[2])})(
+        None, other_sid, other, {}, None)
+
+    assert observed == [(True, False)], "A's notice was in B's hands: A must stay protected"
+    assert sid in server._sessions and not popped
+    assert [e["owner_task_id"] for e in registry.completion_queue.queue] == [session["session_key"]]
+    assert routed == [[]], "B routes none of A's events"
+
+
+def test_post_turn_drain_routes_its_own_notice_and_releases_every_hold(server, registry, detached, monkeypatch):
+    from tools import process_registry_awaiting as awaiting
+
+    sid, session, _ = detached
+    for owner in (session["session_key"], "someone-else"):
+        proc = _spawn(registry, owner, notify_on_complete=True)
+        proc.mark_exited(0)
+        registry._move_to_finished(proc)
+    monkeypatch.setattr(registry, "restore_completions", lambda: 0)
+    routed = []
+    _post_turn_followups(server, {"_notif_handle_ready": lambda *a, **k: routed.append(
+        [(e["owner_task_id"], server._session_is_reapable(sid, session)) for e in a[2]])})(None, sid, session, {}, None)
+
+    assert routed == [[(session["session_key"], False)]], "own notice routed while its session is still held"
+    assert [e["owner_task_id"] for e in registry.completion_queue.queue] == ["someone-else"]
+    assert not awaiting._in_flight, "every hold released once routing returned"
+    assert server._session_is_reapable(sid, session)
+
+
+def test_post_turn_drain_releases_holds_when_routing_raises(server, registry, detached, monkeypatch):
+    from tools import process_registry_awaiting as awaiting
+
+    sid, session, _ = detached
+    proc = _spawn(registry, session["session_key"], notify_on_complete=True)
+    proc.mark_exited(0)
+    registry._move_to_finished(proc)
+    monkeypatch.setattr(registry, "restore_completions", lambda: 0)
+
+    def fail(*_a, **_k):
+        raise RuntimeError("routing failed")
+
+    _post_turn_followups(server, {"_notif_handle_ready": fail})(None, sid, session, {}, None)
+    assert not awaiting._in_flight
+
+
+def test_agent_rebuild_keeps_pre_compression_process_owners(server, registry, detached, monkeypatch):
+    """A Bot Chat capability refresh replaces the agent; a process the old agent started under its pre-compression
+    id must still pin the session (and still be killed by the new agent's close)."""
+    sid, session, _ = detached
+    session["agent"] = SimpleNamespace(_process_owner_task_ids={"pre-compression-id", session["session_key"]},
+                                       _session_db=None, _owns_session_db=False)
+    _spawn(registry, "pre-compression-id", notify_on_complete=True)
+    assert not server._session_is_reapable(sid, session)
+    monkeypatch.setattr(server, "_make_agent", lambda *a, **k: SimpleNamespace(
+        _process_owner_task_ids={session["session_key"]}, _session_db=None))
+    monkeypatch.setattr(server, "_release_build_profile_scopes", lambda *_: None)
+
+    server._rebuild_session_agent(sid, session)
+
+    assert "pre-compression-id" in session["agent"]._process_owner_task_ids
+    assert not server._session_is_reapable(sid, session)

tools/process_registry.py

@@ -2546,14 +2546,11 @@ class ProcessRegistry(ProcessTerminationMixin, ProcessCheckpointMixin, ProcessOu
     # ----- Cleanup / Pruning -----
 
     def _prune_if_needed(self):
-        """Drop expired finished sessions, then the oldest survivor while over
-        MAX_PROCESSES. Must hold _lock. Kept either way: a session whose scope stop is still
-        pending (kill_all's only handle on a scope that may hold live descendants), and one
-        whose completion is not yet published (its owner's reaper protection reads it)."""
+        """Drop expired, then oldest-over-cap finished sessions (hold _lock); keep pending scope stops and unpublished notices."""
+        from tools.process_registry_awaiting import retained_for_delivery
         now = time.time()
         prunable = {sid: s for sid, s in self._finished.items()
-                    if not getattr(s, "_scope_stop_pending", False)
-                    and not (getattr(s, "_finish_claimed", False) and not s._completion_event.is_set())}
+                    if not getattr(s, "_scope_stop_pending", False) and not retained_for_delivery(s)}
         expired = [sid for sid, s in prunable.items() if (now - s.started_at) > FINISHED_TTL_SECONDS]
         over_cap = len(self._running) + len(self._finished) - len(expired) >= MAX_PROCESSES
         if over_cap and (survivors := [sid for sid in prunable if sid not in expired]):

tools/process_registry_awaiting.py

@@ -66,14 +66,26 @@ def release(events) -> None:
 
 
 @contextlib.contextmanager
-def holding(owner_task_ids):
-    """Hold ``owner_task_ids`` for a drain that dequeues through another API (``drain_notifications``)."""
-    owners = [o for o in owner_task_ids if o]
-    _adjust(owners, +1)
+def drained_for(registry, owns):
+    """Dequeue everything queued for a caller that consumes only the events ``owns`` accepts (the TUI's post-turn
+    safety net). Every dequeued event's owner is held — foreign ones too — until the rejected events are back on
+    the queue and the accepted ones (yielded) have been routed, so no owner is ever between queue and hold."""
+    registry.restore_completions()
+    ready = take_ready(registry.completion_queue, 0)
     try:
-        yield
+        mine = [evt for evt in ready if owns(evt)]
+        for evt in ready:
+            if not any(evt is m for m in mine):
+                registry.completion_queue.put(evt)
+        yield mine
     finally:
-        _adjust(owners, -1)
+        release(ready)
+
+
+def retained_for_delivery(s) -> bool:
+    """A finished entry whose completion is not yet published: pruning it would erase the only record that a notice
+    is coming, so the registry keeps it until publication."""
+    return bool(getattr(s, "_finish_claimed", False)) and not s._completion_event.is_set()
 
 
 def _notifies(s) -> bool:

tui_gateway/agent_callbacks.py

@@ -574,6 +574,10 @@ def _rebuild_session_agent(sid: str, session: dict, **kwargs):
         # already closed the agent it saw, so one installed now is never closed (#49852).
         closed_midbuild = bool(session.get("_closing"))
         if not closed_midbuild:
+            # The replacement inherits the raw spawning ids: processes started before a compression rotation keep
+            # their old owner, and the reapers' notify keepalive (and close's kill set) match on it.
+            if inherited := set(getattr(old_agent, "_process_owner_task_ids", ())):
+                agent._process_owner_task_ids = {*inherited, *getattr(agent, "_process_owner_task_ids", ())}
             session.update(agent=agent, config_model_seen=config_model_seen)
             owned = opened or bool(getattr(old_agent, "_owns_session_db", False))
             if owned and _transfer_db_to_agent(agent, session_db):

tui_gateway/prompt_turn.py

@@ -463,18 +463,17 @@ def _run_post_turn_followups(
         from tools.process_registry import process_registry
         # _release_turn_scopes / _post_turn_housekeeping have released the worker's runtime scopes. Queue ownership,
         # notification policy and nested dispatch still belong to this session.
-        from tools.process_registry_awaiting import holding
-        # This drain dequeues this session's own events while it is idle; hold its owners until routing finishes so
-        # the reapers cannot close it between dequeue and turn admission.
-        with _session_profile_runtime_scope(session), holding(_session_process_owners(session)):
-            drained = process_registry.drain_notifications(
-                session_key=session.get("session_key", ""),
-                owns_event=lambda e: _session_owns_notification_event(sid, session, e),
-                skip_poll_observed=False)
-            from tools.process_registry_notifications import format_process_notification
+        from tools.process_registry_awaiting import drained_for
+        from tools.process_registry_notifications import format_process_notification, should_surface_notification
+        key = session.get("session_key", "")
+        owns = lambda e: process_registry._owns_event(
+            e, key, lambda evt: _session_owns_notification_event(sid, session, evt), e.get("type") == "async_delegation")
+        # Every dequeued event's owner — this session's and any other session's — stays held until the foreign events
+        # are requeued and ours are routed, so the reapers never see an owner whose notice is in nobody's hands.
+        with _session_profile_runtime_scope(session), drained_for(process_registry, owns) as mine:
             deferred = []
             _notif_handle_ready(
-                sid, session, [event for event, _text in drained],
+                sid, session, [evt for evt in mine if should_surface_notification(evt)],
                 session.setdefault("_notification_emitted", set()), process_registry,
                 format_process_notification, deferred, owned=True)
             for event in deferred: