Newer
Older
navi-1 / clients / terminal / tui / ws_bridge.py
"""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

import asyncio
import json
from pathlib import Path

from textual.app import App

from clients.terminal.tui.events import ConnectionStatusChanged, WsEvent
from clients.terminal.ws_client import NaviWebSocketClient


class WsBridge:
    """Wraps NaviWebSocketClient and forwards events to a Textual App.

    ``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
        # ``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:
        return self._client

    @property
    def connected(self) -> bool:
        return self._connected

    async def start(self) -> None:
        """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:
        """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()
        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 _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:
                    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))