Merge 45339ad26e into f2118fc767

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(<TeamVirtualKeysTable {...defaultProps} organization={mockOrganization} />);
+
+    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({

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<SortingState>) => {

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(
+      <KeyInfoView keyId="tok_123" onClose={() => {}} keyData={keyListRow} onKeyDataUpdate={() => {}} teams={[]} />,
+    );
+    fireEvent.click(screen.getByText("Settings"));
+  };
+
+  const submitEditForm = (formValues: Record<string, unknown>) => {
+    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");
+  });
+});

ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.ts

@@ -9,6 +9,9 @@ const WORD_FORM_BUDGET_DURATIONS: Record<string, string> = {
 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";

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),

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(
+        <KeyEditView
+          keyData={keyFromList}
+          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}
+        />,
+      );
+
+      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(
+        <KeyEditView
+          keyData={{ ...MOCK_KEY_DATA, organization_id: null, org_id: "org-1" }}
+          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={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(
+        <KeyEditView keyData={{ ...MOCK_KEY_DATA, organization_id: null, org_id: "org-1" }} {...sharedProps} />,
+      );
+
+      await waitFor(() => {
+        expect(screen.getByLabelText("Organization")).toHaveValue("Engineering");
+      });
+
+      rerender(<KeyEditView keyData={{ ...MOCK_KEY_DATA, organization_id: null, org_id: "org-2" }} {...sharedProps} />);
+
+      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(
+        <KeyEditView
+          keyData={{ ...MOCK_KEY_DATA, organization_id: null, org_id: null }}
+          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={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,

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<string | null>(keyData.organization_id || null);
   const [autoRotationEnabled, setAutoRotationEnabled] = useState<boolean>(keyData.auto_rotate || false);
   const [rotationInterval, setRotationInterval] = useState<string>(keyData.rotation_interval || "");
   const [neverExpire, setNeverExpire] = useState<boolean>(!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)}
               />
             )}

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({
 
                   <div>
                     <p className="text-sm font-medium">Organization</p>
-                    <p className="text-sm">{(currentKeyData.organization_id ?? currentKeyData.org_id) || "Not Set"}</p>
+                    <p className="text-sm">{keyOrganizationId(currentKeyData) || "Not Set"}</p>
                   </div>
 
                   <div>