fix(claude-code-agent): end the turn on the terminal ResultMessage, not stream EOF

src/coder_eval/agents/claude_code_agent.py

@@ -17,6 +17,8 @@
     ClaudeSDKClient,
     Message,
     ProcessError,
+    ResultMessage,
+    SystemMessage,
     TaskNotificationMessage,
     query,
 )
@@ -168,9 +170,20 @@ def _is_task_notification(message: Any) -> bool:
 def _is_sdk_result_message(message: Any) -> bool:
     """Check if message is the SDK's final ResultMessage (with usage/cost data).
 
-    Distinct from ToolResultBlock which has tool_use_id, and from
-    TaskNotificationMessage which also carries session_id + usage (excluded).
+    Real SDK instances are identified positively by type. The duck-typed
+    fallback (session_id + usage) exists ONLY for test mocks, and it must
+    never match a SystemMessage subclass: the sub-agent lifecycle family
+    (TaskStartedMessage / TaskProgressMessage / TaskNotificationMessage) also
+    carries session_id — and TaskProgressMessage carries usage too, so a
+    sub-agent progress tick would otherwise be misread as the terminal
+    result. That misread silently corrupted session-id advance and token
+    backfill, and — once the turn started ENDING on the terminal result —
+    truncated any turn that spawned a sub-agent.
     """
+    if isinstance(message, ResultMessage):
+        return True
+    if isinstance(message, SystemMessage):
+        return False
     return hasattr(message, "session_id") and hasattr(message, "usage") and not _is_task_notification(message)
 
 
@@ -899,15 +912,8 @@ def _on_turn_timeout() -> None:
                 label=f"Turn timeout ({timeout:g}s)" if timeout else "turn_timeout",
             ):
                 async for message in query(**query_kwargs):
-                    # Wall-clock guard at the TOP of the loop: it breaks BEFORE
-                    # the message is dispatched (and appended), so the
-                    # over-deadline message is DISCARDED — no append, no events.
-                    # Do NOT relocate this to a post-loop check.
-                    if deadline is not None and time.monotonic() > deadline:
-                        state.timeout_hit = True
-                        self._log.warning("Turn timeout reached mid-stream; breaking out of message loop")
+                    if await self._pump_stream_message(state, message, transport):
                         break
-                    state.dispatch(message)
 
             self._log.debug("Agent query stream ended")
 
@@ -1103,6 +1109,67 @@ def _timed_out(timeout_hit: bool, deadline: float | None) -> bool:
             return True
         return deadline is not None and time.monotonic() > deadline
 
+    async def _pump_stream_message(
+        self, state: "_ClaudeTurnState", message: Message, transport: SubprocessCLITransport | None
+    ) -> bool:
+        """Dispatch one SDK stream message; return True when the loop must stop.
+
+        Wall-clock guard FIRST, before the message is dispatched (and appended),
+        so an over-deadline message is DISCARDED — no append, no events. Do NOT
+        relocate that check to a post-loop position.
+
+        After dispatch, the turn ends on the terminal ResultMessage of the
+        one-shot query(): nothing follows it except process exit. Ending HERE
+        rather than at stream EOF matters — a CLI process that lingers after
+        emitting its result (stray child holding stdio, slow shutdown flush)
+        must not keep the stream open until the watchdog converts a finished,
+        successful turn into a TurnTimeoutError. Breaking alone would not reap
+        the subprocess (generator finalization is GC-scheduled, and the SDK's
+        anyio scopes swallow cooperative cancellation), so the CLI gets a short
+        grace to flush its session file and exit, then a hard kill.
+        """
+        if state.deadline is not None and time.monotonic() > state.deadline:
+            state.timeout_hit = True
+            self._log.warning("Turn timeout reached mid-stream; breaking out of message loop")
+            return True
+        state.dispatch(message)
+        if state.sdk_result_summary is not None:
+            self._log.debug("Terminal ResultMessage received; ending turn without waiting for process exit")
+            await self._reap_transport_after_result(transport)
+            return True
+        return False
+
+    @staticmethod
+    async def _reap_transport_after_result(
+        transport: SubprocessCLITransport | None, grace_seconds: float = 10.0
+    ) -> None:
+        """Bounded reap of the CLI subprocess once the terminal ResultMessage is in hand.
+
+        Give the CLI a short grace to exit on its own — it flushes its session
+        file during shutdown, and killing it instantly can lose the resume
+        transcript's tail (same reason the SDK transport ``close()`` waits
+        before signaling). If it is still alive after the grace, SIGKILL: the
+        turn's result is already captured, so nothing of value can be lost,
+        and a lingering process must not hold the turn (or leak into the next
+        one). No-op when the process already exited — the common case.
+        """
+        if transport is None:
+            return
+        proc = getattr(transport, "_process", None)
+        if proc is None:
+            return
+        reap_deadline = time.monotonic() + grace_seconds
+        while proc.returncode is None and time.monotonic() < reap_deadline:
+            await asyncio.sleep(0.1)
+        if proc.returncode is None:
+            logger.warning(
+                "Claude CLI subprocess (pid=%s) still alive %.0fs after its terminal ResultMessage; hard-killing",
+                getattr(proc, "pid", "?"),
+                grace_seconds,
+            )
+            with suppress(OSError):
+                proc.kill()
+
     @staticmethod
     def _kill_transport(transport: SubprocessCLITransport | None) -> None:
         """SIGKILL the subprocess behind `transport`, if any.

tests/test_agent_telemetry.py

@@ -661,7 +661,13 @@ async def mock_query(prompt, options):
 
     @pytest.mark.asyncio
     async def test_multiple_assistant_turns(self, tmp_path):
-        """Verify multiple assistant turns are tracked separately."""
+        """Verify multiple assistant turns are tracked separately.
+
+        Stream shape mirrors the real one-shot CLI protocol: the ResultMessage
+        is TERMINAL (exactly one, at the end of the session). communicate()
+        ends the turn on it, so a mid-stream ResultMessage would truncate the
+        transcript — that shape does not occur in one-shot mode.
+        """
         tool_use_block_cls, assistant_message_cls, user_message_cls, text_block_cls, _, result_message_cls = (
             create_mock_sdk_messages()
         )
@@ -671,14 +677,13 @@ async def test_multiple_assistant_turns(self, tmp_path):
         tool1 = tool_use_block_cls("toolu_turn1", "Read", {"file_path": "file1.txt"})
         msg1 = assistant_message_cls([text1, tool1])
         user_msg1 = user_message_cls("toolu_turn1", False, "content1")
-        result1 = result_message_cls(usage={"input_tokens": 50, "output_tokens": 30})
 
         # Second turn
         text2 = text_block_cls("Second response")
         tool2 = tool_use_block_cls("toolu_turn2", "Write", {"file_path": "file2.txt", "content": "new content"})
         msg2 = assistant_message_cls([text2, tool2])
         user_msg2 = user_message_cls("toolu_turn2", False, "ok")
-        result2 = result_message_cls(usage={"input_tokens": 60, "output_tokens": 40})
+        result_final = result_message_cls(usage={"input_tokens": 60, "output_tokens": 40})
 
         import coder_eval.agents.claude_code_agent as agent_module
 
@@ -688,10 +693,9 @@ async def test_multiple_assistant_turns(self, tmp_path):
         async def mock_query(prompt, options):
             yield msg1
             yield user_msg1
-            yield result1
             yield msg2
             yield user_msg2
-            yield result2
+            yield result_final
 
         original_query = agent_module.query
         agent_module.query = mock_query
@@ -705,15 +709,18 @@ async def mock_query(prompt, options):
             assistant_msgs = [m for m in turn.messages if isinstance(m, AssistantMessage)]
             assert len(assistant_msgs) == 2
 
-            # First turn
+            # First turn: id-less and not last, so the terminal ResultMessage's
+            # usage does NOT retro-populate it (only the LAST id-less assistant
+            # message is backfilled).
             aturn1 = assistant_msgs[0]
             assert isinstance(aturn1, AssistantMessage)
             assert aturn1.role == "assistant"
-            assert aturn1.input_tokens == 50
-            assert aturn1.output_tokens == 30
+            assert aturn1.input_tokens == 0
+            assert aturn1.output_tokens == 0
             assert len(aturn1.content_blocks) == 2
 
-            # Second turn
+            # Second turn: last id-less assistant message — backfilled from the
+            # terminal ResultMessage's usage.
             aturn2 = assistant_msgs[1]
             assert isinstance(aturn2, AssistantMessage)
             assert aturn2.role == "assistant"

tests/test_agent_timeout.py

@@ -258,3 +258,142 @@ async def mock_query(prompt, options):
         ):
             await agent.communicate("prompt")  # no timeout
             mock_transport_cls.assert_not_called()
+
+
+def _fake_result_message():
+    """Duck-typed SDK ResultMessage: session_id + usage (and no `content`,
+    so it can't false-match the assistant/user branches of dispatch)."""
+    from types import SimpleNamespace
+
+    return SimpleNamespace(
+        session_id="sess-terminal",
+        usage={"input_tokens": 10, "output_tokens": 5},
+        model_usage=None,
+        total_cost_usd=0.01,
+        num_turns=3,
+        subtype="success",
+        stop_reason="end_turn",
+        result="all done",
+        is_error=False,
+    )
+
+
+@pytest.mark.asyncio
+async def test_communicate_ends_turn_on_result_message_when_cli_lingers():
+    """A CLI that emits its terminal ResultMessage but never exits must yield
+    a clean TurnRecord — NOT run to the watchdog and raise TurnTimeoutError.
+
+    Regression for the sweep-killing class: the agent finished (end_turn,
+    subtype=success), the CLI process lingered, the stream never hit EOF, and
+    the watchdog converted a successful run into a 0-score timeout ERROR.
+    """
+    agent = _make_agent()
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        await agent.start(tmpdir)
+
+        async def mock_query(prompt, options, transport=None):
+            yield _fake_result_message()
+            # Simulate the lingering CLI: the stream never ends. Without the
+            # break-on-result fix this sleeps until the watchdog fires.
+            await asyncio.sleep(30)
+
+        fake_transport = MagicMock()
+        fake_transport._process.returncode = 0  # CLI already exited cleanly
+
+        with (
+            patch("coder_eval.agents.claude_code_agent.SubprocessCLITransport") as mock_transport_cls,
+            patch("coder_eval.agents.claude_code_agent.query", mock_query),
+        ):
+            mock_transport_cls.return_value = fake_transport
+
+            # Watchdog at 5s: with the fix communicate() returns immediately
+            # after the ResultMessage; without it, TurnTimeoutError.
+            record = await agent.communicate("prompt", timeout=5.0)
+
+        assert record is not None
+        assert record.result_summary is not None
+        assert record.result_summary.subtype == "success"
+
+
+@pytest.mark.asyncio
+async def test_reap_after_result_kills_lingering_process():
+    """_reap_transport_after_result must SIGKILL a process that outlives its
+    grace period, and leave an already-exited process alone."""
+    lingering = MagicMock()
+    lingering._process.returncode = None
+    lingering._process.pid = 4242
+    await ClaudeCodeAgent._reap_transport_after_result(lingering, grace_seconds=0.3)
+    lingering._process.kill.assert_called_once()
+
+    exited = MagicMock()
+    exited._process.returncode = 0
+    await ClaudeCodeAgent._reap_transport_after_result(exited, grace_seconds=0.3)
+    exited._process.kill.assert_not_called()
+
+    # None transport / None process are no-ops
+    await ClaudeCodeAgent._reap_transport_after_result(None)
+    no_proc = MagicMock()
+    no_proc._process = None
+    await ClaudeCodeAgent._reap_transport_after_result(no_proc)
+
+
+@pytest.mark.asyncio
+async def test_subagent_progress_message_does_not_end_turn():
+    """A sub-agent TaskProgressMessage must NOT be mistaken for the terminal
+    ResultMessage.
+
+    TaskProgressMessage carries BOTH session_id and usage — the duck-typed
+    result check used to match it. That misread silently corrupted session-id
+    advance, and once the turn began ending on the terminal result it
+    truncated any turn that spawned a sub-agent: the agent was killed
+    mid-run ~20s in with zero authored output (observed live: 0-turn,
+    empty-sandbox failures on every sub-agent-spawning task).
+    """
+    from types import SimpleNamespace
+
+    from claude_agent_sdk.types import TaskProgressMessage
+
+    progress = TaskProgressMessage(
+        subtype="task_progress",
+        data={},
+        task_id="task-1",
+        description="sub-agent working",
+        usage={"total_tokens": 10, "tool_uses": 1, "duration_ms": 500},
+        uuid="uuid-1",
+        session_id="sess-sub",
+    )
+    # Classifier level: never a result, even though it has session_id + usage.
+    from coder_eval.agents.claude_code_agent import _is_sdk_result_message
+
+    assert _is_sdk_result_message(progress) is False
+
+    # End-to-end: a stream [assistant-ish, progress, terminal result] must
+    # deliver the REAL result — not end the turn at the progress message.
+    agent = _make_agent()
+    with tempfile.TemporaryDirectory() as tmpdir:
+        await agent.start(tmpdir)
+
+        terminal = _fake_result_message()
+
+        async def mock_query(prompt, options, transport=None):
+            yield SimpleNamespace(content=[], model="mock-model")  # assistant-ish
+            yield progress
+            yield terminal
+            await asyncio.sleep(30)  # linger; the terminal result ends the turn
+
+        fake_transport = MagicMock()
+        fake_transport._process.returncode = 0
+
+        with (
+            patch("coder_eval.agents.claude_code_agent.SubprocessCLITransport") as mock_transport_cls,
+            patch("coder_eval.agents.claude_code_agent.query", mock_query),
+        ):
+            mock_transport_cls.return_value = fake_transport
+            record = await agent.communicate("prompt", timeout=5.0)
+
+        assert record is not None
+        assert record.result_summary is not None
+        assert record.result_summary.subtype == "success"
+        # The assistant message BEFORE the progress tick must have been kept.
+        assert record.agent_output is not None