diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index c23c899332..886b70e32f 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1609,13 +1609,15 @@ async def _run_centralized_common_checks( model-access, budgets, guardrails, org, and vector-store checks. Invariants: - - ``PROXY_ADMIN`` tokens short-circuit (admins bypass these checks - today; preserving that preserves behavior and avoids five DB - fetches per admin request). - ``user_custom_auth`` with ``custom_auth_run_common_checks`` unset skips the gate — matches the existing custom-auth RPS guarantee. Custom-auth deployments don't use OAuth2 / DB-fallback paths, so the skip does not re-open any bypass. + - ``PROXY_ADMIN`` tokens still run through ``common_checks`` so + team-blocked / team-budget / end-user-budget / tag-budget / + vector-store / tool-allowlist enforcement applies to admin keys + too. Admin status is honored where the underlying check exempts it + (``_is_api_route_allowed``, ``organization_role_based_access_check``). """ from litellm.proxy.proxy_server import ( general_settings, @@ -1756,27 +1758,56 @@ async def _run_centralized_common_checks( ) ) - try: - ( - team_object, - user_object, - project_object, - end_user_object, - global_proxy_spend, - ) = await asyncio.gather(*fetch_coros, return_exceptions=False) - except HTTPException: - # Any of the five gathered fetches can raise HTTPException. Only - # reconstruct from the token when a team_id is known — otherwise - # the exception came from a different fetch and the assert in - # _team_obj_from_token would fire. - if user_api_key_auth_obj.team_id is not None: - team_object = _team_obj_from_token(user_api_key_auth_obj) - else: - team_object = None - user_object = None - project_object = None - end_user_object = None - global_proxy_spend = None + # Per-fetch error isolation. ``_safe_fetch`` lets HTTPException, + # ProxyException, and BudgetExceededError escape (everything else is + # already swallowed to None). A bare ``except`` over ``gather`` would + # let one fetch's HTTPException null out every other context — e.g. + # a 404 from ``get_team_object`` (token references a deleted team) + # would silently skip the user, end-user, project, and global-spend + # checks. Use ``return_exceptions=True`` and apply per-fetch fallback + # so a missing team only zeros out the team object. + ( + team_result, + user_result, + project_result, + end_user_result, + global_spend_result, + ) = await asyncio.gather(*fetch_coros, return_exceptions=True) + + # ProxyException / BudgetExceededError are authorization failures — + # propagate so the wrapper renders them. HTTPException is fallback + # material (404 from get_team_object is the only known producer). + for r in ( + team_result, + user_result, + project_result, + end_user_result, + global_spend_result, + ): + if isinstance(r, (ProxyException, litellm.BudgetExceededError)): + raise r + + if isinstance(team_result, HTTPException): + # Token-derived fallback only valid when a team_id is set; + # _team_obj_from_token asserts that precondition. + team_object = ( + _team_obj_from_token(user_api_key_auth_obj) + if user_api_key_auth_obj.team_id is not None + else None + ) + else: + team_object = team_result + + user_object = None if isinstance(user_result, HTTPException) else user_result + project_object = ( + None if isinstance(project_result, HTTPException) else project_result + ) + end_user_object = ( + None if isinstance(end_user_result, HTTPException) else end_user_result + ) + global_proxy_spend = ( + None if isinstance(global_spend_result, HTTPException) else global_spend_result + ) # common_checks identifies admin via user_object, not the token # (non_proxy_admin_allowed_routes_check). JWT admin shortcut and diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 8130e0a942..ca8c7344ab 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -2289,3 +2289,190 @@ async def test_centralized_common_checks_http_exception_without_team_id(): finally: for k, v in originals.items(): setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_centralized_common_checks_team_404_does_not_zero_other_contexts(): + """Per-fetch isolation: an HTTPException(404) from get_team_object + (token references a deleted team) must reconstruct the team from the + token but leave user_object / end_user_object / project_object intact. + Pre-fix a bare ``except HTTPException`` over ``asyncio.gather`` zeroed + every context, silently skipping user-budget, end-user-budget, and + project enforcement whenever the token's team_id was stale.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import HTTPException, Request + from starlette.datastructures import URL + + from litellm.proxy._types import ( + LiteLLM_EndUserTable, + LiteLLM_ProjectTableCachedObj, + ) + + token = UserAPIKeyAuth( + api_key="sk-test", + user_id="u", + team_id="deleted-team", + team_max_budget=5.0, + team_models=["gpt-4o"], + project_id="proj-1", + end_user_id="alice", + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + request._body = json.dumps({"user": "alice", "model": "gpt-4o"}).encode() + + fetched_user = LiteLLM_UserTable( + user_id="u", + user_role=LitellmUserRoles.INTERNAL_USER.value, + max_budget=10.0, + spend=2.0, + ) + fetched_end_user = LiteLLM_EndUserTable(user_id="alice", blocked=False, spend=1.0) + fetched_project = LiteLLM_ProjectTableCachedObj( + project_id="proj-1", + project_alias="Proj 1", + metadata={}, + created_by="admin", + updated_by="admin", + ) + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + side_effect=HTTPException(status_code=404, detail="team-not-found"), + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new_callable=AsyncMock, + return_value=fetched_user, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_project_object", + new_callable=AsyncMock, + return_value=fetched_project, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_end_user_object", + new_callable=AsyncMock, + return_value=fetched_end_user, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_checks, + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"user": "alice", "model": "gpt-4o"}, + route="/chat/completions", + ) + + mock_checks.assert_awaited_once() + kwargs = mock_checks.call_args.kwargs + # team reconstructed from the token + assert kwargs["team_object"] is not None + assert kwargs["team_object"].team_id == "deleted-team" + assert kwargs["team_object"].max_budget == 5.0 + # other contexts must NOT be zeroed by the team fetch failure + assert kwargs["user_object"] is fetched_user + assert kwargs["end_user_object"] is fetched_end_user + assert kwargs["project_object"] is fetched_project + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_centralized_common_checks_user_http_exception_isolates_to_user_only(): + """Per-fetch isolation, mirror of the team case: an HTTPException + from get_user_object must zero only ``user_object``. The successfully + fetched team / end_user / project / global_spend must reach + common_checks intact so their enforcement still runs.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import HTTPException, Request + from starlette.datastructures import URL + + from litellm.proxy._types import ( + LiteLLM_EndUserTable, + LiteLLM_ProjectTableCachedObj, + LiteLLM_TeamTableCachedObj, + ) + + token = UserAPIKeyAuth( + api_key="sk-test", + user_id="u", + team_id="t1", + project_id="proj-1", + end_user_id="alice", + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + request._body = json.dumps({"user": "alice", "model": "gpt-4o"}).encode() + + fetched_team = LiteLLM_TeamTableCachedObj( + team_id="t1", max_budget=20.0, models=["gpt-4o"] + ) + fetched_end_user = LiteLLM_EndUserTable(user_id="alice", blocked=False, spend=1.0) + fetched_project = LiteLLM_ProjectTableCachedObj( + project_id="proj-1", + project_alias="Proj 1", + metadata={}, + created_by="admin", + updated_by="admin", + ) + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + return_value=fetched_team, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new_callable=AsyncMock, + side_effect=HTTPException(status_code=404, detail="user-not-found"), + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_project_object", + new_callable=AsyncMock, + return_value=fetched_project, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_end_user_object", + new_callable=AsyncMock, + return_value=fetched_end_user, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_checks, + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"user": "alice", "model": "gpt-4o"}, + route="/chat/completions", + ) + + mock_checks.assert_awaited_once() + kwargs = mock_checks.call_args.kwargs + assert kwargs["team_object"] is fetched_team + assert kwargs["end_user_object"] is fetched_end_user + assert kwargs["project_object"] is fetched_project + # only the user_object is zeroed by its own fetch failing + assert kwargs["user_object"] is None + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v)