PR feedback

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

Author: dicksontsai

src/claude_code_sdk/_internal/message_parser.py

@@ -1,5 +1,6 @@
 """Message parser for Claude Code SDK responses."""
 
+import logging
 from typing import Any
 
 from ..types import (
@@ -14,6 +15,8 @@
     UserMessage,
 )
 
+logger = logging.getLogger(__name__)
+
 
 def parse_message(data: dict[str, Any]) -> Message | None:
     """
@@ -23,55 +26,82 @@ def parse_message(data: dict[str, Any]) -> Message | None:
         data: Raw message dictionary from CLI output
 
     Returns:
-        Parsed Message object or None if type is unrecognized
+        Parsed Message object or None if type is unrecognized or parsing fails
     """
-    match data["type"]:
+    try:
+        message_type = data.get("type")
+        if not message_type:
+            logger.warning("Message missing 'type' field: %s", data)
+            return None
+
+    except AttributeError:
+        logger.error("Invalid message data type (expected dict): %s", type(data))
+        return None
+
+    match message_type:
         case "user":
-            return UserMessage(content=data["message"]["content"])
+            try:
+                return UserMessage(content=data["message"]["content"])
+            except KeyError as e:
+                logger.error("Missing required field in user message: %s", e)
+                return None
 
         case "assistant":
-            content_blocks: list[ContentBlock] = []
-            for block in data["message"]["content"]:
-                match block["type"]:
-                    case "text":
-                        content_blocks.append(TextBlock(text=block["text"]))
-                    case "tool_use":
-                        content_blocks.append(
-                            ToolUseBlock(
-                                id=block["id"],
-                                name=block["name"],
-                                input=block["input"],
+            try:
+                content_blocks: list[ContentBlock] = []
+                for block in data["message"]["content"]:
+                    match block["type"]:
+                        case "text":
+                            content_blocks.append(TextBlock(text=block["text"]))
+                        case "tool_use":
+                            content_blocks.append(
+                                ToolUseBlock(
+                                    id=block["id"],
+                                    name=block["name"],
+                                    input=block["input"],
+                                )
                             )
-                        )
-                    case "tool_result":
-                        content_blocks.append(
-                            ToolResultBlock(
-                                tool_use_id=block["tool_use_id"],
-                                content=block.get("content"),
-                                is_error=block.get("is_error"),
+                        case "tool_result":
+                            content_blocks.append(
+                                ToolResultBlock(
+                                    tool_use_id=block["tool_use_id"],
+                                    content=block.get("content"),
+                                    is_error=block.get("is_error"),
+                                )
                             )
-                        )
 
-            return AssistantMessage(content=content_blocks)
+                return AssistantMessage(content=content_blocks)
+            except KeyError as e:
+                logger.error("Missing required field in assistant message: %s", e)
+                return None
 
         case "system":
-            return SystemMessage(
-                subtype=data["subtype"],
-                data=data,
-            )
+            try:
+                return SystemMessage(
+                    subtype=data["subtype"],
+                    data=data,
+                )
+            except KeyError as e:
+                logger.error("Missing required field in system message: %s", e)
+                return None
 
         case "result":
-            return ResultMessage(
-                subtype=data["subtype"],
-                duration_ms=data["duration_ms"],
-                duration_api_ms=data["duration_api_ms"],
-                is_error=data["is_error"],
-                num_turns=data["num_turns"],
-                session_id=data["session_id"],
-                total_cost_usd=data.get("total_cost_usd"),
-                usage=data.get("usage"),
-                result=data.get("result"),
-            )
+            try:
+                return ResultMessage(
+                    subtype=data["subtype"],
+                    duration_ms=data["duration_ms"],
+                    duration_api_ms=data["duration_api_ms"],
+                    is_error=data["is_error"],
+                    num_turns=data["num_turns"],
+                    session_id=data["session_id"],
+                    total_cost_usd=data.get("total_cost_usd"),
+                    usage=data.get("usage"),
+                    result=data.get("result"),
+                )
+            except KeyError as e:
+                logger.error("Missing required field in result message: %s", e)
+                return None
 
         case _:
+            logger.debug("Unknown message type: %s", message_type)
             return None

src/claude_code_sdk/_internal/transport/subprocess_cli.py

@@ -306,7 +306,12 @@ async def receive_messages(self) -> AsyncIterator[dict[str, Any]]:
                             yield data
                         except GeneratorExit:
                             return
-                    except json.JSONDecodeError:
+                    except json.JSONDecodeError as e:
+                        logger.warning(
+                            f"Failed to parse JSON from CLI output: {e}. Buffer content: {json_buffer[:200]}..."
+                        )
+                        # Clear buffer to avoid repeated parse attempts on malformed data
+                        json_buffer = ""
                         continue
 
         except anyio.ClosedResourceError:

src/claude_code_sdk/client.py

@@ -128,19 +128,37 @@ async def receive_messages(self) -> AsyncIterator[Message]:
             if message:
                 yield message
 
-    async def send_message(self, content: str, session_id: str = "default") -> None:
-        """Send a new message in streaming mode."""
+    async def send_message(self, prompt: str | AsyncIterable[dict[str, Any]], session_id: str = "default") -> None:
+        """
+        Send a new message in streaming mode.
+
+        Args:
+            prompt: Either a string message or an async iterable of message dictionaries
+            session_id: Session identifier for the conversation
+        """
         if not self._transport:
             raise CLIConnectionError("Not connected. Call connect() first.")
 
-        message = {
-            "type": "user",
-            "message": {"role": "user", "content": content},
-            "parent_tool_use_id": None,
-            "session_id": session_id,
-        }
-
-        await self._transport.send_request([message], {"session_id": session_id})
+        # Handle string prompts
+        if isinstance(prompt, str):
+            message = {
+                "type": "user",
+                "message": {"role": "user", "content": prompt},
+                "parent_tool_use_id": None,
+                "session_id": session_id,
+            }
+            await self._transport.send_request([message], {"session_id": session_id})
+        else:
+            # Handle AsyncIterable prompts
+            messages = []
+            async for msg in prompt:
+                # Ensure session_id is set on each message
+                if "session_id" not in msg:
+                    msg["session_id"] = session_id
+                messages.append(msg)
+
+            if messages:
+                await self._transport.send_request(messages, {"session_id": session_id})
 
     async def interrupt(self) -> None:
         """Send interrupt signal (only works with streaming mode)."""
@@ -150,19 +168,24 @@ async def interrupt(self) -> None:
 
     async def receive_response(self) -> AsyncIterator[Message]:
         """
-        Receive messages from Claude until a ResultMessage is received.
+        Receive messages from Claude until and including a ResultMessage.
 
-        This is an async iterator that yields all messages including the final ResultMessage.
-        It's a convenience method over receive_messages() that automatically stops iteration
-        after receiving a ResultMessage.
+        This async iterator yields all messages in sequence and automatically terminates
+        after yielding a ResultMessage (which indicates the response is complete).
+        It's a convenience method over receive_messages() for single-response workflows.
+
+        **Stopping Behavior:**
+        - Yields each message as it's received
+        - Terminates immediately after yielding a ResultMessage
+        - The ResultMessage IS included in the yielded messages
+        - If no ResultMessage is received, the iterator continues indefinitely
 
         Yields:
             Message: Each message received (UserMessage, AssistantMessage, SystemMessage, ResultMessage)
 
         Example:
             ```python
             async with ClaudeSDKClient() as client:
-                # Send message and process response
                 await client.send_message("What's the capital of France?")
 
                 async for msg in client.receive_response():
@@ -172,14 +195,12 @@ async def receive_response(self) -> AsyncIterator[Message]:
                                 print(f"Claude: {block.text}")
                     elif isinstance(msg, ResultMessage):
                         print(f"Cost: ${msg.total_cost_usd:.4f}")
+                        # Iterator will terminate after this message
             ```
 
         Note:
-            The iterator will automatically stop after yielding a ResultMessage.
-            If you need to collect all messages into a list, use:
-            ```python
-            messages = [msg async for msg in client.receive_response()]
-            ```
+            To collect all messages: `messages = [msg async for msg in client.receive_response()]`
+            The final message in the list will always be a ResultMessage.
         """
         async for message in self.receive_messages():
             yield message