refactor(redis): migrate complexity tier model affinity into a typed store

litellm/caching/affinity_cache.py

@@ -1,72 +1,12 @@
-"""Atomic affinity claims shared by deployment and tier-model selection."""
+"""Pod-local affinity policy shared by deployment and tier-model selection."""
 
-import json
 from collections.abc import Mapping
-from typing import (
-    Final,
-    cast,  # noqa: TID251  # Heterogeneous cache results are narrowed only to object, then validated
-)
+from typing import Final
 
-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
-  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)
-  if ARGV[4] then
-    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
-  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 set_local_affinity_pin(cache: DualCache, cache_key: str, value: object, ttl_seconds: int) -> None:
     """Replace the entry because InMemoryCache.set_cache preserves a live key's expiry."""
@@ -103,90 +43,3 @@ def claim_affinity_pin_in_memory(
     winner: Final = existing if existing is not None and eligible_existing else pin_value
     set_local_affinity_pin(cache, cache_key, winner, ttl_seconds)
     return winner
-
-
-def _decode_pin(value: str) -> object:
-    try:
-        return _PIN_JSON_ADAPTER.validate_json(value)
-    except ValidationError:
-        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,
-    pin_value: Mapping[str, str],
-    ttl_seconds: int,
-    *,
-    eligible_values: tuple[Mapping[str, str], ...] | None = None,
-    preserve_legacy_model_ids: bool = False,
-) -> object:
-    """Return the authoritative first writer, replacing it only when it becomes ineligible.
-
-    Eligible claims refresh the returned winner. Legacy deployment claims only refresh
-    a matching candidate. Resolve Redis per call because the proxy attaches it lazily.
-    """
-    redis_cache: Final = cache.redis_cache
-    if redis_cache is not None:
-        try:
-            claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT)
-            args: Final = (
-                json.dumps(dict(pin_value)),
-                int(ttl_seconds),
-                *(
-                    (json.dumps(tuple(dict(value) for value in eligible_values)),)
-                    if eligible_values is not None
-                    else ()
-                ),
-                *(("1",) if preserve_legacy_model_ids and eligible_values is not None else ()),
-            )
-            raw: Final = cast(  # cast-ok: Redis scripts return heterogeneous values; only object is asserted here
-                object, await claim_script(keys=(cache_key,), args=args)
-            )
-            decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw
-            if not isinstance(decoded, str):
-                return pin_value
-            winner: Final = _decode_pin(decoded)
-            set_local_affinity_pin(cache, cache_key, winner, ttl_seconds)
-            return winner
-        except Exception as error:  # noqa: BLE001  # Redis/Lua faults retain same-pod affinity through local claims
-            verbose_router_logger.debug("Affinity cache: Redis claim failed, using pod-local claim. error=%s", error)
-    return claim_affinity_pin_in_memory(
-        cache,
-        cache_key,
-        pin_value,
-        ttl_seconds,
-        eligible_values=eligible_values,
-        preserve_legacy_model_ids=preserve_legacy_model_ids,
-    )

litellm/redis/stores/complexity_tier_model_affinity.py

@@ -0,0 +1,103 @@
+import hashlib
+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)
+_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
+local decoded, stored = pcall(cjson.decode, current)
+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]
+"""
+
+
+def configuration_fingerprint(config_json: str) -> str:
+    return hashlib.sha256(config_json.encode()).hexdigest()
+
+
+@dataclass(frozen=True, slots=True)
+class TierModelIdentity:
+    model_name: str
+    configuration: str
+    caller: str | None
+    session_id: str
+    tier: str
+    replay_fallback: str | None = None
+
+
+def affinity_key(identity: TierModelIdentity) -> str:
+    parts: Final = (identity.model_name, identity.configuration, identity.caller, identity.session_id, identity.tier)
+    cache_identity: Final = (
+        (*parts, ("replay_fallback", identity.replay_fallback)) if identity.replay_fallback is not None else parts
+    )
+    return "complexity_router_tier_model_affinity:v1:" + hashlib.sha256(json.dumps(cache_identity).encode()).hexdigest()
+
+
+@dataclass(frozen=True, slots=True)
+class TierModelPin:
+    value: object
+
+
+@dataclass(frozen=True, slots=True)
+class UnconfirmedTierModelPin:
+    pass
+
+
+class ComplexityTierModelAffinityStore(Protocol):
+    async def claim(
+        self, identity: TierModelIdentity, model: str, eligible_models: tuple[str, ...], *, ttl_seconds: int
+    ) -> TierModelPin | UnconfirmedTierModelPin | RedisFailure: ...
+
+
+@dataclass(frozen=True, slots=True)
+class TransportComplexityTierModelAffinityStore:
+    transport: RedisTransport
+
+    async def claim(
+        self, identity: TierModelIdentity, model: str, eligible_models: tuple[str, ...], *, ttl_seconds: int
+    ) -> TierModelPin | UnconfirmedTierModelPin | RedisFailure:
+        raw: Final = await self.transport.eval(
+            _CLAIM,
+            (affinity_key(identity),),
+            (
+                json.dumps({"model": model}),
+                int(ttl_seconds),
+                json.dumps(tuple({"model": value} for value in eligible_models)),
+            ),
+        )
+        if isinstance(raw, RedisFailure):
+            return raw
+        if not isinstance(raw, bytes):
+            return UnconfirmedTierModelPin()
+        try:
+            decoded: Final = raw.decode("utf-8")
+        except UnicodeError as error:
+            return RedisFailure(kind="decode", message=str(error))
+        try:
+            return TierModelPin(_JSON.validate_json(decoded))
+        except ValidationError:
+            return TierModelPin(decoded)

litellm/router_strategy/complexity_router/complexity_router.py

@@ -17,7 +17,6 @@ Inspired by ClawRouter: https://github.com/BlockRunAI/ClawRouter
 from __future__ import annotations
 
 import asyncio
-import hashlib
 import json
 import random
 import re
@@ -36,7 +35,11 @@ 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, set_local_affinity_pin
+from litellm.caching.affinity_cache import (
+    ROUTER_SESSION_PINS_TARGET,
+    claim_affinity_pin_in_memory,
+    set_local_affinity_pin,
+)
 from litellm.caching.redis_batch import active_request_redis_batch
 from litellm.constants import (
     COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY,
@@ -76,6 +79,18 @@ from litellm.redis.stores.complexity_session_affinity import (
 from litellm.redis.stores.complexity_session_affinity import (
     pin_value as session_pin_value,
 )
+from litellm.redis.stores.complexity_tier_model_affinity import (
+    ComplexityTierModelAffinityStore,
+    TierModelIdentity,
+    TierModelPin,
+    TransportComplexityTierModelAffinityStore,
+    UnconfirmedTierModelPin,
+    configuration_fingerprint,
+)
+from litellm.redis.stores.complexity_tier_model_affinity import (
+    affinity_key as tier_model_affinity_key,
+)
+from litellm.redis.transport.contract import RedisFailure
 from litellm.router_strategy.adaptive_router.classifier import classify_prompt
 from litellm.router_strategy.complexity_router.context_compaction import compaction_pending
 from litellm.router_strategy.complexity_router.tier_predictor import (
@@ -1515,9 +1530,9 @@ class ComplexityRouter(CustomLogger):
             case _:
                 self._jev_client: JevClassifierClient | None = jev_client_result
 
-        self._tier_affinity_config = hashlib.sha256(
-            self.config.model_dump_json(include=MappingProxyType({"tiers": True, "tier_model_configs": True})).encode()
-        ).hexdigest()
+        self._tier_affinity_config = configuration_fingerprint(
+            self.config.model_dump_json(include=MappingProxyType({"tiers": True, "tier_model_configs": True}))
+        )
 
         # Checked here rather than on the config model because the deployment's
         # complexity_router_default_model arrives outside complexity_router_config and is
@@ -3125,27 +3140,44 @@ class ComplexityRouter(CustomLogger):
         if session_id is None:
             return model
         caller: Final = DeploymentAffinityCheck.get_user_key_from_request_kwargs(request_kwargs)
-        identity: Final = (self.model_name, self._tier_affinity_config, caller, session_id, _tier_name(tier))
-        cache_identity: Final = (
-            (*identity, ("replay_fallback", retained_model)) if retained_model is not None else identity
-        )
-        cache_key: Final = (
-            "complexity_router_tier_model_affinity:v1:"
-            + hashlib.sha256(json.dumps(cache_identity).encode()).hexdigest()
-        )
-        winner: Final = await claim_affinity_pin(
-            self.litellm_router_instance.cache,
-            cache_key,
-            MappingProxyType({"model": model}),
-            self.config.session_affinity_ttl_seconds,
-            eligible_values=tuple(MappingProxyType({"model": candidate}) for candidate in candidates),
+        identity: Final = TierModelIdentity(
+            self.model_name, self._tier_affinity_config, caller, session_id, _tier_name(tier), retained_model
         )
+        winner: Final = await self._claim_tier_model_pin(identity, model, candidates)
         pinned: Final[object] = winner.get("model") if isinstance(winner, Mapping) else None
         if not isinstance(pinned, str) or pinned not in candidates:
             return model
         self._restamp_adaptive_choice(request_kwargs, model, pinned)
         return pinned
 
+    async def _claim_tier_model_pin(
+        self, identity: TierModelIdentity, model: str, candidates: tuple[str, ...]
+    ) -> object:
+        cache: Final = self.litellm_router_instance.cache
+        cache_key: Final = tier_model_affinity_key(identity)
+        ttl: Final = self.config.session_affinity_ttl_seconds
+        redis_cache: Final = cache.redis_cache
+        if redis_cache is not None:
+            store: Final[ComplexityTierModelAffinityStore] = TransportComplexityTierModelAffinityStore(
+                redis_cache.transport()
+            )
+            result: Final = await store.claim(identity, model, candidates, ttl_seconds=ttl)
+            match result:
+                case TierModelPin(value=value):
+                    set_local_affinity_pin(cache, cache_key, value, ttl)
+                    return value
+                case UnconfirmedTierModelPin():
+                    return MappingProxyType({"model": model})
+                case RedisFailure():
+                    pass
+        return claim_affinity_pin_in_memory(
+            cache,
+            cache_key,
+            MappingProxyType({"model": model}),
+            ttl,
+            eligible_values=tuple(MappingProxyType({"model": candidate}) for candidate in candidates),
+        )
+
     async def _pick_model_for_tier(
         self,
         tier: ComplexityTier | str,

tests/unit/redis/stores/test_complexity_tier_model_affinity.py

@@ -0,0 +1,234 @@
+import asyncio
+import hashlib
+import json
+from collections.abc import Iterator
+from typing import Final
+from unittest.mock import patch
+
+import fakeredis
+import pytest
+from redis.exceptions import TimeoutError as RedisTimeoutError
+
+from litellm import Router
+from litellm.caching.redis_batch import request_redis_batch_scope
+from litellm.router_strategy.complexity_router.complexity_router import (
+    ComplexityRouter,
+    _SessionAffinityPin,  # pyright: ignore[reportPrivateUsage]  # builds the retained pin the strategy carries between tiers
+)
+from tests._support.fake_redis import (
+    fake_redis_cache,
+    keys_matching,
+    local_get,
+    local_set,
+    seed_async_client,
+)
+
+
+async def _pin(
+    strategy: ComplexityRouter,
+    tier: str,
+    model: str,
+    candidates: tuple[str, ...],
+    request_kwargs: dict[str, object],
+    retained_pin: _SessionAffinityPin | None = None,
+) -> object:
+    return await strategy._pin_model_for_tier(tier, model, candidates, request_kwargs, retained_pin)  # 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 _strategy(server: fakeredis.FakeServer | None, *, name: str = "router", simple: str = "tiny") -> ComplexityRouter:
+    parent: Final = Router(
+        model_list=[
+            {"model_name": name, "litellm_params": {"model": "openai/test"}, "model_info": {"id": name}}
+            for name in ("tiny", "middle", "large", "top")
+        ]
+    )
+    if server is not None:
+        parent.cache.redis_cache = fake_redis_cache(server, namespace="tier")
+    return ComplexityRouter(
+        model_name=name,
+        litellm_router_instance=parent,
+        derive_savings_baseline=False,
+        complexity_router_config={
+            "tiers": {"SIMPLE": [simple, "middle"], "MEDIUM": "middle", "COMPLEX": "large", "REASONING": "top"},
+            "session_affinity": True,
+            "session_affinity_ttl_seconds": 60,
+        },
+    )
+
+
+def _key(strategy: ComplexityRouter, caller: str | None = "caller", *, fallback: str | None = None) -> str:
+    config: Final = strategy.config.model_dump_json(include={"tiers": True, "tier_model_configs": True})
+    identity: Final = (strategy.model_name, hashlib.sha256(config.encode()).hexdigest(), caller, "same", "SIMPLE")
+    encoded: Final = json.dumps((*identity, ("replay_fallback", fallback)) if fallback is not None else identity)
+    return "tier:complexity_router_tier_model_affinity:v1:" + hashlib.sha256(encoded.encode()).hexdigest()
+
+
+def _kwargs(caller: str | None = "caller") -> dict[str, object]:
+    return {"metadata": {"session_id": "same", "user_api_key_hash": caller}}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "name,caller,simple,fallback",
+    [
+        ("router", "caller", "tiny", None),
+        ("路由", "caller", "tiny", None),
+        ("router", "另一个", "tiny", None),
+        ("router", None, "tiny", None),
+        ("router", "caller", "模型", None),
+        ("router", "caller", "tiny", "removed"),
+    ],
+)
+async def test_identity_and_unicode_wire_match_legacy_hash(
+    name: str, caller: str | None, simple: str, fallback: str | None
+) -> None:
+    server: Final = fakeredis.FakeServer()
+    strategy: Final = _strategy(server, name=name, simple=simple)
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    retained: Final = _SessionAffinityPin(fallback, "SIMPLE") if fallback is not None else None
+    assert await _pin(strategy, "SIMPLE", "模型", ("模型", "tiny"), _kwargs(caller), retained) == "模型"
+    assert await raw.get(_key(strategy, caller, fallback=fallback)) == b'{"model": "\\u6a21\\u578b"}'
+    assert await raw.pttl(_key(strategy, caller, fallback=fallback)) == 60000
+    assert await keys_matching(raw, "tier:tier:*") == ()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "wire,expected",
+    [
+        (b'{"model": "tiny"}', "tiny"),
+        (b'{"model": "removed"}', "middle"),
+        (b'{"model": "tiny", "extra": true}', "middle"),
+        (b'"tiny"', "middle"),
+        (b"tiny", "middle"),
+        (b"{broken", "middle"),
+    ],
+)
+async def test_claim_refreshes_only_complete_eligible_winners(wire: bytes, expected: str) -> None:
+    server: Final = fakeredis.FakeServer()
+    strategy: Final = _strategy(server)
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    await raw.set(_key(strategy), wire, ex=20)
+    assert await _pin(strategy, "SIMPLE", "middle", ("tiny", "middle"), _kwargs()) == expected
+    assert await raw.get(_key(strategy)) == json.dumps({"model": expected}).encode()
+    assert await raw.pttl(_key(strategy)) == 60000
+
+
+@pytest.mark.asyncio
+async def test_concurrent_pods_converge_then_replace_expired_pin() -> None:
+    server: Final = fakeredis.FakeServer()
+    left: Final = _strategy(server)
+    right: Final = _strategy(server)
+    winners: Final = await asyncio.gather(
+        _pin(left, "SIMPLE", "tiny", ("tiny", "middle"), _kwargs()),
+        _pin(right, "SIMPLE", "middle", ("tiny", "middle"), _kwargs()),
+    )
+    assert winners[0] == winners[1]
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    assert await raw.get(_key(left)) == json.dumps({"model": winners[0]}).encode()
+    with patch("time.time", return_value=1700000061.0):
+        assert await _pin(left, "SIMPLE", "middle", ("tiny", "middle"), _kwargs()) == "middle"
+        assert await raw.get(_key(left)) == b'{"model": "middle"}'
+        assert await raw.pttl(_key(left)) == 60000
+
+
+@pytest.mark.asyncio
+async def test_no_session_and_eligible_retained_pin_do_not_claim() -> None:
+    server: Final = fakeredis.FakeServer()
+    strategy: Final = _strategy(server)
+    assert await _pin(strategy, "SIMPLE", "tiny", ("tiny", "middle"), {}) == "tiny"
+    assert (
+        await _pin(strategy, "SIMPLE", "middle", ("tiny", "middle"), _kwargs(), _SessionAffinityPin("tiny", "SIMPLE"))
+        == "tiny"
+    )
+    assert await fakeredis.FakeAsyncRedis(server=server).keys("tier:*") == []
+
+
+@pytest.mark.asyncio
+async def test_late_redis_attachment_replaces_local_winner() -> None:
+    server: Final = fakeredis.FakeServer()
+    strategy: Final = _strategy(None)
+    assert await _pin(strategy, "SIMPLE", "tiny", ("tiny", "middle"), _kwargs()) == "tiny"
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    await raw.set(_key(strategy), b'{"model": "middle"}', ex=20)
+    strategy.litellm_router_instance.cache.redis_cache = fake_redis_cache(server, namespace="tier")
+    assert await _pin(strategy, "SIMPLE", "tiny", ("tiny", "middle"), _kwargs()) == "middle"
+    assert local_get(strategy.litellm_router_instance.cache, _key(strategy).removeprefix("tier:")) == {
+        "model": "middle"
+    }
+
+
+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 claim")
+
+
+class ConsumedClaimReply(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
+        return 1
+
+
+@pytest.mark.asyncio
+async def test_consumed_claim_returns_candidate_without_installing_local_winner() -> None:
+    server: Final = fakeredis.FakeServer()
+    strategy: Final = _strategy(server)
+    cache: Final = strategy.litellm_router_instance.cache
+    redis_cache: Final = cache.redis_cache
+    assert redis_cache is not None
+    seed_async_client(redis_cache, ConsumedClaimReply(server=server))
+    local_set(cache, _key(strategy).removeprefix("tier:"), {"model": "tiny"}, ttl=20)
+    assert await _pin(strategy, "SIMPLE", "middle", ("tiny", "middle"), _kwargs()) == "middle"
+    assert await fakeredis.FakeAsyncRedis(server=server).get(_key(strategy)) == b'{"model": "middle"}'
+    assert local_get(cache, _key(strategy).removeprefix("tier:")) == {"model": "tiny"}
+
+
+@pytest.mark.asyncio
+async def test_claim_leaves_pending_request_and_post_call_work_unflushed() -> None:
+    server: Final = fakeredis.FakeServer()
+    strategy: Final = _strategy(server)
+    redis_cache: Final = strategy.litellm_router_instance.cache.redis_cache
+    assert redis_cache is not None
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    with request_redis_batch_scope(post_call_deadline=3600) as batches:
+        request: Final = batches.batch(redis_cache)
+        post_call: Final = batches.post_call(redis_cache)
+        request.set("pending-request", "ready", 60)
+        post_call.set("pending-post-call", "later", 60)
+        assert await _pin(strategy, "SIMPLE", "tiny", ("tiny", "middle"), _kwargs()) == "tiny"
+        assert await raw.get(_key(strategy)) == b'{"model": "tiny"}'
+        assert await raw.get("tier:pending-request") is None
+        assert await raw.get("tier:pending-post-call") is None
+        await batches.flush_post_call()
+        await request.flush()
+    assert await raw.get("tier:pending-request") == b'"ready"'
+    assert await raw.get("tier:pending-post-call") == b'"later"'
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("applied", [False, True])
+async def test_outage_and_unknown_claim_keep_local_winner(applied: bool) -> None:
+    server: Final = fakeredis.FakeServer()
+    strategy: Final = _strategy(server)
+    cache: Final = strategy.litellm_router_instance.cache
+    redis_cache: Final = cache.redis_cache
+    assert redis_cache is not None
+    local_set(cache, _key(strategy).removeprefix("tier:"), {"model": "tiny"}, ttl=20)
+    if applied:
+        seed_async_client(redis_cache, LostClaimReply(server=server))
+    else:
+        server.connected = False
+    assert await _pin(strategy, "SIMPLE", "middle", ("tiny", "middle"), _kwargs()) == "tiny"
+    server.connected = True
+    assert await fakeredis.FakeAsyncRedis(server=server).get(_key(strategy)) == (
+        b'{"model": "middle"}' if applied else None
+    )
+    with patch("time.time", return_value=1700000021.0):
+        assert local_get(cache, _key(strategy).removeprefix("tier:")) == {"model": "tiny"}

tests/unit/router_strategy/test_complexity_router.py

@@ -14,7 +14,7 @@ from collections.abc import AsyncIterator, Mapping, Sequence
 from copy import deepcopy
 from functools import partial
 from types import MappingProxyType
-from typing import Dict, Final, List, Literal
+from typing import Final, Literal
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
@@ -39,7 +39,6 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 from litellm.router import as_output_cap, stamped_client_ceiling
 from litellm.router_strategy.auto_router.auto_router_model_naming import (
     CUSTOMIZATION_CAPABILITY,
-    GATED_AUTO_ROUTER_CAPABILITIES,
     HEURISTIC_V2_CAPABILITY,
     count_capability_routers,
 )
@@ -128,7 +127,7 @@ def mock_router_instance():
 
 
 @pytest.fixture
-def basic_config() -> Dict:
+def basic_config() -> dict:
     """Basic configuration with tier mappings."""
     return {
         "tiers": {
@@ -1498,7 +1497,6 @@ class TestSingletonMutation:
     def test_default_config_not_mutated(self, mock_router_instance):
         """Test that creating routers without config doesn't mutate defaults."""
         from litellm.router_strategy.complexity_router.config import (
-            DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE,
             ComplexityRouterConfig,
         )
 
@@ -1554,7 +1552,7 @@ class TestKeywordFalsePositives:
         tier, score, signals = complexity_router.classify(prompt)
         # 'entry' contains 'try' but should not trigger code detection
         # Note: 'application' might trigger something, but 'try' should not
-        pass  # Just ensure no crash; false positive check is the main goal
+        # Just ensure no crash; false positive check is the main goal
 
     def test_error_not_in_terrorism(self, complexity_router):
         """'error' should not match in 'terrorism'."""
@@ -1696,17 +1694,19 @@ class TestRouterComplexityDeploymentMethods:
     def test_custom_default_resolution_preserves_explicit_precedence(self, tier_models, explicit, expected):
         from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig
 
-        config: Final = ComplexityRouterConfig.model_validate({
-            "tiers": {"CUSTOM": tier_models, "OTHER": "other-model"},
-            "tier_definitions": [
-                {"name": "CUSTOM", "description": "Custom work"},
-                {"name": "OTHER", "description": "Other work"},
-            ],
-            "fallback_tier": " CUSTOM ",
-            "classifier_type": "llm",
-            "classifier_llm_config": {"model": "classifier"},
-            "default_model": "configured-model",
-        })
+        config: Final = ComplexityRouterConfig.model_validate(
+            {
+                "tiers": {"CUSTOM": tier_models, "OTHER": "other-model"},
+                "tier_definitions": [
+                    {"name": "CUSTOM", "description": "Custom work"},
+                    {"name": "OTHER", "description": "Other work"},
+                ],
+                "fallback_tier": " CUSTOM ",
+                "classifier_type": "llm",
+                "classifier_llm_config": {"model": "classifier"},
+                "default_model": "configured-model",
+            }
+        )
         assert config.resolve_default_model(explicit) == expected
         inferred: Final = config.model_copy(update={"default_model": None})
         assert inferred.resolve_default_model() == "custom-model"
@@ -2142,9 +2142,7 @@ class TestRouterComplexityDeploymentMethods:
 
     @pytest.mark.parametrize("classifier", ("heuristic_first", "hybrid"))
     @pytest.mark.parametrize("held_capability", ("heuristic_v2", "custom_prompt"))
-    def test_v2_chain_registration_enforces_each_claimed_allowance(
-        self, classifier: str, held_capability: str
-    ) -> None:
+    def test_v2_chain_registration_enforces_each_claimed_allowance(self, classifier: str, held_capability: str) -> None:
         held: Final = (
             self._router_row("existing", "existing-id", "heuristic_v2")
             if held_capability == "heuristic_v2"
@@ -2160,7 +2158,8 @@ class TestRouterComplexityDeploymentMethods:
                     "local_heuristic": "heuristic_v2",
                     **(
                         {"heuristic_first_max_tier": "MEDIUM"}
-                        if classifier == "heuristic_first" else {"hybrid_boundary_margin": 0.1}
+                        if classifier == "heuristic_first"
+                        else {"hybrid_boundary_margin": 0.1}
                     ),
                     "classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "Grade by difficulty"},
                     "tiers": {"SIMPLE": "gpt-4o-mini"},
@@ -2785,7 +2784,7 @@ def _wrapped_reply(shape: str, verdict: str) -> str:
 
 
 @pytest.fixture
-def llm_classifier_config() -> Dict:
+def llm_classifier_config() -> dict:
     """Config with an LLM-based classifier wired to a 'haiku-classifier' model."""
     return {
         "tiers": {
@@ -2839,7 +2838,7 @@ class TestLLMClassifierConfig:
             )
 
 
-CAPABILITY_TIERS: Dict[str, str] = {
+CAPABILITY_TIERS: dict[str, str] = {
     "SIMPLE": "efficient-model",
     "REASONING": "capable-model",
 }
@@ -3415,7 +3414,7 @@ class TestCapabilityClassifier:
         mock_router_instance.acompletion.assert_awaited_once()
 
 
-CUSTOM_TIER_LABELS: Dict[str, str] = {
+CUSTOM_TIER_LABELS: dict[str, str] = {
     "SIMPLE": "Cheap",
     "MEDIUM": "Standard",
     "COMPLEX": "Premium",
@@ -4856,7 +4855,7 @@ class TestRouterPreRoutingAliasOverrides:
         reach the outbound request even though the tier deployment is what
         actually gets called."""
         router = self._make_router()
-        request_kwargs: Dict = {}
+        request_kwargs: dict = {}
 
         result = await router.async_pre_routing_hook(
             model="smart-router",
@@ -4889,7 +4888,7 @@ class TestRouterPreRoutingAliasOverrides:
                 {"model_name": "gpt-5-mini", "litellm_params": {"model": "openai/gpt-5-mini"}},
             ]
         )
-        request_kwargs: Dict = {"reasoning_effort": "low"}
+        request_kwargs: dict = {"reasoning_effort": "low"}
 
         deployment = await router.async_get_available_deployment(
             model="smart-router",
@@ -4900,7 +4899,7 @@ class TestRouterPreRoutingAliasOverrides:
         assert deployment["model_name"] == "gpt-5-mini"
         assert request_kwargs["reasoning_effort"] == "xhigh"
 
-    def _make_effort_pinned_router(self, tier_litellm_params: Dict) -> Router:
+    def _make_effort_pinned_router(self, tier_litellm_params: dict) -> Router:
         return Router(
             model_list=[
                 {
@@ -4951,7 +4950,7 @@ class TestRouterPreRoutingAliasOverrides:
         carrier precedence over the reasoning_effort alias, so the pin only
         reaches the wire if those carriers are dropped at the merge."""
         router = self._make_effort_pinned_router({"reasoning_effort": "xhigh"})
-        request_kwargs: Dict = dict(client_carriers)
+        request_kwargs: dict = dict(client_carriers)
 
         await router.async_get_available_deployment(
             model="smart-router",
@@ -4989,7 +4988,7 @@ class TestRouterPreRoutingAliasOverrides:
                 },
             ]
         )
-        request_kwargs: Dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}}
+        request_kwargs: dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}}
 
         await router.async_get_available_deployment_for_pass_through(
             model="smart-router",
@@ -5002,15 +5001,15 @@ class TestRouterPreRoutingAliasOverrides:
         assert "output_config" not in request_kwargs
 
     def test_drop_client_effort_carriers_helper_edge_shapes(self):
-        no_pin: Dict = {"thinking": {"type": "adaptive"}}
+        no_pin: dict = {"thinking": {"type": "adaptive"}}
         Router._drop_client_carriers_a_tier_pin_supersedes(no_pin, {"temperature": 0.1})
         assert no_pin == {"thinking": {"type": "adaptive"}}
 
-        non_dict_carriers: Dict = {"output_config": "max", "reasoning": 3}
+        non_dict_carriers: dict = {"output_config": "max", "reasoning": 3}
         Router._drop_client_carriers_a_tier_pin_supersedes(non_dict_carriers, {"reasoning_effort": "low"})
         assert non_dict_carriers == {"output_config": "max", "reasoning": 3}
 
-        effort_only: Dict = {"output_config": {"effort": "max"}, "reasoning": {"effort": "high"}}
+        effort_only: dict = {"output_config": {"effort": "max"}, "reasoning": {"effort": "high"}}
         Router._pop_effort_from_nested_carrier(effort_only, "output_config")
         Router._pop_effort_from_nested_carrier(effort_only, "reasoning")
         assert effort_only == {}
@@ -5038,7 +5037,7 @@ class TestRouterPreRoutingAliasOverrides:
                 {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}},
             ]
         )
-        request_kwargs: Dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}}
+        request_kwargs: dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}}
 
         await router.async_get_available_deployment(
             model="smart-router",
@@ -5053,7 +5052,7 @@ class TestRouterPreRoutingAliasOverrides:
     @pytest.mark.asyncio
     async def test_client_effort_carriers_survive_when_tier_pins_no_effort(self):
         router = self._make_effort_pinned_router({"temperature": 0.2})
-        request_kwargs: Dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}}
+        request_kwargs: dict = {"thinking": {"type": "adaptive"}, "output_config": {"effort": "max"}}
 
         await router.async_get_available_deployment(
             model="smart-router",
@@ -5098,7 +5097,7 @@ class TestRouterPreRoutingAliasOverrides:
             ]
         )
         real_get_llm_provider = litellm.get_llm_provider
-        copilot_resolutions: List = []
+        copilot_resolutions: list = []
 
         def _guarded(*args, **kwargs):
             target = str(kwargs.get("model") or (args[0] if args else "")) + str(
@@ -5110,7 +5109,7 @@ class TestRouterPreRoutingAliasOverrides:
             return real_get_llm_provider(*args, **kwargs)
 
         monkeypatch.setattr(litellm, "get_llm_provider", _guarded)
-        request_kwargs: Dict = {}
+        request_kwargs: dict = {}
 
         deployment = await router.async_get_available_deployment(
             model="smart-router",
@@ -5178,7 +5177,7 @@ class TestRouterPreRoutingAliasOverrides:
         test_router_init_only_params_are_never_sent_to_a_provider for the
         guard on that downstream filter."""
         router = self._make_router()
-        request_kwargs: Dict = {}
+        request_kwargs: dict = {}
 
         await router.async_pre_routing_hook(
             model="smart-router",
@@ -5232,7 +5231,7 @@ class TestRouterPreRoutingAliasOverrides:
         """A value the caller already passed for this request takes
         precedence over the alias's configured default."""
         router = self._make_router()
-        request_kwargs: Dict = {"drop_params": False}
+        request_kwargs: dict = {"drop_params": False}
 
         await router.async_pre_routing_hook(
             model="smart-router",
@@ -5247,7 +5246,7 @@ class TestRouterPreRoutingAliasOverrides:
         """A plain (non-router-alias) model name is not affected by the
         alias-override merge at all."""
         router = self._make_router()
-        request_kwargs: Dict = {}
+        request_kwargs: dict = {}
 
         result = await router.async_pre_routing_hook(
             model="gpt-4o-mini",
@@ -5282,7 +5281,7 @@ class TestRouterPreRoutingAliasOverrides:
         router.set_model_list(model_list)
         assert "smart-router" in router.adaptive_routers
 
-        request_kwargs: Dict = {}
+        request_kwargs: dict = {}
         await router.async_pre_routing_hook(
             model="smart-router",
             request_kwargs=request_kwargs,
@@ -5344,7 +5343,7 @@ class TestRouterPreRoutingSharedAliasName:
             else [self._marker_entry(), self._plain_entry()]
         )
         router = Router(model_list=[*shared_name_entries, self._tier_entry()])
-        request_kwargs: Dict = {}
+        request_kwargs: dict = {}
 
         result = await router.async_pre_routing_hook(
             model="gpt4o",
@@ -5374,7 +5373,7 @@ class TestRouterPreRoutingSharedAliasName:
             },
         }
         router = Router(model_list=[marker_with_connection_params, self._tier_entry()])
-        request_kwargs: Dict = {}
+        request_kwargs: dict = {}
 
         result = await router.async_pre_routing_hook(
             model="smart",
@@ -5414,7 +5413,7 @@ class TestRouterPreRoutingSharedAliasName:
             ]
         )
 
-        us_kwargs: Dict = {"metadata": {"tags": ["us"]}}
+        us_kwargs: dict = {"metadata": {"tags": ["us"]}}
         us_result = await router.async_pre_routing_hook(
             model="smart",
             request_kwargs=us_kwargs,
@@ -5423,7 +5422,7 @@ class TestRouterPreRoutingSharedAliasName:
         assert us_result is not None and us_result.model == "gpt-us"
         assert us_kwargs["drop_params"] is True
 
-        cn_kwargs: Dict = {"metadata": {"tags": ["cn"]}}
+        cn_kwargs: dict = {"metadata": {"tags": ["cn"]}}
         cn_result = await router.async_pre_routing_hook(
             model="smart",
             request_kwargs=cn_kwargs,
@@ -5597,7 +5596,7 @@ class TestRouterPreRoutingSharedAliasName:
             await asyncio.sleep(0.01)
             return await healthy_deployments(*args, **kwargs)
 
-        sent: Dict[str, str | None] = {}
+        sent: dict[str, str | None] = {}
 
         async def record(**kwargs):
             sent[kwargs["model"]] = kwargs.get("aws_region_name")
@@ -5682,7 +5681,7 @@ class TestAdaptiveSoftFloors:
         return router
 
     @pytest.fixture
-    def hybrid_config(self) -> Dict:
+    def hybrid_config(self) -> dict:
         return {
             "adaptive": True,
             "adaptive_weights": {"quality": 0.7, "cost": 0.3},
@@ -5712,7 +5711,7 @@ class TestAdaptiveSoftFloors:
                 },
             },
         )
-        request_kwargs: Dict = {"metadata": {}}
+        request_kwargs: dict = {"metadata": {}}
 
         with patch(
             "litellm.router_strategy.complexity_router.complexity_router.random.choice",
@@ -5811,7 +5810,7 @@ class TestAdaptiveSoftFloors:
         assert adaptive is not None
         for model in ("cheap", "premium"):
             adaptive._cells[(RequestType.GENERAL, model)] = BanditCell(alpha=6.0, beta=5.0)
-        request_kwargs: Dict = {"metadata": {}}
+        request_kwargs: dict = {"metadata": {}}
 
         with patch(
             "litellm.router_strategy.adaptive_router.bandit.thompson_sample",
@@ -5832,7 +5831,7 @@ class TestAdaptiveSoftFloors:
             litellm_router_instance=adaptive_router_instance,
             complexity_router_config=hybrid_config,
         )
-        request_kwargs: Dict = {"metadata": {}}
+        request_kwargs: dict = {"metadata": {}}
         result = await cr.async_pre_routing_hook(
             model="hybrid",
             request_kwargs=request_kwargs,
@@ -5854,7 +5853,7 @@ class TestLexicalKeywordTierRules:
     """Test deterministic (literal) keyword_tier_rules overrides."""
 
     @pytest.fixture
-    def rule_config(self, basic_config) -> Dict:
+    def rule_config(self, basic_config) -> dict:
         return {
             **basic_config,
             "keyword_tier_rules": [
@@ -5994,7 +5993,7 @@ class TestLexicalKeywordTierRules:
 class TestCjkKeywordTierRules:
     """CJK keyword_tier_rules must fire mid-sentence, where regex word boundaries cannot."""
 
-    def _router(self, mock_router_instance, basic_config, keywords: List[str]) -> ComplexityRouter:
+    def _router(self, mock_router_instance, basic_config, keywords: list[str]) -> ComplexityRouter:
         return ComplexityRouter(
             model_name="test-router",
             litellm_router_instance=mock_router_instance,
@@ -6069,7 +6068,7 @@ class TestCjkKeywordTierRules:
         assert complexity_router._keyword_matches("appelle l' api maintenant", "api") is True
 
 
-def _make_embedding_response(vectors: List[List[float]]) -> "litellm.EmbeddingResponse":
+def _make_embedding_response(vectors: list[list[float]]) -> "litellm.EmbeddingResponse":
     return litellm.EmbeddingResponse(
         model="fake-embed",
         data=[{"embedding": vec, "index": idx, "object": "embedding"} for idx, vec in enumerate(vectors)],
@@ -6088,22 +6087,22 @@ class FakeEmbeddingRouter:
     _CLUSTER_MARKERS = ("k8s", "kube", "container", "cluster", "orchestrat")
 
     def __init__(self):
-        self.async_embedding_calls: List[List[str]] = []
-        self.async_embedding_kwargs: List[Dict] = []
+        self.async_embedding_calls: list[list[str]] = []
+        self.async_embedding_kwargs: list[dict] = []
         # Every embedded batch (sync route-index build AND async query), so tests can count
         # builds independently of which embedding path the library happens to use.
-        self.embedded_batches: List[List[str]] = []
+        self.embedded_batches: list[list[str]] = []
         # Thread ids of the synchronous (route-index build) embedding calls, so a test can
         # assert the build is offloaded off the event-loop thread.
-        self.sync_embedding_thread_ids: List[int] = []
+        self.sync_embedding_thread_ids: list[int] = []
 
-    def _vectors(self, docs: List[str]) -> List[List[float]]:
+    def _vectors(self, docs: list[str]) -> list[list[float]]:
         return [
             [1.0, 0.0] if any(marker in doc.lower() for marker in self._CLUSTER_MARKERS) else [0.0, 1.0] for doc in docs
         ]
 
     @staticmethod
-    def _as_list(text) -> List[str]:
+    def _as_list(text) -> list[str]:
         return text if isinstance(text, list) else [text]
 
     def embedding(self, input, model, **kwargs):
@@ -6598,7 +6597,7 @@ class _StubEncoder:
     """Minimal stand-in for LiteLLMRouterEncoder.aencode_queries, capturing the kwargs it was called with."""
 
     def __init__(self):
-        self.aencode_queries_calls: List[Dict] = []
+        self.aencode_queries_calls: list[dict] = []
 
     async def aencode_queries(self, docs, **kwargs):
         self.aencode_queries_calls.append(kwargs)
@@ -7198,11 +7197,11 @@ class TestSessionAffinity:
     SIMPLE_MESSAGE = [{"role": "user", "content": "Hello!"}]
 
     @pytest.fixture
-    def session_affinity_config(self, basic_config) -> Dict:
+    def session_affinity_config(self, basic_config) -> dict:
         return {**basic_config, "session_affinity": True}
 
     @staticmethod
-    def _request_kwargs(session_id: str) -> Dict:
+    def _request_kwargs(session_id: str) -> dict:
         return {"metadata": {"session_id": session_id}}
 
     @pytest.mark.asyncio
@@ -7609,9 +7608,9 @@ class TestSessionAffinity:
                 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)
+                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:
@@ -7737,8 +7736,12 @@ class TestChildSessionTierPins:
                         "model_info": {"id": f"{group}-{index}"},
                     }
                     for group, index in (
-                        ("simple", 0), ("simple", 1), ("reasoning", 0),
-                        ("reasoning", 1), ("reasoning-peer", 0), ("reasoning-peer", 1),
+                        ("simple", 0),
+                        ("simple", 1),
+                        ("reasoning", 0),
+                        ("reasoning", 1),
+                        ("reasoning-peer", 0),
+                        ("reasoning-peer", 1),
                     )
                 ],
             ],
@@ -7754,9 +7757,9 @@ class TestChildSessionTierPins:
             metadata_variable_name_for_route,
         )
 
-        route: Final = {
-            "chat": "/v1/chat/completions", "responses": "/v1/responses", "messages": "/v1/messages"
-        }[surface]
+        route: Final = {"chat": "/v1/chat/completions", "responses": "/v1/responses", "messages": "/v1/messages"}[
+            surface
+        ]
         bucket: Final = metadata_variable_name_for_route(route)
         data: Final[dict[str, object]] = {bucket: {"user_api_key_hash": caller, "user_api_key_request_route": route}}
         LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
@@ -7770,54 +7773,74 @@ class TestChildSessionTierPins:
         return data, bucket
 
     async def _route(
-        self, router: Router, surface: str, session_id: str | None, prompt: str | None,
-        expected_tier: str, expected_cause: str, caller: str = "caller-a", child: bool = True,
+        self,
+        router: Router,
+        surface: str,
+        session_id: str | None,
+        prompt: str | None,
+        expected_tier: str,
+        expected_cause: str,
+        caller: str = "caller-a",
+        child: bool = True,
         expected_deployment: str | None = None,
     ) -> str:
         from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck
 
         data, bucket = self._normalized_request(surface, session_id, caller, child)
         turn: Final[dict[str, object]] = (
-            {"role": "user", "content": prompt} if prompt is not None else
-            {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "done"}]}
-            if surface == "messages" else
-            {"type": "function_call_output", "call_id": "call_1", "output": "done"}
-            if surface == "responses" else
-            {"role": "tool", "tool_call_id": "call_1", "content": "done"}
+            {"role": "user", "content": prompt}
+            if prompt is not None
+            else {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "done"}]}
+            if surface == "messages"
+            else {"type": "function_call_output", "call_id": "call_1", "output": "done"}
+            if surface == "responses"
+            else {"role": "tool", "tool_call_id": "call_1", "content": "done"}
         )
         tool_call: Final = (
             {"type": "function_call", "call_id": "call_1", "name": "read_file", "arguments": "{}"}
-            if surface == "responses" else
-            {"role": "assistant", "content": [{"type": "tool_use", "id": "call_1", "name": "read_file", "input": {}}]}
-            if surface == "messages" else TestClassificationMode.TOOL_CALL_1
+            if surface == "responses"
+            else {
+                "role": "assistant",
+                "content": [{"type": "tool_use", "id": "call_1", "name": "read_file", "input": {}}],
+            }
+            if surface == "messages"
+            else TestClassificationMode.TOOL_CALL_1
         )
         turns: Final = [turn] if prompt is not None else [TestClassificationMode.SIMPLE_ASK, tool_call, turn]
         if surface == "responses":
             data["input"] = turns
         selected: Final = await router.async_get_available_deployment(
-            model="child-tier-router", request_kwargs=data,
+            model="child-tier-router",
+            request_kwargs=data,
             messages=None if surface == "responses" else turns,
             input=turns if surface == "responses" else None,
         )
         metadata: Final = data[bucket]
         assert isinstance(metadata, dict)
         assert (metadata["session_id"], metadata["trace_id"], data["litellm_session_id"], data["litellm_trace_id"]) == (
-            "parent-session", "parent-session", "parent-session", "parent-session"
+            "parent-session",
+            "parent-session",
+            "parent-session",
+            "parent-session",
         )
         assert (metadata["routing_decision"]["tier"], metadata["routing_decision"]["cause"]) == (
-            expected_tier, expected_cause
+            expected_tier,
+            expected_cause,
         )
         affinity: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck))
         if expected_deployment is not None:
             filtered: Final = await affinity.async_filter_deployments(
-                model=selected["model_name"], healthy_deployments=router.get_model_list(selected["model_name"]) or [],
-                messages=None, request_kwargs=data,
+                model=selected["model_name"],
+                healthy_deployments=router.get_model_list(selected["model_name"]) or [],
+                messages=None,
+                request_kwargs=data,
             )
             assert [deployment["model_info"]["id"] for deployment in filtered] == [expected_deployment]
             assert selected["model_info"]["id"] == expected_deployment
         await affinity.async_pre_call_deployment_hook(
             kwargs={
-                **data, "model_info": selected["model_info"],
+                **data,
+                "model_info": selected["model_info"],
                 bucket: {**metadata, "deployment_model_name": selected["model_name"]},
             },
             call_type=None,
@@ -7845,13 +7868,29 @@ class TestChildSessionTierPins:
                 (None, None, "SIMPLE", "heuristic_scorer", "caller-a", True),
             ):
                 await self._route(
-                    router, surface, session, prompt, tier, cause, caller, child,
+                    router,
+                    surface,
+                    session,
+                    prompt,
+                    tier,
+                    cause,
+                    caller,
+                    child,
                     expected_deployment=parent if tier == "REASONING" and caller == "caller-a" else None,
                 )
-            assert await self._route(
-                router, surface, "parent-session", None, "REASONING", "user_turn_continuation",
-                child=False, expected_deployment=parent,
-            ) == parent
+            assert (
+                await self._route(
+                    router,
+                    surface,
+                    "parent-session",
+                    None,
+                    "REASONING",
+                    "user_turn_continuation",
+                    child=False,
+                    expected_deployment=parent,
+                )
+                == parent
+            )
         finally:
             router.discard()
 
@@ -7862,10 +7901,18 @@ class TestChildSessionTierPins:
             parent: Final = await self._route(
                 router, "chat", "parent-session", "reason carefully", "REASONING", "literal_keyword_match", child=False
             )
-            assert await self._route(
-                router, "chat", "child-session-a", "Hello!", "REASONING", "session_affinity_pin",
-                expected_deployment=parent,
-            ) == parent
+            assert (
+                await self._route(
+                    router,
+                    "chat",
+                    "child-session-a",
+                    "Hello!",
+                    "REASONING",
+                    "session_affinity_pin",
+                    expected_deployment=parent,
+                )
+                == parent
+            )
         finally:
             router.discard()
 
@@ -8804,7 +8851,7 @@ class TestEscalationKeywords:
     one step higher so a user can force a stronger model when unhappy with results."""
 
     @staticmethod
-    def _request_kwargs(session_id: str) -> Dict:
+    def _request_kwargs(session_id: str) -> dict:
         return {"metadata": {"session_id": session_id}}
 
     def test_default_escalation_keyword(self, complexity_router):
@@ -9140,7 +9187,7 @@ class TestEscalationKeywords:
         assert result.model == "o1-b"  # unchanged: no random hop to o1-a / o1-c
 
 
-def _stalled_tool_history(repeats: int = 3) -> List[Dict]:
+def _stalled_tool_history(repeats: int = 3) -> list[dict]:
     """`repeats` identical bash tool calls in a row, the automatic counterpart to a user
     typing an escalation keyword: the assistant, not the human, is the one stuck."""
     return [
@@ -9692,7 +9739,7 @@ class TestRoutingDecisionSurvivesToSpendLogOnEveryMetadataShape:
         # Mirror function_setup: it copies `litellm_metadata` by value into
         # litellm_params AFTER the router hook has run, so the copy must carry
         # the decision. Reading the stash any earlier would lose it.
-        litellm_params: Dict = {}
+        litellm_params: dict = {}
         if "metadata" in request_kwargs:
             litellm_params["metadata"] = request_kwargs["metadata"]
         if isinstance(request_kwargs.get("litellm_metadata"), dict):
@@ -9708,11 +9755,15 @@ class TestRoutingDecisionSurvivesToSpendLogOnEveryMetadataShape:
         assert persisted is not None, f"routing_decision dropped for {expected_bucket}"
         assert persisted["router_model_name"] == "smart-router"
         if classifier_type == "heuristic_v2":
-            assert persisted["heuristic_v2_forecast"] == request_kwargs[expected_bucket]["routing_decision"][
-                "heuristic_v2_forecast"
-            ]
+            assert (
+                persisted["heuristic_v2_forecast"]
+                == request_kwargs[expected_bucket]["routing_decision"]["heuristic_v2_forecast"]
+            )
             assert set(persisted["heuristic_v2_forecast"]["probabilities"]) == {
-                "SIMPLE", "MEDIUM", "COMPLEX", "REASONING"
+                "SIMPLE",
+                "MEDIUM",
+                "COMPLEX",
+                "REASONING",
             }
         else:
             assert "heuristic_v2_forecast" not in persisted
@@ -9747,7 +9798,7 @@ class TestRoutingDecisionIsPerAttempt:
     @pytest.mark.asyncio
     async def test_fallback_to_plain_model_group_clears_the_earlier_decision(self, seed, bucket):
         router = Router(model_list=self.MODEL_LIST)
-        request_kwargs: Dict = dict(seed)
+        request_kwargs: dict = dict(seed)
         messages = [{"role": "user", "content": "Hello!"}]
 
         await router.async_pre_routing_hook(model="smart-router", request_kwargs=request_kwargs, messages=messages)
@@ -9767,7 +9818,7 @@ class TestRoutingDecisionIsPerAttempt:
         Skipping the write there would drop provenance on a successfully routed
         request with no error, so the shared bucket owner replaces the value."""
         router = Router(model_list=self.MODEL_LIST)
-        request_kwargs: Dict = {"litellm_metadata": unusable_bucket}
+        request_kwargs: dict = {"litellm_metadata": unusable_bucket}
 
         response = await router.async_pre_routing_hook(
             model="smart-router",
@@ -9788,7 +9839,7 @@ class TestRecordRoutingDecision:
     DECISION = {"router_model_name": "smart-router", "router_type": "complexity", "routed_model": "gpt-4o-mini"}
 
     def test_none_clears_a_previous_decision_from_both_buckets(self):
-        request_kwargs: Dict = {
+        request_kwargs: dict = {
             "metadata": {"routing_decision": self.DECISION, "keep": 1},
             "litellm_metadata": {"routing_decision": self.DECISION},
         }
@@ -9798,7 +9849,7 @@ class TestRecordRoutingDecision:
         assert request_kwargs["metadata"]["keep"] == 1
 
     def test_none_creates_no_bucket_on_a_request_that_had_none(self):
-        request_kwargs: Dict = {}
+        request_kwargs: dict = {}
         Router._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None)
         assert request_kwargs == {}
 
@@ -9957,7 +10008,7 @@ class TestRedactedLoggingDropsPromptText:
 
     MESSAGES = [{"role": "user", "content": "LITELLM ESCALATE please deploy to k8s now"}]
 
-    async def _decision(self, request_kwargs: Dict) -> Dict:
+    async def _decision(self, request_kwargs: dict) -> dict:
         router = Router(model_list=self.MODEL_LIST)
         response = await router.async_pre_routing_hook(
             model="smart-router", request_kwargs=request_kwargs, messages=self.MESSAGES
@@ -10026,7 +10077,7 @@ class TestRedactedLoggingDropsPromptText:
 
     @pytest.mark.asyncio
     async def test_redaction_via_request_header_is_honored(self):
-        request_kwargs: Dict = {"metadata": {"headers": {"x-litellm-enable-message-redaction": True}}}
+        request_kwargs: dict = {"metadata": {"headers": {"x-litellm-enable-message-redaction": True}}}
         decision = await self._decision(request_kwargs)
         assert "matched_keyword" not in decision
         assert decision["cause"] == "literal_keyword_match"
@@ -10042,8 +10093,10 @@ class TestRedactedLoggingDropsPromptText:
             {
                 "classifier_type": "capability",
                 "capability_classifier_config": {
-                    "efficient_tier": "SIMPLE", "capable_tier": "REASONING",
-                    "base_threshold": 0.5, "threshold_step": 0.1,
+                    "efficient_tier": "SIMPLE",
+                    "capable_tier": "REASONING",
+                    "base_threshold": 0.5,
+                    "threshold_step": 0.1,
                 },
             },
             "capability_classifier_fallback",
@@ -10052,8 +10105,10 @@ class TestRedactedLoggingDropsPromptText:
             {
                 "classifier_type": "llm_v2",
                 "llm_v2_config": {
-                    "efficient_profile": "Routine tasks", "capable_profile": "Complex tasks",
-                    "harness": "Read files and run checks", "max_quality_gap": 0.05,
+                    "efficient_profile": "Routine tasks",
+                    "capable_profile": "Complex tasks",
+                    "harness": "Read files and run checks",
+                    "max_quality_gap": 0.05,
                 },
             },
             "llm_v2_fallback",
@@ -10064,36 +10119,46 @@ async def test_classifier_failure_survives_real_router_decision_and_redaction(
     overrides: Mapping[str, object], expected_cause: str, monkeypatch: pytest.MonkeyPatch
 ) -> None:
     monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
-    router: Final = Router(model_list=[
-        {
-            "model_name": "auto",
-            "litellm_params": {
-                "model": "auto_router/complexity_router",
-                "complexity_router_config": {
-                    "tiers": {"SIMPLE": "cheap", "REASONING": "capable"},
-                    "default_model": "cheap",
-                    "classifier_llm_config": {"model": "classifier", "timeout_ms": 2000},
-                    **overrides,
+    router: Final = Router(
+        model_list=[
+            {
+                "model_name": "auto",
+                "litellm_params": {
+                    "model": "auto_router/complexity_router",
+                    "complexity_router_config": {
+                        "tiers": {"SIMPLE": "cheap", "REASONING": "capable"},
+                        "default_model": "cheap",
+                        "classifier_llm_config": {"model": "classifier", "timeout_ms": 2000},
+                        **overrides,
+                    },
                 },
             },
-        },
-        {
-            "model_name": "classifier",
-            "litellm_params": {
-                "model": "openai/test-classifier", "api_key": "test-key",
-                "api_base": "https://classifier.test/v1",
+            {
+                "model_name": "classifier",
+                "litellm_params": {
+                    "model": "openai/test-classifier",
+                    "api_key": "test-key",
+                    "api_base": "https://classifier.test/v1",
+                },
             },
-        },
-        {"model_name": "cheap", "litellm_params": {"model": "openai/test-cheap", "api_key": "test-key"}},
-        {"model_name": "capable", "litellm_params": {"model": "openai/test-capable", "api_key": "test-key"}},
-    ])
+            {"model_name": "cheap", "litellm_params": {"model": "openai/test-cheap", "api_key": "test-key"}},
+            {"model_name": "capable", "litellm_params": {"model": "openai/test-capable", "api_key": "test-key"}},
+        ]
+    )
     with respx.mock:
         upstream: Final = respx.post("https://classifier.test/v1/chat/completions").respond(
             200,
             json={
-                "id": "classifier-result", "object": "chat.completion", "model": "test-classifier",
-                "choices": [{"index": 0, "message": {"role": "assistant", "content": "PRIVATE_INVALID_REPLY"},
-                             "finish_reason": "stop"}],
+                "id": "classifier-result",
+                "object": "chat.completion",
+                "model": "test-classifier",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {"role": "assistant", "content": "PRIVATE_INVALID_REPLY"},
+                        "finish_reason": "stop",
+                    }
+                ],
             },
         )
         metadata: Final[dict[str, object]] = {"headers": {"x-litellm-enable-message-redaction": True}}
@@ -10110,11 +10175,21 @@ async def test_classifier_failure_survives_real_router_decision_and_redaction(
         assert "PRIVATE_INVALID_REPLY" not in json.dumps(decision)
         assert "signals" not in decision
         if overrides["classifier_type"] == "llm":
-            upstream.respond(200, json={
-                "id": "classifier-recovered", "object": "chat.completion", "model": "test-classifier",
-                "choices": [{"index": 0, "message": {"role": "assistant", "content": '{"tier":"SIMPLE"}'},
-                             "finish_reason": "stop"}],
-            })
+            upstream.respond(
+                200,
+                json={
+                    "id": "classifier-recovered",
+                    "object": "chat.completion",
+                    "model": "test-classifier",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "message": {"role": "assistant", "content": '{"tier":"SIMPLE"}'},
+                            "finish_reason": "stop",
+                        }
+                    ],
+                },
+            )
             recovered: Final = await router.async_pre_routing_hook(
                 model="auto", request_kwargs=request, messages=[{"role": "user", "content": "hi"}]
             )
@@ -12528,7 +12603,7 @@ class TestClassifierFallbackChoice:
             },
         )
         mock_router_instance.cache = DualCache()
-        request_kwargs: Dict = {"metadata": {"session_id": "session-flaky"}}
+        request_kwargs: dict = {"metadata": {"session_id": "session-flaky"}}
 
         mock_router_instance.acompletion = AsyncMock(side_effect=TimeoutError("classifier timed out"))
         first = await router.async_pre_routing_hook(
@@ -12567,7 +12642,7 @@ class TestClassifierFallbackChoice:
             },
         )
         mock_router_instance.cache = DualCache()
-        request_kwargs: Dict = {"metadata": {"session_id": "session-steady"}}
+        request_kwargs: dict = {"metadata": {"session_id": "session-steady"}}
 
         mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "REASONING"}'))
         first = await router.async_pre_routing_hook(
@@ -13045,7 +13120,7 @@ class TestClassificationRubrics:
         assert config.classifier_llm_config.system_prompt == "Grade the data sensitivity of the request."
 
 
-def _custom_tier_config(**overrides) -> Dict:
+def _custom_tier_config(**overrides) -> dict:
     """A valid operator-defined tier set: two built-in names plus one custom tier."""
     return {
         "tiers": {"SIMPLE": "gpt-4o-mini", "COMPLEX": "claude-sonnet-4-20250514", "SECURITY_REVIEW": "o1-preview"},
@@ -14069,8 +14144,8 @@ async def test_tier_params_mask_credentials_in_routing_decision(route, mock_rout
         complexity_router_config=config,
     )
     request_kwargs = {"metadata": {"session_id": "masked-params-session"}}
+    mock_router_instance.cache = DualCache()
     if route == "session":
-        mock_router_instance.cache = DualCache()
         await mock_router_instance.cache.async_set_cache(
             key=router._get_session_affinity_cache_key("masked-params-session", request_kwargs),
             value={"model": "opus", "tier": "REASONING"},
@@ -14816,7 +14891,9 @@ class TestContextWindowEscalation:
         router = ComplexityRouter(
             model_name="test-router",
             litellm_router_instance=_windowed_router(_SMALL, ("mid-model", "openai/gpt-4o-mini", 200000), _BIG),
-            complexity_router_config=_tier_config(tiers={"SIMPLE": ["small-model", "mid-model"], "COMPLEX": "big-model"}),
+            complexity_router_config=_tier_config(
+                tiers={"SIMPLE": ["small-model", "mid-model"], "COMPLEX": "big-model"}
+            ),
         )
 
         result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS)
@@ -14923,9 +15000,7 @@ class TestContextWindowEscalation:
             else MappingProxyType({})
         )
         raw_config: Final = RequestComplexityRouterConfig.model_validate(
-            MappingProxyType(
-                {"tiers": MappingProxyType({"SIMPLE": "small-model", "COMPLEX": "big-model"}), **setting}
-            )
+            MappingProxyType({"tiers": MappingProxyType({"SIMPLE": "small-model", "COMPLEX": "big-model"}), **setting})
         )
         config: Final = (
             RequestComplexityRouterConfig.model_validate_json(raw_config.model_dump_json())
@@ -15086,7 +15161,7 @@ class TestContextWindowEscalation:
             complexity_router_config=_tier_config(tiers={"SIMPLE": "cop-pool", "COMPLEX": "big-model"}),
         )
         real_get_llm_provider = litellm.get_llm_provider
-        copilot_resolutions: List = []
+        copilot_resolutions: list = []
 
         def _guarded(*args, **kwargs):
             target = str(kwargs.get("model") or (args[0] if args else "")) + str(
@@ -15171,9 +15246,7 @@ class TestModalityRouting:
             self.BASE_VISION,
         )
         original: Final = await router.aclassify("What color is this?")
-        result: Final = await router.async_pre_routing_hook(
-            model="m", request_kwargs={}, messages=self.IMAGE_MESSAGE
-        )
+        result: Final = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self.IMAGE_MESSAGE)
 
         assert original.heuristic_v2_forecast is not None
         assert result is not None and result.routing_decision is not None
@@ -17367,16 +17440,18 @@ class TestMaxTokensFromTierModel:
             tier_litellm_params={tier_carrier: 4321},
             max_tokens_from_tier_model="clamp",
             marker_params={marker_carrier: 2000},
-            simple_deployments=[{
-                "model_name": "small",
-                "litellm_params": {
-                    "model": "openai/gpt-5.6",
-                    "api_key": "test-only",
-                    "api_base": "https://tier.test/v1",
-                    surface_key: 9000,
-                },
-                "model_info": {"max_output_tokens": 16384},
-            }],
+            simple_deployments=[
+                {
+                    "model_name": "small",
+                    "litellm_params": {
+                        "model": "openai/gpt-5.6",
+                        "api_key": "test-only",
+                        "api_base": "https://tier.test/v1",
+                        surface_key: 9000,
+                    },
+                    "model_info": {"max_output_tokens": 16384},
+                }
+            ],
         )
 
         with respx.mock(assert_all_mocked=True) as upstream:
@@ -17426,10 +17501,14 @@ class TestMaxTokensFromTierModel:
             max_tokens_from_tier_model="clamp",
             marker_params={marker_carrier: 2000},
             tier_litellm_params={"max_tokens": tier_cap} if tier_cap else None,
-            simple_deployments=None if known_ceiling else [{
-                "model_name": "small",
-                "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"},
-            }],
+            simple_deployments=None
+            if known_ceiling
+            else [
+                {
+                    "model_name": "small",
+                    "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"},
+                }
+            ],
         )
         request: Final[dict[str, object]] = {} if caller_cap is None else {"max_output_tokens": caller_cap}
         await router.async_get_available_deployment(
@@ -17437,14 +17516,15 @@ class TestMaxTokensFromTierModel:
         )
         first: Final = {key: request[key] for key in OUTPUT_TOKEN_CEILING_PARAMS if key in request}
         await router.async_get_available_deployment(
-            model="smart-router", request_kwargs=request,
+            model="smart-router",
+            request_kwargs=request,
             messages=[{"role": "user", "content": self.COMPLEX_PROMPT}],
         )
 
         expected_first: Final = (
             {"max_tokens": tier_cap or caller_cap or 8192}
-            if known_ceiling or tier_cap else
-            {marker_carrier: 2000, **({"max_output_tokens": caller_cap} if caller_cap else {})}
+            if known_ceiling or tier_cap
+            else {marker_carrier: 2000, **({"max_output_tokens": caller_cap} if caller_cap else {})}
         )
         assert first == expected_first
         assert request["max_tokens"] == (caller_cap or 64000)
@@ -17464,10 +17544,12 @@ class TestMaxTokensFromTierModel:
         router: Final = self._router(
             max_tokens_from_tier_model=mode,
             marker_params={marker_carrier: 2000},
-            simple_deployments=[{
-                "model_name": "small",
-                "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"},
-            }],
+            simple_deployments=[
+                {
+                    "model_name": "small",
+                    "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"},
+                }
+            ],
         )
         sent: Final = await self._routed(router)
 
@@ -17696,7 +17778,11 @@ class TestMaxTokensFromTierModel:
             ({"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {}}}, {}),
             ({"litellm_metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {"max_tokens": "2000"}}}, {"max_tokens": 2000}),
             (
-                {"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: {"max_output_tokens": 5, "api_key": "x", "max_tokens": True}}},
+                {
+                    "metadata": {
+                        CLIENT_OUTPUT_CEILING_METADATA_KEY: {"max_output_tokens": 5, "api_key": "x", "max_tokens": True}
+                    }
+                },
                 {"max_output_tokens": 5},
             ),
             ({"metadata": {CLIENT_OUTPUT_CEILING_METADATA_KEY: "not-a-dict"}}, None),

tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py

@@ -11,7 +11,7 @@ from fakeredis.aioredis import FakeConnection
 from redis.asyncio import BlockingConnectionPool, Redis
 
 import litellm
-from litellm.caching.affinity_cache import claim_affinity_pin
+from litellm.caching.affinity_cache import claim_affinity_pin_in_memory
 from litellm.caching.dual_cache import DualCache
 from litellm.caching.in_memory_cache import InMemoryCache
 from litellm.caching.redis_cache import RedisCache
@@ -901,7 +901,6 @@ async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down():
     assert second == "our-deployment"
 
 
-@pytest.mark.asyncio
 @pytest.mark.parametrize(
     ("stored", "expected"),
     [
@@ -914,13 +913,13 @@ async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down():
         (None, {"model": "second"}),
     ],
 )
-async def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl(stored: object, expected: object) -> None:
+def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl(stored: object, expected: object) -> None:
     clock: Final = MagicMock(return_value=100.0)
     cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock))
     cache.in_memory_cache.set_cache("tier-pin", stored, ttl=10)
     clock.return_value = 105.0
 
-    winner: Final = await claim_affinity_pin(
+    winner: Final = claim_affinity_pin_in_memory(
         cache,
         "tier-pin",
         {"model": "second"},
@@ -936,24 +935,21 @@ async def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl(stored
     assert cache.in_memory_cache.get_cache("tier-pin") is None
 
 
-@pytest.mark.asyncio
-async def test_concurrent_eligible_claims_return_one_winner() -> None:
+def test_eligible_local_claims_keep_first_winner() -> None:
     cache: Final = DualCache()
     candidates: Final = ({"model": "first"}, {"model": "second"})
-    winners: Final = await asyncio.gather(
-        *(
-            claim_affinity_pin(
-                cache,
-                "tier-pin",
-                candidates[index % 2],
-                30,
-                eligible_values=candidates,
-            )
-            for index in range(20)
+    winners: Final = tuple(
+        claim_affinity_pin_in_memory(
+            cache,
+            "tier-pin",
+            candidates[index % 2],
+            30,
+            eligible_values=candidates,
         )
+        for index in range(20)
     )
 
-    assert winners == [{"model": "first"}] * 20
+    assert winners == ({"model": "first"},) * 20
     assert cache.in_memory_cache.get_cache("tier-pin") == {"model": "first"}