fix(gateway): ignore stale compression session splits

This commit is contained in:
Jake Present 2026-06-30 11:13:01 -04:00 committed by kshitij
parent d5b4879d4a
commit 00ec3b1884
2 changed files with 125 additions and 28 deletions

View file

@ -10891,13 +10891,17 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
}
await self.hooks.emit("agent:start", hook_ctx)
# Run the agent
# Run the agent. Capture the session id that this run was launched
# against so post-run compression publication can be identity-guarded
# below; a /new or another lifecycle transition may move
# session_entry.session_id while the old run is still unwinding.
_run_start_session_id = session_entry.session_id
agent_result = await self._run_agent(
message=message_text,
context_prompt=context_prompt,
history=history,
source=source,
session_id=session_entry.session_id,
session_id=_run_start_session_id,
session_key=session_key,
run_generation=run_generation,
event_message_id=self._reply_anchor_for_event(event),
@ -10999,17 +11003,27 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
# If the agent's session_id changed during compression, update
# session_entry so transcript writes below go to the right session.
if agent_result.get("session_id") and agent_result["session_id"] != session_entry.session_id:
session_entry.session_id = agent_result["session_id"]
self.session_store._save()
self.session_store._record_gateway_session_peer(
session_entry.session_id,
session_key,
source,
)
await asyncio.to_thread(
self._sync_telegram_topic_binding,
source, session_entry, reason="agent-result-compression",
)
if session_entry.session_id == _run_start_session_id:
session_entry.session_id = agent_result["session_id"]
self.session_store._save()
self.session_store._record_gateway_session_peer(
session_entry.session_id,
session_key,
source,
)
await asyncio.to_thread(
self._sync_telegram_topic_binding,
source, session_entry, reason="agent-result-compression",
)
else:
logger.info(
"Skipping agent-result session split sync for %s because "
"the session binding moved from %s to %s before "
"compression finished",
session_key or "?",
_run_start_session_id,
session_entry.session_id,
)
# Prepend reasoning/thinking if display is enabled (per-platform).
# Mattermost requires explicit per-platform opt-in because this is
@ -17714,14 +17728,36 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
session_id, agent_session_id,
)
entry = self.session_store._entries.get(session_key)
_session_split_entry_persisted = False
if entry:
entry.session_id = agent_session_id
self.session_store._save()
self.session_store._record_gateway_session_peer(
agent_session_id,
session_key,
source,
)
entry_session_id = getattr(entry, "session_id", None)
if not _run_still_current():
logger.info(
"Skipping session split sync for stale run %s"
"generation %s is no longer current",
session_key or "?",
run_generation,
)
elif entry_session_id == agent_session_id:
_session_split_entry_persisted = True
elif entry_session_id != session_id:
logger.info(
"Skipping session split sync for %s because the "
"session binding moved from %s to %s before "
"compression finished",
session_key or "?",
session_id,
entry_session_id,
)
else:
entry.session_id = agent_session_id
self.session_store._save()
self.session_store._record_gateway_session_peer(
agent_session_id,
session_key,
source,
)
_session_split_entry_persisted = True
# If this is a Telegram DM and source.thread_id was lost during
# the session split (synthetic / recovered event), restore it
@ -17729,8 +17765,10 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
# correct message_thread_id instead of routing to the General
# thread. Failure here is non-fatal — we log and continue;
# worst case the message lands in General, which is the
# pre-fix behaviour.
if (
# pre-fix behaviour. Only do this after this run successfully
# published its session split; a stale /stop→/new predecessor
# must not mutate routing/binding state for the fresh session.
if _session_split_entry_persisted and (
getattr(source, "platform", None) == Platform.TELEGRAM
and getattr(source, "chat_type", None) == "dm"
and getattr(source, "thread_id", None) is None
@ -17754,7 +17792,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
"Failed to restore thread_id from binding after session split",
exc_info=True,
)
if entry:
if _session_split_entry_persisted:
self._sync_telegram_topic_binding(
source, entry, reason="agent-run-compression",
)

View file

@ -119,7 +119,7 @@ def _runner(session_store):
return runner
def test_failed_turn_still_syncs_compression_session_split(monkeypatch):
def _install_compression_failure_agent(monkeypatch):
fake_run_agent = types.ModuleType("run_agent")
fake_run_agent.AIAgent = _CompressionThenFailureAgent
monkeypatch.setitem(sys.modules, "run_agent", fake_run_agent)
@ -132,11 +132,9 @@ def test_failed_turn_still_syncs_compression_session_split(monkeypatch):
monkeypatch.setattr(tools_config, "_get_platform_tools", lambda *_args, **_kwargs: {"core"})
session_store = _SessionStore()
runner = _runner(session_store)
source = SessionSource(platform=Platform.TELEGRAM, chat_id="12345", chat_type="dm", user_id="user-1")
result = asyncio.run(
def _run_compression_failure_turn(runner, source, *, run_generation=None):
return asyncio.run(
asyncio.wait_for(
runner._run_agent(
message="continue",
@ -145,11 +143,22 @@ def test_failed_turn_still_syncs_compression_session_split(monkeypatch):
source=source,
session_id="session-before-compression",
session_key=SESSION_KEY,
run_generation=run_generation,
),
timeout=2,
)
)
def test_failed_turn_still_syncs_compression_session_split(monkeypatch):
_install_compression_failure_agent(monkeypatch)
session_store = _SessionStore()
runner = _runner(session_store)
source = SessionSource(platform=Platform.TELEGRAM, chat_id="12345", chat_type="dm", user_id="user-1")
result = _run_compression_failure_turn(runner, source)
assert result["failed"] is True
assert result["session_id"] == "session-after-compression"
assert result["history_offset"] == 0
@ -158,3 +167,53 @@ def test_failed_turn_still_syncs_compression_session_split(monkeypatch):
runner._sync_telegram_topic_binding.assert_called_once_with(
source, session_store.entry, reason="agent-run-compression"
)
def test_stale_run_does_not_overwrite_new_session_after_compression(monkeypatch):
"""A /stop + /new can invalidate a run while its compression is still unwinding.
The stale run may still return with a rotated agent.session_id, but it must
not publish that old compressed child back into the channel's active session
binding. The outer gateway stale-result check will discard the response too;
this regression covers the earlier side effect inside _run_agent().
"""
_install_compression_failure_agent(monkeypatch)
session_store = _SessionStore()
runner = _runner(session_store)
runner._session_run_generation[SESSION_KEY] = 2
source = SessionSource(platform=Platform.TELEGRAM, chat_id="12345", chat_type="dm", user_id="user-1")
result = _run_compression_failure_turn(runner, source, run_generation=1)
assert result["failed"] is True
assert result["session_id"] == "session-after-compression"
assert result["history_offset"] == 0
assert session_store.entry.session_id == "session-before-compression"
assert session_store.save_calls == 0
assert getattr(runner._sync_telegram_topic_binding, "call_count") == 0
def test_session_split_sync_skips_when_binding_already_moved(monkeypatch):
"""A live session binding is identity-guarded, not blindly overwritten.
This catches the exact race where an old run starts with session A, /new
moves the binding to fresh session B, and the old run finishes compression
into child C. C must not replace B.
"""
_install_compression_failure_agent(monkeypatch)
session_store = _SessionStore()
session_store.entry.session_id = "fresh-session-after-new"
runner = _runner(session_store)
runner._session_run_generation[SESSION_KEY] = 1
source = SessionSource(platform=Platform.TELEGRAM, chat_id="12345", chat_type="dm", user_id="user-1")
result = _run_compression_failure_turn(runner, source, run_generation=1)
assert result["failed"] is True
assert result["session_id"] == "session-after-compression"
assert result["history_offset"] == 0
assert session_store.entry.session_id == "fresh-session-after-new"
assert session_store.save_calls == 0
assert getattr(runner._sync_telegram_topic_binding, "call_count") == 0