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) + )