diff --git a/agent/model_metadata.py b/agent/model_metadata.py index 636034f558..dedb1ed9a0 100644 --- a/agent/model_metadata.py +++ b/agent/model_metadata.py @@ -21,7 +21,7 @@ import yaml if TYPE_CHECKING: # pragma: no cover — runtime import is lazy (see below) import requests -from utils import atomic_json_write, base_url_host_matches, base_url_hostname +from utils import atomic_json_write, atomic_yaml_write, base_url_host_matches, base_url_hostname from hermes_constants import OPENROUTER_MODELS_URL @@ -1476,9 +1476,13 @@ def save_context_length(model: str, base_url: str, length: int) -> None: cache[key] = length path = _get_context_cache_path() try: - path.parent.mkdir(parents=True, exist_ok=True) - with open(path, "w", encoding="utf-8") as f: - yaml.dump({"context_lengths": cache}, f, default_flow_style=False) + # Atomic write (temp file + fsync + os.replace): a plain truncating + # ``open(path, "w")`` leaves the file empty/partial if the process is + # killed mid-dump, and the next _load_context_cache() swallows the + # resulting YAML error and returns {} — silently wiping EVERY cached + # context length. It also exposes torn reads to a concurrent process + # reading between truncate and dump-complete. + atomic_yaml_write(path, {"context_lengths": cache}) logger.info("Cached context length %s -> %s tokens", key, f"{length:,}") except Exception as e: logger.debug("Failed to save context length cache: %s", e) @@ -1525,9 +1529,9 @@ def _invalidate_cached_context_length(model: str, base_url: str) -> None: cache.pop(k, None) path = _get_context_cache_path() try: - path.parent.mkdir(parents=True, exist_ok=True) - with open(path, "w", encoding="utf-8") as f: - yaml.dump({"context_lengths": cache}, f, default_flow_style=False) + # Atomic write — see save_context_length() for why a plain truncating + # open() here risks wiping the entire cache on an interrupted dump. + atomic_yaml_write(path, {"context_lengths": cache}) except Exception as e: logger.debug("Failed to invalidate context length cache entry %s: %s", key, e) diff --git a/tests/agent/test_model_metadata.py b/tests/agent/test_model_metadata.py index 78ba56c167..726bd4e68f 100644 --- a/tests/agent/test_model_metadata.py +++ b/tests/agent/test_model_metadata.py @@ -1206,6 +1206,41 @@ class TestContextLengthCache: assert get_model_context_length("unknown/model", base_url="http://local") == 65536 + def test_write_failure_leaves_existing_cache_intact(self, tmp_path, monkeypatch): + """An interrupted write must not corrupt or wipe the existing cache. + + The old non-atomic ``open(path, "w")`` truncated the file before + dumping, so a crash/kill mid-write left empty or partial YAML — and + the next load swallowed the error and returned ``{}``, silently + wiping EVERY persisted context length. The atomic temp-file + + ``os.replace`` write leaves the previous file byte-for-byte intact + when the swap fails. + """ + import utils + import agent.model_metadata as mm + + cache_file = tmp_path / "cache.yaml" + monkeypatch.setattr(mm, "_get_context_cache_path", lambda: cache_file) + + # Seed a valid, populated cache. + save_context_length("model-a", "http://a", 64000) + original_bytes = cache_file.read_bytes() + + # Simulate a crash during the atomic swap step. + def _boom(*_args, **_kwargs): + raise OSError("simulated crash during atomic replace") + + monkeypatch.setattr(utils, "atomic_replace", _boom) + + # save_context_length is best-effort and swallows the error. + save_context_length("model-b", "http://b", 128000) + + # Original file survives untouched — not truncated or emptied. + assert cache_file.read_bytes() == original_bytes + assert get_cached_context_length("model-a", "http://a") == 64000 + # The failed write must not leave a stray temp file behind. + assert list(cache_file.parent.glob(".cache_*.tmp")) == [] + class TestGrok43StaleCacheGuard: """Pre-catalog builds resolved grok-4.3 via the generic 'grok-4' catch-all