From c329d15f7f5bd37fd3a393888904c68ae897b29e Mon Sep 17 00:00:00 2001 From: Ryan Chen Date: Tue, 13 Oct 2026 03:40:09 -0400 Subject: [PATCH] Prune LangGraph checkpoints per thread and delete them with the conversation Every graph step stores a full snapshot of the conversation, nothing pruned old ones, and deleting a conversation left its thread behind (about 1.6 MB per checkpoint at 200 turns). After each turn, only the newest CHECKPOINT_KEEP_PER_THREAD (default 20) checkpoints of that thread are kept, plus their pending writes; state reads only use the newest, and pruning a 200-checkpoint thread takes about 2 ms. Deleting a conversation also deletes its thread, without failing the request if the checkpoint store errors. Co-Authored-By: Claude Opus 5.5 --- .env.example | 3 + blueprints/conversation/__init__.py | 7 + blueprints/conversation/checkpoints.py | 87 ++++++++++ blueprints/conversation/runtime.py | 15 ++ docker-compose.yml | 1 + tests/unit/test_checkpoints.py | 226 +++++++++++++++++++++++++ 6 files changed, 339 insertions(+) create mode 100644 blueprints/conversation/checkpoints.py create mode 100644 tests/unit/test_checkpoints.py diff --git a/.env.example b/.env.example index 95aecb0..39e53a6 100644 --- a/.env.example +++ b/.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. diff --git a/blueprints/conversation/__init__.py b/blueprints/conversation/__init__.py index a688149..6a5ed56 100644 --- a/blueprints/conversation/__init__.py +++ b/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}) diff --git a/blueprints/conversation/checkpoints.py b/blueprints/conversation/checkpoints.py new file mode 100644 index 0000000..93c0a79 --- /dev/null +++ b/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) diff --git a/blueprints/conversation/runtime.py b/blueprints/conversation/runtime.py index 740293c..0f937fb 100644 --- a/blueprints/conversation/runtime.py +++ b/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) diff --git a/docker-compose.yml b/docker-compose.yml index 18afa6b..fdca63c 100644 --- a/docker-compose.yml +++ b/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} diff --git a/tests/unit/test_checkpoints.py b/tests/unit/test_checkpoints.py new file mode 100644 index 0000000..119eb62 --- /dev/null +++ b/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