refactor(redis): migrate deployment session affinity into a typed store

litellm/redis/stores/deployment_session_affinity.py

@@ -0,0 +1,135 @@
+import ast
+import json
+from dataclasses import dataclass
+from typing import Final, Protocol
+
+from pydantic import JsonValue, TypeAdapter, ValidationError
+
+from litellm.redis.transport.contract import RedisFailure, RedisTransport
+
+_JSON: Final = TypeAdapter[JsonValue](JsonValue)
+_OBJECT: Final = TypeAdapter(object)
+_CLAIM: Final = """
+local current = redis.call('GET', KEYS[1])
+if current == false then
+  redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])
+  return ARGV[1]
+end
+if ARGV[3] then
+  local decoded, stored = pcall(cjson.decode, current)
+  local legacy = current
+  if decoded and type(stored) == 'string' then legacy = stored end
+  for _, eligible in ipairs(cjson.decode(ARGV[3])) do
+    if eligible['model_id'] == legacy then
+      redis.call('EXPIRE', KEYS[1], ARGV[2])
+      return current
+    end
+  end
+  if decoded and type(stored) == 'table' then
+    for _, eligible in ipairs(cjson.decode(ARGV[3])) do
+      local matches = true
+      for key, value in pairs(eligible) do
+        if stored[key] ~= value then matches = false; break end
+      end
+      for key, _ in pairs(stored) do
+        if eligible[key] == nil then matches = false; break end
+      end
+      if matches then
+        redis.call('EXPIRE', KEYS[1], ARGV[2])
+        return current
+      end
+    end
+  end
+  redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])
+  return ARGV[1]
+end
+if current == ARGV[1] then
+  redis.call('EXPIRE', KEYS[1], ARGV[2])
+end
+return current
+"""
+
+
+def affinity_key(model_group: str, caller_scope: str, session_id: str) -> str:
+    return f"deployment_affinity:v1:session:{model_group}:{caller_scope}:{session_id}"
+
+
+@dataclass(frozen=True, slots=True)
+class SessionPin:
+    value: object
+
+
+@dataclass(frozen=True, slots=True)
+class UnconfirmedSessionPin:
+    pass
+
+
+class DeploymentSessionAffinityStore(Protocol):
+    async def read(self, model_group: str, caller_scope: str, session_id: str) -> SessionPin | RedisFailure: ...
+
+    async def claim(
+        self,
+        model_group: str,
+        caller_scope: str,
+        session_id: str,
+        model_id: str,
+        *,
+        ttl_seconds: int,
+        eligible_model_ids: frozenset[str] | None = None,
+    ) -> SessionPin | UnconfirmedSessionPin | RedisFailure: ...
+
+
+@dataclass(frozen=True, slots=True)
+class TransportDeploymentSessionAffinityStore:
+    transport: RedisTransport
+
+    async def read(self, model_group: str, caller_scope: str, session_id: str) -> SessionPin | RedisFailure:
+        raw: Final = await self.transport.get(affinity_key(model_group, caller_scope, session_id))
+        if isinstance(raw, RedisFailure):
+            return raw
+        if raw is None:
+            return SessionPin(None)
+        try:
+            decoded: Final = raw.decode("utf-8")
+            try:
+                return SessionPin(_JSON.validate_json(decoded))
+            except ValidationError:
+                return SessionPin(_OBJECT.validate_python(ast.literal_eval(decoded)))
+        except (UnicodeError, ValueError, SyntaxError) as error:
+            return RedisFailure(kind="decode", message=str(error))
+
+    async def claim(
+        self,
+        model_group: str,
+        caller_scope: str,
+        session_id: str,
+        model_id: str,
+        *,
+        ttl_seconds: int,
+        eligible_model_ids: frozenset[str] | None = None,
+    ) -> SessionPin | UnconfirmedSessionPin | RedisFailure:
+        raw: Final = await self.transport.eval(
+            _CLAIM,
+            (affinity_key(model_group, caller_scope, session_id),),
+            (
+                json.dumps({"model_id": model_id}),
+                int(ttl_seconds),
+                *(
+                    (json.dumps(tuple({"model_id": value} for value in eligible_model_ids)),)
+                    if eligible_model_ids is not None
+                    else ()
+                ),
+            ),
+        )
+        if isinstance(raw, RedisFailure):
+            return raw
+        if not isinstance(raw, bytes):
+            return UnconfirmedSessionPin()
+        try:
+            decoded: Final = raw.decode("utf-8")
+        except UnicodeError as error:
+            return RedisFailure(kind="decode", message=str(error))
+        try:
+            return SessionPin(_JSON.validate_json(decoded))
+        except ValidationError:
+            return SessionPin(decoded)

litellm/router_utils/pre_call_checks/deployment_affinity_check.py

@@ -23,7 +23,6 @@ from litellm._internal_context import with_service_target
 from litellm._logging import verbose_router_logger
 from litellm.caching.affinity_cache import (
     ROUTER_SESSION_PINS_TARGET,
-    claim_affinity_pin,
     claim_affinity_pin_in_memory,
     set_local_affinity_pin,
 )
@@ -31,6 +30,15 @@ from litellm.caching.dual_cache import DualCache
 from litellm.caching.redis_batch import active_request_redis_batch
 from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
 from litellm.integrations.custom_logger import CustomLogger, Span
+from litellm.redis.stores.deployment_session_affinity import (
+    DeploymentSessionAffinityStore,
+    SessionPin,
+    TransportDeploymentSessionAffinityStore,
+    UnconfirmedSessionPin,
+)
+from litellm.redis.stores.deployment_session_affinity import (
+    affinity_key as session_affinity_key,
+)
 from litellm.redis.stores.deployment_user_affinity import (
     DeploymentUserAffinityStore,
     TransportDeploymentUserAffinityStore,
@@ -280,7 +288,7 @@ class DeploymentAffinityCheck(CustomLogger):
         `"unscoped"` covers direct Router usage with no authenticated caller, matching
         the complexity router's own session pin key."""
         hashed_user_key: Final = cls._hash_user_key(user_key) if user_key is not None else "unscoped"
-        return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{hashed_user_key}:{session_id}"
+        return session_affinity_key(model_group, hashed_user_key, session_id)
 
     @staticmethod
     def _get_session_id_from_metadata_dict(metadata: Mapping[object, object]) -> str | None:
@@ -399,8 +407,45 @@ class DeploymentAffinityCheck(CustomLogger):
                 return self._pinned_model_id(result.value)
         return self._claim_pin_in_memory(cache_key, DeploymentAffinityCacheValue(model_id=model_id), self.ttl_seconds)
 
-    async def _claim_pin(self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int) -> str | None:
-        winner: Final = await claim_affinity_pin(self.cache, cache_key, pin_value, ttl_seconds)
+    async def _claim_session_pin(
+        self,
+        model_group: str,
+        session_id: str,
+        user_key: str | None,
+        model_id: str,
+        ttl_seconds: int,
+        *,
+        eligible_model_ids: frozenset[str] | None = None,
+    ) -> str | None:
+        cache_key: Final = self.get_session_affinity_cache_key(model_group, session_id, user_key)
+        redis_cache: Final = self.cache.redis_cache
+        if redis_cache is not None:
+            store: Final[DeploymentSessionAffinityStore] = TransportDeploymentSessionAffinityStore(
+                redis_cache.transport()
+            )
+            result: Final = await store.claim(
+                model_group,
+                self._hash_user_key(user_key) if user_key is not None else "unscoped",
+                session_id,
+                model_id,
+                ttl_seconds=ttl_seconds,
+                eligible_model_ids=eligible_model_ids,
+            )
+            if isinstance(result, UnconfirmedSessionPin):
+                return model_id
+            if isinstance(result, SessionPin):
+                self._set_local_pin(cache_key, result.value, ttl_seconds)
+                return self._pinned_model_id(result.value)
+        winner: Final = claim_affinity_pin_in_memory(
+            self.cache,
+            cache_key,
+            DeploymentAffinityCacheValue(model_id=model_id),
+            ttl_seconds,
+            eligible_values=tuple(DeploymentAffinityCacheValue(model_id=value) for value in eligible_model_ids)
+            if eligible_model_ids is not None
+            else None,
+            preserve_legacy_model_ids=eligible_model_ids is not None,
+        )
         return self._pinned_model_id(winner)
 
     def _claim_pin_in_memory(
@@ -487,17 +532,12 @@ class DeploymentAffinityCheck(CustomLogger):
         if session_affinity_active:
             session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs=request_kwargs)
             if session_id is not None:
-                session_cache_key: Final = self.get_session_affinity_cache_key(
-                    model_group=stable_model_map_key, session_id=session_id, user_key=user_key
-                )
-                session_cache_result: Final = (
-                    await self._read_shared_session_pin(
-                        session_cache_key,
-                        self._get_marker_session_affinity_ttl(request_kwargs) or self.ttl_seconds,
-                        parent_otel_span,
-                    )
-                    if self.session_affinity_repin_on_failover
-                    else await self.cache.async_get_cache(key=session_cache_key)
+                session_cache_result: Final = await self._read_session_pin(
+                    stable_model_map_key,
+                    session_id,
+                    user_key,
+                    self._get_marker_session_affinity_ttl(request_kwargs) or self.ttl_seconds,
+                    parent_otel_span,
                 )
 
                 session_model_id: str | None = None
@@ -561,18 +601,36 @@ class DeploymentAffinityCheck(CustomLogger):
         )
         return [deployment]
 
-    async def _read_shared_session_pin(self, cache_key: str, ttl_seconds: int, parent_otel_span: Span | None) -> object:
+    async def _read_session_pin(
+        self,
+        model_group: str,
+        session_id: str,
+        user_key: str | None,
+        ttl_seconds: int,
+        parent_otel_span: Span | None,
+    ) -> object:
+        cache_key: Final = self.get_session_affinity_cache_key(model_group, session_id, user_key)
         redis_cache: Final = self.cache.redis_cache
-        if redis_cache is None:
-            return TypeAdapter(object).validate_python(await self.cache.async_get_cache(key=cache_key))
+        local: Final = TypeAdapter(object).validate_python(
+            await self.cache.async_get_cache(key=cache_key, local_only=True)
+        )
+        if redis_cache is None or (not self.session_affinity_repin_on_failover and local is not None):
+            return local
         local_before: Final = TypeAdapter(object).validate_python(self.cache.in_memory_cache.get_cache(cache_key))
-        try:
-            shared: Final[object] = await redis_cache.async_get_cache(
-                key=cache_key, parent_otel_span=parent_otel_span, raise_on_error=True
-            )
-        except Exception:  # noqa: BLE001  # Optional affinity falls back locally on Redis admission or read failure
-            verbose_router_logger.debug("Session pin read failed, using pod-local affinity")
+        batch: Final = active_request_redis_batch(redis_cache)
+        if not self.session_affinity_repin_on_failover and batch is not None and batch.read_as_missing(cache_key):
+            return None
+        store: Final[DeploymentSessionAffinityStore] = TransportDeploymentSessionAffinityStore(redis_cache.transport())
+        result: Final = await store.read(
+            model_group, self._hash_user_key(user_key) if user_key is not None else "unscoped", session_id
+        )
+        if not isinstance(result, SessionPin):
             return TypeAdapter(object).validate_python(await self.cache.async_get_cache(key=cache_key, local_only=True))
+        shared: Final = result.value
+        if not self.session_affinity_repin_on_failover:
+            if shared is not None:
+                await self.cache.async_set_cache(key=cache_key, value=shared, local_only=True)
+            return shared
         if self.cache.in_memory_cache.get_cache(cache_key) is local_before:
             if shared is None:
                 self.cache.in_memory_cache.delete_cache(cache_key)
@@ -614,16 +672,13 @@ class DeploymentAffinityCheck(CustomLogger):
         model_id: Final = model_info.get("id")
         if model_id is None or str(model_id) not in eligible_deployment_ids:
             return
-        cache_key: Final = self.get_session_affinity_cache_key(
-            model_group, session_id, self.get_user_key_from_request_kwargs(request_kwargs)
-        )
-        await claim_affinity_pin(
-            self.cache,
-            cache_key,
-            DeploymentAffinityCacheValue(model_id=str(model_id)),
+        await self._claim_session_pin(
+            model_group,
+            session_id,
+            self.get_user_key_from_request_kwargs(request_kwargs),
+            str(model_id),
             marker_ttl if marker_ttl is not None else self.ttl_seconds,
-            eligible_values=tuple(DeploymentAffinityCacheValue(model_id=value) for value in eligible_deployment_ids),
-            preserve_legacy_model_ids=True,
+            eligible_model_ids=eligible_deployment_ids,
         )
 
     async def async_pre_call_deployment_hook(
@@ -733,19 +788,18 @@ class DeploymentAffinityCheck(CustomLogger):
         if session_id is not None:
             try:
                 session_affinity_ttl: Final = marker_session_ttl if marker_session_ttl is not None else self.ttl_seconds
-                session_cache_key: Final = self.get_session_affinity_cache_key(
-                    model_group=deployment_model_name, session_id=session_id, user_key=user_key
-                )
-                claimed_session_pin: Final = await self._claim_pin(
-                    cache_key=session_cache_key,
-                    pin_value=pin_value,
+                claimed_session_pin: Final = await self._claim_session_pin(
+                    model_group=deployment_model_name,
+                    session_id=session_id,
+                    user_key=user_key,
+                    model_id=str(model_id),
                     ttl_seconds=session_affinity_ttl,
                 )
                 if claimed_session_pin == pin_value["model_id"]:
                     verbose_router_logger.debug(
                         "DeploymentAffinityCheck: set session affinity mapping model_map_key=%s deployment=%s ttl=%s session_id=%s",
                         deployment_model_name,
-                        model_id,
+                        pin_value["model_id"],
                         session_affinity_ttl,
                         session_id,
                     )
@@ -754,7 +808,7 @@ class DeploymentAffinityCheck(CustomLogger):
                         "DeploymentAffinityCheck: session pin already claimed model_map_key=%s existing=%s ours=%s session_id=%s",
                         deployment_model_name,
                         claimed_session_pin,
-                        model_id,
+                        pin_value["model_id"],
                         session_id,
                     )
             except Exception as e:

tests/_support/fake_redis.py

@@ -3,6 +3,7 @@ from typing import Final, TypeVar
 
 import fakeredis
 from pydantic import TypeAdapter
+from redis.asyncio import Redis
 
 import litellm
 from litellm._service_logger import ServiceLogging
@@ -93,6 +94,12 @@ def async_client(cache: RedisCache) -> fakeredis.FakeAsyncRedis:
     return client
 
 
+def raw_client(cache: RedisCache) -> "Redis[bytes]":
+    client: Final = cache.init_async_client()  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # legacy RedisCache API is untyped
+    assert isinstance(client, Redis)
+    return client  # pyright: ignore[reportUnknownVariableType]  # redis-py erases the reply type at runtime
+
+
 def pin_async_client(cache: RedisCache) -> Callable[[], None]:
     client: Final = async_client(cache)
 

tests/unit/redis/stores/test_deployment_session_affinity.py

@@ -0,0 +1,181 @@
+import hashlib
+import json
+from collections.abc import Callable, Iterator
+from typing import Final
+from unittest.mock import patch
+
+import fakeredis
+import pytest
+from redis.exceptions import TimeoutError as RedisTimeoutError
+
+from litellm.caching.dual_cache import DualCache
+from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck
+from tests._support.fake_redis import (
+    fake_redis_cache,
+    keys_matching,
+    local_get,
+    local_set,
+    seed_async_client,
+)
+
+
+async def _pre_call(checker: DeploymentAffinityCheck, request: dict[str, object], call_type: None = None) -> None:
+    await checker.async_pre_call_deployment_hook(request, call_type)  # pyright: ignore[reportUnknownMemberType]  # legacy cache API is untyped
+
+
+async def _read_session(
+    checker: DeploymentAffinityCheck,
+    model_group: str,
+    session_id: str,
+    user_key: str | None,
+    ttl_seconds: int,
+    parent_otel_span: None,
+) -> object:
+    return await checker._read_session_pin(model_group, session_id, user_key, ttl_seconds, parent_otel_span)  # pyright: ignore[reportPrivateUsage]  # exercises the pin path without a full router request
+
+
+@pytest.fixture(autouse=True)
+def fixed_clock() -> Iterator[None]:
+    with patch("time.time", return_value=1700000000.0):
+        yield
+
+
+def _checker(server: fakeredis.FakeServer, *, repin: bool = True) -> DeploymentAffinityCheck:
+    return DeploymentAffinityCheck(
+        DualCache(redis_cache=fake_redis_cache(server, namespace="session")),
+        60,
+        False,
+        False,
+        enable_session_id_affinity=True,
+        session_affinity_repin_on_failover=repin,
+    )
+
+
+def _key(caller: str = "caller") -> str:
+    return "session:deployment_affinity:v1:session:group:" + hashlib.sha256(caller.encode()).hexdigest() + ":same"
+
+
+def _request(model_id: str, caller: str = "caller") -> dict[str, object]:
+    return {
+        "model_info": {"id": model_id},
+        "metadata": {
+            "deployment_model_name": "group",
+            "model_group": "group",
+            "session_id": "same",
+            "user_api_key_hash": caller,
+        },
+    }
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("repin", [False, True])
+@pytest.mark.parametrize("wire", [b'"chosen"', b'{"model_id": "chosen"}', b"{'model_id': 'chosen'}"])
+async def test_session_reads_legacy_values_without_refresh(wire: bytes, repin: bool) -> None:
+    server: Final = fakeredis.FakeServer()
+    checker: Final = _checker(server, repin=repin)
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    await raw.set(_key(), wire, ex=20)
+    assert await _read_session(checker, "group", "same", "caller", 60, None) in ("chosen", {"model_id": "chosen"})
+    assert await raw.get(_key()) == wire
+    assert await raw.pttl(_key()) == 20000
+
+
+@pytest.mark.asyncio
+async def test_session_claim_bytes_scope_and_claim_time_ttl() -> None:
+    server: Final = fakeredis.FakeServer()
+    checker: Final = _checker(server)
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    await _pre_call(checker, _request("部署"), None)
+    assert await raw.get(_key()) == b'{"model_id": "\\u90e8\\u7f72"}'
+    assert await raw.pttl(_key()) == 60000
+    await raw.pexpire(_key(), 20000)
+    await _pre_call(checker, _request("other"), None)
+    assert await raw.pttl(_key()) == 20000
+    await _pre_call(checker, _request("部署"), None)
+    assert await raw.pttl(_key()) == 60000
+    await _pre_call(checker, _request("other", "second"), None)
+    assert await raw.get(_key("second")) == b'{"model_id": "other"}'
+    assert await keys_matching(raw, "session:session:*") == ()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "wire,expected",
+    [
+        (b'{"model_id": "chosen"}', b'{"model_id": "chosen"}'),
+        (b'"chosen"', b'"chosen"'),
+        (b"chosen", b"chosen"),
+        (b'{"model_id": "removed"}', b'{"model_id": "replacement"}'),
+        (b'{"model_id": "chosen", "extra": true}', b'{"model_id": "replacement"}'),
+        (b"invalid{", b'{"model_id": "replacement"}'),
+    ],
+)
+async def test_success_claim_preserves_only_eligible_complete_values(wire: bytes, expected: bytes) -> None:
+    server: Final = fakeredis.FakeServer()
+    checker: Final = _checker(server)
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    await raw.set(_key(), wire, ex=20)
+    await checker.async_repin_session_after_success(
+        _request("replacement"), "group", frozenset(("chosen", "replacement"))
+    )
+    assert await raw.get(_key()) == expected
+    assert await raw.pttl(_key()) == 60000
+
+
+@pytest.mark.asyncio
+async def test_shared_miss_clears_local_while_outage_retains_it() -> None:
+    server: Final = fakeredis.FakeServer()
+    checker: Final = _checker(server)
+    local_key: Final = _key().removeprefix("session:")
+    local_set(checker.cache, local_key, {"model_id": "stale"}, ttl=60)
+    assert await _read_session(checker, "group", "same", "caller", 60, None) is None
+    assert local_get(checker.cache, local_key) is None
+    local_set(checker.cache, local_key, {"model_id": "local"}, ttl=60)
+    server.connected = False
+    assert await _read_session(checker, "group", "same", "caller", 60, None) == {"model_id": "local"}
+
+
+class ConcurrentRead(fakeredis.FakeAsyncRedis):
+    after_read: Callable[[], None]
+
+    async def get(self, name: str) -> object:  # pyright: ignore[reportIncompatibleMethodOverride]  # narrows fakeredis for the test
+        result = await super().get(name)  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # fakeredis is untyped
+        self.after_read()
+        return result  # pyright: ignore[reportUnknownVariableType]  # fakeredis is untyped
+
+
+@pytest.mark.asyncio
+async def test_shared_read_does_not_overwrite_a_newer_local_claim() -> None:
+    server: Final = fakeredis.FakeServer()
+    checker: Final = _checker(server)
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    await raw.set(_key(), json.dumps({"model_id": "shared"}), ex=20)
+    local_key: Final = _key().removeprefix("session:")
+    client: Final = ConcurrentRead(server=server)
+    client.after_read = lambda: local_set(checker.cache, local_key, {"model_id": "newer"}, ttl=60)
+    redis_cache: Final = checker.cache.redis_cache
+    assert redis_cache is not None
+    seed_async_client(redis_cache, client)
+    assert await _read_session(checker, "group", "same", "caller", 60, None) == {"model_id": "shared"}
+    assert local_get(checker.cache, local_key) == {"model_id": "newer"}
+
+
+class LostClaimReply(fakeredis.FakeAsyncRedis):
+    async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | int) -> object:
+        await super().eval(script, numkeys, *keys_and_args)  # pyright: ignore[reportUnknownMemberType]  # fakeredis is untyped
+        raise RedisTimeoutError("reply lost after apply")
+
+
+@pytest.mark.asyncio
+async def test_session_unknown_write_retains_local_continuity() -> None:
+    server: Final = fakeredis.FakeServer()
+    checker: Final = _checker(server)
+    redis_cache: Final = checker.cache.redis_cache
+    assert redis_cache is not None
+    seed_async_client(redis_cache, LostClaimReply(server=server))
+    local_key: Final = _key().removeprefix("session:")
+    local_set(checker.cache, local_key, {"model_id": "local"}, ttl=60)
+    await _pre_call(checker, _request("shared"), None)
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    assert await raw.get(_key()) == b'{"model_id": "shared"}'
+    assert local_get(checker.cache, local_key) == {"model_id": "local"}

tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py

@@ -1,8 +1,8 @@
 import asyncio
 import json
-from typing import Final
 from collections.abc import AsyncIterator, Mapping
 from datetime import datetime, timezone
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -19,6 +19,7 @@ from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESS
 from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
     DeploymentAffinityCheck,
 )
+from tests._support.fake_redis import async_client, fake_redis_cache, raw_client, seed_async_client
 
 
 def _repinning_router(enabled: bool | None = True) -> litellm.Router:
@@ -121,6 +122,17 @@ async def test_unsuccessful_responses_do_not_repin(status: str, wrapped: bool) -
         router.discard()
 
 
+async def _claim_session(
+    callback: DeploymentAffinityCheck,
+    model_group: str,
+    session_id: str,
+    user_key: str | None,
+    model_id: str,
+    ttl_seconds: int,
+) -> object:
+    return await callback._claim_session_pin(model_group, session_id, user_key, model_id, ttl_seconds)  # pyright: ignore[reportPrivateUsage]  # exercises the pin path without a full router request
+
+
 async def test_concurrent_successful_replacements_keep_first_winner() -> None:
     router: Final = _repinning_router()
     try:
@@ -288,8 +300,8 @@ async def test_delayed_shared_pin_read_preserves_a_newer_local_write(
             await resume.wait()
             return value
 
-    reader: Final = DelayedReader(connection_pool=redis.init_async_client().connection_pool)
-    litellm.in_memory_llm_clients_cache.set_cache(key=redis._get_async_client_cache_key(), value=reader)
+    reader: Final = DelayedReader(connection_pool=raw_client(redis).connection_pool)
+    seed_async_client(redis, reader)
     pending: Final = asyncio.create_task(_selected_session_ids(router))
     try:
         await observed.wait()
@@ -817,12 +829,12 @@ async def test_in_memory_reclaim_slides_idle_window_only_for_the_stored_deployme
     cache.in_memory_cache.set_cache(pin_key, stored_pin, ttl=10)
     first_expiry = cache.in_memory_cache.ttl_dict[pin_key]
 
-    reclaimed = await callback._claim_pin(cache_key=pin_key, pin_value={"model_id": "deployment-1"}, ttl_seconds=777)
+    reclaimed = await _claim_session(callback, "target-group", "slide-session", "key-1", "deployment-1", 777)
     assert reclaimed == "deployment-1"
     assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5)
     assert cache.in_memory_cache.ttl_dict[pin_key] > first_expiry
 
-    lost = await callback._claim_pin(cache_key=pin_key, pin_value={"model_id": "deployment-2"}, ttl_seconds=10)
+    lost = await _claim_session(callback, "target-group", "slide-session", "key-1", "deployment-2", 10)
     assert lost == "deployment-1"
     assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5)
 
@@ -840,32 +852,21 @@ async def test_claim_pin_uses_redis_attached_after_construction():
         enable_responses_api_affinity=False,
     )
 
-    captured = {}
-
-    async def fake_runner(keys, args, client=None):
-        captured["keys"] = keys
-        captured["args"] = args
-        return b'{"model_id": "other-pod-winner"}'
-
-    late_redis = MagicMock()
-    late_redis.async_register_script = MagicMock(return_value=fake_runner)
+    server = FakeServer()
+    late_redis = fake_redis_cache(server, namespace="late")
     cache.redis_cache = late_redis
 
     import time as time_module
 
     pin_key = _session_pin_key("late-redis-session", "key-1")
+    await late_redis.async_set_cache(pin_key, {"model_id": "other-pod-winner"}, ttl=10)
     cache.in_memory_cache.set_cache(pin_key, {"model_id": "other-pod-winner"}, ttl=10)
 
-    claimed = await callback._claim_pin(
-        cache_key=pin_key,
-        pin_value={"model_id": "our-deployment"},
-        ttl_seconds=777,
-    )
+    claimed = await _claim_session(callback, "target-group", "late-redis-session", "key-1", "our-deployment", 777)
 
     assert claimed == "other-pod-winner"
     assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5)
-    assert captured["keys"] == (pin_key,)
-    assert captured["args"] == ('{"model_id": "our-deployment"}', 777)
+    assert await late_redis.async_get_cache(pin_key) == {"model_id": "other-pod-winner"}
     assert cache.in_memory_cache.get_cache(_session_pin_key("late-redis-session", "key-1")) == {
         "model_id": "other-pod-winner"
     }
@@ -885,20 +886,18 @@ async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down():
         enable_responses_api_affinity=False,
     )
 
-    async def exploding_runner(keys, args, client=None):
-        raise ConnectionError("redis is down")
-
-    down_redis = MagicMock()
-    down_redis.async_register_script = MagicMock(return_value=exploding_runner)
+    server = FakeServer()
+    server.connected = False
+    down_redis = fake_redis_cache(server, namespace="down")
     cache.redis_cache = down_redis
 
     key = _session_pin_key("outage-session", "key-1")
-    claimed = await callback._claim_pin(cache_key=key, pin_value={"model_id": "our-deployment"}, ttl_seconds=777)
+    claimed = await _claim_session(callback, "target-group", "outage-session", "key-1", "our-deployment", 777)
 
     assert claimed == "our-deployment"
     assert cache.in_memory_cache.get_cache(key) == {"model_id": "our-deployment"}
 
-    second = await callback._claim_pin(cache_key=key, pin_value={"model_id": "another-deployment"}, ttl_seconds=777)
+    second = await _claim_session(callback, "target-group", "outage-session", "key-1", "another-deployment", 777)
     assert second == "our-deployment"
 
 
@@ -979,14 +978,15 @@ async def test_legacy_deployment_claim_retains_decoder_and_keepalive(
         enable_user_key_affinity=False,
         enable_responses_api_affinity=False,
     )
-    cache.in_memory_cache.set_cache("deployment-pin", stored, ttl=10)
+    pin_key: Final = callback.get_session_affinity_cache_key("target-group", "decoder", None)
+    cache.in_memory_cache.set_cache(pin_key, stored, ttl=10)
     clock.return_value = 105.0
 
-    winner: Final = await callback._claim_pin("deployment-pin", {"model_id": "7"}, 30)
+    winner: Final = await _claim_session(callback, "target-group", "decoder", None, "7", 30)
 
     assert winner == expected
-    assert cache.in_memory_cache.ttl_dict["deployment-pin"] == (135.0 if refresh else 110.0)
-    assert cache.in_memory_cache.get_cache("deployment-pin") == ({"model_id": "7"} if refresh else stored)
+    assert cache.in_memory_cache.ttl_dict[pin_key] == (135.0 if refresh else 110.0)
+    assert cache.in_memory_cache.get_cache(pin_key) == ({"model_id": "7"} if refresh else stored)
 
 
 @pytest.mark.asyncio
@@ -998,16 +998,13 @@ async def test_legacy_deployment_claim_retains_decoder_and_keepalive(
         ("winner", "winner", "winner"),
         (b"winner", "winner", "winner"),
         ('{"model": "winner"}', None, {"model": "winner"}),
-        (None, "candidate", None),
-        (123, "candidate", None),
-        ({"model_id": "winner"}, "candidate", None),
+        (None, "candidate", {"model_id": "candidate"}),
     ],
 )
 async def test_redis_deployment_claim_preserves_legacy_result_decoding(
-    raw: object, expected: str | None, stored: object
+    raw: str | bytes | None, expected: str | None, stored: object
 ) -> None:
-    redis: Final = MagicMock()
-    redis.async_register_script.return_value = AsyncMock(return_value=raw)
+    redis: Final = fake_redis_cache(FakeServer(), namespace="decode")
     cache: Final = DualCache(redis_cache=redis)
     callback: Final = DeploymentAffinityCheck(
         cache=cache,
@@ -1016,10 +1013,13 @@ async def test_redis_deployment_claim_preserves_legacy_result_decoding(
         enable_responses_api_affinity=False,
     )
 
-    winner: Final = await callback._claim_pin("deployment-pin", {"model_id": "candidate"}, 30)
+    pin_key: Final = callback.get_session_affinity_cache_key("target-group", "decoder", None)
+    if raw is not None:
+        await async_client(redis).set(redis.check_and_fix_namespace(pin_key), raw, ex=10)
+    winner: Final = await _claim_session(callback, "target-group", "decoder", None, "candidate", 30)
 
     assert winner == expected
-    assert cache.in_memory_cache.get_cache("deployment-pin") == stored
+    assert cache.in_memory_cache.get_cache(pin_key) == stored
 
 
 @pytest.mark.asyncio