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: