fix(streaming): preserve usage fields when start omits usage

Author: chenlichao

src/anthropic/lib/streaming/_beta_messages.py

@@ -8,10 +8,10 @@
 import httpx2 as httpx
 from pydantic import BaseModel
 
+from anthropic.types.beta.beta_usage import BetaUsage
 from anthropic.types.beta.beta_tool_use_block import BetaToolUseBlock
 from anthropic.types.beta.beta_mcp_tool_use_block import BetaMCPToolUseBlock
 from anthropic.types.beta.beta_server_tool_use_block import BetaServerToolUseBlock
-from anthropic.types.usage import Usage
 
 from ..._types import NotGiven, not_given
 from ..._utils import consume_sync_iterator, consume_async_iterator
@@ -562,10 +562,10 @@ def accumulate_event(
             # case the snapshot has no usage yet. Initialize it from the delta
             # so the final message still carries token counts, and tolerate
             # streams that never supply usage.
-            current_snapshot.usage = Usage(
-                input_tokens=event.usage.input_tokens or 0,
-                output_tokens=event.usage.output_tokens,
-            )
+            usage = event.usage.to_dict()
+            if event.usage.input_tokens is None:
+                usage["input_tokens"] = 0
+            current_snapshot.usage = construct_type(type_=BetaUsage, value=usage)
         else:
             current_snapshot.usage.output_tokens = event.usage.output_tokens
 

src/anthropic/lib/streaming/_messages.py

@@ -7,9 +7,9 @@
 import httpx2 as httpx
 from pydantic import BaseModel
 
+from anthropic.types.usage import Usage
 from anthropic.types.tool_use_block import ToolUseBlock
 from anthropic.types.server_tool_use_block import ServerToolUseBlock
-from anthropic.types.usage import Usage
 
 from ._types import (
     TextEvent,
@@ -526,10 +526,10 @@ def accumulate_event(
             # case the snapshot has no usage yet. Initialize it from the delta
             # so the final message still carries token counts, and tolerate
             # streams that never supply usage.
-            current_snapshot.usage = Usage(
-                input_tokens=event.usage.input_tokens or 0,
-                output_tokens=event.usage.output_tokens,
-            )
+            usage = event.usage.to_dict()
+            if event.usage.input_tokens is None:
+                usage["input_tokens"] = 0
+            current_snapshot.usage = construct_type(type_=Usage, value=usage)
         else:
             current_snapshot.usage.output_tokens = event.usage.output_tokens
 

tests/lib/streaming/fixtures/beta_usage_omitted_with_optional_fields_response.txt

@@ -0,0 +1,17 @@
+event: message_start
+data: {"type":"message_start","message":{"id":"msg_beta_usage_omitted_optional","type":"message","role":"assistant","content":[],"model":"claude-test","stop_reason":null,"stop_sequence":null}}
+
+event: content_block_start
+data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
+
+event: content_block_delta
+data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}
+
+event: content_block_stop
+data: {"type":"content_block_stop","index":0}
+
+event: message_delta
+data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":12,"cache_creation_input_tokens":4,"cache_read_input_tokens":3,"output_tokens":6,"output_tokens_details":{"thinking_tokens":2},"server_tool_use":{"web_search_requests":1,"web_fetch_requests":0},"iterations":[{"type":"message","model":"claude-test","input_tokens":12,"cache_creation_input_tokens":4,"cache_read_input_tokens":3,"output_tokens":6}],"fallback_credit":{"status":{"type":"redeemed"}}}}
+
+event: message_stop
+data: {"type":"message_stop"}

tests/lib/streaming/fixtures/usage_omitted_with_optional_fields_response.txt

@@ -0,0 +1,17 @@
+event: message_start
+data: {"type":"message_start","message":{"id":"msg_usage_omitted_optional","type":"message","role":"assistant","content":[],"model":"claude-test","stop_reason":null,"stop_sequence":null}}
+
+event: content_block_start
+data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
+
+event: content_block_delta
+data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}
+
+event: content_block_stop
+data: {"type":"content_block_stop","index":0}
+
+event: message_delta
+data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":12,"cache_creation_input_tokens":4,"cache_read_input_tokens":3,"output_tokens":6,"output_tokens_details":{"thinking_tokens":2},"server_tool_use":{"web_search_requests":1,"web_fetch_requests":0}}}
+
+event: message_stop
+data: {"type":"message_stop"}

tests/lib/streaming/test_beta_messages.py

@@ -11,6 +11,7 @@
 from anthropic import Anthropic, AsyncAnthropic
 from anthropic._utils import assert_overloads_in_sync, assert_signatures_in_sync
 from anthropic._compat import PYDANTIC_V1
+from anthropic.types.beta.beta_usage import BetaUsage
 from anthropic.types.beta.beta_message import BetaMessage
 from anthropic.lib.streaming._beta_types import (
     BetaInputJsonEvent,
@@ -536,6 +537,37 @@ def test_context_management_propagated(self, respx_mock: MockRouter) -> None:
         ) as stream:
             assert_context_management_response(stream.get_final_message())
 
+    @pytest.mark.respx(base_url=base_url)
+    def test_usage_omitted_at_message_start_uses_beta_usage_and_preserves_delta_optional_usage_fields(
+        self, respx_mock: MockRouter
+    ) -> None:
+        respx_mock.post("/v1/messages").mock(
+            return_value=httpx.Response(200, content=get_response("beta_usage_omitted_with_optional_fields_response.txt"))
+        )
+
+        client = Anthropic(base_url=base_url, api_key=api_key)
+
+        with client.beta.messages.stream(
+            max_tokens=1024,
+            messages=[{"role": "user", "content": "Say hello there!"}],
+            model="claude-test",
+        ) as stream:
+            message = stream.get_final_message()
+
+        assert isinstance(message.usage, BetaUsage)
+        assert message.usage.input_tokens == 12
+        assert message.usage.output_tokens == 6
+        assert message.usage.cache_creation_input_tokens == 4
+        assert message.usage.cache_read_input_tokens == 3
+        assert message.usage.output_tokens_details is not None
+        assert message.usage.output_tokens_details.thinking_tokens == 2
+        assert message.usage.server_tool_use is not None
+        assert message.usage.server_tool_use.web_search_requests == 1
+        assert message.usage.iterations is not None
+        assert message.usage.iterations[0].type == "message"
+        assert message.usage.fallback_credit is not None
+        assert message.usage.fallback_credit.status.type == "redeemed"
+
 
 class TestAsyncMessages:
     @pytest.mark.asyncio

tests/lib/streaming/test_messages.py

@@ -409,6 +409,32 @@ def test_usage_omitted_at_message_start(self, respx_mock: MockRouter) -> None:
         assert message.usage.input_tokens == 12
         assert message.usage.output_tokens == 6
 
+    @pytest.mark.respx(base_url=base_url)
+    def test_usage_omitted_at_message_start_preserves_delta_optional_usage_fields(
+        self, respx_mock: MockRouter
+    ) -> None:
+        respx_mock.post("/v1/messages").mock(
+            return_value=httpx.Response(200, content=get_response("usage_omitted_with_optional_fields_response.txt"))
+        )
+
+        client = Anthropic(base_url=base_url, api_key=api_key)
+
+        with client.messages.stream(
+            max_tokens=1024,
+            messages=[{"role": "user", "content": "Say hello there!"}],
+            model="claude-test",
+        ) as stream:
+            message = stream.get_final_message()
+
+        assert message.usage.input_tokens == 12
+        assert message.usage.output_tokens == 6
+        assert message.usage.cache_creation_input_tokens == 4
+        assert message.usage.cache_read_input_tokens == 3
+        assert message.usage.output_tokens_details is not None
+        assert message.usage.output_tokens_details.thinking_tokens == 2
+        assert message.usage.server_tool_use is not None
+        assert message.usage.server_tool_use.web_search_requests == 1
+
 
 class TestAsyncMessages:
     @pytest.mark.asyncio