fix(key budgets): restrict end user budget lookups to proxy admins

end_user_id was free-form and the lookup was a global find_unique with nothing tying the end
user to the key, its team or the caller. Any valid key could read any end user's alias, budget,
live spend and reset times, and a null entity_label told the caller whether an id existed at
all. /key/budgets made that reachable for every role, since a caller trivially satisfies the
key check by inspecting their own key.

End users are a proxy-global namespace, LiteLLM_EndUserTable carries no team or organization
column, so there is nothing to scope a non-admin against. Gate the parameter on proxy admin and
admin viewer, matching /customer/info and /customer/list. The gate is on the parameter, not the
route, so a non-admin reading their own key's budgets is unaffected.
This commit is contained in:
ryan-crabbe-berri 2026-08-19 17:15:26 -07:00
parent 1282a6d232
commit 2c608b05a6
2 changed files with 71 additions and 1 deletions

View file

@ -3753,6 +3753,12 @@ async def info_key_fn(
raise handle_exception_on_proxy(e)
_END_USER_BUDGET_READER_ROLES: Final = (
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
)
def _key_not_found_error() -> ProxyException:
return ProxyException(
message="Key not found in database",
@ -3800,7 +3806,8 @@ async def key_budgets_fn(
hash. Defaults to the key in the Authorization header when omitted (`GET /key/budgets`).
- end_user_id: str | None (query parameter) - Also report the budgets that would apply to this
end user. Omitted end users produce no `end_user` rows, because nothing binds an end user to
a key outside a request.
a key outside a request. Proxy admins only, since end users are a proxy-global namespace with
no key, team or organization scoping to check a caller against.
Returns:
- key: str - The key that was looked up, echoed back as it was passed in
@ -3839,6 +3846,7 @@ async def key_budgets_fn(
prisma_client,
proxy_logging_obj,
user_api_key_cache,
user_custom_auth,
)
try:
@ -3847,6 +3855,15 @@ async def key_budgets_fn(
"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
)
if end_user_id is not None and user_api_key_dict.user_role not in _END_USER_BUDGET_READER_ROLES:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=(
"Only proxy admins can resolve end user budgets, because end users are not scoped to a key, "
f"a team or an organization. Your role={user_api_key_dict.user_role}"
),
)
key: Final = key_id or user_api_key_dict.api_key
if key is None:
raise _key_not_found_error()
@ -3886,6 +3903,12 @@ async def key_budgets_fn(
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
general_settings=general_settings,
custom_auth_enabled=user_custom_auth is not None,
),
token_end_user_max_budget=(
user_api_key_dict.end_user_max_budget
if end_user_id is not None and end_user_id == user_api_key_dict.end_user_id
else None
),
)
return KeyBudgetsResponse(key=key, budgets=budgets)

View file

@ -17195,3 +17195,50 @@ async def test_key_budgets_treat_a_non_positive_cap_as_unset(max_budget):
assert by_scope["organization"].status == "unlimited"
assert by_scope["key_model"].max_budget is None
assert by_scope["key_model"].status == "unlimited"
@pytest.mark.parametrize(
"role",
[
LitellmUserRoles.INTERNAL_USER.value,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value,
LitellmUserRoles.TEAM.value,
LitellmUserRoles.CUSTOMER.value,
None,
],
)
@pytest.mark.asyncio
async def test_key_budgets_refuse_to_resolve_an_end_user_for_a_non_admin(role):
"""End users are proxy-global, so without this any key holder could read any customer's spend."""
caller = UserAPIKeyAuth(api_key="sk-caller", user_role=role)
with _budgets_route_world(key_row=_budgets_key_row(), caller=caller) as resolver:
response = client.get("/key/budgets?end_user_id=someone-elses-customer")
assert response.status_code == 403
assert "end user" in json.dumps(response.json()).lower()
resolver.assert_not_awaited()
@pytest.mark.parametrize(
"role",
[LitellmUserRoles.PROXY_ADMIN.value, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value],
)
@pytest.mark.asyncio
async def test_key_budgets_resolve_an_end_user_for_an_admin(role):
"""The gate has to keep admins working, or it has just deleted the parameter."""
caller = UserAPIKeyAuth(api_key="sk-admin", user_role=role)
with _budgets_route_world(key_row=_budgets_key_row(), caller=caller) as resolver:
response = client.get("/key/budgets?end_user_id=a-customer")
assert response.status_code == 200
assert resolver.await_args.kwargs["end_user_id"] == "a-customer"
@pytest.mark.asyncio
async def test_key_budgets_without_an_end_user_stay_open_to_non_admins():
"""The gate is on the parameter, not the route: a caller reading their own key must be unaffected."""
caller = UserAPIKeyAuth(api_key="sk-caller", user_role=LitellmUserRoles.INTERNAL_USER.value)
with _budgets_route_world(key_row=_budgets_key_row(), caller=caller):
response = client.get("/key/budgets")
assert response.status_code == 200