"""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"