From bf35f64a7f14161933dfa608577a977d107b1569 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Sun, 26 Jul 2026 22:45:11 -0400 Subject: [PATCH] refac --- backend/open_webui/socket/main.py | 89 +++++++++++++++++++------------ backend/open_webui/tasks.py | 16 +++--- 2 files changed, 61 insertions(+), 44 deletions(-) diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index 78e374d126..b5a0b53e38 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -170,25 +170,31 @@ YDOC_MANAGER = YdocManager( async def periodic_session_pool_cleanup(): """Reap orphaned SESSION_POOL entries that missed heartbeats (e.g. crashed instance).""" - if not session_aquire_func(): - log.debug('Session cleanup lock held by another node. Skipping.') - return + retry_delay = random.uniform(WEBSOCKET_REDIS_LOCK_TIMEOUT / 2, WEBSOCKET_REDIS_LOCK_TIMEOUT) + while True: + if not session_aquire_func(): + log.debug('Session cleanup lock held by another node. Retrying.') + await asyncio.sleep(retry_delay) + continue - try: - while True: - if not session_renew_func(): - log.error('Unable to renew session cleanup lock. Exiting.') - return + try: + while True: + if not session_renew_func(): + log.warning('Unable to renew session cleanup lock. Retrying cleanup ownership.') + break - now = int(time.time()) - for sid in list(SESSION_POOL.keys()): - entry = SESSION_POOL.get(sid) - if entry and now - entry.get('last_seen_at', 0) > SESSION_POOL_TIMEOUT: - log.warning(f'Reaping orphaned session {sid} (user {entry.get("id")})') - del SESSION_POOL[sid] - await asyncio.sleep(SESSION_POOL_TIMEOUT) - finally: - session_release_func() + now = int(time.time()) + for sid in list(SESSION_POOL.keys()): + entry = SESSION_POOL.get(sid) + if entry and now - entry.get('last_seen_at', 0) > SESSION_POOL_TIMEOUT: + log.warning(f'Reaping orphaned session {sid} (user {entry.get("id")})') + try: + del SESSION_POOL[sid] + except KeyError: + pass + await asyncio.sleep(SESSION_POOL_TIMEOUT) + finally: + session_release_func() async def periodic_usage_pool_cleanup(): @@ -349,7 +355,7 @@ async def connect(sid, environ, auth): user = await get_verified_user_by_token(auth['token'], redis) if user: - SESSION_POOL[sid] = { + socket_user = { **user.model_dump( exclude=[ 'profile_image_url', @@ -361,6 +367,8 @@ async def connect(sid, environ, auth): ), 'last_seen_at': int(time.time()), } + SESSION_POOL[sid] = socket_user + await sio.save_session(sid, {'user': socket_user}) await sio.enter_room(sid, f'user:{user.id}') @@ -378,23 +386,21 @@ async def user_join(sid, data): if not user: return - existing = SESSION_POOL.get(sid) - if existing and existing.get('id') == user.id: - SESSION_POOL[sid] = {**existing, 'last_seen_at': int(time.time())} - else: - SESSION_POOL[sid] = { - **user.model_dump( - exclude=[ - 'profile_image_url', - 'profile_banner_image_url', - 'date_of_birth', - 'bio', - 'gender', - ] - ), - 'last_seen_at': int(time.time()), - } + socket_user = { + **user.model_dump( + exclude=[ + 'profile_image_url', + 'profile_banner_image_url', + 'date_of_birth', + 'bio', + 'gender', + ] + ), + 'last_seen_at': int(time.time()), + } + SESSION_POOL[sid] = socket_user + await sio.save_session(sid, {'user': socket_user}) await sio.enter_room(sid, f'user:{user.id}') # Join all the channels only if user has channels permission @@ -504,7 +510,15 @@ async def channel_events(sid, data): @sio.on('events:chat') async def chat_events(sid, data): - user = SESSION_POOL.get(sid) + try: + session = await sio.get_session(sid) + user = session.get('user') + except KeyError: + user = None + + if not user: + user = SESSION_POOL.get(sid) + if not user: return @@ -1069,6 +1083,11 @@ async def get_event_call(request_info): ) except TimeoutError: log.warning(f'Event caller timed out for session {session_id}') + if SESSION_POOL.get(session_id) == session: + try: + del SESSION_POOL[session_id] + except KeyError: + pass return {'error': 'Event call timed out. The browser tab may be inactive or closed.'} if 'session_id' in request_info and 'chat_id' in request_info and 'message_id' in request_info: diff --git a/backend/open_webui/tasks.py b/backend/open_webui/tasks.py index d6b2c00419..e5ff754297 100644 --- a/backend/open_webui/tasks.py +++ b/backend/open_webui/tasks.py @@ -2,10 +2,8 @@ import asyncio import json import logging -from typing import Dict, List, Optional from uuid import uuid4 -from fastapi import Request from redis.asyncio import Redis from open_webui.env import REDIS_KEY_PREFIX @@ -13,7 +11,7 @@ from open_webui.env import REDIS_KEY_PREFIX log = logging.getLogger(__name__) # A dictionary to keep track of active tasks -tasks: Dict[str, asyncio.Task] = {} +tasks: dict[str, asyncio.Task] = {} item_tasks = {} @@ -46,7 +44,7 @@ async def redis_task_command_listener(app): ### ------------------------------ -async def redis_save_task(redis: Redis, task_id: str, item_id: Optional[str]): +async def redis_save_task(redis: Redis, task_id: str, item_id: str | None): pipe = redis.pipeline() pipe.hset(REDIS_TASKS_KEY, task_id, item_id or '') if item_id: @@ -54,7 +52,7 @@ async def redis_save_task(redis: Redis, task_id: str, item_id: Optional[str]): await pipe.execute() -async def redis_cleanup_task(redis: Redis, task_id: str, item_id: Optional[str]): +async def redis_cleanup_task(redis: Redis, task_id: str, item_id: str | None): pipe = redis.pipeline() pipe.hdel(REDIS_TASKS_KEY, task_id) if item_id: @@ -67,11 +65,11 @@ async def redis_cleanup_task(redis: Redis, task_id: str, item_id: Optional[str]) await pipe.execute() -async def redis_list_tasks(redis: Redis) -> List[str]: +async def redis_list_tasks(redis: Redis) -> list[str]: return list(await redis.hkeys(REDIS_TASKS_KEY)) -async def redis_list_item_tasks(redis: Redis, item_id: str) -> List[str]: +async def redis_list_item_tasks(redis: Redis, item_id: str) -> list[str]: return list(await redis.smembers(f'{REDIS_ITEM_TASKS_KEY}:{item_id}')) @@ -101,11 +99,11 @@ async def cleanup_task(redis, task_id: str, id=None): item_tasks.pop(id, None) -async def create_task(redis, coroutine, id=None): +async def create_task(redis, coroutine, id=None, task_id=None): """ Create a new asyncio task and add it to the global task dictionary. """ - task_id = str(uuid4()) # Generate a unique ID for the task + task_id = task_id or str(uuid4()) # Generate a unique ID for the task task = asyncio.create_task(coroutine) # Create the task # Add a done callback for cleanup