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 9897bb4d47a..f9e03564ac9 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -30,10 +30,14 @@ vi.mock("../molecules/notifications_manager", () => { return { default: Notifications }; }); -// Roles: ensure 'admin' has write access -vi.mock("../../utils/roles", () => ({ - rolesWithWriteAccess: ["admin"], -})); +// Roles: ensure 'admin' has write access and include all role helper functions +vi.mock("../../utils/roles", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + rolesWithWriteAccess: ["admin"], + }; +}); // Helpers used in rendering vi.mock("@/utils/dataUtils", () => ({ @@ -218,6 +222,40 @@ vi.mock("../common_components/AutoRotationView", async () => { return { __esModule: true, default: AutoRotationView }; }); +// Mock Next.js router to avoid "invariant expected app router to be mounted" error +vi.mock("next/navigation", () => ({ + useRouter: () => ({ + push: vi.fn(), + replace: vi.fn(), + refresh: vi.fn(), + back: vi.fn(), + forward: vi.fn(), + prefetch: vi.fn(), + }), +})); + +// Mock useTeams hook +vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({ + default: vi.fn(() => ({ + teams: [], + setTeams: vi.fn(), + })), +})); + +// Mock useAuthorized hook to avoid Next.js router dependency +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: vi.fn(() => ({ + accessToken: "access_abc", + userId: "user_1", + userRole: "admin", + premiumUser: true, + token: "token_123", + userEmail: "test@example.com", + disabledPersonalKeyCreation: false, + showSSOBanner: false, + })), +})); + // KeyEditView mock: triggers onSubmit with our injected form values vi.mock("./key_edit_view", async () => { const React = await import("react"); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx index 8137a70124f..219c3cc515d 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx @@ -1,9 +1,20 @@ -import { render, waitFor } from "@testing-library/react"; -import { describe, expect, it } from "vitest"; -import { KeyResponse } from "../key_team_helpers/key_list"; +import { render, waitFor, screen } from "@testing-library/react"; +import { describe, expect, it, vi, beforeEach } from "vitest"; +import { KeyResponse, Team } from "../key_team_helpers/key_list"; import KeyInfoView from "./key_info_view"; +import useTeams from "@/app/(dashboard)/hooks/useTeams"; + +vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({ + default: vi.fn(), +})); describe("KeyInfoView", () => { + beforeEach(() => { + vi.mocked(useTeams).mockReturnValue({ + teams: [], + setTeams: vi.fn(), + }); + }); const MOCK_KEY_DATA: KeyResponse = { token: "test-token-123", token_id: "test-token-123", @@ -114,4 +125,134 @@ describe("KeyInfoView", () => { expect(metadataBlock?.textContent?.trim()).toBe("{}"); }); }); + + it("should allow proxy admin to modify key", async () => { + vi.mocked(useTeams).mockReturnValue({ + teams: [], + setTeams: vi.fn(), + }); + + const keyData = { ...MOCK_KEY_DATA, user_id: "other-user-id" }; + render( + {}} + keyId={"test-key-id"} + onKeyDataUpdate={() => {}} + accessToken={"test-token"} + userID={"proxy-admin-user"} + userRole={"proxy_admin"} + premiumUser={true} + teams={[]} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Regenerate Key")).toBeInTheDocument(); + expect(screen.getByText("Delete Key")).toBeInTheDocument(); + }); + }); + + it("should allow team admin to modify key", async () => { + const teamId = "test-team-id"; + const teamAdminUserId = "team-admin-user"; + const mockTeam: Team = { + team_id: teamId, + team_alias: "Test Team", + models: [], + max_budget: null, + budget_duration: null, + tpm_limit: null, + rpm_limit: null, + organization_id: "org-1", + created_at: "2025-01-01T00:00:00Z", + keys: [], + members_with_roles: [ + { + user_id: teamAdminUserId, + role: "admin", + }, + ], + }; + + vi.mocked(useTeams).mockReturnValue({ + teams: [mockTeam], + setTeams: vi.fn(), + }); + + const keyData = { ...MOCK_KEY_DATA, team_id: teamId, user_id: "other-user-id" }; + render( + {}} + keyId={"test-key-id"} + onKeyDataUpdate={() => {}} + accessToken={"test-token"} + userID={teamAdminUserId} + userRole={"user"} + premiumUser={true} + teams={[]} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Regenerate Key")).toBeInTheDocument(); + expect(screen.getByText("Delete Key")).toBeInTheDocument(); + }); + }); + + it("should allow owner to modify their own key", async () => { + vi.mocked(useTeams).mockReturnValue({ + teams: [], + setTeams: vi.fn(), + }); + + const ownerUserId = "owner-user-id"; + const keyData = { ...MOCK_KEY_DATA, user_id: ownerUserId }; + render( + {}} + keyId={"test-key-id"} + onKeyDataUpdate={() => {}} + accessToken={"test-token"} + userID={ownerUserId} + userRole={"user"} + premiumUser={true} + teams={[]} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Regenerate Key")).toBeInTheDocument(); + expect(screen.getByText("Delete Key")).toBeInTheDocument(); + }); + }); + + it("should not allow other user to modify key", async () => { + vi.mocked(useTeams).mockReturnValue({ + teams: [], + setTeams: vi.fn(), + }); + + const keyData = { ...MOCK_KEY_DATA, user_id: "owner-user-id" }; + render( + {}} + keyId={"test-key-id"} + onKeyDataUpdate={() => {}} + accessToken={"test-token"} + userID={"other-user-id"} + userRole={"user"} + premiumUser={true} + teams={[]} + />, + ); + + await waitFor(() => { + expect(screen.queryByText("Regenerate Key")).not.toBeInTheDocument(); + expect(screen.queryByText("Delete Key")).not.toBeInTheDocument(); + }); + }); }); 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 dcf6eb8b7eb..8b7f9e77879 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -1,3 +1,4 @@ +import useTeams from "@/app/(dashboard)/hooks/useTeams"; import { formatNumberWithCommas, copyToClipboard as utilCopyToClipboard } from "@/utils/dataUtils"; import { mapEmptyStringToNull } from "@/utils/keyUpdateUtils"; import { ArrowLeftIcon, RefreshIcon, TrashIcon } from "@heroicons/react/outline"; @@ -5,7 +6,7 @@ import { Badge, Button, Card, Grid, Tab, TabGroup, TabList, TabPanel, TabPanels, import { Button as AntdButton, Form, Tooltip } from "antd"; import { CheckIcon, CopyIcon } from "lucide-react"; import { useEffect, useState } from "react"; -import { rolesWithWriteAccess } from "../../utils/roles"; +import { isProxyAdminRole, isUserTeamAdminForSingleTeam, rolesWithWriteAccess } from "../../utils/roles"; import { mapDisplayToInternalNames, mapInternalToDisplayNames } from "../callback_info_helpers"; import AutoRotationView from "../common_components/AutoRotationView"; import { extractLoggingSettings, formatMetadataForDisplay, stripTagsFromMetadata } from "../key_info_utils"; @@ -54,6 +55,7 @@ export default function KeyInfoView({ setAccessToken, backButtonText = "Back to Keys", }: KeyInfoViewProps) { + const { teams: teamsData } = useTeams(); const [isEditing, setIsEditing] = useState(false); const [form] = Form.useForm(); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); @@ -301,6 +303,15 @@ export default function KeyInfoView({ return `${dateStr} at ${timeStr}`; }; + const canModifyKey = + isProxyAdminRole(userRole || "") || + (teamsData && + isUserTeamAdminForSingleTeam( + teamsData?.filter((team) => team.team_id === currentKeyData.team_id)[0], + userID || "", + )) || + userID === currentKeyData.user_id; + return (
@@ -349,7 +360,7 @@ export default function KeyInfoView({ )}
- {userRole && rolesWithWriteAccess.includes(userRole) && ( + {canModifyKey && (