simbarag/tests/unit/test_checkpoints.py

"""LangGraph checkpoint pruning/deletion: bounds the on-disk checkpoint DB."""

from types import SimpleNamespace
from unittest.mock import AsyncMock

import aiosqlite
import pytest
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.base import empty_checkpoint
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from langgraph.graph import END, START, MessagesState, StateGraph
from quart import Quart
from quart_jwt_extended import JWTManager, create_access_token
from tortoise import Tortoise

from blueprints.conversation import conversation_blueprint
from blueprints.conversation.checkpoints import (
    DEFAULT_KEEP_PER_THREAD,
    delete_thread_checkpoints,
    keep_per_thread,
    prune_thread_checkpoints,
)
from blueprints.conversation.models import Conversation
from blueprints.users.models import User


async def _checkpoint_ids(saver, thread_id):
    async with saver.conn.execute(
        "SELECT checkpoint_id FROM checkpoints WHERE thread_id = ? ORDER BY checkpoint_id",
        (str(thread_id),),
    ) as cur:
        return [row[0] async for row in cur]


async def _write_rows(saver, thread_id):
    async with saver.conn.execute(
        "SELECT checkpoint_id FROM writes WHERE thread_id = ?",
        (str(thread_id),),
    ) as cur:
        return [row[0] async for row in cur]


async def _add_checkpoint(saver, thread_id):
    checkpoint = empty_checkpoint()
    config = {"configurable": {"thread_id": str(thread_id), "checkpoint_ns": ""}}
    await saver.aput(config, checkpoint, {}, {})
    return checkpoint["id"]


@pytest.fixture
async def saver(tmp_path):
    async with aiosqlite.connect(tmp_path / "checkpoints.db") as conn:
        s = AsyncSqliteSaver(conn)
        await s.setup()
        yield s


async def test_prune_keeps_newest_n_and_drops_orphaned_writes_only_for_pruned_thread(
    saver,
):
    kept_ids = []
    pruned_ids = []
    for i in range(5):
        cid = await _add_checkpoint(saver, "pruned-thread")
        (pruned_ids if i < 3 else kept_ids).append(cid)
        await saver.aput_writes(
            {
                "configurable": {
                    "thread_id": "pruned-thread",
                    "checkpoint_ns": "",
                    "checkpoint_id": cid,
                }
            },
            [("some_channel", "value")],
            f"task-{i}",
        )
    other_id = await _add_checkpoint(saver, "other-thread")
    await saver.aput_writes(
        {
            "configurable": {
                "thread_id": "other-thread",
                "checkpoint_ns": "",
                "checkpoint_id": other_id,
            }
        },
        [("some_channel", "value")],
        "task-other",
    )

    await prune_thread_checkpoints(saver, "pruned-thread", keep=2)

    remaining = await _checkpoint_ids(saver, "pruned-thread")
    assert sorted(remaining) == sorted(kept_ids)
    remaining_writes = await _write_rows(saver, "pruned-thread")
    assert sorted(remaining_writes) == sorted(kept_ids)
    for cid in pruned_ids:
        assert cid not in remaining
        assert cid not in remaining_writes

    # Untouched: pruning one thread must not affect another.
    assert await _checkpoint_ids(saver, "other-thread") == [other_id]
    assert await _write_rows(saver, "other-thread") == [other_id]


async def test_prune_preserves_full_latest_messages(saver):
    def echo(state):
        last = state["messages"][-1].content
        return {"messages": [AIMessage(content=f"reply {last}")]}

    graph = StateGraph(MessagesState)
    graph.add_node("answer", echo)
    graph.add_edge(START, "answer")
    graph.add_edge("answer", END)
    agent = graph.compile(checkpointer=saver)

    thread_id = "many-turns"
    config = {"configurable": {"thread_id": thread_id}}
    turn_count = 25
    for i in range(turn_count):
        await agent.ainvoke(
            {"messages": [HumanMessage(content=f"msg {i}")]}, config=config
        )

    assert len(await _checkpoint_ids(saver, thread_id)) > DEFAULT_KEEP_PER_THREAD

    await prune_thread_checkpoints(saver, thread_id, keep=5)

    assert len(await _checkpoint_ids(saver, thread_id)) == 5
    state = await agent.aget_state(config)
    # Every turn's human + AI message is still present: pruning dropped old
    # checkpoint rows, but the newest row's channel_values already holds the
    # full accumulated message list (not a delta).
    assert len(state.values["messages"]) == turn_count * 2
    assert state.values["messages"][0].content == "msg 0"
    assert state.values["messages"][-1].content == f"reply msg {turn_count - 1}"


@pytest.fixture
async def conv_app(tmp_path):
    await Tortoise.init(
        db_url=f"sqlite://{tmp_path / 'conv.sqlite'}",
        modules={
            "models": ["blueprints.users.models", "blueprints.conversation.models"]
        },
    )
    await Tortoise.generate_schemas()
    app = Quart(__name__)
    app.config.update(
        TESTING=True,
        JWT_SECRET_KEY="checkpoint-test-secret-at-least-thirty-two-characters",
    )
    JWTManager(app)
    app.register_blueprint(conversation_blueprint)
    user = await User.create(username="owner", email="owner@test")
    async with app.app_context():
        headers = {
            "Authorization": f"Bearer {create_access_token(identity=str(user.id))}"
        }
    try:
        yield app, headers, user
    finally:
        await Tortoise.close_connections()


async def test_delete_conversation_via_api_removes_its_thread_and_leaves_others(
    conv_app, tmp_path
):
    app, headers, user = conv_app
    conv1 = await Conversation.create(user=user, name="one")
    conv2 = await Conversation.create(user=user, name="two")

    async with aiosqlite.connect(tmp_path / "checkpoints.db") as conn:
        s = AsyncSqliteSaver(conn)
        await s.setup()
        await _add_checkpoint(s, str(conv1.id))
        await _add_checkpoint(s, str(conv2.id))
        app.main_agent = SimpleNamespace(checkpointer=s)

        response = await app.test_client().delete(
            f"/api/conversation/{conv1.id}", headers=headers
        )
        assert response.status_code == 200
        assert (await response.get_json())["success"] is True
        assert await Conversation.get_or_none(id=conv1.id) is None

        assert await _checkpoint_ids(s, str(conv1.id)) == []
        assert len(await _checkpoint_ids(s, str(conv2.id))) == 1


async def test_delete_conversation_succeeds_even_when_checkpoint_delete_fails(
    conv_app,
):
    app, headers, user = conv_app
    conv = await Conversation.create(user=user, name="doomed")
    app.main_agent = SimpleNamespace(
        checkpointer=SimpleNamespace(
            adelete_thread=AsyncMock(side_effect=RuntimeError("sqlite is gone"))
        )
    )

    response = await app.test_client().delete(
        f"/api/conversation/{conv.id}", headers=headers
    )
    assert response.status_code == 200
    assert (await response.get_json())["success"] is True
    assert await Conversation.get_or_none(id=conv.id) is None


async def test_delete_thread_checkpoints_swallows_errors(caplog):
    checkpointer = SimpleNamespace(
        adelete_thread=AsyncMock(side_effect=RuntimeError("boom"))
    )
    await delete_thread_checkpoints(checkpointer, "t")
    assert "Failed to delete checkpoint thread" in caplog.text


@pytest.mark.parametrize("raw", ["0", "-5", "not-a-number"])
def test_keep_per_thread_rejects_invalid_values(monkeypatch, caplog, raw):
    monkeypatch.setenv("CHECKPOINT_KEEP_PER_THREAD", raw)
    assert keep_per_thread() == DEFAULT_KEEP_PER_THREAD
    assert "invalid" in caplog.text


def test_keep_per_thread_accepts_valid_override(monkeypatch):
    monkeypatch.setenv("CHECKPOINT_KEEP_PER_THREAD", "7")
    assert keep_per_thread() == 7