fix(gateway): re-baseline agent cache count after first-turn session_meta

The cross-process cache-coherence guard (#45966) compares a session's
on-disk message_count against a snapshot stored next to the cached agent,
rebuilding the agent on a mismatch so a foreign writer (e.g. the dashboard
backend) can't leave the in-memory transcript stale.

On a fresh gateway conversation the post-turn re-baseline
(_refresh_agent_cache_message_count) ran BEFORE the first-turn `session_meta`
marker row was appended to the transcript. That append goes through
append_to_transcript -> append_message, which increments message_count
unconditionally. So the snapshot was left exactly one short of the live
count, and on turn 2 of every fresh conversation the guard mistook this
process's own session_meta write for a foreign write, evicting and rebuilding
the cached agent — silently busting the per-conversation prompt cache the
cache exists to protect.

Move the re-baseline to after the turn's full transcript persistence block
(including the session_meta append and the compression session_id swap). The
snapshot now matches the live count, so the guard fires only on genuinely
foreign writes. This also makes the call honor its own documented contract of
using the compaction-updated session_id.

Adds a regression test that drives the real _handle_message_with_agent
against a real SessionDB and asserts the invariant: after a fresh first turn,
snapshot == live message_count, so the next turn's guard reuses the cached
agent. Fails before this change, passes after.
This commit is contained in:
Jason 2026-06-29 03:02:29 +08:00 committed by Teknium
parent 6bc0a7ce80
commit aa4731598c
2 changed files with 289 additions and 14 deletions

View file

@ -10679,19 +10679,12 @@ 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.
await self._refresh_agent_cache_message_count(
session_key, session_entry.session_id
)
# NOTE: the cross-process cache-coherence re-baseline
# (_refresh_agent_cache_message_count) is intentionally deferred
# until AFTER this turn's transcript persistence block below — it
# must include the first-turn `session_meta` marker row and the
# compression session_id swap, both of which happen later. See
# the call site after the `update_session(...)` write.
# Successful turn — clear any stuck-loop counter for this session.
# This ensures the counter only accumulates across CONSECUTIVE
@ -11077,6 +11070,28 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
last_prompt_tokens=agent_result.get("last_prompt_tokens", 0),
)
# Re-baseline the cached agent's message_count snapshot now that
# ALL of this turn's transcript writes are done — the agent's
# flushed user/assistant/tool rows AND the first-turn `session_meta`
# marker appended above. 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 the
# process's own turn grows message_count and the next turn sees a
# mismatch and rebuilds the agent — destroying prompt caching.
#
# This MUST run after the `session_meta` append: that row also
# increments message_count, so re-baselining before it (the old
# position) left the snapshot one short and the guard mis-fired on
# turn 2 of EVERY fresh gateway conversation, rebuilding the cached
# agent and busting the prompt cache. Running here also uses the
# compaction-updated session_id (the agent_result session_id swap
# above), matching this function's documented contract. Refreshing
# here makes the guard fire only on a DIFFERENT process's writes.
# Fail-safe inside the helper.
await self._refresh_agent_cache_message_count(
session_key, session_entry.session_id
)
# Intentional silence is a delivery decision, not a transcript
# mutation. The agent's [SILENT]/NO_REPLY assistant turn above is
# still persisted in session history so later turns keep normal
@ -18220,7 +18235,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
# follow-up. Use the same (session_key, session_id) the
# recursive call runs under so the snapshot matches exactly
# what the follow-up's guard will consult. Fail-safe in helper.
self._refresh_agent_cache_message_count(session_key, session_id)
await self._refresh_agent_cache_message_count(session_key, session_id)
followup_result = await self._run_agent(
message=next_message,

View file

@ -0,0 +1,260 @@
"""Regression: first-turn ``session_meta`` row must be re-baselined into the
agent cache's message_count snapshot.
Bug
---
On a fresh gateway conversation the post-turn re-baseline
(``_refresh_agent_cache_message_count``) runs *before* the first-turn
``session_meta`` marker row is appended to the transcript:
gateway/run.py:
... agent run completes ...
self._refresh_agent_cache_message_count(...) # snapshot taken HERE
...
if not history: # session_meta written LATER
append_to_transcript({"role": "session_meta", ...})
``append_to_transcript`` (no ``skip_db``) increments the session's
``message_count`` unconditionally (hermes_state.append_message), so the
snapshot ends up exactly +1 below the live on-disk count.
The cross-process coherence guard (#45966) compares the live count against
that snapshot on the *next* inbound message, sees ``live != snapshot``,
mistakes this process's own ``session_meta`` write for a foreign write, and
**rebuilds the cached agent on turn 2 of every fresh conversation** silently
busting the per-conversation prompt cache the cache exists to protect.
The fix re-baselines AFTER all of this turn's transcript writes (including the
first-turn ``session_meta`` row), so the snapshot matches the live count and
the guard fires only on genuinely foreign writes.
This drives the REAL ``_handle_message_with_agent`` against a REAL SessionDB
(the ``session_meta`` write actually increments ``message_count``) and asserts
the invariant: after a first turn, snapshot == live count next turn reuses.
"""
import sys
import threading
import types
from datetime import datetime
from unittest.mock import AsyncMock, MagicMock
import pytest
import gateway.run as gateway_run
from gateway.config import GatewayConfig, Platform
from gateway.platforms.base import MessageEvent
from gateway.session import SessionEntry, SessionSource
SESSION_KEY = "agent:main:telegram:group:-1001:12345"
SESSION_ID = "sess-first-turn"
def _bootstrap(monkeypatch, tmp_path, db):
"""GatewayRunner wired to a REAL SessionDB for count reads, mirroring the
proven #42039 harness but with a live cache + real transcript counter."""
fake_dotenv = types.ModuleType("dotenv")
fake_dotenv.load_dotenv = lambda *args, **kwargs: None
monkeypatch.setitem(sys.modules, "dotenv", fake_dotenv)
config = GatewayConfig()
runner = gateway_run.GatewayRunner(config)
runner.adapters = {}
runner._running_agents = {}
runner._running_agents_ts = {}
runner._pending_messages = {}
runner._pending_approvals = {}
runner._is_user_authorized = lambda _source: True
runner._set_session_env = lambda _context: None
runner._handle_active_session_busy_message = AsyncMock(return_value=False)
# REAL SessionDB so the guard's get_session(...).message_count is live.
runner._session_db = db
runner._recover_telegram_topic_thread_id = lambda _source: None
runner._cache_session_source = lambda _key, _source: None
runner._is_session_run_current = lambda _key, _gen: True
runner._begin_session_run_generation = lambda _key: 1
runner._reply_anchor_for_event = lambda _event: None
runner._get_guild_id = lambda _event: None
runner._should_send_voice_reply = lambda *_a, **_kw: False
runner.hooks = MagicMock()
runner.hooks.emit = AsyncMock()
# Live agent cache (not a MagicMock) so the re-baseline actually rewrites
# the snapshot tuple in place.
runner._agent_cache = {}
runner._agent_cache_lock = threading.Lock()
runner.session_store = MagicMock()
runner.session_store.get_or_create_session.return_value = SessionEntry(
session_key=SESSION_KEY,
session_id=SESSION_ID,
created_at=datetime.now(),
updated_at=datetime.now(),
platform=Platform.TELEGRAM,
chat_type="group",
)
# Empty history → triggers the first-turn ``session_meta`` write path.
runner.session_store.load_transcript.return_value = []
runner.session_store.has_platform_message_id.return_value = False
runner.session_store.update_session = MagicMock()
# Mirror the real SessionStore.append_to_transcript: forward non-skip_db
# writes to the real DB so a ``session_meta`` row genuinely increments
# message_count, exactly as in production. skip_db=True writes (the agent
# already persisted them via _flush_messages_to_session_db) are no-ops here.
def _append(session_id, message, skip_db=False):
if not skip_db:
db.append_message(
session_id=session_id,
role=message.get("role", "unknown"),
content=message.get("content"),
)
runner.session_store.append_to_transcript = MagicMock(side_effect=_append)
monkeypatch.setattr(gateway_run, "_hermes_home", tmp_path)
monkeypatch.setattr(
gateway_run, "_resolve_runtime_agent_kwargs", lambda: {"api_key": "fake"}
)
monkeypatch.setattr(
"agent.model_metadata.get_model_context_length",
lambda *_args, **_kwargs: 100_000,
)
return runner
def _event():
return MessageEvent(
text="hello world",
source=SessionSource(
platform=Platform.TELEGRAM,
chat_id="-1001",
chat_type="group",
user_id="12345",
),
message_id="msg-1",
)
def _source():
return SessionSource(
platform=Platform.TELEGRAM,
chat_id="-1001",
chat_type="group",
user_id="12345",
)
def _live_count(db, session_id):
row = db.get_session(session_id)
return (row.get("message_count", 0) if row else 0)
@pytest.mark.asyncio
async def test_first_turn_session_meta_is_captured_by_rebaseline(
monkeypatch, tmp_path
):
"""After a fresh first turn, the cache snapshot must equal the live
message_count including the first-turn ``session_meta`` row.
WITHOUT the fix the re-baseline snapshots the count *before* the
session_meta append, leaving the snapshot one short; the cross-process
guard then rebuilds the cached agent on turn 2 (prompt-cache churn).
"""
from hermes_state import SessionDB
db = SessionDB(db_path=tmp_path / "sessions.db")
db.create_session(SESSION_ID, source="telegram")
runner = _bootstrap(monkeypatch, tmp_path, db)
# Cache snapshot taken at agent-BUILD time = count before this turn's
# writes (a fresh session → 0). This is what the #45966 guard stores.
build_count = _live_count(db, SESSION_ID)
agent_obj = object()
with runner._agent_cache_lock:
runner._agent_cache[SESSION_KEY] = (agent_obj, "sig", build_count)
# Stubbed agent run: the gateway's own user/assistant rows are persisted
# by the agent (skip_db=True downstream); only the session_meta marker is
# written by the gateway with skip_db=False.
runner._run_agent = AsyncMock(
return_value={
"final_response": "Hi there!",
"messages": [
{"role": "user", "content": "hello world"},
{"role": "assistant", "content": "Hi there!"},
],
"tools": [{"name": "noop"}],
"history_offset": 0,
"last_prompt_tokens": 0,
}
)
await runner._handle_message_with_agent(_event(), _source(), SESSION_KEY, 1)
# The first-turn session_meta row was written → live count advanced.
live = _live_count(db, SESSION_ID)
assert live == build_count + 1, (
"first-turn session_meta should increment message_count by exactly 1"
)
# THE INVARIANT: the cache snapshot must now equal the live count, so the
# next turn's cross-process guard reuses the cached agent.
with runner._agent_cache_lock:
cached = runner._agent_cache[SESSION_KEY]
snapshot = cached[2]
assert snapshot == live, (
f"cache snapshot {snapshot} != live count {live}: the first-turn "
f"session_meta write was not re-baselined, so the #45966 guard will "
f"rebuild the cached agent on turn 2 and bust the prompt cache."
)
# And the cached agent instance must be untouched (never rebuilt).
assert cached[0] is agent_obj
@pytest.mark.asyncio
async def test_next_turn_guard_reuses_cached_agent_after_first_turn(
monkeypatch, tmp_path
):
"""End-to-end consequence: with the snapshot correctly re-baselined, the
production cross-process guard's reuse condition (live == snapshot) holds
on turn 2 no rebuild, prompt cache preserved."""
from hermes_state import SessionDB
db = SessionDB(db_path=tmp_path / "sessions.db")
db.create_session(SESSION_ID, source="telegram")
runner = _bootstrap(monkeypatch, tmp_path, db)
with runner._agent_cache_lock:
runner._agent_cache[SESSION_KEY] = (
object(), "sig", _live_count(db, SESSION_ID),
)
runner._run_agent = AsyncMock(
return_value={
"final_response": "Hi there!",
"messages": [
{"role": "user", "content": "hello world"},
{"role": "assistant", "content": "Hi there!"},
],
"tools": [{"name": "noop"}],
"history_offset": 0,
"last_prompt_tokens": 0,
}
)
await runner._handle_message_with_agent(_event(), _source(), SESSION_KEY, 1)
# Replicate the production cache-hit guard's reuse decision exactly:
# reuse iff live on-disk count == snapshot stored next to the agent.
live = _live_count(db, SESSION_ID)
with runner._agent_cache_lock:
snapshot = runner._agent_cache[SESSION_KEY][2]
would_reuse = (live == snapshot)
assert would_reuse, (
"turn-2 cross-process guard would rebuild the cached agent because "
"the first-turn session_meta write was not re-baselined into the "
"snapshot — this is the prompt-cache regression under test."
)