agent-lab/tests/test_sqlite_store.py

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={"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)

    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]["total_tokens"] == 10
    assert usage["turns"] == [
        {
            "turn_index": 1,
            "prompt_tokens": 3,
            "completion_tokens": 7,
            "total_tokens": 10,
            "cached_tokens": 2,
            "elapsed_ms": 450,
        }
    ]
    assert usage["session"]["total_tokens"] == 10


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