Merge pull request #22678 from Harshit28j/litellm_custom_auth_opt_in

fix(proxy): make common_checks opt-in for custom auth

docs/my-website/docs/proxy/guardrails/quick_start.md

@@ -358,13 +358,13 @@ response = client.chat.completions.create(
         }
     ],
     extra_body={
-      "guardrails": [
+      "guardrails": {
         "aporia-pre-guard": {
           "extra_body": {
             "success_threshold": 0.9
           }
         }
-      ]
+      }
     }
 
 )
@@ -387,13 +387,13 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
         "content": "what llm are you"
         }
     ],
-    "guardrails": [
+    "guardrails": {
       "aporia-pre-guard": {
         "extra_body": {
           "success_threshold": 0.9
         }
       }
-    ]
+    }
 }'
 ```
 </TabItem>
@@ -451,7 +451,6 @@ curl -X POST 'http://0.0.0.0:4000/key/generate' \
     -H 'Content-Type: application/json' \
     -d '{
             "guardrails": ["aporia-pre-guard", "aporia-post-guard"]
-        }
     }'
 ```
 
@@ -465,7 +464,6 @@ curl --location 'http://0.0.0.0:4000/key/update' \
     --data '{
         "key": "sk-jNm1Zar7XfNdZXp49Z1kSQ",
         "guardrails": ["aporia-pre-guard", "aporia-post-guard"]
-        }
 }'
 ```
 

litellm/proxy/auth/auth_checks.py

@@ -279,7 +279,6 @@ async def common_checks(  # noqa: PLR0915
     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.
@@ -499,21 +498,18 @@ async def common_checks(  # noqa: PLR0915
         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(

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

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()

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()