diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f10a62f328b..4cdbfefc9be 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -13532,6 +13532,14 @@ async def _stream_mcp_asgi_response( handler_task = asyncio.create_task(handle_fn(scope, receive, bridging_send)) + # If the handler task dies (exception or cancellation) without sending the EOF + # sentinel, body_iter() would block forever on body_queue.get(). The callback + # below guarantees the queue gets unblocked regardless of how the task ends. + def _ensure_eof(task: asyncio.Task) -> None: + body_queue.put_nowait(None) + + handler_task.add_done_callback(_ensure_eof) + try: status, raw_headers = await asyncio.wait_for( asyncio.shield(headers_ready), timeout=30.0 diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 111bb51315b..9712ec7aa55 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -170,13 +170,14 @@ class LiteLLM_Proxy_MCP_Handler: from litellm.proxy._experimental.mcp_server.server import ( _apply_toolset_scope, ) - from litellm.proxy._experimental.mcp_server.toolset_db import ( - get_mcp_toolset_by_name, - ) from litellm.proxy.proxy_server import prisma_client if prisma_client is not None and user_api_key_auth is not None: - toolset = await get_mcp_toolset_by_name(prisma_client, name) + toolset = ( + await global_mcp_server_manager.get_toolset_by_name_cached( + prisma_client, name + ) + ) if toolset is not None: user_api_key_auth = await _apply_toolset_scope( user_api_key_auth, toolset.toolset_id