From 1fe238b2311d679903fc7583d30e18106357c0a6 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 10 Mar 2026 23:15:14 -0700 Subject: [PATCH] Add team-scoped MCP server filtering for key creation and fix UnboundLocalError When creating a key, the MCP server list now filters by the selected team's allowed servers. Also fixes UnboundLocalError on `is_restricted_virtual_key` when `team_id` query param was provided to GET /v1/mcp/server. Co-Authored-By: Claude Opus 4.6 --- .../key_management_endpoints.py | 78 +++++++++++++++++++ .../mcp_management_endpoints.py | 4 + .../object_permission_utils.py | 8 +- .../hooks/mcpServers/useMCPServers.ts | 2 +- 4 files changed, 87 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 1c0c212b60b..ff4d861b1b2 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2052,6 +2052,84 @@ async def update_key_fn( user_api_key_cache=user_api_key_cache, ) + # Only check team limits if key has a team_id + team_obj: Optional[LiteLLM_TeamTableCachedObj] = None + if data.team_id is not None: + team_obj = await get_team_object( + team_id=data.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + + if team_obj is not None: + await _check_team_key_limits( + team_table=team_obj, + data=data, + prisma_client=prisma_client, + ) + + # Validate key against project limits if project_id is being set + _project_id_to_check = getattr(data, "project_id", None) or getattr( + existing_key_row, "project_id", None + ) + if _project_id_to_check is not None and ( + data.models is not None or data.max_budget is not None + ): + await _check_project_key_limits( + project_id=_project_id_to_check, + data=data, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + # if team change - check if this is possible + if is_different_team(data=data, existing_key_row=existing_key_row): + if llm_router is None: + raise HTTPException( + status_code=400, + detail={ + "error": "LLM router not found. Please set it up by passing in a valid config.yaml or adding models via the UI." + }, + ) + # team_obj should be set since is_different_team() returns True only when data.team_id is not None + if team_obj is None: + raise HTTPException( + status_code=500, + detail={ + "error": "Team object not found for team change validation" + }, + ) + await validate_key_team_change( + key=existing_key_row, + team=team_obj, + change_initiated_by=user_api_key_dict, + llm_router=llm_router, + ) + + # Set Management Endpoint Metadata Fields + + # Validate MCP servers in object_permission against the effective team + if data.object_permission is not None: + effective_team_obj = team_obj + # If team_id isn't being changed, resolve the existing key's team + if effective_team_obj is None and existing_key_row.team_id: + effective_team_obj = await get_team_object( + team_id=existing_key_row.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + object_permission_dict = ( + data.object_permission.model_dump() + if hasattr(data.object_permission, "model_dump") + else data.object_permission + ) + await validate_key_mcp_servers_against_team( + object_permission=object_permission_dict, + team_obj=effective_team_obj, + ) + non_default_values = await prepare_key_update_data( data=data, existing_key_row=existing_key_row ) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 3e5b729cea6..57c1786b610 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -778,6 +778,10 @@ if MCP_AVAILABLE: aggregated_servers.values() ) + redacted_mcp_servers = _redact_mcp_credentials_list( + aggregated_servers.values() + ) + # augment the mcp servers with public status if litellm.public_mcp_servers is not None: for server in redacted_mcp_servers: diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 8aba8307b9d..164d65c7e64 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -208,10 +208,10 @@ async def _resolve_team_allowed_mcp_servers( ) direct_servers: List[str] = team_object_permission.mcp_servers or [] - access_group_servers: List[ - str - ] = await MCPRequestHandler._get_mcp_servers_from_access_groups( - team_object_permission.mcp_access_groups or [] + access_group_servers: List[str] = ( + await MCPRequestHandler._get_mcp_servers_from_access_groups( + team_object_permission.mcp_access_groups or [] + ) ) raw_tool_perms = team_object_permission.mcp_tool_permissions or {} if isinstance(raw_tool_perms, str): diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts index 9210e25e1a8..a00a16cbeee 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts @@ -9,7 +9,7 @@ const mcpServersKeys = createQueryKeys("mcpServers"); export const useMCPServers = (teamId?: string | null) => { const { accessToken } = useAuthorized(); return useQuery({ - queryKey: mcpServersKeys.list(teamId ? { filters: { teamId } } : undefined), + queryKey: mcpServersKeys.list({ teamId: teamId ?? undefined }), queryFn: async () => await fetchMCPServers(accessToken!, teamId), enabled: !!accessToken, });