diff --git a/gateway/shutdown_flush.py b/gateway/shutdown_flush.py index 3ffa74ea25644..e144c6ef43e7a 100644 --- a/gateway/shutdown_flush.py +++ b/gateway/shutdown_flush.py @@ -55,29 +55,16 @@ def _fsync_directory(path: Path) -> None: def _write_payload(flush_dir: Path, payload: Dict[str, Any]) -> None: """Atomically write one private, uniquely named recovery payload.""" + from utils import atomic_json_write + file_id = uuid.uuid4().hex final_path = flush_dir / f"pending-{file_id}.json" - temp_path = flush_dir / f".pending-{file_id}.tmp" - file_descriptor = -1 - - try: - file_descriptor = os.open( - temp_path, - os.O_WRONLY | os.O_CREAT | os.O_EXCL, - 0o600, - ) - with os.fdopen(file_descriptor, "w", encoding="utf-8") as handle: - file_descriptor = -1 # The file object now owns the descriptor. - json.dump(payload, handle, ensure_ascii=False, default=str) - handle.flush() - os.fsync(handle.fileno()) - - os.replace(temp_path, final_path) - except Exception: - if file_descriptor >= 0: - os.close(file_descriptor) - temp_path.unlink(missing_ok=True) - raise + atomic_json_write( + final_path, + payload, + mode=0o600, + default=str, + ) try: _fsync_directory(flush_dir) diff --git a/tests/gateway/test_shutdown_flush.py b/tests/gateway/test_shutdown_flush.py index 92b070964142f..a6f57b8137b0b 100644 --- a/tests/gateway/test_shutdown_flush.py +++ b/tests/gateway/test_shutdown_flush.py @@ -74,7 +74,7 @@ def test_flush_write_failure_leaves_no_recovery_file(tmp_path, monkeypatch): def fail_replace(source, destination): raise OSError("simulated replace failure") - monkeypatch.setattr("gateway.shutdown_flush.os.replace", fail_replace) + monkeypatch.setattr("utils.os.replace", fail_replace) assert flush_pending_to_file({"session": "message"}, reason="test") == 0 assert list(flush_dir.iterdir()) == []