diff --git a/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts b/tests/e2e/ui/tests/proxy-admin/keyBudgetWindow.spec.ts index 4e4d0a395c3..399d0947dae 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); + } }, }); @@ -34,15 +37,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/templates/key_edit_view.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx index 9e3e19f1e1b..37e66d78953 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 @@ -1004,6 +1004,73 @@ 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))); + }); + + 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 }))); + }); + it("should persist a canonical budget_duration value, not a word-form the backend cannot parse", async () => { const onSubmitMock = vi.fn().mockResolvedValue(undefined); renderWithProviders( @@ -1545,7 +1612,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(); @@ -1575,17 +1642,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 ID" })).toBeDisabled(); + view.rerender(<>); view.rerender(renderEditor({ ...key, project_id: null })); expect(screen.getByRole("combobox", { name: "Team ID" })).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 1342b97d1a0..577b17a5dd5 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -88,7 +88,11 @@ interface KeyEditViewProps { premiumUser?: boolean; } -export function KeyEditView({ +export function KeyEditView(props: KeyEditViewProps) { + return ; +} + +function KeyEditSession({ keyData, onCancel, onSubmit, @@ -104,7 +108,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,21 +193,17 @@ 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); + form.setValue("rotation_interval", rotationInterval, { shouldDirty: true }); } }, [rotationInterval, form]); @@ -327,7 +329,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) => { @@ -335,16 +337,16 @@ export function KeyEditView({ const selectedTeam = teams?.find((t) => t.team_id === teamId) || null; if (selectedTeam?.organization_id) { setSelectedOrganizationId(selectedTeam.organization_id); - form.setValue("organization_id", selectedTeam.organization_id); + form.setValue("organization_id", selectedTeam.organization_id, { shouldDirty: true }); } else if (!teamId) { setSelectedOrganizationId(null); - form.setValue("organization_id", null); + form.setValue("organization_id", null, { shouldDirty: true }); } }; const handleDisabledCallbacksChange = (internalValues: string[]) => { setDisabledCallbacks(mapInternalToDisplayNames(internalValues)); - form.setValue("disabled_callbacks", internalValues); + form.setValue("disabled_callbacks", internalValues, { shouldDirty: true }); }; const modelOptions = [ @@ -409,14 +411,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; } }} @@ -766,7 +768,7 @@ export function KeyEditView({ selectedAccessGroups={mcpSelection?.accessGroups || []} selectedToolsets={mcpSelection?.toolsets || []} toolPermissions={(mcpToolPermissions as Record | undefined) || {}} - onChange={(toolPerms) => form.setValue("mcp_tool_permissions", toolPerms)} + onChange={(toolPerms) => form.setValue("mcp_tool_permissions", toolPerms, { shouldDirty: true })} /> @@ -828,7 +830,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 }) + } /> )}