Merge branch 'simp/r2-state-b' into simp/integration2

# Conflicts:
#	hermes_state_messages.py
This commit is contained in:
Teknium
2026-09-02 17:26:25 -07:00
4 changed files with 282 additions and 510 deletions
+50 -97
View File
@@ -9,7 +9,7 @@ import json
import logging
import sqlite3
import time
from typing import Any, Dict, List, Optional
from typing import Any, Dict, List, Optional, Tuple
from hermes_state_common import _sql_session_last_active, is_automatic_end_reason
@@ -27,6 +27,24 @@ def _ended_by_compression(row) -> bool:
return row is not None and row["ended_at"] is not None and row["end_reason"] == "compression"
def _claim_lease_row(conn, table: str, key_col: str, key: str, holder: str, now: float, expires_at: float,
stale) -> Tuple[bool, Optional[str]]:
"""Single-transaction lease claim: DELETE a stale holder's row (``stale(holder,
expires_at)``), INSERT OR IGNORE ours, then SELECT to confirm ownership (INSERT OR
IGNORE gives no rowcount signal). Returns ``(acquired, reclaimed_holder)``."""
reclaimed_holder = None
row = conn.execute(f"SELECT holder, expires_at FROM {table} WHERE {key_col} = ?", (key,)).fetchone()
if row is not None and stale(row["holder"], row["expires_at"]):
conn.execute(f"DELETE FROM {table} WHERE {key_col} = ? AND holder = ?", (key, row["holder"]))
reclaimed_holder = row["holder"]
conn.execute(
f"INSERT OR IGNORE INTO {table} ({key_col}, holder, acquired_at, expires_at) VALUES (?, ?, ?, ?)",
(key, holder, now, expires_at),
)
owner = conn.execute(f"SELECT holder FROM {table} WHERE {key_col} = ?", (key,)).fetchone()
return owner is not None and owner["holder"] == holder, reclaimed_holder
class SessionCompressionMixin:
"""Compression lineage, cooldown/streak counters, locks and turn leases."""
@@ -102,8 +120,7 @@ class SessionCompressionMixin:
return False
updated = conn.execute(
"UPDATE sessions SET ended_at = NULL, end_reason = NULL "
"WHERE id = ? AND ended_at IS NOT NULL "
"AND end_reason = 'compression'",
"WHERE id = ? AND ended_at IS NOT NULL AND end_reason = 'compression'",
(session_id,),
)
# rowcount==1 is guaranteed by the parent SELECT in this same txn. A False
@@ -115,7 +132,10 @@ class SessionCompressionMixin:
def _publish_child_session_row(self, conn, parent, *, parent_session_id, child_session_id, source,
model, model_config, system_prompt, cwd, profile_name) -> None:
"""INSERT the compression child's ``sessions`` row copied from *parent*."""
"""INSERT the compression child's ``sessions`` row copied from *parent*. Same contract
as _insert_session_row's compression-fork backfill: the child stays on the parent's
profile and keeps gateway routing/origin columns; no owner on either side -> this
store's profile."""
system_prompt_hash = self._store_system_prompt(conn, system_prompt)
conn.execute(
"""INSERT INTO sessions (
@@ -126,27 +146,12 @@ class SessionCompressionMixin:
thread_id, display_name, origin_json, started_at
) VALUES (?, ?, ?, ?, NULL, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
(
child_session_id,
source,
model,
json.dumps(model_config) if model_config else None,
system_prompt_hash,
parent_session_id,
cwd or parent["cwd"],
parent["git_branch"],
child_session_id, source, model, json.dumps(model_config) if model_config else None,
system_prompt_hash, parent_session_id, cwd or parent["cwd"], parent["git_branch"],
parent["git_repo_root"],
# Same contract as _insert_session_row's compression-fork backfill: the
# child stays on the parent's profile and keeps gateway routing/origin
# columns; no owner on either side -> this store's profile.
profile_name or parent["profile_name"] or self._own_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(),
parent["user_id"], parent["session_key"], parent["chat_id"], parent["chat_type"],
parent["thread_id"], parent["display_name"], parent["origin_json"], time.time(),
),
)
@@ -209,8 +214,7 @@ class SessionCompressionMixin:
# Deliberate boundaries still fail closed.
if is_automatic_end_reason(parent["end_reason"]):
conn.execute(
"UPDATE sessions SET ended_at = NULL, end_reason = NULL "
"WHERE id = ?",
"UPDATE sessions SET ended_at = NULL, end_reason = NULL WHERE id = ?",
(parent_session_id,),
)
else:
@@ -232,11 +236,9 @@ class SessionCompressionMixin:
_ceiling_clause = " AND id <= ?"
_params.append(int(watermark_ceiling))
tail_ids, tail_tool_calls = self._tail_rows_after_watermark(
conn,
"SELECT id, tool_calls FROM messages "
conn, "SELECT id, tool_calls FROM messages "
"WHERE session_id = ? AND active = 1 AND id > ?"
f"{_ceiling_clause} ORDER BY id",
_params,
f"{_ceiling_clause} ORDER BY id", _params,
)
if tail_ids:
self._clone_message_rows(conn, tail_ids, session_id=child_session_id)
@@ -248,8 +250,7 @@ class SessionCompressionMixin:
)
updated = conn.execute(
"UPDATE sessions SET ended_at = ?, end_reason = 'compression' "
"WHERE id = ? AND ended_at IS NULL",
(time.time(), parent_session_id),
"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}")
@@ -318,8 +319,7 @@ class SessionCompressionMixin:
def _do(conn):
cursor = conn.execute(
"UPDATE sessions SET compression_failure_cooldown_until = ?, "
"compression_failure_error = ? WHERE id = ?",
(deadline, error, session_id),
"compression_failure_error = ? WHERE id = ?", (deadline, error, session_id),
)
if cursor.rowcount != 1:
raise RuntimeError(f"compression cooldown rollback session missing: {session_id}")
@@ -344,8 +344,7 @@ class SessionCompressionMixin:
self._write_sql_logged(
"clear_compression_failure_cooldown", session_id,
"UPDATE sessions SET compression_failure_cooldown_until = NULL, "
"compression_failure_error = NULL WHERE id = ?",
(session_id,),
"compression_failure_error = NULL WHERE id = ?", (session_id,),
)
def _read_session_number(self, column: str, session_id: str, cast: type, zero: Any) -> Any:
@@ -419,8 +418,7 @@ class SessionCompressionMixin:
expires_at = time.time() + ttl_seconds
try:
return self._write_rowcount(
"UPDATE compression_locks SET expires_at = ? "
"WHERE session_id = ? AND holder = ?",
"UPDATE compression_locks SET expires_at = ? WHERE session_id = ? AND holder = ?",
(expires_at, session_id, holder),
) > 0
except sqlite3.Error as exc:
@@ -442,25 +440,10 @@ class SessionCompressionMixin:
expires_at = now + ttl_seconds
def _do(conn):
reclaimed_holder = None
row = conn.execute(_LOCK_ROW_SQL, (session_id,)).fetchone()
if row is not None:
current_holder, current_expires_at = row[0], row[1]
if current_expires_at < now or _compression_lock_holder_process_is_dead(current_holder):
conn.execute(
"DELETE FROM compression_locks "
"WHERE session_id = ? AND holder = ?",
(session_id, current_holder),
)
reclaimed_holder = current_holder
conn.execute(
"INSERT OR IGNORE INTO compression_locks "
"(session_id, holder, acquired_at, expires_at) "
"VALUES (?, ?, ?, ?)",
(session_id, holder, now, expires_at),
return _claim_lease_row(
conn, "compression_locks", "session_id", session_id, holder, now, expires_at,
lambda h, e: e < now or _compression_lock_holder_process_is_dead(h),
)
row = conn.execute("SELECT holder FROM compression_locks WHERE session_id = ?", (session_id,)).fetchone()
return row is not None and row[0] == holder, reclaimed_holder
try:
acquired, reclaimed_holder = self._execute_write(_do)
@@ -480,8 +463,7 @@ class SessionCompressionMixin:
return
self._write_sql_logged(
"release_compression_lock", session_id,
"DELETE FROM compression_locks "
"WHERE session_id = ? AND holder = ?",
"DELETE FROM compression_locks WHERE session_id = ? AND holder = ?",
(session_id, holder),
)
@@ -497,8 +479,7 @@ class SessionCompressionMixin:
def _row(sid: str):
row = conn.execute(
"SELECT id, parent_session_id, source, model_config, end_reason "
"FROM sessions WHERE id = ?",
(sid,),
"FROM sessions WHERE id = ?", (sid,),
).fetchone()
return dict(row) if row else None
@@ -537,29 +518,10 @@ class SessionCompressionMixin:
def _do(conn):
conversation_id = self._session_turn_lease_key_on_conn(conn, session_id)
row = conn.execute(
"SELECT holder, expires_at FROM session_turn_leases "
"WHERE conversation_id = ?",
(conversation_id,),
).fetchone()
if row is not None:
current_holder = row["holder"]
if float(row["expires_at"]) <= now or _compression_lock_holder_process_is_dead(current_holder):
conn.execute(
"DELETE FROM session_turn_leases "
"WHERE conversation_id = ? AND holder = ?",
(conversation_id, current_holder),
)
conn.execute(
"INSERT OR IGNORE INTO session_turn_leases "
"(conversation_id, holder, acquired_at, expires_at) "
"VALUES (?, ?, ?, ?)",
(conversation_id, holder, now, expires_at),
)
owner = conn.execute(
"SELECT holder FROM session_turn_leases WHERE conversation_id = ?", (conversation_id,),
).fetchone()
return owner is not None and owner["holder"] == holder
return _claim_lease_row(
conn, "session_turn_leases", "conversation_id", conversation_id, holder, now, expires_at,
lambda h, e: float(e) <= now or _compression_lock_holder_process_is_dead(h),
)[0]
return bool(self._execute_write(_do, patience_s=patience_s))
@@ -622,8 +584,7 @@ class SessionCompressionMixin:
conversation_id = self._session_turn_lease_key_on_conn(conn, session_id)
cursor = conn.execute(
"UPDATE session_turn_leases SET expires_at = ? "
"WHERE conversation_id = ? AND holder = ?",
(expires_at, conversation_id, holder),
"WHERE conversation_id = ? AND holder = ?", (expires_at, conversation_id, holder),
)
return cursor.rowcount > 0
@@ -637,8 +598,7 @@ class SessionCompressionMixin:
def _do(conn):
conversation_id = self._session_turn_lease_key_on_conn(conn, session_id)
conn.execute(
"DELETE FROM session_turn_leases "
"WHERE conversation_id = ? AND holder = ?",
"DELETE FROM session_turn_leases WHERE conversation_id = ? AND holder = ?",
(conversation_id, holder),
)
@@ -649,8 +609,7 @@ class SessionCompressionMixin:
if not session_id:
return None
row = self._read_one(
"SELECT holder FROM compression_locks "
"WHERE session_id = ? AND expires_at >= ?",
"SELECT holder FROM compression_locks WHERE session_id = ? AND expires_at >= ?",
(session_id, time.time()),
)
return None if row is None else row[0]
@@ -660,10 +619,7 @@ class SessionCompressionMixin:
has messages, no end_reason/ended_at, api_call_count=0, older than 7 days) as
``orphaned_compression``. Non-destructive."""
cutoff = time.time() - 604800 # 7 days
def _do(conn):
now = time.time()
result = conn.execute(
return self._write_rowcount(
"""
UPDATE sessions
SET ended_at = ?,
@@ -684,11 +640,8 @@ class SessionCompressionMixin:
WHERE m.session_id = sessions.id
)
""",
(now, cutoff),
)
return result.rowcount
return self._execute_write(_do) or 0
(time.time(), cutoff),
) or 0
def get_compression_chain(self, session_id: str) -> List[str]:
"""Walk the compression-continuation chain forward: root-first through the tip
+224 -397
View File
File diff suppressed because it is too large Load Diff
+5 -10
View File
@@ -79,8 +79,7 @@ class SessionTitlesMixin:
nothing overwrites a user name, re-running the titler on an llm row is a no-op).
No writer may move a hidden canonical Bot Chat off its title. Read and write are
one compare-and-swap transaction, so a manual ``/title`` racing an in-flight
generation is not clobbered.
"""
generation is not clobbered."""
title = self.sanitize_title(title)
is_user = source == self.TITLE_SOURCE_USER
new_rank = self._title_rank(source) if not is_user else None
@@ -166,20 +165,16 @@ class SessionTitlesMixin:
if source not in self._TITLE_SOURCE_RANK:
raise ValueError(f"invalid title source: {source!r}")
return self._write_rowcount(
"UPDATE sessions SET title_source = ? "
"WHERE id = ? AND title IS NOT NULL",
"UPDATE sessions SET title_source = ? WHERE id = ? AND title IS NOT NULL",
(source, session_id),
) > 0
def get_session_by_title(self, title: str) -> Optional[Dict[str, Any]]:
"""Look up a session by exact title. Returns session dict or None."""
row = self._read_one(
"SELECT s.*, "
"COALESCE(sp.prompt, s.system_prompt) AS _system_prompt_resolved "
"FROM sessions s "
"LEFT JOIN system_prompts sp ON sp.hash = s.system_prompt_hash "
"WHERE s.title = ?",
(title,),
"SELECT s.*, COALESCE(sp.prompt, s.system_prompt) AS _system_prompt_resolved "
"FROM sessions s LEFT JOIN system_prompts sp ON sp.hash = s.system_prompt_hash "
"WHERE s.title = ?", (title,),
)
return self._session_row_dict(row) if row else None
+3 -6
View File
@@ -245,8 +245,7 @@ class SessionUsageMixin:
# Writer stuck mid-apply: leave deltas unapplied rather than race it.
logger.warning(
"async token accounting: writer did not stop within %.0fs; "
"%d queued delta(s) not persisted",
join_timeout, len(self._token_queue),
"%d queued delta(s) not persisted", join_timeout, len(self._token_queue),
)
return
# Writer gone: apply leftovers synchronously under the same busy protocol. Wait
@@ -362,12 +361,10 @@ class SessionUsageMixin:
caller's write txn after the ``sessions`` UPDATE. A missing model/provider falls
back to the session row (same COALESCE behaviour as the summary update) — except
for aux rows (``task`` set), which must NOT inherit the main-loop route (vision
on gemini while the main loop runs anthropic): missing info stays 'unknown'/empty.
"""
on gemini while the main loop runs anthropic): missing info stays 'unknown'/empty."""
row = conn.execute(
"SELECT model, billing_provider, billing_base_url, billing_mode "
"FROM sessions WHERE id = ?",
(session_id,),
"FROM sessions WHERE id = ?", (session_id,),
).fetchone()
sess = dict(row) if (row is not None and not task) else {}
eff_model = model or sess.get("model") or "unknown"