fix(mcp): prevent body_iter deadlock + use cached toolset lookup in responses API

- _stream_mcp_asgi_response: add done callback to handler_task that puts
  the EOF sentinel on body_queue when the task exits, preventing body_iter
  from hanging forever if the handler raises after headers are sent.
- litellm_proxy_mcp_handler: replace raw get_mcp_toolset_by_name() DB call
  with global_mcp_server_manager.get_toolset_by_name_cached() so toolset
  resolution uses the 60s TTL cache added for this purpose instead of
  hitting the DB on every responses-API request.
This commit is contained in:
Ishaan Jaffer 2026-03-23 10:06:54 -07:00
parent 399f3f81ba
commit a23b5d5466
2 changed files with 13 additions and 4 deletions

View file

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

View file

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