fix(router): read a Bedrock ARN's base_model window only in context-window escalation

A Bedrock ARN tier with no window of its own now escalates on its base_model's
exact cost-map window. An own window, declared or registered, still wins, and
compaction, pre-call checks and get_router_model_info keep main's behavior.
Covers application inference profile, provisioned and imported model ARNs,
and adds e2e metadata to the new e2e tests.

litellm/llms/bedrock/common_utils.py

@@ -981,6 +981,10 @@ def extract_model_name_from_bedrock_arn(model: str) -> str:
     return model
 
 
+def is_bedrock_arn(model: str) -> bool:
+    return re.search(r"\barn:aws[\w-]*:bedrock:", model) is not None
+
+
 def is_bedrock_application_inference_profile_arn(model: str) -> bool:
     """
     An application inference profile ARN ends in an opaque id with no provider

litellm/router.py

@@ -135,7 +135,7 @@ from litellm.llms.base_llm.vector_store.transformation import (
     RouterVectorStoreEmbeddingExecutor,
     vector_store_request_metadata,
 )
-from litellm.llms.bedrock.common_utils import get_bedrock_base_model, is_bedrock_application_inference_profile_arn
+from litellm.llms.bedrock.common_utils import get_bedrock_base_model
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
 from litellm.llms.openai_like.json_loader import JSONProviderRegistry
 from litellm.llms.openai_like.model_info import (
@@ -1104,7 +1104,7 @@ def as_output_cap(value: object) -> int | None:
     return cap if cap >= 0 else None
 
 
-def _exact_base_model_max_input_tokens(base_model: str) -> int | None:
+def exact_base_model_max_input_tokens(base_model: str) -> int | None:
     try:
         model_info: Final = litellm.get_model_info(model=base_model)
     except Exception:  # noqa: BLE001  # get_model_info raises ModelNotMappedError or a bare Exception for an unmapped model
@@ -1133,22 +1133,15 @@ def _declared_base_model(user_model_info: Mapping[str, object]) -> str | None:
             return None
 
 
-def _base_model_window(
+def _base_model_widened_window(
     route_model_info: ModelMapInfo, base_model: str | None, custom_llm_provider: str
 ) -> Mapping[str, int]:
-    if base_model is None:
-        return MappingProxyType({})
-    if is_bedrock_application_inference_profile_arn(route_model_info["key"]):
-        profile_max_input_tokens: Final = _exact_base_model_max_input_tokens(base_model)
-        return MappingProxyType(
-            {} if profile_max_input_tokens is None else {"max_input_tokens": profile_max_input_tokens}
-        )
     route_max_input_tokens: Final = route_model_info.get("max_input_tokens")
-    if route_max_input_tokens is None:
+    if base_model is None or route_max_input_tokens is None:
         return MappingProxyType({})
     if not _names_route_model(base_model, route_model_info["key"], custom_llm_provider):
         return MappingProxyType({})
-    base_max_input_tokens: Final = _exact_base_model_max_input_tokens(base_model)
+    base_max_input_tokens: Final = exact_base_model_max_input_tokens(base_model)
     if base_max_input_tokens is None or base_max_input_tokens <= route_max_input_tokens:
         return MappingProxyType({})
     return MappingProxyType({"max_input_tokens": base_max_input_tokens})
@@ -11993,8 +11986,7 @@ 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.
-        For every other provider, a base_model naming the same model as `model:` can only widen its context window,
-        and a Bedrock application inference profile ARN, which names no model, takes its base_model's window.
+        For every other provider, a base_model naming the same model as `model:` can only widen its context window.
 
         Returns
         - ModelInfo - If found -> typed dict with max tokens, input cost, etc.
@@ -12101,7 +12093,7 @@ class Router:
         # values are skipped or Deployment's None pricing defaults would erase the map's
         merged_model_info: Final[ModelMapInfo] = {
             **copy.deepcopy(model_info),
-            **_base_model_window(
+            **_base_model_widened_window(
                 model_info,
                 _declared_base_model(user_model_info) or _declared_base_model(litellm_params.model_extra or {}),
                 custom_llm_provider,

litellm/router_strategy/complexity_router/complexity_router.py

@@ -1273,6 +1273,28 @@ def _window_can_hold(window: int | None, needed: int, buffer: float) -> bool:
     return window is None or needed <= int(window * buffer)
 
 
+def _declared_base_model(deployment: Mapping[str, object]) -> str | None:
+    match deployment:
+        case {"model_info": {"base_model": str() as base_model}} | {
+            "litellm_params": {"base_model": str() as base_model}
+        }:
+            return base_model
+        case _:
+            return None
+
+
+def _bedrock_arn_base_model_window(deployment: Mapping[str, object]) -> int | None:
+    from litellm.llms.bedrock.common_utils import is_bedrock_arn
+    from litellm.router import exact_base_model_max_input_tokens
+
+    base_model: Final = _declared_base_model(deployment)
+    match deployment:
+        case {"litellm_params": {"model": str() as model}} if base_model is not None and is_bedrock_arn(model):
+            return exact_base_model_max_input_tokens(base_model)
+        case _:
+            return None
+
+
 def _group_provably_fits(facts: tuple[int | None, bool], needed: int, buffer: float) -> bool:
     window, has_unknown = facts
     return window is not None and not has_unknown and needed <= int(window * buffer)
@@ -3391,7 +3413,8 @@ class ComplexityRouter(CustomLogger):
         return name if self.config.has_custom_tiers else ComplexityTier(name)
 
     def _deployment_window(self, group: str, deployment: Mapping[str, object]) -> int | None:
-        return self._deployment_limit(group, deployment, "max_input_tokens")
+        window: Final = self._deployment_limit(group, deployment, "max_input_tokens")
+        return window if window is not None else _bedrock_arn_base_model_window(deployment)
 
     def _deployment_limit(
         self, group: str, deployment: Mapping[str, object], key: Literal["max_input_tokens", "max_output_tokens"]

tests/code_coverage_tests/router_code_coverage.py

@@ -138,8 +138,8 @@ ignored_function_names = [
     "_pre_call_check_requested_chat_params",  # Tested through the unsupported-param pre-call check tests in test_router.py
     "_deployment_supports_param",  # Tested through the deployment-capability pre-call check tests in test_router.py
     "_capability_model_and_provider",  # Tested through the authenticating-provider pre-call check test in test_router.py
-    "_exact_base_model_max_input_tokens",  # Tested through get_router_model_info in test_router_model_cost_isolation.py
-    "_base_model_window",  # Tested through get_router_model_info in test_router_model_cost_isolation.py
+    "exact_base_model_max_input_tokens",  # Tested through get_router_model_info in test_router_model_cost_isolation.py
+    "_base_model_widened_window",  # Tested through get_router_model_info in test_router_model_cost_isolation.py
     "_provider_model_id",  # Tested through get_router_model_info in test_router_model_cost_isolation.py
     "_names_route_model",  # Tested through get_router_model_info in test_router_model_cost_isolation.py
     "_declared_base_model",  # Tested through get_router_model_info in test_router_model_cost_isolation.py

tests/e2e/router/test_complexity_router_context_window_base_model_e2e.py

@@ -18,9 +18,9 @@ from dataclasses import dataclass
 from typing import Final
 
 import pytest
-
 from e2e_config import unique_marker
 from e2e_http import unwrap
+from e2e_metadata import Domain, Mode, Provider, Subject, meta
 from lifecycle import ResourceManager
 from models import (
     ChatBody,
@@ -159,6 +159,14 @@ def routers(proxy: ProxyClient, router_stack: ExitStack) -> EscalationRouters:
 
 class TestContextWindowEscalationFromBaseModel:
     @pytest.mark.covers("reliability.routing.context_window_escalation.escalates_from_base_model_window")
+    @meta(
+        Subject(
+            domain=Domain.ROUTING,
+            providers=(Provider.BEDROCK,),
+            models=(SIMPLE_BACKEND, UPPER_BACKEND),
+            mode=Mode.NONSTREAM,
+        )
+    )
     def test_base_model_window_escalates(
         self, proxy: ProxyClient, resources: ResourceManager, routers: EscalationRouters
     ) -> None:
@@ -174,6 +182,14 @@ class TestContextWindowEscalationFromBaseModel:
         _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")
+    @meta(
+        Subject(
+            domain=Domain.ROUTING,
+            providers=(Provider.BEDROCK,),
+            models=(EXPLICIT_WINDOW_BACKEND, UPPER_BACKEND),
+            mode=Mode.NONSTREAM,
+        )
+    )
     def test_explicit_window_escalates(
         self, proxy: ProxyClient, resources: ResourceManager, routers: EscalationRouters
     ) -> None:
@@ -189,6 +205,14 @@ class TestContextWindowEscalationFromBaseModel:
         _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")
+    @meta(
+        Subject(
+            domain=Domain.ROUTING,
+            providers=(Provider.BEDROCK,),
+            models=(SIMPLE_BACKEND,),
+            mode=Mode.NONSTREAM,
+        )
+    )
     def test_fitting_prompt_stays_on_the_simple_tier(
         self, proxy: ProxyClient, resources: ResourceManager, routers: EscalationRouters
     ) -> None:

tests/unit/router_strategy/complexity_router/test_context_compaction.py

@@ -312,6 +312,36 @@ async def test_undersized_native_compactor_does_not_block_explicit_escalation(
     assert json.loads(answer.calls[0].request.content)["model"] == "wide-answer"
 
 
+@pytest.mark.parametrize("escalate", [False, True])
+async def test_a_bedrock_arn_tier_windowed_only_through_base_model_is_never_compacted(
+    wire: tuple[respx.Route, respx.Route], escalate: bool,
+) -> None:
+    base_model: Final = "meta.llama3-8b-instruct-v1:0"
+    base_window: Final = litellm.get_model_info(model=base_model)["max_input_tokens"]
+    assert base_window is not None
+    router: Final = make_router(settings={
+        "enable_context_window_escalation": escalate,
+        "tiers": {"SIMPLE": "profile", "MEDIUM": "large", "COMPLEX": "large", "REASONING": "large"},
+    })
+    router.add_deployment(Deployment(
+        model_name="profile", litellm_params={
+            "model": "bedrock/converse/arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/abc123xyz",
+            "mock_response": "profile answer",
+        }, model_info={"id": "profile-answer", "base_model": base_model},
+    ))
+    payload: Final = {"messages": [
+        {"role": "user", "content": "Project code MAPLE-47. Background detail. " * (base_window // 14)},
+        {"role": "assistant", "content": "Recorded"},
+        {"role": "user", "content": "Answer with the project code"},
+    ]}
+    output: Final = base_window // 4
+    trigger: Final = int(base_window * ContextCompactionConfig().trigger_ratio)
+    assert trigger - output < router._count_pre_call_check_tokens(payload["messages"], None, payload) < trigger
+    response: Final = await router.acompletion(model="auto", max_tokens=output, **payload)
+    assert all(route.call_count == 0 for route in wire)
+    assert response.choices[0].message.content == "profile answer"
+
+
 @pytest.mark.parametrize(
     ("surface", "failure"),
     [(surface, failure) for surface in ("chat", "messages", "responses") for failure in ("unsigned", "oversized", "provider")]

tests/unit/router_strategy/test_complexity_router.py

@@ -14354,6 +14354,24 @@ _CJK_TURNS = [
 ]
 
 
+_BEDROCK_PROFILE_ARN: Final = (
+    "bedrock/converse/arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/abc123xyz"
+)
+_LLAMA_BASE_MODEL: Final = "meta.llama3-8b-instruct-v1:0"
+
+
+def _turns_past(window: int) -> list[dict[str, str]]:
+    return [
+        {
+            "role": "user",
+            "content": "Here is everything discussed so far. "
+            + "The meeting notes were saved to the shared folder for later review this week. " * (window // 4),
+        },
+        {"role": "assistant", "content": "Noted, I have read all of it."},
+        {"role": "user", "content": "ok continue"},
+    ]
+
+
 def _tier_config(**overrides: object) -> dict[str, object]:
     return {
         "tiers": {"SIMPLE": "small-model", "COMPLEX": "big-model"},
@@ -14393,43 +14411,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_an_inference_profile_window_declared_only_through_base_model_escalates_too(self):
-        base_model: Final = "meta.llama3-8b-instruct-v1:0"
-        base_window: Final = litellm.get_model_info(model=base_model)["max_input_tokens"]
-        assert base_window is not None
+    @staticmethod
+    def _base_model_tier_router(small: dict[str, object], base_window: int) -> ComplexityRouter:
         litellm_router: Final = 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": "small-model", **small},
                 {
                     "model_name": "big-model",
                     "litellm_params": {"model": "openai/gpt-4o-mini", "mock_response": "ok"},
-                    "model_info": {"max_input_tokens": base_window * 8},
+                    "model_info": {"max_input_tokens": base_window * 16},
                 },
             ]
         )
-        router: Final = ComplexityRouter(
-            model_name="test-router",
-            litellm_router_instance=litellm_router,
-            complexity_router_config=_tier_config(),
+        return ComplexityRouter(
+            model_name="test-router", litellm_router_instance=litellm_router, complexity_router_config=_tier_config()
+        )
+
+    @pytest.mark.asyncio
+    @pytest.mark.usefixtures("local_model_cost_map")
+    @pytest.mark.parametrize(
+        "arn",
+        [
+            _BEDROCK_PROFILE_ARN,
+            "bedrock/converse/arn:aws:bedrock:us-west-2:000000000000:provisioned-model/abc123xyz",
+            "bedrock/llama/arn:aws:bedrock:us-west-2:000000000000:imported-model/abc123xyz",
+        ],
+    )
+    @pytest.mark.parametrize("declared_in", ["model_info", "litellm_params"])
+    async def test_a_bedrock_arn_window_declared_only_through_base_model_escalates_too(
+        self, arn: str, declared_in: str
+    ):
+        base_window: Final = litellm.get_model_info(model=_LLAMA_BASE_MODEL)["max_input_tokens"]
+        assert base_window is not None
+        router: Final = self._base_model_tier_router(
+            {
+                "litellm_params": {
+                    "model": arn,
+                    "mock_response": "ok",
+                    **({"base_model": _LLAMA_BASE_MODEL} if declared_in == "litellm_params" else {}),
+                },
+                "model_info": {"base_model": _LLAMA_BASE_MODEL} if declared_in == "model_info" else {},
+            },
+            base_window,
         )
-        filler_repeats: Final = (base_window * 4) // len(_CONTEXT_FILLER) + 2
-        oversized_turns: Final = [
-            {"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: Final = await router.async_pre_routing_hook(
-            model="test-router", request_kwargs={}, messages=oversized_turns
+            model="test-router", request_kwargs={}, messages=_turns_past(base_window)
         )
 
         assert result is not None
@@ -14437,6 +14464,67 @@ class TestContextWindowEscalation:
         assert result.routing_decision["context_escalation_original_tier"] == "SIMPLE"
         assert "context_escalation" in result.routing_decision["signals"]
 
+    @pytest.mark.asyncio
+    @pytest.mark.usefixtures("local_model_cost_map")
+    @pytest.mark.parametrize("own_window", ["declared", "registered", "not_a_bedrock_arn"])
+    async def test_base_model_never_overrides_an_own_window_or_reaches_a_non_arn_route(
+        self, own_window: str, monkeypatch: pytest.MonkeyPatch
+    ):
+        base_window: Final = litellm.get_model_info(model=_LLAMA_BASE_MODEL)["max_input_tokens"]
+        assert base_window is not None
+        if own_window == "registered":
+            monkeypatch.setitem(
+                litellm.model_cost,
+                _BEDROCK_PROFILE_ARN,
+                {"max_input_tokens": base_window * 8, "litellm_provider": "bedrock", "mode": "chat"},
+            )
+        router: Final = self._base_model_tier_router(
+            {
+                "litellm_params": {
+                    "model": (
+                        "openai/a-model-the-cost-map-does-not-know"
+                        if own_window == "not_a_bedrock_arn"
+                        else _BEDROCK_PROFILE_ARN
+                    ),
+                    "mock_response": "ok",
+                },
+                "model_info": {
+                    "base_model": _LLAMA_BASE_MODEL,
+                    **({"max_input_tokens": base_window * 8} if own_window == "declared" else {}),
+                },
+            },
+            base_window,
+        )
+
+        result: Final = await router.async_pre_routing_hook(
+            model="test-router", request_kwargs={}, messages=_turns_past(base_window)
+        )
+
+        assert result is not None
+        assert result.model == "small-model"
+        assert "context_escalated" not in result.routing_decision
+
+    @pytest.mark.asyncio
+    @pytest.mark.usefixtures("local_model_cost_map")
+    async def test_an_unmapped_base_model_leaves_a_bedrock_arn_tier_in_place(self):
+        base_window: Final = litellm.get_model_info(model=_LLAMA_BASE_MODEL)["max_input_tokens"]
+        assert base_window is not None
+        router: Final = self._base_model_tier_router(
+            {
+                "litellm_params": {"model": _BEDROCK_PROFILE_ARN, "mock_response": "ok"},
+                "model_info": {"base_model": "anthropic.claude-not-in-the-cost-map-v1:0"},
+            },
+            base_window,
+        )
+
+        result: Final = await router.async_pre_routing_hook(
+            model="test-router", request_kwargs={}, messages=_turns_past(base_window)
+        )
+
+        assert result is not None
+        assert result.model == "small-model"
+        assert "context_escalated" not in result.routing_decision
+
     @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

@@ -16263,55 +16263,6 @@ class TestAzureBaseModelFallbackLogging:
         )
         assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"]
 
-
-INFERENCE_PROFILE_ARN: Final = (
-    "bedrock/converse/arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/abc123xyz"
-)
-INFERENCE_PROFILE_BASE_MODEL: Final = "global.anthropic.claude-haiku-4-5-20251001-v1:0"
-
-
-@pytest.mark.usefixtures("local_model_cost_map")
-class TestApplicationInferenceProfileWindowFromBaseModel:
-    def _profile_window(self, model_info: dict[str, object], litellm_params: dict[str, object] | None = None) -> object:
-        router: Final = litellm.Router(
-            model_list=[
-                {
-                    "model_name": "haiku",
-                    "litellm_params": {"model": INFERENCE_PROFILE_ARN, **(litellm_params or {})},
-                    "model_info": {"id": "haiku-profile", **model_info},
-                }
-            ]
-        )
-        return router.get_router_model_info(deployment=None, received_model_name="haiku", id="haiku-profile")[
-            "max_input_tokens"
-        ]
-
-    def test_model_info_base_model_sets_the_profile_window(self):
-        base_window: Final = litellm.get_model_info(model=INFERENCE_PROFILE_BASE_MODEL)["max_input_tokens"]
-        assert base_window is not None
-
-        assert self._profile_window({"base_model": INFERENCE_PROFILE_BASE_MODEL}) == base_window
-
-    def test_litellm_params_base_model_sets_the_profile_window(self):
-        base_window: Final = litellm.get_model_info(model=INFERENCE_PROFILE_BASE_MODEL)["max_input_tokens"]
-        assert base_window is not None
-
-        assert self._profile_window({}, {"base_model": INFERENCE_PROFILE_BASE_MODEL}) == base_window
-
-    def test_explicit_max_input_tokens_still_wins_over_base_model(self):
-        base_window: Final = litellm.get_model_info(model=INFERENCE_PROFILE_BASE_MODEL)["max_input_tokens"]
-        assert base_window is not None
-
-        declared_window: Final = self._profile_window(
-            {"base_model": INFERENCE_PROFILE_BASE_MODEL, "max_input_tokens": base_window // 2}
-        )
-
-        assert declared_window == base_window // 2
-
-    def test_an_unmapped_base_model_leaves_the_profile_window_unknown(self):
-        assert self._profile_window({"base_model": "anthropic.claude-not-in-the-cost-map-v1:0"}) is None
-
-
 def test_model_group_info_intersects_supported_reasoning_efforts():
     router = litellm.Router(
         model_list=[