Skip to content

fix(auth): gate post-custom-auth DB lookups behind opt-in flag#25634

Merged
krrish-berri-2 merged 1 commit intoBerriAI:mainfrom
michelligabriele:fix/post-custom-auth-gate
Apr 13, 2026
Merged

fix(auth): gate post-custom-auth DB lookups behind opt-in flag#25634
krrish-berri-2 merged 1 commit intoBerriAI:mainfrom
michelligabriele:fix/post-custom-auth-gate

Conversation

@michelligabriele
Copy link
Copy Markdown
Collaborator

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

  • I have Added testing in the tests/test_litellm/ directory, Adding at least 1 test is a hard requirement - see details
  • My PR passes all unit tests on make test-unit
  • My PR's scope is as isolated as possible, it only solves 1 specific problem
  • I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

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

CI status guideline:

  • 50-55 passing tests: main is stable with minor issues.
  • 45-49 passing tests: acceptable but needs attention
  • <= 40 passing tests: unstable; be careful with your merges and assess the risk.
  • 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_auth handler returning a static trusted token:

Version RPS p95 latency Failures
v1.81.13 358 1300 ms 0
v1.82.6+ 249 6300 ms 0

~44% RPS drop and ~5x p95 increase, entirely explained by the two unconditional _run_post_custom_auth_checks call sites firing per-request DB lookups on already-trusted tokens.

Type

🐛 Bug Fix

Changes

Problem

Since v1.82.6, _user_api_key_auth_builder unconditionally calls _run_post_custom_auth_checks() on both the enterprise_custom_auth and user_custom_auth return paths. That helper issues per-request get_user_object / get_team_object / get_project_object DB 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_checks invocations behind a new module-level flag litellm.enable_post_custom_auth_checks, default False. 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:

litellm_settings:
  enable_post_custom_auth_checks: true

Same getattr(litellm, ...) pattern and same flag name as the stalled #24589deliberately 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:

  1. The branch carried ~99 files of unrelated CI / workflow / dependabot churn from an older base, most of which conflicts with intervening main.
  2. A CLA identity issue after the v1.82.7/v1.82.8 security incident blocks the original author's commits from being accepted.

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_builder end-to-end (not just the inner helper) so the outer gate is actually proven:

  1. test_user_custom_auth_skips_post_custom_auth_checks_by_default — default gate closed, helper not called.
  2. test_user_custom_auth_runs_post_custom_auth_checks_when_opt_inenable_post_custom_auth_checks=True, helper awaited.
  3. test_enterprise_custom_auth_skips_post_custom_auth_checks_by_default — mirror for the enterprise branch.
  4. 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 set litellm_settings.enable_post_custom_auth_checks: true to restore v1.82.6 behavior.

Also worth flagging: users who previously set general_settings.custom_auth_run_common_checks: true on v1.82.6+ will need to also set enable_post_custom_auth_checks: true after 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:

  • No docs / migration note (flagged as P2 on fix(auth): remove unconditional post-custom-auth DB lookups #24589 — follow-up).
  • No consolidation of custom_auth_run_common_checks (in general_settings) and enable_post_custom_auth_checks (on the litellm module) into a single namespace — worth doing, but as a follow-up issue.
  • No refactor of _run_post_custom_auth_checks itself — only the callers are gated.
  • No CI / workflow / dependabot changes, no dep bumps.

Rollback

Trivial: single-commit revert. No schema changes, no config migrations, no new dependencies.

@vercel
Copy link
Copy Markdown

vercel bot commented Apr 13, 2026

The latest updates on your projects. Learn more about Vercel for GitHub.

Project Deployment Actions Updated (UTC)
litellm Ready Ready Preview, Comment Apr 13, 2026 2:53pm

Request Review

@codspeed-hq
Copy link
Copy Markdown
Contributor

codspeed-hq bot commented Apr 13, 2026

Merging this PR will not alter performance

✅ 16 untouched benchmarks


Comparing michelligabriele:fix/post-custom-auth-gate (531c9ec) with main (5e80e07)

Open in CodSpeed

@codecov
Copy link
Copy Markdown

codecov bot commented Apr 13, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@greptile-apps
Copy link
Copy Markdown
Contributor

greptile-apps bot commented Apr 13, 2026

Greptile Summary

This PR gates the two _run_post_custom_auth_checks call sites (on both the enterprise_custom_auth and user_custom_auth return paths) behind a new litellm.enable_post_custom_auth_checks flag (default False), restoring the pre-v1.82.6 fast path and fixing a reported ~44% RPS regression. The proxy config catch-all (setattr(litellm, key, value)) already handles litellm_settings: enable_post_custom_auth_checks: true without any additional wiring, so the implementation is minimal and correct.

Confidence Score: 5/5

Safe 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.

Important Files Changed

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]
Loading

Reviews (1): Last reviewed commit: "fix(auth): gate post-custom-auth DB look..." | Re-trigger Greptile

Comment on lines +171 to +179
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(
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 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.

@krrish-berri-2 krrish-berri-2 merged commit 0eae9f1 into BerriAI:main Apr 13, 2026
50 of 51 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants