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), ]