diff --git a/litellm/caching/affinity_cache.py b/litellm/caching/affinity_cache.py index d9a40c44d2b..2b8dc67da7a 100644 --- a/litellm/caching/affinity_cache.py +++ b/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, diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index f5fd4ae9c98..82aaa221d31 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/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, diff --git a/tests/integration/sdk/test_dual_cache_redis.py b/tests/integration/sdk/test_dual_cache_redis.py index 3527dbf91ee..7d05eb2795b 100644 --- a/tests/integration/sdk/test_dual_cache_redis.py +++ b/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() diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index 78dc2eaf314..cb3faf89358 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/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")