mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-11-06 05:17:15 +00:00
73 lines
3.1 KiB
Python
73 lines
3.1 KiB
Python
"""``context_length_cache.yaml`` writes and their cross-process serialization.
|
|
|
|
Every agent (CLI, gateway sessions, cron jobs, subagents) persists probed context windows with a
|
|
load -> mutate -> write cycle; unserialized, one writer rewrites a stale snapshot and drops the
|
|
other's entry (or resurrects an invalidated, too-large length).
|
|
"""
|
|
|
|
import contextlib
|
|
import functools
|
|
import logging
|
|
import threading
|
|
|
|
from utils import atomic_yaml_write
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_CONTEXT_CACHE_LOCK_HOLDER = threading.local()
|
|
_CONTEXT_CACHE_LOCK_TIMEOUT_SECONDS = 5.0
|
|
|
|
|
|
def _write_context_cache(cache: dict[str, int], bedrock_confirmed: dict | None = None) -> None:
|
|
"""Atomic write (a truncating write killed mid-dump leaves a partial file that
|
|
_load_context_cache() swallows as {}, wiping EVERY cached length). Raises on failure."""
|
|
from agent.model_metadata import _get_context_cache_path # facade seam: tests patch the path there
|
|
document = {"context_lengths": cache}
|
|
if bedrock_confirmed:
|
|
document["bedrock_confirmed_v1"] = bedrock_confirmed
|
|
atomic_yaml_write(_get_context_cache_path(), document)
|
|
|
|
|
|
def locked_context_cache(fn):
|
|
"""Run ``fn``'s load -> mutate -> write cycle under the cache file's cross-process lock.
|
|
|
|
The cache is best-effort: when the lock cannot be taken in time the update is skipped, never raised.
|
|
"""
|
|
@functools.wraps(fn)
|
|
def wrapper(*args, **kwargs):
|
|
from agent.model_metadata import _get_context_cache_path # facade seam: tests patch the path there
|
|
from hermes_cli.auth import _file_lock
|
|
path = _get_context_cache_path()
|
|
lock_path = path.with_name(path.name + ".lock")
|
|
with contextlib.ExitStack() as stack:
|
|
try:
|
|
stack.enter_context(_file_lock(
|
|
lock_path, _CONTEXT_CACHE_LOCK_HOLDER, _CONTEXT_CACHE_LOCK_TIMEOUT_SECONDS,
|
|
f"Timed out waiting for the context length cache lock ({lock_path})"))
|
|
except (OSError, TimeoutError):
|
|
logger.debug("Skipping context length cache update", exc_info=True)
|
|
return None
|
|
return fn(*args, **kwargs)
|
|
return wrapper
|
|
|
|
|
|
@locked_context_cache
|
|
def drop_context_cache_rows(stale_keys: set[str], label: str) -> None:
|
|
"""Remove every key shape of one cached entry (and its Bedrock provenance) under the cache lock.
|
|
|
|
Only the file update is locked: callers evict their in-process probe caches unconditionally."""
|
|
from agent.model_metadata import _load_context_cache_document # facade seam: tests patch it there
|
|
document = _load_context_cache_document()
|
|
cache = document.get("context_lengths") or {}
|
|
confirmed = document.get("bedrock_confirmed_v1") or {}
|
|
confirmed = confirmed if isinstance(confirmed, dict) else {}
|
|
if not any(k in cache or k in confirmed for k in stale_keys):
|
|
return
|
|
for k in stale_keys:
|
|
cache.pop(k, None)
|
|
confirmed.pop(k, None)
|
|
try:
|
|
_write_context_cache(cache, confirmed)
|
|
except Exception:
|
|
logger.debug("Failed to invalidate context length cache entry %s", label, exc_info=True)
|