Mirrors/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-07-10 12:35:55 +02:00
Code Issues Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-04-24 15:29:36 +09:00
parent b73538ece7
commit b87c755574
2 changed files with 31 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
backend/open_webui/main.py
View File

@@ -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

22
backend/open_webui/utils/mcp/client.py
View File

@@ -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)