diff --git a/gateway/slash_commands.py b/gateway/slash_commands.py index fb3cec080accc..e719eb2ac88a0 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 69b5323b5616a..5ec16a5fe4df0 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 68c8ba36af0f4..b32b6ffe27c1b 100644 --- a/hermes_state.py +++ b/hermes_state.py @@ -5883,6 +5883,39 @@ def _encode_display_metadata(display_metadata: Any) -> Optional[str]: ) return None + 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. + + 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). + """ + 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]]: """Decode a ``display_metadata`` column into the dict every reader expects. @@ -5995,28 +6028,9 @@ 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 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) + 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, @@ -6072,6 +6086,79 @@ 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, + chunk_rows: Optional[int] = None, + ) -> int: + """Append multiple messages atomically in ONE write transaction. + + ``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. + + Atomicity contract: all rows land or none do (the caller re-flushes + 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 0 + + 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, + ) + 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 tool_calls_total > 0: + conn.execute( + """UPDATE sessions SET message_count = message_count + ?, + tool_call_count = tool_call_count + ? WHERE id = ?""", + (inserted, tool_calls_total, session_id), + ) + else: + conn.execute( + "UPDATE sessions SET message_count = message_count + ? WHERE id = ?", + (inserted, session_id), + ) + return inserted + + # 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 953a127323fac..e2bc5f8660bd6 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,48 @@ 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/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": ( "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 c08b012aca3dd..afb8e967bc9da 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/agent/test_verification_stop_caching.py b/tests/agent/test_verification_stop_caching.py index 7620d9eb51706..83bde23206eec 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/hermes_state/test_append_messages_batch.py b/tests/hermes_state/test_append_messages_batch.py new file mode 100644 index 0000000000000..a65436ee93e51 --- /dev/null +++ b/tests/hermes_state/test_append_messages_batch.py @@ -0,0 +1,194 @@ +"""Tests for SessionDB.append_messages_batch (#23254 salvage). + +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 +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_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( + "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_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", []) == 0 + 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_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() + + 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 ff8b9cf3122dd..70e262fbf289b 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 b1a40e78dc912..e169f1e06821d 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_run_agent.py b/tests/run_agent/test_run_agent.py index 4bc1c69b76b75..6161e266b0719 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: diff --git a/tests/run_agent/test_tool_name_db_persistence.py b/tests/run_agent/test_tool_name_db_persistence.py index 3fcf7f33c3ad9..29596c04ddb79 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" diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index e8aaa16e05725..18dae86177d8c 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 @@ -16053,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. diff --git a/tui_gateway/methods_session.py b/tui_gateway/methods_session.py index d84fedcd4314e..1a3e09f16e49c 100644 --- a/tui_gateway/methods_session.py +++ b/tui_gateway/methods_session.py @@ -2649,15 +2649,23 @@ 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 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, + [ + { + "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 + ], + chunk_rows=500, + ) 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 a3475fc11ecb3..8144189ef41e4 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -2749,16 +2749,26 @@ 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"), - ) + # 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, + [ + { + "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 + ], + chunk_rows=500, + ) session["_branch_seed_persisted"] = True except Exception as exc: from hermes_state import is_disk_full_error