Merge pull request 'Prune LangGraph checkpoints per thread and delete them with the conversation' (#85) from perf/prune-checkpoints into main

.env.example

@@ -204,6 +204,9 @@ OBSIDIAN_VAULT_NAME=
 # LangGraph checkpoint store (SQLite). Must be on a persistent volume.
 # Both blue and green containers share this volume via docker-compose.
 CHECKPOINT_DB_PATH=/app/checkpoints/checkpoints.db
+# Checkpoints keep a full message-list snapshot per turn; this bounds how many
+# are kept per conversation thread (older ones are pruned after each turn).
+CHECKPOINT_KEEP_PER_THREAD=20
 
 # Beli restaurant recommendations
 # Beli account credentials for each profile are stored encrypted in the database.

blueprints/conversation/__init__.py

@@ -32,6 +32,7 @@ from utils.s3_client import get_image as s3_get_image
 from utils.s3_client import upload_file as s3_upload_file
 from utils.s3_client import upload_image as s3_upload_image
 
+from .checkpoints import delete_thread_checkpoints
 from .logic import (
     add_message_to_conversation,
     persist_agent_tool_messages,
@@ -758,5 +759,11 @@ async def delete_conversation(conversation_id: str):
     store = getattr(current_app, "conversation_summary_store", None)
     if store is not None:
         await store.delete(str(user_uuid), str(conversation.id))
+    thread_id = str(conversation.id)
     await conversation.delete()
+    checkpointer = getattr(
+        getattr(current_app, "main_agent", None), "checkpointer", None
+    )
+    if checkpointer is not None:
+        await delete_thread_checkpoints(checkpointer, thread_id)
     return jsonify({"success": True})

blueprints/conversation/checkpoints.py

@@ -0,0 +1,87 @@
+"""Keeps the LangGraph checkpoint DB from growing without bound.
+
+Every `AsyncSqliteSaver` row stores a full snapshot of the graph's channel
+values (the complete message list so far), not a delta — so an unpruned
+thread's checkpoints grow roughly linearly with turn count.
+"""
+
+import logging
+import os
+
+DEFAULT_KEEP_PER_THREAD = 20
+
+
+def keep_per_thread() -> int:
+    raw = os.getenv("CHECKPOINT_KEEP_PER_THREAD", str(DEFAULT_KEEP_PER_THREAD))
+    try:
+        value = int(raw)
+    except ValueError:
+        value = -1
+    if value < 1:
+        logging.warning(
+            "CHECKPOINT_KEEP_PER_THREAD=%r is invalid; falling back to %d",
+            raw,
+            DEFAULT_KEEP_PER_THREAD,
+        )
+        return DEFAULT_KEEP_PER_THREAD
+    return value
+
+
+async def prune_thread_checkpoints(
+    checkpointer, thread_id: str, keep: int | None = None
+) -> None:
+    """Keep only the newest `keep` checkpoints (and their writes) for a thread.
+
+    Safe to call after every turn: each checkpoint row is a self-contained
+    channel snapshot rather than a delta, and `aget_tuple`/`aget_state` only
+    ever read the single newest row for a (thread_id, checkpoint_ns) pair, so
+    dropping older rows never affects state reads. Reuses the saver's own
+    lock so this never races its aput/aget_tuple calls.
+    """
+    keep = keep_per_thread() if keep is None else keep
+    thread_id = str(thread_id)
+    conn = checkpointer.conn
+    async with checkpointer.lock, conn.cursor() as cur:
+        await cur.execute(
+            "SELECT DISTINCT checkpoint_ns FROM checkpoints WHERE thread_id = ?",
+            (thread_id,),
+        )
+        namespaces = [row[0] for row in await cur.fetchall()]
+        for ns in namespaces:
+            await cur.execute(
+                """
+                DELETE FROM checkpoints
+                WHERE thread_id = ? AND checkpoint_ns = ?
+                AND checkpoint_id NOT IN (
+                    SELECT checkpoint_id FROM checkpoints
+                    WHERE thread_id = ? AND checkpoint_ns = ?
+                    ORDER BY checkpoint_id DESC
+                    LIMIT ?
+                )
+                """,
+                (thread_id, ns, thread_id, ns, keep),
+            )
+            await cur.execute(
+                """
+                DELETE FROM writes
+                WHERE thread_id = ? AND checkpoint_ns = ?
+                AND checkpoint_id NOT IN (
+                    SELECT checkpoint_id FROM checkpoints
+                    WHERE thread_id = ? AND checkpoint_ns = ?
+                )
+                """,
+                (thread_id, ns, thread_id, ns),
+            )
+        await conn.commit()
+
+
+async def delete_thread_checkpoints(checkpointer, thread_id: str) -> None:
+    """Delete a conversation's checkpoint thread; never raises.
+
+    Called from conversation delete, which must still succeed even if the
+    checkpoint store is unreachable or the thread is already gone.
+    """
+    try:
+        await checkpointer.adelete_thread(str(thread_id))
+    except Exception:
+        logging.exception("Failed to delete checkpoint thread %s", thread_id)

blueprints/conversation/runtime.py

@@ -13,6 +13,7 @@ from langchain_core.messages import convert_to_messages
 from quart import current_app
 
 from metrics import AGENT_REQUESTS_TOTAL, AGENT_RESPONSE_SECONDS
+from .checkpoints import prune_thread_checkpoints
 from .logic import get_agent_for
 from .memory import build_system_prompt_with_memories
 
@@ -119,6 +120,18 @@ def _record(turn: PreparedTurn, started: float, result: str) -> None:
     )
 
 
+async def _prune_after_turn(turn: PreparedTurn) -> None:
+    checkpointer = getattr(turn.agent, "checkpointer", None)
+    if checkpointer is None:
+        return
+    try:
+        await prune_thread_checkpoints(
+            checkpointer, turn.config["configurable"]["thread_id"]
+        )
+    except Exception:
+        logging.exception("Failed to prune checkpoints")
+
+
 async def invoke_prepared(turn: PreparedTurn) -> dict:
     started = time.monotonic()
     result = "success"
@@ -132,6 +145,7 @@ async def invoke_prepared(turn: PreparedTurn) -> dict:
         raise
     finally:
         _record(turn, started, result)
+        await _prune_after_turn(turn)
 
 
 async def invoke_agent_for_user(
@@ -169,3 +183,4 @@ async def stream_prepared(turn: PreparedTurn):
         raise
     finally:
         _record(turn, started, result)
+        await _prune_after_turn(turn)

docker-compose.yml

@@ -11,6 +11,7 @@ x-raggr-common: &raggr-common
     - PROMPT_BUDGET_TOKENS=${PROMPT_BUDGET_TOKENS:-16000}
     - RESPONSE_RESERVE_TOKENS=${RESPONSE_RESERVE_TOKENS:-4096}
     - CHECKPOINT_DB_PATH=${CHECKPOINT_DB_PATH:-/app/checkpoints/checkpoints.db}
+    - CHECKPOINT_KEEP_PER_THREAD=${CHECKPOINT_KEEP_PER_THREAD:-20}
     - CONVERSATION_SUMMARY_DB_PATH=${CONVERSATION_SUMMARY_DB_PATH:-${CHECKPOINT_DB_PATH:-/app/checkpoints/checkpoints.db}.summaries.db}
     - CHAT_RUN_LEASE_SECONDS=${CHAT_RUN_LEASE_SECONDS:-30}
     - CHAT_RUN_MAX_SECONDS=${CHAT_RUN_MAX_SECONDS:-1800}

tests/unit/test_checkpoints.py

@@ -0,0 +1,226 @@
+"""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