feat(bedrock): route token counts from model metadata (#47441)

* feat(bedrock): route token counts from model metadata

* fix(bedrock): keep counting metadata compatible with older proxies

litellm/llms/bedrock/count_tokens/bedrock_token_counter.py

@@ -59,8 +59,7 @@ def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model:
 
 
 class BedrockTokenCounter(BaseTokenCounter):
-    """Token counter for AWS Bedrock: bedrock-runtime CountTokens first, and Anthropic's count_tokens
-    on bedrock-mantle for Claude models bedrock-runtime answers 400 or 404 for"""
+    """Counts on the configured Bedrock endpoint, with Mantle retry when Runtime rejects a Claude model"""
 
     def __init__(
         self,
@@ -127,17 +126,22 @@ class BedrockTokenCounter(BaseTokenCounter):
             **({"tools": tools} if tools else {}),
             **({"system": system} if system else {}),
         }
-        resolved_model: Final = get_bedrock_base_model(model_to_use)
+        use_mantle: Final = model_to_use.startswith("bedrock_mantle/")
+        resolved_model: Final = get_bedrock_base_model(model_to_use.removeprefix("bedrock_mantle/"))
 
-        runtime: Final = await self._count_with(
-            self._runtime_handler, RUNTIME_TOKENIZER_TYPE, request_data, litellm_params, resolved_model
+        primary: Final = await self._count_with(
+            self._mantle_handler if use_mantle else self._runtime_handler,
+            MANTLE_TOKENIZER_TYPE if use_mantle else RUNTIME_TOKENIZER_TYPE,
+            request_data,
+            litellm_params,
+            resolved_model,
         )
         outcome: Final = (
             await self._count_with(
                 self._mantle_handler, MANTLE_TOKENIZER_TYPE, request_data, litellm_params, resolved_model
             )
-            if _runtime_rejected_claude_model(runtime, resolved_model)
-            else runtime
+            if not use_mantle and _runtime_rejected_claude_model(primary, resolved_model)
+            else primary
         )
         match outcome:
             case CountedTokens():

litellm/proxy/common_utils/token_count_settings.py

@@ -9,7 +9,7 @@ import litellm
 from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping
 from litellm.litellm_core_utils.core_helpers import as_str_mapping
 from litellm.proxy._types import ProxyException
-from litellm.types.token_count import TokenCountSourceFields, TokenCountTokenizer
+from litellm.types.token_count import TokenCountProvider, TokenCountSourceFields, TokenCountTokenizer
 from litellm.types.utils import TokenCountResponse
 
 CANNOT_COUNT_MODEL_STATUSES: Final = frozenset({400, 404})
@@ -46,6 +46,7 @@ def _with_hint(message: str, hint: str) -> str:
 class TokenCountModel:
     model: str
     region: str | None
+    litellm_provider: TokenCountProvider | None = None
 
 
 @dataclass(frozen=True, slots=True)
@@ -100,9 +101,27 @@ def _non_blank(value: object) -> str | None:
     return value if isinstance(value, str) and value.strip() else None
 
 
+def _count_provider(value: object) -> TokenCountProvider | None:
+    match value:
+        case "bedrock":
+            return "bedrock"
+        case "bedrock_mantle":
+            return "bedrock_mantle"
+        case _:
+            return None
+
+
 def _count_model(value: Mapping[str, object]) -> TokenCountModel | None:
     model: Final = _non_blank(value.get("model"))
-    return None if model is None else TokenCountModel(model=model, region=_non_blank(value.get("region")))
+    return (
+        None
+        if model is None
+        else TokenCountModel(
+            model=model,
+            region=_non_blank(value.get("region")),
+            litellm_provider=_count_provider(value.get("litellm_provider")),
+        )
+    )
 
 
 def _geography(region: str) -> str:
@@ -112,7 +131,10 @@ def _geography(region: str) -> str:
 def _within_deployment_geography(
     count_model: TokenCountModel | None, deployment_region: Callable[[], str]
 ) -> TokenCountModel | None:
-    if count_model is None or count_model.region is None or count_model.model.startswith("global."):
+    if count_model is None or count_model.region is None:
+        return count_model
+    model: Final = count_model.model.removeprefix("bedrock/").removeprefix("bedrock_mantle/")
+    if model.startswith("global."):
         return count_model
     if _geography(deployment_region()) == _geography(count_model.region):
         return count_model

litellm/proxy/proxy_server.py

@@ -14851,9 +14851,14 @@ async def _provider_token_count(
     Returns None when the count should be estimated locally, or the provider's failure when the framed local
     estimate stands in for it and that failure is the answer if the estimate cannot read the request."""
     count_model: Final = settings.count_model
-    counted: Final = await (
-        count_with(model_to_use, None) if count_model is None else count_with(count_model.model, count_model.region)
+    counting_model: Final = (
+        model_to_use
+        if count_model is None
+        else count_model.model
+        if count_model.litellm_provider is None
+        else f"{count_model.litellm_provider}/{count_model.model.removeprefix('bedrock/').removeprefix('bedrock_mantle/')}"
     )
+    counted: Final = await count_with(counting_model, None if count_model is None else count_model.region)
     if counted is None:
         return None
     if isinstance(counted, TokenCountResponse):

litellm/types/token_count.py

@@ -4,11 +4,13 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
 
 TokenCountSource: TypeAlias = Literal["provider", "token_count_model", "local_estimate"]
 TokenCountTokenizer: TypeAlias = Literal["kimi_k3"]
+TokenCountProvider: TypeAlias = Literal["bedrock", "bedrock_mantle"]
 
 
 class TokenCountModelInfo(TypedDict):
     model: ReadOnly[str]
     region: ReadOnly[NotRequired[str]]
+    litellm_provider: ReadOnly[NotRequired[TokenCountProvider]]
 
 
 class TokenCountSourceFields(TypedDict, total=False):

model_prices_and_context_window.schema.json

@@ -1229,9 +1229,16 @@
             "region": {
               "type": "string",
               "pattern": "\\S"
+            },
+            "litellm_provider": {
+              "type": "string",
+              "enum": [
+                "bedrock",
+                "bedrock_mantle"
+              ]
             }
           },
-          "description": "Model that counts tokens for this one, and the AWS region its count API runs in."
+          "description": "Model that counts tokens for this one, its counting provider, and the AWS region its count API runs in."
         },
         "token_count_tokenizer": {
           "type": "string",

tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py

@@ -174,3 +174,50 @@ async def test_claude_runtime_404_retries_with_system_and_tools_on_mantle() -> N
         "system": "You are a terse assistant.",
         "tools": tools,
     }
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("model", "host", "tokenizer", "expected_tokens"),
+    (
+        ("bedrock_mantle/anthropic.claude-opus-5-5", MANTLE_HOST, "bedrock_mantle_api", 271),
+        ("bedrock/anthropic.claude-3-haiku-20240307-v1:0", RUNTIME_HOST, "bedrock_api", 193),
+    ),
+)
+async def test_qualified_counting_model_uses_its_endpoint_without_a_failed_probe(
+    model: str, host: str, tokenizer: str, expected_tokens: int
+) -> None:
+    bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 193}), mantle=(200, {"input_tokens": 271}))
+
+    result: Final = await _counter(bedrock).count_tokens(
+        model_to_use=model,
+        messages=MESSAGES,
+        contents=None,
+        deployment=DEPLOYMENT,
+        request_model="served-model",
+    )
+
+    assert result is not None
+    assert result.total_tokens == expected_tokens
+    assert result.tokenizer_type == tokenizer
+    assert bedrock.posted_hosts() == (host,)
+    if host == MANTLE_HOST:
+        assert _JSON_BODY.validate_json(bedrock.requests[0].content)["model"] == model.split("/", 1)[1]
+
+
+@pytest.mark.asyncio
+async def test_explicit_mantle_failure_does_not_probe_runtime() -> None:
+    bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 193}), mantle=(404, UNSUPPORTED))
+
+    result: Final = await _counter(bedrock).count_tokens(
+        model_to_use="bedrock_mantle/anthropic.claude-opus-5-5",
+        messages=MESSAGES,
+        contents=None,
+        deployment=DEPLOYMENT,
+    )
+
+    assert result is not None
+    assert result.error is True
+    assert result.status_code == 404
+    assert result.tokenizer_type == "bedrock_mantle_api"
+    assert bedrock.posted_hosts() == (MANTLE_HOST,)

tests/unit/proxy/common_utils/test_token_count_settings.py

@@ -4,6 +4,7 @@ from typing import Final
 
 import pytest
 
+from litellm.llms.bedrock.common_utils import get_bedrock_base_model
 from litellm.proxy.common_utils.token_count_settings import (
     ProviderCountFailure,
     TokenCountModel,
@@ -33,6 +34,13 @@ def _in_us_west_2() -> str:
             {"token_count_framing": 5, "token_count_model": {"model": "custom-model", "region": "ap-southeast-4"}},
             TokenCountSettings(framing=5, count_model=TokenCountModel(model="custom-model", region="ap-southeast-4")),
         ),
+        (
+            {"token_count_model": {"model": "custom-model", "litellm_provider": "bedrock_mantle"}},
+            TokenCountSettings(
+                framing=97,
+                count_model=TokenCountModel(model="custom-model", region=None, litellm_provider="bedrock_mantle"),
+            ),
+        ),
         (
             {"token_count_model": {"model": "custom-model"}},
             TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)),
@@ -94,13 +102,15 @@ def test_shipped_count_region_never_takes_the_prompt_out_of_the_deployment_geogr
     assert settings.count_model == TokenCountModel(model="anthropic.claude-opus-5-5", region=expected_region)
 
 
-def test_shipped_global_count_model_keeps_its_region_for_a_deployment_in_any_geography() -> None:
+@pytest.mark.parametrize("prefix", ("", "bedrock/", "bedrock_mantle/"))
+def test_shipped_global_count_model_keeps_its_region_for_a_deployment_in_any_geography(prefix: str) -> None:
+    model: Final = f"{prefix}global.anthropic.claude-opus-5-5"
     settings: Final = resolve_token_count_settings(
         configured={},
-        default={"token_count_model": {"model": "global.anthropic.claude-opus-5-5", "region": "us-east-1"}},
+        default={"token_count_model": {"model": model, "region": "us-east-1"}},
         deployment_region=lambda: "eu-west-1",
     )
-    assert settings.count_model == TokenCountModel(model="global.anthropic.claude-opus-5-5", region="us-east-1")
+    assert settings.count_model == TokenCountModel(model=model, region="us-east-1")
 
 
 def test_shipped_australian_count_region_stays_with_an_australian_deployment() -> None:
@@ -157,21 +167,24 @@ def test_count_source_fields_name_the_counting_model_only_when_token_count_model
     assert count_source_fields(_counted(source)) == expected
 
 
-def test_every_shipped_token_count_model_is_another_mapped_model_of_the_same_provider_in_a_named_region() -> None:
+def test_every_shipped_token_count_model_resolves_to_a_mapped_aws_model() -> None:
     cost_map: Final = json.loads(_COST_MAP.read_text())
+    aws_models: Final = frozenset(
+        get_bedrock_base_model(name)
+        for name, entry in cost_map.items()
+        if isinstance(entry, dict) and entry.get("litellm_provider") in ("bedrock", "bedrock_converse")
+    )
     count_models: Final = {
         name: entry["token_count_model"]
         for name, entry in cost_map.items()
         if isinstance(entry, dict) and "token_count_model" in entry
     }
     assert count_models
-    assert {
-        name: count_model
-        for name, count_model in count_models.items()
-        if cost_map.get(count_model["model"], {}).get("litellm_provider") != cost_map[name]["litellm_provider"]
-        or count_model["model"] == name
-        or not count_model.get("region")
-    } == {}
+    for name, count_model in count_models.items():
+        model: Final = count_model["model"]
+        assert get_bedrock_base_model(model) in aws_models
+        assert cost_map[name]["litellm_provider"] in ("bedrock", "bedrock_converse")
+        assert "region" not in count_model or count_model["region"].strip()
 
 
 _UPSTREAM_404: Final = {

tests/unit/proxy/test_proxy_token_counter.py

@@ -67,6 +67,7 @@ from litellm.proxy.common_utils.token_count_settings import (
     TokenCountModel,
     TokenCountSettings,
     cost_map_entry,
+    resolve_token_count_settings,
 )
 from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request
 from litellm.proxy.proxy_server import (
@@ -1636,6 +1637,65 @@ async def test_model_without_token_count_model_counts_itself_in_the_deployment_r
     assert ([call[0] for call in counter.calls], counter.regions) == (["us.anthropic.claude-sonnet-5-5"], ["us-west-2"])
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("provider", "expected_host", "expected_tokens"),
+    (
+        ("bedrock", "bedrock-runtime.us-east-1.amazonaws.com", 193),
+        ("bedrock_mantle", "bedrock-mantle.us-east-1.api.aws", 271),
+    ),
+)
+async def test_token_count_model_provider_selects_the_endpoint_and_preserves_the_public_model(
+    provider: str, expected_host: str, expected_tokens: int, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
+    model: Final = "anthropic.claude-opus-5-5"
+
+    def peer(request: httpx.Request) -> httpx.Response:
+        assert request.url.host == expected_host
+        if provider == "bedrock_mantle":
+            assert json.loads(request.content)["model"] == model
+            return httpx.Response(200, json={"input_tokens": expected_tokens})
+        assert request.url.path == f"/model/{model}/count-tokens"
+        return httpx.Response(200, json={"inputTokens": expected_tokens})
+
+    settings: Final = resolve_token_count_settings(
+        configured={},
+        default={"token_count_model": {"model": model, "litellm_provider": provider, "region": "us-east-1"}},
+        deployment_region=lambda: "us-west-2",
+    )
+    counter: Final = BedrockTokenCounter(client=AsyncHTTPHandler(transport=httpx.MockTransport(peer)))
+    result: Final = await _provider_token_count(
+        count_with=partial(
+            _try_provider_token_count,
+            counter,
+            "bedrock",
+            messages=_COUNT_REQUEST.messages,
+            contents=_COUNT_REQUEST.contents,
+            request_model=_COUNT_REQUEST.model,
+            deployment={
+                "litellm_params": {
+                    "aws_access_key_id": "AKIATESTACCESSKEY",
+                    "aws_secret_access_key": "test-secret",
+                    "aws_region_name": "us-west-2",
+                }
+            },
+            tools=_COUNT_REQUEST.tools,
+            system=_COUNT_REQUEST.system,
+        ),
+        model_to_use="us.anthropic.claude-sonnet-5-5",
+        settings=settings,
+        local_estimate_allowed=False,
+    )
+
+    assert isinstance(result, TokenCountResponse)
+    assert (result.total_tokens, result.model_used, result.count_source) == (
+        expected_tokens,
+        model,
+        "token_count_model",
+    )
+
+
 @pytest.mark.parametrize(
     ("count_model", "region"),
     [
@@ -2658,7 +2718,7 @@ async def test_uncountable_model_keeps_the_provider_error_when_local_counting_is
 @pytest.mark.parametrize(
     ("model_info", "counted_with"),
     [
-        ({}, ("us.anthropic.claude-opus-5-5", "us-east-1", 447)),
+        ({}, ("bedrock_mantle/us.anthropic.claude-opus-5-5", "us-east-1", 447)),
         (
             {"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "ap-southeast-4"}},
             ("us.anthropic.claude-opus-5-5", "ap-southeast-4", 447),
@@ -2674,6 +2734,7 @@ async def test_configured_token_count_model_overrides_the_shipped_one_and_an_emp
         {
             "us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", total_tokens=300),
             "us.anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447),
+            "bedrock_mantle/us.anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447),
         }
     )
     monkeypatch.setattr(
@@ -2712,7 +2773,7 @@ async def test_configured_token_count_model_overrides_the_shipped_one_and_an_emp
 async def test_shipped_token_count_model_keeps_the_prompt_in_the_deployment_geography(
     monkeypatch: pytest.MonkeyPatch, aws_region_name: str, counted_in: str
 ):
-    counter: Final = _ScriptedCounter({"anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)})
+    counter: Final = _ScriptedCounter({"bedrock_mantle/anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)})
     monkeypatch.setattr(
         litellm.proxy.proxy_server,
         "llm_router",
@@ -2736,7 +2797,7 @@ async def test_shipped_token_count_model_keeps_the_prompt_in_the_deployment_geog
         request=TokenCountRequest(model="sonnet", messages=[{"role": "user", "content": "hi"}]), call_endpoint=True
     )
 
-    assert (counter.calls[0][0], counter.regions[0]) == ("anthropic.claude-opus-5-5", counted_in)
+    assert (counter.calls[0][0], counter.regions[0]) == ("bedrock_mantle/anthropic.claude-opus-5-5", counted_in)
 
 
 class _SlowVault:
@@ -2756,7 +2817,7 @@ async def _longest_stall(done: asyncio.Event) -> float:
 async def test_region_lookup_for_a_shipped_count_model_leaves_the_event_loop_free(monkeypatch: pytest.MonkeyPatch):
     from litellm.proxy import proxy_server
 
-    counter: Final = _ScriptedCounter({"anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)})
+    counter: Final = _ScriptedCounter({"bedrock_mantle/anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)})
 
     def bedrock_counter(deployment: object, model: str) -> tuple[_ScriptedCounter, str, str]:
         return counter, model, "bedrock"