fix(gateway): preserve peer routing across compression recovery

This commit is contained in:
João Vitor Cunha 2026-06-29 21:21:40 -03:00 committed by kshitij
parent e2ffbf0cf4
commit d5b4879d4a
3 changed files with 103 additions and 2 deletions

View file

@ -11001,6 +11001,11 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
if agent_result.get("session_id") and agent_result["session_id"] != session_entry.session_id: if agent_result.get("session_id") and agent_result["session_id"] != session_entry.session_id:
session_entry.session_id = agent_result["session_id"] session_entry.session_id = agent_result["session_id"]
self.session_store._save() self.session_store._save()
self.session_store._record_gateway_session_peer(
session_entry.session_id,
session_key,
source,
)
await asyncio.to_thread( await asyncio.to_thread(
self._sync_telegram_topic_binding, self._sync_telegram_topic_binding,
source, session_entry, reason="agent-result-compression", source, session_entry, reason="agent-result-compression",
@ -17712,6 +17717,11 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
if entry: if entry:
entry.session_id = agent_session_id entry.session_id = agent_session_id
self.session_store._save() self.session_store._save()
self.session_store._record_gateway_session_peer(
agent_session_id,
session_key,
source,
)
# If this is a Telegram DM and source.thread_id was lost during # If this is a Telegram DM and source.thread_id was lost during
# the session split (synthetic / recovered event), restore it # the session split (synthetic / recovered event), restore it

View file

@ -965,6 +965,7 @@ class SessionStore:
return return
stale_keys: list = [] stale_keys: list = []
recovered_keys = 0
try: try:
for key, entry in self._entries.items(): for key, entry in self._entries.items():
row = db.get_session(entry.session_id) row = db.get_session(entry.session_id)
@ -972,6 +973,43 @@ class SessionStore:
# end_reason is None -> session alive — keep # end_reason is None -> session alive — keep
# end_reason not None -> session ended — prune # end_reason not None -> session ended — prune
if row is not None and row.get("end_reason") is not None: if row is not None and row.get("end_reason") is not None:
recovered_entry = None
if entry.origin is not None:
try:
recovered_entry = self._recover_session_from_db(
session_key=key,
source=entry.origin,
now=_now(),
)
except Exception as exc:
logger.debug(
"gateway.session: recovery lookup failed for stale "
"sessions.json entry %r -> %s: %s",
key,
entry.session_id,
exc,
)
# If the stale entry points at a compression-ended parent but
# a newer live child session exists for the exact same gateway
# peer, repoint the routing index instead of dropping it. A
# hard restart between compression rotation and the next clean
# save otherwise leaves Telegram with no resumable mapping, so
# queued/resume-pending work disappears until the user sends a
# fresh message.
if recovered_entry is not None and recovered_entry.session_id != entry.session_id:
logger.warning(
"gateway.session: repointing stale sessions.json entry "
"%r from ended %s (end_reason=%r) to recovered %s",
key,
entry.session_id,
row["end_reason"],
recovered_entry.session_id,
)
self._entries[key] = recovered_entry
recovered_keys += 1
continue
logger.warning( logger.warning(
"gateway.session: pruning stale sessions.json entry " "gateway.session: pruning stale sessions.json entry "
"%r -> %s (end_reason=%r); left by a crashed gateway", "%r -> %s (end_reason=%r); left by a crashed gateway",
@ -988,7 +1026,7 @@ class SessionStore:
for key in stale_keys: for key in stale_keys:
del self._entries[key] del self._entries[key]
if stale_keys: if stale_keys or recovered_keys:
self._save() self._save()
def _save(self) -> None: def _save(self) -> None:

View file

@ -12,7 +12,7 @@ from datetime import datetime, timedelta
from unittest.mock import MagicMock, patch from unittest.mock import MagicMock, patch
from gateway.config import GatewayConfig, Platform, SessionResetPolicy from gateway.config import GatewayConfig, Platform, SessionResetPolicy
from gateway.session import SessionEntry, SessionStore from gateway.session import SessionEntry, SessionSource, SessionStore
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@ -31,6 +31,18 @@ def _make_entry(key: str, session_id: str) -> SessionEntry:
) )
def _make_entry_with_origin(key: str, session_id: str) -> SessionEntry:
entry = _make_entry(key, session_id)
entry.origin = SessionSource(
platform=Platform.TELEGRAM,
chat_id="5140768830",
chat_type="dm",
user_id="5140768830",
user_name="João",
)
return entry
def _make_store_with_db(tmp_path, db_mock) -> SessionStore: def _make_store_with_db(tmp_path, db_mock) -> SessionStore:
"""Build a SessionStore with a mock SessionDB, bypassing disk load.""" """Build a SessionStore with a mock SessionDB, bypassing disk load."""
config = GatewayConfig(default_reset_policy=SessionResetPolicy(mode="none")) config = GatewayConfig(default_reset_policy=SessionResetPolicy(mode="none"))
@ -98,6 +110,47 @@ class TestPruneStaleSessionsLocked:
assert "key_b" not in store._entries assert "key_b" not in store._entries
assert "key_c" in store._entries assert "key_c" in store._entries
def test_repoints_stale_compression_parent_to_latest_live_child(self, tmp_path):
"""Compression-ended parents should recover their live child mapping.
A gateway crash can leave sessions.json pointing at the pre-compression
parent (end_reason='compression') even though the agent already rotated
into a live child session. If the child has gateway peer metadata, the
startup prune pass must repoint the route instead of deleting it, or
restart auto-resume and queued follow-ups have no session to continue.
"""
key = "agent:main:telegram:dm:5140768830"
db = _db_returning({
"sid_parent": {"end_reason": "compression", "id": "sid_parent"},
})
db.find_latest_gateway_session_for_peer.return_value = {
"id": "sid_child",
"started_at": 1782744974.0,
}
store = _make_store_with_db(tmp_path, db)
store._entries[key] = _make_entry_with_origin(key, "sid_parent")
store._prune_stale_sessions_locked()
assert key in store._entries
assert store._entries[key].session_id == "sid_child"
db.find_latest_gateway_session_for_peer.assert_called_once()
db.reopen_session.assert_called_once_with("sid_child")
def test_prunes_stale_entry_when_recovery_only_finds_same_ended_session(self, tmp_path):
key = "agent:main:telegram:dm:5140768830"
db = _db_returning({"sid_parent": {"end_reason": "agent_close", "id": "sid_parent"}})
db.find_latest_gateway_session_for_peer.return_value = {
"id": "sid_parent",
"started_at": 1782744974.0,
}
store = _make_store_with_db(tmp_path, db)
store._entries[key] = _make_entry_with_origin(key, "sid_parent")
store._prune_stale_sessions_locked()
assert key not in store._entries
def test_noop_when_db_is_none(self, tmp_path): def test_noop_when_db_is_none(self, tmp_path):
config = GatewayConfig(default_reset_policy=SessionResetPolicy(mode="none")) config = GatewayConfig(default_reset_policy=SessionResetPolicy(mode="none"))
with patch("gateway.session.SessionStore._ensure_loaded"): with patch("gateway.session.SessionStore._ensure_loaded"):