diff --git a/examples/streaming_mode_trio.py b/examples/streaming_mode_trio.py new file mode 100644 index 00000000..50366f7f --- /dev/null +++ b/examples/streaming_mode_trio.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 +""" +Example of multi-turn conversation using trio with the Claude SDK. + +This demonstrates how to use the ClaudeSDKClient with trio for interactive, +stateful conversations where you can send follow-up messages based on +Claude's responses. +""" + +import trio + +from claude_code_sdk import ( + AssistantMessage, + ClaudeCodeOptions, + ClaudeSDKClient, + ResultMessage, + SystemMessage, + TextBlock, + UserMessage, +) + + +def display_message(msg): + """Standardized message display function. + + - UserMessage: "User: " + - AssistantMessage: "Claude: " + - SystemMessage: ignored + - ResultMessage: "Result ended" + cost if available + """ + if isinstance(msg, UserMessage): + for block in msg.content: + if isinstance(block, TextBlock): + print(f"User: {block.text}") + elif isinstance(msg, AssistantMessage): + for block in msg.content: + if isinstance(block, TextBlock): + print(f"Claude: {block.text}") + elif isinstance(msg, SystemMessage): + # Ignore system messages + pass + elif isinstance(msg, ResultMessage): + print("Result ended") + + +async def multi_turn_conversation(): + """Example of a multi-turn conversation using trio.""" + async with ClaudeSDKClient( + options=ClaudeCodeOptions(model="claude-3-5-sonnet-20241022") + ) as client: + print("=== Multi-turn Conversation with Trio ===\n") + + # First turn: Simple math question + print("User: What's 15 + 27?") + await client.query("What's 15 + 27?") + + async for message in client.receive_response(): + display_message(message) + print() + + # Second turn: Follow-up calculation + print("User: Now multiply that result by 2") + await client.query("Now multiply that result by 2") + + async for message in client.receive_response(): + display_message(message) + print() + + # Third turn: One more operation + print("User: Divide that by 7 and round to 2 decimal places") + await client.query("Divide that by 7 and round to 2 decimal places") + + async for message in client.receive_response(): + display_message(message) + + print("\nConversation complete!") + + +if __name__ == "__main__": + trio.run(multi_turn_conversation) diff --git a/src/claude_code_sdk/_internal/transport/subprocess_cli.py b/src/claude_code_sdk/_internal/transport/subprocess_cli.py index 34b70342..56a2de9a 100644 --- a/src/claude_code_sdk/_internal/transport/subprocess_cli.py +++ b/src/claude_code_sdk/_internal/transport/subprocess_cli.py @@ -45,6 +45,7 @@ def __init__( self._pending_control_responses: dict[str, dict[str, Any]] = {} self._request_counter = 0 self._close_stdin_after_prompt = close_stdin_after_prompt + self._task_group: anyio.abc.TaskGroup | None = None def _find_cli(self) -> str: """Find Claude Code CLI binary.""" @@ -160,9 +161,9 @@ async def connect(self) -> None: if self._process.stdin: self._stdin_stream = TextSendStream(self._process.stdin) # Start streaming messages to stdin in background - import asyncio - - asyncio.create_task(self._stream_to_stdin()) + self._task_group = anyio.create_task_group() + await self._task_group.__aenter__() + self._task_group.start_soon(self._stream_to_stdin) else: # String mode: close stdin immediately (backward compatible) if self._process.stdin: @@ -183,6 +184,12 @@ async def disconnect(self) -> None: if not self._process: return + # Cancel task group if it exists + if self._task_group: + self._task_group.cancel_scope.cancel() + await self._task_group.__aexit__(None, None, None) + self._task_group = None + if self._process.returncode is None: try: self._process.terminate() diff --git a/src/claude_code_sdk/client.py b/src/claude_code_sdk/client.py index 8e86ba7c..cf668fdc 100644 --- a/src/claude_code_sdk/client.py +++ b/src/claude_code_sdk/client.py @@ -63,7 +63,7 @@ class ClaudeSDKClient: await client.query("Count to 1000") # Interrupt after 2 seconds - await asyncio.sleep(2) + await anyio.sleep(2) await client.interrupt() # Send new instruction