From 28bca76dd99a66212a0c4fd1b60153a120f8394b Mon Sep 17 00:00:00 2001 From: mukunda katta Date: Thu, 30 Apr 2026 16:22:36 -0700 Subject: [PATCH] Add timestamps to message models --- .../_internal/message_parser.py | 9 ++++ src/claude_agent_sdk/types.py | 4 ++ tests/test_message_parser.py | 53 +++++++++++++++++++ tests/test_types.py | 15 +++++- 4 files changed, 79 insertions(+), 2 deletions(-) diff --git a/src/claude_agent_sdk/_internal/message_parser.py b/src/claude_agent_sdk/_internal/message_parser.py index 757c5ceb2..5e8c4c8c1 100644 --- a/src/claude_agent_sdk/_internal/message_parser.py +++ b/src/claude_agent_sdk/_internal/message_parser.py @@ -85,12 +85,14 @@ def parse_message(data: dict[str, Any]) -> Message | None: return UserMessage( content=user_content_blocks, uuid=uuid, + timestamp=data.get("timestamp"), parent_tool_use_id=parent_tool_use_id, tool_use_result=tool_use_result, ) return UserMessage( content=data["message"]["content"], uuid=uuid, + timestamp=data.get("timestamp"), parent_tool_use_id=parent_tool_use_id, tool_use_result=tool_use_result, ) @@ -155,6 +157,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"), + timestamp=data.get("timestamp"), ) except KeyError as e: raise MessageParseError( @@ -175,6 +178,7 @@ def parse_message(data: dict[str, Any]) -> Message | None: session_id=data["session_id"], tool_use_id=data.get("tool_use_id"), task_type=data.get("task_type"), + timestamp=data.get("timestamp"), ) case "task_progress": return TaskProgressMessage( @@ -187,6 +191,7 @@ def parse_message(data: dict[str, Any]) -> Message | None: session_id=data["session_id"], tool_use_id=data.get("tool_use_id"), last_tool_name=data.get("last_tool_name"), + timestamp=data.get("timestamp"), ) case "task_notification": return TaskNotificationMessage( @@ -200,6 +205,7 @@ def parse_message(data: dict[str, Any]) -> Message | None: session_id=data["session_id"], tool_use_id=data.get("tool_use_id"), usage=data.get("usage"), + timestamp=data.get("timestamp"), ) case "mirror_error": # SDK-synthesized via report_mirror_error — never emitted by the CLI subprocess. @@ -208,11 +214,13 @@ def parse_message(data: dict[str, Any]) -> Message | None: data=data, key=data.get("key"), error=data.get("error", ""), + timestamp=data.get("timestamp"), ) case _: return SystemMessage( subtype=subtype, data=data, + timestamp=data.get("timestamp"), ) except KeyError as e: raise MessageParseError( @@ -237,6 +245,7 @@ def parse_message(data: dict[str, Any]) -> Message | None: permission_denials=data.get("permission_denials"), errors=data.get("errors"), uuid=data.get("uuid"), + timestamp=data.get("timestamp"), ) except KeyError as e: raise MessageParseError( diff --git a/src/claude_agent_sdk/types.py b/src/claude_agent_sdk/types.py index 9c2be63fe..6814b3717 100644 --- a/src/claude_agent_sdk/types.py +++ b/src/claude_agent_sdk/types.py @@ -964,6 +964,7 @@ class UserMessage: content: str | list[ContentBlock] uuid: str | None = None + timestamp: str | None = None parent_tool_use_id: str | None = None tool_use_result: dict[str, Any] | None = None @@ -981,6 +982,7 @@ class AssistantMessage: stop_reason: str | None = None session_id: str | None = None uuid: str | None = None + timestamp: str | None = None @dataclass @@ -989,6 +991,7 @@ class SystemMessage: subtype: str data: dict[str, Any] + timestamp: str | None = field(default=None, kw_only=True) class TaskUsage(TypedDict): @@ -1093,6 +1096,7 @@ class ResultMessage: permission_denials: list[Any] | None = None errors: list[str] | None = None uuid: str | None = None + timestamp: str | None = None @dataclass diff --git a/tests/test_message_parser.py b/tests/test_message_parser.py index 69863bb45..839ed35f6 100644 --- a/tests/test_message_parser.py +++ b/tests/test_message_parser.py @@ -52,6 +52,17 @@ def test_parse_user_message_with_uuid(self): assert message.uuid == "msg-abc123-def456" assert len(message.content) == 1 + def test_parse_user_message_with_timestamp(self): + """Test parsing the transcript timestamp field for user messages.""" + data = { + "type": "user", + "timestamp": "2026-04-30T12:00:00.000Z", + "message": {"content": [{"type": "text", "text": "Hello"}]}, + } + message = parse_message(data) + assert isinstance(message, UserMessage) + assert message.timestamp == "2026-04-30T12:00:00.000Z" + def test_parse_user_message_with_tool_use(self): """Test parsing a user message with tool_use block.""" data = { @@ -215,6 +226,48 @@ def test_parse_user_message_with_tool_use_result(self): assert message.tool_use_result["structuredPatch"][0]["oldStart"] == 33 assert message.uuid == "2ace3375-1879-48a0-a421-6bce25a9295a" + def test_parse_assistant_system_and_result_timestamps(self): + """Test parsing transcript timestamps for non-user message types.""" + assistant = parse_message( + { + "type": "assistant", + "timestamp": "2026-04-30T12:00:01.000Z", + "session_id": "session-123", + "message": { + "id": "msg_123", + "model": "claude-opus-4-1-20250805", + "content": [{"type": "text", "text": "Hi"}], + }, + } + ) + assert isinstance(assistant, AssistantMessage) + assert assistant.timestamp == "2026-04-30T12:00:01.000Z" + + system = parse_message( + { + "type": "system", + "subtype": "init", + "timestamp": "2026-04-30T12:00:02.000Z", + } + ) + assert isinstance(system, SystemMessage) + assert system.timestamp == "2026-04-30T12:00:02.000Z" + + result = parse_message( + { + "type": "result", + "subtype": "success", + "duration_ms": 10, + "duration_api_ms": 8, + "is_error": False, + "num_turns": 1, + "session_id": "session-123", + "timestamp": "2026-04-30T12:00:03.000Z", + } + ) + assert isinstance(result, ResultMessage) + assert result.timestamp == "2026-04-30T12:00:03.000Z" + def test_parse_user_message_with_string_content_and_tool_use_result(self): """Test parsing a user message with string content and tool_use_result.""" tool_result_data = {"filePath": "/path/to/file.py", "userModified": True} diff --git a/tests/test_types.py b/tests/test_types.py index fbd07509f..6f1be5539 100644 --- a/tests/test_types.py +++ b/tests/test_types.py @@ -27,15 +27,24 @@ class TestMessageTypes: def test_user_message_creation(self): """Test creating a UserMessage.""" - msg = UserMessage(content="Hello, Claude!") + msg = UserMessage( + content="Hello, Claude!", + timestamp="2026-04-30T12:00:00.000Z", + ) assert msg.content == "Hello, Claude!" + assert msg.timestamp == "2026-04-30T12:00:00.000Z" def test_assistant_message_with_text(self): """Test creating an AssistantMessage with text content.""" text_block = TextBlock(text="Hello, human!") - msg = AssistantMessage(content=[text_block], model="claude-opus-4-1-20250805") + msg = AssistantMessage( + content=[text_block], + model="claude-opus-4-1-20250805", + timestamp="2026-04-30T12:00:01.000Z", + ) assert len(msg.content) == 1 assert msg.content[0].text == "Hello, human!" + assert msg.timestamp == "2026-04-30T12:00:01.000Z" def test_assistant_message_with_thinking(self): """Test creating an AssistantMessage with thinking content.""" @@ -75,10 +84,12 @@ def test_result_message(self): num_turns=1, session_id="session-123", total_cost_usd=0.01, + timestamp="2026-04-30T12:00:02.000Z", ) assert msg.subtype == "success" assert msg.total_cost_usd == 0.01 assert msg.session_id == "session-123" + assert msg.timestamp == "2026-04-30T12:00:02.000Z" class TestOptions: