fix(auth): stop re-caching stale auth objects so key updates propagate across replicas

After every successful virtual-key auth the proxy unconditionally wrote the
in-memory auth object back into DualCache. With enable_redis_auth_cache and
multiple replicas, a replica holding a stale in-memory token republished it to
shared Redis with a full TTL on every request, so /key/update and /key/delete
never took effect while the key kept calling. Both load paths already cache
the pristine object at load time (IdentityStore._resolve_key, get_key_object),
so the post-auth write only served to persist request-mutated state.

Two more writers of the same state carried the same defect. The per-request
key-spend writeback in update_cache republished the full auth object after
every priced request, and even scoped to the local pod it could race an
invalidation and re-insert a revoked key with a fresh TTL; it is removed, as
spend is tracked through the spend🔑* counters that budget enforcement
reads first. The DualCache Redis-to-memory backfill wrote values into the
in-memory cache without a ttl, so replicas hydrated from Redis held auth
objects for InMemoryCache's 600s default instead of user_api_key_cache_ttl;
the backfill now injects default_in_memory_ttl like async_set_cache does.

Resolves LIT-4350
This commit is contained in:
Yassin Kortam 2026-07-16 13:24:51 -07:00
parent 2162da5015
commit d1bf57e535
6 changed files with 223 additions and 31 deletions

View file

@ -160,6 +160,8 @@ class DualCache(BaseCache):
if redis_result is not None:
# Update in-memory cache with the value from Redis
if "ttl" not in kwargs and self.default_in_memory_ttl is not None:
kwargs["ttl"] = self.default_in_memory_ttl
self.in_memory_cache.set_cache(key, redis_result, **kwargs)
result = redis_result
@ -226,6 +228,8 @@ class DualCache(BaseCache):
if redis_result is not None:
# Update in-memory cache with the value from Redis
if "ttl" not in kwargs and self.default_in_memory_ttl is not None:
kwargs["ttl"] = self.default_in_memory_ttl
await self.in_memory_cache.async_set_cache(key, redis_result, **kwargs)
result = redis_result
@ -313,6 +317,9 @@ class DualCache(BaseCache):
# Pre-compute key-to-index mapping for O(1) lookup
key_to_index = {key: i for i, key in enumerate(keys)}
if "ttl" not in kwargs and self.default_in_memory_ttl is not None:
kwargs["ttl"] = self.default_in_memory_ttl
# Update both result and in-memory cache in a single loop
for key, value in redis_result.items():
result[key_to_index[key]] = value

View file

@ -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

View file

@ -2814,21 +2814,6 @@ async def update_cache(
)
# 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,14 @@ async def update_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,
if values_to_update_in_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,
)
)
)
def run_ollama_serve():

View file

@ -88,6 +88,62 @@ 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():
"""The Redis-to-memory backfill must honor default_in_memory_ttl.
On an in-memory miss with a Redis hit, async_get_cache writes the Redis
value into the in-memory cache. Without a ttl this write fell back to
InMemoryCache's own default_ttl (600s), so a replica whose copy of a
management object came from Redis held it ten times longer than the
configured TTL; a deleted or updated virtual key kept working on that
replica for up to 10 minutes instead of converging within
user_api_key_cache_ttl.
"""
in_memory_cache = InMemoryCache(default_ttl=600)
mock_redis = MagicMock(spec=RedisCache)
mock_redis.async_get_cache = AsyncMock(return_value="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_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():
"""

View file

@ -1,3 +1,4 @@
import asyncio
import json
import os
import sys
@ -220,6 +221,103 @@ async def test_should_not_reuse_cached_key_object_for_request_state():
assert second_request_key.request_route 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
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.2": 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")
)
mock_proxy_logging_obj = MagicMock()
mock_proxy_logging_obj.internal_usage_cache = MagicMock()
mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock()
mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
import litellm.proxy.proxy_server as _proxy_server_mod
_attrs_to_set = {
"prisma_client": MagicMock(),
"user_api_key_cache": key_cache,
"proxy_logging_obj": mock_proxy_logging_obj,
"master_key": "sk-master-key",
"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",
}
_original_values = {
attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set
}
try:
for attr, val in _attrs_to_set.items():
setattr(_proxy_server_mod, attr, val)
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.2": 3}}
finally:
for attr, val in _original_values.items():
setattr(_proxy_server_mod, attr, val)
@pytest.mark.asyncio
async def test_custom_auth_does_not_enforce_key_model_access_by_default():
valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"])

View file

@ -4528,6 +4528,61 @@ async def test_tag_cache_update_multiple_tags():
assert tag_updates["tag:tag2"]["spend"] == 25.0
@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._types import UserAPIKeyAuth
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_pipeline_honors_user_api_key_cache_ttl():
"""