From b71f02dbcf5f349794faa2c8cfc264f07bdc1718 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 29 Sep 2026 22:49:20 -0700 Subject: [PATCH] fix(ui): keep MCP permissions visible after key, team and MCP server saves (#43810) * fix(ui): keep MCP permissions visible after key, team and MCP server saves Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): type the object_permission include as a prisma TypedDict Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): do not block key save confirmation on cache refetch Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: ryan Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 2 + litellm/proxy/utils.py | 4 +- .../test_key_management_endpoints.py | 13 +++++- .../test_prisma_client_writes.py | 10 ++++- .../hooks/mcpServers/useMCPServers.ts | 2 +- .../app/(dashboard)/hooks/teams/useTeams.ts | 16 +++++++- .../_components/mcp_server_view.test.tsx | 41 +++++++++++++++++-- .../_components/mcp_server_view.tsx | 5 +++ .../team/TeamInfo.integration.test.tsx | 3 +- .../src/components/team/TeamInfo.test.tsx | 25 ++++++++++- .../src/components/team/TeamInfo.tsx | 4 +- .../KeyInfoView.handleKeyUpdate.test.tsx | 23 ++++++++++- .../components/templates/key_info_view.tsx | 2 +- 13 files changed, 133 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 7e159ec90e7..d37dfe87ad5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2442,11 +2442,13 @@ async def _update_key_row_with_soft_budget( existing_key_row=existing_key_row, changed_by=changed_by, ) + include_object_permission: Final[prisma.types.LiteLLM_VerificationTokenInclude] = {"object_permission": True} updated_row: Final = await tx.litellm_verificationtoken.update( where=key_where, data=with_settings_updated_at( prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token})) ), + include=include_object_permission, ) updated_data: Final[Mapping[str, object]] = ( updated_row.model_dump() if updated_row is not None else MappingProxyType({}) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ea294b76e92..43ad433c19b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -258,7 +258,7 @@ if TYPE_CHECKING: from prisma.actions import LiteLLM_DeprecatedVerificationTokenActions from prisma.client import TransactionManager from prisma.models import LiteLLM_DeprecatedVerificationToken - from prisma.types import HttpConfig + from prisma.types import HttpConfig, LiteLLM_VerificationTokenInclude from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation @@ -5439,9 +5439,11 @@ class PrismaClient: # check if plain text or hash token = _hash_token_if_needed(token=token) db_data["token"] = token + include_object_permission: Final[LiteLLM_VerificationTokenInclude] = {"object_permission": True} response: Final = await VerificationTokenRepository(self).table.update( where={"token": token}, data=with_settings_updated_at(db_data), + include=include_object_permission, ) verbose_proxy_logger.debug("\033[91m" + f"DB Token Table update succeeded {response}" + "\033[0m") _data: dict = {} diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index aa6be328f4a..a5d2828dd9c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -20291,7 +20291,11 @@ async def test_update_key_row_with_soft_budget_updates_budget_and_key_in_transac existing_key = LiteLLM_VerificationToken(token="test-token", budget_id=None) created_row = MagicMock(budget_id="budget-new") updated_row = MagicMock() - updated_row.model_dump.return_value = {"token": "hashed", "budget_id": "budget-new"} + updated_row.model_dump.return_value = { + "token": "hashed", + "budget_id": "budget-new", + "object_permission": {"mcp_servers": ["srv-1"], "mcp_tool_permissions": {"srv-1": ["read"]}}, + } tx = MagicMock() tx.litellm_budgettable.create = AsyncMock(return_value=created_row) tx.litellm_verificationtoken.update = AsyncMock(return_value=updated_row) @@ -20312,10 +20316,15 @@ async def test_update_key_row_with_soft_budget_updates_budget_and_key_in_transac ) assert set(result) == {"token", "data"} - assert result["data"] == {"token": "hashed", "budget_id": "budget-new"} + assert result["data"] == { + "token": "hashed", + "budget_id": "budget-new", + "object_permission": {"mcp_servers": ["srv-1"], "mcp_tool_permissions": {"srv-1": ["read"]}}, + } tx.litellm_verificationtoken.update.assert_awaited_once() update_call = tx.litellm_verificationtoken.update.await_args assert update_call.kwargs["where"] == {"token": result["token"]} + assert update_call.kwargs["include"] == {"object_permission": True} assert update_call.kwargs["data"]["budget_id"] == "budget-new" assert "soft_budget" not in update_call.kwargs["data"] diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py index dd241397e87..6e69444a1b5 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_writes.py @@ -155,6 +155,7 @@ async def test_update_data_token_hashes_and_updates( "token": hashlib.sha256(token.encode()).hexdigest(), "spend": 1.0, "user_id": "u1", + "object_permission": {"mcp_servers": ["srv-1"]}, }, ) prisma_client.db.litellm_verificationtoken.update = AsyncMock(return_value=response) @@ -167,15 +168,22 @@ async def test_update_data_token_hashes_and_updates( actual = { "result": result, "where": update_kwargs["where"], + "include": update_kwargs["include"], "data_token": update_kwargs["data"]["token"], "data_spend": update_kwargs["data"]["spend"], } assert actual == { "result": { "token": hashed, - "data": {"token": hashed, "spend": 1.0, "user_id": "u1"}, + "data": { + "token": hashed, + "spend": 1.0, + "user_id": "u1", + "object_permission": {"mcp_servers": ["srv-1"]}, + }, }, "where": {"token": hashed}, + "include": {"object_permission": True}, "data_token": hashed, "data_spend": 1.0, } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts index 9210e25e1a8..597c5f7b2da 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts @@ -4,7 +4,7 @@ import { fetchMCPServers } from "@/components/networking"; import { MCPServer } from "@/components/mcp_tools/types"; import useAuthorized from "../useAuthorized"; -const mcpServersKeys = createQueryKeys("mcpServers"); +export const mcpServersKeys = createQueryKeys("mcpServers"); export const useMCPServers = (teamId?: string | null) => { const { accessToken } = useAuthorized(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts index 05025adc5e6..7d1d035b4d4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts @@ -1,4 +1,11 @@ -import { keepPreviousData, useInfiniteQuery, useQuery, useQueryClient, UseQueryResult } from "@tanstack/react-query"; +import { + keepPreviousData, + QueryClient, + useInfiniteQuery, + useQuery, + useQueryClient, + UseQueryResult, +} from "@tanstack/react-query"; import { Team } from "@/components/key_team_helpers/key_list"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { fetchTeams } from "@/app/(dashboard)/networking"; @@ -110,7 +117,7 @@ export const useTeamsTable = ( }); }; -const teamKeys = createQueryKeys("teams"); +export const teamKeys = createQueryKeys("teams"); export const useTeams = (): UseQueryResult => { const { accessToken, userId, userRole } = useAuthorized(); return useQuery({ @@ -179,6 +186,11 @@ export const useTeam = (teamId?: string) => { const infiniteTeamKeys = createQueryKeys("infiniteTeams"); +export const invalidateTeamQueries = (queryClient: QueryClient) => + Promise.all( + [teamsTableKeys, teamKeys, infiniteTeamKeys].map((keys) => queryClient.invalidateQueries({ queryKey: keys.all })), + ); + export const useInfiniteTeams = (pageSize: number = 50, search?: string, organizationId?: string | null) => { const { accessToken, userId, userRole } = useAuthorized(); const isAdmin = userRole === "Admin" || userRole === "Admin Viewer"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.test.tsx index 2f7f989c099..7e1eda9143f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.test.tsx @@ -7,13 +7,21 @@ import * as networking from "@/components/networking"; import { setSecureItem } from "@/utils/secureStorage"; import { EDIT_OAUTH_UI_STATE_KEY } from "./mcp_server_edit"; import type { MCPServer } from "@/components/mcp_tools/types"; +import { mcpServersKeys } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; vi.mock(".", () => ({ MCPToolsViewer: () =>
tools viewer
, })); vi.mock("./mcp_server_edit", () => ({ - default: () =>
edit form
, + default: ({ mcpServer, onSuccess }: { mcpServer: MCPServer; onSuccess: (server: MCPServer) => void }) => ( +
+ edit form + +
+ ), EDIT_OAUTH_UI_STATE_KEY: "litellm-mcp-oauth-edit-state", })); @@ -33,9 +41,15 @@ const baseServer = { auth_type: "api_key", } as MCPServer; -const renderView = (overrides: Partial = {}, props: Record = {}) => +const newQueryClient = () => new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } }); + +const renderView = ( + overrides: Partial = {}, + props: Record = {}, + queryClient: QueryClient = newQueryClient(), +) => render( - + { expect(await screen.findByText("edit form")).toBeInTheDocument(); }); + it("drops the cached server list and tool catalog once the edit form saves", async () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: Infinity } } }); + const serversKey = mcpServersKeys.list(); + const toolsKey = ["mcpTools", "srv-1", {}, null]; + const otherToolsKey = ["mcpTools", "srv-2", {}, null]; + queryClient.setQueryData(serversKey, [baseServer]); + queryClient.setQueryData(toolsKey, { tools: [] }); + queryClient.setQueryData(otherToolsKey, { tools: [] }); + const onBack = vi.fn(); + renderView({}, { onBack }, queryClient); + + await userEvent.click(screen.getByRole("tab", { name: "Settings" })); + await userEvent.click(await screen.findByRole("button", { name: "Edit Settings" })); + await userEvent.click(await screen.findByRole("button", { name: "save edit" })); + + expect(queryClient.getQueryState(serversKey)?.isInvalidated).toBe(true); + expect(queryClient.getQueryState(toolsKey)?.isInvalidated).toBe(true); + expect(queryClient.getQueryState(otherToolsKey)?.isInvalidated).toBe(false); + expect(onBack).toHaveBeenCalledTimes(1); + }); + it("opens straight into the edit form when isEditing is set", async () => { renderView({}, { isEditing: true }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.tsx index c97596ce0f6..278a98a2fa6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_view.tsx @@ -1,4 +1,6 @@ import React, { useState } from "react"; +import { useQueryClient } from "@tanstack/react-query"; +import { mcpServersKeys } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import { ArrowLeft, Eye, EyeOff } from "lucide-react"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; @@ -64,6 +66,7 @@ export const MCPServerView: React.FC = ({ }) => { // Open the editing Settings tab on first render when returning from the edit OAuth // redirect, so the "token fetched" feedback shows where the user left off (Settings=2). + const queryClient = useQueryClient(); const canEdit = isProxyAdmin && !isViewOnly && !mcpServer.is_config; const returningFromEditOAuth = isReturningFromEditOAuth(canEdit, mcpServer.server_id); const [editing, setEditing] = useState(isEditing || returningFromEditOAuth); @@ -75,6 +78,8 @@ export const MCPServerView: React.FC = ({ const canRevokeUserCredentials = userRole !== null && isProxyAdminRole(userRole) && !isViewOnly; const handleSuccess = (updated: MCPServer) => { + void queryClient.invalidateQueries({ queryKey: mcpServersKeys.all }); + void queryClient.invalidateQueries({ queryKey: ["mcpTools", updated.server_id] }); setEditing(false); onBack(); }; diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx index 2d84d78d56c..456dc91b13f 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx @@ -77,7 +77,8 @@ vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ useAllProxyModels: vi.fn(), })); -vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", async (importOriginal) => ({ + ...(await importOriginal()), useTeam: vi.fn(), })); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index a693ee971d4..4650f4b6987 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -83,7 +83,8 @@ vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ useAllProxyModels: vi.fn(), })); -vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", async (importOriginal) => ({ + ...(await importOriginal()), useTeam: vi.fn(), })); @@ -233,7 +234,7 @@ vi.mock("../key_team_helpers/filter_helpers", () => ({ import { useAllProxyModels } from "@/app/(dashboard)/hooks/models/useModels"; import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; -import { useTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { teamKeys, teamsTableKeys, useTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; @@ -1146,6 +1147,26 @@ describe("TeamInfoView", () => { }); }); + it("invalidates the cached team list and team detail queries after saving team settings", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: ["gpt-4"] })); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + const tableKey = teamsTableKeys.list({ page: 1, limit: 10 }); + const detailKey = teamKeys.detail("123"); + testQueryClient.setQueryData(tableKey, { teams: [], total: 0 }); + testQueryClient.setQueryData(detailKey, { team_id: "123" }); + + renderWithProviders(); + + await user.click(await screen.findByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await screen.findByLabelText("Team Name"); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(testQueryClient.getQueryState(tableKey)?.isInvalidated).toBe(true)); + expect(testQueryClient.getQueryState(detailKey)?.isInvalidated).toBe(true); + }); + const openSettingsEditorForTeam = async ( user: ReturnType, teamOverrides: Record, diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 3845f94593d..5dfcf1d8e35 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -2,6 +2,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import type { components } from "@/lib/http/schema"; import useCan from "@/app/(dashboard)/hooks/useCan"; import { organizationKeys, useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; +import { invalidateTeamQueries } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useQueryClient } from "@tanstack/react-query"; import UserSearchModal from "@/components/common_components/user_search_modal"; import { @@ -915,7 +916,8 @@ const TeamInfoView: React.FC = ({ const persistTeamUpdate = async (token: string, updateData: Record) => { await teamUpdateCall(token, updateData); - queryClient.invalidateQueries({ queryKey: organizationKeys.all }); + void queryClient.invalidateQueries({ queryKey: organizationKeys.all }); + void invalidateTeamQueries(queryClient); setIsEditing(false); }; diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index 522af5a85ad..09042dea930 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -7,10 +7,11 @@ vi.mock("@/lib/toast", () => ({ })); // ---- Hoisted shared mocks (safe to use inside vi.mock factories) ---- -const { keyUpdateCallMock, keyDeleteCallMock, mockUseAuthorized } = vi.hoisted(() => { +const { keyUpdateCallMock, keyDeleteCallMock, invalidateQueriesMock, mockUseAuthorized } = vi.hoisted(() => { return { keyUpdateCallMock: vi.fn().mockResolvedValue({}), keyDeleteCallMock: vi.fn().mockResolvedValue({}), + invalidateQueriesMock: vi.fn().mockResolvedValue(undefined), mockUseAuthorized: vi.fn(), }; }); @@ -170,7 +171,7 @@ vi.mock("@tanstack/react-query", async (importOriginal) => { const actual = await importOriginal(); return { ...actual, - useQueryClient: () => ({ invalidateQueries: vi.fn() }), + useQueryClient: () => ({ invalidateQueries: invalidateQueriesMock }), }; }); @@ -366,6 +367,24 @@ describe("KeyInfoView handleKeyUpdate mcp_toolsets", () => { }); }); +describe("KeyInfoView handleKeyUpdate cache sync", () => { + it("should invalidate every cached key query so the list and detail views re-read the saved key", async () => { + keyUpdateCallMock.mockResolvedValueOnce({ + object_permission: { mcp_servers: ["srv-1"], mcp_tool_permissions: { "srv-1": ["read_wiki"] } }, + }); + renderView(true); + + fireEvent.click(screen.getByText("Settings")); + fireEvent.click(screen.getByText("Edit Settings")); + (globalThis as any).__TEST_FORM_VALUES = { token: "tok_123", metadata: {} }; + + fireEvent.click(screen.getByText("Mock Submit")); + + await waitFor(() => expect(toast.success).toHaveBeenCalledWith("Key updated successfully")); + expect(invalidateQueriesMock).toHaveBeenCalledWith({ queryKey: ["keys"] }); + }); +}); + describe("KeyInfoView handleKeyUpdate skills", () => { it("should forward the skills the edit form supplies into object_permission and drop the form key", async () => { renderView(true); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 7eb09926caf..cfb5e9fa1f8 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -381,8 +381,8 @@ export default function KeyInfoView({ const newKeyValues = await keyUpdateCall(accessToken, formValues); - // Update local state setCurrentKeyData((prevData) => (prevData ? { ...prevData, ...newKeyValues } : undefined)); + void queryClient.invalidateQueries({ queryKey: keyKeys.all }); if (onKeyDataUpdate) { onKeyDataUpdate(newKeyValues);