diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index 94aae1804..255c9ace5 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -1847,13 +1847,28 @@ 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'): - for client in reversed(mcp_clients.values()): - await client.disconnect() + + 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}') + + 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: {e}') - pass + log.debug(f'Error cleaning up MCP clients: {e}') # Emit chat:active=false when task completes try: if metadata.get('chat_id'): diff --git a/backend/open_webui/utils/mcp/client.py b/backend/open_webui/utils/mcp/client.py index 13745aa9f..beb2f1507 100644 --- a/backend/open_webui/utils/mcp/client.py +++ b/backend/open_webui/utils/mcp/client.py @@ -1,7 +1,10 @@ import asyncio +import logging from typing import Optional from contextlib import AsyncExitStack +log = logging.getLogger(__name__) + import anyio from mcp import ClientSession @@ -136,8 +139,39 @@ class MCPClient: return result_dict async def disconnect(self): - # Clean up and close the session - await self.exit_stack.aclose() + """Clean up and close the session. + + This method is idempotent — calling it multiple times or on a + client that was never connected is safe. It shields the close + operation from CancelledError and adds a timeout so a hung MCP + server cannot block the event loop indefinitely. + """ + exit_stack = self.exit_stack + if exit_stack is None: + return + + # Prevent double-close from concurrent callers + self.exit_stack = None + self.session = None + + try: + await asyncio.wait_for( + asyncio.shield(exit_stack.aclose()), + timeout=5.0, + ) + except asyncio.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) async def __aenter__(self): await self.exit_stack.__aenter__()