diff --git a/e2e-tests/test_message_fields.py b/e2e-tests/test_message_fields.py new file mode 100644 index 00000000..397b5495 --- /dev/null +++ b/e2e-tests/test_message_fields.py @@ -0,0 +1,180 @@ +"""End-to-end tests for preserved message fields with real Claude API calls. + +These tests verify that AssistantMessage and ResultMessage correctly preserve +fields from the raw CLI JSON output that were previously dropped during parsing. +See issue #562. +""" + +from typing import Any + +import pytest + +from claude_agent_sdk import ClaudeSDKClient +from claude_agent_sdk.types import ( + AssistantMessage, + ClaudeAgentOptions, + ResultMessage, + StreamEvent, + SystemMessage, +) + + +@pytest.mark.e2e +@pytest.mark.asyncio +async def test_assistant_message_preserves_fields(): + """Test that AssistantMessage preserves id, usage, stop_reason, session_id, uuid.""" + + options = ClaudeAgentOptions( + model="claude-sonnet-4-5", + max_turns=1, + ) + + collected_messages: list[Any] = [] + + async with ClaudeSDKClient(options) as client: + await client.query("Say hi") + + async for message in client.receive_response(): + collected_messages.append(message) + + # Find AssistantMessage + assistant_messages = [ + msg for msg in collected_messages if isinstance(msg, AssistantMessage) + ] + assert len(assistant_messages) >= 1, "No AssistantMessage received" + + msg = assistant_messages[0] + + # model should be a real model string (not empty or synthetic) + assert msg.model is not None + assert len(msg.model) > 0 + assert "claude" in msg.model, f"Unexpected model: {msg.model}" + + # id should be an Anthropic message ID + assert msg.id is not None, "AssistantMessage.id should not be None" + assert msg.id.startswith("msg_"), f"Unexpected message id format: {msg.id}" + + # usage should contain token counts + assert msg.usage is not None, "AssistantMessage.usage should not be None" + assert "input_tokens" in msg.usage, "usage missing input_tokens" + assert "output_tokens" in msg.usage, "usage missing output_tokens" + assert msg.usage["input_tokens"] >= 0 + assert msg.usage["output_tokens"] >= 0 + + # session_id should be present + assert msg.session_id is not None, "AssistantMessage.session_id should not be None" + assert len(msg.session_id) > 0 + + # uuid should be present + assert msg.uuid is not None, "AssistantMessage.uuid should not be None" + assert len(msg.uuid) > 0 + + +@pytest.mark.e2e +@pytest.mark.asyncio +async def test_result_message_preserves_model_usage(): + """Test that ResultMessage preserves modelUsage, permission_denials, and uuid.""" + + options = ClaudeAgentOptions( + model="claude-sonnet-4-5", + max_turns=1, + ) + + collected_messages: list[Any] = [] + + async with ClaudeSDKClient(options) as client: + await client.query("Say hi") + + async for message in client.receive_response(): + collected_messages.append(message) + + # Find ResultMessage + result_messages = [ + msg for msg in collected_messages if isinstance(msg, ResultMessage) + ] + assert len(result_messages) == 1, "Expected exactly one ResultMessage" + + result = result_messages[0] + + # model_usage should contain per-model breakdown + assert result.model_usage is not None, ( + "ResultMessage.model_usage should not be None" + ) + assert len(result.model_usage) >= 1, "model_usage should have at least one model" + + # The model key should be a real model identifier + model_names = list(result.model_usage.keys()) + assert any( + "claude" in name for name in model_names + ), f"No claude model in model_usage keys: {model_names}" + + # Each model entry should have token and cost fields + for model_name, model_data in result.model_usage.items(): + assert "inputTokens" in model_data, f"{model_name} missing inputTokens" + assert "outputTokens" in model_data, f"{model_name} missing outputTokens" + assert "costUSD" in model_data, f"{model_name} missing costUSD" + + # total_cost_usd should match modelUsage costUSD sum + if result.total_cost_usd is not None: + model_cost_sum = sum( + data.get("costUSD", 0) for data in result.model_usage.values() + ) + assert abs(result.total_cost_usd - model_cost_sum) < 0.0001, ( + f"total_cost_usd ({result.total_cost_usd}) doesn't match " + f"modelUsage costUSD sum ({model_cost_sum})" + ) + + # permission_denials should be present (empty list for simple queries) + assert result.permission_denials is not None, ( + "ResultMessage.permission_denials should not be None" + ) + assert isinstance(result.permission_denials, list) + + # uuid should be present + assert result.uuid is not None, "ResultMessage.uuid should not be None" + assert len(result.uuid) > 0 + + +@pytest.mark.e2e +@pytest.mark.asyncio +async def test_streaming_mode_preserves_fields(): + """Test that fields are preserved in streaming mode with include_partial_messages.""" + + options = ClaudeAgentOptions( + include_partial_messages=True, + model="claude-sonnet-4-5", + max_turns=1, + ) + + collected_messages: list[Any] = [] + + async with ClaudeSDKClient(options) as client: + await client.query("Say hi") + + async for message in client.receive_response(): + collected_messages.append(message) + + # AssistantMessage fields should still be preserved in streaming mode + assistant_messages = [ + msg for msg in collected_messages if isinstance(msg, AssistantMessage) + ] + assert len(assistant_messages) >= 1, "No AssistantMessage received" + + msg = assistant_messages[0] + assert msg.id is not None, "AssistantMessage.id missing in streaming mode" + assert msg.usage is not None, "AssistantMessage.usage missing in streaming mode" + assert msg.session_id is not None, ( + "AssistantMessage.session_id missing in streaming mode" + ) + + # ResultMessage model_usage should also be preserved + result_messages = [ + msg for msg in collected_messages if isinstance(msg, ResultMessage) + ] + assert len(result_messages) == 1 + assert result_messages[0].model_usage is not None, ( + "ResultMessage.model_usage missing in streaming mode" + ) + assert result_messages[0].uuid is not None, ( + "ResultMessage.uuid missing in streaming mode" + ) diff --git a/src/claude_agent_sdk/_internal/message_parser.py b/src/claude_agent_sdk/_internal/message_parser.py index f91081c3..c4163559 100644 --- a/src/claude_agent_sdk/_internal/message_parser.py +++ b/src/claude_agent_sdk/_internal/message_parser.py @@ -127,6 +127,12 @@ def parse_message(data: dict[str, Any]) -> Message: model=data["message"]["model"], parent_tool_use_id=data.get("parent_tool_use_id"), error=data.get("error"), + id=data["message"].get("id"), + usage=data["message"].get("usage"), + stop_reason=data["message"].get("stop_reason"), + stop_sequence=data["message"].get("stop_sequence"), + session_id=data.get("session_id"), + uuid=data.get("uuid"), ) except KeyError as e: raise MessageParseError( @@ -157,6 +163,10 @@ def parse_message(data: dict[str, Any]) -> Message: usage=data.get("usage"), result=data.get("result"), structured_output=data.get("structured_output"), + model_usage=data.get("modelUsage"), + stop_reason=data.get("stop_reason"), + permission_denials=data.get("permission_denials"), + uuid=data.get("uuid"), ) except KeyError as e: raise MessageParseError( diff --git a/src/claude_agent_sdk/types.py b/src/claude_agent_sdk/types.py index 1b544a4d..ab204576 100644 --- a/src/claude_agent_sdk/types.py +++ b/src/claude_agent_sdk/types.py @@ -657,6 +657,12 @@ class AssistantMessage: model: str parent_tool_use_id: str | None = None error: AssistantMessageError | None = None + id: str | None = None + usage: dict[str, Any] | None = None + stop_reason: str | None = None + stop_sequence: str | None = None + session_id: str | None = None + uuid: str | None = None @dataclass @@ -681,6 +687,10 @@ class ResultMessage: usage: dict[str, Any] | None = None result: str | None = None structured_output: Any = None + model_usage: dict[str, Any] | None = None + stop_reason: str | None = None + permission_denials: list[Any] | None = None + uuid: str | None = None @dataclass diff --git a/tests/test_message_parser.py b/tests/test_message_parser.py index dc21ee6a..380daeb9 100644 --- a/tests/test_message_parser.py +++ b/tests/test_message_parser.py @@ -437,3 +437,173 @@ def test_parse_assistant_message_with_rate_limit_error(self): message = parse_message(data) assert isinstance(message, AssistantMessage) assert message.error == "rate_limit" + + def test_parse_assistant_message_preserves_all_fields(self): + """Test that AssistantMessage preserves id, usage, stop_reason, session_id, uuid. + + These fields are present in the raw CLI JSON output but were previously + dropped during parsing. See issue #562. + """ + data = { + "type": "assistant", + "message": { + "model": "claude-sonnet-4-5-20250929", + "id": "msg_01HRq7YZE3apPqSHydvG77Ve", + "type": "message", + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Hi! I'm ready to help.", + } + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": 3, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 20012, + "output_tokens": 1, + "service_tier": "standard", + }, + }, + "parent_tool_use_id": None, + "session_id": "fdf2d90a-fd9e-4736-ae35-806edd13643f", + "uuid": "0dbd2453-1209-4fe9-bd51-4102f64e33df", + } + message = parse_message(data) + assert isinstance(message, AssistantMessage) + assert message.model == "claude-sonnet-4-5-20250929" + assert message.id == "msg_01HRq7YZE3apPqSHydvG77Ve" + assert message.usage is not None + assert message.usage["input_tokens"] == 3 + assert message.usage["cache_read_input_tokens"] == 20012 + assert message.usage["output_tokens"] == 1 + assert message.stop_reason == "end_turn" + assert message.stop_sequence is None + assert message.session_id == "fdf2d90a-fd9e-4736-ae35-806edd13643f" + assert message.uuid == "0dbd2453-1209-4fe9-bd51-4102f64e33df" + + def test_parse_assistant_message_optional_fields_default_to_none(self): + """Test that new optional fields default to None when not present.""" + data = { + "type": "assistant", + "message": { + "content": [{"type": "text", "text": "Hello"}], + "model": "claude-opus-4-1-20250805", + }, + } + message = parse_message(data) + assert isinstance(message, AssistantMessage) + assert message.id is None + assert message.usage is None + assert message.stop_reason is None + assert message.stop_sequence is None + assert message.session_id is None + assert message.uuid is None + + def test_parse_result_message_preserves_all_fields(self): + """Test that ResultMessage preserves model_usage, stop_reason, permission_denials, uuid. + + These fields are present in the raw CLI JSON output but were previously + dropped during parsing. See issue #562. + """ + data = { + "type": "result", + "subtype": "success", + "is_error": False, + "duration_ms": 2995, + "duration_api_ms": 2190, + "num_turns": 1, + "result": "Hi! I'm ready to help.", + "stop_reason": None, + "session_id": "fdf2d90a-fd9e-4736-ae35-806edd13643f", + "total_cost_usd": 0.010620999999999998, + "usage": { + "input_tokens": 3, + "output_tokens": 24, + }, + "modelUsage": { + "claude-sonnet-4-5-20250929": { + "inputTokens": 3, + "outputTokens": 24, + "cacheReadInputTokens": 20012, + "cacheCreationInputTokens": 0, + "costUSD": 0.010620999999999998, + "contextWindow": 200000, + "maxOutputTokens": 64000, + } + }, + "permission_denials": [], + "uuid": "d379c496-f33a-4ea4-b920-3c5483baa6f7", + } + message = parse_message(data) + assert isinstance(message, ResultMessage) + assert message.subtype == "success" + assert message.total_cost_usd == 0.010620999999999998 + # New fields + assert message.model_usage is not None + assert "claude-sonnet-4-5-20250929" in message.model_usage + model_data = message.model_usage["claude-sonnet-4-5-20250929"] + assert model_data["inputTokens"] == 3 + assert model_data["outputTokens"] == 24 + assert model_data["cacheReadInputTokens"] == 20012 + assert model_data["costUSD"] == 0.010620999999999998 + assert message.stop_reason is None + assert message.permission_denials == [] + assert message.uuid == "d379c496-f33a-4ea4-b920-3c5483baa6f7" + + def test_parse_result_message_optional_fields_default_to_none(self): + """Test that new optional fields default to None when not present.""" + data = { + "type": "result", + "subtype": "success", + "duration_ms": 1000, + "duration_api_ms": 500, + "is_error": False, + "num_turns": 2, + "session_id": "session_123", + } + message = parse_message(data) + assert isinstance(message, ResultMessage) + assert message.model_usage is None + assert message.stop_reason is None + assert message.permission_denials is None + assert message.uuid is None + + def test_parse_result_message_model_usage_multiple_models(self): + """Test ResultMessage with modelUsage containing multiple models. + + When subagents use different models, modelUsage has multiple entries. + """ + data = { + "type": "result", + "subtype": "success", + "is_error": False, + "duration_ms": 16025, + "duration_api_ms": 24144, + "num_turns": 9, + "result": "Done.", + "session_id": "session_456", + "total_cost_usd": 0.25, + "modelUsage": { + "claude-sonnet-4-5-20250929": { + "inputTokens": 100, + "outputTokens": 200, + "costUSD": 0.15, + }, + "claude-haiku-4-5-20251001": { + "inputTokens": 50, + "outputTokens": 100, + "costUSD": 0.10, + }, + }, + } + message = parse_message(data) + assert isinstance(message, ResultMessage) + assert message.model_usage is not None + assert len(message.model_usage) == 2 + assert "claude-sonnet-4-5-20250929" in message.model_usage + assert "claude-haiku-4-5-20251001" in message.model_usage + assert message.model_usage["claude-sonnet-4-5-20250929"]["costUSD"] == 0.15 + assert message.model_usage["claude-haiku-4-5-20251001"]["costUSD"] == 0.10