diff --git a/evals/shutdown_flush_adapter_key_strand.py b/evals/shutdown_flush_adapter_key_strand.py new file mode 100644 index 0000000000000..82a540e5d21ef --- /dev/null +++ b/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()) diff --git a/gateway/platforms/base.py b/gateway/platforms/base.py index 73270e9bd4a41..b0e59a2965687 100644 --- a/gateway/platforms/base.py +++ b/gateway/platforms/base.py @@ -4764,8 +4764,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, diff --git a/gateway/run_shutdown.py b/gateway/run_shutdown.py index e1f7285b9148b..18b4bb937a161 100644 --- a/gateway/run_shutdown.py +++ b/gateway/run_shutdown.py @@ -1996,14 +1996,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() diff --git a/gateway/shutdown_flush.py b/gateway/shutdown_flush.py index 9700bd71d2135..d6a4d98013635 100644 --- a/gateway/shutdown_flush.py +++ b/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 from hermes_dir_acl import make_user_dir logger = logging.getLogger(__name__) @@ -88,25 +88,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 @@ -116,7 +252,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, diff --git a/tests/gateway/test_recall_notice_flush_identity.py b/tests/gateway/test_recall_notice_flush_identity.py new file mode 100644 index 0000000000000..adc8b0df35fe0 --- /dev/null +++ b/tests/gateway/test_recall_notice_flush_identity.py @@ -0,0 +1,232 @@ +"""A sender-less Yuanbao recall notice must survive the shutdown flush into its own session. + +Reported in review of this PR: a value whose ``source`` names no participant is queued by a +platform producer straight into the slot of the turn it belongs to, so a flush that re-derives +the key from that source alone loses the participant segment the slot carries and the notice +strands in the spool instead of replaying. The pre-PR flush recorded the slot key verbatim. + +The recall guard is that producer: it rebuilds the notice's source from the recall callback, which +``DecodeMiddleware.parse_json_push`` admits without ``from_account``. This test drives the real +decode, the real recall producer, the real adapter shutdown flush and the real +SessionStore/SQLite reopen, and asserts the exact notice text lands in the ORIGINAL session with +the spool drained. ``from_account=""`` is the supported partial-source case; ``from_account="u1"`` +is a positive control proving the fixture and the recovery seam are sound. +""" + +import asyncio +import json +from pathlib import Path + +import pytest + +from gateway.config import GatewayConfig, PlatformConfig +from plugins.platforms.yuanbao.adapter import ( + DecodeMiddleware, + InboundContext, + RecallGuardMiddleware, + YuanbaoAdapter, +) +from gateway.session import SessionStore +from gateway.shutdown_flush import recover_pending_to_db + +RECALLED_MSG_ID = "m-original" +GROUP_CODE = "G" +ORIGINAL_SENDER = "u1" +# Both the adapter (PlatformConfig.extra) and the store use the gateway-wide isolation flags, so +# the slot key the notice is queued under IS the key the store routes the original turn by. +FLAGS = dict(group_sessions_per_user=True, thread_sessions_per_user=True) + + +def _config() -> GatewayConfig: + return GatewayConfig(**FLAGS) + + +def _open_store(tmp_path: Path) -> SessionStore: + return SessionStore(sessions_dir=tmp_path / "sessions", config=_config()) + + +@pytest.fixture +def hermes_home(tmp_path, monkeypatch): + """Real SQLite under an isolated HERMES_HOME (same seam as the other reviewer probes).""" + import hermes_constants + import hermes_state + + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + monkeypatch.setattr(hermes_constants, "get_hermes_home", lambda: tmp_path) + monkeypatch.setattr(hermes_state, "DEFAULT_DB_PATH", hermes_state._IMPORT_DEFAULT_DB_PATH) + return tmp_path + + +def _wire(tmp_path: Path): + store = _open_store(tmp_path) + adapter = YuanbaoAdapter(PlatformConfig(enabled=True, token="***")) + adapter.config.extra.update(**FLAGS) + adapter.set_session_store(store) + return store, adapter + + +def _original_source(adapter: YuanbaoAdapter): + """The real group source of the message that is being processed when the recall arrives.""" + return adapter.build_source( + chat_id=f"group:{GROUP_CODE}", chat_type="group", + user_id=ORIGINAL_SENDER, thread_id="main", + ) + + +async def _noop() -> None: + return None + + +async def _emit_recall(adapter: YuanbaoAdapter, from_account: str) -> str: + """Drive the real decode + real recall producer; return the notice text queued for the slot.""" + raw = { + "callback_command": "Group.CallbackAfterRecallMsg", + "group_code": GROUP_CODE, + "recall_msg_seq_list": [{"msg_id": RECALLED_MSG_ID}], + } + if from_account: # otherwise exercise the supported callback without this key + raw["from_account"] = from_account + push = DecodeMiddleware.parse_json_push(raw) + assert push is not None, "the decoder must admit this recall callback" + assert push["callback_command"] == "Group.CallbackAfterRecallMsg" + assert push["from_account"] == from_account + + slot = adapter._source_session_key(_original_source(adapter)) + adapter._processing_msg_ids[slot] = RECALLED_MSG_ID + adapter._processing_msg_texts[slot] = "" # blank: no content-redact poll + adapter._active_sessions[slot] = asyncio.Event() + + await RecallGuardMiddleware().handle(InboundContext(adapter=adapter, push=push), _noop) + + notice = adapter._pending_messages.get(slot) + assert notice is not None, "the recall interrupt must queue its notice on the matched slot" + return notice.text + + +def _spool_keys(flush_dir: Path) -> list[str]: + return sorted(json.loads(p.read_text(encoding="utf-8"))["session_key"] + for p in flush_dir.glob("pending-*.json")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("from_account", ["", ORIGINAL_SENDER]) +async def test_recall_notice_replays_into_the_original_session(tmp_path, hermes_home, from_account): + """The queued notice must survive a real adapter shutdown flush and land in the original row.""" + store, adapter = _wire(tmp_path) + try: + entry = store.get_or_create_session(_original_source(adapter)) + slot = adapter._source_session_key(_original_source(adapter)) + # Precondition: the slot key and the store key are the same string (flags agree). + assert slot == entry.session_key == store._generate_session_key(_original_source(adapter)) + + notice_text = await _emit_recall(adapter, from_account) + + # Real production flush seam: the exact changed line in base.py. + await adapter.cancel_background_tasks() + + store.close_all_db_handles() + reopened = _open_store(tmp_path) + try: + db = reopened._db_for_key(entry.session_key) + recovered = recover_pending_to_db( + db, session_resolver=reopened.resolve_session_id_for_key) + contents = [m["content"] for m in db.get_messages(entry.session_id)] + finally: + reopened.close_all_db_handles() + + assert recovered == 1, ( + "the recall notice was not recovered; spool keys=" + f"{_spool_keys(tmp_path / 'pending_messages')}" + ) + assert contents == [notice_text], f"notice text not in the original session: {contents!r}" + assert not list((tmp_path / "pending_messages").glob("*.json")), "spool not drained" + finally: + store.close_all_db_handles() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("from_account", ["", ORIGINAL_SENDER]) +async def test_flushed_key_is_the_slot_the_notice_was_queued_under(tmp_path, hermes_home, + from_account): + """The spooled key must be the slot the notice was queued under, so recovery can resolve it.""" + store, adapter = _wire(tmp_path) + try: + store.get_or_create_session(_original_source(adapter)) + slot = adapter._source_session_key(_original_source(adapter)) + await _emit_recall(adapter, from_account) + + await adapter.cancel_background_tasks() + + assert _spool_keys(tmp_path / "pending_messages") == [slot] + finally: + store.close_all_db_handles() + + +# -- the DM branch of the same producer (C2C.CallbackAfterMsgWithDraw) ------------------------- + + +def _original_dm_source(adapter: YuanbaoAdapter): + """The real direct-chat source of the message being processed when the withdrawal arrives.""" + return adapter.build_source( + chat_id=f"direct:{ORIGINAL_SENDER}", chat_type="dm", user_id=ORIGINAL_SENDER, + ) + + +async def _emit_c2c_recall(adapter: YuanbaoAdapter, from_account: str) -> str: + """Drive the real decode + real C2C withdrawal producer; return the queued notice text.""" + raw = { + "callback_command": "C2C.CallbackAfterMsgWithDraw", + "msg_id": RECALLED_MSG_ID, + "msg_seq": 7, + } + if from_account: + raw["from_account"] = from_account + push = DecodeMiddleware.parse_json_push(raw) + assert push is not None, "the decoder must admit this withdrawal callback" + + slot = adapter._source_session_key(_original_dm_source(adapter)) + adapter._processing_msg_ids[slot] = RECALLED_MSG_ID + adapter._processing_msg_texts[slot] = "" # blank: no content-redact poll + adapter._active_sessions[slot] = asyncio.Event() + + await RecallGuardMiddleware().handle(InboundContext(adapter=adapter, push=push), _noop) + + notice = adapter._pending_messages.get(slot) + assert notice is not None, "the withdrawal interrupt must queue its notice on the matched slot" + return notice.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("from_account", ["", ORIGINAL_SENDER]) +async def test_c2c_withdrawal_notice_replays_into_the_original_dm_session(tmp_path, hermes_home, + from_account): + """A direct-chat notice carries the chat it interrupts, so a shutdown cannot move it.""" + store, adapter = _wire(tmp_path) + try: + entry = store.get_or_create_session(_original_dm_source(adapter)) + slot = adapter._source_session_key(_original_dm_source(adapter)) + assert slot == entry.session_key == store._generate_session_key(_original_dm_source(adapter)) + + notice_text = await _emit_c2c_recall(adapter, from_account) + assert _spool_keys(tmp_path / "pending_messages") == [] + + await adapter.cancel_background_tasks() + + assert _spool_keys(tmp_path / "pending_messages") == [slot], ( + "the withdrawal notice was spooled under a key the store does not route by" + ) + + store.close_all_db_handles() + reopened = _open_store(tmp_path) + try: + db = reopened._db_for_key(entry.session_key) + recovered = recover_pending_to_db( + db, session_resolver=reopened.resolve_session_id_for_key) + contents = [m["content"] for m in db.get_messages(entry.session_id)] + finally: + reopened.close_all_db_handles() + + assert recovered == 1 + assert contents == [notice_text] + finally: + store.close_all_db_handles() diff --git a/tests/gateway/test_shutdown_flush_store_key.py b/tests/gateway/test_shutdown_flush_store_key.py new file mode 100644 index 0000000000000..5afae5d07f1bd --- /dev/null +++ b/tests/gateway/test_shutdown_flush_store_key.py @@ -0,0 +1,332 @@ +"""A shutdown flush must record a key the next boot's recovery can resolve (#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``. A platform-block override reaches +``PlatformConfig.extra`` (``PlatformConfig.from_dict`` promotes untyped platform keys) and the +runner only ``setdefault``s the gateway-wide value, so one platform can legitimately key a +conversation differently from the store — and a standalone gateway hands that same key to the +busy handler, so the queue head and the FIFO tail inherit it. ``recover_pending_to_db`` resolves +a payload's key through the store alone — routing map first, then the durable row under the exact +key — so a flush written under the adapter key matched neither: the message stranded in the spool +with a "Cannot recover pending message" warning on every boot. + +These tests drive the real shutdown entry points (``adapter.cancel_background_tasks()``, +``runner.stop()``) and the real recovery pass, with the two derivations disagreeing. +""" + +import asyncio +import json +from unittest.mock import patch + +import pytest + +from gateway.config import GatewayConfig +from gateway.platforms.event import MessageEvent, MessageType +from plugins.platforms.yuanbao.adapter import RecallGuardMiddleware +from gateway.session import SessionStore +from gateway.shutdown_flush import flush_pending_to_file, recover_pending_to_db +from tests.gateway.restart_test_helpers import make_restart_runner, make_restart_source + +# One platform overrides per-user thread isolation; the store keeps the gateway-wide default +# (threads shared unless ``thread_sessions_per_user``). Same conversation, two keys. +ADAPTER_KEY = "agent:main:telegram:group:900:900:u1" +STORE_KEY = "agent:main:telegram:group:900:900" + + +class _FakeGatewayDB: + """Minimal SessionDB stand-in: the exact-key peer finder the resolver uses, plus row appends.""" + + def __init__(self, rows_by_key=None): + self.rows_by_key = rows_by_key or {} + self.appended = [] + + def find_latest_gateway_session_for_peer(self, *, source, session_key=None, **kwargs): + return self.rows_by_key.get(session_key) + + def append_message(self, **kwargs): + self.appended.append(kwargs) + + +def _store(tmp_path, db, *, config=None): + with patch("gateway.session.SessionStore._ensure_loaded"): + store = SessionStore(sessions_dir=tmp_path / "sessions", config=config or GatewayConfig()) + store._loaded = True + store._db = db # pin: _db_for_key returns this db for every key (no profile I/O) + return store + + +@pytest.fixture +def flush_dir(tmp_path, monkeypatch): + directory = tmp_path / "pending_messages" + directory.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr("gateway.shutdown_flush._get_flush_dir", lambda: directory) + return directory + + +def _queued_event() -> MessageEvent: + return MessageEvent( + text="queued while busy", + message_type=MessageType.TEXT, + source=make_restart_source(chat_id="900", chat_type="group", thread_id="900"), + message_id="queued-1", + ) + + +def _payload_keys(flush_dir) -> list[str]: + return sorted(json.loads(p.read_text())["session_key"] + for p in flush_dir.glob("pending-*.json")) + + +def _spooled_files(flush_dir) -> list: + return list(flush_dir.glob("pending-*.json")) + + +def _wire(tmp_path, flush_dir, *, store_db=None): + """A restart-test runner + adapter whose platform flags disagree with the store's.""" + runner, adapter = make_restart_runner() + store = _store(tmp_path, store_db if store_db is not None else _FakeGatewayDB()) + runner.session_store = store + adapter.set_session_store(store) + adapter.config.extra["thread_sessions_per_user"] = True + return runner, adapter, store + + +def test_the_two_derivations_disagree_for_the_same_source(tmp_path, flush_dir): + """Precondition of the whole bug: the slot key and the store key are different strings.""" + _runner, adapter, store = _wire(tmp_path, flush_dir) + event = _queued_event() + assert adapter._event_session_key(event) == ADAPTER_KEY + assert store._generate_session_key(event.source) == STORE_KEY + + +@pytest.mark.asyncio +async def test_adapter_shutdown_flush_records_the_store_key(tmp_path, flush_dir): + _runner, adapter, _ = _wire(tmp_path, flush_dir) + event = _queued_event() + adapter._pending_messages[ADAPTER_KEY] = event + + await adapter.cancel_background_tasks() + + assert _payload_keys(flush_dir) == [STORE_KEY] + + +@pytest.mark.asyncio +async def test_adapter_keyed_flush_is_recovered_on_the_next_boot(tmp_path, flush_dir): + """End to end: the spool a restart reads back lands in the session the store routes to.""" + db = _FakeGatewayDB({STORE_KEY: {"id": "sid-900", "started_at": 0.0}}) + _runner, adapter, store = _wire(tmp_path, flush_dir, store_db=db) + adapter._pending_messages[ADAPTER_KEY] = _queued_event() + + await adapter.cancel_background_tasks() + recovered = recover_pending_to_db(db, session_resolver=store.resolve_session_id_for_key) + + assert recovered == 1 + assert [call["session_id"] for call in db.appended] == ["sid-900"] + assert [call["content"] for call in db.appended] == ["queued while busy"] + assert _spooled_files(flush_dir) == [] + + +@pytest.mark.asyncio +async def test_runner_shutdown_flush_records_the_store_key(tmp_path, flush_dir): + """The runner's own stop() flushes the queue head and the FIFO tail the same way.""" + runner, adapter, _ = _wire(tmp_path, flush_dir) + event = _queued_event() + adapter_key = adapter._event_session_key(event) + runner._pending_messages[adapter_key] = event + runner._queued_events = {adapter_key: [event]} + + with patch("gateway.status.remove_pid_file"), patch( + "gateway.status.publish_runtime_status" + ): + await runner.stop() + + assert _payload_keys(flush_dir) == [STORE_KEY, STORE_KEY] + + +def test_a_slot_without_a_source_keeps_its_key(tmp_path, flush_dir): + """Raw-str slots (and agent-history payloads) have no source to re-derive from.""" + def _unusable(_source): + raise RuntimeError("no session store wired") + + assert flush_pending_to_file({"agent:main:telegram:dm:7": "raw text"}, + session_key_for=_unusable) == 1 + assert _payload_keys(flush_dir) == ["agent:main:telegram:dm:7"] + + +def test_a_failing_derivation_keeps_the_slot_key(tmp_path, flush_dir): + """A derivation that raises must never cost the payload its place in the spool.""" + def _unusable(_source): + raise RuntimeError("no session store wired") + + assert flush_pending_to_file({ADAPTER_KEY: _queued_event()}, + session_key_for=_unusable) == 1 + assert _payload_keys(flush_dir) == [ADAPTER_KEY] + + +@pytest.mark.parametrize("derived", ["", 0, None, ["agent:main:telegram:group:900:900"]]) +def test_a_derivation_result_that_names_no_session_keeps_the_slot_key(tmp_path, flush_dir, + derived): + """Only a non-empty str is a key the store can route by; anything else keeps the slot key. + + A derivation that returns a list, a number, ``None``, or an empty string would otherwise + record a payload the next boot's resolution cannot look up — the strand this PR fixes. + """ + assert flush_pending_to_file({ADAPTER_KEY: _queued_event()}, + session_key_for=lambda _source: derived) == 1 + assert _payload_keys(flush_dir) == [ADAPTER_KEY] + + +# --- A supported synthetic producer whose value keeps the slot key's owner -------------------- +# +# ``RecallGuardMiddleware._interrupt_for_recall`` (yuanbao.py's branch C) writes its notice +# straight into the matched in-flight slot, and a recall callback may arrive with no +# ``from_account`` (``DecodeMiddleware.parse_json_push`` admits it). The notice's ``source`` is +# then partial, so its own derivation drops the participant the slot key holds and names a key no +# session was ever created under — re-deriving there strands the notice. The slot itself was +# filled by the ordinary dispatch, so its key is the one the store routes and must be kept. +RECALL_SLOT_KEY = "agent:main:telegram:group:group:G:main:u1" +RECALL_PARTIAL_KEY = "agent:main:telegram:group:group:G:main" + + +def _wire_aligned(tmp_path, db): + """Adapter and store agree on per-user isolation: the matched slot key is a routed key.""" + runner, adapter = make_restart_runner() + store = _store(tmp_path, db, config=GatewayConfig(thread_sessions_per_user=True)) + runner.session_store = store + adapter.set_session_store(store) + adapter.config.extra["group_sessions_per_user"] = True + adapter.config.extra["thread_sessions_per_user"] = True + return runner, adapter, store + + +def _queue_recall_notice(adapter, recalled_id: str = "recalled-1") -> MessageEvent: + """Drive the real recall producer against the maps the ordinary dispatch would have filled.""" + matched = make_restart_source(chat_id="group:G", chat_type="group", thread_id="main") + assert adapter._event_session_key(MessageEvent( + text="", message_type=MessageType.TEXT, source=matched)) == RECALL_SLOT_KEY + adapter._active_sessions[RECALL_SLOT_KEY] = asyncio.Event() + adapter._processing_msg_ids = {RECALL_SLOT_KEY: recalled_id} + adapter._processing_msg_texts = {} # blank: no redaction polling for a notice under test + RecallGuardMiddleware._interrupt_for_recall(adapter, RECALL_SLOT_KEY, recalled_id, "G", "") + return adapter._pending_messages[RECALL_SLOT_KEY] + + +@pytest.mark.asyncio +async def test_a_recall_notice_with_a_partial_source_keeps_its_slot_key(tmp_path, flush_dir): + """The notice keeps the key it was queued under, and recovery lands it in that session.""" + db = _FakeGatewayDB({RECALL_SLOT_KEY: {"id": "sid-recall", "started_at": 0.0}}) + _runner, adapter, store = _wire_aligned(tmp_path, db) + notice = _queue_recall_notice(adapter) + assert notice.source.user_id is None + assert store._generate_session_key(notice.source) == RECALL_PARTIAL_KEY + + await adapter.cancel_background_tasks() + + assert _payload_keys(flush_dir) == [RECALL_SLOT_KEY] + recovered = recover_pending_to_db(db, session_resolver=store.resolve_session_id_for_key) + assert recovered == 1 + assert [call["session_id"] for call in db.appended] == ["sid-recall"] + + +# --- The two directions the keying rule must not get wrong -------------------------------- +# +# A key the store can resolve is not by itself proof that the payload belongs to that session, +# and a lookup the store cannot answer is not proof that the key is wrong. Both directions are +# covered here as the paired regressions of the rule in ``_recoverable_key``. +SHARED_SLOT_KEY = "agent:main:telegram:group:900" +SENDER_KEY = "agent:main:telegram:group:900:u1" + + +def _group_event() -> MessageEvent: + """A turn in a group chat, no thread: only the per-user flag separates the two keys.""" + return MessageEvent( + text="queued while busy", + message_type=MessageType.TEXT, + source=make_restart_source(chat_id="900", chat_type="group"), + message_id="queued-2", + ) + + +def _wire_shared_slot(tmp_path, db): + """The inverse override: the platform shares group keys, the store isolates them.""" + runner, adapter = make_restart_runner() + store = _store(tmp_path, db) + runner.session_store = store + adapter.set_session_store(store) + adapter.config.extra["group_sessions_per_user"] = False + return runner, adapter, store + + +@pytest.mark.asyncio +async def test_a_resolvable_slot_key_naming_another_session_does_not_capture_the_payload( + tmp_path, flush_dir): + """The payload is recorded under its own sender's key, not the userless session's. + + Both keys have a session behind them here: a callback-shaped turn already created one under + the shared slot key, and the payload's sender has their own. Resolving the slot key then + says only that *a* session exists under that string — the payload's source carries the + participant, so the store's own derivation is the session its turn is admitted in. + """ + db = _FakeGatewayDB({SHARED_SLOT_KEY: {"id": "sid-userless", "started_at": 0.0}, + SENDER_KEY: {"id": "sid-sender", "started_at": 0.0}}) + _runner, adapter, store = _wire_shared_slot(tmp_path, db) + event = _group_event() + assert adapter._event_session_key(event) == SHARED_SLOT_KEY + assert store._generate_session_key(event.source) == SENDER_KEY + adapter._pending_messages[SHARED_SLOT_KEY] = event + + await adapter.cancel_background_tasks() + + assert _payload_keys(flush_dir) == [SENDER_KEY] + recovered = recover_pending_to_db(db, session_resolver=store.resolve_session_id_for_key) + assert recovered == 1 + assert [call["session_id"] for call in db.appended] == ["sid-sender"] + assert _spooled_files(flush_dir) == [] + + +@pytest.mark.asyncio +async def test_an_unavailable_lookup_does_not_move_the_notice_off_its_slot_key( + tmp_path, flush_dir): + """A store that cannot answer for a key leaves the notice where its producer queued it. + + ``SessionStore._db_for_key`` fails closed (no reachable home for the key), and the resolver + then answers ``None`` for both keys — a miss it never reported. Reading that as "the slot key + is wrong" re-keys the notice onto a derivation the store cannot land either, and the notice + strands with the spool file kept. The lookup comes back before recovery runs, which is the + state the flush must not destroy. + """ + db = _FakeGatewayDB({RECALL_SLOT_KEY: {"id": "sid-recall", "started_at": 0.0}}) + _runner, adapter, store = _wire_aligned(tmp_path, db) + notice = _queue_recall_notice(adapter) + + with patch.object(store, "_db_for_key", lambda *_args, **_kwargs: None): + await adapter.cancel_background_tasks() + + assert _payload_keys(flush_dir) == [RECALL_SLOT_KEY] + recovered = recover_pending_to_db(db, session_resolver=store.resolve_session_id_for_key) + assert recovered == 1 + assert [call["session_id"] for call in db.appended] == ["sid-recall"] + assert [call["content"] for call in db.appended] == [notice.text] + assert _spooled_files(flush_dir) == [] + + +@pytest.mark.asyncio +async def test_the_runner_flush_keeps_a_routed_slot_key_too(tmp_path, flush_dir): + """``runner.stop()`` keys its queue head the same way: routed slot key, not a re-derivation.""" + db = _FakeGatewayDB({RECALL_SLOT_KEY: {"id": "sid-recall", "started_at": 0.0}}) + runner, adapter, store = _wire_aligned(tmp_path, db) + notice = _queue_recall_notice(adapter) + adapter._pending_messages.pop(RECALL_SLOT_KEY) # moved into the runner's queue head + runner.session_store = store + runner._pending_messages[RECALL_SLOT_KEY] = notice + + with patch("gateway.status.remove_pid_file"), patch( + "gateway.status.publish_runtime_status" + ): + await runner.stop() + + assert _payload_keys(flush_dir) == [RECALL_SLOT_KEY] + recovered = recover_pending_to_db(db, session_resolver=store.resolve_session_id_for_key) + assert recovered == 1 + assert [call["session_id"] for call in db.appended] == ["sid-recall"]