fix: reuse shared atomic session log helper

This commit is contained in:
teknium1
2026-03-14 02:56:13 -07:00
parent f685741481
commit cbbba87099
4 changed files with 53 additions and 15 deletions

View File

@@ -28,7 +28,6 @@ import json
import logging import logging
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
import os import os
import tempfile
import random import random
import re import re
import sys import sys
@@ -101,6 +100,7 @@ from agent.trajectory import (
convert_scratchpad_to_think, has_incomplete_scratchpad, convert_scratchpad_to_think, has_incomplete_scratchpad,
save_trajectory as _save_trajectory_to_file, save_trajectory as _save_trajectory_to_file,
) )
from utils import atomic_json_write
HONCHO_TOOL_NAMES = { HONCHO_TOOL_NAMES = {
"honcho_context", "honcho_context",
@@ -1399,17 +1399,12 @@ class AIAgent:
"messages": cleaned, "messages": cleaned,
} }
tmp_dir = str(self.session_log_file.parent) if hasattr(self.session_log_file, 'parent') else os.path.dirname(str(self.session_log_file)) atomic_json_write(
fd, tmp_path = tempfile.mkstemp(dir=tmp_dir, suffix='.tmp', prefix='.session_') self.session_log_file,
try: entry,
with os.fdopen(fd, 'w', encoding='utf-8') as f: indent=2,
json.dump(entry, f, indent=2, ensure_ascii=False, default=str) default=str,
f.flush() )
os.fsync(f.fileno())
os.replace(tmp_path, str(self.session_log_file))
except:
os.unlink(tmp_path)
raise
except Exception as e: except Exception as e:
if self.verbose_logging: if self.verbose_logging:

View File

@@ -97,6 +97,17 @@ class TestAtomicJsonWrite:
text = target.read_text() text = target.read_text()
assert ' "a"' in text # 4-space indent assert ' "a"' in text # 4-space indent
def test_accepts_json_dump_default_hook(self, tmp_path):
class CustomValue:
def __str__(self):
return "custom-value"
target = tmp_path / "custom_default.json"
atomic_json_write(target, {"value": CustomValue()}, default=str)
result = json.loads(target.read_text(encoding="utf-8"))
assert result == {"value": "custom-value"}
def test_unicode_content(self, tmp_path): def test_unicode_content(self, tmp_path):
target = tmp_path / "unicode.json" target = tmp_path / "unicode.json"
data = {"emoji": "🎉", "japanese": "日本語"} data = {"emoji": "🎉", "japanese": "日本語"}

View File

@@ -1928,6 +1928,24 @@ class TestSafeWriter:
assert inner.getvalue() == "test" assert inner.getvalue() == "test"
class TestSaveSessionLogAtomicWrite:
def test_uses_shared_atomic_json_helper(self, agent, tmp_path):
agent.session_log_file = tmp_path / "session.json"
messages = [{"role": "user", "content": "hello"}]
with patch("run_agent.atomic_json_write", create=True) as mock_atomic_write:
agent._save_session_log(messages)
mock_atomic_write.assert_called_once()
call_args = mock_atomic_write.call_args
assert call_args.args[0] == agent.session_log_file
payload = call_args.args[1]
assert payload["session_id"] == agent.session_id
assert payload["messages"] == messages
assert call_args.kwargs["indent"] == 2
assert call_args.kwargs["default"] is str
# =================================================================== # ===================================================================
# Anthropic adapter integration fixes # Anthropic adapter integration fixes
# =================================================================== # ===================================================================

View File

@@ -9,17 +9,25 @@ from typing import Any, Union
import yaml import yaml
def atomic_json_write(path: Union[str, Path], data: Any, *, indent: int = 2) -> None: def atomic_json_write(
path: Union[str, Path],
data: Any,
*,
indent: int = 2,
**dump_kwargs: Any,
) -> None:
"""Write JSON data to a file atomically. """Write JSON data to a file atomically.
Uses temp file + fsync + os.replace to ensure the target file is never Uses temp file + fsync + os.replace to ensure the target file is never
left in a partially-written state. If the process crashes mid-write, left in a partially-written state. If the process crashes mid-write,
the previous version of the file remains intact. the previous version of the file remains intact.
Args: Args:
path: Target file path (will be created or overwritten). path: Target file path (will be created or overwritten).
data: JSON-serializable data to write. data: JSON-serializable data to write.
indent: JSON indentation (default 2). indent: JSON indentation (default 2).
**dump_kwargs: Additional keyword args forwarded to json.dump(), such
as default=str for non-native types.
""" """
path = Path(path) path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True) path.parent.mkdir(parents=True, exist_ok=True)
@@ -31,7 +39,13 @@ def atomic_json_write(path: Union[str, Path], data: Any, *, indent: int = 2) ->
) )
try: try:
with os.fdopen(fd, "w", encoding="utf-8") as f: with os.fdopen(fd, "w", encoding="utf-8") as f:
json.dump(data, f, indent=indent, ensure_ascii=False) json.dump(
data,
f,
indent=indent,
ensure_ascii=False,
**dump_kwargs,
)
f.flush() f.flush()
os.fsync(f.fileno()) os.fsync(f.fileno())
os.replace(tmp_path, path) os.replace(tmp_path, path)