Handle background tasks in agent response loop

agents/bug-fix/hackbot_agents/bug_fix/agent.py

@@ -19,12 +19,15 @@
     ClaudeAgentOptions,
     ClaudeSDKClient,
     McpServerConfig,
-    ResultMessage,
 )
 from hackbot_runtime import ActionsRecorder, AgentError, HackbotAgentResult
 from hackbot_runtime.actions import ACTIONS_SERVER_NAME
 from hackbot_runtime.actions.claude_sdk import actions_server_for, actions_to_tool_names
-from hackbot_runtime.claude import Reporter
+from hackbot_runtime.claude import (
+    Reporter,
+    UnsettledResponseError,
+    receive_settled_response,
+)
 
 from .config import (
     BUGZILLA_READ_TOOLS,
@@ -98,6 +101,7 @@ async def run_bug_fix(
     verbose: bool = False,
     log: Path | None = None,
     actions_recorder: ActionsRecorder | None = None,
+    background_task_timeout_s: float = 3 * 60 * 60,
 ) -> BugFixResult:
     """Triage and fix a single Bugzilla bug with a claude-agent-sdk agent.
 
@@ -165,18 +169,19 @@ async def run_bug_fix(
         setting_sources=[],
     )
 
-    result_msg: ResultMessage | None = None
     with Reporter(verbose=verbose, log_path=log) as reporter:
         reporter.header(f"bug {bug}")
         async with ClaudeSDKClient(options=options) as client:
             await client.query(user_prompt)
-            async for msg in client.receive_response():
-                reporter.message(msg)
-                if isinstance(msg, ResultMessage):
-                    result_msg = msg
+            try:
+                result_msg = await receive_settled_response(
+                    client,
+                    on_message=reporter.message,
+                    timeout_s=background_task_timeout_s,
+                )
+            except UnsettledResponseError as exc:
+                raise AgentError(f"bug {bug}: agent run did not settle: {exc}") from exc
 
-    if result_msg is None:
-        raise AgentError(f"bug {bug}: agent produced no result message")
     if result_msg.is_error:
         raise AgentError(
             f"bug {bug} triage failed: {result_msg.result or result_msg.subtype}"

libs/agent-tools/pyproject.toml

@@ -15,7 +15,7 @@ bugzilla = [
     "six",
 ]
 firefox = ["grizzly-framework", "prefpicker"]
-claude-sdk = ["claude-agent-sdk>=0.1.30"]
+claude-sdk = ["claude-agent-sdk>=0.2.30"]
 searchfox = ["searchfox>=0.20.3"]
 vcs = ["httpx"]
 

libs/hackbot-runtime/hackbot_runtime/claude.py

@@ -10,13 +10,21 @@
 
 from __future__ import annotations
 
+import asyncio
 import json
+from collections.abc import Callable
 from pathlib import Path
 
 from claude_agent_sdk import (
+    TERMINAL_TASK_STATUSES,
     AssistantMessage,
+    ClaudeSDKClient,
+    Message,
     ResultMessage,
     SystemMessage,
+    TaskNotificationMessage,
+    TaskStartedMessage,
+    TaskUpdatedMessage,
     TextBlock,
     ThinkingBlock,
     ToolResultBlock,
@@ -25,6 +33,33 @@
 )
 
 
+class UnsettledResponseError(RuntimeError):
+    """``receive_settled_response`` gave up before the agent's turn settled.
+
+    ``pending`` is the ``task_id -> description`` map of deferring tasks
+    still open when this was raised (empty if none were ever pending — the
+    connection just ended with no result at all). Stored as an attribute
+    for callers that want to inspect it programmatically, e.g. to name the
+    stuck task(s) rather than just log the message.
+    """
+
+    def __init__(self, reason: str, pending: dict[str, str]):
+        self.reason = reason
+        self.pending = pending
+
+    def __str__(self) -> str:
+        if not self.pending:
+            return self.reason
+        return f"{self.reason} ({len(self.pending)} task(s) still pending)"
+
+
+# Task types whose completion the CLI itself resumes the turn for — mirrors
+# claude_agent_sdk._internal.query.DEFERRING_TASK_TYPES (not public API, so
+# duplicated here rather than imported).
+# https://github.com/anthropics/claude-agent-sdk-python/blob/bc0c9af676d9a63ac20a98cf1b7ba4794382c3cc/src/claude_agent_sdk/_internal/query.py#L38-L52
+_DEFERRING_TASK_TYPES = frozenset({"local_agent", "local_workflow"})
+
+
 def _truncate(s: str, n: int = 500) -> str:
     return s if len(s) <= n else s[:n] + f"... [{len(s) - n} more chars]"
 
@@ -123,3 +158,66 @@ def message(self, msg) -> None:
             self._emit(line, always=True)
             if msg.is_error:
                 self._emit(f"[done] ERROR: {msg.result}", always=True)
+
+
+async def receive_settled_response(
+    client: ClaudeSDKClient,
+    on_message: Callable[["Message"], None] | None = None,
+    *,
+    timeout_s: float = 3600,
+) -> ResultMessage:
+    """Drive ``client`` to a *settled* :class:`ResultMessage`.
+
+    ``client.receive_response()`` stops at the first ``ResultMessage``, but
+    the CLI can emit one while a task the agent backgrounded is still
+    running, reporting the turn "done" prematurely (see
+    anthropics/claude-agent-sdk-python#1138). This drains
+    ``client.receive_messages()`` instead (it doesn't stop at a
+    ``ResultMessage``) and only returns once one arrives with no *deferring*
+    task (``local_agent``/``local_workflow``, see ``_DEFERRING_TASK_TYPES``)
+    still open — backgrounded shells and Monitor watches run forever by
+    design and are never waited on. A task's terminal state can arrive as
+    either a ``TaskNotificationMessage`` or a ``TaskUpdatedMessage``, so both
+    clear it.
+
+    Args:
+        client: A connected client with a query already sent.
+        on_message: Called with each message as it streams in, before this
+            function's own bookkeeping. Optional.
+        timeout_s: Bounds the wait so a task that never settles raises
+            ``UnsettledResponseError`` instead of hanging. Defaults to an
+            hour (a full Firefox build); pass a larger value for
+            longer-running work.
+
+    Raises:
+        UnsettledResponseError: timed out, or the connection ended before
+            any ``ResultMessage`` arrived.
+    """
+    pending: dict[str, str] = {}
+    result_msg: ResultMessage | None = None
+
+    try:
+        async with asyncio.timeout(timeout_s):
+            async for msg in client.receive_messages():
+                if on_message is not None:
+                    on_message(msg)
+
+                if isinstance(msg, TaskStartedMessage):
+                    if msg.task_type in _DEFERRING_TASK_TYPES:
+                        pending[msg.task_id] = msg.description
+                elif isinstance(msg, (TaskNotificationMessage, TaskUpdatedMessage)):
+                    if msg.status in TERMINAL_TASK_STATUSES:
+                        pending.pop(msg.task_id, None)
+                elif isinstance(msg, ResultMessage):
+                    result_msg = msg
+                    if not pending:
+                        return result_msg
+    except TimeoutError as exc:
+        raise UnsettledResponseError(
+            f"timed out after {timeout_s:.0f}s waiting for the response to settle",
+            pending,
+        ) from exc
+
+    raise UnsettledResponseError(
+        "connection ended before a settled ResultMessage arrived", pending
+    )

libs/hackbot-runtime/tests/test_claude.py

@@ -1,6 +1,21 @@
-"""Tests for the shared claude-agent-sdk Reporter (hackbot_runtime.claude)."""
+"""Tests for the shared claude-agent-sdk helpers (hackbot_runtime.claude)."""
 
-from hackbot_runtime.claude import Reporter, _truncate
+import asyncio
+
+import pytest
+from claude_agent_sdk import (
+    ResultMessage,
+    SystemMessage,
+    TaskNotificationMessage,
+    TaskStartedMessage,
+    TaskUpdatedMessage,
+)
+from hackbot_runtime.claude import (
+    Reporter,
+    UnsettledResponseError,
+    _truncate,
+    receive_settled_response,
+)
 
 
 def test_truncate_short_string_unchanged():
@@ -34,3 +49,136 @@ def test_no_log_file_when_path_is_none(tmp_path):
     with Reporter(verbose=True, log_path=None) as reporter:
         reporter.header("section")
     assert not list(tmp_path.iterdir())
+
+
+class _FakeClient:
+    """Replays a fixed message list from ``receive_messages()``.
+
+    If ``hang_after`` is set, blocks forever once the list is exhausted,
+    simulating a connection left open with no further messages — the
+    situation ``timeout_s`` in ``receive_settled_response`` is meant to
+    bound.
+    """
+
+    def __init__(self, messages, hang_after: bool = False):
+        self._messages = messages
+        self._hang_after = hang_after
+
+    async def receive_messages(self):
+        for msg in self._messages:
+            yield msg
+        if self._hang_after:
+            await asyncio.Event().wait()
+
+
+def _result(is_error: bool = False, num_turns: int = 1) -> ResultMessage:
+    return ResultMessage(
+        subtype="success",
+        duration_ms=1,
+        duration_api_ms=1,
+        is_error=is_error,
+        num_turns=num_turns,
+        session_id="s1",
+    )
+
+
+def _task_started(task_id: str, task_type: str = "local_agent") -> TaskStartedMessage:
+    return TaskStartedMessage(
+        subtype="task_started",
+        data={},
+        task_id=task_id,
+        description="do a thing",
+        uuid="u1",
+        session_id="s1",
+        task_type=task_type,
+    )
+
+
+def _task_notification(
+    task_id: str, status: str = "completed"
+) -> TaskNotificationMessage:
+    return TaskNotificationMessage(
+        subtype="task_notification",
+        data={},
+        task_id=task_id,
+        status=status,
+        output_file="",
+        summary="done",
+        uuid="u2",
+        session_id="s1",
+    )
+
+
+def _task_updated(task_id: str, status: str = "completed") -> TaskUpdatedMessage:
+    return TaskUpdatedMessage(
+        subtype="task_updated",
+        data={},
+        task_id=task_id,
+        patch={"status": status},
+        status=status,
+    )
+
+
+async def test_receive_settled_response_returns_immediately_when_nothing_pending():
+    result = _result()
+    client = _FakeClient([result])
+    seen = []
+
+    got = await receive_settled_response(client, on_message=seen.append)
+
+    assert got is result
+    assert seen == [result]
+
+
+async def test_receive_settled_response_keeps_draining_past_early_result():
+    started = _task_started("t1")
+    early_result = _result(num_turns=1)
+    notification = _task_notification("t1")
+    final_result = _result(num_turns=2)
+    client = _FakeClient([started, early_result, notification, final_result])
+
+    got = await receive_settled_response(client)
+
+    # The first ResultMessage arrived while "t1" was still open — it must be
+    # ignored in favor of the one that follows the task's terminal message.
+    assert got is final_result
+
+
+async def test_receive_settled_response_task_updated_also_clears_pending():
+    started = _task_started("t1")
+    early_result = _result(num_turns=1)
+    updated = _task_updated("t1")
+    final_result = _result(num_turns=2)
+    client = _FakeClient([started, early_result, updated, final_result])
+
+    got = await receive_settled_response(client)
+
+    assert got is final_result
+
+
+async def test_receive_settled_response_ignores_non_deferring_task_types():
+    # A backgrounded shell (task_type="local_bash") can run indefinitely by
+    # design — the CLI itself never holds the result frame back for one, so
+    # neither should we. Settling immediately (rather than waiting on "t1")
+    # is the correct behavior here, not a race we need to rescue.
+    started = _task_started("t1", task_type="local_bash")
+    result = _result()
+    client = _FakeClient([started, result])
+
+    got = await receive_settled_response(client)
+
+    assert got is result
+
+
+async def test_receive_settled_response_raises_on_timeout_when_task_never_settles():
+    client = _FakeClient([_task_started("t1"), _result()], hang_after=True)
+
+    with pytest.raises(UnsettledResponseError):
+        await receive_settled_response(client, timeout_s=0.05)
+
+
+async def test_receive_settled_response_raises_when_stream_ends_without_result():
+    client = _FakeClient([SystemMessage(subtype="init", data={})])
+
+    with pytest.raises(UnsettledResponseError):
+        await receive_settled_response(client)