fix(count_tokens): fall back to Mantle on Bedrock 404 (Bedrock part of #46430)

Backport of #46430 to stable/1.104.x.
Cherry-picked from 08a24d391a (main), Bedrock changes only.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/llms/bedrock/count_tokens/bedrock_token_counter.py

@@ -20,6 +20,7 @@ from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCo
 
 RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api"
 MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api"
+_RUNTIME_FALLBACK_STATUSES: Final = frozenset({400, 404})
 
 
 class _CountTokensReply(LiteLLMPydanticObjectBase):
@@ -46,14 +47,14 @@ CountTokensOutcome = CountedTokens | CountTokensFailure
 def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool:
     return (
         isinstance(outcome, CountTokensFailure)
-        and outcome.status_code == 400
+        and outcome.status_code in _RUNTIME_FALLBACK_STATUSES
         and is_mantle_claude_model(resolved_model)
     )
 
 
 class BedrockTokenCounter(BaseTokenCounter):
     """Token counter for AWS Bedrock: bedrock-runtime CountTokens first, and Anthropic's count_tokens
-    on bedrock-mantle for the Claude models bedrock-runtime answers 400 for"""
+    on bedrock-mantle for the Claude models bedrock-runtime answers 400 or 404 for"""
 
     def __init__(
         self,

tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py

@@ -451,6 +451,52 @@ def test_gemini_count_tokens_route_reports_the_runtime_total_for_a_model_bedrock
         assert mantle.drain() == ()
 
 
+def test_runtime_404_on_claude_falls_back_to_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None:
+    gateway: Final = counting_proxy.gateway
+    system: Final = f"{_SYSTEM} status=404"
+    with (
+        wire_server(_runtime) as runtime,
+        wire_server(_mantle(), port=mantle_port) as mantle,
+        gateway.scenario() as scenario,
+    ):
+        model: Final = _deployment(scenario, runtime.url)
+        response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS})
+        assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text
+        assert len(_runtime_count_targets(runtime)) == 1
+        assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),)
+
+
+def test_runtime_404_and_mantle_404_surface_provider_error(counting_proxy: OwnedProxy, mantle_port: int) -> None:
+    gateway: Final = counting_proxy.gateway
+    system: Final = f"{_SYSTEM} status=404"
+    with (
+        wire_server(_runtime) as runtime,
+        wire_server(_mantle(_rejecting(404)), port=mantle_port) as mantle,
+        gateway.scenario() as scenario,
+    ):
+        model: Final = _deployment(scenario, runtime.url)
+        response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS})
+        assert response.status_code == 404, response.text
+        assert _REJECTION in response.text
+        assert len(_runtime_count_targets(runtime)) == 1
+        assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),)
+
+
+def test_runtime_404_on_nova_does_not_call_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None:
+    gateway: Final = counting_proxy.gateway
+    with (
+        wire_server(_runtime) as runtime,
+        wire_server(_mantle(), port=mantle_port) as mantle,
+        gateway.scenario() as scenario,
+    ):
+        model: Final = _deployment(scenario, runtime.url, _NOVA)
+        body: Final = {**_bare(model), "messages": [{"role": "user", "content": "status=404"}]}
+        response: Final = _count(gateway, body)
+        assert response.status_code == 404, response.text
+        assert len(_runtime_count_targets(runtime)) == 1
+        assert mantle.drain() == ()
+
+
 def test_responses_input_tokens_counts_through_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None:
     gateway: Final = counting_proxy.gateway
     with (

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

@@ -93,8 +93,10 @@ async def test_bedrock_runtime_count_is_kept_when_it_answers() -> None:
 @pytest.mark.parametrize(
     ("model_to_use", "runtime"),
     (
+        ("global.anthropic.claude-opus-4-8", (401, {"Message": "not authorized to perform: bedrock:CountTokens"})),
         ("global.anthropic.claude-opus-4-8", (403, {"Message": "not authorized to perform: bedrock:CountTokens"})),
         ("amazon.nova-pro-v1:0", (400, UNSUPPORTED)),
+        ("amazon.nova-pro-v1:0", (404, UNSUPPORTED)),
     ),
 )
 async def test_other_bedrock_runtime_failures_are_not_retried_on_mantle(
@@ -120,7 +122,7 @@ async def test_other_bedrock_runtime_failures_are_not_retried_on_mantle(
 @pytest.mark.asyncio
 async def test_mantle_failure_is_reported_with_its_status() -> None:
     bedrock: Final = _Bedrock(
-        runtime=(400, UNSUPPORTED),
+        runtime=(404, UNSUPPORTED),
         mantle=(404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}}),
     )
 
@@ -138,3 +140,37 @@ async def test_mantle_failure_is_reported_with_its_status() -> None:
     assert result.tokenizer_type == "bedrock_mantle_api"
     assert "does not exist" in (result.error_message or "")
     assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST)
+
+
+@pytest.mark.asyncio
+async def test_claude_runtime_404_retries_with_system_and_tools_on_mantle() -> None:
+    bedrock: Final = _Bedrock(runtime=(404, UNSUPPORTED), mantle=(200, {"input_tokens": 2177}))
+    tools: Final = [
+        {
+            "name": "lookup",
+            "description": "Find a record",
+            "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}},
+        }
+    ]
+
+    result: Final = await _counter(bedrock).count_tokens(
+        model_to_use="global.anthropic.claude-opus-4-8",
+        messages=MESSAGES,
+        contents=None,
+        deployment=DEPLOYMENT,
+        request_model="claude-opus-4-8",
+        system="You are a terse assistant.",
+        tools=tools,
+    )
+
+    assert result is not None
+    assert result.error is False
+    assert result.total_tokens == 2177
+    assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST)
+    mantle_body: Final = _JSON_BODY.validate_json(bedrock.requests[1].content)
+    assert mantle_body == {
+        "model": "anthropic.claude-opus-4-8",
+        "messages": MESSAGES,
+        "system": "You are a terse assistant.",
+        "tools": tools,
+    }