scriptoriumadmin/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-08-28 05:27:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Merge 75fe65c2f8 into db22d89cab

Browse source
This commit is contained in:
Classic298 2026-08-26 06:23:13 +08:00 • committed by GitHub
commit f4c14a51db
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 104 additions and 80 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
backend/open_webui/routers/channels.py
View file

@ -1242,7 +1242,7 @@ async def post_new_message(
except Exception as e:
log.debug(e)
active_user_ids = get_user_ids_from_room(f'channel:{channel.id}')
active_user_ids = await get_user_ids_from_room(f'channel:{channel.id}')
# NOTE: We intentionally do NOT pass db to background_handler.
# Background tasks should manage their own short-lived sessions to avoid

164
backend/open_webui/socket/main.py
View file

@ -175,48 +175,65 @@ YDOC_MANAGER = YdocManager(
)
def session_pool_batches():
"""All session pool entries, in bounded batches for the Redis backing."""
if WEBSOCKET_MANAGER == 'redis':
return SESSION_POOL.scan_batches()
return [list(SESSION_POOL.items())]
async def periodic_session_pool_cleanup():
"""Reap orphaned SESSION_POOL entries that missed heartbeats (e.g. crashed instance)."""
retry_delay = random.uniform(WEBSOCKET_REDIS_LOCK_TIMEOUT / 2, WEBSOCKET_REDIS_LOCK_TIMEOUT)
renew_interval = max(WEBSOCKET_REDIS_LOCK_TIMEOUT / 2, 0.5)
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.warning('Unable to renew session cleanup lock. Retrying cleanup ownership.')
break
if not session_aquire_func():
log.debug('Session cleanup lock held by another node. Retrying.')
await asyncio.sleep(retry_delay)
continue
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
next_cleanup_at = time.monotonic() + SESSION_POOL_TIMEOUT
lock_lost = False
try:
while True:
sleep_for = min(renew_interval, next_cleanup_at - time.monotonic())
if sleep_for <= 0:
break
await asyncio.sleep(sleep_for)
if not session_renew_func():
log.warning('Unable to renew session cleanup lock. Retrying cleanup ownership.')
lock_lost = True
break
if lock_lost:
break
finally:
session_release_func()
now = int(time.time())
for batch in session_pool_batches():
expired = {
sid: entry.get('id')
for sid, entry in batch
if now - entry.get('last_seen_at', 0) > SESSION_POOL_TIMEOUT
}
if expired:
log.warning('Reaping %d orphaned session(s) (sid: user): %s', len(expired), expired)
if WEBSOCKET_MANAGER == 'redis':
SESSION_POOL.discard(*expired)
else:
for sid in expired:
SESSION_POOL.pop(sid, None)
await asyncio.sleep(0) # don't hold the loop for the whole sweep
next_cleanup_at = time.monotonic() + SESSION_POOL_TIMEOUT
lock_lost = False
while True:
sleep_for = min(renew_interval, next_cleanup_at - time.monotonic())
if sleep_for <= 0:
break
await asyncio.sleep(sleep_for)
if not session_renew_func():
log.warning('Unable to renew session cleanup lock. Retrying cleanup ownership.')
lock_lost = True
break
if lock_lost:
break
finally:
session_release_func()
except Exception:
log.exception('Session pool cleanup failed. Retrying.')
await asyncio.sleep(retry_delay)
async def periodic_usage_pool_cleanup():
@ -283,32 +300,32 @@ def get_user_id_from_session_pool(sid):
return None
async def get_socket_session_user(sid):
"""Session user from this worker's local Socket.IO store; only locally connected sids are ever looked up."""
try:
return (await sio.get_session(sid)).get('user')
except KeyError:
return None
def get_session_ids_from_room(room):
"""Get all session IDs from a specific room."""
members = get_room_sid_map(sio.manager, '/', room)
return list(members) if members else []
def get_session_ids_by_user_id(user_id: str) -> list[str]:
async def get_session_ids_by_user_id(user_id: str) -> list[str]:
"""Get known session IDs for a user across the local rooms and shared session pool."""
session_ids = set(get_session_ids_from_room(f'user:{user_id}'))
session_ids.update(sid for sid, entry in SESSION_POOL.items() if entry and entry.get('id') == user_id)
for batch in session_pool_batches():
session_ids.update(sid for sid, entry in batch if entry.get('id') == user_id)
await asyncio.sleep(0) # don't hold the loop for the whole pool
return list(session_ids)
def get_user_ids_from_room(room):
active_session_ids = get_session_ids_from_room(room)
# Single pool lookup per session (each .get is a Redis round trip
# when the session pool is Redis-backed).
active_user_ids = list(
{
entry['id']
for entry in (SESSION_POOL.get(session_id) for session_id in active_session_ids)
if entry is not None
}
)
return active_user_ids
async def get_user_ids_from_room(room):
users = [await get_socket_session_user(session_id) for session_id in get_session_ids_from_room(room)]
return {user['id'] for user in users if user}
async def emit_to_users(event: str, data: dict, user_ids: list[str]):
@ -351,7 +368,7 @@ async def disconnect_user_sessions(user_id: str):
The client will automatically reconnect and re-authenticate with
fresh data from the database.
"""
session_ids = get_session_ids_by_user_id(user_id)
session_ids = await get_session_ids_by_user_id(user_id)
for sid in session_ids:
try:
await sio.disconnect(sid)
@ -364,7 +381,7 @@ async def disconnect_user_sessions(user_id: str):
@sio.on('usage')
async def usage(sid, data):
if sid in SESSION_POOL:
if await get_socket_session_user(sid):
model_id = data['model']
# Record the timestamp for the last update
current_time = int(time.time())
@ -446,7 +463,7 @@ async def user_join(sid, data):
@sio.on('heartbeat')
async def heartbeat(sid, data):
user = SESSION_POOL.get(sid)
user = await get_socket_session_user(sid)
if user:
SESSION_POOL[sid] = {**user, 'last_seen_at': int(time.time())}
await Users.update_last_active_by_id(user['id'])
@ -519,7 +536,7 @@ async def channel_events(sid, data):
event_data = data['data']
event_type = event_data['type']
user = SESSION_POOL.get(sid)
user = await get_socket_session_user(sid)
if not user:
return
@ -559,15 +576,7 @@ async def get_folder_unread_counts(user_id: str) -> dict[str, int]:
@sio.on('events:chat')
async def chat_events(sid, data):
try:
session = await sio.get_session(sid)
user = session.get('user')
except KeyError:
user = None
if not user:
user = SESSION_POOL.get(sid)
user = await get_socket_session_user(sid)
if not user:
return
@ -621,7 +630,7 @@ def normalize_document_id(document_id: str) -> str:
@sio.on('ydoc:document:join')
async def ydoc_document_join(sid, data):
"""Handle user joining a document"""
user = SESSION_POOL.get(sid)
user = await get_socket_session_user(sid)
if not user:
return
@ -781,7 +790,7 @@ async def yjs_document_update(sid, data):
return
# Verify write permission — room membership only proves read access
user = SESSION_POOL.get(sid)
user = await get_socket_session_user(sid)
if not user:
return
@ -851,7 +860,7 @@ async def yjs_document_update(sid, data):
@sio.on('ydoc:document:leave')
async def yjs_document_leave(sid, data):
"""Handle user leaving a document"""
user = SESSION_POOL.get(sid)
user = await get_socket_session_user(sid)
if not user: # authenticated session required (parity with sibling handlers)
return
try:
@ -883,7 +892,7 @@ async def yjs_document_leave(sid, data):
@sio.on('ydoc:awareness:update')
async def yjs_awareness_update(sid, data):
"""Handle awareness updates (cursors, selections, etc.)"""
user = SESSION_POOL.get(sid)
user = await get_socket_session_user(sid)
if not user: # authenticated session required (parity with sibling handlers)
return
try:
@ -907,23 +916,24 @@ async def yjs_awareness_update(sid, data):
@sio.event
async def disconnect(sid, reason=None):
if sid in SESSION_POOL:
try:
del SESSION_POOL[sid]
# Clean up USAGE_POOL entries for this session
for model_id in list(USAGE_POOL.keys()):
connections = USAGE_POOL.get(model_id)
if connections and sid in connections:
del connections[sid]
if not connections:
del USAGE_POOL[model_id]
else:
USAGE_POOL[model_id] = connections
await YDOC_MANAGER.remove_user_from_all_documents(sid)
else:
except KeyError:
pass
# print(f"Unknown session ID {sid} disconnected")
# Clean up USAGE_POOL entries for this session
for model_id, connections in list(USAGE_POOL.items()):
if sid in connections:
del connections[sid]
if not connections:
try:
del USAGE_POOL[model_id]
except KeyError:
pass
else:
USAGE_POOL[model_id] = connections
await YDOC_MANAGER.remove_user_from_all_documents(sid)
async def _make_channel_emitter(request_info):

18
backend/open_webui/socket/utils.py
View file

@ -112,6 +112,21 @@ class RedisDict:
def items(self):
return [(k, JSONCodec.loads(v)) for k, v in self.redis.hgetall(self.name).items()]
def scan_batches(self):
"""Yield lists of (key, value) pairs via incremental HSCAN; a field may repeat across batches."""
cursor = 0
while True:
cursor, batch = self.redis.hscan(self.name, cursor, count=200)
if batch:
yield [(k, JSONCodec.loads(v)) for k, v in batch.items()]
if cursor == 0:
break
def discard(self, *keys):
"""Delete fields in one HDEL; no keys is a no-op (HDEL rejects an empty field list)."""
if keys:
self.redis.hdel(self.name, *keys)
def set(self, mapping: dict):
if not mapping:
self.clear()
@ -140,8 +155,7 @@ class RedisDict:
# We never DELETE the whole hash — this eliminates the race window
# where concurrent readers would see an empty models dict.
self.redis.hset(self.name, mapping=serialized)
if keys_to_remove:
self.redis.hdel(self.name, *keys_to_remove)
self.discard(*keys_to_remove)
if self._signature_name:
self.redis.set(self._signature_name, signature)