fix: use UniqueViolationError import; guard _ensure_eof for error/cancel only

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,
         )

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)