Fix ClaudeSDKClient.receive_response() returning before background agent work finishes

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.
 

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:

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."""