diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2ab147043d9..3d6f5a5b6f6 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -43,6 +43,7 @@ from litellm.proxy._experimental.mcp_server.db import ( from litellm.proxy._types import * from litellm.proxy._types import LiteLLM_VerificationToken from litellm.proxy.auth.auth_checks import ( + _cache_key_object, _delete_cache_key_object, can_team_access_model, get_org_object, @@ -916,6 +917,30 @@ async def _common_key_generation_helper( # noqa: PLR0915 response.token_id ) # remap token to use the hash, and leave the key in the `key` field [TODO]: clean up generate_key_helper_fn to do this + # Warm the shared auth cache so the very first authenticated request with this + # key does not depend on instant DB visibility from /key/generate to the next + # replica's auth lookup. Best-effort: a Redis outage must not break key creation. + try: + from litellm.proxy.proxy_server import ( + proxy_logging_obj, + user_api_key_cache, + ) + + if response.token is not None and user_api_key_cache is not None: + await _cache_key_object( + hashed_token=response.token, + user_api_key_obj=UserAPIKeyAuth( + **response.model_dump(exclude_none=True) + ), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: + verbose_proxy_logger.exception( + "generate_key_fn: failed to warm user_api_key_cache for newly created key (best-effort): %s", + str(e), + ) + asyncio.create_task( KeyManagementEventHooks.async_key_generated_hook( data=data, diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 26f04a4abcb..feb1528749b 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -3016,3 +3016,32 @@ async def test_team_member_budget_check_zero_per_member_row_still_blocks(): proxy_logging_obj=proxy_logging_obj, ) assert exc_info.value.max_budget == 0.0 + + +@pytest.mark.asyncio +async def test_get_key_object_returns_pre_warmed_entry_without_db_lookup(): + """ + Regression for case 2026-05-13-zurich-invalid-proxy-token: when + /key/generate has warmed the cache, get_key_object must serve the + key from cache without touching the DB. This is the path that + survives DB-side replication lag / pool isolation. + """ + from litellm.proxy.auth.auth_checks import get_key_object + + pre_warmed_key = UserAPIKeyAuth(token="hashed-warmed-token") + + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=pre_warmed_key) + mock_cache.async_set_cache = AsyncMock() + + mock_prisma_client = MagicMock() + mock_prisma_client.get_data = AsyncMock() + + result = await get_key_object( + hashed_token="hashed-warmed-token", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + + assert result.token == "hashed-warmed-token" + mock_prisma_client.get_data.assert_not_called() diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index c2ac941e737..3083e816a5c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -11201,3 +11201,101 @@ async def test_ghsa_q775_admin_bypasses_budget_ceiling(): litellm_changed_by=None, ) assert result is not None + + +@pytest.mark.asyncio +async def test_common_key_generation_helper_warms_auth_cache(): + """ + /key/generate must warm user_api_key_cache so the first authenticated + request after creation does not depend on instant DB visibility. + + Regression for case 2026-05-13-zurich-invalid-proxy-token: + multi-replica deployments behind pooled/replicated Postgres saw + token_not_found_in_db on the very first request after key creation. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + mock_cache = MagicMock() + mock_cache.async_set_cache = AsyncMock() + mock_proxy_logging = MagicMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.llm_router") as mock_router, + patch("litellm.proxy.proxy_server.premium_user", False), + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn" + ) as mock_generate_key, + ): + mock_prisma.return_value = AsyncMock() + mock_router.return_value = None + mock_generate_key.return_value = { + "key": "sk-test-warm-cache", + "token": "sk-test-warm-cache", + "token_id": "hashed-token-warm-cache", + "expires": None, + "user_id": "test-user", + "team_id": None, + } + + await _common_key_generation_helper( + data=GenerateKeyRequest(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" + ), + litellm_changed_by=None, + team_table=None, + ) + + mock_cache.async_set_cache.assert_awaited() + call_kwargs = mock_cache.async_set_cache.await_args.kwargs + assert call_kwargs["key"] == "hashed-token-warm-cache" + assert isinstance(call_kwargs["value"], UserAPIKeyAuth) + assert call_kwargs["value"].token == "hashed-token-warm-cache" + + +@pytest.mark.asyncio +async def test_common_key_generation_helper_cache_warm_is_best_effort(): + """ + A Redis (or any cache-layer) failure during the post-create cache warm + must NOT cause /key/generate to fail. The key is already in Postgres. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + mock_cache = MagicMock() + mock_cache.async_set_cache = AsyncMock(side_effect=ConnectionError("redis down")) + mock_proxy_logging = MagicMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.llm_router") as mock_router, + patch("litellm.proxy.proxy_server.premium_user", False), + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn" + ) as mock_generate_key, + ): + mock_prisma.return_value = AsyncMock() + mock_router.return_value = None + mock_generate_key.return_value = { + "key": "sk-test-cache-fail", + "token": "sk-test-cache-fail", + "token_id": "hashed-token-cache-fail", + "expires": None, + "user_id": "test-user", + "team_id": None, + } + + # Must not raise — cache failure is logged, response is still returned. + result = await _common_key_generation_helper( + data=GenerateKeyRequest(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" + ), + litellm_changed_by=None, + team_table=None, + ) + assert result.token == "hashed-token-cache-fail"