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