fix(auth): gate post-custom-auth DB lookups behind opt-in flag#25634
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Greptile SummaryThis PR gates the two Confidence Score: 5/5Safe to merge — the change is minimal, correctly restores the pre-v1.82.6 default, and is fully covered by the four new tests. Both production changes are guarded by the existing getattr(litellm, ..., False) pattern used throughout the proxy; the proxy config catch-all already propagates litellm_settings: enable_post_custom_auth_checks: true without any extra wiring. No P0/P1 findings; the only comment is a P2 style suggestion about in-function imports in tests. No files require special attention.
|
| Filename | Overview |
|---|---|
| litellm/proxy/auth/user_api_key_auth.py | Adds an opt-in gate (getattr(litellm, "enable_post_custom_auth_checks", False)) around both post-custom-auth DB lookup call sites; logic is correct and consistent with the existing pattern. |
| tests/test_litellm/proxy/auth/test_user_api_key_auth.py | Adds four end-to-end async tests covering default (gate closed) and opt-in (gate open) for both the user and enterprise custom-auth branches; cleanup is properly handled in finally blocks. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[_user_api_key_auth_builder] --> B{enterprise_custom_auth set?}
B -- Yes --> C[await enterprise_custom_auth]
C --> D{response is UserAPIKeyAuth?}
D -- Yes --> E[model_validate response]
E --> F{enable_post_custom_auth_checks?\ndefault: False}
F -- True --> G[_run_post_custom_auth_checks\nuser/team/project DB lookups]
G --> H[return validated]
F -- False --> H
D -- No --> I[use as raw api_key\ncontinue to standard auth]
B -- No --> J{user_custom_auth set?}
J -- Yes --> K[await user_custom_auth]
K --> L[model_validate response]
L --> M{enable_post_custom_auth_checks?\ndefault: False}
M -- True --> N[_run_post_custom_auth_checks\nuser/team/project DB lookups]
N --> O[return validated]
M -- False --> O
J -- No --> P[Standard LiteLLM auth path]
Reviews (1): Last reviewed commit: "fix(auth): gate post-custom-auth DB look..." | Re-trigger Greptile
| from fastapi import Request | ||
| from starlette.datastructures import URL | ||
|
|
||
| import litellm | ||
| import litellm.proxy.proxy_server as _proxy_server_mod | ||
| from litellm.proxy._types import LitellmUserRoles | ||
| from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder | ||
|
|
||
| trusted_token = UserAPIKeyAuth( |
There was a problem hiding this comment.
In-function imports repeated across all four tests
Per the project style guide, imports belong at the module level. The same six imports appear verbatim inside every new test function body — they can be hoisted to the top of the file (or at minimum deduplicated), and _user_api_key_auth_builder can be added to the existing import block at line 25.
Move these six lines to the top of the file (alongside the existing imports) and remove the duplicated copies in all four test functions.
Relevant issues
Follow-up to #24589 (stalled). Fixes a ~44% RPS regression on the custom-auth hot path reported on v1.82.6+.
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
tests/test_litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewDelays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
CI (LiteLLM team)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Screenshots / Proof of Fix
Benchmark comparing v1.81.13 vs v1.82.6+ under a 750-user / 100-spawn-rate / 60s locust shape against a mocked model, with a
user_custom_authhandler returning a static trusted token:~44% RPS drop and ~5x p95 increase, entirely explained by the two unconditional
_run_post_custom_auth_checkscall sites firing per-request DB lookups on already-trusted tokens.Type
🐛 Bug Fix
Changes
Problem
Since v1.82.6,
_user_api_key_auth_builderunconditionally calls_run_post_custom_auth_checks()on both theenterprise_custom_authanduser_custom_authreturn paths. That helper issues per-requestget_user_object/get_team_object/get_project_objectDB lookups on responses that the custom-auth handler already vouched for — a pure perf hit for deployments whose custom-auth handler is the source of truth, which is the common case.The customer impact is the regression table above: trusted-token proxy throughput drops from 358 → 249 RPS and p95 latency blows out from 1.3 s → 6.3 s at identical load.
Fix
Gate both
_run_post_custom_auth_checksinvocations behind a new module-level flaglitellm.enable_post_custom_auth_checks, defaultFalse. This restores v1.81.13 behavior for everyone by default, and users who rely on the v1.82.6 post-custom-auth DB lookups (e.g. budget enforcement on custom-auth-supplied end-users) can opt back in with:Same
getattr(litellm, ...)pattern and same flag name as the stalled #24589 — deliberately identical to keep review overhead minimal.Why re-submit instead of rebasing #24589
#24589 is currently unmergeable for two reasons unrelated to the fix itself:
main.This PR contains only the two gate changes and their tests, branched fresh from current
main. #24589 can be closed once this merges.Tests added
Four new async tests in
tests/test_litellm/proxy/auth/test_user_api_key_auth.py, all driving_user_api_key_auth_builderend-to-end (not just the inner helper) so the outer gate is actually proven:test_user_custom_auth_skips_post_custom_auth_checks_by_default— default gate closed, helper not called.test_user_custom_auth_runs_post_custom_auth_checks_when_opt_in—enable_post_custom_auth_checks=True, helper awaited.test_enterprise_custom_auth_skips_post_custom_auth_checks_by_default— mirror for the enterprise branch.test_enterprise_custom_auth_runs_post_custom_auth_checks_when_opt_in— opt-in mirror for the enterprise branch.The enterprise-branch coverage specifically addresses Greptile's feedback on #24589 asking for an outer-gate test that short-circuits the call before it ever runs, rather than only testing the inner helper in isolation.
Backwards compatibility
This is a silent behavior change for the narrow set of users who (a) upgraded through v1.82.6+, (b) use
custom_auth, and (c) actively depend on the post-custom-auth DB lookups populating user/team/project state on trusted tokens. Those users must setlitellm_settings.enable_post_custom_auth_checks: trueto restore v1.82.6 behavior.Also worth flagging: users who previously set
general_settings.custom_auth_run_common_checks: trueon v1.82.6+ will need to also setenable_post_custom_auth_checks: trueafter this change — otherwise the inner flag becomes a no-op, because the outer gate short-circuits before the inner flag is consulted.The tradeoff is deliberate — leaving the regression in place penalizes every custom-auth deployment, which is a much larger population than the users who relied on the v1.82.6 behavior.
Scope explicitly out
To keep this PR minimal and reviewable in one pass:
custom_auth_run_common_checks(ingeneral_settings) andenable_post_custom_auth_checks(on thelitellmmodule) into a single namespace — worth doing, but as a follow-up issue._run_post_custom_auth_checksitself — only the callers are gated.Rollback
Trivial: single-commit revert. No schema changes, no config migrations, no new dependencies.