Fix examples

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

Author: dicksontsai

examples/streaming_mode.py

@@ -13,6 +13,7 @@
     AssistantMessage,
     ClaudeCodeOptions,
     ClaudeSDKClient,
+    CLIConnectionError,
     ResultMessage,
     TextBlock,
 )
@@ -27,18 +28,13 @@ async def example_basic_streaming():
         await client.send_message("What is 2+2?")
 
         # Receive complete response using the helper method
-        messages, result = await client.receive_response()
-
-        # Extract text from assistant's response
-        for msg in messages:
+        async for msg in client.receive_response():
             if isinstance(msg, AssistantMessage):
                 for block in msg.content:
                     if isinstance(block, TextBlock):
                         print(f"Claude: {block.text}")
-
-        # Print cost if available
-        if result and result.total_cost_usd:
-            print(f"Cost: ${result.total_cost_usd:.4f}")
+            elif isinstance(msg, ResultMessage) and msg.total_cost_usd:
+                print(f"Cost: ${msg.total_cost_usd:.4f}")
 
     print("Session ended\n")
 
@@ -52,32 +48,22 @@ async def example_multi_turn_conversation():
         print("User: What's the capital of France?")
         await client.send_message("What's the capital of France?")
 
-        messages, _ = await client.receive_response()
-
         # Extract and print response
-        for msg in messages:
-            if isinstance(msg, AssistantMessage):
-                for block in msg.content:
-                    if isinstance(block, TextBlock):
-                        print(f"Claude: {block.text}")
+        async for msg in client.receive_response():
+            content_blocks = getattr(msg, 'content', [])
+            for block in content_blocks:
+                if isinstance(block, TextBlock):
+                    print(f"{block.text}")
 
         # Second turn - follow-up
         print("\nUser: What's the population of that city?")
         await client.send_message("What's the population of that city?")
 
-        messages, _ = await client.receive_response()
-
-        for msg in messages:
-            if isinstance(msg, AssistantMessage):
-                for block in msg.content:
-                    if isinstance(block, TextBlock):
-                        print(f"Claude: {block.text}")
-
-        for msg in messages:
-            if isinstance(msg, AssistantMessage):
-                for block in msg.content:
-                    if isinstance(block, TextBlock):
-                        print(f"Claude: {block.text}")
+        async for msg in client.receive_response():
+            content_blocks = getattr(msg, 'content', [])
+            for block in content_blocks:
+                if isinstance(block, TextBlock):
+                    print(f"{block.text}")
 
     print("\nConversation ended\n")
 
@@ -102,7 +88,7 @@ async def receive_messages():
         questions = [
             "What is 2 + 2?",
             "What is the square root of 144?",
-            "What is 15% of 80?",
+            "What is 10% of 80?",
         ]
 
         for question in questions:
@@ -168,9 +154,7 @@ async def consume_messages():
         await client.send_message("Never mind, just tell me a quick joke")
 
         # Get the joke
-        messages, result = await client.receive_response()
-
-        for msg in messages:
+        async for msg in client.receive_response():
             if isinstance(msg, AssistantMessage):
                 for block in msg.content:
                     if isinstance(block, TextBlock):
@@ -233,10 +217,8 @@ async def example_with_options():
             "Create a simple hello.txt file with a greeting message"
         )
 
-        messages, result = await client.receive_response()
-
         tool_uses = []
-        for msg in messages:
+        async for msg in client.receive_response():
             if isinstance(msg, AssistantMessage):
                 for block in msg.content:
                     if isinstance(block, TextBlock):
@@ -257,48 +239,34 @@ async def example_error_handling():
     client = ClaudeSDKClient()
 
     try:
-        # Connect with custom stream
-        async def message_stream():
-            yield {
-                "type": "user",
-                "message": {"role": "user", "content": "Hello"},
-                "parent_tool_use_id": None,
-                "session_id": "error-demo",
-            }
-
-        await client.connect(message_stream())
+        await client.connect()
 
-        # Create a background task to consume messages (required for interrupt to work)
-        consume_task = None
+        # Send a message that will take time to process
+        await client.send_message("Run a bash sleep command for 60 seconds")
 
-        async def consume_messages():
-            """Background message consumer."""
-            async for msg in client.receive_messages():
-                if isinstance(msg, AssistantMessage):
-                    print("Received response from Claude")
-
-        # Receive messages with timeout
+        # Try to receive response with a short timeout
         try:
-            # Start consuming messages in background
-            consume_task = asyncio.create_task(consume_messages())
-
-            # Wait for response with timeout
-            await asyncio.wait_for(consume_task, timeout=30.0)
+            messages = []
+            async with asyncio.timeout(10.0):
+                async for msg in client.receive_response():
+                    messages.append(msg)
+                    if isinstance(msg, AssistantMessage):
+                        for block in msg.content:
+                            if isinstance(block, TextBlock):
+                                print(f"Claude: {block.text[:50]}...")
+                    elif isinstance(msg, ResultMessage):
+                        print("Received complete response")
+                        break
 
         except asyncio.TimeoutError:
-            print("Response timeout - sending interrupt")
-            # Note: interrupt requires active message consumption
-            # Since we're already consuming in the background task, interrupt will work
-            await client.interrupt()
+            print("\nResponse timeout after 10 seconds - demonstrating graceful handling")
+            print(f"Received {len(messages)} messages before timeout")
 
-            # Cancel the consume task
-            if consume_task:
-                consume_task.cancel()
-                with contextlib.suppress(asyncio.CancelledError):
-                    await consume_task
+    except CLIConnectionError as e:
+        print(f"Connection error: {e}")
 
     except Exception as e:
-        print(f"Error: {e}")
+        print(f"Unexpected error: {e}")
 
     finally:
         # Always disconnect

examples/streaming_mode_ipython.py

@@ -10,17 +10,18 @@
 # BASIC STREAMING
 # ============================================================================
 
-from claude_code_sdk import ClaudeSDKClient, AssistantMessage, TextBlock
+from claude_code_sdk import ClaudeSDKClient, AssistantMessage, TextBlock, ResultMessage
 
 async with ClaudeSDKClient() as client:
     await client.send_message("What is 2+2?")
-    messages, result = await client.receive_response()
 
-    for msg in messages:
+    async for msg in client.receive_response():
         if isinstance(msg, AssistantMessage):
             for block in msg.content:
                 if isinstance(block, TextBlock):
                     print(f"Claude: {block.text}")
+        elif isinstance(msg, ResultMessage) and msg.total_cost_usd:
+            print(f"Cost: ${msg.total_cost_usd:.4f}")
 
 
 # ============================================================================
@@ -31,18 +32,17 @@
 from claude_code_sdk import ClaudeSDKClient, AssistantMessage, TextBlock
 
 async with ClaudeSDKClient() as client:
-    async def receive_response():
-        messages, _ = await client.receive_response()
-        for msg in messages:
+    async def send_and_receive(prompt):
+        await client.send_message(prompt)
+        async for msg in client.receive_response():
             if isinstance(msg, AssistantMessage):
                 for block in msg.content:
                     if isinstance(block, TextBlock):
                         print(f"Claude: {block.text}")
 
-    await client.send_message("Tell me a short joke")
-    await receive_response()
-    await client.send_message("Now tell me a fun fact")
-    await receive_response()
+    await send_and_receive("Tell me a short joke")
+    print("\n---\n")
+    await send_and_receive("Now tell me a fun fact")
 
 
 # ============================================================================
@@ -58,8 +58,7 @@ async def receive_response():
 
 # Helper to get response
 async def get_response():
-    messages, result = await client.receive_response()
-    for msg in messages:
+    async for msg in client.receive_response():
         if isinstance(msg, AssistantMessage):
             for block in msg.content:
                 if isinstance(block, TextBlock):
@@ -123,9 +122,8 @@ async def consume_messages():
     # Send a new message after interrupt
     print("\n--- After interrupt, sending new message ---\n")
     await client.send_message("Just say 'Hello! I was interrupted.'")
-    messages, result = await client.receive_response()
 
-    for msg in messages:
+    async for msg in client.receive_response():
         if isinstance(msg, AssistantMessage):
             for block in msg.content:
                 if isinstance(block, TextBlock):
@@ -142,12 +140,41 @@ async def consume_messages():
     async with ClaudeSDKClient() as client:
         await client.send_message("Run a bash sleep command for 60 seconds")
 
-        # Timeout after 30 seconds
-        messages, result = await asyncio.wait_for(
-            client.receive_response(), timeout=20.0
-        )
+        # Timeout after 20 seconds
+        messages = []
+        async with asyncio.timeout(20.0):
+            async for msg in client.receive_response():
+                messages.append(msg)
+                if isinstance(msg, AssistantMessage):
+                    for block in msg.content:
+                        if isinstance(block, TextBlock):
+                            print(f"Claude: {block.text}")
 
 except asyncio.TimeoutError:
-    print("Request timed out")
+    print("Request timed out after 20 seconds")
 except Exception as e:
-    print(f"Error: {e}")
+    print(f"Error: {e}")
+
+
+# ============================================================================
+# COLLECTING ALL MESSAGES INTO A LIST
+# ============================================================================
+
+from claude_code_sdk import ClaudeSDKClient, AssistantMessage, TextBlock, ResultMessage
+
+async with ClaudeSDKClient() as client:
+    await client.send_message("What are the primary colors?")
+
+    # Collect all messages into a list
+    messages = [msg async for msg in client.receive_response()]
+
+    # Process them afterwards
+    for msg in messages:
+        if isinstance(msg, AssistantMessage):
+            for block in msg.content:
+                if isinstance(block, TextBlock):
+                    print(f"Claude: {block.text}")
+        elif isinstance(msg, ResultMessage):
+            print(f"Total messages: {len(messages)}")
+            if msg.total_cost_usd:
+                print(f"Cost: ${msg.total_cost_usd:.4f}")

src/claude_code_sdk/client.py

@@ -146,50 +146,43 @@ async def interrupt(self) -> None:
             raise CLIConnectionError("Not connected. Call connect() first.")
         await self._transport.interrupt()
 
-    async def receive_response(self) -> tuple[list[Message], ResultMessage | None]:
+    async def receive_response(self) -> AsyncIterator[Message]:
         """
-        Receive a complete response from Claude, collecting all messages until ResultMessage.
+        Receive messages from Claude until a ResultMessage is received.
 
-        Compared to receive_messages(), this is a convenience method that
-        handles the common pattern of receiving messages until Claude completes
-        its response. It collects all messages and returns them along with the
-        final 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.
 
-        Returns:
-            tuple: A tuple of (messages, result) where:
-                - messages: List of all messages received (UserMessage, AssistantMessage, SystemMessage)
-                - result: The final ResultMessage if received, None if stream ended without result
+        Yields:
+            Message: Each message received (UserMessage, AssistantMessage, SystemMessage, ResultMessage)
 
         Example:
             ```python
             async with ClaudeSDKClient() as client:
-                # First turn
+                # Send message and process response
                 await client.send_message("What's the capital of France?")
-                messages, result = await client.receive_response()
 
-                # Extract assistant's response
-                for msg in messages:
+                async for msg in client.receive_response():
                     if isinstance(msg, AssistantMessage):
                         for block in msg.content:
                             if isinstance(block, TextBlock):
                                 print(f"Claude: {block.text}")
+                    elif isinstance(msg, ResultMessage):
+                        print(f"Cost: ${msg.total_cost_usd:.4f}")
+            ```
 
-                # Second turn
-                await client.send_message("What's the population?")
-                messages, result = await client.receive_response()
-                # ... process response
+        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()]
             ```
         """
-        from .types import ResultMessage
-
-        messages = []
         async for message in self.receive_messages():
-            messages.append(message)
+            yield message
             if isinstance(message, ResultMessage):
-                return messages, message
-
-        # Stream ended without ResultMessage
-        return messages, None
+                return
 
     async def disconnect(self) -> None:
         """Disconnect from Claude."""