Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions 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
backend/open_webui/main.py
+21 -29
View File
@@ -1909,37 +1909,29 @@ async def chat_completion(
except Exception: except Exception:
pass pass
finally: finally:
# Clean up MCP clients and emit chat:active=false. # MCP cleanup — MUST run in the SAME asyncio task as
# Shield the entire block from CancelledError so cleanup # connect() because the MCP SDK's streamablehttp_client
# can finish even when the task is being stopped. # uses anyio task groups whose cancel scopes enforce
async def cleanup_process_chat(): # same-task exit. Do NOT wrap in asyncio.shield() or
try: # asyncio.wait_for() — both create a new task.
if mcp_clients := metadata.get('mcp_clients'): # MCPClient.disconnect() self-shields via
# anyio.CancelScope(shield=True).
async def cleanup_mcp_clients(): try:
for client in reversed(list(mcp_clients.values())): if mcp_clients := metadata.get('mcp_clients'):
try: for client in reversed(list(mcp_clients.values())):
await client.disconnect() try:
except Exception as e: await client.disconnect()
log.debug(f'Error disconnecting MCP client: {e}') except Exception as e:
log.debug(f'Error disconnecting MCP client: {e}')
await asyncio.wait_for(cleanup_mcp_clients(), timeout=10.0) except Exception as e:
except asyncio.TimeoutError: log.debug(f'Error cleaning up MCP clients: {e}')
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}')
try: try:
await asyncio.shield(cleanup_process_chat()) if metadata.get('chat_id'):
except (asyncio.CancelledError, Exception): 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 pass
# Fan out: one task per model # Fan out: one task per model
backend/open_webui/utils/mcp/client.py
+10 -12
View File
@@ -155,20 +155,18 @@ class MCPClient:
self.session = None self.session = None
try: try:
await asyncio.wait_for( # IMPORTANT: Do NOT use asyncio.shield() or asyncio.wait_for()
asyncio.shield(exit_stack.aclose()), # here — both create a new asyncio task. The MCP SDK's
timeout=5.0, # streamablehttp_client uses anyio task groups / cancel scopes
) # that MUST be exited in the same task they were entered in.
except asyncio.TimeoutError: # 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') log.warning('MCPClient.disconnect() timed out after 5 s')
except RuntimeError as exc: 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) log.debug('MCPClient.disconnect() suppressed RuntimeError: %s', exc)
except Exception as exc: except Exception as exc:
log.debug('MCPClient.disconnect() error: %s', exc) log.debug('MCPClient.disconnect() error: %s', exc)