refactor(redis): move Claude session router bindings behind a typed store

litellm/redis/stores/claude_code_router_binding.py

@@ -0,0 +1,69 @@
+import ast
+import json
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Final, Protocol
+
+from pydantic import JsonValue, TypeAdapter
+
+from litellm.redis.transport.contract import RedisFailure, RedisTransport
+
+_PREFIX: Final = "claude_code_session_router:v1:"
+_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
+_OBJECT: Final = TypeAdapter(object)
+
+
+@dataclass(frozen=True, slots=True)
+class SessionRouterScope:
+    caller: str
+    session_id: str
+
+
+def binding_key(scope: SessionRouterScope) -> str:
+    return f"{_PREFIX}{scope.caller}:{scope.session_id}"
+
+
+def binding_scope(key: str) -> SessionRouterScope | None:
+    if not key.startswith(_PREFIX):
+        return None
+    caller, separator, session_id = key[len(_PREFIX) :].rpartition(":")
+    return SessionRouterScope(caller, session_id) if separator and caller and session_id else None
+
+
+def _decode(raw: bytes) -> object:
+    text: Final = raw.decode("utf-8")
+    try:
+        return _JSON.validate_json(text)
+    except ValueError:
+        return _OBJECT.validate_python(ast.literal_eval(text))
+
+
+class ClaudeCodeRouterBindingStore(Protocol):
+    async def read_binding(self, scope: SessionRouterScope) -> str | None | RedisFailure: ...
+
+    async def bind_router(self, scope: SessionRouterScope, model: str, *, ttl: timedelta) -> None | RedisFailure: ...
+
+    async def remove_binding(self, scope: SessionRouterScope) -> None | RedisFailure: ...
+
+
+@dataclass(frozen=True, slots=True)
+class TransportClaudeCodeRouterBindingStore:
+    transport: RedisTransport
+
+    async def read_binding(self, scope: SessionRouterScope) -> str | None | RedisFailure:
+        raw: Final = await self.transport.get(binding_key(scope))
+        if raw is None or isinstance(raw, RedisFailure):
+            return raw
+        try:
+            value: Final = _decode(raw)
+        except (ValueError, SyntaxError, UnicodeError) as error:
+            return RedisFailure(kind="decode", message=str(error))
+        return value if isinstance(value, str) else None
+
+    async def bind_router(self, scope: SessionRouterScope, model: str, *, ttl: timedelta) -> None | RedisFailure:
+        result: Final = await self.transport.set_seconds(binding_key(scope), json.dumps(model), ttl=ttl)
+        return result if isinstance(result, RedisFailure) else None
+
+    async def remove_binding(self, scope: SessionRouterScope) -> None | RedisFailure:
+        result: Final = await self.transport.delete(binding_key(scope))
+        return result if isinstance(result, RedisFailure) else None

litellm/router.py

@@ -83,7 +83,6 @@ from litellm.caching.caching import (
     RedisCache,
     RedisClusterCache,
 )
-from litellm.caching.redis_cache import log_redis_failure
 from litellm.constants import (
     BAGGAGE_MAPPED_TAGS_METADATA_KEY,
     CLIENT_OUTPUT_CEILING_METADATA_KEY,
@@ -182,6 +181,13 @@ from litellm.llms.openai_like.model_info import (
     MODEL_INFO_REFRESH_SECONDS,
     get_openai_compatible_model_info,
 )
+from litellm.redis.stores.claude_code_router_binding import (
+    ClaudeCodeRouterBindingStore,
+    SessionRouterScope,
+    TransportClaudeCodeRouterBindingStore,
+    binding_key,
+    binding_scope,
+)
 from litellm.redis.stores.router_global_usage import (
     GlobalUsageCounter,
     RouterGlobalUsageStore,
@@ -16238,6 +16244,10 @@ class Router(DeprecatedRouterMembers):
         )
 
     def _claude_code_session_router_cache_key(self, request_kwargs: Mapping[str, object]) -> str | None:
+        scope: Final = self._claude_code_session_router_scope(request_kwargs)
+        return None if scope is None else binding_key(scope)
+
+    def _claude_code_session_router_scope(self, request_kwargs: Mapping[str, object]) -> SessionRouterScope | None:
         session_id: Final = self._request_header(request_kwargs, "x-claude-code-session-id")
         if session_id is None or _CLAUDE_CODE_SESSION_ID_RE.fullmatch(session_id) is None:
             return None
@@ -16248,33 +16258,52 @@ class Router(DeprecatedRouterMembers):
         caller_scope: Final = metadata.get("user_api_key_hash")
         if not isinstance(caller_scope, str) or not caller_scope:
             return None
-        return f"claude_code_session_router:v1:{caller_scope}:{session_id}"
+        return SessionRouterScope(caller_scope, session_id)
 
     @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET)
-    async def _delete_claude_code_session_router_binding(self, cache_key: str) -> None:
+    async def _delete_claude_code_session_router_binding(self, cache_key: SessionRouterScope | str) -> None:
+        scope: Final = binding_scope(cache_key) if isinstance(cache_key, str) else cache_key
+        if scope is None:
+            return
         try:
-            await self._claude_code_session_router_cache.async_delete_cache(key=cache_key)
-        except Exception as e:  # noqa: BLE001  # cache cleanup must not fail an otherwise routable request
-            verbose_router_logger.warning(
-                "Failed to delete Claude Code session router binding; the binding may remain until its TTL expires: %s",
-                e,
-            )
+            session_cache: Final = self._claude_code_session_router_cache
+            session_cache.in_memory_cache.delete_cache(binding_key(scope))
+            redis: Final = session_cache.redis_cache
+            if redis is not None:
+                store: Final[ClaudeCodeRouterBindingStore] = TransportClaudeCodeRouterBindingStore(redis.transport())
+                await store.remove_binding(scope)
+        except Exception as e:  # noqa: BLE001  # local cache cleanup must not fail an otherwise routable request
+            verbose_router_logger.warning("Failed to delete Claude Code session router binding: %s", e)
 
     @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET)
-    async def _get_claude_code_session_router_binding(self, cache_key: str) -> object:
-        session_cache: Final = self._claude_code_session_router_cache
-        try:
-            if session_cache.redis_cache is None:
-                return await session_cache.async_get_cache(key=cache_key)
-            return await session_cache.redis_cache.async_get_cache(key=cache_key)
-        except Exception as e:  # noqa: BLE001  # an optional binding must not make routing depend on Redis
-            log_redis_failure(
-                verbose_router_logger,
-                logging.WARNING,
-                "Failed to read Claude Code session router binding; using the requested model",
-                e,
-            )
+    async def _get_claude_code_session_router_binding(self, cache_key: SessionRouterScope | str) -> str | None:
+        scope: Final = binding_scope(cache_key) if isinstance(cache_key, str) else cache_key
+        if scope is None:
             return None
+        session_cache: Final = self._claude_code_session_router_cache
+        redis: Final = session_cache.redis_cache
+        if redis is None:
+            local: Final = _OPAQUE_VALUE_ADAPTER.validate_python(
+                await session_cache.async_get_cache(key=binding_key(scope), local_only=True)
+            )
+            return local if isinstance(local, str) else None
+        store: Final[ClaudeCodeRouterBindingStore] = TransportClaudeCodeRouterBindingStore(redis.transport())
+        shared: Final = await store.read_binding(scope)
+        return None if isinstance(shared, RedisFailure) else shared
+
+    @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET)
+    async def _bind_claude_code_session_router(self, scope: SessionRouterScope, model: str) -> None:
+        session_cache: Final = self._claude_code_session_router_cache
+        await session_cache.async_set_cache(
+            key=binding_key(scope),
+            value=model,
+            ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS,
+            local_only=True,
+        )
+        redis: Final = session_cache.redis_cache
+        if redis is not None:
+            store: Final[ClaudeCodeRouterBindingStore] = TransportClaudeCodeRouterBindingStore(redis.transport())
+            await store.bind_router(scope, model, ttl=timedelta(seconds=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS))
 
     @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET)
     async def _resolve_claude_code_session_router(
@@ -16294,36 +16323,28 @@ class Router(DeprecatedRouterMembers):
             )
         ):
             return registered_model_name
-        cache_key: Final = self._claude_code_session_router_cache_key(request_kwargs)
-        if cache_key is None or not isinstance(request_kwargs, dict):
+        scope: Final = self._claude_code_session_router_scope(request_kwargs)
+        if scope is None or not isinstance(request_kwargs, dict):
             return registered_model_name
         if request_kwargs.get("fallback_depth") not in (None, 0):
             return registered_model_name
 
         agent_id: Final = self._request_header(request_kwargs, "x-claude-code-agent-id")
         if agent_id is not None:
-            bound_model: Final = await self._get_claude_code_session_router_binding(cache_key)
+            bound_model: Final = await self._get_claude_code_session_router_binding(scope)
             if not isinstance(bound_model, str):
                 return registered_model_name
             bound_registered_model: Final = self._core_get_model_from_alias(model=bound_model) or bound_model
             if self._select_pre_routing_strategy(bound_registered_model, request_kwargs) is None:
-                await self._delete_claude_code_session_router_binding(cache_key)
+                await self._delete_claude_code_session_router_binding(scope)
                 return registered_model_name
-            await self._claude_code_session_router_cache.async_set_cache(
-                key=cache_key,
-                value=bound_model,
-                ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS,
-            )
+            await self._bind_claude_code_session_router(scope, bound_model)
             self._stamp_or_clear_metadata_key(request_kwargs, "model_group", bound_model)
             return bound_registered_model
 
         if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None:
             return registered_model_name
-        await self._claude_code_session_router_cache.async_set_cache(
-            key=cache_key,
-            value=model,
-            ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS,
-        )
+        await self._bind_claude_code_session_router(scope, model)
         return registered_model_name
 
     async def _core_async_pre_routing_hook(

tests/code_coverage_tests/router_code_coverage.py

@@ -177,6 +177,8 @@ ignored_function_names = [
     "_responses_continuation_ceilings",  # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py
     "_read_global_usage_remote",  # Tested through the shared Redis usage header tests in test_router_helper_utils.py
     "_set_global_usage_local",  # Tested through the shared Redis usage header tests in test_router_helper_utils.py
+    "_claude_code_session_router_scope",  # Tested through the Claude Code session router binding tests in test_router.py
+    "_bind_claude_code_session_router",  # Tested through the Claude Code session router binding tests in test_router.py
     "_team_deployments_named",  # Tested through the team lookups in the concurrent delete tests in test_router.py
     "_served_openai_models",  # Cached as served_openai_models, tested through the pre-call check token count tests in test_router.py
     "priority_header_update",  # Tested through the priority header tests in test_common_request_processing.py

tests/unit/redis/stores/test_claude_code_router_binding.py

@@ -0,0 +1,223 @@
+import json
+from collections.abc import Iterator, Mapping
+from datetime import timedelta
+from typing import Final, Protocol, cast
+from unittest.mock import patch
+
+import fakeredis
+import pytest
+from pydantic import TypeAdapter
+from redis.exceptions import ConnectionError as RedisConnectionError
+
+from litellm import Router
+from litellm.caching.redis_cache import RedisCache
+from litellm.redis.stores.claude_code_router_binding import SessionRouterScope, TransportClaudeCodeRouterBindingStore
+from litellm.redis.transport.contract import Arg, RedisFailure
+from litellm.types.router import PreRoutingHookResponse
+from tests._support.fake_redis import fake_redis_cache
+
+KEY: Final = "claude_code_session_router:v1:key:scope:session-1234"
+MODEL: Final = 'smart-雪 "router"'
+SCOPE: Final = SessionRouterScope("key:scope", "session-1234")
+_METADATA: Final = TypeAdapter(dict[str, object])
+
+
+class BindingRouter(Protocol):
+    def update_redis_cache(self, cache: RedisCache) -> None: ...
+
+    async def async_pre_routing_hook(
+        self, model: str, request_kwargs: Mapping[str, object]
+    ) -> PreRoutingHookResponse | None: ...
+
+
+class LegacyBindingCache(Protocol):
+    async def async_get_cache(self, key: str) -> object: ...
+
+    async def async_set_cache(self, key: str, value: str, *, ttl: int) -> None: ...
+
+
+@pytest.fixture(autouse=True)
+def clock() -> Iterator[None]:
+    with patch("time.time", return_value=1800000000.0):
+        yield
+
+
+def _router(cache: RedisCache | None) -> BindingRouter:
+    router: Final = Router(
+        model_list=[
+            {"model_name": "plain", "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "fake"}},
+            *(
+                {
+                    "model_name": model,
+                    "litellm_params": {
+                        "model": "auto_router/complexity_router",
+                        "complexity_router_config": {
+                            "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), "plain")
+                        },
+                    },
+                }
+                for model in (MODEL, "premium-router")
+            ),
+        ]
+    )
+    typed: Final = cast(BindingRouter, router)  # cast-ok: typed view of the actual binding caller, without replacing it
+    if cache is not None:
+        typed.update_redis_cache(cache)
+    return typed
+
+
+def _kwargs(
+    *, agent: bool = False, caller: str = "key:scope", session: str = "session-1234", depth: int = 0
+) -> Mapping[str, object]:
+    return {
+        "metadata": {"user_api_key_hash": caller},
+        "proxy_server_request": {
+            "headers": {
+                "X-Claude-Code-Session-Id": session,
+                **({"x-claude-code-agent-id": "agent-1234"} if agent else {}),
+            }
+        },
+        "fallback_depth": depth,
+    }
+
+
+async def _route(router: BindingRouter, model: str, kwargs: Mapping[str, object]) -> str:
+    await router.async_pre_routing_hook(model, dict(kwargs))
+    bound: Final = _METADATA.validate_python(kwargs["metadata"]).get("model_group")
+    return bound if isinstance(bound, str) else model
+
+
+@pytest.mark.asyncio
+async def test_binding_old_new_wire_namespace_and_read_refresh(server: fakeredis.FakeServer) -> None:
+    cache: Final = fake_redis_cache(server, namespace="tenant")
+    legacy: Final = cast(LegacyBindingCache, cache)  # cast-ok: rolling-deploy writer and reader
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    await legacy.async_set_cache(KEY, MODEL, ttl=3600)
+    router: Final = _router(cache)
+    assert await _route(router, "plain", _kwargs(agent=True)) == MODEL
+    assert await raw.get("tenant:" + KEY) == json.dumps(MODEL).encode()
+    assert await raw.exists(KEY) == 0
+    with patch("time.time", return_value=1800000100.0):
+        assert await raw.pttl("tenant:" + KEY) == 3500000
+        assert await _route(router, "plain", _kwargs(agent=True)) == MODEL
+        assert await raw.pttl("tenant:" + KEY) == 3600000
+        assert await _route(router, "premium-router", _kwargs()) == "premium-router"
+    assert await legacy.async_get_cache(KEY) == "premium-router"
+    assert await raw.get("tenant:" + KEY) == b'"premium-router"'
+
+
+@pytest.mark.asyncio
+async def test_shared_miss_and_remote_update_override_a_fresh_local_binding(server: fakeredis.FakeServer) -> None:
+    cache: Final = fake_redis_cache(server)
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    main: Final = _router(cache)
+    agent: Final = _router(cache)
+    await _route(main, MODEL, _kwargs())
+    assert await _route(agent, "plain", _kwargs(agent=True)) == MODEL
+    await _route(main, "premium-router", _kwargs())
+    assert await _route(agent, "plain", _kwargs(agent=True)) == "premium-router"
+    await raw.delete(KEY)
+    assert await _route(agent, "plain", _kwargs(agent=True)) == "plain"
+
+
+@pytest.mark.asyncio
+async def test_local_bindings_survive_without_redis_but_are_not_backfilled_on_attach(
+    server: fakeredis.FakeServer,
+) -> None:
+    router: Final = _router(None)
+    await _route(router, MODEL, _kwargs())
+    assert await _route(router, "plain", _kwargs(agent=True)) == MODEL
+    router.update_redis_cache(fake_redis_cache(server))
+    assert await _route(router, "plain", _kwargs(agent=True)) == "plain"
+    assert await fakeredis.FakeAsyncRedis(server=server).exists(KEY) == 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("raw_value", [b"null", b"4", b"[]", b'{"router": "other"}', b"not-json", b"\xff"])
+async def test_non_string_and_malformed_bindings_are_ignored_not_rewritten(
+    server: fakeredis.FakeServer, raw_value: bytes
+) -> None:
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    await raw.set(KEY, raw_value, ex=600)
+    assert await _route(_router(fake_redis_cache(server)), "plain", _kwargs(agent=True)) == "plain"
+    assert await raw.get(KEY) == raw_value
+    assert await raw.pttl(KEY) == 600000
+
+
+@pytest.mark.asyncio
+async def test_legacy_python_literal_binding_refreshes_as_json(server: fakeredis.FakeServer) -> None:
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    await raw.set(KEY, b"'premium-router'", ex=600)
+    assert await _route(_router(fake_redis_cache(server)), "plain", _kwargs(agent=True)) == "premium-router"
+    assert await raw.get(KEY) == b'"premium-router"'
+    assert await raw.pttl(KEY) == 3600000
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("value", [b'"removed-router"', b'""'])
+async def test_unknown_strategy_deletes_the_binding_instead_of_tombstoning(
+    server: fakeredis.FakeServer, value: bytes
+) -> None:
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    await raw.set(KEY, value, ex=3600)
+    assert await _route(_router(fake_redis_cache(server)), "plain", _kwargs(agent=True)) == "plain"
+    assert await raw.exists(KEY) == 0
+
+
+@pytest.mark.asyncio
+async def test_invalid_scope_other_caller_and_fallback_cannot_reuse_the_binding(server: fakeredis.FakeServer) -> None:
+    router: Final = _router(fake_redis_cache(server))
+    await _route(router, MODEL, _kwargs())
+    assert await _route(router, "plain", _kwargs(agent=True, caller="other")) == "plain"
+    assert await _route(router, "plain", _kwargs(agent=True, caller="")) == "plain"
+    assert await _route(router, "plain", _kwargs(agent=True, session="short")) == "plain"
+    assert await _route(router, "plain", _kwargs(agent=True, depth=1)) == "plain"
+    assert await fakeredis.FakeAsyncRedis(server=server).get(KEY) == json.dumps(MODEL).encode()
+
+
+@pytest.mark.asyncio
+async def test_binding_outage_never_uses_stale_local_and_shares_the_breaker() -> None:
+    server: Final = fakeredis.FakeServer()
+    cache: Final = fake_redis_cache(server)
+    router: Final = _router(cache)
+    await _route(router, MODEL, _kwargs())
+    server.connected = False
+    assert await _route(router, "plain", _kwargs(agent=True)) == "plain"
+    assert await _route(router, "premium-router", _kwargs()) == "premium-router"
+    store: Final = TransportClaudeCodeRouterBindingStore(cache.transport())
+    for _ in range(5):
+        await store.read_binding(SCOPE)
+    result: Final = await store.bind_router(SCOPE, MODEL, ttl=timedelta(seconds=3600))
+    assert isinstance(result, RedisFailure) and result.kind == "breaker_open"
+
+
+class RedisWriter(Protocol):
+    async def set(
+        self, name: str, value: Arg, *, ex: timedelta | None = None, px: int | None = None, nx: bool = False
+    ) -> object: ...
+
+
+@pytest.mark.asyncio
+async def test_unknown_applied_binding_write_remains_routable_without_replay(server: fakeredis.FakeServer) -> None:
+    cache: Final = fake_redis_cache(server)
+    router: Final = _router(cache)
+    original_set: Final = cast(type[RedisWriter], fakeredis.FakeAsyncRedis).set  # cast-ok: typed redis-py boundary
+
+    async def apply_then_disconnect(
+        client: RedisWriter,
+        name: str,
+        value: Arg,
+        *,
+        ex: timedelta | None = None,
+        px: int | None = None,
+        nx: bool = False,
+    ) -> object:
+        await original_set(client, name, value, ex=ex, px=px, nx=nx)
+        raise RedisConnectionError("response lost after SET applied")
+
+    with patch.object(fakeredis.FakeAsyncRedis, "set", new=apply_then_disconnect):
+        assert await _route(router, MODEL, _kwargs()) == MODEL
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    assert await raw.get(KEY) == json.dumps(MODEL).encode()
+    assert await raw.pttl(KEY) == 3600000
+    assert await _route(_router(cache), "plain", _kwargs(agent=True)) == MODEL