scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 117

fix(mcp): handle RecordNotFoundError in update_mcp_toolset; union direct servers with toolset servers

Browse source
This commit is contained in:
Ishaan Jaffer 2026-03-23 11:05:32 -07:00
parent 97961237f1
commit 936f6b8ad5
2 changed files with 20 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
litellm/proxy/_experimental/mcp_server/toolset_db.py
View file

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

10
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

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