From 00ec3b1884c5ee0fbd4e042e5ad7aef7141a0a17 Mon Sep 17 00:00:00 2001 From: Jake Present Date: Tue, 30 Jun 2026 11:13:01 -0400 Subject: [PATCH] fix(gateway): ignore stale compression session splits --- gateway/run.py | 84 ++++++++++++++----- .../test_compression_failure_session_sync.py | 69 +++++++++++++-- 2 files changed, 125 insertions(+), 28 deletions(-) diff --git a/gateway/run.py b/gateway/run.py index b6afc605ec0..ce6c8950f13 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -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", ) diff --git a/tests/gateway/test_compression_failure_session_sync.py b/tests/gateway/test_compression_failure_session_sync.py index 5e093eb0bf5..8874c016cce 100644 --- a/tests/gateway/test_compression_failure_session_sync.py +++ b/tests/gateway/test_compression_failure_session_sync.py @@ -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