From 6777ffa4bb67310bc5a7044892cc7cf48bdb685e Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:46:15 -0700 Subject: [PATCH] feat: reuse transactional rewind with generation and receipt fences --- hermes_state_messages.py | 35 +++--------------- hermes_state_mutation_transcript.py | 38 ++++++++++++++++++++ hermes_state_mutations.py | 19 +++++++++- hermes_state_runtime.py | 2 +- tests/state/test_runtime_rewind_mutations.py | 34 ++++++++++++++++++ 5 files changed, 95 insertions(+), 33 deletions(-) create mode 100644 hermes_state_mutation_transcript.py create mode 100644 tests/state/test_runtime_rewind_mutations.py diff --git a/hermes_state_messages.py b/hermes_state_messages.py index 796ad49a10..01940aa3e2 100644 --- a/hermes_state_messages.py +++ b/hermes_state_messages.py @@ -1010,38 +1010,11 @@ class SessionMessagesMixin: ``expected_active_ids`` / ``expected_target_content`` pin the active set and canonical live payload in-txn before any mutation (presentation-only metadata changes don't invalidate a rewind). A live turn lease refuses; expired/dead holders are reclaimed. ``rewind_count`` always increments.""" + from hermes_state_mutation_transcript import rewind_in_transaction def _do(conn): - self._check_transcript_write_guards( - conn, session_id, None, reject_active_turn_lease=True, reject_active_compression_lock=True) - if expected_active_ids is not None: - active_rows = conn.execute(_ACTIVE_IDS_SQL, (session_id,)).fetchall() - if [int(r[0]) for r in active_rows] != expected_active_ids: - raise RuntimeError("active transcript changed before the rewind could be persisted") - row = conn.execute( - "SELECT * FROM messages WHERE id = ? AND session_id = ?", (target_message_id, session_id)).fetchone() - if row is None: - raise ValueError(f"message {target_message_id} not found in session {session_id}") - target_row = dict(row) - if target_row.get("role") != "user": - raise ValueError( - f"rewind target must be a 'user' message (got role={target_row.get('role')!r}, id={target_message_id})") - replacement_message_id = replacement = None - if preserve_compaction_handoff or expected_target_content is not None: - replacement = self._split_rewind_target(target_row, expected_target_content, preserve_compaction_handoff) - ids = [r[0] for r in conn.execute("SELECT id FROM messages WHERE session_id = ? AND id >= ? AND active = 1", - (session_id, target_message_id)).fetchall()] - if ids: - conn.execute(f"UPDATE messages SET active = 0 WHERE id IN ({_placeholders(ids)})", ids) - if replacement is not None: - self._insert_message_rows(conn, session_id, [replacement]) - replacement_message_id = int(conn.execute("SELECT last_insert_rowid()").fetchone()[0]) - conn.execute( - "UPDATE sessions SET rewind_count = COALESCE(rewind_count, 0) + 1 WHERE id = ?", (session_id,)) - message_count, tool_call_count = self._active_transcript_counts(conn, session_id) - conn.execute(f"{_SET_COUNTERS_SQL} WHERE id = ?", (message_count, tool_call_count, session_id)) - head_id = conn.execute( - "SELECT MAX(id) FROM messages WHERE session_id = ? AND active = 1", (session_id,)).fetchone()[0] - return target_row, ids, head_id, replacement_message_id + return rewind_in_transaction(self, conn, session_id, target_message_id, + preserve_compaction_handoff=preserve_compaction_handoff, + expected_active_ids=expected_active_ids, expected_target_content=expected_target_content) target_row, rewound, new_head_id, replacement_message_id = self._execute_write(_do) # Decode for the prompt-buffer prefill without a second fallible DB operation. target_row["content"] = self._decode_content(target_row.get("content")) diff --git a/hermes_state_mutation_transcript.py b/hermes_state_mutation_transcript.py new file mode 100644 index 0000000000..b8187430fd --- /dev/null +++ b/hermes_state_mutation_transcript.py @@ -0,0 +1,38 @@ +"""Transaction-scoped rewind shared by storage and canonical mutations.""" +from hermes_state_messages import _ACTIVE_IDS_SQL, _SET_COUNTERS_SQL, _placeholders + + +def rewind_in_transaction(self, conn, session_id, target_message_id, *, + preserve_compaction_handoff=False, expected_active_ids=None, + expected_target_content=None): + self._check_transcript_write_guards( + conn, session_id, None, reject_active_turn_lease=True, reject_active_compression_lock=True) + if expected_active_ids is not None: + active_rows = conn.execute(_ACTIVE_IDS_SQL, (session_id,)).fetchall() + if [int(r[0]) for r in active_rows] != expected_active_ids: + raise RuntimeError("active transcript changed before the rewind could be persisted") + row = conn.execute( + "SELECT * FROM messages WHERE id = ? AND session_id = ?", (target_message_id, session_id)).fetchone() + if row is None: + raise ValueError(f"message {target_message_id} not found in session {session_id}") + target_row = dict(row) + if target_row.get("role") != "user": + raise ValueError( + f"rewind target must be a 'user' message (got role={target_row.get('role')!r}, id={target_message_id})") + replacement_message_id = replacement = None + if preserve_compaction_handoff or expected_target_content is not None: + replacement = self._split_rewind_target(target_row, expected_target_content, preserve_compaction_handoff) + ids = [r[0] for r in conn.execute("SELECT id FROM messages WHERE session_id = ? AND id >= ? AND active = 1", + (session_id, target_message_id)).fetchall()] + if ids: + conn.execute(f"UPDATE messages SET active = 0 WHERE id IN ({_placeholders(ids)})", ids) + if replacement is not None: + self._insert_message_rows(conn, session_id, [replacement]) + replacement_message_id = int(conn.execute("SELECT last_insert_rowid()").fetchone()[0]) + conn.execute( + "UPDATE sessions SET rewind_count = COALESCE(rewind_count, 0) + 1 WHERE id = ?", (session_id,)) + message_count, tool_call_count = self._active_transcript_counts(conn, session_id) + conn.execute(f"{_SET_COUNTERS_SQL} WHERE id = ?", (message_count, tool_call_count, session_id)) + head_id = conn.execute( + "SELECT MAX(id) FROM messages WHERE session_id = ? AND active = 1", (session_id,)).fetchone()[0] + return target_row, ids, head_id, replacement_message_id diff --git a/hermes_state_mutations.py b/hermes_state_mutations.py index 73e058ba55..758bb07262 100644 --- a/hermes_state_mutations.py +++ b/hermes_state_mutations.py @@ -12,6 +12,11 @@ def validate_action(operation, payload): return if operation == 'import' and set(payload) == {'sessions'} and isinstance(payload['sessions'], list): return + if (operation == 'rewind' and 'target_message_id' in payload + and not set(payload) - {'target_message_id', 'preserve_compaction_handoff'} + and type(payload['target_message_id']) is int and payload['target_message_id'] > 0 + and type(payload.get('preserve_compaction_handoff', False)) is bool): + return required = {'rename': {'title'}, 'archive': {'archived'}} if operation == 'sidebar': valid = bool(payload) and not set(payload) - METADATA_FIELDS.keys() @@ -22,7 +27,7 @@ def validate_action(operation, payload): def apply_action(db, conn, session_id, operation, payload): - handlers = {'delete': _delete, 'import': _import} + handlers = {'delete': _delete, 'import': _import, 'rewind': _rewind} if operation in handlers: return handlers[operation](db, conn, session_id, payload) affected = set() @@ -58,6 +63,18 @@ def _delete(db, conn, session_id, payload): return set(), {'deleted_ids': targets} +def _rewind(db, conn, session_id, payload): + from hermes_state_mutation_guards import require_idle + from hermes_state_mutation_transcript import rewind_in_transaction + require_idle(db, conn, [session_id]) + target, ids, head, replacement = rewind_in_transaction(db, conn, session_id, + payload['target_message_id'], preserve_compaction_handoff=payload.get('preserve_compaction_handoff', False)) + target['content'] = db._decode_content(target.get('content')) + conn.execute('UPDATE sessions SET runtime_generation=runtime_generation+1 WHERE id=?', (session_id,)) + return {session_id}, {'rewound_count': len(ids), 'target_message': target, + 'new_head_id': head, 'replacement_message_id': replacement} + + def _import(db, conn, session_id, payload): rows = payload['sessions'] if len(rows) > db._IMPORT_MAX_SESSIONS: diff --git a/hermes_state_runtime.py b/hermes_state_runtime.py index ca0454c0f5..7cfbb26f48 100644 --- a/hermes_state_runtime.py +++ b/hermes_state_runtime.py @@ -188,7 +188,7 @@ def mutate_runtime_session(db, *, epoch: int, principal_id: str, session_id: str validate_action(operation, payload) if expected_generation is not None and (type(expected_generation) is not int or expected_generation < 0): raise RuntimeStoreError('invalid_params') - if operation == 'delete' and expected_generation is None: + if operation in {'delete', 'rewind'} and expected_generation is None: raise RuntimeStoreError('invalid_params') # Snapshot caller data before waiting for the writer lock. payload = json.loads(_json(payload)) diff --git a/tests/state/test_runtime_rewind_mutations.py b/tests/state/test_runtime_rewind_mutations.py new file mode 100644 index 0000000000..967859dde5 --- /dev/null +++ b/tests/state/test_runtime_rewind_mutations.py @@ -0,0 +1,34 @@ +"""Rewind fences the generation and preserves archived history.""" +import pytest +from hermes_state import SessionDB +import hermes_state_runtime as rt + + +def test_rewind_has_existing_semantics_with_atomic_retry(tmp_path): + with SessionDB(db_path=tmp_path / 'state.db') as db: + db.create_session('s', source='test') + db.append_messages_batch('s', [{'role': 'user', 'content': 'first'}, {'role': 'assistant', 'content': 'reply'}, {'role': 'user', 'content': 'again'}]) + target = db.get_messages('s')[-1]['id'] + epoch = rt.begin_runtime_epoch(db, instance_id='owner') + args = dict(epoch=epoch, principal_id='human', session_id='s', request_id='rewind', + expected_revision=0, expected_generation=0, operation='rewind', payload={'target_message_id': target}) + receipt = rt.mutate_runtime_session(db, **args) + assert receipt['rewound_count'] == 1 and receipt['target_message']['content'] == 'again' + assert [m['content'] for m in db.get_messages('s')] == ['first', 'reply'] + assert db.message_count('s') == 3 + assert db.get_session('s')['runtime_generation'] == 1 + assert rt.mutate_runtime_session(db, **args) == receipt + assert db.get_session('s')['rewind_count'] == 1 + + +def test_rewind_refuses_queued_work(tmp_path): + with SessionDB(db_path=tmp_path / 'state.db') as db: + db.create_session('s', source='test') + db.append_message('s', role='user', content='keep') + target = db.get_messages('s')[0]['id'] + epoch = rt.begin_runtime_epoch(db, instance_id='owner') + rt.admit_session_input(db, epoch=epoch, principal_id='human', session_id='s', request_id='input', payload={'text': 'queued'}) + with pytest.raises(rt.RuntimeStoreError, match='session_busy'): + rt.mutate_runtime_session(db, epoch=epoch, principal_id='human', session_id='s', request_id='rewind', + expected_revision=0, expected_generation=0, operation='rewind', payload={'target_message_id': target}) + assert db.get_messages('s')[0]['content'] == 'keep'