diff --git a/gateway/session.py b/gateway/session.py index c14e9bd0301a..f69e94c22572 100644 --- a/gateway/session.py +++ b/gateway/session.py @@ -1002,11 +1002,28 @@ def rewrite_transcript(self, session_id: str, messages: List[Dict[str, Any]]) -> except Exception as e: logger.debug("Failed to rewrite transcript in DB: %s", e) - # JSONL: overwrite the file + # JSONL: write to a temp file first so crashes never leave partial output transcript_path = self.get_transcript_path(session_id) - with open(transcript_path, "w", encoding="utf-8") as f: - for msg in messages: - f.write(json.dumps(msg, ensure_ascii=False) + "\n") + transcript_path.parent.mkdir(parents=True, exist_ok=True) + + import tempfile + + fd, tmp_path = tempfile.mkstemp( + dir=str(transcript_path.parent), suffix=".tmp", prefix=".transcript_" + ) + try: + with os.fdopen(fd, "w", encoding="utf-8") as f: + for msg in messages: + f.write(json.dumps(msg, ensure_ascii=False) + "\n") + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_path, transcript_path) + except BaseException: + try: + os.unlink(tmp_path) + except OSError as e: + logger.debug("Could not remove temp file %s: %s", tmp_path, e) + raise def load_transcript(self, session_id: str) -> List[Dict[str, Any]]: """Load all messages from a session's transcript.""" diff --git a/tests/gateway/test_session.py b/tests/gateway/test_session.py index 39e4aad3d652..1df0770e77f7 100644 --- a/tests/gateway/test_session.py +++ b/tests/gateway/test_session.py @@ -420,6 +420,40 @@ def test_rewrite_with_empty_list(self, store): reloaded = store.load_transcript(session_id) assert reloaded == [] + def test_rewrite_failure_leaves_existing_jsonl_intact(self, store, monkeypatch): + session_id = "test_session_atomic_failure" + original_messages = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + ] + replacement_messages = [ + {"role": "user", "content": "rewrite me"}, + {"role": "assistant", "content": "this should not land"}, + ] + + for msg in original_messages: + store.append_to_transcript(session_id, msg) + + transcript_path = store.get_transcript_path(session_id) + original_contents = transcript_path.read_text(encoding="utf-8") + real_dumps = json.dumps + call_count = 0 + + def fail_mid_rewrite(*args, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 2: + raise RuntimeError("boom during rewrite") + return real_dumps(*args, **kwargs) + + monkeypatch.setattr("gateway.session.json.dumps", fail_mid_rewrite) + + with pytest.raises(RuntimeError, match="boom during rewrite"): + store.rewrite_transcript(session_id, replacement_messages) + + assert transcript_path.read_text(encoding="utf-8") == original_contents + assert store.load_transcript(session_id) == original_messages + class TestLoadTranscriptCorruptLines: """Regression: corrupt JSONL lines (e.g. from mid-write crash) must be