Skip to content

fix(compression): recover rotated session lineage #70121

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
ruizanthony wants to merge 1 commit into from
Closed

fix(compression): recover rotated session lineage #70121

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
296 changes: 220 additions & 76 deletions agent/conversation_compression.py
View file
Open in desktop

Large diffs are not rendered by default.

8 changes: 8 additions & 0 deletions agent/turn_context.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
IDLE_COMPACTION_STATUS_TEMPLATE,
PREFLIGHT_COMPRESSION_STATUS_TEMPLATE,
conversation_history_after_compression,
recover_rotated_compression_session,
)
from agent.context_engine import automatic_compaction_status_message
from agent.iteration_budget import IterationBudget
Expand Down Expand Up @@ -352,6 +353,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
Expand Down
98 changes: 94 additions & 4 deletions gateway/session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1054,6 +1054,11 @@ def __init__(self, sessions_dir: Path, config: GatewayConfig,
self._inflight_lock = threading.Lock()
self._inflight_sessions: Dict[str, _SessionFlight] = {}
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
Expand Down Expand Up @@ -2598,6 +2603,29 @@ def peek_session_id(self, session_key: str) -> Optional[str]:
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:
Expand All @@ -2606,8 +2634,6 @@ def append_to_transcript(self, session_id: str, message: Dict[str, Any], skip_db
_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))
Expand All @@ -2623,12 +2649,76 @@ def append_to_transcript(self, session_id: str, message: Dict[str, Any], skip_db
# 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)
Expand All @@ -2639,7 +2729,7 @@ def append_to_transcript(self, session_id: str, message: Dict[str, Any], skip_db
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:
Expand All @@ -2656,7 +2746,7 @@ def append_to_transcript(self, session_id: str, message: Dict[str, Any], skip_db
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]
Expand Down
178 changes: 178 additions & 0 deletions hermes_state.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1505,6 +1505,21 @@ def load_fts5_cjk_extension(conn: sqlite3.Connection) -> bool:
"""


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.
Expand Down Expand Up @@ -3785,6 +3800,136 @@ def find_latest_gateway_session_for_peer(
).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
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, 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"],
profile_name,
parent["user_id"],
parent["session_key"],
parent["chat_id"],
parent["chat_type"],
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.

Expand Down Expand Up @@ -5598,6 +5743,7 @@ def append_message(
effect_disposition: Optional[str] = None,
timestamp: Any = None,
api_content: Optional[str] = None,
compression_lock_holder: Optional[str] = None,
) -> int:
"""
Append a message to a session. Returns the message row ID.
Expand Down Expand Up @@ -5661,6 +5807,28 @@ def append_message(
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,
Expand Down Expand Up @@ -5832,6 +6000,16 @@ def replace_messages(
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,),
Expand Down
3 changes: 3 additions & 0 deletions run_agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2042,6 +2042,9 @@ def _flush_messages_to_session_db_unlocked(
codex_message_items=msg.get("codex_message_items") if role == "assistant" else None,
timestamp=_row_timestamp,
api_content=_row_api_content,
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
Expand Down
Loading