From ee9f0db1cff3d3694a9e246351be785366c0058a Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 23 Jul 2026 23:37:40 +0000 Subject: [PATCH] fix(bedrock-mantle): backfill usage on non-streaming /v1/messages responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock/messages/mantle_transformation.py | 23 +++++ .../test_litellm/llms/bedrock/test_mantle.py | 97 +++++++++++++++++++ 2 files changed, 120 insertions(+) diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py index da7b8697a6b..65a2ab3b9a7 100644 --- a/litellm/llms/bedrock/messages/mantle_transformation.py +++ b/litellm/llms/bedrock/messages/mantle_transformation.py @@ -17,6 +17,10 @@ from litellm.llms.bedrock.common_utils import build_mantle_messages_url from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeMessagesConfig, ) +from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, + AnthropicUsage, +) from litellm.types.router import GenericLiteLLMParams if TYPE_CHECKING: @@ -103,6 +107,25 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): ) return {**request, "model": model_id, **stream_fields} + def transform_anthropic_messages_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + ) -> AnthropicMessagesResponse: + response = super().transform_anthropic_messages_response( + model=model, + raw_response=raw_response, + logging_obj=logging_obj, + ) + existing_usage: AnthropicUsage = response.get("usage") or AnthropicUsage() + normalized_usage: AnthropicUsage = { + "input_tokens": 0, + "output_tokens": 0, + **existing_usage, + } + return {**response, "usage": normalized_usage} + def get_async_streaming_response_iterator( self, model: str, diff --git a/tests/test_litellm/llms/bedrock/test_mantle.py b/tests/test_litellm/llms/bedrock/test_mantle.py index c5ce0d5aba7..d34517f61f6 100644 --- a/tests/test_litellm/llms/bedrock/test_mantle.py +++ b/tests/test_litellm/llms/bedrock/test_mantle.py @@ -399,6 +399,103 @@ async def test_mantle_anthropic_messages_sends_workspace_header_and_clean_body() assert "aws_bedrock_project_id" not in requests[0]["body"] +def _usageless_anthropic_response(url: str) -> httpx.Response: + return httpx.Response( + status_code=200, + json={ + "id": "msg_classifier", + "type": "message", + "role": "assistant", + "model": "anthropic.claude-opus-4-8", + "content": [{"type": "text", "text": "safe"}], + "stop_reason": "end_turn", + "stop_sequence": None, + }, + request=httpx.Request("POST", url), + ) + + +@pytest.mark.asyncio +async def test_mantle_anthropic_messages_backfills_missing_usage(): + """ + Regression for LIT-4758: a Mantle non-streaming response with no `usage` + object must not reach the client usage-less, or Claude Code's auto-mode + classifier crashes on `usage.input_tokens`. + """ + import litellm + + async def mock_post(self, url, data=None, headers=None, **kwargs): + return _usageless_anthropic_response(str(url)) + + try: + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new=mock_post, + ): + response = await litellm.anthropic_messages( + model="bedrock/mantle/anthropic.claude-opus-4-8", + messages=[{"role": "user", "content": "is `Bash(ls)` safe?"}], + max_tokens=10, + aws_access_key_id="fake-key", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + ) + finally: + await litellm.close_litellm_async_clients() + + assert response["usage"]["input_tokens"] == 0 + assert response["usage"]["output_tokens"] == 0 + + +@pytest.mark.asyncio +async def test_mantle_anthropic_messages_preserves_upstream_usage(): + """Backfill must not clobber a usage object the upstream did return.""" + import litellm + + def _response_with_usage(url: str) -> httpx.Response: + return httpx.Response( + status_code=200, + json={ + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "anthropic.claude-opus-4-8", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": 42, + "output_tokens": 7, + "cache_read_input_tokens": 5, + }, + }, + request=httpx.Request("POST", url), + ) + + async def mock_post(self, url, data=None, headers=None, **kwargs): + return _response_with_usage(str(url)) + + try: + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new=mock_post, + ): + response = await litellm.anthropic_messages( + model="bedrock/mantle/anthropic.claude-opus-4-8", + messages=[{"role": "user", "content": "hello"}], + max_tokens=10, + aws_access_key_id="fake-key", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + ) + finally: + await litellm.close_litellm_async_clients() + + assert response["usage"]["input_tokens"] == 42 + assert response["usage"]["output_tokens"] == 7 + assert response["usage"]["cache_read_input_tokens"] == 5 + + @pytest.mark.asyncio async def test_mantle_anthropic_messages_routes_to_vpc_api_base(): import litellm