From 38e2fc45d54d329c7410dea4c0b5b54ead4f7b7a Mon Sep 17 00:00:00 2001 From: Brian Krabach Date: Wed, 1 Apr 2026 06:56:21 -0700 Subject: [PATCH] fix: WS proxy checks ttyd liveness before accepting + reset counter on message MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause (4th iteration): The WS proxy called websocket.accept() before verifying ttyd was alive. The browser 'open' event fired immediately, resetting _reconnectAttempts to 0. Counter bounced 0→1→0→1 forever — the client-side /connect POST at >= 2 attempts never fired. Server fix: _ttyd_is_listening() TCP probe (socket.create_connection, <1ms) before websocket.accept(). If ttyd is dead, auto-spawns it from active_session in state (kill_ttyd → spawn_ttyd → sleep 0.8s) THEN accepts. Browser 'open' only fires when ttyd is confirmed reachable. Client fix: _reconnectAttempts reset moved from 'open' handler to 'message' handler. First data message proves ttyd is alive and relaying — not just that the proxy accepted. Belt-and-suspenders: client-side /connect fetch still fires at attempt >= 2 as a fallback. Tests added: - test_ttyd_is_listening_function_exists: function importable from main - test_ws_proxy_checks_ttyd_before_accepting: source inspection ensures _ttyd_is_listening() call appears before await websocket.accept() - test_ws_proxy_auto_spawns_ttyd_when_dead: verifies spawn_ttyd called with active_session when _ttyd_is_listening returns False - terminal.mjs: _reconnectAttempts = 0 NOT in open handler, IS in message --- muxplex/frontend/terminal.js | 12 +++- muxplex/frontend/tests/test_terminal.mjs | 51 ++++++++++++++++ muxplex/main.py | 48 ++++++++++++++- muxplex/tests/test_ws_proxy.py | 76 ++++++++++++++++++++++++ 4 files changed, 184 insertions(+), 3 deletions(-) diff --git a/muxplex/frontend/terminal.js b/muxplex/frontend/terminal.js index 11be612..38c94c0 100644 --- a/muxplex/frontend/terminal.js +++ b/muxplex/frontend/terminal.js @@ -71,7 +71,11 @@ function connectWebSocket(name, sourceUrl) { ws.addEventListener('open', function() { if (ws !== _ws) return; // stale connection — superseded by a newer one, ignore - _reconnectAttempts = 0; // reset backoff counter on successful connection + // NOTE: do NOT reset _reconnectAttempts here. The server-side proxy accepts + // the WS before confirming ttyd is alive (auto-spawning if needed), but the + // browser 'open' event fires as soon as the proxy accepts — not when ttyd + // is actually ready. Resetting here caused the 0→1→0→1 bounce. Instead, + // reset on first data message (proves ttyd is alive and relaying). if (reconnectOverlay) reconnectOverlay.classList.add('hidden'); // Step 1: TEXT frame auth handshake — ttyd checks AuthToken before starting PTY ws.send(JSON.stringify({ AuthToken: '' })); @@ -86,6 +90,12 @@ function connectWebSocket(name, sourceUrl) { ws.addEventListener('message', function(e) { if (ws !== _ws) return; // stale connection — superseded by a newer one, ignore if (!_term) return; + // First data message proves ttyd is alive and relaying — safe to reset counter. + // We deliberately do NOT reset in the 'open' handler: the server-side proxy + // accepts the browser WS before ttyd is fully confirmed alive, so 'open' + // firing alone doesn't mean data will flow. Resetting here prevents the + // 0→1→0→1 bounce that kept the reconnect loop from escalating to /connect. + if (_reconnectAttempts > 0) _reconnectAttempts = 0; if (e.data instanceof ArrayBuffer) { var msg = new Uint8Array(e.data); if (msg.length < 1) return; diff --git a/muxplex/frontend/tests/test_terminal.mjs b/muxplex/frontend/tests/test_terminal.mjs index 4216a96..618e700 100644 --- a/muxplex/frontend/tests/test_terminal.mjs +++ b/muxplex/frontend/tests/test_terminal.mjs @@ -831,6 +831,57 @@ test('terminal.js WebSocket reconnect awaits /connect before creating WS', () => assert.ok(connectFn.includes('return;'), 'connect() must return after fetch to prevent falling through to immediate WS creation'); }); +// --- Reconnect counter: must reset on message, not on open --- + +test('terminal.js resets _reconnectAttempts on first message, not on open', () => { + const source = fs.readFileSync(new URL('../terminal.js', import.meta.url), 'utf8'); + + // Find the open handler body (between "addEventListener('open'" and its closing "})") + const openStart = source.indexOf("addEventListener('open'"); + assert.ok(openStart !== -1, "must have an open handler"); + // Find the matching closing "})" for the open handler — walk from openStart + let depth = 0; + let openBodyEnd = -1; + for (let i = openStart; i < source.length - 1; i++) { + if (source[i] === '{') depth++; + else if (source[i] === '}') { + depth--; + if (depth === 0) { openBodyEnd = i; break; } + } + } + assert.ok(openBodyEnd !== -1, "must find the end of the open handler"); + const openBody = source.substring(openStart, openBodyEnd + 1); + + // _reconnectAttempts = 0 must NOT appear in the open handler + // (the proxy accepts before ttyd is alive, so open doesn't prove ttyd is up) + assert.ok( + !openBody.includes('_reconnectAttempts = 0'), + '_reconnectAttempts must NOT be reset in the open handler — ' + + 'the proxy accepts the WS before confirming ttyd is alive; ' + + 'reset must happen on first message (proves ttyd is sending data)', + ); + + // _reconnectAttempts reset must appear in the message handler instead + const msgStart = source.indexOf("addEventListener('message'"); + assert.ok(msgStart !== -1, "must have a message handler"); + let msgDepth = 0; + let msgBodyEnd = -1; + for (let i = msgStart; i < source.length - 1; i++) { + if (source[i] === '{') msgDepth++; + else if (source[i] === '}') { + msgDepth--; + if (msgDepth === 0) { msgBodyEnd = i; break; } + } + } + assert.ok(msgBodyEnd !== -1, "must find the end of the message handler"); + const msgBody = source.substring(msgStart, msgBodyEnd + 1); + assert.ok( + msgBody.includes('_reconnectAttempts'), + '_reconnectAttempts must be reset inside the message handler ' + + '(first data message proves ttyd is alive and relaying)', + ); +}); + // --- Issue 4: setTerminalFontSize --- test('terminal.js exposes window._setTerminalFontSize function', () => { diff --git a/muxplex/main.py b/muxplex/main.py index 4f2eb36..23b32a0 100644 --- a/muxplex/main.py +++ b/muxplex/main.py @@ -556,12 +556,36 @@ async def instance_info() -> dict: # --------------------------------------------------------------------------- +def _ttyd_is_listening() -> bool: + """Return True if something is accepting TCP connections on TTYD_PORT. + + Uses a raw socket connect (no WebSocket handshake, no PTY spawned). + Takes < 1 ms on localhost when ttyd is running; fails immediately with + ConnectionRefusedError when it's not. OSError/TimeoutError are also + caught so the caller always gets a bool. + """ + import socket as _sock + + try: + with _sock.create_connection(("127.0.0.1", TTYD_PORT), timeout=0.5): + return True + except (ConnectionRefusedError, OSError, TimeoutError): + return False + + @app.websocket("/terminal/ws") async def terminal_ws_proxy(websocket: WebSocket) -> None: """Proxy WebSocket frames between the browser and ttyd. - Accepts with subprotocol 'tty' (required by ttyd), then opens a connection - to ws://localhost:{TTYD_PORT}/ws and relays frames bidirectionally. + Checks that ttyd is alive BEFORE accepting the browser WebSocket. If ttyd + is not listening (e.g. after a service restart), auto-spawns it using the + active_session from state, then waits briefly for it to bind its port. + + Only after ttyd is confirmed reachable does the function call + websocket.accept() — so the browser's 'open' event only fires once a real + relay is possible. This prevents the reconnect-counter bounce bug where + the proxy accepted immediately (resetting _reconnectAttempts to 0) and + then closed as soon as it couldn't reach the dead ttyd. """ # Auth check before accepting — BaseHTTPMiddleware doesn't cover WebSocket scope host = websocket.client.host if websocket.client else "" @@ -573,6 +597,26 @@ async def terminal_ws_proxy(websocket: WebSocket) -> None: await websocket.close(code=4001) return + # Ensure ttyd is reachable BEFORE accepting the browser WS. + # After a service restart ttyd is dead but clients reconnect immediately. + # Auto-spawn from active_session so the browser's 'open' event only fires + # when a real relay is possible — eliminates the 0→1→0→1 counter bounce. + if not _ttyd_is_listening(): + try: + async with state_lock: + state = load_state() + session_name = state.get("active_session") + if session_name: + _log.info( + "WS proxy: ttyd not listening, auto-spawning for '%s'", + session_name, + ) + await kill_ttyd() + await spawn_ttyd(session_name) + await asyncio.sleep(0.8) # wait for ttyd to bind its port + except Exception as exc: + _log.warning("WS proxy: failed to auto-spawn ttyd: %s", exc) + await websocket.accept(subprotocol="tty") ttyd_url = f"ws://localhost:{TTYD_PORT}/ws" diff --git a/muxplex/tests/test_ws_proxy.py b/muxplex/tests/test_ws_proxy.py index 48177df..d9d6cf1 100644 --- a/muxplex/tests/test_ws_proxy.py +++ b/muxplex/tests/test_ws_proxy.py @@ -125,6 +125,82 @@ class FakeTtydWs: # --------------------------------------------------------------------------- +# --------------------------------------------------------------------------- +# Tests: ttyd liveness check before websocket.accept +# --------------------------------------------------------------------------- + + +def test_ttyd_is_listening_function_exists(): + """_ttyd_is_listening() must exist in main.py (TCP probe helper).""" + # Import will fail if function doesn't exist — that IS the failing test + from muxplex.main import _ttyd_is_listening # noqa: F401 + + assert callable(_ttyd_is_listening) + + +def test_ws_proxy_checks_ttyd_before_accepting(): + """terminal_ws_proxy must check _ttyd_is_listening BEFORE websocket.accept. + + Root cause of the reconnect loop: the proxy called websocket.accept() before + checking if ttyd was alive. The browser's 'open' event fired immediately, + resetting _reconnectAttempts to 0. The counter bounced 0→1→0→1 forever so + the client-side /connect POST (at >= 2 attempts) never fired. + + Fix: check _ttyd_is_listening() first. If not listening, auto-spawn ttyd + THEN accept — so the browser only gets 'open' when ttyd is actually ready. + """ + source = inspect.getsource(terminal_ws_proxy) + # Use "await websocket.accept" to avoid matching the docstring mention + accept_idx = source.index("await websocket.accept") + ttyd_check_idx = source.index("_ttyd_is_listening") + assert ttyd_check_idx < accept_idx, ( + "_ttyd_is_listening() must be checked BEFORE await websocket.accept() — " + "proxy must not accept the browser WS until ttyd is confirmed alive" + ) + + +def test_ws_proxy_auto_spawns_ttyd_when_dead(monkeypatch): + """WS proxy must call spawn_ttyd when _ttyd_is_listening returns False.""" + import asyncio + + spawn_calls = [] + + async def mock_spawn_ttyd(name: str): + spawn_calls.append(name) + + async def mock_kill_ttyd(): + pass + + async def mock_sleep(_delay: float): + pass # no-op so tests don't actually wait + + # Patch _ttyd_is_listening to report ttyd as dead + monkeypatch.setattr("muxplex.main._ttyd_is_listening", lambda: False) + # Patch spawn_ttyd / kill_ttyd so tests don't touch real processes + monkeypatch.setattr("muxplex.main.spawn_ttyd", mock_spawn_ttyd) + monkeypatch.setattr("muxplex.main.kill_ttyd", mock_kill_ttyd) + # asyncio.sleep is called after spawn — patch to be a no-op + monkeypatch.setattr(asyncio, "sleep", mock_sleep) + + # Provide a fake websockets.connect that immediately closes (no real ttyd) + fake_ws = FakeTtydWs(responses=[]) + monkeypatch.setattr("muxplex.main.websockets.connect", lambda *a, **kw: fake_ws) + + # Patch load_state to return state with active_session + monkeypatch.setattr( + "muxplex.main.load_state", + lambda: {"active_session": "test-session", "sessions": {}, "session_order": []}, + ) + + with _make_authed_client() as c: + with c.websocket_connect("/terminal/ws") as _: + pass + + assert spawn_calls == ["test-session"], ( + "spawn_ttyd must be called with active_session when ttyd is not listening" + ) + + def test_terminal_ws_proxy_does_not_use_receive_bytes(): """Regression: receive_bytes() silently drops TEXT frames (like the ttyd auth token).