fix(count_tokens): keep a local conversion 400 as a 400 on the Claude count routes

The anthropic, azure_ai and Bedrock Mantle count handlers re-raise a BadRequestError raised before the provider call, and their token counters mark it rejected_before_provider_call, so a framing deployment answers the 400 instead of an estimate that leaves the attachment out, and Mantle no longer turns it into a 500
This commit is contained in:
mateo-berri 2026-10-30 04:57:00 -07:00
parent 3f3f9bd451
commit f18056a85b
9 changed files with 92 additions and 35 deletions

View file

@ -116,8 +116,8 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
except AnthropicError:
# Re-raise Anthropic exceptions as-is
raise
except litellm.BadRequestError as e:
raise AnthropicError(status_code=400, message=e.message) from e
except litellm.BadRequestError:
raise
except httpx.HTTPStatusError as e:
# HTTP errors - preserve the actual status code
verbose_logger.error("HTTP error in CountTokens handler: %s", e)

View file

@ -143,6 +143,17 @@ class AnthropicTokenCounter(BaseTokenCounter):
tokenizer_type="anthropic_api",
original_response=result,
)
except litellm.BadRequestError as e:
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="anthropic_api",
error=True,
error_message=e.message,
status_code=400,
rejected_before_provider_call=True,
)
except (AnthropicError, AuthenticationError) as e:
verbose_logger.warning("Anthropic CountTokens error: status=%s, message=%s", e.status_code, e.message)
return TokenCountResponse(

View file

@ -128,8 +128,8 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
except AnthropicError:
# Re-raise Anthropic exceptions as-is
raise
except litellm.BadRequestError as e:
raise AnthropicError(status_code=400, message=e.message) from e
except litellm.BadRequestError:
raise
except httpx.HTTPStatusError as e:
# HTTP errors - preserve the actual status code
verbose_logger.error("HTTP error in CountTokens handler: %s", e)

View file

@ -147,6 +147,17 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
"original_response": result,
}
)
except litellm.BadRequestError as e:
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="azure_ai_anthropic_api",
error=True,
error_message=e.message,
status_code=400,
rejected_before_provider_call=True,
)
except AnthropicError as e:
verbose_logger.warning(
"Azure AI Anthropic CountTokens API error: status=%s, message=%s", e.status_code, e.message

View file

@ -49,6 +49,7 @@ class CountTokensFailure:
status_code: int
message: str
tokenizer_type: str
rejected_before_provider_call: bool = False
CountTokensOutcome = CountedTokens | CountTokensFailure
@ -120,6 +121,10 @@ class BedrockTokenCounter(BaseTokenCounter):
client=self._client,
)
reply: Final = _CountTokensReply.model_validate(result)
except litellm.BadRequestError as e:
return CountTokensFailure(
status_code=400, message=e.message, tokenizer_type=tokenizer_type, rejected_before_provider_call=True
)
except BedrockError as e:
verbose_logger.debug(
"%s CountTokens API error: status=%s, message=%s", tokenizer_type, e.status_code, e.message
@ -214,6 +219,7 @@ class BedrockTokenCounter(BaseTokenCounter):
error=True,
error_message=outcome.message,
status_code=outcome.status_code,
rejected_before_provider_call=outcome.rejected_before_provider_call,
)
case _:
assert_never(outcome)

View file

@ -99,7 +99,7 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM):
response=response,
)
return _COUNT_RESPONSE.validate_json(response.content)
except BedrockError:
except (BedrockError, litellm.BadRequestError):
raise
except httpx.HTTPStatusError as e:
raise BedrockError(

View file

@ -14,7 +14,6 @@ from pydantic import JsonValue, TypeAdapter
import litellm
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.anthropic.common_utils import AnthropicError
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
from litellm.llms.anthropic.count_tokens.token_counter import AnthropicTokenCounter
from litellm.llms.anthropic.count_tokens.transformation import (
@ -526,35 +525,6 @@ def test_count_drops_a_leading_system_message_without_text() -> None:
assert result == {"model": "claude-opus-5-5", "messages": [{"role": "user", "content": "hi"}]}
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport_clients")
@pytest.mark.parametrize("handler_type", (AnthropicCountTokensHandler, AzureAIAnthropicCountTokensHandler))
async def test_handler_answers_400_for_a_text_attachment_it_cannot_decode(
handler_type: type[AnthropicCountTokensHandler] | type[AzureAIAnthropicCountTokensHandler],
) -> None:
text_attachment: Final = "data:text/plain;charset=bogus-cs;base64," + b64encode(b"hello").decode()
messages: Final[list[dict[str, JsonValue]]] = [
{"role": "user", "content": [{"type": "image_url", "image_url": {"url": text_attachment}}]}
]
api_base: Final = "https://gateway.example"
handler: Final = handler_type()
with respx.mock:
route: Final = respx.post(f"{api_base}/v1/messages/count_tokens").mock(
return_value=httpx.Response(200, json={"input_tokens": 7})
)
with pytest.raises(AnthropicError) as exc_info:
await (
handler.handle_count_tokens_request(
model="claude-opus-5-5", messages=messages, api_base=api_base, auth_header={"x-api-key": "test-key"}
) if isinstance(handler, AnthropicCountTokensHandler) else handler.handle_count_tokens_request(
model="claude-opus-5-5", messages=messages, api_base=api_base, api_key="test-key"
)
)
assert exc_info.value.status_code == 400
assert "bogus-cs" in exc_info.value.message
assert not route.called
@pytest.mark.asyncio
async def test_handler_sends_the_leading_system_run_as_system_not_as_a_message(httpx_transport_clients):
"""The wire body is what the provider judges: ``system`` carries the prompt and no message has

View file

@ -293,6 +293,24 @@ async def test_claude_count_rejects_a_lone_responses_tool_without_calling_bedroc
assert bedrock.posted_hosts() == ()
@pytest.mark.asyncio
async def test_claude_count_rejects_a_text_attachment_it_cannot_decode_without_calling_bedrock() -> None:
bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 412}), mantle=(200, {"input_tokens": 2177}))
text_attachment: Final = "data:text/plain;charset=bogus-cs;base64," + base64.b64encode(b"hello").decode()
result: Final = await _counter(bedrock).count_tokens(
model_to_use="bedrock_mantle/anthropic.claude-opus-5-5",
messages=[{"role": "user", "content": [{"type": "image_url", "image_url": {"url": text_attachment}}]}],
contents=None,
deployment=DEPLOYMENT,
)
assert result is not None
assert (result.error, result.status_code, result.rejected_before_provider_call) == (True, 400, True)
assert "bogus-cs" in (result.error_message or "")
assert bedrock.posted_hosts() == ()
_CLAUDE_REPLY: Final = {
"id": "msg_1",
"type": "message",

View file

@ -2040,6 +2040,47 @@ async def test_framing_keeps_the_400_for_blocks_the_provider_counter_cannot_conv
assert "search_result" in raised.value.message
@pytest.mark.parametrize(
"litellm_params",
[
{"model": "anthropic/claude-opus-5-5", "api_key": "sk-unused-framed"},
{
"model": "azure_ai/claude-opus-5-5",
"api_key": "sk-unused-framed",
"api_base": "https://framed.services.ai.azure.com/anthropic",
},
],
ids=["anthropic", "azure_ai"],
)
@pytest.mark.asyncio
async def test_framing_keeps_the_400_for_a_text_attachment_the_claude_counter_cannot_decode(
monkeypatch: pytest.MonkeyPatch, litellm_params: dict[str, str]
):
monkeypatch.setattr(
litellm.proxy.proxy_server,
"llm_router",
Router(
model_list=[
{"model_name": "framed", "litellm_params": litellm_params, "model_info": {"token_count_framing": 100}}
]
),
)
text_attachment: Final = "data:text/plain;charset=bogus-cs;base64," + base64.b64encode(b"hello").decode()
with pytest.raises(ProxyException) as raised:
await count_request_tokens(
request=TokenCountRequest(
model="framed",
messages=[{"role": "user", "content": [{"type": "image_url", "image_url": {"url": text_attachment}}]}],
),
call_endpoint=True,
messages_format="anthropic",
)
assert raised.value.code == "400"
assert "bogus-cs" in raised.value.message
@pytest.mark.asyncio
@pytest.mark.parametrize(
"unreadable_body",