Merge pull request #13364 from NovoNordisk-OpenSource/websocket-behind-auth

refactor: hiding websocket initalization behind authentication

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 (expires_at is not None) and int(time.time()) > expires_at:

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

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']);
+    });
+};

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 } });
 

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