From ded69f8d0341ac474f33425fc0ba9b880f5022c2 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 22 Sep 2026 10:53:10 -0700 Subject: [PATCH] fix(ui): surface the owner's user budget on keys without their own budget (#38220) * fix(ui): surface the owner's user budget on keys without their own budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): apply the owner's budget hint to team keys when apply_user_budget_to_team_keys is on Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): read only the apply_user_budget_to_team_keys flag from general_settings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): mark the general_settings cast as cast-ok Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): load the owner's budget for keys opened outside the current page Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: jesus --- .../proxy_setting_endpoints.py | 20 +++- .../test_proxy_setting_endpoints.py | 85 +++++++++++++++ .../(dashboard)/hooks/keys/useKeyInfo.test.ts | 101 +++++++++++++++++ .../app/(dashboard)/hooks/keys/useKeyInfo.ts | 25 ++++- .../useApplyUserBudgetToTeamKeys.ts | 6 ++ .../VirtualKeysPage/VirtualKeysTable.test.tsx | 4 + .../VirtualKeysPage/VirtualKeysTable.tsx | 15 ++- .../VirtualKeysPage/keyTableColumns.tsx | 10 +- .../components/key_team_helpers/key_list.tsx | 4 +- .../shared/InheritedBudgetHint.test.tsx | 54 +++++++++- .../components/shared/InheritedBudgetHint.tsx | 29 ++++- .../key_info_view.budget_display.test.tsx | 102 +++++++++++++++++- .../components/templates/key_info_view.tsx | 7 +- 13 files changed, 447 insertions(+), 15 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyInfo.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useApplyUserBudgetToTeamKeys.ts diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index f4d4ccf5851..f37b82f3b54 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -364,11 +364,18 @@ ALLOWED_UI_SETTINGS_FIELDS: Final = { } ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: Final = "enable_ptu_cost_attribution" +APPLY_USER_BUDGET_TO_TEAM_KEYS_UI_SETTING: Final = "apply_user_budget_to_team_keys" # UI settings derived from the deployment environment. Deliberately kept out of # ALLOWED_UI_SETTINGS_FIELDS: they are read-only, never persisted, and PATCH # rejects them so an admin cannot flip an env-gated feature at runtime. -_DERIVED_UI_SETTINGS_FIELDS: Final[frozenset[str]] = frozenset({ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING}) +_DERIVED_UI_SETTINGS_FIELDS: Final[frozenset[str]] = frozenset( + {ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING, APPLY_USER_BUDGET_TO_TEAM_KEYS_UI_SETTING} +) + + +def _apply_user_budget_to_team_keys_enabled(settings: Mapping[str, object]) -> bool: + return settings.get(APPLY_USER_BUDGET_TO_TEAM_KEYS_UI_SETTING) is True def _derived_ui_setting_value(key: str) -> object: @@ -381,6 +388,12 @@ def _derived_ui_setting_value(key: str) -> object: """ if key == ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: return is_ptu_cost_attribution_enabled() + if key == APPLY_USER_BUDGET_TO_TEAM_KEYS_UI_SETTING: + from litellm.proxy.proxy_server import general_settings + + return _apply_user_budget_to_team_keys_enabled( + cast(Mapping[str, object], general_settings) # cast-ok: proxy_server declares general_settings as bare dict + ) return None @@ -776,6 +789,8 @@ def _ui_setting_source( if key == ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: configured_value: Final = get_secret(PTU_COST_ATTRIBUTION_ENV_VAR, None) return "config" if configured_value is not None or value is True else "default" + if key == APPLY_USER_BUDGET_TO_TEAM_KEYS_UI_SETTING: + return "config" if value is True else "default" return source_for(settings, key, _model_field_default(settings_class, key)) @@ -1782,6 +1797,9 @@ async def get_ui_settings(): { **resolved_settings.values, ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: is_ptu_cost_attribution_enabled(), + APPLY_USER_BUDGET_TO_TEAM_KEYS_UI_SETTING: _derived_ui_setting_value( + APPLY_USER_BUDGET_TO_TEAM_KEYS_UI_SETTING + ), } ) source: Final[Mapping[str, FieldSource]] = MappingProxyType( diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index eecee2fd0f1..9ba4975cb19 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -3670,6 +3670,91 @@ class TestPtuCostAttributionUISetting: assert not mock_prisma.db.litellm_uisettings.upsert.called +class TestApplyUserBudgetToTeamKeysUISetting: + """``apply_user_budget_to_team_keys`` mirrors general_settings on every GET. + + The proxy enforces the key owner's user budget on team keys only when + ``general_settings.apply_user_budget_to_team_keys`` is on, so the dashboard + shows the owner's budget gate on a team key iff this derived value is true. + Like the other derived settings it is read-only and never persisted. + """ + + @staticmethod + def _mock_prisma(monkeypatch, stored=None): + from unittest.mock import AsyncMock, MagicMock + + mock_prisma = MagicMock() + mock_record = None + if stored is not None: + mock_record = MagicMock() + mock_record.ui_settings = stored + mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=mock_record) + mock_prisma.db.litellm_uisettings.upsert = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + return mock_prisma + + def test_reported_false_when_general_settings_lacks_the_flag(self, mock_auth, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + self._mock_prisma(monkeypatch) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["apply_user_budget_to_team_keys"] is False + + def test_reported_true_when_general_settings_enables_the_flag(self, mock_auth, monkeypatch): + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"apply_user_budget_to_team_keys": True}, + ) + self._mock_prisma(monkeypatch) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["apply_user_budget_to_team_keys"] is True + + def test_non_json_values_elsewhere_in_general_settings_do_not_break_get(self, mock_auth, monkeypatch): + """general_settings holds non-JSON values at runtime (e.g. RoleBasedPermissions + instances under "role_permissions"); only the flag itself may be inspected.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"role_permissions": [object()], "apply_user_budget_to_team_keys": True}, + ) + self._mock_prisma(monkeypatch) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["apply_user_budget_to_team_keys"] is True + + def test_reads_only_the_flag_key(self): + from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( + _apply_user_budget_to_team_keys_enabled, + ) + + assert _apply_user_budget_to_team_keys_enabled({"apply_user_budget_to_team_keys": True}) is True + assert _apply_user_budget_to_team_keys_enabled({}) is False + assert _apply_user_budget_to_team_keys_enabled({"apply_user_budget_to_team_keys": "true"}) is False + + def test_a_persisted_true_cannot_forge_the_derived_value(self, mock_auth, monkeypatch): + """A row written before the allowlist existed must not be able to turn the feature on.""" + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + self._mock_prisma(monkeypatch, stored={"apply_user_budget_to_team_keys": True}) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["apply_user_budget_to_team_keys"] is False + + def test_is_not_an_allowlisted_persisted_setting(self): + from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( + ALLOWED_UI_SETTINGS_FIELDS, + ) + + assert "apply_user_budget_to_team_keys" not in ALLOWED_UI_SETTINGS_FIELDS + + class TestTeamAdminEditableTeamFieldsSetting: """team_admin_editable_team_fields: the proxy-wide allow-list update_team applies to team admins.""" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyInfo.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyInfo.test.ts new file mode 100644 index 00000000000..96e9cff6ac7 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyInfo.test.ts @@ -0,0 +1,101 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { renderHook, waitFor } from "@testing-library/react"; +import React, { ReactNode } from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { keyInfoV1Call, userGetInfoV2 } from "@/components/networking"; + +import { useKeyInfo } from "./useKeyInfo"; + +vi.mock("@/components/networking", () => ({ + keyInfoV1Call: vi.fn(), + userGetInfoV2: vi.fn(), +})); + +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +const mockKeyInfoV1Call = vi.mocked(keyInfoV1Call); +const mockUserGetInfoV2 = vi.mocked(userGetInfoV2); + +const createWrapper = () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + return { queryClient, wrapper }; +}; + +const KEY_ID = "sk-key-1"; +const ACCESS_TOKEN = "sk-access"; + +const OWNER = { + user_id: "user-1", + user_email: "owner@example.com", + user_alias: "Budget Owner", + user_role: "user", + spend: 0, + max_budget: 1500, + models: [], + budget_duration: "1mo", + budget_reset_at: null, + metadata: null, + created_at: null, + updated_at: null, + sso_user_id: null, + teams: [], +}; + +const keyInfoResponse = (info: Record) => ({ info }); + +describe("useKeyInfo", () => { + beforeEach(() => { + mockKeyInfoV1Call.mockReset(); + mockUserGetInfoV2.mockReset(); + mockUseAuthorized.mockReturnValue({ accessToken: ACCESS_TOKEN }); + }); + + it("attaches the owner's budget fields when the key has a user_id", async () => { + mockKeyInfoV1Call.mockResolvedValue(keyInfoResponse({ user_id: "user-1", key_alias: "team-key" })); + mockUserGetInfoV2.mockResolvedValue(OWNER); + const { wrapper } = createWrapper(); + + const { result } = renderHook(() => useKeyInfo(KEY_ID), { wrapper }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(mockUserGetInfoV2).toHaveBeenCalledWith(ACCESS_TOKEN, "user-1"); + const expectedUser = { + user_id: "user-1", + user_email: "owner@example.com", + user_alias: "Budget Owner", + max_budget: 1500, + budget_duration: "1mo", + }; + expect(result.current.data?.user).toEqual(expectedUser); + }); + + it("does not fetch an owner when the key has no user_id", async () => { + mockKeyInfoV1Call.mockResolvedValue(keyInfoResponse({ user_id: null, key_alias: "service-key" })); + const { wrapper } = createWrapper(); + + const { result } = renderHook(() => useKeyInfo(KEY_ID), { wrapper }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(mockUserGetInfoV2).not.toHaveBeenCalled(); + expect(result.current.data?.user).toBeUndefined(); + }); + + it("still resolves the key data when the owner lookup fails", async () => { + mockKeyInfoV1Call.mockResolvedValue(keyInfoResponse({ user_id: "user-1" })); + mockUserGetInfoV2.mockRejectedValue(new Error("403")); + const { wrapper } = createWrapper(); + + const { result } = renderHook(() => useKeyInfo(KEY_ID), { wrapper }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(result.current.data?.token).toBe(KEY_ID); + expect(result.current.data?.api_key).toBe(KEY_ID); + expect(result.current.data?.user).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyInfo.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyInfo.ts index 57950953d74..be248e9aafb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyInfo.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeyInfo.ts @@ -2,10 +2,25 @@ import { useQuery, UseQueryResult } from "@tanstack/react-query"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { KeyResponse } from "@/components/key_team_helpers/key_list"; -import { keyInfoV1Call } from "@/components/networking"; +import { keyInfoV1Call, userGetInfoV2 } from "@/components/networking"; import { keyKeys } from "./useKeys"; +const fetchOwner = async (accessToken: string, userId: string): Promise => { + try { + const owner = await userGetInfoV2(accessToken, userId); + return { + user_id: owner.user_id, + user_email: owner.user_email, + user_alias: owner.user_alias, + max_budget: owner.max_budget, + budget_duration: owner.budget_duration, + }; + } catch { + return undefined; + } +}; + export function useKeyInfo(keyId: string | null, options?: { enabled?: boolean }): UseQueryResult { const { accessToken } = useAuthorized(); @@ -14,10 +29,16 @@ export function useKeyInfo(keyId: string | null, options?: { enabled?: boolean } queryFn: async () => { if (!accessToken || !keyId) throw new Error("Missing access token or key id"); const keyData = await keyInfoV1Call(accessToken, keyId); + const info = keyData["info"]; + const owner = + typeof info.user_id === "string" && info.user_id !== "" + ? await fetchOwner(accessToken, info.user_id) + : undefined; return { - ...keyData["info"], + ...info, token: keyId, api_key: keyId, + ...(owner ? { user: owner } : {}), }; }, enabled: Boolean(accessToken && keyId) && (options?.enabled ?? true), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useApplyUserBudgetToTeamKeys.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useApplyUserBudgetToTeamKeys.ts new file mode 100644 index 00000000000..3aeed5e4338 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useApplyUserBudgetToTeamKeys.ts @@ -0,0 +1,6 @@ +import { useUISettings } from "./useUISettings"; + +export const APPLY_USER_BUDGET_TO_TEAM_KEYS_SETTING_KEY = "apply_user_budget_to_team_keys"; + +export const useApplyUserBudgetToTeamKeys = (): boolean => + useUISettings().data?.values?.[APPLY_USER_BUDGET_TO_TEAM_KEYS_SETTING_KEY] === true; diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 4d963a2f603..7ef7f1fcb09 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -58,6 +58,10 @@ vi.mock("@/app/(dashboard)/hooks/keys/useKeyInfo", () => ({ useKeyInfo: vi.fn(), })); +vi.mock("@/app/(dashboard)/hooks/uiSettings/useApplyUserBudgetToTeamKeys", () => ({ + useApplyUserBudgetToTeamKeys: vi.fn(() => false), +})); + vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({ default: vi.fn(), })); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index 39dd2cc5ab2..ec405ee9cad 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 { useApplyUserBudgetToTeamKeys } from "@/app/(dashboard)/hooks/uiSettings/useApplyUserBudgetToTeamKeys"; import { useAllTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { @@ -139,10 +140,18 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { const keyList = useMemo(() => keys?.keys ?? [], [keys]); const rowCount = keys?.total_count ?? 0; - const columns = useMemo( - () => getKeyTableColumns({ allTeams, organizations, onSelectKey: (key) => void setSelectedKeyId(key.token) }), - [allTeams, organizations, setSelectedKeyId], + const applyUserBudgetToTeamKeys = useApplyUserBudgetToTeamKeys(); + + const columnDeps = useMemo( + () => ({ + allTeams, + organizations, + onSelectKey: (key: KeyResponse) => void setSelectedKeyId(key.token), + applyUserBudgetToTeamKeys, + }), + [allTeams, organizations, setSelectedKeyId, applyUserBudgetToTeamKeys], ); + const columns = useMemo(() => getKeyTableColumns(columnDeps), [columnDeps]); const selectedKeyFromList = useMemo( () => keyList.find((key) => key.token === selectedKeyId), diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx index c9ee0fa1126..33625c0e1eb 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx @@ -4,7 +4,7 @@ import { Info } from "lucide-react"; import { ColumnDef } from "@tanstack/react-table"; import { DataTableMultiSortHeader, DataTableSortHeader, type DataTableSortField } from "@/components/shared/DataTable"; -import { inheritedBudgetGates } from "@/components/shared/InheritedBudgetHint"; +import { inheritedBudgetGates, keyOwnerBudgetSource } from "@/components/shared/InheritedBudgetHint"; import { HoverCard, HoverCardContent, HoverCardTrigger } from "@/components/ui/hover-card"; import { Skeleton } from "@/components/ui/skeleton"; import { @@ -86,12 +86,14 @@ interface KeyTableColumnsDeps { allTeams: Team[]; organizations: Organization[]; onSelectKey: (key: KeyResponse) => void; + applyUserBudgetToTeamKeys: boolean; } export const getKeyTableColumns = ({ allTeams, organizations, onSelectKey, + applyUserBudgetToTeamKeys, }: KeyTableColumnsDeps): ColumnDef[] => [ { id: "key_alias", @@ -277,7 +279,11 @@ export const getKeyTableColumns = ({ ); }, 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 ad357f5925d..8d39a0fd22b 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 @@ -115,8 +115,10 @@ export interface KeyResponse { next_rotation_at?: string; user?: { user_id: string; - user_email: string; + user_email: string | null; user_alias: string | null; + max_budget?: number | null; + budget_duration?: string | null; }; created_by_user?: { user_id: string; diff --git a/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.test.tsx b/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.test.tsx index 657e859aef9..8faae9c9660 100644 --- a/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.test.tsx @@ -2,7 +2,7 @@ import { render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, expect, it } from "vitest"; -import { InheritedBudgetHint, inheritedBudgetGates } from "./InheritedBudgetHint"; +import { InheritedBudgetHint, inheritedBudgetGates, keyOwnerBudgetSource } from "./InheritedBudgetHint"; const team = { team_id: "team-1", team_alias: "Platform", max_budget: 1200, budget_duration: "30d" }; const organization = { @@ -10,6 +10,13 @@ const organization = { organization_alias: "Acme", litellm_budget_table: { max_budget: 5000, budget_duration: null }, }; +const user = { + user_id: "user-1", + user_email: "owner@example.com", + user_alias: "Key Owner", + max_budget: 1500, + budget_duration: "1mo", +}; describe("inheritedBudgetGates", () => { it("returns team then org gates when both have budgets", () => { @@ -42,6 +49,45 @@ describe("inheritedBudgetGates", () => { ), ).toEqual(["team-1", "org-1"]); }); + + it("returns the owner's user budget as a gate", () => { + expect(inheritedBudgetGates(null, null, user)).toEqual([ + { scope: "User", alias: "Key Owner", maxBudget: 1500, budgetDuration: "1mo" }, + ]); + }); + + it("skips the user gate when the owner has no budget", () => { + expect(inheritedBudgetGates(null, null, { ...user, max_budget: null })).toEqual([]); + expect(inheritedBudgetGates(null, null, null)).toEqual([]); + }); + + it("falls back to email then id for the user alias", () => { + expect(inheritedBudgetGates(null, null, { ...user, user_alias: null })[0].alias).toBe("owner@example.com"); + expect(inheritedBudgetGates(null, null, { ...user, user_alias: null, user_email: null })[0].alias).toBe("user-1"); + }); + + it("lists team, org, and user gates together", () => { + expect(inheritedBudgetGates(team, organization, user).map((g) => g.scope)).toEqual([ + "Team", + "Organization", + "User", + ]); + }); +}); + +describe("keyOwnerBudgetSource", () => { + it("returns the owner on a personal key regardless of the flag", () => { + expect(keyOwnerBudgetSource({ team_id: null, user }, false)).toBe(user); + expect(keyOwnerBudgetSource({ team_id: null, user }, true)).toBe(user); + }); + + it("hides the owner on a team key when the flag is off", () => { + expect(keyOwnerBudgetSource({ team_id: "team-1", user }, false)).toBeNull(); + }); + + it("returns the owner on a team key when the flag is on", () => { + expect(keyOwnerBudgetSource({ team_id: "team-1", user }, true)).toBe(user); + }); }); describe("InheritedBudgetHint", () => { @@ -57,4 +103,10 @@ describe("InheritedBudgetHint", () => { expect(screen.getByTestId("inherited-budget-hint")).toHaveTextContent("Organization Acme: $5,000.00"); expect(screen.getByTestId("inherited-budget-hint")).not.toHaveTextContent("Organization Acme: $5,000.00 /"); }); + + it("shows the owner's user budget on hover", async () => { + render(); + await userEvent.setup().hover(screen.getByLabelText("question-circle")); + expect(screen.getByTestId("inherited-budget-hint")).toHaveTextContent("User Key Owner: $1,500.00 / 1mo"); + }); }); diff --git a/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.tsx b/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.tsx index 5fc80758cfd..4c02efa66ed 100644 --- a/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.tsx +++ b/ui/litellm-dashboard/src/components/shared/InheritedBudgetHint.tsx @@ -6,7 +6,7 @@ import type { Organization } from "@/components/networking"; import { formatNumberWithCommas } from "@/utils/dataUtils"; export interface InheritedBudgetGate { - scope: "Team" | "Organization"; + scope: "Team" | "Organization" | "User"; alias: string; maxBudget: number; budgetDuration: string | null; @@ -15,6 +15,14 @@ export interface InheritedBudgetGate { type TeamBudgetSource = Pick; type OrganizationBudgetSource = Pick; +export interface UserBudgetSource { + user_id: string; + user_alias?: string | null; + user_email?: string | null; + max_budget?: number | null; + budget_duration?: string | null; +} + const teamGate = (team: TeamBudgetSource | null | undefined): InheritedBudgetGate | null => team && team.max_budget != null ? { @@ -38,10 +46,27 @@ const organizationGate = (organization: OrganizationBudgetSource | null | undefi : null; }; +const userGate = (user: UserBudgetSource | null | undefined): InheritedBudgetGate | null => + user && user.max_budget != null + ? { + scope: "User", + alias: user.user_alias || user.user_email || user.user_id, + maxBudget: user.max_budget, + budgetDuration: user.budget_duration ?? null, + } + : null; + export const inheritedBudgetGates = ( team: TeamBudgetSource | null | undefined, organization: OrganizationBudgetSource | null | undefined, -): readonly InheritedBudgetGate[] => [teamGate(team), organizationGate(organization)].filter((gate) => gate !== null); + user?: UserBudgetSource | null, +): readonly InheritedBudgetGate[] => + [teamGate(team), organizationGate(organization), userGate(user)].filter((gate) => gate !== null); + +export const keyOwnerBudgetSource = ( + key: { team_id?: string | null; user?: UserBudgetSource | null }, + applyUserBudgetToTeamKeys: boolean, +): UserBudgetSource | null => (!key.team_id || applyUserBudgetToTeamKeys ? key.user ?? null : null); const formatGate = (gate: InheritedBudgetGate): string => `${gate.scope} ${gate.alias}: $${formatNumberWithCommas(gate.maxBudget, 2)}${gate.budgetDuration ? ` / ${gate.budgetDuration}` : ""}`; diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx index a02f80f7800..965fb3e8054 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.budget_display.test.tsx @@ -7,6 +7,7 @@ import KeyInfoView from "./key_info_view"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import useTeams from "@/app/(dashboard)/hooks/useTeams"; import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; +import { useApplyUserBudgetToTeamKeys } from "@/app/(dashboard)/hooks/uiSettings/useApplyUserBudgetToTeamKeys"; import type { Organization } from "../networking"; // IMPORTANT: do not mock `@/utils/dataUtils` here. We want to exercise the @@ -26,6 +27,9 @@ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: vi.fn() })); vi.mock("@/app/(dashboard)/hooks/projects/useProjects", () => ({ useProjects: vi.fn().mockReturnValue({ data: [], isLoading: false }), })); +vi.mock("@/app/(dashboard)/hooks/uiSettings/useApplyUserBudgetToTeamKeys", () => ({ + useApplyUserBudgetToTeamKeys: vi.fn(() => false), +})); vi.mock("@/app/(dashboard)/hooks/keys/useResetKeySpend", () => ({ useResetKeySpend: vi.fn(() => ({ mutate: vi.fn(), isPending: false })), })); @@ -114,6 +118,26 @@ const baseAuthorized = { isAuthorized: true, }; +const TEST_BUDGET_TEAM_FIELDS = { + team_id: "team-123", + team_alias: "Test Budget", + max_budget: 1200, + budget_duration: "30d", +}; + +const TEAM_KEY_WITH_BUDGETED_OWNER = { + ...MOCK_KEY_DATA, + max_budget: null, + team_id: "team-123", + user: { + user_id: "user-1", + user_email: "owner@example.com", + user_alias: "Budget Owner", + max_budget: 1500, + budget_duration: "1mo", + }, +} as unknown as KeyResponse; + const makeTeam = (overrides: Partial): Team => ({ team_id: "team-default", team_alias: "Default Team", @@ -158,6 +182,7 @@ describe("KeyInfoView overview budget display (LIT-2845)", () => { vi.mocked(useTeams).mockReturnValue({ teams: [], setTeams: vi.fn() }); vi.mocked(useAuthorized).mockReturnValue(baseAuthorized); mockOrganizations([]); + vi.mocked(useApplyUserBudgetToTeamKeys).mockReturnValue(false); }); it("renders a sub-dollar max_budget ($0.10) with 2-decimal precision in the overview Spend card", async () => { @@ -214,7 +239,7 @@ describe("KeyInfoView overview budget display (LIT-2845)", () => { it("never pairs key spend with the team budget: shows Unlimited plus an inherited-budget hint", async () => { vi.mocked(useTeams).mockReturnValue({ - teams: [makeTeam({ team_id: "team-123", team_alias: "Test Budget", max_budget: 1200, budget_duration: "30d" })], + teams: [makeTeam(TEST_BUDGET_TEAM_FIELDS)], setTeams: vi.fn(), }); renderWithProviders( @@ -258,6 +283,81 @@ describe("KeyInfoView overview budget display (LIT-2845)", () => { expect(screen.getByTestId("inherited-budget-hint")).not.toHaveTextContent("Team Org Team"); }); + it("lists the owner's user budget in the hint for a personal key with no budget of its own", async () => { + renderWithProviders( + {}} + keyId={"test-key-id"} + onKeyDataUpdate={() => {}} + teams={[]} + />, + ); + await waitFor(() => { + expect(screen.getByText(/of Unlimited/)).toBeInTheDocument(); + }); + await userEvent.setup().hover(screen.getByLabelText("question-circle")); + expect(screen.getByTestId("inherited-budget-hint")).toHaveTextContent("User Budget Owner: $1,500.00 / 1mo"); + }); + + it("omits the owner's user budget from the hint for a team key when apply_user_budget_to_team_keys is off", async () => { + vi.mocked(useTeams).mockReturnValue({ + teams: [makeTeam(TEST_BUDGET_TEAM_FIELDS)], + setTeams: vi.fn(), + }); + renderWithProviders( + {}} + keyId={"test-key-id"} + onKeyDataUpdate={() => {}} + teams={[]} + />, + ); + await waitFor(() => { + expect(screen.getByText(/of Unlimited/)).toBeInTheDocument(); + }); + await userEvent.setup().hover(screen.getByLabelText("question-circle")); + expect(screen.getByTestId("inherited-budget-hint")).toHaveTextContent("Team Test Budget: $1,200.00 / 30d"); + expect(screen.getByTestId("inherited-budget-hint")).not.toHaveTextContent("User Budget Owner"); + }); + + it("lists the owner's user budget in the hint for a team key when apply_user_budget_to_team_keys is on", async () => { + vi.mocked(useApplyUserBudgetToTeamKeys).mockReturnValue(true); + vi.mocked(useTeams).mockReturnValue({ + teams: [makeTeam(TEST_BUDGET_TEAM_FIELDS)], + setTeams: vi.fn(), + }); + renderWithProviders( + {}} + keyId={"test-key-id"} + onKeyDataUpdate={() => {}} + teams={[]} + />, + ); + await waitFor(() => { + expect(screen.getByText(/of Unlimited/)).toBeInTheDocument(); + }); + await userEvent.setup().hover(screen.getByLabelText("question-circle")); + expect(screen.getByTestId("inherited-budget-hint")).toHaveTextContent("Team Test Budget: $1,200.00 / 30d"); + expect(screen.getByTestId("inherited-budget-hint")).toHaveTextContent("User Budget Owner: $1,500.00 / 1mo"); + }); + it("renders 'Unlimited' with no hint when neither key, team, nor org has a budget", async () => { vi.mocked(useTeams).mockReturnValue({ teams: [makeTeam({ team_id: "team-789", team_alias: "Free Team" })], 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 06e088f956b..3c2c9bb352d 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -1,6 +1,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useProjects } from "@/app/(dashboard)/hooks/projects/useProjects"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; +import { useApplyUserBudgetToTeamKeys } from "@/app/(dashboard)/hooks/uiSettings/useApplyUserBudgetToTeamKeys"; import useTeams from "@/app/(dashboard)/hooks/useTeams"; import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { formatNumberWithCommas } from "@/utils/dataUtils"; @@ -45,7 +46,7 @@ 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 } from "../shared/InheritedBudgetHint"; +import { InheritedBudgetHint, inheritedBudgetGates, keyOwnerBudgetSource } from "../shared/InheritedBudgetHint"; import { KeyEditView } from "./key_edit_view"; interface KeyInfoViewProps { @@ -95,6 +96,7 @@ export default function KeyInfoView({ const { data: organizations } = useOrganizations(); const { data: projects } = useProjects(); const { data: uiSettingsData } = useUISettings(); + const applyUserBudgetToTeamKeys = useApplyUserBudgetToTeamKeys(); const { data: allMcpServers } = useMCPServers(); const { data: allMcpToolsets } = useMCPToolsets(); const enableProjectsUI = Boolean(uiSettingsData?.values?.enable_projects_ui); @@ -507,7 +509,8 @@ export default function KeyInfoView({ const hasOwnBudget = currentKeyData.max_budget !== null; const budgetDisplay = hasOwnBudget ? `$${formatNumberWithCommas(currentKeyData.max_budget, 2)}` : "Unlimited"; - const inheritedGates = hasOwnBudget ? [] : inheritedBudgetGates(parentTeam, parentOrg); + const ownerUser = keyOwnerBudgetSource(currentKeyData, applyUserBudgetToTeamKeys); + const inheritedGates = hasOwnBudget ? [] : inheritedBudgetGates(parentTeam, parentOrg, ownerUser); return (