scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(mcp): add mcp_toolsets to ObjectPermissionBase, fix multi-toolset overwrite, fix delete 404, allow standalone key toolsets

Browse source
This commit is contained in:
Ishaan Jaffer 2026-03-23 10:45:22 -07:00
parent 65ae86b407
commit e9c4c71a1f
4 changed files with 63 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -97,9 +97,15 @@ async def delete_mcp_toolset(
prisma_client: PrismaClient,
toolset_id: str,
) -> Optional[MCPToolset]:
row = await prisma_client.db.litellm_mcptoolsettable.delete(
where={"toolset_id": toolset_id}
)
if row is None:
return None
try:
row = await prisma_client.db.litellm_mcptoolsettable.delete(
where={"toolset_id": toolset_id}
)
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)

1
litellm/proxy/_types.py
View file

@ -855,6 +855,7 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase):
mcp_servers: Optional[List[str]] = None
mcp_access_groups: Optional[List[str]] = None
mcp_tool_permissions: Optional[Dict[str, List[str]]] = None
mcp_toolsets: Optional[List[str]] = None
vector_stores: Optional[List[str]] = None
agents: Optional[List[str]] = None
agent_access_groups: Optional[List[str]] = None

28
litellm/proxy/management_helpers/object_permission_utils.py
View file

@ -380,30 +380,26 @@ async def validate_key_mcp_servers_against_team(
detail={"error": detail},
)
# Validate requested toolsets (must be subset of team's toolsets)
if requested_toolsets:
# Validate requested toolsets against team's allowed toolsets.
# Only enforce the team-based restriction when a team is present — standalone
# keys (no team) can freely be granted any toolset by an admin.
if requested_toolsets and team_obj is not None:
team_toolsets: Set[str] = set()
if (
team_obj is not None
and team_obj.object_permission is not None
team_obj.object_permission is not None
and team_obj.object_permission.mcp_toolsets
):
team_toolsets = set(team_obj.object_permission.mcp_toolsets)
disallowed_toolsets = requested_toolsets - team_toolsets
if disallowed_toolsets:
if team_obj is not None:
detail = (
f"Key requests MCP toolsets not allowed by team '{team_obj.team_id}': "
f"{sorted(disallowed_toolsets)}. "
f"Team allows: {sorted(team_toolsets)}."
)
else:
detail = (
f"Key is not in a team. MCP toolsets cannot be assigned to "
f"keys outside of a team. Disallowed toolsets: {sorted(disallowed_toolsets)}."
)
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={"error": detail},
detail={
"error": (
f"Key requests MCP toolsets not allowed by team '{team_obj.team_id}': "
f"{sorted(disallowed_toolsets)}. "
f"Team allows: {sorted(team_toolsets)}."
)
},
)

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

@ -160,28 +160,24 @@ class LiteLLM_Proxy_MCP_Handler:
):
mcp_servers.append(server_url.split("/")[-1])
# Resolve toolset names: if any name in mcp_servers is a toolset (not a real
# server name), apply toolset scope to user_api_key_auth so that only the
# toolset's servers and tools are visible. Non-toolset names are kept as-is.
# Resolve toolset names: collect all toolset IDs first, then apply their
# combined permissions in a single pass so multiple toolsets are unioned
# rather than the last one overwriting the others.
resolved_mcp_servers: List[str] = []
resolved_toolset_ids: List[str] = []
for name in mcp_servers:
if not global_mcp_server_manager.get_mcp_server_by_name(name):
try:
from litellm.proxy._experimental.mcp_server.server import (
_apply_toolset_scope,
)
from litellm.proxy.proxy_server import prisma_client
if prisma_client is not None and user_api_key_auth is not None:
if prisma_client is not None:
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
)
resolved_toolset_ids.append(toolset.toolset_id)
# Don't add to resolved_mcp_servers — toolset scope
# restricts via object_permission, not server name filter.
continue
@ -189,6 +185,39 @@ class LiteLLM_Proxy_MCP_Handler:
verbose_logger.debug(f"Could not resolve '{name}' as toolset: {_e}")
resolved_mcp_servers.append(name)
# Apply all resolved toolsets at once (union), avoiding permission overwrite.
if resolved_toolset_ids and user_api_key_auth is not None:
try:
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
tool_permissions = (
await global_mcp_server_manager.resolve_toolset_tool_permissions(
toolset_ids=resolved_toolset_ids
)
)
server_ids = list(tool_permissions.keys())
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_tool_permissions": tool_permissions,
"mcp_toolsets": [],
"mcp_access_groups": [],
}
)
else:
updated_op = LiteLLM_ObjectPermissionTable(
object_permission_id="toolset-scope",
mcp_servers=server_ids,
mcp_tool_permissions=tool_permissions,
)
user_api_key_auth = user_api_key_auth.model_copy(
update={"object_permission": updated_op}
)
except Exception as _e:
verbose_logger.debug(f"Could not apply toolset permissions: {_e}")
tools = await _get_tools_from_mcp_servers(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,