diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index ca030284bf..8b2ece07b8 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -5300,7 +5300,7 @@ def test_close_transport_rebinds_session_to_remaining_viewer(monkeypatch): The rebind #83716 added is gone; multi-client fan-out subsumes it. Both windows are attached to the slot at once, so the pop-out is a fan-out peer rather than a viewer waiting to be promoted, and closing it detaches that - peer and collapses the slot back onto the main window. This pins the same + peer while retaining the surviving ordered mailbox. This pins the same guarantee through the mechanism that replaced the rebind: the session is not parked, not reaped, not handed to the orphan reaper, and the surviving window keeps receiving frames. @@ -5311,9 +5311,11 @@ def test_close_transport_rebinds_session_to_remaining_viewer(monkeypatch): class _LiveTransport: def __init__(self): self.frames = [] + self.received = threading.Event() def write(self, obj=None, *a, **k): self.frames.append(obj) + self.received.set() return True main = _LiveTransport() @@ -5332,14 +5334,14 @@ def test_close_transport_rebinds_session_to_remaining_viewer(monkeypatch): reaped, detached = server._close_sessions_for_transport(popout) assert reaped == 0 and detached == 0 - # One peer left, so the fan-out collapses back to the bare transport — - # the slot is indistinguishable from a session that never fanned out. - assert session["transport"] is main + assert server._session_transport_contains(session, main) + assert not server._session_transport_contains(session, popout) assert "multi-sid" not in reap_calls assert server._ws_session_is_orphaned(session) is False # And it is still a working stream, not just a surviving reference. server._emit("message.delta", "multi-sid", {"text": "still here"}) + assert main.received.wait(timeout=5) assert [(f.get("params") or {}).get("type") for f in main.frames] == [ "message.delta" ] diff --git a/tests/tui_gateway/test_multi_client_fanout.py b/tests/tui_gateway/test_multi_client_fanout.py index 90a60e1e7c..0942468973 100644 --- a/tests/tui_gateway/test_multi_client_fanout.py +++ b/tests/tui_gateway/test_multi_client_fanout.py @@ -1,1036 +1,241 @@ -"""A session streams to EVERY attached client, not just the newest one. - -Before this, ``session["transport"]`` held exactly one client and every -``prompt.submit`` / ``session.resume`` / ``session.activate`` / queued-prompt -drain rebound it — so a second client either saw nothing, stole the stream from -the first, or silenced the turn the first was reading. The slot now holds a -``FanoutTransport`` as soon as a second client attaches, and the disconnect path -detaches instead of parking whenever another client is still there. - -Single-client behaviour is the control condition throughout: with one client -attached the slot holds the bare transport and every path behaves exactly as it -did before fan-out existed. -""" - +"""Shared-session routing and backpressure, exercised through real OS pipes.""" +import asyncio +import json +import os +import queue +import socket import threading -import types +from contextlib import ExitStack, suppress + +import pytest from tui_gateway import server -from tui_gateway.transport import FanoutTransport - - -class _FakeClient: - """A connected client: records the frames it receives. - - ``ok`` False models a peer that has gone away (``write`` returns False, the - gateway's peer-gone signal); ``boom`` models a wedged peer whose write - raises. Both must be pruned without disturbing the healthy clients. - """ - - def __init__(self, name: str, *, ok: bool = True, boom: bool = False) -> None: - self.name = name - self.frames: list[dict] = [] - self.closed = False - self._ok = ok - self._boom = boom - - def write(self, obj: dict) -> bool: - if self._boom: - raise RuntimeError(f"{self.name} is wedged") - self.frames.append(obj) - return self._ok - - def close(self) -> None: - self.closed = True - - def types(self) -> list[str]: - return [(f.get("params") or {}).get("type") for f in self.frames] - - -def _session(**extra) -> dict: - return { - "agent": None, - "session_key": "session-key", - "history": [], - "history_lock": threading.Lock(), - "history_version": 0, - "running": False, - "attached_images": [], - "transport": None, - **extra, - } - - -# ── FanoutTransport ──────────────────────────────────────────────────────── - - -def test_fanout_delivers_to_every_attached_transport(): - a, b = _FakeClient("a"), _FakeClient("b") - fan = FanoutTransport(a, b) - - assert fan.write({"frame": 1}) is True - assert a.frames == b.frames == [{"frame": 1}] - assert fan.transports() == [a, b] - - -def test_fanout_attach_is_idempotent_and_detach_is_by_identity(): - a, b = _FakeClient("a"), _FakeClient("b") - fan = FanoutTransport(a) - - assert fan.attach(a) is False # already attached - assert fan.attach(b) is True - assert fan.contains(a) and fan.contains(b) - assert fan.has_transports(excluding=a) is True - - assert fan.detach(a) is True - assert fan.detach(a) is False - assert fan.transports() == [b] - assert fan.has_transports(excluding=b) is False - - -def test_fanout_prunes_a_wedged_peer_without_disturbing_the_rest(): - """(g) A raising peer is dropped; the healthy client keeps its stream.""" - healthy, wedged = _FakeClient("healthy"), _FakeClient("wedged", boom=True) - fan = FanoutTransport(wedged, healthy) - - assert fan.write({"frame": 1}) is True - assert healthy.frames == [{"frame": 1}] - assert fan.transports() == [healthy] - - # And the pruning is permanent — no retry storm against the wedged peer. - assert fan.write({"frame": 2}) is True - assert len(healthy.frames) == 2 - - -def test_fanout_prunes_a_peer_that_reports_gone_and_reports_all_dead(): - gone = _FakeClient("gone", ok=False) - fan = FanoutTransport(gone) - - # write() returned False -> peer gone -> pruned, and an empty fan-out - # reports peer-gone exactly like a single dead transport would. - assert fan.write({"frame": 1}) is False - assert fan.transports() == [] - assert fan.write({"frame": 2}) is False - - -def test_fanout_close_releases_peers_without_closing_their_sockets(): - a = _FakeClient("a") - fan = FanoutTransport(a) - - fan.close() - - assert fan.transports() == [] - # Each client owns its own socket; the WS handler closes it on disconnect. - assert a.closed is False - - -# ── attach ladder ────────────────────────────────────────────────────────── - - -def test_attach_replaces_an_empty_or_stdio_slot_without_wrapping(): - """The single-client shape is unchanged: no FanoutTransport in sight.""" - a = _FakeClient("a") - - empty = _session(transport=None) - assert server._attach_session_transport(empty, a) is True - assert empty["transport"] is a - - stdio = _session(transport=server._stdio_transport) - assert server._attach_session_transport(stdio, a) is True - assert stdio["transport"] is a - - parked = _session(transport=server._detached_ws_transport) - assert server._attach_session_transport(parked, a) is True - assert parked["transport"] is a - - -def test_attach_of_the_same_transport_is_a_noop(): - a = _FakeClient("a") - session = _session(transport=a) - - assert server._attach_session_transport(session, a) is True - assert session["transport"] is a - - -def test_attach_of_a_second_client_wraps_both_and_a_third_joins_the_fanout(): - a, b, c = _FakeClient("a"), _FakeClient("b"), _FakeClient("c") - session = _session(transport=a) - - server._attach_session_transport(session, b) - assert isinstance(session["transport"], FanoutTransport) - assert session["transport"].transports() == [a, b] - - server._attach_session_transport(session, c) - assert session["transport"].transports() == [a, b, c] - - -def test_attach_never_lets_stdio_displace_a_live_client(): - """An unbound-context activate must not silence the websocket that owns it.""" - a = _FakeClient("a") - session = _session(transport=a) - - assert server._attach_session_transport(session, server._stdio_transport) is False - assert session["transport"] is a - - -def test_attach_flattens_a_fanout_argument_instead_of_nesting_it(): - a, b, c = _FakeClient("a"), _FakeClient("b"), _FakeClient("c") - session = _session(transport=a) - server._attach_session_transport(session, b) - - server._attach_session_transport(session, FanoutTransport(b, c)) - - assert session["transport"].transports() == [a, b, c] - - -def test_attach_flattens_a_fanout_argument_onto_a_single_client_slot(monkeypatch): - """The nesting the queued-prompt paths can actually reach. - - A busy submit pins ``t or session["transport"]`` to the queued prompt, so - the envelope can hold the FAN-OUT that was in the slot, and the drain hands - that straight back to attach. Flattening only when the slot ALREADY fans out - leaves the leaf-slot case nesting one fan-out inside another, and every - reader of the slot scans a single level — ``FanoutTransport.contains``, the - steer-authority check, detach — so a peer in the inner fan-out is invisible - to all of them, and a peer in BOTH levels is written to twice per frame. - """ - a, b = _FakeClient("a"), _FakeClient("b") - - # (1) A fan-out arriving at a leaf slot is flattened into it, not wrapped. - session = _session(transport=a) - assert server._attach_session_transport(session, FanoutTransport(a, b)) is True - slot = session["transport"] - assert isinstance(slot, FanoutTransport) - assert slot.transports() == [a, b] - assert not any(isinstance(t, FanoutTransport) for t in slot.transports()) - assert server._session_transport_contains(session, b) is True - - # (2) The sequence that produces it. Two clients attach, so the slot holds a - # fan-out; a busy submit captures that slot in the queued envelope; the - # second client disconnects and the slot collapses back to the first; the - # drain then attaches the captured fan-out to a leaf slot. - monkeypatch.setattr(server, "_run_prompt_submit", lambda *a_, **k: None) - - drained = _session(transport=a) - server._attach_session_transport(drained, b) - captured = drained["transport"] - assert isinstance(captured, FanoutTransport) - server._enqueue_prompt(drained, "queued while busy", captured) - - assert server._detach_session_transport(drained, b) is True - assert drained["transport"] is a - - server._sessions["flatten-sid"] = drained - try: - assert server._drain_queued_prompt("drain", "flatten-sid", drained) is True - # Flat: the captured fan-out lost b to the same detach that collapsed - # the slot, so flattening it re-attaches only a, which is already there. - assert drained["transport"] is a - server._emit("message.delta", "flatten-sid", {"text": "once"}) - finally: - server._sessions.pop("flatten-sid", None) - - # Nesting would have put a inside the slot AND under it: one frame, two - # writes to the same client. - assert a.types() == ["message.delta"] - - -def test_detach_collapses_back_to_the_single_remaining_client(): - a, b = _FakeClient("a"), _FakeClient("b") - session = _session(transport=a) - server._attach_session_transport(session, b) - - assert server._detach_session_transport(session, a) is True - assert session["transport"] is b - assert server._detach_session_transport(session, b) is False - - -def test_detach_of_a_single_client_leaves_the_slot_for_the_caller_to_park(): - a = _FakeClient("a") - session = _session(transport=a) - - # False == "no live client remains" — the disconnect path parks the sentinel. - assert server._detach_session_transport(session, a) is False - assert session["transport"] is a - - -def test_live_transport_predicates_ignore_stdio_and_the_drop_sentinel(): - a = _FakeClient("a") - - assert server._session_has_live_transport(_session(transport=a)) is True - assert server._session_has_live_transport(_session(transport=a), excluding=a) is False - assert server._session_has_live_transport(_session(transport=None)) is False - assert ( - server._session_has_live_transport(_session(transport=server._stdio_transport)) - is False - ) - assert ( - server._session_has_live_transport( - _session(transport=server._detached_ws_transport) - ) - is False - ) - - -def test_a_closed_socket_is_not_a_live_peer(): - """A transport that already latched ``_closed`` is a departed client. - - ``_transport_is_live_peer`` ended on a bare ``return True``, so a WSTransport - whose socket had gone away still counted as a live peer: a session whose only - remaining peer was that dead socket answered "another client is still here" - and escaped both the park and the reap. ``_transport_is_dead`` is the - module's deadness predicate; the ladder now defers to it. - """ - dead = _FakeClient("dead") - dead._closed = True # the flag WSTransport latches when its socket goes away - - assert server._transport_is_live_peer(dead) is False - assert server._session_has_live_transport(_session(transport=dead)) is False - - # And a fan-out that collapses onto a dead peer is parked, not skipped: the - # live client leaving was the last real client. - live_client = _FakeClient("live") - session = _session(transport=FanoutTransport(live_client, dead)) - server._sessions["dead-peer-sid"] = session - try: - assert server._close_sessions_for_transport(live_client) == (0, 1) - assert session["transport"] is server._detached_ws_transport - finally: - server._sessions.pop("dead-peer-sid", None) - - -def test_a_fanout_with_no_live_peer_is_dead(): - """``_transport_is_dead`` is the reapers' gate, and a fan-out could not fail it. - - A ``FanoutTransport`` is never the drop sentinel and its ``__slots__`` give - it no ``_closed``, so a session that fanned out once read as alive forever: - the TTL reaper, the LRU cap and the disconnect revalidation all ask this one - predicate. A fan-out reaches the no-live-peer state without any disconnect - passing through ``_close_sessions_for_transport`` — a write that returns - False or raises prunes that peer — so the empty case is reachable, and it is - dead. - """ - a, b = _FakeClient("a"), _FakeClient("b") - session = _session(transport=a) - server._attach_session_transport(session, b) - assert isinstance(session["transport"], FanoutTransport) - assert server._transport_is_dead(session["transport"]) is False - - # One peer gone, one still reading: the session keeps its client. - a._closed = True # the flag WSTransport latches when its socket goes away - assert server._transport_is_dead(session["transport"]) is False - - b._closed = True - assert server._transport_is_dead(session["transport"]) is True - - # Pruning every peer leaves the fan-out empty, which is nobody attached. - assert server._transport_is_dead(FanoutTransport()) is True - - -def test_steer_authority_recognizes_the_exact_client_inside_a_fanout(): - """Wrapping attached peers must not revoke the commissioning peer's authority.""" - owner, watcher, stranger = ( - _FakeClient("owner"), - _FakeClient("watcher"), - _FakeClient("stranger"), - ) - session = _session(transport=FanoutTransport(owner, watcher)) - server._sessions["sid"] = session - try: - token = server.bind_transport(owner) +from tui_gateway.transport import FanoutTransport, StdioTransport +from tui_gateway.ws import WSTransport + + +class PipeClient: + def __init__(self, stack, *, reading=True): + read_fd, write_fd = os.pipe() + self.reader = stack.enter_context(os.fdopen(read_fd, "r", encoding="utf-8")) + self.writer = os.fdopen(write_fd, "w", encoding="utf-8") + stack.callback(self._cleanup) + self.transport = StdioTransport(lambda: self.writer, threading.Lock()) + self._closed = False + self.writes = 0 + self.frames = queue.Queue() + if reading: + self.thread = threading.Thread(target=self._read, daemon=True) + self.thread.start() + + def _cleanup(self): + with suppress(BrokenPipeError): + self.writer.close() + if hasattr(self, "thread"): + self.thread.join(timeout=5) + assert not self.thread.is_alive() + + def _read(self): + for line in self.reader: + self.frames.put(json.loads(line)) + + def write(self, obj): + self.writes += 1 + return self.transport.write(obj) + + def close(self): + self._closed = True + + def receive(self): + return self.frames.get(timeout=5) + + +class SocketClient(WSTransport): + """Real WSTransport with its ASGI send backed by a kernel socketpair.""" + def __init__(self, stack, *, reading=True): + self.reader, self.writer = socket.socketpair() + self.writer.setblocking(False) + loop = asyncio.new_event_loop() + super().__init__(self, loop) + self.writes = 0 + self.loop_thread = threading.Thread(target=loop.run_forever, daemon=True) + self.loop_thread.start() + self.frames = queue.Queue() + self.read_thread = None + if reading: + self.read_thread = threading.Thread(target=self._read, daemon=True) + self.read_thread.start() + stack.callback(self._cleanup) + + async def send_text(self, payload): + self.writes += 1 + await self._loop.sock_sendall(self.writer, (payload + "\n").encode()) + + def _read(self): + with self.reader.makefile("r", encoding="utf-8") as stream: + for line in stream: + self.frames.put(json.loads(line)) + + def receive(self): + return self.frames.get(timeout=5) + + def _cleanup(self): + async def cancel_sends(): + self.close() + tasks = [task for task in asyncio.all_tasks() if task is not asyncio.current_task()] + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + asyncio.run_coroutine_threadsafe(cancel_sends(), self._loop).result(timeout=5) + self.writer.close() + if self.read_thread: + self.read_thread.join(timeout=5) + assert not self.read_thread.is_alive() + self.reader.close() + self._loop.call_soon_threadsafe(self._loop.stop) + self.loop_thread.join(timeout=5) + self._loop.close() + + +def _session(transport): + return dict(transport=transport, agent=None, session_key="fanout-invariant", + history=[], history_lock=threading.Lock(), history_version=0, + running=False, attached_images=[]) + + +@pytest.mark.parametrize("attachment", ["direct", "flattened"]) +def test_membership_preserves_terminal_delivery_and_revokes_departed_peers(monkeypatch, attachment): + with ExitStack() as stack: + a, b, stranger = [PipeClient(stack) for _ in range(3)] + session = _session(a) + monkeypatch.setitem(server._sessions, "fanout-invariant", session) + newcomer = b if attachment == "direct" else FanoutTransport(a, b) + assert server._attach_session_transport(session, newcomer) + assert server._attach_session_transport(session, b) + for kind in ("message.start", "message.delta", "message.complete"): + server._emit(kind, "fanout-invariant", {"text": "α"}) + first, second = a.receive(), b.receive() + assert first == second + assert first["params"]["type"] == kind + for client, allowed in ((a, True), (b, True), (stranger, False)): + token = server.bind_transport(client) + try: + assert (server._current_session_steer_authority("fanout-invariant")[0] is client) == allowed + finally: + server.reset_transport(token) + # RPC replies stay on their request transport, never the subscriber set. + token = server.bind_transport(a) try: - assert server._current_session_steer_authority("sid") == (owner, session) + assert server.write_json({"jsonrpc": "2.0", "id": "private", "result": "owner only"}) finally: server.reset_transport(token) + assert a.receive()["id"] == "private" + assert b.frames.empty() + from gateway import browser_control_broker as broker_module + monkeypatch.setattr(broker_module, "browser_control_enabled", lambda: True) + for client in (a, b, stranger): + client.auth_identity = {"user_id": "fanout-owner", "provider": "fixture"} + session["profile"] = "default" - token = server.bind_transport(stranger) + def controller(client, action, **params): + return server.dispatch({"jsonrpc": "2.0", "id": 1, + "method": "browser.controller." + action, + "params": {"session_id": "fanout-invariant", **params}}, client) + + registered = controller(a, "register", controller_id="invariant", + browser_profile_id="fixture", capabilities=["controller.noop"], + protocol_version=broker_module.BROWSER_CONTROL_PROTOCOL_VERSION) + assert "result" in registered, registered try: - assert server._current_session_steer_authority("sid") == (None, None) + assert controller(a, "heartbeat")["result"] == {"ok": True} + assert controller(b, "heartbeat")["error"]["message"] == "controller is not owned by this transport" + assert controller(stranger, "heartbeat")["error"]["message"] == "session is not owned by this transport" finally: - server.reset_transport(token) - finally: - server._sessions.pop("sid", None) - - -def test_steer_authority_is_granted_to_an_attached_watcher(): - """A client that attached to watch a session may also steer its subagents. - - This is INTENTIONAL and wider than the pre-fan-out rule, which admitted only - whichever client happened to hold the transport slot. A mirrored session has - no single owner in that slot, so authority is membership in it. Narrowing - this back to the commissioning peer would need a per-subagent record of who - commissioned it, which this change does not add; the widening is stated in - the pull request description, and this test pins it so it cannot be changed - silently in either direction. - """ - owner, watcher, stranger = ( - _FakeClient("owner"), - _FakeClient("watcher"), - _FakeClient("stranger"), - ) - session = _session(transport=owner) - server._attach_session_transport(session, watcher) - solo = _session(transport=owner) - server._sessions["sid"] = session - server._sessions["solo"] = solo - try: - token = server.bind_transport(watcher) - try: - assert server._current_session_steer_authority("sid") == (watcher, session) - # Control: on a single-client session the same watcher is a stranger, - # so the widening reaches attached clients and nobody else. - assert server._current_session_steer_authority("solo") == (None, None) - finally: - server.reset_transport(token) - - token = server.bind_transport(stranger) - try: - assert server._current_session_steer_authority("sid") == (None, None) - finally: - server.reset_transport(token) - finally: - server._sessions.pop("sid", None) - server._sessions.pop("solo", None) - - -# ── browser-control session ownership ────────────────────────────────────── -# -# The four browser.controller.* handlers gate on the session slot exactly as -# subagent.steer did, so fan-out breaks them the same way and the fix is the -# same predicate. These tests pin the gate itself: they assert on the ownership -# error message and stop at the NEXT gate ("no controller registered for this -# session"), so a later broker change cannot make them pass vacuously. - -_CONTROLLER_IDENTITY = {"user_id": "user-fixture", "provider": "provider-fixture"} -_NOT_OWNED = "session is not owned by this transport" -_NO_CONTROLLER = "no controller registered for this session" - - -def _controller_client(name: str) -> _FakeClient: - """A fan-out client that also carries a server-authenticated identity.""" - client = _FakeClient(name) - client.auth_identity = dict(_CONTROLLER_IDENTITY) - return client - - -def _controller_rpc(transport, method_name: str, **params) -> dict: - return server.dispatch( - { - "jsonrpc": "2.0", - "id": 1, - "method": method_name, - "params": params, - }, - transport, - ) - - -def _error_message(response: dict) -> str | None: - return (response.get("error") or {}).get("message") - - -def test_browser_control_ownership_gate_admits_a_peer_inside_a_fanout(): - """Wrapping attached peers must not revoke browser control for all of them.""" - owner, watcher, stranger = ( - _controller_client("owner"), - _controller_client("watcher"), - _controller_client("stranger"), - ) - session = _session(transport=FanoutTransport(owner, watcher), profile="default") - server._sessions["sid"] = session - try: - # The peer that would have registered the controller clears the - # ownership gate and stops at the next one. - assert _error_message(_controller_rpc(owner, "browser.controller.heartbeat", session_id="sid")) == _NO_CONTROLLER - # Widened, exactly as the steer conversion widened steer: any attached - # peer clears the session gate. The broker's is_owner check below it is - # what still refuses a peer that did not register the controller. - assert _error_message(_controller_rpc(watcher, "browser.controller.heartbeat", session_id="sid")) == _NO_CONTROLLER - # An unattached client is still refused at the ownership gate. - assert _error_message(_controller_rpc(stranger, "browser.controller.heartbeat", session_id="sid")) == _NOT_OWNED - finally: - server._sessions.pop("sid", None) - - -def test_browser_control_ownership_gate_is_unchanged_for_a_single_client(): - """Control: a bare slot still compares by identity on all four handlers.""" - owner, stranger = _controller_client("owner"), _controller_client("stranger") - session = _session(transport=owner, profile="default") - server._sessions["solo"] = session - try: - for method_name in ( - "browser.controller.heartbeat", - "browser.controller.result", - "browser.controller.detach", - ): - assert _error_message(_controller_rpc(stranger, method_name, session_id="solo")) == _NOT_OWNED - assert _error_message(_controller_rpc(owner, "browser.controller.heartbeat", session_id="solo")) == _NO_CONTROLLER - finally: - server._sessions.pop("solo", None) - - -def test_browser_controller_registers_and_detaches_inside_a_fanout(monkeypatch): - """End to end: registration on a mirrored session, then a clean detach.""" - from gateway import browser_control_broker - - monkeypatch.setattr( - "gateway.browser_control_broker.browser_control_enabled", lambda: True - ) - owner, watcher, stranger = ( - _controller_client("owner"), - _controller_client("watcher"), - _controller_client("stranger"), - ) - session = _session(transport=FanoutTransport(owner, watcher), profile="default") - server._sessions["sid"] = session - registration = {} - try: - registration = _controller_rpc( - owner, - "browser.controller.register", - session_id="sid", - controller_id="controller-fixture", - browser_profile_id="browser-profile-fixture", - capabilities=["controller.noop"], - protocol_version=browser_control_broker.BROWSER_CONTROL_PROTOCOL_VERSION, - ) - assert registration.get("error") is None, registration - assert registration["result"]["scope"]["session_id"] == "sid" - - # The registering peer owns the controller; the fan-out peer that did - # not register clears the session gate and is refused by the broker. - assert _controller_rpc(owner, "browser.controller.heartbeat", session_id="sid")["result"] == {"ok": True} - assert _error_message(_controller_rpc(watcher, "browser.controller.heartbeat", session_id="sid")) == "controller is not owned by this transport" - assert _error_message(_controller_rpc(stranger, "browser.controller.heartbeat", session_id="sid")) == _NOT_OWNED - - detached = _controller_rpc(owner, "browser.controller.detach", session_id="sid") - assert detached.get("error") is None, detached - finally: - if registration.get("result"): - broker = browser_control_broker.get_browser_control_broker() - scope = broker.scope_for_session( - session_id="sid", - principal_id=registration["result"]["scope"]["principal_id"], - transport_family=registration["result"]["scope"]["transport_family"], - ) - if scope is not None: - broker.detach(scope, notify_controller=False) - server._sessions.pop("sid", None) - - -# ── event delivery ───────────────────────────────────────────────────────── - - -def test_two_attached_clients_both_receive_a_session_event(monkeypatch): - """(a) The headline behaviour: one emit, two clients.""" - a, b = _FakeClient("a"), _FakeClient("b") - session = _session(transport=a) - server._attach_session_transport(session, b) - server._sessions["sid"] = session - try: - server._emit("message.delta", "sid", {"text": "hi"}) - finally: - server._sessions.pop("sid", None) - - assert a.types() == ["message.delta"] - assert b.types() == ["message.delta"] - assert a.frames[0] == b.frames[0] - - -def test_a_single_client_session_writes_through_the_bare_transport(monkeypatch): - """(j) Control: nothing about the one-client path changed.""" - a = _FakeClient("a") - server._sessions["sid"] = _session(transport=a) - try: - assert server._sessions["sid"]["transport"] is a - server._emit("message.delta", "sid", {"text": "hi"}) - finally: - server._sessions.pop("sid", None) - - # The replay contract stamps a monotonic ``seq`` on every WS event frame. - # Strip it: this control test is about ROUTING, not numbering. - frames = [ - {**f, "params": {k: v for k, v in f["params"].items() if k != "seq"}} - for f in a.frames - ] - assert frames == [ - { - "jsonrpc": "2.0", - "method": "event", - "params": { - "type": "message.delta", - "session_id": "sid", - "payload": {"text": "hi"}, - }, - } - ] - - -def test_both_attached_clients_receive_the_terminal_message_complete(): - """The frame that ENDS a turn fans out too, not just the streaming deltas. - - ``message.complete`` is how a client knows the turn is over. A mirrored - session that delivered deltas but dropped the terminal frame would leave the - second client rendering a turn that never finishes. - """ - a, b = _FakeClient("a"), _FakeClient("b") - session = _session(transport=a) - server._attach_session_transport(session, b) - server._sessions["sid"] = session - try: - server._emit("message.delta", "sid", {"text": "par"}) - server._emit("message.complete", "sid", {"text": "part", "status": "ok"}) - finally: - server._sessions.pop("sid", None) - - assert a.types() == ["message.delta", "message.complete"] - assert b.types() == ["message.delta", "message.complete"] - assert a.frames[-1] == b.frames[-1] - assert a.frames[-1]["params"]["payload"] == {"text": "part", "status": "ok"} - - # Control: the one-client path delivers the same terminal frame through the - # bare transport, with no fan-out in the slot. - solo = _FakeClient("solo") - server._sessions["solo-sid"] = _session(transport=solo) - try: - assert server._sessions["solo-sid"]["transport"] is solo - server._emit("message.complete", "solo-sid", {"text": "part", "status": "ok"}) - finally: - server._sessions.pop("solo-sid", None) - - assert solo.types() == ["message.complete"] - assert solo.frames[-1]["params"]["payload"] == {"text": "part", "status": "ok"} - - -# ── prompt.submit / queued drain ─────────────────────────────────────────── - - -def test_second_clients_submit_attaches_and_the_first_keeps_streaming(monkeypatch): - """(c) A second client's prompt.submit must not cut the first one out.""" - monkeypatch.setattr(server, "_run_prompt_submit", lambda *a, **k: None) - monkeypatch.setattr(server, "_ensure_active_session_slot", lambda *a, **k: None) - - watcher, submitter = _FakeClient("watcher"), _FakeClient("submitter") - session = _session(transport=watcher, agent=types.SimpleNamespace()) - server._sessions["sid"] = session - token = server.bind_transport(submitter) - try: - server._methods["prompt.submit"]("r1", {"session_id": "sid", "text": "hello"}) - server._emit("message.delta", "sid", {"text": "answer"}) - finally: - server.reset_transport(token) - server._sessions.pop("sid", None) - - assert isinstance(session["transport"], FanoutTransport) - assert watcher.types() == ["message.delta"] - assert submitter.types() == ["message.delta"] - - -def test_queued_prompt_drain_keeps_both_clients_attached(monkeypatch): - """(d) The hole every competing patch missed: the drain used to rebind. - - Client A is streaming; client B submits mid-turn, so B's prompt is queued - with B's transport pinned to it. When the drain fires it must ATTACH B — - rebinding pinned the whole drained turn to B and silenced A. - """ - 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) - 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")] - assert isinstance(session["transport"], FanoutTransport) - assert a.types() == ["message.delta"] - 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) - - b = _FakeClient("b") - session = _session(transport=server._detached_ws_transport) - server._enqueue_prompt(session, "from B", b) - - assert server._drain_queued_prompt("drain", "sid", session) is True - assert session["transport"] is b - - -# ── disconnect ───────────────────────────────────────────────────────────── - - -def test_watcher_disconnect_leaves_the_other_client_streaming(): - """(e) One client leaving must not park or reap a session someone is reading.""" - a, watcher = _FakeClient("a"), _FakeClient("watcher") - session = _session(transport=a) - server._attach_session_transport(session, watcher) - server._sessions["sid"] = session - try: - assert server._close_sessions_for_transport(watcher) == (0, 0) - assert session["transport"] is a - assert server._ws_session_is_orphaned(session) is False - server._emit("message.delta", "sid", {"text": "still here"}) - finally: - server._sessions.pop("sid", None) - - assert a.types() == ["message.delta"] - assert watcher.types() == [] - - -def test_last_client_disconnect_parks_exactly_as_before(monkeypatch): - """(f) Once the last client goes, today's park + grace-reap path runs.""" - monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0) - - a, watcher = _FakeClient("a"), _FakeClient("watcher") - session = _session(transport=a) - server._attach_session_transport(session, watcher) - server._sessions["sid"] = session - try: - assert server._close_sessions_for_transport(watcher) == (0, 0) + controller(a, "detach") + session["viewers"] = {b: object(), stranger: object()} + assert server._detach_session_transport(session, b) + assert not session["viewers"] + assert not server._session_transport_contains(session, b) + server._emit("message.complete", "fanout-invariant", {"text": "only A"}) + assert a.receive()["params"]["payload"]["text"] == "only A" + assert b.frames.empty() + assert server._attach_session_transport(session, b) + server._emit("message.complete", "fanout-invariant", {"text": "reattached"}) + assert a.receive() == b.receive() + b.close() + # A stale queued envelope must not restore a dead peer's authority. + assert not server._attach_session_transport(session, b) + assert not server._session_transport_contains(session, b) + assert not server._attach_session_transport(session, server._stdio_transport) + assert server._close_sessions_for_transport(b) == (0, 0) assert server._close_sessions_for_transport(a) == (0, 1) assert session["transport"] is server._detached_ws_transport - assert server._ws_session_is_orphaned(session) is True - finally: - server._sessions.pop("sid", None) -def test_last_client_disconnect_reaps_a_close_on_disconnect_session(monkeypatch): - """(f) close_on_disconnect still fires — but only for the LAST client.""" - # The disconnect path claims the session with _pop_session_by_id and finishes - # the close through _teardown_popped_session, so that is what this stubs; the - # sid is asserted through the pop rather than through the call arguments. - closed = [] - - def _fake_teardown(session, *, end_reason: str = "tui_close") -> bool: - closed.append((session, end_reason)) - return True - - monkeypatch.setattr(server, "_teardown_popped_session", _fake_teardown) - - a, watcher = _FakeClient("a"), _FakeClient("watcher") - session = _session(transport=a, close_on_disconnect=True) - server._attach_session_transport(session, watcher) - server._sessions["sid"] = session - try: - assert server._close_sessions_for_transport(watcher) == (0, 0) - assert closed == [] - assert server._close_sessions_for_transport( - a, end_reason="ws_disconnect" - ) == (1, 0) - assert "sid" not in server._sessions - finally: - server._sessions.pop("sid", None) - - assert closed == [(session, "ws_disconnect")] - - -def test_orphan_check_spares_a_session_that_still_has_a_fanout_peer(): - # _ws_session_is_orphaned stayed slot-identity based on main: a fan-out is - # never the drop sentinel, and the disconnect path only parks the sentinel - # once the last peer is gone, so these hold without a liveness rewrite. - a, watcher = _FakeClient("a"), _FakeClient("watcher") - session = _session(transport=a) - server._attach_session_transport(session, watcher) - - assert server._ws_session_is_orphaned(session) is False - - # Losing one of two clients collapses the slot but keeps the session live. - assert server._detach_session_transport(session, a) is True - assert server._ws_session_is_orphaned(session) is False - - # Losing the last one reports "no client left"; the caller then parks the - # drop sentinel, which is what makes the session orphaned. - assert server._detach_session_transport(session, watcher) is False - session["transport"] = server._detached_ws_transport - assert server._ws_session_is_orphaned(session) is True - - -# --------------------------------------------------------------------------- -# Concurrent fan-out delivery -# -# A SERIAL ``FanoutTransport.write`` — each peer written in turn — would let a -# client whose write parks for the full ``tui_gateway.ws._WS_WRITE_TIMEOUT_S`` -# (10s, the non-streaming path's ``fut.result`` wait for a stalled event loop) -# hold the same frame back from every healthy client behind it in the list. -# That is the property these tests pin, so the writes run concurrently, with two -# cases deliberately kept inline: a lone peer (single-client sessions must not -# change at all) and a caller that is already on an event loop (handing that -# write to a pool thread would make the pool thread block on the loop it just -# left, freezing both). -# -# Extra imports for this block: ``asyncio`` below, plus a function-local -# ``tui_gateway.transport`` in the deadline test (it monkeypatches a module -# constant). Everything else — ``threading``, ``FanoutTransport``, -# ``_FakeClient`` — comes from the top of this file. -# --------------------------------------------------------------------------- - -import asyncio - - -class _SlowPeer: - """A peer whose write parks until the test releases it. - - Models the real stall: a non-streaming ``WSTransport.write`` blocks on - ``fut.result(timeout=_WS_WRITE_TIMEOUT_S)`` while the owning event loop is - busy, so the emitting thread sits inside that one peer's write for seconds. - """ - - def __init__(self, name: str = "slow") -> None: - self.name = name - self.frames: list[dict] = [] - self.entered = threading.Event() - self.release = threading.Event() - self.closed = False - - def write(self, obj: dict) -> bool: - self.entered.set() - # Bounded: a regression must fail the assertion below, never hang the - # suite waiting for a release that the failing path never reaches. - self.release.wait(timeout=5.0) - self.frames.append(obj) - return True - - def close(self) -> None: - self.closed = True - - -class _SignallingPeer: - """A healthy peer that announces each frame the moment it lands.""" - - def __init__(self, name: str = "healthy") -> None: - self.name = name - self.frames: list[dict] = [] - self.got_frame = threading.Event() - self.closed = False - - def write(self, obj: dict) -> bool: - self.frames.append(obj) - self.got_frame.set() - return True - - def close(self) -> None: - self.closed = True - - -class _ThreadRecordingPeer: - """A healthy peer that records which thread each of its writes ran on.""" - - def __init__(self, name: str = "peer") -> None: - self.name = name - self.frames: list[dict] = [] - self.threads: list[threading.Thread] = [] - self.closed = False - - def write(self, obj: dict) -> bool: - self.threads.append(threading.current_thread()) - self.frames.append(obj) - return True - - def close(self) -> None: - self.closed = True - - -def _healthy_peer_is_not_held_behind(slow_first: bool) -> None: - """One wedged peer, one healthy peer, one ``write`` from a worker thread. - - The healthy peer must hold the frame while the wedged peer is still parked - inside its own write. Asserted for both list orders so the fix cannot be a - special case that only ever hurries the first (or the last) peer along. - """ - slow = _SlowPeer() - healthy = _SignallingPeer() - peers = (slow, healthy) if slow_first else (healthy, slow) - fanout = FanoutTransport(*peers) - frame = {"jsonrpc": "2.0", "method": "event", "params": {"type": "message.complete"}} - - returned = threading.Event() - result: dict = {} - - def emit() -> None: +@pytest.mark.linux_only +@pytest.mark.parametrize("client_type", [PipeClient, SocketClient]) +@pytest.mark.parametrize("slow_first", [True, False]) +@pytest.mark.parametrize("on_loop", [True, False]) +def test_backpressure_never_blocks_later_frames_or_other_subscribers(slow_first, on_loop, client_type, monkeypatch): + monkeypatch.setattr("tui_gateway.ws._WS_WRITE_TIMEOUT_S", 0.01) + monkeypatch.setattr("tui_gateway.ws._TOKEN_COALESCE_S", 0) + with ExitStack() as stack: + healthy, slow = client_type(stack), client_type(stack, reading=False) + # Fill the actual kernel pipe, not a fake wait in a transport.write(). + fd = slow.writer.fileno() + os.set_blocking(fd, False) try: - result["ok"] = fanout.write(frame) + while True: + os.write(fd, b"x" * 4096) + except BlockingIOError: + pass finally: - returned.set() + os.set_blocking(fd, client_type is PipeClient) + fan = FanoutTransport(*((slow, healthy) if slow_first else (healthy, slow))) + returned = threading.Event() + errors = [] - # Emit from a worker thread, which is where the gateway's event writes come - # from: ``handle_ws`` runs ``server.dispatch`` on ``asyncio.to_thread``. - worker = threading.Thread(target=emit, name="fanout-emitter", daemon=True) - worker.start() - try: - assert slow.entered.wait(timeout=2.0), "the wedged peer never got the frame" - assert healthy.got_frame.wait(timeout=2.0), ( - "the healthy peer did not get the frame while a wedged peer held its " - "write open — the fan-out is still serial" - ) - assert healthy.frames == [frame] - assert slow.frames == [] - finally: - slow.release.set() - assert returned.wait(timeout=5.0), "the fan-out write never returned" - worker.join(timeout=5.0) + def emit(): + try: + for kind in ("message.start", "message.delta", "message.complete"): + assert fan.write({"params": {"type": kind}}) + except BaseException as exc: + errors.append(exc) + finally: + returned.set() - # The fan-out still collects: the wedged peer's frame landed before the - # call returned, and both peers stay attached. - assert slow.frames == [frame] - assert result["ok"] is True - assert fanout.transports() == [*peers] + async def loop_emit(): + emit() - -def test_a_wedged_peer_does_not_hold_the_frame_from_a_healthy_peer_behind_it(): - _healthy_peer_is_not_held_behind(slow_first=True) - - -def test_a_wedged_peer_does_not_hold_the_frame_from_a_healthy_peer_ahead_of_it(): - _healthy_peer_is_not_held_behind(slow_first=False) - - -def test_a_write_from_inside_an_event_loop_stays_on_the_loop_thread(): - """The deadlock the concurrent path must never introduce. - - ``WSTransport.write`` fires and forgets when it can see it is running on - its own loop, and BLOCKS on ``fut.result`` when it cannot. Hand a - loop-thread write to a pool thread and that pool thread takes the blocking - path, waiting on a loop that is itself waiting on the pool — the whole - gateway stalls for the write timeout. So an on-loop caller stays inline. - """ - a, b = _ThreadRecordingPeer("a"), _ThreadRecordingPeer("b") - fanout = FanoutTransport(a, b) - frame = {"params": {"type": "message.complete"}} - outcome: dict = {} - - async def emit() -> None: - outcome["loop_thread"] = threading.current_thread() - outcome["ok"] = fanout.write(frame) - - # Driven from a side thread with a bounded join so a real deadlock fails - # this test instead of hanging the run. - runner = threading.Thread(target=lambda: asyncio.run(emit()), daemon=True) - runner.start() - runner.join(timeout=5.0) - assert not runner.is_alive(), "fan-out write from the event loop did not return" - - assert outcome["ok"] is True - assert a.frames == [frame] and b.frames == [frame] - assert a.threads == [outcome["loop_thread"]] - assert b.threads == [outcome["loop_thread"]] - - -def test_a_lone_peer_is_written_on_the_calling_thread(): - """Single-client sessions keep the exact behaviour they had before.""" - only = _ThreadRecordingPeer("only") - fanout = FanoutTransport(only) - - frame = {"params": {"type": "message.complete"}} - assert fanout.write(frame) is True - - assert only.frames == [frame] - assert only.threads == [threading.current_thread()] - - -def test_the_concurrent_path_prunes_dead_peers_and_still_delivers(): - """Pruning is unchanged: ``False`` and a raise both detach the peer.""" - gone = _FakeClient("gone", ok=False) - wedged = _FakeClient("wedged", boom=True) - healthy = _FakeClient("healthy") - fanout = FanoutTransport(gone, wedged, healthy) - - frame = {"params": {"type": "message.complete"}} - assert fanout.write(frame) is True - - assert healthy.frames == [frame] - assert fanout.transports() == [healthy] - - -def test_the_concurrent_path_reports_peer_gone_when_every_peer_is_dead(): - gone = _FakeClient("gone", ok=False) - wedged = _FakeClient("wedged", boom=True) - fanout = FanoutTransport(gone, wedged) - - assert fanout.write({"params": {"type": "message.complete"}}) is False - assert fanout.transports() == [] - - -def test_consecutive_frames_reach_every_peer_in_order(): - """The fan-out collects before returning, so frame A is on every peer - before frame B is dispatched to any of them.""" - a, b = _ThreadRecordingPeer("a"), _ThreadRecordingPeer("b") - fanout = FanoutTransport(a, b) - first = {"params": {"type": "message.delta", "text": "1"}} - second = {"params": {"type": "message.complete"}} - - assert fanout.write(first) is True - # Collect-before-return: nothing is still in flight when write() answers. - assert a.frames == [first] and b.frames == [first] - - assert fanout.write(second) is True - assert a.frames == [first, second] - assert b.frames == [first, second] - - -def test_a_peer_that_misses_the_deadline_is_kept_and_counts_as_delivered(monkeypatch): - """Slow is not dead. - - A peer still writing when the fan-out deadline expires stays attached and - the frame counts as in flight — which is what ``WSTransport.write`` itself - reports when its own write times out. Here it is the ONLY live peer, so the - ``True`` return rests entirely on it. - """ - from tui_gateway import transport as transport_module - - monkeypatch.setattr(transport_module, "_FANOUT_WRITE_DEADLINE_S", 0.05) - - slow = _SlowPeer() - gone = _FakeClient("gone", ok=False) - fanout = FanoutTransport(slow, gone) - frame = {"params": {"type": "message.complete"}} - - assert fanout.write(frame) is True - # The dead peer is pruned; the slow one is not. - assert fanout.transports() == [slow] - assert slow.frames == [] - - slow.release.set() + worker = threading.Thread(target=(lambda: asyncio.run(loop_emit())) if on_loop else emit, daemon=True) + worker.start() + try: + assert returned.wait(3), "slow subscriber blocked the emitting turn" + assert not errors + assert [healthy.receive()["params"]["type"] for _ in range(3)] == [ + "message.start", "message.delta", "message.complete"] + # Exhaust only the slow peer's bounded backlog; pace the healthy + # reader by receipts so scheduler latency cannot make it overflow. + for n in range(1024): + frame = {"params": {"type": "message.delta", "n": n}} + assert fan.write(frame) + assert healthy.receive() == frame + if not fan.contains(slow): + break + assert not fan.contains(slow), "slow backlog grew without bound" + for _ in range(16): + assert fan.attach(slow) + assert fan.write({"reattach": True}) + assert healthy.receive() == {"reattach": True} + fan.detach(slow) + assert slow.writes == 1, "reattach spawned more writers behind blocked I/O" + assert fan.contains(healthy) + assert fan.write({"params": {"type": "message.complete"}}) + assert healthy.receive()["params"]["type"] == "message.complete" + finally: + # Closing the real reader releases any blocked writer even on RED. + slow.reader.close() + worker.join(timeout=15) + fan.close() + assert not worker.is_alive() + assert not fan.write({"after": "close"}) diff --git a/tests/tui_gateway/test_shared_session_delivery.py b/tests/tui_gateway/test_shared_session_delivery.py index 6d5cf6aa29..fef62c34a3 100644 --- a/tests/tui_gateway/test_shared_session_delivery.py +++ b/tests/tui_gateway/test_shared_session_delivery.py @@ -7,10 +7,12 @@ from tui_gateway import server class Peer: def __init__(self): self.frames = [] + self.received = threading.Event() self._closed = False def write(self, frame): self.frames.append(frame) + self.received.set() return not self._closed def close(self): @@ -24,10 +26,14 @@ def test_reattach_preserves_terminal_delivery(monkeypatch): with session["history_lock"]: server._rebind_live_transport("shared", session, second) server._emit("message.complete", "shared", {"text": "finished"}) + assert first.received.wait(timeout=5) + assert second.received.wait(timeout=5) assert first.frames == second.frames assert len(first.frames) == 1 second.close() assert server._close_sessions_for_transport(second) == (0, 0) assert second not in session.get("viewers", {}) + first.received.clear() server._emit("message.complete", "shared", {"text": "still attached"}) + assert first.received.wait(timeout=5) assert len(first.frames) == 2 diff --git a/tui_gateway/session_transports.py b/tui_gateway/session_transports.py index dfe8ae1a35..12ee7bc046 100644 --- a/tui_gateway/session_transports.py +++ b/tui_gateway/session_transports.py @@ -4,146 +4,71 @@ from __future__ import annotations import threading from tui_gateway.method_ctx import bind_module -# ── multi-client fan-out ───────────────────────────────────────────────────── -# -# A session's ``transport`` slot holds either ONE transport — the historical, -# single-client shape, byte-identical to pre-fan-out behaviour — or a -# ``FanoutTransport`` wrapping several. Because the fan-out satisfies the same -# Transport protocol, ``write_json`` and every other reader of the slot are -# untouched; only the attach/detach ladder below knows the difference. -# -# ``_session_transport_lock`` serializes the read-modify-write of that one slot. -# It is a LEAF lock: nothing under it acquires ``_sessions_lock`` or a session's -# ``history_lock`` (the fan-out's own lock is likewise a leaf), so it is safe to -# take while holding either of those. -_session_transport_lock = threading.Lock() +# Leaf lock: callers may hold sessions/history locks, never acquire them here. +_session_transport_lock = threading.RLock() def _transport_is_live_peer(transport) -> bool: - """True when *transport* is a real, currently attached client. - - Excluded: the parked drop sentinel (by definition clientless) and stdio, - which is the process-wide fallback sink — a standalone ``hermes --tui`` - writes there, but nothing "attaches" to it and a second client can never - share it. - """ - if transport is None: - return False - if transport is _detached_ws_transport or transport is _stdio_transport: - return False - if isinstance(transport, (_DropTransport, StdioTransport)): - return False - # A socket that already latched ``_closed`` is a departed client, not a live - # peer: without this, a stale WSTransport keeps a session out of the park / - # reap path and keeps answering "another client is still here". - # ``_transport_is_dead`` is the deadness predicate this module shares with - # the reaper (defined in ``session_reaper``, published onto this namespace - # by ``method_ctx.bind_module``; module globals resolve at call time). - return not _transport_is_dead(transport) + """Exclude the process fallback sink, parked sentinel, and closed peers.""" + return (transport is not None + and transport is not _detached_ws_transport + and transport is not _stdio_transport + and not isinstance(transport, (_DropTransport, StdioTransport)) + and not _transport_is_dead(transport)) def _session_transport_contains(session: dict | None, transport) -> bool: - """True when *transport* is attached to *session*, directly or via fan-out.""" - if not session or transport is None: + if not session or transport is None or _transport_is_dead(transport): return False existing = session.get("transport") - if existing is transport: - return True - if isinstance(existing, FanoutTransport): - return existing.contains(transport) - return False + return existing is transport or ( + isinstance(existing, FanoutTransport) and existing.contains(transport)) def _session_live_transports(session: dict | None) -> list: - """Every live client attached to *session* (empty for a parked/stdio slot).""" existing = (session or {}).get("transport") - if isinstance(existing, FanoutTransport): - return [t for t in existing.transports() if _transport_is_live_peer(t)] - return [existing] if _transport_is_live_peer(existing) else [] + peers = existing.transports() if isinstance(existing, FanoutTransport) else [existing] + return [peer for peer in peers if _transport_is_live_peer(peer)] def _session_has_live_transport(session: dict | None, *, excluding=None) -> bool: - """True when *session* still has a live client attached, ignoring *excluding*. - - ``excluding`` answers whether another live client remains once one peer is - ignored, which is what the detach and orphan paths need: whether anything - survives the departing transport. - """ - return any(t is not excluding for t in _session_live_transports(session)) + return any(peer is not excluding for peer in _session_live_transports(session)) def _attach_session_transport(session: dict | None, transport) -> bool: - """Attach *transport* to *session* ADDITIVELY — never steal the slot. - - A ``FanoutTransport`` argument is FLATTENED before the ladder runs, whatever - the slot holds: each of its peers is attached as a leaf, so the slot never - nests one fan-out inside another. The queued-prompt paths make that - reachable — a busy submit records ``session["transport"]`` as the prompt's - transport, which is the fan-out itself when two clients are attached, and - the drain hands it back here after a disconnect has collapsed the slot to a - single client. Nesting would blind every single-level reader of the slot - (``FanoutTransport.contains``, the steer-authority scan, detach) to the - inner peers, so a client inside the inner fan-out could never be found, - parked, reaped, or granted authority over its own turn. - - The ladder, for the leaf newcomer that always reaches it: - - * same object already in the slot → no-op; - * the slot already fans out → attach into it; - * the slot is empty / stdio / the parked drop sentinel → take the slot, - which is exactly what the old rebind did; - * otherwise → wrap the incumbent and the newcomer in a ``FanoutTransport`` - so both clients keep streaming. - - A non-peer newcomer (stdio, drop sentinel) never displaces a live client: - an activate/resume dispatched without a bound websocket must not silence the - socket that owns the session. - - Returns ``True`` when *transport* is attached afterwards. - """ + """Add live peers; flatten captured queued fanouts without nesting authority.""" if not session or transport is None: return False - if isinstance(transport, FanoutTransport): - # Flatten OUTSIDE the lock — _session_transport_lock is not reentrant. - # Each peer then walks the ladder on its own, so a peer whose socket - # died while the prompt sat in the queue is dropped by the liveness rung - # instead of being pinned back into the slot. - attached = False - for peer in transport.transports(): - if _attach_session_transport(session, peer): - attached = True - return attached with _session_transport_lock: + if isinstance(transport, FanoutTransport): + # Snapshot and attach share detach's lock: a queued fanout cannot + # resurrect a still-open peer removed during flattening. + attached = [_attach_session_transport(session, peer) for peer in transport.transports()] + return any(attached) existing = session.get("transport") - if existing is transport: - return True - if isinstance(existing, FanoutTransport): - if isinstance(transport, FanoutTransport): - for member in transport.transports(): - existing.attach(member) - else: - existing.attach(transport) - return True + if _transport_is_dead(transport): + if isinstance(existing, FanoutTransport): + existing.detach(transport) + return False if not _transport_is_live_peer(transport): - if _transport_is_live_peer(existing): + if _session_has_live_transport(session): return False session["transport"] = transport return True - if not _transport_is_live_peer(existing): - session["transport"] = transport + if existing is transport: return True - session["transport"] = FanoutTransport(existing, transport) + if isinstance(existing, FanoutTransport): + existing.attach(transport) + return existing.contains(transport) + elif _transport_is_live_peer(existing): + session["transport"] = FanoutTransport(existing, transport) + else: + session["transport"] = transport return True def _detach_session_transport(session: dict | None, transport) -> bool: - """Detach *transport* from *session*'s slot. - - Returns ``True`` when a live client OTHER than *transport* remains — i.e. - the session must keep streaming and must NOT be parked or reaped. A - single-client session leaves the departing transport in the slot, exactly as - before fan-out existed; its caller parks the drop sentinel over it. - """ + """Remove membership; return whether another live client prevents parking.""" if not session: return False with _session_transport_lock: @@ -151,33 +76,27 @@ def _detach_session_transport(session: dict | None, transport) -> bool: existing = session.get("transport") if isinstance(existing, FanoutTransport): existing.detach(transport) - remaining = existing.transports() - if len(remaining) == 1: - # Collapse: a session back down to one client is indistinguishable - # from one that never fanned out. - session["transport"] = remaining[0] - return _session_has_live_transport(session, excluding=transport) + viewers = session.get("viewers") or {} + for viewer in list(viewers): + if not existing.contains(viewer) or _transport_is_dead(viewer): + viewers.pop(viewer, None) + # Keep the surviving mailbox: collapsing to a bare transport lets + # new frames overtake its already queued terminal/control events. + return _session_has_live_transport(session, excluding=transport) def _detach_transport_from_sessions(transport) -> list[tuple[str, dict]]: - """Detach *transport* from every session holding it. - - Returns the ``(sid, session)`` pairs left with NO live client — the ones the - disconnect path must park or reap. Sessions that retain another attached - client keep streaming and are not returned. - """ + """Remove even closed/pruned peers' viewer entries; return clientless slots.""" with _sessions_lock: - attached = [ - (sid, s) - for sid, s in _sessions.items() - if _session_transport_contains(s, transport) - ] - return [ - (sid, session) - for sid, session in attached - if not _detach_session_transport(session, transport) - ] - + attached = [] + for sid, session in _sessions.items(): + existing = session.get("transport") + if (existing is transport + or isinstance(existing, FanoutTransport) and existing.contains(transport) + or transport in (session.get("viewers") or {})): + attached.append((sid, session)) + return [(sid, session) for sid, session in attached + if not _detach_session_transport(session, transport)] def register(server) -> None: diff --git a/tui_gateway/transport.py b/tui_gateway/transport.py index 26ba3fd12b..00c23bbb84 100644 --- a/tui_gateway/transport.py +++ b/tui_gateway/transport.py @@ -9,8 +9,8 @@ A :class:`Transport` forwards a JSON-serialisable dict to its peer, so one dispa from __future__ import annotations -import asyncio -import concurrent.futures +from collections import deque +from dataclasses import dataclass, field import contextlib import contextvars import errno @@ -18,7 +18,6 @@ import json import logging import os import threading -import time from typing import Any, Callable, Optional, Protocol, runtime_checkable # Errno values that mean "the peer is gone" rather than "the host has a real I/O problem". Anything @@ -35,65 +34,6 @@ logger = logging.getLogger(__name__) # fully buffered on a pipe, so this ONLY makes sense with ``-u``/``PYTHONUNBUFFERED=1``; otherwise the TUI hangs. _DISABLE_FLUSH = (os.environ.get("HERMES_TUI_GATEWAY_NO_FLUSH", "") or "").strip().lower() in {"1", "true", "yes", "on"} -# Worker pool behind FanoutTransport.write. A non-streaming WebSocket write -# blocks the calling thread while the owning event loop flushes the frame — up -# to ``tui_gateway.ws._WS_WRITE_TIMEOUT_S`` (10s) when that loop is stalled — -# so walking a session's peers one at a time lets one wedged client hold the -# frame back from every healthy client behind it. The pool is created lazily -# and only ever used when a session has more than one peer, so the ordinary -# single-client gateway never starts a thread. -# -# Sizing: these workers are idle except while a peer is actually wedged, so the -# cap only has to cover the wedged peers of all sessions at once. If it were -# ever saturated the excess writes queue and run in submission order, which is -# the pre-pool behaviour — degraded, not broken. -_FANOUT_POOL_MAX_WORKERS = 32 - -# Wall-clock bound on ONE fan-out write, measured from before the first peer is -# dispatched. Deliberately larger than ``tui_gateway.ws._WS_WRITE_TIMEOUT_S`` -# (10.0) so a merely-slow peer reaches its own timeout and reports for itself; -# this deadline only exists so a peer that never returns at all cannot pin the -# emitting thread forever. Not imported from ``tui_gateway.ws``: that module -# imports ``tui_gateway.server``, which imports this one, so the import would -# be circular. Drift is harmless — a peer that misses this deadline is treated -# as in-flight, not as dead (see ``FanoutTransport.write``). -_FANOUT_WRITE_DEADLINE_S = 12.0 - -_fanout_pool: "concurrent.futures.ThreadPoolExecutor | None" = None -_fanout_pool_lock = threading.Lock() - - -def _fanout_write_pool() -> "concurrent.futures.ThreadPoolExecutor": - """The shared fan-out pool, created on first multi-peer write.""" - global _fanout_pool - pool = _fanout_pool - if pool is not None: - return pool - with _fanout_pool_lock: - if _fanout_pool is None: - _fanout_pool = concurrent.futures.ThreadPoolExecutor( - max_workers=_FANOUT_POOL_MAX_WORKERS, - thread_name_prefix="tui-fanout", - ) - return _fanout_pool - - -def _caller_is_on_event_loop() -> bool: - """True when the CALLING thread is running an asyncio event loop. - - Mirrors the probe in ``tui_gateway.ws.WSTransport.write``, which takes a - fire-and-forget path when it can see it is on its own loop and a BLOCKING - ``fut.result`` path when it cannot. ``FanoutTransport`` owns no loop and - its peers may sit on different ones, so the test here is the conservative - one: if this thread is running any loop at all, keep the peer writes on it. - """ - try: - asyncio.get_running_loop() - except RuntimeError: - return False - return True - - @runtime_checkable class Transport(Protocol): """Minimal interface every transport implements.""" @@ -176,206 +116,139 @@ class StdioTransport: return None +@dataclass(eq=False) +class _FanoutPeer: + transport: Transport + pending: deque = field(default_factory=deque) + pending_bytes: int = 0 + writing: bool = False + attached: bool = True + generation: int = 0 + + class FanoutTransport: - """Delivers one JSON frame to every client attached to a session. + """Ordered, bounded session-event mailboxes; RPC replies remain request-local. - A session's ``transport`` slot used to hold exactly one client, so every - ``prompt.submit`` / ``session.resume`` / ``session.activate`` / queued-prompt - drain REBOUND it and whichever client was there before went silent. This - object goes in the same slot and satisfies the same :class:`Transport` - protocol, so ``server.write_json`` — and every other reader of that slot — - is unchanged; the only difference is that N clients receive the frame - instead of the most recent one. - - Scope: this carries a session's ASYNC EVENT stream only. Request/response - RPCs still answer on the request's context-bound transport - (``dispatch`` → ``current_transport()``), so a client only ever sees replies - to its own calls. - - Peers are pruned on failure: a transport that returns ``False`` (peer gone) - or raises is dropped from the fan-out instead of being written to forever. - Pruning is about dead peers, not slow ones — a client that is merely wedged - is kept, and :meth:`write` runs the peer writes concurrently so its stall - stays its own. + One slow socket must not stop the emitting turn or any healthy subscriber. + Each peer has at most one daemon writer and a bounded backlog. On overflow + it loses its subscription (history/replay is the recovery path), not other + sessions sharing its socket. A write already in the OS cannot be revoked. """ - __slots__ = ("_lock", "_transports") + _MAX_PENDING_FRAMES = 256 + _MAX_PENDING_BYTES = 4 * 1024 * 1024 - def __init__(self, *transports: "Transport") -> None: + def __init__(self, *transports: Transport) -> None: self._lock = threading.Lock() - self._transports: list["Transport"] = [] + self._peers: list[_FanoutPeer] = [] for transport in transports: self.attach(transport) - def attach(self, transport: "Transport") -> bool: - """Add *transport*. Returns ``True`` when it was not already attached.""" + def attach(self, transport: Transport) -> bool: if transport is None or transport is self: return False with self._lock: - for existing in self._transports: - if existing is transport: - return False - self._transports.append(transport) + for peer in self._peers: + if peer.transport is transport: + if peer.attached: + return False + # Reuse the in-flight writer: reconnect cannot spawn more + # threads or overtake a write already inside this socket. + peer.attached = True + peer.generation += 1 + return True + self._peers.append(_FanoutPeer(transport)) return True - def detach(self, transport: "Transport") -> bool: - """Remove *transport*. Returns ``True`` when it was attached.""" + def _remove(self, peer: _FanoutPeer) -> None: + # Membership lock held; identity fences a stale writer from removing + # a later attachment of the same transport. + peer.attached = False + peer.pending.clear() + peer.pending_bytes = 0 + if not peer.writing and peer in self._peers: + self._peers.remove(peer) + + def detach(self, transport: Transport) -> bool: with self._lock: - for idx, existing in enumerate(self._transports): - if existing is transport: - del self._transports[idx] + for peer in self._peers: + if peer.attached and peer.transport is transport: + self._remove(peer) return True return False - def contains(self, transport: "Transport") -> bool: + def contains(self, transport: Transport) -> bool: with self._lock: - return any(existing is transport for existing in self._transports) + return any(peer.attached and peer.transport is transport for peer in self._peers) - def transports(self) -> list["Transport"]: - """A snapshot of the attached transports (safe to iterate).""" + def transports(self) -> list[Transport]: with self._lock: - return list(self._transports) + return [peer.transport for peer in self._peers if peer.attached] - def has_transports(self, *, excluding: "Transport | None" = None) -> bool: - """True when any transport other than *excluding* is still attached.""" - with self._lock: - return any(existing is not excluding for existing in self._transports) + def has_transports(self, *, excluding: Transport | None = None) -> bool: + return any(peer is not excluding for peer in self.transports()) - @staticmethod - def _deliver(transport: "Transport", obj: dict, dead: list) -> bool: - """Write *obj* to one peer inline. Records a dead peer in *dead*. - - Returns ``True`` when the peer took the frame. - """ - try: - ok = transport.write(obj) - except Exception: - logger.debug("fanout write failed; pruning peer", exc_info=True) - dead.append(transport) - return False - if ok: - return True - dead.append(transport) - return False - - def write(self, obj: dict) -> bool: - """Deliver *obj* to every attached transport. - - Iterates a SNAPSHOT and never holds the lock across a peer write, so a - peer write can never block an attach or a detach. - - With more than one peer the writes run CONCURRENTLY: every peer but one - is handed to a shared worker pool and the last runs on the calling - thread, which is going to block here anyway. Serial delivery would let - a client wedged for the full WS write timeout hold the frame back from - every healthy client behind it in the list; this way each peer's stall - is its own. Two cases stay inline: - - * exactly one peer, so a single-client session is untouched; - * a caller already running on an event loop. ``WSTransport.write`` - fires and forgets when it sees its own loop and BLOCKS on - ``fut.result`` when it does not, so moving a loop-thread write onto a - pool thread would make that thread wait on the loop it just left - while the loop waits on it — a full write-timeout freeze. - - Pruning is unchanged: a peer that returns ``False`` or raises is - detached, whichever path it took. A peer that has not answered by - ``_FANOUT_WRITE_DEADLINE_S`` is NOT pruned — slow is not dead — and - counts as delivered, which is what ``WSTransport.write`` itself reports - when its own write times out: the frame is queued on that peer's loop - and flushes when the loop breathes. If that peer really is gone, its - next write returns ``False`` promptly and prunes it then. - - The collect happens before the return, so frame A is on every peer that - answered within the deadline before frame B is dispatched to any of - them. A peer whose write was still queued on the pool when the deadline - passed is the exception: the next frame can reach that peer's own - ``_token_lock`` while the previous one is still waiting for it, and the - lock — not this method — decides the order they land in. Within a peer, - ordering is that peer's own business (``WSTransport`` queues under its - token lock, ``StdioTransport`` writes under its stream lock); - ``server.write_json`` takes no lock, so concurrent emitters interleave - here exactly as they already did. - - Returns ``True`` when at least one client accepted the frame or still - has it in flight — an all-dead fan-out reports peer-gone exactly like a - single dead transport would. - """ - targets = self.transports() - if not targets: - return False - - dead: list["Transport"] = [] - delivered = False - - if len(targets) == 1 or _caller_is_on_event_loop(): - for transport in targets: - if self._deliver(transport, obj, dead): - delivered = True - for transport in dead: - self.detach(transport) - return delivered - - deadline = time.monotonic() + _FANOUT_WRITE_DEADLINE_S - pool = _fanout_write_pool() - dispatched: list[tuple["Transport", "concurrent.futures.Future | None"]] = [] - for transport in targets[:-1]: + def _drain(self, peer: _FanoutPeer) -> None: + while True: + with self._lock: + if not peer.attached or not peer.pending: + peer.writing = False + if not peer.attached: + self._remove(peer) + return + generation = peer.generation + frame, size = peer.pending.popleft() + peer.pending_bytes -= size try: - dispatched.append((transport, pool.submit(transport.write, obj))) - except RuntimeError: - # Pool refused the work (interpreter shutting down). Deliver - # inline below rather than dropping the frame. - dispatched.append((transport, None)) - - # The last peer runs here: one fewer worker per write, and a two-client - # session costs a single pool thread. - if self._deliver(targets[-1], obj, dead): - delivered = True - - # One deadline for the whole batch, then read only the futures that - # finished. Waiting first and calling ``result()`` on a settled future - # keeps our deadline distinguishable from a peer that raised: on 3.11+ - # ``concurrent.futures.TimeoutError`` IS the builtin ``TimeoutError``, - # so a ``result(timeout=...)`` cannot tell the two apart. - futures = [fut for _, fut in dispatched if fut is not None] - if futures: - concurrent.futures.wait( - futures, timeout=max(0.0, deadline - time.monotonic()) - ) - - for transport, fut in dispatched: - if fut is None: - if self._deliver(transport, obj, dead): - delivered = True - continue - if not fut.done(): - # Slow, not dead: leave it attached and count the frame as in - # flight. Cancelling would not stop a write already running. - logger.warning( - "fanout write still pending after %ss; peer left attached", - _FANOUT_WRITE_DEADLINE_S, - ) - delivered = True - continue - try: - ok = fut.result() + from tui_gateway.ws import WSTransport + if isinstance(peer.transport, WSTransport): + # write() acknowledges buffered tokens/timeouts, not socket + # progress. Await the real send so WS cannot move an + # unbounded backlog underneath this bounded mailbox. + from agent.async_utils import safe_schedule_threadsafe + future = safe_schedule_threadsafe( + peer.transport.write_async(frame), peer.transport._loop) + ok = future is not None and future.result() + else: + ok = peer.transport.write(frame) except Exception: logger.debug("fanout write failed; pruning peer", exc_info=True) - dead.append(transport) - continue - if ok: - delivered = True - else: - dead.append(transport) + ok = False + if not ok: + with self._lock: + if peer.generation != generation: + continue + peer.writing = False + self._remove(peer) + return - for transport in dead: - self.detach(transport) - return delivered + def write(self, obj: dict) -> bool: + # Freeze the queued frame so a caller cannot mutate it after admission. + encoded = json.dumps(obj, ensure_ascii=False) + size = len(encoded.encode("utf-8", errors="surrogatepass")) + frame = json.loads(encoded) + with self._lock: + for peer in list(self._peers): + if not peer.attached: + continue + if (len(peer.pending) >= self._MAX_PENDING_FRAMES + or peer.pending_bytes + size > self._MAX_PENDING_BYTES): + logger.warning("fanout subscriber backlog full; detaching peer") + self._remove(peer) + continue + peer.pending.append((frame, size)) + peer.pending_bytes += size + if not peer.writing: + peer.writing = True + threading.Thread(target=self._drain, args=(peer,), + name="tui-fanout", daemon=True).start() + return any(peer.attached for peer in self._peers) def close(self) -> None: - """Detach every peer. Does NOT close them — each owns its own socket.""" + """Detach without closing sockets owned by the connection handlers.""" with self._lock: - self._transports = [] + for peer in list(self._peers): + self._remove(peer) class TeeTransport: