diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index db1a089ff7..715a395e67 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2301,23 +2301,16 @@ async def validate_key_team_change( # Check if the team has access to the key's models if len(key.models) > 0: for model in key.models: + # Skip special sentinel values — "all-team-models" means + # "use whatever the team allows", so it's always valid. + if model == SpecialModelNames.all_team_models.value: + continue await can_team_access_model( model=model, team_object=team, llm_router=llm_router, ) - # Check if the key's user_id is a member of the team - member_object = _get_user_in_team( - team_table=cast(LiteLLM_TeamTableCachedObj, team), user_id=key.user_id - ) - if key.user_id is not None: - if not member_object: - raise HTTPException( - status_code=403, - detail=f"User={key.user_id} is not a member of the team={team.team_id}. Check team members via `/team/info`.", - ) - # Check if the key's tpm/rpm limit is less than the team's tpm/rpm limit if key.tpm_limit is not None: if team.tpm_limit and key.tpm_limit > team.tpm_limit: @@ -2331,6 +2324,17 @@ async def validate_key_team_change( detail=f"Key={key.token} has a rpm_limit={key.rpm_limit} which is greater than the team's rpm_limit={team.rpm_limit}.", ) + # Check if the key's user_id is a member of the team + member_object = _get_user_in_team( + team_table=cast(LiteLLM_TeamTableCachedObj, team), user_id=key.user_id + ) + if key.user_id is not None: + if not member_object: + raise HTTPException( + status_code=403, + detail=f"User={key.user_id} is not a member of the team={team.team_id}. Check team members via `/team/info`.", + ) + # Check if the person initiating the change is a Proxy Admin or Team Admin if change_initiated_by.user_role == LitellmUserRoles.PROXY_ADMIN.value: return diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index e496cf373e..cfc16808af 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -1502,6 +1502,57 @@ async def test_validate_key_team_change_with_member_permissions(): ) +@pytest.mark.asyncio +async def test_validate_key_team_change_skips_all_team_models_sentinel(): + """ + Test that validate_key_team_change skips the 'all-team-models' sentinel + value when checking if the target team can access the key's models. + + Keys with models=["all-team-models"] mean "use whatever models the team + allows", so moving them to any team should not fail model validation. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + mock_key = MagicMock() + mock_key.user_id = "test-user-123" + mock_key.models = ["all-team-models"] + mock_key.tpm_limit = None + mock_key.rpm_limit = None + + mock_team = MagicMock() + mock_team.team_id = "test-team-456" + mock_team.models = ["gpt-4", "claude-3"] + mock_team.members_with_roles = [] + mock_team.tpm_limit = None + mock_team.rpm_limit = None + + mock_change_initiator = MagicMock() + mock_change_initiator.user_id = "test-user-123" + mock_change_initiator.user_role = LitellmUserRoles.PROXY_ADMIN.value + + mock_router = MagicMock() + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.can_team_access_model", + new_callable=AsyncMock, + ) as mock_can_access: + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._get_user_in_team" + ) as mock_get_user: + mock_get_user.return_value = MagicMock() + + await validate_key_team_change( + key=mock_key, + team=mock_team, + change_initiated_by=mock_change_initiator, + llm_router=mock_router, + ) + + # can_team_access_model should NOT have been called since + # "all-team-models" is a sentinel that should be skipped + mock_can_access.assert_not_called() + + def test_key_rotation_fields_helper(): """ Test the key data update logic for rotation fields. diff --git a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx new file mode 100644 index 0000000000..1f6a61f39f --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx @@ -0,0 +1,69 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; +import OrganizationDropdown from "./OrganizationDropdown"; + +const MOCK_ORGS = [ + { + organization_id: "org-1", + organization_alias: "Engineering", + budget_id: "", + metadata: {}, + models: [], + spend: 0, + model_spend: {}, + created_at: "", + created_by: "", + updated_at: "", + }, + { + organization_id: "org-2", + organization_alias: "Sales", + budget_id: "", + metadata: {}, + models: [], + spend: 0, + model_spend: {}, + created_at: "", + created_by: "", + updated_at: "", + }, +]; + +describe("OrganizationDropdown", () => { + it("should render", () => { + render(); + expect(screen.getByRole("combobox")).toBeInTheDocument(); + }); + + it("should display organization options when opened", async () => { + const user = userEvent.setup(); + render(); + + await user.click(screen.getByRole("combobox")); + + expect(await screen.findByText("Engineering")).toBeInTheDocument(); + expect(screen.getByText("Sales")).toBeInTheDocument(); + }); + + it("should call onChange with the org id when an organization is selected", async () => { + const onChange = vi.fn(); + const user = userEvent.setup(); + render(); + + await user.click(screen.getByRole("combobox")); + await user.click(await screen.findByText("Engineering")); + + expect(onChange).toHaveBeenCalledWith("org-1", expect.anything()); + }); + + it("should add ant-select-disabled class when disabled prop is true", () => { + const { container } = render(); + expect(container.querySelector(".ant-select-disabled")).toBeTruthy(); + }); + + it("should render with empty organizations list", () => { + render(); + expect(screen.getByRole("combobox")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx new file mode 100644 index 0000000000..ac93041f3d --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx @@ -0,0 +1,52 @@ +import React from "react"; +import { Select } from "antd"; +import { Organization } from "../networking"; + +interface OrganizationDropdownProps { + organizations?: Organization[] | null; + value?: string; + onChange?: (value: string) => void; + disabled?: boolean; + loading?: boolean; +} + +const OrganizationDropdown: React.FC = ({ + organizations, + value, + onChange, + disabled, + loading, +}) => { + return ( + + ); +}; + +export default OrganizationDropdown; diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx index 3ed4c80aea..eef7292dac 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx @@ -213,7 +213,22 @@ vi.mock("../common_components/PassThroughRoutesSelector", () => ({ default: () = vi.mock("../common_components/PremiumLoggingSettings", () => ({ default: () => null })); vi.mock("../common_components/RateLimitTypeFormItem", () => ({ default: () => null })); vi.mock("../common_components/RouterSettingsAccordion", () => ({ default: () => null })); -vi.mock("../common_components/team_dropdown", () => ({ default: () => null })); +vi.mock("../common_components/team_dropdown", () => ({ + default: ({ teams, onChange, disabled }: { teams?: any[]; onChange?: (v: string) => void; disabled?: boolean }) => ( + + ), +})); vi.mock("../CreateUserButton", () => ({ CreateUserButton: () => null })); vi.mock("../mcp_server_management/MCPServerSelector", () => ({ default: () => null })); vi.mock("../mcp_server_management/MCPToolPermissions", () => ({ default: () => null })); @@ -227,6 +242,31 @@ vi.mock("@/app/(dashboard)/hooks/projects/useProjects", () => ({ useProjects: vi.fn().mockReturnValue({ data: [], isLoading: false }), })); +vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ + useOrganizations: vi.fn().mockReturnValue({ + data: [ + { organization_id: "org-1", organization_alias: "Engineering" }, + { organization_id: "org-2", organization_alias: "Sales" }, + ], + isLoading: false, + }), +})); + +vi.mock("../common_components/OrganizationDropdown", () => ({ + default: ({ value, onChange, disabled }: { value?: string; onChange?: (v: string) => void; disabled?: boolean }) => ( + + ), +})); + vi.mock("../common_components/ProjectDropdown", () => ({ default: ({ value, onChange }: { value?: string; onChange?: (v: string) => void }) => ( { expect(setFieldsValueMock).toHaveBeenCalledWith({ key_type: "management" }); }); }); + + describe("organization dropdown", () => { + it("should render the organization dropdown when modal is open", async () => { + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create new key/i })); + }); + + await waitFor(() => { + expect(screen.getByTestId("org-dropdown")).toBeInTheDocument(); + }); + }); + + it("should disable the organization dropdown for non-admin users", async () => { + authorizedState = { ...defaultAuthorizedState, userRole: "Internal User" }; + + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create new key/i })); + }); + + await waitFor(() => { + expect(screen.getByTestId("org-dropdown")).toBeDisabled(); + }); + }); + + it("should enable the organization dropdown for admin users", async () => { + authorizedState = { ...defaultAuthorizedState, userRole: "Admin" }; + + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create new key/i })); + }); + + await waitFor(() => { + expect(screen.getByTestId("org-dropdown")).not.toBeDisabled(); + }); + }); + + it("should render team dropdown alongside organization dropdown", async () => { + const teamsWithOrg = [ + { team_id: "team-1", team_alias: "Team Alpha", organization_id: "org-1", models: [] }, + ]; + + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create new key/i })); + }); + + await waitFor(() => { + expect(screen.getByTestId("org-dropdown")).toBeInTheDocument(); + expect(screen.getByTestId("team-dropdown")).toBeInTheDocument(); + }); + }); + + it("should set organization_id in form state when org is selected", async () => { + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create new key/i })); + }); + + await waitFor(() => { + expect(screen.getByTestId("org-dropdown")).toBeInTheDocument(); + }); + + act(() => { + fireEvent.change(screen.getByTestId("org-dropdown"), { target: { value: "org-1" } }); + }); + + expect(formStateRef.current["organization_id"]).toBe("org-1"); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 0926b11fe0..71e882db3f 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -1,5 +1,6 @@ "use client"; import { keyKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { useProjects } from "@/app/(dashboard)/hooks/projects/useProjects"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; @@ -23,6 +24,7 @@ import PremiumLoggingSettings from "../common_components/PremiumLoggingSettings" import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem"; import RouterSettingsAccordion, { RouterSettingsAccordionValue } from "../common_components/RouterSettingsAccordion"; import TeamDropdown from "../common_components/team_dropdown"; +import OrganizationDropdown from "../common_components/OrganizationDropdown"; import ProjectDropdown from "../common_components/ProjectDropdown"; import { CreateUserButton } from "../CreateUserButton"; import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; @@ -160,6 +162,7 @@ export const fetchUserModels = async ( const CreateKey: React.FC = ({ team, teams, data, addKey, autoOpenCreate, prefillData }) => { const { accessToken, userId: userID, userRole, premiumUser } = useAuthorized(); const canEditGuardrails = premiumUser || (userRole != null && rolesWithWriteAccess.includes(userRole)); + const { data: organizations, isLoading: isOrganizationsLoading } = useOrganizations(); const { data: projects, isLoading: isProjectsLoading } = useProjects(); const { data: uiSettingsData } = useUISettings(); const enableProjectsUI = Boolean(uiSettingsData?.values?.enable_projects_ui); @@ -179,6 +182,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp const [promptsList, setPromptsList] = useState([]); const [loggingSettings, setLoggingSettings] = useState([]); const [selectedCreateKeyTeam, setSelectedCreateKeyTeam] = useState(team); + const [selectedOrganizationId, setSelectedOrganizationId] = useState(null); const [selectedProjectId, setSelectedProjectId] = useState(null); const [isCreateUserModalVisible, setIsCreateUserModalVisible] = useState(false); const [newlyCreatedUserId, setNewlyCreatedUserId] = useState(null); @@ -207,6 +211,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp setRouterSettings(null); setRouterSettingsKey((prev) => prev + 1); setSelectedAgentId(null); + setSelectedOrganizationId(null); setSelectedProjectId(null); }; @@ -224,6 +229,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp setRouterSettings(null); setRouterSettingsKey((prev) => prev + 1); setSelectedAgentId(null); + setSelectedOrganizationId(null); setSelectedProjectId(null); }; @@ -752,6 +758,32 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp )} + + Organization{" "} + + + + + } + name="organization_id" + className="mt-4" + > + { + setSelectedOrganizationId(orgId || null); + // Clear team and project when org changes + setSelectedCreateKeyTeam(null); + setSelectedProjectId(null); + form.setFieldValue("team_id", undefined); + form.setFieldValue("project_id", undefined); + }} + /> + @@ -773,7 +805,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp help={keyOwner === "service_account" ? "required" : ""} > t.organization_id === selectedOrganizationId) : teams} disabled={selectedProjectId !== null} loading={!teams} onChange={(teamId) => { @@ -781,6 +813,14 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp setSelectedCreateKeyTeam(selectedTeam); setSelectedProjectId(null); form.setFieldValue("project_id", undefined); + // Auto-populate org from team for non-admin users + if (selectedTeam?.organization_id) { + setSelectedOrganizationId(selectedTeam.organization_id); + form.setFieldValue("organization_id", selectedTeam.organization_id); + } else if (!teamId) { + setSelectedOrganizationId(null); + form.setFieldValue("organization_id", undefined); + } }} /> @@ -1531,6 +1571,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp excludedFields={[ "key_alias", "team_id", + "organization_id", "models", "duration", "metadata", diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index b00a8d1e3f..2e4d0d97e4 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -53,6 +53,16 @@ vi.mock("../organisms/create_key_button", () => ({ fetchTeamModels: vi.fn().mockResolvedValue(["team-model-1", "team-model-2"]), })); +vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ + useOrganizations: vi.fn().mockReturnValue({ + data: [ + { organization_id: "org-1", organization_alias: "Engineering" }, + { organization_id: "org-2", organization_alias: "Sales" }, + ], + isLoading: false, + }), +})); + vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ useAccessGroups: vi.fn().mockReturnValue({ data: [ @@ -576,4 +586,91 @@ describe("KeyEditView", () => { resolveSubmit(); } }); + + describe("organization dropdown", () => { + it("should render the organization dropdown", async () => { + renderWithProviders( + {}} + onSubmit={async () => {}} + accessToken="" + userID="" + userRole="Admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Organization")).toBeInTheDocument(); + }); + }); + + it("should disable the organization dropdown for non-admin users", async () => { + const { container } = renderWithProviders( + {}} + onSubmit={async () => {}} + accessToken="" + userID="" + userRole="Internal User" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Organization")).toBeInTheDocument(); + }); + + const orgFormItem = screen.getByText("Organization").closest(".ant-form-item"); + const disabledSelect = orgFormItem?.querySelector(".ant-select-disabled"); + expect(disabledSelect).toBeTruthy(); + }); + + it("should not disable the organization dropdown for admin users", async () => { + const { container } = renderWithProviders( + {}} + onSubmit={async () => {}} + accessToken="" + userID="" + userRole="Admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Organization")).toBeInTheDocument(); + }); + + const orgFormItem = screen.getByText("Organization").closest(".ant-form-item"); + const disabledSelect = orgFormItem?.querySelector(".ant-select-disabled"); + expect(disabledSelect).toBeFalsy(); + }); + + it("should initialize organization from keyData", async () => { + const keyWithOrg = { + ...MOCK_KEY_DATA, + organization_id: "org-1", + }; + + renderWithProviders( + {}} + onSubmit={async () => {}} + accessToken="" + userID="" + userRole="Admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Engineering")).toBeInTheDocument(); + }); + }); + }); }); 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 b6c00577c9..cf431d1024 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -1,4 +1,5 @@ import GuardrailSelector from "@/components/guardrails/GuardrailSelector"; +import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { useProjects } from "@/app/(dashboard)/hooks/projects/useProjects"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import PolicySelector from "@/components/policies/PolicySelector"; @@ -13,6 +14,7 @@ import { mapInternalToDisplayNames } from "../callback_info_helpers"; import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings"; import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector"; import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem"; +import OrganizationDropdown from "../common_components/OrganizationDropdown"; import { extractLoggingSettings, formatMetadataForDisplay, stripTagsFromMetadata } from "../key_info_utils"; import { KeyResponse } from "../key_team_helpers/key_list"; import MCPServerSelector from "../mcp_server_management/MCPServerSelector"; @@ -96,10 +98,12 @@ export function KeyEditView({ ? 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); const [isKeySaving, setIsKeySaving] = useState(false); + const { data: organizations, isLoading: isOrganizationsLoading } = useOrganizations(); const { data: projects } = useProjects(); const { data: uiSettingsData } = useUISettings(); const enableProjectsUI = Boolean(uiSettingsData?.values?.enable_projects_ui); @@ -610,6 +614,28 @@ export function KeyEditView({ /> + + Organization{" "} + + + + + } + name="organization_id" + > + { + setSelectedOrganizationId(orgId || null); + form.setFieldValue("team_id", undefined); + }} + /> + + { + const selectedTeam = teams?.find((t) => t.team_id === teamId) || null; + if (selectedTeam?.organization_id) { + setSelectedOrganizationId(selectedTeam.organization_id); + form.setFieldValue("organization_id", selectedTeam.organization_id); + } else if (!teamId) { + setSelectedOrganizationId(null); + form.setFieldValue("organization_id", undefined); + } + }} filterOption={(input, option) => { - const team = teams?.find((t) => t.team_id === option?.value); + const filteredTeams = selectedOrganizationId + ? teams?.filter((t) => t.organization_id === selectedOrganizationId) + : teams; + const team = filteredTeams?.find((t) => t.team_id === option?.value); if (!team) return false; return team.team_alias?.toLowerCase().includes(input.toLowerCase()) ?? false; }} > - {teams?.map((team) => ( + {(selectedOrganizationId + ? teams?.filter((t) => t.organization_id === selectedOrganizationId) + : teams + )?.map((team) => ( {`${team.team_alias} (${team.team_id})`}