fix(router): resolve a deployment's window from base_model for every provider

litellm/router.py

@@ -11035,7 +11035,9 @@ class Router:
 
         Augment litellm info with additional params set in `model_info`.
 
-        For azure models, ignore the `model:`. Only set max tokens, cost values if base_model is set.
+        A deployment's `base_model` (model_info first, then litellm_params) is the cost-map key for every
+        provider, the way /model/info resolves it. Without one, azure needs the deployment name to be a
+        known azure key and every other provider resolves its own `model:`.
 
         Returns
         - ModelInfo - If found -> typed dict with max tokens, input cost, etc.
@@ -11077,7 +11079,7 @@ class Router:
         )
 
         ## SET MODEL TO 'model=' - if base_model is None + not azure
-        if custom_llm_provider == "azure" and base_model is None:
+        if base_model is None and custom_llm_provider == "azure":
             # Router init auto-registers every deployment name into
             # litellm.model_cost as a zeroed stub, so membership alone can't
             # tell a resolvable name apart; require usable limits/costs.
@@ -11099,7 +11101,7 @@ class Router:
                     "Could not identify azure model '%s'. Set azure 'base_model' for accurate max tokens, cost tracking, etc.- https://docs.litellm.ai/docs/proxy/cost_tracking#spend-tracking-for-azure-openai-models",
                     _model,
                 )
-        elif custom_llm_provider != "azure":
+        elif base_model is None:
             model = _model
 
             if "*" in model:  # only call pattern_router for wildcard models

tests/e2e/coverage_registry/reliability.yaml

@@ -30,6 +30,9 @@
 - {id: reliability.routing.semantic_auto_router.responses_input_routed, module: reliability, tier: P0, behavior: routing, variant: semantic_auto_router, assertions: [responses_input_routed], exercised_on: [responses], source: "litellm/router_strategy/auto_router/auto_router.py:131", fail_before_fix: proven, rationale: "/v1/responses input is resolved into messages for the semantic auto-router pre-routing hook instead of failing 400 Unmapped LLM provider auto_router (GitHub PR #37333)"}
 - {id: reliability.routing.strategy_alias.custom_pricing_ignored, module: reliability, tier: P1, behavior: routing, variant: strategy_alias, assertions: [custom_pricing_ignored], exercised_on: [chat_completions], source: "litellm/router.py:11489", rationale: "Custom pricing on a strategy-router alias never prices the routed request; spend logs at the routed tier deployment's own rate (GitHub PR #36691)"}
 - {id: reliability.routing.complexity_heuristic.scores_current_ask_only, module: reliability, tier: P1, behavior: routing, variant: complexity_heuristic, assertions: [scores_current_ask_only], exercised_on: [chat_completions], source: "router_strategy/complexity_router/complexity_router.py:942", rationale: "The heuristic complexity classifier scores the caller's current ask only, so a keyword-heavy agent system prompt cannot inflate the tier (GitHub PR #36721)"}
+- {id: reliability.routing.context_window_escalation.escalates_from_base_model_window, module: reliability, tier: P0, behavior: routing, variant: context_window_escalation, assertions: [escalates_from_base_model_window], exercised_on: [chat_completions], source: "litellm/router.py:11081", fail_before_fix: proven, rationale: "An oversized SIMPLE request moves to the next tier that fits when the SIMPLE deployment's window comes only from model_info.base_model (a Bedrock application inference profile ARN), instead of staying put and 400ing at the provider"}
+- {id: reliability.routing.context_window_escalation.escalates_from_explicit_window, module: reliability, tier: P1, behavior: routing, variant: context_window_escalation, assertions: [escalates_from_explicit_window], exercised_on: [chat_completions], source: "router_strategy/complexity_router/complexity_router.py:1238", rationale: "An oversized SIMPLE request moves to the next tier that fits when the SIMPLE deployment declares max_input_tokens explicitly"}
+- {id: reliability.routing.context_window_escalation.fitting_prompt_stays_on_tier, module: reliability, tier: P1, behavior: routing, variant: context_window_escalation, assertions: [fitting_prompt_stays_on_tier], exercised_on: [chat_completions], source: "router_strategy/complexity_router/complexity_router.py:1238", rationale: "A request that fits the SIMPLE deployment's base_model window is served by that deployment, so resolving the window never over-escalates"}
 - {id: reliability.cache.exact.returns_cached, module: reliability, tier: P1, behavior: cache, variant: exact, assertions: [returns_cached], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/caching.py", rationale: "Response cache returns cached on exact match"}
 - {id: reliability.cache.prompt_caching_model_select.returns_cached, module: reliability, tier: P1, behavior: cache, variant: prompt_caching_model_select, assertions: [returns_cached], exercised_on: [chat_completions], source: "router_utils/prompt_caching_cache.py", rationale: "Selects model supporting prompt caching for cacheable prefix; runs only on a stack with the prompt_caching pre-call check enabled (E2E_PROMPT_CACHING_STACK)"}
 - {id: reliability.circuit_breaker.redis.trips_then_recovers, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/redis_cache.py:99", rationale: "Redis breaker CLOSED->OPEN->HALF_OPEN; guards all cache/rate-limit ops"}

tests/e2e/models.py

@@ -1319,6 +1319,8 @@ class ModelInfoBody(BaseModel):
     team_id: str | None = None
     allowed_fails: int | None = None
     allowed_fails_policy: dict[str, int] | None = None
+    base_model: str | None = None
+    max_input_tokens: int | None = None
 
 
 class ModelNewBody(BaseModel):

tests/e2e/router/test_complexity_router_context_window_base_model_e2e.py

@@ -0,0 +1,219 @@
+"""Live e2e pins for the complexity router's context-window escalation when a tier's
+window is declared only through ``model_info.base_model``.
+
+A Bedrock application inference profile ARN has no cost-map row of its own, so the
+SIMPLE tier's window exists only through ``base_model``. ``Router.get_router_model_info``
+read base_model for azure deployments alone and resolved every other provider's
+deployment by its own ``model:``, so the ARN came back with no window, the escalation
+gate treated an unknown window as able to hold anything, and an oversized request the
+classifier scored SIMPLE stayed on the 200k model and 400'd at Bedrock with "prompt is
+too long" instead of moving to the 1M MEDIUM tier.
+
+The explicit-window control declares ``max_input_tokens`` below its model's own row and
+sends a request that fits the row but not the declaration, so it escalates only because
+a declared window is honored. It lives on its own backend model: the router registers a
+deployment's ``model_info`` under the shared cost-map key of its ``model:``, so a declared
+window on the ARN would leak to the base_model deployment and hide the bug. A short
+request on the base_model router proves the gate only moves what does not fit.
+
+Every deployment is registered via /model/new (stage has no static config for these)
+and the served deployment is read back from the spend log's ``model``, which stores
+either the registered alias or the provider-prefixed form. The prompts are plain prose
+with a "how many" ask, which keeps the heuristic classifier on SIMPLE, and the router
+estimates their size as chars/4.
+"""
+
+import os
+from collections.abc import Iterator
+from contextlib import ExitStack
+from dataclasses import dataclass
+from typing import Final
+
+import pytest
+
+from e2e_config import unique_marker
+from e2e_http import unwrap
+from lifecycle import ResourceManager
+from models import (
+    ChatBody,
+    ChatMessage,
+    ChatResponse,
+    KeyGenerateBody,
+    LiteLLMParamsBody,
+    ModelInfoBody,
+    ModelNewBody,
+    SpendLogRow,
+)
+from proxy_client import ProxyClient
+
+pytestmark = pytest.mark.e2e
+
+SIMPLE_BACKEND: Final = os.environ.get(
+    "E2E_BEDROCK_INFERENCE_PROFILE_ARN",
+    "bedrock/converse/arn:aws:bedrock:us-west-2:888602223428:application-inference-profile/sccejsk6r44v",
+)
+SIMPLE_BASE_MODEL: Final = "global.anthropic.claude-haiku-4-5-20251001-v1:0"
+EXPLICIT_WINDOW_BACKEND: Final = "bedrock/converse/us.anthropic.claude-haiku-4-5-20251001-v1:0"
+DECLARED_WINDOW_TOKENS: Final = 60_000
+UPPER_BACKEND: Final = "bedrock/converse/global.amazon.nova-2-lite-v1:0"
+MAX_TOKENS: Final = 8
+OVERSIZED_PROMPT_CHARS: Final = 1_000_000
+BEYOND_DECLARED_WINDOW_PROMPT_CHARS: Final = 400_000
+FILLER: Final = (
+    "On the morning of day {n}, the garden club met beside the old mill and talked about which flowers "
+    "to grow along the river walk. Everyone brought warm bread and shared stories about the cold winter "
+    "and the first robin of spring.\n"
+)
+ASK: Final = "How many days are written above? Reply with the single word OK."
+
+
+@dataclass(frozen=True, slots=True)
+class EscalationRouters:
+    base_model_router: str
+    explicit_window_router: str
+    simple_from_base_model: str
+    simple_with_explicit_window: str
+    upper: str
+
+
+def _bedrock_params(model: str) -> LiteLLMParamsBody:
+    return LiteLLMParamsBody(
+        model=model,
+        aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID",
+        aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY",
+        aws_region_name="us-west-2",
+    )
+
+
+def _router_params(simple: str, upper: str) -> LiteLLMParamsBody:
+    return LiteLLMParamsBody(
+        model="auto_router/complexity_router",
+        complexity_router_config={
+            "classifier_type": "heuristic",
+            "enable_context_window_escalation": True,
+            "tiers": {"SIMPLE": simple, "MEDIUM": upper, "COMPLEX": upper, "REASONING": upper},
+        },
+    )
+
+
+def _prompt_of(chars: int) -> str:
+    marker: Final = unique_marker()
+    line_count: Final = -(-chars // len(FILLER.format(n=f"1 ({marker})")))
+    lines: Final = tuple(FILLER.format(n=f"{n} ({marker})") for n in range(1, line_count + 1))
+    return "".join(lines) + "\n" + ASK
+
+
+def _chat_body(model: str, content: str) -> ChatBody:
+    return ChatBody(model=model, messages=[ChatMessage(role="user", content=content)], max_tokens=MAX_TOKENS)
+
+
+def _key_for(proxy: ProxyClient, resources: ResourceManager, models: list[str]) -> str:
+    key: Final = proxy.generate_key(KeyGenerateBody(models=models, user_id="e2e-context-window-base-model"))
+    resources.defer(lambda: proxy.delete_key(key))
+    return key
+
+
+def _served_by(rows: list[SpendLogRow], alias: str, backend: str, context: str) -> None:
+    served: Final = tuple(row.model for row in rows)
+    assert served and all(model in {alias, backend} for model in served), (
+        f"{context}: expected every request to be served by {alias} ({backend}), spend logs show {served}"
+    )
+
+
+def _answered(response: ChatResponse) -> None:
+    message: Final = response.choices[0].message if response.choices else None
+    assert message is not None and message.content, f"expected an answer, got {response}"
+
+
+@pytest.fixture(scope="class")
+def router_stack() -> Iterator[ExitStack]:
+    with ExitStack() as stack:
+        yield stack
+
+
+@pytest.fixture(scope="class")
+def routers(proxy: ProxyClient, router_stack: ExitStack) -> EscalationRouters:
+    marker: Final = unique_marker()
+    named: Final = EscalationRouters(
+        base_model_router=f"e2e-cw-router-base-{marker}",
+        explicit_window_router=f"e2e-cw-router-explicit-{marker}",
+        simple_from_base_model=f"e2e-cw-simple-base-{marker}",
+        simple_with_explicit_window=f"e2e-cw-simple-explicit-{marker}",
+        upper=f"e2e-cw-upper-{marker}",
+    )
+    registrations: Final = (
+        ModelNewBody(
+            model_name=named.simple_from_base_model,
+            litellm_params=_bedrock_params(SIMPLE_BACKEND),
+            model_info=ModelInfoBody(base_model=SIMPLE_BASE_MODEL),
+        ),
+        ModelNewBody(
+            model_name=named.simple_with_explicit_window,
+            litellm_params=_bedrock_params(EXPLICIT_WINDOW_BACKEND),
+            model_info=ModelInfoBody(max_input_tokens=DECLARED_WINDOW_TOKENS),
+        ),
+        ModelNewBody(model_name=named.upper, litellm_params=_bedrock_params(UPPER_BACKEND), model_info=ModelInfoBody()),
+        ModelNewBody(
+            model_name=named.base_model_router,
+            litellm_params=_router_params(named.simple_from_base_model, named.upper),
+            model_info=ModelInfoBody(),
+        ),
+        ModelNewBody(
+            model_name=named.explicit_window_router,
+            litellm_params=_router_params(named.simple_with_explicit_window, named.upper),
+            model_info=ModelInfoBody(),
+        ),
+    )
+    for body in registrations:
+        router_stack.callback(proxy.delete_model, proxy.register_model(body))
+    return named
+
+
+class TestContextWindowEscalationFromBaseModel:
+    @pytest.mark.covers("reliability.routing.context_window_escalation.escalates_from_base_model_window")
+    def test_base_model_window_escalates(
+        self, proxy: ProxyClient, resources: ResourceManager, routers: EscalationRouters
+    ) -> None:
+        key: Final = _key_for(
+            proxy, resources, [routers.base_model_router, routers.simple_from_base_model, routers.upper]
+        )
+
+        response: Final = unwrap(
+            proxy.chat(key, _chat_body(routers.base_model_router, _prompt_of(OVERSIZED_PROMPT_CHARS)))
+        )
+
+        _answered(response)
+        _served_by(proxy.poll_logs_for_key(key), routers.upper, UPPER_BACKEND, "oversized ask, window from base_model")
+
+    @pytest.mark.covers("reliability.routing.context_window_escalation.escalates_from_explicit_window")
+    def test_explicit_window_escalates(
+        self, proxy: ProxyClient, resources: ResourceManager, routers: EscalationRouters
+    ) -> None:
+        key: Final = _key_for(
+            proxy, resources, [routers.explicit_window_router, routers.simple_with_explicit_window, routers.upper]
+        )
+
+        response: Final = unwrap(
+            proxy.chat(key, _chat_body(routers.explicit_window_router, _prompt_of(BEYOND_DECLARED_WINDOW_PROMPT_CHARS)))
+        )
+
+        _answered(response)
+        _served_by(proxy.poll_logs_for_key(key), routers.upper, UPPER_BACKEND, "ask beyond the declared window")
+
+    @pytest.mark.covers("reliability.routing.context_window_escalation.fitting_prompt_stays_on_tier")
+    def test_fitting_prompt_stays_on_the_simple_tier(
+        self, proxy: ProxyClient, resources: ResourceManager, routers: EscalationRouters
+    ) -> None:
+        key: Final = _key_for(
+            proxy, resources, [routers.base_model_router, routers.simple_from_base_model, routers.upper]
+        )
+
+        response: Final = unwrap(proxy.chat(key, _chat_body(routers.base_model_router, f"say hello {unique_marker()}")))
+
+        _answered(response)
+        _served_by(
+            proxy.poll_logs_for_key(key),
+            routers.simple_from_base_model,
+            SIMPLE_BACKEND,
+            "short ask, window from base_model",
+        )

tests/unit/router_strategy/test_complexity_router.py

@@ -14225,6 +14225,52 @@ class TestContextWindowEscalation:
         assert result.routing_decision["tier"] == "COMPLEX"
         assert "context_escalation" in result.routing_decision["signals"]
 
+    @pytest.mark.asyncio
+    @pytest.mark.usefixtures("local_model_cost_map")
+    async def test_a_window_declared_only_through_base_model_escalates_too(self):
+        """A Bedrock application inference profile ARN has no cost-map row; its window exists
+        only through `model_info.base_model`. The gate used to resolve that deployment to no
+        window at all (base_model was read for azure only), treat it as infinite, and leave an
+        oversized SIMPLE prompt on it to 400 at the provider."""
+        base_model = "meta.llama3-8b-instruct-v1:0"
+        base_window = litellm.get_model_info(model=f"bedrock/{base_model}")["max_input_tokens"]
+        assert base_window is not None
+        litellm_router = Router(
+            model_list=[
+                {
+                    "model_name": "small-model",
+                    "litellm_params": {
+                        "model": "bedrock/converse/arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/abc123xyz",
+                        "mock_response": "ok",
+                    },
+                    "model_info": {"base_model": base_model},
+                },
+                {
+                    "model_name": "big-model",
+                    "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"},
+                    "model_info": {"max_input_tokens": base_window * 8},
+                },
+            ]
+        )
+        router = ComplexityRouter(
+            model_name="test-router",
+            litellm_router_instance=litellm_router,
+            complexity_router_config=_tier_config(),
+        )
+        filler_repeats = (base_window * 4) // len(_CONTEXT_FILLER) + 2
+        oversized_turns = [
+            {"role": "user", "content": "Here is everything discussed so far. " + _CONTEXT_FILLER * filler_repeats},
+            {"role": "assistant", "content": "Noted, I have read all of it."},
+            {"role": "user", "content": "ok continue"},
+        ]
+
+        result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=oversized_turns)
+
+        assert result is not None
+        assert result.model == "big-model"
+        assert result.routing_decision["context_escalation_original_tier"] == "SIMPLE"
+        assert "context_escalation" in result.routing_decision["signals"]
+
     @pytest.mark.asyncio
     async def test_a_prompt_that_fits_routes_exactly_as_before(self):
         """The gate must be invisible for normal traffic: same model, no escalation facts."""

tests/unit/test_router/test_router.py

@@ -13650,6 +13650,73 @@ class TestAzureBaseModelFallbackLogging:
         )
         assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"]
 
+
+INFERENCE_PROFILE_ARN = "bedrock/converse/arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/abc123xyz"
+INFERENCE_PROFILE_BASE_MODEL = "global.anthropic.claude-haiku-4-5-20251001-v1:0"
+
+
+@pytest.mark.usefixtures("local_model_cost_map")
+class TestBaseModelResolvesWindowForEveryProvider:
+    """A Bedrock application inference profile ARN is not in the cost map, so its window and
+    prices only exist through `base_model`. get_router_model_info used to read base_model for
+    azure alone and resolved every other provider's deployment by its own `model:`, so the
+    ARN came back unmapped and the complexity router's context-window escalation, which
+    treats an unknown window as infinite, never moved an oversized request off that tier."""
+
+    def _router(self, model_info: dict[str, object], litellm_params: dict[str, object] | None = None) -> litellm.Router:
+        return litellm.Router(
+            model_list=[
+                {
+                    "model_name": "haiku",
+                    "litellm_params": {"model": INFERENCE_PROFILE_ARN, **(litellm_params or {})},
+                    "model_info": {"id": "haiku-profile", **model_info},
+                }
+            ]
+        )
+
+    def test_model_info_base_model_resolves_a_non_azure_deployment(self):
+        router = self._router({"base_model": INFERENCE_PROFILE_BASE_MODEL})
+
+        model_info = router.get_router_model_info(deployment=None, received_model_name="haiku", id="haiku-profile")
+
+        base = litellm.get_model_info(model=f"bedrock/{INFERENCE_PROFILE_BASE_MODEL}")
+        assert model_info["max_input_tokens"] == base["max_input_tokens"]
+        assert model_info["input_cost_per_token"] == base["input_cost_per_token"]
+
+    def test_litellm_params_base_model_resolves_a_non_azure_deployment(self):
+        router = self._router({}, {"base_model": INFERENCE_PROFILE_BASE_MODEL})
+
+        model_info = router.get_router_model_info(deployment=None, received_model_name="haiku", id="haiku-profile")
+
+        base = litellm.get_model_info(model=f"bedrock/{INFERENCE_PROFILE_BASE_MODEL}")
+        assert model_info["max_input_tokens"] == base["max_input_tokens"]
+
+    def test_explicit_max_input_tokens_still_wins_over_base_model(self):
+        base_window = litellm.get_model_info(model=f"bedrock/{INFERENCE_PROFILE_BASE_MODEL}")["max_input_tokens"]
+        assert base_window is not None
+        router = self._router({"base_model": INFERENCE_PROFILE_BASE_MODEL, "max_input_tokens": base_window // 2})
+
+        model_info = router.get_router_model_info(deployment=None, received_model_name="haiku", id="haiku-profile")
+
+        assert model_info["max_input_tokens"] == base_window // 2
+
+    def test_a_deployment_without_base_model_still_resolves_its_own_model(self):
+        router = litellm.Router(
+            model_list=[
+                {
+                    "model_name": "nova",
+                    "litellm_params": {"model": "bedrock/converse/global.amazon.nova-2-lite-v1:0"},
+                    "model_info": {"id": "nova-plain"},
+                }
+            ]
+        )
+
+        model_info = router.get_router_model_info(deployment=None, received_model_name="nova", id="nova-plain")
+
+        own = litellm.get_model_info(model="bedrock/converse/global.amazon.nova-2-lite-v1:0")
+        assert model_info["max_input_tokens"] == own["max_input_tokens"]
+
+
 def test_model_group_info_intersects_supported_reasoning_efforts():
     router = litellm.Router(
         model_list=[