diff --git a/gateway/run.py b/gateway/run.py index a5fe58f97a5..b6afc605ec0 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -11001,6 +11001,11 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew 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", @@ -17712,6 +17717,11 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew 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, + ) # If this is a Telegram DM and source.thread_id was lost during # the session split (synthetic / recovered event), restore it diff --git a/gateway/session.py b/gateway/session.py index 20ac65c7e4f..c741cfba50c 100644 --- a/gateway/session.py +++ b/gateway/session.py @@ -965,6 +965,7 @@ class SessionStore: return stale_keys: list = [] + recovered_keys = 0 try: for key, entry in self._entries.items(): row = db.get_session(entry.session_id) @@ -972,6 +973,43 @@ class SessionStore: # end_reason is None -> session alive — keep # end_reason not None -> session ended — prune 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( "gateway.session: pruning stale sessions.json entry " "%r -> %s (end_reason=%r); left by a crashed gateway", @@ -988,7 +1026,7 @@ class SessionStore: for key in stale_keys: del self._entries[key] - if stale_keys: + if stale_keys or recovered_keys: self._save() def _save(self) -> None: diff --git a/tests/gateway/test_session_store_stale_prune.py b/tests/gateway/test_session_store_stale_prune.py index dac5a3e02b8..9b8cf1fe441 100644 --- a/tests/gateway/test_session_store_stale_prune.py +++ b/tests/gateway/test_session_store_stale_prune.py @@ -12,7 +12,7 @@ from datetime import datetime, timedelta from unittest.mock import MagicMock, patch 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: """Build a SessionStore with a mock SessionDB, bypassing disk load.""" config = GatewayConfig(default_reset_policy=SessionResetPolicy(mode="none")) @@ -98,6 +110,47 @@ class TestPruneStaleSessionsLocked: assert "key_b" not 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): config = GatewayConfig(default_reset_policy=SessionResetPolicy(mode="none")) with patch("gateway.session.SessionStore._ensure_loaded"):