fix: preserve session pins during image capability checks

litellm/router_strategy/complexity_router/README.md

@@ -433,8 +433,8 @@ continue serving text requests but are excluded from image requests.
 A same-tier re-pick keeps the decision's cause and adds `modality:image` to `signals`; a tier
 change or default takeover records `cause: modality_escalation` with the displaced placement
 (`modality_escalated_from:<TIER>` or `modality_displaced_default_model`). Escalations are never
-pinned by session affinity, and by default a KEPT session pin bypasses the gate: a session pinned
-to a text-only or unknown model fails before provider dispatch when an image arrives.
+pinned by session affinity. Without `modality_pin_override`, an image request fails before provider
+dispatch when the pinned group has no compatible deployment.
 
 Context-window and modality recovery take priority over the default model. If a compatible tier
 cannot serve, the router checks the remaining compatible recovery tiers before using `default_model`.

litellm/router_strategy/complexity_router/complexity_router.py

@@ -42,6 +42,7 @@ from litellm.constants import (
     RETURN_RAW_MODEL_NAME_METADATA_KEY,
     SESSION_ID_GENERATED_METADATA_KEY,
 )
+from litellm.exceptions import BadRequestError
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.integrations.otel.runtime import phase_event
 from litellm.litellm_core_utils.classifier_logging import masked_originating_request
@@ -3719,34 +3720,29 @@ class ComplexityRouter(CustomLogger):
         context_fit: _RequestContextFit | None = None,
         retained_pin: _SessionAffinityPin | None = None,
     ) -> PreRoutingHookResponse:
-        """Replace a routed model that cannot accept this request's image input.
-
-        The single modality owner, applied to the decided response at the hook's exits so every
-        routing path is covered uniformly. A KEPT session pin is exempt by design (its cause)
-        unless modality_pin_override is set, in which case the image turn is re-placed and reported
-        as modality_pin_override while the stored pin, written upstream from the session's own
-        model, is left for the next text turn; replacement picks and every other path are just
-        responses. The re-placement walks
-        UPWARD-ONLY from the decision's tier (so a plan-mode floor can never be undercut), picks
-        through `_pick_model_for_tier` so routing plugins still apply, then falls to
-        default_model (never on plugin routers, and never on a plan-floored decision, since
-        default_model carries no tier guarantee), else raises the clear 400. The rewritten
-        decision keeps its cause on a same-tier repick and becomes modality_escalation when the
-        tier moved or default_model took over, with the displaced placement in signals.
-        """
+        """Keep compatible placements, reject incompatible pins, and repick permitted image routes."""
         decision: Final = response.routing_decision
         if (
             not resolved_messages
             or response.model is None
-            or (
-                decision is not None
-                and decision.get("cause") == "session_affinity_pin"
-                and not self.config.modality_pin_override
-            )
             or not request_contains_image_content(resolved_messages)
             or self._model_accepts_image_input(response.model, _ROUTING_KWARGS_ADAPTER.validate_python(request_kwargs))
         ):
             return response
+        if (
+            decision is not None
+            and decision.get("cause") == "session_affinity_pin"
+            and not self.config.modality_pin_override
+        ):
+            raise BadRequestError(
+                message=(
+                    f"Auto-router {self.model_name} is pinned to '{response.model}', which has no deployment "
+                    "with confirmed image support. Configure model_info.supports_vision for a vision-capable "
+                    "deployment, enable modality_pin_override to permit relocation, or remove the image content."
+                ),
+                model=self.model_name,
+                llm_provider="",
+            )
         fit: Final = context_fit or await self._request_context_fit(resolved_messages, request_kwargs)
         eligible: Final = frozenset(
             name
@@ -3775,9 +3771,7 @@ class ComplexityRouter(CustomLogger):
             new_tier = None
             new_model = self._placed_default_model()
         else:
-            import litellm
-
-            raise litellm.BadRequestError(
+            raise BadRequestError(
                 message=(
                     f"Auto-router {self.model_name} received a request with image input, but no model "
                     f"at or above the decided tier has confirmed image support. "
@@ -3791,7 +3785,7 @@ class ComplexityRouter(CustomLogger):
         same_tier: Final = capable is not None and decided == capable
         base_cause: Final = (decision.get("cause") if decision is not None else None) or "default_fallback"
         # Reaching here on a kept pin means modality_pin_override is on, since the guard above
-        # returns otherwise. The model moved off the pin even on a same-tier repick, so reporting
+        # rejects otherwise. The model moved off the pin even on a same-tier repick, so reporting
         # the pin's own cause would claim the session's model served a request it did not.
         displaced_pin: Final = base_cause == "session_affinity_pin"
         displaced_default: Final = decided is None and response.model == self.config.default_model

tests/unit/router_strategy/test_complexity_router.py

@@ -14853,7 +14853,7 @@ class TestModalityRouting:
             ("no_ask_capable_default_kept", "vision-default", "default_fallback"),
             ("no_ask_text_default_displaced", "vision-mid", "modality_escalation"),
             ("custom_tiers_walk", "premium-model", "modality_escalation"),
-            ("pin_kept_bypasses", "text-cheap", "session_affinity_pin"),
+            ("pin_kept_rejects", None, None),
             ("pin_replacement_gated", "vision-big", "modality_escalation"),
             ("pin_override_escalates", "vision-mid", "modality_pin_override"),
             ("pin_override_same_tier", "vision-cheap", "modality_pin_override"),
@@ -14932,6 +14932,10 @@ class TestModalityRouting:
             mock_router_instance.model_list = []
             mock_router_instance.model_name_to_deployment_indices = {}
         router = self._router(mock_router_instance, config, vision)
+        if path == "pin_kept_rejects":
+            with pytest.raises(litellm.BadRequestError, match="modality_pin_override"):
+                await router.async_pre_routing_hook(model="m", request_kwargs=request_kwargs, messages=messages)
+            return
         result = await router.async_pre_routing_hook(model="m", request_kwargs=request_kwargs, messages=messages)
         assert result.model == expected_model
         assert result.routing_decision["cause"] == expected_cause

tests/unit/test_router/test_router.py

@@ -29993,3 +29993,70 @@ async def test_image_dispatch_and_fallback_never_call_text_deployment(
         assert vision_route.called
         assert backup_route.called is primary_fails
         assert not text_route.called
+
+
+@pytest.mark.parametrize("allow_pin_override", (False, True))
+async def test_image_request_cannot_relocate_retained_pin_through_health_fallback(allow_pin_override: bool) -> None:
+    router: Final = Router(
+        model_list=[
+            {
+                "model_name": "automatic",
+                "litellm_params": {
+                    "model": "auto_router/complexity_router",
+                    "complexity_router_config": {
+                        "tiers": {"SIMPLE": "text-pool", "MEDIUM": "vision-pool"},
+                        "default_model": "vision-pool",
+                        "session_affinity": True,
+                        "modality_pin_override": allow_pin_override,
+                    },
+                },
+            },
+            {
+                "model_name": "text-pool",
+                "litellm_params": {"model": "openai/custom-text", "api_key": "test"},
+                "model_info": {"id": "text", "supports_vision": False},
+            },
+            {
+                "model_name": "vision-pool",
+                "litellm_params": {"model": "openai/custom-vision", "api_key": "test"},
+                "model_info": {"id": "vision", "supports_vision": True},
+            },
+        ]
+    )
+    metadata: Final = {"session_id": "image-session", "user_api_key_hash": "test-key"}
+    text_messages: Final = [{"role": "user", "content": "Hi"}]
+    first: Final = await router.async_get_available_deployment(
+        model="automatic",
+        messages=text_messages,
+        request_kwargs={"metadata": dict(metadata)},
+    )
+    assert first["model_info"]["id"] == "text"
+    image_messages: Final = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "What color is this?"},
+                {"type": "image_url", "image_url": {"url": "https://example.invalid/image.png"}},
+            ],
+        }
+    ]
+    if allow_pin_override:
+        image: Final = await router.async_get_available_deployment(
+            model="automatic",
+            messages=image_messages,
+            request_kwargs={"metadata": dict(metadata)},
+        )
+        assert image["model_info"]["id"] == "vision"
+    else:
+        with pytest.raises(litellm.BadRequestError, match="modality_pin_override"):
+            await router.async_get_available_deployment(
+                model="automatic",
+                messages=image_messages,
+                request_kwargs={"metadata": dict(metadata)},
+            )
+    last: Final = await router.async_get_available_deployment(
+        model="automatic",
+        messages=text_messages,
+        request_kwargs={"metadata": dict(metadata)},
+    )
+    assert last["model_info"]["id"] == "text"