From b44755db96fa54c6d2e4c38960ea20cbac0a93fc Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Tue, 3 Mar 2026 22:50:34 +0530 Subject: [PATCH] fix(proxy): make common_checks opt-in for custom auth via custom_auth_run_common_checks Replaces the skip_route_check approach from PR #22662 with a configurable opt-in flag. By default, common_checks() is not run for custom auth flows, preserving backwards compatibility with pre-#22164 behavior. Users who want budget/team/route enforcement on custom auth can enable it: general_settings: custom_auth_run_common_checks: true Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/auth/auth_checks.py | 28 +++---- litellm/proxy/auth/user_api_key_auth.py | 32 ++++---- .../proxy/auth/test_auth_checks.py | 77 +++++++++---------- .../auth/test_custom_auth_end_user_budget.py | 21 ++++- 4 files changed, 86 insertions(+), 72 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 91ac58215a..ef6b0ac462 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -234,7 +234,6 @@ async def common_checks( request: Request, skip_budget_checks: bool = False, project_object: Optional[LiteLLM_ProjectTableCachedObj] = None, - skip_route_check: bool = False, ) -> bool: """ Common checks across jwt + key-based auth. @@ -454,21 +453,18 @@ async def common_checks( user_object=user_object, route=route, request_body=request_body ) - if not skip_route_check: - token_team = getattr(valid_token, "team_id", None) - token_type: Literal["ui", "api"] = ( - "ui" - if token_team is not None and token_team == "litellm-dashboard" - else "api" - ) - _is_route_allowed = _is_allowed_route( - route=route, - token_type=token_type, - user_obj=user_object, - request=request, - request_data=request_body, - valid_token=valid_token, - ) + token_team = getattr(valid_token, "team_id", None) + token_type: Literal["ui", "api"] = ( + "ui" if token_team is not None and token_team == "litellm-dashboard" else "api" + ) + _is_route_allowed = _is_allowed_route( + route=route, + token_type=token_type, + user_obj=user_object, + request=request, + request_data=request_body, + valid_token=valid_token, + ) # 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store await vector_store_access_check( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 1d575fb513..7b52f6bb96 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1752,21 +1752,21 @@ async def _run_post_custom_auth_checks( if _project_obj is not None: valid_token.project_metadata = _project_obj.metadata - _ = await common_checks( - request=request, - request_body=request_data, - team_object=_team_obj, - user_object=user_object, - end_user_object=end_user_object, - general_settings=general_settings, - global_proxy_spend=None, - route=route, - llm_router=llm_router, - proxy_logging_obj=proxy_logging_obj, - valid_token=valid_token, - skip_budget_checks=False, - project_object=_project_obj, - skip_route_check=True, - ) + if general_settings.get("custom_auth_run_common_checks", False): + _ = await common_checks( + request=request, + request_body=request_data, + team_object=_team_obj, + user_object=user_object, + end_user_object=end_user_object, + general_settings=general_settings, + global_proxy_spend=None, + route=route, + llm_router=llm_router, + proxy_logging_obj=proxy_logging_obj, + valid_token=valid_token, + skip_budget_checks=False, + project_object=_project_obj, + ) return valid_token diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 501c2285d1..ff1bc5b258 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1522,52 +1522,51 @@ async def test_get_fuzzy_user_object_case_insensitive_email(): @pytest.mark.asyncio -async def test_common_checks_skip_route_check_for_custom_auth(): +async def test_custom_auth_common_checks_opt_in(): """ - Test that custom routes (e.g. /ldap/ngs/ready) pass common_checks when - skip_route_check=True, which is the case for custom auth flows. + Test that _run_post_custom_auth_checks only runs common_checks when + custom_auth_run_common_checks is explicitly set to True in general_settings. - Regression test for: custom user-added routes being rejected as admin-only - after _run_post_custom_auth_checks was introduced. + By default (False), common_checks is skipped for backwards compatibility + with custom auth flows that existed before PR #22164. """ - from fastapi import Request + from litellm.proxy.auth.user_api_key_auth import _run_post_custom_auth_checks - from litellm.proxy.auth.auth_checks import common_checks - - mock_request = MagicMock(spec=Request) valid_token = UserAPIKeyAuth(token="test-token") + mock_request = MagicMock() - # Without skip_route_check, a custom route with unknown user should fail - with pytest.raises(Exception): - await common_checks( - request_body={}, - team_object=None, - user_object=None, - end_user_object=None, - global_proxy_spend=None, - general_settings={}, - route="/ldap/ngs/ready", - llm_router=None, - proxy_logging_obj=MagicMock(), + # Default (no flag) — common_checks should NOT be called + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_common, patch( + "litellm.proxy.proxy_server.general_settings", + {}, + ): + mock_common.return_value = True + result = await _run_post_custom_auth_checks( valid_token=valid_token, request=mock_request, - skip_route_check=False, + request_data={}, + route="/ldap/ngs/ready", + parent_otel_span=None, ) + mock_common.assert_not_called() - # With skip_route_check=True (custom auth path), the same route should pass - result = await common_checks( - request_body={}, - team_object=None, - user_object=None, - end_user_object=None, - global_proxy_spend=None, - general_settings={}, - route="/ldap/ngs/ready", - llm_router=None, - proxy_logging_obj=MagicMock(), - valid_token=valid_token, - request=mock_request, - skip_route_check=True, - ) - - assert result is True + # With flag=True — common_checks SHOULD be called + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_common, patch( + "litellm.proxy.proxy_server.general_settings", + {"custom_auth_run_common_checks": True}, + ): + mock_common.return_value = True + result = await _run_post_custom_auth_checks( + valid_token=valid_token, + request=mock_request, + request_data={}, + route="/chat/completions", + parent_otel_span=None, + ) + mock_common.assert_called_once() diff --git a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py index 73a9718842..18816dcec4 100644 --- a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py @@ -10,9 +10,10 @@ from litellm.proxy._types import UserAPIKeyAuth @pytest.mark.asyncio async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id(): - # Test backwards compatibility + # Test backwards compatibility — common_checks only runs when opt-in flag is set valid_token = UserAPIKeyAuth(token="test_token") + # Default: common_checks should NOT be called with patch( "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock ) as mock_common: @@ -26,6 +27,24 @@ async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id(): ) assert result.token == "test_token" assert getattr(result, "end_user_id", None) is None + mock_common.assert_not_awaited() + + # With opt-in flag: common_checks SHOULD be called + with patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock + ) as mock_common, patch( + "litellm.proxy.proxy_server.general_settings", + {"custom_auth_run_common_checks": True}, + ): + mock_common.return_value = True + result = await _run_post_custom_auth_checks( + valid_token=valid_token, + request=None, + request_data={}, + route="/v1/chat/completions", + parent_otel_span=None, + ) + assert result.token == "test_token" mock_common.assert_awaited_once()