fix(query): implement structured concurrency to prevent query() deadlocks (fixes #1136)

src/claude_agent_sdk/_internal/client.py

@@ -6,6 +6,9 @@
 from dataclasses import asdict, replace
 from typing import Any
 
+import anyio
+
+from .._errors import ProcessError
 from ..types import (
     ClaudeAgentOptions,
     HookEvent,
@@ -201,33 +204,63 @@ async def _on_mirror_error(key: Any, error: str) -> None:
             )
 
         try:
-            # Start reading messages
-            await query.start()
+            read_done = anyio.Event()
 
-            # Always initialize to send agents via stdin (matching TypeScript SDK)
-            await query.initialize()
+            async def _run_read_messages() -> None:
+                try:
+                    await query._read_messages()
+                finally:
+                    read_done.set()
 
-            # Handle prompt input
-            if isinstance(prompt, str):
-                # For string prompts, write user message to stdin after initialize
-                # (matching TypeScript SDK behavior)
-                user_message = {
-                    "type": "user",
-                    "session_id": "",
-                    "message": {"role": "user", "content": prompt},
-                    "parent_tool_use_id": None,
-                }
-                await chosen_transport.write(json.dumps(user_message) + "\n")
-                query.spawn_task(query.wait_for_result_and_end_input())
-            elif isinstance(prompt, AsyncIterable):
-                # Stream input in background for async iterables
-                query.spawn_task(query.stream_input(prompt))
-
-            # Yield parsed messages, skipping unknown message types
-            async for data in query.receive_messages():
-                message = parse_message(data)
-                if message is not None:
-                    yield message
+            async def _monitor_process() -> None:
+                if hasattr(chosen_transport, "wait"):
+                    returncode = await chosen_transport.wait()
+                    if (
+                        type(returncode) is int
+                        and returncode != 0
+                        and not read_done.is_set()
+                    ):
+                        raise ProcessError(
+                            f"Claude Code process exited unexpectedly with code {returncode}",
+                            exit_code=returncode,
+                            stderr="",
+                        )
+
+            try:
+                async with anyio.create_task_group() as tg:
+                    tg.start_soon(_run_read_messages)
+                    tg.start_soon(_monitor_process)
+
+                    # Always initialize to send agents via stdin (matching TypeScript SDK)
+                    await query.initialize()
+
+                    if isinstance(prompt, str):
+                        user_message = {
+                            "type": "user",
+                            "session_id": "",
+                            "message": {"role": "user", "content": prompt},
+                            "parent_tool_use_id": None,
+                        }
+                        await chosen_transport.write(json.dumps(user_message) + "\n")
+                        tg.start_soon(query.wait_for_result_and_end_input)
+                    elif isinstance(prompt, AsyncIterable):
+                        tg.start_soon(query.stream_input, prompt)
+
+                    async for data in query.receive_messages():
+                        message = parse_message(data)
+                        if message is not None:
+                            yield message
+
+                    tg.cancel_scope.cancel()
+            except BaseException as e:
+                if isinstance(e, GeneratorExit):
+                    raise
+                if hasattr(e, "exceptions") and any(
+                    isinstance(exc, GeneratorExit)
+                    for exc in getattr(e, "exceptions", ())
+                ):
+                    raise GeneratorExit() from None
+                raise
 
         finally:
             await query.close()

src/claude_agent_sdk/_internal/query.py

@@ -832,15 +832,12 @@ async def stream_input(self, stream: AsyncIterable[dict[str, Any]]) -> None:
         If SDK MCP servers or hooks are present, waits for the first result
         before closing stdin to allow bidirectional control protocol communication.
         """
-        try:
-            async for message in stream:
-                if self._closed:
-                    break
-                await self.transport.write(json.dumps(message) + "\n")
+        async for message in stream:
+            if self._closed:
+                break
+            await self.transport.write(json.dumps(message) + "\n")
 
-            await self.wait_for_result_and_end_input()
-        except Exception as e:
-            logger.debug(f"Error streaming input: {e}")
+        await self.wait_for_result_and_end_input()
 
     async def receive_messages(self) -> AsyncIterator[dict[str, Any]]:
         """Receive SDK messages (not control messages)."""

src/claude_agent_sdk/_internal/transport/__init__.py

@@ -83,5 +83,13 @@ async def end_input(self) -> None:
         """End the input stream (close stdin for process transports)."""
         pass
 
+    async def wait(self) -> int | None:
+        """Wait for transport process to exit.
+
+        Returns:
+            Process exit code, or None if transport has no subprocess.
+        """
+        return None
+
 
 __all__ = ["Transport"]

src/claude_agent_sdk/_internal/transport/subprocess_cli.py

@@ -15,6 +15,7 @@
 from typing import Any, cast
 
 import anyio
+import anyio.to_thread
 from anyio.abc import Process
 from anyio.streams.text import TextReceiveStream, TextSendStream
 
@@ -433,7 +434,7 @@ def _build_command(self) -> list[str]:
                 cmd.append(f"--{flag}")
             else:
                 # Flag with value
-                cmd.extend([f"--{flag}", str(value)])
+                cmd.extend([f"--{flag}", value])
 
         # Resolve thinking config -> --thinking / --max-thinking-tokens
         # `thinking` takes precedence over the deprecated `max_thinking_tokens`
@@ -828,6 +829,16 @@ def guard(length: int) -> None:
             )
             raise self._exit_error
 
+    async def wait(self) -> int | None:
+        """Wait for the transport process to exit.
+
+        Returns:
+            Process exit code, or None if transport has no subprocess.
+        """
+        if self._process is not None:
+            return await self._process.wait()
+        return None
+
     async def _check_claude_version(self) -> None:
         """Check Claude Code version and warn if below minimum."""
         if self._cli_path is None:

tests/test_client.py

@@ -34,7 +34,9 @@ async def mock_generator():
 
                 assert len(messages) == 1
                 assert isinstance(messages[0], AssistantMessage)
-                assert messages[0].content[0].text == "4"
+                block = messages[0].content[0]
+                assert isinstance(block, TextBlock)
+                assert block.text == "4"
 
         anyio.run(_test)
 
@@ -227,13 +229,6 @@ async def _test():
                 mock_query.initialize = AsyncMock()
                 mock_query.close = AsyncMock()
                 mock_query.close_receive_stream = Mock()
-                mock_query._tg = None
-
-                def _consume_coro(coro):
-                    coro.close()
-                    return Mock()
-
-                mock_query.spawn_task = Mock(side_effect=_consume_coro)
 
                 async def mock_receive():
                     yield {
@@ -251,11 +246,7 @@ async def mock_receive():
                 async for _ in query(prompt="test", options=ClaudeAgentOptions()):
                     pass
 
-                mock_query.spawn_task.assert_called_once()
-                assert not mock_query.wait_for_result_and_end_input.await_args_list, (
-                    "wait_for_result_and_end_input should be spawned as a task, "
-                    "not awaited directly"
-                )
+                mock_query.wait_for_result_and_end_input.assert_called_once()
 
         anyio.run(_test)
 

tests/test_query.py

@@ -175,7 +175,7 @@ async def _test():
             server = _make_greet_server()
             mock_transport = _make_mock_transport(messages=_ASSISTANT_AND_RESULT)
 
-            call_order = []
+            call_order: list[tuple[str, ...]] = []
             original_write = mock_transport.write
 
             async def tracking_write(data):
@@ -316,7 +316,7 @@ def test_string_prompt_with_hooks_waits_for_result(self):
         async def _test():
             mock_transport = _make_mock_transport(messages=_ASSISTANT_AND_RESULT)
 
-            call_order = []
+            call_order: list[tuple[str, ...]] = []
 
             async def tracking_write(data):
                 call_order.append(("write", data))
@@ -371,7 +371,7 @@ async def _test():
             server = _make_greet_server()
             mock_transport = _make_mock_transport(messages=_ASSISTANT_AND_RESULT)
 
-            call_order = []
+            call_order: list[tuple[str, ...]] = []
             original_write = mock_transport.write
 
             async def tracking_write(data):
@@ -1240,3 +1240,75 @@ async def _test():
             assert isinstance(q.pending_control_results["req_1"], ProcessError)
 
         anyio.run(_test)
+
+
+class TestStructuredConcurrency:
+    """Tests for structured concurrency in query() (Issue #1136)."""
+
+    @pytest.mark.anyio
+    async def test_async_iterable_exception_propagates_fast_without_timeout(self):
+        mock_transport = AsyncMock()
+        mock_transport.is_ready.return_value = True
+        mock_transport.connect = AsyncMock()
+        mock_transport.close = AsyncMock()
+        mock_transport.end_input = AsyncMock()
+        mock_transport.write = AsyncMock()
+        mock_transport.wait = AsyncMock(return_value=0)
+
+        async def faulty_prompt():
+            yield {"type": "user", "message": "hi"}
+            raise ValueError("AsyncIterable prompt iteration failed")
+
+        async def mock_read_messages():
+            await anyio.sleep(10)
+            if False:
+                yield {}
+
+        mock_transport.read_messages = mock_read_messages
+
+        # Fixes SIM117 and prevents the initialization timeout
+        with (
+            patch(
+                "claude_agent_sdk._internal.query.Query.initialize",
+                new_callable=AsyncMock,
+            ),
+            pytest.raises(Exception) as exc_info,
+        ):
+            async for _ in query(prompt=faulty_prompt(), transport=mock_transport):
+                pass
+
+        assert "AsyncIterable prompt iteration failed" in repr(exc_info.value)
+
+    @pytest.mark.anyio
+    async def test_subprocess_crash_propagates_fast_without_timeout(self):
+        """When the transport subprocess crashes (wait() resolves non-zero), query()
+
+        must fail fast and propagate ProcessError without hanging on stdout.
+        """
+        mock_transport = AsyncMock()
+        mock_transport.is_ready.return_value = True
+        mock_transport.connect = AsyncMock()
+        mock_transport.close = AsyncMock()
+        mock_transport.end_input = AsyncMock()
+        mock_transport.write = AsyncMock()
+        mock_transport.wait = AsyncMock(return_value=1)
+
+        async def hanging_read_messages():
+            await anyio.sleep(10)
+            if False:
+                yield {}
+
+        mock_transport.read_messages = hanging_read_messages
+
+        # Fixes SIM117 and prevents the initialization timeout
+        with (
+            patch(
+                "claude_agent_sdk._internal.query.Query.initialize",
+                new_callable=AsyncMock,
+            ),
+            pytest.raises(Exception) as exc_info,
+        ):
+            async for _ in query(prompt="test prompt", transport=mock_transport):
+                pass
+
+        assert "1" in str(exc_info.value) or "exited" in str(exc_info.value)