mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
fix(ui): show the team member budget on keys without their own budget (#43927)
Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
77b6d048df
commit
b72bfa735b
19 changed files with 1038 additions and 44 deletions
|
|
@ -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<typeof import("@/components/networking")>();
|
||||
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 });
|
||||
});
|
||||
});
|
||||
|
|
@ -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<Record<string, TeamMemberBudgetSource>> => {
|
||||
const { accessToken, userId } = useAuthorized();
|
||||
const uniqueTeamIds = useMemo(() => [...new Set(teamIds)], [teamIds]);
|
||||
const [lastExpiryCheckMs, setLastExpiryCheckMs] = useState<number | null>(null);
|
||||
const combine = useMemo(
|
||||
() =>
|
||||
(results: UseQueryResult<TeamMemberBudgetSource>[]): Readonly<Record<string, TeamMemberBudgetSource>> =>
|
||||
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<TeamMemberBudgetSource> =>
|
||||
(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]);
|
||||
};
|
||||
|
|
@ -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<typeof import("../networking")>()),
|
||||
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<OnUrlUpdateFunction>) => 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(<VirtualKeysTable />);
|
||||
|
||||
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(<VirtualKeysTable />);
|
||||
|
||||
|
|
|
|||
|
|
@ -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]);
|
||||
|
||||
|
|
|
|||
|
|
@ -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<Record<string, TeamMemberBudgetSource>>;
|
||||
onSelectKey: (key: KeyResponse) => void;
|
||||
applyUserBudgetToTeamKeys: boolean;
|
||||
}
|
||||
|
|
@ -92,6 +98,7 @@ interface KeyTableColumnsDeps {
|
|||
export const getKeyTableColumns = ({
|
||||
allTeams,
|
||||
organizations,
|
||||
teamMemberBudgets,
|
||||
onSelectKey,
|
||||
applyUserBudgetToTeamKeys,
|
||||
}: KeyTableColumnsDeps): ColumnDef<KeyResponse>[] => [
|
||||
|
|
@ -269,19 +276,33 @@ export const getKeyTableColumns = ({
|
|||
accessorKey: "spend",
|
||||
meta: { title: "Spend / Budget", skeleton: "meter", numeric: true },
|
||||
header: ({ table }) => <DataTableMultiSortHeader table={table} fields={SPEND_BUDGET_SORT_FIELDS} />,
|
||||
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 (
|
||||
<SpendBudgetCell
|
||||
spend={row.original.spend}
|
||||
maxBudget={row.original.max_budget}
|
||||
teamMemberGate={memberGate}
|
||||
inheritedGates={
|
||||
row.original.max_budget == null
|
||||
? inheritedBudgetGates(team, organization, keyOwnerBudgetSource(row.original, applyUserBudgetToTeamKeys))
|
||||
? inheritedBudgetGates(
|
||||
team,
|
||||
organization,
|
||||
keyOwnerBudgetSource(row.original, applyUserBudgetToTeamKeys),
|
||||
memberGate,
|
||||
)
|
||||
: []
|
||||
}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -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<string, unknown>;
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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", () => {
|
||||
|
|
|
|||
|
|
@ -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<Team, "team_id" | "team_alias" | "max_budget" | "budget_duration">;
|
||||
|
|
@ -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 },
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
<SpendBudgetCell
|
||||
spend={5}
|
||||
maxBudget={null}
|
||||
teamMemberGate={teamMemberGate}
|
||||
inheritedGates={[teamMemberGate]}
|
||||
spendDecimals={2}
|
||||
/>,
|
||||
);
|
||||
|
||||
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(<SpendBudgetCell spend={10} maxBudget={100} teamMemberGate={teamMemberGate} />);
|
||||
|
||||
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(<SpendBudgetCell spend={10} maxBudget={null} inheritedGates={[]} />);
|
||||
expect(screen.queryByLabelText("question-circle")).not.toBeInTheDocument();
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
<div className="flex min-w-[130px] flex-col gap-1">
|
||||
<div className="whitespace-nowrap text-xs">
|
||||
<span className="font-medium tabular-nums text-foreground">{spendText}</span>{" "}
|
||||
<span className="text-muted-foreground">{budgetLabel}</span>
|
||||
{budget === null && <InheritedBudgetHint gates={inheritedGates} />}
|
||||
{ownBudget === null && <InheritedBudgetHint gates={inheritedGates} />}
|
||||
</div>
|
||||
{hasBudget && (
|
||||
<Meter
|
||||
value={spendValue}
|
||||
max={budget}
|
||||
aria-valuetext={`${spendText} of $${formatNumberWithCommas(budget, budgetDecimals)}`}
|
||||
>
|
||||
<Meter value={meterSpendValue} max={budget} aria-valuetext={meterValueText}>
|
||||
<MeterTrack>
|
||||
<MeterIndicator tone={meterTone(pct)} />
|
||||
</MeterTrack>
|
||||
|
|
|
|||
|
|
@ -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 }) => (
|
||||
<div>
|
||||
<button onClick={() => setIsAddMemberModalVisible(true)}>Add Member</button>
|
||||
</div>
|
||||
)),
|
||||
default: vi.fn(
|
||||
({
|
||||
setIsAddMemberModalVisible,
|
||||
setSelectedEditMember,
|
||||
setIsEditMemberModalVisible,
|
||||
}: {
|
||||
setIsAddMemberModalVisible: (visible: boolean) => void;
|
||||
setSelectedEditMember: (member: Member) => void;
|
||||
setIsEditMemberModalVisible: (visible: boolean) => void;
|
||||
}) => (
|
||||
<div>
|
||||
<button onClick={() => setIsAddMemberModalVisible(true)}>Add Member</button>
|
||||
<button
|
||||
onClick={() => {
|
||||
setSelectedEditMember({ user_email: "edit@test.com", user_id: "edit-user", role: "user" });
|
||||
setIsEditMemberModalVisible(true);
|
||||
}}
|
||||
>
|
||||
Edit Member
|
||||
</button>
|
||||
</div>
|
||||
),
|
||||
),
|
||||
}));
|
||||
|
||||
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(<TeamInfoView {...defaultProps} />);
|
||||
|
||||
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(
|
||||
|
|
|
|||
|
|
@ -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<TeamInfoProps> = ({
|
|||
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<TeamInfoProps> = ({
|
|||
};
|
||||
|
||||
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<TeamInfoProps> = ({
|
|||
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<TeamInfoProps> = ({
|
|||
setIsDeleting(true);
|
||||
try {
|
||||
await teamMemberDeleteCall(accessToken, teamId, memberToDelete);
|
||||
void queryClient.invalidateQueries({ queryKey: teamMemberBudgetQueryKey(teamId) });
|
||||
|
||||
toast.success("Team member removed successfully");
|
||||
|
||||
|
|
|
|||
|
|
@ -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<typeof useKeys>;
|
||||
const mockUseTeamMemberBudgets = useTeamMemberBudgets as MockedFunction<typeof useTeamMemberBudgets>;
|
||||
|
||||
const KEY_HASH = "88a145505dd6e87e2ea166fcef1e4b53948dbdb32af6431dfd05ec06b571ee52";
|
||||
|
||||
|
|
@ -51,6 +58,18 @@ const createMockKey = (overrides: Partial<KeyResponse> = {}): KeyResponse =>
|
|||
...overrides,
|
||||
}) as KeyResponse;
|
||||
|
||||
const createKeysQueryResult = (data: KeysResponse): ReturnType<typeof useKeys> => {
|
||||
const queryClient = new QueryClient();
|
||||
const queryOptions = {
|
||||
queryKey: ["team-virtual-keys-table-test"],
|
||||
queryFn: async () => data,
|
||||
initialData: data,
|
||||
enabled: false,
|
||||
};
|
||||
const observer = new QueryObserver<KeysResponse>(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(<TeamVirtualKeysTable {...defaultProps} />);
|
||||
|
||||
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: {
|
||||
|
|
|
|||
|
|
@ -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<Record<string, boolean>>({});
|
||||
|
||||
|
|
@ -296,11 +305,24 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
|
|||
accessorKey: "max_budget",
|
||||
meta: { title: "Budget (USD)", numeric: true },
|
||||
header: ({ column }) => <DataTableSortHeader column={column} title="Budget (USD)" variant="header-cycle" />,
|
||||
size: 110,
|
||||
size: 160,
|
||||
enableSorting: true,
|
||||
cell: (info) => (
|
||||
<MoneyCell value={info.getValue() as number | null} decimals={0} emptyText="Unlimited" showZero />
|
||||
),
|
||||
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 ? (
|
||||
<span className="block w-full whitespace-nowrap text-right tabular-nums">
|
||||
${formatNumberWithCommas(memberGate.maxBudget, 0)}{" "}
|
||||
<span className="text-muted-foreground">(team member)</span>
|
||||
</span>
|
||||
) : (
|
||||
<MoneyCell value={info.getValue() as number | null} decimals={0} emptyText="Unlimited" showZero />
|
||||
);
|
||||
},
|
||||
},
|
||||
{
|
||||
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<SortingState>) => {
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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<typeof useAuthorized>;
|
||||
|
||||
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(
|
||||
<KeyInfoView
|
||||
keyData={
|
||||
{
|
||||
...MOCK_KEY_DATA,
|
||||
max_budget: null,
|
||||
team_id: "member-budget-info-team",
|
||||
user_id: "owner-user-id",
|
||||
user: { user_id: "owner-user-id", user_email: "owner@example.com", user_alias: null },
|
||||
} as unknown as KeyResponse
|
||||
}
|
||||
onClose={() => {}}
|
||||
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(
|
||||
<KeyInfoView
|
||||
keyData={{ ...MOCK_KEY_DATA, max_budget: 20, team_id: "own-budget-info-team", user_id: "owner-user-id" }}
|
||||
onClose={() => {}}
|
||||
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);
|
||||
|
||||
|
|
|
|||
|
|
@ -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<KeyResponse | undefined>(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<Date | null>(null);
|
||||
const [keyDataUpdateHeldUntilModalClose, setKeyDataUpdateHeldUntilModalClose] = useState<Partial<KeyResponse> | 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 (
|
||||
<div className="w-full h-full overflow-y-auto p-4">
|
||||
|
|
@ -701,7 +726,9 @@ export default function KeyInfoView({
|
|||
<div className="mt-2">
|
||||
<h3 className="text-lg font-medium">${formatNumberWithCommas(currentKeyData.spend, 4)}</h3>
|
||||
<p className="text-sm">
|
||||
of {budgetDisplay}
|
||||
{memberGate
|
||||
? `Team member spend $${formatNumberWithCommas(memberGate.spend ?? 0, 4)} of $${formatNumberWithCommas(memberGate.maxBudget, 2)} (team member budget)`
|
||||
: `of ${budgetDisplay}`}
|
||||
<InheritedBudgetHint gates={inheritedGates} />
|
||||
</p>
|
||||
{currentKeyData.budget_reset_at && (
|
||||
|
|
@ -999,11 +1026,7 @@ export default function KeyInfoView({
|
|||
|
||||
<div>
|
||||
<p className="text-sm font-medium">Budget</p>
|
||||
<p className="text-sm">
|
||||
{currentKeyData.max_budget !== null
|
||||
? `$${formatNumberWithCommas(currentKeyData.max_budget, 2)}`
|
||||
: "Unlimited"}
|
||||
</p>
|
||||
<p className="text-sm">{budgetDisplay}</p>
|
||||
</div>
|
||||
|
||||
<div>
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue