From b87c7555740dc87c69eea9704750bbbfdcb3db0f Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Fri, 24 Apr 2026 15:29:36 +0900
Subject: [PATCH] refac

---
 backend/open_webui/main.py             | 50 +++++++++++---------------
 backend/open_webui/utils/mcp/client.py | 22 ++++++------
 2 files changed, 31 insertions(+), 41 deletions(-)

diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py
index d75f1af4f4..b56659e721 100644
--- a/backend/open_webui/main.py
+++ b/backend/open_webui/main.py
@@ -1909,37 +1909,29 @@ async def chat_completion(
                 except Exception:
                     pass
         finally:
-            # 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_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}')
+            # MCP cleanup — MUST run in the SAME asyncio task as
+            # connect() because the MCP SDK's streamablehttp_client
+            # uses anyio task groups whose cancel scopes enforce
+            # same-task exit.  Do NOT wrap in asyncio.shield() or
+            # asyncio.wait_for() — both create a new task.
+            # MCPClient.disconnect() self-shields via
+            # anyio.CancelScope(shield=True).
+            try:
+                if mcp_clients := metadata.get('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}')
+            except Exception as e:
+                log.debug(f'Error cleaning up MCP clients: {e}')
 
             try:
-                await asyncio.shield(cleanup_process_chat())
-            except (asyncio.CancelledError, Exception):
+                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:
                 pass
 
     # Fan out: one task per model
diff --git a/backend/open_webui/utils/mcp/client.py b/backend/open_webui/utils/mcp/client.py
index effe4b1637..759bcc0a31 100644
--- a/backend/open_webui/utils/mcp/client.py
+++ b/backend/open_webui/utils/mcp/client.py
@@ -155,20 +155,18 @@ class MCPClient:
         self.session = None
 
         try:
-            await asyncio.wait_for(
-                asyncio.shield(exit_stack.aclose()),
-                timeout=5.0,
-            )
-        except asyncio.TimeoutError:
+            # IMPORTANT: Do NOT use asyncio.shield() or asyncio.wait_for()
+            # here — both create a new asyncio task.  The MCP SDK's
+            # streamablehttp_client uses anyio task groups / cancel scopes
+            # that MUST be exited in the same task they were entered in.
+            # Using anyio.CancelScope(shield=True) protects from
+            # CancelledError while staying in the current task.
+            with anyio.CancelScope(shield=True):
+                with anyio.fail_after(5.0):
+                    await exit_stack.aclose()
+        except TimeoutError:
             log.warning('MCPClient.disconnect() timed out after 5 s')
         except RuntimeError as exc:
-            # The MCP SDK's streamable_http transport uses anyio task
-            # groups and async generators internally.  When we close
-            # a session that was interrupted mid-flight these can
-            # raise RuntimeError ("aclose(): asynchronous generator is
-            # already running" or "Attempted to exit cancel scope in a
-            # different task").  Swallowing the error here prevents the
-            # orphaned coroutines from spinning at 100 % CPU.
             log.debug('MCPClient.disconnect() suppressed RuntimeError: %s', exc)
         except Exception as exc:
             log.debug('MCPClient.disconnect() error: %s', exc)