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,161 @@ 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,
+                     key_lands: Optional[Callable[[str], Optional[bool]]] = None) -> str:
+    """The key to record for *value*: the store's own key for it, while the store can land that.
+
+    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, in either
+    direction. Recovery resolves a key through the store only (routing map, then the durable row
+    under the exact key), so a key that matches neither strands the message in the spool forever.
+
+    ``session_key_for`` is the question worth asking of that store: it maps the value's ``source``
+    to the key the store routes *this payload* by (``SessionStore._generate_session_key``, or
+    ``GatewayRunner._session_key_for_source``) — the key a turn for this payload is admitted
+    under. ``key_lands`` asks the same store whether a key has a session behind it: ``True`` (a
+    session the store can reach), ``False`` (the store answered: no session under this key), or
+    ``None`` (the store could not answer for this key — no reachable home for it).
+
+    The slot key is the producer's own key, the one the queue re-dispatches the payload by. It is
+    kept when the derivation is unconfirmed (``None``: an unavailable lookup decides nothing) and
+    when the store knows a session under the slot key and none under the derivation — a synthetic
+    producer such as a recall notice may carry a partial ``source`` whose own derivation drops the
+    participant the slot key holds, and moving a payload that already has a home onto a key no
+    session was ever created under is the strand this re-keying exists to prevent. The derivation
+    is recorded otherwise, including when the store knows no session under either key: there the
+    payload's own key is the one recovery names in its warning, and nothing is lost — the spool
+    file is kept. A value without a ``source`` (a raw str slot, an agent-history payload) and any
+    derivation failure keep the slot key unchanged.
+    """
+    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
+    if not (isinstance(key, str) and key) or key == slot_key:
+        return slot_key
+    if not callable(key_lands):
+        return key  # nothing to ask: the derivation stands, as it always has with no store
+    lands = _key_lands(key_lands, key)
+    if lands is None:
+        return slot_key
+    if lands:
+        return key
+    return slot_key if _key_lands(key_lands, slot_key) is True else key
+
+
+def _key_lands(key_lands: Callable[[str], Optional[bool]],
+               session_key: str) -> Optional[bool]:
+    """``key_lands(session_key)``, guarded: a failing check answers "unknown" (``None``)."""
+    try:
+        answer = key_lands(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 lands", session_key, exc_info=True)
+        return None
+    return None if answer is None else bool(answer)
+
+
+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_key_lands(store: Any) -> Optional[Callable[[str], Optional[bool]]]:
+    """``key_lands`` for *store*: does it have a session behind this key?
+
+    Uses the resolver ``recover_pending_to_db`` hands payloads to (routing map first, then the
+    durable row under the exact key), so a key it cannot land is a key recovery would strand.
+    ``True`` and ``False`` are the store's own answers; ``None`` reports that it could not answer
+    for that key at all — ``SessionStore._db_for_key`` failing closed (the seam
+    ``resolve_session_id_for_key`` returns ``None`` on), i.e. no reachable home for the key — which
+    is not a verdict on the key. ``None`` (no store, or one without the resolver) leaves every key
+    as its derivation gives it.
+    """
+    resolve = getattr(store, "resolve_session_id_for_key", None)
+    db_for_key = getattr(store, "_db_for_key", None)
+    if not callable(resolve):
+        return None
+
+    def lands(session_key: str) -> Optional[bool]:
+        if resolve(session_key):
+            return True
+        # Nothing resolved: only now is it worth asking whether the store could look at all, so
+        # "no session under this key" is never read out of an unreachable one.
+        if callable(db_for_key) and db_for_key(session_key) is None:
+            return None
+        return False
+
+    return lands
+
+
+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),
+            "key_lands": store_key_lands(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),
+            "key_lands": store_key_lands(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,
+                          key_lands: Optional[Callable[[str], Optional[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; ``key_lands`` is that store's own verdict on a key (``True``/``False``, ``None`` when it
+    could not answer for that key), which decides whether the derivation is recorded (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, key_lands)
+            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,
+                           key_lands: Optional[Callable[[str], Optional[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 ``key_lands`` key each event exactly as in
+    :func:`flush_pending_to_file`.
     """
     if not overflow_by_session:
         return 0
@@ -115,7 +251,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, key_lands)
+                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,