diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 229a8140389..493422a67b0 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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) 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 3318392e174..e9d08c090f2 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 @@ -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