From b58b3adb9beefce2dd5bed5106a7476743c2f111 Mon Sep 17 00:00:00 2001
From: kshitij <kshitijkapoor0611@gmail.com>
Date: Mon, 3 Aug 2026 16:43:44 +0530
Subject: [PATCH] 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 e8aaa16e05..df245ad873 100644
--- a/tests/test_tui_gateway_server.py
+++ b/tests/test_tui_gateway_server.py
@@ -11438,6 +11438,11 @@ def test_session_branch_writes_to_parent_profile_db(monkeypatch, tmp_path):
         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 test_session_branch_installs_parent_profile_secret_scope(monkeypatch, tmp_pa
         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 d84fedcd43..acbfe7a8b6 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 b2dfc34e2a..a97fa2353c 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