diff --git a/gateway/agent_cache_pressure.py b/gateway/agent_cache_pressure.py new file mode 100644 index 0000000000000..7be52e1bc9260 --- /dev/null +++ b/gateway/agent_cache_pressure.py @@ -0,0 +1,281 @@ +"""Memory-pressure bounds for the gateway's per-session AIAgent cache. + +The gateway caches one ``AIAgent`` per session so a long-lived conversation +reuses its prompt prefix instead of rebuilding the system prompt every turn. +Each cached agent also pins ``_session_messages`` — the full live transcript, +tool outputs included, which is tens of MB on a tool-heavy session. + +``gateway/run.py`` bounds that cache two ways, and both are blind to how much +memory it actually holds: + +* the LRU cap counts *entries*, not bytes, and 128 warm transcripts is + several GB; +* the idle TTL only sheds agents that went quiet for an hour, and it + deliberately defers eviction for a finalizable session that has not expired + yet, so a busy gateway hoards every transcript all day. + +This module supplies the missing signal: the process's own anonymous RSS, +compared against a budget derived from the cgroup limit the gateway actually +runs under. ``GatewayRunner._sweep_agent_cache_under_pressure`` uses it to +shed LRU transcripts through the existing soft-eviction path, which rebuilds +from the persisted session on the next turn (#80764). + +Everything here is pure or read-only so it can be tested without a gateway. +Config lives under ``agent.agent_cache`` in ``config.yaml``. +""" + +from __future__ import annotations + +import os +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Iterable, List, Optional, Tuple + +# Fraction of the resolved memory limit at which we start shedding +# transcripts. Deliberately well under the limit: on the reported incident +# the gateway hit cgroup ``memory.high`` throttling with swap full, and a +# SIGTERM flush from there could not finish inside systemd's stop timeout. +# Eviction has to happen while the process still has room to breathe. +_AUTO_BUDGET_FRACTION = 0.65 +# Below this a "budget" is noise — small containers would evict on every pass +# and never keep a warm prefix. +_AUTO_BUDGET_FLOOR_MB = 512 + +_DEFAULT_MAX_EVICTIONS_PER_PASS = 16 +# Never let a pressure pass touch the hottest sessions: they are the ones +# whose prompt cache is worth the most, and shedding them just moves the cost +# to the next turn instead of removing it. +_DEFAULT_PROTECT_RECENT = 8 + +_BYTES_PER_MB = 1024 * 1024 + + +@dataclass(frozen=True) +class AgentCacheBounds: + """Operator-facing bounds for the per-session agent cache. + + ``max_size`` and ``idle_ttl_secs`` are ``None`` when the operator did not + set them, so ``gateway/run.py`` keeps using its module-level defaults. + ``memory_high_mb`` is ``None`` when pressure eviction is switched off. + """ + + max_size: Optional[int] = None + idle_ttl_secs: Optional[float] = None + memory_high_mb: Optional[int] = None + max_evictions_per_pass: int = _DEFAULT_MAX_EVICTIONS_PER_PASS + protect_recent: int = _DEFAULT_PROTECT_RECENT + + +def _positive_int(value: Any) -> Optional[int]: + if isinstance(value, bool) or value is None: + return None + try: + parsed = int(value) + except (TypeError, ValueError): + return None + return parsed if parsed > 0 else None + + +def _positive_float(value: Any) -> Optional[float]: + if isinstance(value, bool) or value is None: + return None + try: + parsed = float(value) + except (TypeError, ValueError): + return None + return parsed if parsed > 0 else None + + +def _cgroup_limit_bytes() -> Optional[int]: + """Return the memory limit this process runs under, if it is cgroup-capped. + + Prefers cgroup v2 ``memory.high`` (the throttling point — passing it is + what stalled the reported shutdown) over ``memory.max``, and falls back to + cgroup v1. ``max`` / absurd sentinel values mean "unlimited". + """ + if sys.platform != "linux": + return None + candidates = ( + "/sys/fs/cgroup/memory.high", + "/sys/fs/cgroup/memory.max", + "/sys/fs/cgroup/memory/memory.limit_in_bytes", + ) + for candidate in candidates: + try: + raw = Path(candidate).read_text(encoding="utf-8").strip() + except OSError: + continue + if not raw or raw == "max": + continue + try: + limit = int(raw) + except ValueError: + continue + # cgroup v1 reports "unlimited" as a near-2^63 sentinel. + if limit <= 0 or limit >= (1 << 62): + continue + return limit + return None + + +def _total_memory_bytes() -> Optional[int]: + try: + return int(os.sysconf("SC_PAGE_SIZE")) * int(os.sysconf("SC_PHYS_PAGES")) + except (OSError, ValueError, AttributeError): + pass + try: + import psutil # type: ignore + + return int(psutil.virtual_memory().total) + except Exception: + return None + + +def resolve_memory_high_mb(setting: Any) -> Optional[int]: + """Resolve the ``memory_high_mb`` setting into an absolute MB budget. + + ``"auto"`` derives a budget from the cgroup limit the gateway runs under + (or total RAM when uncapped), which is what makes this fix work out of the + box on the containerised/systemd deployments where the leak bites. A + positive number is taken literally; anything falsy disables the pass. + """ + if isinstance(setting, str): + normalized = setting.strip().lower() + if normalized != "auto": + return ( + None + if normalized in ("", "off", "none", "false", "disabled") + else _positive_int(normalized) + ) + elif isinstance(setting, bool): + if not setting: + return None + else: + return _positive_int(setting) + + limit = _cgroup_limit_bytes() or _total_memory_bytes() + if not limit: + return None + budget = int(limit * _AUTO_BUDGET_FRACTION / _BYTES_PER_MB) + return budget if budget >= _AUTO_BUDGET_FLOOR_MB else None + + +def resolve_agent_cache_bounds(config: Any) -> AgentCacheBounds: + """Read ``agent.agent_cache`` out of a raw config mapping. + + Reads the *raw* user config (the gateway's loader does not deep-merge + ``DEFAULT_CONFIG``), so an absent key stays absent and the caller can tell + "operator chose 128" from "operator said nothing". + """ + section: Any = None + if isinstance(config, dict): + agent_cfg = config.get("agent") + if isinstance(agent_cfg, dict): + section = agent_cfg.get("agent_cache") + if not isinstance(section, dict): + section = {} + + max_evictions = _positive_int(section.get("max_evictions_per_pass")) + protect_recent = section.get("protect_recent") + protect_parsed = _positive_int(protect_recent) + if protect_parsed is None and protect_recent == 0: + protect_parsed = 0 + + return AgentCacheBounds( + max_size=_positive_int(section.get("max_size")), + idle_ttl_secs=_positive_float(section.get("idle_ttl_secs")), + memory_high_mb=resolve_memory_high_mb(section.get("memory_high_mb", "auto")), + max_evictions_per_pass=( + max_evictions if max_evictions is not None else _DEFAULT_MAX_EVICTIONS_PER_PASS + ), + protect_recent=( + protect_parsed if protect_parsed is not None else _DEFAULT_PROTECT_RECENT + ), + ) + + +def read_anon_rss_mb() -> Optional[int]: + """Return the process's anonymous resident memory in MB, or None. + + Anonymous pages are the ones cached transcripts live in — the reported + incident measured 11.0 GB of anon out of 11.0 GB total, so file-backed + pages are noise here. ``collect_memory_snapshot`` already reads + ``/proc/self/status`` without a dependency; psutil covers everything else, + where only total RSS is available. + """ + try: + from hermes_cli.mem_trim import collect_memory_snapshot + + snapshot = collect_memory_snapshot() + anon_kib = snapshot.get("rss_anon_kib") + if isinstance(anon_kib, int) and anon_kib > 0: + return anon_kib // 1024 + rss_kib = snapshot.get("rss_kib") + if isinstance(rss_kib, int) and rss_kib > 0: + return rss_kib // 1024 + except Exception: + pass + + try: + import psutil # type: ignore + + return int(psutil.Process(os.getpid()).memory_info().rss / _BYTES_PER_MB) + except Exception: + return None + + +def transcript_persistence_caught_up(agent: Any) -> bool: + """True when the agent's live transcript is fully on disk. + + Soft eviction drops ``_session_messages`` and rebuilds it from the + persisted session next turn, so it is only safe once persistence has + caught up. ``_last_flushed_db_idx`` is advanced to ``len(messages)`` by + ``AIAgent._flush_messages_to_session_db`` and only on a fully successful + write — the same divergence the FTS write-corruption guard reacts to when + it preserves live history over a lagging transcript. Unknown shapes are + treated as *not* caught up: a skipped eviction costs memory, a wrong one + costs the user their conversation. + """ + messages = getattr(agent, "_session_messages", None) + if not isinstance(messages, list): + return False + flushed = getattr(agent, "_last_flushed_db_idx", None) + if not isinstance(flushed, int) or isinstance(flushed, bool): + return False + return flushed >= len(messages) + + +def plan_pressure_evictions( + ordered_entries: Iterable[Tuple[str, Any]], + *, + is_evictable: Callable[[str, Any], bool], + max_evictions: int, + protect_recent: int = 0, +) -> List[Tuple[str, Any]]: + """Choose which cached sessions to shed, least-recently-used first. + + ``ordered_entries`` must be in LRU→MRU order (the cache is an + ``OrderedDict`` kept in that order by ``move_to_end`` on every hit). The + batch is capped so one pass cannot stall the gateway tearing down clients. + + ``protect_recent`` is an upper bound, clamped to half the cache: a handful + of sessions can be big enough to exhaust the budget on their own (a single + tool-heavy transcript runs to hundreds of MB), and a fixed guard would + then protect the entire cache and leave the gateway climbing toward the + OOM killer with nothing it is willing to shed. + """ + entries = list(ordered_entries) + if max_evictions <= 0 or not entries: + return [] + protect = min(max(protect_recent, 0), len(entries) // 2) + if protect: + entries = entries[:-protect] + + plan: List[Tuple[str, Any]] = [] + for key, agent in entries: + if len(plan) >= max_evictions: + break + if is_evictable(key, agent): + plan.append((key, agent)) + return plan diff --git a/gateway/run.py b/gateway/run.py index 125185aeaa816..4cc5773e80063 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -71,6 +71,12 @@ from hermes_cli.fallback_config import get_fallback_chain # long-lived gateways (each AIAgent holds LLM clients, tool schemas, # memory providers, etc.). LRU order + idle TTL eviction are enforced # from _enforce_agent_cache_cap() and _session_expiry_watcher() below. +# +# These are the defaults; `agent.agent_cache.max_size` / +# `agent.agent_cache.idle_ttl_secs` in config.yaml override them per +# deployment. Neither bound knows how many BYTES a cached agent holds, so +# _sweep_agent_cache_under_pressure() adds the missing memory-pressure valve +# (see gateway/agent_cache_pressure.py). _AGENT_CACHE_MAX_SIZE = 128 _AGENT_CACHE_IDLE_TTL_SECS = 3600.0 # evict agents idle for >1h _PLATFORM_CONNECT_TIMEOUT_SECS_DEFAULT = 30.0 @@ -12180,6 +12186,17 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew except Exception as _e: logger.debug("Idle agent sweep failed: %s", _e) + # Neither the LRU cap nor the idle TTL is aware of how much + # memory a cached transcript costs, so a busy gateway keeps + # every warm session's tool output resident until RSS hits the + # cgroup limit (#80764). Shed LRU transcripts once the heap is + # over budget; they reload from the persisted session on the + # next turn. + try: + self._sweep_agent_cache_under_pressure() + except Exception as _e: + logger.debug("Agent cache pressure sweep failed: %s", _e) + # Periodically prune stale SessionStore entries. The # in-memory dict (and sessions.json) would otherwise grow # unbounded in gateways serving many rotating chats / @@ -23731,8 +23748,156 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew if hasattr(agent, "_session_messages"): agent._session_messages = [] + def _agent_cache_bounds(self): + """Operator-configured agent-cache bounds, resolved once per process. + + Resolved lazily rather than in ``__init__`` so it also works for the + ``__new__``-constructed runners used by tests and by the slash-command + mixin. + """ + bounds = getattr(self, "_agent_cache_bounds_cache", None) + if bounds is None: + from gateway.agent_cache_pressure import resolve_agent_cache_bounds + + try: + bounds = resolve_agent_cache_bounds(_load_gateway_config()) + except Exception as _e: + from gateway.agent_cache_pressure import AgentCacheBounds + + logger.debug("Agent cache bounds config read failed: %s", _e) + bounds = AgentCacheBounds() + self._agent_cache_bounds_cache = bounds + return bounds + + def _agent_cache_cap(self) -> int: + """Effective LRU cap — the configured override, else the default.""" + configured = self._agent_cache_bounds().max_size + return configured if configured else _AGENT_CACHE_MAX_SIZE + + def _agent_cache_idle_ttl(self) -> float: + """Effective idle TTL in seconds — configured override, else default.""" + configured = self._agent_cache_bounds().idle_ttl_secs + return configured if configured else _AGENT_CACHE_IDLE_TTL_SECS + + def _sweep_agent_cache_under_pressure(self) -> int: + """Shed cached transcripts once the gateway's own heap nears its budget. + + The LRU cap counts entries and the idle sweep counts seconds; neither + knows that one cached agent pins a full ``_session_messages`` + transcript — tens of MB on a session with 100+ tool calls. A gateway + serving many chats therefore holds every warm transcript indefinitely: + agents that took a turn within the TTL are never idle-swept, and the + sweep additionally defers finalizable sessions until they expire. RSS + climbs until the cgroup throttles and SIGTERM can no longer flush + inside systemd's stop timeout (#80764). + + This is the missing valve. Above the configured anonymous-RSS budget + it evicts LRU agents through the same soft path the cap enforcer uses, + so the transcript is dropped and rebuilt from the persisted session on + the next turn. Three things are never touched: agents mid-turn (their + clients and sandboxes are in use), the most recently used sessions + (whose prompt cache is worth the most), and any session whose live + transcript has not finished reaching disk. + + Returns the number of entries evicted (0 when memory is fine). + """ + from gateway.agent_cache_pressure import ( + plan_pressure_evictions, + read_anon_rss_mb, + transcript_persistence_caught_up, + ) + + bounds = self._agent_cache_bounds() + if not bounds.memory_high_mb: + return 0 + _cache = getattr(self, "_agent_cache", None) + _lock = getattr(self, "_agent_cache_lock", None) + if not _cache or _lock is None: + # Nothing cached — whatever is using the heap, it isn't us, and + # warning about it every tick would point at the wrong subsystem. + return 0 + + rss_mb = read_anon_rss_mb() + if rss_mb is None or rss_mb < bounds.memory_high_mb: + return 0 + + running_ids = { + id(a) + for _, a in self._running_agent_items() + if a is not None and a is not _AGENT_PENDING_SENTINEL + } + + def _is_evictable(key: str, agent: Any) -> bool: + if agent is None or agent is _AGENT_PENDING_SENTINEL: + return False + if id(agent) in running_ids: + return False + return transcript_persistence_caught_up(agent) + + with _lock: + ordered = [ + (key, entry[0] if isinstance(entry, tuple) and entry else entry) + for key, entry in _cache.items() + ] + plan = plan_pressure_evictions( + ordered, + is_evictable=_is_evictable, + max_evictions=bounds.max_evictions_per_pass, + protect_recent=bounds.protect_recent, + ) + for key, _ in plan: + _cache.pop(key, None) + + if not plan: + _mid_turn = sum(1 for _, a in ordered if a is not None and id(a) in running_ids) + logger.warning( + "Agent cache pressure: anon RSS %dMB over budget %dMB but no " + "evictable session (%d cached, %d mid-turn) — memory will keep " + "climbing until those turns finish.", + rss_mb, bounds.memory_high_mb, len(ordered), _mid_turn, + ) + return 0 + + logger.warning( + "Agent cache pressure: anon RSS %dMB over budget %dMB — evicting " + "%d LRU session(s): %s", + rss_mb, bounds.memory_high_mb, len(plan), + ", ".join(key for key, _ in plan), + ) + try: + threading.Thread( + target=self._release_pressure_batch, + args=(plan,), + daemon=True, + name="agent-cache-pressure", + ).start() + except Exception: + self._release_pressure_batch(plan) + return len(plan) + + def _release_pressure_batch(self, plan: List[tuple]) -> None: + """Release a pressure-evicted batch, then return the heap to the OS. + + Sequential on one daemon thread rather than a thread per agent: the + batch is already capped, and the point of the pass is to reclaim + memory, not to race N teardowns. The trailing ``malloc_trim`` is what + turns "Python dropped the transcript" into "RSS actually fell" — + without it glibc keeps the freed arenas and the cgroup never notices. + """ + for key, agent in plan: + try: + self._commit_then_release_soft(agent, key) + except Exception as _e: + logger.debug("Pressure release failed for %s: %s", key, _e) + try: + from hermes_cli.mem_trim import trim_memory + + trim_memory(force=True, reason="agent_cache_pressure") + except Exception: + pass + def _enforce_agent_cache_cap(self) -> None: - """Evict oldest cached agents when cache exceeds _AGENT_CACHE_MAX_SIZE. + """Evict oldest cached agents when cache exceeds the LRU cap. Must be called with _agent_cache_lock held. Resource cleanup (memory provider shutdown, tool resource close) is scheduled @@ -23772,7 +23937,8 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew # already-cached long-running one. The cache may therefore stay # temporarily over cap; it will re-check on the next insert, # after active turns have finished. - excess = max(0, len(_cache) - _AGENT_CACHE_MAX_SIZE) + cap = self._agent_cache_cap() + excess = max(0, len(_cache) - cap) evict_plan: List[tuple] = [] # [(key, agent), ...] if excess > 0: ordered_keys = list(_cache.keys()) @@ -23786,12 +23952,12 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew for key, _ in evict_plan: _cache.pop(key, None) - remaining_over_cap = len(_cache) - _AGENT_CACHE_MAX_SIZE + remaining_over_cap = len(_cache) - cap if remaining_over_cap > 0: logger.warning( "Agent cache over cap (%d > %d); %d excess slot(s) held by " "mid-turn agents — will re-check on next insert.", - len(_cache), _AGENT_CACHE_MAX_SIZE, remaining_over_cap, + len(_cache), cap, remaining_over_cap, ) for key, agent in evict_plan: @@ -23814,7 +23980,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew ).start() def _sweep_idle_cached_agents(self) -> int: - """Evict cached agents whose AIAgent has been idle > _AGENT_CACHE_IDLE_TTL_SECS. + """Evict cached agents whose AIAgent has been idle past the idle TTL. Safe to call from the session expiry watcher without holding the cache lock — acquires it internally. Returns the number of entries @@ -23829,6 +23995,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew if _cache is None or _lock is None: return 0 now = time.time() + idle_ttl = self._agent_cache_idle_ttl() to_evict: List[tuple] = [] running_ids = { id(a) @@ -23845,7 +24012,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew last_activity = getattr(agent, "_last_activity_ts", None) if last_activity is None: continue - if (now - last_activity) > _AGENT_CACHE_IDLE_TTL_SECS: + if (now - last_activity) > idle_ttl: # Check whether the session has actually expired in the # session store. If it hasn't (e.g. daily-reset mode # where the reset fires hours after the user's last diff --git a/hermes_cli/config_defaults.py b/hermes_cli/config_defaults.py index 0b14c8db0dc8a..ca665b0dfac97 100644 --- a/hermes_cli/config_defaults.py +++ b/hermes_cli/config_defaults.py @@ -40,6 +40,29 @@ DEFAULT_CONFIG = { # rejected with a resend notice rather than run without serialization. # Non-positive values fall back to 1800 seconds. "gateway_turn_lease_timeout": 1800, + # Per-session AIAgent cache in the gateway. Each cached agent keeps a + # warm prompt prefix AND the session's full transcript, so the cache + # trades memory for cost: too small and every turn re-pays an uncached + # prompt, too large and tool-heavy transcripts fill the heap. + "agent_cache": { + # LRU entry cap. + "max_size": 128, + # Evict an agent that has been idle this long (seconds). + "idle_ttl_secs": 3600, + # Anonymous-RSS budget (MB) above which the gateway starts shedding + # least-recently-used transcripts, which reload from the persisted + # session on the next turn. "auto" derives the budget from the + # cgroup memory limit the gateway runs under (or total RAM when + # uncapped); a number sets it explicitly; 0/off disables the pass + # and lets memory grow to whatever the two bounds above allow. + "memory_high_mb": "auto", + # Upper bound on how many sessions one pressure pass sheds, so a + # burst of teardowns cannot stall the gateway. + "max_evictions_per_pass": 16, + # Most-recently-used sessions the pressure pass never touches — + # they are the ones actively paying for a warm prompt cache. + "protect_recent": 8, + }, # Force-interrupt budget once gateway stop()/drain has begun # (seconds). Applies to SIGTERM/external stop and to the final # phase of in-band restart after any after-turn wait. 0 = interrupt diff --git a/tests/gateway/test_agent_cache_pressure.py b/tests/gateway/test_agent_cache_pressure.py new file mode 100644 index 0000000000000..52571a9b0d844 --- /dev/null +++ b/tests/gateway/test_agent_cache_pressure.py @@ -0,0 +1,395 @@ +"""Memory-pressure eviction for the gateway agent cache (#80764). + +The LRU cap counts entries and the idle sweep counts seconds, so a gateway +serving many warm sessions holds every full transcript resident until the +cgroup kills it. These tests pin the pressure valve that sheds them, and the +three things it must never shed: a mid-turn agent, the most-recently-used +sessions, and a session whose transcript has not finished reaching disk. +""" + +import threading +from collections import OrderedDict +from unittest.mock import MagicMock + +import pytest + +from gateway.agent_cache_pressure import ( + AgentCacheBounds, + plan_pressure_evictions, + resolve_agent_cache_bounds, + resolve_memory_high_mb, + transcript_persistence_caught_up, +) + + +class TestBoundsResolution: + """Absent config must stay absent so gateway/run.py keeps its defaults.""" + + def test_absent_section_leaves_lru_bounds_unset(self): + bounds = resolve_agent_cache_bounds({}) + assert bounds.max_size is None + assert bounds.idle_ttl_secs is None + + def test_configured_values_are_honoured(self): + bounds = resolve_agent_cache_bounds( + { + "agent": { + "agent_cache": { + "max_size": 32, + "idle_ttl_secs": 600, + "memory_high_mb": 2048, + "max_evictions_per_pass": 4, + "protect_recent": 2, + } + } + } + ) + assert bounds.max_size == 32 + assert bounds.idle_ttl_secs == 600.0 + assert bounds.memory_high_mb == 2048 + assert bounds.max_evictions_per_pass == 4 + assert bounds.protect_recent == 2 + + def test_garbage_values_fall_back_to_defaults(self): + """A typo in config.yaml must not disable the cache or crash startup.""" + bounds = resolve_agent_cache_bounds( + {"agent": {"agent_cache": {"max_size": "lots", "idle_ttl_secs": -5}}} + ) + assert bounds.max_size is None + assert bounds.idle_ttl_secs is None + assert bounds.max_evictions_per_pass > 0 + + def test_protect_recent_zero_is_respected(self): + """0 means "shed anything", which is distinct from "unset".""" + bounds = resolve_agent_cache_bounds( + {"agent": {"agent_cache": {"protect_recent": 0}}} + ) + assert bounds.protect_recent == 0 + + +class TestMemoryBudgetResolution: + @pytest.mark.parametrize("setting", [0, False, None, "off", "none", ""]) + def test_falsy_settings_disable_the_pass(self, setting): + assert resolve_memory_high_mb(setting) is None + + @pytest.mark.parametrize("setting", [4096, "4096", 4096.0]) + def test_explicit_budget_is_taken_literally(self, setting): + assert resolve_memory_high_mb(setting) == 4096 + + def test_auto_derives_a_budget_below_the_cgroup_limit(self, monkeypatch): + """The budget must leave headroom: hitting memory.high is what makes + the shutdown flush time out in the first place.""" + import gateway.agent_cache_pressure as acp + + limit_mb = 10 * 1024 + monkeypatch.setattr(acp, "_cgroup_limit_bytes", lambda: limit_mb * 1024 * 1024) + + budget = resolve_memory_high_mb("auto") + + assert budget is not None + assert 0 < budget < limit_mb + + def test_auto_is_disabled_when_no_limit_is_discoverable(self, monkeypatch): + import gateway.agent_cache_pressure as acp + + monkeypatch.setattr(acp, "_cgroup_limit_bytes", lambda: None) + monkeypatch.setattr(acp, "_total_memory_bytes", lambda: None) + + assert resolve_memory_high_mb("auto") is None + + +class TestPersistenceGuard: + """Soft eviction drops the transcript, so it may only run once the + transcript is durable. Exercised against the real AIAgent flush.""" + + def _agent(self, tmp_path, session_id): + from hermes_state import SessionDB + from run_agent import AIAgent + + db = SessionDB(db_path=tmp_path / "sessions.db") + agent = AIAgent( + model="anthropic/claude-sonnet-4", + api_key="test", + base_url="https://openrouter.ai/api/v1", + provider="openrouter", + max_iterations=5, + quiet_mode=True, + skip_context_files=True, + skip_memory=True, + session_id=session_id, + session_db=db, + ) + db.create_session(session_id, source="telegram") + agent._session_db_created = True + return agent + + def test_fresh_agent_holds_nothing_to_lose(self, tmp_path): + agent = self._agent(tmp_path, "fresh") + try: + assert transcript_persistence_caught_up(agent) is True + finally: + agent.close() + + def test_unflushed_turn_blocks_eviction_then_flush_unblocks_it(self, tmp_path): + agent = self._agent(tmp_path, "lagging") + try: + messages = [ + {"role": "user", "content": "read the logs"}, + {"role": "assistant", "content": "done"}, + ] + agent._session_messages = messages + + assert transcript_persistence_caught_up(agent) is False, ( + "a transcript that never reached disk must not be dropped — " + "the session would come back with amnesia" + ) + + assert agent._flush_messages_to_session_db(messages) is True + assert transcript_persistence_caught_up(agent) is True + finally: + agent.close() + + def test_unknown_shapes_are_treated_as_unsafe(self): + assert transcript_persistence_caught_up(object()) is False + assert transcript_persistence_caught_up(None) is False + + +class TestEvictionPlanner: + def _entries(self, n): + return [(f"s{i}", MagicMock()) for i in range(n)] + + def test_evicts_least_recently_used_first(self): + entries = self._entries(6) + plan = plan_pressure_evictions( + entries, is_evictable=lambda k, a: True, max_evictions=2, protect_recent=0 + ) + assert [key for key, _ in plan] == ["s0", "s1"] + + def test_never_touches_the_protected_tail(self): + entries = self._entries(10) + plan = plan_pressure_evictions( + entries, is_evictable=lambda k, a: True, max_evictions=10, protect_recent=3 + ) + assert [key for key, _ in plan] == ["s0", "s1", "s2", "s3", "s4", "s5", "s6"] + + @pytest.mark.parametrize("size", [1, 2, 3, 5]) + def test_a_small_cache_of_large_transcripts_is_still_shedable(self, size): + """A fixed MRU guard would protect the whole cache when a couple of + sessions are big enough to blow the budget on their own — the gateway + would then climb toward the OOM killer with nothing it would shed.""" + plan = plan_pressure_evictions( + self._entries(size), + is_evictable=lambda k, a: True, + max_evictions=10, + protect_recent=8, + ) + assert plan, f"nothing evictable with {size} cached session(s)" + assert len(plan) <= size + + def test_protection_still_keeps_the_hottest_session(self): + plan = plan_pressure_evictions( + self._entries(4), + is_evictable=lambda k, a: True, + max_evictions=10, + protect_recent=8, + ) + assert "s3" not in [key for key, _ in plan] + + def test_skipped_candidates_do_not_consume_the_batch(self): + """Skipping a protected session must not shrink the batch — otherwise + one wedged session throttles the whole pass.""" + entries = self._entries(6) + plan = plan_pressure_evictions( + entries, + is_evictable=lambda k, a: k != "s0", + max_evictions=2, + protect_recent=0, + ) + assert [key for key, _ in plan] == ["s1", "s2"] + + +class TestGatewayPressureSweep: + """End-to-end against the real GatewayRunner method.""" + + def _runner(self, bounds=None): + from gateway.run import GatewayRunner + + runner = GatewayRunner.__new__(GatewayRunner) + runner._agent_cache = OrderedDict() + runner._agent_cache_lock = threading.Lock() + runner._running_agents = {} + runner._agent_cache_bounds_cache = bounds or AgentCacheBounds( + memory_high_mb=1000, max_evictions_per_pass=8, protect_recent=1 + ) + return runner + + def _cached_agent(self, *, persisted=True, messages=2): + agent = MagicMock() + agent._session_messages = [{"role": "user", "content": "x"}] * messages + agent._last_flushed_db_idx = messages if persisted else 0 + return agent + + def _at_rss(self, monkeypatch, mb): + import gateway.agent_cache_pressure as acp + + monkeypatch.setattr(acp, "read_anon_rss_mb", lambda: mb) + + def test_no_eviction_below_budget(self, monkeypatch): + runner = self._runner() + self._at_rss(monkeypatch, 400) + for i in range(5): + runner._agent_cache[f"s{i}"] = (self._cached_agent(), "sig") + + assert runner._sweep_agent_cache_under_pressure() == 0 + assert len(runner._agent_cache) == 5 + + def test_over_budget_sheds_lru_and_frees_the_transcript(self, monkeypatch): + runner = self._runner() + self._at_rss(monkeypatch, 4000) + released: list = [] + runner._commit_then_release_soft = lambda agent, key: ( + released.append(key), + setattr(agent, "_session_messages", []), + ) + + for i in range(4): + runner._agent_cache[f"s{i}"] = (self._cached_agent(), "sig") + oldest = runner._agent_cache["s0"][0] + + evicted = runner._sweep_agent_cache_under_pressure() + + assert evicted == 3 # protect_recent=1 keeps the newest + assert "s0" not in runner._agent_cache + assert "s3" in runner._agent_cache + _wait_for(lambda: released == ["s0", "s1", "s2"]) + assert oldest._session_messages == [] + + def test_mid_turn_session_is_never_evicted(self, monkeypatch): + runner = self._runner() + self._at_rss(monkeypatch, 4000) + runner._commit_then_release_soft = lambda agent, key: None + + active = self._cached_agent() + runner._agent_cache["s-active"] = (active, "sig") + runner._agent_cache["s-idle"] = (self._cached_agent(), "sig") + runner._agent_cache["s-new"] = (self._cached_agent(), "sig") + runner._running_agents["s-active"] = active + + runner._sweep_agent_cache_under_pressure() + + assert "s-active" in runner._agent_cache, ( + "evicting a mid-turn agent tears down the clients and sandbox the " + "running request is using" + ) + assert "s-idle" not in runner._agent_cache + + def test_lagging_persistence_blocks_eviction(self, monkeypatch): + runner = self._runner() + self._at_rss(monkeypatch, 4000) + runner._commit_then_release_soft = lambda agent, key: None + + runner._agent_cache["s-lagging"] = ( + self._cached_agent(persisted=False), "sig", + ) + runner._agent_cache["s-durable"] = (self._cached_agent(), "sig") + runner._agent_cache["s-new"] = (self._cached_agent(), "sig") + + runner._sweep_agent_cache_under_pressure() + + assert "s-lagging" in runner._agent_cache, ( + "dropping a transcript that never reached disk loses the " + "conversation the FTS guard exists to protect" + ) + assert "s-durable" not in runner._agent_cache + + def test_empty_cache_is_a_no_op(self, monkeypatch): + """Heap pressure with nothing cached is somebody else's problem.""" + runner = self._runner() + self._at_rss(monkeypatch, 999_999) + + assert runner._sweep_agent_cache_under_pressure() == 0 + + def test_all_candidates_skipped_reports_zero_without_raising(self, monkeypatch): + runner = self._runner() + self._at_rss(monkeypatch, 4000) + runner._commit_then_release_soft = lambda agent, key: None + for i in range(3): + runner._agent_cache[f"s{i}"] = ( + self._cached_agent(persisted=False), "sig", + ) + + assert runner._sweep_agent_cache_under_pressure() == 0 + assert len(runner._agent_cache) == 3 + + def test_disabled_budget_is_a_no_op(self, monkeypatch): + runner = self._runner(bounds=AgentCacheBounds(memory_high_mb=None)) + self._at_rss(monkeypatch, 999_999) + runner._agent_cache["s0"] = (self._cached_agent(), "sig") + + assert runner._sweep_agent_cache_under_pressure() == 0 + assert "s0" in runner._agent_cache + + +class TestConfiguredBoundsReachTheCache: + """The two existing bounds must be operator-tunable, and must keep their + built-in values when config.yaml says nothing.""" + + def _runner(self, bounds): + from gateway.run import GatewayRunner + + runner = GatewayRunner.__new__(GatewayRunner) + runner._agent_cache_bounds_cache = bounds + return runner + + def test_unset_config_keeps_the_built_in_defaults(self): + from gateway import run as gw_run + + runner = self._runner(AgentCacheBounds()) + assert runner._agent_cache_cap() == gw_run._AGENT_CACHE_MAX_SIZE + assert runner._agent_cache_idle_ttl() == gw_run._AGENT_CACHE_IDLE_TTL_SECS + + def test_configured_cap_bounds_the_real_enforcer(self): + """A configured cap must actually shrink the cache, not just report.""" + runner = self._runner(AgentCacheBounds(max_size=2)) + runner._agent_cache = OrderedDict() + runner._agent_cache_lock = threading.Lock() + runner._running_agents = {} + runner._release_evicted_agent_soft = lambda agent: None + runner._commit_then_release_soft = lambda agent, key: None + + with runner._agent_cache_lock: + for i in range(5): + runner._agent_cache[f"s{i}"] = (MagicMock(), "sig") + runner._enforce_agent_cache_cap() + + assert len(runner._agent_cache) == 2 + assert list(runner._agent_cache) == ["s3", "s4"] + + def test_configured_idle_ttl_drives_the_real_sweep(self): + import time as _t + + runner = self._runner(AgentCacheBounds(idle_ttl_secs=0.01)) + runner._agent_cache = OrderedDict() + runner._agent_cache_lock = threading.Lock() + runner._running_agents = {} + runner._release_evicted_agent_soft = lambda agent: None + runner.session_store = None + + stale = MagicMock() + stale._last_activity_ts = _t.time() - 5.0 + runner._agent_cache["s-stale"] = (stale, "sig") + + assert runner._sweep_idle_cached_agents() == 1 + assert "s-stale" not in runner._agent_cache + + +def _wait_for(predicate, timeout: float = 3.0) -> None: + """Wait for a background release thread to finish its work.""" + import time as _t + + deadline = _t.time() + timeout + while _t.time() < deadline: + if predicate(): + return + _t.sleep(0.02) + assert predicate(), "background release did not complete in time"