diff --git a/src/claude_agent_sdk/_internal/query.py b/src/claude_agent_sdk/_internal/query.py index 66dfde9e7..25af0573f 100644 --- a/src/claude_agent_sdk/_internal/query.py +++ b/src/claude_agent_sdk/_internal/query.py @@ -55,6 +55,16 @@ # or it will hang the query (see Query._track_task_lifecycle). DEFERRING_TASK_TYPES = frozenset({"local_agent", "local_workflow"}) +# Cap on Query._deferred_result_ids. is_deferred_result() is the only thing +# that retires an entry, and only ClaudeSDKClient.receive_response() calls +# it — a caller reading raw frames via receive_messages() instead never +# consults it, so nothing would otherwise bound this set for a long-lived +# Query that sees many deferred results over its lifetime. Evicting the +# oldest entry past this cap trades a vanishingly unlikely stale-turn +# false negative for a hard memory bound; the real fix is calling +# is_deferred_result(), this is just a backstop for callers who don't. +_MAX_DEFERRED_RESULT_IDS = 128 + def _error_result_text(message: dict[str, Any]) -> str: """Pick the most informative text from a ``result`` frame with ``is_error``. @@ -188,6 +198,16 @@ 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. + # + # A plain dict, not a set: insertion order gives FIFO eviction in + # _read_messages once _MAX_DEFERRED_RESULT_IDS is exceeded. + self._deferred_result_ids: dict[str, None] = {} # 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 +399,16 @@ 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[result_uuid] = None + if ( + len(self._deferred_result_ids) + > _MAX_DEFERRED_RESULT_IDS + ): + self._deferred_result_ids.pop( + next(iter(self._deferred_result_ids)) + ) else: self._first_result_event.set() if message.get("is_error"): @@ -957,6 +987,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.pop(result_uuid, None) + 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..8530ed5bc 100644 --- a/tests/test_streaming_client.py +++ b/tests/test_streaming_client.py @@ -488,6 +488,217 @@ 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_deferred_result_ids_bounded_via_receive_messages(self): + """Query._deferred_result_ids is only retired by is_deferred_result(), + which only receive_response() calls. A caller reading raw frames via + receive_messages() instead never retires an entry, so a long-lived + Query that sees many deferred results must still bound the set + itself rather than leaking one entry per deferred result forever.""" + from claude_agent_sdk._internal.query import _MAX_DEFERRED_RESULT_IDS + + num_deferred_results = _MAX_DEFERRED_RESULT_IDS + 50 + + 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 + + # One background task that never settles, so every result + # frame below arrives with _inflight_tasks non-empty and is + # marked deferred. + yield { + "type": "system", + "subtype": "task_started", + "task_id": "task-1", + "task_type": "local_agent", + "description": "background subagent", + "uuid": "uuid-ts1", + "session_id": "test", + } + for i in range(num_deferred_results): + yield { + "type": "result", + "subtype": "success", + "duration_ms": 1, + "duration_api_ms": 1, + "is_error": False, + "num_turns": 1, + "session_id": "test", + "total_cost_usd": 0.0, + "uuid": f"uuid-r{i}", + } + + mock_transport.read_messages = mock_receive + + async with ClaudeSDKClient() as client: + result_count = 0 + async for msg in client.receive_messages(): + if isinstance(msg, ResultMessage): + result_count += 1 + if result_count == num_deferred_results: + break + + assert ( + len(client._query._deferred_result_ids) <= _MAX_DEFERRED_RESULT_IDS + ) + @pytest.mark.anyio async def test_interrupt(self): """Test interrupt functionality."""