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}