253 lines
11 KiB
Python
253 lines
11 KiB
Python
"""A bad thread must not prevent healthy conversations from being archived."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import errno
|
|
import sqlite3
|
|
from contextlib import asynccontextmanager
|
|
from datetime import datetime
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
from langgraph.checkpoint.base import empty_checkpoint
|
|
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from app.gateway.routers import chat_archive as archive_api
|
|
from deerflow.persistence.base import Base
|
|
from deerflow.persistence.thread_meta.model import ThreadMetaRow
|
|
from deerflow.runtime.checkpointer.chat_archive import (
|
|
archive_db_path,
|
|
ensure_catalog,
|
|
try_acquire_archive_job_lock,
|
|
)
|
|
|
|
MONTH = "2025-02"
|
|
THREADS = ("a-first", "b-middle", "c-last")
|
|
|
|
|
|
async def checkpoint_ids(saver, thread_id):
|
|
return [
|
|
item.config["configurable"]["checkpoint_id"]
|
|
async for item in saver.alist({"configurable": {"thread_id": thread_id}})
|
|
]
|
|
|
|
|
|
async def add_checkpoint(saver, thread_id):
|
|
checkpoint = empty_checkpoint()
|
|
await saver.aput(
|
|
{"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}},
|
|
checkpoint, {"source": "archive-failure-test"}, checkpoint["channel_versions"],
|
|
)
|
|
return checkpoint["id"]
|
|
|
|
|
|
@asynccontextmanager
|
|
async def archive_fixture(tmp_path):
|
|
root = tmp_path / "archives"
|
|
ensure_catalog(root)
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'metadata.db'}")
|
|
try:
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all, tables=[ThreadMetaRow.__table__])
|
|
sessions = async_sessionmaker(engine, expire_on_commit=False)
|
|
created_at = datetime(2025, 2, 12, 12)
|
|
async with sessions() as session:
|
|
for thread_id in THREADS:
|
|
session.add(ThreadMetaRow(
|
|
thread_id=thread_id, assistant_id="assistant", user_id="owner",
|
|
display_name=thread_id, status="idle", metadata_json={},
|
|
created_at=created_at, updated_at=created_at,
|
|
))
|
|
await session.commit()
|
|
async with AsyncSqliteSaver.from_conn_string(str(tmp_path / "primary.db")) as primary:
|
|
await primary.setup()
|
|
for thread_id in THREADS:
|
|
await add_checkpoint(primary, thread_id)
|
|
yield SimpleNamespace(
|
|
root=root, primary=primary,
|
|
runner=archive_api.ChatArchiveRunner(root, primary, sessions),
|
|
)
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
async def run_job(fixture, action):
|
|
job = await fixture.runner.start(MONTH, action)
|
|
await fixture.runner._tasks[job["job_id"]]
|
|
return archive_api._job_dict(fixture.root, job["job_id"])
|
|
|
|
|
|
def assert_partial(job, failed_thread, total=3):
|
|
assert job["status"] == "partial"
|
|
assert job["processed_threads"] == total
|
|
assert job["failed_threads"] == 1
|
|
assert failed_thread in job["error"]
|
|
assert job["current_thread_id"] is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("failed_thread", THREADS[:2])
|
|
async def test_copy_failure_continues_and_never_exposes_unverified_thread_to_purge(tmp_path, monkeypatch, failed_thread):
|
|
real_copy = archive_api.copy_thread_state
|
|
|
|
async def fail_one(source, target, thread_id):
|
|
if thread_id == failed_thread:
|
|
raise ValueError("unsupported checkpoint payload")
|
|
return await real_copy(source, target, thread_id)
|
|
|
|
monkeypatch.setattr(archive_api, "copy_thread_state", fail_one)
|
|
monkeypatch.setattr(archive_api, "_PURGE_THREAD_DELAY_SECONDS", 0)
|
|
async with archive_fixture(tmp_path) as fixture:
|
|
result = await run_job(fixture, "archive")
|
|
assert_partial(result, failed_thread)
|
|
assert "unsupported checkpoint payload" in result["error"]
|
|
with sqlite3.connect(fixture.root / "catalog.db") as conn:
|
|
rows = conn.execute("SELECT thread_id FROM archive_threads").fetchall()
|
|
assert {row[0] for row in rows} == set(THREADS) - {failed_thread}
|
|
async with AsyncSqliteSaver.from_conn_string(str(archive_db_path(fixture.root, MONTH))) as archived:
|
|
assert await checkpoint_ids(archived, "c-last")
|
|
|
|
purge = await run_job(fixture, "purge")
|
|
assert purge["status"] == "completed"
|
|
assert purge["processed_threads"] == 2
|
|
for thread_id in THREADS:
|
|
assert bool(await checkpoint_ids(fixture.primary, thread_id)) is (thread_id == failed_thread)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("failure_mode", ["copy_error", "verification_error"])
|
|
async def test_purge_keeps_unverified_pg_data_and_continues_other_threads(tmp_path, monkeypatch, failure_mode):
|
|
monkeypatch.setattr(archive_api, "_PURGE_THREAD_DELAY_SECONDS", 0)
|
|
async with archive_fixture(tmp_path) as fixture:
|
|
assert (await run_job(fixture, "archive"))["status"] == "completed"
|
|
await add_checkpoint(fixture.primary, "a-first")
|
|
real_copy = archive_api.copy_thread_state
|
|
|
|
async def fail_one(source, target, thread_id):
|
|
if thread_id == "a-first":
|
|
if failure_mode == "copy_error":
|
|
raise ValueError("cannot copy new checkpoint")
|
|
return False # Deliberately leave the new PG checkpoint absent from SQLite.
|
|
return await real_copy(source, target, thread_id)
|
|
|
|
monkeypatch.setattr(archive_api, "copy_thread_state", fail_one)
|
|
purge = await run_job(fixture, "purge")
|
|
assert_partial(purge, "a-first")
|
|
assert len(await checkpoint_ids(fixture.primary, "a-first")) == 2
|
|
assert await checkpoint_ids(fixture.primary, "b-middle") == []
|
|
assert await checkpoint_ids(fixture.primary, "c-last") == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rollback_continues_after_one_thread_fails(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(archive_api, "_PURGE_THREAD_DELAY_SECONDS", 0)
|
|
async with archive_fixture(tmp_path) as fixture:
|
|
await run_job(fixture, "archive")
|
|
await run_job(fixture, "purge")
|
|
real_copy = archive_api.copy_thread_state
|
|
|
|
async def fail_one(source, target, thread_id):
|
|
if thread_id == "a-first":
|
|
raise ValueError("cannot restore this checkpoint")
|
|
return await real_copy(source, target, thread_id)
|
|
|
|
monkeypatch.setattr(archive_api, "copy_thread_state", fail_one)
|
|
result = await run_job(fixture, "rollback")
|
|
assert_partial(result, "a-first")
|
|
assert await checkpoint_ids(fixture.primary, "a-first") == []
|
|
assert await checkpoint_ids(fixture.primary, "b-middle")
|
|
assert await checkpoint_ids(fixture.primary, "c-last")
|
|
async with AsyncSqliteSaver.from_conn_string(str(archive_db_path(fixture.root, MONTH))) as archived:
|
|
assert await checkpoint_ids(archived, "a-first")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_disk_exhaustion_stops_instead_of_attempting_more_threads(tmp_path, monkeypatch):
|
|
visited = []
|
|
|
|
async def disk_full(source, target, thread_id):
|
|
visited.append(thread_id)
|
|
raise OSError(errno.ENOSPC, "disk full")
|
|
|
|
monkeypatch.setattr(archive_api, "copy_thread_state", disk_full)
|
|
async with archive_fixture(tmp_path) as fixture:
|
|
job = await run_job(fixture, "archive")
|
|
assert job["status"] == "failed"
|
|
assert visited == ["a-first"]
|
|
for thread_id in THREADS:
|
|
assert await checkpoint_ids(fixture.primary, thread_id)
|
|
|
|
|
|
def persisted_job(job_id, *, status="running"):
|
|
return {
|
|
"job_id": job_id, "month": MONTH, "action": "archive", "status": status,
|
|
"total_threads": 3, "processed_threads": 0, "failed_threads": 0,
|
|
"current_thread_id": None, "message": "", "error": None,
|
|
"created_at": "2026-10-25T00:00:00+08:00", "updated_at": "2026-10-25T00:00:00+08:00",
|
|
}
|
|
|
|
|
|
def test_recovery_leaves_live_older_process_alone_then_recovers_after_lock_release(tmp_path):
|
|
root = tmp_path / "recovery"
|
|
archive_api._save_job(root, persisted_job("older-live-process"))
|
|
archive_api._upsert_month(root, MONTH, status="copying", total=3, job_id="older-live-process")
|
|
operation_lock = try_acquire_archive_job_lock(root)
|
|
assert operation_lock is not None
|
|
try:
|
|
archive_api._recover_jobs(root, started_before="2026-10-26T00:00:00+08:00")
|
|
assert archive_api._job_dict(root, "older-live-process")["status"] == "running"
|
|
finally:
|
|
operation_lock.close()
|
|
archive_api._recover_jobs(root, started_before="2026-10-26T00:00:00+08:00")
|
|
recovered = archive_api._job_dict(root, "older-live-process")
|
|
assert recovered["status"] == "interrupted"
|
|
assert recovered["error"]
|
|
assert archive_api._catalog_months(root)[MONTH]["status"] == "interrupted"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_another_process_operation_lock_prevents_start(tmp_path):
|
|
async with archive_fixture(tmp_path) as fixture:
|
|
lock = try_acquire_archive_job_lock(fixture.root)
|
|
assert lock is not None
|
|
try:
|
|
with pytest.raises(HTTPException) as error:
|
|
await fixture.runner.start(MONTH, "archive")
|
|
assert error.value.status_code == 409
|
|
finally:
|
|
lock.close()
|
|
|
|
|
|
def test_latest_job_uses_insertion_order_when_timestamps_tie(tmp_path):
|
|
root = tmp_path / "tie"
|
|
archive_api._save_job(root, persisted_job("z-earlier", status="completed"))
|
|
archive_api._save_job(root, persisted_job("a-newer", status="partial"))
|
|
assert archive_api._catalog_months(root)[MONTH]["latest_job"]["job_id"] == "a-newer"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_inventory_rechecks_owner_after_previously_live_process_exits(tmp_path):
|
|
root = tmp_path / "later-exit"
|
|
archive_api._save_job(root, persisted_job("later-exit"))
|
|
request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace()))
|
|
lock = try_acquire_archive_job_lock(root)
|
|
assert lock is not None
|
|
try:
|
|
await archive_api._ensure_archive_recovery_async(request, root)
|
|
assert archive_api._job_dict(root, "later-exit")["status"] == "running"
|
|
finally:
|
|
lock.close()
|
|
await archive_api._ensure_archive_recovery_async(request, root)
|
|
assert archive_api._job_dict(root, "later-exit")["status"] == "interrupted"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_new_action_recovers_unowned_stale_job_without_gateway_restart(tmp_path):
|
|
async with archive_fixture(tmp_path) as fixture:
|
|
archive_api._save_job(fixture.root, persisted_job("orphaned"))
|
|
result = await run_job(fixture, "archive")
|
|
assert result["status"] == "completed"
|
|
assert archive_api._job_dict(fixture.root, "orphaned")["status"] == "interrupted"
|