fix: add ASGI middleware to prevent WS reaching StaticFiles

_WSInterceptMiddleware intercepts /ws WebSocket scopes at the outermost
ASGI layer before Starlette routing is consulted, so Mount('/') can never
hand a WS connection to _SPAStaticFiles regardless of route ordering.

Also downgrade the _SPAStaticFiles non-HTTP fallback log from WARNING to
DEBUG — the graceful close still fires as a safety net, but no longer
spams the log since the middleware handles the normal case.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Khoa (Revenovich) Tran Gia
2026-03-03 11:19:30 +07:00
parent 1c852f26f8
commit e83703aff0

View File

@@ -73,6 +73,26 @@ class _SecurityHeadersMiddleware(BaseHTTPMiddleware):
response.headers["Content-Security-Policy"] = self._CSP
return response
class _WSInterceptMiddleware:
"""Short-circuit /ws WebSocket connections before Starlette routing.
In Starlette 0.52+, Mount('/') returns Match.FULL for every WebSocket
scope, so a WS connection can reach _SPAStaticFiles if the dedicated /ws
WebSocketRoute is somehow not matched first. Wrapping the whole ASGI app
here guarantees /ws is handled before any route or mount is consulted.
"""
def __init__(self, app, ws_handler) -> None:
self._app = app
self._ws = ws_handler
async def __call__(self, scope, receive, send) -> None:
if scope["type"] == "websocket" and scope.get("path") == "/ws":
await self._ws(scope, receive, send)
else:
await self._app(scope, receive, send)
class _SPAStaticFiles(StaticFiles):
"""StaticFiles with SPA fallback: serve index.html for unknown paths.
@@ -83,15 +103,13 @@ class _SPAStaticFiles(StaticFiles):
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.
# Match.FULL for ALL WebSocket scopes; _WSInterceptMiddleware should
# prevent any WS connection from reaching here, but close gracefully
# as a last-resort safety net rather than raising AssertionError.
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.",
logger.debug(
"non-HTTP scope reached StaticFiles (path=%r, type=%r) — closing gracefully",
scope.get("path"),
scope.get("type"),
)
@@ -170,14 +188,11 @@ def create_app() -> FastAPI:
app.include_router(workflow_router, prefix="/api/workflow", tags=["workflow"])
app.include_router(ws_router, tags=["ws"])
from web.routers.ws_router import websocket_endpoint as _ws_endpoint
# 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
# StaticFiles catch-all mount below.
app.add_websocket_route("/ws", _ws_endpoint)
# Serve frontend static files (if built)
@@ -185,6 +200,12 @@ def create_app() -> FastAPI:
app.mount("/", _SPAStaticFiles(directory=str(_WEB_STATIC), html=True), name="static")
logger.info("Serving frontend from %s", _WEB_STATIC)
# Wrap the entire ASGI app so /ws WebSocket connections are intercepted
# before Starlette routing can hand them to the StaticFiles Mount('/').
# add_middleware places this as the outermost layer — it runs before any
# route or mount is consulted.
app.add_middleware(_WSInterceptMiddleware, ws_handler=_ws_endpoint)
@app.on_event("startup")
async def _startup():
asyncio.create_task(_status_ticker())