diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 43646a11c63..beb679ec316 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -512,7 +512,14 @@ async def _auto_register_jwt_mapping( """ Auto-register: create a new virtual key + mapping for an unrecognised JWT claim value. The new key carries no model/budget restrictions; admins can tighten it later. + + Race safety: if two concurrent requests both reach here simultaneously (both saw + no mapping in the DB), one will win the unique-constraint race on + litellm_jwtkeymapping. The loser catches the conflict, fetches the winner's + mapping, and proceeds — no orphaned keys and no error surfaced to the caller. """ + # Inline import required: key_management_endpoints imports user_api_key_auth + # (line 51) so a module-level import here would create a circular dependency. from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_helper_fn, ) @@ -527,15 +534,38 @@ async def _auto_register_jwt_mapping( ) token_hash = key_data["token"] - await prisma_client.db.litellm_jwtkeymapping.create( - data={ - "jwt_claim_name": virtual_key_claim_field, - "jwt_claim_value": claim_value, - "token": token_hash, - "created_by": "auto_register", - "updated_by": "auto_register", - } - ) + try: + await prisma_client.db.litellm_jwtkeymapping.create( + data={ + "jwt_claim_name": virtual_key_claim_field, + "jwt_claim_value": claim_value, + "token": token_hash, + "created_by": "auto_register", + "updated_by": "auto_register", + } + ) + except Exception as e: + error_str = str(e).lower() + if "unique" in error_str or "p2002" in error_str: + # A concurrent request won the race — fetch the winning mapping and + # use its token. The key we just generated is orphaned but harmless; + # it will be excluded from spend tracking since nothing maps to it. + verbose_proxy_logger.debug( + "JWT Key Mapping (auto_register): unique conflict on create — " + "fetching winner's mapping for %s='%s'.", + virtual_key_claim_field, + claim_value, + ) + token_hash = await get_jwt_key_mapping_object( + jwt_claim_name=virtual_key_claim_field, + jwt_claim_value=claim_value, + prisma_client=prisma_client, + ) + if token_hash is None: + # Should not happen, but guard against a delete racing our fetch. + return None + else: + raise await user_api_key_cache.async_set_cache( key=cache_key, @@ -544,8 +574,9 @@ async def _auto_register_jwt_mapping( ) verbose_proxy_logger.info( - f"JWT Key Mapping (auto_register): created new virtual key for " - f"{virtual_key_claim_field}='{claim_value}'." + "JWT Key Mapping (auto_register): created new virtual key for %s='%s'.", + virtual_key_claim_field, + claim_value, ) return await get_key_object( @@ -591,6 +622,20 @@ async def _resolve_jwt_to_virtual_key( status_code=403, detail=f"JWT Key Mapping: No registered mapping for {virtual_key_claim_field}='{claim_value}'. Access denied.", ) + if behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER and prisma_client is not None: + # Stale sentinel written under a prior fallback_team_mapping config — + # evict it and auto-register now that the policy has changed. + await user_api_key_cache.async_delete_cache(cache_key) + return await _auto_register_jwt_mapping( + virtual_key_claim_field=virtual_key_claim_field, + claim_value=str(claim_value), + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + cache_key=cache_key, + ) return None elif cached_mapping is not None: return await get_key_object( diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/proxy_unit_tests/test_jwt_key_mapping.py index 1b9d72d1b50..2456731889b 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/proxy_unit_tests/test_jwt_key_mapping.py @@ -558,6 +558,123 @@ async def test_auto_register_creates_key_and_mapping(): assert cached == "hashed_auto_key" +@pytest.mark.asyncio +async def test_auto_register_triggers_on_stale_no_mapping_sentinel(): + """ + If the cache holds a stale __NO_MAPPING__ sentinel (written under a prior + fallback_team_mapping config) and behavior is now AUTO_REGISTER, the sentinel + must be evicted and auto-registration must run — not silently return None. + """ + from litellm.proxy._types import UnregisteredJWTClientBehavior + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + virtual_key_claim_field="email", + unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.AUTO_REGISTER, + virtual_key_mapping_cache_ttl=300, + ) + jwt_claims = {"email": "alice@corp.com"} + + prisma_client = MagicMock() + prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None) + prisma_client.db.litellm_jwtkeymapping.create = AsyncMock() + + user_api_key_cache = DualCache() + # Seed the stale sentinel + await user_api_key_cache.async_set_cache( + "jwt_key_mapping:email:alice@corp.com", "__NO_MAPPING__" + ) + + mock_key_obj = UserAPIKeyAuth(token="hashed_auto_key", team_id=None) + + with patch( + "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + ) as mock_get_key, patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + new_callable=AsyncMock, + ) as mock_gen_key: + mock_gen_key.return_value = {"token": "hashed_auto_key", "key": "sk-auto-key"} + mock_get_key.return_value = mock_key_obj + + result = await _resolve_jwt_to_virtual_key( + jwt_claims=jwt_claims, + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + # Must have auto-registered, not returned None + assert result == mock_key_obj + prisma_client.db.litellm_jwtkeymapping.create.assert_called_once() + + +@pytest.mark.asyncio +async def test_auto_register_race_condition_unique_conflict(): + """ + If two concurrent requests both call _auto_register_jwt_mapping and the + second hits a unique-constraint violation on create, it must fall back to + fetching the winner's mapping — no error surfaced to the caller. + """ + from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping + from litellm.proxy._types import UnregisteredJWTClientBehavior + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + virtual_key_claim_field="sub", + unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.AUTO_REGISTER, + virtual_key_mapping_cache_ttl=300, + ) + + prisma_client = MagicMock() + prisma_client.db.litellm_jwtkeymapping.create = AsyncMock( + side_effect=Exception("Unique constraint failed (P2002)") + ) + # Simulate the winner's mapping already in DB after the conflict + winner_mapping = MagicMock() + winner_mapping.token = "winner_token_hash" + winner_mapping.is_active = True + prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock( + return_value=winner_mapping + ) + + user_api_key_cache = DualCache() + mock_key_obj = UserAPIKeyAuth(token="winner_token_hash", team_id=None) + + with patch( + "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock + ) as mock_get_key, patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + new_callable=AsyncMock, + return_value={"token": "loser_token_hash", "key": "sk-loser"}, + ): + mock_get_key.return_value = mock_key_obj + + result = await _auto_register_jwt_mapping( + virtual_key_claim_field="sub", + claim_value="user-42", + jwt_handler=jwt_handler, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + cache_key="jwt_key_mapping:sub:user-42", + ) + + assert result == mock_key_obj + # Cache should hold the winner's token, not the loser's + cached = await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:user-42") + assert cached == "winner_token_hash" + mock_get_key.assert_called_once_with( + hashed_token="winner_token_hash", + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + # ────────────────────────────────────────────── # Tests: backward-compat alias jwt_client_id_field # ──────────────────────────────────────────────