From 0b07f628ffb2d4fe8d2c271fe578a9a2796382ec Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 10:30:03 -0700 Subject: [PATCH 01/26] [Test] UI: Add vitest coverage for 10 previously untested components Add unit tests for: - SimpleToolCallBlock, SimpleMessageBlock, CollapsibleMessage, HistoryTree (log details drawer) - OnboardingForm (onboarding flow) - TeamsHeaderTabs, TeamsTable (teams page) - transform_key_info, filter_helpers (key/team helpers) - queryKeysFactory (query key generation utility) 47 new tests covering conditional rendering, user interactions, data transformation, and error handling. Co-Authored-By: Claude Opus 4.6 --- .../hooks/common/queryKeysFactory.test.ts | 34 +++++ .../teams/components/TeamsHeaderTabs.test.tsx | 54 ++++++++ .../components/TeamsTable/TeamsTable.test.tsx | 129 ++++++++++++++++++ .../app/onboarding/OnboardingForm.test.tsx | 95 +++++++++++++ .../key_team_helpers/filter_helpers.test.ts | 90 ++++++++++++ .../transform_key_info.test.ts | 62 +++++++++ .../CollapsibleMessage.test.tsx | 54 ++++++++ .../LogDetailsDrawer/HistoryTree.test.tsx | 50 +++++++ .../SimpleMessageBlock.test.tsx | 55 ++++++++ .../SimpleToolCallBlock.test.tsx | 51 +++++++ 10 files changed, 674 insertions(+) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsHeaderTabs.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/TeamsTable.test.tsx create mode 100644 ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.test.ts create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/transform_key_info.test.ts create mode 100644 ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/CollapsibleMessage.test.tsx create mode 100644 ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/HistoryTree.test.tsx create mode 100644 ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SimpleMessageBlock.test.tsx create mode 100644 ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SimpleToolCallBlock.test.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.test.ts new file mode 100644 index 0000000000..39afd04409 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.test.ts @@ -0,0 +1,34 @@ +import { describe, it, expect } from "vitest"; +import { createQueryKeys } from "./queryKeysFactory"; + +describe("createQueryKeys", () => { + const keys = createQueryKeys("books"); + + it("should return the resource name as the base key", () => { + expect(keys.all).toEqual(["books"]); + }); + + it("should generate a lists key", () => { + expect(keys.lists()).toEqual(["books", "list"]); + }); + + it("should generate a list key with params", () => { + expect(keys.list({ page: 1, limit: 10 })).toEqual([ + "books", + "list", + { params: { page: 1, limit: 10 } }, + ]); + }); + + it("should generate a list key with undefined params when none provided", () => { + expect(keys.list()).toEqual(["books", "list", { params: undefined }]); + }); + + it("should generate a details key", () => { + expect(keys.details()).toEqual(["books", "detail"]); + }); + + it("should generate a detail key for a specific ID", () => { + expect(keys.detail("123")).toEqual(["books", "detail", "123"]); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsHeaderTabs.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsHeaderTabs.test.tsx new file mode 100644 index 0000000000..50a7f10f04 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsHeaderTabs.test.tsx @@ -0,0 +1,54 @@ +import { render, screen } from "@testing-library/react"; +import React from "react"; +import { describe, expect, it, vi } from "vitest"; +import TeamsHeaderTabs from "./TeamsHeaderTabs"; + +vi.mock("@tremor/react", () => ({ + TabGroup: ({ children, ...props }: any) =>
{children}
, + TabList: ({ children, ...props }: any) =>
{children}
, + Tab: ({ children, ...props }: any) => , + TabPanels: ({ children, ...props }: any) =>
{children}
, + Text: ({ children, ...props }: any) => {children}, + Icon: ({ onClick, ...props }: any) => + + ); + } + + return ( + + columns={teamColumns} + dataSource={displayTeams} + rowKey="team_id" + pagination={false} + onChange={handleTableSort} + locale={{ + emptyText: ( +
+ +
+ No teams yet +
+
+ + Create your first team to organize members and manage access to models. + +
+ {canCreateOrManageTeams(userRole, userID, organizations) && ( + + )} +
+ ), + }} + scroll={{ x: 1000 }} + size="middle" + /> + ); + }; + + const tabItems = [ + { + key: "your-teams", + label: "Your Teams", + children: ( + <> + + + + } + suffix={isSearching ? : null} + placeholder="Search teams by name..." + onChange={(e) => handleSearchChange(e.target.value)} + allowClear + style={{ maxWidth: 400 }} + /> + handleFilterChange("organization_id", value || "")} + loading={isLoading} + /> + + { + setCurrentPage(page); + setPageSize(size); + fetchTeamsV2({ page, size }); + }} + size="small" + showTotal={(total) => `${total} teams`} + showSizeChanger + pageSizeOptions={["10", "20", "50"]} + /> + + + {renderTeamsContent()} + + + + + ), + }, + { + key: "available-teams", + label: "Available Teams", + children: , + }, + ...(isProxyAdminRole(userRole || "") + ? [ + { + key: "default-settings", + label: "Default Team Settings", + children: , + }, + ] + : []), + ]; + return ( -
- - - {canCreateOrManageTeams(userRole, userID, organizations) && ( - - )} - {selectedTeamId ? ( - { - setTeams((teams) => { - if (teams == null) { - return teams; - } - const updated = teams.map((team) => { - if (data.team_id === team.team_id) { - return updateExistingKeys(team, data); - } - return team; - }); - // Minimal fix: refresh the full team list after an update - if (accessToken) { - fetchTeams(accessToken, userID, userRole, currentOrg, setTeams); - } - return updated; - }); - }} - onClose={() => { - setSelectedTeamId(null); - setEditTeam(false); - }} - accessToken={accessToken} - is_team_admin={is_team_admin(teams?.find((team) => team.team_id === selectedTeamId))} - is_proxy_admin={userRole == "Admin"} - userModels={userModels} - editTeam={editTeam} - premiumUser={premiumUser} - /> - ) : ( - - -
- Your Teams - Available Teams - {isProxyAdminRole(userRole || "") && Default Team Settings} -
-
- {lastRefreshed && Last Refreshed: {lastRefreshed}} - -
-
- - - - Click on “Team ID” to view team details and manage team members. - - - - -
-
- {/* Search and Filter Controls */} -
- {/* Team Alias Search */} - handleFilterChange("team_alias", value)} - icon={Search} - /> + + {selectedTeamId ? ( + { + setTeams((teams) => { + if (teams == null) { + return teams; + } + return teams.map((team) => { + if (data.team_id === team.team_id) { + return updateExistingKeys(team, data); + } + return team; + }); + }); + fetchTeamsV2(); + }} + onClose={() => { + setSelectedTeamId(null); + setEditTeam(false); + }} + accessToken={accessToken} + is_team_admin={is_team_admin(teams?.find((team) => team.team_id === selectedTeamId))} + is_proxy_admin={userRole == "Admin"} + userModels={userModels} + editTeam={editTeam} + premiumUser={premiumUser} + /> + ) : ( + <> + + + + <TeamOutlined style={{ marginRight: 8 }} /> + Teams + + + Manage teams, members, and their access to models and budgets + + + {canCreateOrManageTeams(userRole, userID, organizations) && ( + + )} + - {/* Filter Button */} - setShowFilters(!showFilters)} - active={showFilters} - hasActiveFilters={!!(filters.team_id || filters.team_alias || filters.organization_id)} - /> + + + )} - {/* Reset Filters Button */} - -
- - {/* Additional Filters */} - {showFilters && ( -
- {/* Team ID Search */} - handleFilterChange("team_id", value)} - icon={User} - /> - - {/* Organization Dropdown */} -
- -
-
- )} -
-
- - - - Team Name - Team ID - Created - Spend (USD) - Budget (USD) - Models - Organization - Info - Actions - - - - - {teams && teams.length > 0 ? ( - teams - .filter((team) => { - if (!currentOrg) return true; - return team.organization_id === currentOrg.organization_id; - }) - .sort((a, b) => new Date(b.created_at).getTime() - new Date(a.created_at).getTime()) - .map((team: any) => ( - - - {team["team_alias"]} - - -
- - - -
-
- - {team.created_at ? new Date(team.created_at).toLocaleDateString() : "N/A"} - - - {formatNumberWithCommas(team["spend"], 4)} - - - {team["max_budget"] !== null && team["max_budget"] !== undefined - ? team["max_budget"] - : "No limit"} - - 3 ? "px-0" : ""} - > -
- {Array.isArray(team.models) ? ( -
- {team.models.length === 0 ? ( - - All Proxy Models - - ) : ( - <> -
- {team.models.length > 3 && ( -
- { - setExpandedAccordions((prev) => ({ - ...prev, - [team.team_id]: !prev[team.team_id], - })); - }} - /> -
- )} -
- {team.models.slice(0, 3).map((model: string, index: number) => - model === "all-proxy-models" ? ( - - All Proxy Models - - ) : ( - - - {model.length > 30 - ? `${getModelDisplayName(model).slice(0, 30)}...` - : getModelDisplayName(model)} - - - ), - )} - {team.models.length > 3 && !expandedAccordions[team.team_id] && ( - - - +{team.models.length - 3}{" "} - {team.models.length - 3 === 1 ? "more model" : "more models"} - - - )} - {expandedAccordions[team.team_id] && ( -
- {team.models.slice(3).map((model: string, index: number) => - model === "all-proxy-models" ? ( - - All Proxy Models - - ) : ( - - - {model.length > 30 - ? `${getModelDisplayName(model).slice(0, 30)}...` - : getModelDisplayName(model)} - - - ), - )} -
- )} -
-
- - )} -
- ) : null} -
-
- - - {getOrganizationAlias(team.organization_id, organizationsData || organizations)} - - - - {perTeamInfo && - team.team_id && - perTeamInfo[team.team_id] && - perTeamInfo[team.team_id].keys && - perTeamInfo[team.team_id].keys.length}{" "} - Keys - - - {perTeamInfo && - team.team_id && - perTeamInfo[team.team_id] && - perTeamInfo[team.team_id].team_info && - perTeamInfo[team.team_id].team_info.members_with_roles && - perTeamInfo[team.team_id].team_info.members_with_roles.length}{" "} - Members - - - - {userRole == "Admin" ? ( - <> - { - setSelectedTeamId(team.team_id); - setEditTeam(true); - }} - dataTestId="edit-team-button" - tooltipText="Edit team" - /> - handleDelete(team)} - dataTestId="delete-team-button" - tooltipText="Delete team" - /> - - ) : null} - -
- )) - ) : ( - - -
- No teams found - Adjust your filters or create a new team -
-
-
- )} -
-
- -
- -
-
- - - - {isProxyAdminRole(userRole || "") && ( - - - - )} -
-
- )} - {canCreateOrManageTeams(userRole, userID, organizations) && ( + {canCreateOrManageTeams(userRole, userID, organizations) && ( = ({ : "" } > - = ({ optionFilterProp="children" > {adminOrgs?.map((org) => ( - + {org.organization_alias}{" "} ({org.organization_id}) - + ))} - + {/* Show message when org admin needs to select organization */} {isOrgAdmin && !isSingleOrg && adminOrgs.length > 1 && (
- + Please select an organization to create a team for. You can only create teams within organizations where you are an admin. @@ -1190,11 +1211,11 @@ const Teams: React.FC = ({ - - daily - weekly - monthly - + @@ -1313,7 +1334,7 @@ const Teams: React.FC = ({ className="mt-8" help="Select existing guardrails or enter new ones" > - = ({ className="mt-8" help="Select existing policies or enter new ones" > - = ({
- + Create custom aliases for models that can be used by team members in API calls. This allows you to create shortcuts for specific models. @@ -1548,14 +1569,12 @@ const Teams: React.FC = ({
- Create Team +
)} - - -
+ ); }; diff --git a/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx b/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx index 488913a734..2f146aab72 100644 --- a/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx +++ b/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx @@ -6,6 +6,7 @@ import { ChevronUpIcon, ChevronDownIcon, ExternalLinkIcon, + ClipboardCopyIcon, } from "@heroicons/react/outline"; import { Tooltip } from "antd"; import BaseActionButton from "../BaseActionButton"; @@ -32,6 +33,7 @@ export const TableIconActionButtonMap: Record void; disabled?: boolean; loading?: boolean; + style?: React.CSSProperties; } const OrganizationDropdown: React.FC = ({ @@ -16,16 +19,18 @@ const OrganizationDropdown: React.FC = ({ onChange, disabled, loading, + style, }) => { return ( diff --git a/ui/litellm-dashboard/src/components/ui/AntDLoadingSpinner.tsx b/ui/litellm-dashboard/src/components/ui/AntDLoadingSpinner.tsx new file mode 100644 index 0000000000..9e90f77584 --- /dev/null +++ b/ui/litellm-dashboard/src/components/ui/AntDLoadingSpinner.tsx @@ -0,0 +1,12 @@ +import { Spin } from "antd"; +import { LoadingOutlined } from "@ant-design/icons"; + +interface AntDLoadingSpinnerProps { + size?: "small" | "default" | "large"; + fontSize?: number; +} + +export function AntDLoadingSpinner({ size, fontSize }: AntDLoadingSpinnerProps) { + const indicator = ; + return ; +} From f4e5fb2b4ae4d388d62e66eaba64c33a162bf8d6 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 22:50:35 -0700 Subject: [PATCH 14/26] fix: exclude team_id from Prisma create payload in _prepare_mcp_server_data team_id is a request-level field, not a DB column. Excluding it prevents Prisma from rejecting the create call. Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/_experimental/mcp_server/db.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index fbef33c32e..0bc102889b 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -40,8 +40,8 @@ def _prepare_mcp_server_data( """ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - # Convert model to dict - data_dict = data.model_dump(exclude_none=True) + # Convert model to dict, excluding fields not in the DB schema + data_dict = data.model_dump(exclude_none=True, exclude={"team_id"}) # Ensure alias is always present in the dict (even if None) if "alias" not in data_dict: data_dict["alias"] = getattr(data, "alias", None) From 6c06e6dc80fd7a2870a3f6fef57a8d651de1a82b Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 22:57:27 -0700 Subject: [PATCH 15/26] fix: link new object_permission_id back to team on first MCP server create When a team has no object_permission_id yet, the auto-assign logic creates an ObjectPermissionTable row but never linked it to the team. Now updates the team's object_permission_id after creation. Co-Authored-By: Claude Opus 4.6 --- .../management_endpoints/mcp_management_endpoints.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 3400f90f48..522a533d43 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1358,7 +1358,7 @@ if MCP_AVAILABLE: set(existing_mcp_servers + [new_mcp_server.server_id]) ) - await handle_update_object_permission_common( + new_permission_id = await handle_update_object_permission_common( data_json={ "object_permission": { "mcp_servers": updated_mcp_servers, @@ -1367,6 +1367,13 @@ if MCP_AVAILABLE: existing_object_permission_id=existing_permission_id, prisma_client=prisma_client, ) + + # If the team had no object_permission_id, link the new one + if existing_permission_id is None and new_permission_id is not None: + await prisma_client.db.litellm_teamtable.update( + where={"team_id": manager_team_id}, + data={"object_permission_id": new_permission_id}, + ) except Exception as e: verbose_proxy_logger.exception(f"Error creating mcp server: {str(e)}") raise HTTPException( From dfc4401e1c305f11c898e989991503122668bede Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 23:00:38 -0700 Subject: [PATCH 16/26] fix: use check_db_only=True in _assert_can_manage_team_mcp_server Cached team objects may have stale object_permission data. Bypassing cache ensures the server-in-team check uses fresh data. Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/management_endpoints/mcp_management_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 522a533d43..a3b0b3ab28 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -179,7 +179,7 @@ if MCP_AVAILABLE: team_id=resolved_team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, - check_db_only=False, + check_db_only=True, # bypass cache to get fresh object_permission ) if not _is_user_team_mcp_manager(user_api_key_dict, team_obj): From e71410470267f64a00f61d7b75094a8628054b0e Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 23:10:44 -0700 Subject: [PATCH 17/26] fix: address Greptile review feedback - P0: Separate try/except for auto-assign so server creation succeeds even if team permission update fails - P1: Clean up team permission entry on MCP server delete - P1: Add MCP_AVAILABLE skip guard to tests - P2: Return team_obj from _assert_can_manage_team_mcp_server to eliminate redundant get_team_object call in create endpoint Co-Authored-By: Claude Opus 4.6 --- .../mcp_management_endpoints.py | 80 +++++++++++++------ .../test_mcp_manager_role.py | 21 +++-- 2 files changed, 72 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index a3b0b3ab28..f137902f1b 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -21,7 +21,7 @@ import json import os from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Iterable, List, Literal, Optional +from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple from fastapi import ( APIRouter, @@ -147,12 +147,12 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth, team_id: Optional[str] = None, server_id: Optional[str] = None, - ) -> str: + ) -> Tuple[str, "LiteLLM_TeamTableCachedObj"]: """ Verify that the caller is an MCP server manager for a team and (for edit/delete) that the target server belongs to that team. - Returns the team_id the caller is managing. + Returns a tuple of (team_id, team_obj) for downstream use. Raises HTTPException(400) if no team_id can be determined. Raises HTTPException(403) if the caller is not an MCP manager or server not in team. """ @@ -165,6 +165,7 @@ if MCP_AVAILABLE: detail={"error": "team_id is required for MCP server manager operations."}, ) + # When the API key is team-scoped, ensure the request team_id matches if ( team_id and user_api_key_dict.team_id @@ -200,7 +201,7 @@ if MCP_AVAILABLE: }, ) - return resolved_team_id + return resolved_team_id, team_obj @dataclass class _TemporaryMCPServerEntry: @@ -1277,6 +1278,7 @@ if MCP_AVAILABLE: is_proxy_admin = LitellmUserRoles.PROXY_ADMIN == user_api_key_dict.user_role manager_team_id: Optional[str] = None + manager_team_obj = None if not is_proxy_admin: # Check if the user is an MCP manager for the specified team if payload.team_id is None: @@ -1286,7 +1288,7 @@ if MCP_AVAILABLE: "error": "team_id is required when creating MCP servers as a team MCP manager." }, ) - manager_team_id = await _assert_can_manage_team_mcp_server( + manager_team_id, manager_team_obj = await _assert_can_manage_team_mcp_server( user_api_key_dict=user_api_key_dict, team_id=payload.team_id, ) @@ -1334,25 +1336,26 @@ if MCP_AVAILABLE: # Ensure registry is up to date by reloading from database await global_mcp_server_manager.reload_servers_from_database() - # If created by an MCP manager, auto-assign the server to their team - if manager_team_id is not None: - from litellm.proxy.proxy_server import user_api_key_cache + except Exception as e: + verbose_proxy_logger.exception(f"Error creating mcp server: {str(e)}") + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={"error": f"Error creating mcp server: {str(e)}"}, + ) - team_obj = await get_team_object( - team_id=manager_team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=False, - ) + # Auto-assign the server to the manager's team (separate from create + # so a failure here doesn't mask the successfully created server). + if manager_team_id is not None and manager_team_obj is not None: + try: existing_permission_id = getattr( - team_obj, "object_permission_id", None + manager_team_obj, "object_permission_id", None ) # Read existing mcp_servers list and append the new server existing_mcp_servers: list = [] - if team_obj.object_permission is not None: + if manager_team_obj.object_permission is not None: existing_mcp_servers = ( - team_obj.object_permission.mcp_servers or [] + manager_team_obj.object_permission.mcp_servers or [] ) updated_mcp_servers = list( set(existing_mcp_servers + [new_mcp_server.server_id]) @@ -1374,12 +1377,10 @@ if MCP_AVAILABLE: where={"team_id": manager_team_id}, data={"object_permission_id": new_permission_id}, ) - except Exception as e: - verbose_proxy_logger.exception(f"Error creating mcp server: {str(e)}") - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Error creating mcp server: {str(e)}"}, - ) + except Exception as e: + verbose_proxy_logger.exception( + f"MCP server created but failed to auto-assign to team {manager_team_id}: {str(e)}" + ) return _redact_mcp_credentials(new_mcp_server) @router.post( @@ -1593,8 +1594,9 @@ if MCP_AVAILABLE: ) # Authz - proxy admins or team MCP managers can delete MCP servers + manager_team_obj = None if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: - await _assert_can_manage_team_mcp_server( + _, manager_team_obj = await _assert_can_manage_team_mcp_server( user_api_key_dict=user_api_key_dict, server_id=server_id, ) @@ -1612,6 +1614,36 @@ if MCP_AVAILABLE: # Ensure registry is up to date by reloading from database await global_mcp_server_manager.reload_servers_from_database() + # Remove server from the manager's team permission list + if manager_team_obj is not None: + try: + existing_permission_id = getattr( + manager_team_obj, "object_permission_id", None + ) + if ( + existing_permission_id is not None + and manager_team_obj.object_permission is not None + ): + existing_mcp_servers = ( + manager_team_obj.object_permission.mcp_servers or [] + ) + updated_mcp_servers = [ + s for s in existing_mcp_servers if s != server_id + ] + await handle_update_object_permission_common( + data_json={ + "object_permission": { + "mcp_servers": updated_mcp_servers, + } + }, + existing_object_permission_id=existing_permission_id, + prisma_client=prisma_client, + ) + except Exception as e: + verbose_proxy_logger.exception( + f"MCP server deleted but failed to remove from team permissions: {str(e)}" + ) + # TODO: Enterprise: Finish audit log trail if litellm.store_audit_logs: pass diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py index 860b255f9c..f4e91157f2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py @@ -9,6 +9,9 @@ from litellm.proxy._types import ( from litellm.proxy.management_endpoints.common_utils import ( _is_user_team_mcp_manager, ) +from litellm.proxy.management_endpoints import ( + mcp_management_endpoints as mgmt_endpoints, +) class TestIsUserTeamMcpManager: @@ -70,6 +73,9 @@ from litellm.proxy._types import LiteLLM_TeamTableCachedObj @pytest.mark.asyncio +@pytest.mark.skipif( + not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed" +) class TestAssertCanManageTeamMcpServer: async def test_mcp_manager_with_team_id_succeeds(self): from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -86,10 +92,11 @@ class TestAssertCanManageTeamMcpServer: "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object", AsyncMock(return_value=mock_team), ): - result = await _assert_can_manage_team_mcp_server( + team_id, team_obj = await _assert_can_manage_team_mcp_server( user_api_key_dict=user_auth, team_id="team1" ) - assert result == "team1" + assert team_id == "team1" + assert team_obj == mock_team async def test_regular_user_gets_403(self): from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -154,10 +161,11 @@ class TestAssertCanManageTeamMcpServer: AsyncMock(return_value={"server1", "server2"}), ), ): - result = await _assert_can_manage_team_mcp_server( + team_id, team_obj = await _assert_can_manage_team_mcp_server( user_api_key_dict=user_auth, server_id="server1" ) - assert result == "team1" + assert team_id == "team1" + assert team_obj == mock_team async def test_mcp_manager_server_not_in_team_gets_403(self): from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -188,6 +196,9 @@ class TestAssertCanManageTeamMcpServer: @pytest.mark.asyncio +@pytest.mark.skipif( + not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed" +) class TestCreateMcpServerAsManager: async def test_create_auto_assigns_to_team(self): """MCP manager creating a server should auto-assign it to their team's ObjectPermissionTable.""" @@ -234,7 +245,7 @@ class TestCreateMcpServerAsManager: ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", - AsyncMock(return_value="team1"), + AsyncMock(return_value=("team1", mock_team)), ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", From 65bfd449e980ac08e6339af7a569a69dfa38b192 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 23:23:18 -0700 Subject: [PATCH 18/26] address greptile review feedback (greploop iteration 2) - Remove dead get_team_object mock in test (now reuses team_obj from _assert) - Add test for existing_permission_id=None branch (team linkage) - Remaining P1s are by-design per spec (admin blocked from MCP, team_id ignored for proxy admin) Co-Authored-By: Claude Opus 4.6 --- .../test_mcp_manager_role.py | 81 ++++++++++++++++++- 1 file changed, 77 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py index f4e91157f2..67e9dda4ce 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py @@ -262,10 +262,6 @@ class TestCreateMcpServerAsManager: reload_servers_from_database=AsyncMock(), ), ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object", - AsyncMock(return_value=mock_team), - ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common", mock_handle_update, @@ -280,3 +276,80 @@ class TestCreateMcpServerAsManager: assert "existing_server" in mcp_servers assert "new_server_id" in mcp_servers assert call_kwargs["existing_object_permission_id"] == "perm1" + + async def test_create_links_new_permission_to_team_when_none_exists(self): + """When team has no object_permission_id, create should link the new one.""" + from litellm.proxy._types import NewMCPServerRequest + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + add_mcp_server, + ) + + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user1", + api_key="sk-test", + team_id="team1", + ) + + payload = NewMCPServerRequest( + server_name="test_server", + url="https://example.com/mcp", + team_id="team1", + ) + + # Team with NO object_permission_id + mock_team = LiteLLM_TeamTableCachedObj( + team_id="team1", + members_with_roles=[ + Member(user_id="user1", role="mcp_server_manager"), + ], + object_permission_id=None, + ) + mock_team.object_permission = None + + created_server = MagicMock() + created_server.server_id = "new_server_id" + created_server.credentials = None + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.update = AsyncMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_prisma, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", + AsyncMock(return_value=("team1", mock_team)), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server", + AsyncMock(return_value=created_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + MagicMock( + add_server=AsyncMock(), + reload_servers_from_database=AsyncMock(), + ), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common", + AsyncMock(return_value="new_perm_id"), + ), + ): + await add_mcp_server(payload=payload, user_api_key_dict=user_auth) + + # Verify the team was updated with the new object_permission_id + mock_prisma.db.litellm_teamtable.update.assert_called_once_with( + where={"team_id": "team1"}, + data={"object_permission_id": "new_perm_id"}, + ) From 34d954f8cb9d02c5c4176368a9baf0dc0dbb5e8d Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 19 Mar 2026 23:29:38 -0700 Subject: [PATCH 19/26] [Fix] UI: Migrate AntD message API to use context-based MessageManager AntD v5 static message API doesn't render without an App wrapper or useMessage() context holder. Mirrors the existing notification pattern by adding message.useMessage() to AntdGlobalProvider and routing all calls through a new MessageManager module. Co-Authored-By: Claude Opus 4.6 --- .../AccessGroupCreateModal.tsx | 5 ++- .../AccessGroupEditModal.tsx | 5 ++- .../src/components/BulkEditUsers.tsx | 4 +- .../CloudZeroCreateModal.tsx | 9 +++-- .../CloudZeroIntegrationSettings.tsx | 15 ++++---- .../CloudZeroUpdateModal.tsx | 9 +++-- .../ProjectModals/CreateProjectModal.tsx | 7 ++-- .../ProjectModals/EditProjectModal.tsx | 7 ++-- .../SearchTools/SearchToolTester.tsx | 5 ++- .../RouterSettings/Fallbacks/AddFallbacks.tsx | 5 ++- .../Fallbacks/FallbackSelectionForm.tsx | 5 ++- .../src/components/agents/add_agent_form.tsx | 11 +++--- .../src/components/agents/agent_info.tsx | 9 +++-- .../src/components/chat/ChatPage.tsx | 5 ++- .../src/components/chat/MCPAppsPanel.tsx | 6 +-- .../src/components/chat/MCPConnectPicker.tsx | 7 ++-- .../src/components/chat/MCPCredentialsTab.tsx | 5 ++- .../claude_code_plugins/add_plugin_form.tsx | 17 +++++---- .../mcp_tools/ByokCredentialModal.tsx | 9 +++-- .../components/molecules/message_manager.tsx | 38 +++++++++++++++++++ .../src/components/networking.tsx | 4 +- .../organisms/create_key_button.tsx | 2 +- .../chat_ui/CodeInterpreterTool.tsx | 5 ++- .../src/components/policies/index.tsx | 29 +++++++------- .../policies/pipeline_flow_builder.tsx | 9 +++-- .../components/policies/policy_templates.tsx | 5 ++- .../shared/CreatedKeyDisplay.test.tsx | 14 +++---- .../components/shared/CreatedKeyDisplay.tsx | 5 ++- .../src/components/team/TeamInfo.tsx | 7 ++-- .../CreateVectorStore.tsx | 19 +++++----- .../DocumentsTable.tsx | 5 ++- .../VectorStoreTester.tsx | 5 ++- .../view_logs/LogDetailsDrawer/InputCard.tsx | 4 +- .../view_logs/LogDetailsDrawer/OutputCard.tsx | 5 ++- .../src/contexts/AntdGlobalProvider.tsx | 14 ++++--- 35 files changed, 190 insertions(+), 125 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/molecules/message_manager.tsx diff --git a/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupCreateModal.tsx b/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupCreateModal.tsx index b51fbd5bd9..5ec9f04f3b 100644 --- a/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupCreateModal.tsx +++ b/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupCreateModal.tsx @@ -1,5 +1,6 @@ import React from "react"; -import { Modal, Form, message } from "antd"; +import { Modal, Form } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { AccessGroupBaseForm, AccessGroupFormValues, @@ -37,7 +38,7 @@ export function AccessGroupCreateModal({ createMutation.mutate(params, { onSuccess: () => { - message.success("Access group created successfully"); + MessageManager.success("Access group created successfully"); form.resetFields(); onSuccess?.(); onCancel(); diff --git a/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupEditModal.tsx b/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupEditModal.tsx index 919295b6f7..f05edb3fe9 100644 --- a/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupEditModal.tsx +++ b/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupEditModal.tsx @@ -1,5 +1,6 @@ import React, { useEffect } from "react"; -import { Modal, Form, message } from "antd"; +import { Modal, Form } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { AccessGroupBaseForm, AccessGroupFormValues, @@ -55,7 +56,7 @@ export function AccessGroupEditModal({ { accessGroupId: accessGroup.access_group_id, params }, { onSuccess: () => { - message.success("Access group updated successfully"); + MessageManager.success("Access group updated successfully"); onSuccess?.(); onCancel(); }, diff --git a/ui/litellm-dashboard/src/components/BulkEditUsers.tsx b/ui/litellm-dashboard/src/components/BulkEditUsers.tsx index 2f3e57ff2a..2ca69065bc 100644 --- a/ui/litellm-dashboard/src/components/BulkEditUsers.tsx +++ b/ui/litellm-dashboard/src/components/BulkEditUsers.tsx @@ -3,7 +3,6 @@ import { Modal, Typography, Divider, - message, Table, Select, InputNumber, @@ -14,6 +13,7 @@ import { import { userBulkUpdateUserCall, teamBulkMemberAddCall, Member } from "./networking"; import { UserEditView } from "./user_edit_view"; import NotificationsManager from "./molecules/notifications_manager"; +import MessageManager from "@/components/molecules/message_manager"; const { Text, Title } = Typography; @@ -188,7 +188,7 @@ const BulkEditUserModal: React.FC = ({ } if (failedTeams.length > 0) { - message.warning(`Failed to add users to ${failedTeams.length} team(s)`); + MessageManager.warning(`Failed to add users to ${failedTeams.length} team(s)`); } } diff --git a/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroCreateModal.tsx b/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroCreateModal.tsx index feb00fc040..5f460cb7cb 100644 --- a/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroCreateModal.tsx +++ b/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroCreateModal.tsx @@ -1,4 +1,5 @@ -import { Form, Modal, Input, message } from "antd"; +import { Form, Modal, Input } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { useEffect } from "react"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useCloudZeroCreate } from "@/app/(dashboard)/hooks/cloudzero/useCloudZeroCreate"; @@ -31,7 +32,7 @@ export default function CloudZeroCreationModal({ open, onOk, onCancel }: CloudZe }, { onSuccess: () => { - message.success("CloudZero integration created successfully"); + MessageManager.success("CloudZero integration created successfully"); form.resetFields(); onOk(); }, @@ -39,7 +40,7 @@ export default function CloudZeroCreationModal({ open, onOk, onCancel }: CloudZe if (error?.errorFields) { return; } - message.error(error?.message || "Failed to create CloudZero integration"); + MessageManager.error(error?.message || "Failed to create CloudZero integration"); }, }, ); @@ -47,7 +48,7 @@ export default function CloudZeroCreationModal({ open, onOk, onCancel }: CloudZe if (error?.errorFields) { return; } - message.error(error?.message || "Failed to create CloudZero integration"); + MessageManager.error(error?.message || "Failed to create CloudZero integration"); } }; diff --git a/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroIntegrationSettings.tsx b/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroIntegrationSettings.tsx index c161d241f7..62709dac52 100644 --- a/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroIntegrationSettings.tsx +++ b/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroIntegrationSettings.tsx @@ -3,7 +3,8 @@ import { useCloudZeroExport } from "@/app/(dashboard)/hooks/cloudzero/useCloudZe import { useCloudZeroDeleteSettings } from "@/app/(dashboard)/hooks/cloudzero/useCloudZeroSettings"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; -import { Alert, Button, Card, Descriptions, Divider, message, Popconfirm, Tag } from "antd"; +import { Alert, Button, Card, Descriptions, Divider, Popconfirm, Tag } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { CheckCircle, Edit, Play, Trash2, Upload } from "lucide-react"; import { useState } from "react"; import CloudZeroUpdateModal from "./CloudZeroUpdateModal"; @@ -30,10 +31,10 @@ export function CloudZeroIntegrationSettings({ settings, onSettingsUpdated }: Cl { limit: 10 }, { onSuccess: (data) => { - message.success("Dry run completed successfully"); + MessageManager.success("Dry run completed successfully"); }, onError: (error) => { - message.error(error?.message || "Failed to perform dry run"); + MessageManager.error(error?.message || "Failed to perform dry run"); }, }, ); @@ -48,10 +49,10 @@ export function CloudZeroIntegrationSettings({ settings, onSettingsUpdated }: Cl { operation: "replace_hourly" }, { onSuccess: () => { - message.success("Data successfully exported to CloudZero"); + MessageManager.success("Data successfully exported to CloudZero"); }, onError: (error) => { - message.error(error?.message || "Failed to export data"); + MessageManager.error(error?.message || "Failed to export data"); }, }, ); @@ -79,12 +80,12 @@ export function CloudZeroIntegrationSettings({ settings, onSettingsUpdated }: Cl deleteMutation.mutate(undefined, { onSuccess: () => { - message.success("CloudZero integration deleted successfully"); + MessageManager.success("CloudZero integration deleted successfully"); setIsDeleteModalOpen(false); onSettingsUpdated(); }, onError: (error) => { - message.error(error?.message || "Failed to delete CloudZero integration"); + MessageManager.error(error?.message || "Failed to delete CloudZero integration"); }, }); }; diff --git a/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroUpdateModal.tsx b/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroUpdateModal.tsx index 0aca6857b8..a04007897c 100644 --- a/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroUpdateModal.tsx +++ b/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroUpdateModal.tsx @@ -1,6 +1,7 @@ import { useCloudZeroUpdateSettings } from "@/app/(dashboard)/hooks/cloudzero/useCloudZeroSettings"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { Form, Input, message, Modal } from "antd"; +import { Form, Input, Modal } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { useEffect } from "react"; import { CloudZeroSettings } from "./types"; @@ -39,7 +40,7 @@ export default function CloudZeroUpdateModal({ open, onOk, onCancel, settings }: }, { onSuccess: () => { - message.success("CloudZero integration updated successfully"); + MessageManager.success("CloudZero integration updated successfully"); form.resetFields(); onOk(); }, @@ -47,7 +48,7 @@ export default function CloudZeroUpdateModal({ open, onOk, onCancel, settings }: if (error?.errorFields) { return; } - message.error(error?.message || "Failed to update CloudZero integration"); + MessageManager.error(error?.message || "Failed to update CloudZero integration"); }, }, ); @@ -55,7 +56,7 @@ export default function CloudZeroUpdateModal({ open, onOk, onCancel, settings }: if (error?.errorFields) { return; } - message.error(error?.message || "Failed to update CloudZero integration"); + MessageManager.error(error?.message || "Failed to update CloudZero integration"); } }; diff --git a/ui/litellm-dashboard/src/components/Projects/ProjectModals/CreateProjectModal.tsx b/ui/litellm-dashboard/src/components/Projects/ProjectModals/CreateProjectModal.tsx index e490f89303..bbf56e4930 100644 --- a/ui/litellm-dashboard/src/components/Projects/ProjectModals/CreateProjectModal.tsx +++ b/ui/litellm-dashboard/src/components/Projects/ProjectModals/CreateProjectModal.tsx @@ -1,5 +1,6 @@ -import { Modal, Form, Button, Typography, message } from "antd"; +import { Modal, Form, Button, Typography } from "antd"; import { FolderAddOutlined } from "@ant-design/icons"; +import MessageManager from "@/components/molecules/message_manager"; import { useCreateProject, ProjectCreateParams, @@ -32,12 +33,12 @@ export function CreateProjectModal({ createMutation.mutate(params, { onSuccess: () => { - message.success("Project created successfully"); + MessageManager.success("Project created successfully"); form.resetFields(); onClose(); }, onError: (error) => { - message.error(error.message || "Failed to create project"); + MessageManager.error(error.message || "Failed to create project"); }, }); } catch (error) { diff --git a/ui/litellm-dashboard/src/components/Projects/ProjectModals/EditProjectModal.tsx b/ui/litellm-dashboard/src/components/Projects/ProjectModals/EditProjectModal.tsx index 75f56b1373..dc3b43ef73 100644 --- a/ui/litellm-dashboard/src/components/Projects/ProjectModals/EditProjectModal.tsx +++ b/ui/litellm-dashboard/src/components/Projects/ProjectModals/EditProjectModal.tsx @@ -1,6 +1,7 @@ import { useEffect } from "react"; -import { Modal, Form, Button, Typography, message } from "antd"; +import { Modal, Form, Button, Typography } from "antd"; import { SaveOutlined } from "@ant-design/icons"; +import MessageManager from "@/components/molecules/message_manager"; import { ProjectResponse } from "@/app/(dashboard)/hooks/projects/useProjects"; import { useUpdateProject, @@ -80,12 +81,12 @@ export function EditProjectModal({ { projectId: project.project_id, params }, { onSuccess: () => { - message.success("Project updated successfully"); + MessageManager.success("Project updated successfully"); onSuccess?.(); onClose(); }, onError: (error) => { - message.error(error.message || "Failed to update project"); + MessageManager.error(error.message || "Failed to update project"); }, }, ); diff --git a/ui/litellm-dashboard/src/components/SearchTools/SearchToolTester.tsx b/ui/litellm-dashboard/src/components/SearchTools/SearchToolTester.tsx index 28b34b0608..d1cb5077b1 100644 --- a/ui/litellm-dashboard/src/components/SearchTools/SearchToolTester.tsx +++ b/ui/litellm-dashboard/src/components/SearchTools/SearchToolTester.tsx @@ -1,5 +1,6 @@ import React, { useState } from "react"; -import { Button, Input, Typography, Spin, message } from "antd"; +import { Button, Input, Typography, Spin } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { SearchOutlined, LoadingOutlined } from "@ant-design/icons"; import { searchToolQueryCall } from "../networking"; import NotificationsManager from "../molecules/notifications_manager"; @@ -39,7 +40,7 @@ export const SearchToolTester: React.FC = ({ searchToolNa const handleSearch = async () => { if (!query.trim()) { - message.warning("Please enter a search query"); + MessageManager.warning("Please enter a search query"); return; } diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx index 34f059516a..e55089d27d 100644 --- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx @@ -5,8 +5,9 @@ */ import { Button as TremorButton } from "@tremor/react"; -import { Button, message } from "antd"; +import { Button } from "antd"; import React, { useEffect, useState } from "react"; +import MessageManager from "@/components/molecules/message_manager"; import NotificationManager from "../../../molecules/notifications_manager"; import { fetchAvailableModels, ModelGroup } from "../../../playground/llm_calls/fetch_models"; import { AddFallbacksModal } from "./AddFallbacksModal"; @@ -90,7 +91,7 @@ export default function AddFallbacks({ (g) => !g.primaryModel || g.fallbackModels.length === 0, ); if (invalidGroups.length > 0) { - message.error( + MessageManager.error( `Please complete configuration for all groups. ${invalidGroups.length} group(s) incomplete.`, ); return; diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackSelectionForm.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackSelectionForm.tsx index 08b031c683..0482161bc6 100644 --- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackSelectionForm.tsx +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackSelectionForm.tsx @@ -5,9 +5,10 @@ */ import { Button } from "@tremor/react"; -import { message, Tabs } from "antd"; +import { Tabs } from "antd"; import { Plus } from "lucide-react"; import React, { useEffect, useState } from "react"; +import MessageManager from "@/components/molecules/message_manager"; import { FallbackGroup, FallbackGroupConfig } from "./FallbackGroupConfig"; interface FallbackSelectionFormProps { @@ -60,7 +61,7 @@ export function FallbackSelectionForm({ const handleRemoveGroup = (targetId: string) => { if (groups.length === 1) { - message.warning("At least one group is required"); + MessageManager.warning("At least one group is required"); return; } const newGroups = groups.filter((g) => g.id !== targetId); diff --git a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx index c5518596b8..b6d96f0445 100644 --- a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx +++ b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx @@ -1,5 +1,6 @@ import React, { useState, useEffect } from "react"; -import { Modal, Form, message, Select, Input, Steps, Radio, Tag, Divider, Switch, InputNumber, Collapse } from "antd"; +import { Modal, Form, Select, Input, Steps, Radio, Tag, Divider, Switch, InputNumber, Collapse } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { Button } from "@tremor/react"; import { CheckCircleFilled, KeyOutlined, RobotOutlined, AppstoreOutlined, InfoCircleOutlined } from "@ant-design/icons"; import CreatedKeyDisplay from "../shared/CreatedKeyDisplay"; @@ -216,7 +217,7 @@ const AddAgentForm: React.FC = ({ const handleCreateAgent = async () => { if (!accessToken) { - message.error("No access token available"); + MessageManager.error("No access token available"); return; } @@ -226,7 +227,7 @@ const AddAgentForm: React.FC = ({ const values = { ...form.getFieldsValue(true) }; const agentData = buildAgentData(values); if (!agentData) { - message.error("Failed to build agent data"); + MessageManager.error("Failed to build agent data"); setIsSubmitting(false); return; } @@ -301,7 +302,7 @@ const AddAgentForm: React.FC = ({ setCreatedKeyValue(keyResponse.key || null); } else if (keyAssignOption === "existing_key") { if (!selectedExistingKey) { - message.error("Please select an existing key to assign"); + MessageManager.error("Please select an existing key to assign"); setIsSubmitting(false); return; } @@ -318,7 +319,7 @@ const AddAgentForm: React.FC = ({ } catch (error) { console.error("Error creating agent:", error); const errorMessage = error instanceof Error ? error.message : String(error); - message.error(errorMessage ? `Failed to create agent: ${errorMessage}` : "Failed to create agent"); + MessageManager.error(errorMessage ? `Failed to create agent: ${errorMessage}` : "Failed to create agent"); } finally { setIsSubmitting(false); } diff --git a/ui/litellm-dashboard/src/components/agents/agent_info.tsx b/ui/litellm-dashboard/src/components/agents/agent_info.tsx index b41e318a76..d543be8356 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_info.tsx +++ b/ui/litellm-dashboard/src/components/agents/agent_info.tsx @@ -1,6 +1,7 @@ import React, { useState, useEffect } from "react"; import { Card, Title, Text, Button as TremorButton, Tab, TabGroup, TabList, TabPanel, TabPanels} from "@tremor/react"; -import { Form, Input, InputNumber, Button as AntButton, message, Spin, Descriptions, Divider } from "antd"; +import { Form, Input, InputNumber, Button as AntButton, Spin, Descriptions, Divider } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { ArrowLeftIcon } from "@heroicons/react/outline"; import { getAgentInfo, patchAgentCall, getAgentCreateMetadata, AgentCreateInfo } from "../networking"; import { Agent } from "./types"; @@ -72,7 +73,7 @@ const AgentInfoView: React.FC = ({ } } catch (error) { console.error("Error fetching agent info:", error); - message.error("Failed to load agent information"); + MessageManager.error("Failed to load agent information"); } finally { setIsLoading(false); } @@ -111,12 +112,12 @@ const AgentInfoView: React.FC = ({ } await patchAgentCall(accessToken, agentId, updateData); - message.success("Agent updated successfully"); + MessageManager.success("Agent updated successfully"); setIsEditing(false); fetchAgentInfo(); } catch (error) { console.error("Error updating agent:", error); - message.error("Failed to update agent"); + MessageManager.error("Failed to update agent"); } finally { setIsSaving(false); } diff --git a/ui/litellm-dashboard/src/components/chat/ChatPage.tsx b/ui/litellm-dashboard/src/components/chat/ChatPage.tsx index ccf39d2147..b547f74917 100644 --- a/ui/litellm-dashboard/src/components/chat/ChatPage.tsx +++ b/ui/litellm-dashboard/src/components/chat/ChatPage.tsx @@ -1,7 +1,8 @@ "use client"; import React, { useCallback, useEffect, useRef, useState, useLayoutEffect } from "react"; -import { Tooltip, Skeleton, Popover, message } from "antd"; +import { Tooltip, Skeleton, Popover } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { SettingOutlined, PlusOutlined, @@ -212,7 +213,7 @@ const ChatPage: React.FC = ({ accessToken, userRole, userId, user localStorage.setItem(LOCALSTORAGE_MODEL_KEY, JSON.stringify([names[0]])); } }) - .catch(() => message.error("Could not load models")) + .catch(() => MessageManager.error("Could not load models")) .finally(() => setIsLoadingModels(false)); }, [accessToken]); diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx index b805cab71e..d25db73ae4 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx @@ -5,7 +5,7 @@ import { Spin, Input, Button, Skeleton } from "antd"; import { SearchOutlined, ArrowLeftOutlined, RightOutlined, ToolOutlined, CheckCircleOutlined } from "@ant-design/icons"; import { deleteMCPOAuthUserCredential, fetchMCPServers, getMCPOAuthUserCredentialStatus, listMCPTools } from "../networking"; import { AUTH_TYPE, MCPServer, MCPTool, handleTransport } from "../mcp_tools/types"; -import { message } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { useUserMcpOAuthFlow } from "@/hooks/useUserMcpOAuthFlow"; // ── OAuth2 connect button ───────────────────────────────────────────────────── @@ -198,7 +198,7 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange const idToFetch = serverId ?? serverName; const result = await listMCPTools(accessToken, idToFetch); if (result?.error) { - message.warning(`Could not load tools for ${serverName}`); + MessageManager.warning(`Could not load tools for ${serverName}`); return; } // Use the ref so we read the most up-to-date list; guard against duplicates @@ -207,7 +207,7 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange onChange([...selectedServersRef.current, serverName]); } } catch { - message.warning(`Could not load tools for ${serverName}`); + MessageManager.warning(`Could not load tools for ${serverName}`); } finally { setTogglingOn((prev) => { const next = new Set(prev); diff --git a/ui/litellm-dashboard/src/components/chat/MCPConnectPicker.tsx b/ui/litellm-dashboard/src/components/chat/MCPConnectPicker.tsx index a52ce18156..6ef3aecc46 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPConnectPicker.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPConnectPicker.tsx @@ -1,5 +1,6 @@ import React, { useEffect, useState } from "react"; -import { Switch, Spin, message } from "antd"; +import { Switch, Spin } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { fetchMCPServers, listMCPTools } from "../networking"; import { MCPServer } from "../mcp_tools/types"; @@ -57,7 +58,7 @@ const MCPConnectPicker: React.FC = ({ accessToken, selectedServers, onCha const result = await listMCPTools(accessToken, serverName); // listMCPTools never throws; it returns { tools, error, message } on failure if (result?.error) { - message.warning( + MessageManager.warning( `Could not load tools for ${serverName} — it will be excluded from this message.` ); // Do not add to selectedServers @@ -65,7 +66,7 @@ const MCPConnectPicker: React.FC = ({ accessToken, selectedServers, onCha } onChange([...selectedServers, serverName]); } catch { - message.warning( + MessageManager.warning( `Could not load tools for ${serverName} — it will be excluded from this message.` ); // Do not add to selectedServers diff --git a/ui/litellm-dashboard/src/components/chat/MCPCredentialsTab.tsx b/ui/litellm-dashboard/src/components/chat/MCPCredentialsTab.tsx index 363ec8c0e4..e2d879b22a 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPCredentialsTab.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPCredentialsTab.tsx @@ -8,7 +8,8 @@ */ import React, { useCallback, useEffect, useState } from "react"; -import { Spin, message } from "antd"; +import { Spin } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { DeleteOutlined, LinkOutlined } from "@ant-design/icons"; import { Badge, Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow } from "@tremor/react"; import { @@ -77,7 +78,7 @@ const MCPCredentialsTab: React.FC = ({ accessToken }) => { await deleteMCPOAuthUserCredential(accessToken, serverId); setCredentials((prev) => prev.filter((c) => c.server_id !== serverId)); } catch { - message.error("Failed to revoke connection. Please try again."); + MessageManager.error("Failed to revoke connection. Please try again."); } finally { setRevoking((prev) => { const n = new Set(prev); n.delete(serverId); return n; }); } diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx index 217851f128..d5e417a9a8 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx @@ -1,5 +1,6 @@ import React, { useState } from "react"; -import { Modal, Form, Input, Select, message } from "antd"; +import { Modal, Form, Input, Select } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { Button } from "@tremor/react"; import { registerClaudeCodePlugin } from "../networking"; import { @@ -43,13 +44,13 @@ const AddPluginForm: React.FC = ({ const handleSubmit = async (values: any) => { if (!accessToken) { - message.error("No access token available"); + MessageManager.error("No access token available"); return; } // Validate plugin name if (!validatePluginName(values.name)) { - message.error( + MessageManager.error( "Plugin name must be kebab-case (lowercase letters, numbers, and hyphens only)" ); return; @@ -57,7 +58,7 @@ const AddPluginForm: React.FC = ({ // Validate semantic version if provided if (values.version && !isValidSemanticVersion(values.version)) { - message.error( + MessageManager.error( "Version must be in semantic versioning format (e.g., 1.0.0)" ); return; @@ -65,13 +66,13 @@ const AddPluginForm: React.FC = ({ // Validate email if provided if (values.authorEmail && !isValidEmail(values.authorEmail)) { - message.error("Invalid email format"); + MessageManager.error("Invalid email format"); return; } // Validate homepage URL if provided if (values.homepage && !isValidUrl(values.homepage)) { - message.error("Invalid homepage URL format"); + MessageManager.error("Invalid homepage URL format"); return; } @@ -119,14 +120,14 @@ const AddPluginForm: React.FC = ({ } await registerClaudeCodePlugin(accessToken, pluginData); - message.success("Plugin registered successfully"); + MessageManager.success("Plugin registered successfully"); form.resetFields(); setSourceType("github"); onSuccess(); onClose(); } catch (error) { console.error("Error registering plugin:", error); - message.error("Failed to register plugin"); + MessageManager.error("Failed to register plugin"); } finally { setIsSubmitting(false); } diff --git a/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx index 473918c126..58c2a965e0 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx @@ -1,7 +1,8 @@ "use client"; import React, { useState } from "react"; -import { Modal, Input, Switch, message } from "antd"; +import { Modal, Input, Switch } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { KeyOutlined, LockOutlined, @@ -46,7 +47,7 @@ export const ByokCredentialModal: React.FC = ({ const handleAuthorize = async () => { if (!apiKey.trim()) { - message.error("Please enter your API key"); + MessageManager.error("Please enter your API key"); return; } setLoading(true); @@ -63,11 +64,11 @@ export const ByokCredentialModal: React.FC = ({ const err = await response.json(); throw new Error(err?.detail?.error || "Failed to save credential"); } - message.success(`Connected to ${serverDisplayName}`); + MessageManager.success(`Connected to ${serverDisplayName}`); onSuccess(server.server_id); handleClose(); } catch (e: any) { - message.error(e.message || "Failed to connect"); + MessageManager.error(e.message || "Failed to connect"); } finally { setLoading(false); } diff --git a/ui/litellm-dashboard/src/components/molecules/message_manager.tsx b/ui/litellm-dashboard/src/components/molecules/message_manager.tsx new file mode 100644 index 0000000000..e4c1552d5e --- /dev/null +++ b/ui/litellm-dashboard/src/components/molecules/message_manager.tsx @@ -0,0 +1,38 @@ +import { message as staticMessage } from "antd"; +import type { MessageInstance } from "antd/es/message/interface"; + +let messageInstance: MessageInstance | null = null; + +export const setMessageInstance = (instance: MessageInstance) => { + messageInstance = instance; +}; + +const getMessageApi = () => messageInstance || staticMessage; + +const MessageManager = { + success(content: string, duration?: number) { + getMessageApi().success(content, duration); + }, + + error(content: string, duration?: number) { + getMessageApi().error(content, duration); + }, + + warning(content: string, duration?: number) { + getMessageApi().warning(content, duration); + }, + + info(content: string, duration?: number) { + getMessageApi().info(content, duration); + }, + + loading(content: string, duration?: number) { + return getMessageApi().loading(content, duration); + }, + + destroy() { + getMessageApi().destroy(); + }, +}; + +export default MessageManager; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index bb9a0c1e01..956968dfee 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -68,7 +68,7 @@ export const getInProductNudgesCall = async (accessToken: string) => { /** * Helper file for calls being made to proxy */ -import { message } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { clearTokenCookies } from "@/utils/cookieUtils"; import { TagNewRequest, TagUpdateRequest, TagListResponse, TagInfoResponse } from "./tag_management/types"; import { Team } from "./key_team_helpers/key_list"; @@ -555,7 +555,7 @@ export const modelCreateCall = async (accessToken: string, formValues: Model) => console.log("API Response:", data); // Close any existing messages before showing new ones - message.destroy(); + MessageManager.destroy(); // Sequential success messages NotificationsManager.success(`Model ${formValues.model_name} created successfully`); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index d0a7a909d6..df5b3a4b32 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -8,7 +8,7 @@ import { formatNumberWithCommas } from "@/utils/dataUtils"; import { InfoCircleOutlined } from "@ant-design/icons"; import { useQueryClient } from "@tanstack/react-query"; import { Accordion, AccordionBody, AccordionHeader, Button, Col, Grid, Text, TextInput, Title } from "@tremor/react"; -import { Button as Button2, Form, Input, message, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd"; +import { Button as Button2, Form, Input, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd"; import debounce from "lodash/debounce"; import React, { useCallback, useEffect, useState } from "react"; import { rolesWithWriteAccess } from "../../utils/roles"; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterTool.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterTool.tsx index 04711ede9f..dd37769556 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterTool.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterTool.tsx @@ -1,5 +1,6 @@ import React from "react"; -import { Switch, Tooltip, message } from "antd"; +import { Switch, Tooltip } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { CodeOutlined, InfoCircleOutlined, ExclamationCircleOutlined } from "@ant-design/icons"; import { Text } from "@tremor/react"; @@ -38,7 +39,7 @@ const CodeInterpreterTool: React.FC = ({ const handleToggle = (checked: boolean) => { if (checked && !isOpenAI) { - message.warning("Code Interpreter is only available for OpenAI models"); + MessageManager.warning("Code Interpreter is only available for OpenAI models"); return; } onEnabledChange(checked); diff --git a/ui/litellm-dashboard/src/components/policies/index.tsx b/ui/litellm-dashboard/src/components/policies/index.tsx index bd77f3ab9b..f47b8d78b9 100644 --- a/ui/litellm-dashboard/src/components/policies/index.tsx +++ b/ui/litellm-dashboard/src/components/policies/index.tsx @@ -1,6 +1,7 @@ import React, { useState, useEffect, useCallback } from "react"; import { Button, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react"; -import { Modal, message, Alert } from "antd"; +import { Modal, Alert } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { ExclamationCircleOutlined, InfoCircleOutlined } from "@ant-design/icons"; import { isAdminRole } from "@/utils/roles"; import PolicyTable from "./policy_table"; @@ -80,7 +81,7 @@ const PoliciesPanel: React.FC = ({ setPoliciesList(response.policies || []); } catch (error) { console.error("Error fetching policies:", error); - message.error("Failed to fetch policies"); + MessageManager.error("Failed to fetch policies"); } finally { setIsLoading(false); } @@ -95,7 +96,7 @@ const PoliciesPanel: React.FC = ({ setAttachmentsList(response.attachments || []); } catch (error) { console.error("Error fetching attachments:", error); - message.error("Failed to fetch attachments"); + MessageManager.error("Failed to fetch attachments"); } finally { setIsAttachmentsLoading(false); } @@ -148,11 +149,11 @@ const PoliciesPanel: React.FC = ({ setIsDeleting(true); try { await deletePolicyCall(accessToken, policyToDelete.policy_id); - message.success(`Policy "${policyToDelete.policy_name}" deleted successfully`); + MessageManager.success(`Policy "${policyToDelete.policy_name}" deleted successfully`); await fetchPolicies(); } catch (error) { console.error("Error deleting policy:", error); - message.error("Failed to delete policy"); + MessageManager.error("Failed to delete policy"); } finally { setIsDeleting(false); setIsDeleteModalOpen(false); @@ -177,11 +178,11 @@ const PoliciesPanel: React.FC = ({ if (!accessToken) return; try { await deletePolicyAttachmentCall(accessToken, attachmentId); - message.success("Attachment deleted successfully"); + MessageManager.success("Attachment deleted successfully"); fetchAttachments(); } catch (error) { console.error("Error deleting attachment:", error); - message.error("Failed to delete attachment"); + MessageManager.error("Failed to delete attachment"); } }, }); @@ -193,7 +194,7 @@ const PoliciesPanel: React.FC = ({ const handleUseTemplate = async (template: any) => { if (!accessToken) { - message.error("Authentication required"); + MessageManager.error("Authentication required"); return; } @@ -221,7 +222,7 @@ const PoliciesPanel: React.FC = ({ setIsGuardrailSelectionModalOpen(true); } catch (error) { console.error("Error fetching guardrails:", error); - message.error("Failed to load guardrails. Please try again."); + MessageManager.error("Failed to load guardrails. Please try again."); } }; @@ -271,7 +272,7 @@ const PoliciesPanel: React.FC = ({ await proceedWithTemplate(enrichedTemplate); } catch (error) { console.error("Error enriching template:", error); - message.error("Failed to configure template. Please try again."); + MessageManager.error("Failed to configure template. Please try again."); setIsEnrichingTemplate(false); } }; @@ -318,15 +319,15 @@ const PoliciesPanel: React.FC = ({ // Show success message if (createdGuardrails.length > 0) { - message.success( + MessageManager.success( `Created ${createdGuardrails.length} guardrail${createdGuardrails.length > 1 ? "s" : ""}! Complete the policy form to save.` ); } else { - message.success("Template ready! Complete the policy form to save."); + MessageManager.success("Template ready! Complete the policy form to save."); } if (failedGuardrails.length > 0) { - message.warning( + MessageManager.warning( `Failed to create ${failedGuardrails.length} guardrail(s): ${failedGuardrails.join(", ")}. You may need to create them manually.` ); } @@ -348,7 +349,7 @@ const PoliciesPanel: React.FC = ({ setTemplateQueue([]); setTemplateQueueProgress(null); console.error("Error creating guardrails:", error); - message.error("Failed to create guardrails. Please try again."); + MessageManager.error("Failed to create guardrails. Please try again."); } }; diff --git a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx index 89665774ee..b1768d5b81 100644 --- a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx +++ b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx @@ -1,5 +1,6 @@ import React, { useState } from "react"; -import { Select, Typography, message, Spin } from "antd"; +import { Select, Typography, Spin } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { Button, TextInput } from "@tremor/react"; import { ArrowLeftIcon, PlusIcon } from "@heroicons/react/outline"; import { DotsVerticalIcon } from "@heroicons/react/solid"; @@ -1385,17 +1386,17 @@ export const FlowBuilderPage: React.FC = ({ const handleSave = async () => { if (!policyName.trim()) { - message.error("Please enter a policy name"); + MessageManager.error("Please enter a policy name"); return; } if (!accessToken) { - message.error("No access token available"); + MessageManager.error("No access token available"); return; } const emptySteps = pipeline.steps.filter((s) => !s.guardrail); if (emptySteps.length > 0) { - message.error("Please select a guardrail for all steps"); + MessageManager.error("Please select a guardrail for all steps"); return; } diff --git a/ui/litellm-dashboard/src/components/policies/policy_templates.tsx b/ui/litellm-dashboard/src/components/policies/policy_templates.tsx index 98ba0acb6a..c4ea22651c 100644 --- a/ui/litellm-dashboard/src/components/policies/policy_templates.tsx +++ b/ui/litellm-dashboard/src/components/policies/policy_templates.tsx @@ -1,5 +1,6 @@ import React, { useState, useEffect, useMemo } from "react"; -import { Card, Button, Spin, message, Checkbox } from "antd"; +import { Card, Button, Spin, Checkbox } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { ShieldCheckIcon, ShieldExclamationIcon, @@ -184,7 +185,7 @@ const PolicyTemplates: React.FC = ({ onUseTemplate, onOpen onTemplatesLoaded?.(data); } catch (error) { console.error("Error fetching policy templates:", error); - message.error("Failed to fetch policy templates"); + MessageManager.error("Failed to fetch policy templates"); } finally { setIsLoading(false); } diff --git a/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx index 013e222c6b..eb6fe5f50c 100644 --- a/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx @@ -3,15 +3,11 @@ import { render, screen, act } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import CreatedKeyDisplay from "./CreatedKeyDisplay"; -vi.mock("antd", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - message: { success: vi.fn() }, - }; -}); +vi.mock("@/components/molecules/message_manager", () => ({ + default: { success: vi.fn(), error: vi.fn(), warning: vi.fn(), info: vi.fn(), loading: vi.fn(), destroy: vi.fn() }, +})); -import { message } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; describe("CreatedKeyDisplay", () => { beforeEach(() => { @@ -52,7 +48,7 @@ describe("CreatedKeyDisplay", () => { await user.click(screen.getByRole("button", { name: /copy virtual key/i })); - expect(message.success).toHaveBeenCalledWith("Key copied to clipboard"); + expect(MessageManager.success).toHaveBeenCalledWith("Key copied to clipboard"); }); it("should revert button text back after 2 seconds", async () => { diff --git a/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx index d75c63c646..00ffb8c6e2 100644 --- a/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx +++ b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx @@ -1,6 +1,7 @@ import React, { useState } from "react"; import { CopyToClipboard } from "react-copy-to-clipboard"; -import { Button, message } from "antd"; +import { Button } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; interface CreatedKeyDisplayProps { apiKey: string; @@ -15,7 +16,7 @@ const CreatedKeyDisplay: React.FC = ({ apiKey }) => { const handleCopy = () => { setCopied(true); - message.success("Key copied to clipboard"); + MessageManager.success("Key copied to clipboard"); setTimeout(() => setCopied(false), 2000); }; diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index d2ce79580d..a4c7ae2bbb 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -20,7 +20,8 @@ import { isProxyAdminRole } from "@/utils/roles"; import { EditOutlined, InfoCircleOutlined, SaveOutlined } from "@ant-design/icons"; import { ArrowLeftIcon } from "@heroicons/react/outline"; import { Badge, Card, Grid, Text, TextInput, Title } from "@tremor/react"; -import { Button, Form, Input, message, Select, Switch, Tabs, Tooltip } from "antd"; +import { Button, Form, Input, Select, Switch, Tabs, Tooltip } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { CheckIcon, CopyIcon } from "lucide-react"; import React, { useEffect, useMemo, useState } from "react"; import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils"; @@ -366,7 +367,7 @@ const TeamInfoView: React.FC = ({ tpm_limit: values.tpm_limit, rpm_limit: values.rpm_limit, }; - message.destroy(); // Remove all existing toasts + MessageManager.destroy(); // Remove all existing toasts await teamMemberUpdateCall(accessToken, teamId, member); @@ -388,7 +389,7 @@ const TeamInfoView: React.FC = ({ } setIsEditMemberModalVisible(false); - message.destroy(); // Remove all existing toasts + MessageManager.destroy(); // Remove all existing toasts NotificationsManager.fromBackend(errMsg); console.error("Error updating team member:", error); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx b/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx index 97162aa132..b189168e1b 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx @@ -1,6 +1,7 @@ import React, { useState } from "react"; import { Card, Title, Text } from "@tremor/react"; -import { Upload, Button, Select, Form, message, Alert, Tooltip, Input } from "antd"; +import { Upload, Button, Select, Form, Alert, Tooltip, Input } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { InboxOutlined, InfoCircleOutlined } from "@ant-design/icons"; import type { UploadProps } from "antd"; import { ragIngestCall } from "../networking"; @@ -47,13 +48,13 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu ].includes(file.type); if (!isValidType) { - message.error(`${file.name} is not a supported file type. Please upload PDF, TXT, DOCX, or MD files.`); + MessageManager.error(`${file.name} is not a supported file type. Please upload PDF, TXT, DOCX, or MD files.`); return Upload.LIST_IGNORE; } const isLt50M = file.size / 1024 / 1024 < 50; if (!isLt50M) { - message.error(`${file.name} must be smaller than 50MB!`); + MessageManager.error(`${file.name} must be smaller than 50MB!`); return Upload.LIST_IGNORE; } @@ -87,12 +88,12 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu const handleCreateVectorStore = async () => { if (documents.length === 0) { - message.warning("Please upload at least one document"); + MessageManager.warning("Please upload at least one document"); return; } if (!selectedProvider) { - message.warning("Please select a provider"); + MessageManager.warning("Please select a provider"); return; } @@ -100,7 +101,7 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu const requiredFields = getProviderSpecificFields(selectedProvider).filter((field) => field.required); for (const field of requiredFields) { if (!providerParams[field.name]) { - message.warning(`Please provide ${field.label}`); + MessageManager.warning(`Please provide ${field.label}`); return; } } @@ -108,17 +109,17 @@ const CreateVectorStore: React.FC = ({ accessToken, onSu // S3 Vectors specific validation if (selectedProvider === "s3_vectors") { if (providerParams.vector_bucket_name && providerParams.vector_bucket_name.length < 3) { - message.warning("Vector bucket name must be at least 3 characters"); + MessageManager.warning("Vector bucket name must be at least 3 characters"); return; } if (providerParams.index_name && providerParams.index_name.length > 0 && providerParams.index_name.length < 3) { - message.warning("Index name must be at least 3 characters if provided"); + MessageManager.warning("Index name must be at least 3 characters if provided"); return; } } if (!accessToken) { - message.error("No access token available"); + MessageManager.error("No access token available"); return; } diff --git a/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.tsx b/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.tsx index aeb4240d36..c1ce57ad33 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.tsx @@ -1,5 +1,6 @@ import React from "react"; -import { Table, Badge, Tooltip, message } from "antd"; +import { Table, Badge, Tooltip } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { EyeOutlined, CopyOutlined, DeleteOutlined } from "@ant-design/icons"; import { DocumentUpload } from "./types"; @@ -11,7 +12,7 @@ interface DocumentsTableProps { const DocumentsTable: React.FC = ({ documents, onRemove }) => { const handleCopyId = (uid: string) => { navigator.clipboard.writeText(uid); - message.success("Document ID copied to clipboard"); + MessageManager.success("Document ID copied to clipboard"); }; const getStatusBadge = (status: DocumentUpload["status"]) => { diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTester.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTester.tsx index f8b5ada0c3..5c59334fa4 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTester.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTester.tsx @@ -1,5 +1,6 @@ import React, { useState } from "react"; -import { Button, Input, Card, Typography, Spin, message, Divider } from "antd"; +import { Button, Input, Card, Typography, Spin, Divider } from "antd"; +import MessageManager from "@/components/molecules/message_manager"; import { SendOutlined, DatabaseOutlined, LoadingOutlined, DownOutlined, RightOutlined } from "@ant-design/icons"; import { vectorStoreSearchCall } from "../networking"; import NotificationsManager from "../molecules/notifications_manager"; @@ -46,7 +47,7 @@ export const VectorStoreTester: React.FC = ({ vectorStor const handleSearch = async () => { if (!query.trim()) { - message.warning("Please enter a search query"); + MessageManager.warning("Please enter a search query"); return; } diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/InputCard.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/InputCard.tsx index 1299f9c168..8fb47274a8 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/InputCard.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/InputCard.tsx @@ -4,7 +4,7 @@ */ import { useState } from 'react'; -import { message } from 'antd'; +import MessageManager from "@/components/molecules/message_manager"; import { ParsedMessage } from './prettyMessagesTypes'; import { SectionHeader } from './SectionHeader'; import { CollapsibleMessage } from './CollapsibleMessage'; @@ -33,7 +33,7 @@ export function InputCard({ messages, promptTokens, inputCost }: InputCardProps) const handleCopy = () => { const content = lastMessage?.content || ''; navigator.clipboard.writeText(content); - message.success('Input copied'); + MessageManager.success('Input copied'); }; return ( diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/OutputCard.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/OutputCard.tsx index eff8d83cbd..22f1708ee4 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/OutputCard.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/OutputCard.tsx @@ -4,7 +4,8 @@ */ import { useState } from 'react'; -import { Typography, message as antdMessage } from 'antd'; +import { Typography } from 'antd'; +import MessageManager from "@/components/molecules/message_manager"; import { ParsedMessage } from './prettyMessagesTypes'; import { SectionHeader } from './SectionHeader'; import { SimpleMessageBlock } from './SimpleMessageBlock'; @@ -25,7 +26,7 @@ export function OutputCard({ message, completionTokens, outputCost }: OutputCard const content = message.content || ''; navigator.clipboard.writeText(content); - antdMessage.success('Output copied'); + MessageManager.success('Output copied'); }; if (!message) { diff --git a/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx b/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx index 5b5c1036fa..cc182ee36d 100644 --- a/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx +++ b/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx @@ -1,23 +1,27 @@ "use client"; import React, { useEffect, useRef } from "react"; -import { notification } from "antd"; +import { notification, message } from "antd"; import { setNotificationInstance } from "@/components/molecules/notifications_manager"; +import { setMessageInstance } from "@/components/molecules/message_manager"; export default function AntdGlobalProvider({ children }: { children: React.ReactNode }) { - const [api, contextHolder] = notification.useNotification(); + const [notificationApi, notificationContextHolder] = notification.useNotification(); + const [messageApi, messageContextHolder] = message.useMessage(); const initialized = useRef(false); useEffect(() => { if (!initialized.current) { - setNotificationInstance(api); + setNotificationInstance(notificationApi); + setMessageInstance(messageApi); initialized.current = true; } - }, [api]); + }, [notificationApi, messageApi]); return ( <> - {contextHolder} + {notificationContextHolder} + {messageContextHolder} {children} ); From b8c9bf7d256ed90b1f8cfdbe8b57d79f118009f4 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 20 Mar 2026 09:05:18 -0700 Subject: [PATCH 20/26] refactor: extract _auto_assign and _remove helpers, use team_endpoints helper - Replace raw prisma_client.db.litellm_teamtable.update with handle_update_object_permission from team_endpoints (follows established helper-function pattern) - Extract _auto_assign_mcp_server_to_team and _remove_mcp_server_from_team helpers for reuse and testability - Update tests to mock at the correct boundaries Co-Authored-By: Claude Opus 4.6 --- .../mcp_management_endpoints.py | 119 ++++++++++-------- .../test_mcp_manager_role.py | 89 ++++--------- 2 files changed, 96 insertions(+), 112 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index f137902f1b..e137ad15f1 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -203,6 +203,69 @@ if MCP_AVAILABLE: return resolved_team_id, team_obj + async def _auto_assign_mcp_server_to_team( + server_id: str, + team_id: str, + team_obj: "LiteLLM_TeamTableCachedObj", + prisma_client: Any, + ) -> None: + """ + Add an MCP server to a team's ObjectPermissionTable and link the + permission back to the team if it didn't have one yet. + + Uses handle_update_object_permission (the team-endpoint helper) so + the object_permission_id linkage follows the same pattern as + team_endpoints.update_team. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + handle_update_object_permission, + ) + + existing_mcp_servers: list = [] + if team_obj.object_permission is not None: + existing_mcp_servers = team_obj.object_permission.mcp_servers or [] + updated_mcp_servers = list(set(existing_mcp_servers + [server_id])) + + # Build the data dict in the same shape team_endpoints uses + data_json: Dict[str, Any] = { + "object_permission": {"mcp_servers": updated_mcp_servers}, + } + data_json = await handle_update_object_permission( + data_json=data_json, + existing_team_row=team_obj, + ) + + # If handle_update_object_permission produced an object_permission_id, + # persist it on the team row (it sets data_json["object_permission_id"]). + if "object_permission_id" in data_json: + await prisma_client.db.litellm_teamtable.update( + where={"team_id": team_id}, + data={"object_permission_id": data_json["object_permission_id"]}, + ) + + async def _remove_mcp_server_from_team( + server_id: str, + team_obj: "LiteLLM_TeamTableCachedObj", + ) -> None: + """Remove a server ID from a team's ObjectPermissionTable.mcp_servers list.""" + from litellm.proxy.proxy_server import prisma_client + + existing_permission_id = getattr(team_obj, "object_permission_id", None) + if ( + existing_permission_id is None + or team_obj.object_permission is None + ): + return + + existing_mcp_servers = team_obj.object_permission.mcp_servers or [] + updated_mcp_servers = [s for s in existing_mcp_servers if s != server_id] + + await handle_update_object_permission_common( + data_json={"object_permission": {"mcp_servers": updated_mcp_servers}}, + existing_object_permission_id=existing_permission_id, + prisma_client=prisma_client, + ) + @dataclass class _TemporaryMCPServerEntry: server: MCPServer @@ -1347,36 +1410,12 @@ if MCP_AVAILABLE: # so a failure here doesn't mask the successfully created server). if manager_team_id is not None and manager_team_obj is not None: try: - existing_permission_id = getattr( - manager_team_obj, "object_permission_id", None - ) - - # Read existing mcp_servers list and append the new server - existing_mcp_servers: list = [] - if manager_team_obj.object_permission is not None: - existing_mcp_servers = ( - manager_team_obj.object_permission.mcp_servers or [] - ) - updated_mcp_servers = list( - set(existing_mcp_servers + [new_mcp_server.server_id]) - ) - - new_permission_id = await handle_update_object_permission_common( - data_json={ - "object_permission": { - "mcp_servers": updated_mcp_servers, - } - }, - existing_object_permission_id=existing_permission_id, + await _auto_assign_mcp_server_to_team( + server_id=new_mcp_server.server_id, + team_id=manager_team_id, + team_obj=manager_team_obj, prisma_client=prisma_client, ) - - # If the team had no object_permission_id, link the new one - if existing_permission_id is None and new_permission_id is not None: - await prisma_client.db.litellm_teamtable.update( - where={"team_id": manager_team_id}, - data={"object_permission_id": new_permission_id}, - ) except Exception as e: verbose_proxy_logger.exception( f"MCP server created but failed to auto-assign to team {manager_team_id}: {str(e)}" @@ -1617,28 +1656,10 @@ if MCP_AVAILABLE: # Remove server from the manager's team permission list if manager_team_obj is not None: try: - existing_permission_id = getattr( - manager_team_obj, "object_permission_id", None + await _remove_mcp_server_from_team( + server_id=server_id, + team_obj=manager_team_obj, ) - if ( - existing_permission_id is not None - and manager_team_obj.object_permission is not None - ): - existing_mcp_servers = ( - manager_team_obj.object_permission.mcp_servers or [] - ) - updated_mcp_servers = [ - s for s in existing_mcp_servers if s != server_id - ] - await handle_update_object_permission_common( - data_json={ - "object_permission": { - "mcp_servers": updated_mcp_servers, - } - }, - existing_object_permission_id=existing_permission_id, - prisma_client=prisma_client, - ) except Exception as e: verbose_proxy_logger.exception( f"MCP server deleted but failed to remove from team permissions: {str(e)}" diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py index 67e9dda4ce..ed601dfffa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py @@ -233,7 +233,7 @@ class TestCreateMcpServerAsManager: created_server.server_id = "new_server_id" created_server.credentials = None - mock_handle_update = AsyncMock() + mock_auto_assign = AsyncMock() with ( patch( @@ -263,38 +263,23 @@ class TestCreateMcpServerAsManager: ), ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common", - mock_handle_update, + "litellm.proxy.management_endpoints.mcp_management_endpoints._auto_assign_mcp_server_to_team", + mock_auto_assign, ), ): await add_mcp_server(payload=payload, user_api_key_dict=user_auth) - # Verify handle_update_object_permission_common was called with merged server list - mock_handle_update.assert_called_once() - call_kwargs = mock_handle_update.call_args.kwargs - mcp_servers = call_kwargs["data_json"]["object_permission"]["mcp_servers"] - assert "existing_server" in mcp_servers - assert "new_server_id" in mcp_servers - assert call_kwargs["existing_object_permission_id"] == "perm1" + # Verify _auto_assign_mcp_server_to_team was called with the right args + mock_auto_assign.assert_called_once() + call_kwargs = mock_auto_assign.call_args.kwargs + assert call_kwargs["server_id"] == "new_server_id" + assert call_kwargs["team_id"] == "team1" + assert call_kwargs["team_obj"] == mock_team - async def test_create_links_new_permission_to_team_when_none_exists(self): - """When team has no object_permission_id, create should link the new one.""" - from litellm.proxy._types import NewMCPServerRequest + async def test_auto_assign_links_new_permission_to_team(self): + """_auto_assign_mcp_server_to_team should create permission and link to team.""" from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - add_mcp_server, - ) - - user_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, - user_id="user1", - api_key="sk-test", - team_id="team1", - ) - - payload = NewMCPServerRequest( - server_name="test_server", - url="https://example.com/mcp", - team_id="team1", + _auto_assign_mcp_server_to_team, ) # Team with NO object_permission_id @@ -307,46 +292,24 @@ class TestCreateMcpServerAsManager: ) mock_team.object_permission = None - created_server = MagicMock() - created_server.server_id = "new_server_id" - created_server.credentials = None - mock_prisma = MagicMock() mock_prisma.db.litellm_teamtable.update = AsyncMock() - with ( - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma, - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", - AsyncMock(return_value=("team1", mock_team)), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(return_value=None), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server", - AsyncMock(return_value=created_server), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - MagicMock( - add_server=AsyncMock(), - reload_servers_from_database=AsyncMock(), - ), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common", - AsyncMock(return_value="new_perm_id"), - ), + # handle_update_object_permission sets object_permission_id in data_json + async def fake_handle(data_json, existing_team_row): + data_json["object_permission_id"] = "new_perm_id" + return data_json + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.handle_update_object_permission", + side_effect=fake_handle, ): - await add_mcp_server(payload=payload, user_api_key_dict=user_auth) + await _auto_assign_mcp_server_to_team( + server_id="new_server_id", + team_id="team1", + team_obj=mock_team, + prisma_client=mock_prisma, + ) # Verify the team was updated with the new object_permission_id mock_prisma.db.litellm_teamtable.update.assert_called_once_with( From 1cd7a48c33cd72239913ccf73ecd87c3be4d5a30 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 20 Mar 2026 09:19:53 -0700 Subject: [PATCH 21/26] Add tests for edit and delete MCP server manager paths Addresses Greptile feedback about missing integration tests for PUT/DELETE when invoked by mcp_server_manager role. Adds tests for edit success/403, delete success with team cleanup/403, and the _remove_mcp_server_from_team helper directly. Co-Authored-By: Claude Opus 4.6 --- .../test_mcp_manager_role.py | 239 ++++++++++++++++++ 1 file changed, 239 insertions(+) diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py index ed601dfffa..36cecde65b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py @@ -69,6 +69,7 @@ class TestIsUserTeamMcpManager: from unittest.mock import AsyncMock, MagicMock, patch +from fastapi import HTTPException from litellm.proxy._types import LiteLLM_TeamTableCachedObj @@ -316,3 +317,241 @@ class TestCreateMcpServerAsManager: where={"team_id": "team1"}, data={"object_permission_id": "new_perm_id"}, ) + + +@pytest.mark.asyncio +@pytest.mark.skipif( + not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed" +) +class TestEditMcpServerAsManager: + async def test_edit_succeeds_for_mcp_manager(self): + """MCP manager should be able to edit a server assigned to their team.""" + from litellm.proxy._types import UpdateMCPServerRequest + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + edit_mcp_server, + ) + + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user1", + api_key="sk-test", + team_id="team1", + ) + + payload = UpdateMCPServerRequest( + server_id="server1", + description="Updated description", + ) + + mock_team = LiteLLM_TeamTableCachedObj( + team_id="team1", + members_with_roles=[ + Member(user_id="user1", role="mcp_server_manager"), + ], + ) + + updated_server = MagicMock() + updated_server.server_id = "server1" + updated_server.credentials = None + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", + AsyncMock(return_value=("team1", mock_team)), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.update_mcp_server", + AsyncMock(return_value=updated_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + MagicMock( + update_server=AsyncMock(), + reload_servers_from_database=AsyncMock(), + ), + ), + ): + result = await edit_mcp_server(payload=payload, user_api_key_dict=user_auth) + # Should not raise — edit succeeded + assert result is not None + + async def test_edit_fails_for_server_not_in_team(self): + """MCP manager should get 403 when editing a server not in their team.""" + from litellm.proxy._types import UpdateMCPServerRequest + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + edit_mcp_server, + ) + + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user1", + api_key="sk-test", + team_id="team1", + ) + + payload = UpdateMCPServerRequest( + server_id="server_not_in_team", + description="Updated description", + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", + AsyncMock(side_effect=HTTPException(status_code=403, detail="Not in team")), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await edit_mcp_server(payload=payload, user_api_key_dict=user_auth) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +@pytest.mark.skipif( + not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed" +) +class TestDeleteMcpServerAsManager: + async def test_delete_succeeds_and_cleans_up_team(self): + """MCP manager deleting a server should also remove it from team permissions.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + remove_mcp_server, + ) + + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user1", + api_key="sk-test", + team_id="team1", + ) + + mock_team = LiteLLM_TeamTableCachedObj( + team_id="team1", + members_with_roles=[ + Member(user_id="user1", role="mcp_server_manager"), + ], + ) + + deleted_server = MagicMock() + deleted_server.server_id = "server1" + + mock_remove_from_team = AsyncMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", + AsyncMock(return_value=("team1", mock_team)), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.delete_mcp_server", + AsyncMock(return_value=deleted_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + MagicMock( + remove_server=MagicMock(), + reload_servers_from_database=AsyncMock(), + ), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._remove_mcp_server_from_team", + mock_remove_from_team, + ), + ): + response = await remove_mcp_server( + server_id="server1", user_api_key_dict=user_auth + ) + assert response.status_code == 202 + + # Verify team cleanup was called + mock_remove_from_team.assert_called_once() + call_kwargs = mock_remove_from_team.call_args.kwargs + assert call_kwargs["server_id"] == "server1" + assert call_kwargs["team_obj"] == mock_team + + async def test_delete_fails_for_server_not_in_team(self): + """MCP manager should get 403 when deleting a server not in their team.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + remove_mcp_server, + ) + + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user1", + api_key="sk-test", + team_id="team1", + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", + AsyncMock(side_effect=HTTPException(status_code=403, detail="Not in team")), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await remove_mcp_server( + server_id="server_not_in_team", user_api_key_dict=user_auth + ) + assert exc_info.value.status_code == 403 + + async def test_remove_mcp_server_from_team_helper(self): + """_remove_mcp_server_from_team should update the permission list without the deleted server.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _remove_mcp_server_from_team, + ) + + mock_team = LiteLLM_TeamTableCachedObj( + team_id="team1", + members_with_roles=[ + Member(user_id="user1", role="mcp_server_manager"), + ], + object_permission_id="perm1", + ) + mock_team.object_permission = MagicMock( + mcp_servers=["server1", "server2", "server3"] + ) + + mock_handle_common = AsyncMock() + + with ( + patch( + "litellm.proxy.proxy_server.prisma_client", + MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common", + mock_handle_common, + ), + ): + await _remove_mcp_server_from_team( + server_id="server2", + team_obj=mock_team, + ) + + mock_handle_common.assert_called_once() + call_kwargs = mock_handle_common.call_args.kwargs + updated_servers = call_kwargs["data_json"]["object_permission"]["mcp_servers"] + assert "server2" not in updated_servers + assert "server1" in updated_servers + assert "server3" in updated_servers + assert call_kwargs["existing_object_permission_id"] == "perm1" From 700fd86de971f0ca50bafafed129ab17af661f3a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 20 Mar 2026 12:13:09 -0700 Subject: [PATCH 22/26] Fix importorskip guard and add LiteLLM_TeamTableCachedObj import - Add pytest.importorskip("mcp") at module level so tests skip cleanly in CI environments without the mcp package (instead of ImportError) - Import LiteLLM_TeamTableCachedObj into MCP_AVAILABLE block so type annotations resolve for static analysis and get_type_hints() - Remove string quotes from type annotations now that the import exists Co-Authored-By: Claude Opus 4.6 --- .../proxy/management_endpoints/mcp_management_endpoints.py | 7 ++++--- .../proxy/management_endpoints/test_mcp_manager_role.py | 2 ++ 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index e137ad15f1..8c8bd62988 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -112,6 +112,7 @@ if MCP_AVAILABLE: ) from litellm.proxy._types import ( LiteLLM_MCPServerTable, + LiteLLM_TeamTableCachedObj, LitellmUserRoles, MakeMCPServersPublicRequest, MCPApprovalStatus, @@ -147,7 +148,7 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth, team_id: Optional[str] = None, server_id: Optional[str] = None, - ) -> Tuple[str, "LiteLLM_TeamTableCachedObj"]: + ) -> Tuple[str, LiteLLM_TeamTableCachedObj]: """ Verify that the caller is an MCP server manager for a team and (for edit/delete) that the target server belongs to that team. @@ -206,7 +207,7 @@ if MCP_AVAILABLE: async def _auto_assign_mcp_server_to_team( server_id: str, team_id: str, - team_obj: "LiteLLM_TeamTableCachedObj", + team_obj: LiteLLM_TeamTableCachedObj, prisma_client: Any, ) -> None: """ @@ -245,7 +246,7 @@ if MCP_AVAILABLE: async def _remove_mcp_server_from_team( server_id: str, - team_obj: "LiteLLM_TeamTableCachedObj", + team_obj: LiteLLM_TeamTableCachedObj, ) -> None: """Remove a server ID from a team's ObjectPermissionTable.mcp_servers list.""" from litellm.proxy.proxy_server import prisma_client diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py index 36cecde65b..86b0419fb2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py @@ -1,5 +1,7 @@ import pytest +pytest.importorskip("mcp", reason="mcp package not installed") + from litellm.proxy._types import ( LiteLLM_TeamTable, LitellmUserRoles, From ba4aae02c75a4087e8041e6781c0501904ea6f1d Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 20 Mar 2026 13:10:10 -0700 Subject: [PATCH 23/26] Fix outdated MCP server auth test for team MCP manager flow The test_create_mcp_server_auth_failure test expected a 403 for non-admin users, but the team MCP manager feature changed the auth flow to first check for team_id (400) before checking permissions. Split into two tests: one for missing team_id (400) and one for non-manager rejection (403). Co-Authored-By: Claude Opus 4.6 --- .../test_mcp_servers.py | 60 +++++++++++++++---- 1 file changed, 49 insertions(+), 11 deletions(-) diff --git a/tests/store_model_in_db_tests/test_mcp_servers.py b/tests/store_model_in_db_tests/test_mcp_servers.py index a369cce83c..68ef977615 100644 --- a/tests/store_model_in_db_tests/test_mcp_servers.py +++ b/tests/store_model_in_db_tests/test_mcp_servers.py @@ -6,7 +6,7 @@ import os import asyncio from unittest import mock from fastapi.testclient import TestClient -from fastapi import FastAPI +from fastapi import FastAPI, HTTPException from starlette import status @@ -259,45 +259,83 @@ async def test_create_duplicate_mcp_server(): @pytest.mark.asyncio -async def test_create_mcp_server_auth_failure(): +async def test_create_mcp_server_auth_failure_no_team_id(): """ - Test that non-admin users cannot create MCP servers. + Test that non-admin users without a team_id get a 400 error + requiring team_id for team MCP manager flow. """ - # Mock the database functions directly with mock.patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.MCP_AVAILABLE", True, ), mock.patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw" ) as mock_get_prisma: - # Import after mocking from litellm.proxy.management_endpoints.mcp_management_endpoints import ( add_mcp_server, ) from fastapi import HTTPException - # Mock database client mock_prisma = mock.Mock() mock_get_prisma.return_value = mock_prisma - # Set up test data server_id = str(uuid.uuid4()) mcp_server_request = generate_mcpserver_create_request(server_id=server_id) - # Create mock user auth without admin role user_auth = UserAPIKeyAuth( api_key=TEST_MASTER_KEY, user_id="test-user", - user_role=LitellmUserRoles.INTERNAL_USER, # Not an admin + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + with pytest.raises(HTTPException) as exc_info: + await add_mcp_server( + payload=mcp_server_request, user_api_key_dict=user_auth + ) + + assert exc_info.value.status_code == 400 + assert "team_id is required" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_create_mcp_server_auth_failure_not_manager(): + """ + Test that non-admin users with a team_id but without MCP manager + permissions get a 403 error. + """ + with mock.patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.MCP_AVAILABLE", + True, + ), mock.patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw" + ) as mock_get_prisma, mock.patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", + side_effect=HTTPException( + status_code=403, + detail={"error": "You do not have permission to manage MCP servers for this team."}, + ), + ): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + add_mcp_server, + ) + + mock_prisma = mock.Mock() + mock_get_prisma.return_value = mock_prisma + + server_id = str(uuid.uuid4()) + mcp_server_request = generate_mcpserver_create_request(server_id=server_id) + mcp_server_request.team_id = "some-team-id" + + user_auth = UserAPIKeyAuth( + api_key=TEST_MASTER_KEY, + user_id="test-user", + user_role=LitellmUserRoles.INTERNAL_USER, ) - # Expect HTTPException to be raised with pytest.raises(HTTPException) as exc_info: await add_mcp_server( payload=mcp_server_request, user_api_key_dict=user_auth ) - # Verify the exception details assert exc_info.value.status_code == 403 assert "permission" in str(exc_info.value.detail) From c9683c6f9751b6877b938a5b2f29937394c8179d Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 20 Mar 2026 15:41:57 -0700 Subject: [PATCH 24/26] Revert "[Feature] Team MCP Server Manager Role" --- litellm/proxy/_experimental/mcp_server/db.py | 4 +- litellm/proxy/_types.py | 7 +- .../management_endpoints/common_utils.py | 12 - .../mcp_management_endpoints.py | 221 +------ .../test_mcp_manager_role.py | 559 ------------------ 5 files changed, 31 insertions(+), 772 deletions(-) delete mode 100644 tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 0bc102889b..fbef33c32e 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -40,8 +40,8 @@ def _prepare_mcp_server_data( """ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - # Convert model to dict, excluding fields not in the DB schema - data_dict = data.model_dump(exclude_none=True, exclude={"team_id"}) + # Convert model to dict + data_dict = data.model_dump(exclude_none=True) # Ensure alias is always present in the dict (even if None) if "alias" not in data_dict: data_dict["alias"] = getattr(data, "alias", None) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 826e5c9a7b..91a953c217 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1144,10 +1144,6 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): None, description="Server-managed: set by the endpoint; caller values are overridden.", ) - team_id: Optional[str] = Field( - None, - description="Team ID to assign the MCP server to. Required for team MCP managers.", - ) @model_validator(mode="before") @classmethod @@ -1612,9 +1608,8 @@ class Member(MemberBase): role: Literal[ "admin", "user", - "mcp_server_manager", ] = Field( - description="The role of the user within the team. 'admin' users can manage team settings and members, 'user' is a regular team member, 'mcp_server_manager' can manage MCP servers for the team" + description="The role of the user within the team. 'admin' users can manage team settings and members, 'user' is a regular team member" ) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 092b22cee6..efc42d3355 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -41,18 +41,6 @@ def _is_user_team_admin( return False -def _is_user_team_mcp_manager( - user_api_key_dict: UserAPIKeyAuth, team_obj: LiteLLM_TeamTable -) -> bool: - for member in team_obj.members_with_roles: - if ( - member.user_id is not None and member.user_id == user_api_key_dict.user_id - ) and member.role == "mcp_server_manager": - return True - - return False - - async def _is_user_org_admin_for_team( user_api_key_dict: UserAPIKeyAuth, team_obj: LiteLLM_TeamTable ) -> bool: diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 883e1a94aa..f29a721ede 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -21,7 +21,7 @@ import json import os from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple +from typing import Any, Dict, Iterable, List, Literal, Optional from fastapi import ( APIRouter, @@ -112,7 +112,6 @@ if MCP_AVAILABLE: ) from litellm.proxy._types import ( LiteLLM_MCPServerTable, - LiteLLM_TeamTableCachedObj, LitellmUserRoles, MakeMCPServersPublicRequest, MCPApprovalStatus, @@ -131,142 +130,11 @@ if MCP_AVAILABLE: ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.http_parsing_utils import _read_request_body - from litellm.proxy.auth.auth_checks import get_team_object - from litellm.proxy.management_endpoints.common_utils import ( - _is_user_team_mcp_manager, - _user_has_admin_view, - ) - from litellm.proxy.management_helpers.object_permission_utils import ( - _get_team_allowed_mcp_servers, - handle_update_object_permission_common, - ) + from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.types.mcp import MCPCredentials from litellm.types.mcp_server.mcp_server_manager import MCPServer - async def _assert_can_manage_team_mcp_server( - user_api_key_dict: UserAPIKeyAuth, - team_id: Optional[str] = None, - server_id: Optional[str] = None, - ) -> Tuple[str, LiteLLM_TeamTableCachedObj]: - """ - Verify that the caller is an MCP server manager for a team and (for edit/delete) - that the target server belongs to that team. - - Returns a tuple of (team_id, team_obj) for downstream use. - Raises HTTPException(400) if no team_id can be determined. - Raises HTTPException(403) if the caller is not an MCP manager or server not in team. - """ - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache - - resolved_team_id = team_id or user_api_key_dict.team_id - if not resolved_team_id: - raise HTTPException( - status_code=400, - detail={"error": "team_id is required for MCP server manager operations."}, - ) - - # When the API key is team-scoped, ensure the request team_id matches - if ( - team_id - and user_api_key_dict.team_id - and team_id != user_api_key_dict.team_id - ): - raise HTTPException( - status_code=403, - detail={"error": "team_id does not match the API key's team."}, - ) - - team_obj = await get_team_object( - team_id=resolved_team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=True, # bypass cache to get fresh object_permission - ) - - if not _is_user_team_mcp_manager(user_api_key_dict, team_obj): - raise HTTPException( - status_code=403, - detail={ - "error": f"User does not have mcp_server_manager role in team {resolved_team_id}." - }, - ) - - if server_id is not None: - team_server_ids = await _get_team_allowed_mcp_servers(team_obj) - if server_id not in team_server_ids: - raise HTTPException( - status_code=403, - detail={ - "error": f"MCP server {server_id} is not assigned to team {resolved_team_id}." - }, - ) - - return resolved_team_id, team_obj - - async def _auto_assign_mcp_server_to_team( - server_id: str, - team_id: str, - team_obj: LiteLLM_TeamTableCachedObj, - prisma_client: Any, - ) -> None: - """ - Add an MCP server to a team's ObjectPermissionTable and link the - permission back to the team if it didn't have one yet. - - Uses handle_update_object_permission (the team-endpoint helper) so - the object_permission_id linkage follows the same pattern as - team_endpoints.update_team. - """ - from litellm.proxy.management_endpoints.team_endpoints import ( - handle_update_object_permission, - ) - - existing_mcp_servers: list = [] - if team_obj.object_permission is not None: - existing_mcp_servers = team_obj.object_permission.mcp_servers or [] - updated_mcp_servers = list(set(existing_mcp_servers + [server_id])) - - # Build the data dict in the same shape team_endpoints uses - data_json: Dict[str, Any] = { - "object_permission": {"mcp_servers": updated_mcp_servers}, - } - data_json = await handle_update_object_permission( - data_json=data_json, - existing_team_row=team_obj, - ) - - # If handle_update_object_permission produced an object_permission_id, - # persist it on the team row (it sets data_json["object_permission_id"]). - if "object_permission_id" in data_json: - await prisma_client.db.litellm_teamtable.update( - where={"team_id": team_id}, - data={"object_permission_id": data_json["object_permission_id"]}, - ) - - async def _remove_mcp_server_from_team( - server_id: str, - team_obj: LiteLLM_TeamTableCachedObj, - ) -> None: - """Remove a server ID from a team's ObjectPermissionTable.mcp_servers list.""" - from litellm.proxy.proxy_server import prisma_client - - existing_permission_id = getattr(team_obj, "object_permission_id", None) - if ( - existing_permission_id is None - or team_obj.object_permission is None - ): - return - - existing_mcp_servers = team_obj.object_permission.mcp_servers or [] - updated_mcp_servers = [s for s in existing_mcp_servers if s != server_id] - - await handle_update_object_permission_common( - data_json={"object_permission": {"mcp_servers": updated_mcp_servers}}, - existing_object_permission_id=existing_permission_id, - prisma_client=prisma_client, - ) - @dataclass class _TemporaryMCPServerEntry: server: MCPServer @@ -1338,27 +1206,16 @@ if MCP_AVAILABLE: # Validate and normalize payload fields validate_and_normalize_mcp_server_payload(payload) - # AuthZ - proxy admins or team MCP managers can create MCP servers - is_proxy_admin = LitellmUserRoles.PROXY_ADMIN == user_api_key_dict.user_role - manager_team_id: Optional[str] = None - - manager_team_obj = None - if not is_proxy_admin: - # Check if the user is an MCP manager for the specified team - if payload.team_id is None: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": "team_id is required when creating MCP servers as a team MCP manager." - }, - ) - manager_team_id, manager_team_obj = await _assert_can_manage_team_mcp_server( - user_api_key_dict=user_api_key_dict, - team_id=payload.team_id, + # AuthZ - restrict only proxy admins to create mcp servers + if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": "User does not have permission to create mcp servers. You can only create mcp servers if you are a PROXY_ADMIN." + }, ) - - # Fail if the MCP server with this id already exists - if payload.server_id is not None: + elif payload.server_id is not None: + # fail if the mcp server with id already exists mcp_server = await get_mcp_server(prisma_client, payload.server_id) if mcp_server is not None: raise HTTPException( @@ -1367,8 +1224,7 @@ if MCP_AVAILABLE: "error": f"MCP Server with id {payload.server_id} already exists. Cannot create another." }, ) - - if ( + elif ( SpecialMCPServerName.all_team_servers == payload.server_id or SpecialMCPServerName.all_proxy_servers == payload.server_id ): @@ -1399,28 +1255,12 @@ if MCP_AVAILABLE: # Ensure registry is up to date by reloading from database await global_mcp_server_manager.reload_servers_from_database() - except Exception as e: verbose_proxy_logger.exception(f"Error creating mcp server: {str(e)}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={"error": f"Error creating mcp server: {str(e)}"}, ) - - # Auto-assign the server to the manager's team (separate from create - # so a failure here doesn't mask the successfully created server). - if manager_team_id is not None and manager_team_obj is not None: - try: - await _auto_assign_mcp_server_to_team( - server_id=new_mcp_server.server_id, - team_id=manager_team_id, - team_obj=manager_team_obj, - prisma_client=prisma_client, - ) - except Exception as e: - verbose_proxy_logger.exception( - f"MCP server created but failed to auto-assign to team {manager_team_id}: {str(e)}" - ) return _redact_mcp_credentials(new_mcp_server) @router.post( @@ -1637,12 +1477,15 @@ if MCP_AVAILABLE: "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - # Authz - proxy admins or team MCP managers can delete MCP servers - manager_team_obj = None + # Authz - restrict only admins to delete mcp servers if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: - _, manager_team_obj = await _assert_can_manage_team_mcp_server( - user_api_key_dict=user_api_key_dict, - server_id=server_id, + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": "Call not allowed to delete MCP server. User is not a proxy admin. route={}".format( + "DELETE /v1/mcp/server" + ) + }, ) # try to delete the mcp server @@ -1658,18 +1501,6 @@ if MCP_AVAILABLE: # Ensure registry is up to date by reloading from database await global_mcp_server_manager.reload_servers_from_database() - # Remove server from the manager's team permission list - if manager_team_obj is not None: - try: - await _remove_mcp_server_from_team( - server_id=server_id, - team_obj=manager_team_obj, - ) - except Exception as e: - verbose_proxy_logger.exception( - f"MCP server deleted but failed to remove from team permissions: {str(e)}" - ) - # TODO: Enterprise: Finish audit log trail if litellm.store_audit_logs: pass @@ -1971,11 +1802,15 @@ if MCP_AVAILABLE: # Validate and normalize payload fields validate_and_normalize_mcp_server_payload(payload) - # Authz - proxy admins or team MCP managers can update MCP servers + # Authz - restrict only admins to delete mcp servers if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: - await _assert_can_manage_team_mcp_server( - user_api_key_dict=user_api_key_dict, - server_id=payload.server_id, + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": "Call not allowed to update MCP server. User is not a proxy admin. route={}".format( + "PUT /v1/mcp/server" + ) + }, ) # try to update the mcp server diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py deleted file mode 100644 index 86b0419fb2..0000000000 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py +++ /dev/null @@ -1,559 +0,0 @@ -import pytest - -pytest.importorskip("mcp", reason="mcp package not installed") - -from litellm.proxy._types import ( - LiteLLM_TeamTable, - LitellmUserRoles, - Member, - UserAPIKeyAuth, -) -from litellm.proxy.management_endpoints.common_utils import ( - _is_user_team_mcp_manager, -) -from litellm.proxy.management_endpoints import ( - mcp_management_endpoints as mgmt_endpoints, -) - - -class TestIsUserTeamMcpManager: - def test_mcp_server_manager_role_returns_true(self): - user_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, - user_id="user1", - api_key="sk-test", - ) - team = LiteLLM_TeamTable( - team_id="team1", - members_with_roles=[ - Member(user_id="user1", role="mcp_server_manager") - ], - ) - assert _is_user_team_mcp_manager(user_auth, team) is True - - def test_regular_user_role_returns_false(self): - user_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, - user_id="user1", - api_key="sk-test", - ) - team = LiteLLM_TeamTable( - team_id="team1", - members_with_roles=[Member(user_id="user1", role="user")], - ) - assert _is_user_team_mcp_manager(user_auth, team) is False - - def test_admin_role_returns_false(self): - user_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, - user_id="user1", - api_key="sk-test", - ) - team = LiteLLM_TeamTable( - team_id="team1", - members_with_roles=[Member(user_id="user1", role="admin")], - ) - assert _is_user_team_mcp_manager(user_auth, team) is False - - def test_user_not_in_team_returns_false(self): - user_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, - user_id="user2", - api_key="sk-test", - ) - team = LiteLLM_TeamTable( - team_id="team1", - members_with_roles=[ - Member(user_id="user1", role="mcp_server_manager") - ], - ) - assert _is_user_team_mcp_manager(user_auth, team) is False - - -from unittest.mock import AsyncMock, MagicMock, patch -from fastapi import HTTPException -from litellm.proxy._types import LiteLLM_TeamTableCachedObj - - -@pytest.mark.asyncio -@pytest.mark.skipif( - not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed" -) -class TestAssertCanManageTeamMcpServer: - async def test_mcp_manager_with_team_id_succeeds(self): - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - _assert_can_manage_team_mcp_server, - ) - user_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test", - ) - mock_team = LiteLLM_TeamTableCachedObj( - team_id="team1", - members_with_roles=[Member(user_id="user1", role="mcp_server_manager")], - ) - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object", - AsyncMock(return_value=mock_team), - ): - team_id, team_obj = await _assert_can_manage_team_mcp_server( - user_api_key_dict=user_auth, team_id="team1" - ) - assert team_id == "team1" - assert team_obj == mock_team - - async def test_regular_user_gets_403(self): - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - _assert_can_manage_team_mcp_server, - ) - user_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test", - ) - mock_team = LiteLLM_TeamTableCachedObj( - team_id="team1", - members_with_roles=[Member(user_id="user1", role="user")], - ) - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object", - AsyncMock(return_value=mock_team), - ): - with pytest.raises(Exception) as exc_info: - await _assert_can_manage_team_mcp_server( - user_api_key_dict=user_auth, team_id="team1" - ) - assert exc_info.value.status_code == 403 - - async def test_admin_gets_403(self): - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - _assert_can_manage_team_mcp_server, - ) - user_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test", - ) - mock_team = LiteLLM_TeamTableCachedObj( - team_id="team1", - members_with_roles=[Member(user_id="user1", role="admin")], - ) - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object", - AsyncMock(return_value=mock_team), - ): - with pytest.raises(Exception) as exc_info: - await _assert_can_manage_team_mcp_server( - user_api_key_dict=user_auth, team_id="team1" - ) - assert exc_info.value.status_code == 403 - - async def test_mcp_manager_server_in_team_succeeds(self): - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - _assert_can_manage_team_mcp_server, - ) - user_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test", team_id="team1", - ) - mock_team = LiteLLM_TeamTableCachedObj( - team_id="team1", - members_with_roles=[Member(user_id="user1", role="mcp_server_manager")], - ) - with ( - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object", - AsyncMock(return_value=mock_team), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_allowed_mcp_servers", - AsyncMock(return_value={"server1", "server2"}), - ), - ): - team_id, team_obj = await _assert_can_manage_team_mcp_server( - user_api_key_dict=user_auth, server_id="server1" - ) - assert team_id == "team1" - assert team_obj == mock_team - - async def test_mcp_manager_server_not_in_team_gets_403(self): - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - _assert_can_manage_team_mcp_server, - ) - user_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test", team_id="team1", - ) - mock_team = LiteLLM_TeamTableCachedObj( - team_id="team1", - members_with_roles=[Member(user_id="user1", role="mcp_server_manager")], - ) - with ( - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object", - AsyncMock(return_value=mock_team), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_allowed_mcp_servers", - AsyncMock(return_value={"server2", "server3"}), - ), - ): - with pytest.raises(Exception) as exc_info: - await _assert_can_manage_team_mcp_server( - user_api_key_dict=user_auth, server_id="server1" - ) - assert exc_info.value.status_code == 403 - - -@pytest.mark.asyncio -@pytest.mark.skipif( - not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed" -) -class TestCreateMcpServerAsManager: - async def test_create_auto_assigns_to_team(self): - """MCP manager creating a server should auto-assign it to their team's ObjectPermissionTable.""" - from litellm.proxy._types import NewMCPServerRequest - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - add_mcp_server, - ) - - user_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, - user_id="user1", - api_key="sk-test", - team_id="team1", - ) - - payload = NewMCPServerRequest( - server_name="test-server", - url="https://example.com/mcp", - team_id="team1", - ) - - mock_team = LiteLLM_TeamTableCachedObj( - team_id="team1", - members_with_roles=[ - Member(user_id="user1", role="mcp_server_manager"), - ], - object_permission_id="perm1", - ) - mock_team.object_permission = MagicMock(mcp_servers=["existing_server"]) - - created_server = MagicMock() - created_server.server_id = "new_server_id" - created_server.credentials = None - - mock_auto_assign = AsyncMock() - - with ( - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=MagicMock(), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", - AsyncMock(return_value=("team1", mock_team)), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(return_value=None), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server", - AsyncMock(return_value=created_server), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - MagicMock( - add_server=AsyncMock(), - reload_servers_from_database=AsyncMock(), - ), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._auto_assign_mcp_server_to_team", - mock_auto_assign, - ), - ): - await add_mcp_server(payload=payload, user_api_key_dict=user_auth) - - # Verify _auto_assign_mcp_server_to_team was called with the right args - mock_auto_assign.assert_called_once() - call_kwargs = mock_auto_assign.call_args.kwargs - assert call_kwargs["server_id"] == "new_server_id" - assert call_kwargs["team_id"] == "team1" - assert call_kwargs["team_obj"] == mock_team - - async def test_auto_assign_links_new_permission_to_team(self): - """_auto_assign_mcp_server_to_team should create permission and link to team.""" - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - _auto_assign_mcp_server_to_team, - ) - - # Team with NO object_permission_id - mock_team = LiteLLM_TeamTableCachedObj( - team_id="team1", - members_with_roles=[ - Member(user_id="user1", role="mcp_server_manager"), - ], - object_permission_id=None, - ) - mock_team.object_permission = None - - mock_prisma = MagicMock() - mock_prisma.db.litellm_teamtable.update = AsyncMock() - - # handle_update_object_permission sets object_permission_id in data_json - async def fake_handle(data_json, existing_team_row): - data_json["object_permission_id"] = "new_perm_id" - return data_json - - with patch( - "litellm.proxy.management_endpoints.team_endpoints.handle_update_object_permission", - side_effect=fake_handle, - ): - await _auto_assign_mcp_server_to_team( - server_id="new_server_id", - team_id="team1", - team_obj=mock_team, - prisma_client=mock_prisma, - ) - - # Verify the team was updated with the new object_permission_id - mock_prisma.db.litellm_teamtable.update.assert_called_once_with( - where={"team_id": "team1"}, - data={"object_permission_id": "new_perm_id"}, - ) - - -@pytest.mark.asyncio -@pytest.mark.skipif( - not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed" -) -class TestEditMcpServerAsManager: - async def test_edit_succeeds_for_mcp_manager(self): - """MCP manager should be able to edit a server assigned to their team.""" - from litellm.proxy._types import UpdateMCPServerRequest - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - edit_mcp_server, - ) - - user_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, - user_id="user1", - api_key="sk-test", - team_id="team1", - ) - - payload = UpdateMCPServerRequest( - server_id="server1", - description="Updated description", - ) - - mock_team = LiteLLM_TeamTableCachedObj( - team_id="team1", - members_with_roles=[ - Member(user_id="user1", role="mcp_server_manager"), - ], - ) - - updated_server = MagicMock() - updated_server.server_id = "server1" - updated_server.credentials = None - - with ( - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=MagicMock(), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", - AsyncMock(return_value=("team1", mock_team)), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.update_mcp_server", - AsyncMock(return_value=updated_server), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - MagicMock( - update_server=AsyncMock(), - reload_servers_from_database=AsyncMock(), - ), - ), - ): - result = await edit_mcp_server(payload=payload, user_api_key_dict=user_auth) - # Should not raise — edit succeeded - assert result is not None - - async def test_edit_fails_for_server_not_in_team(self): - """MCP manager should get 403 when editing a server not in their team.""" - from litellm.proxy._types import UpdateMCPServerRequest - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - edit_mcp_server, - ) - - user_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, - user_id="user1", - api_key="sk-test", - team_id="team1", - ) - - payload = UpdateMCPServerRequest( - server_id="server_not_in_team", - description="Updated description", - ) - - with ( - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=MagicMock(), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", - AsyncMock(side_effect=HTTPException(status_code=403, detail="Not in team")), - ), - ): - with pytest.raises(HTTPException) as exc_info: - await edit_mcp_server(payload=payload, user_api_key_dict=user_auth) - assert exc_info.value.status_code == 403 - - -@pytest.mark.asyncio -@pytest.mark.skipif( - not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed" -) -class TestDeleteMcpServerAsManager: - async def test_delete_succeeds_and_cleans_up_team(self): - """MCP manager deleting a server should also remove it from team permissions.""" - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - remove_mcp_server, - ) - - user_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, - user_id="user1", - api_key="sk-test", - team_id="team1", - ) - - mock_team = LiteLLM_TeamTableCachedObj( - team_id="team1", - members_with_roles=[ - Member(user_id="user1", role="mcp_server_manager"), - ], - ) - - deleted_server = MagicMock() - deleted_server.server_id = "server1" - - mock_remove_from_team = AsyncMock() - - with ( - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=MagicMock(), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", - AsyncMock(return_value=("team1", mock_team)), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.delete_mcp_server", - AsyncMock(return_value=deleted_server), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - MagicMock( - remove_server=MagicMock(), - reload_servers_from_database=AsyncMock(), - ), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._remove_mcp_server_from_team", - mock_remove_from_team, - ), - ): - response = await remove_mcp_server( - server_id="server1", user_api_key_dict=user_auth - ) - assert response.status_code == 202 - - # Verify team cleanup was called - mock_remove_from_team.assert_called_once() - call_kwargs = mock_remove_from_team.call_args.kwargs - assert call_kwargs["server_id"] == "server1" - assert call_kwargs["team_obj"] == mock_team - - async def test_delete_fails_for_server_not_in_team(self): - """MCP manager should get 403 when deleting a server not in their team.""" - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - remove_mcp_server, - ) - - user_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, - user_id="user1", - api_key="sk-test", - team_id="team1", - ) - - with ( - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=MagicMock(), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", - AsyncMock(side_effect=HTTPException(status_code=403, detail="Not in team")), - ), - ): - with pytest.raises(HTTPException) as exc_info: - await remove_mcp_server( - server_id="server_not_in_team", user_api_key_dict=user_auth - ) - assert exc_info.value.status_code == 403 - - async def test_remove_mcp_server_from_team_helper(self): - """_remove_mcp_server_from_team should update the permission list without the deleted server.""" - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - _remove_mcp_server_from_team, - ) - - mock_team = LiteLLM_TeamTableCachedObj( - team_id="team1", - members_with_roles=[ - Member(user_id="user1", role="mcp_server_manager"), - ], - object_permission_id="perm1", - ) - mock_team.object_permission = MagicMock( - mcp_servers=["server1", "server2", "server3"] - ) - - mock_handle_common = AsyncMock() - - with ( - patch( - "litellm.proxy.proxy_server.prisma_client", - MagicMock(), - ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common", - mock_handle_common, - ), - ): - await _remove_mcp_server_from_team( - server_id="server2", - team_obj=mock_team, - ) - - mock_handle_common.assert_called_once() - call_kwargs = mock_handle_common.call_args.kwargs - updated_servers = call_kwargs["data_json"]["object_permission"]["mcp_servers"] - assert "server2" not in updated_servers - assert "server1" in updated_servers - assert "server3" in updated_servers - assert call_kwargs["existing_object_permission_id"] == "perm1" From 6862930538ed74ecd9411b24b11fd5928af3fa5a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 20 Mar 2026 15:44:00 -0700 Subject: [PATCH 25/26] Revert test to match reverted team MCP manager feature The team MCP manager feature was reverted in PR #24255, so the test needs to go back to the original single auth failure test that expects a 403 for non-admin users. Co-Authored-By: Claude Opus 4.6 --- .../test_mcp_servers.py | 60 ++++--------------- 1 file changed, 11 insertions(+), 49 deletions(-) diff --git a/tests/store_model_in_db_tests/test_mcp_servers.py b/tests/store_model_in_db_tests/test_mcp_servers.py index 68ef977615..a369cce83c 100644 --- a/tests/store_model_in_db_tests/test_mcp_servers.py +++ b/tests/store_model_in_db_tests/test_mcp_servers.py @@ -6,7 +6,7 @@ import os import asyncio from unittest import mock from fastapi.testclient import TestClient -from fastapi import FastAPI, HTTPException +from fastapi import FastAPI from starlette import status @@ -259,83 +259,45 @@ async def test_create_duplicate_mcp_server(): @pytest.mark.asyncio -async def test_create_mcp_server_auth_failure_no_team_id(): +async def test_create_mcp_server_auth_failure(): """ - Test that non-admin users without a team_id get a 400 error - requiring team_id for team MCP manager flow. + Test that non-admin users cannot create MCP servers. """ + # Mock the database functions directly with mock.patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.MCP_AVAILABLE", True, ), mock.patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw" ) as mock_get_prisma: + # Import after mocking from litellm.proxy.management_endpoints.mcp_management_endpoints import ( add_mcp_server, ) from fastapi import HTTPException + # Mock database client mock_prisma = mock.Mock() mock_get_prisma.return_value = mock_prisma + # Set up test data server_id = str(uuid.uuid4()) mcp_server_request = generate_mcpserver_create_request(server_id=server_id) + # Create mock user auth without admin role user_auth = UserAPIKeyAuth( api_key=TEST_MASTER_KEY, user_id="test-user", - user_role=LitellmUserRoles.INTERNAL_USER, - ) - - with pytest.raises(HTTPException) as exc_info: - await add_mcp_server( - payload=mcp_server_request, user_api_key_dict=user_auth - ) - - assert exc_info.value.status_code == 400 - assert "team_id is required" in str(exc_info.value.detail) - - -@pytest.mark.asyncio -async def test_create_mcp_server_auth_failure_not_manager(): - """ - Test that non-admin users with a team_id but without MCP manager - permissions get a 403 error. - """ - with mock.patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.MCP_AVAILABLE", - True, - ), mock.patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw" - ) as mock_get_prisma, mock.patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", - side_effect=HTTPException( - status_code=403, - detail={"error": "You do not have permission to manage MCP servers for this team."}, - ), - ): - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - add_mcp_server, - ) - - mock_prisma = mock.Mock() - mock_get_prisma.return_value = mock_prisma - - server_id = str(uuid.uuid4()) - mcp_server_request = generate_mcpserver_create_request(server_id=server_id) - mcp_server_request.team_id = "some-team-id" - - user_auth = UserAPIKeyAuth( - api_key=TEST_MASTER_KEY, - user_id="test-user", - user_role=LitellmUserRoles.INTERNAL_USER, + user_role=LitellmUserRoles.INTERNAL_USER, # Not an admin ) + # Expect HTTPException to be raised with pytest.raises(HTTPException) as exc_info: await add_mcp_server( payload=mcp_server_request, user_api_key_dict=user_auth ) + # Verify the exception details assert exc_info.value.status_code == 403 assert "permission" in str(exc_info.value.detail) From e678ddea4310754a001659e66c170439562cf2ae Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 20 Mar 2026 16:12:48 -0700 Subject: [PATCH 26/26] Fix unreachable special MCP server name guard in add_mcp_server The special name check (all_team_servers, all_proxy_servers) was an elif after the server_id-is-not-None check, making it unreachable since special names are non-None strings. Split into separate if blocks so the special name guard runs before the duplicate-ID check. Co-Authored-By: Claude Opus 4.6 --- .../mcp_management_endpoints.py | 25 +++++++++++-------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index f29a721ede..e4bb288cda 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1214,17 +1214,9 @@ if MCP_AVAILABLE: "error": "User does not have permission to create mcp servers. You can only create mcp servers if you are a PROXY_ADMIN." }, ) - elif payload.server_id is not None: - # fail if the mcp server with id already exists - mcp_server = await get_mcp_server(prisma_client, payload.server_id) - if mcp_server is not None: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": f"MCP Server with id {payload.server_id} already exists. Cannot create another." - }, - ) - elif ( + + # Block reserved special server IDs + if ( SpecialMCPServerName.all_team_servers == payload.server_id or SpecialMCPServerName.all_proxy_servers == payload.server_id ): @@ -1235,6 +1227,17 @@ if MCP_AVAILABLE: }, ) + if payload.server_id is not None: + # fail if the mcp server with id already exists + mcp_server = await get_mcp_server(prisma_client, payload.server_id) + if mcp_server is not None: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": f"MCP Server with id {payload.server_id} already exists. Cannot create another." + }, + ) + # TODO: audit log for create # Admin-created servers are always active — clear any submission lifecycle