fix(gateway): re-baseline agent-cache message_count after each turn
The #45966 cross-process coherence guard snapshots a session's on-disk message_count next to the cached agent and rebuilds the agent when the count changes. But the snapshot is taken at agent-BUILD time — before the turn writes its own user + assistant (+ tool) rows — and the cache entry is never rewritten on a reuse. So this process's OWN turn grows message_count, and the very next turn sees a mismatch and rebuilds the agent. That happens every turn, for every conversation, silently destroying the per-conversation prompt caching the cache exists to protect (AGENTS.md: prompt caching is sacred). Add _refresh_agent_cache_message_count(): after a turn completes and the agent has flushed its rows to the SessionDB, re-baseline the stored count to the now-current value. The guard then fires ONLY when a DIFFERENT process changes the transcript — preserving the #45966 fix while keeping the cache warm for normal single-process operation. Tests drive the real SessionDB + the real guard condition: 5 consecutive same-process turns now all REUSE the cached agent (0 before the fix); a cross-process append still invalidates; and the re-baseline is fail-safe (no DB, falsy session_id, raising probe, legacy 2-tuple, pending sentinel all no-op).
This commit is contained in:
@@ -8875,6 +8875,20 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
||||
_response_time, _api_calls, _resp_len,
|
||||
)
|
||||
|
||||
# Re-baseline the cached agent's message_count snapshot now that
|
||||
# this turn has completed and the agent has flushed its rows to
|
||||
# the SessionDB. The cross-process coherence guard (#45966)
|
||||
# snapshots the count at agent-BUILD time (before this turn's own
|
||||
# writes) and never refreshes it on reuse — so without this, this
|
||||
# process's own turn would grow the count and the next turn would
|
||||
# see a mismatch and rebuild the agent every turn, destroying
|
||||
# prompt caching. Refreshing here makes the guard fire only on a
|
||||
# DIFFERENT process's writes. Uses the (possibly compaction-
|
||||
# updated) live session_id. Fail-safe inside the helper.
|
||||
self._refresh_agent_cache_message_count(
|
||||
session_key, session_entry.session_id
|
||||
)
|
||||
|
||||
# Successful turn — clear any stuck-loop counter for this session.
|
||||
# This ensures the counter only accumulates across CONSECUTIVE
|
||||
# restarts where the session was active (never completed).
|
||||
@@ -12784,6 +12798,57 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
|
||||
if release_running_state:
|
||||
self._release_running_agent_state(session_key)
|
||||
|
||||
def _refresh_agent_cache_message_count(
|
||||
self, session_key: str, session_id: Optional[str]
|
||||
) -> None:
|
||||
"""Re-baseline a cached agent's stored message_count after THIS turn.
|
||||
|
||||
The cross-process coherence guard (#45966) compares the session's
|
||||
on-disk ``message_count`` against the count snapshotted next to the
|
||||
cached agent, and rebuilds the agent on a mismatch. But the snapshot
|
||||
is taken at agent-BUILD time — before this turn writes its own user +
|
||||
assistant (+ tool) rows — and the cache entry is never rewritten on a
|
||||
reuse. So without this re-baseline, THIS process's own turn would
|
||||
grow ``message_count`` and the very next turn would see a mismatch
|
||||
and rebuild the agent — every turn, for every conversation — silently
|
||||
destroying the per-conversation prompt caching the cache exists to
|
||||
protect.
|
||||
|
||||
Call this once a turn has completed and the agent has flushed its
|
||||
rows to the SessionDB. It snapshots the now-current count (which
|
||||
includes this process's own writes) so the guard only fires when a
|
||||
DIFFERENT process changes the transcript out from under us. The
|
||||
``_sig`` is left untouched; only the count element is refreshed, and
|
||||
only when the same agent is still cached (no rebuild/eviction raced
|
||||
in between). Fail-safe: any DB error leaves the snapshot as-is, which
|
||||
at worst costs one unnecessary rebuild on the next turn.
|
||||
"""
|
||||
if self._session_db is None or not session_id:
|
||||
return
|
||||
_cache_lock = getattr(self, "_agent_cache_lock", None)
|
||||
_cache = getattr(self, "_agent_cache", None)
|
||||
if not _cache_lock or _cache is None:
|
||||
return
|
||||
try:
|
||||
_sess_row = self._session_db.get_session(session_id)
|
||||
_live = _sess_row.get("message_count", 0) if _sess_row else None
|
||||
except Exception:
|
||||
return
|
||||
if _live is None:
|
||||
return
|
||||
with _cache_lock:
|
||||
cached = _cache.get(session_key)
|
||||
# Only re-baseline a live 3-tuple entry; skip pending sentinels,
|
||||
# legacy 2-tuples (they intentionally opt out of the guard), and
|
||||
# the case where the entry was evicted/rebuilt mid-turn.
|
||||
if (
|
||||
isinstance(cached, tuple)
|
||||
and len(cached) > 2
|
||||
and cached[0] is not _AGENT_PENDING_SENTINEL
|
||||
):
|
||||
if cached[2] != _live:
|
||||
_cache[session_key] = (cached[0], cached[1], _live)
|
||||
|
||||
def _evict_cached_agent(self, session_key: str) -> None:
|
||||
"""Remove a cached agent for a session (called on /new, /model, etc).
|
||||
|
||||
|
||||
Reference in New Issue
Block a user