From 0b79cb2c9a6c8e686f976e0ecccb445f6eafe1fb Mon Sep 17 00:00:00 2001 From: itxaiohanglover <1531137510@qq.com> Date: Sat, 20 Jun 2026 01:42:31 +0800 Subject: [PATCH] fix: preserve wire fields in parse_message via data field Add `data: dict[str, Any] | None = None` to AssistantMessage and ResultMessage, mirroring the existing SystemMessage.data pattern. The parser now passes the raw frame dict as `data=data`, so newer CLI wire fields (ttft_ms, terminal_reason, fast_mode_state, stop_details, diagnostics, context_management, caller, request_id) are reachable before they are explicitly modeled. Fixes #1026 --- .../_internal/message_parser.py | 2 + src/claude_agent_sdk/types.py | 20 +++- tests/test_message_parser.py | 108 ++++++++++++++++++ 3 files changed, 128 insertions(+), 2 deletions(-) diff --git a/src/claude_agent_sdk/_internal/message_parser.py b/src/claude_agent_sdk/_internal/message_parser.py index 574816c6f..15b4f7d53 100644 --- a/src/claude_agent_sdk/_internal/message_parser.py +++ b/src/claude_agent_sdk/_internal/message_parser.py @@ -179,6 +179,7 @@ def parse_message(data: dict[str, Any]) -> Message | None: stop_reason=data["message"].get("stop_reason"), session_id=data.get("session_id"), uuid=data.get("uuid"), + data=data, ) except KeyError as e: raise MessageParseError( @@ -270,6 +271,7 @@ def parse_message(data: dict[str, Any]) -> Message | None: errors=data.get("errors"), api_error_status=data.get("api_error_status"), uuid=data.get("uuid"), + data=data, ) except KeyError as e: raise MessageParseError( diff --git a/src/claude_agent_sdk/types.py b/src/claude_agent_sdk/types.py index ee925b351..b94d8fb58 100644 --- a/src/claude_agent_sdk/types.py +++ b/src/claude_agent_sdk/types.py @@ -1023,7 +1023,14 @@ class UserMessage: @dataclass class AssistantMessage: - """Assistant message with content blocks.""" + """Assistant message with content blocks. + + The ``data`` field retains the full raw wire frame from the CLI so that + newer, unmodeled fields (e.g. ``stop_details``, ``diagnostics``, + ``context_management``, ``caller``, ``request_id``) are reachable before + they are explicitly modeled. This mirrors the ``SystemMessage.data`` + pattern. Defaults to ``None`` for manually constructed messages. + """ content: list[ContentBlock] model: str @@ -1034,6 +1041,7 @@ class AssistantMessage: stop_reason: str | None = None session_id: str | None = None uuid: str | None = None + data: dict[str, Any] | None = None @dataclass @@ -1143,7 +1151,14 @@ class DeferredToolUse: @dataclass class ResultMessage: - """Result message with cost and usage information.""" + """Result message with cost and usage information. + + The ``data`` field retains the full raw wire frame from the CLI so that + newer, unmodeled fields (e.g. ``ttft_ms``, ``terminal_reason``, + ``fast_mode_state``, ``request_id``) are reachable before they are + explicitly modeled. This mirrors the ``SystemMessage.data`` pattern. + Defaults to ``None`` for manually constructed messages. + """ subtype: str duration_ms: int @@ -1165,6 +1180,7 @@ class ResultMessage: # Emitted by the CLI since v2.1.110. Safe to log (no message content). api_error_status: int | None = None uuid: str | None = None + data: dict[str, Any] | None = None @dataclass diff --git a/tests/test_message_parser.py b/tests/test_message_parser.py index 7ce2990ca..8bbd25698 100644 --- a/tests/test_message_parser.py +++ b/tests/test_message_parser.py @@ -1075,3 +1075,111 @@ def test_parse_hook_event_message_minimal(self): assert message.hook_event_name == "Stop" assert message.session_id is None assert message.uuid is None + + # -- data field: raw frame retention (issue #1026) ----------------------- + + def test_assistant_message_retains_raw_frame(self): + """AssistantMessage retains the full raw wire frame in ``data``. + + Newer CLI versions emit fields like ``stop_details``, ``diagnostics``, + ``context_management``, and ``request_id`` that are not yet modeled + on the dataclass. ``data`` is the forward-compatible escape hatch + so consumers can reach them without re-parsing the wire. + """ + data = { + "type": "assistant", + "message": { + "content": [{"type": "text", "text": "hi"}], + "model": "claude-sonnet-4-5", + }, + "request_id": "req_01ABC", + "stop_details": None, + "diagnostics": None, + "context_management": None, + } + message = parse_message(data) + assert isinstance(message, AssistantMessage) + assert message.data is not None + assert message.data == data + assert message.data["request_id"] == "req_01ABC" + assert message.data["stop_details"] is None + assert message.data["diagnostics"] is None + assert message.data["context_management"] is None + + def test_assistant_message_data_defaults_none(self): + """Manually constructed AssistantMessage has data=None (backward compat).""" + msg = AssistantMessage( + content=[TextBlock(text="hi")], + model="claude-sonnet-4-5", + ) + assert msg.data is None + + def test_result_message_retains_raw_frame(self): + """ResultMessage retains the full raw wire frame in ``data``. + + Newer CLI versions emit fields like ``ttft_ms``, ``terminal_reason``, + ``fast_mode_state``, and ``request_id`` that are not yet modeled on + the dataclass. ``data`` is the forward-compatible escape hatch so + consumers can reach them without re-parsing the wire. + """ + data = { + "type": "result", + "subtype": "success", + "duration_ms": 8127, + "duration_api_ms": 7903, + "is_error": False, + "num_turns": 4, + "session_id": "s", + "total_cost_usd": 0.02, + "usage": {}, + "ttft_ms": 2806, + "terminal_reason": "completed", + "fast_mode_state": "off", + "request_id": "req_01XYZ", + } + message = parse_message(data) + assert isinstance(message, ResultMessage) + assert message.data is not None + assert message.data == data + assert message.data["ttft_ms"] == 2806 + assert message.data["terminal_reason"] == "completed" + assert message.data["fast_mode_state"] == "off" + assert message.data["request_id"] == "req_01XYZ" + + def test_result_message_data_defaults_none(self): + """Manually constructed ResultMessage has data=None (backward compat).""" + msg = ResultMessage( + subtype="success", + duration_ms=1000, + duration_api_ms=500, + is_error=False, + num_turns=1, + session_id="s", + ) + assert msg.data is None + + def test_assistant_message_caller_preserved_in_data(self): + """Per-content-block ``caller`` is reachable via ``data``. + + The ``caller`` field lives inside individual content blocks on the + wire; it is not modeled on ``ContentBlock`` dataclasses, but the + raw frame in ``data`` preserves it. + """ + data = { + "type": "assistant", + "message": { + "content": [ + { + "type": "text", + "text": "Hello", + "caller": {"type": "direct"}, + }, + ], + "model": "claude-sonnet-4-5", + }, + } + message = parse_message(data) + assert isinstance(message, AssistantMessage) + assert message.data is not None + block = message.data["message"]["content"][0] + assert block["caller"] == {"type": "direct"}