bhetherman/open-webui
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-07-26 22:45:11 -04:00
parent fc4906c9e9
commit bf35f64a7f
2 changed files with 61 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/socket/main.py
+54 -35
View File
@@ -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:
backend/open_webui/tasks.py
+7 -9
View File
@@ -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