From 936f6b8ad5280fcd2eca4b34bd63f812d5a261ce Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Mon, 23 Mar 2026 11:05:32 -0700 Subject: [PATCH] fix(mcp): handle RecordNotFoundError in update_mcp_toolset; union direct servers with toolset servers --- .../_experimental/mcp_server/toolset_db.py | 18 +++++++++++++----- .../responses/mcp/litellm_proxy_mcp_handler.py | 10 +++++++--- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index c424f51774e..adea36d2dea 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -81,15 +81,23 @@ async def update_mcp_toolset( prisma_client: PrismaClient, data: UpdateMCPToolsetRequest, touched_by: str, -) -> MCPToolset: +) -> Optional[MCPToolset]: data_dict = data.model_dump(exclude_none=True, exclude={"toolset_id"}) if "tools" in data_dict: data_dict["tools"] = json.dumps(data_dict["tools"]) data_dict["updated_by"] = touched_by - row = await prisma_client.db.litellm_mcptoolsettable.update( - where={"toolset_id": data.toolset_id}, - data=data_dict, - ) + try: + row = await prisma_client.db.litellm_mcptoolsettable.update( + where={"toolset_id": data.toolset_id}, + data=data_dict, + ) + except Exception as e: + if ( + "RecordNotFoundError" in type(e).__name__ + or "record was not found" in str(e).lower() + ): + return None + raise return _toolset_from_row(row) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 8cf195494c1..6d0b6b87ce6 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -216,12 +216,16 @@ class LiteLLM_Proxy_MCP_Handler: toolset_ids=resolved_toolset_ids ) ) - server_ids = list(tool_permissions.keys()) + # Union toolset server IDs with direct servers the user also requested, + # so explicitly-selected servers aren't dropped by downstream permission filtering. + all_server_ids = list( + set(tool_permissions.keys()) | set(resolved_mcp_servers) + ) existing_op = user_api_key_auth.object_permission if existing_op is not None: updated_op = existing_op.model_copy( update={ - "mcp_servers": server_ids, + "mcp_servers": all_server_ids, "mcp_tool_permissions": tool_permissions, "mcp_toolsets": [], "mcp_access_groups": [], @@ -230,7 +234,7 @@ class LiteLLM_Proxy_MCP_Handler: else: updated_op = LiteLLM_ObjectPermissionTable( object_permission_id="toolset-scope", - mcp_servers=server_ids, + mcp_servers=all_server_ids, mcp_tool_permissions=tool_permissions, ) user_api_key_auth = user_api_key_auth.model_copy(