From c96cdcd2f30a791ba6e5bd313157df14d2f14eec Mon Sep 17 00:00:00 2001 From: songjianglin <270247441+2017955725@users.noreply.github.com> Date: Fri, 9 Oct 2026 01:45:24 +0800 Subject: [PATCH] fix(ui): support editing virtual key model aliases --- .../components/key_team_helpers/key_list.tsx | 2 +- .../templates/KeyEditViewControls.tsx | 40 +++++++++ .../templates/keyEditFormValues.test.ts | 30 +++++++ .../components/templates/keyEditFormValues.ts | 7 +- .../key_edit_view.integration.test.tsx | 82 ++++++++++++++++++- .../components/templates/key_edit_view.tsx | 20 ++--- 6 files changed, 166 insertions(+), 15 deletions(-) diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index 8d39a0fd22b..11379716185 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -41,7 +41,7 @@ export interface KeyResponse { max_budget: number; expires: string; models: string[]; - aliases: Record; + aliases: Record; config: Record; user_id: string; team_id: string | null; diff --git a/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx b/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx index 99220af7c6f..bb3612a2ed7 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx @@ -5,9 +5,12 @@ import { Textarea } from "@/components/ui/textarea"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import { CircleHelp } from "lucide-react"; import { FormField } from "@/components/shared/form/FormField"; +import { Field, FieldLabel } from "@/components/ui/field"; import { toast } from "@/lib/toast"; import AgentSelector from "../agent_management/AgentSelector"; import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem"; +import ModelAliasManager from "../common_components/ModelAliasManager"; +import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/BudgetWindowsEditor"; import NumericalInput from "../shared/numerical_input"; import SkillSelector from "../skills/SkillSelector"; import { moveTagsOutOfMetadataJson } from "./keyEditFieldNormalizers"; @@ -29,6 +32,43 @@ const KEY_TYPE_OPTIONS = [ { value: "management", label: "Management", hint: "Can call only management routes (user/team/key management)" }, ]; +export const KeyModelAliasField = ({ + control, + accessToken, +}: { + control: Control; + accessToken: string; +}) => ( + + {({ value, onChange }) => ( + + )} + +); + +export const KeyBudgetWindowsField = ({ + value, + onChange, +}: { + value: BudgetWindowEntry[]; + onChange: (value: BudgetWindowEntry[]) => void; +}) => ( + + + {labelWithHint( + "Budget Windows", + "Set multiple independent budget windows (e.g., hourly $10 AND monthly $200). Each window tracks spend separately and resets on its own schedule.", + )} + + + +); + export const KeyTypeSelect = ({ id, value, diff --git a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.test.ts b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.test.ts index f12088bea19..1a52d7177b5 100644 --- a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.test.ts +++ b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.test.ts @@ -4,6 +4,36 @@ import { keyEditFormSchema, toKeyEditFormValues, toSubmittedValues } from "./key const parse = (values: Record) => keyEditFormSchema.safeParse(values); +describe("key model aliases", () => { + it("loads key aliases independently of team aliases", () => { + const keyData = { + aliases: { chat: "chat-dev" }, + team_model_aliases: { chat: "team-chat" }, + } as unknown as KeyResponse; + expect(toKeyEditFormValues(keyData).aliases).toEqual(keyData.aliases); + }); + + it("preserves aliases through schema validation", () => { + expect(keyEditFormSchema.parse({ aliases: { chat: "chat-dev" } })).toEqual({ aliases: { chat: "chat-dev" } }); + }); + + it("rejects alias targets that are not model names", () => { + expect(parse({ aliases: { chat: 123 } }).success).toBe(false); + }); + + it("submits a changed mapping including an explicitly emptied one", () => { + const gates = { canViewPolicies: false, canViewPrompts: false, aliasesChanged: true }; + expect(toSubmittedValues({ aliases: { chat: "chat-prod" } }, gates).aliases).toEqual({ chat: "chat-prod" }); + expect(toSubmittedValues({ aliases: {} }, gates).aliases).toEqual({}); + }); + + it("omits aliases when the field was not changed", () => { + expect( + toSubmittedValues({ aliases: { chat: "chat-dev" } }, { canViewPolicies: false, canViewPrompts: false }), + ).not.toHaveProperty("aliases"); + }); +}); + describe("tpd_limit round trip", () => { const keyData = { token: "tok", models: [], rpm_limit: 5, tpd_limit: 250000 } as unknown as KeyResponse; diff --git a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts index 0ae60bbcc81..f62267b3e20 100644 --- a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts +++ b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts @@ -20,6 +20,7 @@ export interface AgentsAndGroups { export interface KeyEditFormValues { key_alias?: string; models?: string[]; + aliases?: Record; allowed_routes?: string; max_budget?: number | string | null; soft_budget?: number | string | null; @@ -68,6 +69,7 @@ const readMetadata = (keyData: KeyResponse, key: string): unknown => export const toKeyEditFormValues = (keyData: KeyResponse): KeyEditFormValues => ({ key_alias: keyData.key_alias, models: keyData.models, + aliases: keyData.aliases ?? {}, allowed_routes: Array.isArray(keyData.allowed_routes) && keyData.allowed_routes.length > 0 ? keyData.allowed_routes.join(", ") : "", max_budget: keyData.max_budget, @@ -124,6 +126,7 @@ export const toKeyEditFormValues = (keyData: KeyResponse): KeyEditFormValues => export const keyEditFormSchema = z.object({ key_alias: z.custom().optional(), models: z.custom().optional(), + aliases: z.record(z.string(), z.string()).optional(), allowed_routes: z.custom().optional(), max_budget: z.custom().optional(), soft_budget: z.custom().optional(), @@ -173,14 +176,16 @@ export const keyEditFormSchema = z.object({ export interface MountedFieldGates { canViewPolicies: boolean; canViewPrompts: boolean; + aliasesChanged?: boolean; } export const toSubmittedValues = ( values: KeyEditFormValues, - { canViewPolicies, canViewPrompts }: MountedFieldGates, + { canViewPolicies, canViewPrompts, aliasesChanged }: MountedFieldGates, ): Record => ({ key_alias: values.key_alias, models: values.models, + ...(aliasesChanged ? { aliases: values.aliases } : {}), allowed_routes: values.allowed_routes, max_budget: values.max_budget, soft_budget: values.soft_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..216f90b37a2 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 @@ -1,4 +1,4 @@ -import { fireEvent, screen, waitFor } from "@testing-library/react"; +import { fireEvent, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { chooseSelectOption, renderWithProviders, testQueryClient } from "../../../tests/test-utils"; @@ -30,6 +30,9 @@ vi.mock("../networking", async () => { modelAvailableCall: vi.fn().mockResolvedValue({ data: [{ id: "gpt-4" }, { id: "gpt-3.5-turbo" }], }), + modelHubCall: vi.fn().mockResolvedValue({ + data: [{ model_group: "chat-dev" }, { model_group: "chat-prod" }], + }), tagListCall: vi.fn().mockResolvedValue({ tag1: { name: "tag1", description: "Test tag 1" }, tag2: { name: "tag2", description: "Test tag 2" }, @@ -2156,6 +2159,83 @@ describe("KeyEditView", () => { expect(onSubmitMock.mock.calls[0][0]).toStrictEqual(UNTOUCHED_SAVE_PAYLOAD); }); + it("loads key model aliases and submits an edited target with its dirty field", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderForPayload(onSubmit, { + ...MOCK_KEY_DATA, + aliases: { chat: "chat-dev" }, + team_model_aliases: { chat: "team-chat" }, + }); + + await userEvent.click(await screen.findByRole("button", { name: "Edit chat" })); + const target = within(screen.getByRole("row", { name: /Save Cancel/ })).getByRole("combobox"); + await userEvent.clear(target); + await userEvent.type(target, "chat-prod"); + await userEvent.click(await screen.findByRole("option", { name: "chat-prod" })); + await userEvent.click(screen.getByRole("button", { name: "Save", exact: true })); + + expect(onSubmit).not.toHaveBeenCalled(); + expect(await screen.findByText("chat-prod")).toBeInTheDocument(); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmit).toHaveBeenCalledWith( + expect.objectContaining({ aliases: { chat: "chat-prod" } }), + expect.arrayContaining(["aliases"]), + ); + }); + }); + + it("submits an empty alias map when the last key alias is deleted", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderForPayload(onSubmit, { ...MOCK_KEY_DATA, aliases: { chat: "chat-dev" } }); + + await userEvent.click(await screen.findByRole("button", { name: "Delete chat" })); + expect(onSubmit).not.toHaveBeenCalled(); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmit).toHaveBeenCalledWith( + expect.objectContaining({ aliases: {} }), + expect.arrayContaining(["aliases"]), + ); + }); + }); + + it("adds a model alias to a key that has none", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderForPayload(onSubmit); + + fireEvent.change(await screen.findByLabelText("Alias Name"), { target: { value: "chat" } }); + await userEvent.type(screen.getByPlaceholderText("Select target model"), "chat-dev"); + await userEvent.click(await screen.findByRole("option", { name: "chat-dev" })); + await userEvent.click(screen.getByRole("button", { name: "Add Alias" })); + expect(onSubmit).not.toHaveBeenCalled(); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmit).toHaveBeenCalledWith( + expect.objectContaining({ aliases: { chat: "chat-dev" } }), + expect.arrayContaining(["aliases"]), + ); + }); + }); + + it("keeps stored aliases out of an unrelated key edit", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderForPayload(onSubmit, { ...MOCK_KEY_DATA, aliases: { chat: "chat-dev" } }); + + expect(await screen.findByText("chat-dev")).toBeInTheDocument(); + fireEvent.change(screen.getByLabelText("Key Alias"), { target: { value: "renamed-key" } }); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmit).toHaveBeenCalledTimes(1); + }); + expect(onSubmit.mock.calls[0][0]).not.toHaveProperty("aliases"); + expect(onSubmit.mock.calls[0][1]).not.toContain("aliases"); + }); + it("drops the policy and prompt keys entirely for a role that cannot see those fields", async () => { const onSubmitMock = vi.fn().mockResolvedValue(undefined); renderWithProviders( 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..65a93b7f295 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -33,7 +33,9 @@ import { import { KeyAgentAndSkillFields, KeyBudgetNumberField, + KeyBudgetWindowsField, KeyMetadataField, + KeyModelAliasField, KeyRateLimitFields, KeyTypeSelect, labelWithHint, @@ -55,7 +57,7 @@ import { } from "../key_team_helpers/endUserBudgetPayload"; import { ModelMaxBudgetField } from "../key_team_helpers/ModelMaxBudgetEditor"; import { useModelMaxBudgetField } from "../key_team_helpers/useModelMaxBudgetField"; -import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/BudgetWindowsEditor"; +import { BudgetWindowEntry } from "../key_team_helpers/BudgetWindowsEditor"; import { TagRateLimitEditor, TagRateLimitEntry, @@ -140,6 +142,7 @@ export function KeyEditView({ const enableProjectsUI = Boolean(uiSettingsData?.values?.enable_projects_ui); const hasProject = Boolean(keyData.project_id); const detachProject = hasProject && form.watch("project_id") === null; + const aliasesChanged = Boolean(form.formState.dirtyFields.aliases); const canDetachProject = canDetachKeyProject(team, organizations, userID, userRole); const allowedRoutesValue = form.watch("allowed_routes"); @@ -235,7 +238,6 @@ export function KeyEditView({ .filter((route: string) => route.length > 0); } } - // If it's already an array (shouldn't happen, but handle it), keep as is // Backend rejects non-empty allowed_routes from non-admins, so re-sending // an unchanged value 403s a team admin. Set compare tolerates reorder. @@ -366,7 +368,7 @@ export function KeyEditView({ onSubmit={(event) => { moveMetadataTagsToTagsField(form); return form.handleSubmit((values) => - handleSubmit(toSubmittedValues(values, { canViewPolicies, canViewPrompts })), + handleSubmit(toSubmittedValues(values, { canViewPolicies, canViewPrompts, aliasesChanged })), )(event); }} > @@ -401,6 +403,8 @@ export function KeyEditView({ )} + + Key Type - - - {labelWithHint( - "Budget Windows", - "Set multiple independent budget windows (e.g., hourly $10 AND monthly $200). Each window tracks spend separately and resets on its own schedule.", - )} - - - +