From ae5abfc53c5b04a9d2a49e755c6e3de153e9c49c Mon Sep 17 00:00:00 2001 From: Classic298 <27028174+Classic298@users.noreply.github.com> Date: Wed, 2 Sep 2026 01:22:34 +0200 Subject: [PATCH] fix: direct connection chat completions hang with multiple workers A chat completion through a direct connection never finishes when Open WebUI runs with more than one worker. The browser does the model call itself and streams the chunks back over its websocket, and that websocket is held by whichever worker it happened to connect to, which is often not the worker running the HTTP request. The chunks then arrive at a process that holds no listener for that request, so nothing is ever written to the response and the client waits forever. Workers now relay those chunks to each other over one shared Redis pub/sub channel: a worker holding a chunk for a completion it is not running publishes it, and the worker that is running it puts the chunk on that completion's queue. A chunk is only queued or relayed when the emitting session's user id matches the user id the channel is named for; binding to the user leaves a browser that reconnects mid-stream able to finish, where binding to the socket session would drop the rest of its chunks. Picking the chunks up needs a Socket.IO catch-all handler, because the event name carries the request id and a fixed name would need a frontend change. Events the server previously left unhandled now reach one, so an unknown event sent with an acknowledgement callback gets an empty ack where it got none. One shared channel spares a subscribe and unsubscribe round trip per stream, at the cost of every worker parsing every relayed chunk. Publishes are serialised behind one lock per worker, which also keeps the relay on a single Redis connection. Past that connection's throughput chunks wait on the lock. A lock per channel orders chunks just as well but allows one publish in flight per stream, which past roughly a hundred concurrent streams on a worker exhausts the Redis client's connection pool and silently drops chunks. The relay only runs when the websocket manager is Redis, so without it the behaviour is unchanged. --- backend/open_webui/main.py | 8 +++ backend/open_webui/socket/main.py | 90 ++++++++++++++++++++++++++++++- backend/open_webui/utils/chat.py | 29 +++------- 3 files changed, 104 insertions(+), 23 deletions(-) diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index a3bdf4f327..c95aad90c3 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -182,7 +182,9 @@ from open_webui.routers.retrieval import ( get_rf, ) from open_webui.socket.main import ( + DIRECT_COMPLETION_RELAY_ENABLED, MODELS, + direct_completion_listener, get_event_emitter, get_models_in_use, get_user_id_from_session_pool, @@ -386,6 +388,9 @@ async def lifespan(app: FastAPI): if app.state.redis is not None: app.state.redis_task_command_listener = asyncio.create_task(redis_task_command_listener(app)) + if DIRECT_COMPLETION_RELAY_ENABLED: + app.state.direct_completion_listener = asyncio.create_task(direct_completion_listener()) + app.state.periodic_usage_pool_cleanup = asyncio.create_task(periodic_usage_pool_cleanup()) app.state.periodic_session_pool_cleanup = asyncio.create_task(periodic_session_pool_cleanup()) @@ -472,6 +477,9 @@ async def lifespan(app: FastAPI): if hasattr(app.state, 'redis_task_command_listener'): app.state.redis_task_command_listener.cancel() + if hasattr(app.state, 'direct_completion_listener'): + app.state.direct_completion_listener.cancel() + app.state.periodic_usage_pool_cleanup.cancel() app.state.periodic_session_pool_cleanup.cancel() app.state.scheduler_worker_loop.cancel() diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index c3678f592d..30ef579c26 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -5,6 +5,7 @@ import logging import random import sys import time +from contextlib import suppress from typing import Any import pycrdt as Y @@ -37,17 +38,23 @@ from open_webui.models.folders import Folders from open_webui.models.notes import Notes, NoteUpdateForm from open_webui.models.users import UserNameResponse, Users from open_webui.socket.utils import RedisDict, RedisLock, YdocManager -from open_webui.tasks import create_task, stop_item_tasks +from open_webui.tasks import ( + REDIS_PUBSUB_MAX_RECONNECT_INTERVAL, + REDIS_PUBSUB_RECONNECT_INTERVAL, + create_task, + stop_item_tasks, +) from open_webui.utils.access_control import has_permission from open_webui.utils.auth import get_verified_user_by_token from open_webui.utils.chat_id import is_saved_chat_id -from open_webui.utils.json_codec import SOCKETIO_JSON +from open_webui.utils.json_codec import SOCKETIO_JSON, JSONCodec, dumps_bytes from open_webui.utils.misc import get_output_text from open_webui.utils.redis import ( build_sentinel_url, get_redis_connection, get_sentinels_from_env, ) +from redis.exceptions import RedisError from socketio.packet import Packet logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL) @@ -189,6 +196,13 @@ YDOC_MANAGER = YdocManager( redis_key_prefix=f'{REDIS_KEY_PREFIX}:ydoc:documents', ) +REDIS_DIRECT_COMPLETION_CHANNEL = f'{REDIS_KEY_PREFIX}:direct_completion' + +DIRECT_COMPLETION_RELAY_ENABLED = WEBSOCKET_MANAGER == 'redis' + +DIRECT_COMPLETION_QUEUES: dict[str, asyncio.Queue] = {} +DIRECT_COMPLETION_PUBLISH_LOCK = asyncio.Lock() + def get_session_pool_batches(): """All session pool entries, in bounded batches for the Redis backing.""" @@ -943,6 +957,78 @@ async def disconnect(sid, reason=None): # print(f"Unknown session ID {sid} disconnected") +async def direct_completion_listener() -> None: + """Hand every relayed chunk to the queue of the direct completion it belongs to.""" + reconnect_interval = REDIS_PUBSUB_RECONNECT_INTERVAL + + while True: + pubsub = None + try: + # RedisCluster can't route a pubsub subscribe until initialize() fills its slot cache. + await REDIS.initialize() + + pubsub = REDIS.pubsub() + await pubsub.subscribe(REDIS_DIRECT_COMPLETION_CHANNEL) + reconnect_interval = REDIS_PUBSUB_RECONNECT_INTERVAL + + async for message in pubsub.listen(): + if message['type'] != 'message': + continue + chunk = JSONCodec.loads(message['data']) + queue = DIRECT_COMPLETION_QUEUES.get(chunk['channel']) + if queue is not None: + await queue.put(chunk['data']) + log.warning('Direct completion relay listener stopped. Retrying.') + except asyncio.CancelledError: + raise + except Exception: + log.exception('Direct completion relay listener failed. Retrying.') + finally: + if pubsub: + with suppress(Exception): + await pubsub.aclose() + + await asyncio.sleep(reconnect_interval) + reconnect_interval = min(reconnect_interval * 2, REDIS_PUBSUB_MAX_RECONNECT_INTERVAL) + + +def open_direct_completion_queue(channel: str) -> asyncio.Queue: + """Open the queue receiving the client's chunk emits for one direct completion.""" + queue = asyncio.Queue() + DIRECT_COMPLETION_QUEUES[channel] = queue + return queue + + +def close_direct_completion_queue(channel: str) -> None: + DIRECT_COMPLETION_QUEUES.pop(channel, None) + + +async def publish_direct_completion_chunk(channel: str, data: Any) -> None: + """Publish one chunk to the relay channel.""" + async with DIRECT_COMPLETION_PUBLISH_LOCK: + await REDIS.publish(REDIS_DIRECT_COMPLETION_CHANNEL, dumps_bytes({'channel': channel, 'data': data})) + + +@sio.on('*') +async def handle_direct_completion_chunk(event: Any, sid: str, *args: Any) -> None: + """Intake for the chunks a browser streams back for its own direct completion.""" + 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 + + queue = DIRECT_COMPLETION_QUEUES.get(event) + if queue is not None: + await queue.put(args[0]) + elif DIRECT_COMPLETION_RELAY_ENABLED: + try: + await publish_direct_completion_chunk(event, args[0]) + except RedisError as e: + log.debug('Failed to relay direct completion chunk on %s: %s', event, e) + + async def _make_channel_emitter(request_info): """Event emitter that routes pipeline output to a channel message. diff --git a/backend/open_webui/utils/chat.py b/backend/open_webui/utils/chat.py index 0b64a52165..8ae5a69599 100644 --- a/backend/open_webui/utils/chat.py +++ b/backend/open_webui/utils/chat.py @@ -1,4 +1,3 @@ -import asyncio import logging import random import sys @@ -23,9 +22,10 @@ from open_webui.routers.pipelines import ( process_pipeline_outlet_filter, ) from open_webui.socket.main import ( + close_direct_completion_queue, get_event_call, get_event_emitter, - sio, + open_direct_completion_queue, ) from open_webui.utils.filter import ( get_filter_functions, @@ -71,19 +71,7 @@ async def generate_direct_chat_completion( logging.info('WebSocket channel: %s', channel) if form_data.get('stream'): - q = asyncio.Queue() - - async def message_listener(sid, data): - """ - Handle received socket messages and push them into the queue. - """ - await q.put(data) - - def remove_message_listener(): - sio.handlers['/'].pop(channel, None) - - # Register the listener - sio.on(channel, message_listener) + queue = open_direct_completion_queue(channel) # Start processing chat completion in background try: @@ -103,16 +91,15 @@ async def generate_direct_chat_completion( status = res.get('status', False) except BaseException: - remove_message_listener() + close_direct_completion_queue(channel) raise if status: # Define a generator to stream responses async def event_generator(): - nonlocal q try: while True: - data = await q.get() # Wait for new messages + data = await queue.get() # Wait for new messages if isinstance(data, dict): if 'done' in data and data['done']: break # Stop streaming when 'done' is received @@ -127,16 +114,16 @@ async def generate_direct_chat_completion( log.debug('Error in event generator: %s', e) pass finally: - remove_message_listener() + close_direct_completion_queue(channel) # Define a background task to run the event generator async def background(): - remove_message_listener() + close_direct_completion_queue(channel) # Return the streaming response return StreamingResponse(event_generator(), media_type='text/event-stream', background=background) else: - remove_message_listener() + close_direct_completion_queue(channel) raise Exception(str(res)) else: res = await event_caller(