From 57fc329f48f30852fad7ed35de6a0a201ed47f1a Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Mon, 23 Mar 2026 11:47:35 -0700 Subject: [PATCH] fix: use UniqueViolationError import; guard _ensure_eof for error/cancel only --- .../mcp_management_endpoints.py | 19 +++++++++---------- litellm/proxy/proxy_server.py | 3 ++- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 4c9102f04c1..ae6f0ae86d2 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -37,9 +37,10 @@ from fastapi import ( from fastapi.responses import JSONResponse try: - from prisma.errors import RecordNotFoundError + from prisma.errors import RecordNotFoundError, UniqueViolationError except ImportError: RecordNotFoundError = Exception # type: ignore + UniqueViolationError = Exception # type: ignore import litellm from litellm._logging import verbose_logger, verbose_proxy_logger @@ -2074,15 +2075,13 @@ if MCP_AVAILABLE: ) try: result = await create_mcp_toolset(prisma_client, payload, touched_by) - except Exception as e: - if "UniqueViolationError" in type(e).__name__ or "unique" in str(e).lower(): - raise HTTPException( - status_code=status.HTTP_409_CONFLICT, - detail={ - "error": f"A toolset named '{payload.toolset_name}' already exists." - }, - ) - raise + except UniqueViolationError: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail={ + "error": f"A toolset named '{payload.toolset_name}' already exists." + }, + ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5f1d7767b6c..57ffb8e0172 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -13536,7 +13536,8 @@ async def _stream_mcp_asgi_response( # 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) + if task.cancelled() or task.exception() is not None: + body_queue.put_nowait(None) handler_task.add_done_callback(_ensure_eof)