Add UUID to UserMessage response type to improve devX for rewind (#418)

Author: noahzweben

src/claude_agent_sdk/_internal/message_parser.py

@@ -48,6 +48,7 @@ def parse_message(data: dict[str, Any]) -> Message:
         case "user":
             try:
                 parent_tool_use_id = data.get("parent_tool_use_id")
+                uuid = data.get("uuid")
                 if isinstance(data["message"]["content"], list):
                     user_content_blocks: list[ContentBlock] = []
                     for block in data["message"]["content"]:
@@ -74,10 +75,12 @@ def parse_message(data: dict[str, Any]) -> Message:
                                 )
                     return UserMessage(
                         content=user_content_blocks,
+                        uuid=uuid,
                         parent_tool_use_id=parent_tool_use_id,
                     )
                 return UserMessage(
                     content=data["message"]["content"],
+                    uuid=uuid,
                     parent_tool_use_id=parent_tool_use_id,
                 )
             except KeyError as e:

src/claude_agent_sdk/client.py

@@ -264,20 +264,25 @@ async def set_model(self, model: str | None = None) -> None:
     async def rewind_files(self, user_message_id: str) -> None:
         """Rewind tracked files to their state at a specific user message.
 
-        Requires file checkpointing to be enabled via the `enable_file_checkpointing` option
-        when creating the ClaudeSDKClient.
+        Requires:
+            - `enable_file_checkpointing=True` to track file changes
+            - `extra_args={"replay-user-messages": None}` to receive UserMessage
+              objects with `uuid` in the response stream
 
         Args:
             user_message_id: UUID of the user message to rewind to. This should be
                 the `uuid` field from a `UserMessage` received during the conversation.
 
         Example:
             ```python
-            options = ClaudeAgentOptions(enable_file_checkpointing=True)
+            options = ClaudeAgentOptions(
+                enable_file_checkpointing=True,
+                extra_args={"replay-user-messages": None},
+            )
             async with ClaudeSDKClient(options) as client:
                 await client.query("Make some changes to my files")
                 async for msg in client.receive_response():
-                    if isinstance(msg, UserMessage):
+                    if isinstance(msg, UserMessage) and msg.uuid:
                         checkpoint_id = msg.uuid  # Save this for later
 
                 # Later, rewind to that point

src/claude_agent_sdk/types.py

@@ -562,6 +562,7 @@ class UserMessage:
     """User message."""
 
     content: str | list[ContentBlock]
+    uuid: str | None = None
     parent_tool_use_id: str | None = None
 
 

tests/test_message_parser.py

@@ -31,6 +31,21 @@ def test_parse_valid_user_message(self):
         assert isinstance(message.content[0], TextBlock)
         assert message.content[0].text == "Hello"
 
+    def test_parse_user_message_with_uuid(self):
+        """Test parsing a user message with uuid field (issue #414).
+
+        The uuid field is needed for file checkpointing with rewind_files().
+        """
+        data = {
+            "type": "user",
+            "uuid": "msg-abc123-def456",
+            "message": {"content": [{"type": "text", "text": "Hello"}]},
+        }
+        message = parse_message(data)
+        assert isinstance(message, UserMessage)
+        assert message.uuid == "msg-abc123-def456"
+        assert len(message.content) == 1
+
     def test_parse_user_message_with_tool_use(self):
         """Test parsing a user message with tool_use block."""
         data = {