hiding websocket initalization behind authentication

This commit is contained in:
shcodenick 2025-04-30 15:05:57 +02:00
parent f7f3ae7cc9
commit 4c81b6ebc5
5 changed files with 66 additions and 50 deletions

View File

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

View File

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

View File

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

View File

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

View File

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