From 4c81b6ebc5702a7d75206a45affb55f2938e7cd2 Mon Sep 17 00:00:00 2001 From: shcodenick Date: Wed, 30 Apr 2025 15:05:57 +0200 Subject: [PATCH] hiding websocket initalization behind authentication --- backend/open_webui/routers/auths.py | 2 +- backend/open_webui/socket/main.py | 3 ++ src/lib/utils/websocket.ts | 56 +++++++++++++++++++++++++++++ src/routes/+layout.svelte | 51 ++------------------------ src/routes/auth/+page.svelte | 4 +++ 5 files changed, 66 insertions(+), 50 deletions(-) create mode 100644 src/lib/utils/websocket.ts diff --git a/backend/open_webui/routers/auths.py b/backend/open_webui/routers/auths.py index 21e79ee91..055ee61ab 100644 --- a/backend/open_webui/routers/auths.py +++ b/backend/open_webui/routers/auths.py @@ -80,8 +80,8 @@ async def get_session_user( auth_header = request.headers.get("Authorization") auth_token = get_http_authorization_cred(auth_header) token = auth_token.credentials - data = decode_token(token) + expires_at = data.get("exp") if int(time.time()) > expires_at: diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index a06e1e3e9..3d81a4e1e 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -192,6 +192,9 @@ async def connect(sid, environ, auth): # print(f"user {user.name}({user.id}) connected with session ID {sid}") await sio.emit("user-list", {"user_ids": list(USER_POOL.keys())}) await sio.emit("usage", {"models": get_models_in_use()}) + return True + + return False @sio.on("user-join") diff --git a/src/lib/utils/websocket.ts b/src/lib/utils/websocket.ts new file mode 100644 index 000000000..41fc21b50 --- /dev/null +++ b/src/lib/utils/websocket.ts @@ -0,0 +1,56 @@ +import { io } from 'socket.io-client'; + +import { + socket, + activeUserIds, + USAGE_POOL, +} from '$lib/stores'; +import { WEBUI_BASE_URL } from '$lib/constants'; + + +export const setupSocket = async (enableWebsocket) => { + const _socket = io(`${WEBUI_BASE_URL}` || undefined, { + reconnection: true, + reconnectionDelay: 1000, + reconnectionDelayMax: 5000, + randomizationFactor: 0.5, + path: '/ws/socket.io', + transports: enableWebsocket ? ['websocket'] : ['polling', 'websocket'], + auth: { token: localStorage.token } + }); + + await socket.set(_socket); + + _socket.on('connect_error', (err) => { + console.log('connect_error', err); + }); + + _socket.on('connect', () => { + console.log('connected', _socket.id); + }); + + _socket.on('reconnect_attempt', (attempt) => { + console.log('reconnect_attempt', attempt); + }); + + _socket.on('reconnect_failed', () => { + console.log('reconnect_failed'); + }); + + _socket.on('disconnect', (reason, details) => { + console.log(`Socket ${_socket.id} disconnected due to ${reason}`); + if (details) { + console.log('Additional details:', details); + } + }); + + _socket.on('user-list', (data) => { + console.log('user-list', data); + activeUserIds.set(data.user_ids); + }); + + _socket.on('usage', (data) => { + console.log('usage', data); + USAGE_POOL.set(data['models']); + }); +}; diff --git a/src/routes/+layout.svelte b/src/routes/+layout.svelte index 851b7c899..218e92e43 100644 --- a/src/routes/+layout.svelte +++ b/src/routes/+layout.svelte @@ -48,6 +48,7 @@ import NotificationToast from '$lib/components/NotificationToast.svelte'; import AppSidebar from '$lib/components/app/AppSidebar.svelte'; import { chatCompletion } from '$lib/apis/openai'; + import { setupSocket } from '$lib/utils/websocket'; setContext('i18n', i18n); @@ -58,53 +59,6 @@ const BREAKPOINT = 768; - const setupSocket = async (enableWebsocket) => { - const _socket = io(`${WEBUI_BASE_URL}` || undefined, { - reconnection: true, - reconnectionDelay: 1000, - reconnectionDelayMax: 5000, - randomizationFactor: 0.5, - path: '/ws/socket.io', - transports: enableWebsocket ? ['websocket'] : ['polling', 'websocket'], - auth: { token: localStorage.token } - }); - - await socket.set(_socket); - - _socket.on('connect_error', (err) => { - console.log('connect_error', err); - }); - - _socket.on('connect', () => { - console.log('connected', _socket.id); - }); - - _socket.on('reconnect_attempt', (attempt) => { - console.log('reconnect_attempt', attempt); - }); - - _socket.on('reconnect_failed', () => { - console.log('reconnect_failed'); - }); - - _socket.on('disconnect', (reason, details) => { - console.log(`Socket ${_socket.id} disconnected due to ${reason}`); - if (details) { - console.log('Additional details:', details); - } - }); - - _socket.on('user-list', (data) => { - console.log('user-list', data); - activeUserIds.set(data.user_ids); - }); - - _socket.on('usage', (data) => { - console.log('usage', data); - USAGE_POOL.set(data['models']); - }); - }; - const executePythonAsWorker = async (id, code, cb) => { let result = null; let stdout = null; @@ -561,8 +515,6 @@ await WEBUI_NAME.set(backendConfig.name); if ($config) { - await setupSocket($config.features?.enable_websocket ?? true); - const currentUrl = `${window.location.pathname}${window.location.search}`; const encodedUrl = encodeURIComponent(currentUrl); @@ -574,6 +526,7 @@ }); if (sessionUser) { + await setupSocket($config.features?.enable_websocket ?? true); // Save Session User to Store $socket.emit('user-join', { auth: { token: sessionUser.token } }); diff --git a/src/routes/auth/+page.svelte b/src/routes/auth/+page.svelte index be8989002..fc7e4c53a 100644 --- a/src/routes/auth/+page.svelte +++ b/src/routes/auth/+page.svelte @@ -12,6 +12,7 @@ import { WEBUI_NAME, config, user, socket } from '$lib/stores'; import { generateInitialsImage, canvasPixelTest } from '$lib/utils'; + import { setupSocket } from '$lib/utils/websocket'; import Spinner from '$lib/components/common/Spinner.svelte'; import OnBoarding from '$lib/components/OnBoarding.svelte'; @@ -41,6 +42,9 @@ if (sessionUser.token) { localStorage.token = sessionUser.token; } + if (!$socket) { + await setupSocket($config.features?.enable_websocket ?? true); + } $socket.emit('user-join', { auth: { token: sessionUser.token } }); await user.set(sessionUser);