refac

backend/open_webui/main.py

@@ -1869,16 +1869,14 @@ async def chat_completion(
         except asyncio.CancelledError:
             log.info('Chat processing was cancelled')
             try:
+                async def emit_cancel_event():
                     event_emitter = await get_event_emitter(metadata)
                     if event_emitter:
-                    await asyncio.shield(
+                        await event_emitter({'type': 'chat:tasks:cancel'})
-                        event_emitter(
+
-                            {'type': 'chat:tasks:cancel'},
+                await asyncio.shield(emit_cancel_event())
-                        )
+            except Exception:
-                    )
-            except Exception as e:
                 pass
-            finally:
             raise  # re-raise to ensure proper task cancellation handling
         except Exception as e:
             error_detail = e.detail if isinstance(e, HTTPException) else str(e)
@@ -1911,29 +1909,26 @@ async def chat_completion(
                 except Exception:
                     pass
         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:
                     if mcp_clients := metadata.get('mcp_clients'):
 
-                    async def _cleanup_mcp():
+                        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(
+                        await asyncio.wait_for(cleanup_mcp_clients(), timeout=10.0)
-                        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)
@@ -1942,6 +1937,11 @@ async def chat_completion(
                 except Exception as e:
                     log.debug(f'Error emitting chat:active: {e}')
 
+            try:
+                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'):
         task_ids = []

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,13 +4973,21 @@ 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 Chats.upsert_message_to_chat_by_id_and_message_id(
-                        await asyncio.shield(
-                            Chats.upsert_message_to_chat_by_id_and_message_id(
                             metadata['chat_id'],
                             metadata['message_id'],
                             {
@@ -4986,16 +4996,16 @@ async def streaming_chat_response_handler(response, ctx):
                                 'output': output,
                             },
                         )
-                        )
                     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['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