scriptoriumadmin/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-10-09 03:18:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

fix: Direct Connection replies are saved scrambled or empty (#31979)

Browse source 
* fix: Direct Connection replies are saved scrambled or empty

With a Direct Connection the browser tab forwards the model's reply to the server piece by piece. Since the server started checking the tab's session again for every one of those pieces, the pieces could overtake each other while the checks ran, so the saved reply came out in the wrong order, or empty when the end of the stream arrived first. Pieces from one tab are now handled one after another in the order they arrived, and each one is still checked.

Fixes #31953

* fix: check Direct Connection reply pieces side by side while keeping them in order

Handling one tab's reply pieces strictly one after another also made each piece wait for the previous piece's session check, so a fast 2000-piece reply took 12 to 22% longer to save than without the ordering. Each piece's check now starts as soon as it arrives and only the hand-over waits its turn, so replies save as fast as before and still in order. Once a check fails, for example after a sign-out, no later piece from that tab gets through either.
This commit is contained in:
Classic298 2026-10-07 14:38:01 +02:00 • committed by GitHub
parent f6cbeb1a1c
commit a1bb3b3923
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 17 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -6,6 +6,7 @@ import logging
import random
import sys
import time
import weakref
from contextlib import suppress
from typing import Any
@ -209,6 +210,7 @@ REDIS_EVENT_CHANNEL = f'{REDIS_KEY_PREFIX}:direct_completion'
EVENT_QUEUES: dict[str, asyncio.Queue] = {}
EVENT_PUBLISH_LOCK = asyncio.Lock()
SESSION_EVENT_LOCKS: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary()
def get_session_pool_batches():
@ -1068,19 +1070,22 @@ async def socket_event_handler(event: Any, sid: str, *args: Any) -> None:
if not isinstance(event, str) or event.count(':') != 2 or not args:
return
user = await get_socket_session_user(sid)
if not user or user.get('id') != event.split(':', 1)[0]:
return
# The lock keeps arrival order; the sid re-check drops every event after a failed check
session_check = asyncio.create_task(get_socket_session_user(sid))
async with SESSION_EVENT_LOCKS.setdefault(sid, asyncio.Lock()):
user = await session_check
if not user or sid not in LOCAL_AUTHENTICATED_SIDS or user.get('id') != event.split(':', 1)[0]:
return
queue = EVENT_QUEUES.get(event)
if queue is not None:
await queue.put(args[0])
elif WEBSOCKET_MANAGER == 'redis':
try:
async with EVENT_PUBLISH_LOCK:
await REDIS.publish(REDIS_EVENT_CHANNEL, dumps_bytes({'channel': event, 'data': args[0]}))
except RedisError as e:
log.debug('Failed to relay socket event %s: %s', event, e)
queue = EVENT_QUEUES.get(event)
if queue is not None:
await queue.put(args[0])
elif WEBSOCKET_MANAGER == 'redis':
try:
async with EVENT_PUBLISH_LOCK:
await REDIS.publish(REDIS_EVENT_CHANNEL, dumps_bytes({'channel': event, 'data': args[0]}))
except RedisError as e:
log.debug('Failed to relay socket event %s: %s', event, e)
async def _make_channel_emitter(request_info):