From aa48106fcab691de1b2f916d6fd68784429cec17 Mon Sep 17 00:00:00 2001 From: Classic298 <27028174+Classic298@users.noreply.github.com> Date: Thu, 3 Sep 2026 20:57:39 +0200 Subject: [PATCH] fix: unregister the direct-connection stream listener on every exit (#29509) A direct-connection streaming request registers a per-request socket.io handler before it asks the browser to start the completion, and that handler was only removed once the response had been fully streamed. Every other way the request could end left it behind for the life of the process: the call to the browser raising, a non-success status, an ack with no arguments or an ack that is not an object, cancellation while waiting for that ack and a response body closed or garbage-collected before it finished. A client that repeatedly hits a failing direct connection grows the server's handler table and the closures it holds without bound, and nothing ever cleans it up. The removal now runs on every exit from the request, through one named helper that pops with a default so it is safe to run twice on the paths where both the generator's finally and the response's background task fire. The guard covers the exchange up to the status read and catches BaseException, because cancellation is not an Exception, and re-raises it unchanged. Measured across thirteen exit paths: twelve leak a handler on dev and none of those leak here. The thirteenth, a response body that is never iterated at all, behaves the same on both. 200 failing requests leave 200 handlers on dev and none on this branch. Streamed bytes on the success path, exception types and cancellation behaviour are unchanged. --- backend/open_webui/utils/chat.py | 43 +++++++++++++++++++------------- 1 file changed, 26 insertions(+), 17 deletions(-) diff --git a/backend/open_webui/utils/chat.py b/backend/open_webui/utils/chat.py index d74c15e7d1..a769d5f1df 100644 --- a/backend/open_webui/utils/chat.py +++ b/backend/open_webui/utils/chat.py @@ -79,25 +79,34 @@ async def generate_direct_chat_completion( """ await q.put(data) + def remove_message_listener(): + sio.handlers['/'].pop(channel, None) + # Register the listener sio.on(channel, message_listener) # Start processing chat completion in background - res = await event_caller( - { - 'type': 'request:chat:completion', - 'data': { - 'form_data': form_data, - 'model': models[form_data['model']], - 'channel': channel, - 'session_id': session_id, - }, - } - ) + try: + res = await event_caller( + { + 'type': 'request:chat:completion', + 'data': { + 'form_data': form_data, + 'model': models[form_data['model']], + 'channel': channel, + 'session_id': session_id, + }, + } + ) - log.info('res: %s', res) + log.info('res: %s', res) - if res.get('status', False): + status = res.get('status', False) + except BaseException: + remove_message_listener() + raise + + if status: # Define a generator to stream responses async def event_generator(): nonlocal q @@ -117,17 +126,17 @@ async def generate_direct_chat_completion( except Exception as e: log.debug('Error in event generator: %s', e) pass + finally: + remove_message_listener() # Define a background task to run the event generator async def background(): - try: - del sio.handlers['/'][channel] - except Exception as e: - pass + remove_message_listener() # Return the streaming response return StreamingResponse(event_generator(), media_type='text/event-stream', background=background) else: + remove_message_listener() raise Exception(str(res)) else: res = await event_caller(