diff --git a/clients/terminal/tui/ws_bridge.py b/clients/terminal/tui/ws_bridge.py index 455bf6b..5baa12c 100644 --- a/clients/terminal/tui/ws_bridge.py +++ b/clients/terminal/tui/ws_bridge.py @@ -1,4 +1,12 @@ -"""Bridge between NaviWebSocketClient and Textual event loop.""" +"""Bridge between NaviWebSocketClient and Textual event loop. + +Owns the WebSocket lifecycle for the TUI: connect, forward events to the app, +and — when the connection drops — reconnect with exponential backoff until the +app shuts down. A dropped server socket no longer leaves the client stranded; +the status panel flips to ``reconnecting…`` and the next message the user sends +is held in the input queue until a fresh socket is live, so a brief blip +doesn't eat their input. +""" from __future__ import annotations @@ -13,15 +21,38 @@ class WsBridge: - """Wraps NaviWebSocketClient and forwards events to a Textual App.""" + """Wraps NaviWebSocketClient and forwards events to a Textual App. - def __init__(self, app: App, session_id: str, cwd: Path | None = None) -> None: + ``start`` launches a supervisor that connects and runs the receive loop; + when the socket dies it backs off and reconnects indefinitely. ``stop`` + cancels everything promptly (it's the only thing that ends the supervisor — + a dropped connection just triggers another reconnect attempt). + """ + + def __init__( + self, + app: App, + session_id: str, + cwd: Path | None = None, + *, + client: NaviWebSocketClient | None = None, + ) -> None: self.app = app self.session_id = session_id - self._client = NaviWebSocketClient(session_id, cwd=cwd) - self._receive_task: asyncio.Task | None = None + # ``client=`` lets tests inject a fake without touching the network; + # production leaves it None and we build the real client. + self._client = client if client is not None else NaviWebSocketClient(session_id, cwd=cwd) + self._supervise_task: asyncio.Task | None = None self._input_task: asyncio.Task | None = None self._connected = False + self._stopping = False + # Set while a live socket exists; the input loop waits on this so sends + # don't fire into a dead socket. Cleared on drop and on stop. + self._connected_event: asyncio.Event | None = None + # Exponential-backoff state. Reset to 0 on a successful connect. + self._attempt = 0 + self._backoff_min = 1.0 + self._backoff_max = 30.0 @property def client(self) -> NaviWebSocketClient: @@ -32,53 +63,142 @@ return self._connected async def start(self) -> None: - try: - await self._client.connect() - self._connected = True - self._notify(True, "") - self._receive_task = asyncio.create_task(self._receive_loop()) - self._input_task = asyncio.create_task(self._client.input_loop()) - except Exception as exc: - self._connected = False - self._notify(False, str(exc)) + """Launch the supervisor + input loop. The supervisor does the first + connect (and retries on failure), so ``start`` itself never raises — + connection errors surface as ``ConnectionStatusChanged`` events.""" + self._stopping = False + self._connected_event = asyncio.Event() + self._supervise_task = asyncio.create_task(self._supervise()) + self._input_task = asyncio.create_task(self._input_loop()) async def stop(self) -> None: - if self._receive_task: - self._receive_task.cancel() - try: - await self._receive_task - except asyncio.CancelledError: - pass - self._receive_task = None - if self._input_task: - self._input_task.cancel() - try: - await self._input_task - except asyncio.CancelledError: - pass - self._input_task = None + """Shut the bridge down: stop reconnecting, cancel the loops, close + the socket. Idempotent — safe to call from ``on_unmount`` and again on + an ``attach_session`` switch.""" + self._stopping = True + # Unblock the input loop if it's waiting on a live connection. + if self._connected_event is not None: + self._connected_event.set() + # Sentinel so the input loop's queue.get() returns promptly. self._client.stop_input() - await self._client.close() + for task in (self._supervise_task, self._input_task): + if task is not None and not task.done(): + task.cancel() + for task in (self._supervise_task, self._input_task): + if task is not None: + try: + await task + except (asyncio.CancelledError, Exception): + pass + self._supervise_task = None + self._input_task = None + try: + await self._client.close() + except Exception: + pass self._connected = False + if self._connected_event is not None: + self._connected_event.clear() self._notify(False, "") - async def _receive_loop(self) -> None: - try: - ws = self._client._ws - if ws is None: - return - async for raw in ws: + async def _supervise(self) -> None: + """Connect → receive until the socket closes → backoff → repeat. + + The only exit is ``stop()`` setting ``_stopping`` (or cancelling this + task). A dropped connection or a failed connect is just a reason to + back off and try again. + """ + while not self._stopping: + try: + await self._client.connect() + except Exception as exc: + if self._stopping: + break + self._set_connected(False) + self._notify(False, f"reconnecting ({self._attempt + 1})" + (f": {exc}" if str(exc) else "")) + await self._backoff() + continue + + # Connected — reset backoff and let the input loop send. + self._attempt = 0 + self._set_connected(True) + self._notify(True, "online") + + try: + await self._receive_until_closed() + except Exception: + pass + + if self._stopping: + break + self._set_connected(False) + self._notify(False, "reconnecting…") + # Drop the dead socket before reconnecting so connect() is clean. + try: + await self._client.close() + except Exception: + pass + await self._backoff() + + async def _receive_until_closed(self) -> None: + ws = self._client._ws + if ws is None: + return + async for raw in ws: + try: + msg = json.loads(raw) + except json.JSONDecodeError: + continue + self._post(WsEvent(msg)) + + async def _input_loop(self) -> None: + """Drain the client's input queue into the live socket. + + Waits for ``_connected_event`` before each send, so a send never fires + into a dead socket. If a send fails mid-flight (the socket died between + the connect check and the write), the message goes back on the queue + and the dead socket is closed — the supervisor reconnects, the event + is set again, and the message is delivered. Nothing is lost across a + brief drop. + """ + queue = self._client._input_queue + while not self._stopping: + content = await queue.get() + if content is None: + break + if self._connected_event is None: + continue + await self._connected_event.wait() + if self._stopping: + break + try: + await self._client.send(content) + except Exception: + # Socket died under us — requeue and force the supervisor to + # reconnect by closing the dead socket (which ends the receive + # loop and triggers a fresh connect). + queue.put_nowait(content) + self._connected_event.clear() try: - msg = json.loads(raw) - except json.JSONDecodeError: - continue - self._post(WsEvent(msg)) - except Exception: - self._connected = False - self._notify(False, "disconnected") + await self._client.close() + except Exception: + pass + + async def _backoff(self) -> None: + delay = min(self._backoff_max, self._backoff_min * (2 ** self._attempt)) + self._attempt += 1 + await asyncio.sleep(delay) + + def _set_connected(self, connected: bool) -> None: + self._connected = connected + if self._connected_event is not None: + if connected: + self._connected_event.set() + else: + self._connected_event.clear() def _post(self, message) -> None: self.app.post_message(message) def _notify(self, connected: bool, detail: str) -> None: - self._post(ConnectionStatusChanged(connected, detail)) + self._post(ConnectionStatusChanged(connected, detail)) \ No newline at end of file diff --git a/tests/clients/test_ws_bridge.py b/tests/clients/test_ws_bridge.py new file mode 100644 index 0000000..8766f38 --- /dev/null +++ b/tests/clients/test_ws_bridge.py @@ -0,0 +1,253 @@ +"""Tests for WsBridge reconnect behaviour. + +Uses a fake WebSocket client (no network) to drive the supervisor: a socket +delivers a few frames then drops, and the bridge must back off and reconnect, +continuing to forward events and not lose the user's queued input. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from clients.terminal.tui.events import ConnectionStatusChanged, WsEvent +from clients.terminal.tui.ws_bridge import WsBridge + + +class _FakeApp: + """Records every posted message — WsBridge only needs ``post_message``.""" + + def __init__(self) -> None: + self.messages: list = [] + + def post_message(self, message) -> None: + self.messages.append(message) + + +class _FakeWs: + """Async iterator standing in for a real socket. + + ``frames`` are delivered one per ``__anext__``; when they run out the + socket ``close``s (the supervisor treats that as a drop and reconnects). + ``idle=True`` makes the socket stay open with no incoming data — it blocks + until ``close()`` — so a connect can hold ``connected`` True long enough + for the input loop to send. + """ + + def __init__(self, frames: list[str] | None = None, *, idle: bool = False, send_fails: bool = False) -> None: + self._frames = list(frames or []) + self.idle = idle + self.send_fails = send_fails + self.sent: list[str] = [] + self._closed = asyncio.Event() + + def __aiter__(self): + return self + + async def __anext__(self): + if self.idle: + # Open socket, no incoming data: block until closed. + await self._closed.wait() + raise StopAsyncIteration + if not self._frames: + raise StopAsyncIteration + return self._frames.pop(0) + + async def close(self) -> None: + self._closed.set() + + +class _FakeClient: + """Stand-in for NaviWebSocketClient: scripted connect() results, records + sends, exposes the same _input_queue / stop_input surface the bridge uses.""" + + def __init__(self) -> None: + self._ws: _FakeWs | None = None + self._input_queue: asyncio.Queue = asyncio.Queue() + self._ws_sequence: list[_FakeWs] = [] + self.connect_calls = 0 + self.sent: list[str] = [] + + def push_ws(self, ws: _FakeWs) -> None: + self._ws_sequence.append(ws) + + async def connect(self) -> None: + self.connect_calls += 1 + self._ws = self._ws_sequence.pop(0) if self._ws_sequence else _FakeWs([]) + + async def close(self) -> None: + if self._ws is not None: + await self._ws.close() + self._ws = None + + async def send(self, content: str) -> None: + if self._ws is None: + raise RuntimeError("not connected") + if self._ws.send_fails: + raise RuntimeError("socket dead") + self._ws.sent.append(content) + self.sent.append(content) + + def stop_input(self) -> None: + self._input_queue.put_nowait(None) + + def enqueue(self, content: str) -> None: + self._input_queue.put_nowait(content) + + +def _bridge(client: _FakeClient) -> WsBridge: + """A bridge with a tiny backoff so reconnects happen within the test.""" + bridge = WsBridge(_FakeApp(), "sess", client=client) + bridge._backoff_min = 0.01 + bridge._backoff_max = 0.05 + return bridge + + +def _ws_payloads(bridge: WsBridge) -> list[dict]: + return [m.payload for m in bridge.app.messages if isinstance(m, WsEvent)] + + +def _status_details(bridge: WsBridge) -> list[str]: + return [m.detail for m in bridge.app.messages if isinstance(m, ConnectionStatusChanged)] + + +@pytest.mark.anyio +async def test_reconnects_after_drop_and_continues_forwarding_events() -> None: + """A socket delivers two frames then closes; the bridge backs off, + reconnects, and frames from the second socket still reach the app.""" + client = _FakeClient() + client.push_ws(_FakeWs(['{"type":"status","content":"a"}', '{"type":"status","content":"b"}'])) + client.push_ws(_FakeWs(['{"type":"status","content":"c"}'])) + bridge = _bridge(client) + + await bridge.start() + # connect1 → a, b → drop → backoff → connect2 → c → drop. 50ms is enough + # for two tiny backoffs. + await asyncio.sleep(0.05) + await bridge.stop() + + assert client.connect_calls >= 2, "supervisor must reconnect after a drop" + payloads = [p.get("content") for p in _ws_payloads(bridge)] + assert "a" in payloads and "b" in payloads and "c" in payloads + + +@pytest.mark.anyio +async def test_initial_connect_failure_keeps_retrying_until_success() -> None: + """If the server is down at startup, the supervisor keeps retrying and + eventually connects when the server comes back.""" + client = _FakeClient() + bridge = _bridge(client) + + # No ws pushed yet → connect() succeeds but yields an empty (immediately + # closing) socket, which also drops. To simulate a refused connection, + # make connect() fail for the first two attempts. + + original_connect = client.connect + calls = {"n": 0} + + async def flaky_connect(): + calls["n"] += 1 + if calls["n"] <= 2: + raise ConnectionRefusedError("server down") + await original_connect() + + client.connect = flaky_connect + client.push_ws(_FakeWs(['{"type":"status","content":"recovered"}'])) + + await bridge.start() + await asyncio.sleep(0.08) + await bridge.stop() + + assert calls["n"] >= 3, "must keep retrying past the initial refusals" + assert "recovered" in [p.get("content") for p in _ws_payloads(bridge)] + # The status panel saw the offline→online transition. + details = _status_details(bridge) + assert any("reconnecting" in d for d in details) + assert any(d == "online" for d in details) + + +@pytest.mark.anyio +async def test_queued_input_is_held_across_disconnect_then_sent_on_reconnect() -> None: + """A message enqueued while disconnected is not lost — it's sent once the + next socket is live.""" + client = _FakeClient() + # ws1 closes immediately (no frames) → drop; ws2 stays open and accepts sends. + client.push_ws(_FakeWs([])) + client.push_ws(_FakeWs(idle=True)) + bridge = _bridge(client) + # Long first backoff so the disconnected window is reliably hit. + bridge._backoff_min = 0.2 + bridge._backoff_max = 0.2 + + await bridge.start() + await asyncio.sleep(0.02) # ws1 connected then dropped; now in the 200ms backoff + assert bridge.connected is False # genuinely disconnected + client.enqueue("hello from the dark") + await asyncio.sleep(0.25) # backoff done → ws2 connects → input loop sends + await bridge.stop() + + assert "hello from the dark" in client.sent + + +@pytest.mark.anyio +async def test_send_failure_requeues_message_and_resent_after_reconnect() -> None: + """If send() fails mid-flight (socket died right after connect), the + message is requeued and delivered after the next reconnect.""" + client = _FakeClient() + # ws1: open but send always fails; ws2: open and accepts sends. + client.push_ws(_FakeWs(idle=True, send_fails=True)) + client.push_ws(_FakeWs(idle=True)) + bridge = _bridge(client) + + await bridge.start() + await asyncio.sleep(0.02) # ws1 connected (send will fail) + client.enqueue("survive the blip") + await asyncio.sleep(0.06) # send fails → requeue → close → reconnect → resend + await bridge.stop() + + assert "survive the blip" in client.sent + + +@pytest.mark.anyio +async def test_stop_cancels_backoff_promptly() -> None: + """Stopping during a backoff sleep ends the supervisor quickly — no hanging + on a long reconnect delay.""" + client = _FakeClient() + bridge = _bridge(client) + bridge._backoff_min = 10.0 # long backoff + bridge._backoff_max = 10.0 + + await bridge.start() + await asyncio.sleep(0.01) # first connect attempt (no ws → empty socket → drop → long backoff) + loop = asyncio.get_event_loop() + t0 = loop.time() + await asyncio.wait_for(bridge.stop(), timeout=2.0) + elapsed = loop.time() - t0 + assert elapsed < 2.0, "stop() must cancel the backoff, not wait for it" + + +@pytest.mark.anyio +async def test_stop_is_idempotent() -> None: + client = _FakeClient() + bridge = _bridge(client) + await bridge.start() + await asyncio.sleep(0.01) + await bridge.stop() + await bridge.stop() # must not raise + assert bridge.connected is False + + +@pytest.mark.anyio +async def test_supervisor_does_not_reconnect_after_stop() -> None: + """Once stopped, a dropped socket must not trigger further connects.""" + client = _FakeClient() + client.push_ws(_FakeWs([])) # drops immediately + bridge = _bridge(client) + + await bridge.start() + await asyncio.sleep(0.02) + await bridge.stop() + calls_after_stop = client.connect_calls + await asyncio.sleep(0.05) + assert client.connect_calls == calls_after_stop, "stop() must end reconnect attempts" \ No newline at end of file