diff --git a/lamin_cli/agents/_common.py b/lamin_cli/agents/_common.py index dc753ff..cd85e40 100644 --- a/lamin_cli/agents/_common.py +++ b/lamin_cli/agents/_common.py @@ -313,6 +313,69 @@ def contains_finish_invocation( return False +def _message_blocks(msg: dict) -> list[dict]: + content = msg.get("content") + if isinstance(content, str): + return [{"type": "text", "text": content}] + if not isinstance(content, list): + return [] + return [block for block in content if isinstance(block, dict)] + + +def _is_reply_text(block: dict) -> bool: + if block.get("type") != "text": + return False + text = block.get("text") + return isinstance(text, str) and bool(text.strip()) + + +def reply_precedes_finish( + entries: list[dict], shell_tool_names: frozenset[str] +) -> bool: + """Whether assistant text, not another tool call, sits before `lamin finish`.""" + if not entries: + return False + window_start = max(0, len(entries) - _RECENT_ENTRIES_WINDOW) + found: tuple[int, int] | None = None + for msg_index in range(len(entries) - 1, window_start - 1, -1): + blocks = _message_blocks(entries[msg_index]) + for block_index, block in enumerate(blocks): + if block.get("type") != "tool_use": + continue + if block.get("name") not in shell_tool_names: + continue + tool_input = block.get("input") + cmd = tool_input.get("command", "") if isinstance(tool_input, dict) else "" + if isinstance(cmd, str) and _is_finish_invocation(cmd): + found = (msg_index, block_index) + break + if found is not None: + break + if found is None: + return False + msg_index, block_index = found + + def meaningful(blocks: list[dict]) -> list[dict]: + return [ + block + for block in blocks + if block.get("type") not in {"thinking", "reasoning", "tool_result"} + ] + + source = entries[msg_index] + preceding = meaningful(_message_blocks(source)[:block_index]) + if not preceding: + if msg_index == 0: + return False + source = entries[msg_index - 1] + preceding = meaningful(_message_blocks(source)) + if source.get("role") not in {"assistant", None}: + return False + if not preceding: + return False + return _is_reply_text(preceding[-1]) + + def wait_for_finish_invocation( read_fn: Callable[[], _T], is_done_fn: Callable[[_T], bool], @@ -320,6 +383,7 @@ def wait_for_finish_invocation( budget_seconds: float = 8.0, poll_interval_seconds: float = 0.3, liveness_window_seconds: float = _LIVENESS_WINDOW_SECONDS, + preceding_reply_fn: Callable[[_T], bool] | None = None, ) -> _T: try: recently_active = ( @@ -333,7 +397,8 @@ def wait_for_finish_invocation( deadline = time.monotonic() + budget_seconds while True: result = read_fn() - if is_done_fn(result): + reply_ready = preceding_reply_fn is None or preceding_reply_fn(result) + if is_done_fn(result) and reply_ready: return result if time.monotonic() >= deadline: return result diff --git a/lamin_cli/agents/claude.py b/lamin_cli/agents/claude.py index 1cd4845..fb96773 100644 --- a/lamin_cli/agents/claude.py +++ b/lamin_cli/agents/claude.py @@ -242,6 +242,9 @@ def finish_claudecode_session() -> None: is_done_fn=lambda entries: _common.contains_finish_invocation( entries, _SHELL_TOOL_NAMES ), + preceding_reply_fn=lambda entries: _common.reply_precedes_finish( + entries, _SHELL_TOOL_NAMES + ), transcript_path=transcript_path, ) html_doc = _common.render_transcript_html( diff --git a/lamin_cli/agents/codex.py b/lamin_cli/agents/codex.py index efff0f6..26dce64 100644 --- a/lamin_cli/agents/codex.py +++ b/lamin_cli/agents/codex.py @@ -313,6 +313,9 @@ def _read() -> tuple[list[dict], dict[str, int]]: is_done_fn=lambda result: _common.contains_finish_invocation( result[0], _SHELL_TOOL_NAMES ), + preceding_reply_fn=lambda result: _common.reply_precedes_finish( + result[0], _SHELL_TOOL_NAMES + ), transcript_path=paths[-1], ) html_doc = _common.render_transcript_html( diff --git a/lamin_cli/agents/copilot.py b/lamin_cli/agents/copilot.py index 591fe7a..d95cf64 100644 --- a/lamin_cli/agents/copilot.py +++ b/lamin_cli/agents/copilot.py @@ -463,6 +463,9 @@ def _read() -> tuple[list[dict], list[dict]]: is_done_fn=lambda result: _common.contains_finish_invocation( result[1], _SHELL_TOOL_NAMES ), + preceding_reply_fn=lambda result: _common.reply_precedes_finish( + result[1], _SHELL_TOOL_NAMES + ), transcript_path=transcript_path, ) report_ids, share_ids = _report_session_ids(session_id, run.uid) diff --git a/lamin_cli/agents/cursor.py b/lamin_cli/agents/cursor.py index 0466361..7636a70 100644 --- a/lamin_cli/agents/cursor.py +++ b/lamin_cli/agents/cursor.py @@ -509,6 +509,9 @@ def finish_cursor_session() -> None: is_done_fn=lambda parsed: _common.contains_finish_invocation( parsed, _SHELL_TOOL_NAMES ), + preceding_reply_fn=lambda parsed: _common.reply_precedes_finish( + parsed, _SHELL_TOOL_NAMES + ), transcript_path=_cursor_db_path(), ) html_doc = _common.render_transcript_html( diff --git a/tests/agents/test_common.py b/tests/agents/test_common.py index 78b6df8..bb514ad 100644 --- a/tests/agents/test_common.py +++ b/tests/agents/test_common.py @@ -154,6 +154,23 @@ def test_contains_finish_invocation_finds_it_within_recent_window(): assert _common.contains_finish_invocation(entries, _SHELL_TOOL_NAMES) +def test_wait_stops_when_no_reply_precedes_finish(tmp_path): + start = time.monotonic() + result = _common.wait_for_finish_invocation( + lambda: [_bash_entry("echo hi"), _bash_entry("lamin finish")], + lambda entries: _common.contains_finish_invocation(entries, _SHELL_TOOL_NAMES), + transcript_path=_fresh_transcript_path(tmp_path), + budget_seconds=0.3, + preceding_reply_fn=lambda entries: _common.reply_precedes_finish( + entries, _SHELL_TOOL_NAMES + ), + poll_interval_seconds=0.05, + ) + elapsed = time.monotonic() - start + assert 0.25 < elapsed < 0.8 + assert result[-1]["content"][0]["input"]["command"] == "lamin finish" + + def test_wait_for_finish_invocation_retries_until_found(tmp_path): calls = {"n": 0} diff --git a/tests/agents/test_track_cursor.py b/tests/agents/test_track_cursor.py index f1a615b..a25a867 100644 --- a/tests/agents/test_track_cursor.py +++ b/tests/agents/test_track_cursor.py @@ -545,5 +545,5 @@ def write_delayed_finish_command(): elapsed = time.monotonic() - start writer.join() - assert 0.4 < elapsed < 3.0 + assert 0.4 < elapsed < 12.0 assert ln.Run.get(uid=uid).finished_at is not None