diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeamMemberBudgets.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeamMemberBudgets.test.ts new file mode 100644 index 00000000000..e945cd5e5f7 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeamMemberBudgets.test.ts @@ -0,0 +1,250 @@ +import { act, renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { createElement, type ReactNode } from "react"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import { teamInfoCall } from "@/components/networking"; +import type { TeamMemberBudgetSource } from "@/components/shared/InheritedBudgetHint"; +import { teamMemberBudgetGate } from "@/components/shared/InheritedBudgetHint"; +import { teamMemberBudgetQueryKey, useTeamMemberBudgets } from "./useTeamMemberBudgets"; + +const authMocks = vi.hoisted(() => ({ + useAuthorized: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: authMocks.useAuthorized, +})); + +vi.mock("@/components/networking", async (importOriginal) => { + const networking = await importOriginal(); + return { ...networking, teamInfoCall: vi.fn() }; +}); + +describe("useTeamMemberBudgets", () => { + const teamId = "team-1"; + const teamIds = [teamId]; + const teamInfo: TeamMemberBudgetSource = { + team_info: { team_id: teamId }, + }; + beforeEach(() => { + vi.clearAllMocks(); + authMocks.useAuthorized.mockReturnValue({ accessToken: "access-token", userId: "viewer-one" }); + vi.mocked(teamInfoCall).mockResolvedValue(teamInfo as never); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + const createQueryWrapper = () => { + const queryClient = new QueryClient({ + defaultOptions: { + queries: { retry: false }, + }, + }); + const wrapper = ({ children }: { children: ReactNode }) => + createElement(QueryClientProvider, { client: queryClient }, children); + return { queryClient, wrapper }; + }; + + it("keeps separate cached results for different viewers", async () => { + const { queryClient, wrapper } = createQueryWrapper(); + const viewerOneKey = [...teamMemberBudgetQueryKey(teamId), "viewer-one"]; + const viewerTwoKey = [...teamMemberBudgetQueryKey(teamId), "viewer-two"]; + + const firstViewer = renderHook(() => useTeamMemberBudgets(teamIds), { wrapper }); + await waitFor(() => expect(teamInfoCall).toHaveBeenCalledTimes(1)); + firstViewer.unmount(); + + authMocks.useAuthorized.mockReturnValue({ accessToken: "access-token", userId: "viewer-two" }); + renderHook(() => useTeamMemberBudgets(teamIds), { wrapper }); + await waitFor(() => expect(teamInfoCall).toHaveBeenCalledTimes(2)); + + const cachedKeys = queryClient + .getQueryCache() + .findAll({ queryKey: teamMemberBudgetQueryKey(teamId) }) + .map((query) => query.queryKey); + expect(cachedKeys).toHaveLength(2); + expect(cachedKeys).toContainEqual(viewerOneKey); + expect(cachedKeys).toContainEqual(viewerTwoKey); + expect(teamInfoCall).toHaveBeenNthCalledWith(1, "access-token", teamId, { keyLimit: 1 }); + expect(teamInfoCall).toHaveBeenNthCalledWith(2, "access-token", teamId, { keyLimit: 1 }); + }); + + it("refetches a viewer-scoped entry when invalidated by the team prefix", async () => { + const { queryClient, wrapper } = createQueryWrapper(); + renderHook(() => useTeamMemberBudgets(teamIds), { wrapper }); + await waitFor(() => expect(teamInfoCall).toHaveBeenCalledTimes(1)); + + await act(async () => { + await queryClient.invalidateQueries({ queryKey: teamMemberBudgetQueryKey(teamId) }); + }); + + expect(teamInfoCall).toHaveBeenCalledTimes(2); + }); + + it("drops an expired member budget increase without refetching", async () => { + vi.useFakeTimers({ now: new Date("2026-06-20T12:00:00Z") }); + const expiry = new Date(Date.now() + 60_000).toISOString(); + vi.mocked(teamInfoCall).mockResolvedValue({ + team_info: { team_id: teamId }, + team_memberships: [ + { + user_id: "viewer-one", + litellm_budget_table: { + max_budget: 50, + temp_budget_increase: 25, + temp_budget_expiry: expiry, + }, + }, + ], + } as never); + + const { wrapper } = createQueryWrapper(); + const { result } = renderHook( + () => teamMemberBudgetGate(useTeamMemberBudgets(teamIds)[teamId], "viewer-one")?.maxBudget, + { wrapper }, + ); + await act(async () => { + await vi.advanceTimersByTimeAsync(0); + }); + expect(result.current).toBe(75); + + act(() => { + vi.advanceTimersByTime(61_000); + }); + + expect(result.current).toBe(50); + expect(teamInfoCall).toHaveBeenCalledTimes(1); + }); + + it("re-arms the timer for the next expiry after one fires", async () => { + vi.useFakeTimers({ now: new Date("2026-06-20T12:00:00Z") }); + const teamTwoId = "team-2"; + const membershipFor = (increase: number, expiryOffsetMs: number) => ({ + team_info: { team_id: "ignored" }, + team_memberships: [ + { + user_id: "viewer-one", + litellm_budget_table: { + max_budget: 50, + temp_budget_increase: increase, + temp_budget_expiry: new Date(Date.now() + expiryOffsetMs).toISOString(), + }, + }, + ], + }); + vi.mocked(teamInfoCall).mockImplementation(async (_token, id) => + id === teamId ? (membershipFor(25, 60_000) as never) : (membershipFor(10, 120_000) as never), + ); + + const { wrapper } = createQueryWrapper(); + const { result } = renderHook( + () => { + const budgets = useTeamMemberBudgets([teamId, teamTwoId]); + return { + teamOne: teamMemberBudgetGate(budgets[teamId], "viewer-one")?.maxBudget, + teamTwo: teamMemberBudgetGate(budgets[teamTwoId], "viewer-one")?.maxBudget, + }; + }, + { wrapper }, + ); + await act(async () => { + await vi.advanceTimersByTimeAsync(0); + }); + expect(result.current).toEqual({ teamOne: 75, teamTwo: 60 }); + + act(() => { + vi.advanceTimersByTime(61_000); + }); + expect(result.current).toEqual({ teamOne: 50, teamTwo: 60 }); + + act(() => { + vi.advanceTimersByTime(60_000); + }); + expect(result.current).toEqual({ teamOne: 50, teamTwo: 50 }); + }); + + it("returns the same object reference when no increase is active", async () => { + vi.useFakeTimers({ now: new Date("2026-06-20T12:00:00Z") }); + vi.mocked(teamInfoCall).mockResolvedValue({ + team_info: { team_id: teamId }, + team_memberships: [{ user_id: "viewer-one", litellm_budget_table: { max_budget: 50 } }], + } as never); + + const { wrapper } = createQueryWrapper(); + const { result } = renderHook(() => useTeamMemberBudgets(teamIds), { wrapper }); + await act(async () => { + await vi.advanceTimersByTimeAsync(0); + }); + const before = result.current; + + act(() => { + vi.advanceTimersByTime(86_400_000); + }); + + expect(result.current).toBe(before); + }); + + it("does not fire early when the next expiry is far in the future", async () => { + vi.useFakeTimers({ now: new Date("2026-06-20T12:00:00Z") }); + const setTimeoutSpy = vi.spyOn(globalThis, "setTimeout"); + const teamTwoId = "team-2"; + const membershipFor = (increase: number, expiryOffsetMs: number) => ({ + team_info: { team_id: "ignored" }, + team_memberships: [ + { + user_id: "viewer-one", + litellm_budget_table: { + max_budget: 50, + temp_budget_increase: increase, + temp_budget_expiry: new Date(Date.now() + expiryOffsetMs).toISOString(), + }, + }, + ], + }); + vi.mocked(teamInfoCall).mockImplementation(async (_token, id) => + id === teamId ? (membershipFor(25, 40 * 86_400_000) as never) : (membershipFor(10, 41 * 86_400_000) as never), + ); + + const { wrapper } = createQueryWrapper(); + const { result } = renderHook( + () => { + const budgets = useTeamMemberBudgets([teamId, teamTwoId]); + return { + budgets, + teamOne: teamMemberBudgetGate(budgets[teamId], "viewer-one")?.maxBudget, + teamTwo: teamMemberBudgetGate(budgets[teamTwoId], "viewer-one")?.maxBudget, + }; + }, + { wrapper }, + ); + await act(async () => { + await vi.advanceTimersByTimeAsync(0); + }); + expect(result.current).toMatchObject({ teamOne: 75, teamTwo: 60 }); + const before = result.current; + + act(() => { + vi.advanceTimersByTime(1_000); + }); + + expect(result.current).toMatchObject({ teamOne: 75, teamTwo: 60 }); + expect(result.current).toBe(before); + const scheduledDelays = setTimeoutSpy.mock.calls.map(([, delay]) => delay ?? 0); + expect(Math.max(...scheduledDelays)).toBeLessThanOrEqual(2_147_483_647); + + act(() => { + vi.advanceTimersByTime(25 * 86_400_000); + }); + act(() => { + vi.advanceTimersByTime(15.2 * 86_400_000); + }); + expect(result.current).toMatchObject({ teamOne: 50, teamTwo: 60 }); + act(() => { + vi.advanceTimersByTime(2 * 86_400_000); + }); + expect(result.current).toMatchObject({ teamOne: 50, teamTwo: 50 }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeamMemberBudgets.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeamMemberBudgets.ts new file mode 100644 index 00000000000..3d2637b906d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeamMemberBudgets.ts @@ -0,0 +1,63 @@ +"use client"; + +import { useQueries, type UseQueryResult } from "@tanstack/react-query"; +import { useEffect, useMemo, useState } from "react"; + +import type { TeamMemberBudgetSource } from "@/components/shared/InheritedBudgetHint"; +import { tempBudgetExpiryMs } from "@/components/shared/InheritedBudgetHint"; +import { teamInfoCall } from "@/components/networking"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { teamKeys } from "@/app/(dashboard)/hooks/teams/useTeams"; + +export const teamMemberBudgetQueryKey = (teamId: string) => [...teamKeys.all, "memberBudget", teamId] as const; + +const MAX_TIMER_DELAY_MS = 2_147_483_647; + +export const useTeamMemberBudgets = (teamIds: readonly string[]): Readonly> => { + const { accessToken, userId } = useAuthorized(); + const uniqueTeamIds = useMemo(() => [...new Set(teamIds)], [teamIds]); + const [lastExpiryCheckMs, setLastExpiryCheckMs] = useState(null); + const combine = useMemo( + () => + (results: UseQueryResult[]): Readonly> => + Object.fromEntries( + results.flatMap((result, index) => + result.isSuccess && result.data ? [[uniqueTeamIds[index], result.data]] : [], + ), + ), + [uniqueTeamIds], + ); + + const budgets = useQueries({ + queries: uniqueTeamIds.map((teamId) => ({ + queryKey: [...teamMemberBudgetQueryKey(teamId), userId], + queryFn: async (): Promise => + (await teamInfoCall(accessToken!, teamId, { keyLimit: 1 })) as TeamMemberBudgetSource, + enabled: Boolean(accessToken && userId), + retry: false, + })), + combine, + }); + + const pendingExpiries = useMemo( + () => + Object.values(budgets).flatMap((teamInfo) => + (teamInfo.team_memberships ?? []) + .map((membership) => tempBudgetExpiryMs(membership.litellm_budget_table)) + .filter((expiry): expiry is number => expiry !== null), + ), + [budgets], + ); + + useEffect(() => { + const nextExpiryMs = pendingExpiries.filter((expiry) => expiry > Date.now()).sort((a, b) => a - b)[0]; + if (nextExpiryMs === undefined) return; + const timer = setTimeout( + () => setLastExpiryCheckMs(Date.now()), + Math.min(nextExpiryMs - Date.now() + 1, MAX_TIMER_DELAY_MS), + ); + return () => clearTimeout(timer); + }, [pendingExpiries, lastExpiryCheckMs]); + + return useMemo(() => (lastExpiryCheckMs === null ? budgets : { ...budgets }), [budgets, lastExpiryCheckMs]); +}; diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 423723a5938..ccf7f453132 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -1,15 +1,15 @@ -import { screen, waitFor, within, fireEvent } from "@testing-library/react"; +import { act, screen, waitFor, within, fireEvent } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import type { OnUrlUpdateFunction } from "nuqs/adapters/testing"; import { vi, it, expect, beforeEach, describe, Mock, MockedFunction } from "vitest"; -import { chooseSelectOption, renderWithProviders } from "../../../tests/test-utils"; +import { chooseSelectOption, renderWithProviders, testQueryClient } from "../../../tests/test-utils"; import { VirtualKeysTable } from "./VirtualKeysTable"; import { KEY_TABLE_HIDDEN_COLUMNS, KEY_TABLE_SORT_FIELDS } from "./keyTableColumns"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import { useKeyInfo } from "@/app/(dashboard)/hooks/keys/useKeyInfo"; import { KeysResponse, useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import useTeams from "@/app/(dashboard)/hooks/useTeams"; -import { regenerateKeyCall } from "../networking"; +import { regenerateKeyCall, teamInfoCall } from "../networking"; // Resolve debounced values synchronously so an applied filter lands in the useKeys query within the test tick. vi.mock("@tanstack/react-pacer/debouncer", async () => { @@ -30,6 +30,7 @@ vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); vi.mock("../networking", async (importOriginal) => ({ ...(await importOriginal()), regenerateKeyCall: vi.fn(), + teamInfoCall: vi.fn(), })); vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ @@ -43,6 +44,7 @@ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ })); vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + teamKeys: { all: ["teams"] }, useAllTeams: vi.fn(() => ({ data: [{ team_id: "team-1", team_alias: "Test Team" }], isLoading: false, @@ -85,6 +87,7 @@ const mockKey: KeyResponse = { spend: 5.5, total_spend: 42.25, max_budget: 100, + key_type: null, expires: "2999-12-31T23:59:59Z", models: ["gpt-3.5-turbo", "gpt-4"], aliases: {}, @@ -103,7 +106,10 @@ const mockKey: KeyResponse = { allowed_routes: [], permissions: {}, model_spend: { "gpt-3.5-turbo": 2.5, "gpt-4": 3.0 }, - model_max_budget: { "gpt-3.5-turbo": 50, "gpt-4": 50 }, + model_max_budget: { + "gpt-3.5-turbo": { budget_limit: 50, time_period: "1m" }, + "gpt-4": { budget_limit: 50, time_period: "1m" }, + }, soft_budget_cooldown: false, blocked: false, litellm_budget_table: {}, @@ -191,6 +197,7 @@ const lastHistoryMode = (onUrlUpdate: Mock) => onUrlUpdate. beforeEach(() => { vi.clearAllMocks(); + vi.mocked(teamInfoCall).mockReset(); localStorage.clear(); mockUseKeys.mockReturnValue(keysResult([mockKey])); @@ -248,6 +255,31 @@ it("should display key information correctly", async () => { }); }); +it("shows the team-member budget for a key without its own budget", async () => { + vi.mocked(teamInfoCall).mockResolvedValue({ + team_info: { + team_id: "member-budget-table-team", + team_alias: "Test Team", + team_member_budget_table: { max_budget: 50, budget_duration: "30d" }, + }, + team_memberships: [{ user_id: "user-1", litellm_budget_table: null }], + } as never); + mockUseKeys.mockReturnValue( + keysResult([{ ...mockKey, team_id: "member-budget-table-team", max_budget: null } as unknown as KeyResponse]), + ); + + renderWithProviders(); + + expect(await screen.findByText("· $0.00 of $50 (team member)")).toBeInTheDocument(); + expect(teamInfoCall).toHaveBeenCalledWith("test-token", "member-budget-table-team", { keyLimit: 1 }); + + await act(async () => { + await testQueryClient.invalidateQueries({ queryKey: ["teams"] }); + }); + + expect(teamInfoCall).toHaveBeenCalledTimes(2); +}); + it("shows lifetime spend in its own column next to the period spend meter", async () => { renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index cc25d22e546..32801b693d4 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -3,6 +3,7 @@ import { useKeyInfo } from "@/app/(dashboard)/hooks/keys/useKeyInfo"; import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; +import { useTeamMemberBudgets } from "@/app/(dashboard)/hooks/teams/useTeamMemberBudgets"; import { useApplyUserBudgetToTeamKeys } from "@/app/(dashboard)/hooks/uiSettings/useApplyUserBudgetToTeamKeys"; import { useAllTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; @@ -141,16 +142,25 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { const keyList = useMemo(() => keys?.keys ?? [], [keys]); const rowCount = keys?.total_count ?? 0; + const teamIds = useMemo( + () => + Array.from( + new Set(keyList.flatMap((key) => (key.max_budget == null && key.team_id && key.user_id ? [key.team_id] : []))), + ), + [keyList], + ); + const teamMemberBudgets = useTeamMemberBudgets(teamIds); const applyUserBudgetToTeamKeys = useApplyUserBudgetToTeamKeys(); const columnDeps = useMemo( () => ({ allTeams, organizations, + teamMemberBudgets, onSelectKey: (key: KeyResponse) => void setSelectedKeyId(key.token), applyUserBudgetToTeamKeys, }), - [allTeams, organizations, setSelectedKeyId, applyUserBudgetToTeamKeys], + [allTeams, organizations, teamMemberBudgets, setSelectedKeyId, applyUserBudgetToTeamKeys], ); const columns = useMemo(() => getKeyTableColumns(columnDeps), [columnDeps]); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx index d69e90b1882..9751b9ebf2c 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx @@ -4,7 +4,12 @@ import { Info } from "lucide-react"; import { ColumnDef } from "@tanstack/react-table"; import { DataTableMultiSortHeader, DataTableSortHeader, type DataTableSortField } from "@/components/shared/DataTable"; -import { inheritedBudgetGates, keyOwnerBudgetSource } from "@/components/shared/InheritedBudgetHint"; +import { + inheritedBudgetGates, + keyOwnerBudgetSource, + teamMemberBudgetGate, + type TeamMemberBudgetSource, +} from "@/components/shared/InheritedBudgetHint"; import { HoverCard, HoverCardContent, HoverCardTrigger } from "@/components/ui/hover-card"; import { Skeleton } from "@/components/ui/skeleton"; import { @@ -85,6 +90,7 @@ const InfoHeader = ({ label, tooltip }: { label: string; tooltip: string }) => ( interface KeyTableColumnsDeps { allTeams: Team[]; organizations: Organization[]; + teamMemberBudgets: Readonly>; onSelectKey: (key: KeyResponse) => void; applyUserBudgetToTeamKeys: boolean; } @@ -92,6 +98,7 @@ interface KeyTableColumnsDeps { export const getKeyTableColumns = ({ allTeams, organizations, + teamMemberBudgets, onSelectKey, applyUserBudgetToTeamKeys, }: KeyTableColumnsDeps): ColumnDef[] => [ @@ -269,19 +276,33 @@ export const getKeyTableColumns = ({ accessorKey: "spend", meta: { title: "Spend / Budget", skeleton: "meter", numeric: true }, header: ({ table }) => , - size: 180, + size: 240, enableSorting: true, cell: ({ row }) => { const team = allTeams.find((t) => t.team_id === row.original.team_id); const orgId = row.original.organization_id || row.original.org_id || team?.organization_id; const organization = organizations.find((o) => o.organization_id === orgId); + const memberGate = + row.original.max_budget == null + ? teamMemberBudgetGate( + teamMemberBudgets[row.original.team_id ?? ""], + row.original.user_id, + row.original.user?.user_email, + ) + : null; return ( 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..2f0b9c9b00e 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 @@ -38,7 +38,7 @@ export interface KeyResponse { key_alias: string; spend: number; total_spend: number; - max_budget: number; + max_budget: number | null; expires: string; models: string[]; aliases: Record; diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index 0fdf7a4799d..e60c0356751 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -442,6 +442,20 @@ describe("teamInfoCall", () => { expect(parsed.searchParams.get("team_id")).toBe(teamID); }); + it("passes key_limit when requested", async () => { + const mockFetch = vi.spyOn(global, "fetch").mockResolvedValue(new Response("{}", { status: 200 })); + + try { + await Networking.teamInfoCall("token", "team-1", { keyLimit: 1 }); + + const [url] = mockFetch.mock.calls[0]; + const parsed = url instanceof Request ? new URL(url.url) : new URL(url.toString(), "http://example.com"); + expect(parsed.searchParams.get("key_limit")).toBe("1"); + } finally { + mockFetch.mockRestore(); + } + }); + it("should not append team_id when teamID is null", async () => { const mockFetch = vi.fn().mockResolvedValue({ ok: true, @@ -455,6 +469,7 @@ describe("teamInfoCall", () => { const [url] = mockFetch.mock.calls[0]; const parsed = typeof url === "string" ? new URL(url, "http://example.com") : new URL((url as Request).url); expect(parsed.searchParams.has("team_id")).toBe(false); + expect(parsed.searchParams.has("key_limit")).toBe(false); }); }); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 81fd3f7b441..2b38da75411 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -1126,9 +1126,12 @@ export const userGetInfoV2 = async (accessToken: string, userId?: string): Promi } }; -export const teamInfoCall = async (accessToken: string, teamID: string | null) => { +export const teamInfoCall = async (accessToken: string, teamID: string | null, options?: { keyLimit?: number }) => { try { - return await apiClient.get(`/team/info`, { accessToken, query: { team_id: teamID || undefined } }); + return await apiClient.get(`/team/info`, { + accessToken, + query: { team_id: teamID || undefined, key_limit: options?.keyLimit }, + }); } catch (error) { console.error("Failed to create key:", error); throw error; diff --git a/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.test.tsx b/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.test.tsx index 8faae9c9660..699db37331c 100644 --- a/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.test.tsx @@ -2,7 +2,12 @@ import { render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, expect, it } from "vitest"; -import { InheritedBudgetHint, inheritedBudgetGates, keyOwnerBudgetSource } from "./InheritedBudgetHint"; +import { + InheritedBudgetHint, + inheritedBudgetGates, + keyOwnerBudgetSource, + teamMemberBudgetGate, +} from "./InheritedBudgetHint"; const team = { team_id: "team-1", team_alias: "Platform", max_budget: 1200, budget_duration: "30d" }; const organization = { @@ -73,6 +78,237 @@ describe("inheritedBudgetGates", () => { "User", ]); }); + + it("places the team-member gate after the team and before organization and user gates", () => { + const memberGate = { + scope: "Team member" as const, + alias: "owner@example.com in Platform", + maxBudget: 50, + budgetDuration: "30d", + }; + + expect(inheritedBudgetGates(team, organization, user, memberGate).map((gate) => gate.scope)).toEqual([ + "Team", + "Team member", + "Organization", + "User", + ]); + }); +}); + +describe("teamMemberBudgetGate", () => { + const teamInfo = { + team_info: { + team_id: "team-1", + team_alias: "Platform", + team_member_budget_table: { max_budget: 50, budget_duration: "30d" }, + }, + team_memberships: [ + { user_id: "custom-user", spend: 14, litellm_budget_table: { max_budget: 25, budget_duration: "7d" } }, + { user_id: "default-user", spend: 35, litellm_budget_table: { max_budget: null, budget_duration: null } }, + ], + }; + + it("uses a custom membership budget instead of the team default", () => { + const expectedGate = { + scope: "Team member", + alias: "custom@example.com in Platform", + maxBudget: 25, + budgetDuration: "7d", + spend: 14, + }; + + expect(teamMemberBudgetGate(teamInfo, "custom-user", "custom@example.com")).toEqual(expectedGate); + }); + + it("falls back to the team default when the membership budget is null", () => { + const expectedGate = { + scope: "Team member", + alias: "default-user in Platform", + maxBudget: 50, + budgetDuration: "30d", + spend: 35, + }; + + expect(teamMemberBudgetGate(teamInfo, "default-user")).toEqual(expectedGate); + }); + + it("uses the team default when there is no membership row", () => { + expect(teamMemberBudgetGate(teamInfo, "other-user")?.maxBudget).toBe(50); + }); + + it("defaults member spend to zero when there is no membership row", () => { + const expectedGate = { + scope: "Team member", + alias: "other-user in Platform", + maxBudget: 50, + budgetDuration: "30d", + spend: 0, + }; + + expect(teamMemberBudgetGate(teamInfo, "other-user")).toEqual(expectedGate); + }); + + it("adds an active membership increase to a custom membership budget", () => { + const now = new Date("2026-06-20T12:00:00.000Z"); + const increasedTeamInfo = { + ...teamInfo, + team_memberships: [ + { + user_id: "custom-user", + spend: 45, + litellm_budget_table: { + max_budget: 25, + budget_duration: "7d", + temp_budget_increase: 10, + temp_budget_expiry: "2026-06-20T12:30:00Z", + }, + }, + ], + }; + const expectedGate = { + scope: "Team member", + alias: "custom-user in Platform", + maxBudget: 35, + budgetDuration: "7d", + spend: 45, + }; + + expect(teamMemberBudgetGate(increasedTeamInfo, "custom-user", null, now)).toEqual(expectedGate); + }); + + it("adds an active membership increase to the team default when the membership budget is null", () => { + const now = new Date("2026-06-20T12:00:00.000Z"); + const increasedTeamInfo = { + ...teamInfo, + team_memberships: [ + { + user_id: "default-user", + spend: 42, + litellm_budget_table: { + max_budget: null, + budget_duration: null, + temp_budget_increase: 10, + temp_budget_expiry: "2026-06-20T12:30:00Z", + }, + }, + ], + }; + const expectedGate = { + scope: "Team member", + alias: "default-user in Platform", + maxBudget: 60, + budgetDuration: "30d", + spend: 42, + }; + + expect(teamMemberBudgetGate(increasedTeamInfo, "default-user", null, now)).toEqual(expectedGate); + }); + + it("ignores an expired membership increase", () => { + const now = new Date("2026-06-20T12:00:00.000Z"); + const expiredTeamInfo = { + ...teamInfo, + team_memberships: [ + { + user_id: "custom-user", + litellm_budget_table: { + max_budget: 25, + budget_duration: "7d", + temp_budget_increase: 10, + temp_budget_expiry: "2026-06-20T11:30:00Z", + }, + }, + ], + }; + + expect(teamMemberBudgetGate(expiredTeamInfo, "custom-user", null, now)?.maxBudget).toBe(25); + }); + + it("treats a naive membership expiry as UTC", () => { + const teamInfoWithNaiveExpiry = { + ...teamInfo, + team_memberships: [ + { + user_id: "default-user", + litellm_budget_table: { + max_budget: null, + temp_budget_increase: 10, + temp_budget_expiry: "2026-06-20T12:00:00", + }, + }, + ], + }; + const beforeExpiry = new Date("2026-06-20T11:59:59.000Z"); + const afterExpiry = new Date("2026-06-20T12:00:01.000Z"); + + expect(teamMemberBudgetGate(teamInfoWithNaiveExpiry, "default-user", null, beforeExpiry)?.maxBudget).toBe(60); + expect(teamMemberBudgetGate(teamInfoWithNaiveExpiry, "default-user", null, afterExpiry)?.maxBudget).toBe(50); + }); + + it("returns no gate when the team default is zero and the member has no budget", () => { + expect( + teamMemberBudgetGate( + { ...teamInfo, team_info: { ...teamInfo.team_info, team_member_budget_table: { max_budget: 0 } } }, + "other-user", + ), + ).toBeNull(); + }); + + it("does not create a gate from an increase when the team default is zero", () => { + const now = new Date("2026-06-20T12:00:00.000Z"); + const noBudgetWithIncrease = { + ...teamInfo, + team_info: { ...teamInfo.team_info, team_member_budget_table: { max_budget: 0 } }, + team_memberships: [ + { + user_id: "default-user", + litellm_budget_table: { + max_budget: null, + temp_budget_increase: 10, + temp_budget_expiry: "2026-06-20T12:30:00Z", + }, + }, + ], + }; + + expect(teamMemberBudgetGate(noBudgetWithIncrease, "default-user", null, now)).toBeNull(); + }); + + it("keeps an explicit zero membership budget as a gate", () => { + const expectedGate = { + scope: "Team member", + alias: "zero-user in Platform", + maxBudget: 0, + budgetDuration: null, + spend: 0, + }; + + expect( + teamMemberBudgetGate( + { + ...teamInfo, + team_memberships: [{ user_id: "zero-user", litellm_budget_table: { max_budget: 0 } }], + }, + "zero-user", + ), + ).toEqual(expectedGate); + }); + + it("returns no gate without a user id", () => { + expect(teamMemberBudgetGate(teamInfo, null)).toBeNull(); + }); + + it("falls back to the user id and team id when labels are absent", () => { + expect( + teamMemberBudgetGate( + { + team_info: { team_id: "team-id", team_member_budget_table: { max_budget: 50 } }, + }, + "user-id", + )?.alias, + ).toBe("user-id in team-id"); + }); }); describe("keyOwnerBudgetSource", () => { diff --git a/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.tsx b/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.tsx index 4c02efa66ed..5c0e01ac20d 100644 --- a/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.tsx +++ b/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.tsx @@ -6,10 +6,36 @@ import type { Organization } from "@/components/networking"; import { formatNumberWithCommas } from "@/utils/dataUtils"; export interface InheritedBudgetGate { - scope: "Team" | "Organization" | "User"; + scope: "Team" | "Team member" | "Organization" | "User"; alias: string; maxBudget: number; budgetDuration: string | null; + spend?: number | null; +} + +export type MemberBudgetRow = + | { + max_budget?: number | null; + budget_duration?: string | null; + temp_budget_increase?: number | null; + temp_budget_expiry?: string | null; + } + | null + | undefined; + +export interface TeamMemberBudgetSource { + team_info: { + team_id: string; + team_alias?: string | null; + team_member_budget_table?: MemberBudgetRow; + }; + team_memberships?: + | readonly { + user_id: string; + spend?: number | null; + litellm_budget_table?: MemberBudgetRow; + }[] + | null; } type TeamBudgetSource = Pick; @@ -56,12 +82,53 @@ const userGate = (user: UserBudgetSource | null | undefined): InheritedBudgetGat } : null; +export const tempBudgetExpiryMs = (row: MemberBudgetRow): number | null => { + if (row?.temp_budget_increase == null || row.temp_budget_expiry == null) return null; + const expiryValue = /(?:Z|[+-]\d{2}:?\d{2})$/i.test(row.temp_budget_expiry) + ? row.temp_budget_expiry + : `${row.temp_budget_expiry}Z`; + return new Date(expiryValue).getTime(); +}; + +const activeIncrease = (row: MemberBudgetRow, now: Date): number => { + const expiry = tempBudgetExpiryMs(row); + if (expiry === null || expiry <= now.getTime()) return 0; + return row?.temp_budget_increase ?? 0; +}; + +export const teamMemberBudgetGate = ( + teamInfo: TeamMemberBudgetSource | null | undefined, + userId: string | null | undefined, + memberLabel?: string | null, + now: Date = new Date(), +): InheritedBudgetGate | null => { + if (!teamInfo || !userId) return null; + + const membership = teamInfo.team_memberships?.find((row) => row.user_id === userId); + const membershipBudget = membership?.litellm_budget_table; + const teamBudget = teamInfo.team_info.team_member_budget_table; + const hasMembershipBudget = membershipBudget?.max_budget != null; + const budgetRow = hasMembershipBudget ? membershipBudget : teamBudget; + if (!budgetRow || budgetRow.max_budget == null || (!hasMembershipBudget && budgetRow.max_budget <= 0)) { + return null; + } + + return { + scope: "Team member", + alias: `${memberLabel || userId} in ${teamInfo.team_info.team_alias || teamInfo.team_info.team_id}`, + maxBudget: budgetRow.max_budget + activeIncrease(membershipBudget, now), + budgetDuration: budgetRow.budget_duration ?? null, + spend: membership?.spend ?? 0, + }; +}; + export const inheritedBudgetGates = ( team: TeamBudgetSource | null | undefined, organization: OrganizationBudgetSource | null | undefined, user?: UserBudgetSource | null, + teamMember?: InheritedBudgetGate | null, ): readonly InheritedBudgetGate[] => - [teamGate(team), organizationGate(organization), userGate(user)].filter((gate) => gate !== null); + [teamGate(team), teamMember ?? null, organizationGate(organization), userGate(user)].filter((gate) => gate !== null); export const keyOwnerBudgetSource = ( key: { team_id?: string | null; user?: UserBudgetSource | null }, diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.test.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.test.tsx index 587d0275c8a..7973752cb31 100644 --- a/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.test.tsx @@ -62,6 +62,48 @@ describe("SpendBudgetCell", () => { expect(screen.getByLabelText("question-circle")).toBeInTheDocument(); }); + it("uses membership spend against the team-member budget when the key has no own budget", () => { + const teamMemberGate = { + scope: "Team member" as const, + alias: "user@example.com in Team A", + maxBudget: 50, + budgetDuration: "30d", + spend: 45, + }; + const { container } = render( + , + ); + + expect(screen.getByText("$5.00")).toBeInTheDocument(); + expect(screen.getByText("· $45.00 of $50 (team member)")).toBeInTheDocument(); + const meter = screen.getByRole("meter"); + expect(meter).toHaveAttribute("aria-valuenow", "45"); + expect(meter).toHaveAttribute("aria-valuemax", "50"); + expect(meter).toHaveAttribute("aria-valuetext", "team member spend $45.00 of $50"); + expect(indicator(container)?.className).toContain("bg-warning"); + expect(screen.getByLabelText("question-circle")).toBeInTheDocument(); + }); + + it("uses the key's own budget instead of a team-member budget", () => { + const teamMemberGate = { + scope: "Team member" as const, + alias: "user@example.com in Team A", + maxBudget: 50, + budgetDuration: "30d", + }; + render(); + + expect(screen.getByText("of $100")).toBeInTheDocument(); + expect(screen.queryByText(/team member/)).not.toBeInTheDocument(); + expect(screen.getByRole("meter")).toHaveAttribute("aria-valuemax", "100"); + }); + it("shows no inherited-budget hint when there is nothing to inherit", () => { render(); expect(screen.queryByLabelText("question-circle")).not.toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx index 127c7202121..914d3a98c81 100644 --- a/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx +++ b/ui/litellm-dashboard/src/components/shared/table_cells/spend_budget_cell.tsx @@ -7,6 +7,7 @@ import { formatNumberWithCommas, getSpendString } from "@/utils/dataUtils"; interface SpendBudgetCellProps { spend: number | null | undefined; maxBudget: number | null | undefined; + teamMemberGate?: InheritedBudgetGate | null; inheritedGates?: readonly InheritedBudgetGate[]; spendDecimals?: number; budgetDecimals?: number; @@ -18,34 +19,62 @@ const meterTone = (pct: number): "default" | "warning" | "over" => { return "default"; }; +const getBudgetLabel = ( + budget: number | null, + usesTeamMemberBudget: boolean, + memberSpendText: string, + budgetDecimals: number, +): string => { + if (budget === null) return "· Unlimited"; + const formattedBudget = formatNumberWithCommas(budget, budgetDecimals); + if (usesTeamMemberBudget) return `· ${memberSpendText} of $${formattedBudget} (team member)`; + return `of $${formattedBudget}`; +}; + +const getMeterValueText = ( + hasBudget: boolean, + budget: number | null, + meterSpendText: string, + budgetDecimals: number, +): string => { + if (!hasBudget || budget === null) return ""; + const formattedBudget = formatNumberWithCommas(budget, budgetDecimals); + return `${meterSpendText} of $${formattedBudget}`; +}; + export function SpendBudgetCell({ spend, maxBudget, + teamMemberGate = null, inheritedGates = [], spendDecimals = 4, budgetDecimals = 0, }: SpendBudgetCellProps) { const spendValue = typeof spend === "number" && !Number.isNaN(spend) ? spend : 0; - const budget = maxBudget ?? null; + const ownBudget = maxBudget ?? null; + const budget = ownBudget ?? teamMemberGate?.maxBudget ?? null; + const usesTeamMemberBudget = ownBudget === null && teamMemberGate !== null; + const memberSpendValue = + typeof teamMemberGate?.spend === "number" && !Number.isNaN(teamMemberGate.spend) ? teamMemberGate.spend : 0; + const meterSpendValue = usesTeamMemberBudget ? memberSpendValue : spendValue; const hasBudget = typeof budget === "number" && budget > 0; - const pct = hasBudget ? (spendValue / budget) * 100 : 0; + const pct = hasBudget ? (meterSpendValue / budget) * 100 : 0; const spendText = spendValue > 0 ? getSpendString(spendValue, spendDecimals) : "$0.00"; - const budgetLabel = budget === null ? "· Unlimited" : `of $${formatNumberWithCommas(budget, budgetDecimals)}`; + const memberSpendText = memberSpendValue > 0 ? getSpendString(memberSpendValue, spendDecimals) : "$0.00"; + const budgetLabel = getBudgetLabel(budget, usesTeamMemberBudget, memberSpendText, budgetDecimals); + const meterSpendText = usesTeamMemberBudget ? `team member spend ${memberSpendText}` : spendText; + const meterValueText = getMeterValueText(hasBudget, budget, meterSpendText, budgetDecimals); return (
{spendText}{" "} {budgetLabel} - {budget === null && } + {ownBudget === null && }
{hasBudget && ( - + diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index 3230ffe5e60..03cee6e5fc1 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -5,6 +5,7 @@ import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { chooseSelectOption, renderWithProviders, testQueryClient } from "../../../tests/test-utils"; import { toast } from "@/lib/toast"; +import type { Member } from "@/components/networking"; import type { EffectiveMcpServer } from "../mcp_server_management/effectiveMcpServers"; import type { MCPServer } from "../mcp_tools/types"; import TeamInfoView, { @@ -133,11 +134,29 @@ vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({ })); vi.mock("@/components/team/TeamMemberTab", () => ({ - default: vi.fn(({ setIsAddMemberModalVisible }) => ( -
- -
- )), + default: vi.fn( + ({ + setIsAddMemberModalVisible, + setSelectedEditMember, + setIsEditMemberModalVisible, + }: { + setIsAddMemberModalVisible: (visible: boolean) => void; + setSelectedEditMember: (member: Member) => void; + setIsEditMemberModalVisible: (visible: boolean) => void; + }) => ( +
+ + +
+ ), + ), })); vi.mock("@/components/common_components/user_search_modal", () => ({ @@ -1076,6 +1095,31 @@ describe("TeamInfoView", () => { }); }); + it("invalidates the member budget query after a successful member update", async () => { + const user = userEvent.setup({ delay: null }); + const teamData = createMockTeamData(); + vi.mocked(networking.teamInfoCall).mockResolvedValue(teamData); + testQueryClient.clear(); + const invalidateQueriesSpy = vi.spyOn(testQueryClient, "invalidateQueries"); + + try { + renderWithProviders(); + + await user.click(await screen.findByRole("tab", { name: "Members" })); + await user.click(await screen.findByRole("button", { name: "Edit Member" })); + await user.click(await screen.findByRole("button", { name: "Submit" })); + + await waitFor(() => expect(networking.teamMemberUpdateCall).toHaveBeenCalled()); + await waitFor(() => + expect(invalidateQueriesSpy).toHaveBeenCalledWith({ + queryKey: ["teams", "memberBudget", defaultProps.teamId], + }), + ); + } finally { + invalidateQueriesSpy.mockRestore(); + } + }); + it("should display soft budget in settings view when present", async () => { const user = userEvent.setup({ delay: null }); vi.mocked(networking.teamInfoCall).mockResolvedValue( diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 3f933c4a4f6..a782112edb2 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -3,6 +3,7 @@ import type { components } from "@/lib/http/schema"; import useCan from "@/app/(dashboard)/hooks/useCan"; import { organizationKeys, useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { invalidateTeamQueries } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { teamMemberBudgetQueryKey } from "@/app/(dashboard)/hooks/teams/useTeamMemberBudgets"; import { useQueryClient } from "@tanstack/react-query"; import UserSearchModal from "@/components/common_components/user_search_modal"; import { @@ -734,6 +735,8 @@ const TeamInfoView: React.FC = ({ setTeamData(await teamInfoCall(accessToken, teamId)); } catch { toast.fromError("Failed to load team information"); + } finally { + void queryClient.invalidateQueries({ queryKey: teamMemberBudgetQueryKey(teamId) }); } }; @@ -820,6 +823,7 @@ const TeamInfoView: React.FC = ({ }; await teamMemberAddCall(accessToken, teamId, member); + void queryClient.invalidateQueries({ queryKey: teamMemberBudgetQueryKey(teamId) }); toast.success("Team member added successfully"); setIsAddMemberModalVisible(false); @@ -866,6 +870,7 @@ const TeamInfoView: React.FC = ({ toast.dismiss(); // Remove all existing toasts await teamMemberUpdateCall(accessToken, teamId, member); + void queryClient.invalidateQueries({ queryKey: teamMemberBudgetQueryKey(teamId) }); toast.success("Team member updated successfully"); setIsEditMemberModalVisible(false); @@ -903,6 +908,7 @@ const TeamInfoView: React.FC = ({ setIsDeleting(true); try { await teamMemberDeleteCall(accessToken, teamId, memberToDelete); + void queryClient.invalidateQueries({ queryKey: teamMemberBudgetQueryKey(teamId) }); toast.success("Team member removed successfully"); diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx index 7561264ea65..5c3af3f6c3f 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx @@ -1,8 +1,10 @@ import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi, MockedFunction } from "vitest"; +import { QueryClient, QueryObserver } from "@tanstack/react-query"; import { fireEvent, renderWithProviders, screen, waitFor, within } from "../../../tests/test-utils"; import { TeamVirtualKeysTable } from "./TeamVirtualKeysTable"; 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"; @@ -10,6 +12,10 @@ vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ useKeys: vi.fn(), })); +vi.mock("@/app/(dashboard)/hooks/teams/useTeamMemberBudgets", () => ({ + useTeamMemberBudgets: vi.fn(), +})); + vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); vi.mock("../key_team_helpers/fetch_available_models_team_key", () => ({ @@ -31,6 +37,7 @@ vi.mock("@tanstack/react-pacer/debouncer", () => ({ })); const mockUseKeys = useKeys as MockedFunction; +const mockUseTeamMemberBudgets = useTeamMemberBudgets as MockedFunction; const KEY_HASH = "88a145505dd6e87e2ea166fcef1e4b53948dbdb32af6431dfd05ec06b571ee52"; @@ -51,6 +58,18 @@ const createMockKey = (overrides: Partial = {}): KeyResponse => ...overrides, }) as KeyResponse; +const createKeysQueryResult = (data: KeysResponse): ReturnType => { + const queryClient = new QueryClient(); + const queryOptions = { + queryKey: ["team-virtual-keys-table-test"], + queryFn: async () => data, + initialData: data, + enabled: false, + }; + const observer = new QueryObserver(queryClient, queryOptions); + return observer.getCurrentResult(); +}; + const mockOrganization: Organization = { organization_id: "org-123", organization_alias: "Test Org", @@ -78,6 +97,7 @@ describe("TeamVirtualKeysTable", () => { beforeEach(() => { vi.clearAllMocks(); + mockUseTeamMemberBudgets.mockReturnValue({}); mockUseKeys.mockReturnValue({ data: { keys: [], total_count: 0, current_page: 1, total_pages: 1 } as KeysResponse, isPending: false, @@ -163,6 +183,33 @@ describe("TeamVirtualKeysTable", () => { expect(screen.getByText("bob_key_team1")).toBeInTheDocument(); }); + it("shows the team member budget for budgetless keys and keeps each key's own budget", async () => { + const keysResponse: KeysResponse = { + keys: [ + createMockKey({ token: "sk-no-budget", token_id: "key-no-budget", max_budget: null }), + createMockKey({ token: "sk-own-budget", token_id: "key-own-budget", max_budget: 75 }), + ], + total_count: 2, + current_page: 1, + total_pages: 1, + }; + mockUseKeys.mockReturnValue(createKeysQueryResult(keysResponse)); + const teamMemberBudgets = { + "team-1": { + team_info: { team_id: "team-1", team_member_budget_table: { max_budget: 50 } }, + team_memberships: [], + }, + }; + mockUseTeamMemberBudgets.mockReturnValue(teamMemberBudgets); + + renderWithProviders(); + + const inheritedBudget = await screen.findByText("$50", { exact: false }); + expect(inheritedBudget).toHaveTextContent("$50 (team member)"); + expect(screen.getByText("$75")).toBeInTheDocument(); + expect(mockUseTeamMemberBudgets).toHaveBeenCalledWith(["team-1"]); + }); + it("should show the current range from total_count when multiple pages exist", async () => { mockUseKeys.mockReturnValue({ data: { diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx index 4e4fa3f4bb2..8ea559990c0 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx @@ -1,5 +1,7 @@ "use client"; import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { useTeamMemberBudgets } from "@/app/(dashboard)/hooks/teams/useTeamMemberBudgets"; +import { teamMemberBudgetGate } from "@/components/shared/InheritedBudgetHint"; import { SimpleTooltip } from "@/components/ui/tooltip"; import { DateCell, @@ -21,6 +23,7 @@ import { Input } from "@/components/ui/input"; import { orgDetailHref, userDetailHref } from "@/utils/entityLinks"; import { DEFAULT_PROXY_ADMIN_USER_ID } from "@/utils/sentinels"; import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; +import { formatNumberWithCommas } from "@/utils/dataUtils"; import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { ColumnDef, ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; import { ChevronDown, ChevronRight } from "lucide-react"; @@ -97,6 +100,12 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi })); }, [keys?.keys, organization?.organization_id]); + const teamMemberBudgetTeamIds = useMemo( + () => (displayKeys.some((key) => key.max_budget == null && Boolean(key.user_id)) ? [teamId] : []), + [displayKeys, teamId], + ); + const teamMemberBudgets = useTeamMemberBudgets(teamMemberBudgetTeamIds); + const rowCount = keys?.total_count ?? 0; const [expandedAccordions, setExpandedAccordions] = useState>({}); @@ -296,11 +305,24 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi accessorKey: "max_budget", meta: { title: "Budget (USD)", numeric: true }, header: ({ column }) => , - size: 110, + size: 160, enableSorting: true, - cell: (info) => ( - - ), + cell: (info) => { + const key = info.row.original; + const memberGate = + key.max_budget == null + ? teamMemberBudgetGate(teamMemberBudgets[teamId], key.user_id, key.user?.user_email) + : null; + + return memberGate ? ( + + ${formatNumberWithCommas(memberGate.maxBudget, 0)}{" "} + (team member) + + ) : ( + + ); + }, }, { id: "budget_reset_at", @@ -420,7 +442,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi }, }, ], - [expandedAccordions], + [expandedAccordions, teamMemberBudgets, teamId], ); 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 88f7fbb9f8c..c82a4823dfd 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -31,6 +31,10 @@ vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ useOrganizations: () => ({ data: [] }), })); +vi.mock("@/app/(dashboard)/hooks/teams/useTeamMemberBudgets", () => ({ + useTeamMemberBudgets: vi.fn().mockReturnValue({}), +})); + // Networking: wire the hoisted fns so we can assert calls later vi.mock("../networking", () => { return { 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 f4e1fe288ec..a1d3e498dd4 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 @@ -5,7 +5,7 @@ import { screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; -import { keyDeleteCall, keyUpdateCall } from "../networking"; +import { keyDeleteCall, keyUpdateCall, teamInfoCall } from "../networking"; import { QueryClient } from "@tanstack/react-query"; import KeyInfoView, { needsLifetimeSpendBackfill } from "./key_info_view"; @@ -92,6 +92,7 @@ vi.mock("../networking", () => ({ serverRootPath: "", keyDeleteCall: vi.fn().mockResolvedValue({}), keyUpdateCall: vi.fn().mockResolvedValue({}), + teamInfoCall: vi.fn(), getPolicyInfoWithGuardrails: vi.fn().mockResolvedValue({ resolved_guardrails: ["guardrail-1", "guardrail-2"], }), @@ -114,6 +115,7 @@ vi.mock("@/utils/dataUtils", () => ({ describe("KeyInfoView", () => { beforeEach(() => { + vi.mocked(teamInfoCall).mockReset(); vi.mocked(useTeams).mockReturnValue({ teams: [], setTeams: vi.fn(), @@ -128,6 +130,7 @@ describe("KeyInfoView", () => { total_spend: 0, max_budget: 0, expires: "null", + key_type: null, models: [], aliases: {}, config: {}, @@ -154,6 +157,7 @@ describe("KeyInfoView", () => { litellm_budget_table: {}, organization_id: null, created_at: "2025-10-29T01:26:41.613000Z", + last_active: null, updated_at: "2025-10-29T01:47:33.980000Z", team_spend: 100, team_alias: "", @@ -196,13 +200,29 @@ describe("KeyInfoView", () => { accessToken: "test-token", userId: "test-user", userRole: "admin", + userRoleLabel: "Admin", premiumUser: true, token: "test-token", userEmail: null, disabledPersonalKeyCreation: null, + isLoading: false, + isAuthorized: true, + isViewOnly: false, + loginMethod: null, + passwordResetRequired: false, showSSOBanner: false, }; + const authorizedForTeamBudget = { + ...baseUseAuthorizedMock, + isLoading: false, + isAuthorized: true, + userRoleLabel: "Admin", + isViewOnly: false, + loginMethod: null, + passwordResetRequired: false, + } satisfies ReturnType; + const openMoreKeyActions = async () => { await userEvent.click(await screen.findByRole("button", { name: /more key actions/i })); }; @@ -296,6 +316,66 @@ describe("KeyInfoView", () => { expect(screen.getByTestId("key-lifetime-spend")).toHaveTextContent("Lifetime spend: $340.5000"); }); + it("shows a team-member budget for a key without its own budget", async () => { + vi.mocked(useAuthorized).mockReturnValue(authorizedForTeamBudget); + vi.mocked(teamInfoCall).mockResolvedValue({ + team_info: { + team_id: "member-budget-info-team", + team_alias: "Test Team", + team_member_budget_table: { max_budget: 50, budget_duration: "30d" }, + }, + team_memberships: [{ user_id: "owner-user-id", spend: 45, litellm_budget_table: null }], + } as never); + + renderWithProviders( + {}} + keyId="test-key-id" + teams={[]} + />, + ); + + expect(await screen.findByText("Team member spend $45.0000 of $50.00 (team member budget)")).toBeInTheDocument(); + expect(screen.queryByText("of Unlimited")).not.toBeInTheDocument(); + expect(teamInfoCall).toHaveBeenCalledWith("test-token", "member-budget-info-team", { keyLimit: 1 }); + + await userEvent.click(screen.getByRole("tab", { name: "Settings" })); + expect(await screen.findByText("$50.00 (team member budget)")).toBeInTheDocument(); + }); + + it("keeps an own key budget in the Overview when a team-member budget exists", async () => { + vi.mocked(useAuthorized).mockReturnValue(authorizedForTeamBudget); + vi.mocked(teamInfoCall).mockResolvedValue({ + team_info: { + team_id: "own-budget-info-team", + team_alias: "Test Team", + team_member_budget_table: { max_budget: 50, budget_duration: "30d" }, + }, + team_memberships: [], + } as never); + + renderWithProviders( + {}} + keyId="test-key-id" + teams={[]} + />, + ); + + expect(await screen.findByText("of $20.00")).toBeInTheDocument(); + expect(teamInfoCall).not.toHaveBeenCalled(); + }); + it("shows the backfill hint when lifetime spend trails the period spend", async () => { vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock); 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 04a792bb5ca..7b924243fb5 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -21,7 +21,7 @@ import { KeyInfoHeader } from "./KeyInfoHeader"; import KeySavingsTab from "./KeySavingsTab"; import KeyAutoRouterUsageTab from "./KeyAutoRouterUsageTab"; import { useActivityDateRange } from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange"; -import { useEffect, useState } from "react"; +import { useEffect, useMemo, useState } from "react"; import { hasProxyWideSpendView, isProxyAdminRole, @@ -44,11 +44,17 @@ import { keyKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useQueryClient } from "@tanstack/react-query"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; +import { useTeamMemberBudgets } from "@/app/(dashboard)/hooks/teams/useTeamMemberBudgets"; import { extractMcpEntitlement } from "../mcp_server_management/mcpEntitlement"; import ObjectPermissionsView from "../object_permissions_view"; import { RegenerateKeyModal } from "../organisms/RegenerateKeyModal"; import { parseErrorMessage } from "../shared/errorUtils"; -import { InheritedBudgetHint, inheritedBudgetGates, keyOwnerBudgetSource } from "../shared/InheritedBudgetHint"; +import { + InheritedBudgetHint, + inheritedBudgetGates, + keyOwnerBudgetSource, + teamMemberBudgetGate, +} from "../shared/InheritedBudgetHint"; import { KeyEditView } from "./key_edit_view"; import { isTeamAdminEditingMemberKey, teamAdminMemberKeyPayload } from "./teamAdminMemberKeyPayload"; @@ -118,6 +124,14 @@ export default function KeyInfoView({ const { mutate: setKeyBlockedState, isPending: blockLoading } = useSetKeyBlockedState(); // Add local state to maintain key data and track regeneration const [currentKeyData, setCurrentKeyData] = useState(keyData); + const memberBudgetTeamIds = useMemo( + () => + currentKeyData?.max_budget === null && currentKeyData.team_id && currentKeyData.user_id + ? [currentKeyData.team_id] + : [], + [currentKeyData], + ); + const teamMemberBudgets = useTeamMemberBudgets(memberBudgetTeamIds); const [lastRegeneratedAt, setLastRegeneratedAt] = useState(null); const [keyDataUpdateHeldUntilModalClose, setKeyDataUpdateHeldUntilModalClose] = useState | null>( null, @@ -535,9 +549,20 @@ export default function KeyInfoView({ const parentOrg = orgId ? organizations?.find((org) => org.organization_id === orgId) : null; const hasOwnBudget = currentKeyData.max_budget !== null; - const budgetDisplay = hasOwnBudget ? `$${formatNumberWithCommas(currentKeyData.max_budget, 2)}` : "Unlimited"; const ownerUser = keyOwnerBudgetSource(currentKeyData, applyUserBudgetToTeamKeys); - const inheritedGates = hasOwnBudget ? [] : inheritedBudgetGates(parentTeam, parentOrg, ownerUser); + const memberGate = hasOwnBudget + ? null + : teamMemberBudgetGate( + teamMemberBudgets[currentKeyData.team_id ?? ""], + currentKeyData.user_id, + currentKeyData.user?.user_email, + ); + const budgetDisplay = (() => { + if (hasOwnBudget) return `$${formatNumberWithCommas(currentKeyData.max_budget, 2)}`; + if (memberGate) return `$${formatNumberWithCommas(memberGate.maxBudget, 2)} (team member budget)`; + return "Unlimited"; + })(); + const inheritedGates = hasOwnBudget ? [] : inheritedBudgetGates(parentTeam, parentOrg, ownerUser, memberGate); return (
@@ -701,7 +726,9 @@ export default function KeyInfoView({

${formatNumberWithCommas(currentKeyData.spend, 4)}

- of {budgetDisplay} + {memberGate + ? `Team member spend $${formatNumberWithCommas(memberGate.spend ?? 0, 4)} of $${formatNumberWithCommas(memberGate.maxBudget, 2)} (team member budget)` + : `of ${budgetDisplay}`}

{currentKeyData.budget_reset_at && ( @@ -999,11 +1026,7 @@ export default function KeyInfoView({

Budget

-

- {currentKeyData.max_budget !== null - ? `$${formatNumberWithCommas(currentKeyData.max_budget, 2)}` - : "Unlimited"} -

+

{budgetDisplay}