From 7cbf417420a1490f3a137edfb1ce33695a04c8c3 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 18 Aug 2026 15:15:27 +0000 Subject: [PATCH] Fix ClaudeSDKClient.receive_response() returning before background agent work finishes receive_response() stopped on the first ResultMessage it saw, but a result frame only marks the end of one turn, not necessarily the run: when a delegated background agent/workflow task is still in flight, the CLI emits that result to close out the current turn and then continues with a follow-up turn once the task completes, ending in a second, later ResultMessage. receive_response() had no way to tell the two apart, so it returned on the first (intermediate) one and silently missed everything from the follow-up turn, including the real final result (#1138). Query already tracks in-flight delegated agent tasks via _track_task_lifecycle()/_inflight_tasks (added for #1088, to avoid closing stdin too early), so this reuses that same signal: each "result" frame sent while a delegated task is in flight has its uuid recorded, and Query.is_deferred_result() lets receive_response() recognize such a frame and keep reading instead of returning early. Fixes #1138. --- src/claude_agent_sdk/_internal/query.py | 32 ++++++ src/claude_agent_sdk/client.py | 19 +++- tests/test_streaming_client.py | 131 ++++++++++++++++++++++++ 3 files changed, 178 insertions(+), 4 deletions(-) diff --git a/src/claude_agent_sdk/_internal/query.py b/src/claude_agent_sdk/_internal/query.py index 66dfde9e7..fa4e1b7f4 100644 --- a/src/claude_agent_sdk/_internal/query.py +++ b/src/claude_agent_sdk/_internal/query.py @@ -188,6 +188,13 @@ def __init__( # a result that arrives while this set is non-empty must not close # stdin. self._inflight_tasks: set[str] = set() + # UUIDs of "result" frames sent while _inflight_tasks was non-empty — + # i.e. a turn boundary, not the end of the run (see #1138). Consulted + # (and cleared) by is_deferred_result() so a higher-level convenience + # API like ClaudeSDKClient.receive_response() can tell such a frame + # apart from the run-ending result and keep reading instead of + # returning early while delegated agent work is still in flight. + self._deferred_result_ids: set[str] = set() # Set to the result payload when the most recent message is a result # with is_error=True. Used to replace the generic "exit code 1" # ProcessError with a ResultError carrying what the CLI already @@ -379,6 +386,9 @@ async def _read_messages(self) -> None: "keeping stdin open", len(self._inflight_tasks), ) + result_uuid = message.get("uuid") + if isinstance(result_uuid, str): + self._deferred_result_ids.add(result_uuid) else: self._first_result_event.set() if message.get("is_error"): @@ -957,6 +967,28 @@ def _track_task_lifecycle(self, message: dict[str, Any]) -> None: if status in TERMINAL_TASK_STATUSES: self._inflight_tasks.discard(task_id) + def is_deferred_result(self, result_uuid: str | None) -> bool: + """Whether a ``result`` frame was an intermediate turn boundary. + + True means the frame arrived while delegated agent work + (``DEFERRING_TASK_TYPES``) was still in flight, so it ended one turn + but not the run — a later ``result`` frame with no tasks in flight is + still coming. Consulted by ``ClaudeSDKClient.receive_response()`` so + it can keep reading instead of returning on this frame and silently + missing the rest of the run (#1138). + + Looks up (and clears) the marker recorded in ``_read_messages``. + Returns ``False`` — "this is the run-ending result" — for ``None`` or + an id this ``Query`` never marked as deferred, which is the correct + default for CLI versions that predate background agent tasks, where + every result already ended the run. + """ + if result_uuid is None: + return False + was_deferred = result_uuid in self._deferred_result_ids + self._deferred_result_ids.discard(result_uuid) + return was_deferred + def _has_bidirectional_needs(self) -> bool: """Whether the CLI may still send control requests that need a reply. diff --git a/src/claude_agent_sdk/client.py b/src/claude_agent_sdk/client.py index 8a78d196b..7f3ab9576 100644 --- a/src/claude_agent_sdk/client.py +++ b/src/claude_agent_sdk/client.py @@ -531,17 +531,23 @@ async def get_server_info(self) -> dict[str, Any] | None: async def receive_response(self) -> AsyncIterator[Message]: """ - Receive messages from Claude until and including a ResultMessage. + Receive messages from Claude until and including the run-ending ResultMessage. This async iterator yields all messages in sequence and automatically terminates - after yielding a ResultMessage (which indicates the response is complete). + after yielding the ResultMessage that ends the run. It's a convenience method over receive_messages() for single-response workflows. **Stopping Behavior:** - Yields each message as it's received - - Terminates immediately after yielding a ResultMessage + - Terminates immediately after yielding the run-ending ResultMessage - The ResultMessage IS included in the yielded messages - If no ResultMessage is received, the iterator continues indefinitely + - A ResultMessage is skipped over (not treated as terminal) when it only + ends one turn while delegated agent work (a background subagent or + workflow) is still in flight — the run continues with a follow-up + turn, which ends in its own, later ResultMessage. This keeps the + iterator from returning early and silently missing the rest of the + run (#1138). Yields: Message: Each message received (UserMessage, AssistantMessage, SystemMessage, ResultMessage) @@ -565,9 +571,14 @@ async def receive_response(self) -> AsyncIterator[Message]: To collect all messages: `messages = [msg async for msg in client.receive_response()]` The final message in the list will always be a ResultMessage. """ + if not self._query: + raise CLIConnectionError("Not connected. Call connect() first.") + query = self._query async for message in self.receive_messages(): yield message - if isinstance(message, ResultMessage): + if isinstance(message, ResultMessage) and not query.is_deferred_result( + message.uuid + ): return async def disconnect(self) -> None: diff --git a/tests/test_streaming_client.py b/tests/test_streaming_client.py index 861ce6ec3..0e6527c83 100644 --- a/tests/test_streaming_client.py +++ b/tests/test_streaming_client.py @@ -488,6 +488,137 @@ async def mock_receive(): assert isinstance(messages[0], AssistantMessage) assert isinstance(messages[1], ResultMessage) + @pytest.mark.anyio + async def test_receive_response_waits_for_deferred_result(self): + """receive_response() must not stop on a result that only ends one + turn while delegated agent work is still in flight (#1138) — it + should keep reading through to the run-ending result.""" + + with patch( + "claude_agent_sdk._internal.transport.subprocess_cli.SubprocessCLITransport" + ) as mock_transport_class: + mock_transport = create_mock_transport() + mock_transport_class.return_value = mock_transport + + async def mock_receive(): + await anyio.sleep(0.01) + written = mock_transport.write.call_args_list + for call in written: + data = call[0][0] + try: + msg = json.loads(data.strip()) + if ( + msg.get("type") == "control_request" + and msg.get("request", {}).get("subtype") == "initialize" + ): + yield { + "type": "control_response", + "response": { + "request_id": msg.get("request_id"), + "subtype": "success", + "commands": [], + "output_style": "default", + }, + } + break + except (json.JSONDecodeError, KeyError, AttributeError): + pass + + yield { + "type": "assistant", + "message": { + "role": "assistant", + "content": [{"type": "text", "text": "First turn"}], + "model": "claude-opus-4-1-20250805", + }, + } + yield { + "type": "system", + "subtype": "task_started", + "task_id": "task-1", + "task_type": "local_agent", + "description": "background subagent", + "uuid": "uuid-ts1", + "session_id": "test", + } + # Turn boundary: this result ends the first turn, but the + # background agent task is still running. + yield { + "type": "result", + "subtype": "success", + "duration_ms": 1000, + "duration_api_ms": 800, + "is_error": False, + "num_turns": 1, + "session_id": "test", + "total_cost_usd": 0.001, + "uuid": "uuid-r1", + } + # The background task settles, waking the parent for a + # follow-up turn. + yield { + "type": "system", + "subtype": "task_notification", + "task_id": "task-1", + "status": "completed", + "output_file": "/tmp/task-1.output", + "summary": "done", + "uuid": "uuid-tn1", + "session_id": "test", + } + yield { + "type": "assistant", + "message": { + "role": "assistant", + "content": [ + {"type": "text", "text": "Follow-up after the task"} + ], + "model": "claude-opus-4-1-20250805", + }, + } + # Run-ending result: no tasks in flight. + yield { + "type": "result", + "subtype": "success", + "duration_ms": 500, + "duration_api_ms": 400, + "is_error": False, + "num_turns": 2, + "session_id": "test", + "total_cost_usd": 0.002, + "uuid": "uuid-r2", + } + # This should not be yielded — receive_response() must have + # already returned after the run-ending result above. + yield { + "type": "assistant", + "message": { + "role": "assistant", + "content": [{"type": "text", "text": "Should not see this"}], + }, + "model": "claude-opus-4-1-20250805", + } + + mock_transport.read_messages = mock_receive + + async with ClaudeSDKClient() as client: + messages = [] + async for msg in client.receive_response(): + messages.append(msg) + + assert [type(m).__name__ for m in messages] == [ + "AssistantMessage", + "TaskStartedMessage", + "ResultMessage", + "TaskNotificationMessage", + "AssistantMessage", + "ResultMessage", + ] + results = [m for m in messages if isinstance(m, ResultMessage)] + assert len(results) == 2 + assert results[0].uuid == "uuid-r1" + assert results[1].uuid == "uuid-r2" + @pytest.mark.anyio async def test_interrupt(self): """Test interrupt functionality."""