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 test_direct_session_db_flushes_share_marker_claim(agent):
                 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 @@ class TestPersistUserMessageOverride:
             "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: