refac
This commit is contained in:
@@ -1847,13 +1847,28 @@ async def chat_completion(
|
|||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
finally:
|
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:
|
try:
|
||||||
if mcp_clients := metadata.get('mcp_clients'):
|
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:
|
except Exception as e:
|
||||||
log.debug(f'Error cleaning up: {e}')
|
log.debug(f'Error cleaning up MCP clients: {e}')
|
||||||
pass
|
|
||||||
# Emit chat:active=false when task completes
|
# Emit chat:active=false when task completes
|
||||||
try:
|
try:
|
||||||
if metadata.get('chat_id'):
|
if metadata.get('chat_id'):
|
||||||
|
|||||||
@@ -1,7 +1,10 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
|
import logging
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
from contextlib import AsyncExitStack
|
from contextlib import AsyncExitStack
|
||||||
|
|
||||||
|
log = logging.getLogger(__name__)
|
||||||
|
|
||||||
import anyio
|
import anyio
|
||||||
|
|
||||||
from mcp import ClientSession
|
from mcp import ClientSession
|
||||||
@@ -136,8 +139,39 @@ class MCPClient:
|
|||||||
return result_dict
|
return result_dict
|
||||||
|
|
||||||
async def disconnect(self):
|
async def disconnect(self):
|
||||||
# Clean up and close the session
|
"""Clean up and close the session.
|
||||||
await self.exit_stack.aclose()
|
|
||||||
|
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):
|
async def __aenter__(self):
|
||||||
await self.exit_stack.__aenter__()
|
await self.exit_stack.__aenter__()
|
||||||
|
|||||||
Reference in New Issue
Block a user