diff --git a/examples/streaming_mode.py b/examples/streaming_mode.py index 73eb4109..da2bc1b8 100755 --- a/examples/streaming_mode.py +++ b/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, } diff --git a/src/claude_code_sdk/_internal/message_parser.py b/src/claude_code_sdk/_internal/message_parser.py index 858e24fa..8477e51a 100644 --- a/src/claude_code_sdk/_internal/message_parser.py +++ b/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( diff --git a/src/claude_code_sdk/types.py b/src/claude_code_sdk/types.py index bd3c7267..46bab99a 100644 --- a/src/claude_code_sdk/types.py +++ b/src/claude_code_sdk/types.py @@ -73,7 +73,7 @@ class ToolResultBlock: class UserMessage: """User message.""" - content: str + content: str | list[ContentBlock] @dataclass diff --git a/tests/test_message_parser.py b/tests/test_message_parser.py index 0eb43542..b6475301 100644 --- a/tests/test_message_parser.py +++ b/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."""