scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(key budgets): withhold the proxy-wide numbers from callers who may not read them

Browse source 
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.
This commit is contained in:
ryan-crabbe-berri 2026-08-20 19:50:34 -07:00
parent 7ae1005109
commit d0082eb644
8 changed files with 173 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
litellm/proxy/management_endpoints/key_budget_resolver.py
View file

@ -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 (

7
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -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,
),
)

9
litellm/types/proxy/management_endpoints/key_management_endpoints.py
View file

@ -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

75
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -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

19
ui/litellm-dashboard/src/components/templates/KeyBudgetsBulletChart.test.tsx
View file

@ -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(<KeyBudgetsBulletChart budgets={[TEAM_HALF, UNREADABLE_TEAM]} />);
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(<KeyBudgetsBulletChart budgets={[TEAM_HALF, RESTRICTED_PROXY]} />);
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(<KeyBudgetsBulletChart budgets={[ORG_NEARLY, MEMBER_OVER]} />);

2
ui/litellm-dashboard/src/components/templates/KeyBudgetsBulletChart.tsx
View file

@ -113,7 +113,7 @@ function Verdict({ budgets }: { budgets: readonly KeyBudgetEntry[] }) {
)}
{unknown.length > 0 && (
<p className="text-xs text-amber-600">
{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(", ")}.
</p>
)}

19
ui/litellm-dashboard/src/components/templates/KeyBudgetsTableColumns.test.ts
View file

@ -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);

22
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -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}