From 0ee8d418784619e5b9cf5aae3b5fdaed764fba4b Mon Sep 17 00:00:00 2001 From: Anthony Ruiz Date: Thu, 23 Jul 2026 10:49:40 +0000 Subject: [PATCH] fix(compression): recover rotated session lineage --- agent/conversation_compression.py | 329 +++++++++++++----- agent/turn_context.py | 8 + gateway/session.py | 98 +++++- hermes_state.py | 188 ++++++++++ run_agent.py | 3 + .../agent/test_compression_concurrent_fork.py | 90 +++++ .../agent/test_compression_rotation_state.py | 34 +- tests/agent/test_turn_context.py | 31 ++ tests/gateway/test_session.py | 101 ++++++ .../test_compression_boundary_hook.py | 12 +- tests/state/test_compression_lineage_guard.py | 196 +++++++++++ tests/test_hermes_state.py | 2 +- 12 files changed, 979 insertions(+), 113 deletions(-) create mode 100644 tests/state/test_compression_lineage_guard.py diff --git a/agent/conversation_compression.py b/agent/conversation_compression.py index 2a2ec1b0b4d1..81a839d1c57d 100644 --- a/agent/conversation_compression.py +++ b/agent/conversation_compression.py @@ -40,7 +40,7 @@ import uuid import threading from datetime import datetime from pathlib import Path -from typing import Any, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple from agent.context_engine import ( automatic_compaction_status_message, @@ -373,6 +373,142 @@ def compression_skipped_due_to_lock(agent: Any) -> bool: return _sig is True or isinstance(_sig, str) +def _durable_history_matches_snapshot( + durable: Any, snapshot: List[Dict[str, Any]] +) -> bool: + """Return whether durable parent rows are the ordered snapshot prefix.""" + if not isinstance(durable, list) or len(durable) > len(snapshot): + return False + identity_fields = ( + "role", "content", "tool_call_id", "tool_calls", "tool_name", "api_content", + ) + for stored, live in zip(durable, snapshot): + if not isinstance(stored, dict) or not isinstance(live, dict): + return False + if any(stored.get(key) != live.get(key) for key in identity_fields): + return False + return True + + +def _adopt_live_compression_child( + agent: Any, + session_db: Any, + parent_session_id: str, +) -> Optional[List[Dict[str, Any]]]: + """Move a stale compression contender onto the unique durable child. + + Resolve and load first, then mutate the live agent. This ordering keeps the + stale contender fail-closed when lineage is ambiguous or the compacted + handoff cannot be read. + """ + finder = getattr(type(session_db), "find_live_compression_child", None) + loader = getattr(type(session_db), "get_messages_as_conversation", None) + if not callable(finder) or not callable(loader): + return None + child = finder(session_db, parent_session_id) + if not child or not child.get("id"): + return None + child_session_id = str(child["id"]) + recovered = loader(session_db, child_session_id) + if not isinstance(recovered, list) or not recovered: + return None + # Revalidate after loading: the child may have rotated or a competing + # continuation may have appeared between the two DB reads. + confirmed = finder(session_db, parent_session_id) + if not confirmed or str(confirmed.get("id") or "") != child_session_id: + return None + + agent.session_id = child_session_id + try: + from gateway.session_context import set_current_session_id + + set_current_session_id(child_session_id) + except Exception: + os.environ["HERMES_SESSION_ID"] = child_session_id + try: + from hermes_logging import set_session_context + + set_session_context(child_session_id) + except Exception: + pass + + agent._session_db_created = True + if child.get("system_prompt"): + agent._cached_system_prompt = child["system_prompt"] + agent._last_flushed_db_idx = len(recovered) + agent._flushed_db_message_session_id = child_session_id + agent._flushed_db_message_ids = { + id(message) for message in recovered if isinstance(message, dict) + } + + on_session_start = getattr(agent.context_compressor, "on_session_start", None) + if callable(on_session_start): + try: + on_session_start( + child_session_id, + boundary_reason="compression", + old_session_id=parent_session_id, + session_db=session_db, + platform=getattr(agent, "platform", None) or "cli", + conversation_id=getattr(agent, "_gateway_session_key", None), + ) + except Exception as exc: + logger.debug("context engine compression-child adoption failed: %s", exc) + else: + bind_state = getattr(agent.context_compressor, "bind_session_state", None) + if callable(bind_state): + try: + bind_state(session_db=session_db, session_id=child_session_id) + except Exception: + pass + try: + if agent._memory_manager: + agent._memory_manager.on_session_switch( + child_session_id, + parent_session_id=parent_session_id, + reset=False, + reason="compression", + ) + except Exception as exc: + logger.debug("memory manager compression-child adoption failed: %s", exc) + + return recovered + + +def recover_rotated_compression_session( + agent: Any, +) -> Optional[List[Dict[str, Any]]]: + """Recover a stale live agent before a new turn writes to its old parent.""" + session_db = getattr(agent, "_session_db", None) + session_id = getattr(agent, "session_id", None) or "" + if session_db is None or not session_id: + return None + try: + if not _session_was_rotated_by_compression(session_db, session_id): + return None + # Rotation publication holds the parent compression lease until the + # child handoff is durable. A concurrent turn waits briefly rather than + # observing the intentional parent-ended/child-empty intermediate state. + holder_getter = getattr(session_db, "get_compression_lock_holder", None) + for attempt in range(21): + recovered = _adopt_live_compression_child(agent, session_db, session_id) + if recovered is not None: + return recovered + holder = holder_getter(session_id) if callable(holder_getter) else None + if not holder or attempt == 20: + return None + time.sleep(0.05) + return None + except Exception as exc: + logger.warning( + "compression session recovery failed for session=%s (%s: %s)", + session_id, + type(exc).__name__, + exc, + ) + return None + + def _compression_lock_holder(agent: Any) -> str: """Build a unique holder id for the lock: pid:tid:agent-instance:uuid. @@ -1461,6 +1597,8 @@ def compress_context( if _lock_released: return _lock_released = True + if getattr(agent, "_active_compression_lock_holder", None) == _lock_holder: + agent._active_compression_lock_holder = None if _lock_refresher is not None: try: _lock_refresher.stop() @@ -1472,6 +1610,9 @@ def compress_context( except Exception as _rel_err: logger.debug("compression lock release failed: %s", _rel_err) + if _lock_holder is not None: + agent._active_compression_lock_holder = _lock_holder + # A delayed contender can acquire the parent lock after the winning path # has released it and completed rotation. The lock serializes work but does # not by itself prove that this stale agent still owns a live parent. @@ -1494,15 +1635,25 @@ def compress_context( _existing_sp = agent._build_system_prompt(system_message) return messages, _existing_sp if _parent_already_rotated: - logger.info( - "compression skipped: session=%s was already rotated by " - "another compression path", - _lock_sid, + recovered_messages = _adopt_live_compression_child( + agent, _lock_db, _lock_sid ) _release_lock() _existing_sp = getattr(agent, "_cached_system_prompt", None) if not _existing_sp: _existing_sp = agent._build_system_prompt(system_message) + if recovered_messages is not None: + logger.warning( + "compression recovery: stale session=%s adopted live child=%s", + _lock_sid, + agent.session_id, + ) + return recovered_messages, _existing_sp + logger.warning( + "compression skipped: session=%s was already rotated by " + "another compression path, but no unique live child could be adopted", + _lock_sid, + ) return messages, _existing_sp # The agent may have been constructed before another path completed an @@ -1536,6 +1687,28 @@ def compress_context( ) _lock_refresher.start() + # The caller's history snapshot predates lease acquisition. Reload the + # durable parent after the lease is live; any row not represented as the + # ordered prefix means a frontend/background writer committed in that + # window, so publishing from this snapshot would omit a durable turn. + if _lock_db is not None and _lock_sid: + durable_loader = getattr( + type(_lock_db), "get_messages_as_conversation", None + ) + if callable(durable_loader): + durable_parent = durable_loader(_lock_db, _lock_sid) + if not _durable_history_matches_snapshot(durable_parent, messages): + logger.warning( + "compression aborted: session=%s changed before lease " + "acquisition; preserving newer durable messages", + _lock_sid, + ) + _release_lock() + existing_prompt = getattr(agent, "_cached_system_prompt", None) + if not existing_prompt: + existing_prompt = agent._build_system_prompt(system_message) + return messages, existing_prompt + # Notify external memory provider before compression discards context. # The provider's on_pre_compress() may return a string of insights it # wants surfaced inside the compression summary; capture and forward it @@ -1880,40 +2053,13 @@ def compress_context( ) except Exception: pass # best-effort — don't block compression on a flush error - # Propagate title to the new session with auto-numbering - old_title = agent._session_db.get_session_title(agent.session_id) - agent._session_db.end_session(agent.session_id, "compression") - old_session_id = agent.session_id - agent.session_id = f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:6]}" - # Ordering contract: the agent thread updates the contextvar here; - # the gateway propagates to SessionEntry after run_in_executor returns. - try: - from gateway.session_context import set_current_session_id - - set_current_session_id(agent.session_id) - except Exception: - os.environ["HERMES_SESSION_ID"] = agent.session_id - # The gateway/tools session context (ContextVar + env) and the - # logging session context are SEPARATE mechanisms. The call above - # moves the former; the ``[session_id]`` tag on log lines comes - # from ``hermes_logging._session_context`` (set once per turn in - # conversation_loop.py). Without this, post-rotation log lines in - # the same turn keep the STALE old id while the message/DB/gateway - # state carry the new one — breaking log correlation exactly at the - # compaction boundary (see #34089). Guarded separately so a logging - # failure can never regress the routing update above. - try: - from hermes_logging import set_session_context - - set_session_context(agent.session_id) - except Exception: - pass - agent._session_db_created = False + # Publish parent closure + child row + compacted handoff in + # one transaction. No reader can observe a missing/empty child. # The rotation child must stay on the parent's profile — # mirror _ensure_db_session's stamp ("default" persists as - # NULL). _insert_session_row's parent backfill additionally - # COALESCEs from the parent row, covering app-global remote - # sessions whose thread lacks the HERMES_HOME context. + # NULL). publish_compression_child additionally COALESCEs + # from the parent row, covering app-global remote sessions + # whose thread lacks the HERMES_HOME context. try: from hermes_cli.profiles import get_active_profile_name @@ -1922,53 +2068,39 @@ def compress_context( _profile_for_child = None except Exception: _profile_for_child = None + old_title = agent._session_db.get_session_title(agent.session_id) + old_session_id = agent.session_id + new_session_id = ( + f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_" + f"{uuid.uuid4().hex[:6]}" + ) + agent._session_db.publish_compression_child( + parent_session_id=old_session_id, + child_session_id=new_session_id, + source=agent.platform + or os.environ.get("HERMES_SESSION_SOURCE", "cli"), + model=agent.model, + model_config=agent._session_init_model_config, + system_prompt=new_system_prompt, + messages=compressed, + cwd=getattr(agent, "working_directory", None), + profile_name=_profile_for_child, + compression_lock_holder=_lock_holder, + require_compression_lease=_lock_holder is not None, + ) + agent.session_id = new_session_id try: - agent._session_db.create_session( - session_id=agent.session_id, - source=agent.platform or os.environ.get("HERMES_SESSION_SOURCE", "cli"), - model=agent.model, - model_config=agent._session_init_model_config, - parent_session_id=old_session_id, - profile_name=_profile_for_child, - ) - except Exception as _cs_err: - # The child row could not be created (e.g. FK constraint, - # contended write). Previously the outer handler simply - # warned and let the agent continue on the NEW id — which - # has no row in state.db, producing an orphan: the parent - # is ended, the child is never indexed, and every - # subsequent message is attributed to a session that - # doesn't exist (#33906/#33907). Roll the live id back to - # the parent so the conversation stays attached to a real, - # indexed session instead of a phantom. - logger.warning( - "Compression child session create failed (%s) — " - "rolling back to parent session %s to avoid an orphan.", - _cs_err, old_session_id, - ) - agent.session_id = old_session_id - try: - from gateway.session_context import set_current_session_id - set_current_session_id(agent.session_id) - except Exception: - os.environ["HERMES_SESSION_ID"] = agent.session_id - try: - from hermes_logging import set_session_context - set_session_context(agent.session_id) - except Exception: - pass - # Re-open the parent: it was ended above, but we're - # continuing on it, so it must not stay closed. - try: - agent._session_db.reopen_session(old_session_id) - except Exception: - pass - old_session_id = None # no rotation happened - # The parent row already exists in state.db, so mark the - # session as created — _ensure_db_session would otherwise - # retry a (harmless INSERT OR IGNORE) create next turn. - agent._session_db_created = True - raise + from gateway.session_context import set_current_session_id + + set_current_session_id(agent.session_id) + except Exception: + os.environ["HERMES_SESSION_ID"] = agent.session_id + try: + from hermes_logging import set_session_context + + set_session_context(agent.session_id) + except Exception: + pass agent._session_db_created = True split_status = "rotated_committed" # Carry a persistent /goal onto the continuation session. @@ -1988,18 +2120,14 @@ def compress_context( except (ValueError, Exception) as e: logger.debug("Could not propagate title on compression: %s", e) - # Shared post-write steps (both modes target agent.session_id, which - # in-place keeps and rotation has already reassigned to the new id): - # refresh the stored system prompt and reset the flush cursor so the - # next turn re-bases its append diff. - agent._session_db.update_system_prompt(agent.session_id, new_system_prompt) + # In-place mode still updates/replaces the current row here. + # Rotation already published prompt + compacted handoff atomically. if in_place: + agent._session_db.update_system_prompt( + agent.session_id, new_system_prompt + ) agent._last_flushed_db_idx = 0 else: - # A headless turn can be killed before its finalizer. Persist - # the rotated child's compacted handoff at the boundary so - # the new session is immediately resumable. - agent._session_db.replace_messages(agent.session_id, compressed) agent._last_flushed_db_idx = len(compressed) agent._flushed_db_message_session_id = agent.session_id agent._flushed_db_message_ids = { @@ -2009,7 +2137,22 @@ def compress_context( } _session_commit_succeeded = True except Exception as e: - split_status = "aborted" if locals().get("old_session_id") is None and not in_place else "failed_not_indexed" + if ( + not in_place + and locals().get("old_session_id") + and agent.session_id == old_session_id + ): + # Atomic publication failed (including lease loss): keep the + # parent live and discard the stale compacted snapshot. + old_session_id = None + messages[:] = copy.deepcopy(messages_before_compression) + compressed = messages + _compression_made_progress = False + split_status = ( + "aborted" + if locals().get("old_session_id") is None and not in_place + else "failed_not_indexed" + ) # If the rotation rolled back to the parent (orphan-avoidance # above), agent.session_id is the still-indexed parent and # old_session_id was cleared — so this is recovery, not an diff --git a/agent/turn_context.py b/agent/turn_context.py index 59914befe0c7..ac9555d012e4 100644 --- a/agent/turn_context.py +++ b/agent/turn_context.py @@ -36,6 +36,7 @@ from agent.conversation_compression import ( PREFLIGHT_COMPRESSION_STATUS_TEMPLATE, compression_skipped_due_to_lock, conversation_history_after_compression, + recover_rotated_compression_session, ) from agent.context_engine import automatic_compaction_status_message from agent.iteration_budget import IterationBudget @@ -353,6 +354,13 @@ def build_turn_context( # Guard stdio against OSError from broken pipes (systemd/headless/daemon). install_safe_stdio() + # Recover a session rotated by another path before binding log/turn ids or + # copying client-supplied history. Everything in this turn must consistently + # belong to the canonical child, including observability metadata. + recovered_history = recover_rotated_compression_session(agent) + if recovered_history is not None: + conversation_history = recovered_history + # NOTE: the DB session row is created later, AFTER the system prompt is # restored/built (see _ensure_db_session() below the system-prompt block). # Creating it here — before _cached_system_prompt is populated — inserts a diff --git a/gateway/session.py b/gateway/session.py index 7f630dd861df..89f4188669e7 100644 --- a/gateway/session.py +++ b/gateway/session.py @@ -1186,6 +1186,11 @@ class SessionStore: self._legacy_slack_claim_lock = threading.Lock() self._claimed_legacy_slack_keys: set[str] = set() self._transcript_retry_lock = threading.Lock() + # Exactly one transcript drainer mutates routing/queues at a time. SQLite + # serializes writes anyway; this outer lock also makes parent->child + # queue migration and routing publication linearizable. + self._transcript_drain_lock = threading.RLock() + self._transcript_reroutes: Dict[str, str] = {} self._dirty_transcripts: Dict[str, List[Dict[str, Any]]] = {} self._transcript_append_failures: Dict[str, int] = {} self._fts_rebuild_attempted = False @@ -2914,6 +2919,29 @@ class SessionStore: return getattr(entry, "session_id", None) if entry else None def append_to_transcript(self, session_id: str, message: Dict[str, Any], skip_db: bool = False) -> None: + """Serialize transcript draining across queue migration boundaries.""" + if not self._db or skip_db: + return + drain_lock = getattr(self, "_transcript_drain_lock", None) + if drain_lock is None: + # Compatibility for old in-memory/test instances created via + # object.__new__ before this field existed. + drain_lock = threading.RLock() + self._transcript_drain_lock = drain_lock + with drain_lock: + reroutes = getattr(self, "_transcript_reroutes", None) + if reroutes is None: + reroutes = {} + self._transcript_reroutes = reroutes + seen = set() + while session_id in reroutes and session_id not in seen: + seen.add(session_id) + session_id = reroutes[session_id] + self._append_to_transcript_serialized(session_id, message) + + def _append_to_transcript_serialized( + self, session_id: str, message: Dict[str, Any] + ) -> None: """Append a message to a session's transcript (SQLite). Args: @@ -2922,8 +2950,6 @@ class SessionStore: _flush_messages_to_session_db(), preventing the duplicate-write bug (#860). """ - if not self._db or skip_db: - return with self._transcript_retry_lock: pending = self._dirty_transcripts.setdefault(session_id, []) pending.append(dict(message)) @@ -2939,12 +2965,76 @@ class SessionStore: # Snapshot the first pending message, then release the lock # before the DB write so other sessions are not blocked. msg = pending[0] + queue_session_id = session_id # DB write outside the retry lock — other sessions can append # concurrently. We re-acquire the lock only to update the queue. while True: try: self._append_transcript_message(session_id, msg) except Exception as exc: + from hermes_state import CompressionSessionClosedError + + if isinstance(exc, CompressionSessionClosedError): + child = self._db.find_live_compression_child(session_id) + child_id = str(child["id"]) if child and child.get("id") else "" + if child_id: + try: + self._append_transcript_message(child_id, msg) + except Exception as reroute_exc: + exc = reroute_exc + else: + with self._transcript_retry_lock: + if pending and pending[0] is msg: + pending.pop(0) + existing_child_pending = self._dirty_transcripts.get( + child_id, [] + ) + if pending: + # Older parent backlog must precede messages + # already queued directly on the child. + pending.extend(existing_child_pending) + self._dirty_transcripts[child_id] = pending + elif existing_child_pending: + pending = existing_child_pending + self._dirty_transcripts.pop(queue_session_id, None) + previous_failures = self._transcript_append_failures.pop( + queue_session_id, 0 + ) + if previous_failures: + self._transcript_append_failures[child_id] = max( + previous_failures, + self._transcript_append_failures.get(child_id, 0), + ) + self._transcript_reroutes[session_id] = child_id + queue_session_id = child_id + # Publish routing only after the retry queue has moved, + # so new child writes cannot bypass older parent backlog. + with self._lock: + for entry in self._entries.values(): + if entry.session_id == session_id: + entry.session_id = child_id + self._save() + if not pending: + return + msg = pending[0] + session_id = child_id + continue + else: + # This is a permanent routing invariant failure, not a + # transient DB outage. Drop it from the retry queue so it + # cannot poison later transcript writes indefinitely. + with self._transcript_retry_lock: + if pending and pending[0] is msg: + pending.pop(0) + if not pending: + self._dirty_transcripts.pop(queue_session_id, None) + self._transcript_append_failures.pop(session_id, None) + logger.error( + "Session DB transcript append rejected for compression-ended " + "%s with no unique live child; not retrying", + session_id, + ) + return if self._is_fts_corruption_error(exc) and self._rebuild_fts_once(): try: self._append_transcript_message(session_id, msg) @@ -2955,7 +3045,7 @@ class SessionStore: if pending and pending[0] is msg: pending.pop(0) if not pending: - self._dirty_transcripts.pop(session_id, None) + self._dirty_transcripts.pop(queue_session_id, None) self._transcript_append_failures.pop(session_id, None) continue with self._transcript_retry_lock: @@ -2972,7 +3062,7 @@ class SessionStore: if pending and pending[0] is msg: pending.pop(0) if not pending: - self._dirty_transcripts.pop(session_id, None) + self._dirty_transcripts.pop(queue_session_id, None) self._transcript_append_failures.pop(session_id, None) return msg = pending[0] diff --git a/hermes_state.py b/hermes_state.py index 64e65d64b34b..6ae4863b1843 100644 --- a/hermes_state.py +++ b/hermes_state.py @@ -1560,6 +1560,21 @@ END; """ +class CompressionSessionClosedError(RuntimeError): + """A durable write targeted a parent already closed by compression.""" + + def __init__(self, session_id: str): + self.session_id = session_id + super().__init__( + f"Session {session_id!r} is closed by compression; " + "adopt its live continuation before appending messages" + ) + + +class CompressionSessionBusyError(RuntimeError): + """A non-owner tried to write while compression owns the session.""" + + class SessionDB: """ SQLite-backed session storage with FTS5 search. @@ -3845,6 +3860,146 @@ class SessionDB: ).fetchone() return dict(row) if row else None + def find_live_compression_child( + self, parent_session_id: str + ) -> Optional[Dict[str, Any]]: + """Return the unique live direct child of a compression-ended session. + + A stale agent may observe that another compression path already rotated + its parent. Recovery is safe only when the durable lineage identifies + exactly one live direct continuation. Multiple children are treated as + ambiguous and fail closed rather than guessing which transcript owns + subsequent messages. + """ + if not parent_session_id: + return None + with self._lock: + parent = self._conn.execute( + "SELECT ended_at, end_reason FROM sessions WHERE id = ?", + (parent_session_id,), + ).fetchone() + if ( + parent is None + or parent["ended_at"] is None + or parent["end_reason"] != "compression" + ): + return None + rows = self._conn.execute( + """ + SELECT * FROM sessions + WHERE parent_session_id = ? + AND ended_at IS NULL + AND json_extract(COALESCE(model_config, '{}'), '$._branched_from') IS NULL + AND json_extract(COALESCE(model_config, '{}'), '$._delegate_from') IS NULL + AND COALESCE(source, '') != 'tool' + ORDER BY started_at ASC + LIMIT 2 + """, + (parent_session_id,), + ).fetchall() + return dict(rows[0]) if len(rows) == 1 else None + + def publish_compression_child( + self, + *, + parent_session_id: str, + child_session_id: str, + source: str, + messages: List[Dict[str, Any]], + model: str = None, + model_config: Dict[str, Any] = None, + system_prompt: str = None, + cwd: str = None, + profile_name: str = None, + compression_lock_holder: str = None, + require_compression_lease: bool = True, + ) -> None: + """Atomically close a parent and publish its durable compression child. + + The parent closure, child row, and compacted handoff become visible in + one transaction. Readers can therefore observe either the live parent or + a complete child, never an ended parent with a missing/empty child. + """ + def _do(conn): + lock_row = conn.execute( + "SELECT holder, expires_at FROM compression_locks WHERE session_id = ?", + (parent_session_id,), + ).fetchone() + if require_compression_lease and ( + lock_row is None + or not compression_lock_holder + or lock_row["holder"] != compression_lock_holder + or float(lock_row["expires_at"]) <= time.time() + ): + raise CompressionSessionBusyError( + f"Compression lease lost before publication: {parent_session_id}" + ) + parent = conn.execute( + """SELECT ended_at, cwd, git_branch, git_repo_root, + user_id, session_key, chat_id, chat_type, + thread_id, display_name, origin_json, profile_name + FROM sessions WHERE id = ?""", + (parent_session_id,), + ).fetchone() + if parent is None: + raise RuntimeError(f"Compression parent not found: {parent_session_id}") + if parent["ended_at"] is not None: + raise RuntimeError(f"Compression parent already ended: {parent_session_id}") + if not messages: + raise RuntimeError("Compression child handoff must not be empty") + + conn.execute( + """INSERT INTO sessions ( + id, source, model, model_config, system_prompt, + parent_session_id, cwd, git_branch, git_repo_root, + profile_name, user_id, session_key, chat_id, chat_type, + thread_id, display_name, origin_json, started_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + ( + child_session_id, + source, + model, + json.dumps(model_config) if model_config else None, + system_prompt, + parent_session_id, + cwd or parent["cwd"], + parent["git_branch"], + parent["git_repo_root"], + # Same inheritance contract as _insert_session_row's + # compression-fork backfill (#59527 / cross-profile jump + # fix): the child stays on the parent's profile and keeps + # the gateway routing/origin columns so peer recovery + # still works after a crash at the boundary. + profile_name or parent["profile_name"], + parent["user_id"], + parent["session_key"], + parent["chat_id"], + parent["chat_type"], + parent["thread_id"], + parent["display_name"], + parent["origin_json"], + time.time(), + ), + ) + total_messages, total_tool_calls = self._insert_message_rows( + conn, child_session_id, messages + ) + conn.execute( + "UPDATE sessions SET message_count = ?, tool_call_count = ? WHERE id = ?", + (total_messages, total_tool_calls, child_session_id), + ) + updated = conn.execute( + "UPDATE sessions SET ended_at = ?, end_reason = 'compression' " + "WHERE id = ? AND ended_at IS NULL", + (time.time(), parent_session_id), + ) + if updated.rowcount != 1: + raise RuntimeError( + f"Compression parent changed during publication: {parent_session_id}" + ) + + self._execute_write(_do) + def end_session(self, session_id: str, end_reason: str) -> None: """Mark a session as ended. @@ -5850,6 +6005,7 @@ class SessionDB: api_content: Optional[str] = None, display_kind: Optional[str] = None, display_metadata: Optional[Dict[str, Any]] = None, + compression_lock_holder: Optional[str] = None, ) -> int: """ Append a message to a session. Returns the message row ID. @@ -5916,6 +6072,28 @@ class SessionDB: num_tool_calls = len(tool_calls) if isinstance(tool_calls, list) else 1 def _do(conn): + active_lock = conn.execute( + "SELECT holder FROM compression_locks " + "WHERE session_id = ? AND expires_at > ?", + (session_id, time.time()), + ).fetchone() + if ( + active_lock is not None + and active_lock["holder"] != compression_lock_holder + ): + raise CompressionSessionBusyError( + f"Session {session_id!r} is being compressed by another writer" + ) + session = conn.execute( + "SELECT ended_at, end_reason FROM sessions WHERE id = ?", + (session_id,), + ).fetchone() + if ( + session is not None + and session["ended_at"] is not None + and session["end_reason"] == "compression" + ): + raise CompressionSessionClosedError(session_id) cursor = conn.execute( """INSERT INTO messages (session_id, role, content, tool_call_id, tool_calls, tool_name, effect_disposition, timestamp, token_count, finish_reason, @@ -6125,6 +6303,16 @@ class SessionDB: active_clause = " AND active = 1" if active_only else "" def _do(conn): + session = conn.execute( + "SELECT ended_at, end_reason FROM sessions WHERE id = ?", + (session_id,), + ).fetchone() + if ( + session is not None + and session["ended_at"] is not None + and session["end_reason"] == "compression" + ): + raise CompressionSessionClosedError(session_id) conn.execute( f"DELETE FROM messages WHERE session_id = ?{active_clause}", (session_id,), diff --git a/run_agent.py b/run_agent.py index c73f167bb902..ebd280c4768f 100644 --- a/run_agent.py +++ b/run_agent.py @@ -2097,6 +2097,9 @@ class AIAgent: and not msg.get("_compressed_summary_has_user_turn") else msg.get("display_kind") ), + compression_lock_holder=getattr( + self, "_active_compression_lock_holder", None + ), ) msg[_DB_PERSISTED_MARKER] = True # The intrinsic markers are now the sole source of truth. Reset the diff --git a/tests/agent/test_compression_concurrent_fork.py b/tests/agent/test_compression_concurrent_fork.py index 2e4624d1bdac..8af78d5ccdb6 100644 --- a/tests/agent/test_compression_concurrent_fork.py +++ b/tests/agent/test_compression_concurrent_fork.py @@ -316,6 +316,35 @@ def test_concurrent_compression_does_not_fork_session(tmp_path: Path) -> None: ) +def test_durable_message_committed_before_lease_aborts_stale_snapshot( + tmp_path: Path, +) -> None: + """A durable row absent from the caller snapshot must survive in the parent.""" + db = SessionDB(db_path=tmp_path / "state.db") + parent_sid = "PRE_LEASE_DURABLE_RACE" + db.create_session(parent_sid, source="webui") + db.append_message(parent_sid, "user", "old durable") + + # Frontend takes its snapshot, then another producer commits before this + # compressor acquires the lease. + stale_snapshot = [{"role": "user", "content": "old durable"}] + db.append_message(parent_sid, "assistant", "late committed before lease") + agent = _build_agent_with_db(db, parent_sid) + + returned, _system_prompt = agent._compress_context( + stale_snapshot, "sys", approx_tokens=120_000 + ) + + assert returned is stale_snapshot + assert agent.session_id == parent_sid + assert db.find_live_compression_child(parent_sid) is None + assert [m["content"] for m in db.get_messages_as_conversation(parent_sid)] == [ + "old durable", + "late committed before lease", + ] + agent.context_compressor.compress.assert_not_called() + + def test_skipped_compression_returns_messages_unchanged(tmp_path: Path) -> None: """The loser of the lock race must return its input messages verbatim. @@ -507,6 +536,67 @@ def test_commit_fence_waits_for_an_active_commit() -> None: assert result["cancelled"] is False +def test_delayed_contender_adopts_unique_rotated_child(tmp_path: Path) -> None: + """A stale agent must continue on the winner's compacted child transcript.""" + db = SessionDB(db_path=tmp_path / "state.db") + parent_sid = "STALE_PARENT" + child_sid = "CANONICAL_CHILD" + db.create_session(parent_sid, source="webui") + db.end_session(parent_sid, "compression") + db.create_session(child_sid, source="webui", parent_session_id=parent_sid) + compacted = [ + {"role": "user", "content": "[CONTEXT COMPACTION] summary"}, + {"role": "assistant", "content": "compacted tail"}, + ] + db.replace_messages(child_sid, compacted) + + agent = _build_agent_with_db(db, parent_sid) + stale_messages = [ + {"role": "user", "content": "stale"}, + {"role": "assistant", "content": "x" * 1000}, + ] + recovered, _system_prompt = agent._compress_context( + stale_messages, "sys", approx_tokens=120_000 + ) + + assert agent.session_id == child_sid + assert [(m["role"], m["content"]) for m in recovered] == [ + ("user", "[CONTEXT COMPACTION] summary"), + ("assistant", "compacted tail"), + ] + assert agent._session_db_created is True + assert agent._flushed_db_message_session_id == child_sid + assert agent._last_flushed_db_idx == len(recovered) + agent.context_compressor.compress.assert_not_called() + lifecycle_args, lifecycle_kwargs = agent.context_compressor.on_session_start.call_args + assert lifecycle_args == (child_sid,) + assert lifecycle_kwargs["boundary_reason"] == "compression" + assert lifecycle_kwargs["old_session_id"] == parent_sid + assert lifecycle_kwargs["session_db"] is db + + +def test_delayed_contender_fails_closed_without_unique_child(tmp_path: Path) -> None: + """Missing or ambiguous lineage must not silently select a continuation.""" + db = SessionDB(db_path=tmp_path / "state.db") + parent_sid = "AMBIGUOUS_PARENT" + db.create_session(parent_sid, source="webui") + db.end_session(parent_sid, "compression") + db.create_session("CHILD_A", source="webui", parent_session_id=parent_sid) + db.create_session("CHILD_B", source="webui", parent_session_id=parent_sid) + db.replace_messages("CHILD_A", [{"role": "user", "content": "a"}]) + db.replace_messages("CHILD_B", [{"role": "user", "content": "b"}]) + agent = _build_agent_with_db(db, parent_sid) + stale_messages = [{"role": "user", "content": "stale"}] + + returned, _system_prompt = agent._compress_context( + stale_messages, "sys", approx_tokens=120_000 + ) + + assert returned is stale_messages or returned == stale_messages + assert agent.session_id == parent_sid + agent.context_compressor.compress.assert_not_called() + + def test_compression_restores_user_turn_when_compressor_drops_all_users(tmp_path: Path) -> None: """Provider chat templates need at least one user message after compaction. diff --git a/tests/agent/test_compression_rotation_state.py b/tests/agent/test_compression_rotation_state.py index 0ba525d62532..11d67f28ca83 100644 --- a/tests/agent/test_compression_rotation_state.py +++ b/tests/agent/test_compression_rotation_state.py @@ -126,21 +126,35 @@ class TestOrphanRollbackOnCreateFailure: db.create_session(parent, source="cli") agent = _build_agent_with_db(db, parent) - # Make the CHILD create_session raise, but let the initial parent - # end_session/reopen work. We patch create_session to blow up. - real_create = db.create_session + # Atomic publication failure must leave the live parent and caller's + # original list untouched even when a plugin compressor mutates in place. + original = _msgs() + + def _mutating_compress(live_messages, **_kwargs): + live_messages[:] = [ + {"role": "user", "content": "mutated compacted snapshot"} + ] + return live_messages + + agent.context_compressor.compress.side_effect = _mutating_compress def _boom(*a, **k): - raise RuntimeError("FOREIGN KEY constraint failed") + raise RuntimeError("simulated atomic publication failure") - with patch.object(db, "create_session", side_effect=_boom): - agent._compress_context(_msgs(), "sys", approx_tokens=120_000) + with patch.object(db, "publish_compression_child", side_effect=_boom): + returned, _system_prompt = agent._compress_context( + original, "sys", approx_tokens=120_000 + ) - # The live id must roll back to the still-indexed parent — NOT a - # phantom child id that has no row in state.db. assert agent.session_id == parent - assert db.get_session(parent) is not None - _ = real_create # silence unused + assert [(m["role"], m["content"]) for m in returned] == [ + (m["role"], m["content"]) for m in _msgs() + ] + assert returned is original + parent_row = db.get_session(parent) + assert parent_row is not None + assert parent_row["ended_at"] is None + assert db.find_live_compression_child(parent) is None class TestWorkspaceMetadataFollowsRotation: diff --git a/tests/agent/test_turn_context.py b/tests/agent/test_turn_context.py index 8c34a25c8ede..fcf94fe39a07 100644 --- a/tests/agent/test_turn_context.py +++ b/tests/agent/test_turn_context.py @@ -208,6 +208,37 @@ def test_returns_turn_context_with_user_message_appended(): assert ctx.active_system_prompt == "SYSTEM" +def test_turn_start_replaces_stale_parent_history_with_compression_child(): + agent = _FakeAgent() + stale_history = [{"role": "user", "content": "stale parent"}] + compacted_history = [ + {"role": "user", "content": "[CONTEXT COMPACTION] summary"}, + {"role": "assistant", "content": "child tail"}, + ] + + def _recover(_agent): + _agent.session_id = "compression-child" + return compacted_history + + log_context = MagicMock() + with patch( + "agent.turn_context.recover_rotated_compression_session", + side_effect=_recover, + ): + ctx = _build( + agent, + conversation_history=stale_history, + set_session_context=log_context, + ) + + assert agent.session_id == "compression-child" + assert agent._current_turn_id.startswith("compression-child:") + log_context.assert_called_once_with("compression-child") + assert ctx.conversation_history == compacted_history + assert ctx.messages == compacted_history + [{"role": "user", "content": "hello"}] + assert all(message.get("content") != "stale parent" for message in ctx.messages) + + def test_applies_agent_side_effects(): agent = _FakeAgent() _build(agent) diff --git a/tests/gateway/test_session.py b/tests/gateway/test_session.py index 2913cc7a1d8a..ca86fc728414 100644 --- a/tests/gateway/test_session.py +++ b/tests/gateway/test_session.py @@ -5,6 +5,7 @@ from dataclasses import replace from datetime import datetime from pathlib import Path from unittest.mock import patch, MagicMock +from hermes_state import SessionDB from gateway.config import Platform, HomeChannel, GatewayConfig, PlatformConfig from gateway.platforms.base import MessageEvent from gateway.session import ( @@ -2272,6 +2273,106 @@ class TestRewriteTranscriptPreservesReasoning: class TestGatewaySessionDbRecovery: + def test_compression_closed_parent_reroutes_without_retry_queue(self, tmp_path): + import threading + from types import SimpleNamespace + + db = SessionDB(db_path=tmp_path / "state.db") + db.create_session("parent", source="telegram") + db.end_session("parent", "compression") + db.create_session("child", source="telegram", parent_session_id="parent") + db.replace_messages("child", [{"role": "user", "content": "summary"}]) + + store = object.__new__(SessionStore) + store._db = db + store._lock = threading.RLock() + store._entries = {"route": SimpleNamespace(session_id="parent")} + store._loaded = True + store._save = lambda: None + store._transcript_retry_lock = threading.Lock() + store._dirty_transcripts = {} + store._transcript_append_failures = {} + store._fts_rebuild_attempted = False + + store.append_to_transcript( + "parent", {"role": "assistant", "content": "routed to child"} + ) + + assert store._entries["route"].session_id == "child" + assert "parent" not in store._dirty_transcripts + assert [m["content"] for m in db.get_messages_as_conversation("parent")] == [] + assert [m["content"] for m in db.get_messages_as_conversation("child")] == [ + "summary", + "routed to child", + ] + db.close() + + def test_transcript_reroute_migrates_remaining_backlog_to_child(self): + import threading + from types import SimpleNamespace + from hermes_state import CompressionSessionClosedError + + class FakeDb: + def find_live_compression_child(self, session_id): + assert session_id == "parent" + return {"id": "child"} + + store = object.__new__(SessionStore) + store._db = FakeDb() + store._lock = threading.RLock() + store._entries = {"route": SimpleNamespace(session_id="parent")} + store._loaded = True + store._save = lambda: None + store._transcript_retry_lock = threading.Lock() + store._dirty_transcripts = { + "parent": [ + {"role": "user", "content": "old-1"}, + {"role": "assistant", "content": "old-2"}, + ] + } + store._transcript_append_failures = {"parent": 2} + store._fts_rebuild_attempted = True + child_attempts = [] + failed_old_2 = False + + def _append(session_id, message): + nonlocal failed_old_2 + if session_id == "parent": + raise CompressionSessionClosedError("parent") + child_attempts.append(message["content"]) + if message["content"] == "old-2" and not failed_old_2: + failed_old_2 = True + raise RuntimeError("transient child failure") + + store._append_transcript_message = _append + store.append_to_transcript( + "parent", {"role": "user", "content": "old-3"} + ) + + assert child_attempts == ["old-1", "old-2"] + assert store._entries["route"].session_id == "child" + assert "parent" not in store._dirty_transcripts + assert [m["content"] for m in store._dirty_transcripts["child"]] == [ + "old-2", + "old-3", + ] + assert store._transcript_append_failures["child"] >= 2 + + # A producer still holding the stale parent id must join and drain the + # child backlog before its newer message; no duplicate old-1 is allowed. + store.append_to_transcript( + "parent", {"role": "assistant", "content": "new-after-reroute"} + ) + assert child_attempts == [ + "old-1", + "old-2", + "old-2", + "old-3", + "new-after-reroute", + ] + assert "parent" not in store._dirty_transcripts + assert "child" not in store._dirty_transcripts + def test_transcript_append_rebuilds_fts_and_retries_dirty_rows_in_order(self): import threading diff --git a/tests/run_agent/test_compression_boundary_hook.py b/tests/run_agent/test_compression_boundary_hook.py index 36e5b97af613..67d0e25a68fb 100644 --- a/tests/run_agent/test_compression_boundary_hook.py +++ b/tests/run_agent/test_compression_boundary_hook.py @@ -120,14 +120,16 @@ class TestCompressionBoundaryHook: ) ) agent.context_compressor = compressor - original_update = db.update_system_prompt + original_publish = db.publish_compression_child - def _record_update(*args, **kwargs): - result = original_update(*args, **kwargs) + def _record_publish(*args, **kwargs): + result = original_publish(*args, **kwargs) events.append("persist") return result - with patch.object(db, "update_system_prompt", side_effect=_record_update): + with patch.object( + db, "publish_compression_child", side_effect=_record_publish + ): agent._compress_context( [{"role": "user", "content": "request"}], "sys", @@ -174,7 +176,7 @@ class TestCompressionBoundaryHook: with patch.object( db, - "update_system_prompt", + "publish_compression_child", side_effect=RuntimeError("synthetic commit failure"), ): agent._compress_context( diff --git a/tests/state/test_compression_lineage_guard.py b/tests/state/test_compression_lineage_guard.py new file mode 100644 index 000000000000..bc41632c98da --- /dev/null +++ b/tests/state/test_compression_lineage_guard.py @@ -0,0 +1,196 @@ +"""Regression tests for stale writes after a compression session split.""" + +from __future__ import annotations + +import pytest + +from hermes_state import SessionDB + + +@pytest.fixture() +def db(tmp_path): + session_db = SessionDB(db_path=tmp_path / "state.db") + try: + yield session_db + finally: + session_db.close() + + +def _compression_parent(db: SessionDB, session_id: str = "parent") -> None: + db.create_session(session_id, source="webui") + db.append_message(session_id, "user", "before split") + db.end_session(session_id, "compression") + + +def test_find_live_compression_child_returns_unique_direct_child(db: SessionDB) -> None: + _compression_parent(db) + db.create_session("child", source="webui", parent_session_id="parent") + + child = db.find_live_compression_child("parent") + + assert child is not None + assert child["id"] == "child" + assert child["parent_session_id"] == "parent" + assert child["ended_at"] is None + + +def test_find_live_compression_child_fails_closed_when_ambiguous(db: SessionDB) -> None: + _compression_parent(db) + db.create_session("child-a", source="webui", parent_session_id="parent") + db.create_session("child-b", source="webui", parent_session_id="parent") + + assert db.find_live_compression_child("parent") is None + + +def test_find_live_compression_child_ignores_ended_children(db: SessionDB) -> None: + _compression_parent(db) + db.create_session("ended-child", source="webui", parent_session_id="parent") + db.end_session("ended-child", "agent_close") + + assert db.find_live_compression_child("parent") is None + + +def test_find_live_compression_child_ignores_non_continuation_children( + db: SessionDB, +) -> None: + _compression_parent(db) + db.create_session("canonical", source="webui", parent_session_id="parent") + db.create_session( + "branch", + source="webui", + parent_session_id="parent", + model_config={"_branched_from": "parent"}, + ) + db.create_session( + "delegate", + source="webui", + parent_session_id="parent", + model_config={"_delegate_from": "parent"}, + ) + db.create_session("tool-child", source="tool", parent_session_id="parent") + + child = db.find_live_compression_child("parent") + + assert child is not None + assert child["id"] == "canonical" + + +def test_append_message_rejects_compression_ended_parent_atomically(db: SessionDB) -> None: + _compression_parent(db) + before = db.get_session("parent")["message_count"] + + with pytest.raises(RuntimeError, match="closed by compression"): + db.append_message("parent", "assistant", "must not land on parent") + + assert db.get_session("parent")["message_count"] == before + assert [m["content"] for m in db.get_messages("parent")] == ["before split"] + + +def test_append_message_preserves_legacy_behavior_for_other_end_reasons(db: SessionDB) -> None: + db.create_session("ended", source="test") + db.end_session("ended", "agent_close") + + message_id = db.append_message("ended", "user", "legacy append") + + assert isinstance(message_id, int) + assert db.get_messages("ended")[-1]["content"] == "legacy append" + + +def test_replace_messages_rejects_compression_ended_parent_atomically( + db: SessionDB, +) -> None: + _compression_parent(db) + + with pytest.raises(RuntimeError, match="closed by compression"): + db.replace_messages("parent", [{"role": "user", "content": "rewrite"}]) + + assert [m["content"] for m in db.get_messages("parent")] == ["before split"] + + +def test_publish_compression_child_is_atomic_on_handoff_failure( + db: SessionDB, monkeypatch +) -> None: + db.create_session("atomic-parent", source="webui") + db.append_message("atomic-parent", "user", "original") + assert db.try_acquire_compression_lock("atomic-parent", "winner", ttl_seconds=60) + + def _boom(*_args, **_kwargs): + raise RuntimeError("handoff insert failed") + + monkeypatch.setattr(db, "_insert_message_rows", _boom) + with pytest.raises(RuntimeError, match="handoff insert failed"): + db.publish_compression_child( + parent_session_id="atomic-parent", + child_session_id="atomic-child", + source="webui", + messages=[{"role": "user", "content": "summary"}], + compression_lock_holder="winner", + ) + + parent = db.get_session("atomic-parent") + assert parent is not None + assert parent["ended_at"] is None + assert db.get_session("atomic-child") is None + + +def test_publish_compression_child_exposes_complete_child(db: SessionDB) -> None: + db.create_session("atomic-parent", source="webui") + db.append_message("atomic-parent", "user", "original") + assert db.try_acquire_compression_lock("atomic-parent", "winner", ttl_seconds=60) + + db.publish_compression_child( + parent_session_id="atomic-parent", + child_session_id="atomic-child", + source="webui", + system_prompt="compressed system", + messages=[{"role": "user", "content": "summary"}], + compression_lock_holder="winner", + ) + + assert db.get_session("atomic-parent")["end_reason"] == "compression" + child = db.find_live_compression_child("atomic-parent") + assert child is not None + assert child["id"] == "atomic-child" + assert child["system_prompt"] == "compressed system" + assert [m["content"] for m in db.get_messages("atomic-child")] == ["summary"] + + +def test_publish_compression_child_rejects_lost_or_expired_lease(db: SessionDB) -> None: + db.create_session("lease-parent", source="webui") + db.append_message("lease-parent", "user", "new durable turn") + assert db.try_acquire_compression_lock("lease-parent", "new-winner", ttl_seconds=60) + + with pytest.raises(RuntimeError, match="lease lost"): + db.publish_compression_child( + parent_session_id="lease-parent", + child_session_id="stale-child", + source="webui", + messages=[{"role": "user", "content": "stale summary"}], + compression_lock_holder="old-loser", + ) + + parent = db.get_session("lease-parent") + assert parent is not None + assert parent["ended_at"] is None + assert db.get_session("stale-child") is None + assert [m["content"] for m in db.get_messages("lease-parent")] == [ + "new durable turn" + ] + + +def test_compression_lease_blocks_non_owner_but_allows_owner_flush( + db: SessionDB, +) -> None: + db.create_session("leased", source="webui") + assert db.try_acquire_compression_lock("leased", "winner", ttl_seconds=60) + + with pytest.raises(RuntimeError, match="being compressed"): + db.append_message("leased", "user", "late stale turn") + + db.append_message( + "leased", + "assistant", + "winner flush", + compression_lock_holder="winner", + ) + assert [m["content"] for m in db.get_messages("leased")] == ["winner flush"] diff --git a/tests/test_hermes_state.py b/tests/test_hermes_state.py index 657938e327e2..a0299df447b2 100644 --- a/tests/test_hermes_state.py +++ b/tests/test_hermes_state.py @@ -4693,13 +4693,13 @@ class TestListSessionsRich: """ t0 = 1709500000.0 db.create_session("root1", "cli") + db.append_message("root1", "user", "old ask") with db._lock: db._conn.execute("UPDATE sessions SET started_at=? WHERE id=?", (t0, "root1")) db._conn.execute( "UPDATE sessions SET ended_at=?, end_reason=? WHERE id=?", (t0 + 100, "compression", "root1"), ) - db.append_message("root1", "user", "old ask") # Continuation tip created after root ended; last activity much later. db.create_session("tip1", "cli", parent_session_id="root1")