mirror of
https://github.com/open-webui/open-webui.git
synced 2026-08-26 23:44:48 -06:00
refac
This commit is contained in:
Timothy Jaeryang Baek
+37
-37
@@ -1869,17 +1869,15 @@ async def chat_completion(
|
||||
except asyncio.CancelledError:
|
||||
log.info('Chat processing was cancelled')
|
||||
try:
|
||||
event_emitter = await get_event_emitter(metadata)
|
||||
if event_emitter:
|
||||
await asyncio.shield(
|
||||
event_emitter(
|
||||
{'type': 'chat:tasks:cancel'},
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
async def emit_cancel_event():
|
||||
event_emitter = await get_event_emitter(metadata)
|
||||
if event_emitter:
|
||||
await event_emitter({'type': 'chat:tasks:cancel'})
|
||||
|
||||
await asyncio.shield(emit_cancel_event())
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
raise # re-raise to ensure proper task cancellation handling
|
||||
raise # re-raise to ensure proper task cancellation handling
|
||||
except Exception as e:
|
||||
error_detail = e.detail if isinstance(e, HTTPException) else str(e)
|
||||
log.error('Error processing chat payload: %s', error_detail)
|
||||
@@ -1911,36 +1909,38 @@ async def chat_completion(
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
# Clean up MCP clients. Shield the entire block from
|
||||
# CancelledError so disconnect() can finish even when the
|
||||
# task is being stopped. Each client is isolated so one
|
||||
# failure doesn't skip the rest.
|
||||
try:
|
||||
if mcp_clients := metadata.get('mcp_clients'):
|
||||
# Clean up MCP clients and emit chat:active=false.
|
||||
# Shield the entire block from CancelledError so cleanup
|
||||
# can finish even when the task is being stopped.
|
||||
async def cleanup_process_chat():
|
||||
try:
|
||||
if mcp_clients := metadata.get('mcp_clients'):
|
||||
|
||||
async def _cleanup_mcp():
|
||||
for client in reversed(list(mcp_clients.values())):
|
||||
try:
|
||||
await client.disconnect()
|
||||
except Exception as e:
|
||||
log.debug(f'Error disconnecting MCP client: {e}')
|
||||
async def cleanup_mcp_clients():
|
||||
for client in reversed(list(mcp_clients.values())):
|
||||
try:
|
||||
await client.disconnect()
|
||||
except Exception as 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:
|
||||
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.shield(cleanup_process_chat())
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass
|
||||
|
||||
# Fan out: one task per model
|
||||
if metadata.get('session_id') and metadata.get('chat_id'):
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user