From a279510e352ffc7736411fe9f22d6216458f445d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 11:25:03 -0700 Subject: [PATCH] fix(bedrock): refuse Mantle-only token counts on a pinned deployment with a clear 400 A token count model on bedrock_mantle used to fall back to the pinned bedrock-runtime endpoint with a model id only Mantle hosts, which failed every time. It now answers a 400 before any call, naming BEDROCK_MANTLE_API_BASE, and the proxy skips its region hint and framed local estimate for it. --- .../count_tokens/bedrock_token_counter.py | 21 +++++- .../bedrock/count_tokens/mantle_handler.py | 5 +- .../test_bedrock_token_counter.py | 38 ++++++---- tests/unit/proxy/test_proxy_token_counter.py | 69 +++++++++++++++++++ 4 files changed, 114 insertions(+), 19 deletions(-) diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 38512d99631..58f45a50fa1 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -19,6 +19,7 @@ from litellm.llms.bedrock.base_aws_llm import run_aws_signing from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.bedrock.count_tokens.mantle_handler import ( + NO_MANTLE_ENDPOINT_MESSAGE, BedrockMantleCountTokensHandler, deployment_mantle_endpoint, ) @@ -82,6 +83,19 @@ async def _has_mantle_endpoint(litellm_params: Mapping[str, object]) -> bool: return await run_aws_signing(deployment_mantle_endpoint, litellm_params) is not None +def _no_mantle_endpoint_refusal(request_model: str, model_used: str) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_used, + tokenizer_type=MANTLE_TOKENIZER_TYPE, + error=True, + error_message=NO_MANTLE_ENDPOINT_MESSAGE, + status_code=400, + rejected_before_provider_call=True, + ) + + def _request_data( model: str, messages: Sequence[Mapping[str, object]], tools: Sequence[Mapping[str, object]] | None, system: object ) -> Mapping[str, object]: @@ -159,8 +173,9 @@ class BedrockTokenCounter(BaseTokenCounter): counted_model: Final = model_to_use.removeprefix("bedrock_mantle/") resolved_model: Final = get_bedrock_base_model(counted_model) litellm_params: Final = (deployment or {}).get("litellm_params", {}) - counts_on_mantle: Final = use_mantle and await _has_mantle_endpoint(litellm_params) - primary_tokenizer: Final = MANTLE_TOKENIZER_TYPE if counts_on_mantle else RUNTIME_TOKENIZER_TYPE + if use_mantle and not await _has_mantle_endpoint(litellm_params): + return _no_mantle_endpoint_refusal(request_model, model_to_use) + primary_tokenizer: Final = MANTLE_TOKENIZER_TYPE if use_mantle else RUNTIME_TOKENIZER_TYPE translated: Final = ( await anthropic_count_fields( counted_model, @@ -187,7 +202,7 @@ class BedrockTokenCounter(BaseTokenCounter): await self._count_with( self._mantle_handler, primary_tokenizer, mantle_request, litellm_params, resolved_model ) - if counts_on_mantle + if use_mantle else await self._count_with( self._runtime_handler, primary_tokenizer, runtime_request, litellm_params, resolved_model ) diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py index ad82604815c..cdc090e5951 100644 --- a/litellm/llms/bedrock/count_tokens/mantle_handler.py +++ b/litellm/llms/bedrock/count_tokens/mantle_handler.py @@ -19,8 +19,9 @@ from litellm.secret_managers.main import get_secret_str MANTLE_COUNT_TOKENS_SUFFIX: Final = "/count_tokens" MANTLE_ANTHROPIC_VERSION: Final = "2023-06-01" NO_MANTLE_ENDPOINT_MESSAGE: Final = ( - "This deployment pins its own bedrock-runtime endpoint, so its prompt is not counted on the public " - "bedrock-mantle host. Set BEDROCK_MANTLE_API_BASE to the bedrock-mantle endpoint it may use." + "This model counts tokens on bedrock-mantle, but the deployment pins its own bedrock-runtime endpoint " + "through api_base, aws_bedrock_runtime_endpoint, or AWS_BEDROCK_RUNTIME_ENDPOINT, so its prompt is not sent " + "to the public bedrock-mantle host. Set BEDROCK_MANTLE_API_BASE to the bedrock-mantle endpoint it may use." ) diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py index 8f90a709ac4..0f37fcb6356 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py @@ -8,6 +8,7 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter +from litellm.llms.bedrock.count_tokens.mantle_handler import NO_MANTLE_ENDPOINT_MESSAGE from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault @@ -502,20 +503,23 @@ async def test_pinned_runtime_endpoint_retries_on_the_mantle_host_the_operator_n @pytest.mark.asyncio @pytest.mark.parametrize( - ("runtime", "expected"), + ("endpoint", "environment"), ( - ((200, {"inputTokens": 193}), (False, 193, None)), - ((400, UNSUPPORTED), (True, 0, 400)), + ({"api_base": PRIVATE_RUNTIME}, {}), + ({"aws_bedrock_runtime_endpoint": PRIVATE_RUNTIME}, {}), + ({}, {"AWS_BEDROCK_RUNTIME_ENDPOINT": PRIVATE_RUNTIME}), ), - ids=("runtime_counts", "runtime_rejects"), + ids=("api_base", "aws_bedrock_runtime_endpoint", "AWS_BEDROCK_RUNTIME_ENDPOINT"), ) -async def test_mantle_counting_model_on_a_pinned_runtime_deployment_counts_on_that_endpoint_only( - monkeypatch: pytest.MonkeyPatch, - runtime: tuple[int, Mapping[str, object]], - expected: tuple[bool, int, int | None], +async def test_mantle_counting_model_on_a_pinned_runtime_deployment_is_refused_before_any_call( + monkeypatch: pytest.MonkeyPatch, endpoint: Mapping[str, str], environment: Mapping[str, str] ) -> None: - vault: Final = host_in_loop_probing_vault(monkeypatch, "BEDROCK_MANTLE_API_BASE", "AWS_BEDROCK_RUNTIME_ENDPOINT") - bedrock: Final = _Bedrock(runtime=runtime, mantle=(200, {"input_tokens": 271}), runtime_host=PRIVATE_RUNTIME_HOST) + for name, value in environment.items(): + monkeypatch.setenv(name, value) + vault: Final = host_in_loop_probing_vault(monkeypatch, "BEDROCK_MANTLE_API_BASE") + bedrock: Final = _Bedrock( + runtime=(200, {"inputTokens": 193}), mantle=(200, {"input_tokens": 271}), runtime_host=PRIVATE_RUNTIME_HOST + ) result: Final = await _counter(bedrock).count_tokens( model_to_use="bedrock_mantle/anthropic.claude-opus-5-5", @@ -525,16 +529,22 @@ async def test_mantle_counting_model_on_a_pinned_runtime_deployment_counts_on_th "litellm_params": { **DEPLOYMENT["litellm_params"], "model": "bedrock/us.anthropic.claude-opus-5-5", - "api_base": PRIVATE_RUNTIME, + **endpoint, } }, request_model="claude-opus-5-5", ) assert result is not None - assert (result.error, result.total_tokens, result.status_code) == expected - assert result.tokenizer_type == "bedrock_api" - assert bedrock.posted_urls() == (f"{PRIVATE_RUNTIME}/model/anthropic.claude-opus-5-5/count-tokens",) + assert (result.error, result.status_code, result.total_tokens, result.tokenizer_type) == ( + True, + 400, + 0, + "bedrock_mantle_api", + ) + assert result.error_message == NO_MANTLE_ENDPOINT_MESSAGE + assert result.rejected_before_provider_call + assert bedrock.requests == () assert vault.read.is_set() assert not vault.read_on_loop.is_set() diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index b21730c64f6..153bf11c590 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -1762,6 +1762,75 @@ async def test_token_count_model_provider_selects_the_endpoint_and_preserves_the ) +class _RejectingBedrockPeer: + def __init__(self) -> None: + self.requests: tuple[httpx.Request, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.requests = (*self.requests, request) + return httpx.Response(400, json={"message": "The provided model identifier is invalid."}) + + +_MANTLE_COUNT_MODEL: Final = MappingProxyType( + { + "token_count_framing": 10, + "token_count_model": { + "model": "anthropic.claude-opus-5-5", + "litellm_provider": "bedrock_mantle", + "region": "us-east-1", + }, + } +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("configured", "default"), + ((_MANTLE_COUNT_MODEL, {}), ({}, _MANTLE_COUNT_MODEL)), + ids=("configured_token_count_model", "cost_map_token_count_model"), +) +async def test_mantle_token_count_model_on_a_pinned_runtime_deployment_answers_a_400_naming_the_fix( + configured: Mapping[str, object], default: Mapping[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + for name in ("AWS_BEARER_TOKEN_BEDROCK", "BEDROCK_MANTLE_API_BASE", "AWS_BEDROCK_RUNTIME_ENDPOINT"): + monkeypatch.delenv(name, raising=False) + peer: Final = _RejectingBedrockPeer() + counter: Final = BedrockTokenCounter(client=AsyncHTTPHandler(transport=httpx.MockTransport(peer))) + with pytest.raises(ProxyException) as raised: + 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": { + "model": "bedrock/us.anthropic.claude-opus-5-5", + "api_base": "https://vpce-0a1b2c3d.bedrock-runtime.us-east-1.vpce.amazonaws.com", + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-east-1", + } + }, + tools=_COUNT_REQUEST.tools, + system=_COUNT_REQUEST.system, + ), + model_to_use="us.anthropic.claude-opus-5-5", + request_model=_COUNT_REQUEST.model, + settings=resolve_token_count_settings( + configured=configured, default=default, deployment_region=lambda: "us-east-1" + ), + local_estimate_allowed=True, + ) + + assert raised.value.code == "400" + assert "BEDROCK_MANTLE_API_BASE" in raised.value.message + assert "If that region is not enabled" not in raised.value.message + assert peer.requests == () + + @pytest.mark.parametrize( ("count_model", "region"), [