From 84adc800f6d3ca191ced374af1bfffc12e43b05e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:53:34 +0000 Subject: [PATCH] refactor(redis): move the auth cache invalidation codec into a store module --- .../auth_cache_invalidation_pubsub.py | 57 +++---------------- .../redis/stores/auth_cache_invalidation.py | 50 ++++++++++++++++ .../stores/test_auth_cache_invalidation.py | 48 ++++++++++++++++ 3 files changed, 107 insertions(+), 48 deletions(-) create mode 100644 litellm/redis/stores/auth_cache_invalidation.py create mode 100644 tests/unit/redis/stores/test_auth_cache_invalidation.py diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index b2640d2a765..1b455d5550e 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -1,15 +1,9 @@ import asyncio -import json -from collections.abc import Mapping, Sequence -from dataclasses import asdict, dataclass -from types import MappingProxyType +from collections.abc import Sequence from typing import TYPE_CHECKING, Final -from pydantic import TypeAdapter - from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.proxy.common_utils.config_sync_pubsub import ( # noqa: F401 # legacy module exports ConfigSyncPubSub, _ConfigSyncPubSub, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -18,6 +12,12 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( # noqa: F401 # leg pubsub_capable_client, ) from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET +from litellm.redis.stores.auth_cache_invalidation import ( + AUTH_CACHE_INVALIDATION_CHANNEL, + CacheInvalidation, + decode_cache_invalidation, + encode_cache_invalidation, +) if TYPE_CHECKING: from litellm.caching.dual_cache import DualCache @@ -25,7 +25,6 @@ if TYPE_CHECKING: from litellm.caching.redis_cache import RedisCache from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -AUTH_CACHE_INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" _POLL_TIMEOUT_SECONDS: Final = 1.0 _MAX_PENDING_PUBLISHES: Final = 1024 _MAX_IN_FLIGHT_PUBLISHES: Final = 16 @@ -34,7 +33,6 @@ _pending_publishes: Final[set[asyncio.Task[None]]] = set() # mutable-ok: strong _in_flight_publishes: Final = asyncio.Semaphore(_MAX_IN_FLIGHT_PUBLISHES) _BACKOFF_INITIAL_SECONDS: Final = 5.0 _BACKOFF_MAX_SECONDS: Final = 60.0 -_DECODED_JSON: Final = TypeAdapter(object) def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str: @@ -43,45 +41,8 @@ def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str: return f"{redis_cache.namespace}:{AUTH_CACHE_INVALIDATION_CHANNEL}" -@dataclass(frozen=True, slots=True) -class _CacheInvalidationMessage: - cache_key: str - new_value: float | None = None - ttl: float | None = None - - def _cache_invalidation_message_json(cache_key: str, new_value: float | None = None, ttl: float | None = None) -> str: - message: Final[Mapping[str, object]] = MappingProxyType( - asdict(_CacheInvalidationMessage(cache_key=cache_key, new_value=new_value, ttl=ttl)) - ) - return json.dumps({field: value for field, value in message.items() if value is not None}) - - -def _finite_number_or_none(value: object) -> float | None: - if isinstance(value, bool) or not isinstance(value, (int, float)): - return None - return float(value) - - -def _message_from_data(data: object) -> _CacheInvalidationMessage | None: - if isinstance(data, bytes): - data = data.decode("utf-8", errors="replace") # rebind-ok: normalizing the wire payload to str - if not isinstance(data, str): - return None - try: - parsed: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(data))) - except json.JSONDecodeError: - return None - if parsed is None: - return None - cache_key: Final = parsed.get("cache_key") - if not isinstance(cache_key, str): - return None - return _CacheInvalidationMessage( - cache_key=cache_key, - new_value=_finite_number_or_none(parsed.get("new_value")), - ttl=_finite_number_or_none(parsed.get("ttl")), - ) + return encode_cache_invalidation(CacheInvalidation(cache_key=cache_key, new_value=new_value, ttl=ttl)) async def _publish_to_redis(redis_cache: "RedisCache", cache_key: str, message: str) -> None: @@ -294,7 +255,7 @@ class AuthCacheInvalidationSubscriber: @with_service_target(AUTH_OBJECTS_TARGET) def _apply_message(self, message: object) -> None: data: Final = message.get("data") if isinstance(message, dict) else None - parsed: Final = _message_from_data(data) + parsed: Final = decode_cache_invalidation(data) if parsed is None: return if parsed.new_value is not None: diff --git a/litellm/redis/stores/auth_cache_invalidation.py b/litellm/redis/stores/auth_cache_invalidation.py new file mode 100644 index 00000000000..ee069221ebb --- /dev/null +++ b/litellm/redis/stores/auth_cache_invalidation.py @@ -0,0 +1,50 @@ +import json +from dataclasses import dataclass +from typing import Final + +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.core_helpers import as_str_mapping + +AUTH_CACHE_INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" +_DECODED_JSON: Final = TypeAdapter(object) + + +@dataclass(frozen=True, slots=True) +class CacheInvalidation: + cache_key: str + new_value: float | None = None + ttl: float | None = None + + +def encode_cache_invalidation(invalidation: CacheInvalidation) -> str: + fields: Final[tuple[tuple[str, str | float | None], ...]] = ( + ("cache_key", invalidation.cache_key), + ("new_value", invalidation.new_value), + ("ttl", invalidation.ttl), + ) + return json.dumps({field: value for field, value in fields if value is not None}) + + +def _number_or_none(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return float(value) + + +def decode_cache_invalidation(data: object) -> CacheInvalidation | None: + text: Final = data.decode("utf-8", errors="replace") if isinstance(data, bytes) else data + if not isinstance(text, str): + return None + try: + parsed: Final = as_str_mapping(_DECODED_JSON.validate_python(json.loads(text))) + except json.JSONDecodeError: + return None + if parsed is None: + return None + cache_key: Final = parsed.get("cache_key") + if not isinstance(cache_key, str): + return None + return CacheInvalidation( + cache_key=cache_key, new_value=_number_or_none(parsed.get("new_value")), ttl=_number_or_none(parsed.get("ttl")) + ) diff --git a/tests/unit/redis/stores/test_auth_cache_invalidation.py b/tests/unit/redis/stores/test_auth_cache_invalidation.py new file mode 100644 index 00000000000..9b9c78dffb6 --- /dev/null +++ b/tests/unit/redis/stores/test_auth_cache_invalidation.py @@ -0,0 +1,48 @@ +import math +from typing import Final + +import pytest + +from litellm.redis.stores.auth_cache_invalidation import ( + CacheInvalidation, + decode_cache_invalidation, + encode_cache_invalidation, +) + +OLD_POD_MESSAGES: Final = ( + (b'{"cache_key": "project_id:p-1"}', CacheInvalidation("project_id:p-1")), + ( + b'{"cache_key": "spend:team_member:u-1:t-1", "new_value": 0.0, "ttl": 60.5}', + CacheInvalidation("spend:team_member:u-1:t-1", new_value=0.0, ttl=60.5), + ), + (b'{"cache_key": "k\\u96ea", "ttl": 30}', CacheInvalidation("k\u96ea", ttl=30.0)), +) + + +@pytest.mark.parametrize(("wire", "invalidation"), OLD_POD_MESSAGES) +def test_codec_round_trips_the_exact_bytes_old_pods_publish(wire: bytes, invalidation: CacheInvalidation) -> None: + assert decode_cache_invalidation(wire) == invalidation + assert decode_cache_invalidation(wire.decode()) == invalidation + if invalidation.ttl != 30.0: + assert encode_cache_invalidation(invalidation).encode() == wire + + +def test_encode_keeps_the_old_field_order_and_drops_unset_fields() -> None: + assert encode_cache_invalidation(CacheInvalidation("k", ttl=5.0)) == '{"cache_key": "k", "ttl": 5.0}' + assert encode_cache_invalidation(CacheInvalidation("k", new_value=1.0)) == '{"cache_key": "k", "new_value": 1.0}' + + +@pytest.mark.parametrize( + "data", + [b"not json", b'{"other": "x"}', b'{"cache_key": 7}', b"[1]", "raw string", None, 7], +) +def test_unusable_messages_decode_to_nothing(data: object) -> None: + assert decode_cache_invalidation(data) is None + + +def test_non_numeric_values_are_dropped_and_non_finite_numbers_pass_through_like_old_pods() -> None: + decoded: Final = decode_cache_invalidation(b'{"cache_key": "k", "new_value": true, "ttl": Infinity}') + + assert decoded is not None + assert (decoded.cache_key, decoded.new_value, math.isinf(decoded.ttl or 0.0)) == ("k", None, True) + assert decode_cache_invalidation(b'{"cache_key": "k", "new_value": "1"}') == CacheInvalidation("k")