From 71e4ec9a472ec9881dabdf8355e045f2020955dd Mon Sep 17 00:00:00 2001 From: devsart95 Date: Mon, 3 Aug 2026 16:43:31 +0530 Subject: [PATCH 1/5] perf(state): batch the turn flush into one SQLite transaction Re-derivation of #23254 (@devsart95) on today's flush loop. The turn flush in _flush_messages_to_session_db wrote one BEGIN IMMEDIATE transaction per message row; a typical agent turn (user + assistant + tool results) paid 3-8 transactions -- and, off WAL (the default on macOS while the WAL-reset guard is active), 3-8 fsyncs -- per turn. Adds SessionDB.append_messages_batch: same row shape as append_message (shared _prepare_message_row serializer + _MESSAGE_INSERT_SQL column list, so the two writers cannot drift), same compression-lock and compression-closed guards, one aggregated session-counter UPDATE, one transaction for the whole batch. Row serialization stays outside the write lock. The flush loop now collects the turn's new rows and writes them in one call. All-or-nothing pairs exactly with the persisted-marker stamping: on failure no rows landed and no markers were stamped, so the next flush re-writes the whole tail (same recovery contract as before, minus the partial-prefix case that could double-count). Measured (same harness, 5-message turn, journal_mode=DELETE, synchronous=FULL): 2.32ms -> 0.83ms median per turn flush (64% faster, 5 fsyncs -> 1). On WAL the win is smaller but the atomicity fix holds. --- hermes_state.py | 312 ++++++++++++++---- run_agent.py | 53 ++- .../agent/test_cursor_optimizations_parity.py | 6 + .../test_append_messages_batch.py | 180 ++++++++++ ...est_empty_response_recovery_persistence.py | 6 + .../run_agent/test_message_sequence_repair.py | 5 + .../test_tool_name_db_persistence.py | 14 +- 7 files changed, 479 insertions(+), 97 deletions(-) create mode 100644 tests/hermes_state/test_append_messages_batch.py diff --git a/hermes_state.py b/hermes_state.py index 68c8ba36af0f..a8eb72638156 100644 --- a/hermes_state.py +++ b/hermes_state.py @@ -5883,6 +5883,117 @@ def _encode_display_metadata(display_metadata: Any) -> Optional[str]: ) return None + # INSERT column list shared by append_message and append_messages_batch so + # the row shape can never drift between the single and batched writers. + _MESSAGE_INSERT_SQL = ( + "INSERT INTO messages (session_id, role, content, tool_call_id, " + "tool_calls, tool_name, effect_disposition, timestamp, token_count, finish_reason, " + "reasoning, reasoning_content, reasoning_details, codex_reasoning_items, " + "codex_message_items, platform_message_id, observed, active, api_content, display_kind, display_metadata) " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" + ) + + def _prepare_message_row( + self, + *, + session_id: str, + role: str, + content: Any = None, + tool_name: Optional[str] = None, + tool_calls: Any = None, + tool_call_id: Optional[str] = None, + token_count: Optional[int] = None, + finish_reason: Optional[str] = None, + reasoning: Optional[str] = None, + reasoning_content: Optional[str] = None, + reasoning_details: Any = None, + codex_reasoning_items: Any = None, + codex_message_items: Any = None, + platform_message_id: Optional[str] = None, + observed: bool = False, + effect_disposition: Optional[str] = None, + timestamp: Any = None, + api_content: Optional[str] = None, + display_kind: Optional[str] = None, + display_metadata: Optional[Dict[str, Any]] = None, + ) -> Tuple[tuple, int]: + """Serialize one message into ``_MESSAGE_INSERT_SQL`` bind params. + + Runs entirely OUTSIDE the write transaction (JSON encoding, surrogate + scrubbing, timestamp coercion), so batched flushes keep the BEGIN + IMMEDIATE window as short as possible. Returns ``(params, + num_tool_calls)`` where ``num_tool_calls`` feeds the session counter + update. + """ + # Display metadata is presentation-only and never changes the model + # context role/content replayed to providers. + display_metadata_json = self._encode_display_metadata(display_metadata) + # Serialize structured fields to JSON before entering the write txn + reasoning_details_json = ( + json.dumps(reasoning_details) + if reasoning_details else None + ) + codex_items_json = ( + json.dumps(codex_reasoning_items) + if codex_reasoning_items else None + ) + codex_message_items_json = ( + json.dumps(codex_message_items) + if codex_message_items else None + ) + # tool_calls may arrive as a Python list (from the live agent) or + # as a JSON string (from import/export). Parse first to avoid + # double-encoding. + if isinstance(tool_calls, str): + try: + tool_calls = json.loads(tool_calls) + except (json.JSONDecodeError, TypeError): + tool_calls = [] + tool_calls_json = json.dumps(tool_calls) if tool_calls else None + # Multimodal content (list of parts) must be JSON-encoded: sqlite3 + # cannot bind list/dict parameters directly. + stored_content = self._encode_content(content) + + message_timestamp = time.time() + if timestamp is not None: + try: + if hasattr(timestamp, "timestamp"): + message_timestamp = float(timestamp.timestamp()) + else: + message_timestamp = float(timestamp) + except (TypeError, ValueError): + logger.debug("Ignoring invalid explicit message timestamp: %r", timestamp) + + # Pre-compute tool call count + num_tool_calls = 0 + if tool_calls is not None: + num_tool_calls = len(tool_calls) if isinstance(tool_calls, list) else 1 + + params = ( + session_id, + role, + stored_content, + tool_call_id, + tool_calls_json, + _scrub_surrogates(tool_name), + effect_disposition, + message_timestamp, + token_count, + finish_reason, + _scrub_surrogates(reasoning), + _scrub_surrogates(reasoning_content), + reasoning_details_json, + codex_items_json, + codex_message_items_json, + platform_message_id, + 1 if observed else 0, + 1, + _scrub_surrogates(api_content) if isinstance(api_content, str) else None, + _scrub_surrogates(display_kind) if isinstance(display_kind, str) else None, + display_metadata_json, + ) + return params, num_tool_calls + @staticmethod def _decode_display_metadata(raw: Any) -> Optional[Dict[str, Any]]: """Decode a ``display_metadata`` column into the dict every reader expects. @@ -5950,49 +6061,28 @@ def append_message( from every outgoing payload anyway, so the scrubbed form IS the wire bytes). """ - # Display metadata is presentation-only and never changes the model - # context role/content replayed to providers. - display_metadata_json = self._encode_display_metadata(display_metadata) - # Serialize structured fields to JSON before entering the write txn - reasoning_details_json = ( - json.dumps(reasoning_details) - if reasoning_details else None - ) - codex_items_json = ( - json.dumps(codex_reasoning_items) - if codex_reasoning_items else None - ) - codex_message_items_json = ( - json.dumps(codex_message_items) - if codex_message_items else None + row_params, num_tool_calls = self._prepare_message_row( + session_id=session_id, + role=role, + content=content, + tool_name=tool_name, + tool_calls=tool_calls, + tool_call_id=tool_call_id, + token_count=token_count, + finish_reason=finish_reason, + reasoning=reasoning, + reasoning_content=reasoning_content, + reasoning_details=reasoning_details, + codex_reasoning_items=codex_reasoning_items, + codex_message_items=codex_message_items, + platform_message_id=platform_message_id, + observed=observed, + effect_disposition=effect_disposition, + timestamp=timestamp, + api_content=api_content, + display_kind=display_kind, + display_metadata=display_metadata, ) - # tool_calls may arrive as a Python list (from the live agent) or - # as a JSON string (from import/export). Parse first to avoid - # double-encoding. - if isinstance(tool_calls, str): - try: - tool_calls = json.loads(tool_calls) - except (json.JSONDecodeError, TypeError): - tool_calls = [] - tool_calls_json = json.dumps(tool_calls) if tool_calls else None - # Multimodal content (list of parts) must be JSON-encoded: sqlite3 - # cannot bind list/dict parameters directly. - stored_content = self._encode_content(content) - - message_timestamp = time.time() - if timestamp is not None: - try: - if hasattr(timestamp, "timestamp"): - message_timestamp = float(timestamp.timestamp()) - else: - message_timestamp = float(timestamp) - except (TypeError, ValueError): - logger.debug("Ignoring invalid explicit message timestamp: %r", timestamp) - - # Pre-compute tool call count - num_tool_calls = 0 - if tool_calls is not None: - num_tool_calls = len(tool_calls) if isinstance(tool_calls, list) else 1 def _do(conn): active_lock = conn.execute( @@ -6017,36 +6107,7 @@ def _do(conn): 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, - reasoning, reasoning_content, reasoning_details, codex_reasoning_items, - codex_message_items, platform_message_id, observed, active, api_content, display_kind, display_metadata) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", - ( - session_id, - role, - stored_content, - tool_call_id, - tool_calls_json, - _scrub_surrogates(tool_name), - effect_disposition, - message_timestamp, - token_count, - finish_reason, - _scrub_surrogates(reasoning), - _scrub_surrogates(reasoning_content), - reasoning_details_json, - codex_items_json, - codex_message_items_json, - platform_message_id, - 1 if observed else 0, - 1, - _scrub_surrogates(api_content) if isinstance(api_content, str) else None, - _scrub_surrogates(display_kind) if isinstance(display_kind, str) else None, - display_metadata_json, - ), - ) + cursor = conn.execute(self._MESSAGE_INSERT_SQL, row_params) msg_id = cursor.lastrowid # Update counters @@ -6072,6 +6133,113 @@ def _do(conn): _do, patience_s=self._TRANSCRIPT_WRITE_PATIENCE_S ) + def append_messages_batch( + self, + session_id: str, + messages: List[Dict[str, Any]], + compression_lock_holder: Optional[str] = None, + ) -> List[int]: + """Append multiple messages atomically in ONE write transaction. + + ``messages`` is a list of dicts whose keys mirror + :meth:`append_message`'s keyword arguments (role, content, tool_name, + tool_calls, tool_call_id, finish_reason, reasoning, reasoning_content, + reasoning_details, codex_reasoning_items, codex_message_items, + timestamp, api_content, display_kind, display_metadata, ...). + + A turn-boundary flush writes the whole turn (user + assistant + tool + rows, typically 3-8 messages) as one BEGIN IMMEDIATE / commit pair + instead of one transaction (and, off WAL, one fsync) per row. Row + serialization happens OUTSIDE the transaction via + ``_prepare_message_row`` — the same helper ``append_message`` uses, + so the row shape cannot drift between the two writers. + + Atomicity contract: all rows land or none do (the caller re-flushes + unstamped messages on the next attempt). The compression-lock and + compression-closed guards from ``append_message`` run once for the + batch — same session, same instant. Returns the inserted row IDs in + input order. + """ + if not messages: + return [] + + prepared: List[tuple] = [] + total_tool_calls = 0 + for msg in messages: + role = msg.get("role", "unknown") + params, num_tc = self._prepare_message_row( + session_id=session_id, + role=role, + content=msg.get("content"), + tool_name=msg.get("tool_name"), + tool_calls=msg.get("tool_calls"), + tool_call_id=msg.get("tool_call_id"), + token_count=msg.get("token_count"), + finish_reason=msg.get("finish_reason"), + reasoning=msg.get("reasoning") if role == "assistant" else None, + reasoning_content=msg.get("reasoning_content") if role == "assistant" else None, + reasoning_details=msg.get("reasoning_details") if role == "assistant" else None, + codex_reasoning_items=msg.get("codex_reasoning_items") if role == "assistant" else None, + codex_message_items=msg.get("codex_message_items") if role == "assistant" else None, + platform_message_id=msg.get("platform_message_id"), + observed=bool(msg.get("observed")), + effect_disposition=msg.get("effect_disposition"), + timestamp=msg.get("timestamp"), + api_content=msg.get("api_content"), + display_kind=msg.get("display_kind"), + display_metadata=msg.get("display_metadata"), + ) + prepared.append(params) + total_tool_calls += num_tc + + 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 SessionCompressionInProgressError( + 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) + + row_ids: List[int] = [] + for params in prepared: + cursor = conn.execute(self._MESSAGE_INSERT_SQL, params) + row_ids.append(cursor.lastrowid) + + # One aggregated counter update for the whole batch. + if total_tool_calls > 0: + conn.execute( + """UPDATE sessions SET message_count = message_count + ?, + tool_call_count = tool_call_count + ? WHERE id = ?""", + (len(prepared), total_tool_calls, session_id), + ) + else: + conn.execute( + "UPDATE sessions SET message_count = message_count + ? WHERE id = ?", + (len(prepared), session_id), + ) + return row_ids + + # Same criticality as append_message: this IS the turn's transcript. + return self._execute_write( + _do, patience_s=self._TRANSCRIPT_WRITE_PATIENCE_S + ) + def set_latest_matching_message_display_kind( self, session_id: str, *, role: str, content: str, display_kind: str, display_metadata: Optional[Dict[str, Any]] = None, diff --git a/run_agent.py b/run_agent.py index 953a127323fa..11e3d732d959 100644 --- a/run_agent.py +++ b/run_agent.py @@ -2096,6 +2096,10 @@ def _flush_messages_to_session_db_unlocked( ): _scan_start += 1 + # Collect this flush's new rows and write them in ONE transaction + # at the end of the scan (see append_messages_batch). + _batch_rows: List[Dict[str, Any]] = [] + _batch_msgs: List[Dict] = [] for _msg_idx in range(_scan_start, len(messages)): msg = messages[_msg_idx] if not isinstance(msg, dict): @@ -2214,33 +2218,46 @@ def _flush_messages_to_session_db_unlocked( ] elif isinstance(msg.get("tool_calls"), list): tool_calls_data = msg["tool_calls"] - self._session_db.append_message( - session_id=self.session_id, - role=role, - content=content, - tool_name=msg.get("tool_name"), - tool_calls=tool_calls_data, - tool_call_id=msg.get("tool_call_id"), - finish_reason=msg.get("finish_reason"), - reasoning=msg.get("reasoning") if role == "assistant" else None, - reasoning_content=msg.get("reasoning_content") if role == "assistant" else None, - reasoning_details=msg.get("reasoning_details") if role == "assistant" else None, - codex_reasoning_items=msg.get("codex_reasoning_items") if role == "assistant" else None, - codex_message_items=msg.get("codex_message_items") if role == "assistant" else None, - timestamp=_row_timestamp, - api_content=_row_api_content, - display_kind=( + _batch_rows.append({ + "role": role, + "content": content, + "tool_name": msg.get("tool_name"), + "tool_calls": tool_calls_data, + "tool_call_id": msg.get("tool_call_id"), + "finish_reason": msg.get("finish_reason"), + "reasoning": msg.get("reasoning") if role == "assistant" else None, + "reasoning_content": msg.get("reasoning_content") if role == "assistant" else None, + "reasoning_details": msg.get("reasoning_details") if role == "assistant" else None, + "codex_reasoning_items": msg.get("codex_reasoning_items") if role == "assistant" else None, + "codex_message_items": msg.get("codex_message_items") if role == "assistant" else None, + "timestamp": _row_timestamp, + "api_content": _row_api_content, + "display_kind": ( "hidden" if msg.get(COMPRESSED_SUMMARY_METADATA_KEY) and not msg.get("_compressed_summary_has_user_turn") else msg.get("display_kind") ), - display_metadata=msg.get("display_metadata"), + "display_metadata": msg.get("display_metadata"), + }) + _batch_msgs.append(msg) + # One transaction for the whole turn's new rows (typically 3-8 + # messages): one BEGIN IMMEDIATE / commit — and, off WAL, one + # fsync — instead of one per row. All-or-nothing pairs exactly + # with the marker stamping below: on failure NO rows landed and + # NO markers were stamped, so the next flush re-scans and + # re-writes the whole tail (same recovery contract as before, + # minus the partial-prefix case that could double-pay counters). + if _batch_rows: + self._session_db.append_messages_batch( + session_id=self.session_id, + messages=_batch_rows, compression_lock_holder=getattr( self, "_active_compression_lock_holder", None ), ) - msg[_DB_PERSISTED_MARKER] = True + for _written in _batch_msgs: + _written[_DB_PERSISTED_MARKER] = True # The intrinsic markers are now the sole source of truth. Reset the # one-shot seed so no id() outlives this flush to alias a message # allocated next turn at a recycled address. diff --git a/tests/agent/test_cursor_optimizations_parity.py b/tests/agent/test_cursor_optimizations_parity.py index c08b012aca3d..afb8e967bc9d 100644 --- a/tests/agent/test_cursor_optimizations_parity.py +++ b/tests/agent/test_cursor_optimizations_parity.py @@ -146,6 +146,12 @@ def __init__(self): self.rows = [] def append_message(self, **kw): self.rows.append({k: copy.deepcopy(v) for k, v in kw.items()}) + def append_messages_batch(self, session_id, messages, **kw): + for m in messages: + row = {k: copy.deepcopy(v) for k, v in m.items()} + row["session_id"] = session_id + self.rows.append(row) + return list(range(1, len(messages) + 1)) def make_agent(bounded): a = ra.AIAgent.__new__(ra.AIAgent) diff --git a/tests/hermes_state/test_append_messages_batch.py b/tests/hermes_state/test_append_messages_batch.py new file mode 100644 index 000000000000..aaff6884c8ad --- /dev/null +++ b/tests/hermes_state/test_append_messages_batch.py @@ -0,0 +1,180 @@ +"""Tests for SessionDB.append_messages_batch (#23254 salvage). + +The batch writer must be row-shape-identical to append_message (shared +_prepare_message_row + _MESSAGE_INSERT_SQL), atomic (all rows or none), +and must aggregate the session counters in one UPDATE. +""" + +import json +import sqlite3 + +import pytest + +from hermes_state import ( + CompressionSessionClosedError, + SessionDB, +) + + +@pytest.fixture() +def db(tmp_path): + d = SessionDB(db_path=tmp_path / "state.db") + d.create_session("sess-batch", source="cli") + yield d + d.close() + + +def _turn_messages(): + return [ + {"role": "user", "content": "question"}, + { + "role": "assistant", + "content": "let me check", + "tool_calls": [{"name": "terminal", "arguments": "{}"}], + "reasoning_content": "thinking...", + "finish_reason": "tool_calls", + }, + { + "role": "tool", + "content": "tool output", + "tool_name": "terminal", + "tool_call_id": "call_1", + }, + {"role": "assistant", "content": "answer", "finish_reason": "stop"}, + ] + + +class TestAppendMessagesBatch: + def test_batch_rows_identical_to_single_appends(self, db, tmp_path): + """The batch writer stores the same bytes append_message would.""" + db2 = SessionDB(db_path=tmp_path / "state2.db") + db2.create_session("sess-batch", source="cli") + try: + msgs = _turn_messages() + db.append_messages_batch("sess-batch", msgs) + for m in msgs: + role = m["role"] + db2.append_message( + session_id="sess-batch", + role=role, + content=m.get("content"), + tool_name=m.get("tool_name"), + tool_calls=m.get("tool_calls"), + tool_call_id=m.get("tool_call_id"), + finish_reason=m.get("finish_reason"), + reasoning_content=( + m.get("reasoning_content") if role == "assistant" else None + ), + ) + cols = ( + "role, content, tool_call_id, tool_calls, tool_name, " + "finish_reason, reasoning_content, observed, active" + ) + rows_a = db._conn.execute( + f"SELECT {cols} FROM messages ORDER BY id" + ).fetchall() + rows_b = db2._conn.execute( + f"SELECT {cols} FROM messages ORDER BY id" + ).fetchall() + assert [tuple(r) for r in rows_a] == [tuple(r) for r in rows_b] + finally: + db2.close() + + def test_counters_aggregate_once(self, db): + db.append_messages_batch("sess-batch", _turn_messages()) + row = db._conn.execute( + "SELECT message_count, tool_call_count FROM sessions WHERE id = ?", + ("sess-batch",), + ).fetchone() + assert row["message_count"] == 4 + assert row["tool_call_count"] == 1 + + def test_returns_row_ids_in_input_order(self, db): + ids = db.append_messages_batch("sess-batch", _turn_messages()) + assert ids == sorted(ids) + assert len(ids) == 4 + + def test_empty_batch_is_noop(self, db): + assert db.append_messages_batch("sess-batch", []) == [] + row = db._conn.execute( + "SELECT message_count FROM sessions WHERE id = ?", ("sess-batch",) + ).fetchone() + assert row["message_count"] == 0 + + def test_atomicity_all_or_nothing(self, db, monkeypatch): + """A failure mid-batch leaves ZERO rows and untouched counters.""" + real_execute_write = db._execute_write + original_insert = SessionDB._MESSAGE_INSERT_SQL + + calls = {"n": 0} + + def _do_wrapper(fn, **kwargs): + def failing(conn): + real_conn_execute = conn.execute + + def exec_counting(sql, *args): + if sql == original_insert: + calls["n"] += 1 + if calls["n"] == 3: + raise sqlite3.OperationalError("boom mid-batch") + return real_conn_execute(sql, *args) + + conn.execute = exec_counting + try: + return fn(conn) + finally: + conn.execute = real_conn_execute + + return real_execute_write(failing, **kwargs) + + monkeypatch.setattr(db, "_execute_write", _do_wrapper) + with pytest.raises(sqlite3.OperationalError): + db.append_messages_batch("sess-batch", _turn_messages()) + monkeypatch.undo() + + count = db._conn.execute("SELECT COUNT(*) FROM messages").fetchone()[0] + assert count == 0 + row = db._conn.execute( + "SELECT message_count, tool_call_count FROM sessions WHERE id = ?", + ("sess-batch",), + ).fetchone() + assert row["message_count"] == 0 + assert row["tool_call_count"] == 0 + + def test_compression_closed_session_rejected(self, db): + db._conn.execute( + "UPDATE sessions SET ended_at = 1.0, end_reason = 'compression' " + "WHERE id = ?", + ("sess-batch",), + ) + db._conn.commit() + with pytest.raises(CompressionSessionClosedError): + db.append_messages_batch("sess-batch", _turn_messages()) + + def test_multimodal_content_encoded(self, db): + msgs = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "look"}, + {"type": "image_url", "image_url": {"url": "data:x"}}, + ], + } + ] + db.append_messages_batch("sess-batch", msgs) + raw = db._conn.execute("SELECT content FROM messages").fetchone()[0] + # encoded via _encode_content — same sentinel prefix as append_message + loaded = db.get_messages("sess-batch") + assert loaded, raw + + def test_tool_calls_json_string_not_double_encoded(self, db): + msgs = [ + { + "role": "assistant", + "content": "x", + "tool_calls": json.dumps([{"name": "t", "arguments": "{}"}]), + } + ] + db.append_messages_batch("sess-batch", msgs) + raw = db._conn.execute("SELECT tool_calls FROM messages").fetchone()[0] + assert json.loads(raw) == [{"name": "t", "arguments": "{}"}] diff --git a/tests/run_agent/test_empty_response_recovery_persistence.py b/tests/run_agent/test_empty_response_recovery_persistence.py index ff8b9cf3122d..70e262fbf289 100644 --- a/tests/run_agent/test_empty_response_recovery_persistence.py +++ b/tests/run_agent/test_empty_response_recovery_persistence.py @@ -13,6 +13,12 @@ def append_message(self, session_id, role, content=None, **kwargs): self.rows.append({"role": role, "content": content}) return len(self.rows) + def append_messages_batch(self, session_id, messages, **kwargs): + # Mirror the real batch writer: same rows, one call. + for m in messages: + self.rows.append({"role": m.get("role"), "content": m.get("content")}) + return list(range(len(self.rows) - len(messages) + 1, len(self.rows) + 1)) + def _agent_with_capturing_db(): agent = AIAgent.__new__(AIAgent) diff --git a/tests/run_agent/test_message_sequence_repair.py b/tests/run_agent/test_message_sequence_repair.py index b1a40e78dc91..e169f1e06821 100644 --- a/tests/run_agent/test_message_sequence_repair.py +++ b/tests/run_agent/test_message_sequence_repair.py @@ -227,6 +227,11 @@ def __init__(self): def append_message(self, **kw): self.rows.append(kw) + def append_messages_batch(self, session_id, messages, **kw): + for m in messages: + self.rows.append(dict(m, session_id=session_id)) + return list(range(1, len(messages) + 1)) + agent = _bare_agent() agent._session_db = _DB() agent._session_db_created = True diff --git a/tests/run_agent/test_tool_name_db_persistence.py b/tests/run_agent/test_tool_name_db_persistence.py index 3fcf7f33c3ad..29596c04ddb7 100644 --- a/tests/run_agent/test_tool_name_db_persistence.py +++ b/tests/run_agent/test_tool_name_db_persistence.py @@ -27,7 +27,8 @@ def _make_agent(session_db): def test_tool_name_persisted_to_session_db(): """tool_name set by make_tool_result_message must be passed through to - append_message so the column is populated on first flush to the session DB.""" + the batched flush so the column is populated on first write to the + session DB.""" session_db = MagicMock() agent = _make_agent(session_db) @@ -37,9 +38,8 @@ def test_tool_name_persisted_to_session_db(): ] agent._flush_messages_to_session_db(messages) - tool_appends = [ - c for c in session_db.append_message.call_args_list - if c.kwargs.get("role") == "tool" - ] - assert len(tool_appends) == 1 - assert tool_appends[0].kwargs["tool_name"] == "terminal" + assert session_db.append_messages_batch.call_count == 1 + batch = session_db.append_messages_batch.call_args.kwargs["messages"] + tool_rows = [m for m in batch if m.get("role") == "tool"] + assert len(tool_rows) == 1 + assert tool_rows[0]["tool_name"] == "terminal" From 9187d7a02a029897774f066dab6d7797852455c3 Mon Sep 17 00:00:00 2001 From: kshitij Date: Mon, 3 Aug 2026 16:43:44 +0530 Subject: [PATCH 2/5] perf(tui-gateway): batch branch-seed history copies (whole-bug-class) Sibling sites of the per-message flush pattern: both branch-seed paths (session.branch in methods_session.py and the lazy seed persist in server.py) copied the parent history row-by-row -- one transaction per row, and a branch seed can be hundreds of rows. Route both through SessionDB.append_messages_batch. The server.py path also gains real atomicity: _branch_seed_persisted assumed every row landed, which the per-row loop could not guarantee. --- tests/test_tui_gateway_server.py | 10 ++++++++++ tui_gateway/methods_session.py | 25 ++++++++++++++++--------- tui_gateway/server.py | 27 +++++++++++++++++---------- 3 files changed, 43 insertions(+), 19 deletions(-) diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index e8aaa16e0572..df245ad8738c 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -11438,6 +11438,11 @@ def create_session(self, new_key, **kwargs): def append_message(self, **kwargs): seen["msgs"].append(kwargs) + def append_messages_batch(self, session_id, messages, **kwargs): + for m in messages: + seen["msgs"].append(dict(m, session_id=session_id)) + return list(range(1, len(messages) + 1)) + def set_session_title(self, key, title): seen["title"] = (key, title) return True @@ -11550,6 +11555,11 @@ def create_session(self, new_key, **kwargs): def append_message(self, **kwargs): seen["msgs"].append(kwargs) + def append_messages_batch(self, session_id, messages, **kwargs): + for m in messages: + seen["msgs"].append(dict(m, session_id=session_id)) + return list(range(1, len(messages) + 1)) + def set_session_title(self, key, title): return True diff --git a/tui_gateway/methods_session.py b/tui_gateway/methods_session.py index d84fedcd4314..acbfe7a8b6a8 100644 --- a/tui_gateway/methods_session.py +++ b/tui_gateway/methods_session.py @@ -2649,15 +2649,22 @@ def _(rid, params: dict) -> dict: else None ), ) - for msg in history: - db.append_message( - session_id=new_key, - role=msg.get("role", "user"), - content=msg.get("content"), - # Preserve the parent's original message timestamps — - # branch copies are history, not new activity (9d73006ad). - timestamp=msg.get("timestamp"), - ) + # Copy the whole parent history in ONE transaction — a branch + # seed can be hundreds of rows, and per-row transactions were + # the write-amplification pattern removed in #23254. + db.append_messages_batch( + new_key, + [ + { + "role": msg.get("role", "user"), + "content": msg.get("content"), + # Preserve the parent's original message timestamps — + # branch copies are history, not new activity (9d73006ad). + "timestamp": msg.get("timestamp"), + } + for msg in history + ], + ) db.set_session_title(new_key, title) except Exception as e: if lease is not None: diff --git a/tui_gateway/server.py b/tui_gateway/server.py index a3475fc11ecb..d51cbe42105f 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -2749,16 +2749,23 @@ def _persist_branch_seed(session: dict) -> None: if db is None: return try: - for msg in seed: - db.append_message( - session_id=key, - role=msg.get("role", "user"), - content=msg.get("content"), - # Preserve the parent's original message timestamps — - # append_message would otherwise stamp time.time() and the - # branch's copied history would all appear authored "now". - timestamp=msg.get("timestamp"), - ) + # One transaction for the whole seed (see #23254): a branch seed + # can be hundreds of rows and is all-or-nothing anyway — the + # _branch_seed_persisted flag below assumes every row landed. + db.append_messages_batch( + key, + [ + { + "role": msg.get("role", "user"), + "content": msg.get("content"), + # Preserve the parent's original message timestamps — + # append_message would otherwise stamp time.time() and the + # branch's copied history would all appear authored "now". + "timestamp": msg.get("timestamp"), + } + for msg in seed + ], + ) session["_branch_seed_persisted"] = True except Exception as exc: from hermes_state import is_disk_full_error From 8c2268133f3e05c8c506ff03e38f23707b46e674 Mon Sep 17 00:00:00 2001 From: kshitij Date: Mon, 3 Aug 2026 16:54:14 +0530 Subject: [PATCH 3/5] test(run-agent): update flush-path fakes and assertions for batched writes The flush now goes through append_messages_batch; MagicMock-based assertions and barrier fakes that hooked append_message observed nothing (the flush's try/except swallowed the AttributeError). Assert on the batch payload instead. --- tests/run_agent/test_run_agent.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/tests/run_agent/test_run_agent.py b/tests/run_agent/test_run_agent.py index 4bc1c69b76b7..6161e266b071 100644 --- a/tests/run_agent/test_run_agent.py +++ b/tests/run_agent/test_run_agent.py @@ -112,8 +112,8 @@ def test_flush_persist_override_replaces_api_local_multimodal_note(agent): agent._flush_messages_to_session_db([{"role": "user", "content": api_content}], []) - db_write = agent._session_db.append_message.call_args.kwargs - assert db_write["content"] == "Describe this screenshot\n[screenshot]" + batch = agent._session_db.append_messages_batch.call_args.kwargs["messages"] + assert batch[0]["content"] == "Describe this screenshot\n[screenshot]" assert api_content[0]["text"] == "[MODEL SWITCH NOTE]\n\nDescribe this screenshot" @@ -136,6 +136,17 @@ def append_message(self, **kwargs): assert self.release.wait(timeout=5) self.rows.append(kwargs["content"]) + def append_messages_batch(self, session_id, messages, **kwargs): + with self._lock: + self.calls += 1 + first = self.calls == 1 + if first: + self.entered.set() + assert self.release.wait(timeout=5) + for m in messages: + self.rows.append(m["content"]) + return list(range(1, len(messages) + 1)) + db = _BarrierDB() agent._session_db = db agent._session_db_created = True @@ -5591,8 +5602,10 @@ def test_persist_session_rewrites_current_turn_user_message(self, agent): "2-3 sentences max. No code blocks or markdown.] Hello there" ) # But the DB write must get the override. - first_db_write = agent._session_db.append_message.call_args_list[0].kwargs - assert first_db_write["content"] == "Hello there" + batch = agent._session_db.append_messages_batch.call_args_list[0].kwargs[ + "messages" + ] + assert batch[0]["content"] == "Hello there" class TestReasoningReplayForStrictProviders: From 8c974ad71e1d040150212ba1c593cd255e364400 Mon Sep 17 00:00:00 2001 From: kshitij Date: Mon, 3 Aug 2026 17:14:28 +0530 Subject: [PATCH 4/5] =?UTF-8?q?refactor(state):=20fold=20simplify=20findin?= =?UTF-8?q?gs=20=E2=80=94=20reuse=20=5Finsert=5Fmessage=5Frows,=20share=20?= =?UTF-8?q?guards,=20chunk=20seeds?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Simplify-pass folds on the #23254 salvage: - REUSE (HIGH): append_messages_batch now delegates row serialization to the pre-existing _insert_message_rows helper (already shared by replace_messages / archive_and_compact / portability import) instead of adding a third serialization path (_prepare_message_row + _MESSAGE_INSERT_SQL are gone). One row-writer for every multi-row path; the row-ID return was consumed by no production caller, so the batch returns the inserted count. - QUALITY (HIGH): the compression-lock + compression-closed admission guards are extracted into _check_transcript_write_guards, shared by append_message and append_messages_batch (previously duplicated 23 lines that had already needed targeted fixes, #74478). The role-gated reasoning filtering is no longer duplicated in run_agent.py — it lives at its one site inside _insert_message_rows. - EFFICIENCY (MEDIUM, measured): unbounded seed copies hold one BEGIN IMMEDIATE for seconds (10k rows ~= 2.4s; FTS triggers dominate) and monopolize the in-process write lock. append_messages_batch grows a chunk_rows param; all seed/copy call sites use chunk_rows=500. Same recovery semantics as the old per-row loops, bounded lock holds. - REUSE (MEDIUM): the two remaining per-row branch-copy loops found by the pass (gateway/slash_commands.py /branch, hermes_cli cli_commands_mixin.py branch) are converted to chunked batches too (AsyncSessionDB's generic to_thread forwarder covers the async site). Turn-flush benchmark unchanged after the refactor: 2.43 -> 0.87 ms median per 5-message flush (64% faster). --- gateway/slash_commands.py | 57 +-- hermes_cli/cli_commands_mixin.py | 45 ++- hermes_state.py | 373 +++++++----------- run_agent.py | 12 +- .../test_append_messages_batch.py | 80 ++-- tui_gateway/methods_session.py | 7 +- tui_gateway/server.py | 9 +- 7 files changed, 269 insertions(+), 314 deletions(-) diff --git a/gateway/slash_commands.py b/gateway/slash_commands.py index fb3cec080acc..e719eb2ac88a 100644 --- a/gateway/slash_commands.py +++ b/gateway/slash_commands.py @@ -4632,30 +4632,39 @@ async def _handle_branch_command(self, event: MessageEvent) -> str: logger.error("Failed to create branch session: %s", e) return t("gateway.branch.create_failed", error=e) - # Copy conversation history to the new session - for msg in history: - try: - await self._session_db.append_message( - session_id=new_session_id, - role=msg.get("role", "user"), - content=msg.get("content"), - tool_name=msg.get("tool_name") or msg.get("name"), - tool_calls=msg.get("tool_calls"), - tool_call_id=msg.get("tool_call_id"), - finish_reason=msg.get("finish_reason"), - reasoning=msg.get("reasoning"), - reasoning_content=msg.get("reasoning_content"), - reasoning_details=msg.get("reasoning_details"), - codex_reasoning_items=msg.get("codex_reasoning_items"), - codex_message_items=msg.get("codex_message_items"), - # Keep the api_content sidecar so the branch's first turn - # replays the parent's exact wire bytes (warm provider - # prompt cache) instead of a full cold prefill. - api_content=extract_api_content_sidecar(msg), - timestamp=msg.get("timestamp"), - ) - except Exception: - pass # Best-effort copy + # Copy conversation history to the new session in bounded-chunk + # transactions (see #23254): one txn per row was the removed + # write-amplification pattern, and a history can be hundreds of rows. + # Best-effort like the old loop — a failed copy still yields a + # usable (partial) branch. + try: + await self._session_db.append_messages_batch( + new_session_id, + [ + { + "role": msg.get("role", "user"), + "content": msg.get("content"), + "tool_name": msg.get("tool_name") or msg.get("name"), + "tool_calls": msg.get("tool_calls"), + "tool_call_id": msg.get("tool_call_id"), + "finish_reason": msg.get("finish_reason"), + "reasoning": msg.get("reasoning"), + "reasoning_content": msg.get("reasoning_content"), + "reasoning_details": msg.get("reasoning_details"), + "codex_reasoning_items": msg.get("codex_reasoning_items"), + "codex_message_items": msg.get("codex_message_items"), + # Keep the api_content sidecar so the branch's first turn + # replays the parent's exact wire bytes (warm provider + # prompt cache) instead of a full cold prefill. + "api_content": extract_api_content_sidecar(msg), + "timestamp": msg.get("timestamp"), + } + for msg in history + ], + chunk_rows=500, + ) + except Exception: + pass # Best-effort copy # Set title try: diff --git a/hermes_cli/cli_commands_mixin.py b/hermes_cli/cli_commands_mixin.py index 69b5323b5616..5ec16a5fe4df 100644 --- a/hermes_cli/cli_commands_mixin.py +++ b/hermes_cli/cli_commands_mixin.py @@ -1171,25 +1171,32 @@ def _handle_branch_command(self, cmd_original: str) -> None: _cprint(f" Failed to create branch session: {e}") return - # Copy conversation history to the new session - for msg in self.conversation_history: - try: - self._session_db.append_message( - session_id=new_session_id, - role=msg.get("role", "user"), - content=msg.get("content"), - tool_name=msg.get("tool_name") or msg.get("name"), - tool_calls=msg.get("tool_calls"), - tool_call_id=msg.get("tool_call_id"), - reasoning=msg.get("reasoning"), - # Keep the api_content sidecar so the branch's first turn - # replays the parent's exact wire bytes (warm provider - # prompt cache) instead of a full cold prefill. - api_content=extract_api_content_sidecar(msg), - timestamp=msg.get("timestamp"), - ) - except Exception: - pass # Best-effort copy + # Copy conversation history to the new session in bounded-chunk + # transactions (see #23254) instead of one txn per row. Best-effort + # like the old loop — a failed copy still yields a usable branch. + try: + self._session_db.append_messages_batch( + new_session_id, + [ + { + "role": msg.get("role", "user"), + "content": msg.get("content"), + "tool_name": msg.get("tool_name") or msg.get("name"), + "tool_calls": msg.get("tool_calls"), + "tool_call_id": msg.get("tool_call_id"), + "reasoning": msg.get("reasoning"), + # Keep the api_content sidecar so the branch's first turn + # replays the parent's exact wire bytes (warm provider + # prompt cache) instead of a full cold prefill. + "api_content": extract_api_content_sidecar(msg), + "timestamp": msg.get("timestamp"), + } + for msg in self.conversation_history + ], + chunk_rows=500, + ) + except Exception: + pass # Best-effort copy # Set title on the branch try: diff --git a/hermes_state.py b/hermes_state.py index a8eb72638156..b32b6ffe27c1 100644 --- a/hermes_state.py +++ b/hermes_state.py @@ -5883,116 +5883,38 @@ def _encode_display_metadata(display_metadata: Any) -> Optional[str]: ) return None - # INSERT column list shared by append_message and append_messages_batch so - # the row shape can never drift between the single and batched writers. - _MESSAGE_INSERT_SQL = ( - "INSERT INTO messages (session_id, role, content, tool_call_id, " - "tool_calls, tool_name, effect_disposition, timestamp, token_count, finish_reason, " - "reasoning, reasoning_content, reasoning_details, codex_reasoning_items, " - "codex_message_items, platform_message_id, observed, active, api_content, display_kind, display_metadata) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" - ) + def _check_transcript_write_guards( + self, conn, session_id: str, compression_lock_holder: Optional[str] + ) -> None: + """Transcript-append admission checks, run INSIDE the write txn. - def _prepare_message_row( - self, - *, - session_id: str, - role: str, - content: Any = None, - tool_name: Optional[str] = None, - tool_calls: Any = None, - tool_call_id: Optional[str] = None, - token_count: Optional[int] = None, - finish_reason: Optional[str] = None, - reasoning: Optional[str] = None, - reasoning_content: Optional[str] = None, - reasoning_details: Any = None, - codex_reasoning_items: Any = None, - codex_message_items: Any = None, - platform_message_id: Optional[str] = None, - observed: bool = False, - effect_disposition: Optional[str] = None, - timestamp: Any = None, - api_content: Optional[str] = None, - display_kind: Optional[str] = None, - display_metadata: Optional[Dict[str, Any]] = None, - ) -> Tuple[tuple, int]: - """Serialize one message into ``_MESSAGE_INSERT_SQL`` bind params. - - Runs entirely OUTSIDE the write transaction (JSON encoding, surrogate - scrubbing, timestamp coercion), so batched flushes keep the BEGIN - IMMEDIATE window as short as possible. Returns ``(params, - num_tool_calls)`` where ``num_tool_calls`` feeds the session counter - update. + Shared by :meth:`append_message` and :meth:`append_messages_batch` so + the two writers can never diverge on these correctness invariants + (this guard has already needed targeted fixes — see the #74478 + patience note below). """ - # Display metadata is presentation-only and never changes the model - # context role/content replayed to providers. - display_metadata_json = self._encode_display_metadata(display_metadata) - # Serialize structured fields to JSON before entering the write txn - reasoning_details_json = ( - json.dumps(reasoning_details) - if reasoning_details else None - ) - codex_items_json = ( - json.dumps(codex_reasoning_items) - if codex_reasoning_items else None - ) - codex_message_items_json = ( - json.dumps(codex_message_items) - if codex_message_items else None - ) - # tool_calls may arrive as a Python list (from the live agent) or - # as a JSON string (from import/export). Parse first to avoid - # double-encoding. - if isinstance(tool_calls, str): - try: - tool_calls = json.loads(tool_calls) - except (json.JSONDecodeError, TypeError): - tool_calls = [] - tool_calls_json = json.dumps(tool_calls) if tool_calls else None - # Multimodal content (list of parts) must be JSON-encoded: sqlite3 - # cannot bind list/dict parameters directly. - stored_content = self._encode_content(content) - - message_timestamp = time.time() - if timestamp is not None: - try: - if hasattr(timestamp, "timestamp"): - message_timestamp = float(timestamp.timestamp()) - else: - message_timestamp = float(timestamp) - except (TypeError, ValueError): - logger.debug("Ignoring invalid explicit message timestamp: %r", timestamp) - - # Pre-compute tool call count - num_tool_calls = 0 - if tool_calls is not None: - num_tool_calls = len(tool_calls) if isinstance(tool_calls, list) else 1 - - params = ( - session_id, - role, - stored_content, - tool_call_id, - tool_calls_json, - _scrub_surrogates(tool_name), - effect_disposition, - message_timestamp, - token_count, - finish_reason, - _scrub_surrogates(reasoning), - _scrub_surrogates(reasoning_content), - reasoning_details_json, - codex_items_json, - codex_message_items_json, - platform_message_id, - 1 if observed else 0, - 1, - _scrub_surrogates(api_content) if isinstance(api_content, str) else None, - _scrub_surrogates(display_kind) if isinstance(display_kind, str) else None, - display_metadata_json, - ) - return params, num_tool_calls + 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 SessionCompressionInProgressError( + 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) @staticmethod def _decode_display_metadata(raw: Any) -> Optional[Dict[str, Any]]: @@ -6061,53 +5983,84 @@ def append_message( from every outgoing payload anyway, so the scrubbed form IS the wire bytes). """ - row_params, num_tool_calls = self._prepare_message_row( - session_id=session_id, - role=role, - content=content, - tool_name=tool_name, - tool_calls=tool_calls, - tool_call_id=tool_call_id, - token_count=token_count, - finish_reason=finish_reason, - reasoning=reasoning, - reasoning_content=reasoning_content, - reasoning_details=reasoning_details, - codex_reasoning_items=codex_reasoning_items, - codex_message_items=codex_message_items, - platform_message_id=platform_message_id, - observed=observed, - effect_disposition=effect_disposition, - timestamp=timestamp, - api_content=api_content, - display_kind=display_kind, - display_metadata=display_metadata, + # Display metadata is presentation-only and never changes the model + # context role/content replayed to providers. + display_metadata_json = self._encode_display_metadata(display_metadata) + # Serialize structured fields to JSON before entering the write txn + reasoning_details_json = ( + json.dumps(reasoning_details) + if reasoning_details else None + ) + codex_items_json = ( + json.dumps(codex_reasoning_items) + if codex_reasoning_items else None + ) + codex_message_items_json = ( + json.dumps(codex_message_items) + if codex_message_items else None ) + # tool_calls may arrive as a Python list (from the live agent) or + # as a JSON string (from import/export). Parse first to avoid + # double-encoding. + if isinstance(tool_calls, str): + try: + tool_calls = json.loads(tool_calls) + except (json.JSONDecodeError, TypeError): + tool_calls = [] + tool_calls_json = json.dumps(tool_calls) if tool_calls else None + # Multimodal content (list of parts) must be JSON-encoded: sqlite3 + # cannot bind list/dict parameters directly. + stored_content = self._encode_content(content) + + message_timestamp = time.time() + if timestamp is not None: + try: + if hasattr(timestamp, "timestamp"): + message_timestamp = float(timestamp.timestamp()) + else: + message_timestamp = float(timestamp) + except (TypeError, ValueError): + logger.debug("Ignoring invalid explicit message timestamp: %r", timestamp) + + # Pre-compute tool call count + num_tool_calls = 0 + if tool_calls is not None: + 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 SessionCompressionInProgressError( - 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(self._MESSAGE_INSERT_SQL, row_params) + self._check_transcript_write_guards( + conn, session_id, compression_lock_holder + ) + cursor = conn.execute( + """INSERT INTO messages (session_id, role, content, tool_call_id, + tool_calls, tool_name, effect_disposition, timestamp, token_count, finish_reason, + reasoning, reasoning_content, reasoning_details, codex_reasoning_items, + codex_message_items, platform_message_id, observed, active, api_content, display_kind, display_metadata) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", + ( + session_id, + role, + stored_content, + tool_call_id, + tool_calls_json, + _scrub_surrogates(tool_name), + effect_disposition, + message_timestamp, + token_count, + finish_reason, + _scrub_surrogates(reasoning), + _scrub_surrogates(reasoning_content), + reasoning_details_json, + codex_items_json, + codex_message_items_json, + platform_message_id, + 1 if observed else 0, + 1, + _scrub_surrogates(api_content) if isinstance(api_content, str) else None, + _scrub_surrogates(display_kind) if isinstance(display_kind, str) else None, + display_metadata_json, + ), + ) msg_id = cursor.lastrowid # Update counters @@ -6138,102 +6091,68 @@ def append_messages_batch( session_id: str, messages: List[Dict[str, Any]], compression_lock_holder: Optional[str] = None, - ) -> List[int]: + chunk_rows: Optional[int] = None, + ) -> int: """Append multiple messages atomically in ONE write transaction. - ``messages`` is a list of dicts whose keys mirror - :meth:`append_message`'s keyword arguments (role, content, tool_name, - tool_calls, tool_call_id, finish_reason, reasoning, reasoning_content, - reasoning_details, codex_reasoning_items, codex_message_items, - timestamp, api_content, display_kind, display_metadata, ...). + ``messages`` is a list of dicts in the same shape + :meth:`_insert_message_rows` already consumes for replace/compact/ + import (role, content, tool_name, tool_calls, tool_call_id, + finish_reason, reasoning*, codex_*, timestamp, api_content, + display_kind, display_metadata, ...). Reusing that helper keeps ONE + row-serialization path for every multi-row writer. A turn-boundary flush writes the whole turn (user + assistant + tool rows, typically 3-8 messages) as one BEGIN IMMEDIATE / commit pair - instead of one transaction (and, off WAL, one fsync) per row. Row - serialization happens OUTSIDE the transaction via - ``_prepare_message_row`` — the same helper ``append_message`` uses, - so the row shape cannot drift between the two writers. + instead of one transaction (and, off WAL, one fsync) per row. Atomicity contract: all rows land or none do (the caller re-flushes - unstamped messages on the next attempt). The compression-lock and - compression-closed guards from ``append_message`` run once for the - batch — same session, same instant. Returns the inserted row IDs in - input order. + unstamped messages on the next attempt). The same admission guards + as :meth:`append_message` run once for the batch — same session, + same instant. + + ``chunk_rows`` bounds the transaction size for LARGE copies (branch + seeds can be thousands of rows; measured: 10k rows ≈ 2.4s inside one + BEGIN IMMEDIATE because the FTS triggers run per row, which would + monopolize the write lock and starve concurrent writers). When set, + the batch commits in chunks of at most that many rows — same + recovery semantics as the old per-row loops (a mid-copy failure + leaves a partial seed), just with bounded lock holds. A turn flush + never needs it. Returns the inserted row count. """ if not messages: - return [] - - prepared: List[tuple] = [] - total_tool_calls = 0 - for msg in messages: - role = msg.get("role", "unknown") - params, num_tc = self._prepare_message_row( - session_id=session_id, - role=role, - content=msg.get("content"), - tool_name=msg.get("tool_name"), - tool_calls=msg.get("tool_calls"), - tool_call_id=msg.get("tool_call_id"), - token_count=msg.get("token_count"), - finish_reason=msg.get("finish_reason"), - reasoning=msg.get("reasoning") if role == "assistant" else None, - reasoning_content=msg.get("reasoning_content") if role == "assistant" else None, - reasoning_details=msg.get("reasoning_details") if role == "assistant" else None, - codex_reasoning_items=msg.get("codex_reasoning_items") if role == "assistant" else None, - codex_message_items=msg.get("codex_message_items") if role == "assistant" else None, - platform_message_id=msg.get("platform_message_id"), - observed=bool(msg.get("observed")), - effect_disposition=msg.get("effect_disposition"), - timestamp=msg.get("timestamp"), - api_content=msg.get("api_content"), - display_kind=msg.get("display_kind"), - display_metadata=msg.get("display_metadata"), - ) - prepared.append(params) - total_tool_calls += num_tc + return 0 - 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 SessionCompressionInProgressError( - f"Session {session_id!r} is being compressed by another writer" + if chunk_rows is not None and len(messages) > chunk_rows: + inserted_total = 0 + for start in range(0, len(messages), chunk_rows): + inserted_total += self.append_messages_batch( + session_id, + messages[start:start + chunk_rows], + compression_lock_holder=compression_lock_holder, ) - 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) - - row_ids: List[int] = [] - for params in prepared: - cursor = conn.execute(self._MESSAGE_INSERT_SQL, params) - row_ids.append(cursor.lastrowid) + return inserted_total + def _do(conn): + self._check_transcript_write_guards( + conn, session_id, compression_lock_holder + ) + inserted, tool_calls_total = self._insert_message_rows( + conn, session_id, messages + ) # One aggregated counter update for the whole batch. - if total_tool_calls > 0: + if tool_calls_total > 0: conn.execute( """UPDATE sessions SET message_count = message_count + ?, tool_call_count = tool_call_count + ? WHERE id = ?""", - (len(prepared), total_tool_calls, session_id), + (inserted, tool_calls_total, session_id), ) else: conn.execute( "UPDATE sessions SET message_count = message_count + ? WHERE id = ?", - (len(prepared), session_id), + (inserted, session_id), ) - return row_ids + return inserted # Same criticality as append_message: this IS the turn's transcript. return self._execute_write( diff --git a/run_agent.py b/run_agent.py index 11e3d732d959..e2bc5f8660bd 100644 --- a/run_agent.py +++ b/run_agent.py @@ -2225,11 +2225,13 @@ def _flush_messages_to_session_db_unlocked( "tool_calls": tool_calls_data, "tool_call_id": msg.get("tool_call_id"), "finish_reason": msg.get("finish_reason"), - "reasoning": msg.get("reasoning") if role == "assistant" else None, - "reasoning_content": msg.get("reasoning_content") if role == "assistant" else None, - "reasoning_details": msg.get("reasoning_details") if role == "assistant" else None, - "codex_reasoning_items": msg.get("codex_reasoning_items") if role == "assistant" else None, - "codex_message_items": msg.get("codex_message_items") if role == "assistant" else None, + # Reasoning/codex fields are role-gated (assistant-only) + # inside _insert_message_rows — pass through untouched. + "reasoning": msg.get("reasoning"), + "reasoning_content": msg.get("reasoning_content"), + "reasoning_details": msg.get("reasoning_details"), + "codex_reasoning_items": msg.get("codex_reasoning_items"), + "codex_message_items": msg.get("codex_message_items"), "timestamp": _row_timestamp, "api_content": _row_api_content, "display_kind": ( diff --git a/tests/hermes_state/test_append_messages_batch.py b/tests/hermes_state/test_append_messages_batch.py index aaff6884c8ad..a65436ee93e5 100644 --- a/tests/hermes_state/test_append_messages_batch.py +++ b/tests/hermes_state/test_append_messages_batch.py @@ -1,8 +1,9 @@ """Tests for SessionDB.append_messages_batch (#23254 salvage). -The batch writer must be row-shape-identical to append_message (shared -_prepare_message_row + _MESSAGE_INSERT_SQL), atomic (all rows or none), -and must aggregate the session counters in one UPDATE. +The batch writer reuses _insert_message_rows (the same row-serialization +path as replace/compact/import), runs the same admission guards as +append_message, is atomic (all rows or none), and aggregates the session +counters in one UPDATE. """ import json @@ -80,6 +81,26 @@ def test_batch_rows_identical_to_single_appends(self, db, tmp_path): finally: db2.close() + def test_reasoning_gated_to_assistant_rows(self, db): + """_insert_message_rows role-gates reasoning fields; a tool row + carrying reasoning keys must not persist them.""" + db.append_messages_batch( + "sess-batch", + [ + { + "role": "tool", + "content": "out", + "tool_name": "t", + "tool_call_id": "c1", + "reasoning_content": "should not persist", + } + ], + ) + row = db._conn.execute( + "SELECT reasoning_content FROM messages" + ).fetchone() + assert row[0] is None + def test_counters_aggregate_once(self, db): db.append_messages_batch("sess-batch", _turn_messages()) row = db._conn.execute( @@ -89,13 +110,11 @@ def test_counters_aggregate_once(self, db): assert row["message_count"] == 4 assert row["tool_call_count"] == 1 - def test_returns_row_ids_in_input_order(self, db): - ids = db.append_messages_batch("sess-batch", _turn_messages()) - assert ids == sorted(ids) - assert len(ids) == 4 + def test_returns_inserted_count(self, db): + assert db.append_messages_batch("sess-batch", _turn_messages()) == 4 def test_empty_batch_is_noop(self, db): - assert db.append_messages_batch("sess-batch", []) == [] + assert db.append_messages_batch("sess-batch", []) == 0 row = db._conn.execute( "SELECT message_count FROM sessions WHERE id = ?", ("sess-batch",) ).fetchone() @@ -103,31 +122,26 @@ def test_empty_batch_is_noop(self, db): def test_atomicity_all_or_nothing(self, db, monkeypatch): """A failure mid-batch leaves ZERO rows and untouched counters.""" - real_execute_write = db._execute_write - original_insert = SessionDB._MESSAGE_INSERT_SQL - - calls = {"n": 0} - - def _do_wrapper(fn, **kwargs): - def failing(conn): - real_conn_execute = conn.execute - - def exec_counting(sql, *args): - if sql == original_insert: - calls["n"] += 1 - if calls["n"] == 3: - raise sqlite3.OperationalError("boom mid-batch") - return real_conn_execute(sql, *args) - - conn.execute = exec_counting - try: - return fn(conn) - finally: - conn.execute = real_conn_execute - - return real_execute_write(failing, **kwargs) - - monkeypatch.setattr(db, "_execute_write", _do_wrapper) + real_insert = SessionDB._insert_message_rows + + def failing_insert(self_db, conn, session_id, messages): + real_conn_execute = conn.execute + calls = {"n": 0} + + def exec_counting(sql, *args): + if sql.lstrip().startswith("INSERT INTO messages"): + calls["n"] += 1 + if calls["n"] == 3: + raise sqlite3.OperationalError("boom mid-batch") + return real_conn_execute(sql, *args) + + conn.execute = exec_counting + try: + return real_insert(self_db, conn, session_id, messages) + finally: + conn.execute = real_conn_execute + + monkeypatch.setattr(SessionDB, "_insert_message_rows", failing_insert) with pytest.raises(sqlite3.OperationalError): db.append_messages_batch("sess-batch", _turn_messages()) monkeypatch.undo() diff --git a/tui_gateway/methods_session.py b/tui_gateway/methods_session.py index acbfe7a8b6a8..1a3e09f16e49 100644 --- a/tui_gateway/methods_session.py +++ b/tui_gateway/methods_session.py @@ -2649,9 +2649,9 @@ def _(rid, params: dict) -> dict: else None ), ) - # Copy the whole parent history in ONE transaction — a branch - # seed can be hundreds of rows, and per-row transactions were - # the write-amplification pattern removed in #23254. + # Copy the whole parent history in bounded-chunk transactions — + # a branch seed can be hundreds of rows, and per-row transactions + # were the write-amplification pattern removed in #23254. db.append_messages_batch( new_key, [ @@ -2664,6 +2664,7 @@ def _(rid, params: dict) -> dict: } for msg in history ], + chunk_rows=500, ) db.set_session_title(new_key, title) except Exception as e: diff --git a/tui_gateway/server.py b/tui_gateway/server.py index d51cbe42105f..8144189ef41e 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -2749,9 +2749,11 @@ def _persist_branch_seed(session: dict) -> None: if db is None: return try: - # One transaction for the whole seed (see #23254): a branch seed - # can be hundreds of rows and is all-or-nothing anyway — the - # _branch_seed_persisted flag below assumes every row landed. + # Bounded-chunk transactions (see #23254): a branch seed can be + # hundreds of rows; chunking keeps each BEGIN IMMEDIATE short so + # concurrent writers aren't starved. Recovery semantics match the + # old per-row loop (mid-copy failure leaves a partial seed with + # _branch_seed_persisted unset). db.append_messages_batch( key, [ @@ -2765,6 +2767,7 @@ def _persist_branch_seed(session: dict) -> None: } for msg in seed ], + chunk_rows=500, ) session["_branch_seed_persisted"] = True except Exception as exc: From 773f4a4250b0d939ed310b3f153ff5d1bb22379d Mon Sep 17 00:00:00 2001 From: kshitij Date: Mon, 3 Aug 2026 20:04:41 +0530 Subject: [PATCH 5/5] fix(tests): update two more append_message.call_args assertions to append_messages_batch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI-caught: test_verification_stop_caching and test_tui_gateway_server::test_native_vision_turn_persists_a_renderable_image_ref both assert on append_message.call_args, but the flush loop now calls append_messages_batch. Same class of test-fake fallout fixed in 5 other files — these two were missed. --- tests/agent/test_verification_stop_caching.py | 5 +++-- tests/test_tui_gateway_server.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/agent/test_verification_stop_caching.py b/tests/agent/test_verification_stop_caching.py index 7620d9eb5170..83bde23206ee 100644 --- a/tests/agent/test_verification_stop_caching.py +++ b/tests/agent/test_verification_stop_caching.py @@ -84,8 +84,9 @@ def test_db_flush_drops_only_nudge_keeps_candidate(tmp_path, monkeypatch): agent._flush_messages_to_session_db(messages, conversation_history=[]) persisted = [ - kwargs.get("content") - for _args, kwargs in agent._session_db.append_message.call_args_list + msg.get("content") + for _args, kwargs in agent._session_db.append_messages_batch.call_args_list + for msg in kwargs["messages"] ] assert "hi" in persisted assert "verified and clean" in persisted diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index df245ad8738c..18dae86177d8 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -16063,7 +16063,7 @@ def test_native_vision_turn_persists_a_renderable_image_ref(tmp_path): agent._flush_messages_to_session_db([{"role": "user", "content": native_parts}], []) - written = agent._session_db.append_message.call_args.kwargs["content"] + written = agent._session_db.append_messages_batch.call_args.kwargs["messages"][0]["content"] assert f"@image:`{img}`" in written assert "what is in this photo?" in written # The model keeps the pixels for the rest of the session.