fix(apodex): offload per-turn session persist off the event loop

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 <id>`` 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
 

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