From d0082eb6447254017f63f709b15f5f81996fcc4e Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 20 Aug 2026 19:50:34 -0700 Subject: [PATCH] fix(key budgets): withhold the proxy-wide numbers from callers who may not read them The proxy row's limit and spend cover the whole deployment, not the key being inspected, and every proxy-wide spend route is admin-only. Reporting them on a route any key holder can call handed one tenant the total spend of all of them. The row still ships, because dropping it would read as "no proxy budget applies", which is the guess this endpoint exists to remove. Its numbers are blanked, `spend_state` gains `restricted` so a blank is never mistaken for a failed read or a zero, and the row reports unknown: a caller who cannot see the numbers cannot rule the scope out. The proxy budget row is not read from the database at all for those callers. --- .../key_budget_resolver.py | 44 ++++++++++- .../key_management_endpoints.py | 7 +- .../key_management_endpoints.py | 9 ++- .../test_key_management_endpoints.py | 75 ++++++++++++++++++- .../templates/KeyBudgetsBulletChart.test.tsx | 19 ++++- .../templates/KeyBudgetsBulletChart.tsx | 2 +- .../templates/KeyBudgetsTableColumns.test.ts | 19 +++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 22 +++--- 8 files changed, 173 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_budget_resolver.py b/litellm/proxy/management_endpoints/key_budget_resolver.py index 72596626da3..27e67ad1c8e 100644 --- a/litellm/proxy/management_endpoints/key_budget_resolver.py +++ b/litellm/proxy/management_endpoints/key_budget_resolver.py @@ -106,6 +106,14 @@ _RESERVATION_NOTE: Final = KeyBudgetNote( "read-time check; requests it cannot price up front are still gated by the read-time check alone" ), ) +_PROXY_RESTRICTED_NOTE: Final = KeyBudgetNote( + code="proxy_spend_restricted", + severity="warning", + text=( + "the proxy-wide budget applies to this key, but its limit and spend cover the whole deployment " + "and are only reported to proxy admins, so this row cannot be ruled out from here" + ), +) _PROJECT_SPEND_NOTE: Final = KeyBudgetNote( code="project_spend_not_tracked", severity="warning", @@ -204,6 +212,7 @@ class KeyBudgetResolverDeps: user_api_key_cache: UserApiKeyCache proxy_logging_obj: ProxyLogging general_settings: Mapping[str, object] + proxy_spend_visible: bool custom_auth_enabled: bool = False read_spend: SpendReader = field(default=_read_counter_spend) @@ -236,7 +245,12 @@ class _RecordedSpend: value: float -_SpendSource = _CounterSpend | _RecordedSpend | _UnknownSpend +@dataclass(frozen=True, slots=True) +class _RestrictedSpend: + """Stands in for a number the caller may not read, which is not a number nobody could read.""" + + +_SpendSource = _CounterSpend | _RecordedSpend | _UnknownSpend | _RestrictedSpend @dataclass(frozen=True, slots=True) @@ -352,6 +366,7 @@ class _KeyBudgetContext: valid_token: UserAPIKeyAuth token_inputs: _TokenBudgetInputs end_user_id: str | None + proxy_spend_visible: bool custom_auth_enabled: bool custom_auth_skips_checks: bool general_settings: Mapping[str, object] @@ -396,6 +411,8 @@ async def _read_spend(plan: _PlannedBudget, deps: KeyBudgetResolverDeps) -> _Spe match source: case _UnknownSpend(): return _SpendReading(value=None, state="unavailable") + case _RestrictedSpend(): + return _SpendReading(value=None, state="restricted") case _RecordedSpend(): return _SpendReading(value=source.value, state="live") case _CounterSpend(): @@ -446,10 +463,11 @@ def _status(plan: _PlannedBudget, spend: float | None, exceeded: bool) -> Budget """ A row nobody could evaluate reports ``unknown``, never ``unlimited``. - Reading it as unset is the failure this endpoint exists to prevent: an unreadable entity and an - unreadable counter both leave a budget that may well be the one blocking the key. + Reading it as unset is the failure this endpoint exists to prevent: an unreadable entity, an + unreadable counter and a number withheld from the caller all leave a budget that may well be the + one blocking the key. """ - if isinstance(plan.spend_source, _UnknownSpend): + if isinstance(plan.spend_source, _UnknownSpend | _RestrictedSpend): return "unknown" if plan.max_budget is None: return "unlimited" @@ -529,6 +547,7 @@ async def _load_context( valid_token=valid_token, token_inputs=token_inputs, end_user_id=end_user_id, + proxy_spend_visible=deps.proxy_spend_visible, custom_auth_enabled=deps.custom_auth_enabled, custom_auth_skips_checks=( deps.custom_auth_enabled and deps.general_settings.get("custom_auth_run_common_checks") is not True @@ -659,6 +678,8 @@ async def _load_budget_meta( async def _load_proxy_budget(deps: KeyBudgetResolverDeps) -> _ProxyBudget | _Unavailable | None: + if not deps.proxy_spend_visible: + return None try: row: Final = await UserRepository(deps.prisma_client).find_by_id(LITELLM_PROXY_BUDGET_NAME) except Exception: # noqa: BLE001 # every entity load degrades to "unknown" rather than failing the report @@ -862,8 +883,23 @@ def _budget_meta(context: _KeyBudgetContext, budget_id: str | None) -> _BudgetRo def _plan_proxy(context: _KeyBudgetContext) -> tuple[_PlannedBudget, ...]: + """The proxy budget is the whole deployment's spend, so its numbers are not every key holder's.""" proxy: Final = context.proxy max_budget: Final = litellm.max_budget if litellm.max_budget > 0 else None + if not context.proxy_spend_visible: + return ( + _PlannedBudget( + scope="proxy", + entity_id=None, + entity_label=None, + enforcement="hard", + max_budget=None, + comparison=">", + source="litellm_settings.max_budget", + spend_source=_RestrictedSpend(), + notes=(_PROXY_RESTRICTED_NOTE,), + ), + ) match proxy: case _Unavailable(): return ( diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 819089eec81..193e63f7ee1 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3880,8 +3880,10 @@ async def key_budgets_fn( but places no limit on it - spend: float | None - Spend as the enforcing check reads it, from the same cross-pod counter, not the periodically-synced database column. `null` only when the read failed - - spend_state: str - Whether `spend` was read (`live`) or is missing because the entity or its - counter could not be read (`unavailable`) + - spend_state: str - Whether `spend` was read (`live`), is missing because the entity or its + counter could not be read (`unavailable`), or is withheld from this caller (`restricted`). + The proxy-wide row is `restricted` for everyone but a proxy admin, since its limit and spend + cover the whole deployment rather than this key - remaining: float | None - `max_budget - spend`, when both are known - comparison: str - The operator the enforcing check uses, which differs per scope - budget_duration / budget_reset_at / window_start: When spend next resets to zero @@ -3985,6 +3987,7 @@ async def key_budgets_fn( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, general_settings=general_settings, + proxy_spend_visible=user_api_key_dict.user_role in _PROXY_WIDE_READER_ROLES, custom_auth_enabled=user_custom_auth is not None, ), ) diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 5a4dfc818ce..e14a14e8528 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -136,6 +136,7 @@ BudgetNoteCode = Literal[ "end_user_route_only", "entity_unavailable", "project_spend_not_tracked", + "proxy_spend_restricted", "request_tags_add_budgets", "reservation_blocks_at_limit", "rolling_window", @@ -145,7 +146,7 @@ BudgetNoteCode = Literal[ BudgetNoteSeverity = Literal["info", "warning"] -BudgetSpendState = Literal["live", "unavailable"] +BudgetSpendState = Literal["live", "unavailable", "restricted"] class KeyBudgetNote(BaseModel): @@ -171,9 +172,9 @@ class KeyBudgetEntry(BaseModel): """ One budget that can gate requests made with a key, with its live spend. - ``status`` is ``unknown`` when the row could not be evaluated, either because the entity behind it - was unreadable or because its spend was, and it is never ``unlimited`` in that case: an unreadable - scope is not a scope the reader may rule out. + ``status`` is ``unknown`` when the row could not be evaluated, whether because the entity behind + it was unreadable, because its spend was, or because the caller may not read the numbers, and it + is never ``unlimited`` in that case: a scope nobody could evaluate is not one to rule out. """ scope: BudgetScope diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 5356da63b8e..b501cd0d110 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -16057,12 +16057,13 @@ class _RecordingSpendReader: return self.spend_by_counter_key.get(counter_key, 0.0) -def _budgets_deps(read_spend=None, general_settings=None, custom_auth_enabled=False): +def _budgets_deps(read_spend=None, general_settings=None, custom_auth_enabled=False, proxy_spend_visible=True): return KeyBudgetResolverDeps( prisma_client=MagicMock(), user_api_key_cache=MagicMock(), proxy_logging_obj=MagicMock(), general_settings=general_settings if general_settings is not None else {}, + proxy_spend_visible=proxy_spend_visible, custom_auth_enabled=custom_auth_enabled, read_spend=read_spend or _RecordingSpendReader({}), ) @@ -16978,8 +16979,10 @@ def test_key_budgets_classify_every_note_code_and_leave_none_to_a_default(): "end_user_route_only": "warning", "project_spend_not_tracked": "warning", "request_tags_add_budgets": "warning", - # `status` carries it too, but only as a value no client built before this endpoint can know + # `status` and `spend_state` carry these too, but only as values no client built before this + # endpoint can know "entity_unavailable": "warning", + "proxy_spend_restricted": "warning", } @@ -17255,3 +17258,71 @@ async def test_key_budgets_never_send_an_admin_through_a_teams_permission_list(r response = client.get(f"/key/{_BUDGETS_KEY_HASH}/budgets") assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_key_budgets_withhold_the_proxy_wide_numbers_from_a_caller_who_may_not_read_them(): + """ + The proxy row's limit and spend are the whole deployment's, not this key's. + + Every proxy-wide spend route is admin-only, so reporting them on a route any key holder can call + would hand one tenant the total spend of all of them. + """ + with _budgets_world(**_fully_populated_world()), patch.object(litellm, "max_budget", 1000.0): + budgets = await resolve_key_budgets( + valid_token=_budgets_token(), end_user_id=None, deps=_budgets_deps(proxy_spend_visible=False) + ) + + entry = next(e for e in budgets if e.scope == "proxy") + assert (entry.max_budget, entry.spend, entry.remaining) == (None, None, None) + assert (entry.spend_state, entry.status) == ("restricted", "unknown") + assert [note.code for note in entry.notes] == ["proxy_spend_restricted"] + + +@pytest.mark.asyncio +async def test_key_budgets_still_report_the_proxy_row_to_a_caller_who_may_not_read_its_numbers(): + """Dropping the row would read as "no proxy budget applies", which is the guess this route removes.""" + with _budgets_world(**_fully_populated_world()), patch.object(litellm, "max_budget", 1000.0): + budgets = await resolve_key_budgets( + valid_token=_budgets_token(), end_user_id=None, deps=_budgets_deps(proxy_spend_visible=False) + ) + + assert [e.scope for e in budgets].count("proxy") == 1 + assert next(e for e in budgets if e.scope == "proxy").status != "unlimited" + + +@pytest.mark.asyncio +async def test_key_budgets_never_read_the_proxy_budget_row_for_a_caller_who_may_not_see_it(): + """Withholding a number the report already fetched still logs the read and still costs the query.""" + user_repository = MagicMock() + user_repository.return_value.find_by_id = AsyncMock(return_value=None) + with ( + _budgets_world(**_fully_populated_world()), + patch(f"{_BUDGETS_RESOLVER}.UserRepository", user_repository), + ): + await resolve_key_budgets( + valid_token=_budgets_token(), end_user_id=None, deps=_budgets_deps(proxy_spend_visible=False) + ) + + user_repository.return_value.find_by_id.assert_not_awaited() + + +@pytest.mark.parametrize( + ("role", "expected"), + [ + (LitellmUserRoles.PROXY_ADMIN.value, True), + (LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, True), + (LitellmUserRoles.INTERNAL_USER.value, False), + (LitellmUserRoles.TEAM.value, False), + (None, False), + ], +) +@pytest.mark.asyncio +async def test_key_budgets_route_shows_the_proxy_numbers_only_to_a_proxy_wide_reader(role, expected): + """The gate has to arrive from the request, or the resolver decides it from nothing.""" + caller = UserAPIKeyAuth(api_key=_BUDGETS_KEY_HASH, user_id="user-budgets", user_role=role) + with _budgets_route_world(key_row=_budgets_key_row(), caller=caller) as resolver: + response = client.get("/key/budgets") + + assert response.status_code == 200 + assert resolver.await_args.kwargs["deps"].proxy_spend_visible is expected diff --git a/ui/litellm-dashboard/src/components/templates/KeyBudgetsBulletChart.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyBudgetsBulletChart.test.tsx index 9b6f9ebc368..23034c328c8 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyBudgetsBulletChart.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyBudgetsBulletChart.test.tsx @@ -51,6 +51,14 @@ const UNREADABLE_TEAM: KeyBudgetEntry = { status: "unknown", notes: [{ code: "entity_unavailable", severity: "warning", text: "could not read" }], }; +const RESTRICTED_PROXY: KeyBudgetEntry = { + ...OK, + scope: "proxy", + spend: null, + spend_state: "restricted", + status: "unknown", + notes: [{ code: "proxy_spend_restricted", severity: "warning", text: "admins only" }], +}; const DEAD_PROJECT: KeyBudgetEntry = { ...OK, scope: "project", @@ -108,13 +116,20 @@ describe("KeyBudgetsBulletChart", () => { expect(screen.queryByTestId("key-budget-bullet-blocking")).not.toBeInTheDocument(); }); - it("names a scope nobody could read, since a verdict that ignores it is a verdict ruling it out", () => { + it("names a scope nobody could evaluate, since a verdict that ignores it is a verdict ruling it out", () => { render(); - expect(verdict()).toHaveTextContent("1 scope could not be read"); + expect(verdict()).toHaveTextContent("1 scope could not be evaluated"); expect(verdict()).toHaveTextContent("Team"); }); + it("names a scope whose numbers this caller may not read, which is no more ruled out than a failed read", () => { + render(); + + expect(verdict()).toHaveTextContent("1 scope could not be evaluated"); + expect(verdict()).toHaveTextContent("Proxy"); + }); + it("draws each bar in proportion to its budget, and never past the end of its track", () => { render(); diff --git a/ui/litellm-dashboard/src/components/templates/KeyBudgetsBulletChart.tsx b/ui/litellm-dashboard/src/components/templates/KeyBudgetsBulletChart.tsx index 58a3cfa46de..c600a486a14 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyBudgetsBulletChart.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyBudgetsBulletChart.tsx @@ -113,7 +113,7 @@ function Verdict({ budgets }: { budgets: readonly KeyBudgetEntry[] }) { )} {unknown.length > 0 && (

- {unknown.length} {unknown.length === 1 ? "scope" : "scopes"} could not be read, so nothing on{" "} + {unknown.length} {unknown.length === 1 ? "scope" : "scopes"} could not be evaluated, so nothing on{" "} {unknown.length === 1 ? "it" : "them"} can be ruled out: {unknown.map(scopeLabel).join(", ")}.

)} diff --git a/ui/litellm-dashboard/src/components/templates/KeyBudgetsTableColumns.test.ts b/ui/litellm-dashboard/src/components/templates/KeyBudgetsTableColumns.test.ts index ca9be5c91d9..e4455062e58 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyBudgetsTableColumns.test.ts +++ b/ui/litellm-dashboard/src/components/templates/KeyBudgetsTableColumns.test.ts @@ -275,6 +275,25 @@ const UNRESOLVED: KeyBudgetEntry = { notes: [ENTITY_UNAVAILABLE_NOTE], }; +// The proxy budget applies, but its limit and spend are the whole deployment's, so a caller who is +// not a proxy admin gets the row without the numbers. Same shape as a failed read, same treatment. +const RESTRICTED: KeyBudgetEntry = { + ...UNRESOLVED, + scope: "proxy", + entity_id: null, + spend_state: "restricted", + source: "litellm_settings.max_budget", + notes: [], +}; + +describe("a scope whose numbers the caller may not read", () => { + it("is treated as unknown rather than as a confident zero or an unlimited scope", () => { + expect(cannotTrip(RESTRICTED)).toBe(false); + expect(isBlockingRow(RESTRICTED)).toBe(false); + expect(rowRank(RESTRICTED)).toStrictEqual(rowRank(UNRESOLVED)); + }); +}); + describe("a scope the server could not resolve", () => { it("is not dead, since the budget it hides can be the one denying every request", () => { expect(cannotTrip(UNRESOLVED)).toBe(false); diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 72a72a18978..cc342551aa6 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -6841,8 +6841,10 @@ export interface paths { * but places no limit on it * - spend: float | None - Spend as the enforcing check reads it, from the same cross-pod * counter, not the periodically-synced database column. `null` only when the read failed - * - spend_state: str - Whether `spend` was read (`live`) or is missing because the entity or its - * counter could not be read (`unavailable`) + * - spend_state: str - Whether `spend` was read (`live`), is missing because the entity or its + * counter could not be read (`unavailable`), or is withheld from this caller (`restricted`). + * The proxy-wide row is `restricted` for everyone but a proxy admin, since its limit and spend + * cover the whole deployment rather than this key * - remaining: float | None - `max_budget - spend`, when both are known * - comparison: str - The operator the enforcing check uses, which differs per scope * - budget_duration / budget_reset_at / window_start: When spend next resets to zero @@ -7541,8 +7543,10 @@ export interface paths { * but places no limit on it * - spend: float | None - Spend as the enforcing check reads it, from the same cross-pod * counter, not the periodically-synced database column. `null` only when the read failed - * - spend_state: str - Whether `spend` was read (`live`) or is missing because the entity or its - * counter could not be read (`unavailable`) + * - spend_state: str - Whether `spend` was read (`live`), is missing because the entity or its + * counter could not be read (`unavailable`), or is withheld from this caller (`restricted`). + * The proxy-wide row is `restricted` for everyone but a proxy admin, since its limit and spend + * cover the whole deployment rather than this key * - remaining: float | None - `max_budget - spend`, when both are known * - comparison: str - The operator the enforcing check uses, which differs per scope * - budget_duration / budget_reset_at / window_start: When spend next resets to zero @@ -26365,9 +26369,9 @@ export interface components { * KeyBudgetEntry * @description One budget that can gate requests made with a key, with its live spend. * - * ``status`` is ``unknown`` when the row could not be evaluated, either because the entity behind it - * was unreadable or because its spend was, and it is never ``unlimited`` in that case: an unreadable - * scope is not a scope the reader may rule out. + * ``status`` is ``unknown`` when the row could not be evaluated, whether because the entity behind + * it was unreadable, because its spend was, or because the caller may not read the numbers, and it + * is never ``unlimited`` in that case: a scope nobody could evaluate is not one to rule out. */ KeyBudgetEntry: { /** Budget Duration */ @@ -26411,7 +26415,7 @@ export interface components { * Spend State * @enum {string} */ - spend_state: "live" | "unavailable"; + spend_state: "live" | "unavailable" | "restricted"; /** * Status * @enum {string} @@ -26436,7 +26440,7 @@ export interface components { * Code * @enum {string} */ - code: "custom_auth_may_override_end_user_cap" | "custom_auth_skips_read_time_checks" | "end_user_route_only" | "entity_unavailable" | "project_spend_not_tracked" | "request_tags_add_budgets" | "reservation_blocks_at_limit" | "rolling_window" | "throttled_instead_of_blocked" | "user_budget_not_applied_to_team_key"; + code: "custom_auth_may_override_end_user_cap" | "custom_auth_skips_read_time_checks" | "end_user_route_only" | "entity_unavailable" | "project_spend_not_tracked" | "proxy_spend_restricted" | "request_tags_add_budgets" | "reservation_blocks_at_limit" | "rolling_window" | "throttled_instead_of_blocked" | "user_budget_not_applied_to_team_key"; /** * Severity * @enum {string}