refactor(redis): move the auth cache invalidation codec into a store module

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:

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"))
+    )

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