From 355e329b9f6f4bf1298c4303515e393a7290bec4 Mon Sep 17 00:00:00 2001 From: zengweihao Date: Sun, 28 Dec 2025 23:36:08 +0800 Subject: [PATCH] fix: preserve per-step usage in AssistantMessage The Python SDK was discarding `id` and `usage` fields from assistant messages during parsing, making it impossible to track per-step token usage. This is inconsistent with the TypeScript SDK which preserves these fields. Changes: - Add `message_id` and `usage` fields to `AssistantMessage` dataclass - Extract these fields in `message_parser.py` when parsing assistant messages Use case: - Monitor per-step token consumption for cost analysis - Detect when API features like `clear_tool_results` trigger by observing input_tokens drops between steps - Parity with TypeScript SDK behavior --- src/claude_agent_sdk/_internal/message_parser.py | 3 +++ src/claude_agent_sdk/types.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/src/claude_agent_sdk/_internal/message_parser.py b/src/claude_agent_sdk/_internal/message_parser.py index 4bfe8145..7f893d54 100644 --- a/src/claude_agent_sdk/_internal/message_parser.py +++ b/src/claude_agent_sdk/_internal/message_parser.py @@ -124,6 +124,9 @@ def parse_message(data: dict[str, Any]) -> Message: model=data["message"]["model"], parent_tool_use_id=data.get("parent_tool_use_id"), error=data["message"].get("error"), + # Per-step fields from Anthropic API response + message_id=data["message"].get("id"), + usage=data["message"].get("usage"), ) except KeyError as e: raise MessageParseError( diff --git a/src/claude_agent_sdk/types.py b/src/claude_agent_sdk/types.py index 9c09345f..244115e6 100644 --- a/src/claude_agent_sdk/types.py +++ b/src/claude_agent_sdk/types.py @@ -574,6 +574,9 @@ class AssistantMessage: model: str parent_tool_use_id: str | None = None error: AssistantMessageError | None = None + # Per-step fields from Anthropic API response + message_id: str | None = None # API message ID (e.g., "msg_...") + usage: dict[str, Any] | None = None # Per-step token usage @dataclass