Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 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
backend/open_webui/main.py
+37 -37
View File
@@ -1869,17 +1869,15 @@ async def chat_completion(
except asyncio.CancelledError: except asyncio.CancelledError:
log.info('Chat processing was cancelled') log.info('Chat processing was cancelled')
try: try:
event_emitter = await get_event_emitter(metadata) async def emit_cancel_event():
if event_emitter: event_emitter = await get_event_emitter(metadata)
await asyncio.shield( if event_emitter:
event_emitter( await event_emitter({'type': 'chat:tasks:cancel'})
{'type': 'chat:tasks:cancel'},
) await asyncio.shield(emit_cancel_event())
) except Exception:
except Exception as e:
pass pass
finally: raise # re-raise to ensure proper task cancellation handling
raise # re-raise to ensure proper task cancellation handling
except Exception as e: except Exception as e:
error_detail = e.detail if isinstance(e, HTTPException) else str(e) error_detail = e.detail if isinstance(e, HTTPException) else str(e)
log.error('Error processing chat payload: %s', error_detail) log.error('Error processing chat payload: %s', error_detail)
@@ -1911,36 +1909,38 @@ async def chat_completion(
except Exception: except Exception:
pass pass
finally: finally:
# Clean up MCP clients. Shield the entire block from # Clean up MCP clients and emit chat:active=false.
# CancelledError so disconnect() can finish even when the # Shield the entire block from CancelledError so cleanup
# task is being stopped. Each client is isolated so one # can finish even when the task is being stopped.
# failure doesn't skip the rest. async def cleanup_process_chat():
try: try:
if mcp_clients := metadata.get('mcp_clients'): if mcp_clients := metadata.get('mcp_clients'):
async def _cleanup_mcp(): async def cleanup_mcp_clients():
for client in reversed(list(mcp_clients.values())): for client in reversed(list(mcp_clients.values())):
try: try:
await client.disconnect() await client.disconnect()
except Exception as e: except Exception as e:
log.debug(f'Error disconnecting MCP client: {e}') log.debug(f'Error disconnecting MCP client: {e}')
await asyncio.wait_for(cleanup_mcp_clients(), timeout=10.0)
except asyncio.TimeoutError:
log.warning('MCP client cleanup timed out after 10 s')
except Exception as e:
log.debug(f'Error cleaning up MCP clients: {e}')
try:
if metadata.get('chat_id'):
event_emitter = await get_event_emitter(metadata, update_db=False)
if event_emitter:
await event_emitter({'type': 'chat:active', 'data': {'active': False}})
except Exception as e:
log.debug(f'Error emitting chat:active: {e}')
await asyncio.wait_for(
asyncio.shield(_cleanup_mcp()),
timeout=10.0,
)
except asyncio.TimeoutError:
log.warning('MCP client cleanup timed out after 10 s')
except Exception as e:
log.debug(f'Error cleaning up MCP clients: {e}')
# Emit chat:active=false when task completes
try: try:
if metadata.get('chat_id'): await asyncio.shield(cleanup_process_chat())
event_emitter = await get_event_emitter(metadata, update_db=False) except (asyncio.CancelledError, Exception):
if event_emitter: pass
await event_emitter({'type': 'chat:active', 'data': {'active': False}})
except Exception as e:
log.debug(f'Error emitting chat:active: {e}')
# Fan out: one task per model # Fan out: one task per model
if metadata.get('session_id') and metadata.get('chat_id'): if metadata.get('session_id') and metadata.get('chat_id'):
backend/open_webui/utils/middleware.py
+30 -20
View File
@@ -4269,6 +4269,8 @@ async def streaming_chat_response_handler(response, ctx):
'data': data, 'data': data,
} }
) )
except (asyncio.CancelledError, KeyboardInterrupt):
raise
except Exception as e: except Exception as e:
done = 'data: [DONE]' in line done = 'data: [DONE]' in line
if done: if done:
@@ -4971,31 +4973,39 @@ async def streaming_chat_response_handler(response, ctx):
await outlet_filter_handler(ctx) await outlet_filter_handler(ctx)
except asyncio.CancelledError: except asyncio.CancelledError:
log.warning('Task was cancelled!') 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: if not ENABLE_REALTIME_CHAT_SAVE:
# Save message in the database await Chats.upsert_message_to_chat_by_id_and_message_id(
await asyncio.shield( metadata['chat_id'],
Chats.upsert_message_to_chat_by_id_and_message_id( metadata['message_id'],
metadata['chat_id'], {
metadata['message_id'], 'done': True,
{ 'content': serialize_output(output),
'done': True, 'output': output,
'content': serialize_output(output), },
'output': output,
},
)
) )
else: else:
await asyncio.shield( await Chats.upsert_message_to_chat_by_id_and_message_id(
Chats.upsert_message_to_chat_by_id_and_message_id( metadata['chat_id'],
metadata['chat_id'], metadata['message_id'],
metadata['message_id'], {'done': True},
{'done': True},
)
) )
except Exception:
try:
await asyncio.shield(save_cancelled_state())
except (asyncio.CancelledError, Exception):
pass pass
raise # re-raise CancelledError for proper propagation raise # re-raise CancelledError for proper propagation