agent-lab/tests/test_sqlite_store.py

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
from agent_lab.infrastructure.sqlite_store import SQLiteSessionStore


def test_sqlite_store_persists_session_messages_audit_and_usage(tmp_path):
    store = SQLiteSessionStore(tmp_path / "agent_lab.sqlite3")

    session_id = store.create_session(
        title="Debug first turn",
        config={
            "tool_invocation_mode": "chat_agent_tools",
            "chat_agent": {"model": "chat-model"},
        },
    )
    store.start_turn(session_id, turn_index=1, user_message="debug this")
    store.append_message(
        session_id,
        turn_index=1,
        message=ChatMessage(role="user", content="debug this"),
    )
    store.append_message(
        session_id,
        turn_index=1,
        message=ChatMessage(role="assistant", content="answer"),
    )
    store.append_audit(
        session_id,
        event="chat_agent_request",
        details={"model": "chat-model"},
        turn_index=1,
        round_index=1,
    )
    store.append_usage(
        session_id,
        turn_index=1,
        round_index=1,
        usage=TokenUsage(
            prompt_tokens=3,
            completion_tokens=7,
            total_tokens=10,
            cached_tokens=2,
        ),
        ttft_ms=120,
        elapsed_ms=450,
    )
    store.complete_turn(session_id, turn_index=1, wall_time_ms=425)

    sessions = store.list_sessions()
    assert sessions == [
        {
            "id": session_id,
            "title": "Debug first turn",
            "created_at": sessions[0]["created_at"],
            "updated_at": sessions[0]["updated_at"],
            "turn_count": 1,
            "total_tokens": 10,
        }
    ]
    assert store.get_session(session_id)["config"]["chat_agent"]["model"] == "chat-model"
    assert [message["content"] for message in store.list_messages(session_id)] == [
        "debug this",
        "answer",
    ]
    assert store.list_audit_logs(session_id)[0]["event"] == "chat_agent_request"
    usage = store.usage_summary(session_id)
    assert usage["calls"][0] | {"created_at": "ignored"} == {
        "id": usage["calls"][0]["id"],
        "turn_index": 1,
        "round_index": 1,
        "prompt_tokens": 3,
        "completion_tokens": 7,
        "total_tokens": 10,
        "cached_tokens": 2,
        "ttft_ms": 120,
        "elapsed_ms": 450,
        "mode": "chat_agent_tools",
        "agent": "chat_agent",
        "call_kind": "chat_completion",
        "event_id": None,
        "event_name": None,
        "used_fallback": False,
        "tool_latency_ms": None,
        "metric_key": None,
        "created_at": "ignored",
    }
    assert usage["turns"] == [
        {
            "turn_index": 1,
            "prompt_tokens": 3,
            "completion_tokens": 7,
            "total_tokens": 10,
            "cached_tokens": 2,
            "elapsed_ms": 450,
            "call_count": 1,
            "fallback_count": 0,
            "tool_count": 0,
            "turn_wall_time_ms": 425,
        }
    ]
    assert usage["session"] == {
        "prompt_tokens": 3,
        "completion_tokens": 7,
        "total_tokens": 10,
        "cached_tokens": 2,
        "elapsed_ms": 450,
        "call_count": 1,
        "fallback_count": 0,
        "tool_count": 0,
        "turn_wall_time_ms": 425,
    }


def test_sqlite_store_persists_typed_calls_without_double_counting_metrics(tmp_path):
    store = SQLiteSessionStore(tmp_path / "agent_lab.sqlite3")
    session_id = store.create_session(
        title="comparable metrics",
        config={"tool_invocation_mode": "dual_agent"},
    )
    store.start_turn(session_id, turn_index=1, user_message="search")

    store.append_usage(
        session_id,
        turn_index=1,
        round_index=1,
        usage=TokenUsage(prompt_tokens=2, completion_tokens=3, total_tokens=5),
        ttft_ms=10,
        elapsed_ms=30,
        metadata={"metric_key": "chat:1:1"},
    )
    fallback_metadata = {
        "agent": "event_agent",
        "call_kind": "argument_fallback",
        "event_id": "event-1",
        "event_name": "mock_search",
        "used_fallback": True,
        "metric_key": "fallback:1:1:event-1",
    }
    for _ in range(2):
        store.append_usage(
            session_id,
            turn_index=1,
            round_index=1,
            usage=TokenUsage(prompt_tokens=4, completion_tokens=6, total_tokens=10),
            ttft_ms=5,
            elapsed_ms=20,
            metadata=fallback_metadata,
        )
    store.append_usage(
        session_id,
        turn_index=1,
        round_index=1,
        usage=TokenUsage(prompt_tokens=100, completion_tokens=100, total_tokens=200),
        ttft_ms=None,
        elapsed_ms=80,
        metadata={
            "agent": "event_kernel",
            "call_kind": "tool_execution",
            "event_id": "event-1",
            "event_name": "mock_search",
            "used_fallback": True,
            "tool_latency_ms": 80,
            "metric_key": "tool:1:1:event-1",
        },
    )
    store.complete_turn(session_id, turn_index=1, wall_time_ms=95)

    usage = store.usage_summary(session_id)

    assert len(usage["calls"]) == 3
    assert [call["mode"] for call in usage["calls"]] == ["dual_agent"] * 3
    assert [call["agent"] for call in usage["calls"]] == [
        "chat_agent",
        "event_agent",
        "event_kernel",
    ]
    assert [call["call_kind"] for call in usage["calls"]] == [
        "chat_completion",
        "argument_fallback",
        "tool_execution",
    ]
    assert usage["calls"][1]["event_id"] == "event-1"
    assert usage["calls"][1]["event_name"] == "mock_search"
    assert usage["calls"][1]["used_fallback"] is True
    assert usage["calls"][2]["tool_latency_ms"] == 80
    assert usage["calls"][2]["total_tokens"] == 0
    assert usage["turns"] == [
        {
            "turn_index": 1,
            "prompt_tokens": 6,
            "completion_tokens": 9,
            "total_tokens": 15,
            "cached_tokens": 0,
            "elapsed_ms": 130,
            "call_count": 3,
            "fallback_count": 1,
            "tool_count": 1,
            "turn_wall_time_ms": 95,
        }
    ]
    assert usage["session"] == {
        "prompt_tokens": 6,
        "completion_tokens": 9,
        "total_tokens": 15,
        "cached_tokens": 0,
        "elapsed_ms": 130,
        "call_count": 3,
        "fallback_count": 1,
        "tool_count": 1,
        "turn_wall_time_ms": 95,
    }
    assert store.list_sessions()[0]["total_tokens"] == 15


def test_sqlite_store_metric_keys_are_unique_within_each_session(tmp_path):
    store = SQLiteSessionStore(tmp_path / "agent_lab.sqlite3")
    for session_id in ("session-a", "session-b"):
        store.create_session(title=session_id, config={}, session_id=session_id)
        store.start_turn(session_id, turn_index=1, user_message="test")
        for _ in range(2):
            store.append_usage(
                session_id,
                turn_index=1,
                round_index=1,
                usage=TokenUsage(total_tokens=1),
                ttft_ms=None,
                elapsed_ms=1,
                metadata={"metric_key": "chat:1:1"},
            )

    assert len(store.usage_summary("session-a")["calls"]) == 1
    assert len(store.usage_summary("session-b")["calls"]) == 1


def test_sqlite_store_migrates_legacy_usage_and_turn_metrics_additively(tmp_path):
    database_path = tmp_path / "legacy-usage.sqlite3"
    connection = sqlite3.connect(database_path)
    connection.executescript(
        """
        CREATE TABLE sessions (
            id TEXT PRIMARY KEY,
            title TEXT NOT NULL,
            created_at TEXT NOT NULL,
            updated_at TEXT NOT NULL,
            config_json TEXT NOT NULL
        );
        CREATE TABLE turns (
            session_id TEXT NOT NULL,
            turn_index INTEGER NOT NULL,
            user_message TEXT NOT NULL,
            started_at TEXT NOT NULL,
            completed_at TEXT,
            PRIMARY KEY (session_id, turn_index)
        );
        CREATE TABLE usage_stats (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id TEXT NOT NULL,
            turn_index INTEGER NOT NULL,
            round_index INTEGER NOT NULL,
            prompt_tokens INTEGER NOT NULL,
            completion_tokens INTEGER NOT NULL,
            total_tokens INTEGER NOT NULL,
            cached_tokens INTEGER NOT NULL,
            ttft_ms INTEGER,
            elapsed_ms INTEGER NOT NULL,
            created_at TEXT NOT NULL
        );
        INSERT INTO sessions VALUES (
            'legacy', 'Legacy', 'now', 'now',
            '{"tool_invocation_mode":"chat_agent_tools"}'
        );
        INSERT INTO turns VALUES ('legacy', 1, 'old', 'now', 'now');
        INSERT INTO usage_stats (
            session_id, turn_index, round_index, prompt_tokens,
            completion_tokens, total_tokens, cached_tokens,
            ttft_ms, elapsed_ms, created_at
        ) VALUES ('legacy', 1, 1, 1, 2, 3, 0, 4, 5, 'now');
        """
    )
    connection.commit()
    connection.close()

    store = SQLiteSessionStore(database_path)
    usage = store.usage_summary("legacy")

    assert usage["calls"][0]["agent"] == "chat_agent"
    assert usage["calls"][0]["call_kind"] == "chat_completion"
    assert usage["calls"][0]["mode"] == "chat_agent_tools"
    assert usage["calls"][0]["used_fallback"] is False
    assert usage["turns"][0]["turn_wall_time_ms"] is None
    migrated = sqlite3.connect(database_path)
    usage_columns = {row[1] for row in migrated.execute("PRAGMA table_info(usage_stats)")}
    turn_columns = {row[1] for row in migrated.execute("PRAGMA table_info(turns)")}
    indexes = {row[1] for row in migrated.execute("PRAGMA index_list(usage_stats)")}
    metric_index_sql = migrated.execute(
        "SELECT sql FROM sqlite_master WHERE name = 'usage_stats_metric_key_uq'"
    ).fetchone()[0]
    migrated.close()
    assert {
        "agent",
        "call_kind",
        "event_id",
        "event_name",
        "used_fallback",
        "tool_latency_ms",
        "metric_key",
    } <= usage_columns
    assert "wall_time_ms" in turn_columns
    assert "usage_stats_metric_key_uq" in indexes
    assert "session_id, metric_key" in metric_index_sql


def test_sqlite_store_deduplicates_metric_keys_before_creating_unique_index(
    tmp_path,
):
    database_path = tmp_path / "duplicate-metric-keys.sqlite3"
    connection = sqlite3.connect(database_path)
    connection.execute(
        """
        CREATE TABLE usage_stats (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id TEXT NOT NULL,
            turn_index INTEGER NOT NULL,
            round_index INTEGER NOT NULL,
            prompt_tokens INTEGER NOT NULL,
            completion_tokens INTEGER NOT NULL,
            total_tokens INTEGER NOT NULL,
            cached_tokens INTEGER NOT NULL,
            ttft_ms INTEGER,
            elapsed_ms INTEGER NOT NULL,
            metric_key TEXT,
            created_at TEXT NOT NULL
        )
        """
    )
    connection.executemany(
        """
        INSERT INTO usage_stats (
            id, session_id, turn_index, round_index, prompt_tokens,
            completion_tokens, total_tokens, cached_tokens, ttft_ms,
            elapsed_ms, metric_key, created_at
        ) VALUES (?, 'legacy', 1, 1, 0, 0, 0, 0, NULL, 1, ?, 'now')
        """,
        [
            (7, "chat:1:1"),
            (11, "chat:1:1"),
            (13, None),
            (17, None),
            (19, None),
        ],
    )
    connection.commit()
    connection.close()

    SQLiteSessionStore(database_path).list_sessions()

    migrated = sqlite3.connect(database_path)
    rows = migrated.execute(
        """
        SELECT id, session_id, metric_key
        FROM usage_stats
        ORDER BY id
        """
    ).fetchall()
    indexes = [row[1] for row in migrated.execute("PRAGMA index_list(usage_stats)")]
    migrated.close()

    assert rows == [
        (7, "legacy", "chat:1:1"),
        (13, "legacy", None),
        (17, "legacy", None),
        (19, "legacy", None),
    ]
    assert indexes.count("usage_stats_metric_key_uq") == 1


def test_sqlite_store_session_list_totals_are_not_multiplied_by_turns(tmp_path):
    store = SQLiteSessionStore(tmp_path / "agent_lab.sqlite3")
    session_id = store.create_session(title="two turns", config={})

    for turn_index, tokens in [(1, 2), (2, 4)]:
        store.start_turn(
            session_id,
            turn_index=turn_index,
            user_message=f"turn {turn_index}",
        )
        store.append_usage(
            session_id,
            turn_index=turn_index,
            round_index=1,
            usage=TokenUsage(total_tokens=tokens),
            ttft_ms=None,
            elapsed_ms=10,
        )

    assert store.list_sessions()[0]["turn_count"] == 2
    assert store.list_sessions()[0]["total_tokens"] == 6


def test_sqlite_store_roundtrips_assistant_tool_calls_and_tool_replies(tmp_path):
    store = SQLiteSessionStore(tmp_path / "agent_lab.sqlite3")
    session_id = store.create_session(title="direct tools", config={})
    store.start_turn(session_id, turn_index=1, user_message="search")
    tool_calls = [
        ToolCallEvent(
            id="call-1",
            name="mock_search",
            arguments={"query": "agent lab"},
            raw_arguments='{"query":"agent lab"}',
        ),
        ToolCallEvent(
            id="call-2",
            name="mock_ticket",
            arguments={"title": "follow up"},
            raw_arguments='{"title":"follow up"}',
        ),
    ]

    store.append_message(
        session_id,
        turn_index=1,
        message=ChatMessage(
            role="assistant",
            content="",
            tool_calls=tool_calls,
        ),
    )
    store.append_message(
        session_id,
        turn_index=1,
        message=ChatMessage(
            role="tool",
            content='{"ok":true}',
            name="mock_search",
            tool_call_id="call-1",
        ),
    )

    messages = store.list_messages(session_id)

    assert messages[0]["tool_calls"] == [call.model_dump() for call in tool_calls]
    assert messages[1]["tool_calls"] == []
    assert messages[1]["tool_call_id"] == "call-1"


def test_sqlite_store_migrates_legacy_messages_with_empty_tool_calls(tmp_path):
    database_path = tmp_path / "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', 'old answer', NULL, NULL, 'now')
        """
    )
    connection.commit()
    connection.close()

    store = SQLiteSessionStore(database_path)

    assert store.list_messages("legacy")[0]["tool_calls"] == []
    migrated = sqlite3.connect(database_path).execute(
        "PRAGMA table_info(messages)"
    ).fetchall()
    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(
        """
        CREATE TABLE turns (
            session_id TEXT NOT NULL,
            turn_index INTEGER NOT NULL,
            user_message TEXT NOT NULL,
            started_at TEXT NOT NULL,
            completed_at TEXT,
            PRIMARY KEY (session_id, turn_index)
        )
        """
    )
    connection.execute(
        """
        CREATE TABLE usage_stats (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id TEXT NOT NULL,
            turn_index INTEGER NOT NULL,
            round_index INTEGER NOT NULL,
            prompt_tokens INTEGER NOT NULL,
            completion_tokens INTEGER NOT NULL,
            total_tokens INTEGER NOT NULL,
            cached_tokens INTEGER NOT NULL,
            ttft_ms INTEGER,
            elapsed_ms INTEGER NOT NULL,
            metric_key TEXT,
            created_at TEXT NOT NULL
        )
        """
    )
    connection.executemany(
        """
        INSERT INTO usage_stats (
            id, session_id, turn_index, round_index, prompt_tokens,
            completion_tokens, total_tokens, cached_tokens, ttft_ms,
            elapsed_ms, metric_key, created_at
        ) VALUES (?, 'legacy', 1, 1, 0, 0, 0, 0, NULL, 1, ?, 'now')
        """,
        [
            (7, "chat:1:1"),
            (11, "chat:1:1"),
            (13, None),
            (17, None),
        ],
    )
    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
    migrated = real_connect(database_path)
    turn_columns = [row[1] for row in migrated.execute("PRAGMA table_info(turns)")]
    usage_columns = [
        row[1] for row in migrated.execute("PRAGMA table_info(usage_stats)")
    ]
    metric_indexes = [
        row[1] for row in migrated.execute("PRAGMA index_list(usage_stats)")
    ]
    usage_rows = migrated.execute(
        """
        SELECT id, session_id, metric_key
        FROM usage_stats
        ORDER BY id
        """
    ).fetchall()
    migrated.close()
    assert turn_columns.count("wall_time_ms") == 1
    for column_name in (
        "agent",
        "call_kind",
        "event_id",
        "event_name",
        "used_fallback",
        "tool_latency_ms",
        "metric_key",
    ):
        assert usage_columns.count(column_name) == 1
    assert metric_indexes.count("usage_stats_metric_key_uq") == 1
    assert usage_rows == [
        (7, "legacy", "chat:1:1"),
        (13, "legacy", None),
        (17, "legacy", None),
    ]