fix(router): keep session tier pins coherent across workers (#48645)

litellm/caching/affinity_cache.py

@@ -4,18 +4,27 @@ import json
 from collections.abc import Mapping
 from typing import (
     Final,
-    cast,  # noqa: TID251  # Redis script results are narrowed only to object, then validated
+    cast,  # noqa: TID251  # Heterogeneous cache results are narrowed only to object, then validated
 )
 
 from pydantic import JsonValue, TypeAdapter, ValidationError
 
 from litellm._logging import verbose_router_logger
 from litellm.caching.dual_cache import DualCache
+from litellm.caching.redis_batch import active_request_redis_batch
 
 ROUTER_SESSION_PINS_TARGET: Final = "router_session_pins"
 
 _PIN_JSON_ADAPTER: Final = TypeAdapter[JsonValue](JsonValue)
 
+_READ_PIN_SCRIPT: Final = """
+local current = redis.call('GET', KEYS[1])
+if current ~= false then
+  redis.call('EXPIRE', KEYS[1], ARGV[1])
+end
+return current
+"""
+
 _CLAIM_PIN_SCRIPT: Final = """
 local current = redis.call('GET', KEYS[1])
 if current == false then
@@ -103,6 +112,37 @@ def _decode_pin(value: str) -> object:
         return value
 
 
+async def read_affinity_pin(cache: DualCache, cache_key: str, ttl_seconds: int) -> object:
+    """Refresh the authoritative pin without writing a previously observed value back."""
+    redis_cache: Final = cache.redis_cache
+    if redis_cache is not None:
+        try:
+            read_script: Final = redis_cache.async_register_script(_READ_PIN_SCRIPT)
+            batch: Final = active_request_redis_batch(redis_cache)
+            raw: Final = cast(  # cast-ok: Redis script returns a serialized pin or a miss; validate before use
+                object,
+                await batch.script(_READ_PIN_SCRIPT, read_script, (cache_key,), (int(ttl_seconds),))
+                if batch is not None and batch.pending
+                else await read_script(keys=(cache_key,), args=(int(ttl_seconds),)),
+            )
+            decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw
+            value: Final = _decode_pin(decoded) if isinstance(decoded, str) else None
+            if value is None:
+                cache.in_memory_cache.delete_cache(cache_key)
+            else:
+                set_local_affinity_pin(cache, cache_key, value, ttl_seconds)
+        except Exception as error:  # noqa: BLE001  # Redis faults retain pod-local continuity, unlike a successful miss
+            verbose_router_logger.debug("Affinity cache: Redis read failed, using pod-local pin. error=%s", error)
+        else:
+            return value
+    local_value: Final = cast(  # cast-ok: the heterogeneous local cache is parsed by the affinity caller
+        object, cache.in_memory_cache.get_cache(cache_key)
+    )
+    if local_value is not None:
+        set_local_affinity_pin(cache, cache_key, local_value, ttl_seconds)
+    return local_value
+
+
 async def claim_affinity_pin(
     cache: DualCache,
     cache_key: str,

litellm/router_strategy/complexity_router/complexity_router.py

@@ -35,7 +35,7 @@ from typing_extensions import ReadOnly, TypedDict
 
 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
+from litellm.caching.affinity_cache import ROUTER_SESSION_PINS_TARGET, claim_affinity_pin, read_affinity_pin
 from litellm.constants import (
     EMPTY_MAPPING,
     INTERNAL_CALL_ORIGIN_METADATA_KEY,
@@ -4473,7 +4473,9 @@ class ComplexityRouter(CustomLogger):
         )
 
         if cache_key is not None and pin_replay_allowed:
-            pinned_value: Final = await self.litellm_router_instance.cache.async_get_cache(key=cache_key)
+            pinned_value: Final = await read_affinity_pin(
+                self.litellm_router_instance.cache, cache_key, self.config.session_affinity_ttl_seconds
+            )
             pinned_pin: Final = _parse_session_affinity_pin(pinned_value, self.config.tier_names())
             if pinned_pin is not None:
                 user_message: Final = (
@@ -4570,13 +4572,12 @@ class ComplexityRouter(CustomLogger):
                         if pin_placement is not None and pin_context_original_tier is not None
                         else floor_model
                     )
-                    # Refresh the TTL on every hit so an active session doesn't lose its
-                    # pin mid-conversation just because it outlives the original write.
-                    await self.litellm_router_instance.cache.async_set_cache(
-                        key=cache_key,
-                        value=_session_affinity_cache_value(session_model, resolved_pin_tier),
-                        ttl=self.config.session_affinity_ttl_seconds,
-                    )
+                    if escalated:
+                        await self.litellm_router_instance.cache.async_set_cache(
+                            key=cache_key,
+                            value=_session_affinity_cache_value(session_model, resolved_pin_tier),
+                            ttl=self.config.session_affinity_ttl_seconds,
+                        )
                     if self.config.adaptive:
                         from litellm.router_strategy.adaptive_router.config import (
                             ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY,

tests/integration/sdk/test_dual_cache_redis.py

@@ -6,7 +6,7 @@ from unittest.mock import patch
 
 import pytest
 
-from litellm.caching.affinity_cache import claim_affinity_pin
+from litellm.caching.affinity_cache import claim_affinity_pin, read_affinity_pin
 from litellm.caching.dual_cache import DualCache
 from litellm.caching.in_memory_cache import InMemoryCache
 from litellm.caching.redis_cache import RedisCache
@@ -156,3 +156,121 @@ async def test_session_repin_lua_preserves_healthy_legacy_pins_and_replaces_unav
         assert await redis.async_get_cache(key) == replacement
     finally:
         await redis.async_delete_cache(key)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("stored", ['{"model":"current","tier":"REASONING"}', '"current"', "current"])
+async def test_tier_pin_reads_shared_state_refreshes_ttl_and_forgets_deleted_pins(stored: str) -> None:
+    redis: Final = _redis_cache()
+    redis.namespace = f"tier-read-{uuid.uuid4().hex}"
+    workers: Final = (DualCache(redis_cache=redis), DualCache(redis_cache=redis))
+    key: Final = "session"
+    namespaced: Final = redis.check_and_fix_namespace(key)
+    try:
+        for worker in workers:
+            worker.in_memory_cache.set_cache(key, {"model": "old", "tier": "SIMPLE"}, ttl=60)
+        redis.redis_client.set(namespaced, stored, ex=10)
+        first: Final = await read_affinity_pin(workers[0], key, 60)
+        second: Final = await read_affinity_pin(workers[1], key, 60)
+        assert first == second and first in ("current", {"model": "current", "tier": "REASONING"})
+        assert all(worker.in_memory_cache.get_cache(key) == first for worker in workers)
+        assert 10 < redis.redis_client.ttl(namespaced) <= 60
+        replacement: Final = {"model": "newer", "tier": "SIMPLE"}
+        await redis.async_set_cache(key, replacement, ttl=10)
+        assert await read_affinity_pin(workers[0], key, 60) == replacement
+        assert await read_affinity_pin(workers[1], key, 60) == replacement
+        await redis.async_delete_cache(key)
+        for worker in workers:
+            assert await read_affinity_pin(worker, key, 60) is None
+            assert worker.in_memory_cache.get_cache(key) is None
+        assert redis.redis_client.exists(namespaced) == 0
+    finally:
+        await redis.async_delete_cache(key)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mode", ["joined", "separate", "script_error"])
+async def test_tier_pin_joins_pending_reads_without_reusing_request_state(mode: str) -> None:
+    import json
+    from collections.abc import Iterable
+
+    import litellm
+    from redis.asyncio import BlockingConnectionPool, Connection, Redis
+
+    from litellm.caching.redis_batch import request_redis_batch_scope
+
+    sends: Final[list[bytes]] = []
+
+    class RecordingConnection(Connection):
+        async def send_packed_command(
+            self, command: bytes | str | Iterable[bytes], check_health: bool = True
+        ) -> None:
+            packed: Final = command.encode() if isinstance(command, str) else (
+                command if isinstance(command, bytes) else b"".join(command)
+            )
+            sends.append(packed)
+            await super().send_packed_command(packed, check_health=check_health)
+
+    namespace: Final = f"tier-batch-{uuid.uuid4().hex}"
+    redis: Final = RedisCache(
+        host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]),
+        namespace=namespace, client_name=namespace,
+    )
+    pool: Final = BlockingConnectionPool(
+        host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]),
+        connection_class=RecordingConnection,
+    )
+    client: Final = Redis(connection_pool=pool)
+    client_key: Final = redis._get_async_client_cache_key()
+    litellm.in_memory_llm_clients_cache.set_cache(key=client_key, value=client)
+    cache: Final = DualCache(redis_cache=redis)
+    key: Final = redis.check_and_fix_namespace("pin")
+    other_key: Final = redis.check_and_fix_namespace("other")
+    stale: Final = {"model": "old", "tier": "SIMPLE"}
+    current: Final = {"model": "current", "tier": "REASONING"}
+    newer: Final = {"model": "newer", "tier": "COMPLEX"}
+    try:
+        redis.redis_client.set(key, json.dumps(current), ex=10)
+        redis.redis_client.set(other_key, json.dumps("other-value"), ex=60)
+        assert await read_affinity_pin(cache, "pin", 60) == current
+        cache.in_memory_cache.set_cache("pin", stale, ttl=60)
+        assert cache.in_memory_cache.get_cache("pin") == stale
+        if mode == "script_error":
+            redis.redis_client.delete(key)
+            redis.redis_client.rpush(key, "wrong-type")
+        start: Final = len(sends)
+        with request_redis_batch_scope() as request:
+            batch: Final = request.batch(redis)
+            pending: Final = batch.mget(("other",))
+            if mode == "separate":
+                assert await pending == {"other": "other-value"}
+            assert await read_affinity_pin(cache, "pin", 60) == (stale if mode == "script_error" else current)
+            assert await pending == {"other": "other-value"}
+            wire: Final = tuple(data for data in sends[start:] if b"MGET" in data or b"EVALSHA" in data)
+            assert len(wire) == (2 if mode == "separate" else 1)
+            assert sum(b"MGET" in data for data in wire) == sum(b"EVALSHA" in data for data in wire) == 1
+            if mode == "script_error":
+                assert redis.redis_client.type(key) == b"list"
+                redis.redis_client.delete(key)
+            else:
+                assert 10 < redis.redis_client.ttl(key) <= 60
+            redis.redis_client.set(key, json.dumps(newer), ex=10)
+            later_pending: Final = batch.mget(("other",))
+            assert await read_affinity_pin(cache, "pin", 60) == newer
+            assert await later_pending == {"other": "other-value"}
+            assert cache.in_memory_cache.get_cache("pin") == newer
+            assert 10 < redis.redis_client.ttl(key) <= 60
+            assert redis.redis_client.get(key) == json.dumps(newer).encode()
+            redis.redis_client.delete(key)
+            assert await read_affinity_pin(cache, "pin", 60) is None
+            assert cache.in_memory_cache.get_cache("pin") is None
+            assert redis.redis_client.exists(key) == 0
+    finally:
+        redis.redis_client.delete(key, other_key)
+        for cached_key in tuple(litellm.in_memory_llm_clients_cache.cache_dict):
+            if client_key in cached_key:
+                litellm.in_memory_llm_clients_cache.delete_cache(cached_key)
+        await client.aclose()
+        await pool.aclose()
+        await redis.disconnect()
+        redis.redis_client.close()

tests/unit/router_strategy/test_complexity_router.py

@@ -35,6 +35,7 @@ from litellm.router_utils.auto_router_model_naming import (
 from litellm._logging import verbose_router_logger
 from litellm.caching.dual_cache import DualCache
 from litellm.caching.in_memory_cache import InMemoryCache
+from litellm.caching.redis_cache import RedisCache
 from litellm.constants import (
     CLIENT_OUTPUT_CEILING_METADATA_KEY,
     OUTPUT_TOKEN_CEILING_PARAMS,
@@ -7455,35 +7456,14 @@ class TestSessionAffinity:
         assert simple.model == "gpt-4o-mini"
 
     @pytest.mark.asyncio
-    async def test_respects_ttl_seconds(self, mock_router_instance, basic_config):
-        cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None)
-        cache.async_get_cache = AsyncMock(return_value=None)
+    @pytest.mark.parametrize("preexisting", [False, True])
+    async def test_session_pin_ttl_slides_on_replay(
+        self, mock_router_instance: MagicMock, basic_config: dict[str, object], preexisting: bool
+    ) -> None:
+        clock: Final = MagicMock(return_value=100.0)
+        cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock))
         mock_router_instance.cache = cache
-        router = ComplexityRouter(
-            model_name="test-router",
-            litellm_router_instance=mock_router_instance,
-            complexity_router_config={
-                **basic_config,
-                "session_affinity": True,
-                "session_affinity_ttl_seconds": 120,
-            },
-        )
-        await router.async_pre_routing_hook(
-            model="test-model", request_kwargs=self._request_kwargs("session-1"), messages=self.SIMPLE_MESSAGE
-        )
-        cache.async_set_cache.assert_called_once()
-        call_kwargs = cache.async_set_cache.call_args.kwargs
-        assert call_kwargs["ttl"] == 120
-        assert call_kwargs["value"] == {"model": "gpt-4o-mini", "tier": "SIMPLE"}
-
-    @pytest.mark.asyncio
-    async def test_ttl_refreshed_on_cache_hit(self, mock_router_instance, basic_config):
-        """Regression: a pinned turn must refresh the TTL, not just the first write --
-        otherwise a session outliving session_affinity_ttl_seconds silently loses its pin."""
-        cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None)
-        cache.async_get_cache = AsyncMock(return_value="o1-preview")
-        mock_router_instance.cache = cache
-        router = ComplexityRouter(
+        router: Final = ComplexityRouter(
             model_name="test-router",
             litellm_router_instance=mock_router_instance,
             complexity_router_config={
@@ -7492,14 +7472,153 @@ class TestSessionAffinity:
                 "session_affinity_ttl_seconds": 90,
             },
         )
-        result = await router.async_pre_routing_hook(
-            model="test-model", request_kwargs=self._request_kwargs("session-1"), messages=self.SIMPLE_MESSAGE
+        request_kwargs: Final = self._request_kwargs("session-1")
+        key: Final = router._get_session_affinity_cache_key("session-1", request_kwargs)
+        if preexisting:
+            cache.in_memory_cache.set_cache(key, "o1-preview", ttl=90)
+        first: Final = await router.async_pre_routing_hook(
+            model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE
         )
-        assert result.model == "o1-preview"
-        cache.async_set_cache.assert_called_once()
-        call_kwargs = cache.async_set_cache.call_args.kwargs
-        assert call_kwargs["value"] == {"model": "o1-preview", "tier": "REASONING"}
-        assert call_kwargs["ttl"] == 90
+        assert first.model == ("o1-preview" if preexisting else "gpt-4o-mini")
+        clock.return_value = 180.0
+        replay: Final = await router.async_pre_routing_hook(
+            model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE
+        )
+        assert replay.model == first.model
+        clock.return_value = 200.0
+        assert cache.in_memory_cache.get_cache(key) == (
+            "o1-preview" if preexisting else {"model": "gpt-4o-mini", "tier": "SIMPLE"}
+        )
+        clock.return_value = 271.0
+        assert cache.in_memory_cache.get_cache(key) is None
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("batch_pending", [False, True])
+    @pytest.mark.parametrize("remote_state", ["latest", "interleaved", "outage"])
+    async def test_continuation_follows_redis_and_never_republishes_a_stale_pin(
+        self, mock_router_instance: MagicMock, basic_config: dict[str, object], remote_state: str, batch_pending: bool
+    ) -> None:
+        from fakeredis import FakeRedis
+        from redis.asyncio import BlockingConnectionPool, Connection, Redis
+        from redis.asyncio.retry import Retry
+        from redis.backoff import NoBackoff
+
+        from litellm.caching.redis_batch import request_redis_batch_scope
+
+        clock: Final = MagicMock(return_value=100.0)
+        simple: Final = {"model": "gpt-4o-mini", "tier": "SIMPLE"}
+        latest: Final = {"model": "o1-preview", "tier": "REASONING"}
+        stored: Final = {"pin": json.dumps(simple if remote_state == "interleaved" else latest)}
+        available: Final = asyncio.Event()
+        read_started: Final = asyncio.Event()
+        release_read: Final = asyncio.Event()
+        if remote_state != "outage":
+            available.set()
+        if remote_state != "interleaved":
+            release_read.set()
+
+        async def execute(*args: object, **kwargs: object) -> object:
+            if args[0] in ("EVALSHA", "GET", "MGET"):
+                if not available.is_set():
+                    raise ConnectionError("Redis unavailable")
+                if args[0] == "MGET":
+                    return [json.dumps("other-value")]
+                observed: Final = stored["pin"]
+                read_started.set()
+                await release_read.wait()
+                return observed
+            if args[0] == "SET":
+                stored["pin"] = str(args[2])
+            return b"PONG" if args[0] == "PING" else b"OK"
+
+        pack_commands: Final = MagicMock(side_effect=tuple)
+
+        async def connection() -> MagicMock:
+            replies: Final[asyncio.Queue[object]] = asyncio.Queue()
+
+            async def send_command(*args: object) -> None:
+                await replies.put(await execute(*args))
+
+            async def send_packed_command(commands: Sequence[Sequence[object]]) -> None:
+                for command in commands:
+                    await send_command(*command)
+
+            return MagicMock(
+                spec=Connection,
+                retry=Retry(NoBackoff(), 0),
+                send_command=AsyncMock(side_effect=send_command),
+                pack_commands=pack_commands,
+                send_packed_command=AsyncMock(side_effect=send_packed_command),
+                read_response=AsyncMock(side_effect=replies.get),
+            )
+
+        sync_client: Final = FakeRedis()
+        pool: Final = MagicMock(
+            spec=BlockingConnectionPool,
+            connection_kwargs={"protocol": 2},
+            get_encoder=sync_client.get_encoder,
+            get_connection=AsyncMock(side_effect=connection),
+            release=AsyncMock(),
+        )
+        client: Final = Redis(connection_pool=pool)
+        redis_cache: Final = RedisCache(host="tier-pin-fixture", connection_pool=sync_client.connection_pool)
+        client_key: Final = redis_cache._get_async_client_cache_key()
+        litellm.in_memory_llm_clients_cache.set_cache(key=client_key, value=client)
+        try:
+            cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock), redis_cache=redis_cache)
+            mock_router_instance.cache = cache
+            router: Final = ComplexityRouter(
+                model_name="test-router",
+                litellm_router_instance=mock_router_instance,
+                complexity_router_config={
+                    **basic_config,
+                    "classification_mode": "user_turn",
+                    "deployment_affinity": False,
+                    "session_affinity_ttl_seconds": 10,
+                },
+            )
+            request: Final = self._request_kwargs("shared-session")
+            key: Final = router._get_session_affinity_cache_key("shared-session", request)
+            cache.in_memory_cache.set_cache(key, simple, ttl=10)
+            clock.return_value = 105.0
+            messages: Final = [*self.SIMPLE_MESSAGE, {"role": "tool", "content": "done"}]
+            with request_redis_batch_scope() as scope:
+                batch: Final = scope.batch(redis_cache)
+                unrelated: Final = batch.mget(("other",)) if batch_pending else None
+                pending: Final = asyncio.create_task(router.async_pre_routing_hook("m", request, messages))
+                if remote_state == "interleaved":
+                    await asyncio.wait_for(read_started.wait(), timeout=1)
+                    stored["pin"] = json.dumps(latest)
+                    release_read.set()
+                result: Final = await pending
+                if unrelated is not None:
+                    if remote_state == "outage":
+                        with pytest.raises(ConnectionError, match="Redis unavailable"):
+                            await unrelated
+                    else:
+                        assert await unrelated == {"other": "other-value"}
+                assert result.model == (latest["model"] if remote_state == "latest" else simple["model"])
+                assert result.routing_decision["cause"] == "user_turn_continuation"
+                assert json.loads(stored["pin"]) == latest
+                clock.return_value = 111.0
+                assert cache.in_memory_cache.get_cache(key) is not None
+                available.set()
+                recovery_read: Final = batch.mget(("other",)) if batch_pending else None
+                recovered: Final = await router.async_pre_routing_hook("m", request, messages)
+                if recovery_read is not None:
+                    assert await recovery_read == {"other": "other-value"}
+                assert recovered.model == latest["model"]
+                assert cache.in_memory_cache.get_cache(key) == latest
+                assert read_started.is_set()
+                assert tuple(
+                    tuple(args[0] for args in sent.args[0]) for sent in pack_commands.call_args_list
+                ) == (("MGET", "EVALSHA"),) * (2 if batch_pending else 0)
+        finally:
+            for cached_key in tuple(litellm.in_memory_llm_clients_cache.cache_dict):
+                if client_key in cached_key:
+                    litellm.in_memory_llm_clients_cache.delete_cache(cached_key)
+            await client.aclose()
+            sync_client.close()
 
     @pytest.mark.asyncio
     async def test_different_api_keys_do_not_share_pin(self, mock_router_instance, session_affinity_config):
@@ -13870,8 +13989,7 @@ async def test_session_pin_uses_recorded_tier_when_model_is_in_multiple_tiers(mo
 
 @pytest.mark.asyncio
 async def test_session_pin_survives_json_list_round_trip(mock_router_instance):
-    cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None)
-    cache.async_get_cache = AsyncMock(return_value=["shared", "SIMPLE"])
+    cache: Final = DualCache()
     mock_router_instance.cache = cache
     router = ComplexityRouter(
         model_name="test-router",
@@ -13885,6 +14003,8 @@ async def test_session_pin_survives_json_list_round_trip(mock_router_instance):
         },
     )
     request_kwargs = {"metadata": {"session_id": "json-round-trip-session"}}
+    key: Final = router._get_session_affinity_cache_key("json-round-trip-session", request_kwargs)
+    cache.in_memory_cache.set_cache(key, ["shared", "SIMPLE"])
 
     response = await router.async_pre_routing_hook(
         model="test-router",
@@ -13895,7 +14015,7 @@ async def test_session_pin_survives_json_list_round_trip(mock_router_instance):
     assert response is not None
     assert response.model == "shared"
     assert response.litellm_params == {"reasoning_effort": "low"}
-    assert cache.async_set_cache.call_args.kwargs["value"] == {"model": "shared", "tier": "SIMPLE"}
+    assert cache.in_memory_cache.get_cache(key) == ["shared", "SIMPLE"]
 
 
 HEURISTIC_FIRST_TIERS: dict[str, str] = {
@@ -15024,8 +15144,7 @@ class TestModalityRouting:
                 {"role": "user", "content": [{"type": "text", "text": "quick lookup: what is this?"}, IMG_PART]}
             ]
         elif path.startswith(("pin_kept", "pin_replacement", "pin_override")):
-            cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None)
-            cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"})
+            cache: Final = DualCache()
             mock_router_instance.cache = cache
             config["session_affinity"] = True
             request_kwargs = {"metadata": {"session_id": "s1"}}
@@ -15047,6 +15166,11 @@ class TestModalityRouting:
             mock_router_instance.model_list = []
             _answer_deployments_named(mock_router_instance)
         router = self._router(mock_router_instance, config, vision)
+        if path.startswith(("pin_kept", "pin_replacement", "pin_override")):
+            mock_router_instance.cache.in_memory_cache.set_cache(
+                router._get_session_affinity_cache_key("s1", request_kwargs),
+                {"model": "text-cheap", "tier": "SIMPLE"},
+            )
         result = await router.async_pre_routing_hook(model="m", request_kwargs=request_kwargs, messages=messages)
         assert result.model == expected_model
         assert result.routing_decision["cause"] == expected_cause
@@ -15214,8 +15338,7 @@ class TestModalityRouting:
     @pytest.mark.asyncio
     async def test_pin_override_serves_the_image_turn_without_repinning(self, mock_router_instance):
         """The override is for one request: the session keeps the model it was pinned to."""
-        cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None)
-        cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"})
+        cache: Final = DualCache()
         mock_router_instance.cache = cache
         router = self._router(
             mock_router_instance,
@@ -15228,6 +15351,8 @@ class TestModalityRouting:
             dict(self.BASE_VISION),
         )
         request_kwargs = {"metadata": {"session_id": "s1"}}
+        key: Final = router._get_session_affinity_cache_key("s1", request_kwargs)
+        cache.in_memory_cache.set_cache(key, {"model": "text-cheap", "tier": "SIMPLE"})
 
         image_turn = await router.async_pre_routing_hook(
             model="m", request_kwargs=request_kwargs, messages=self.IMAGE_MESSAGE
@@ -15236,7 +15361,7 @@ class TestModalityRouting:
         assert image_turn.routing_decision["cause"] == "modality_pin_override"
         assert "modality_escalated_from:SIMPLE" in image_turn.routing_decision["signals"]
 
-        assert cache.async_set_cache.await_args.kwargs["value"] == {"model": "text-cheap", "tier": "SIMPLE"}
+        assert cache.in_memory_cache.get_cache(key) == {"model": "text-cheap", "tier": "SIMPLE"}
 
         text_turn = await router.async_pre_routing_hook(
             model="m", request_kwargs={"metadata": {"session_id": "s1"}}, messages=[{"role": "user", "content": "hi"}]
@@ -15247,8 +15372,7 @@ class TestModalityRouting:
     @pytest.mark.asyncio
     async def test_pin_override_with_no_capable_model_rejects_and_keeps_the_pin(self, mock_router_instance):
         """The clear 400 replaces the provider's, and a rejected turn must not cost the session its pin."""
-        cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None)
-        cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"})
+        cache: Final = DualCache()
         mock_router_instance.cache = cache
         router = self._router(
             mock_router_instance,
@@ -15260,11 +15384,13 @@ class TestModalityRouting:
             },
             {"text-cheap": False, "text-big": False},
         )
+        key: Final = router._get_session_affinity_cache_key("s1", {"metadata": {"session_id": "s1"}})
+        cache.in_memory_cache.set_cache(key, {"model": "text-cheap", "tier": "SIMPLE"})
         with pytest.raises(litellm.BadRequestError, match="no model"):
             await router.async_pre_routing_hook(
                 model="m", request_kwargs={"metadata": {"session_id": "s1"}}, messages=self.IMAGE_MESSAGE
             )
-        assert cache.async_set_cache.await_args.kwargs["value"] == {"model": "text-cheap", "tier": "SIMPLE"}
+        assert cache.in_memory_cache.get_cache(key) == {"model": "text-cheap", "tier": "SIMPLE"}
 
 
 @pytest.mark.usefixtures("local_model_cost_map")