refac

backend/open_webui/utils/mcp/client.py

@@ -82,7 +82,7 @@ class MCPClient:
                     await self.session.initialize()
                 self.exit_stack = exit_stack.pop_all()
             except Exception as e:
-                await asyncio.shield(self.disconnect())
+                await self.disconnect()
                 raise e
 
     async def list_tool_specs(self) -> Optional[dict]:
@@ -150,9 +150,7 @@ class MCPClient:
         """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.
+        client that was never connected is safe.
         """
         exit_stack = self.exit_stack
         if exit_stack is None:
@@ -172,8 +170,11 @@ class MCPClient:
             # We simply call aclose() directly. If the task is cancelled, the
             # sockets will eventually be cleaned up by garbage collection.
             await exit_stack.aclose()
-        except TimeoutError:
-            log.warning('MCPClient.disconnect() timed out after 5 s')
+        except asyncio.CancelledError as exc:
+            task = asyncio.current_task()
+            if task is not None and task.cancelling():
+                raise
+            log.debug('MCPClient.disconnect() suppressed internal cancellation: %s', exc)
         except RuntimeError as exc:
             log.debug('MCPClient.disconnect() suppressed RuntimeError: %s', exc)
         except Exception as exc: