diff --git a/src/coder_eval/agents/claude_code_agent.py b/src/coder_eval/agents/claude_code_agent.py index eb5b506e..da4e1b8e 100644 --- a/src/coder_eval/agents/claude_code_agent.py +++ b/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. diff --git a/tests/test_agent_telemetry.py b/tests/test_agent_telemetry.py index 7f6fe08b..1ac1c005 100644 --- a/tests/test_agent_telemetry.py +++ b/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" diff --git a/tests/test_agent_timeout.py b/tests/test_agent_timeout.py index bd9fff7b..eac86c9a 100644 --- a/tests/test_agent_timeout.py +++ b/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