diff --git a/gateway/run.py b/gateway/run.py index f59aa5dec796e..a98d8297b7dfa 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -1651,35 +1651,6 @@ def _multiplex_profile_homes(config: object) -> list[tuple[str, "Path"]]: return list(profiles_to_serve(multiplex=True)) -def _recover_pending_flushes(runner) -> int: - """Replay every ``pending_messages`` spool this gateway owns into state.db; return the count. - - ``_get_flush_dir`` follows the active HERMES_HOME, so a routed turn on a multiplexed gateway spools - its stalled transcript backlog under ``profiles//`` and the runtime drain forgets it on - restart. After the launch home, replay each served profile inside its own home so the default - store ``recover_pending_to_db`` opens is that profile's state.db (#123584). - """ - from gateway.shutdown_flush import recover_pending_to_db - from hermes_constants import get_hermes_home, reset_hermes_home_override, set_hermes_home_override - - resolver = runner.session_store.resolve_session_id_for_key - recovered = recover_pending_to_db(session_resolver=resolver) - if not getattr(runner.config, "multiplex_profiles", False): - return recovered - launch_home = Path(get_hermes_home()).resolve() - for name, home in _multiplex_profile_homes(runner.config): - if Path(home).resolve() == launch_home or not (Path(home) / "pending_messages").is_dir(): - continue - token = set_hermes_home_override(str(home)) - try: - recovered += recover_pending_to_db(session_resolver=resolver) - except Exception: # one profile's unreadable spool must not strand the others' - logger.warning("Pending-message recovery failed for profile %s", name, exc_info=True) - finally: - reset_hermes_home_override(token) - return recovered - - def _cron_tick_profile_homes(config: object) -> list[tuple[str, "Path"]]: """Profile homes the in-process ticker visits: the served set PLUS the process-active profile: ``profiles_to_serve`` lists default + every live named profile, but a ``--profile @@ -5968,13 +5939,6 @@ def _start_keepalive() -> None: if not success: _shutdown_gateway_health_export(runner) return False - - def _recover_pending() -> None: - recovered = _recover_pending_flushes(runner) - if recovered: - logger.info("Recovered %d pending message(s) from shutdown flush", recovered) - - _best_effort(_recover_pending) if runner.should_exit_cleanly: _shutdown_gateway_health_export(runner) if runner.exit_reason: diff --git a/gateway/run_startup.py b/gateway/run_startup.py index a20ba923f33cd..94b941728881a 100644 --- a/gateway/run_startup.py +++ b/gateway/run_startup.py @@ -1508,6 +1508,7 @@ async def _start_post_connect_services(self, connected_count: int) -> None: async def _start_finish_wiring(self, connected_count: int) -> None: """Post-connect wiring: services, boot notifications, startup restore, recovered watchers.""" from gateway.run import _planned_restart_notification_pending, _restart_notification_pending + from gateway.shutdown_flush import recover_gateway_pending await self._start_post_connect_services(connected_count) # Let fresh adapters settle before lifecycle sends (helps Discord thread deliveries). if connected_count > 0: @@ -1523,6 +1524,11 @@ async def _start_finish_wiring(self, connected_count: int) -> None: await self._await_startup_boot_sends( planned_restart_notification_pending=_planned_restart_notification_pending(), ) + # Replay the previous run's pending spool before resume turns and queued inbound write live + # rows: only this pass tells the store which sessions have spooled rows, and a live row + # written first would get a lower row id than them for good. + with _log_suppressed(logging.WARNING, "Pending-message recovery failed: %s"): + recover_gateway_pending(self) # Auto-resume restart-interrupted sessions (ledger-answered ones were cleared above); a failed # auto-resume stays visible on the next user message. self._schedule_resume_pending_sessions() diff --git a/gateway/session_transcript.py b/gateway/session_transcript.py index 1d3c9f5e05610..9d33ba85b1cdc 100644 --- a/gateway/session_transcript.py +++ b/gateway/session_transcript.py @@ -5,6 +5,7 @@ from __future__ import annotations import contextlib +import json import logging import threading import time @@ -41,6 +42,55 @@ def _spool_dropped(session_id: str, message: Dict[str, Any]): "reasoning", "reasoning_content", "reasoning_details", "codex_reasoning_items", "codex_message_items", ) +# Message keys bound straight to a TEXT column. A spooled message is JSON from disk, and a dict or list +# here makes sqlite refuse the whole row. +_BOUND_TEXT_KEYS = ("role", "tool_name", "tool_call_id", "reasoning", "reasoning_content", "platform_message_id") + + +def _bindable_text(value: Any) -> Any: + """*value* when sqlite can bind it as is, else its JSON text (``str`` if that fails too). An int + outside sqlite's signed 64-bit range raises OverflowError even for a TEXT column.""" + if value is None or isinstance(value, (str, float)) or ( + isinstance(value, int) and -2**63 <= value < 2**63): + return value + try: + return json.dumps(value, default=str) + except (TypeError, ValueError): + return str(value) + + +def transcript_append_kwargs(session_id: str, message: Dict[str, Any], *, fallback_ts: Any = None) -> Dict[str, Any]: + """``SessionDB.append_message`` kwargs for one transcript row. The live writer and the restart + replay of the transcript spool (gateway/shutdown_flush.py) both build rows here, so a replayed + message lands as the row the live drain would have written. Fields are listed rather than + splatted: a spooled message is arbitrary JSON from disk, and an unexpected key would raise and + abort the replay. *fallback_ts* applies only to a message with no timestamp at all, because a + truthiness test would rewrite epoch 0. Values bound straight to a TEXT column are made bindable + (:func:`_bindable_text`), so one malformed field cannot make the row unwritable.""" + is_assistant = message.get("role") == "assistant" + timestamp = message.get("timestamp") + kwargs = { + "session_id": session_id, + "role": message.get("role", "unknown"), + "content": message.get("content"), + "tool_name": message.get("tool_name"), + "tool_calls": message.get("tool_calls"), + "tool_call_id": message.get("tool_call_id"), + **{k: message.get(k) if is_assistant else None for k in _ASSISTANT_ONLY_KEYS}, + "platform_message_id": message.get("platform_message_id") or message.get("message_id"), + "observed": bool(message.get("observed")), + "timestamp": fallback_ts if timestamp is None else timestamp, + # Exact bytes sent to the API (prompt-cache-stable replay); must survive every + # persistence path or the next replay diverges. + "api_content": extract_api_content_sidecar(message), + # Presentation typing ("internal_notification" for self-injected async-delegation/background + # notification turns, #82888). DB-only; stripped from provider-bound payloads. + "display_kind": message.get("display_kind"), + "display_metadata": message.get("display_metadata"), + } + for key in _BOUND_TEXT_KEYS: + kwargs[key] = _bindable_text(kwargs[key]) + return kwargs class SessionTranscriptMixin: @@ -244,10 +294,15 @@ def _ack_head() -> bool: # DB write outside the retry lock so other sessions can append. while True: - # Spooled backlog (cap eviction or a stalled session) is older than ``msg``: replay it - # first so recovery keeps transcript order; a still-dead DB just fails both. - self._drain_spooled_drops(session_id) + spool_exc = None try: + # Spooled backlog (cap eviction, a stalled session, a boot hold-back) is older than + # ``msg``: replay it first. While any of it stays on disk ``msg`` stays queued, since + # writing it now would give it a lower row id than that backlog for good. The replay's + # own error is raised so the repair/divert handling below still classifies it. + spool_exc = self._drain_spooled_drops(session_id) + if spool_exc is not None: + raise spool_exc self._append_transcript_message(session_id, msg) except Exception as exc: from hermes_state import StateDbCorruptError, StateDbReplacedError @@ -289,6 +344,8 @@ def _ack_head() -> bool: "no unique live child; not retrying", session_id) return if self._is_fts_corruption_error(exc) and self._rebuild_fts_once(): + if spool_exc is not None: + continue # repaired: drain the older spool again before ``msg`` try: self._append_transcript_message(session_id, msg) except Exception as retry_exc: @@ -334,19 +391,29 @@ def _spool_stalled_backlog(self, session_id: str, queue_session_id: str) -> int: break spooled += 1 if spooled: - self._lazy("_spooled_drop_sessions", set).add(session_id) with self._transcript_retry_lock: + self._lazy("_spooled_drop_sessions", set).add(session_id) del pending[:spooled] if not pending: self._dirty_transcripts.pop(queue_session_id, None) return spooled - def _drain_spooled_drops(self, session_id: str) -> None: - """Replay cap-dropped spooled transcript messages after DB recovery. Best-effort: replay - failures keep the spool files for the next successful flush; nothing here may raise.""" + def _drain_spooled_drops(self, session_id: str) -> Exception | None: + """Replay cap-dropped spooled transcript messages after DB recovery; return the replay + failure while some of them are still on disk, else None. Best-effort: replay failures keep + the spool files for the next successful flush; nothing here may raise.""" spooled_sessions = getattr(self, "_spooled_drop_sessions", None) if not spooled_sessions or session_id not in spooled_sessions: - return + return None + failures: list[Exception] = [] + + def replay(message: Dict[str, Any]) -> None: + try: + self._append_transcript_message(session_id, message) + except Exception as exc: + failures.append(exc) + raise + try: from gateway.shutdown_flush import drain_transcript_spool # Inside an outage the append that follows logs/escalates the same failure; the @@ -354,13 +421,24 @@ def _drain_spooled_drops(self, session_id: str) -> None: with self._transcript_retry_lock: known_failing = bool(self._transcript_append_failures.get(session_id)) _replayed, remaining = drain_transcript_spool( - session_id, lambda message: self._append_transcript_message(session_id, message), - db_known_failing=known_failing, - ) + session_id, replay, db_known_failing=known_failing) if not remaining: - spooled_sessions.discard(session_id) + with self._transcript_retry_lock: + spooled_sessions.discard(session_id) + return None + return failures[-1] if failures else RuntimeError( + f"older spooled transcript rows for {session_id} still pending") except Exception as exc: logger.warning("Failed to drain transcript spool for %s: %s", session_id, exc) + return None + + def mark_spooled_drop_sessions(self, session_ids) -> None: + """Have each session's next transcript write drain the on-disk spool first. Boot recovery marks + the sessions it held back, so their next live row cannot land ahead of those files.""" + if not session_ids: + return + with self._transcript_retry_lock: + self._lazy("_spooled_drop_sessions", set).update(session_ids) def _append_transcript_message(self, session_id: str, message: Dict[str, Any]) -> None: """Write one transcript row. Caller handles retry queuing.""" @@ -370,27 +448,7 @@ def _append_transcript_message(self, session_id: str, message: Dict[str, Any]) - # into the ambient store. raise RuntimeError( f"no owning session store for {session_id}; deferring transcript write") - is_assistant = message.get("role") == "assistant" - _db.append_message( - session_id=session_id, - role=message.get("role", "unknown"), - content=message.get("content"), - tool_name=message.get("tool_name"), - tool_calls=message.get("tool_calls"), - tool_call_id=message.get("tool_call_id"), - **{k: message.get(k) if is_assistant else None for k in _ASSISTANT_ONLY_KEYS}, - platform_message_id=(message.get("platform_message_id") or message.get("message_id")), - observed=bool(message.get("observed")), - timestamp=message.get("timestamp"), - # Exact bytes sent to the API (prompt-cache-stable replay); must survive every - # persistence path or the next replay diverges. - api_content=extract_api_content_sidecar(message), - # Presentation typing (e.g. "internal_notification"); DB-only. - # "internal_notification" for self-injected async-delegation/background notification turns, - # #82888). DB-only; stripped from provider-bound payloads. - display_kind=message.get("display_kind"), - display_metadata=message.get("display_metadata"), - ) + _db.append_message(**transcript_append_kwargs(session_id, message)) @staticmethod def _is_fts_corruption_error(exc: Exception) -> bool: diff --git a/gateway/shutdown_flush.py b/gateway/shutdown_flush.py index 74af44744f5a6..e395fc5bd9979 100644 --- a/gateway/shutdown_flush.py +++ b/gateway/shutdown_flush.py @@ -4,7 +4,7 @@ ``agent._session_messages`` are the only surviving copies; shutdown ``.clear()`` would drop them. All hooks write atomic JSON payloads under ``/pending_messages/``: ``flush_pending_to_file`` / ``flush_overflow_to_file`` (queue head / FIFO tail, before clear), -``recover_pending_to_db`` (after ``runner.start()``; replays via ``SessionDB.append_message``, +``recover_pending_to_db`` (at startup, before resume turns; replays via ``SessionDB.append_message``, deletes each file on success), ``flush_agent_history_to_file`` (DB flush raised), ``spool_dropped_transcript_message`` / ``drain_transcript_spool``. """ @@ -15,7 +15,9 @@ import itertools import json import logging +import math import os +import sqlite3 import time import uuid from pathlib import Path @@ -27,8 +29,15 @@ # operation. Payloads carry the full transcript message dict for verbatim replay. # See #78182. TRANSCRIPT_CAP_DROP_REASON = "transcript_cap_drop" +# Reason tag for agent-history snapshots: kept for manual operator recovery, never replayed. +AGENT_HISTORY_REASON = "shutdown-with-unpersisted-agent-history" +# Suffix for a spool file whose row the database rejected; no replay glob matches it. +QUARANTINE_SUFFIX = ".bad" # Monotonic tiebreaker so same-second spool files replay in drop order. _TRANSCRIPT_SPOOL_SEQ = itertools.count() +# Spool files whose row was written but which could not be deleted. Replaying one again would write +# its row twice, so both replay passes skip it for the rest of this process. +_REPLAYED_UNREMOVABLE: set[Path] = set() def _get_flush_dir(): @@ -136,9 +145,11 @@ def drain_transcript_spool(session_id: str, replay, *, db_known_failing: bool = """Replay cap-dropped transcript messages spooled for *session_id*; return ``(replayed, remaining)``. ``replay(message_dict)`` runs per message in drop order; a spool file is deleted only after its replay succeeds. The first failure stops the drain (the DB is likely still - unhealthy) and keeps the rest for retry. With ``db_known_failing`` (the caller's last write - already failed and is being logged/escalated) a replay failure is expected and logs at DEBUG, - so a stalled session does not add one WARNING per append on top of its ERROR (#114266). + unhealthy) and keeps the rest for retry, unless the DB rejected that row's own values + (:func:`is_row_rejection`): it is quarantined and the drain goes on. With ``db_known_failing`` + (the caller's last write already failed and is being logged/escalated) a replay failure is + expected and logs at DEBUG, so a stalled session does not add one WARNING per append on top of + its ERROR (#114266). """ try: candidates = list(_get_flush_dir().glob("pending-*.json")) @@ -147,6 +158,8 @@ def drain_transcript_spool(session_id: str, replay, *, db_known_failing: bool = return 0, 0 entries = [] for path in candidates: + if path in _REPLAYED_UNREMOVABLE: + continue try: payload = json.loads(path.read_text(encoding="utf-8-sig")) except Exception: @@ -162,18 +175,21 @@ def drain_transcript_spool(session_id: str, replay, *, db_known_failing: bool = logger.warning("Removing structurally invalid transcript spool file %s", path) path.unlink(missing_ok=True) continue - entries.append((payload.get("ts", 0), payload.get("seq", 0), path.name, path, message)) - ordered, replayed, remaining = sorted(entries, key=lambda e: e[:3]), 0, 0 - for idx, (_ts, _seq, _name, path, message) in enumerate(ordered): + entries.append((_spool_sort_key(payload, path.name), path, message)) + ordered, replayed, remaining = sorted(entries, key=lambda e: e[0]), 0, 0 + for idx, (_key, path, message) in enumerate(ordered): try: replay(message) except Exception as exc: + if is_row_rejection(exc): # no retry can write it, so it must not hold the session back + _quarantine_spool_file(path, session_id, exc) + continue (logger.debug if db_known_failing else logger.warning)( "Replay of spooled transcript message %s for %s failed; " "keeping spool file for retry: %s", path, session_id, exc) remaining = len(ordered) - idx break - path.unlink(missing_ok=True) + _remove_replayed(path, session_id) replayed += 1 if replayed: logger.info("Replayed %d spooled transcript message(s) for %s after DB recovery", replayed, @@ -206,7 +222,104 @@ def _serialise_value(value: Any) -> Optional[dict]: return {"text": str(value)} +def _sort_number(value: Any) -> float: + """Coerce a spool ordering field to a float; unusable values sort first. + + Ordering fields are read back from JSON on disk and may be missing or + corrupt. A sort key that mixes ``str`` and ``int`` raises ``TypeError`` + and would abort the entire recovery pass, so anything non-numeric is + normalised to ``0.0``. So is a JSON integer too large for a float + (``float(10**400)`` raises ``OverflowError``) and a NaN/Infinity literal + (NaN compares false both ways and would scramble the sort). + """ + if isinstance(value, bool) or not isinstance(value, (int, float)): + return 0.0 + try: + number = float(value) + except OverflowError: + return 0.0 + return number if math.isfinite(number) else 0.0 + + +def _spool_sort_key(payload: Dict[str, Any], name: str) -> tuple: + """Drop-order key ``(ts, seq, filename)`` for one spool payload. The live drain and the restart + pass both sort on it, so they replay a session's files in the same order.""" + # seq is per process (_TRANSCRIPT_SPOOL_SEQ) and ts has one-second resolution, so files two + # processes spool in the same second can interleave; within one process the order is exact. + # Slot heads have no seq and overflow tails start at zero, so a missing seq sorts as -1. + return _sort_number(payload.get("ts")), _sort_number(payload.get("seq", -1)), name + + +def is_row_rejection(exc: BaseException) -> bool: + """True when the database refused this row's own values, so no retry can succeed: a value sqlite + cannot bind (an unsupported type, or an int outside 64 bits), or a malformed field that fails + before the write. Lock, I/O, corruption and routing failures stay retryable, and so do constraint + failures, which can be session-wide. Only a binding message counts among sqlite errors: + ``InterfaceError("no more rows available")`` is WAL contention that ``SessionDB`` retries and + re-raises unchanged once its patience runs out.""" + if isinstance(exc, sqlite3.Error): + return (isinstance(exc, (sqlite3.InterfaceError, sqlite3.ProgrammingError)) + and "binding parameter" in str(exc).lower()) + return isinstance(exc, (TypeError, ValueError, AttributeError, OverflowError)) + + +def _remove_replayed(path: Path, session_id: str) -> None: + """Delete a spool file whose row was just written. If it cannot be deleted, this process never + replays it again, and replay goes on with the files after it, which are still older than any + live row. Nothing on disk marks it, so a restart writes the row a second time.""" + try: + path.unlink(missing_ok=True) + except OSError as exc: + _REPLAYED_UNREMOVABLE.add(path) + logger.error("Spooled transcript row for %s was written, but %s could not be deleted (%s); " + "delete it by hand, or the next start writes that row again", session_id, path, exc) + + +def _quarantine_spool_file(path: Path, session_id: str, exc: BaseException) -> None: + """Move a spool file whose row the database rejected out of every replay, so it cannot hold its + session's later rows back on every write and every restart. It stays on disk for manual recovery.""" + target = path.with_name(path.name + QUARANTINE_SUFFIX) + try: + path.replace(target) + except OSError as move_exc: + logger.error("Spooled transcript row for %s in %s was rejected by the database (%s) and could " + "not be quarantined: %s", session_id, path, exc, move_exc) + return + logger.warning("Spooled transcript row for %s was rejected by the database (%s); moved it to %s " + "for manual recovery and continuing with the session's later rows", + session_id, exc, target) + + +def _order_flush_files(paths) -> list[Path]: + """Order recovery payloads by :func:`_spool_sort_key`, as :func:`drain_transcript_spool` does. + ``SessionDB`` restores a conversation by AUTOINCREMENT id and spool files have random names, so + any other order permanently scrambles the recovered transcript. + + Keeps only the key and path, not decoded payloads; the caller re-reads each file to replay it. + Unparseable files sort last by name. + """ + entries = [] + for path in paths: + try: + payload = json.loads(path.read_text(encoding="utf-8-sig")) + # OSError: unreadable; ValueError: bad JSON or bytes; RecursionError: nesting too deep. + except (OSError, ValueError, RecursionError): + payload = None + if isinstance(payload, dict): + entries.append(((0, *_spool_sort_key(payload, path.name)), path)) + else: + entries.append(((1, 0.0, 0.0, path.name), path)) + entries.sort(key=lambda entry: entry[0]) + return [path for _key, path in entries] + + def recover_pending_to_db(session_db=None, *, session_resolver=None) -> int: + """Replay flush-dir ``*.json`` files into state.db; return the number of messages recovered. + See :func:`recover_pending_spool`, which also reports the sessions it held back.""" + return recover_pending_spool(session_db, session_resolver=session_resolver)[0] + + +def recover_pending_spool(session_db=None, *, session_resolver=None) -> tuple[int, set[str]]: """Replay flush-dir ``*.json`` files via ``SessionDB.append_message``, deleting each on success. ``session_db=None`` opens (and afterwards releases) the shared default ``state.db``. @@ -214,32 +327,40 @@ def recover_pending_to_db(session_db=None, *, session_resolver=None) -> int: ``SessionStore.resolve_session_id_for_key``) is required for real flush files: adapter ``MessageEvent`` objects carry no ``session_id``, so without it every recovery lands in the skip branch. A returned ``db`` routes the append to the profile store owning the key (multiplexed - gateways); ``None`` falls back to ``session_db``. Returns the number of messages recovered. + gateways); ``None`` falls back to ``session_db``. Returns ``(recovered, held_back)``: the number of + messages recovered, and the session ids whose spool files this pass held back after a failed + replay, so the live writer can drain them before that session's next row. """ - flush_files = sorted(_get_flush_dir().glob("*.json")) + flush_files = _order_flush_files(_get_flush_dir().glob("*.json")) if not flush_files: - return 0 + return 0, set() own_db = session_db is None if own_db: from hermes_state_registry import acquire session_db = acquire() recovered = 0 + # Sessions whose spool replay already failed this pass -> files held back for them. Ordering is a + # per-session property, so one unhealthy session must not hold back the others. + blocked_sessions: Dict[str, int] = {} try: for path in flush_files: + if path in _REPLAYED_UNREMOVABLE: + continue # One unparseable payload or rejected append must only skip THIS file: the file is # never unlinked, so aborting the pass would re-poison every later boot. - # utf-8-sig: our BOM-tolerant read fix for flush files. try: payload = json.loads(path.read_text(encoding="utf-8-sig")) # Agent-history snapshots use a different schema (reason + # messages list) and are meant for manual operator recovery, # not automatic DB insertion. Skip them silently. - if payload.get("reason") == "shutdown-with-unpersisted-agent-history": + if payload.get("reason") == AGENT_HISTORY_REASON: continue if _recover_one_payload(session_db, path, payload, - session_resolver=session_resolver): + session_resolver=session_resolver, + blocked_sessions=blocked_sessions): recovered += 1 - path.unlink(missing_ok=True) + _remove_replayed(path, payload.get("session_key", "")) + # health: allow BLE001 -- per-file boundary moved unchanged from recover_pending_to_db; a traceback per locked-DB file would only add noise except Exception as exc: logger.warning("Failed to recover pending message from %s: %s", path, exc) finally: @@ -249,25 +370,44 @@ def recover_pending_to_db(session_db=None, *, session_resolver=None) -> int: release_or_close(session_db) if recovered: logger.info("Recovered %d pending message(s) from shutdown flush", recovered) - return recovered + if blocked_sessions: + # Otherwise a spool dir that never empties gives an operator no reason why. + logger.info("Held back %d spooled transcript file(s) for the next start after a failed replay: %s", + sum(blocked_sessions.values()), + ", ".join(f"{sid} ({count})" for sid, count in blocked_sessions.items())) + return recovered, set(blocked_sessions) def _recover_one_payload(session_db, path: Path, payload: Dict[str, Any], *, - session_resolver=None) -> bool: - """Append one flush payload to ``session_db``; False (file kept) when structurally invalid.""" - # Cap-dropped transcript payloads carry the full message dict keyed by session_id — replay directly - # (#78182). This handles spool files that were never drained before a restart. + session_resolver=None, blocked_sessions: Dict[str, int]) -> bool: + """Append one flush payload to ``session_db``; False when it was not replayed (the file is kept, + or quarantined when the database rejected the row itself).""" + # Cap-dropped transcript payloads carry the full message dict keyed by session_id: replay it + # directly (#78182). This handles spool files that were never drained before a restart. if payload.get("reason") == TRANSCRIPT_CAP_DROP_REASON: - # Cap-dropped payloads carry the full message dict keyed by session_id — replay directly. + from gateway.session_transcript import transcript_append_kwargs data = payload.get("data", {}) or {} spooled_sid, message = data.get("session_id", ""), data.get("message") if not spooled_sid or not isinstance(message, dict): logger.warning("Cannot recover structurally invalid transcript spool " "file %s; preserved for manual inspection", path) return False - session_db.append_message(session_id=spooled_sid, role=message.get("role", "unknown"), - content=message.get("content") or "", - timestamp=message.get("timestamp") or payload.get("ts")) + if spooled_sid in blocked_sessions: + # An older message for this session could not be replayed. Writing this one now would give + # it a lower row id than the message it follows, permanently inverting the transcript, so + # leave it for the next start. + blocked_sessions[spooled_sid] += 1 + return False + try: + session_db.append_message(**transcript_append_kwargs(spooled_sid, message, fallback_ts=payload.get("ts"))) + except Exception as exc: + if is_row_rejection(exc): # no retry can write it, so it must not hold the session back + _quarantine_spool_file(path, spooled_sid, exc) + return False + # Same contract as drain_transcript_spool: stop this session's replay on the first failure + # and keep the remaining spool files for the next attempt. + blocked_sessions[spooled_sid] = 1 + raise return True session_key, data = payload.get("session_key", ""), payload.get("data", {}) text = data.get("text", "") @@ -299,6 +439,40 @@ def _recover_one_payload(session_db, path: Path, payload: Dict[str, Any], *, return True +def recover_gateway_pending(runner) -> int: + """Replay every ``pending_messages`` spool this gateway owns into state.db; return the count. + + ``_get_flush_dir`` follows the active HERMES_HOME, so a routed turn on a multiplexed gateway spools + its stalled transcript backlog under ``profiles//`` and the runtime drain forgets it on + restart. After the launch home, replay each served profile inside its own home so the default + store ``recover_pending_spool`` opens is that profile's state.db (#123584). The sessions a pass + holds back are marked on the session store, so their next live row drains the spool first. + """ + from gateway.run import _multiplex_profile_homes + from hermes_constants import get_hermes_home, reset_hermes_home_override, set_hermes_home_override + + store = runner.session_store + recovered, held_back = recover_pending_spool(session_resolver=store.resolve_session_id_for_key) + store.mark_spooled_drop_sessions(held_back) + if not getattr(runner.config, "multiplex_profiles", False): + return recovered + launch_home = Path(get_hermes_home()).resolve() + for name, home in _multiplex_profile_homes(runner.config): + if Path(home).resolve() == launch_home or not (Path(home) / "pending_messages").is_dir(): + continue + token = set_hermes_home_override(str(home)) + try: + count, held_back = recover_pending_spool(session_resolver=store.resolve_session_id_for_key) + except Exception: # one profile's unreadable spool must not strand the others' + logger.warning("Pending-message recovery failed for profile %s", name, exc_info=True) + continue + finally: + reset_hermes_home_override(token) + recovered += count + store.mark_spooled_drop_sessions(held_back) + return recovered + + def flush_agent_history_to_file(session_id: Optional[str], history: list) -> None: """Best-effort dump of an agent's in-memory transcript before teardown. Used when ``_flush_messages_to_session_db`` raises (e.g. FTS/SQLite corruption): the transcript is written @@ -316,7 +490,7 @@ def flush_agent_history_to_file(session_id: Optional[str], history: list) -> Non except Exception: continue _write_payload(flush_dir, { - "reason": "shutdown-with-unpersisted-agent-history", "issue": "#72680", + "reason": AGENT_HISTORY_REASON, "issue": "#72680", "session_id": session_id, "count": len(snapshot), "messages": snapshot, }) logger.warning("Preserved %d in-memory message(s) for session %s " diff --git a/tests/gateway/test_profile_spool_recovery.py b/tests/gateway/test_profile_spool_recovery.py index 0afa8d6d583fb..de0da4e4702e4 100644 --- a/tests/gateway/test_profile_spool_recovery.py +++ b/tests/gateway/test_profile_spool_recovery.py @@ -16,9 +16,8 @@ import agent.secret_scope as ss from gateway.config import GatewayConfig from gateway.platforms.base import Platform, SessionSource -from gateway.run import _recover_pending_flushes from gateway.session import SessionEntry, SessionStore -from gateway.shutdown_flush import spool_dropped_transcript_message +from gateway.shutdown_flush import recover_gateway_pending, spool_dropped_transcript_message from hermes_constants import get_hermes_home, reset_hermes_home_override, set_hermes_home_override @@ -73,7 +72,7 @@ def test_boot_recovery_replays_a_routed_profiles_spooled_backlog(multiplex_homes _spool_under(root, default_sid, "default while locked") _spool_under(profile, work_sid, "work while locked") - recovered = _recover_pending_flushes(SimpleNamespace(config=store.config, session_store=store)) + recovered = recover_gateway_pending(SimpleNamespace(config=store.config, session_store=store)) assert recovered == 2 assert [m["content"] for m in default_db.get_messages(default_sid)] == ["default while locked"] diff --git a/tests/gateway/test_shutdown_flush_recovery.py b/tests/gateway/test_shutdown_flush_recovery.py new file mode 100644 index 0000000000000..08f906f5af6d1 --- /dev/null +++ b/tests/gateway/test_shutdown_flush_recovery.py @@ -0,0 +1,471 @@ +"""Cross-restart recovery of cap-dropped transcript spool files (#78182). + +``recover_pending_to_db`` is the restart-time consumer of the same spool +``drain_transcript_spool`` drains during live operation. These tests pin the +properties the live drain already guarantees for that spool. +""" + +import json +import math +import sqlite3 +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +from gateway.shutdown_flush import ( + QUARANTINE_SUFFIX, + TRANSCRIPT_CAP_DROP_REASON, + _order_flush_files, + drain_transcript_spool, + is_row_rejection, + recover_gateway_pending, + recover_pending_to_db, +) + + +@pytest.fixture +def flush_dir(tmp_path, monkeypatch): + """A temp spool directory wired into the module under test.""" + directory = tmp_path / "pending_messages" + directory.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr( + "gateway.shutdown_flush._get_flush_dir", lambda: directory + ) + monkeypatch.setattr("gateway.shutdown_flush._REPLAYED_UNREMOVABLE", set()) + return directory + + +@pytest.fixture +def make_store(tmp_path): + """Build a real ``SessionStore`` whose transcript writes go to *db*. It goes through the + constructor, so new store state does not break these tests.""" + from gateway.config import GatewayConfig + from gateway.session import SessionStore + + def build(db): + with patch("gateway.session.SessionStore._ensure_loaded"): + store = SessionStore(sessions_dir=tmp_path / "sessions", config=GatewayConfig()) + store._loaded = True + store._db = db + return store + + return build + + +@pytest.fixture +def boot_db(monkeypatch): + """Make boot recovery's default state.db the given handle.""" + import hermes_state_registry + + def use(db): + monkeypatch.setattr(hermes_state_registry, "acquire", lambda: db) + monkeypatch.setattr(hermes_state_registry, "release_or_close", lambda _db: None) + + return use + + +def _boot(store) -> int: + runner = SimpleNamespace(config=SimpleNamespace(multiplex_profiles=False), session_store=store) + return recover_gateway_pending(runner) + + +def _write_spool( + flush_dir: Path, + name: str, + session_id: str, + message: dict, + *, + ts: Any, + seq: Any, +) -> Path: + """Write one cap-drop spool payload under an explicit file name. + + Production names these ``pending-.json``; the tests choose the + names so that filename order and drop order can be made to disagree. + """ + path = flush_dir / name + path.write_text( + json.dumps( + { + "session_key": session_id, + "reason": TRANSCRIPT_CAP_DROP_REASON, + "ts": ts, + "seq": seq, + "data": {"session_id": session_id, "message": message}, + } + ), + encoding="utf-8", + ) + return path + + +def _contents(mock_db) -> list: + return [c.kwargs["content"] for c in mock_db.append_message.call_args_list] + + +def _fail_unlink_for(monkeypatch, name: str) -> None: + """Make deleting the spool file *name* fail, as on a directory that stopped being writable.""" + real_unlink = Path.unlink + + def unlink(self, missing_ok=False): + if self.name == name: + raise PermissionError(13, "Permission denied", str(self)) + return real_unlink(self, missing_ok=missing_ok) + + monkeypatch.setattr(Path, "unlink", unlink) + + +def test_replays_in_drop_order_not_file_name_order(flush_dir): + """SessionDB restores by AUTOINCREMENT id, so append order IS the order the user sees + after recovery. Production names are ``pending-.json``; these names sort opposite + to drop order, and a burst inside one second shares ``ts``, so ``seq`` breaks the tie.""" + _write_spool(flush_dir, "pending-zzz.json", "sess-1", + {"role": "user", "content": "first"}, ts=100, seq=0) + _write_spool(flush_dir, "pending-mmm.json", "sess-1", + {"role": "assistant", "content": "second"}, ts=100, seq=1) + _write_spool(flush_dir, "pending-aaa.json", "sess-1", + {"role": "user", "content": "third"}, ts=101, seq=2) + + mock_db = MagicMock() + assert recover_pending_to_db(mock_db) == 3 + assert _contents(mock_db) == ["first", "second", "third"] + + +def test_an_out_of_range_ordering_field_does_not_abort_recovery(flush_dir): + """Ordering runs before the per-file error handling, so one file whose ``ts`` is a JSON + integer no float can hold must not stop every other file from being recovered.""" + _write_spool(flush_dir, "pending-bad.json", "sess-1", + {"role": "user", "content": "huge ts"}, ts=10**400, seq=0) + _write_spool(flush_dir, "pending-ok.json", "sess-2", + {"role": "user", "content": "valid"}, ts=100, seq=1) + + mock_db = MagicMock() + assert recover_pending_to_db(mock_db) == 2 + assert sorted(_contents(mock_db)) == ["huge ts", "valid"] + + +def test_live_drain_and_restart_pass_order_malformed_files_alike(flush_dir): + """Every live write waits on the live drain, so a ``ts`` it cannot compare must not raise there, + and both passes must replay a session's files in the same order.""" + _write_spool(flush_dir, "pending-a.json", "sess-1", + {"role": "user", "content": "string ts"}, ts="later", seq=0) + _write_spool(flush_dir, "pending-b.json", "sess-1", + {"role": "user", "content": "nan ts"}, ts=math.nan, seq=1) + _write_spool(flush_dir, "pending-c.json", "sess-1", + {"role": "user", "content": "ts 5"}, ts=5, seq=2) + + restart = [json.loads(path.read_text())["data"]["message"]["content"] + for path in _order_flush_files(flush_dir.glob("*.json"))] + live = [] + assert drain_transcript_spool("sess-1", lambda message: live.append(message["content"])) == (3, 0) + + assert live == restart == ["string ts", "nan ts", "ts 5"] + + +def test_a_replayed_row_is_the_row_the_live_writer_writes(flush_dir, make_store): + """A replayed spool row must be the row the live writer would have written for the same + message: losing tool_call_id orphans a tool result, losing api_content makes the next replay + diverge. Only a message with no timestamp differs: it takes the payload clock, and epoch 0 is + a real timestamp.""" + tool_calls = [{"id": "call-1", "type": "function", + "function": {"name": "send_payment", "arguments": "{}"}}] + assistant = { + "role": "assistant", "content": None, "tool_calls": tool_calls, + "reasoning": "deliberating", "reasoning_content": "chain", + "reasoning_details": [{"type": "text"}], "codex_reasoning_items": [{"id": "r1"}], + "codex_message_items": [{"id": "m1"}], "platform_message_id": "tg-42", + "observed": True, "timestamp": 0, "api_content": "exact bytes sent to the API", + "display_kind": "internal_notification", + } + tool = {"role": "tool", "content": "receipt-1", "tool_call_id": "call-1", "tool_name": "send_payment"} + user = {"role": "user", "content": "hi", "reasoning": "leaked", "message_id": "tg-7"} + _write_spool(flush_dir, "pending-ccc.json", "sess-1", assistant, ts=100, seq=0) + _write_spool(flush_dir, "pending-bbb.json", "sess-1", tool, ts=100, seq=1) + _write_spool(flush_dir, "pending-aaa.json", "sess-1", user, ts=999, seq=2) + + recovery_db, live_db = MagicMock(), MagicMock() + assert recover_pending_to_db(recovery_db) == 3 + store = make_store(live_db) + for message in (assistant, tool, user): + store._append_transcript_message("sess-1", message) + + replayed = [c.kwargs for c in recovery_db.append_message.call_args_list] + written = [c.kwargs for c in live_db.append_message.call_args_list] + assert [row.pop("timestamp") for row in replayed] == [0, 100, 999] + assert [row.pop("timestamp") for row in written] == [0, None, None] + assert replayed == written + assert replayed[0]["tool_calls"] == tool_calls + assert replayed[0]["api_content"] == "exact bytes sent to the API" + assert (replayed[1]["tool_call_id"], replayed[1]["tool_name"]) == ("call-1", "send_payment") + assert (replayed[2]["reasoning"], replayed[2]["platform_message_id"]) == (None, "tg-7") + + +def test_a_failed_replay_holds_back_that_sessions_later_messages_only(flush_dir, caplog): + """Writing "second" after "first" failed would give it a lower row id than "first" once + "first" is retried on a later start: the inversion lands on disk for good. Another + session is unaffected, and the pass says what it held back.""" + first = _write_spool(flush_dir, "pending-aaa.json", "sess-1", + {"role": "user", "content": "first"}, ts=100, seq=0) + second = _write_spool(flush_dir, "pending-bbb.json", "sess-1", + {"role": "user", "content": "second"}, ts=101, seq=1) + other = _write_spool(flush_dir, "pending-ccc.json", "sess-2", + {"role": "user", "content": "other-session"}, ts=102, seq=2) + + mock_db = MagicMock() + + def append_message(**kwargs): + if kwargs["content"] == "first": + raise RuntimeError("controlled database outage") + return 1 + + mock_db.append_message.side_effect = append_message + + with caplog.at_level("INFO", logger="gateway.shutdown_flush"): + assert recover_pending_to_db(mock_db) == 1 + + assert _contents(mock_db) == ["first", "other-session"] + assert first.exists() and second.exists() + assert not other.exists() + assert "Held back 2 spooled transcript file(s)" in caplog.text and "sess-1 (2)" in caplog.text + + +def test_held_back_files_drain_before_that_sessions_next_live_write(flush_dir, make_store, boot_db): + """The live writer drains a session's spool only when the store knows it has one. If boot + recovery does not say which sessions it held back, the next live row lands ahead of them.""" + _write_spool(flush_dir, "pending-bbb.json", "sess-1", + {"role": "user", "content": "old0"}, ts=100, seq=0) + _write_spool(flush_dir, "pending-aaa.json", "sess-1", + {"role": "user", "content": "old1"}, ts=100, seq=1) + mock_db = MagicMock() + mock_db.append_message.side_effect = [RuntimeError("still locked at boot"), 1, 1, 1] + boot_db(mock_db) + store = make_store(mock_db) + + assert _boot(store) == 0 + store.append_to_transcript("sess-1", {"role": "user", "content": "new-live"}) + + assert _contents(mock_db) == ["old0", "old0", "old1", "new-live"] + assert not list(flush_dir.glob("*.json")) + + +def test_live_write_waits_while_held_back_files_still_fail_to_drain(flush_dir, make_store, boot_db): + """If the pre-write drain fails again, the live row must stay queued in memory: writing it + while the older files are still on disk gives it a lower row id than they will get.""" + _write_spool(flush_dir, "pending-bbb.json", "sess-1", + {"role": "user", "content": "old0"}, ts=100, seq=0) + _write_spool(flush_dir, "pending-aaa.json", "sess-1", + {"role": "user", "content": "old1"}, ts=100, seq=1) + mock_db = MagicMock() + mock_db.append_message.side_effect = [ + RuntimeError("still locked at boot"), RuntimeError("still locked at drain"), 1, 1, 1, 1] + boot_db(mock_db) + store = make_store(mock_db) + + assert _boot(store) == 0 + store.append_to_transcript("sess-1", {"role": "user", "content": "new-live"}) + assert _contents(mock_db) == ["old0", "old0"] + assert len(list(flush_dir.glob("*.json"))) == 2 + + store.append_to_transcript("sess-1", {"role": "user", "content": "new-live-2"}) + assert _contents(mock_db) == ["old0", "old0", "old0", "old1", "new-live", "new-live-2"] + assert not list(flush_dir.glob("*.json")) + assert not store._dirty_transcripts + + +def test_a_poisoned_spool_row_does_not_stop_later_live_rows(flush_dir, tmp_path, make_store, boot_db): + """A spooled field sqlite cannot bind used to fail on every replay. With live writes waiting + behind the spool, the session then stopped persisting for good, across restarts too.""" + from hermes_state import SessionDB + + db = SessionDB(db_path=tmp_path / "state.db") + try: + db.create_session("sess-1", source="telegram") + _write_spool(flush_dir, "pending-bbb.json", "sess-1", + {"role": "tool", "content": "result", "tool_call_id": {"bad": 1}}, ts=100, seq=0) + _write_spool(flush_dir, "pending-aaa.json", "sess-1", + {"role": "user", "content": "older"}, ts=100, seq=1) + # sqlite cannot bind an int past 64 bits, even into a TEXT column. + _write_spool(flush_dir, "pending-ccc.json", "sess-1", + {"role": "tool", "content": "big id", "tool_call_id": 2**100}, ts=100, seq=2) + boot_db(db) + store = make_store(db) + + assert _boot(store) == 3 + store.append_to_transcript("sess-1", {"role": "user", "content": "live"}) + + assert [m["content"] for m in db.get_messages("sess-1")] == ["result", "older", "big id", "live"] + assert not list(flush_dir.iterdir()) + assert not store._dirty_transcripts + finally: + db.close() + + +def test_a_row_the_database_rejects_is_quarantined_not_retried(flush_dir, make_store, boot_db): + """No retry can write a row the database refuses on its own values, so at boot and in the + live drain it is moved aside instead of holding back the session's later rows.""" + def append_message(**kwargs): + if kwargs["content"].startswith("poison"): + raise sqlite3.InterfaceError("Error binding parameter 3: unsupported type") + return 1 + + mock_db = MagicMock() + mock_db.append_message.side_effect = append_message + _write_spool(flush_dir, "pending-aaa.json", "sess-1", + {"role": "user", "content": "poison at boot"}, ts=100, seq=0) + _write_spool(flush_dir, "pending-bbb.json", "sess-1", + {"role": "user", "content": "after"}, ts=100, seq=1) + boot_db(mock_db) + store = make_store(mock_db) + assert _boot(store) == 1 + + _write_spool(flush_dir, "pending-ccc.json", "sess-1", + {"role": "user", "content": "poison live"}, ts=200, seq=2) + _write_spool(flush_dir, "pending-ddd.json", "sess-1", + {"role": "user", "content": "after live"}, ts=200, seq=3) + store.mark_spooled_drop_sessions({"sess-1"}) + store.append_to_transcript("sess-1", {"role": "user", "content": "live"}) + + assert _contents(mock_db) == ["poison at boot", "after", "poison live", "after live", "live"] + assert sorted(p.name for p in flush_dir.iterdir()) == [ + f"pending-aaa.json{QUARANTINE_SUFFIX}", f"pending-ccc.json{QUARANTINE_SUFFIX}"] + assert not store._dirty_transcripts + + +def test_only_errors_about_the_rows_own_values_are_rejections(): + """A rejection is quarantined for good, so it must be an error no retry can clear.""" + assert is_row_rejection(sqlite3.ProgrammingError("Error binding parameter 4: type 'dict' is not supported")) + assert is_row_rejection(sqlite3.InterfaceError("Error binding parameter 3 - probably unsupported type")) + assert is_row_rejection(OverflowError("Python int too large to convert to SQLite INTEGER")) + assert not is_row_rejection(sqlite3.InterfaceError("no more rows available")) + assert not is_row_rejection(sqlite3.OperationalError("database is locked")) + assert not is_row_rejection(sqlite3.IntegrityError("FOREIGN KEY constraint failed")) + assert not is_row_rejection(sqlite3.DatabaseError("database disk image is malformed")) + + +def test_transient_no_more_rows_error_keeps_the_row_for_retry(flush_dir, make_store, boot_db): + """SessionDB retries 'no more rows available' as WAL contention and re-raises it unchanged once + its patience runs out. The row itself is valid, so at boot and in the live drain it must stay on + disk and be written once the contention clears, not be quarantined.""" + transient = sqlite3.InterfaceError("no more rows available") + path = _write_spool(flush_dir, "pending-aaa.json", "sess-1", + {"role": "user", "content": "old"}, ts=100, seq=0) + mock_db = MagicMock() + mock_db.append_message.side_effect = [transient, transient, 1, 1, 1] + boot_db(mock_db) + store = make_store(mock_db) + + assert _boot(store) == 0 + assert path.exists() + store.append_to_transcript("sess-1", {"role": "user", "content": "live"}) + assert _contents(mock_db) == ["old", "old"] + assert path.exists() + + store.append_to_transcript("sess-1", {"role": "user", "content": "live-2"}) + assert _contents(mock_db) == ["old", "old", "old", "live", "live-2"] + assert not list(flush_dir.iterdir()) + assert not store._dirty_transcripts + + +def test_a_replayed_file_that_cannot_be_deleted_is_not_replayed_again(flush_dir, make_store, monkeypatch): + """Its row is already written, so replaying it before the next live write duplicates that row. + The files after it are still older than the live row, so they must land first.""" + _write_spool(flush_dir, "pending-aaa.json", "sess-1", + {"role": "user", "content": "old0"}, ts=100, seq=0) + second = _write_spool(flush_dir, "pending-bbb.json", "sess-1", + {"role": "user", "content": "old1"}, ts=100, seq=1) + _fail_unlink_for(monkeypatch, "pending-aaa.json") + mock_db = MagicMock() + store = make_store(mock_db) + store.mark_spooled_drop_sessions({"sess-1"}) + + store.append_to_transcript("sess-1", {"role": "user", "content": "live"}) + store.mark_spooled_drop_sessions({"sess-1"}) + store.append_to_transcript("sess-1", {"role": "user", "content": "live-2"}) + + assert _contents(mock_db) == ["old0", "old1", "live", "live-2"] + assert not second.exists() + assert not store._dirty_transcripts + + +def test_boot_recovery_counts_and_never_repeats_a_file_it_could_not_delete( + flush_dir, make_store, boot_db, monkeypatch): + """Boot recovery wrote the row, so it is recovered; the live drain must not write it again.""" + _write_spool(flush_dir, "pending-aaa.json", "sess-1", + {"role": "user", "content": "old0"}, ts=100, seq=0) + _write_spool(flush_dir, "pending-bbb.json", "sess-1", + {"role": "user", "content": "old1"}, ts=100, seq=1) + _fail_unlink_for(monkeypatch, "pending-aaa.json") + mock_db = MagicMock() + boot_db(mock_db) + store = make_store(mock_db) + + assert _boot(store) == 2 + store.mark_spooled_drop_sessions({"sess-1"}) + store.append_to_transcript("sess-1", {"role": "user", "content": "live"}) + + assert _contents(mock_db) == ["old0", "old1", "live"] + + +def test_repairable_spool_replay_failure_still_reaches_fts_rebuild(flush_dir, make_store): + """A spool replay failure must keep its own type, so the FTS rebuild still runs. After the + repair the older spooled row is written before the live one.""" + _write_spool(flush_dir, "pending-aaa.json", "sess-1", + {"role": "user", "content": "old"}, ts=100, seq=0) + + class RepairableDb: + def __init__(self): + self.broken, self.rows, self.repairs = True, [], 0 + + def append_message(self, **kwargs): + if self.broken: + raise RuntimeError("no such table: messages_fts") + self.rows.append(kwargs["content"]) + + def rebuild_fts(self): + self.repairs += 1 + self.broken = False + return 1 + + db = RepairableDb() + store = make_store(db) + store.mark_spooled_drop_sessions({"sess-1"}) + + store.append_to_transcript("sess-1", {"role": "user", "content": "live"}) + + assert db.repairs == 1 + assert db.rows == ["old", "live"] + assert not list(flush_dir.glob("*.json")) + assert not store._dirty_transcripts + + +def test_boot_recovery_runs_before_resume_turns_and_queued_inbound(monkeypatch): + """Resume turns and queued inbound write live rows. The store knows nothing about the previous + run's spool until recovery has run, so a live row written first lands ahead of it for good.""" + import asyncio + + import gateway.run as gateway_run + + order = [] + monkeypatch.setattr("gateway.shutdown_flush.recover_gateway_pending", + lambda runner: order.append("recover") or 0) + monkeypatch.setattr(gateway_run, "_restart_notification_pending", lambda: False) + monkeypatch.setattr(gateway_run, "_planned_restart_notification_pending", lambda: False) + + async def noop(*_args, **_kwargs): + return None + + async def finish_startup_restore(): + order.append("drain inbound") + + runner = SimpleNamespace( + _start_post_connect_services=noop, _await_startup_boot_sends=noop, + _schedule_resume_pending_sessions=lambda: order.append("resume"), + _finish_startup_restore=finish_startup_restore, + _send_session_db_warning_notifications=noop, _spawn_supervised=lambda *a, **k: None, + ) + asyncio.run(gateway_run.GatewayRunner._start_finish_wiring(runner, 0)) + + assert order == ["recover", "resume", "drain inbound"] diff --git a/tests/gateway/test_shutdown_spool_limits.py b/tests/gateway/test_shutdown_spool_limits.py new file mode 100644 index 0000000000000..125407d577418 --- /dev/null +++ b/tests/gateway/test_shutdown_spool_limits.py @@ -0,0 +1,54 @@ +"""Recovery must not retain every decoded spool payload while sorting.""" + +import tracemalloc +from pathlib import Path + +import pytest + +from gateway import shutdown_flush + + +def test_recovery_payload_memory_does_not_scale_with_backlog( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(shutdown_flush, "_get_flush_dir", lambda: tmp_path) + content = "x" * (256 * 1024) + for seq in range(24): + shutdown_flush.spool_dropped_transcript_message( + "sess", {"role": "user", "content": content, "timestamp": seq}, + ) + + class Sink: + def append_message(self, **kwargs: object) -> None: + assert kwargs["content"] == content + + tracemalloc.start() + try: + assert shutdown_flush.recover_pending_to_db(Sink()) == 24 + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + # Six MiB of payloads on disk; allow ample room for one decoded payload, + # its source text and the small ordering index, but not the entire backlog. + assert peak < 4 * 1024 * 1024, f"recovery retained {peak} bytes" + assert not list(tmp_path.glob("*.json")) + + +def test_same_second_shutdown_slot_precedes_overflow( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, +) -> None: + from types import SimpleNamespace + from unittest.mock import MagicMock + + monkeypatch.setattr(shutdown_flush, "_get_flush_dir", lambda: tmp_path) + monkeypatch.setattr(shutdown_flush.time, "time", lambda: 100) + names = iter(["zzz", "aaa", "bbb"]) + monkeypatch.setattr(shutdown_flush.uuid, "uuid4", lambda: SimpleNamespace(hex=next(names))) + assert shutdown_flush.flush_pending_to_file({"session": {"text": "A", "session_id": "sess"}}) == 1 + assert shutdown_flush.flush_overflow_to_file({"session": [ + {"text": "B", "session_id": "sess"}, {"text": "C", "session_id": "sess"}, + ]}) == 2 + db = MagicMock() + assert shutdown_flush.recover_pending_to_db(db) == 3 + assert [call.kwargs["content"] for call in db.append_message.call_args_list] == ["A", "B", "C"] + assert not list(tmp_path.glob("*.json"))