From e482d7d3214031e944c6208f42cf3eed41d9aafd Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 16 Jul 2026 13:58:54 -0700 Subject: [PATCH] fix(proxy): stop stale auth cache re-publish so key updates and deletes propagate across replicas With enable_redis_auth_cache and multiple replicas, /key/update and /key/delete delete the Redis auth blob and the handling pod's in-memory entry, but two read-path writers re-published the stale blob from any other replica's per-pod memory back to Redis with a fresh 60s TTL on every request: the post-auth re-cache in user_api_key_auth and the spend writeback in update_cache. Replicas whose in-memory entries expired then re-primed themselves from the poisoned Redis entry, so key limit and access changes never took effect fleet-wide while traffic continued, and a deleted key kept authenticating. The auth object is now written only by the DB-load paths (IdentityStore._resolve_key, get_key_object): the post-auth re-cache is removed outright (even a local-only write could race an invalidation and resurrect a revoked key on this worker) and spend tracking no longer writes the auth object back at all; spend is tracked through the spend:key:* counters. The remaining spend writebacks for user, team, end-user, and tag objects become local-only so they cannot republish stale management objects either, with one deliberate exception: the proxy-wide {litellm_proxy_admin_name}:spend scalar keeps its shared Redis write because the global max_budget check reads it between authoritative DB reloads, and it carries no limits or permissions so sharing it cannot resurrect an invalidated auth blob. DualCache's redis-to-memory read backfill also ignored default_in_memory_ttl, pinning backfilled entries for InMemoryCache's 600s default instead of the configured 60s auth TTL; the backfill now injects the configured default like every write path already does, so a replica primed from Redis converges within the auth cache TTL as well. Consolidates the sibling stale-auth-recache branch; the delete-propagation case is the duplicate ticket LIT-4350. Resolves LIT-4219 --- litellm/caching/dual_cache.py | 18 ++- litellm/proxy/auth/user_api_key_auth.py | 10 -- litellm/proxy/proxy_server.py | 41 +++--- tests/test_litellm/caching/test_dual_cache.py | 54 ++++++++ .../proxy/auth/test_auth_checks.py | 43 ++++++ .../proxy/auth/test_user_api_key_auth.py | 94 ++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 122 ++++++++++++++++++ 7 files changed, 348 insertions(+), 34 deletions(-) diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index be618815a537..0e3c93946fdd 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -103,6 +103,18 @@ def attach_redis_cache( if default_redis_ttl is not None: self.default_redis_ttl = default_redis_ttl + def _backfill_kwargs(self, kwargs: "dict[str, object]") -> "dict[str, object]": + """ + Kwargs for writing a Redis read result into the in-memory tier. + + Applies ``default_in_memory_ttl`` exactly like the write paths do; + without it, backfilled entries fall to ``InMemoryCache``'s own default + TTL and can outlive the TTL this cache was configured with. + """ + if "ttl" not in kwargs and self.default_in_memory_ttl is not None: + return {**kwargs, "ttl": self.default_in_memory_ttl} + return kwargs + def set_cache(self, key, value, local_only: bool = False, **kwargs): # Update both Redis and in-memory cache try: @@ -160,7 +172,7 @@ def get_cache( if redis_result is not None: # Update in-memory cache with the value from Redis - self.in_memory_cache.set_cache(key, redis_result, **kwargs) + self.in_memory_cache.set_cache(key, redis_result, **self._backfill_kwargs(kwargs)) result = redis_result @@ -226,7 +238,7 @@ async def async_get_cache( if redis_result is not None: # Update in-memory cache with the value from Redis - await self.in_memory_cache.async_set_cache(key, redis_result, **kwargs) + await self.in_memory_cache.async_set_cache(key, redis_result, **self._backfill_kwargs(kwargs)) result = redis_result @@ -318,7 +330,7 @@ async def async_batch_get_cache( result[key_to_index[key]] = value if value is not None and self.in_memory_cache is not None: - await self.in_memory_cache.async_set_cache(key, value, **kwargs) + await self.in_memory_cache.async_set_cache(key, value, **self._backfill_kwargs(kwargs)) return result except Exception: diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 0519b5ef0b64..4d07d4c043cc 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1996,16 +1996,6 @@ async def _user_api_key_auth_builder( raise HTTPException(401, detail="Invalid API key, no token associated") api_key = valid_token.token - # Add hashed token to cache - asyncio.create_task( - _cache_key_object( - hashed_token=api_key, - user_api_key_obj=valid_token, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - valid_token_dict = valid_token.model_dump(exclude_none=True) valid_token_dict.pop("token", None) # budget_throttle_pct is excluded from model_dump (it must not leak diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7e6ca5a01088..fad63cb01c35 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2814,21 +2814,6 @@ async def _update_key_cache(token: str, response_cost: float): ) # set cooldown on alert - if existing_spend_obj is not None and getattr(existing_spend_obj, "team_spend", None) is not None: - existing_team_spend = existing_spend_obj.team_spend or 0 - # Calculate the new cost by adding the existing cost and response_cost - existing_spend_obj.team_spend = existing_team_spend + response_cost - - if existing_spend_obj is not None and getattr(existing_spend_obj, "team_member_spend", None) is not None: - existing_team_member_spend = existing_spend_obj.team_member_spend or 0 - # Calculate the new cost by adding the existing cost and response_cost - existing_spend_obj.team_member_spend = existing_team_member_spend + response_cost - - # Existing spend_obj is mutated; UserApiKeyCache.async_set_cache_pipeline turns - # BaseModel values into dicts for Redis (same Codec path as async_set_cache). - existing_spend_obj.spend = new_spend - values_to_update_in_cache.append((hashed_token, existing_spend_obj)) - ### UPDATE USER SPEND ### async def _update_user_cache(): ## UPDATE CACHE FOR USER ID + GLOBAL PROXY @@ -3032,13 +3017,27 @@ async def _update_tag_cache(): if tags is not None: await _update_tag_cache() - asyncio.create_task( - user_api_key_cache.async_set_cache_pipeline( - cache_list=values_to_update_in_cache, - ttl=get_management_object_ttl(user_api_key_cache), - litellm_parent_otel_span=parent_otel_span, + global_proxy_spend_key = "{}:spend".format(litellm_proxy_admin_name) + local_object_updates = tuple((k, v) for k, v in values_to_update_in_cache if k != global_proxy_spend_key) + shared_scalar_updates = tuple((k, v) for k, v in values_to_update_in_cache if k == global_proxy_spend_key) + + if local_object_updates: + asyncio.create_task( + user_api_key_cache.async_set_cache_pipeline( + cache_list=list(local_object_updates), + ttl=get_management_object_ttl(user_api_key_cache), + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + ) + if shared_scalar_updates: + asyncio.create_task( + user_api_key_cache.async_set_cache_pipeline( + cache_list=list(shared_scalar_updates), + ttl=get_management_object_ttl(user_api_key_cache), + litellm_parent_otel_span=parent_otel_span, + ) ) - ) def run_ollama_serve(): diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index f4f88def78d3..47be139eb5ee 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -88,6 +88,60 @@ async def test_dual_cache_async_set_cache_injects_default_in_memory_ttl(): assert expiry <= after + 60 +@pytest.mark.asyncio +async def test_dual_cache_redis_backfill_injects_default_in_memory_ttl(): + """ + A Redis-hit backfill into the in-memory tier must honor + default_in_memory_ttl the same way the write paths do. Without it, the + backfilled entry falls to InMemoryCache's own default_ttl (600s), so a + replica that primed a management object (e.g. a virtual key's auth blob) + from Redis keeps serving it for 10 minutes after the object was updated + and invalidated, instead of re-reading within the configured TTL. + """ + in_memory_cache = InMemoryCache(default_ttl=600) + redis_cache = MagicMock() + redis_cache.async_get_cache = AsyncMock(return_value="redis_value") + dual_cache = DualCache( + in_memory_cache=in_memory_cache, + redis_cache=redis_cache, + default_in_memory_ttl=60, + ) + + before = time.time() + result = await dual_cache.async_get_cache(key="backfill_key") + after = time.time() + + assert result == "redis_value" + expiry = in_memory_cache.ttl_dict["backfill_key"] + assert expiry >= before + 60 + assert expiry <= after + 60 + + +@pytest.mark.asyncio +async def test_dual_cache_batch_redis_backfill_injects_default_in_memory_ttl(): + """async_batch_get_cache's Redis-to-memory backfill must honor + default_in_memory_ttl, same as the single-key path.""" + in_memory_cache = InMemoryCache(default_ttl=600) + mock_redis = MagicMock(spec=RedisCache) + mock_redis.async_batch_get_cache = AsyncMock( + return_value={"batch_backfill_key": "redis_value"} + ) + dual_cache = DualCache( + in_memory_cache=in_memory_cache, + redis_cache=mock_redis, + default_in_memory_ttl=60, + ) + + before = time.time() + result = await dual_cache.async_batch_get_cache(keys=["batch_backfill_key"]) + after = time.time() + + assert result == ["redis_value"] + expiry = in_memory_cache.ttl_dict["batch_backfill_key"] + assert expiry >= before + 60 + assert expiry <= after + 60 + + @pytest.mark.asyncio async def test_dual_cache_async_set_cache_respects_explicit_ttl(): """ diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 8365909f3141..cc4a7d5bfb4e 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -522,6 +522,49 @@ async def test_get_key_object_should_raise_if_reconnect_fails_on_db_connection_e assert mock_prisma_client.get_data.await_count == 1 +def _fake_redis_cache(): + fake_redis = MagicMock() + fake_redis.async_get_cache = AsyncMock(return_value=None) + fake_redis.async_set_cache = AsyncMock() + fake_redis.async_set_cache_pipeline = AsyncMock() + fake_redis.async_delete_cache = AsyncMock() + return fake_redis + + +class TestAuthCacheRedisWritePolicy: + """Redis auth-cache entries may only be written from fresh DB loads. + + With ``enable_redis_auth_cache`` and multiple replicas, a pod that re-publishes + a cache-derived key object to Redis can resurrect a stale auth blob after + ``/key/update`` or ``/key/delete`` already deleted it, so limit changes never + propagate fleet-wide while traffic keeps refreshing the stale entry's TTL. + """ + + @pytest.mark.asyncio + async def test_get_key_object_db_load_publishes_to_redis(self): + mock_prisma_client = MagicMock() + mock_prisma_client.get_data = AsyncMock( + return_value=UserAPIKeyAuth(token="hashed-token-db") + ) + + fake_redis = _fake_redis_cache() + cache = UserApiKeyCache() + cache.redis_cache = fake_redis + + key_obj = await get_key_object( + hashed_token="hashed-token-db", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) + + assert key_obj.token == "hashed-token-db" + fake_redis.async_set_cache.assert_awaited_once() + assert ( + fake_redis.async_set_cache.await_args.kwargs.get("key") + or fake_redis.async_set_cache.await_args.args[0] + ) == "hashed-token-db" + + def test_get_cli_jwt_auth_token_default_expiration(valid_sso_user_defined_values): """Test generating CLI JWT token with default 24-hour expiration""" token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index a0248963cf1a..9ac22086d92f 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -1,3 +1,4 @@ +import asyncio import json import os import sys @@ -4161,6 +4162,99 @@ async def test_auth_path_caches_team_object_under_canonical_team_id_key(): assert cache.get_cache(key=None) is None +@pytest.mark.asyncio +async def test_auth_does_not_rewrite_cached_key_object_back_into_cache(): + """A cache-hit auth must not write the token back into the cache. + + Re-writing on every auth let a replica holding a stale in-memory token + republish it to shared Redis with a fresh TTL on each request, so + /key/update and /key/delete never propagated across replicas or regional + Redis while the key kept calling (stale auth re-cache feedback loop). + Only the DB-load paths (IdentityStore._resolve_key / get_key_object) may + populate the cache. + """ + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.proxy_server import hash_token + + api_key = "sk-lit-cached-key-no-rewrite" + hashed_key = hash_token(api_key) + + key_cache = UserApiKeyCache() + stale_token = UserAPIKeyAuth( + api_key=api_key, + token=hashed_key, + metadata={"model_rpm_limit": {"gpt-5.4-mini": 3}}, + last_refreshed_at=1000.0, + ) + await key_cache.async_set_cache( + key=hashed_key, value=stale_token, model_type=UserAPIKeyAuth + ) + + fetch_from_db = AsyncMock( + side_effect=AssertionError("cache-hit auth must not touch the DB") + ) + + proxy_logging_obj = MagicMock() + proxy_logging_obj.internal_usage_cache = MagicMock() + proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + attrs = { + "prisma_client": MagicMock(), + "user_api_key_cache": key_cache, + "proxy_logging_obj": proxy_logging_obj, + "master_key": "sk-test-master", + "general_settings": {}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + with patch( + "litellm.proxy.auth.resolvers.store._fetch_key_object_from_db_with_reconnect", + fetch_from_db, + ): + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()] + if pending: + await asyncio.wait(pending, timeout=5) + + assert result.token == hashed_key + fetch_from_db.assert_not_called() + + cached_after = await key_cache.async_get_cache( + key=hashed_key, model_type=UserAPIKeyAuth + ) + assert cached_after is not None + assert cached_after.last_refreshed_at == 1000.0 + assert cached_after.metadata == {"model_rpm_limit": {"gpt-5.4-mini": 3}} + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + class TestCheckKeyModelBudgetWithFallback: """`_check_key_model_budget_with_fallback` must reroute a request to the first configured `budget_fallbacks` entry still within its own budget, diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 079b844638ce..7efce05c1deb 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -4568,6 +4568,128 @@ async def test_update_cache_pipeline_honors_user_api_key_cache_ttl(): setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache) +@pytest.mark.asyncio +async def test_spend_tracking_never_writes_the_auth_object_back(): + """Spend tracking must never write the auth object back into the cache. + + Writing the mutated auth object back after every priced request let a + stale copy be re-published with a fresh TTL: to shared Redis it defeated + /key/update and /key/delete across replicas, and even a local-only write + could race an invalidation and resurrect a revoked key on this worker. + Spend is tracked through the spend:key:* counters, so the auth object is + only ever written by the DB-load paths. + """ + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + original_cache = litellm.proxy.proxy_server.user_api_key_cache + cache = UserApiKeyCache() + setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache) + try: + hashed_token = "spend-tracking-no-writeback-token" + await cache.async_set_cache( + key=hashed_token, + value=UserAPIKeyAuth(token=hashed_token, spend=1.0), + model_type=UserAPIKeyAuth, + ) + with ( + patch.object( + cache, "async_set_cache_pipeline", new=AsyncMock() + ) as mock_pipeline, + patch.object(cache, "async_set_cache", new=AsyncMock()) as mock_set, + ): + await litellm.proxy.proxy_server.update_cache( + token=hashed_token, + user_id=None, + end_user_id=None, + team_id=None, + response_cost=5.0, + parent_otel_span=None, + ) + pending = [ + t for t in asyncio.all_tasks() if t is not asyncio.current_task() + ] + if pending: + await asyncio.wait(pending, timeout=5) + + key_pipeline_writes = [ + call + for call in mock_pipeline.call_args_list + if any(k == hashed_token for k, _ in call.kwargs["cache_list"]) + ] + assert key_pipeline_writes == [] + mock_set.assert_not_called() + finally: + setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache) + + +@pytest.mark.asyncio +async def test_update_cache_global_proxy_spend_scalar_stays_shared(): + """ + The proxy-wide spend estimate must keep flowing to Redis when the spend + writeback goes per-pod: the global max_budget check reads the + ``{litellm_proxy_admin_name}:spend`` cache entry between authoritative DB + reloads, so keeping it pod-local would let traffic spread across replicas + exceed the proxy budget by roughly a factor of the replica count within a + cache TTL. Sharing this scalar is safe because it carries no limits or + permissions, so it cannot resurrect an invalidated auth blob. + """ + from litellm.caching.caching import DualCache + + admin_name = litellm.proxy.proxy_server.litellm_proxy_admin_name + global_key = "{}:spend".format(admin_name) + + async def fake_get(key, **kwargs): + if key == "user-lit": + return {"user_id": "user-lit", "spend": 1.0} + if key == global_key: + return 10.0 + return None + + original_cache = litellm.proxy.proxy_server.user_api_key_cache + cache = DualCache(default_in_memory_ttl=300) + setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache) + try: + with patch.object( + cache, "async_get_cache", new=AsyncMock(side_effect=fake_get) + ): + with patch.object( + cache, "async_set_cache_pipeline", new=AsyncMock() + ) as mock_set_cache: + await litellm.proxy.proxy_server.update_cache( + token=None, + user_id="user-lit", + end_user_id=None, + team_id=None, + response_cost=5.0, + parent_otel_span=None, + ) + + pending = [ + t for t in asyncio.all_tasks() if t is not asyncio.current_task() + ] + if pending: + await asyncio.wait(pending, timeout=5) + + calls = mock_set_cache.await_args_list + local_keys = [ + k + for c in calls + if c.kwargs.get("local_only") is True + for k, _ in c.kwargs["cache_list"] + ] + shared_keys = [ + k + for c in calls + if c.kwargs.get("local_only") is not True + for k, _ in c.kwargs["cache_list"] + ] + assert "user-lit" in local_keys + assert global_key not in local_keys + assert shared_keys == [global_key] + finally: + setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache) + + @pytest.mark.asyncio async def test_init_sso_settings_in_db(): """