fix(gateway): skip cross-process guard on session_id switch under same session_key (#54947)

The cross-process coherence guard (#45966) compares the session's
on-disk message_count against the snapshot stored next to the cached
agent, and rebuilds the agent on a mismatch.  The guard is correct
when the cache snapshot and the live count both refer to the same
DB row.  But the agent cache is keyed by session_key, which can
group multiple conversation threads (different session_ids) under
the same key — and the message_count values belong to DIFFERENT
DB rows.

When the user switches from session A to session B under the same
session_key, the cache hit returns A's cached agent.  The guard then
compares A's snapshot count (A.message_count) against B's live count
(B.message_count) — they are NEVER equal because they track
different conversations — and invalidates the cache.  Every session
switch busts the prompt cache and forces a fresh agent build.  The
post-turn re-baseline (#46237) made it worse: it reads the live
count from the CURRENT session_entry.session_id, so each switch
overwrites the original snapshot with the new session's count,
causing the very next switch BACK to the original session to fire
the guard again.

This is the bug from #54947 (P0, sweeper:risk-session-state,
sweeper:risk-caching).

Fix:
  * Record the snapshot's session_id alongside the message_count in
    the cache tuple: (agent, sig, mc, session_id) — a 4-tuple.  The
    cache build at the AIAgent construction site stores the active
    session_id.
  * The cache-hit guard skips the cross-process count comparison
    when the active session_id differs from the snapshot's
    session_id — the comparison is meaningless across different DB
    rows, so the agent is REUSED without invalidation.  The cross-
    process guard still fires when the session_id matches and the
    live count differs (genuine cross-process write on the SAME
    session).
  * _refresh_agent_cache_message_count checks the snapshot's
    session_id: when it differs from the current session_id, the
    snapshot is intentionally left untouched (overwriting it would
    corrupt the original conversation's baseline and cause the
    switch-back to fire the guard).  The legacy 3-tuple shape (no
    session_id) is still re-baselined as before.
  * Backward-compat:
      - 2-tuple (agent, sig) — unchanged, opts out of the guard.
      - 3-tuple (agent, sig, mc) — unchanged behavior, standard
        cross-process check.
      - pending sentinel — unchanged, untouched by re-baseline.
      - new 4-tuple (agent, sig, mc, session_id) — full session_id-
        aware guard with skip on mismatch.

Tests:
  * tests/gateway/test_session_id_cache_coherence.py — 7 tests
    covering L1-L5 from LAYERS.md:
      - L1 session_id switch must REUSE
      - L2 cache tuple records snapshot's session_id
      - L3 re-baseline skips when session_id differs
      - L4 same-session_id turns still re-baseline (#46237 holds)
      - L5 legacy 2-tuples and pending sentinels untouched
      - legacy 3-tuple (no session_id) still guarded (#45966 holds)
      - 3-tuple transitions to 3-tuple (not 4-tuple) on re-baseline

No regressions in 70 existing tests in test_agent_cache.py or 137
related session tests.  Co-authored with #52197 (deferred cleanup
of evicted agents); both fixes compose cleanly.
This commit is contained in:
Tranquil-Flow 2026-06-29 18:13:50 +02:00 committed by Teknium
parent aa4731598c
commit e7562c394f
2 changed files with 357 additions and 4 deletions

View file

@ -14917,6 +14917,15 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
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.
When the cache entry records a ``session_id`` (4-tuple form, #54947)
that differs from the current ``session_id`` meaning the cache
was built for a DIFFERENT conversation under the same ``session_key``
the snapshot is intentionally left untouched. Overwriting it with
the current session's count would corrupt the original conversation's
baseline and cause the next switch back to fire the cross-process
guard spuriously. Fail-safe: the legacy 3-tuple shape (no
``session_id``) is still re-baselined as before.
"""
if self._session_db is None or not session_id:
return
@ -14941,8 +14950,23 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
and len(cached) > 2
and cached[0] is not _AGENT_PENDING_SENTINEL
):
# If the snapshot was taken for a different session_id
# (same session_key, different conversation), leave the
# snapshot alone — the current session_id's count belongs
# to a different DB row (#54947).
_snapshot_sid = cached[3] if len(cached) > 3 else None
if _snapshot_sid is not None and _snapshot_sid != session_id:
return
if cached[2] != _live:
_cache[session_key] = (cached[0], cached[1], _live)
if _snapshot_sid is None:
# Legacy 3-tuple: preserve the original 3-element
# shape so existing entries stay compatible with
# callers that index ``cached[2]`` directly.
_cache[session_key] = (cached[0], cached[1], _live)
else:
_cache[session_key] = (
cached[0], cached[1], _live, _snapshot_sid,
)
def _evict_cached_agent(self, session_key: str) -> None:
"""Remove a cached agent for a session (called on /new, /model, etc).
@ -16649,9 +16673,25 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
if cached and cached[1] == _sig:
# cached[2] is the message_count at cache time;
# stale when a second process appended rows.
# cached[3] (when present) is the session_id the
# snapshot was taken for — used to skip the guard
# when the active session_id differs (#54947).
_cached_mc = cached[2] if len(cached) > 2 else None
_cached_sid = cached[3] if len(cached) > 3 else None
# If the snapshot belongs to a different session_id
# (same session_key, different conversation), the
# message_count comparison is meaningless — the
# counts track DIFFERENT DB rows. REUSE the cached
# agent rather than rebuild and bust the prompt cache
# on every session switch (#54947).
_session_id_mismatch = (
_cached_sid is not None
and session_id is not None
and _cached_sid != session_id
)
if (
_cached_mc is not None
not _session_id_mismatch
and _cached_mc is not None
and _current_msg_count is not None
and _current_msg_count != _cached_mc
):
@ -16682,7 +16722,6 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
_xproc_evicted_agent = _ev_agent
else:
agent = cached[0]
reused_cached_agent = True
# Refresh LRU order so the cap enforcement evicts
# truly-oldest entries, not the one we just used.
if hasattr(_cache, "move_to_end"):
@ -16695,6 +16734,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
# (cached agent may have been created with old config)
agent.max_iterations = max_iterations
logger.debug("Reusing cached agent for session %s", session_key)
reused_cached_agent = True
# Lock released — now schedule cleanup of any cross-process-evicted
# agent on a daemon thread so memory-provider shutdown / socket
@ -16752,7 +16792,14 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
)
if _cache_lock and _cache is not None:
with _cache_lock:
_cache[session_key] = (agent, _sig, _current_msg_count)
# Record the session_id the snapshot was taken for
# alongside the message_count, so the cross-process
# guard can skip the (meaningless) count comparison
# when the active session_id later switches under
# the same session_key (#54947).
_cache[session_key] = (
agent, _sig, _current_msg_count, session_id,
)
self._enforce_agent_cache_cap()
logger.debug("Created new agent for session %s (sig=%s)", session_key, _sig)