diff --git a/litellm/proxy/management_endpoints/key_budget_resolver.py b/litellm/proxy/management_endpoints/key_budget_resolver.py index 0428de74221..72596626da3 100644 --- a/litellm/proxy/management_endpoints/key_budget_resolver.py +++ b/litellm/proxy/management_endpoints/key_budget_resolver.py @@ -1136,6 +1136,13 @@ def _plan_tag(context: _KeyBudgetContext, tag_name: str, tag: LiteLLM_TagTable | def _plan_end_user(context: _KeyBudgetContext) -> tuple[_PlannedBudget, ...]: + """ + An end user with no row still gets `max_end_user_budget_id`; one that could not be read gets nothing. + + `get_end_user_object` swallows its own read failures into the same None auth falls back on, so the + default cap below covers a failed read too. Reaching `unresolved` takes a raise, and both auth call + sites wrap that fallback in the `try` a raise escapes, leaving the request with no end user cap. + """ end_user_id: Final = context.end_user_id if end_user_id is None or "end_user" in context.unresolved: return () diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index e99e49ddf28..819089eec81 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3780,12 +3780,40 @@ async def info_key_fn( raise handle_exception_on_proxy(e) -_END_USER_BUDGET_READER_ROLES: Final = ( +_PROXY_WIDE_READER_ROLES: Final = ( LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, ) +async def _enforce_team_member_budget_permission( + user_api_key_dict: UserAPIKeyAuth, + key: str, + key_info: LiteLLM_VerificationToken, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> None: + """ + Belonging to a key's team is what grants /key/info; these budgets are their own permission. + + A team key's budgets carry the team's, the caller's own membership, the owning user's, the + organization's and the key's tag spend, so a member the team has not granted this route cannot + read another member's key through it. Proxy admins, the calling key itself and keys the caller + owns never reach the team's permission list. + """ + if user_api_key_dict.user_role in _PROXY_WIDE_READER_ROLES: + return + if user_api_key_dict.api_key == key or key_info.user_id == user_api_key_dict.user_id: + return + await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( + user_api_key_dict=user_api_key_dict, + route=KeyManagementRoutes.KEY_BUDGETS, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + existing_key_row=key_info, + ) + + def _key_not_found_error() -> ProxyException: return ProxyException( message="Key not found in database", @@ -3891,7 +3919,7 @@ 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: + if end_user_id is not None and user_api_key_dict.user_role not in _PROXY_WIDE_READER_ROLES: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=( @@ -3932,6 +3960,14 @@ async def key_budgets_fn( detail=f"You are not allowed to access this key's info. Your role={user_api_key_dict.user_role}", ) + await _enforce_team_member_budget_permission( + user_api_key_dict=user_api_key_dict, + key=key, + key_info=key_info, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + # The same object auth resolves the key to, so a stale cached limit is reported as the limit # that will actually be enforced rather than the database value that will not be. resolved_key: Final = await get_key_object( 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 75e8dd2414c..5356da63b8e 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 @@ -18,6 +18,7 @@ import inspect from litellm.proxy._types import ( GenerateKeyRequest, + KeyManagementRoutes, NewUserRequest, LiteLLM_BudgetTable, LiteLLM_OrganizationTable, @@ -16670,6 +16671,51 @@ async def test_key_budgets_warn_that_custom_auth_can_hide_an_end_user_cap(): assert "custom_auth_may_override_end_user_cap" in [note.code for note in warned.notes] +@pytest.mark.asyncio +async def test_key_budgets_report_the_default_end_user_cap_when_the_end_user_has_no_row(): + """ + `max_end_user_budget_id` caps an end user auth never found a row for, and auth reaches that + fallback whenever the lookup comes back empty, including when it came back empty because the + read failed. Leaving it out would report an uncapped end user that requests still get blocked on. + """ + world = _fully_populated_world( + end_user=None, + default_end_user_budget=LiteLLM_BudgetTable(budget_id="budget-default-end-user", max_budget=42.0), + ) + with _budgets_world(**world): + budgets = await resolve_key_budgets( + valid_token=_budgets_token(), end_user_id="unseen-end-user", deps=_budgets_deps() + ) + + entry = next(e for e in budgets if e.scope == "end_user") + assert entry.max_budget == 42.0 + assert entry.source == "litellm_settings.max_end_user_budget_id" + assert entry.status == "ok" + + +@pytest.mark.asyncio +async def test_key_budgets_withhold_a_default_end_user_cap_that_auth_would_not_have_applied(): + """ + Both auth call sites read the end user inside a `try` that swallows, and the branch falling back + to `max_end_user_budget_id` sits inside that same `try`. A lookup that raises therefore leaves + the request with no end user cap at all, so naming the default here would invent an enforcer. + """ + + async def _explode(**kwargs): + raise RuntimeError("end user lookup is down") + + world = _fully_populated_world( + default_end_user_budget=LiteLLM_BudgetTable(budget_id="budget-default-end-user", max_budget=42.0), + ) + with _budgets_world(**world), patch(f"{_BUDGETS_RESOLVER}.get_end_user_object", _explode): + budgets = await resolve_key_budgets( + valid_token=_budgets_token(), end_user_id="end-user-budgets", deps=_budgets_deps() + ) + + entry = next(e for e in budgets if e.scope == "end_user") + assert (entry.status, entry.max_budget, entry.spend) == ("unknown", None, None) + + @pytest.mark.parametrize("max_budget", [0, 0.0, -5.0]) @pytest.mark.asyncio async def test_key_budgets_treat_a_non_positive_cap_as_unset(max_budget): @@ -17098,3 +17144,114 @@ async def test_key_budgets_without_an_end_user_stay_open_to_non_admins(): response = client.get("/key/budgets") assert response.status_code == 200 + + +@_budgets_contextlib.contextmanager +def _budgets_team_world(*, granted_routes, member_role="user", member_id="team-mate"): + team_table = LiteLLM_TeamTableCachedObj( + team_id="team-budgets", + members_with_roles=[Member(user_id=member_id, role=member_role)], + team_member_permissions=granted_routes, + ) + with patch( + "litellm.proxy.management_helpers.team_member_permission_checks.get_team_object", + new_callable=AsyncMock, + return_value=team_table, + ): + yield + + +@pytest.mark.asyncio +async def test_key_budgets_refuse_a_team_mates_key_to_a_member_without_the_budgets_permission(): + """ + Belonging to the team is what /key/info grants, and this route reports far more than /key/info. + + A member the team has not granted the route would otherwise read the team's, the organization's, + the key owner's and the key's tag spend off someone else's key. + """ + caller = UserAPIKeyAuth( + api_key="sk-team-mate", + user_id="team-mate", + team_id="team-budgets", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + with ( + _budgets_team_world(granted_routes=[]), + _budgets_route_world(key_row=_budgets_key_row(), caller=caller) as resolver, + ): + response = client.get(f"/key/{_BUDGETS_KEY_HASH}/budgets") + + assert response.status_code == 401 + assert KeyManagementRoutes.KEY_BUDGETS.value in json.dumps(response.json()) + resolver.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_key_budgets_allow_a_team_mates_key_once_the_team_grants_the_route(): + """Denying a granted member would leave the permission listed but impossible to use.""" + caller = UserAPIKeyAuth( + api_key="sk-team-mate", + user_id="team-mate", + team_id="team-budgets", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + with ( + _budgets_team_world(granted_routes=[KeyManagementRoutes.KEY_BUDGETS.value]), + _budgets_route_world(key_row=_budgets_key_row(), caller=caller) as resolver, + ): + response = client.get(f"/key/{_BUDGETS_KEY_HASH}/budgets") + + assert response.status_code == 200 + resolver.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_key_budgets_let_a_team_admin_read_a_team_key_without_an_explicit_grant(): + caller = UserAPIKeyAuth( + api_key="sk-team-admin", + user_id="team-mate", + team_id="team-budgets", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + with ( + _budgets_team_world(granted_routes=[], member_role="admin"), + _budgets_route_world(key_row=_budgets_key_row(), caller=caller), + ): + response = client.get(f"/key/{_BUDGETS_KEY_HASH}/budgets") + + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_key_budgets_keep_a_member_reading_the_key_they_own_without_the_permission(): + """The permission is about reading someone else's key, so it must not lock a member out of theirs.""" + caller = UserAPIKeyAuth( + api_key="sk-owner", + user_id="user-budgets", + team_id="team-budgets", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + with ( + _budgets_team_world(granted_routes=[], member_id="user-budgets"), + _budgets_route_world(key_row=_budgets_key_row(), caller=caller), + ): + response = client.get(f"/key/{_BUDGETS_KEY_HASH}/budgets") + + assert response.status_code == 200 + + +@pytest.mark.parametrize( + "role", + [LitellmUserRoles.PROXY_ADMIN.value, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value], +) +@pytest.mark.asyncio +async def test_key_budgets_never_send_an_admin_through_a_teams_permission_list(role): + """An admin is in no team, so a membership-derived permission list would deny every admin.""" + caller = UserAPIKeyAuth(api_key="sk-admin", user_id="an-admin", user_role=role) + with ( + _budgets_team_world(granted_routes=[]), + _budgets_route_world(key_row=_budgets_key_row(), caller=caller), + ): + response = client.get(f"/key/{_BUDGETS_KEY_HASH}/budgets") + + assert response.status_code == 200