mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
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:
parent
399f3f81ba
commit
a23b5d5466
2 changed files with 13 additions and 4 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue