litellm/tests/test_litellm/proxy
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
..
_experimental/mcp_server Litellm oss staging 040626 (#29671) 2026-06-04 11:07:20 -07:00
a2a feat(a2a): well-known agent-card discovery + LangGraph Platform mode (#28860) 2026-05-29 20:50:42 -07:00
agent_endpoints Fix : a2a bugs 030626 (#29566) 2026-06-03 11:14:15 -07:00
anthropic_endpoints test: isolate proxy master_key/prisma_client module globals between tests 2026-04-23 15:31:16 -07:00
auth Litellm jwt mapping virtualkeys (#28510) 2026-06-04 19:00:36 -07:00
client fix: scope CLI stored token to base_url to prevent cross-domain credential leakage (#26945) 2026-05-01 12:11:32 -07:00
common_utils fix(reset_budget): write only {spend, budget_reset_at} and stop pre-zeroing counter (#29358) 2026-05-30 17:48:16 -07:00
db feat: add componentized proxy deployment with gateway, backend, ui, and migrations (#27557) 2026-05-16 09:25:17 -07:00
discovery_endpoints style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
experimental/mcp_server Litellm fix GitHub action testing (#11163) 2025-05-26 14:41:42 -07:00
google_endpoints Gemini managed agents support (#28270) 2026-05-19 16:02:03 -07:00
guardrails Litellm oss staging 040626 (#29671) 2026-06-04 11:07:20 -07:00
health_endpoints feat(proxy): native /health/drain preStop hook for graceful shutdown (#29439) 2026-06-02 16:30:44 -07:00
hooks feat(proxy): add per-MCP-server RPM rate limiting for keys and teams (#29482) 2026-06-02 12:52:10 -07:00
image_endpoints style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
management_endpoints fix(key_generate): scope session-token team-key budget exemption to caller-supplied team_id (#29641) 2026-06-03 19:11:53 -07:00
management_helpers feat(mcp/auth): additive key access-group grants + opt-in member assignment (#29313) 2026-05-30 17:35:31 -07:00
memory Litellm memory improvements v2 (#26541) 2026-04-25 19:03:43 -07:00
middleware fix(proxy): point /metrics 401 at the opt-out flag 2026-05-08 18:09:14 -07:00
openai_files_endpoint Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_yj_apr17 2026-04-17 17:36:40 -07:00
pass_through_endpoints fix(proxy): passthrough 404 when SERVER_ROOT_PATH is set (#29658) 2026-06-04 07:44:51 -07:00
policy_engine fix(guardrails): return HTTP 400 for litellm content filter blocks (#28418) 2026-05-30 14:28:46 -07:00
prompts style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
proxy_server tests(proxy_server): surface current behavior in tests (#29309) 2026-05-29 23:17:24 -07:00
public_endpoints fix(mcp): gate /public/mcp_hub strictly on litellm.public_mcp_servers (#27764) 2026-06-04 17:26:59 +05:30
rag_endpoints fix: harden /key/update authorization checks (#27878) 2026-05-14 04:16:04 +00:00
realtime_endpoints test: isolate proxy master_key/prisma_client module globals between tests 2026-04-23 15:31:16 -07:00
response_api_endpoints Litellm websocket improvements (#29563) 2026-06-03 11:48:35 -07:00
shutdown feat(proxy): native /health/drain preStop hook for graceful shutdown (#29439) 2026-06-02 16:30:44 -07:00
spend_tracking fix(proxy): strip NUL bytes from spend log payloads to prevent PostgreSQL 22P05 (#29515) 2026-06-02 12:07:11 -07:00
test_configs Litellm fix GitHub action testing (#11163) 2025-05-26 14:41:42 -07:00
types_utils chore(proxy): also scrub guardrail callbacks / module paths from DB overlay 2026-05-14 01:24:51 +00:00
ui_crud_endpoints chore(ci): merge dev branch (#28657) 2026-05-25 13:48:47 -07:00
utils test(proxy/utils): pin ProxyLogging behavior (#29485) 2026-06-02 17:45:39 -07:00
vector_store_endpoints fix(proxy): restrict vector store index create/delete to proxy admins (#29202) 2026-05-30 15:10:21 -07:00
__init__.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
conftest.py feat(proxy): native /health/drain preStop hook for graceful shutdown (#29439) 2026-06-02 16:30:44 -07:00
test_aiohttp_cleanup_closed.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_aiohttp_session_recovery.py fix: use AsyncMock for concurrent test consistency 2026-03-18 00:54:23 +00:00
test_api_key_masking_in_errors.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_audio_speech_prometheus_hooks.py fix req changes 2026-02-28 21:32:57 +05:30
test_batch_expiry.py fix(proxy): strip LiteLLM policy tracking from OpenAI batch metadata (#28425) 2026-05-26 11:35:42 -07:00
test_batch_metadata_none_fix.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_batch_retrieve_bedrock.py fix(proxy/batches): forward model to retrieve_batch for bedrock 2026-04-29 22:48:03 +02:00
test_budget_reservation.py fix(proxy): gate image-gen reservation strictly on model mode 2026-05-09 09:16:27 -07:00
test_caching_routes.py Litellm oss staging 040626 (#29671) 2026-06-04 11:07:20 -07:00
test_chat_completion_metadata.py fix: propagate JWT auth metadata to OTEL spans (#19627) 2026-01-23 21:21:23 -08:00
test_common_request_processing.py fix(proxy): disable proxy buffering on streaming SSE responses (#29557) 2026-06-04 17:23:14 +05:30
test_component_allowlists.py fix(anthropic/adapter): emit thinking block for reasoning_content-only streaming chunks (#29600) 2026-06-03 11:26:13 -07:00
test_cors_config.py refactor: extract _get_cors_config() for testability, fix no-op CORS tests 2026-04-11 22:24:04 +05:30
test_custom_proxy.py fix(ui/): fix routing for custom server root path (#15701) 2025-10-23 13:59:29 -07:00
test_dynamic_mcp_route.py Litellm oss staging 040626 (#29671) 2026-06-04 11:07:20 -07:00
test_empty_model_list.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_enforce_user_param.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_fallback_management_endpoints.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_fastapi_offline_routes.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_filter_models_by_team_access_group.py fix: resolve access group names in _filter_models_by_team_id (#25224) 2026-05-01 11:55:36 +05:30
test_health_check_functions.py feat(proxy): skip disable_background_health_check models on GET /health when flag set (#27716) 2026-05-13 09:49:05 -07:00
test_health_check_max_tokens.py Litellm agent oss staging 05 11 2026 (#27733) 2026-05-13 14:09:12 -07:00
test_langfuse_passthrough_security.py chore(callbacks): guard dynamic integration hosts 2026-04-30 14:27:19 -07:00
test_lazy_openapi_snapshot.py [Fix] Tests: Move Misplaced Import in Lazy OpenAPI Snapshot Test 2026-05-01 14:24:53 -07:00
test_litellm_pre_call_utils.py fix(proxy): enforce tag budgets for key-level tags (#29108) 2026-05-28 11:39:02 -07:00
test_max_budget_env_var.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_mcp_asgi_response.py fix(mcp): surface upstream 401 for token-forwarding MCP servers (#27847) 2026-05-13 12:03:36 -07:00
test_model_dump_with_preserved_fields.py Fix_mapped tests part 2 2026-02-26 12:43:39 +05:30
test_model_id_header_propagation.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_model_info_default_limits.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_model_level_guardrails.py fix(proxy): run model-level post_call guardrails on streaming requests (#26922) 2026-05-07 11:53:03 -07:00
test_openapi_schema_validation.py fix(proxy): preserve HTTP operations when injecting WebSocket stubs into OpenAPI schema 2026-05-06 00:28:42 +02:00
test_pricing_field_strip.py fix(proxy): move pricing strip below the litellm_metadata JSON-string parse 2026-05-03 02:04:05 +00:00
test_prometheus_cleanup.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_provider_url_destination_guard.py chore(proxy): move URL-valued model/file_id guard from SDK to proxy 2026-05-02 02:17:46 +00:00
test_proxy_cli.py Add granian as a ASGI compliant web server. Provider better throughput stability, (#26027) 2026-05-21 19:08:37 -07:00
test_proxy_logging_hook_detection.py perf: reduce per-request and per-chunk overhead across Anthropic streaming hot paths (#28289) 2026-05-23 12:15:59 -07:00
test_proxy_server.py fix(proxy): omit OpenAI [DONE] on google-genai streamGenerateContent (#29426) 2026-06-01 14:38:19 -07:00
test_proxy_types.py fix(proxy): add default=None to LiteLLM_TeamMembership.litellm_budget_table (#29684) 2026-06-04 12:13:11 -07:00
test_proxy_utils.py feat(otel): typed semconv-aligned OpenTelemetry instrumentation (#28909) 2026-05-29 23:15:27 -07:00
test_pyroscope.py Implement normalize_nonempty_secret_str function to trim whitespace from secrets and treat empty values as unset. Update proxy_server to use this function for Grafana credentials. Enhance tests to validate the new normalization behavior. 2026-05-04 18:17:31 +00:00
test_redis_auth_cache_flag.py refactor(proxy): replace DualCache with UserApiKeyCache for user API key management 2026-04-28 19:15:03 +00:00
test_response_model_sanitization.py perf: eliminate per-request callback scanning on proxy hot path (#27858) 2026-05-14 09:28:31 -07:00
test_route_a2a_models.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_route_llm_request.py feat(proxy): fix vector store retrieve/list/update/delete without model (#27929) 2026-05-15 04:37:59 -07:00
test_sensitive_route_auth.py chore(proxy): guard sensitive public endpoints 2026-04-30 11:52:47 -07:00
test_shared_health_check.py feat(proxy): skip disable_background_health_check models on GET /health when flag set (#27716) 2026-05-13 09:49:05 -07:00
test_spend_log_cleanup.py fix(proxy): keep spend log cleanup running after batch failures and surface DB errors (#27303) 2026-05-06 18:39:15 +00:00
test_swagger_chat_completions.py fix(proxy): avoid misleading multi-method operation ids 2026-04-30 20:44:14 -07:00
test_team_member_update.py fix mapped tests (#12320) 2025-07-04 10:04:43 -07:00
test_team_org_move.py fix(team_endpoints): auto-add SSO team members to org on move (proxy admin only) (#26377) 2026-04-24 08:36:25 -07:00
test_tools_allowlist_enforcement.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_update_llm_router_resilience.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00