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

@@ -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()

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.
 

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(