Merge pull request #5270 from cheahjs/fix/websocket-take-2

fix: socket.io connections failing when websockets are not available

backend/open_webui/main.py

@@ -812,6 +812,24 @@ async def update_embedding_function(request: Request, call_next):
     return response
 
 
+@app.middleware("http")
+async def inspect_websocket(request: Request, call_next):
+    if (
+        "/ws/socket.io" in request.url.path
+        and request.query_params.get("transport") == "websocket"
+    ):
+        upgrade = (request.headers.get("Upgrade") or "").lower()
+        connection = (request.headers.get("Connection") or "").lower().split(",")
+        # Check that there's the correct headers for an upgrade, else reject the connection
+        # This is to work around this upstream issue: https://github.com/miguelgrinberg/python-engineio/issues/367
+        if upgrade != "websocket" or "upgrade" not in connection:
+            return JSONResponse(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                content={"detail": "Invalid WebSocket upgrade request"},
+            )
+    return await call_next(request)
+
+
 app.mount("/ws", socket_app)
 
 app.mount("/ollama", ollama_app)