diff --git a/gateway/session.py b/gateway/session.py index db90d3121727..eab07fb03448 100644 --- a/gateway/session.py +++ b/gateway/session.py @@ -19,6 +19,8 @@ from dataclasses import dataclass from typing import Dict, List, Optional, Any +from utils import atomic_jsonl_write + logger = logging.getLogger(__name__) @@ -1187,11 +1189,12 @@ 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: overwrite the file atomically. A crash between truncate and + # final flush would otherwise leave the transcript empty or partial, + # losing history for pre-DB sessions (see GH-1193 for the read-side + # counterpart that skips corrupt lines after the fact). 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") + atomic_jsonl_write(transcript_path, messages) 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 539b12a5e18e..81cd2939ced2 100644 --- a/tests/gateway/test_session.py +++ b/tests/gateway/test_session.py @@ -443,6 +443,83 @@ def test_rewrite_with_empty_list(self, store): assert reloaded == [] +class TestSessionStoreRewriteTranscriptAtomicity: + """Regression: rewrite_transcript must never leave the JSONL file in a + partially-written state. A crash between truncate and final flush would + otherwise wipe out transcript history — catastrophic for pre-DB sessions + where the JSONL is the only source of truth, and harmful for /retry, + /undo, /compress, /reset flows even when SQLite is present. + + Pairs with TestLoadTranscriptCorruptLines (GH-1193) which handles the + read-side fallback; this class verifies the write-side guarantee.""" + + @pytest.fixture() + def store(self, tmp_path): + config = GatewayConfig() + with patch("gateway.session.SessionStore._ensure_loaded"): + s = SessionStore(sessions_dir=tmp_path, config=config) + s._db = None # exercise the JSONL path directly + s._loaded = True + return s + + def test_mid_write_crash_preserves_previous_transcript(self, store): + session_id = "atomic_crash" + original = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + ] + store.rewrite_transcript(session_id, original) + + # Simulate a crash during json serialization of the *new* payload. + with patch("utils.json.dumps", side_effect=IOError("disk full")): + with pytest.raises(IOError): + store.rewrite_transcript( + session_id, + [{"role": "user", "content": "replacement"}], + ) + + # Transcript must still hold the pre-crash content. + reloaded = store.load_transcript(session_id) + assert reloaded == original + + def test_no_stray_temp_files_after_successful_rewrite(self, store, tmp_path): + session_id = "atomic_cleanup" + store.rewrite_transcript( + session_id, + [{"role": "user", "content": "ok"}], + ) + + session_dir = store.get_transcript_path(session_id).parent + stray = [p for p in session_dir.iterdir() if ".tmp" in p.name] + assert stray == [] + + def test_no_stray_temp_files_after_failed_rewrite(self, store): + session_id = "atomic_failure_cleanup" + store.rewrite_transcript(session_id, [{"role": "user", "content": "seed"}]) + + with patch("utils.json.dumps", side_effect=IOError("disk full")): + with pytest.raises(IOError): + store.rewrite_transcript( + session_id, + [{"role": "user", "content": "attempt"}], + ) + + session_dir = store.get_transcript_path(session_id).parent + stray = [p for p in session_dir.iterdir() if ".tmp" in p.name] + assert stray == [] + + def test_unicode_roundtrip(self, store): + session_id = "atomic_unicode" + messages = [ + {"role": "user", "content": "hello 🎉"}, + {"role": "assistant", "content": "日本語"}, + ] + store.rewrite_transcript(session_id, messages) + + reloaded = store.load_transcript(session_id) + assert reloaded == messages + + class TestLoadTranscriptCorruptLines: """Regression: corrupt JSONL lines (e.g. from mid-write crash) must be skipped instead of crashing the entire transcript load. GH-1193.""" diff --git a/tests/hermes_cli/test_atomic_jsonl_write.py b/tests/hermes_cli/test_atomic_jsonl_write.py new file mode 100644 index 000000000000..ca45fbb85a61 --- /dev/null +++ b/tests/hermes_cli/test_atomic_jsonl_write.py @@ -0,0 +1,153 @@ +"""Tests for utils.atomic_jsonl_write — crash-safe JSONL file writes.""" + +import json +from pathlib import Path +from unittest.mock import patch + +import pytest + +from utils import atomic_jsonl_write + + +def _read_lines(path: Path) -> list: + return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line] + + +class TestAtomicJsonlWrite: + """Core atomic write behavior.""" + + def test_writes_valid_jsonl(self, tmp_path): + target = tmp_path / "lines.jsonl" + items = [{"i": 1}, {"i": 2}, {"i": 3}] + atomic_jsonl_write(target, items) + + assert _read_lines(target) == items + + def test_each_item_on_its_own_line(self, tmp_path): + target = tmp_path / "lines.jsonl" + atomic_jsonl_write(target, [{"a": 1}, {"b": 2}]) + + lines = target.read_text(encoding="utf-8").splitlines() + assert lines == ['{"a": 1}', '{"b": 2}'] + + def test_creates_parent_directories(self, tmp_path): + target = tmp_path / "deep" / "nested" / "dir" / "lines.jsonl" + atomic_jsonl_write(target, [{"ok": True}]) + + assert target.exists() + assert _read_lines(target) == [{"ok": True}] + + def test_overwrites_existing_file(self, tmp_path): + target = tmp_path / "lines.jsonl" + target.write_text('{"old": true}\n', encoding="utf-8") + + atomic_jsonl_write(target, [{"new": True}]) + assert _read_lines(target) == [{"new": True}] + + def test_empty_iterable_produces_empty_file(self, tmp_path): + target = tmp_path / "empty.jsonl" + atomic_jsonl_write(target, []) + + assert target.exists() + assert target.read_text(encoding="utf-8") == "" + + def test_preserves_original_on_serialization_error(self, tmp_path): + target = tmp_path / "lines.jsonl" + original = [{"preserved": True}] + target.write_text(json.dumps(original[0]) + "\n", encoding="utf-8") + + with pytest.raises(TypeError): + atomic_jsonl_write(target, [{"bad": object()}]) + + assert _read_lines(target) == original + + def test_no_leftover_temp_files_on_success(self, tmp_path): + target = tmp_path / "lines.jsonl" + atomic_jsonl_write(target, [{"i": 1}]) + + tmp_files = [f for f in tmp_path.iterdir() if ".tmp" in f.name] + assert tmp_files == [] + assert target.exists() + + def test_no_leftover_temp_files_on_failure(self, tmp_path): + target = tmp_path / "lines.jsonl" + + with pytest.raises(TypeError): + atomic_jsonl_write(target, [{"bad": object()}]) + + tmp_files = [f for f in tmp_path.iterdir() if ".tmp" in f.name] + assert tmp_files == [] + + def test_cleans_up_temp_file_on_baseexception(self, tmp_path): + """KeyboardInterrupt/SystemExit must not leave stray .tmp files.""" + + class SimulatedAbort(BaseException): + pass + + target = tmp_path / "lines.jsonl" + original = [{"preserved": True}] + target.write_text(json.dumps(original[0]) + "\n", encoding="utf-8") + + with patch("utils.json.dumps", side_effect=SimulatedAbort): + with pytest.raises(SimulatedAbort): + atomic_jsonl_write(target, [{"new": True}]) + + tmp_files = [f for f in tmp_path.iterdir() if ".tmp" in f.name] + assert tmp_files == [] + assert _read_lines(target) == original + + def test_mid_write_failure_preserves_prior_file(self, tmp_path): + """A crash after some lines are written must not clobber the target.""" + target = tmp_path / "lines.jsonl" + original = [{"i": 0}, {"i": 1}] + target.write_text( + "\n".join(json.dumps(item) for item in original) + "\n", + encoding="utf-8", + ) + + call_count = {"n": 0} + real_dumps = json.dumps + + def fail_after_two(obj, **kwargs): + call_count["n"] += 1 + if call_count["n"] > 2: + raise IOError("simulated mid-write crash") + return real_dumps(obj, **kwargs) + + with patch("utils.json.dumps", side_effect=fail_after_two): + with pytest.raises(IOError): + atomic_jsonl_write(target, [{"j": i} for i in range(5)]) + + # Target file still holds the original content — tempfile discarded. + assert _read_lines(target) == original + tmp_files = [f for f in tmp_path.iterdir() if ".tmp" in f.name] + assert tmp_files == [] + + def test_accepts_string_path(self, tmp_path): + target = str(tmp_path / "string_path.jsonl") + atomic_jsonl_write(target, [{"string": True}]) + + assert _read_lines(Path(target)) == [{"string": True}] + + def test_unicode_content(self, tmp_path): + target = tmp_path / "unicode.jsonl" + items = [{"emoji": "🎉"}, {"japanese": "日本語"}] + atomic_jsonl_write(target, items) + + assert _read_lines(target) == items + + def test_accepts_generator_of_items(self, tmp_path): + target = tmp_path / "gen.jsonl" + atomic_jsonl_write(target, ({"i": i} for i in range(3))) + + assert _read_lines(target) == [{"i": 0}, {"i": 1}, {"i": 2}] + + def test_forwards_dump_kwargs(self, tmp_path): + class CustomValue: + def __str__(self): + return "custom-value" + + target = tmp_path / "custom.jsonl" + atomic_jsonl_write(target, [{"value": CustomValue()}], default=str) + + assert _read_lines(target) == [{"value": "custom-value"}] diff --git a/utils.py b/utils.py index f3d38006d145..e964eab2065f 100644 --- a/utils.py +++ b/utils.py @@ -162,6 +162,52 @@ def atomic_yaml_write( raise +def atomic_jsonl_write( + path: Union[str, Path], + items: Any, + **dump_kwargs: Any, +) -> None: + """Write an iterable of JSON-serializable items to a file atomically as JSONL. + + Each item is serialized on its own line via ``json.dumps``. Uses temp + file + fsync + os.replace so the target file is never observable in a + partially-written state — if the process crashes mid-write, the + previous version of the file remains intact. + + Args: + path: Target file path (will be created or overwritten). + items: Iterable of JSON-serializable items, one per line. + **dump_kwargs: Additional keyword args forwarded to ``json.dumps``. + ``ensure_ascii`` defaults to ``False`` to preserve non-ASCII text. + """ + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + + dump_kwargs.setdefault("ensure_ascii", False) + + original_mode = _preserve_file_mode(path) + + fd, tmp_path = tempfile.mkstemp( + dir=str(path.parent), + prefix=f".{path.stem}_", + suffix=".tmp", + ) + try: + with os.fdopen(fd, "w", encoding="utf-8") as f: + for item in items: + f.write(json.dumps(item, **dump_kwargs) + "\n") + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_path, path) + _restore_file_mode(path, original_mode) + except BaseException: + try: + os.unlink(tmp_path) + except OSError: + pass + raise + + # ─── JSON Helpers ─────────────────────────────────────────────────────────────