From fa8938a89f152916f35f09b4c81c31e0bc9f57d4 Mon Sep 17 00:00:00 2001 From: Leon Date: Sat, 15 Aug 2026 15:29:01 +0800 Subject: [PATCH] fix(autopilot): offload stage-result writes `_stage_loop` is async, but it persisted each stage's result with a synchronous `mkdir` + `write_text` on the gateway event loop. Every other task on that loop -- chat streaming, WebSocket frames, cron dispatch -- is stalled for the duration of the filesystem work, once per stage boundary. Split the filesystem half into `_write_stage_result` and hand it to `asyncio.to_thread`. The extraction stays on the loop thread: it walks `slot.messages`, which is live mutable state the loop owns, so only the finished text and the slot key cross into the worker. Path format, message ordering, redaction and failure semantics are unchanged. The publication `os.replace` moves into that worker too. An earlier revision kept it on the loop deliberately, reasoning that a rename is metadata-only and that publishing after an uncancelled await made the orphan-writer class structurally unreachable. The first half of that does not hold: the session directory can be network-backed, where a rename is a round trip like any other call, so it is the same stall the offload exists to remove and the same `no-blocking-call-on-event-loop` anchor covers it. The orphan-writer guarantee is kept by a flag rather than by placement. `asyncio.to_thread` cannot interrupt a running worker, so the worker re-reads an `abandoned` event immediately before publishing and the caller sets it the moment its await is cancelled. A capture abandoned during the payload write -- the long half, and effectively all of the window on the slow filesystem this offload targets -- unlinks its temp file and leaves `stage_N_result.md` untouched, which the existing cancellation test still pins. The residual window is the gap between that check and the rename; it is not zero, and it is stated rather than claimed away. Co-authored-by: Kiro Crew Co-Authored-By: Claude Opus 5 --- src/kiro_crew/dashboard/chat_orchestrator.py | 103 ++++- test/test_completion_result_read_off_loop.py | 4 +- test/test_display_time_redaction.py | 9 +- test/test_stage_result_off_loop.py | 375 +++++++++++++++++++ 4 files changed, 480 insertions(+), 11 deletions(-) create mode 100644 test/test_stage_result_off_loop.py diff --git a/src/kiro_crew/dashboard/chat_orchestrator.py b/src/kiro_crew/dashboard/chat_orchestrator.py index f5d45a3e8ba..39c77dac274 100644 --- a/src/kiro_crew/dashboard/chat_orchestrator.py +++ b/src/kiro_crew/dashboard/chat_orchestrator.py @@ -3,7 +3,10 @@ from __future__ import annotations import asyncio +import contextlib import logging +import os +import threading import uuid from datetime import datetime, timezone from pathlib import Path @@ -126,13 +129,69 @@ async def _previous_result_paths( return await asyncio.to_thread(_read_previous_results, recorded) -def _capture_stage_result( +def _write_stage_result( + slot_key: str, + stage_num: int, + result_text: str, + abandoned: threading.Event, +) -> str | None: + """Write *result_text* and publish it, entirely on the worker thread. + + Returns the published path, or ``None`` when the capture was abandoned + before publication. Takes only immutable values plus *abandoned* — never + the slot — so nothing the loop owns is reachable from the worker. + + EVERY filesystem call lives here, ``os.replace`` included. The session + directory can be network-backed, where a rename is a round trip and not the + "metadata-only, therefore free" syscall it looks like locally; running it on + the loop stalls chat streaming, WebSocket frames and cron dispatch for that + duration, which is what ``no-blocking-call-on-event-loop`` forbids. + + Publication stays SAFE without being on the loop. The payload still goes to + a ``uuid``-named sibling that no other writer knows, and the canonical path + is touched only after re-reading *abandoned* — which the caller sets when + its await is cancelled (stop button, slot close, slot-key reuse). So a + worker abandoned during the write — the long half, and the whole of the + window on a slow filesystem — unlinks its temp file and publishes nothing. + + What this does NOT claim: the check and the rename are two operations, so a + cancellation landing between them still publishes. That window is one + already-resolved rename rather than the entire payload write, and it cannot + be closed from the loop side at all while the rename runs on the loop — + there the loop is not even free to observe the cancellation. + """ + session_dir = config_dir() / "sessions" / slot_key + session_dir.mkdir(parents=True, exist_ok=True) + final = session_dir / f"stage_{stage_num}_result.md" + tmp = session_dir / f".stage_{stage_num}_result.{uuid.uuid4().hex}.tmp" + tmp.write_text(result_text, encoding="utf-8") + if abandoned.is_set(): + with contextlib.suppress(OSError): + tmp.unlink() + return None + try: + os.replace(str(tmp), str(final)) + except OSError: + with contextlib.suppress(OSError): + tmp.unlink() + raise + return str(final) + + +async def _capture_stage_result( slot: "_ChatSlot", stage_num: int, ) -> str: """Extract assistant messages since stage start and write to disk. Returns the path to the result file. + + The extraction stays on the event-loop thread because ``slot.messages`` is + live mutable state the loop owns; only the finished text crosses into the + worker, so a concurrent append cannot be read from another thread. The + filesystem half runs off-loop: ``_stage_loop`` is async and this runs at every + stage boundary, so a slow ``mkdir``/``write_text`` would stall chat streaming, + WebSocket frames and cron dispatch for its duration. """ # Collect assistant text from the most recent messages (since last stage separator) result_parts: list[str] = [] @@ -154,11 +213,33 @@ def _capture_stage_result( result_parts.reverse() result_text = "\n\n".join(result_parts) - session_dir = config_dir() / "sessions" / slot.key - session_dir.mkdir(parents=True, exist_ok=True) - path = session_dir / f"stage_{stage_num}_result.md" - path.write_text(result_text, encoding="utf-8") - return str(path) + # The worker owns the whole filesystem half: payload write AND publication. + # Nothing here touches the disk, because the session directory can be + # network-backed and a rename there is a round trip, not a free + # metadata-only syscall. + # + # ``asyncio.to_thread`` cannot interrupt a running worker, so the orphan + # writer is held off with a flag instead of with placement: the worker + # re-reads ``abandoned`` immediately before it publishes, and this handler + # sets it the moment the await is cancelled. A capture abandoned during the + # payload write — the long half, and effectively all of the window on the + # slow filesystem this offload exists for — therefore unlinks its temp file + # and leaves ``stage_N_result.md`` alone. + abandoned = threading.Event() + try: + published = await asyncio.to_thread( + _write_stage_result, slot.key, stage_num, result_text, abandoned + ) + except asyncio.CancelledError: + abandoned.set() + raise + if published is None: + # Only reachable when the worker saw ``abandoned``, which is set solely + # in the handler above -- so the await raised and this line did not run. + # Kept so the signature stays honest rather than asserting the + # invariant with a cast. + raise asyncio.CancelledError() + return published def _completion_excerpts(result_paths: tuple[tuple[int, str], ...]) -> dict[int, str]: @@ -560,11 +641,19 @@ async def _stage_loop( # Capture result to disk try: - result_path = _capture_stage_result(slot, stage_num) + result_path = await _capture_stage_result(slot, stage_num) tracker.record_stage_result(stage_num, result_path) except OSError: logger.warning("Failed to capture stage %d result to disk", stage_num, exc_info=True) + # The offloaded write above is a suspension point between the + # loop's last stop check and the next stage. A stop that lands + # while the worker is writing (user cancel sets tracker.stopped + # and slot._auto_run, but `auto_run` below is a call-time + # snapshot) must not let the run advance to another stage. + if slot._stopping or tracker.stopped: + break + # Gate: if not auto_run, wait for user approval if not auto_run: # Emit completion message — user must click Go for next stage diff --git a/test/test_completion_result_read_off_loop.py b/test/test_completion_result_read_off_loop.py index 8a614ee4a74..6f738fdd247 100644 --- a/test/test_completion_result_read_off_loop.py +++ b/test/test_completion_result_read_off_loop.py @@ -213,8 +213,8 @@ async def test_completion_summary_survives_a_deleted_result_file(monkeypatch, tm real_capture = None - def _capture_then_delete_stage_1(s, stage_num): - path = real_capture(s, stage_num) + async def _capture_then_delete_stage_1(s, stage_num): + path = await real_capture(s, stage_num) if stage_num == 1: (tmp_path / "sessions" / s.key / "stage_1_result.md").unlink() return path diff --git a/test/test_display_time_redaction.py b/test/test_display_time_redaction.py index 43bbcf56241..70102ebf211 100644 --- a/test/test_display_time_redaction.py +++ b/test/test_display_time_redaction.py @@ -162,11 +162,16 @@ def test_side_chat_parent_snapshot_keeps_user_text() -> None: assert "my own question" in _format_parent_snapshot(slot) -def test_stage_result_capture_redacts_before_writing_to_disk(tmp_path, monkeypatch) -> None: +@pytest.mark.asyncio +async def test_stage_result_capture_redacts_before_writing_to_disk(tmp_path, monkeypatch) -> None: """_capture_stage_result writes assistant text to a NEW file on disk. A gateway restart mid-orchestration leaves restored (now unredacted) turns in the window, so without redaction here those bytes would be written out. + + Redaction runs on the event-loop thread and only the finished text is handed + to the worker that writes it, so the credential cannot reach disk even though + the write itself is off-loop. """ monkeypatch.setenv("KIROCREW_HOME", str(tmp_path)) monkeypatch.setattr("kiro_crew.dashboard.chat_orchestrator.config_dir", lambda: tmp_path) @@ -176,7 +181,7 @@ def test_stage_result_capture_redacts_before_writing_to_disk(tmp_path, monkeypat slot = _ChatSlot("chat-1-stage") slot.append("assistant", f"result with {SECRET}", "msg msg-a", broadcast=False) - path = _capture_stage_result(slot, 1) + path = await _capture_stage_result(slot, 1) written = pathlib.Path(path).read_text() assert SECRET not in written, "stage result persisted an unredacted credential" diff --git a/test/test_stage_result_off_loop.py b/test/test_stage_result_off_loop.py new file mode 100644 index 00000000000..616d29cb163 --- /dev/null +++ b/test/test_stage_result_off_loop.py @@ -0,0 +1,375 @@ +"""The autopilot stage-result write must not run on the gateway event loop. + +``_stage_loop`` is async and drives the whole autopilot run, but it persists each +stage's result with a synchronous ``mkdir`` + ``write_text``. Every other task on +the loop — chat streaming, WebSocket frames, cron dispatch — is stalled for the +duration of that filesystem work, once per stage boundary. + +Only the filesystem mutation belongs off-loop. Reading ``slot.messages`` and +redacting stays on the loop thread: it is live mutable slot state, and handing it +to a worker would buy nothing while widening the change into cross-thread access. +These tests pin both halves of that boundary. +""" + +from __future__ import annotations + +import inspect +import os +import pathlib +import threading + +import pytest + + +async def _capture(slot, stage_num): + """Drive the real production capture, sync or async. + + Deliberately tolerant of both shapes so the thread assertions below are what + fails on an unfixed tree. A bare ``await`` would raise "can't be used in + 'await' expression" against the synchronous version, which proves only that + the symbol changed — not that the write was ever on the wrong thread. + """ + from kiro_crew.dashboard.chat_orchestrator import _capture_stage_result + + result = _capture_stage_result(slot, stage_num) + if inspect.isawaitable(result): + result = await result + return result + + +def _slot_with(*assistant_texts: str): + from kiro_crew.dashboard.state import _ChatSlot + + slot = _ChatSlot("chat-1-stage") + for i, text in enumerate(assistant_texts): + slot.append("assistant", text, f"msg msg-a{i}", broadcast=False) + return slot + + +@pytest.mark.asyncio +async def test_stage_result_write_runs_off_the_loop_thread(tmp_path, monkeypatch): + """The bytes reach disk on some thread other than the loop's.""" + monkeypatch.setattr("kiro_crew.dashboard.chat_orchestrator.config_dir", lambda: tmp_path) + slot = _slot_with("stage one output") + + seen_threads: list[int] = [] + real_write = pathlib.Path.write_text + target_dir = tmp_path / "sessions" / "chat-1-stage" + + def recording_write(self, *args, **kwargs): + # Scoped to the stage-result payload write (canonical or temp name, + # so the assertion tracks the write wherever the implementation puts + # the bytes) -- unrelated writes cannot decide this. + if self.parent == target_dir and "stage_1_result" in self.name: + seen_threads.append(threading.get_ident()) + return real_write(self, *args, **kwargs) + + monkeypatch.setattr(pathlib.Path, "write_text", recording_write) + + await _capture(slot, 1) + + assert seen_threads, ( + "the stage result file was never written -- this test no longer " + "exercises the write and would pass vacuously" + ) + assert threading.get_ident() not in seen_threads, ( + "the stage result was written on the event-loop thread; the filesystem " + "work must be handed to asyncio.to_thread" + ) + + +@pytest.mark.asyncio +async def test_stage_result_mkdir_runs_off_the_loop_thread(tmp_path, monkeypatch): + """``mkdir`` is the other blocking syscall, and it is first. + + Offloading only the write would leave a directory creation — which on a cold + or networked session directory is the slower of the two — still on the loop. + """ + monkeypatch.setattr("kiro_crew.dashboard.chat_orchestrator.config_dir", lambda: tmp_path) + slot = _slot_with("stage one output") + + seen_threads: list[int] = [] + real_mkdir = pathlib.Path.mkdir + target = tmp_path / "sessions" / "chat-1-stage" + + def recording_mkdir(self, *args, **kwargs): + # Scoped to the stage-result directory: unrelated mkdir traffic from + # fixtures or lazily-created config dirs runs on the loop legitimately + # and would make a global assertion fail for the wrong reason. + if self == target: + seen_threads.append(threading.get_ident()) + return real_mkdir(self, *args, **kwargs) + + monkeypatch.setattr(pathlib.Path, "mkdir", recording_mkdir) + + await _capture(slot, 1) + + assert seen_threads, ( + "the session directory was never created -- this test no longer " + "exercises the stage-result mkdir and would pass vacuously" + ) + assert threading.get_ident() not in seen_threads, ( + "the session directory was created on the event-loop thread; it must be " + "handed to asyncio.to_thread with the write" + ) + + +@pytest.mark.asyncio +async def test_slot_messages_are_read_on_the_loop_thread(tmp_path, monkeypatch): + """The offload stops at the filesystem — live slot state is not shared. + + ``slot.messages`` is mutable and owned by the loop. Moving the traversal into + a worker would make an unrelated append during the capture a cross-thread + read, so the boundary is asserted from the safe side too, not just the fast + one. + """ + monkeypatch.setattr("kiro_crew.dashboard.chat_orchestrator.config_dir", lambda: tmp_path) + slot = _slot_with("stage one output") + + seen_threads: list[int] = [] + + class RecordingList(list): + def __reversed__(self): + seen_threads.append(threading.get_ident()) + return super().__reversed__() + + slot.messages = RecordingList(slot.messages) + + await _capture(slot, 1) + + assert seen_threads, ( + "slot.messages was never traversed -- this test no longer exercises the " + "extraction and would pass vacuously" + ) + assert seen_threads == [threading.get_ident()] * len(seen_threads), ( + "slot.messages was traversed off the event-loop thread; only the " + "filesystem write may cross into a worker" + ) + + +@pytest.mark.asyncio +async def test_capture_preserves_path_ordering_and_redaction(tmp_path, monkeypatch): + """The offload changes scheduling only — every output contract holds.""" + monkeypatch.setattr("kiro_crew.dashboard.chat_orchestrator.config_dir", lambda: tmp_path) + secret = "AKIAIOSFODNN7EXAMPLE" + slot = _slot_with("first part", f"second part with {secret}") + + path = await _capture(slot, 3) + + assert path == str(tmp_path / "sessions" / "chat-1-stage" / "stage_3_result.md") + written = pathlib.Path(path).read_text(encoding="utf-8") + # Oldest first: the extraction walks backwards then reverses. + assert written.index("first part") < written.index("second part") + assert secret not in written, "stage result persisted an unredacted credential" + + +@pytest.mark.asyncio +async def test_stop_landing_during_the_offloaded_write_does_not_advance(tmp_path, monkeypatch): + """A stop that lands while the worker is writing must halt the run. + + The offloaded write is a suspension point between the loop's last stop + check and the next stage. The user cancel path sets ``tracker.stopped`` + (and ``slot._auto_run``), not ``slot._stopping`` -- and ``auto_run`` is a + call-time snapshot -- so without a re-check after the await the loop + resumes and executes the next stage against a revoked approval. + """ + from unittest.mock import MagicMock + + monkeypatch.setattr("kiro_crew.dashboard.state.config_dir", lambda: tmp_path) + monkeypatch.setattr("kiro_crew.dashboard.chat_orchestrator.config_dir", lambda: tmp_path) + from kiro_crew.dashboard.chat_orchestrator import _stage_loop + from kiro_crew.dashboard.state import _ChatSlot + + state = MagicMock() + state.broadcast_ws = MagicMock() + state.push_slots_update = MagicMock() + state.subagents = MagicMock() + state.subagents.running_agents_for = MagicMock(return_value=[]) + + slot = _ChatSlot("stop-mid-write", mode="orchestrator") + slot._stage_titles = ["A", "B"] + slot._orch_tracker = None + + executed: list[int] = [] + + async def _turn(s, sl, msg, **kw): + executed.append(len(executed) + 1) + sl.append("assistant", f"stage {len(executed)} output", "msg msg-a", broadcast=False) + + monkeypatch.setattr("kiro_crew.dashboard.chat_orchestrator._run_chat", _turn) + + # The cancel lands while the stage-1 result is being written: the worker + # is mid-write when the user's stop is processed on the loop thread. + real_write = pathlib.Path.write_text + + def stopping_write(self, *args, **kwargs): + if "stage_1_result" in self.name: + slot._orch_tracker.stop() + slot._auto_run = False + return real_write(self, *args, **kwargs) + + monkeypatch.setattr(pathlib.Path, "write_text", stopping_write) + + await _stage_loop(state, slot, auto_run=True) + + assert executed == [1], ( + "stage 2 executed after the user's stop landed during the stage-1 " + "result write -- the loop must re-check the stop flags after the " + "offloaded write instead of advancing on a stale snapshot" + ) + + +@pytest.mark.asyncio +async def test_cancelled_capture_never_publishes_the_stage_file(tmp_path, monkeypatch): + """A cancelled capture must leave the canonical stage file untouched. + + ``asyncio.to_thread`` cannot interrupt a worker mid-syscall, so a + cancelled await abandons a still-running writer. If that writer owned + the canonical ``stage_N_result.md``, its bytes could land AFTER a + resumed plan (or a new slot reusing the key) wrote its own result — + silent corruption. The structural guarantee pinned here: the worker + only ever writes a uniquely-named temp file, and publication to the + canonical path happens on the loop thread strictly after an uncancelled + return. An abandoned worker's entire blast radius is an orphan temp + file. + """ + import asyncio + import threading + + monkeypatch.setattr("kiro_crew.dashboard.chat_orchestrator.config_dir", lambda: tmp_path) + slot = _slot_with("slow stage output") + + session_dir = tmp_path / "sessions" / slot.key + final = session_dir / "stage_1_result.md" + + started = threading.Event() + release = threading.Event() + real_write = pathlib.Path.write_text + + def slow_write(self, *args, **kwargs): + if self.parent == session_dir: + started.set() + release.wait(timeout=10) + return real_write(self, *args, **kwargs) + + monkeypatch.setattr(pathlib.Path, "write_text", slow_write) + + task = asyncio.ensure_future(_capture(slot, 1)) + for _ in range(500): + if started.is_set(): + break + await asyncio.sleep(0.01) + assert started.is_set(), "the worker never reached the write -- test scaffold broke" + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + # The new plan's result lands on the canonical path... + session_dir.mkdir(parents=True, exist_ok=True) + real_write(final, "NEW PLAN RESULT", encoding="utf-8") + + # ...then the abandoned worker finishes. It must not clobber the file. + release.set() + for _ in range(500): + await asyncio.sleep(0.01) + if threading.active_count() == 1: + break + await asyncio.sleep(0.05) + + assert final.read_text(encoding="utf-8") == "NEW PLAN RESULT", ( + "an abandoned stage-result writer overwrote the canonical stage file " + "written by a resumed plan -- the worker must only ever write a " + "uniquely-named temp file, with publication happening on the loop " + "thread after an uncancelled return" + ) + + +@pytest.mark.asyncio +async def test_stage_result_publication_runs_off_the_loop_thread(tmp_path, monkeypatch): + """The rename that PUBLISHES the result must not run on the loop either. + + ``os.replace`` reads as a free metadata-only syscall, which is why it is easy + to leave behind on the loop after the payload write has been offloaded. On a + network-backed session directory it is a round trip like any other, so it + stalls chat streaming, WebSocket frames and cron dispatch exactly as the + write did — the same ``no-blocking-call-on-event-loop`` anchor covers both. + + Pinned separately from the write because the two moved at different times: + an implementation can satisfy the write assertion above while still + publishing on the loop, which is precisely the state this test was added for. + """ + monkeypatch.setattr("kiro_crew.dashboard.chat_orchestrator.config_dir", lambda: tmp_path) + slot = _slot_with("stage one output") + + seen_threads: list[int] = [] + real_replace = os.replace + target_dir = tmp_path / "sessions" / "chat-1-stage" + + def recording_replace(src, dst, *args, **kwargs): + dst_path = pathlib.Path(dst) + if dst_path.parent == target_dir and "stage_1_result" in dst_path.name: + seen_threads.append(threading.get_ident()) + return real_replace(src, dst, *args, **kwargs) + + monkeypatch.setattr(os, "replace", recording_replace) + + final = await _capture(slot, 1) + + assert seen_threads, ( + "the stage result was never published through os.replace -- this test no " + "longer exercises publication and would pass vacuously" + ) + assert threading.get_ident() not in seen_threads, ( + "the stage result was published on the event-loop thread; the rename is " + "filesystem I/O and belongs in the same worker as the payload write" + ) + assert ( + pathlib.Path(final).read_text(encoding="utf-8") == "stage one output" + ), "publication moved off the loop but stopped producing the canonical file" + + +@pytest.mark.asyncio +async def test_no_filesystem_call_reaches_the_loop_thread(tmp_path, monkeypatch): + """Whole-boundary check: mkdir, write and replace are all off the loop. + + The per-call tests above each pin one syscall, so a future fourth one could + be added on the loop without failing any of them. This asserts the boundary + itself rather than its current members. + """ + monkeypatch.setattr("kiro_crew.dashboard.chat_orchestrator.config_dir", lambda: tmp_path) + slot = _slot_with("stage one output") + + loop_thread = threading.get_ident() + offenders: list[str] = [] + + real_mkdir = pathlib.Path.mkdir + real_write = pathlib.Path.write_text + real_replace = os.replace + target_dir = tmp_path / "sessions" / "chat-1-stage" + + def recording_mkdir(self, *args, **kwargs): + if threading.get_ident() == loop_thread and self == target_dir: + offenders.append("mkdir") + return real_mkdir(self, *args, **kwargs) + + def recording_write(self, *args, **kwargs): + if threading.get_ident() == loop_thread and self.parent == target_dir: + offenders.append("write_text") + return real_write(self, *args, **kwargs) + + def recording_replace(src, dst, *args, **kwargs): + if threading.get_ident() == loop_thread and pathlib.Path(dst).parent == target_dir: + offenders.append("os.replace") + return real_replace(src, dst, *args, **kwargs) + + monkeypatch.setattr(pathlib.Path, "mkdir", recording_mkdir) + monkeypatch.setattr(pathlib.Path, "write_text", recording_write) + monkeypatch.setattr(os, "replace", recording_replace) + + await _capture(slot, 1) + + assert offenders == [], ( + "stage-result capture performed filesystem work on the event-loop " + "thread: %s" % ", ".join(offenders) + )