Commit bd3b7a6
authored
fix: spawn wait_for_result_and_end_input as background task for string prompts (#780)
## Problem
`query()` with a string prompt and hooks/MCP servers deadlocks once the
internal 100-slot anyio message buffer fills up (~50 tool calls). Each
tool call produces ~2 messages, so the buffer fills after about 50 tool
calls.
### Root cause
For string prompts, `client.py:141` awaited
`wait_for_result_and_end_input()` **before** `receive_messages()`
started draining the buffer:
```python
if isinstance(prompt, str):
await chosen_transport.write(json.dumps(user_message) + "\n")
await query.wait_for_result_and_end_input() # blocks until "result" arrives
async for data in query.receive_messages(): # buffer drain starts here
```
Meanwhile `_read_messages()` keeps reading CLI stdout and pushing into
the 100-slot channel. After ~50 tool calls the channel is full and
`_message_send.send()` blocks. Now `_read_messages` can't read anything
else from stdout, including the "result" message that
`wait_for_result_and_end_input` needs — deadlock.
## Fix
Spawn `wait_for_result_and_end_input()` as a background task instead of
awaiting it inline. This matches the existing `AsyncIterable` path which
already uses `spawn_task(stream_input())`, and allows
`receive_messages()` to start draining the buffer immediately.
```python
# Before (deadlocks)
await query.wait_for_result_and_end_input()
# After (concurrent)
query.spawn_task(query.wait_for_result_and_end_input())
```
## Testing
- Added regression test verifying `spawn_task` is called instead of
direct `await`
- Fixed existing test warnings from unawaited mock coroutines
- All 425 tests pass
- Lint, format, and mypy all clean
Fixes #7791 parent 566e41f commit bd3b7a6
2 files changed
Lines changed: 64 additions & 1 deletion
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
138 | 138 | | |
139 | 139 | | |
140 | 140 | | |
141 | | - | |
| 141 | + | |
142 | 142 | | |
143 | 143 | | |
144 | 144 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
155 | 155 | | |
156 | 156 | | |
157 | 157 | | |
| 158 | + | |
| 159 | + | |
| 160 | + | |
| 161 | + | |
| 162 | + | |
| 163 | + | |
158 | 164 | | |
159 | 165 | | |
160 | 166 | | |
| |||
192 | 198 | | |
193 | 199 | | |
194 | 200 | | |
| 201 | + | |
| 202 | + | |
| 203 | + | |
| 204 | + | |
| 205 | + | |
| 206 | + | |
| 207 | + | |
| 208 | + | |
| 209 | + | |
| 210 | + | |
| 211 | + | |
| 212 | + | |
| 213 | + | |
| 214 | + | |
| 215 | + | |
| 216 | + | |
| 217 | + | |
| 218 | + | |
| 219 | + | |
| 220 | + | |
| 221 | + | |
| 222 | + | |
| 223 | + | |
| 224 | + | |
| 225 | + | |
| 226 | + | |
| 227 | + | |
| 228 | + | |
| 229 | + | |
| 230 | + | |
| 231 | + | |
| 232 | + | |
| 233 | + | |
| 234 | + | |
| 235 | + | |
| 236 | + | |
| 237 | + | |
| 238 | + | |
| 239 | + | |
| 240 | + | |
| 241 | + | |
| 242 | + | |
| 243 | + | |
| 244 | + | |
| 245 | + | |
| 246 | + | |
| 247 | + | |
| 248 | + | |
| 249 | + | |
| 250 | + | |
| 251 | + | |
| 252 | + | |
| 253 | + | |
| 254 | + | |
| 255 | + | |
| 256 | + | |
| 257 | + | |
0 commit comments