deerflow-code/offline-backend-20260512/backend/tests/test_thread_clear.py

"""Clear-all chat history keeps other users and system threads."""

from __future__ import annotations

import pytest
from sqlalchemy import select
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from sqlalchemy.pool import StaticPool

from app.gateway.routers.threads import router
from app.gateway.thread_clear import clear_owned_chat_threads
from deerflow.persistence.base import Base
from deerflow.persistence.run.model import RunRow
from deerflow.persistence.thread_meta.model import ThreadMetaRow
from deerflow.persistence.thread_meta.sql import ThreadMetaRepository


class _Cursor:
    def __init__(self, statements: list[tuple[str, tuple[str, ...]]]) -> None:
        self.statements = statements

    async def execute(self, sql: str, params: tuple[str, ...] = ()) -> None:
        self.statements.append((sql, params))

    async def executemany(self, sql: str, rows: list[tuple[str, ...]]) -> None:
        for row in rows:
            self.statements.append((sql, row))

    async def __aenter__(self) -> _Cursor:
        return self

    async def __aexit__(self, *_args: object) -> bool:
        return False


class _Conn:
    def __init__(self) -> None:
        self.statements: list[tuple[str, tuple[str, ...]]] = []
        self.commits = 0

    def cursor(self) -> _Cursor:
        return _Cursor(self.statements)

    async def commit(self) -> None:
        self.commits += 1


class _Lock:
    async def __aenter__(self) -> _Lock:
        return self

    async def __aexit__(self, *_args: object) -> bool:
        return False


class _SqliteSaver:
    def __init__(self) -> None:
        self.conn = _Conn()
        self.lock = _Lock()
        self.deleted: list[str] = []

    async def adelete_thread(self, thread_id: str) -> None:
        self.deleted.append(thread_id)


_SqliteSaver.__module__ = "langgraph.checkpoint.sqlite.aio"


class _FallbackSaver:
    def __init__(self) -> None:
        self.deleted: list[str] = []

    async def adelete_thread(self, thread_id: str) -> None:
        self.deleted.append(thread_id)


class _Paths:
    def __init__(self) -> None:
        self.removed: list[tuple[str, str | None]] = []

    def delete_thread_dir(self, thread_id: str, *, user_id: str | None = None) -> None:
        self.removed.append((thread_id, user_id))


async def _engine():
    engine = create_async_engine("sqlite+aiosqlite://", poolclass=StaticPool)
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    return engine, async_sessionmaker(engine, expire_on_commit=False)


@pytest.mark.asyncio
async def test_clear_removes_agent_chats_and_keeps_system_threads():
    engine, factory = await _engine()
    repo = ThreadMetaRepository(factory)
    await repo.create("chat-1", user_id="user-a", display_name="普通对话")
    await repo.create("agent-1", user_id="user-a", display_name="智能体", metadata={"agent_id": "writer"})
    await repo.create(
        "system-1",
        user_id="user-a",
        display_name="会商",
        metadata={"system": True, "thread_type": "roundtable"},
    )
    await repo.create("other-1", user_id="user-b", display_name="别人的对话")
    async with factory() as session:
        session.add(RunRow(run_id="run-chat", thread_id="chat-1", user_id="user-a"))
        session.add(RunRow(run_id="run-agent", thread_id="agent-1", user_id="user-a"))
        session.add(RunRow(run_id="run-system", thread_id="system-1", user_id="user-a"))
        await session.commit()

    saver = _SqliteSaver()
    paths = _Paths()
    deleted = await clear_owned_chat_threads(
        store=repo,
        user_id="user-a",
        checkpointer=saver,
        session_factory=factory,
        paths=paths,  # type: ignore[arg-type]
        wait_for_files=True,
    )

    assert deleted == 2
    assert saver.deleted == []
    assert saver.conn.commits == 1
    deleted_ids = {thread_id for _sql, params in saver.conn.statements for thread_id in params}
    assert deleted_ids == {"chat-1", "agent-1"}
    assert {thread_id for thread_id, _user_id in paths.removed} == {"chat-1", "agent-1"}

    remaining = await repo.search(user_id="user-a", exclude_system=False, limit=20)
    assert {row["thread_id"] for row in remaining} == {"system-1"}
    others = await repo.search(user_id="user-b", exclude_system=True, limit=20)
    assert [row["thread_id"] for row in others] == ["other-1"]

    async with factory() as session:
        run_ids = set((await session.execute(select(RunRow.run_id))).scalars().all())
        meta_ids = set((await session.execute(select(ThreadMetaRow.thread_id))).scalars().all())
    assert run_ids == {"run-system"}
    assert meta_ids == {"system-1", "other-1"}
    await engine.dispose()


@pytest.mark.asyncio
async def test_clear_falls_back_to_per_thread_checkpoint_delete():
    engine, factory = await _engine()
    repo = ThreadMetaRepository(factory)
    await repo.create("chat-1", user_id="user-a")
    saver = _FallbackSaver()
    deleted = await clear_owned_chat_threads(
        store=repo,
        user_id="user-a",
        checkpointer=saver,
        session_factory=factory,
        paths=_Paths(),  # type: ignore[arg-type]
        wait_for_files=True,
    )
    assert deleted == 1
    assert saver.deleted == ["chat-1"]
    remaining = await repo.search(user_id="user-a", exclude_system=True, limit=20)
    assert remaining == []
    await engine.dispose()


def test_clear_route_is_registered():
    matches = [
        route
        for route in router.routes
        if getattr(route, "path", None) == "/api/threads/clear" and "POST" in getattr(route, "methods", set())
    ]
    assert matches