diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index d6f4f4c7af..d75f1af4f4 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -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'): diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index 8d3b6dd267..fa6c65f36d 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -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