fix(tui): route bg process notifications to owning session, drop orphaned events

Two complementary fixes for cross-session background-process notification
leakage in the TUI/Desktop multi-session path (#42674, #35652).

1. Poller orphan guard: after _notification_event_belongs_elsewhere
   returns False, check whether the event has a non-empty session_key
   that differs from the current session.  If so the owner session is
   gone — drop the event instead of hijacking it into an unrelated
   session transcript.

2. Post-turn drain filter: the existing drain_notifications() pops every
   event from the global queue regardless of ownership.  Added
   _drain_owned_notifications() which applies the same ownership routing
   used by the poller (consume own, requeue foreign-live, drop orphan),
   and wired it into the post-turn safety drain.

Complementary to PR #42731 which addresses a separate code path in the
same bug class.  Together they close #42674.
This commit is contained in:
Yingliang Zhang 2026-06-29 17:46:29 +08:00 committed by Teknium
parent 2ea39daeb1
commit 81fc24862c
2 changed files with 257 additions and 0 deletions

View file

@ -2284,6 +2284,153 @@ def test_prompt_submit_rejects_negative_truncate_ordinal(monkeypatch):
server._sessions.pop("trunc-sid", None)
def test_notification_poller_drops_orphaned_events(monkeypatch):
"""Completion events whose owner is gone are dropped, not hijacked."""
import queue as _queue_mod
from tools.process_registry import process_registry
emitted = []
sess = _session(session_key="session-a")
server._sessions["sid_a"] = sess
monkeypatch.setattr(server, "_emit", lambda *a, **kw: emitted.append(a))
monkeypatch.setattr(server, "make_stream_renderer", lambda cols: None)
monkeypatch.setattr(server, "render_message", lambda raw, cols: None)
# Isolate the completion queue so no other test/poller can interfere.
isolated_queue: _queue_mod.Queue = _queue_mod.Queue()
monkeypatch.setattr(process_registry, "completion_queue", isolated_queue)
process_registry._completion_consumed.discard("proc_ghost")
# An event owned by session-b — but session-b is NOT in _sessions (gone).
isolated_queue.put({
"type": "completion",
"session_id": "proc_ghost",
"session_key": "session-b",
"command": "echo from ghost",
"exit_code": 0,
"output": "ghost output",
})
stop = threading.Event()
stop.set()
try:
server._notification_poller_loop(stop, "sid_a", sess)
# No status.update emitted — the orphaned event was dropped.
status_calls = [a for a in emitted if a[0] == "status.update"]
assert len(status_calls) == 0, (
f"orphaned event should be dropped, got {len(status_calls)} status.update calls"
)
finally:
server._sessions.pop("sid_a", None)
while not process_registry.completion_queue.empty():
process_registry.completion_queue.get_nowait()
def test_notification_poller_delivers_owned_events(monkeypatch):
"""Events owned by *this* session are delivered normally (regression guard)."""
import queue as _queue_mod
from tools.process_registry import process_registry
turns = []
emitted = []
class _Agent:
def run_conversation(self, prompt, conversation_history=None, stream_callback=None):
turns.append(prompt)
return {"final_response": "ok", "messages": [{"role": "assistant", "content": "ok"}]}
class _ImmediateThread:
def __init__(self, target=None, daemon=None):
self._target = target
def start(self):
self._target()
sess = _session(session_key="session-a", agent=_Agent())
server._sessions["sid_a"] = sess
monkeypatch.setattr(server.threading, "Thread", _ImmediateThread)
monkeypatch.setattr(server, "_emit", lambda *a, **kw: emitted.append(a))
monkeypatch.setattr(server, "make_stream_renderer", lambda cols: None)
monkeypatch.setattr(server, "render_message", lambda raw, cols: None)
isolated_queue: _queue_mod.Queue = _queue_mod.Queue()
monkeypatch.setattr(process_registry, "completion_queue", isolated_queue)
process_registry._completion_consumed.discard("proc_mine")
# An event owned by THIS session.
isolated_queue.put({
"type": "completion",
"session_id": "proc_mine",
"session_key": "session-a",
"command": "echo mine",
"exit_code": 0,
"output": "mine",
})
stop = threading.Event()
stop.set()
try:
server._notification_poller_loop(stop, "sid_a", sess)
status_calls = [a for a in emitted if a[0] == "status.update"]
assert len(status_calls) == 1
assert status_calls[0][2]["kind"] == "process"
assert len(turns) == 1
assert "proc_mine" in turns[0]
finally:
server._sessions.pop("sid_a", None)
while not process_registry.completion_queue.empty():
process_registry.completion_queue.get_nowait()
def test_drain_owned_notifications_routes_by_session_key(monkeypatch):
"""_drain_owned_notifications filters events by session ownership."""
import queue as _queue_mod
sess_a = _session(session_key="session-a")
sess_b = _session(session_key="session-b")
server._sessions["sid_a"] = sess_a
server._sessions["sid_b"] = sess_b
q: _queue_mod.Queue = _queue_mod.Queue()
# Mix of events: ours, foreign (live), orphan, global
q.put({"type": "completion", "session_id": "1", "session_key": "session-a",
"command": "a", "exit_code": 0, "output": "a"})
q.put({"type": "completion", "session_id": "2", "session_key": "session-b",
"command": "b", "exit_code": 0, "output": "b"})
q.put({"type": "completion", "session_id": "3", "session_key": "ghost",
"command": "c", "exit_code": 0, "output": "c"})
q.put({"type": "completion", "session_id": "4",
"command": "d", "exit_code": 0, "output": "d"})
results = server._drain_owned_notifications(q, sess_a)
# Only our event (session-a) + global (no session_key) should be returned.
assert len(results) == 2, f"expected 2 events (ours + global), got {len(results)}"
owned_sids = {evt.get("session_id") for evt, _ in results}
assert "1" in owned_sids # ours
assert "4" in owned_sids # global
# session-b's event should be re-queued for its poller.
# Orphan should be dropped (not in results, not re-queued).
assert q.qsize() == 1, f"1 event (session-b) should be re-queued, got {q.qsize()}"
# Now drain as session-b — should get its re-queued event.
results_b = server._drain_owned_notifications(q, sess_b)
assert len(results_b) == 1, f"session-b should pick up its re-queued event, got {len(results_b)}"
assert results_b[0][0].get("session_id") == "2"
assert q.empty(), f"queue should be empty after both drains, got {q.qsize()}"
# Cleanup
server._sessions.pop("sid_a", None)
server._sessions.pop("sid_b", None)
def test_session_create_does_not_persist_empty_row(monkeypatch):
"""session.create must NOT eagerly write a DB row.

View file

@ -8716,6 +8716,83 @@ def _notification_event_dedup_key(evt: dict) -> tuple:
return (evt_sid, evt_type)
def _drain_owned_notifications(
completion_queue: "queue.Queue",
session: dict,
) -> "list[tuple[dict, str]]":
"""Pop all pending notification events, keeping only those owned by this session.
``process_registry.drain_notifications()`` pops every event from the global
queue regardless of ``session_key`` a turn finishing in session B would
consume an event started by session A. This wrapper applies the same
ownership check used by ``_notification_poller_loop``, so only events that
belong to *this* session (plus global/system events with no ``session_key``)
are dispatched here. Events owned by another live session are re-queued;
orphaned events (owner gone) are dropped (#42674, #35652).
Returns the same ``[(raw_event, formatted_text)]`` shape as
``drain_notifications()``, filtered to this session's events only.
"""
from tools.process_registry import format_process_notification, process_registry
_my_key = str(session.get("session_key") or "")
# Snapshot live session keys (excluding our own) for ownership routing.
# Must be computed fresh so a just-closed session isn't treated as live.
try:
with _sessions_lock:
snapshot = list(_sessions.values())
except Exception:
snapshot = []
live_owner_keys = {
str(s.get("session_key") or "")
for s in snapshot
if s is not session and str(s.get("session_key") or "")
}
owned: list[tuple[dict, str]] = []
requeue: list[dict] = []
while not completion_queue.empty():
try:
evt = completion_queue.get_nowait()
except Exception:
break
_evt_key = str(evt.get("session_key") or "")
if not _evt_key:
# Global/system event with no owner — handle here.
pass
elif _evt_key == _my_key:
# Owned by this session — handle here.
pass
elif _evt_key in live_owner_keys:
# Owned by another live session — requeue.
requeue.append(evt)
continue
else:
# Orphaned event (owner gone) — drop silently.
logger.debug(
"Dropping orphaned background notification for "
"session_key=%s (owner gone, current=%s)",
_evt_key, _my_key,
)
continue
_evt_sid = evt.get("session_id", "")
if evt.get("type") == "completion" and process_registry.is_completion_consumed(_evt_sid):
continue
text = format_process_notification(evt)
if text:
owned.append((evt, text))
# Re-queue events for live sessions so their pollers can handle them.
for evt in requeue:
completion_queue.put(evt)
return owned
def _notification_poller_loop(
stop_event: threading.Event, sid: str, session: dict
) -> None:
@ -8732,6 +8809,8 @@ def _notification_poller_loop(
"""
from tools.process_registry import process_registry, format_process_notification
_my_key = str(session.get("session_key") or "")
_emitted = set() # dedup re-queued events so same completion isn't emitted 50 times while session is busy
while not stop_event.is_set() and not session.get("_finalized"):
try:
@ -8775,6 +8854,20 @@ def _notification_poller_loop(
)
continue
# Orphan guard: _notification_event_belongs_elsewhere returns False for
# events whose owner is no longer live (session closed / /new'd away).
# Previously these orphans were consumed by whichever poller dequeued
# them — injecting an unrelated background-process notification into the
# wrong session's transcript (#42674, #35652). Drop them instead.
_evt_key = str(evt.get("session_key") or "")
if _evt_key and _evt_key != _my_key:
logger.debug(
"Dropping orphaned background notification for "
"session_key=%s (owner gone, current=%s)",
_evt_key, _my_key,
)
continue
_evt_sid = evt.get("session_id", "")
if evt.get("type") == "completion" and process_registry.is_completion_consumed(_evt_sid):
continue
@ -8830,6 +8923,7 @@ def _notification_poller_loop(
# Drain any remaining events after stop signal (process all pending
# before exiting so nothing is lost on shutdown). Events owned by other
# live sessions are set aside and re-queued so their poller still sees them.
# Orphaned events (owner gone) are dropped — same guard as the main loop.
deferred: list = []
while not process_registry.completion_queue.empty():
try:
@ -8848,6 +8942,15 @@ def _notification_poller_loop(
):
deferred.append(evt)
continue
# Orphan guard: same check as the main loop above.
_evt_key = str(evt.get("session_key") or "")
if _evt_key and _evt_key != _my_key:
logger.debug(
"Dropping orphaned background notification for "
"session_key=%s (owner gone, current=%s)",
_evt_key, _my_key,
)
continue
_evt_sid = evt.get("session_id", "")
if evt.get("type") == "completion" and process_registry.is_completion_consumed(_evt_sid):
continue
@ -9411,6 +9514,13 @@ def _run_prompt_submit(rid, sid: str, session: dict, text: Any) -> None:
# Drain completion notifications that arrived during this turn.
# The background poller handles between-turn delivery; this is
# the safety net for events that arrived mid-turn.
#
# Ownership filter (#42674, #35652): drain_notifications() pops
# every event from the global queue regardless of session_key.
# A turn finishing in session B must not consume an event that
# belongs to session A. Owned events are re-queued so the
# owning session's poller can handle them; orphaned events
# (owner gone) are dropped silently.
try:
from tools.process_registry import process_registry