Improve UserMessage types to include ToolResultBlock (anthropics#101)

Fixes anthropics#90

Signed-off-by: Rushil Patel <[email protected]>

Author: dicksontsai

examples/streaming_mode.py

@@ -27,6 +27,8 @@
     ResultMessage,
     SystemMessage,
     TextBlock,
+    ToolResultBlock,
+    ToolUseBlock,
     UserMessage,
 )
 
@@ -303,6 +305,50 @@ async def create_message_stream():
     print("\n")
 
 
+async def example_bash_command():
+    """Example showing tool use blocks when running bash commands."""
+    print("=== Bash Command Example ===")
+    
+    async with ClaudeSDKClient() as client:
+        print("User: Run a bash echo command")
+        await client.query("Run a bash echo command that says 'Hello from bash!'")
+        
+        # Track all message types received
+        message_types = []
+        
+        async for msg in client.receive_messages():
+            message_types.append(type(msg).__name__)
+            
+            if isinstance(msg, UserMessage):
+                # User messages can contain tool results
+                for block in msg.content:
+                    if isinstance(block, TextBlock):
+                        print(f"User: {block.text}")
+                    elif isinstance(block, ToolResultBlock):
+                        print(f"Tool Result (id: {block.tool_use_id}): {block.content[:100] if block.content else 'None'}...")
+                        
+            elif isinstance(msg, AssistantMessage):
+                # Assistant messages can contain tool use blocks
+                for block in msg.content:
+                    if isinstance(block, TextBlock):
+                        print(f"Claude: {block.text}")
+                    elif isinstance(block, ToolUseBlock):
+                        print(f"Tool Use: {block.name} (id: {block.id})")
+                        if block.name == "Bash":
+                            command = block.input.get("command", "")
+                            print(f"  Command: {command}")
+                            
+            elif isinstance(msg, ResultMessage):
+                print("Result ended")
+                if msg.total_cost_usd:
+                    print(f"Cost: ${msg.total_cost_usd:.4f}")
+                break
+        
+        print(f"\nMessage types received: {', '.join(set(message_types))}")
+    
+    print("\n")
+
+
 async def example_error_handling():
     """Demonstrate proper error handling."""
     print("=== Error Handling Example ===")
@@ -359,6 +405,7 @@ async def main():
         "manual_message_handling": example_manual_message_handling,
         "with_options": example_with_options,
         "async_iterable_prompt": example_async_iterable_prompt,
+        "bash_command": example_bash_command,
         "error_handling": example_error_handling,
     }
 

src/claude_code_sdk/_internal/message_parser.py

@@ -45,6 +45,31 @@ def parse_message(data: dict[str, Any]) -> Message:
     match message_type:
         case "user":
             try:
+                if isinstance(data["message"]["content"], list):
+                    user_content_blocks: list[ContentBlock] = []
+                    for block in data["message"]["content"]:
+                        match block["type"]:
+                            case "text":
+                                user_content_blocks.append(
+                                    TextBlock(text=block["text"])
+                                )
+                            case "tool_use":
+                                user_content_blocks.append(
+                                    ToolUseBlock(
+                                        id=block["id"],
+                                        name=block["name"],
+                                        input=block["input"],
+                                    )
+                                )
+                            case "tool_result":
+                                user_content_blocks.append(
+                                    ToolResultBlock(
+                                        tool_use_id=block["tool_use_id"],
+                                        content=block.get("content"),
+                                        is_error=block.get("is_error"),
+                                    )
+                                )
+                    return UserMessage(content=user_content_blocks)
                 return UserMessage(content=data["message"]["content"])
             except KeyError as e:
                 raise MessageParseError(

src/claude_code_sdk/types.py

@@ -73,7 +73,7 @@ class ToolResultBlock:
 class UserMessage:
     """User message."""
 
-    content: str
+    content: str | list[ContentBlock]
 
 
 @dataclass

tests/test_message_parser.py

@@ -9,6 +9,7 @@
     ResultMessage,
     SystemMessage,
     TextBlock,
+    ToolResultBlock,
     ToolUseBlock,
     UserMessage,
 )
@@ -25,6 +26,108 @@ def test_parse_valid_user_message(self):
         }
         message = parse_message(data)
         assert isinstance(message, UserMessage)
+        assert len(message.content) == 1
+        assert isinstance(message.content[0], TextBlock)
+        assert message.content[0].text == "Hello"
+
+    def test_parse_user_message_with_tool_use(self):
+        """Test parsing a user message with tool_use block."""
+        data = {
+            "type": "user",
+            "message": {
+                "content": [
+                    {"type": "text", "text": "Let me read this file"},
+                    {
+                        "type": "tool_use",
+                        "id": "tool_456",
+                        "name": "Read",
+                        "input": {"file_path": "/example.txt"},
+                    },
+                ]
+            },
+        }
+        message = parse_message(data)
+        assert isinstance(message, UserMessage)
+        assert len(message.content) == 2
+        assert isinstance(message.content[0], TextBlock)
+        assert isinstance(message.content[1], ToolUseBlock)
+        assert message.content[1].id == "tool_456"
+        assert message.content[1].name == "Read"
+        assert message.content[1].input == {"file_path": "/example.txt"}
+
+    def test_parse_user_message_with_tool_result(self):
+        """Test parsing a user message with tool_result block."""
+        data = {
+            "type": "user",
+            "message": {
+                "content": [
+                    {
+                        "type": "tool_result",
+                        "tool_use_id": "tool_789",
+                        "content": "File contents here",
+                    }
+                ]
+            },
+        }
+        message = parse_message(data)
+        assert isinstance(message, UserMessage)
+        assert len(message.content) == 1
+        assert isinstance(message.content[0], ToolResultBlock)
+        assert message.content[0].tool_use_id == "tool_789"
+        assert message.content[0].content == "File contents here"
+
+    def test_parse_user_message_with_tool_result_error(self):
+        """Test parsing a user message with error tool_result block."""
+        data = {
+            "type": "user",
+            "message": {
+                "content": [
+                    {
+                        "type": "tool_result",
+                        "tool_use_id": "tool_error",
+                        "content": "File not found",
+                        "is_error": True,
+                    }
+                ]
+            },
+        }
+        message = parse_message(data)
+        assert isinstance(message, UserMessage)
+        assert len(message.content) == 1
+        assert isinstance(message.content[0], ToolResultBlock)
+        assert message.content[0].tool_use_id == "tool_error"
+        assert message.content[0].content == "File not found"
+        assert message.content[0].is_error is True
+
+    def test_parse_user_message_with_mixed_content(self):
+        """Test parsing a user message with mixed content blocks."""
+        data = {
+            "type": "user",
+            "message": {
+                "content": [
+                    {"type": "text", "text": "Here's what I found:"},
+                    {
+                        "type": "tool_use",
+                        "id": "use_1",
+                        "name": "Search",
+                        "input": {"query": "test"},
+                    },
+                    {
+                        "type": "tool_result",
+                        "tool_use_id": "use_1",
+                        "content": "Search results",
+                    },
+                    {"type": "text", "text": "What do you think?"},
+                ]
+            },
+        }
+        message = parse_message(data)
+        assert isinstance(message, UserMessage)
+        assert len(message.content) == 4
+        assert isinstance(message.content[0], TextBlock)
+        assert isinstance(message.content[1], ToolUseBlock)
+        assert isinstance(message.content[2], ToolResultBlock)
+        assert isinstance(message.content[3], TextBlock)
 
     def test_parse_valid_assistant_message(self):
         """Test parsing a valid assistant message."""