From 1c852f26f88f1e99424e77818884d8d91d02f6f6 Mon Sep 17 00:00:00 2001
From: "Khoa (Revenovich) Tran Gia" <khoa.tran-gia@ubisoft.com>
Date: Tue, 3 Mar 2026 10:56:41 +0700
Subject: [PATCH] fix: prevent WebSocket AssertionError when StaticFiles
 catch-all intercepts WS

In Starlette 0.52+, Mount('/') returns Match.FULL for every WebSocket
scope. If APIWebSocketRoute('/ws') is somehow not matched first, the
StaticFiles mount catches the connection and crashes with:
  assert scope["type"] == "http"  # AssertionError

Two-layer fix:
- _SPAStaticFiles.__call__: gracefully close non-HTTP connections with
  WebSocketClose() and log a warning with the actual path/type so the
  routing issue can be diagnosed.
- app.add_websocket_route('/ws', websocket_endpoint): belt-and-suspenders
  registration using Starlette's base WebSocketRoute (simpler than
  FastAPI's APIWebSocketRoute) right before the StaticFiles mount. If
  include_router's APIWebSocketRoute doesn't match, this fallback will.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 web/app.py | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/web/app.py b/web/app.py
index 7bde104..107fafa 100644
--- a/web/app.py
+++ b/web/app.py
@@ -81,6 +81,24 @@ class _SPAStaticFiles(StaticFiles):
     matching file, so client-side routes like /generate work on refresh.
     """
 
+    async def __call__(self, scope, receive, send) -> None:
+        # StaticFiles only handles HTTP. In Starlette 0.52+, Mount('/') returns
+        # Match.FULL for ALL WebSocket scopes, so a WebSocket connection can
+        # reach here if the dedicated /ws route somehow doesn't match first.
+        # Close gracefully instead of asserting.
+        if scope.get("type") != "http":
+            from starlette.websockets import WebSocketClose
+            logger.warning(
+                "WebSocket or non-HTTP scope reached StaticFiles "
+                "(path=%r, type=%r) — closing gracefully. "
+                "This indicates a routing issue; /ws route did not match.",
+                scope.get("path"),
+                scope.get("type"),
+            )
+            await WebSocketClose()(scope, receive, send)
+            return
+        await super().__call__(scope, receive, send)
+
     async def get_response(self, path: str, scope):
         try:
             return await super().get_response(path, scope)
@@ -152,6 +170,16 @@ def create_app() -> FastAPI:
     app.include_router(workflow_router, prefix="/api/workflow", tags=["workflow"])
     app.include_router(ws_router, tags=["ws"])
 
+    # Belt-and-suspenders: register /ws directly on the app so it sits at the
+    # top of the route list and is guaranteed to be checked before the
+    # StaticFiles catch-all mount below.  In Starlette 0.52+, Mount('/')
+    # returns Match.FULL for every WebSocket scope, which means if the route
+    # from include_router is somehow not matched first, the mount wins and
+    # StaticFiles crashes with AssertionError.  Having the route registered
+    # twice is harmless — the first Match.FULL in the route list wins.
+    from web.routers.ws_router import websocket_endpoint as _ws_endpoint
+    app.add_websocket_route("/ws", _ws_endpoint)
+
     # Serve frontend static files (if built)
     if _WEB_STATIC.exists() and any(_WEB_STATIC.iterdir()):
         app.mount("/", _SPAStaticFiles(directory=str(_WEB_STATIC), html=True), name="static")