From 4d32503f86d6a1eb6dffd2147a9ba77595bc42b1 Mon Sep 17 00:00:00 2001 From: Ryan Tucker Date: Sat, 15 Aug 2026 18:56:38 -0500 Subject: [PATCH] fix(tui_gateway): skip a dead peer's transport when draining its queued prompt MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _drain_queued_prompt attaches the transport pinned to the queued envelope, but the client that queued the prompt may have disconnected while the prompt sat in the queue. Attaching it then pins a dead peer into the session's fan-out, where it costs one failed write before the fan-out prunes it. That is not a regression — the line this replaced rebound the whole slot to that same dead transport, which cost the session for the entire drained turn rather than one write — but the guard is one condition, so it goes in. _transport_is_dead is the predicate already used by the reaper and the orphan check: the parked drop sentinel, or a transport carrying _closed. Drain semantics are otherwise unchanged. The prompt still runs, and the clients already attached keep their stream; only the dead peer's attachment is skipped. tests/tui_gateway/test_multi_client_fanout.py pins that the drained prompt still dispatches while the disconnected queuer stays out of the slot, alongside the existing case where a live queuer is attached and the single-client control where the queuer takes an unoccupied slot. Raised by the automated review on #86784. --- tests/tui_gateway/test_multi_client_fanout.py | 32 +++++++++++++++++++ tui_gateway/session_auto_continue.py | 12 ++++--- 2 files changed, 39 insertions(+), 5 deletions(-) diff --git a/tests/tui_gateway/test_multi_client_fanout.py b/tests/tui_gateway/test_multi_client_fanout.py index d1a38656bb..90a60e1e7c 100644 --- a/tests/tui_gateway/test_multi_client_fanout.py +++ b/tests/tui_gateway/test_multi_client_fanout.py @@ -656,6 +656,38 @@ def test_queued_prompt_drain_keeps_both_clients_attached(monkeypatch): assert b.types() == ["message.delta"] +def test_queued_prompt_drain_skips_a_queuer_that_disconnected(monkeypatch): + """B goes away while its prompt waits: the prompt runs, the dead pin does not. + + Attaching a transport whose client already left would pin a dead peer into + the slot until the first failed write prunes it. A keeps its stream and + stays the only attached client. + """ + dispatched = [] + monkeypatch.setattr( + server, + "_run_prompt_submit", + lambda rid, sid, _session, text, **kw: dispatched.append((rid, text)), + ) + + a, b = _FakeClient("a"), _FakeClient("b") + session = _session(transport=a) + server._enqueue_prompt(session, "from B", b) + b._closed = True # what _transport_is_dead reads: B's socket went away + server._sessions["sid"] = session + try: + assert server._drain_queued_prompt("drain", "sid", session) is True + server._emit("message.delta", "sid", {"text": "drained answer"}) + finally: + server._sessions.pop("sid", None) + + assert dispatched == [("drain", "from B")] # drain semantics unchanged + assert session["transport"] is a + assert server._session_transport_contains(session, b) is False + assert a.types() == ["message.delta"] + assert b.types() == [] + + def test_queued_prompt_drain_still_rebinds_a_single_client_session(monkeypatch): """(j) Control: with one client the drain lands on the queuer's transport.""" monkeypatch.setattr(server, "_run_prompt_submit", lambda *a, **k: None) diff --git a/tui_gateway/session_auto_continue.py b/tui_gateway/session_auto_continue.py index 858a1f736f..e1c7e79110 100644 --- a/tui_gateway/session_auto_continue.py +++ b/tui_gateway/session_auto_continue.py @@ -287,11 +287,13 @@ def _drain_queued_prompt(rid, sid: str, session: dict) -> bool: queue_generation = int(session.get("_queued_prompt_generation", 0)) _ac_set_queue(session, session.get("queued_prompts") or []) session["running"] = True - if queued.get("transport") is not None: - # The queuer's transport is pinned so the drained turn reaches the client that sent it — but - # ATTACHED, not rebound: a mid-turn prompt from a second client used to silence the first for the - # whole drained turn. - _attach_session_transport(session, queued["transport"]) + queued_transport = queued.get("transport") + # The queuer's transport is pinned so the drained turn reaches the client that sent it — but + # ATTACHED, not rebound: a mid-turn prompt from a second client used to silence the first for the + # whole drained turn. A peer that disconnected while its prompt sat in the queue is skipped: the + # prompt still runs, only the dead pin is dropped. + if queued_transport is not None and not _transport_is_dead(queued_transport): + _attach_session_transport(session, queued_transport) use_compute_host = _session_uses_compute_host(session) with session["history_lock"]: if int(session.get("_queued_prompt_generation", 0)) != queue_generation: