From 3cb1306ee3d7fd4f5a107e0211fd2a3d3e1a3809 Mon Sep 17 00:00:00 2001 From: Tanush Jain Date: Mon, 20 Jul 2026 22:49:04 +0530 Subject: [PATCH] fix(query): implement structured concurrency with task groups (fixes #1136) --- src/claude_agent_sdk/_internal/client.py | 83 +++++++++++++------ src/claude_agent_sdk/_internal/query.py | 13 ++- .../_internal/transport/__init__.py | 8 ++ .../_internal/transport/subprocess_cli.py | 13 ++- tests/test_client.py | 17 +--- tests/test_query.py | 78 ++++++++++++++++- 6 files changed, 162 insertions(+), 50 deletions(-) diff --git a/src/claude_agent_sdk/_internal/client.py b/src/claude_agent_sdk/_internal/client.py index 7b0a59387..453e514aa 100644 --- a/src/claude_agent_sdk/_internal/client.py +++ b/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() diff --git a/src/claude_agent_sdk/_internal/query.py b/src/claude_agent_sdk/_internal/query.py index 7aded9384..192c13029 100644 --- a/src/claude_agent_sdk/_internal/query.py +++ b/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).""" diff --git a/src/claude_agent_sdk/_internal/transport/__init__.py b/src/claude_agent_sdk/_internal/transport/__init__.py index 67388d78d..62c5d1dbd 100644 --- a/src/claude_agent_sdk/_internal/transport/__init__.py +++ b/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"] diff --git a/src/claude_agent_sdk/_internal/transport/subprocess_cli.py b/src/claude_agent_sdk/_internal/transport/subprocess_cli.py index faaec2d5e..37b671005 100644 --- a/src/claude_agent_sdk/_internal/transport/subprocess_cli.py +++ b/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: diff --git a/tests/test_client.py b/tests/test_client.py index 3d9c0293c..72cc0b661 100644 --- a/tests/test_client.py +++ b/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) diff --git a/tests/test_query.py b/tests/test_query.py index 3ae65093d..66af5f3c3 100644 --- a/tests/test_query.py +++ b/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)