diff --git a/apodex/session.py b/apodex/session.py index d58fcf0..b76e098 100644 --- a/apodex/session.py +++ b/apodex/session.py @@ -11,6 +11,7 @@ import asyncio import json import os +import threading from pathlib import Path from typing import Any @@ -189,6 +190,10 @@ def __init__( # plugins.tools._path_auth._authorized_local_path). Without this they # only allow a few default dirs and deny the user's repo. self._authorize_workspace(cwd) + # _persist() now runs both on the main thread (start_new_session, + # rename_session) and off-thread (_on_turn's asyncio.to_thread), so + # concurrent writers must serialize on the same checkpoint file. + self._persist_lock = threading.Lock() @staticmethod def _active_spill_workspace() -> Path | None: @@ -476,7 +481,11 @@ async def _on_turn(self, turn: int, messages: list, metadata: dict) -> None: after each completed turn — keep history current and persist.""" self.history = list(messages) self.display_history = list(messages) - self._persist() + # _persist() does synchronous file I/O over the full history; run it + # off the event loop so long sessions don't stall on every turn. + # Awaited between turns; _persist_lock also serializes snapshots and + # writes if cancellation leaves this worker running in the background. + await asyncio.to_thread(self._persist) # ── persistence (interrupt-safe resume) ─────────────────────────────── def _enrich_task(self, task: str) -> str: @@ -591,21 +600,29 @@ def replay_history(self) -> list[Message]: def _persist(self) -> None: """Checkpoint session state so ``--resume `` can continue it. - Best-effort; a failed write never disrupts the session.""" + Best-effort; a failed write never disrupts the session. + + Serialized via ``_persist_lock`` and written atomically (tmp file + + ``os.replace``) because this runs from both the main thread + (``start_new_session`` / ``rename_session``) and a worker thread + (``_on_turn``'s ``asyncio.to_thread``) — without both, concurrent + writers can interleave and corrupt the checkpoint file.""" try: import json from apodex.todo import get_todos - snapshot = getattr(self.r, "snapshot_state", None) - if callable(snapshot): - raw_tui_state = snapshot() - self.tui_state = raw_tui_state if isinstance(raw_tui_state, dict) else {} + # Snapshot under the same lock as the write: a cancelled + # to_thread worker can otherwise overwrite a newer checkpoint + # with a payload it captured before waiting for this lock. + with self._persist_lock: + snapshot = getattr(self.r, "snapshot_state", None) + if callable(snapshot): + raw_tui_state = snapshot() + self.tui_state = raw_tui_state if isinstance(raw_tui_state, dict) else {} - path = _session_state_path(self.session_id) - os.makedirs(os.path.dirname(path), exist_ok=True) - with open(path, "w", encoding="utf-8") as f: - json.dump({ + path = _session_state_path(self.session_id) + payload = { "session_id": self.session_id, "created_at": self.created_at, "local_timezone": self.local_timezone, @@ -633,7 +650,12 @@ def _persist(self) -> None: {"content": item.content, "status": item.status} for item in get_todos() ], - }, f, ensure_ascii=False) + } + os.makedirs(os.path.dirname(path), exist_ok=True) + tmp_path = f"{path}.{os.getpid()}.{threading.get_ident()}.tmp" + with open(tmp_path, "w", encoding="utf-8") as f: + json.dump(payload, f, ensure_ascii=False) + os.replace(tmp_path, path) except Exception: pass diff --git a/apodex/tests/test_changes.py b/apodex/tests/test_changes.py index 41de6a8..265f9a2 100644 --- a/apodex/tests/test_changes.py +++ b/apodex/tests/test_changes.py @@ -748,6 +748,89 @@ def test_session_persist_and_resume(tmp_path, monkeypatch): assert "edited.py" in s2.journal.to_dict() or True # journal restored shape + +@pytest.mark.asyncio +async def test_cancelled_checkpoint_cannot_overwrite_newer_save(tmp_path, monkeypatch): + import threading + from types import SimpleNamespace + + from apodex import session as session_module + from apodex.session import TerminalSession + + checkpoint = tmp_path / "state.json" + monkeypatch.setattr(session_module, "_session_state_path", lambda _: str(checkpoint)) + snapshot_started = threading.Event() + release_snapshot = threading.Event() + old_save_finished = threading.Event() + old_thread_id = None + + class CheckpointLock: + def __init__(self): + self.lock = threading.Lock() + + def __enter__(self): + # Let the old snapshot finish once a newer save is waiting for + # its lock. Before the fix the old snapshot owns no lock, so the + # newer save completes first and releases it in the test below. + if threading.get_ident() != old_thread_id and self.lock.locked(): + release_snapshot.set() + self.lock.acquire() + + def __exit__(self, *_): + self.lock.release() + + def journal_snapshot(): + if threading.get_ident() == old_thread_id: + snapshot_started.set() + if not release_snapshot.wait(5): + raise TimeoutError("old checkpoint was never released") + return {} + + session = TerminalSession.__new__(TerminalSession) + session.__dict__.update( + r=SimpleNamespace(), session_id="test", created_at="", local_timezone="", + session_name="old name", mode="coding", cwd=str(tmp_path), + cfg=SimpleNamespace(model="fake"), history=[], display_history=[], + workflow_turns=[], usage=SimpleNamespace(to_dict=lambda: {}), tui_state={}, + journal=SimpleNamespace(to_dict=journal_snapshot, + observed_paths=lambda: [], revert_bases=lambda: {}), + plan_state=SimpleNamespace(active=False), _persist_lock=CheckpointLock(), + ) + persist = session._persist + + def tracked_persist(): + nonlocal old_thread_id + is_old = old_thread_id is None + if is_old: + old_thread_id = threading.get_ident() + try: + persist() + finally: + if is_old: + old_save_finished.set() + + monkeypatch.setattr(session, "_persist", tracked_persist) + turn = asyncio.create_task(session._on_turn(1, [{"role": "user", "content": "old"}], {})) + try: + assert await asyncio.to_thread(snapshot_started.wait, 5) + turn.cancel() + with pytest.raises(asyncio.CancelledError): + await turn + session.history = [{"role": "user", "content": "new"}] + session.display_history = list(session.history) + await asyncio.to_thread(session.rename_session, "new name") + finally: + release_snapshot.set() + assert await asyncio.to_thread(old_save_finished.wait, 5) + if not turn.done(): + await turn + + state = json.loads(checkpoint.read_text()) + assert state["name"] == "new name" + assert state["history"] == session.history + assert state["display_history"] == session.display_history + + def test_follow_up_receives_exact_agent_and_host_deliverable_paths(tmp_path, monkeypatch): from apodex.config import ModelConfig from apodex.render import Renderer