feat(ws): implement full duplex bi-directional WebSocket proxy bridge supporting subdomains and path-based routing

app/proxy.py

@@ -1,8 +1,11 @@
 import httpx
 import re
+import asyncio
 from typing import AsyncGenerator, Dict, Any, Optional, List, Tuple
-from fastapi import Request, Response
+from fastapi import Request, Response, WebSocket, WebSocketDisconnect
 from fastapi.responses import StreamingResponse
+import websockets
+from websockets.exceptions import ConnectionClosed
 import logging
 
 logger = logging.getLogger("hermes_hub.proxy")
@@ -19,6 +22,16 @@ HOP_BY_HOP_HEADERS = {
     "content-length"
 }
 
+WS_HOP_BY_HOP_HEADERS = {
+    "connection",
+    "upgrade",
+    "sec-websocket-key",
+    "sec-websocket-version",
+    "sec-websocket-extensions",
+    "sec-websocket-accept",
+    "host"
+}
+
 _http_transport: Optional[httpx.AsyncHTTPTransport] = None
 
 def get_http_transport() -> httpx.AsyncHTTPTransport:
@@ -158,3 +171,88 @@ async def proxy_request(
             status_code=500,
             media_type="text/plain"
         )
+
+async def proxy_websocket(
+    client_ws: WebSocket,
+    backend_url: str,
+    path: str,
+    universe_id: Optional[str] = None
+):
+    """
+    Bi-directional full duplex WebSocket proxy bridge.
+    Forwards connection from client browser to backend Hermes universe instance.
+    """
+    ws_base = backend_url.replace("https://", "wss://").replace("http://", "ws://").rstrip("/")
+    sub_path = path.lstrip("/")
+    upstream_url = f"{ws_base}/{sub_path}" if sub_path else ws_base
+    
+    if client_ws.url.query:
+        upstream_url = f"{upstream_url}?{client_ws.url.query}"
+        
+    upstream_headers = {}
+    for k, v in client_ws.headers.items():
+        if k.lower() not in WS_HOP_BY_HOP_HEADERS:
+            upstream_headers[k] = v
+            
+    upstream_headers["x-forwarded-host"] = client_ws.headers.get("host", "")
+    upstream_headers["x-forwarded-proto"] = client_ws.url.scheme or "http"
+    if client_ws.client:
+        upstream_headers["x-forwarded-for"] = client_ws.client.host
+        
+    # Check subprotocols
+    subprotocols_raw = client_ws.headers.get("sec-websocket-protocol", "")
+    subprotocols = [s.strip() for s in subprotocols_raw.split(",") if s.strip()] or None
+    
+    try:
+        async with websockets.connect(
+            upstream_url,
+            extra_headers=upstream_headers,
+            subprotocols=subprotocols,
+            ping_interval=20,
+            ping_timeout=20,
+            max_size=10 * 1024 * 1024
+        ) as upstream_ws:
+            # Accept client connection with negotiated subprotocol
+            await client_ws.accept(subprotocol=upstream_ws.subprotocol)
+            
+            async def client_to_upstream():
+                try:
+                    while True:
+                        msg = await client_ws.receive()
+                        if "text" in msg:
+                            await upstream_ws.send(msg["text"])
+                        elif "bytes" in msg:
+                            await upstream_ws.send(msg["bytes"])
+                        elif msg.get("type") == "websocket.disconnect":
+                            break
+                except (WebSocketDisconnect, ConnectionClosed, asyncio.CancelledError):
+                    pass
+                    
+            async def upstream_to_client():
+                try:
+                    async for msg in upstream_ws:
+                        if isinstance(msg, str):
+                            await client_ws.send_text(msg)
+                        elif isinstance(msg, bytes):
+                            await client_ws.send_bytes(msg)
+                except (WebSocketDisconnect, ConnectionClosed, asyncio.CancelledError):
+                    pass
+                    
+            done, pending = await asyncio.wait(
+                [
+                    asyncio.create_task(client_to_upstream()),
+                    asyncio.create_task(upstream_to_client())
+                ],
+                return_when=asyncio.FIRST_COMPLETED
+            )
+            for task in pending:
+                task.cancel()
+                
+    except (WebSocketDisconnect, ConnectionClosed):
+        pass
+    except Exception as e:
+        logger.error(f"WebSocket proxy bridge error for {upstream_url}: {e}")
+        try:
+            await client_ws.close(code=1011)
+        except Exception:
+            pass