fix: make session transcript restoration atomic

Problem:
- Session mode mismatches could overwrite persisted architecture config, and direct tool transcripts could be stored or restored as invalid partial protocol groups.

Risk:
- SQLite transaction and migration locking change persistence timing; mode, cancellation, repair, concurrency, replay, focused, and full-suite tests cover compatibility and failure boundaries.

src/agent_lab/application/runtime.py

@@ -552,6 +552,7 @@ class DebugRuntime:
             cached_tokens = 0
             saw_event = False
             assistant_content: list[str] = []
+            assistant_message: ChatMessage | None = None
             events: list[ToolCallEvent] = []
             batch_result: EventBatchResult | None = None
             decision: EventBatchDecision | None = None
@@ -715,11 +716,14 @@ class DebugRuntime:
                     ),
                 )
                 messages.append(assistant_message)
-                self._append_persisted_message(
-                    session_id,
-                    turn_index,
-                    assistant_message,
-                )
+                if not (
+                    events and request.tool_invocation_mode == "chat_agent_tools"
+                ):
+                    self._append_persisted_message(
+                        session_id,
+                        turn_index,
+                        assistant_message,
+                    )
 
             if events:
                 batch_scope = self._event_scope(
@@ -751,12 +755,12 @@ class DebugRuntime:
                         scope=batch_scope,
                     )
                     messages.extend(tool_replies)
-                    for reply in tool_replies:
-                        self._append_persisted_message(
-                            session_id,
-                            turn_index,
-                            reply,
-                        )
+                    assert assistant_message is not None
+                    self._append_persisted_messages(
+                        session_id,
+                        turn_index,
+                        [assistant_message, *tool_replies],
+                    )
                 else:
                     response = asyncio.get_running_loop().create_future()
                     await queues.events.put(
@@ -1278,7 +1282,7 @@ class DebugRuntime:
     ) -> list[ChatMessage]:
         if self.session_store is None or session_id is None:
             return []
-        return [
+        restored = [
             ChatMessage.model_validate(
                 {
                     "role": message["role"],
@@ -1290,6 +1294,39 @@ class DebugRuntime:
             )
             for message in self.session_store.list_messages(session_id)
         ]
+        return self._repair_restored_messages(restored)
+
+    def _repair_restored_messages(
+        self,
+        messages: list[ChatMessage],
+    ) -> list[ChatMessage]:
+        repaired: list[ChatMessage] = []
+        index = 0
+        while index < len(messages):
+            message = messages[index]
+            if message.role == "tool":
+                index += 1
+                continue
+            if not message.tool_calls:
+                repaired.append(message)
+                index += 1
+                continue
+
+            expected_ids = {tool_call.id for tool_call in message.tool_calls}
+            replies: list[ChatMessage] = []
+            cursor = index + 1
+            while (
+                cursor < len(messages)
+                and messages[cursor].role == "tool"
+                and len(replies) < len(expected_ids)
+            ):
+                replies.append(messages[cursor])
+                cursor += 1
+            reply_ids = {reply.tool_call_id for reply in replies}
+            if len(replies) == len(expected_ids) and reply_ids == expected_ids:
+                repaired.extend([message, *replies])
+            index = cursor
+        return repaired
 
     def _start_persisted_turn(
         self,
@@ -1333,6 +1370,20 @@ class DebugRuntime:
             message=message,
         )
 
+    def _append_persisted_messages(
+        self,
+        session_id: str | None,
+        turn_index: int,
+        messages: list[ChatMessage],
+    ) -> None:
+        if self.session_store is None or session_id is None:
+            return
+        self.session_store.append_messages(
+            session_id,
+            turn_index=turn_index,
+            messages=messages,
+        )
+
     def _append_persisted_usage(
         self,
         session_id: str | None,

src/agent_lab/application/session_store.py

@@ -52,6 +52,15 @@ class SessionStore(Protocol):
     ) -> None:
         ...
 
+    def append_messages(
+        self,
+        session_id: str,
+        *,
+        turn_index: int,
+        messages: list[ChatMessage],
+    ) -> None:
+        ...
+
     def append_audit(
         self,
         session_id: str,

src/agent_lab/infrastructure/sqlite_store.py

@@ -42,7 +42,19 @@ class SQLiteSessionStore:
         title: str,
         config: dict[str, Any],
     ) -> str:
-        if session_id and self.get_session(session_id) is not None:
+        existing_session = self.get_session(session_id) if session_id else None
+        if existing_session is not None:
+            persisted_mode = existing_session["config"].get(
+                "tool_invocation_mode",
+                "dual_agent",
+            )
+            requested_mode = config.get("tool_invocation_mode", "dual_agent")
+            if persisted_mode != requested_mode:
+                raise ValueError(
+                    "session tool invocation mode mismatch: "
+                    f"persisted={persisted_mode} requested={requested_mode}"
+                )
+            assert session_id is not None
             self.update_session_config(session_id, config=config)
             return session_id
         return self.create_session(title=title, config=config, session_id=session_id)
@@ -168,34 +180,61 @@ class SQLiteSessionStore:
         turn_index: int,
         message: ChatMessage,
     ) -> None:
+        self.append_messages(
+            session_id,
+            turn_index=turn_index,
+            messages=[message],
+        )
+
+    def append_messages(
+        self,
+        session_id: str,
+        *,
+        turn_index: int,
+        messages: list[ChatMessage],
+    ) -> None:
+        if not messages:
+            return
         now = self._now()
         with self._lock:
-            self._connect().execute(
-                """
-                INSERT INTO messages
-                    (
-                        session_id, turn_index, role, content, name,
-                        tool_call_id, tool_calls_json, created_at
-                    )
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
-                """,
-                (
-                    session_id,
-                    turn_index,
-                    message.role,
-                    message.content,
-                    message.name,
-                    message.tool_call_id,
-                    json.dumps(
-                        [tool_call.model_dump() for tool_call in message.tool_calls],
-                        ensure_ascii=False,
-                        separators=(",", ":"),
-                    ),
-                    now,
-                ),
-            )
-            self._touch_session_locked(session_id, now)
-            self._connect().commit()
+            connection = self._connect()
+            try:
+                connection.execute("BEGIN")
+                connection.executemany(
+                    """
+                    INSERT INTO messages
+                        (
+                            session_id, turn_index, role, content, name,
+                            tool_call_id, tool_calls_json, created_at
+                        )
+                    VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+                    """,
+                    [
+                        (
+                            session_id,
+                            turn_index,
+                            message.role,
+                            message.content,
+                            message.name,
+                            message.tool_call_id,
+                            json.dumps(
+                                [
+                                    tool_call.model_dump()
+                                    for tool_call in message.tool_calls
+                                ],
+                                ensure_ascii=False,
+                                separators=(",", ":"),
+                            ),
+                            now,
+                        )
+                        for message in messages
+                    ],
+                )
+                self._touch_session_locked(session_id, now)
+                connection.commit()
+            except BaseException:
+                connection.rollback()
+                raise
 
     def append_audit(
         self,
@@ -477,18 +516,33 @@ class SQLiteSessionStore:
             );
             """
         )
-        message_columns = {
+        try:
+            self._connection.execute("BEGIN IMMEDIATE")
+            if "tool_calls_json" not in self._message_columns():
+                try:
+                    self._connection.execute(
+                        """
+                        ALTER TABLE messages
+                        ADD COLUMN tool_calls_json TEXT NOT NULL DEFAULT '[]'
+                        """
+                    )
+                except sqlite3.OperationalError as exc:
+                    if (
+                        "duplicate column name" not in str(exc).lower()
+                        or "tool_calls_json" not in self._message_columns()
+                    ):
+                        raise
+            self._connection.commit()
+        except BaseException:
+            self._connection.rollback()
+            raise
+
+    def _message_columns(self) -> set[str]:
+        assert self._connection is not None
+        return {
             row[1]
             for row in self._connection.execute("PRAGMA table_info(messages)")
         }
-        if "tool_calls_json" not in message_columns:
-            self._connection.execute(
-                """
-                ALTER TABLE messages
-                ADD COLUMN tool_calls_json TEXT NOT NULL DEFAULT '[]'
-                """
-            )
-        self._connection.commit()
 
     def _touch_session_locked(self, session_id: str, now: str) -> None:
         self._connect().execute(

src/agent_lab/presentation/static/app.js

@@ -992,6 +992,14 @@ function renderReplayMessages(messages) {
   }
 
   messages.forEach((message) => {
+    if (
+      message.role === "assistant" &&
+      !message.content &&
+      Array.isArray(message.tool_calls) &&
+      message.tool_calls.length > 0
+    ) {
+      return;
+    }
     const item = document.createElement("div");
     const kind = ["user", "assistant", "tool"].includes(message.role)
       ? message.role

tests/test_debug_runtime.py

@@ -2090,6 +2090,155 @@ async def test_direct_mode_reconnect_restores_provider_valid_session_transcript(
     ) == 1
 
 
+@pytest.mark.asyncio
+async def test_cancelled_direct_batch_persists_no_partial_protocol_and_reconnects(tmp_path):
+    store = SQLiteSessionStore(tmp_path / "agent_lab.sqlite3")
+    handler_started = asyncio.Event()
+    release_handler = asyncio.Event()
+
+    async def blocking_handler(event: ToolCallEvent) -> dict[str, Any]:
+        handler_started.set()
+        await release_handler.wait()
+        return {"ok": True, "call_id": event.id}
+
+    registry = _single_tool_registry(blocking_handler)
+    first_client = ScriptedChatClient(
+        [[StreamItem.provider_tool_call(_tool_call("cancelled-call"))]]
+    )
+    first_runtime = DebugRuntime(
+        first_client,
+        registry=registry,
+        session_store=store,
+    )
+    request = _direct_request(
+        enabled_tools=["safe_tool"],
+        session_id="cancelled-direct",
+    )
+
+    first_runtime.start_session(request)
+    await asyncio.wait_for(handler_started.wait(), timeout=1)
+    await first_runtime.aclose()
+
+    persisted_after_cancel = store.list_messages("cancelled-direct")
+    assert [message["role"] for message in persisted_after_cancel] == ["user"]
+    assert not any(message["tool_calls"] for message in persisted_after_cancel)
+
+    reconnect_client = TranscriptValidatingScriptedChatClient(
+        [[StreamItem.message_delta("reconnected")]]
+    )
+    reconnect_runtime = DebugRuntime(
+        reconnect_client,
+        registry=registry,
+        session_store=store,
+    )
+    reconnect_request = request.model_copy(
+        update={"user_message": "after cancellation"}
+    )
+    queues = reconnect_runtime.start_session(reconnect_request)
+    outputs: list[dict[str, Any]] = []
+    while not any(
+        message["type"] in {"turn_completed", "error"} for message in outputs
+    ):
+        outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
+    await reconnect_runtime.aclose()
+
+    assert not any(message["type"] == "error" for message in outputs)
+    assert [
+        (message.role, message.content)
+        for message in reconnect_client.messages_by_call[0]
+        if message.role != "system"
+    ] == [
+        ("user", "debug direct tools"),
+        ("user", "after cancellation"),
+    ]
+
+
+@pytest.mark.asyncio
+async def test_direct_restore_drops_invalid_protocol_fragments_and_keeps_visible_history(
+    tmp_path,
+):
+    store = SQLiteSessionStore(tmp_path / "agent_lab.sqlite3")
+    session_id = store.create_session(
+        title="repair transcript",
+        config={"tool_invocation_mode": "chat_agent_tools"},
+        session_id="repair-transcript",
+    )
+    store.start_turn(session_id, turn_index=1, user_message="visible user")
+    complete_calls = [_tool_call("complete-1"), _tool_call("complete-2")]
+    store.append_messages(
+        session_id,
+        turn_index=1,
+        messages=[
+            ChatMessage(role="user", content="visible user"),
+            ChatMessage(
+                role="assistant",
+                content="using tools",
+                tool_calls=complete_calls,
+            ),
+            ChatMessage(
+                role="tool",
+                content='{"id":1}',
+                name="safe_tool",
+                tool_call_id="complete-1",
+            ),
+            ChatMessage(
+                role="tool",
+                content='{"id":2}',
+                name="safe_tool",
+                tool_call_id="complete-2",
+            ),
+            ChatMessage(role="assistant", content="visible answer"),
+            ChatMessage(
+                role="tool",
+                content='{"orphan":true}',
+                name="safe_tool",
+                tool_call_id="orphan-call",
+            ),
+            ChatMessage(role="assistant", content="visible after orphan"),
+            ChatMessage(
+                role="assistant",
+                content="",
+                tool_calls=[_tool_call("incomplete-call")],
+            ),
+        ],
+    )
+    store.complete_turn(session_id, turn_index=1)
+    client = TranscriptValidatingScriptedChatClient(
+        [[StreamItem.message_delta("repaired")]]
+    )
+    runtime = DebugRuntime(
+        client,
+        registry=_single_tool_registry(lambda event: {"ok": True}),
+        session_store=store,
+    )
+    request = _direct_request(
+        enabled_tools=["safe_tool"],
+        session_id=session_id,
+    ).model_copy(update={"user_message": "current user"})
+
+    queues = runtime.start_session(request)
+    outputs: list[dict[str, Any]] = []
+    while not any(
+        message["type"] in {"turn_completed", "error"} for message in outputs
+    ):
+        outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
+    await runtime.aclose()
+
+    assert not any(message["type"] == "error" for message in outputs)
+    transcript = client.messages_by_call[0]
+    assert [(message.role, message.content) for message in transcript] == [
+        ("system", "You are a debugger."),
+        ("user", "visible user"),
+        ("assistant", "using tools"),
+        ("tool", '{"id":1}'),
+        ("tool", '{"id":2}'),
+        ("assistant", "visible answer"),
+        ("assistant", "visible after orphan"),
+        ("user", "current user"),
+    ]
+    assert transcript[2].tool_calls == complete_calls
+
+
 @pytest.mark.asyncio
 async def test_dual_mode_reconnect_restores_visible_history_without_tool_protocol(tmp_path):
     store = SQLiteSessionStore(tmp_path / "agent_lab.sqlite3")

tests/test_sqlite_store.py

@@ -1,4 +1,8 @@
 import sqlite3
+from concurrent.futures import ThreadPoolExecutor
+from threading import Barrier
+
+import pytest
 
 from agent_lab.domain.events import ToolCallEvent
 from agent_lab.domain.messages import ChatMessage, TokenUsage
@@ -182,3 +186,173 @@ def test_sqlite_store_migrates_legacy_messages_with_empty_tool_calls(tmp_path):
     tool_calls_column = next(column for column in migrated if column[1] == "tool_calls_json")
     assert tool_calls_column[3] == 1
     assert tool_calls_column[4] == "'[]'"
+
+
+@pytest.mark.parametrize(
+    ("persisted_config", "requested_mode", "persisted_mode"),
+    [
+        ({"tool_invocation_mode": "chat_agent_tools"}, "dual_agent", "chat_agent_tools"),
+        ({}, "chat_agent_tools", "dual_agent"),
+    ],
+)
+def test_sqlite_store_rejects_session_mode_mismatch_without_mutation(
+    tmp_path,
+    persisted_config,
+    requested_mode,
+    persisted_mode,
+):
+    store = SQLiteSessionStore(tmp_path / "agent_lab.sqlite3")
+    session_id = store.create_session(
+        title="mode locked",
+        config=persisted_config,
+        session_id="mode-locked",
+    )
+    store.start_turn(session_id, turn_index=1, user_message="past user")
+    store.append_message(
+        session_id,
+        turn_index=1,
+        message=ChatMessage(role="user", content="past user"),
+    )
+    before_session = store.get_session(session_id)
+    before_messages = store.list_messages(session_id)
+
+    with pytest.raises(
+        ValueError,
+        match=(
+            "session tool invocation mode mismatch: "
+            f"persisted={persisted_mode} requested={requested_mode}"
+        ),
+    ):
+        store.ensure_session(
+            session_id,
+            title="ignored",
+            config={"tool_invocation_mode": requested_mode, "changed": True},
+        )
+
+    assert store.get_session(session_id) == before_session
+    assert store.list_messages(session_id) == before_messages
+
+
+def test_sqlite_store_append_messages_rolls_back_the_whole_protocol_group(tmp_path):
+    database_path = tmp_path / "agent_lab.sqlite3"
+    store = SQLiteSessionStore(database_path)
+    session_id = store.create_session(title="atomic group", config={})
+    store.start_turn(session_id, turn_index=1, user_message="use tool")
+    connection = sqlite3.connect(database_path)
+    connection.execute(
+        """
+        CREATE TRIGGER reject_tool_message
+        BEFORE INSERT ON messages
+        WHEN NEW.role = 'tool'
+        BEGIN
+            SELECT RAISE(ABORT, 'tool blocked');
+        END
+        """
+    )
+    connection.commit()
+    connection.close()
+    tool_call = ToolCallEvent(
+        id="atomic-call",
+        name="mock_search",
+        arguments={"query": "atomic"},
+        raw_arguments='{"query":"atomic"}',
+    )
+
+    with pytest.raises(sqlite3.IntegrityError, match="tool blocked"):
+        store.append_messages(
+            session_id,
+            turn_index=1,
+            messages=[
+                ChatMessage(role="assistant", content="", tool_calls=[tool_call]),
+                ChatMessage(
+                    role="tool",
+                    content='{"ok":true}',
+                    name="mock_search",
+                    tool_call_id="atomic-call",
+                ),
+            ],
+        )
+
+    assert store.list_messages(session_id) == []
+
+
+def test_sqlite_store_migration_is_safe_across_concurrent_initialization(
+    tmp_path,
+    monkeypatch,
+):
+    database_path = tmp_path / "concurrent-legacy.sqlite3"
+    connection = sqlite3.connect(database_path)
+    connection.execute(
+        """
+        CREATE TABLE messages (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            session_id TEXT NOT NULL,
+            turn_index INTEGER NOT NULL,
+            role TEXT NOT NULL,
+            content TEXT NOT NULL,
+            name TEXT,
+            tool_call_id TEXT,
+            created_at TEXT NOT NULL
+        )
+        """
+    )
+    connection.execute(
+        """
+        INSERT INTO messages
+            (session_id, turn_index, role, content, name, tool_call_id, created_at)
+        VALUES ('legacy', 1, 'assistant', 'visible', NULL, NULL, 'now')
+        """
+    )
+    connection.commit()
+    connection.close()
+    worker_count = 8
+    barrier = Barrier(worker_count)
+    real_connect = sqlite3.connect
+
+    class CoordinatedConnection:
+        def __init__(self, connection):
+            self._connection = connection
+
+        @property
+        def row_factory(self):
+            return self._connection.row_factory
+
+        @row_factory.setter
+        def row_factory(self, value):
+            self._connection.row_factory = value
+
+        def execute(self, sql, *args):
+            cursor = self._connection.execute(sql, *args)
+            if (
+                sql.strip() == "PRAGMA table_info(messages)"
+                and not self._connection.in_transaction
+            ):
+                rows = cursor.fetchall()
+                barrier.wait(timeout=5)
+                return rows
+            return cursor
+
+        def __getattr__(self, name):
+            return getattr(self._connection, name)
+
+    def coordinated_connect(*args, **kwargs):
+        return CoordinatedConnection(real_connect(*args, **kwargs))
+
+    monkeypatch.setattr(
+        "agent_lab.infrastructure.sqlite_store.sqlite3.connect",
+        coordinated_connect,
+    )
+
+    def initialize_store() -> list[dict]:
+        store = SQLiteSessionStore(database_path)
+        barrier.wait()
+        return store.list_messages("legacy")
+
+    with ThreadPoolExecutor(max_workers=worker_count) as executor:
+        results = list(executor.map(lambda _: initialize_store(), range(worker_count)))
+
+    assert all(messages[0]["tool_calls"] == [] for messages in results)
+    columns = real_connect(database_path).execute(
+        "PRAGMA table_info(messages)"
+    ).fetchall()
+    assert [column[1] for column in columns].count("tool_calls_json") == 1

tests/test_websocket_api.py

@@ -1554,6 +1554,21 @@ def test_static_app_loads_session_replay_and_usage():
     assert "Finish current turn before switching sessions" in js
 
 
+def test_static_replay_hides_empty_assistant_tool_protocol_records():
+    js = Path("src/agent_lab/presentation/static/app.js").read_text()
+    render_replay = js[
+        js.index("function renderReplayMessages(") :
+        js.index("function appendAuditEntry(")
+    ]
+
+    assert 'message.role === "assistant"' in render_replay
+    assert "!message.content" in render_replay
+    assert "Array.isArray(message.tool_calls)" in render_replay
+    assert "message.tool_calls.length > 0" in render_replay
+    assert "return;" in render_replay
+    assert '["user", "assistant", "tool"].includes(message.role)' in render_replay
+
+
 def test_static_session_switch_guard_uses_active_turn_not_open_socket():
     js = Path("src/agent_lab/presentation/static/app.js").read_text()
     create_session = js[js.index("async function createSession()") : js.index("async function loadSelectedSession()")]