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 <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user