From 37b3a697c7f051eb00b0a57914159adc4ec3d5d8 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 16 Oct 2026 17:20:46 -0700 Subject: [PATCH] fix(ui): preserve key settings drafts during background refreshes (#41350) * fix(ui): preserve key settings drafts during background refreshes * fix(ui): refresh untouched key settings while preserving edits * fix(ui): keep team-change MCP clears across key refreshes clearMcpGrants now marks the MCP fields dirty, so a background refresh with keepDirtyValues no longer restores the grants a team change just cleared. Also adapts the draft tests to the onSubmit(values, dirtyFields) signature on main and asserts a refresh does not mark untouched fields dirty, since the team admin member-key path rejects saves with non-budget dirty fields. KeyEditView becomes a one-line arrow to stay under the max-lines limit. * fix(ui): keep MCP tool defaults out of a key's edited fields MCPToolPermissions writes non-delete tool defaults for unrestricted servers on mount. Marking that write dirty made a team admin's budget-only save on a member key get rejected for editing mcp_tool_permissions. onChange now says whether a write is a default or a user edit, and the key editor dirties only edits. The edit session is now keyed by token or token_id, matching the form's own fallback, so switching between keys that only carry token_id starts fresh. * fix(ui): keep a new MCP server's default tool limits across refreshes Default tool writes for a server the user just added in this edit session are part of that edit, so they are marked dirty and survive a background refresh. Defaults written for an already saved server stay clean, so a team admin's budget-only save on a member key is still accepted. --- .../tests/proxy-admin/keyBudgetWindow.spec.ts | 50 +++- .../MCPToolPermissions.test.tsx | 53 +++-- .../MCPToolPermissions.tsx | 11 +- .../components/templates/keyEditFormValues.ts | 14 +- .../key_edit_view.integration.test.tsx | 222 +++++++++++++++++- .../components/templates/key_edit_view.tsx | 43 ++-- 6 files changed, 335 insertions(+), 58 deletions(-) diff --git a/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts b/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts index 50850a5cd58..d0b1bb9ecfb 100644 --- a/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts +++ b/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts @@ -18,8 +18,11 @@ const test = base.extend<{ scopedKey: ScopedKey }>({ team_id: E2E_TEAM_CRUD_ID, models: [CHAT_MODEL_A], }); - await use({ alias, token: created.token }); - await deleteVirtualKey(page.request, created.token); + try { + await use({ alias, token: created.token }); + } finally { + await deleteVirtualKey(page.request, created.token); + } }, }); @@ -36,15 +39,44 @@ test.describe("Proxy Admin - Key budget window", () => { await navigateToPage(page, Page.ApiKeys); await dismissFeedbackPopup(page); - await openKeyDetail(page, alias); - - await page.getByRole("tab", { name: "Settings" }).click(); - await page.getByRole("button", { name: "Edit Settings" }).click(); - - await page.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("12.5"); + // The initial, unfiltered list must already contain this newly created key. + // useKeys keeps that list visible while the filtered search is pending. + await expect(page.getByRole("row").filter({ hasText: alias })).toBeVisible({ timeout: 15_000 }); + // Deliver the real search result after editing starts: replacing the list's + // key object must not reset the open form and erase its unsaved budget. + const searchReceived = Promise.withResolvers(); + const releaseSearch = Promise.withResolvers(); + const isAliasSearch = (url: string) => { + const parsed = new URL(url); + return parsed.pathname.endsWith("/key/list") && parsed.searchParams.get("search") === alias; + }; + await page.route("**/key/list?**", async (route) => { + if (!isAliasSearch(route.request().url())) return route.continue(); + const response = await route.fetch(); + searchReceived.resolve(); + await releaseSearch.promise; + await route.fulfill({ response }); + }); + try { + await openKeyDetail(page, alias); + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + await page.getByRole("spinbutton", { name: "Max Budget (USD)" }).fill("12.5"); + await searchReceived.promise; + const refreshed = page.waitForResponse((response) => isAliasSearch(response.url())); + releaseSearch.resolve(); + await refreshed; + } finally { + releaseSearch.resolve(); + await page.unrouteAll({ behavior: "wait" }); + } await page.getByLabel("Reset Budget", { exact: true }).click(); await page.getByRole("option", { name: "monthly", exact: true }).click(); - await page.getByRole("button", { name: "Save Changes" }).click(); + const saved = await captureRequestBody(page, { method: "POST", urlIncludes: "/key/update" }, async () => { + await page.getByRole("button", { name: "Save Changes" }).click(); + }); + expect(saved.max_budget, "a delayed search response must not erase the budget draft").toBe("12.5"); + expect(saved.budget_duration).toBe("30d"); await expect .poll(async () => (await readKeyInfo(page.request, token)).max_budget, { diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx index 3757720c699..bad4356aafb 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx @@ -79,7 +79,7 @@ describe("MCPToolPermissions", () => { const expectedToolPermissions = { [mockServerId]: ["read_wiki_contents", "ask_question"], }; - expect(mockOnChange).toHaveBeenCalledWith(expectedToolPermissions); + expect(mockOnChange).toHaveBeenCalledWith(expectedToolPermissions, "edit"); // Verify API calls // Note: useMCPServers uses useAuthorized() internally, which returns "123" from global mock @@ -135,9 +135,12 @@ describe("MCPToolPermissions", () => { // Selecting every displayed tool writes the wildcard, which also covers tools the // server adds later. - expect(mockOnChange).toHaveBeenCalledWith({ - [mockServerId]: ["*"], - }); + expect(mockOnChange).toHaveBeenCalledWith( + { + [mockServerId]: ["*"], + }, + "edit", + ); }); it("should deselect all tools when Deselect All button is clicked", async () => { @@ -186,9 +189,12 @@ describe("MCPToolPermissions", () => { await userEvent.click(deselectAllButton); // Verify onChange was called with no tools selected - expect(mockOnChange).toHaveBeenCalledWith({ - [mockServerId]: [], - }); + expect(mockOnChange).toHaveBeenCalledWith( + { + [mockServerId]: [], + }, + "edit", + ); }); describe("wildcard all-tools grant", () => { @@ -238,7 +244,7 @@ describe("MCPToolPermissions", () => { expect(await screen.findByText("read_wiki_structure")).toBeInTheDocument(); await userEvent.click(screen.getByRole("button", { name: "Select All" })); - expect(mockOnChange).toHaveBeenCalledWith({ [wildcardServerId]: ["*"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [wildcardServerId]: ["*"] }, "edit"); }); it("converts back to an enumerated list when one tool is unchecked from a wildcard grant", async () => { @@ -256,9 +262,12 @@ describe("MCPToolPermissions", () => { await userEvent.click(screen.getByText("Flat List")); await userEvent.click(screen.getByRole("checkbox", { name: "ask_question" })); - expect(mockOnChange).toHaveBeenCalledWith({ - [wildcardServerId]: ["read_wiki_structure", "read_wiki_contents"], - }); + expect(mockOnChange).toHaveBeenCalledWith( + { + [wildcardServerId]: ["read_wiki_structure", "read_wiki_contents"], + }, + "edit", + ); }); }); @@ -320,7 +329,7 @@ describe("MCPToolPermissions", () => { expect(deleteIssue).toBeChecked(); await userEvent.click(listIssues); - expect(mockOnChange).toHaveBeenCalledWith({ [groupServer.server_id]: ["delete_issue"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [groupServer.server_id]: ["delete_issue"] }, "edit"); }); it("marks an access-group server as inherited and leaves a directly selected one unmarked", async () => { @@ -469,7 +478,7 @@ describe("MCPToolPermissions", () => { // First checkbox is the header toggle of the group holding list_issues. await userEvent.click(screen.getAllByRole("checkbox")[0]); - expect(mockOnChange).toHaveBeenCalledWith({ [toolsetServer.server_id]: ["list_issues"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [toolsetServer.server_id]: ["list_issues"] }, "edit"); }); // Copying the toolset's tools into the entry would outlive the toolset, so a write keeps only @@ -502,7 +511,7 @@ describe("MCPToolPermissions", () => { // A toolset-sourced server never writes the wildcard: that would create a standing direct // grant outliving the toolset. The write keeps only the tools this level grants itself. - expect(mockOnChange).toHaveBeenCalledWith({ [toolsetServer.server_id]: ["delete_issue"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [toolsetServer.server_id]: ["delete_issue"] }, "edit"); }); // The default narrows an unrestricted server; against a toolset-restricted one it would widen @@ -597,7 +606,7 @@ describe("MCPToolPermissions", () => { expect(await screen.findByText("Payments")).toBeInTheDocument(); expect(screen.queryByText("srv-collide")).not.toBeInTheDocument(); await waitFor(() => { - expect(mockOnChange).toHaveBeenCalledWith({ "srv-collide": ["list_issues"] }); + expect(mockOnChange).toHaveBeenCalledWith({ "srv-collide": ["list_issues"] }, "default"); }); expect(mockOnChange.mock.calls.every(([written]) => !Object.hasOwn(written, "srv-twin"))).toBe(true); expect(networking.listMCPTools).not.toHaveBeenCalledWith(mockAccessToken, "srv-twin"); @@ -640,7 +649,7 @@ describe("MCPToolPermissions", () => { ); await waitFor(() => { - expect(mockOnChange).toHaveBeenCalledWith({ [directServer.server_id]: ["list_issues"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [directServer.server_id]: ["list_issues"] }, "default"); }); }); @@ -804,7 +813,7 @@ describe("MCPToolPermissions", () => { expect(await screen.findByText("list_issues")).toBeInTheDocument(); await userEvent.click(screen.getByRole("button", { name: "Deselect All" })); - expect(mockOnChange).toHaveBeenCalledWith({ github_mcp: [] }); + expect(mockOnChange).toHaveBeenCalledWith({ github_mcp: [] }, "edit"); }); }); @@ -984,7 +993,7 @@ describe("MCPToolPermissions", () => { await userEvent.click(readGroupToggle); - expect(mockOnChange).toHaveBeenCalledWith({ [mockServerId]: ["delete_document"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [mockServerId]: ["delete_document"] }, "edit"); }); it("adds the rest of a partially-allowed risk group when its mixed toggle is clicked", async () => { @@ -996,7 +1005,7 @@ describe("MCPToolPermissions", () => { await userEvent.click(readGroupToggle); - expect(mockOnChange).toHaveBeenCalledWith({ [mockServerId]: ["list_documents", "get_document"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [mockServerId]: ["list_documents", "get_document"] }, "edit"); }); it("toggles a single tool exactly once when its checkbox is clicked inside the clickable row", async () => { @@ -1006,7 +1015,7 @@ describe("MCPToolPermissions", () => { await userEvent.click(await screen.findByRole("checkbox", { name: "delete_document" })); expect(mockOnChange).toHaveBeenCalledTimes(1); - expect(mockOnChange).toHaveBeenCalledWith({ [mockServerId]: ["list_documents", "get_document"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [mockServerId]: ["list_documents", "get_document"] }, "edit"); }); it("toggles a single tool when the row around its checkbox is clicked", async () => { @@ -1016,7 +1025,7 @@ describe("MCPToolPermissions", () => { await userEvent.click(await screen.findByText("Destroy a document")); expect(mockOnChange).toHaveBeenCalledTimes(1); - expect(mockOnChange).toHaveBeenCalledWith({ [mockServerId]: ["list_documents", "get_document"] }); + expect(mockOnChange).toHaveBeenCalledWith({ [mockServerId]: ["list_documents", "get_document"] }, "edit"); }); it("re-renders each checkbox from the permissions it emitted", async () => { @@ -1103,7 +1112,7 @@ describe("MCPToolPermissions", () => { />, ); - await waitFor(() => expect(onChange).toHaveBeenCalledWith({ "srv-files": ["read_file"] })); + await waitFor(() => expect(onChange).toHaveBeenCalledWith({ "srv-files": ["read_file"] }, "default")); }); }); }); diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx index 8067936e715..b12b31601b8 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx @@ -20,13 +20,15 @@ import { resolveEffectiveMcpServers, } from "./effectiveMcpServers"; +export type ToolPermissionsChangeOrigin = "default" | "edit"; + interface MCPToolPermissionsProps { accessToken: string; selectedServers: readonly string[]; selectedAccessGroups?: readonly string[]; selectedToolsets?: readonly string[]; toolPermissions: Record; - onChange: (toolPermissions: Record) => void; + onChange: (toolPermissions: Record, origin: ToolPermissionsChangeOrigin) => void; disabled?: boolean; preserveUnrestrictedServers?: readonly string[]; } @@ -141,7 +143,10 @@ const MCPToolPermissions: React.FC = ({ const nonDeleteTools = fetchedTools .filter((t) => classifyToolOp(t.name, t.description || "") !== "delete") .map((t) => t.name); - onChange(applyToolPermissionWrite({ toolPermissions: latestPermissions, entry, allowed: nonDeleteTools })); + onChange( + applyToolPermissionWrite({ toolPermissions: latestPermissions, entry, allowed: nonDeleteTools }), + "default", + ); } } } catch (err) { @@ -175,7 +180,7 @@ const MCPToolPermissions: React.FC = ({ entry.source.kind !== "toolset" && names.length > 0 && names.every((n) => allowed.includes(n)) ? [MCP_ALL_TOOLS_WILDCARD] : allowed; - onChange(applyToolPermissionWrite({ toolPermissions, entry, allowed: next })); + onChange(applyToolPermissionWrite({ toolPermissions, entry, allowed: next }), "edit"); }; const handleSelectAll = (entry: EffectiveMcpServer) => { diff --git a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts index 35f09011ab9..ffbdd7ae086 100644 --- a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts +++ b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts @@ -1,7 +1,8 @@ import { z } from "zod"; -import { UseFormSetValue } from "react-hook-form"; +import { UseFormReturn, UseFormSetValue } from "react-hook-form"; import { KeyResponse } from "../key_team_helpers/key_list"; +import { ToolPermissionsChangeOrigin } from "../mcp_server_management/MCPToolPermissions"; import { extractLoggingSettings, formatMetadataForDisplay, stripTagsFromMetadata } from "../key_info_utils"; import { mapInternalToDisplayNames } from "../callback_info_helpers"; import { estimateChecks, estimateFields } from "./estimatedOutputTokens"; @@ -222,6 +223,13 @@ export const toSubmittedValues = ( }); export const clearMcpGrants = (setValue: UseFormSetValue) => { - setValue("mcp_servers_and_groups", { servers: [], accessGroups: [], toolsets: [] }); - setValue("mcp_tool_permissions", {}); + setValue("mcp_servers_and_groups", { servers: [], accessGroups: [], toolsets: [] }, { shouldDirty: true }); + setValue("mcp_tool_permissions", {}, { shouldDirty: true }); }; + +export const writeMcpToolPermissions = + (form: Pick, "getFieldState" | "setValue">) => + (toolPermissions: Record, origin: ToolPermissionsChangeOrigin) => + form.setValue("mcp_tool_permissions", toolPermissions, { + shouldDirty: origin === "edit" || form.getFieldState("mcp_servers_and_groups").isDirty, + }); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx index ead17fa0c95..4dd98b31781 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx @@ -10,6 +10,7 @@ import { getPoliciesList, getUiSettings, getPromptsList, + listMCPTools, modelAvailableCall, vectorStoreListCall, } from "../networking"; @@ -1011,6 +1012,179 @@ describe("KeyEditView", () => { }); }); + it("should preserve a budget draft when the same key refreshes before saving", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + const props = { + keyData: MOCK_KEY_DATA, + onCancel: vi.fn(), + onSubmit, + accessToken: "test-token", + userID: "test-user", + userRole: "admin", + }; + const view = renderWithProviders(); + + fireEvent.change(await screen.findByRole("spinbutton", { name: "Max Budget (USD)" }), { + target: { value: "12.5" }, + }); + view.rerender( + , + ); + await chooseSelectOption(userEvent, screen.getByLabelText("Reset Budget"), "monthly"); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + const expected = { + max_budget: "12.5", + budget_duration: "30d", + rpm_limit: 75, + models: ["gpt-4"], + metadata: JSON.stringify({ note: "new" }, null, 2), + vector_stores: ["store-new"], + }; + await waitFor(() => expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining(expected), expect.any(Array))); + const [, dirtyFields] = onSubmit.mock.calls[0]; + expect(dirtyFields).toContain("max_budget"); + expect(dirtyFields).not.toContain("rpm_limit"); + }); + + const mcpServer = (server_id: string) => ({ + server_id, + server_name: server_id, + url: `https://${server_id}.example`, + transport: "http", + auth_type: "none", + created_at: "", + created_by: "", + updated_at: "", + updated_by: "", + }); + const mcpTools = (names: string[]) => ({ + tools: names.map((name) => ({ name })), + error: null, + message: null, + stack_trace: null, + }); + + it("should not mark MCP tool defaults as an edit when saving only a budget", async () => { + vi.mocked(fetchMCPServers).mockResolvedValue([mcpServer("srv-a")]); + vi.mocked(listMCPTools).mockResolvedValue(mcpTools(["read_page", "delete_page"])); + testQueryClient.clear(); + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderWithProviders( + , + ); + await waitFor(() => expect(listMCPTools).toHaveBeenCalledWith("test-token", "srv-a")); + + fireEvent.change(await screen.findByRole("spinbutton", { name: "Max Budget (USD)" }), { + target: { value: "12.5" }, + }); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(onSubmit).toHaveBeenCalled()); + const [values, dirtyFields] = onSubmit.mock.calls[0]; + expect(values.mcp_tool_permissions).toHaveProperty("srv-a"); + expect(dirtyFields).toEqual(["max_budget"]); + vi.mocked(fetchMCPServers).mockResolvedValue([]); + vi.mocked(listMCPTools).mockResolvedValue(mcpTools([])); + }); + + it("should keep default tool limits for a newly added MCP server when the same key refreshes", async () => { + vi.mocked(fetchMCPServers).mockResolvedValue([mcpServer("mcp-1")]); + vi.mocked(listMCPTools).mockResolvedValue(mcpTools(["read-page", "delete-page"])); + testQueryClient.clear(); + const onSubmit = vi.fn().mockResolvedValue(undefined); + const props = { + keyData: MOCK_KEY_DATA, + onCancel: vi.fn(), + onSubmit, + accessToken: "test-token", + userID: "test-user", + userRole: "admin", + }; + const view = renderWithProviders(); + + await userEvent.click(await screen.findByTestId("mcp-server-selector")); + await waitFor(() => expect(listMCPTools).toHaveBeenCalledWith("test-token", "mcp-1")); + view.rerender(); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(onSubmit).toHaveBeenCalled()); + expect(onSubmit.mock.calls[0][0]).toMatchObject({ + rpm_limit: 75, + mcp_tool_permissions: { "mcp-1": ["read-page"] }, + }); + vi.mocked(fetchMCPServers).mockResolvedValue([]); + vi.mocked(listMCPTools).mockResolvedValue(mcpTools([])); + }); + + it("should start a fresh budget draft when switching between keys identified only by token_id", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + const props = { + keyData: { ...MOCK_KEY_DATA, token: "", token_id: "key-one" }, + onCancel: vi.fn(), + onSubmit, + accessToken: "test-token", + userID: "test-user", + userRole: "admin", + }; + const view = renderWithProviders(); + fireEvent.change(await screen.findByRole("spinbutton", { name: "Max Budget (USD)" }), { + target: { value: "12.5" }, + }); + view.rerender( + , + ); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => + expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ max_budget: 25 }), expect.any(Array)), + ); + }); + + it("should start a fresh budget draft when switching to another key", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + const props = { + keyData: MOCK_KEY_DATA, + onCancel: vi.fn(), + onSubmit, + accessToken: "test-token", + userID: "test-user", + userRole: "admin", + }; + const view = renderWithProviders(); + fireEvent.change(await screen.findByRole("spinbutton", { name: "Max Budget (USD)" }), { + target: { value: "12.5" }, + }); + view.rerender( + , + ); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => + expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ max_budget: 25 }), expect.any(Array)), + ); + }); + it("should persist a canonical budget_duration value, not a word-form the backend cannot parse", async () => { const onSubmitMock = vi.fn().mockResolvedValue(undefined); renderWithProviders( @@ -1615,6 +1789,46 @@ describe("KeyEditView", () => { }); }); + it("keeps the team change's cleared MCP grants when the same key refreshes before saving", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + const keyData = { + ...MOCK_KEY_DATA, + organization_id: "org-1", + team_id: "group-maple", + object_permission: { mcp_servers: ["srv-a"], mcp_tool_permissions: { "srv-a": ["read"] } }, + }; + const renderEditor = (data: KeyResponse) => ( + {}} + onSubmit={onSubmit} + accessToken="" + userID="" + userRole="Admin" + premiumUser={false} + /> + ); + const view = renderWithProviders(renderEditor(keyData)); + + await waitFor(() => { + expect(screen.getAllByRole("button", { name: "Clear" })).toHaveLength(2); + }); + await userEvent.click(screen.getAllByRole("button", { name: "Clear" })[1]); + await waitFor(() => { + expect(screen.getByTestId("mcp-server-selector")).toHaveAttribute("data-team-id", ""); + }); + view.rerender(renderEditor({ ...keyData, rpm_limit: 75 })); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmit).toHaveBeenCalled(); + }); + expect(onSubmit.mock.calls[0][0]).toMatchObject({ + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: {}, + }); + }); + it("clears the organization and its dependent team in the update payload", async () => { const onSubmit = vi.fn().mockResolvedValue(undefined); renderWithProviders( @@ -1649,7 +1863,7 @@ describe("KeyEditView", () => { }); }); - it("should save an explicit project detach while keeping parents locked until the saved key changes", async () => { + it("should save an explicit project detach while keeping parents locked until reopening the editor", async () => { vi.mocked(getUiSettings).mockResolvedValue({ values: { enable_projects_ui: true } }); const onSubmit = vi.fn().mockResolvedValue(undefined); const onCancel = vi.fn(); @@ -1679,17 +1893,23 @@ describe("KeyEditView", () => { await userEvent.click(screen.getByRole("button", { name: "Cancel" })); expect(onCancel).toHaveBeenCalledOnce(); expect(onSubmit).not.toHaveBeenCalled(); + // KeyInfoView unmounts the editor on cancel and after a successful save. + view.rerender(<>); view.rerender(renderEditor({ ...key })); await userEvent.click(await screen.findByRole("button", { name: "Detach from project" })); + view.rerender(renderEditor({ ...key, spend: 1 })); + expect(screen.getByRole("button", { name: "Keep project" })).toBeVisible(); await userEvent.click(screen.getByRole("button", { name: /save changes/i })); const expectedDetach = { project_id: null, organization_id: "org-1", team_id: "group-maple", models: key.models }; await waitFor(() => expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining(expectedDetach), expect.any(Array)), ); expect(screen.getByRole("combobox", { name: "Team" })).toBeDisabled(); + view.rerender(<>); view.rerender(renderEditor({ ...key, project_id: null })); expect(screen.getByRole("combobox", { name: "Team" })).toBeEnabled(); expect(screen.queryByRole("button", { name: "Detach from project" })).not.toBeInTheDocument(); + view.rerender(<>); view.rerender(renderEditor(key, "Internal User")); expect(screen.queryByRole("button", { name: "Detach from project" })).not.toBeInTheDocument(); view.rerender(renderEditor(key, "Org Admin")); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index 14724fcf6ef..44e9af0d316 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -46,6 +46,7 @@ import { toKeyEditFormValues, toSubmittedValues, clearMcpGrants, + writeMcpToolPermissions, } from "./keyEditFormValues"; import { BudgetFallbacksEditor } from "../key_team_helpers/BudgetFallbacksEditor"; import { END_USER_BUDGET_HINT, EndUserBudgetSelect } from "../key_team_helpers/EndUserBudgetSelect"; @@ -89,7 +90,11 @@ interface KeyEditViewProps { premiumUser?: boolean; } -export function KeyEditView({ +const sessionKey = (keyData: KeyResponse) => keyData.token || keyData.token_id; + +export const KeyEditView = (props: KeyEditViewProps) => ; + +function KeyEditSession({ keyData, onCancel, onSubmit, @@ -105,7 +110,9 @@ export function KeyEditView({ const canEditEstimates = userRole != null && isProxyAdminRole(userRole); const estimateTooltip = estimateTooltips(canEditEstimates); const form = useZodForm(keyEditFormSchema, { - defaultValues: toKeyEditFormValues(keyData), + // Refresh untouched fields without erasing the user's in-progress edits. + values: toKeyEditFormValues(keyData), + resetOptions: { keepDirtyValues: true }, }); const [promptsList, setPromptsList] = useState([]); const [tagsList, setTagsList] = useState>({}); @@ -187,22 +194,16 @@ export function KeyEditView({ // Sync disabled callbacks with form when component mounts useEffect(() => { - form.setValue("disabled_callbacks", disabledCallbacks); + form.setValue("disabled_callbacks", disabledCallbacks, { shouldDirty: true }); }, [form, disabledCallbacks]); - useEffect(() => { - form.reset(toKeyEditFormValues(keyData)); - }, [keyData, form]); - // Sync auto-rotation state with form values useEffect(() => { - form.setValue("auto_rotate", autoRotationEnabled); + form.setValue("auto_rotate", autoRotationEnabled, { shouldDirty: true }); }, [autoRotationEnabled, form]); useEffect(() => { - if (rotationInterval) { - form.setValue("rotation_interval", rotationInterval); - } + if (rotationInterval) form.setValue("rotation_interval", rotationInterval, { shouldDirty: true }); }, [rotationInterval, form]); // Fetch tags for selector @@ -327,7 +328,7 @@ export function KeyEditView({ const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | null) => { setField(orgId); setSelectedOrganizationId(orgId); - form.setValue("team_id", null); + form.setValue("team_id", null, { shouldDirty: true }); }; const handleTeamChange = (setField: (value: string | null) => void, teamId: string | null) => { @@ -336,12 +337,12 @@ export function KeyEditView({ const organizationId = teams?.find((t) => t.team_id === teamId)?.organization_id; if (!organizationId) return; setSelectedOrganizationId(organizationId); - form.setValue("organization_id", organizationId); + form.setValue("organization_id", organizationId, { shouldDirty: true }); }; const handleDisabledCallbacksChange = (internalValues: string[]) => { setDisabledCallbacks(mapInternalToDisplayNames(internalValues)); - form.setValue("disabled_callbacks", internalValues); + form.setValue("disabled_callbacks", internalValues, { shouldDirty: true }); }; const modelOptions = [ @@ -417,14 +418,14 @@ export function KeyEditView({ onChange={(value) => { switch (value) { case "default": - form.setValue("allowed_routes", ""); + form.setValue("allowed_routes", "", { shouldDirty: true }); break; case "llm_api": - form.setValue("allowed_routes", "llm_api_routes"); + form.setValue("allowed_routes", "llm_api_routes", { shouldDirty: true }); break; case "management": - form.setValue("allowed_routes", "management_routes"); - form.setValue("models", []); + form.setValue("allowed_routes", "management_routes", { shouldDirty: true }); + form.setValue("models", [], { shouldDirty: true }); break; } }} @@ -775,7 +776,7 @@ export function KeyEditView({ selectedAccessGroups={mcpSelection?.accessGroups || []} selectedToolsets={mcpSelection?.toolsets || []} toolPermissions={(mcpToolPermissions as Record | undefined) || {}} - onChange={(toolPerms) => form.setValue("mcp_tool_permissions", toolPerms)} + onChange={writeMcpToolPermissions(form)} /> @@ -827,7 +828,9 @@ export function KeyEditView({ canDetach={canDetachProject} pending={detachProject} disabled={isKeySaving} - onToggle={() => form.setValue("project_id", detachProject ? keyData.project_id : null)} + onToggle={() => + form.setValue("project_id", detachProject ? keyData.project_id : null, { shouldDirty: true }) + } /> )}