github-starred/open-webui
2
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-05-03 10:49:21 -05:00
Code Issues 915 Packages Projects Releases 126 Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-04-24 15:21:37 +09:00
parent 58bc254809
commit 90584ab6f3
2 changed files with 67 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
backend/open_webui/utils/middleware.py
View File

@@ -4269,6 +4269,8 @@ async def streaming_chat_response_handler(response, ctx):
'data': data,
}
)
except (asyncio.CancelledError, KeyboardInterrupt):
raise
except Exception as e:
done = 'data: [DONE]' in line
if done:
@@ -4971,31 +4973,39 @@ async def streaming_chat_response_handler(response, ctx):
await outlet_filter_handler(ctx)
except asyncio.CancelledError:
log.warning('Task was cancelled!')
try:
await asyncio.shield(event_emitter({'type': 'chat:tasks:cancel'}))
# Close the response body iterator to trigger cleanup
# in stream_wrapper's finally block and release the
# upstream connection. Without this, the async
# generator is orphaned and may spin in anyio internals.
if hasattr(response, 'body_iterator') and hasattr(response.body_iterator, 'aclose'):
try:
await asyncio.shield(response.body_iterator.aclose())
except (asyncio.CancelledError, Exception):
pass
async def save_cancelled_state():
await event_emitter({'type': 'chat:tasks:cancel'})
if not ENABLE_REALTIME_CHAT_SAVE:
# Save message in the database
await asyncio.shield(
Chats.upsert_message_to_chat_by_id_and_message_id(
metadata['chat_id'],
metadata['message_id'],
{
'done': True,
'content': serialize_output(output),
'output': output,
},
)
await Chats.upsert_message_to_chat_by_id_and_message_id(
metadata['chat_id'],
metadata['message_id'],
{
'done': True,
'content': serialize_output(output),
'output': output,
},
)
else:
await asyncio.shield(
Chats.upsert_message_to_chat_by_id_and_message_id(
metadata['chat_id'],
metadata['message_id'],
{'done': True},
)
await Chats.upsert_message_to_chat_by_id_and_message_id(
metadata['chat_id'],
metadata['message_id'],
{'done': True},
)
except Exception:
try:
await asyncio.shield(save_cancelled_state())
except (asyncio.CancelledError, Exception):
pass
raise # re-raise CancelledError for proper propagation