litellm/tests
Shivam Rawat 3bd89f209e
Litellm jwt mapping virtualkeys (#28510)
* restore an explicit no-match policy

* fix(jwt): fix AUTO_REGISTER sentinel bypass, race condition, and inline import comment

- AUTO_REGISTER now evicts stale __NO_MAPPING__ sentinel instead of silently
  returning None when cached under a prior fallback_team_mapping config
- Race condition in _auto_register_jwt_mapping: catch P2002 unique-constraint
  violation on concurrent creates, fetch the winning mapping, proceed cleanly
- Added comment on inline generate_key_helper_fn import explaining the circular
  dependency (key_management_endpoints imports user_api_key_auth at line 51)
- 3 new tests: stale sentinel eviction, race condition winner fallback, and the
  existing auto_register happy path

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix(jwt): cache __NO_MAPPING__ sentinel before raising 403 in REJECT mode

REJECT mode was raising HTTPException immediately on a DB miss without writing
the __NO_MAPPING__ sentinel, causing every subsequent rejected request to
re-query the DB. Write the sentinel first so repeated rejections are served
from cache within virtual_key_mapping_cache_ttl.

Adds test asserting DB is not hit on the second reject after a cache-warm miss.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix(jwt): enforce no-match policy when prisma_client is None

The early `if prisma_client is None: return None` guard ran before the
no-match policy check, silently bypassing REJECT and AUTO_REGISTER — every
JWT client fell through to team auth regardless of configuration.

Fix: treat prisma_client=None as a definitive DB miss and fall through to the
same policy block as a real miss. REJECT now raises 403, AUTO_REGISTER raises
500 with a clear message (can't create keys without a DB), FALLBACK_TEAM_MAPPING
returns None unchanged.

Adds three tests: REJECT/403 with no DB, FALLBACK returns None with no DB,
AUTO_REGISTER/500 with no DB.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix(jwt): consistent AUTO_REGISTER on cached sentinel; clean up race orphans

Addresses Greptile review on PR #25570 cherry-pick.

1. Inconsistent AUTO_REGISTER when __NO_MAPPING__ sentinel is cached:
   The cached-sentinel branch silently returned None when prisma_client was
   None, while the fresh path raised HTTP 500 under the same config. Same
   request, different access-control outcome depending on cache state. Both
   paths now raise the same 500.

2. Orphaned virtual keys from race-condition losers:
   On unique-constraint conflict, generate_key_helper_fn had already persisted
   an unrestricted virtual key in LiteLLM_VerificationToken with the cleartext
   in request memory. Under sustained concurrency these accumulated
   indefinitely. The loser now deletes its orphan before falling back to the
   winner's mapping; failure to delete is logged but does not fail the request.

Also corrects a latent FK bug surfaced while fixing #2: the mapping row was
storing the plaintext key in LiteLLM_JWTKeyMapping.token, but that column FKs
to the hashed LiteLLM_VerificationToken.token — now hashed at the call site.

Tests:
- updated test_auto_register_creates_key_and_mapping to assert the hashed
  token is stored, not the plaintext
- updated test_auto_register_race_condition_unique_conflict to assert the
  orphan is deleted with the correct hashed token
- added test_auto_register_raises_500_when_sentinel_cached_and_no_db
- added test_auto_register_race_conflict_tolerates_delete_failure

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix(jwt): close REJECT bypass when JWT omits the configured claim field

A JWT presented without the configured `virtual_key_claim_field` previously
returned None at the `claim_value is None` guard before the
`unregistered_jwt_client_behavior` check ran. A caller who knows the configured
claim-field name could bypass REJECT by simply omitting that field and falling
through to team-based JWT auth.

Apply the no-match policy on a missing claim:
  - REJECT          → 403
  - AUTO_REGISTER   → 403 (no stable identity to map; refuse rather than
                     create a sentinel-keyed record)
  - FALLBACK_TEAM_MAPPING → return None (unchanged, backward-compatible)

Adds three tests covering each branch of the missing-claim path.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix(jwt): AUTO_REGISTER inherits team_id so keys are bounded by team limits

Auto-registered virtual keys were created with no team, model, route, rate, or
budget constraints — broader access than the standard team-based JWT auth path
the same client would have taken. Under AUTO_REGISTER, resolve the team_id
from the JWT (via the operator-configured team_id_jwt_field / team_id_default)
and stamp it on the new key. Downstream auth then applies the team's
budget/models/tpm/rpm/allowed_routes via the existing virtual-key flow.

Policy when team_id_jwt_field is configured:
  - JWT carries team claim → stamp resolved team_id
  - JWT lacks claim + team_id_default set → stamp default
  - JWT lacks claim + no default → 403 (refuse to create an unbounded key)

When neither team_id_jwt_field nor team_id_default is configured, the
operator has explicitly opted out of team-based limits — the auto-created
key has no team_id (matches what team-auth would do in the same config).

Adds 4 tests covering each branch.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix(jwt): make AUTO_REGISTER functional in prod; raise on missing winner

Two correctness fixes flagged by Greptile on the AUTO_REGISTER path:

1. generate_key_helper_fn was called without table_name="key". Without that,
   the helper falls into the user-upsert branch (table_name in (None, "user"))
   and tries to insert into LiteLLM_UserTable with user_id=None, which hits
   the NOT NULL @id constraint. AUTO_REGISTER would never have succeeded in
   production. Now passes table_name="key" explicitly, matching the
   /key/generate caller.

2. When the race loser refetches the winner's mapping and gets None (winner
   row concurrently deleted), the previous code returned None — and the
   caller in _resolve_jwt_to_virtual_key then fell through to less-
   restrictive team-based JWT auth, silently bypassing the configured
   AUTO_REGISTER policy. Now raises HTTP 503 so the caller retries against
   a stable state rather than getting unintended fallback access.

Adds one test for the 503 winner-vanishes path.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix(jwt): defer AUTO_REGISTER until JWT policy is enforced by auth_builder

Closes the JWT policy bypass on the AUTO_REGISTER path flagged by veria-ai.

Before: when unregistered_jwt_client_behavior=auto_register and the JWT's
claim was unmapped, _resolve_jwt_to_virtual_key validated the JWT signature
and then immediately created a virtual key + mapping. JWTAuthManager.auth_builder
never ran for the first request (the new key short-circuited the team-auth
path), and every subsequent request hit the cached mapping — so custom_validate,
RBAC, scope_mappings, and user_allowed_email_domain were never enforced for
auto-registered clients.

After: _resolve_jwt_to_virtual_key returns a _PendingAutoRegister signal
instead of creating the key. The caller in _user_api_key_auth_builder runs
JWTAuthManager.auth_builder, then — only on a validated, policy-passing
result — calls _auto_register_jwt_mapping with the team_id / user_id from
that result. The created key inherits team + user limits from the validated
identity, and future cache hits load that already-policy-checked key.

Also drops the interim _resolve_inherited_team_id helper that pulled team_id
from raw JWT claims — same bypass risk; team_id now comes exclusively from
auth_builder.

Tests:
  - Rewrote two existing tests to assert _resolve_jwt_to_virtual_key returns
    _PendingAutoRegister (no key created yet) for both the fresh-DB-miss
    and stale-sentinel branches
  - Added a contract test that _auto_register_jwt_mapping stamps the
    validated team_id/user_id onto generate_key_helper_fn
  - Removed four stale team-binding tests that exercised the prior
    raw-claim helper

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Update user_api_key_auth.py

* fix(jwt): cache proxy-admin AUTO_REGISTER path to avoid repeated DB lookups

Cache-miss regression introduced by the deferred-auto-register refactor:
when a JWT under AUTO_REGISTER resolved to a proxy admin, the is_proxy_admin
early-return in _user_api_key_auth_builder ran *before* the pending
auto-register cache-write block. Result: no cache entry, so every
subsequent proxy-admin request re-queried get_jwt_key_mapping_object
indefinitely.

Fix: write a __JWT_PROXY_ADMIN__ sentinel to user_api_key_cache before the
early return when a pending auto-register existed. _resolve_jwt_to_virtual_key
treats that sentinel as "skip mapping, fall through to auth_builder", so
future requests from the same JWT identity hit the cache instead of the DB.
auth_builder still runs full JWT policy on every request — only the
mapping DB lookup is short-circuited.

Adds one test asserting the sentinel cache-hit returns None without
hitting prisma_client.db.litellm_jwtkeymapping.find_first.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix(proxy): stamp org context on JWT auto-registered keys

AUTO_REGISTER keys were created with team_id and user_id only, so org budget checks were skipped after switching to the key-scoped path.

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-04 19:00:36 -07:00
..
agent_tests Revert "chore(tests): migrate Bedrock CI to AWS account 941277531214 (#28728)" (#29326) 2026-05-30 11:26:24 -07:00
audio_tests fix(tests): stabilize image-edit VCR cassettes to stop live gpt-image-1 spend (#28110) 2026-05-18 09:15:39 -07:00
basic_proxy_startup_tests
batches_tests test: stabilize batch VCR coverage and stop live upload/network leaks (#29477) 2026-06-02 16:11:52 -07:00
benchmarks
code_coverage_tests Litellm oss staging (#29492) 2026-06-02 08:48:10 -07:00
documentation_tests
enterprise fix(bedrock): support tool search results + chat annotations (#29120) 2026-05-29 20:48:36 -07:00
guardrails_tests Revert "chore(tests): migrate Bedrock CI to AWS account 941277531214 (#28728)" (#29326) 2026-05-30 11:26:24 -07:00
image_gen_tests Revert "chore(tests): migrate Bedrock CI to AWS account 941277531214 (#28728)" (#29326) 2026-05-30 11:26:24 -07:00
integration CI: copy of #25177 (OCI GenAI: embeddings, streaming/reasoning fixes, model catalog) (#28223) 2026-05-23 12:15:41 -07:00
litellm fix(gemini): googleSearch + server-side tools and googleMaps JSON schema (#29582) 2026-06-04 07:43:30 -07:00
litellm_core_utils
litellm_utils_tests test(vcr): close out the remaining VCR live-call leaks (#29603) 2026-06-03 13:46:43 -07:00
litellm-proxy-extras
llm_responses_api_testing test(responses): bump deprecated gemini-3-pro-preview to gemini-3.1-pro-preview (#29433) 2026-06-01 09:54:30 -07:00
llm_translation Litellm oss staging 040626 (#29671) 2026-06-04 11:07:20 -07:00
load_tests
local_testing Litellm oss staging 040626 (#29671) 2026-06-04 11:07:20 -07:00
logging_callback_tests test(vcr): close out the remaining VCR live-call leaks (#29603) 2026-06-03 13:46:43 -07:00
mcp_tests [internal copy of #28008] Support MCP OAuth passthrough and issuer-scoped JWT auth (#28356) 2026-06-02 12:22:04 -07:00
multi_instance_e2e_tests
ocr_tests test(vcr): close out the remaining VCR live-call leaks (#29603) 2026-06-03 13:46:43 -07:00
old_proxy_tests/tests
openai_endpoints_tests chore(ci): modernize model references in tests and configs (#27856) 2026-05-15 15:44:28 -07:00
otel_tests feat(prometheus): add user_email and user_alias to user budget metrics (#28155) 2026-05-18 16:28:14 -07:00
pass_through_tests test(vcr): close out the remaining VCR live-call leaks (#29603) 2026-06-03 13:46:43 -07:00
pass_through_unit_tests test(vcr): close out the remaining VCR live-call leaks (#29603) 2026-06-03 13:46:43 -07:00
proxy_admin_ui_tests fix(guardrails): persist disable_global_guardrails on keys (#29233) 2026-05-28 21:19:04 -07:00
proxy_behavior test(proxy): phase-4 payload behavior pinning for tier-2/3 key + team management endpoints (#28681) 2026-05-23 12:16:29 -07:00
proxy_e2e_anthropic_messages_tests chore(ci): modernize model references in tests and configs (#27856) 2026-05-15 15:44:28 -07:00
proxy_migration_tests test(proxy): stop running real-DB tests in GitHub Actions unit jobs (#29700) 2026-06-04 14:56:02 -07:00
proxy_security_tests test(proxy): stop running real-DB tests in GitHub Actions unit jobs (#29700) 2026-06-04 14:56:02 -07:00
proxy_unit_tests Litellm jwt mapping virtualkeys (#28510) 2026-06-04 19:00:36 -07:00
router_unit_tests fix(router): enforce deployment budgets for dynamically added models (#29273) 2026-05-29 19:43:14 -07:00
scim_tests
search_tests fix(tests): stabilize image-edit VCR cassettes to stop live gpt-image-1 spend (#28110) 2026-05-18 09:15:39 -07:00
spend_tracking_tests chore(ci): modernize model references in tests and configs (#27856) 2026-05-15 15:44:28 -07:00
store_model_in_db_tests
test_litellm Litellm jwt mapping virtualkeys (#28510) 2026-06-04 19:00:36 -07:00
unified_google_tests fix(tests): stabilize image-edit VCR cassettes to stop live gpt-image-1 spend (#28110) 2026-05-18 09:15:39 -07:00
vector_store_tests Revert "chore(tests): migrate Bedrock CI to AWS account 941277531214 (#28728)" (#29326) 2026-05-30 11:26:24 -07:00
windows_tests ci: reproduce default-Windows wheel install to guard MAX_PATH (#29597) 2026-06-03 11:28:08 -07:00
__init__.py
_flush_vcr_cache.py tests(vcr): isolate cassette redis to CASSETTE_REDIS_URL 2026-05-01 12:32:59 -07:00
_live_test_helpers.py test(vcr): close out the remaining VCR live-call leaks (#29603) 2026-06-03 13:46:43 -07:00
_vcr_conftest_common.py test(vcr): close out the remaining VCR live-call leaks (#29603) 2026-06-03 13:46:43 -07:00
_vcr_redis_persister.py fix(tests/vcr): make Redis cassette cache replay deterministically (zero VCR misses on consecutive runs) (#28826) 2026-05-26 11:30:44 -07:00
eval_swe_bench.py
gettysburg.wav
large_text.py
openai_batch_completions.jsonl
README.MD
test_budget_management.py
test_callbacks_on_proxy.py test(callbacks): harden flaky proxy callback-leak detector (#28195) 2026-05-18 16:39:02 -07:00
test_config.py
test_debug_warning.py
test_default_encoding_non_root.py
test_end_users.py chore(ci): modernize model references in tests and configs (#27856) 2026-05-15 15:44:28 -07:00
test_entrypoint.py
test_fallbacks.py
test_gpt5_azure_temperature_support.py
test_health.py fix(tests): swap dall-e to gpt-image-1 after openai deprecation 2026-05-12 16:55:18 -07:00
test_keys.py fix(tests): swap dall-e to gpt-image-1 after openai deprecation 2026-05-12 16:55:18 -07:00
test_litellm_proxy_responses_config.py chore(ci): modernize model references in tests and configs (#27856) 2026-05-15 15:44:28 -07:00
test_logging.conf
test_models.py
test_new_vector_store_endpoints.py
test_openai_endpoints.py fix(ci): keep coverage rename green when a parallel node runs no tests (#29608) 2026-06-03 13:37:53 -07:00
test_organizations.py
test_otel_thread_leak.py
test_passthrough_endpoints.py
test_presidio_latency.py
test_proxy_server_non_root.py
test_ratelimit.py chore(ci): modernize model references in tests and configs (#27856) 2026-05-15 15:44:28 -07:00
test_resource_cleanup.py
test_service_logger_otel.py
test_spend_logs.py Litellm oss staging 04 21 2026 2 (#26569) 2026-05-20 21:25:19 -07:00
test_team_logging.py
test_team_members.py Litellm oss staging 04 21 2026 2 (#26569) 2026-05-20 21:25:19 -07:00
test_team.py
test_users.py Fix: tag budget reset must drop stale management-cache entry (#27568) 2026-05-10 00:18:55 +00:00

In total litellm runs 1000+ tests

[02/20/2025] Update:

To make it easier to contribute and map what behavior is tested,

we've started mapping the litellm directory in tests/test_litellm

This folder can only run mock tests.