Merge pull request #41174 from BerriAI/litellm_tier_model_affinity

fix(router): preserve session model choice within each complexity tier

litellm/caching/affinity_cache.py

@@ -0,0 +1,125 @@
+"""Atomic affinity claims shared by deployment and tier-model selection."""
+
+import json
+from collections.abc import Mapping
+from typing import (
+    Final,
+    cast,  # noqa: TID251  # Redis script 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
+
+_PIN_JSON_ADAPTER: Final = TypeAdapter[JsonValue](JsonValue)
+
+_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 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."""
+    cache.in_memory_cache.delete_cache(cache_key)
+    cache.in_memory_cache.set_cache(cache_key, value, ttl=ttl_seconds)
+
+
+def _legacy_pin_matches(stored: object, pin_value: Mapping[str, str]) -> bool:
+    if isinstance(stored, dict):
+        return all(stored.get(key) is not None and str(stored[key]) == value for key, value in pin_value.items())
+    return isinstance(stored, str) and len(pin_value) == 1 and stored in pin_value.values()
+
+
+def claim_affinity_pin_in_memory(
+    cache: DualCache,
+    cache_key: str,
+    pin_value: Mapping[str, str],
+    ttl_seconds: int,
+    *,
+    eligible_values: tuple[Mapping[str, str], ...] | None = None,
+) -> object:
+    """No await between read and write, so same-loop claims agree during a Redis outage."""
+    existing: Final[object] = cache.in_memory_cache.get_cache(cache_key)
+    if existing is not None and eligible_values is None:
+        if _legacy_pin_matches(existing, pin_value):
+            set_local_affinity_pin(cache, cache_key, pin_value, ttl_seconds)
+        return existing
+    winner: Final = existing if existing is not None and existing in (eligible_values or ()) 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 claim_affinity_pin(
+    cache: DualCache,
+    cache_key: str,
+    pin_value: Mapping[str, str],
+    ttl_seconds: int,
+    *,
+    eligible_values: tuple[Mapping[str, str], ...] | None = None,
+) -> 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)),  # mutable-ok: JSON serialization requires dict, not a generic Mapping
+                int(ttl_seconds),
+                *(
+                    (json.dumps(tuple(dict(value) for value in eligible_values)),)  # mutable-ok: JSON requires dict
+                    if 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)

litellm/router_strategy/complexity_router/complexity_router.py

@@ -16,6 +16,8 @@ Inspired by ClawRouter: https://github.com/BlockRunAI/ClawRouter
 from __future__ import annotations
 
 import asyncio
+import hashlib
+import json
 import random
 import re
 import time
@@ -28,6 +30,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
 from pydantic import BaseModel, TypeAdapter, ValidationError, create_model
 
 from litellm._logging import verbose_router_logger
+from litellm.caching.affinity_cache import claim_affinity_pin
 from litellm.constants import (
     EMPTY_MAPPING,
     INTERNAL_CALL_ORIGIN_METADATA_KEY,
@@ -55,6 +58,7 @@ from litellm.router_strategy.complexity_router.tier_predictor import (
     TierSuccessPredictor,
     resolve_tier_artifact,
 )
+from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck
 from litellm.types.llms.openai import (
     AllMessageValues,
     ChatCompletionImageObject,
@@ -1119,10 +1123,10 @@ class _ContextWindowPlacement(NamedTuple):
 
 class _SessionAffinityPin(NamedTuple):
     model: str
-    tier: ComplexityTier | None
+    tier: ComplexityTier | str | None
 
 
-def _parse_session_affinity_pin(value: object) -> _SessionAffinityPin | None:
+def _parse_session_affinity_pin(value: object, active_tiers: tuple[str, ...]) -> _SessionAffinityPin | None:
     if isinstance(value, str):
         return _SessionAffinityPin(model=value, tier=None)
     parts: Final[tuple[object, object] | None] = (
@@ -1137,8 +1141,11 @@ def _parse_session_affinity_pin(value: object) -> _SessionAffinityPin | None:
     model, tier_value = parts
     if not isinstance(model, str):
         return None
-    tier: Final = ComplexityTier(tier_value) if isinstance(tier_value, str) else None
-    return _SessionAffinityPin(model=model, tier=tier)
+    if tier_value is None:
+        return _SessionAffinityPin(model=model, tier=None)
+    if not isinstance(tier_value, str) or tier_value not in active_tiers:
+        return None
+    return _SessionAffinityPin(model=model, tier=_built_in_tier_or_none(tier_value) or tier_value)
 
 
 def _session_affinity_cache_value(model: str, tier: ComplexityTier | str | None) -> Mapping[str, str | None]:
@@ -1195,6 +1202,10 @@ class ComplexityRouter(CustomLogger):
         if default_model:
             self.config.default_model = default_model
 
+        self._tier_affinity_config = hashlib.sha256(
+            self.config.model_dump_json(include=MappingProxyType({"tiers": True, "tier_model_configs": True})).encode()
+        ).hexdigest()
+
         # Checked here rather than on the config model because the deployment's
         # complexity_router_default_model arrives outside complexity_router_config and is
         # applied just above, so a validator on the model would reject a deployment that
@@ -2259,6 +2270,51 @@ class ComplexityRouter(CustomLogger):
     def _tier_pools(self) -> dict[str, list[str]]:
         return {tier: (models if isinstance(models, list) else [models]) for tier, models in self.config.tiers.items()}
 
+    async def _pin_model_for_tier(
+        self,
+        tier: ComplexityTier | str,
+        model: str,
+        candidates: tuple[str, ...],
+        request_kwargs: dict[str, object],  # mutable-ok: adaptive feedback metadata must follow the selected model
+        retained_pin: _SessionAffinityPin | None = None,
+    ) -> str:
+        if not self._uses_deployment_pin or model not in candidates:
+            return model
+        retained_model: Final = (
+            retained_pin.model
+            if retained_pin is not None
+            and retained_pin.tier is not None
+            and _tier_name(retained_pin.tier) == _tier_name(tier)
+            else None
+        )
+        if retained_model is not None and retained_model in candidates:
+            self._restamp_adaptive_choice(request_kwargs, model, retained_model)
+            return retained_model
+        session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs)
+        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),
+        )
+        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 _pick_model_for_tier(
         self,
         tier: ComplexityTier | str,
@@ -2266,11 +2322,18 @@ class ComplexityRouter(CustomLogger):
         resolved_messages: list[dict[str, Any]] | None,
         request_kwargs: dict,
         allowed_models: tuple[str, ...] | None = None,
+        retained_pin: _SessionAffinityPin | None = None,
     ) -> str:
         if not self.config.plugins:
-            if allowed_models is not None:
-                return self._pick_from_tier_value(allowed_models, _tier_name(tier))
-            return self.get_model_for_tier(tier)
+            candidates: Final = (
+                allowed_models if allowed_models is not None else tuple(self._tier_pools().get(_tier_name(tier), ()))
+            )
+            selected: Final = (
+                self._pick_from_tier_value(allowed_models, _tier_name(tier))
+                if allowed_models is not None
+                else self.get_model_for_tier(tier)
+            )
+            return await self._pin_model_for_tier(tier, selected, candidates, request_kwargs, retained_pin)
 
         from litellm.types.router import RoutingContext
 
@@ -2369,6 +2432,40 @@ class ComplexityRouter(CustomLogger):
         self._adaptive_chosen_model_key = ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY
         return self.adaptive_router
 
+    def _adaptive_candidate_models(
+        self,
+        classified_tier: ComplexityTier | str,
+        hard_floor: ComplexityTier | str | None = None,
+        hard_ceiling: ComplexityTier | str | None = None,
+        fit_filter: frozenset[str] | None = None,
+    ) -> tuple[str, ...]:
+        pools: Final = self._tier_pools()
+        candidates: Final = (
+            tuple(pools.get(_tier_name(classified_tier), ()))
+            if self.config.adaptive_eligible == "classified_tier"
+            else tuple(dict.fromkeys(chain.from_iterable(pools.values())))
+        )
+        floor: Final = self._active_tier_severity(hard_floor) if hard_floor is not None else None
+        ceiling: Final = self._active_tier_severity(hard_ceiling) if hard_ceiling is not None else None
+        return tuple(
+            model
+            for model in _allowed(candidates, fit_filter)
+            if (
+                floor is None
+                or any(
+                    self._active_tier_severity(tier) >= floor
+                    for tier in self._model_tiers.get(model, (classified_tier,))
+                )
+            )
+            and (
+                ceiling is None
+                or any(
+                    self._active_tier_severity(tier) <= ceiling
+                    for tier in self._model_tiers.get(model, (classified_tier,))
+                )
+            )
+        )
+
     def _soft_floor_pick(
         self,
         classified_tier: ComplexityTier | str,
@@ -2436,34 +2533,17 @@ class ComplexityRouter(CustomLogger):
                         ],
                     }
             return chosen_model
-        if self.config.adaptive_eligible == "classified_tier":
-            candidates = list(classified_candidates)
-            if not candidates:
-                return self._fitting_tier_fallback(classified_tier, fit_filter)
-        else:
-            candidates = list(_allowed(tuple(adaptive.config.available_models), fit_filter))
+        candidates: Final = self._adaptive_candidate_models(classified_tier, fit_filter=fit_filter)
 
         all_costs: Final = [adaptive.model_to_cost.get(m, 0.0) for m in candidates]
         quality_weight: Final = self.config.adaptive_weights.quality
         cost_weight: Final = self.config.adaptive_weights.cost
         penalty_weight: Final = self.config.tier_distance_penalty
 
-        floor_severity: Final = self._active_tier_severity(hard_floor) if hard_floor is not None else None
-        ceiling_severity: Final = self._active_tier_severity(hard_ceiling) if hard_ceiling is not None else None
         best_model: str | None = None
         best_score = float("-inf")
         candidate_scores: Final[list[dict[str, object]]] = []
-        for model in candidates:
-            if floor_severity is not None and all(
-                self._active_tier_severity(model_tier) < floor_severity
-                for model_tier in self._model_tiers.get(model, (classified_tier,))
-            ):
-                continue
-            if ceiling_severity is not None and all(
-                self._active_tier_severity(model_tier) > ceiling_severity
-                for model_tier in self._model_tiers.get(model, (classified_tier,))
-            ):
-                continue
+        for model in self._adaptive_candidate_models(classified_tier, hard_floor, hard_ceiling, fit_filter):
             cell = adaptive._cells[(request_type, model)]
             quality_sample = thompson_sample(cell)
             cost_score = normalized_cost(adaptive.model_to_cost.get(model, 0.0), all_costs)
@@ -2644,8 +2724,6 @@ class ComplexityRouter(CustomLogger):
         """Prompt content the resolved message list never carries: the Responses API's
         `instructions`, the /v1/messages top-level `system` block, and tool definitions.
         A coding agent's context is dominated by these."""
-        import json
-
         instructions: Final = request_kwargs.get("instructions")
         proxy_request: Final = request_kwargs.get("proxy_server_request")
         body: Final = proxy_request.get("body") if isinstance(proxy_request, Mapping) else None
@@ -2831,19 +2909,21 @@ class ComplexityRouter(CustomLogger):
         )
         return higher_tiers[0] if higher_tiers else tier
 
-    def _escalated_pin(self, pinned_model: str) -> str | None:
+    def _escalated_pin(self, pinned_model: str, tier: ComplexityTier | str | None = None) -> _SessionAffinityPin | None:
         """Bump a session's pinned model to the next-higher configured tier.
 
         Returns None when the pin no longer maps to any configured tier, signalling
         a full reclassification instead.
         """
-        pinned_tier: Final = self._tier_for_model(pinned_model)
+        pinned_tier: Final = tier if tier is not None else self._tier_for_model(pinned_model)
         if pinned_tier is None:
             return None
         escalated_tier: Final = self._escalate_tier(pinned_tier)
         if escalated_tier == pinned_tier:
-            return pinned_model
-        return self.get_model_for_tier(escalated_tier)
+            return _SessionAffinityPin(pinned_model, pinned_tier)
+        return _SessionAffinityPin(
+            self.get_model_for_tier(escalated_tier), _built_in_tier_or_none(_tier_name(escalated_tier))
+        )
 
     def _vision_verdicts(self, model_name: str) -> tuple[bool | None, ...]:
         """Declared vision support per deployment serving the name: True, False, or None when
@@ -2907,6 +2987,7 @@ class ComplexityRouter(CustomLogger):
         resolved_messages: Sequence[Mapping[str, object]] | None,
         request_kwargs: dict,  # mutable-ok: same shape the hook receives
         context_fit: _RequestContextFit | None = None,
+        retained_pin: _SessionAffinityPin | None = None,
     ) -> PreRoutingHookResponse:
         """Replace a routed model that cannot accept this request's image input.
 
@@ -2955,6 +3036,7 @@ class ComplexityRouter(CustomLogger):
                 repick_messages,  # pyright: ignore[reportArgumentType]  # hook-resolved message dicts; the pick only reads them
                 request_kwargs,
                 allowed_models=tuple(entry for entry in pools.get(capable, ()) if entry in eligible),
+                retained_pin=retained_pin,
             )
         elif self._modality_default_model_usable(request_kwargs, resolved_messages, eligible):
             new_tier = None
@@ -3098,6 +3180,7 @@ class ComplexityRouter(CustomLogger):
         resolved_messages: Sequence[Mapping[str, object]] | None,
         request_kwargs: dict,  # mutable-ok: same shape the hook receives
         context_fit: _RequestContextFit | None = None,
+        retained_pin: _SessionAffinityPin | None = None,
     ) -> PreRoutingHookResponse:
         """Try compatible tier recovery before the default, preserving request policy and fit."""
         decision: Final = response.routing_decision
@@ -3155,6 +3238,7 @@ class ComplexityRouter(CustomLogger):
                         repick_messages,  # pyright: ignore[reportArgumentType]  # hook-resolved message dicts; the pick only reads them
                         request_kwargs,
                         allowed_models=live,
+                        retained_pin=retained_pin,
                     )
                 except ValueError as exc:
                     verbose_router_logger.debug(
@@ -3247,8 +3331,13 @@ class ComplexityRouter(CustomLogger):
         """The adaptive feedback loop reads its chosen-model marker from request metadata; a
         gate rewrite must move the marker with the model or rewards land on the displaced one."""
         metadata: Final = request_kwargs.get("metadata")
-        if isinstance(metadata, dict) and metadata.get("adaptive_router_chosen_model") == old_model:
+        if not isinstance(metadata, dict):
+            return
+        if metadata.get("adaptive_router_chosen_model") == old_model:
             metadata["adaptive_router_chosen_model"] = new_model
+        decision: Final = metadata.get("adaptive_router_decision")
+        if isinstance(decision, dict) and decision.get("chosen_model") == old_model:
+            decision["chosen_model"] = new_model
 
     def _lexical_tier_override(self, user_message: str) -> KeywordOverride | None:
         """When keyword_tier_rules match literally, the most-severe matched tier wins.
@@ -3561,25 +3650,42 @@ 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_pin: Final = _parse_session_affinity_pin(pinned_value)
+            pinned_pin: Final = _parse_session_affinity_pin(pinned_value, self.config.tier_names())
             if pinned_pin is not None:
-                routed_model: str | None = pinned_pin.model
-                pin_escalation_keyword: str | None = None
-                if self.escalation_keywords:
-                    user_message: Final = (
-                        _newest_turn_ask(resolved_messages, marker_pairs) if resolved_messages else None
+                user_message: Final = _newest_turn_ask(resolved_messages, marker_pairs) if resolved_messages else None
+                pin_escalation_keyword: Final = (
+                    self._matched_escalation_keyword(user_message) if user_message is not None else None
+                )
+                selected_pin: Final = (
+                    self._escalated_pin(pinned_pin.model, pinned_pin.tier)
+                    if pin_escalation_keyword is not None
+                    else _SessionAffinityPin(
+                        pinned_pin.model,
+                        pinned_pin.tier if pinned_pin.tier is not None else self._tier_for_model(pinned_pin.model),
                     )
-                    if user_message is not None:
-                        pin_escalation_keyword = self._matched_escalation_keyword(user_message)
-                    if pin_escalation_keyword is not None:
-                        routed_model = self._escalated_pin(pinned_pin.model)
-                if routed_model is not None:
-                    escalated: Final = routed_model != pinned_pin.model
-                    resolved_pin_tier: Final = (
-                        pinned_pin.tier
-                        if not escalated and pinned_pin.tier is not None
-                        else self._tier_for_model(routed_model)
+                )
+                if selected_pin is not None:
+                    escalated: Final = selected_pin.model != pinned_pin.model or (
+                        pin_escalation_keyword is not None
+                        and pinned_pin.tier is not None
+                        and selected_pin.tier != pinned_pin.tier
                     )
+                    resolved_pin_tier: Final = selected_pin.tier
+                    session_model: Final = (
+                        await self._pin_model_for_tier(
+                            resolved_pin_tier,
+                            selected_pin.model,
+                            tuple(self._tier_pools().get(_tier_name(resolved_pin_tier), ())),
+                            request_kwargs,
+                        )
+                        if escalated and resolved_pin_tier is not None
+                        else selected_pin.model
+                    )
+                    retained_pin: Final = _SessionAffinityPin(session_model, resolved_pin_tier)
+                    if resolved_pin_tier is not None:
+                        await self._pin_model_for_tier(
+                            resolved_pin_tier, session_model, (session_model,), request_kwargs
+                        )
                     # The floor outranks the pin because plan mode is a transient state of the
                     # session, not a request to move it: the turns carrying the sentinel route at
                     # the floor, and the stored pin deliberately keeps the session's own model so
@@ -3590,16 +3696,28 @@ class ComplexityRouter(CustomLogger):
                     plan_floored: Final = (
                         pinned_tier is not None and self._apply_plan_mode_floor(pinned_tier) != pinned_tier
                     )
-                    session_model: Final = routed_model
-                    if plan_floored and pinned_tier is not None:
-                        routed_model = self.get_model_for_tier(self._apply_plan_mode_floor(pinned_tier))
-                    pin_source_tier: Final = self._tier_for_model(routed_model)
+                    floor_model: Final = (
+                        await self._pick_model_for_tier(
+                            self._apply_plan_mode_floor(pinned_tier),
+                            messages,
+                            resolved_messages,
+                            request_kwargs,
+                            retained_pin=retained_pin,
+                        )
+                        if plan_floored and pinned_tier is not None
+                        else session_model
+                    )
+                    pin_source_tier: Final = (
+                        self._apply_plan_mode_floor(pinned_tier)
+                        if plan_floored and pinned_tier is not None
+                        else resolved_pin_tier
+                    )
                     pin_placement: Final = (
                         await self._context_window_placement(
                             pin_source_tier,
                             resolved_messages,
                             request_kwargs,
-                            pool_override=(routed_model,),
+                            pool_override=(floor_model,),
                             context_fit=context_fit,
                         )
                         if pin_source_tier is not None
@@ -3612,11 +3730,18 @@ class ComplexityRouter(CustomLogger):
                         and _tier_name(pin_placement.tier) != _tier_name(pin_source_tier)
                         else None
                     )
-                    if pin_placement is not None and pin_context_original_tier is not None:
-                        # The stored pin below keeps the session's own model on purpose.
-                        routed_model = self._pick_from_tier_value(
-                            pin_placement.allowed_models, _tier_name(pin_placement.tier)
+                    routed_model: Final = (
+                        await self._pick_model_for_tier(
+                            pin_placement.tier,
+                            messages,
+                            resolved_messages,
+                            request_kwargs,
+                            allowed_models=pin_placement.allowed_models,
+                            retained_pin=retained_pin,
                         )
+                        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(
@@ -3644,7 +3769,7 @@ class ComplexityRouter(CustomLogger):
                     routed_pin_tier: Final = (
                         pin_placement.tier
                         if pin_placement is not None and pin_context_original_tier is not None
-                        else (self._tier_for_model(routed_model) if plan_floored else resolved_pin_tier)
+                        else pin_source_tier
                     )
                     session_tier_litellm_params: Final = self._litellm_params_for_model(routed_pin_tier, routed_model)
                     has_original_messages: Final = messages is not None and len(messages) > 0
@@ -3671,12 +3796,14 @@ class ComplexityRouter(CustomLogger):
                                 resolved_messages,
                                 request_kwargs,
                                 context_fit,
+                                retained_pin,
                             ),
                             messages,
                             input,
                             resolved_messages,
                             request_kwargs,
                             context_fit,
+                            retained_pin,
                         )
                     )
 
@@ -3961,13 +4088,21 @@ class ComplexityRouter(CustomLogger):
             housekeeping_ceiling: Final = tier if outcome.cause == "housekeeping" else None
             # A context-escalated tier becomes the hard floor: a floor the bandit can slide
             # under is not a floor.
-            routed_model = self._soft_floor_pick(
+            adaptive_floor: Final = tier if context_original_tier is not None else plan_floor
+            adaptive_fit: Final = context_placement.holdable_models if context_placement is not None else None
+            sampled_model: Final = self._soft_floor_pick(
                 tier,
                 ask,
                 request_kwargs,
-                hard_floor=tier if context_original_tier is not None else plan_floor,
+                hard_floor=adaptive_floor,
                 hard_ceiling=housekeeping_ceiling,
-                fit_filter=context_placement.holdable_models if context_placement is not None else None,
+                fit_filter=adaptive_fit,
+            )
+            routed_model = await self._pin_model_for_tier(  # rebind-ok: reuse the eligible tier winner
+                tier,
+                sampled_model,
+                self._adaptive_candidate_models(tier, adaptive_floor, housekeeping_ceiling, adaptive_fit),
+                request_kwargs,
             )
             adaptive: Final = self._ensure_adaptive_router()
             if adaptive is not None:

litellm/router_strategy/complexity_router/config.py

@@ -1256,20 +1256,16 @@ class ComplexityRouterConfig(BaseModel):
     deployment_affinity: bool = Field(
         default=True,
         description=(
-            "When True and a session_id is resolvable on the request, pin the deployment chosen "
-            "inside each routed model group and reuse it whenever the session returns to that "
-            "group, without pinning which group the session routes to. Independent of "
-            "session_affinity, which pins the model group instead (and always carries this "
-            "deployment pin with it): with session_affinity off, "
-            "every turn is still classified on its own merits while a session that escalates to a "
-            "stronger tier and comes back still lands on the deployment it used before, which is "
-            "what keeps a provider prompt cache warm. Pins are held per model group, so switching "
-            "tiers does not disturb the pin left behind in the previous group. On by default "
-            "because re-shuffling a conversation across deployments of the same model discards "
-            "that cache for no benefit; set False to keep every turn load-balanced across the "
-            "group, which is what a deployment set with tight per-deployment rate limits wants. "
-            "Inert when no session_id is resolvable, since there is nothing to key a pin on, and "
-            "suppressed when plugins are configured, for the same reason session_affinity is."
+            "When True and a client session_id is resolvable, reuse the session's chosen model "
+            "for each classified tier and its deployment within each model group. With "
+            "session_affinity off, every turn is still classified: moving to another tier leaves "
+            "the previous tier's model pin intact for a later return. Pins yield to current "
+            "candidate, context, modality, and availability constraints. Adaptive selection chooses "
+            "the initial model from its eligible pool, then reuses that choice per tier. This "
+            "reduces avoidable provider prompt-cache misses; it does not guarantee cache hits. "
+            "Set False to select models and load-balance deployments on every turn, unless "
+            "session_affinity or user_turn classification requires a pin. Inert without a client "
+            "session_id and suppressed when plugins are configured."
         ),
     )
     session_affinity_ttl_seconds: int = Field(
@@ -1277,7 +1273,7 @@ class ComplexityRouterConfig(BaseModel):
         gt=0,
         description=(
             "TTL for the session affinity pin; refreshed on every cache hit. Bounds both the "
-            "session_affinity model pin and the deployment_affinity deployment pin, so it measures "
+            "session_affinity model pin and the deployment_affinity per-tier model and deployment pins, so it measures "
             "idle time for the session's routing decisions rather than total session length"
         ),
     )

litellm/router_utils/pre_call_checks/deployment_affinity_check.py

@@ -13,13 +13,13 @@ where routing to a consistent deployment is still beneficial.
 """
 
 import hashlib
-import json
 from collections.abc import Mapping, Sequence
 from typing import Any, Final, cast
 
-from typing_extensions import TypedDict
+from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_router_logger
+from litellm.caching.affinity_cache import claim_affinity_pin, claim_affinity_pin_in_memory, set_local_affinity_pin
 from litellm.caching.dual_cache import DualCache
 from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
 from litellm.integrations.custom_logger import CustomLogger, Span
@@ -28,8 +28,8 @@ from litellm.types.llms.openai import AllMessageValues
 from litellm.types.utils import CallTypes
 
 
-class DeploymentAffinityCacheValue(TypedDict):
-    model_id: str
+class DeploymentAffinityCacheValue(TypedDict, closed=True):
+    model_id: ReadOnly[str]
 
 
 VALID_MODEL_GROUP_AFFINITY_FLAGS: Final = frozenset(
@@ -60,19 +60,6 @@ def warn_on_unknown_model_group_affinity_flags(model_group_affinity_config: Mapp
             )
 
 
-_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 current == ARGV[1] then
-  redis.call('EXPIRE', KEYS[1], ARGV[2])
-end
-return current
-"""
-
-
 class DeploymentAffinityCheck(CustomLogger):
     """
     Router deployment affinity callback.
@@ -255,34 +242,33 @@ class DeploymentAffinityCheck(CustomLogger):
         return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{hashed_user_key}:{session_id}"
 
     @staticmethod
-    def _get_session_id_from_metadata_dict(metadata: dict) -> str | None:
+    def _get_session_id_from_metadata_dict(metadata: Mapping[object, object]) -> str | None:
         session_id: Final = metadata.get("session_id")
         if session_id is None or metadata.get(SESSION_ID_GENERATED_METADATA_KEY):
             return None
         return str(session_id)
 
     @staticmethod
-    def _iter_metadata_dicts(request_kwargs: dict) -> list[dict]:
+    def _iter_metadata_dicts(request_kwargs: Mapping[str, object]) -> tuple[Mapping[object, object], ...]:
         """
         Return all metadata dicts available on the request.
 
         Depending on the endpoint, Router may populate `metadata` or `litellm_metadata`.
         Users may also send one or both, so we check both (rather than using `or`).
         """
-        metadata_dicts: Final[list[dict]] = []
-        for key in ("litellm_metadata", "metadata"):
-            md = request_kwargs.get(key)
-            if isinstance(md, dict):
-                metadata_dicts.append(md)
-        return metadata_dicts
+        return tuple(
+            cast(Mapping[object, object], metadata)  # cast-ok: isinstance proves mapping shape; values remain opaque
+            for key in ("litellm_metadata", "metadata")
+            if isinstance(metadata := request_kwargs.get(key), dict)
+        )
 
     @staticmethod
-    def _first_metadata_value(metadata_dicts: Sequence[dict], key: str) -> str | None:
+    def _first_metadata_value(metadata_dicts: Sequence[Mapping[object, object]], key: str) -> str | None:
         value: Final = next((metadata[key] for metadata in metadata_dicts if metadata.get(key) is not None), None)
         return None if value is None else str(value)
 
     @classmethod
-    def _get_user_key_from_request_kwargs(cls, request_kwargs: dict) -> str | None:
+    def get_user_key_from_request_kwargs(cls, request_kwargs: Mapping[str, object]) -> str | None:
         """
         Extract a stable affinity key from request kwargs.
 
@@ -334,74 +320,17 @@ class DeploymentAffinityCheck(CustomLogger):
         return None
 
     def _set_local_pin(self, cache_key: str, value: object, ttl_seconds: int) -> None:
-        """The one owner of authoritative local pin writes: a plain set keeps a live
-        key's original expiry (`allow_ttl_override`), so the entry is replaced to make
-        the TTL real. Every local pin write goes through here so the redis-winner sync
-        and the pod-local claim can never disagree about expiry again."""
-        self.cache.in_memory_cache.delete_cache(cache_key)
-        self.cache.in_memory_cache.set_cache(cache_key, value, ttl=ttl_seconds)
+        set_local_affinity_pin(self.cache, cache_key, value, ttl_seconds)
 
     async def _claim_pin(self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int) -> str | None:
-        """First-writer-wins pin write: store `pin_value` only when the key is absent and
-        return the deployment id the key holds afterwards, so a caller learns whether it won
-        by comparing against its own id, and None when the stored value is one no reader can
-        interpret. Concurrent claimers converge on the
-        first write instead of the last. Re-claiming with the stored value refreshes its
-        TTL, the same keepalive the complexity router's model pin documents: an active
-        session must not lose its pin mid-conversation just because it outlives the
-        original write, so the affinity TTL (the Router's
-        `deployment_affinity_ttl_seconds`, or a pre-routing hook's per-request
-        `session_affinity_ttl_seconds` override) bounds idle time, not total
-        session length. On Redis one Lua script does the get-or-set-or-refresh
-        atomically (same registration seam the rate limiters use) and the in-memory
-        tier is synchronized to the winner; without Redis, and whenever Redis is
-        unreachable, the pod-local check-and-set below stands in and is atomic because it
-        runs synchronously on the event loop. Degrading to a pod-local claim rather than
-        propagating the fault is what keeps same-pod stickiness through a Redis blip: the
-        caller only logs this result, so an escaping error would leave the session with no
-        pin at all and reshuffle every turn for the outage, which is worse than losing
-        cross-pod agreement. The redis tier is
-        resolved per call because the proxy attaches it after Router construction
-        (`Router._update_redis_cache`); the compiled script is cached per event loop
-        underneath the registration seam.
-        """
-        redis_cache: Final = self.cache.redis_cache
-        if redis_cache is not None:
-            try:
-                claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT)
-                raw: Final = await claim_script(keys=(cache_key,), args=(json.dumps(pin_value), int(ttl_seconds)))
-                decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw
-                if not isinstance(decoded, str):
-                    return pin_value["model_id"]
-                try:
-                    winner: object = json.loads(decoded)
-                except json.JSONDecodeError:
-                    winner = decoded
-                self._set_local_pin(cache_key=cache_key, value=winner, ttl_seconds=ttl_seconds)
-                return self._pinned_model_id(winner)
-            except Exception as e:  # noqa: BLE001  # any Redis/Lua failure degrades to the pod-local claim, never unpins
-                verbose_router_logger.debug(
-                    "DeploymentAffinityCheck: redis pin claim failed, falling back to pod-local claim. error=%s", e
-                )
-
-        return self._claim_pin_in_memory(cache_key=cache_key, pin_value=pin_value, ttl_seconds=ttl_seconds)
+        winner: Final = await claim_affinity_pin(self.cache, cache_key, pin_value, ttl_seconds)
+        return self._pinned_model_id(winner)
 
     def _claim_pin_in_memory(
         self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int
     ) -> str | None:
-        """Pod-local half of the claim, used when no Redis tier is attached and as the
-        fallback when the Redis claim fails. Mirrors the Lua script exactly, including
-        the keepalive: re-claiming with the stored value slides the idle window through
-        `_set_local_pin`. Both branches stay synchronous, hence atomic on the event
-        loop."""
-        existing: Final = self.cache.in_memory_cache.get_cache(cache_key)
-        if existing is not None:
-            existing_model_id: Final = self._pinned_model_id(existing)
-            if existing_model_id == pin_value["model_id"]:
-                self._set_local_pin(cache_key=cache_key, value=pin_value, ttl_seconds=ttl_seconds)
-            return existing_model_id
-        self._set_local_pin(cache_key=cache_key, value=pin_value, ttl_seconds=ttl_seconds)
-        return pin_value["model_id"]
+        winner: Final = claim_affinity_pin_in_memory(self.cache, cache_key, pin_value, ttl_seconds)
+        return self._pinned_model_id(winner)
 
     @staticmethod
     def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None:
@@ -465,7 +394,7 @@ class DeploymentAffinityCheck(CustomLogger):
             enable_session_id or self._get_marker_session_affinity_ttl(request_kwargs=request_kwargs) is not None
         )
         user_key: Final = (
-            self._get_user_key_from_request_kwargs(request_kwargs=request_kwargs)
+            self.get_user_key_from_request_kwargs(request_kwargs=request_kwargs)
             if (session_affinity_active or enable_user_key)
             else None
         )
@@ -580,7 +509,7 @@ class DeploymentAffinityCheck(CustomLogger):
             return None
 
         user_key: Final = (
-            self._get_user_key_from_request_kwargs(request_kwargs=kwargs)
+            self.get_user_key_from_request_kwargs(request_kwargs=kwargs)
             if (enable_user_key or session_affinity_active)
             else None
         )

tests/test_litellm/router_strategy/test_complexity_router.py

@@ -9,7 +9,7 @@ import json
 import logging
 import sys
 import time
-from collections.abc import AsyncIterator, Mapping
+from collections.abc import AsyncIterator, Mapping, Sequence
 from copy import deepcopy
 from functools import partial
 from typing import Dict, Final, List, Literal
@@ -31,6 +31,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.constants import (
     OUTPUT_TOKEN_CEILING_PARAMS,
     RETURN_RAW_MODEL_NAME_METADATA_KEY,
@@ -70,6 +71,7 @@ from litellm.router_strategy.complexity_router.tier_predictor import (
 from litellm.types.router import (
     Deployment,
     LiteLLM_Params,
+    PreRoutingHookResponse,
     RouterErrors,
     TaggedPreRoutingStrategy,
 )
@@ -5574,6 +5576,387 @@ class TestRoutingDecisionCauseLogging:
         assert "cause=semantic_keyword_match" not in router_log_capture.text
 
 
+class TestTierModelAffinity:
+    @staticmethod
+    async def _route(
+        router: ComplexityRouter,
+        metadata: Mapping[str, object],
+        proposed_model: str,
+        prompt: str = "compact",
+        messages: list[dict[str, object]] | None = None,
+    ) -> PreRoutingHookResponse:
+        def choose(candidates: Sequence[str]) -> str:
+            return proposed_model if proposed_model in candidates else candidates[0]
+
+        request_metadata: Final = dict(metadata)
+        with patch(  # test-quality-ok: [TQ008] alternate proposals make affinity reuse deterministic
+            "litellm.router_strategy.complexity_router.complexity_router.random.choice",
+            side_effect=choose,
+        ):
+            result: Final = await router.async_pre_routing_hook(
+                model="affinity-router",
+                request_kwargs={"metadata": request_metadata},
+                messages=messages if messages is not None else [{"role": "user", "content": prompt}],
+            )
+        assert result is not None
+        if router.config.adaptive:
+            assert request_metadata["adaptive_router_chosen_model"] == result.model
+        return result
+
+    @staticmethod
+    def _router(
+        mock_router_instance: MagicMock,
+        adaptive: bool = False,
+        deployment_affinity: bool = True,
+        plugins: bool = False,
+    ) -> ComplexityRouter:
+        mock_router_instance.cache = DualCache()
+        mock_router_instance.model_list = []
+        mock_router_instance.model_name_to_deployment_indices = {}
+        return ComplexityRouter(
+            model_name="affinity-router",
+            litellm_router_instance=mock_router_instance,
+            complexity_router_config={
+                "tiers": {
+                    tier: [
+                        {"model_name": model, "litellm_params": {"temperature": temperature}}
+                        for model in ("model-a", "model-b")
+                    ]
+                    for tier, temperature in (("SIMPLE", 0.1), ("REASONING", 0.9))
+                },
+                "adaptive": adaptive,
+                "deployment_affinity": deployment_affinity,
+                "session_affinity": False,
+                **({"plugins": [_DummyPlugin()]} if plugins else {}),
+            },
+        )
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("adaptive", [False, True])
+    async def test_reuses_model_per_tier_without_pinning_classification(
+        self, mock_router_instance: MagicMock, adaptive: bool
+    ) -> None:
+        router: Final = self._router(mock_router_instance, adaptive=adaptive)
+        metadata: Final = {"session_id": "same-session"}
+        first: Final = await self._route(router, metadata, "model-a")
+        if adaptive:
+            from litellm.router_strategy.adaptive_router.bandit import BanditCell
+            from litellm.router_strategy.adaptive_router.classifier import classify_prompt
+
+            bandit: Final = router._ensure_adaptive_router()
+            assert bandit is not None
+            bandit._cells[(classify_prompt("compact"), "model-a")] = BanditCell(alpha=5.0, beta=5.0)
+        repeated: Final = await self._route(router, metadata, "model-b")
+        reasoning: Final = await self._route(
+            router, metadata, "model-b", "Let's think step by step and reason through this problem carefully."
+        )
+        returned: Final = await self._route(router, metadata, "model-b")
+
+        assert (first.model, repeated.model, reasoning.model, returned.model) == (
+            "model-a", "model-a", "model-b", "model-a"
+        )
+        assert tuple(result.routing_decision["tier"] for result in (first, repeated, reasoning, returned)) == (
+            "SIMPLE", "SIMPLE", "REASONING", "SIMPLE"
+        )
+        assert returned.litellm_params == {"temperature": 0.1}
+        assert reasoning.litellm_params == {"temperature": 0.9}
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("identity_key", ["user_api_key_hash", "user_api_key_user_id"])
+    async def test_isolates_sessions_and_authenticated_callers(
+        self, mock_router_instance: MagicMock, identity_key: str
+    ) -> None:
+        router: Final = self._router(mock_router_instance)
+        first_caller: Final = {"session_id": "shared", identity_key: "caller-a"}
+        other_caller: Final = {"session_id": "shared", identity_key: "caller-b"}
+        other_session: Final = {"session_id": "separate", identity_key: "caller-a"}
+
+        assert (await self._route(router, first_caller, "model-a")).model == "model-a"
+        assert (await self._route(router, other_caller, "model-b")).model == "model-b"
+        assert (await self._route(router, other_session, "model-b")).model == "model-b"
+        assert (await self._route(router, first_caller, "model-b")).model == "model-a"
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "metadata,deployment_affinity,plugins",
+        [
+            ({}, True, False),
+            ({"session_id": "generated", SESSION_ID_GENERATED_METADATA_KEY: True}, True, False),
+            ({"session_id": "provided"}, False, False),
+            ({"session_id": "provided"}, True, True),
+        ],
+        ids=["absent-session", "generated-session", "disabled", "plugin-policy"],
+    )
+    async def test_does_not_pin_without_eligible_session(
+        self,
+        mock_router_instance: MagicMock,
+        metadata: Mapping[str, object],
+        deployment_affinity: bool,
+        plugins: bool,
+    ) -> None:
+        router: Final = self._router(
+            mock_router_instance, deployment_affinity=deployment_affinity, plugins=plugins
+        )
+        assert (await self._route(router, metadata, "model-a")).model == "model-a"
+        assert (await self._route(router, metadata, "model-b")).model == "model-b"
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("adaptive", [False, True])
+    async def test_replaces_pin_outside_the_context_candidate_domain(self, adaptive: bool) -> None:
+        router: Final = ComplexityRouter(
+            model_name="affinity-router",
+            litellm_router_instance=_windowed_router(_SMALL, _BIG),
+            complexity_router_config={
+                "tiers": {"SIMPLE": ["small-model", "big-model"]},
+                "adaptive": adaptive,
+                "deployment_affinity": True,
+                "session_affinity": False,
+            },
+        )
+        metadata: Final = {"session_id": "growing-context"}
+        assert (await self._route(router, metadata, "small-model")).model == "small-model"
+        oversized: Final = await router.async_pre_routing_hook(
+            model="affinity-router",
+            request_kwargs={"metadata": dict(metadata)},
+            messages=_OVERSIZED_TURNS,
+        )
+        assert oversized is not None
+        assert oversized.model == "big-model"
+        assert oversized.routing_decision["tier"] == "SIMPLE"
+        assert (await self._route(router, metadata, "small-model")).model == "big-model"
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("session_affinity", [False, True], ids=["user-turn", "session-affinity"])
+    @pytest.mark.parametrize("gate", ["image", "health"])
+    async def test_temporary_replay_gate_keeps_the_held_tiers_model_preference(
+        self, mock_router_instance: MagicMock, session_affinity: bool, gate: Literal["image", "health"]
+    ) -> None:
+        async def get_healthy_deployments(
+            model: str,
+            request_kwargs: Mapping[str, object],
+            messages: Sequence[Mapping[str, object]] | None = None,
+            input: object = None,
+            parent_otel_span: object = None,
+            health_check_probe: bool = False,
+        ) -> list[dict[str, object]]:
+            unavailable: Final = (
+                gate == "health"
+                and model == "model-a"
+                and messages is not None
+                and bool(messages)
+                and messages[-1].get("role") == "tool"
+            )
+            return [] if unavailable else [{"model_name": model, "model_info": {"id": f"deployment-{model}"}}]
+
+        cache: Final = DualCache()
+        mock_router_instance.cache = cache
+        mock_router_instance.async_get_healthy_deployments = get_healthy_deployments
+        router: Final = TestModalityRouting._router(
+            mock_router_instance,
+            {
+                "tiers": {"SIMPLE": ["model-a", "model-b"]},
+                "deployment_affinity": True,
+                "session_affinity": session_affinity,
+                "classification_mode": "every_request" if session_affinity else "user_turn",
+                "modality_routing": True,
+                "modality_pin_override": True,
+            },
+            {"model-a": False, "model-b": True},
+        )
+        metadata: Final = {"session_id": "replay-session"}
+        continuation: Final[list[dict[str, object]]] = [
+            {"role": "user", "content": "compact"},
+            {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [
+                    {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}
+                ],
+            },
+            {"role": "tool", "tool_call_id": "call_1", "content": [IMG_PART] if gate == "image" else "done"},
+        ]
+        assert (await self._route(router, metadata, "model-a")).model == "model-a"
+
+        replayed: Final = await self._route(router, metadata, "model-b", messages=continuation)
+        assert replayed.model == "model-b"
+        assert replayed.routing_decision["tier"] == "SIMPLE"
+        assert replayed.routing_decision["cause"] == (
+            "health_failover"
+            if gate == "health"
+            else ("modality_pin_override" if session_affinity else "user_turn_continuation")
+        )
+        cache_key: Final = router._get_session_affinity_cache_key("replay-session", {"metadata": metadata})
+        assert await cache.async_get_cache(cache_key) == {"model": "model-a", "tier": "SIMPLE"}
+
+        next_ask: Final = await self._route(router, metadata, "model-b")
+        assert next_ask.model == "model-a"
+        assert next_ask.routing_decision["tier"] == "SIMPLE"
+        assert next_ask.routing_decision["cause"] == (
+            "session_affinity_pin" if session_affinity else "heuristic_scorer"
+        )
+
+    @pytest.mark.asyncio
+    async def test_user_turn_replay_refreshes_the_model_used_within_its_tier(
+        self, mock_router_instance: MagicMock
+    ) -> None:
+        clock: Final = MagicMock(return_value=100.0)
+        mock_router_instance.cache = DualCache(in_memory_cache=InMemoryCache(clock=clock))
+        router: Final = ComplexityRouter(
+            model_name="affinity-router",
+            litellm_router_instance=mock_router_instance,
+            complexity_router_config={
+                "tiers": {"SIMPLE": ["model-a", "model-b"]},
+                "classification_mode": "user_turn",
+                "session_affinity_ttl_seconds": 10,
+            },
+        )
+        metadata: Final = {"session_id": "same-session"}
+        continuation: Final[list[dict[str, object]]] = [
+            {"role": "user", "content": "compact"},
+            {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [
+                    {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}
+                ],
+            },
+            {"role": "tool", "tool_call_id": "call_1", "content": "done"},
+        ]
+        assert (await self._route(router, metadata, "model-a")).model == "model-a"
+        clock.return_value = 105.0
+        replayed: Final = await self._route(router, metadata, "model-b", messages=continuation)
+        assert replayed.model == "model-a"
+        assert replayed.routing_decision["cause"] == "user_turn_continuation"
+
+        clock.return_value = 111.0
+        next_ask: Final = await self._route(router, metadata, "model-b")
+        assert next_ask.model == "model-a"
+        assert next_ask.routing_decision["tier"] == "SIMPLE"
+        assert next_ask.routing_decision["cause"] == "heuristic_scorer"
+
+    @pytest.mark.asyncio
+    async def test_session_escalation_keeps_the_selected_tier_when_models_overlap(
+        self, mock_router_instance: MagicMock
+    ) -> None:
+        cache: Final = DualCache()
+        mock_router_instance.cache = cache
+        router: Final = ComplexityRouter(
+            model_name="affinity-router",
+            litellm_router_instance=mock_router_instance,
+            complexity_router_config={
+                "tiers": {
+                    "SIMPLE": "base",
+                    **{
+                        tier: [
+                            {"model_name": model, "litellm_params": {"temperature": temperature}}
+                            for model in models
+                        ]
+                        for tier, models, temperature in (
+                            ("MEDIUM", ("shared", "middle"), 0.4),
+                            ("COMPLEX", ("shared", "higher"), 0.8),
+                        )
+                    },
+                },
+                "session_affinity": True,
+                "keyword_tier_rules": [{"keywords": ["visit_complex"], "tier": "COMPLEX"}],
+            },
+        )
+        metadata: Final = {"session_id": "same-session"}
+        assert (await self._route(router, metadata, "higher", "visit_complex")).model == "higher"
+        cache_key: Final = router._get_session_affinity_cache_key("same-session", {"metadata": metadata})
+        await cache.async_set_cache(cache_key, {"model": "base", "tier": "SIMPLE"}, ttl=600)
+
+        result: Final = await self._route(router, metadata, "shared", "LITELLM ESCALATE")
+        assert result.model == "shared"
+        assert result.routing_decision["tier"] == "MEDIUM"
+        assert result.routing_decision["cause"] == "session_affinity_escalation"
+        assert result.litellm_params == {"temperature": 0.4}
+        assert await cache.async_get_cache(cache_key) == {"model": "shared", "tier": "MEDIUM"}
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "stale_tier",
+        ["NON_REASONING", "REMOVED_TIER", 7, []],
+        ids=["inactive-tier", "unknown-tier", "integer-tier", "list-tier"],
+    )
+    @pytest.mark.parametrize(
+        "prompt,expected_model,expected_tier",
+        [("compact", "model-a", "SIMPLE"), ("LITELLM ESCALATE", "model-b", "MEDIUM")],
+        ids=["ordinary-replay", "escalation"],
+    )
+    async def test_reclassifies_session_pin_outside_the_active_tier_ladder(
+        self,
+        mock_router_instance: MagicMock,
+        stale_tier: object,
+        prompt: str,
+        expected_model: str,
+        expected_tier: str,
+    ) -> None:
+        cache: Final = DualCache()
+        mock_router_instance.cache = cache
+        router: Final = ComplexityRouter(
+            model_name="affinity-router",
+            litellm_router_instance=mock_router_instance,
+            complexity_router_config={
+                "tiers": {"SIMPLE": "model-a", "MEDIUM": "model-b"},
+                "session_affinity": True,
+            },
+        )
+        metadata: Final = {"session_id": "same-session"}
+        cache_key: Final = router._get_session_affinity_cache_key("same-session", {"metadata": metadata})
+        await cache.async_set_cache(cache_key, {"model": "model-a", "tier": stale_tier}, ttl=600)
+
+        result: Final = await self._route(router, metadata, expected_model, prompt)
+
+        assert result.model == expected_model
+        assert result.routing_decision["tier"] == expected_tier
+        assert result.routing_decision["cause"] == "heuristic_scorer"
+        assert await cache.async_get_cache(cache_key) == {"model": expected_model, "tier": expected_tier}
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("classification_mode", ["every_request", "user_turn"])
+    async def test_custom_tier_keeps_its_own_model(
+        self, mock_router_instance: MagicMock, classification_mode: Literal["every_request", "user_turn"]
+    ) -> None:
+        mock_router_instance.cache = DualCache()
+        router: Final = ComplexityRouter(
+            model_name="affinity-router",
+            litellm_router_instance=mock_router_instance,
+            complexity_router_config=_custom_tier_config(
+                tiers={"SIMPLE": ["model-a", "model-b"], "SECURITY_REVIEW": ["model-a", "model-b"], "COMPLEX": "model-a"},
+                deployment_affinity=True,
+                classification_mode=classification_mode,
+                keyword_tier_rules=[
+                    {"keywords": ["compact"], "tier": "SIMPLE"},
+                    {"keywords": ["audit"], "tier": "SECURITY_REVIEW"},
+                ],
+            ),
+        )
+        metadata: Final = {"session_id": "custom-session"}
+        assert (await self._route(router, metadata, "model-a")).model == "model-a"
+        assert (await self._route(router, metadata, "model-b", "audit")).model == "model-b"
+        assert (await self._route(router, metadata, "model-b")).model == "model-a"
+        retained: Final = await self._route(router, metadata, "model-a", "audit")
+        assert retained.model == "model-b"
+        assert retained.routing_decision["tier"] == "SECURITY_REVIEW"
+        if classification_mode == "user_turn":
+            continuation: Final[list[dict[str, object]]] = [
+                {"role": "user", "content": "audit"},
+                {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": [
+                        {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}
+                    ],
+                },
+                {"role": "tool", "tool_call_id": "call_1", "content": "done"},
+            ]
+            replayed: Final = await self._route(router, metadata, "model-a", messages=continuation)
+            assert replayed.model == "model-b"
+            assert replayed.routing_decision["tier"] == "SECURITY_REVIEW"
+            assert replayed.routing_decision["cause"] == "user_turn_continuation"
+
+
 class TestSessionAffinity:
     """Test the session_affinity sticky-routing behavior (off by default)."""
 
@@ -5638,11 +6021,8 @@ class TestSessionAffinity:
         tier_pinned,
         deployment_pinned,
     ):
-        """deployment_affinity pins the deployment inside each routed group without pinning which
-        group the session routes to, so with session_affinity off the tier must still reclassify
-        on every turn while the marker the Router stamps is still emitted. Turn 1 classifies
-        REASONING and turn 2 SIMPLE, so a reclassified turn 2 moves model while a tier-pinned one
-        does not. plugins suppress both pins, since a stale pin would bypass the plugin pipeline."""
+        """Deployment affinity retains a model per tier while classification continues.
+        Session affinity keeps the first tier too; plugins suppress both affinity policies."""
         mock_router_instance.cache = DualCache()
         router = ComplexityRouter(
             model_name="test-router",
@@ -5692,8 +6072,7 @@ class TestSessionAffinity:
 
     @pytest.mark.asyncio
     async def test_disabled_by_default_reclassifies_every_turn(self, mock_router_instance, basic_config):
-        """Regression: session_affinity defaults to False, so a shared session_id must NOT
-        pin the first turn's model; every turn is classified on its own merits."""
+        """With session_affinity off, a shared session can move from REASONING to SIMPLE."""
         assert "session_affinity" not in basic_config
         mock_router_instance.cache = DualCache()
         router = ComplexityRouter(
@@ -5848,7 +6227,7 @@ class TestSessionAffinity:
 
     @pytest.mark.asyncio
     async def test_respects_ttl_seconds(self, mock_router_instance, basic_config):
-        cache = AsyncMock()
+        cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None)
         cache.async_get_cache = AsyncMock(return_value=None)
         mock_router_instance.cache = cache
         router = ComplexityRouter(
@@ -5872,7 +6251,7 @@ class TestSessionAffinity:
     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 = AsyncMock()
+        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(
@@ -7112,7 +7491,8 @@ class TestEscalationKeywords:
             complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": ["o1-a", "o1-b", "o1-c"]}},
         )
         for pinned in ("o1-a", "o1-b", "o1-c"):
-            assert router._escalated_pin(pinned) == pinned
+            escalated: Final = router._escalated_pin(pinned)
+            assert (escalated.model, escalated.tier) == (pinned, "REASONING")
 
     @pytest.mark.asyncio
     async def test_session_escalation_at_ceiling_keeps_multi_model_pin(self, mock_router_instance):
@@ -12009,7 +12389,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 = AsyncMock()
+    cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None)
     cache.async_get_cache = AsyncMock(return_value=["shared", "SIMPLE"])
     mock_router_instance.cache = cache
     router = ComplexityRouter(
@@ -12988,7 +13368,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 = AsyncMock()
+            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"})
             mock_router_instance.cache = cache
             config["session_affinity"] = True
@@ -13178,7 +13558,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 = AsyncMock()
+        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"})
         mock_router_instance.cache = cache
         router = self._router(
@@ -13211,7 +13591,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 = AsyncMock()
+        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"})
         mock_router_instance.cache = cache
         router = self._router(
@@ -14322,18 +14702,37 @@ class TestTierHealthFailover:
             cooling=("id-a1",),
             raises_for={"exhausted-b": raised},
         )
-        key = router._get_session_affinity_cache_key("sess-exhausted", {})
-        await router.litellm_router_instance.cache.async_set_cache(
-            key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
-        )
-        results = [
-            await router.async_pre_routing_hook(
-                model="m", request_kwargs={"metadata": {"session_id": "sess-exhausted"}}, messages=self.SIMPLE_MESSAGE
+        sessions: Final = tuple(f"sess-exhausted-{sample}" for sample in range(20))
+        await asyncio.gather(
+            *(
+                router.litellm_router_instance.cache.async_set_cache(
+                    key=router._get_session_affinity_cache_key(session_id, {}),
+                    value={"model": "dead-a", "tier": "SIMPLE"},
+                    ttl=600,
+                )
+                for session_id in sessions
             )
-            for _ in range(20)
+        )
+        results: Final = [
+            await router.async_pre_routing_hook(
+                model="m", request_kwargs={"metadata": {"session_id": session_id}}, messages=self.SIMPLE_MESSAGE
+            )
+            for session_id in sessions
         ]
         assert {r.model for r in results} == expected
 
+        def choose_other(candidates: Sequence[str]) -> str:
+            return next((model for model in candidates if model != results[0].model), candidates[0])
+
+        with patch(  # test-quality-ok: [TQ008] an alternate healthy proposal proves retained affinity across failover
+            "litellm.router_strategy.complexity_router.complexity_router.random.choice",
+            side_effect=choose_other,
+        ):
+            retained: Final = await router.async_pre_routing_hook(
+                model="m", request_kwargs={"metadata": {"session_id": sessions[0]}}, messages=self.SIMPLE_MESSAGE
+            )
+        assert retained.model == results[0].model
+
     @pytest.mark.asyncio
     async def test_a_group_the_router_has_no_deployment_for_is_not_a_failover_target(self, mock_router_instance):
         """The owner answers an unconfigured group with BadRequestError. Reading that as live

tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py

@@ -1,3 +1,5 @@
+import asyncio
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -6,7 +8,9 @@ import pytest
 import json
 
 import litellm
+from litellm.caching.affinity_cache import claim_affinity_pin
 from litellm.caching.dual_cache import DualCache
+from litellm.caching.in_memory_cache import InMemoryCache
 from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
 from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
     DeploymentAffinityCheck,
@@ -558,6 +562,124 @@ 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"),
+    [
+        ({"model": "first"}, {"model": "first"}),
+        ('{ "model" : "first" }', {"model": "first"}),
+        ({"model": "removed"}, {"model": "second"}),
+        ({"model": "first", "extra": "stale"}, {"model": "second"}),
+        ({"model_id": "first"}, {"model": "second"}),
+        ("first", {"model": "second"}),
+        (None, {"model": "second"}),
+    ],
+)
+async 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(
+        cache, "tier-pin", {"model": "second"}, 30,
+        eligible_values=({"model": "first"}, {"model": "second"}),
+    )
+
+    assert winner == expected
+    assert cache.in_memory_cache.ttl_dict["tier-pin"] == 135.0
+    clock.return_value = 111.0
+    assert cache.in_memory_cache.get_cache("tier-pin") == expected
+    clock.return_value = 136.0
+    assert cache.in_memory_cache.get_cache("tier-pin") is None
+
+
+@pytest.mark.asyncio
+async def test_concurrent_eligible_claims_return_one_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)
+    ))
+
+    assert winners == [{"model": "first"}] * 20
+    assert cache.in_memory_cache.get_cache("tier-pin") == {"model": "first"}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("stored", "expected", "refresh"),
+    [
+        ({"model_id": 7}, "7", True),
+        ({"model_id": "other"}, "other", False),
+        ({"model": "7"}, None, False),
+        (["7"], None, False),
+    ],
+)
+async def test_legacy_deployment_claim_retains_decoder_and_keepalive(
+    stored: object, expected: str | None, refresh: bool
+) -> None:
+    clock: Final = MagicMock(return_value=100.0)
+    cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock))
+    callback: Final = DeploymentAffinityCheck(
+        cache=cache, ttl_seconds=30,
+        enable_user_key_affinity=False, enable_responses_api_affinity=False,
+    )
+    cache.in_memory_cache.set_cache("deployment-pin", stored, ttl=10)
+    clock.return_value = 105.0
+
+    winner: Final = await callback._claim_pin(
+        "deployment-pin", {"model_id": "7"}, 30
+    )
+
+    assert winner == expected
+    assert cache.in_memory_cache.ttl_dict["deployment-pin"] == (
+        135.0 if refresh else 110.0
+    )
+    assert cache.in_memory_cache.get_cache("deployment-pin") == (
+        {"model_id": "7"} if refresh else stored
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("raw", "expected", "stored"),
+    [
+        (b'{"model_id": "winner"}', "winner", {"model_id": "winner"}),
+        ('"winner"', "winner", "winner"),
+        ("winner", "winner", "winner"),
+        (b"winner", "winner", "winner"),
+        ('{"model": "winner"}', None, {"model": "winner"}),
+        (None, "candidate", None),
+        (123, "candidate", None),
+        ({"model_id": "winner"}, "candidate", None),
+    ],
+)
+async def test_redis_deployment_claim_preserves_legacy_result_decoding(
+    raw: object, expected: str | None, stored: object
+) -> None:
+    redis: Final = MagicMock()
+    redis.async_register_script.return_value = AsyncMock(return_value=raw)
+    cache: Final = DualCache(redis_cache=redis)
+    callback: Final = DeploymentAffinityCheck(
+        cache=cache, ttl_seconds=30,
+        enable_user_key_affinity=False, enable_responses_api_affinity=False,
+    )
+
+    winner: Final = await callback._claim_pin(
+        "deployment-pin", {"model_id": "candidate"}, 30
+    )
+
+    assert winner == expected
+    assert cache.in_memory_cache.get_cache("deployment-pin") == stored
+
+
 @pytest.mark.asyncio
 async def test_marker_session_affinity_read_and_write_agree_for_wildcard_groups():
     """Wildcard deployments keep the literal pattern as model_name on both the read

ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx

@@ -28,13 +28,13 @@ export const AffinityControls: React.FC<{
         <Switch
           checked={value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY}
           onCheckedChange={(deploymentAffinity) => onChange({ ...value, deployment_affinity: deploymentAffinity })}
-          aria-label="Pin a session to one deployment per model group"
+          aria-label="Pin one model deployment per tier"
         />
-        <strong className="font-semibold">Pin a session to one deployment per model group</strong>
+        <strong className="font-semibold">Pin one model deployment per tier</strong>
       </div>
       <span className="block text-xs mb-3 text-muted-foreground">
-        Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to
-        load-balance every turn.
+        Reuses the model chosen for each tier and its deployment when available. Requests can still move between tiers.
+        Turn off to select models and load-balance deployments every turn.
       </span>
       <div style={{ maxWidth: 320 }}>
         <label className="block text-sm font-medium mb-1" htmlFor="session-affinity-ttl">

ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx

@@ -959,7 +959,7 @@ describe("ComplexityRouterConfig affinity panel", () => {
     renderWithProviders(<ComplexityRouterConfig {...baseProps} />);
     fireEvent.click(screen.getByText("Advanced: Affinity"));
 
-    expect(screen.getByRole("switch", { name: "Pin a session to one deployment per model group" })).toBeChecked();
+    expect(screen.getByRole("switch", { name: "Pin one model deployment per tier" })).toBeChecked();
     expect(screen.queryByRole("switch", { name: "Pin a session to its first model" })).not.toBeInTheDocument();
   });
 
@@ -968,7 +968,7 @@ describe("ComplexityRouterConfig affinity panel", () => {
     renderWithProviders(<ComplexityRouterConfig {...baseProps} onChange={onChange} />);
     fireEvent.click(screen.getByText("Advanced: Affinity"));
 
-    fireEvent.click(screen.getByRole("switch", { name: "Pin a session to one deployment per model group" }));
+    fireEvent.click(screen.getByRole("switch", { name: "Pin one model deployment per tier" }));
 
     expect(onChange).toHaveBeenCalledWith({ ...defaultValue, deployment_affinity: false });
   });
@@ -979,7 +979,7 @@ describe("ComplexityRouterConfig affinity panel", () => {
     );
     fireEvent.click(screen.getByText("Advanced: Affinity"));
 
-    expect(screen.getByRole("switch", { name: "Pin a session to one deployment per model group" })).not.toBeChecked();
+    expect(screen.getByRole("switch", { name: "Pin one model deployment per tier" })).not.toBeChecked();
   });
 
   it("writes an idle TTL on blur and keeps the partial input as a draft while typing", () => {

ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx

@@ -715,9 +715,7 @@ describe("AddAutoRouterTab", () => {
     await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router");
     expandDetailedConfiguration();
     await user.click(screen.getByText("Advanced: Affinity"));
-    expect(
-      await screen.findByRole("switch", { name: "Pin a session to one deployment per model group" }),
-    ).toBeChecked();
+    expect(await screen.findByRole("switch", { name: "Pin one model deployment per tier" })).toBeChecked();
 
     await user.click(screen.getByRole("button", { name: /add auto router/i }));
 
@@ -736,7 +734,7 @@ describe("AddAutoRouterTab", () => {
     await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router");
     expandDetailedConfiguration();
     await user.click(screen.getByText("Advanced: Affinity"));
-    await user.click(await screen.findByRole("switch", { name: "Pin a session to one deployment per model group" }));
+    await user.click(await screen.findByRole("switch", { name: "Pin one model deployment per tier" }));
 
     await user.click(screen.getByRole("button", { name: /add auto router/i }));
 

ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx

@@ -519,9 +519,7 @@ describe("EditAutoRouterModal deployment affinity", () => {
     renderWithStoredConfig(STORED_CONFIG);
 
     await user.click(await screen.findByText("Advanced: Affinity"));
-    expect(
-      await screen.findByRole("switch", { name: "Pin a session to one deployment per model group" }),
-    ).toBeChecked();
+    expect(await screen.findByRole("switch", { name: "Pin one model deployment per tier" })).toBeChecked();
 
     await user.click(screen.getByRole("button", { name: /save changes/i }));
 
@@ -534,9 +532,7 @@ describe("EditAutoRouterModal deployment affinity", () => {
     renderWithStoredConfig({ ...STORED_CONFIG, deployment_affinity: false });
 
     await user.click(await screen.findByText("Advanced: Affinity"));
-    expect(
-      await screen.findByRole("switch", { name: "Pin a session to one deployment per model group" }),
-    ).not.toBeChecked();
+    expect(await screen.findByRole("switch", { name: "Pin one model deployment per tier" })).not.toBeChecked();
 
     await user.click(screen.getByRole("button", { name: /save changes/i }));
 
@@ -549,7 +545,7 @@ describe("EditAutoRouterModal deployment affinity", () => {
     renderWithStoredConfig(STORED_CONFIG);
 
     await user.click(await screen.findByText("Advanced: Affinity"));
-    await user.click(await screen.findByRole("switch", { name: "Pin a session to one deployment per model group" }));
+    await user.click(await screen.findByRole("switch", { name: "Pin one model deployment per tier" }));
 
     await user.click(screen.getByRole("button", { name: /save changes/i }));