diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx index 5c3af3f6c3f..c890501bb2c 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx @@ -7,6 +7,7 @@ import { KeysResponse, useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useTeamMemberBudgets } from "@/app/(dashboard)/hooks/teams/useTeamMemberBudgets"; import { KeyResponse } from "../key_team_helpers/key_list"; import { Organization } from "../networking"; +import KeyInfoView from "../templates/key_info_view"; vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ useKeys: vi.fn(), @@ -442,6 +443,22 @@ describe("TeamVirtualKeysTable", () => { }); }); + it("should open a key without its own organization without the team's organization", async () => { + const keyWithoutOrganization = { + keys: [createMockKey({ token: "sk-no-org", key_alias: "no_org_key", organization_id: null })], + total_count: 1, + current_page: 1, + total_pages: 1, + }; + mockUseKeys.mockReturnValue(createKeysQueryResult(keyWithoutOrganization)); + renderWithProviders(); + + await userEvent.click(await screen.findByRole("button", { name: /sk-no-org/ })); + + await screen.findByText("Key Info View"); + expect(vi.mocked(KeyInfoView).mock.lastCall?.[0].keyData).toMatchObject({ organization_id: null }); + }); + describe("entity links out of the key rows", () => { const renderRow = async (key: KeyResponse, organization: Organization | null = null) => { mockUseKeys.mockReturnValue({ diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx index 8ea559990c0..2e7d2c001a1 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx @@ -90,15 +90,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi const { data: keys, isPending: isLoading, isFetching, refetch } = useKeys(pageIndex + 1, pageSize, keyListOptions); - const displayKeys = useMemo(() => { - const kList = keys?.keys || []; - const orgId = organization?.organization_id; - if (!orgId) return kList; - return kList.map((k: KeyResponse) => ({ - ...k, - organization_id: (k.organization_id ?? k.org_id) || orgId, - })); - }, [keys?.keys, organization?.organization_id]); + const displayKeys = useMemo(() => keys?.keys ?? [], [keys?.keys]); const teamMemberBudgetTeamIds = useMemo( () => (displayKeys.some((key) => key.max_budget == null && Boolean(key.user_id)) ? [teamId] : []), @@ -180,7 +172,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi }, { id: "organization_id", - accessorKey: "organization_id", + accessorFn: (key) => (key.organization_id ?? key.org_id) || organization?.organization_id || null, header: "Organization ID", size: 140, enableSorting: false, @@ -442,7 +434,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi }, }, ], - [expandedAccordions, teamMemberBudgets, teamId], + [expandedAccordions, teamMemberBudgets, teamId, organization?.organization_id], ); const handleSortingChange = useCallback((updaterOrValue: React.SetStateAction) => { 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 c82a4823dfd..9e48c4e3aba 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -1,6 +1,7 @@ import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { toast } from "@/lib/toast"; +import type { KeyResponse } from "../key_team_helpers/key_list"; vi.mock("@/lib/toast", () => ({ toast: { error: vi.fn(), success: vi.fn() }, @@ -655,3 +656,60 @@ describe("KeyInfoView handleKeyUpdate soft_budget", () => { expect(keyUpdateCallMock).not.toHaveBeenCalled(); }); }); + +describe("KeyInfoView handleKeyUpdate organization", () => { + const keyListRow = { ...baseKeyData, org_id: "org-1" } as KeyResponse; + + const adminAuth = { + accessToken: "access_abc", + userId: "user_1", + userRole: "Admin", + premiumUser: true, + token: "token_123", + userEmail: "test@example.com", + disabledPersonalKeyCreation: false, + showSSOBanner: false, + }; + + const renderKeyListRowSettings = () => { + mockUseAuthorized.mockReturnValue(adminAuth); + render( + {}} keyData={keyListRow} onKeyDataUpdate={() => {}} teams={[]} />, + ); + fireEvent.click(screen.getByText("Settings")); + }; + + const submitEditForm = (formValues: Record) => { + fireEvent.click(screen.getByText("Edit Settings")); + Object.assign(globalThis, { __TEST_FORM_VALUES: formValues }); + fireEvent.click(screen.getByText("Mock Submit")); + }; + + const shownOrganization = () => screen.getByText("Organization", { selector: "p" }).nextElementSibling; + + it("should show no organization after a save clears the one a key list row carried as org_id", async () => { + keyUpdateCallMock.mockResolvedValueOnce({ ...baseKeyData, organization_id: null }); + renderKeyListRowSettings(); + expect(shownOrganization()).toHaveTextContent("org-1"); + + submitEditForm({ token: "tok_123", organization_id: null }); + + await waitFor(() => { + expect(shownOrganization()).toHaveTextContent("Not Set"); + }); + expect(keyUpdateCallMock.mock.calls[0][1]).toHaveProperty("organization_id", null); + }); + + it("should leave organization_id out of a save that keeps the key's organization", async () => { + keyUpdateCallMock.mockResolvedValueOnce({ ...baseKeyData, organization_id: "org-1" }); + renderKeyListRowSettings(); + + submitEditForm({ token: "tok_123", organization_id: "org-1" }); + + await waitFor(() => { + expect(keyUpdateCallMock).toHaveBeenCalledTimes(1); + }); + expect(keyUpdateCallMock.mock.calls[0][1]).not.toHaveProperty("organization_id"); + expect(shownOrganization()).toHaveTextContent("org-1"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.ts b/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.ts index 1cfd1ae35f7..4ddc2c60bd3 100644 --- a/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.ts +++ b/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.ts @@ -9,6 +9,9 @@ const WORD_FORM_BUDGET_DURATIONS: Record = { export const canonicalBudgetDuration = (duration: string | null | undefined): string | null => duration ? WORD_FORM_BUDGET_DURATIONS[duration] ?? duration : null; +export const keyOrganizationId = (key: { organization_id?: string | null; org_id?: string | null }): string | null => + key.organization_id || key.org_id || null; + // Determine the key_type display value from allowed_routes export const keyTypeFromRoutes = (allowedRoutes: string[] | null | undefined): string => { if (!allowedRoutes || allowedRoutes.length === 0) return "default"; diff --git a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts index ffbdd7ae086..c23b9e76957 100644 --- a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts +++ b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts @@ -6,7 +6,7 @@ import { ToolPermissionsChangeOrigin } from "../mcp_server_management/MCPToolPer import { extractLoggingSettings, formatMetadataForDisplay, stripTagsFromMetadata } from "../key_info_utils"; import { mapInternalToDisplayNames } from "../callback_info_helpers"; import { estimateChecks, estimateFields } from "./estimatedOutputTokens"; -import { canonicalBudgetDuration } from "./keyEditFieldNormalizers"; +import { canonicalBudgetDuration, keyOrganizationId } from "./keyEditFieldNormalizers"; export interface McpServersAndGroups { servers: string[]; @@ -109,7 +109,7 @@ export const toKeyEditFormValues = (keyData: KeyResponse): KeyEditFormValues => accessGroups: keyData.object_permission?.agent_access_groups || [], }, skills: keyData.object_permission?.skills || [], - organization_id: keyData.organization_id, + organization_id: keyOrganizationId(keyData), team_id: keyData.team_id, project_id: keyData.project_id, logging_settings: extractLoggingSettings(keyData.metadata), 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 4dd98b31781..901b98407c5 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 @@ -1671,6 +1671,148 @@ describe("KeyEditView", () => { expect(await screen.findByText("Engineering")).toBeInTheDocument(); }); + it("should initialize organization from a key list row that only carries org_id", async () => { + const keyFromList = { + ...MOCK_KEY_DATA, + organization_id: null, + org_id: "org-1", + }; + + renderWithProviders( + {}} + onSubmit={async () => {}} + accessToken="" + userID="" + userRole="Admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByLabelText("Organization")).toHaveValue("Engineering"); + }); + + await userEvent.click(screen.getByRole("combobox", { name: "Team" })); + + expect(await screen.findByRole("option", { name: /Alpha/ })).toBeInTheDocument(); + expect(screen.queryByRole("option", { name: /Beta/ })).not.toBeInTheDocument(); + }); + + it("should let a non-admin move a key to their team in another organization", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + + renderWithProviders( + {}} + onSubmit={onSubmitMock} + accessToken="" + userID="" + userRole="Internal User" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByLabelText("Organization")).toHaveValue("Engineering"); + }); + + await userEvent.click(screen.getByRole("combobox", { name: "Team" })); + await userEvent.click(await screen.findByRole("option", { name: /Beta/ })); + + await waitFor(() => { + expect(screen.getByLabelText("Organization")).toHaveValue("Sales"); + }); + + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalled(); + }); + expect(onSubmitMock.mock.calls[0][0]).toMatchObject({ organization_id: "org-2", team_id: "team-b" }); + }); + + it("should re-scope teams when the edited key switches to another organization", async () => { + const sharedProps = { + teams: [ + { team_id: "team-a", team_alias: "Alpha", organization_id: "org-1" }, + { team_id: "team-b", team_alias: "Beta", organization_id: "org-2" }, + ], + onCancel: () => {}, + onSubmit: async () => {}, + accessToken: "", + userID: "", + userRole: "Admin", + premiumUser: false, + }; + + const { rerender } = renderWithProviders( + , + ); + + await waitFor(() => { + expect(screen.getByLabelText("Organization")).toHaveValue("Engineering"); + }); + + rerender(); + + await waitFor(() => { + expect(screen.getByLabelText("Organization")).toHaveValue("Sales"); + }); + + await userEvent.click(screen.getByRole("combobox", { name: "Team" })); + + expect(await screen.findByRole("option", { name: /Beta/ })).toBeInTheDocument(); + expect(screen.queryByRole("option", { name: /Alpha/ })).not.toBeInTheDocument(); + }); + + it("should adopt the organization of a team picked in the form", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + + renderWithProviders( + {}} + onSubmit={onSubmitMock} + accessToken="" + userID="" + userRole="Admin" + premiumUser={false} + />, + ); + + await screen.findByRole("button", { name: /save changes/i }); + expect(screen.getByLabelText("Organization")).toHaveValue(""); + + await userEvent.click(screen.getByRole("combobox", { name: "Team" })); + await userEvent.click(await screen.findByRole("option", { name: /Beta/ })); + + await waitFor(() => { + expect(screen.getByLabelText("Organization")).toHaveValue("Sales"); + }); + + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalled(); + }); + expect(onSubmitMock.mock.calls[0][0].organization_id).toBe("org-2"); + }); + it("should initialize organization from keyData", async () => { const keyWithOrg = { ...MOCK_KEY_DATA, 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 44e9af0d316..3e844d38d68 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -123,7 +123,6 @@ function KeyEditSession({ ? mapInternalToDisplayNames(keyData.metadata.litellm_disabled_callbacks) : [], ); - const [selectedOrganizationId, setSelectedOrganizationId] = useState(keyData.organization_id || null); const [autoRotationEnabled, setAutoRotationEnabled] = useState(keyData.auto_rotate || false); const [rotationInterval, setRotationInterval] = useState(keyData.rotation_interval || ""); const [neverExpire, setNeverExpire] = useState(!keyData.expires); @@ -147,9 +146,11 @@ function KeyEditSession({ const { data: uiSettingsData } = useUISettings(); const enableProjectsUI = Boolean(uiSettingsData?.values?.enable_projects_ui); const hasProject = Boolean(keyData.project_id); + const canChangeOrganization = userRole === "Admin" && !hasProject; const detachProject = hasProject && form.watch("project_id") === null; const canDetachProject = canDetachKeyProject(team, organizations, userID, userRole); + const selectedOrganizationId = form.watch("organization_id") ?? null; const allowedRoutesValue = form.watch("allowed_routes"); const selectedModels = (form.watch("models") as string[] | undefined) ?? []; const allowedRoutes = parseAllowedRoutes(allowedRoutesValue); @@ -327,7 +328,6 @@ function KeyEditSession({ const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | null) => { setField(orgId); - setSelectedOrganizationId(orgId); form.setValue("team_id", null, { shouldDirty: true }); }; @@ -336,7 +336,6 @@ function KeyEditSession({ clearMcpGrants(form.setValue); const organizationId = teams?.find((t) => t.team_id === teamId)?.organization_id; if (!organizationId) return; - setSelectedOrganizationId(organizationId); form.setValue("organization_id", organizationId, { shouldDirty: true }); }; @@ -354,9 +353,10 @@ function KeyEditSession({ })), ]; - const visibleTeams = selectedOrganizationId - ? teams?.filter((t) => t.organization_id === selectedOrganizationId) - : teams; + const visibleTeams = + canChangeOrganization && selectedOrganizationId + ? teams?.filter((t) => t.organization_id === selectedOrganizationId) + : teams; const teamOptionsFor = (selectedTeamId: string | null): SearchSelectOption[] => { const selectedTeamOutsideFilter = teams?.find((t) => t.team_id === selectedTeamId && !visibleTeams?.includes(t)); @@ -797,7 +797,7 @@ function KeyEditSession({ value={value} organizations={organizations} loading={isOrganizationsLoading} - disabled={userRole !== "Admin" || hasProject} + disabled={!canChangeOrganization} onChange={(orgId) => handleOrganizationChange(onChange, orgId)} /> )} 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 7b924243fb5..523c14d378d 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -56,6 +56,7 @@ import { teamMemberBudgetGate, } from "../shared/InheritedBudgetHint"; import { KeyEditView } from "./key_edit_view"; +import { keyOrganizationId } from "./keyEditFieldNormalizers"; import { isTeamAdminEditingMemberKey, teamAdminMemberKeyPayload } from "./teamAdminMemberKeyPayload"; export function needsLifetimeSpendBackfill(spend: number, totalSpend: number | null | undefined): boolean { @@ -239,6 +240,10 @@ export default function KeyInfoView({ delete formValues.disable_global_guardrails; } + if ((formValues.organization_id || null) === keyOrganizationId(currentKeyData)) { + delete formValues.organization_id; + } + // Handle max budget empty string formValues.max_budget = mapEmptyStringToNull(formValues.max_budget); @@ -397,7 +402,9 @@ export default function KeyInfoView({ const newKeyValues = await keyUpdateCall(accessToken, formValues); - setCurrentKeyData((prevData) => (prevData ? { ...prevData, ...newKeyValues } : undefined)); + setCurrentKeyData((prevData) => + prevData ? { ...prevData, ...newKeyValues, org_id: newKeyValues.organization_id } : undefined, + ); void queryClient.invalidateQueries({ queryKey: keyKeys.all }); if (onKeyDataUpdate) { @@ -545,7 +552,7 @@ export default function KeyInfoView({ const lastConfiguredAt = currentKeyData.settings_updated_at || currentKeyData.created_at; const parentTeam = currentKeyData.team_id ? teamsData?.find((team) => team.team_id === currentKeyData.team_id) : null; - const orgId = currentKeyData.organization_id || currentKeyData.org_id || parentTeam?.organization_id || ""; + const orgId = keyOrganizationId(currentKeyData) || parentTeam?.organization_id || ""; const parentOrg = orgId ? organizations?.find((org) => org.organization_id === orgId) : null; const hasOwnBudget = currentKeyData.max_budget !== null; @@ -970,7 +977,7 @@ export default function KeyInfoView({ Organization - {(currentKeyData.organization_id ?? currentKeyData.org_id) || "Not Set"} + {keyOrganizationId(currentKeyData) || "Not Set"}
Organization
{(currentKeyData.organization_id ?? currentKeyData.org_id) || "Not Set"}
{keyOrganizationId(currentKeyData) || "Not Set"}