Recover from a btcdata restart without reloading the page

htmx's WebSocket extension reconnects only after close codes 1006/1011/1012/1013.
_ws_proxy closed the browser socket with a normal 1000 on every upstream exit, so
htmx read a btcdata restart as a deliberate shutdown and left both panels frozen
until the visitor reloaded. New browser connections always worked, which hid the
defect from non-browser checks.

The proxy now reconnects upstream behind the still-open browser socket on a bounded
schedule, resets that schedule after a connection that lasted long enough to count
as healthy, and closes with 1012 only once the schedule runs out, so recovery falls
back to htmx's own backoff. A concurrent reader detects a departed browser, which a
dead upstream would otherwise be indistinguishable from, and an unrenderable payload
is logged and skipped rather than ending the stream.

tests/test_ws_proxy.py covers this against a real loopback WebSocket server: same
socket recovery across an upstream restart for both streams, the retryable close
code when the upstream stays down, and skip-on-unrenderable.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

app/main.py

@@ -1,6 +1,8 @@
+import asyncio
 import json
 import logging
 import os
+import time
 from datetime import datetime
 
 import httpx
@@ -367,36 +369,112 @@ def strike_json_to_html(text: str) -> str:
 
 # --- WebSocket proxy ---
 
+# htmx's WebSocket extension reconnects a browser socket only after an abnormal
+# close (1006/1011/1012/1013). A normal 1000 close means "the server is done",
+# so the extension never reconnects and the panel stays frozen until the visitor
+# reloads. A btcdata restart must therefore never reach the browser as a normal
+# close: the proxy reconnects upstream behind the still-open browser socket, and
+# only if that keeps failing does it close with 1012 (Service Restart) and hand
+# recovery to the page's own reconnect.
+BROWSER_RETRY_CLOSE_CODE = 1012
+UPSTREAM_OPEN_TIMEOUT_SECONDS = 5.0
+UPSTREAM_RETRY_DELAYS_SECONDS = (0.5, 1.0, 2.0, 4.0, 8.0, 8.0)
+UPSTREAM_STABLE_SECONDS = 30.0
+
 
 def _btcdata_ws_url(path: str) -> str:
     return BTCDATA_URL.replace("http", "ws", 1) + path
 
 
-async def _ws_proxy(browser_ws: WebSocket, upstream_path: str, transform):
+async def _await_browser_close(browser_ws: WebSocket) -> None:
+    """Return as soon as the browser side goes away.
+
+    The browser sends nothing, but reading is the only way to notice that it
+    left while the proxy is waiting on a dead upstream rather than sending.
+    """
+    try:
+        while True:
+            if (await browser_ws.receive())["type"] == "websocket.disconnect":
+                return
+    except (WebSocketDisconnect, RuntimeError):
+        return
+
+
+async def _forward_upstream(browser_ws: WebSocket, upstream_url: str, transform) -> None:
+    """Forward one upstream connection to the browser until that connection ends."""
+    async with websockets.connect(
+        upstream_url, ping_interval=None, open_timeout=UPSTREAM_OPEN_TIMEOUT_SECONDS
+    ) as upstream:
+        unrenderable = 0
+        async for message in upstream:
+            if not message:
+                continue
+            try:
+                html = transform(message)
+            except Exception:
+                # One unexpected payload must not take down a working page. Log
+                # the first per connection, then stay quiet about the rest.
+                unrenderable += 1
+                logger.log(
+                    logging.WARNING if unrenderable == 1 else logging.DEBUG,
+                    "Discarding an upstream message that could not be rendered",
+                    exc_info=True,
+                )
+                continue
+            await browser_ws.send_text(html)
+
+
+async def _ws_proxy(browser_ws: WebSocket, upstream_path: str, transform) -> None:
     await browser_ws.accept()
     upstream_url = _btcdata_ws_url(upstream_path)
     logger.info("Browser connected, proxying to %s", upstream_url)
 
+    browser_closed = asyncio.create_task(_await_browser_close(browser_ws))
+    attempt = 0
     try:
-        async with websockets.connect(upstream_url, ping_interval=None) as upstream:
-            async for message in upstream:
-                if not message:
-                    continue
-                try:
-                    html = transform(message)
-                    await browser_ws.send_text(html)
-                except WebSocketDisconnect:
-                    break
-                except Exception:
-                    logger.debug("Error forwarding to browser", exc_info=True)
-                    break
-    except WebSocketDisconnect:
-        pass
-    except Exception:
-        logger.debug("Upstream connection error", exc_info=True)
+        while True:
+            started = time.monotonic()
+            forwarding = asyncio.create_task(
+                _forward_upstream(browser_ws, upstream_url, transform)
+            )
+            done, _ = await asyncio.wait(
+                {forwarding, browser_closed}, return_when=asyncio.FIRST_COMPLETED
+            )
+            if browser_closed in done:
+                # The browser left; a healthy upstream would otherwise keep this
+                # task forwarding into a socket nobody is reading.
+                forwarding.cancel()
+                await asyncio.gather(forwarding, return_exceptions=True)
+                logger.info("Browser disconnected from %s", upstream_url)
+                return
+            try:
+                forwarding.result()
+                logger.info("Upstream %s closed; reconnecting", upstream_url)
+            except Exception:
+                logger.info("Upstream %s unavailable; reconnecting", upstream_url)
+                logger.debug("Upstream connection error", exc_info=True)
+            if time.monotonic() - started >= UPSTREAM_STABLE_SECONDS:
+                # A connection that worked for a while means a restart, not a
+                # persistently broken upstream: start the schedule over.
+                attempt = 0
+            if attempt >= len(UPSTREAM_RETRY_DELAYS_SECONDS):
+                logger.warning(
+                    "Upstream %s still unavailable; asking the browser to reconnect",
+                    upstream_url,
+                )
+                return
+            await asyncio.wait({browser_closed}, timeout=UPSTREAM_RETRY_DELAYS_SECONDS[attempt])
+            attempt += 1
+            if browser_closed.done():
+                logger.info("Browser disconnected from %s", upstream_url)
+                return
     finally:
+        browser_closed.cancel()
+        await asyncio.gather(browser_closed, return_exceptions=True)
         try:
-            await browser_ws.close()
+            # Every server-side exit, expected or not, closes with a code the
+            # page retries; a browser that already left ignores it.
+            await browser_ws.close(code=BROWSER_RETRY_CLOSE_CODE)
         except Exception:
             pass