Skip to content

Commit 5f8351f

Browse files
authored
Make streaming implementation trio-compatible (#84)
## Summary - Replace asyncio.create_task() with anyio task group for trio compatibility - Update client.py docstring example to use anyio.sleep - Add trio example demonstrating multi-turn conversation ## Details The SDK already uses anyio for most async operations, but one line was using asyncio.create_task() which broke trio compatibility. This PR fixes that by using anyio's task group API with proper lifecycle management. ### Changes: 1. **subprocess_cli.py**: Replace asyncio.create_task() with anyio task group, ensuring proper cleanup on disconnect 2. **client.py**: Update docstring example to use anyio.sleep instead of asyncio.sleep 3. **streaming_mode_trio.py**: Add new example showing how to use the SDK with trio ## Test plan - [x] All existing tests pass - [x] Manually tested with trio runtime (created test script that successfully runs multi-turn conversation) - [x] Linting and type checking pass 🤖 Generated with [Claude Code](https://claude.ai/code)
1 parent b25cdf8 commit 5f8351f

File tree

3 files changed

+91
-4
lines changed

3 files changed

+91
-4
lines changed

examples/streaming_mode_trio.py

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
#!/usr/bin/env python3
2+
"""
3+
Example of multi-turn conversation using trio with the Claude SDK.
4+
5+
This demonstrates how to use the ClaudeSDKClient with trio for interactive,
6+
stateful conversations where you can send follow-up messages based on
7+
Claude's responses.
8+
"""
9+
10+
import trio
11+
12+
from claude_code_sdk import (
13+
AssistantMessage,
14+
ClaudeCodeOptions,
15+
ClaudeSDKClient,
16+
ResultMessage,
17+
SystemMessage,
18+
TextBlock,
19+
UserMessage,
20+
)
21+
22+
23+
def display_message(msg):
24+
"""Standardized message display function.
25+
26+
- UserMessage: "User: <content>"
27+
- AssistantMessage: "Claude: <content>"
28+
- SystemMessage: ignored
29+
- ResultMessage: "Result ended" + cost if available
30+
"""
31+
if isinstance(msg, UserMessage):
32+
for block in msg.content:
33+
if isinstance(block, TextBlock):
34+
print(f"User: {block.text}")
35+
elif isinstance(msg, AssistantMessage):
36+
for block in msg.content:
37+
if isinstance(block, TextBlock):
38+
print(f"Claude: {block.text}")
39+
elif isinstance(msg, SystemMessage):
40+
# Ignore system messages
41+
pass
42+
elif isinstance(msg, ResultMessage):
43+
print("Result ended")
44+
45+
46+
async def multi_turn_conversation():
47+
"""Example of a multi-turn conversation using trio."""
48+
async with ClaudeSDKClient(
49+
options=ClaudeCodeOptions(model="claude-3-5-sonnet-20241022")
50+
) as client:
51+
print("=== Multi-turn Conversation with Trio ===\n")
52+
53+
# First turn: Simple math question
54+
print("User: What's 15 + 27?")
55+
await client.query("What's 15 + 27?")
56+
57+
async for message in client.receive_response():
58+
display_message(message)
59+
print()
60+
61+
# Second turn: Follow-up calculation
62+
print("User: Now multiply that result by 2")
63+
await client.query("Now multiply that result by 2")
64+
65+
async for message in client.receive_response():
66+
display_message(message)
67+
print()
68+
69+
# Third turn: One more operation
70+
print("User: Divide that by 7 and round to 2 decimal places")
71+
await client.query("Divide that by 7 and round to 2 decimal places")
72+
73+
async for message in client.receive_response():
74+
display_message(message)
75+
76+
print("\nConversation complete!")
77+
78+
79+
if __name__ == "__main__":
80+
trio.run(multi_turn_conversation)

src/claude_code_sdk/_internal/transport/subprocess_cli.py

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ def __init__(
4545
self._pending_control_responses: dict[str, dict[str, Any]] = {}
4646
self._request_counter = 0
4747
self._close_stdin_after_prompt = close_stdin_after_prompt
48+
self._task_group: anyio.abc.TaskGroup | None = None
4849

4950
def _find_cli(self) -> str:
5051
"""Find Claude Code CLI binary."""
@@ -160,9 +161,9 @@ async def connect(self) -> None:
160161
if self._process.stdin:
161162
self._stdin_stream = TextSendStream(self._process.stdin)
162163
# Start streaming messages to stdin in background
163-
import asyncio
164-
165-
asyncio.create_task(self._stream_to_stdin())
164+
self._task_group = anyio.create_task_group()
165+
await self._task_group.__aenter__()
166+
self._task_group.start_soon(self._stream_to_stdin)
166167
else:
167168
# String mode: close stdin immediately (backward compatible)
168169
if self._process.stdin:
@@ -183,6 +184,12 @@ async def disconnect(self) -> None:
183184
if not self._process:
184185
return
185186

187+
# Cancel task group if it exists
188+
if self._task_group:
189+
self._task_group.cancel_scope.cancel()
190+
await self._task_group.__aexit__(None, None, None)
191+
self._task_group = None
192+
186193
if self._process.returncode is None:
187194
try:
188195
self._process.terminate()

src/claude_code_sdk/client.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ class ClaudeSDKClient:
6363
await client.query("Count to 1000")
6464
6565
# Interrupt after 2 seconds
66-
await asyncio.sleep(2)
66+
await anyio.sleep(2)
6767
await client.interrupt()
6868
6969
# Send new instruction

0 commit comments

Comments
 (0)