fix(proxy/key_management): warm user_api_key_cache on /key/generate

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,

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

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"