Merge pull request #34446 from BerriAI/litellm_fix_mantle_missing_usage

fix(bedrock-mantle): backfill usage on non-streaming /v1/messages responses

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,

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