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.
This commit is contained in:
Classic298 2026-09-03 20:57:39 +02:00 committed by GitHub
parent 8600b0564b
commit aa48106fca
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

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