devsart95
06ae5b6faa
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.
46 lines
1.6 KiB
Python
46 lines
1.6 KiB
Python
"""Test that tool_name is correctly persisted to the session DB for tool-result messages.
|
|
|
|
make_tool_result_message() sets tool_name on every tool-result dict at construction
|
|
time. This test verifies that the value survives the flush path into the session DB.
|
|
"""
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from run_agent import AIAgent
|
|
from agent.tool_dispatch_helpers import make_tool_result_message
|
|
|
|
|
|
def _make_agent(session_db):
|
|
with (
|
|
patch("run_agent.get_tool_definitions", return_value=[]),
|
|
patch("run_agent.check_toolset_requirements", return_value={}),
|
|
patch("run_agent.OpenAI"),
|
|
):
|
|
return AIAgent(
|
|
api_key="test-key",
|
|
base_url="https://openrouter.ai/api/v1",
|
|
quiet_mode=True,
|
|
skip_context_files=True,
|
|
skip_memory=True,
|
|
session_db=session_db,
|
|
)
|
|
|
|
|
|
def test_tool_name_persisted_to_session_db():
|
|
"""tool_name set by make_tool_result_message must be passed through to
|
|
the batched flush so the column is populated on first write to the
|
|
session DB."""
|
|
session_db = MagicMock()
|
|
agent = _make_agent(session_db)
|
|
|
|
messages = [
|
|
{"role": "user", "content": "run a command"},
|
|
make_tool_result_message("terminal", "$ ls\nfile.txt", "c1"),
|
|
]
|
|
agent._flush_messages_to_session_db(messages)
|
|
|
|
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"
|