UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-01-18 00:30:51 +00:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

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

Browse Source 
fix: socket.io connections failing when websockets are not available
This commit is contained in:
Timothy Jaeryang Baek 2024-09-09 23:19:05 +01:00 committed by GitHub
commit 1023ff8454
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 33 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
backend/open_webui/apps/socket/main.py
View File

@ -2,9 +2,16 @@ import asyncio
import socketio
from open_webui.apps.webui.models.users import Users
from open_webui.env import ENABLE_WEBSOCKET_SUPPORT
from open_webui.utils.utils import decode_token
sio = socketio.AsyncServer(cors_allowed_origins=[], async_mode="asgi")
sio = socketio.AsyncServer(
cors_allowed_origins=[],
async_mode="asgi",
transports=(["polling", "websocket"] if ENABLE_WEBSOCKET_SUPPORT else ["polling"]),
allow_upgrades=ENABLE_WEBSOCKET_SUPPORT,
always_connect=True,
)
app = socketio.ASGIApp(sio, socketio_path="/ws/socket.io")
# Dictionary to maintain the user pool

4
backend/open_webui/env.py
View File

@ -273,3 +273,7 @@ WEBUI_SESSION_COOKIE_SECURE = os.environ.get(
if WEBUI_AUTH and WEBUI_SECRET_KEY == "":
raise ValueError(ERROR_MESSAGES.ENV_VAR_NOT_FOUND)
ENABLE_WEBSOCKET_SUPPORT = (
os.environ.get("ENABLE_WEBSOCKET_SUPPORT", "True").lower() == "true"
)

18
backend/open_webui/main.py
View File

@ -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)

13
src/routes/+layout.svelte
View File

@ -38,27 +38,20 @@
let loaded = false;
const BREAKPOINT = 768;
const setupSocket = (websocket = true) => {
const setupSocket = () => {
const _socket = io(`${WEBUI_BASE_URL}` || undefined, {
reconnection: true,
reconnectionDelay: 1000,
reconnectionDelayMax: 5000,
randomizationFactor: 0.5,
path: '/ws/socket.io',
auth: { token: localStorage.token },
transports: websocket ? ['websocket'] : ['polling']
auth: { token: localStorage.token }
});
socket.set(_socket);
_socket.on('connect_error', (err) => {
if (err.message.includes('websocket')) {
console.log('WebSocket connection failed, falling back to polling');
_socket.close();
setupSocket(false);
} else {
console.log('connect_error', err);
}
console.log('connect_error', err);
});
_socket.on('connect', () => {