fix(gateway): flush a spooled pending message under the session store's key so a restart can replay it

evals/shutdown_flush_adapter_key_strand.py

@@ -0,0 +1,121 @@
+"""Shutdown flush spooled under the adapter's session key strands on the next boot (#136189).
+
+``_pending_messages`` slots are keyed by the adapter's own isolation flags
+(``PlatformConfig.extra``), while ``SessionStore`` routes by the gateway-wide
+``group_sessions_per_user`` / ``thread_sessions_per_user``. ``recover_pending_to_db`` resolves a
+payload's key through the store alone (routing map, then the durable row under the exact key), so
+a spool file written under the adapter key matches neither and is preserved forever.
+
+This probe drives the real entry points only — ``adapter.cancel_background_tasks()`` then
+``recover_pending_to_db`` — so it runs unchanged against a fixed and an unfixed checkout:
+
+    python evals/shutdown_flush_adapter_key_strand.py [ROOT]
+
+Prints the two derivations, the key actually spooled, the replay count and a ``VERDICT=`` line
+(``STRANDED`` before the fix, ``RECOVERED`` after).
+"""
+
+import asyncio
+import json
+import os
+import sys
+import tempfile
+from pathlib import Path
+
+ROOT = Path(sys.argv[1]).resolve() if len(sys.argv) > 1 else Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+HOME = Path(tempfile.mkdtemp(prefix="flush-strand-"))
+os.environ["HERMES_HOME"] = str(HOME)
+os.environ["HERMES_NIX_BUILD"] = "1"
+
+from gateway.config import GatewayConfig, Platform, PlatformConfig
+from gateway.platforms.base import BasePlatformAdapter, SendResult
+from gateway.platforms.event import MessageEvent, MessageType
+from gateway.session import SessionSource, SessionStore
+from gateway.shutdown_flush import recover_pending_to_db
+
+
+class _Adapter(BasePlatformAdapter):
+    """The minimum a real adapter needs to reach the shutdown flush."""
+
+    def __init__(self):
+        super().__init__(PlatformConfig(enabled=True, token="***"), Platform.TELEGRAM)
+
+    async def connect(self, *, is_reconnect: bool = False):
+        return True
+
+    async def disconnect(self):
+        return None
+
+    async def send(self, chat_id, content, reply_to=None, metadata=None):
+        return SendResult(success=True, message_id="1")
+
+    async def send_typing(self, chat_id, metadata=None):
+        return None
+
+    async def get_chat_info(self, chat_id):
+        return {"id": chat_id}
+
+
+class _Row:
+    """The durable row the resolver's exact-key peer finder would return."""
+
+    def __init__(self, session_id: str, session_key: str):
+        self.session_id = session_id
+        self.session_key = session_key
+        self.appended: list[dict] = []
+
+    def find_latest_gateway_session_for_peer(self, *, source, session_key=None, **kwargs):
+        # The real finder runs only the ``s.session_key = ?`` branch when no chat tuple is given:
+        # a row matches one exact key, never "any key".
+        if session_key != self.session_key:
+            return None
+        return {"id": self.session_id, "started_at": 0.0}
+
+    def append_message(self, **kwargs):
+        self.appended.append(kwargs)
+
+
+def _spooled_keys() -> list[str]:
+    directory = HOME / "pending_messages"
+    return sorted(json.loads(path.read_text())["session_key"]
+                  for path in directory.glob("*.json"))
+
+
+def main() -> int:
+    store = SessionStore(sessions_dir=HOME / "sessions", config=GatewayConfig())
+
+    adapter = _Adapter()
+    adapter.set_session_store(store)
+    # One platform overrides per-user thread isolation; the store keeps the gateway-wide default.
+    adapter.config.extra["thread_sessions_per_user"] = True
+
+    source = SessionSource(platform=Platform.TELEGRAM, chat_id="900", chat_type="group",
+                           thread_id="900", user_id="u1")
+    event = MessageEvent(text="queued while busy", message_type=MessageType.TEXT,
+                         source=source, message_id="queued-1")
+    slot_key = adapter._event_session_key(event)
+    store_key = store._generate_session_key(source)
+    db = _Row("sid-900", store_key)
+    store._db = db  # every key resolves against this handle (no profile I/O)
+    adapter._pending_messages[slot_key] = event
+
+    asyncio.run(adapter.cancel_background_tasks())
+    spooled = _spooled_keys()
+    recovered = recover_pending_to_db(db, session_resolver=store.resolve_session_id_for_key)
+    remaining = len(_spooled_keys())
+
+    print(f"slot_key={slot_key}")
+    print(f"store_key={store_key}")
+    print(f"spooled_keys={spooled}")
+    print(f"recovered={recovered}")
+    print(f"replayed_content={[call.get('content') for call in db.appended]}")
+    print(f"spool_remaining={remaining}")
+    verdict = "RECOVERED" if recovered == 1 and not remaining else "STRANDED"
+    print(f"VERDICT={verdict}")
+    return 0 if verdict == "RECOVERED" else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

gateway/platforms/base.py

@@ -4793,8 +4793,8 @@ async def cancel_background_tasks(self) -> None:
                                self.name, sum(not t.done() for t in tasks))
                 break
         with contextlib.suppress(Exception):  # flush pending messages to disk before clearing
-            from gateway.shutdown_flush import flush_pending_to_file
-            flush_pending_to_file(self._pending_messages, reason="adapter_shutdown")
+            from gateway.shutdown_flush import adapter_flush_keying, flush_pending_to_file
+            flush_pending_to_file(self._pending_messages, reason="adapter_shutdown", **adapter_flush_keying(self))
         for state in self._text_debounce_store().values():
             state.cancel_timer()
         for bucket in (self._background_tasks, self._expected_cancelled_tasks, self._session_tasks,

gateway/run_shutdown.py

@@ -2009,14 +2009,14 @@ async def _stop_release_runtime_state(self, ctx: GatewayShutdownMixin._StopConte
             self._release_running_agent_state(_session_key)
         # Flush pending messages before clearing: under FTS5 corruption they are the only surviving copy.
         with suppress(Exception):
-            from gateway.shutdown_flush import flush_pending_to_file
-            flush_pending_to_file(dict(self._pending_messages), reason="shutdown")
+            from gateway.shutdown_flush import flush_pending_to_file, runner_flush_keying
+            flush_pending_to_file(dict(self._pending_messages), reason="shutdown", **runner_flush_keying(self))
         # The overflow FIFO tail lives in SessionState.conversation.queued_events — flush it too.
         with suppress(Exception):
-            from gateway.shutdown_flush import flush_overflow_to_file
+            from gateway.shutdown_flush import flush_overflow_to_file, runner_flush_keying
             flush_overflow_to_file(
                 {_k: list(_v) for _k, _v in dict(getattr(self, "_queued_events", None) or {}).items() if _v},
-                reason="shutdown",
+                reason="shutdown", **runner_flush_keying(self),
             )
         # Live SessionState views: clear() resets one field per session (never a wholesale dict swap).
         self._running_agents.clear()

gateway/shutdown_flush.py

@@ -21,7 +21,7 @@
 import time
 import uuid
 from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any, Callable, Dict, Optional
 
 logger = logging.getLogger(__name__)
 
@@ -87,25 +87,135 @@ def _flush_value(flush_dir: Path, kind: str, session_key: str, value: Any, **ext
         return False
 
 
-def flush_pending_to_file(pending: dict[str, Any], *, reason: str = "shutdown") -> int:
-    """Serialise non-empty ``_pending_messages`` slots (``MessageEvent`` or str); return count."""
+def _recoverable_key(value: Any, slot_key: str,
+                     session_key_for: Optional[Callable[[Any], Any]] = None,
+                     slot_key_routes: Optional[Callable[[str], bool]] = None) -> str:
+    """The key to record for *value*: the slot key when the store routes it, else its store key.
+
+    Slots are keyed by the adapter's own isolation flags (``PlatformConfig.extra``), which may
+    differ from the gateway-wide ones ``SessionStore`` routes by: an operator can override
+    ``group_sessions_per_user`` / ``thread_sessions_per_user`` on one platform. Recovery resolves
+    a key through the store only (routing map, then the durable row under the exact key), so a
+    slot key that matches neither strands the message in the spool forever. That is the only case
+    in which ``session_key_for`` replaces it: the store's key derivation
+    (``SessionStore._generate_session_key``, or ``GatewayRunner._session_key_for_source``).
+
+    A slot key the store *can* route is recorded unchanged. It is the key the producing adapter
+    queued the value under and the key its live drain re-dispatches it with, so it is the
+    payload's established owner; re-deriving can lose that owner, because a synthetic producer
+    (a recall notice written straight into the slot) may carry a partial ``source`` whose own
+    derivation drops the participant the slot key holds — the notice then strands on a key no
+    session was ever created under. ``slot_key_routes`` is the store's exact-key resolver; a
+    value without a ``source`` (a raw str slot, an agent-history payload) and any derivation
+    failure keep the slot key unchanged.
+    """
+    if callable(slot_key_routes):
+        try:
+            if slot_key_routes(slot_key):
+                return slot_key
+        except Exception:  # a routability check must never cost the payload its key
+            logger.debug("Could not check whether %s is routable; deriving its store key",
+                         slot_key, exc_info=True)
+    source = getattr(value, "source", None)
+    if source is None or not callable(session_key_for):
+        return slot_key
+    try:
+        key = session_key_for(source)
+    except Exception:  # a store/profile lookup can fail any number of ways; the slot key still works
+        logger.debug("Could not resolve the store key for %s; recording the slot key",
+                     slot_key, exc_info=True)
+        return slot_key
+    return key if isinstance(key, str) and key else slot_key
+
+
+def adapter_store_key_for(adapter: Any) -> Optional[Callable[[Any], Any]]:
+    """``session_key_for`` for an adapter's flush: its session store's key derivation, or None.
+
+    ``None`` (no store wired, or an object without the derivation) makes the flushers record each
+    slot key verbatim, which is the pre-#136189 behaviour.
+    """
+    derive = getattr(getattr(adapter, "_session_store", None), "_generate_session_key", None)
+    return derive if callable(derive) else None
+
+
+def store_slot_key_routes(store: Any) -> Optional[Callable[[str], bool]]:
+    """``slot_key_routes`` for *store*: does its own exact-key resolver land this key?
+
+    The resolver is the one ``recover_pending_to_db`` hands payloads to (routing map first, then
+    the durable row under the exact key), so a slot key it cannot land is a slot key recovery
+    would strand — the only case the flush re-keys. ``None`` (no store, or one without the
+    resolver) leaves every key as its derivation gives it, the head's behaviour.
+    """
+    resolve = getattr(store, "resolve_session_id_for_key", None)
+    if not callable(resolve):
+        return None
+
+    def routes(session_key: str) -> bool:
+        try:
+            return bool(resolve(session_key))
+        except Exception:  # a profile lookup can fail; an unconfirmed key is not worth recording
+            logger.debug("Could not ask the store whether %s routes; deriving its store key",
+                         session_key, exc_info=True)
+            return False
+
+    return routes
+
+
+def adapter_flush_keying(adapter: Any) -> dict[str, Any]:
+    """The keying kwargs for an adapter's shutdown flush.
+
+    Bundled as the gateway's other ``**_kwargs`` builders are, and because both flush call sites
+    live in files the repo's file-length ratchet already holds over target (``base.py``,
+    ``run_shutdown.py``): the bundle keeps each call site at its current line count.
+    """
+    store = getattr(adapter, "_session_store", None)
+    return {"session_key_for": adapter_store_key_for(adapter),
+            "slot_key_routes": store_slot_key_routes(store)}
+
+
+def runner_flush_keying(runner: Any) -> dict[str, Any]:
+    """The keying kwargs for the runner's own flush (queue head and overflow tail).
+
+    Same bundle, same reason; the derivation is the runner's own
+    (``GatewayRunner._session_key_for_source``), which falls back to the config when no store is
+    wired.
+    """
+    return {"session_key_for": getattr(runner, "_session_key_for_source", None),
+            "slot_key_routes": store_slot_key_routes(getattr(runner, "session_store", None))}
+
+
+def flush_pending_to_file(pending: dict[str, Any], *, reason: str = "shutdown",
+                          session_key_for: Optional[Callable[[Any], Any]] = None,
+                          slot_key_routes: Optional[Callable[[str], bool]] = None) -> int:
+    """Serialise non-empty ``_pending_messages`` slots (``MessageEvent`` or str); return count.
+
+    ``session_key_for`` maps a slot's ``source`` to the key the session store routes it by, so
+    the payload replays after a restart even when the adapter that queued it derived a different
+    key; ``slot_key_routes`` says whether the slot key itself is one recovery can resolve, in
+    which case it is recorded as it stands (see :func:`_recoverable_key`).
+    """
     if not pending:
         return 0
     flush_dir, ts, flushed = _get_flush_dir(), int(time.time()), 0
     for session_key, value in list(pending.items()):
         if value is not None:
-            flushed += _flush_value(flush_dir, "pending", session_key, value, reason=reason, ts=ts)
+            key = _recoverable_key(value, session_key, session_key_for, slot_key_routes)
+            flushed += _flush_value(flush_dir, "pending", key, value, reason=reason, ts=ts)
     if flushed:
         logger.info("Flushed %d pending message(s) to %s (reason=%s)", flushed, flush_dir, reason)
     return flushed
 
 
-def flush_overflow_to_file(overflow_by_session: dict[str, Any], *, reason: str = "shutdown") -> int:
+def flush_overflow_to_file(overflow_by_session: dict[str, Any], *, reason: str = "shutdown",
+                           session_key_for: Optional[Callable[[Any], Any]] = None,
+                           slot_key_routes: Optional[Callable[[str], bool]] = None) -> int:
     """Serialise the FIFO overflow tails (``queued_events``) to disk; return events flushed.
 
     The adapter slot holds the queue head and ``SessionState.conversation.queued_events`` the
     tail; both must survive restart. Each event is its own payload in the slot-flush shape so
     ``recover_pending_to_db`` replays them unchanged; ``seq`` preserves arrival order per session.
+    ``session_key_for`` and ``slot_key_routes`` key each event exactly as in
+    :func:`flush_pending_to_file`.
     """
     if not overflow_by_session:
         return 0
@@ -115,7 +225,8 @@ def flush_overflow_to_file(overflow_by_session: dict[str, Any], *, reason: str =
             continue
         for seq, value in enumerate(list(events)):
             if value is not None:
-                flushed += _flush_value(flush_dir, "overflow", session_key, value, reason=reason,
+                key = _recoverable_key(value, session_key, session_key_for, slot_key_routes)
+                flushed += _flush_value(flush_dir, "overflow", key, value, reason=reason,
                                         ts=ts, seq=seq)
     if flushed:
         logger.info("Flushed %d queued overflow message(s) to %s (reason=%s)", flushed, flush_dir,