"""Focused persistence tests for project-scoped cross-session memory.""" from __future__ import annotations from datetime import UTC, datetime, timedelta import pytest from sqlalchemy import delete, event, func, select from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from deerflow.persistence.analysis_project_memory import ( AnalysisProjectMemoryItemRow, AnalysisProjectMemoryJobRow, AnalysisProjectMemoryStateRow, AnalysisProjectMemoryStore, ) from deerflow.persistence.analysis_projects import AnalysisProjectRow async def _sql_store(tmp_path, *, project_id: str = "project-1", owner: str = "user-1"): engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'project-memory.db'}") @event.listens_for(engine.sync_engine, "connect") def _foreign_keys(dbapi_connection, _connection_record): # noqa: ANN001 cursor = dbapi_connection.cursor() cursor.execute("PRAGMA foreign_keys=ON") cursor.close() async with engine.begin() as connection: await connection.run_sync(AnalysisProjectRow.__table__.create) await connection.run_sync(AnalysisProjectMemoryStateRow.__table__.create) await connection.run_sync(AnalysisProjectMemoryItemRow.__table__.create) await connection.run_sync(AnalysisProjectMemoryJobRow.__table__.create) sf = async_sessionmaker(engine, expire_on_commit=False) now = datetime.now(UTC) async with sf() as session: session.add( AnalysisProjectRow( id=project_id, user_id=owner, name="Project", mode="qa", background="", knowledge_base_ids=[], assistant_knowledge_base_ids=[], created_at=now, updated_at=now, ) ) await session.commit() return AnalysisProjectMemoryStore(sf), sf, engine @pytest.mark.asyncio async def test_memory_fallback_supersedes_items_and_builds_snapshot() -> None: store = AnalysisProjectMemoryStore() state = await store.get_state("p1", "u1") assert state is not None assert state["enabled"] is True assert await store.get_state("p1", "another-user") is None first = await store.upsert_item( "p1", "u1", logical_key="deadline", kind="constraint", content="Deadline is Friday", source_thread_id="thread-1", source_message_ids=["message-1"], ) second = await store.upsert_item( "p1", "u1", logical_key="deadline", kind="constraint", content="Deadline is Monday", source_thread_id="thread-2", ) assert first["id"] != second["id"] all_items = await store.list_items("p1", "u1", status=None) assert {item["status"] for item in all_items} == {"active", "superseded"} old = next(item for item in all_items if item["id"] == first["id"]) assert old["superseded_by"] == second["id"] state = await store.get_state("p1", "u1") assert state is not None assert state["item_count"] == 1 assert state["snapshot"]["summary"] == "已沉淀 1 条项目记忆。" assert state["snapshot"]["statistics"]["by_kind"]["constraint"] == 1 assert state["snapshot"]["constraints"] == ["Deadline is Monday"] @pytest.mark.asyncio async def test_recall_uses_lexical_fallback_and_vector_identity_gate() -> None: store = AnalysisProjectMemoryStore() await store.upsert_item( "p1", "u1", logical_key="budget", kind="fact", content="Budget is 500 yuan", importance=0.2, embedding=[1.0, 0.0], embedding_model="encoder-a", embedding_fingerprint="a-v1", ) await store.upsert_item( "p1", "u1", logical_key="venue", kind="fact", content="Venue is Shanghai", importance=0.9, embedding=[0.0, 1.0], embedding_model="encoder-b", embedding_fingerprint="b-v1", ) lexical = await store.recall("p1", "u1", query="What is the budget?") assert lexical["items"][0]["logical_key"] == "budget" assert all(item["recall_signals"]["vector"] == 0 for item in lexical["items"]) wrong_model = await store.recall( "p1", "u1", query="unrelated", query_embedding=[1.0, 0.0], query_embedding_model="encoder-b", query_embedding_fingerprint="b-v1", ) budget = next(item for item in wrong_model["items"] if item["logical_key"] == "budget") assert budget["recall_signals"]["vector"] == 0 assert budget["recall_signals"]["vector_compatible"] is False matching = await store.recall( "p1", "u1", query="unrelated", query_embedding=[1.0, 0.0], query_embedding_model="encoder-a", query_embedding_fingerprint="a-v1", ) budget = next(item for item in matching["items"] if item["logical_key"] == "budget") assert budget["recall_signals"]["vector"] == 1.0 @pytest.mark.asyncio async def test_apply_extraction_snapshot_is_order_independent() -> None: first = AnalysisProjectMemoryStore() second = AnalysisProjectMemoryStore() deltas = [ { "logical_key": "decision-b", "kind": "decision", "content": "Use option B", "importance": 0.7, }, { "logical_key": "fact-a", "kind": "fact", "content": "The audience is internal", "importance": 0.8, }, ] for item in deltas: result = await first.apply_extraction( "p1", "u1", items=[item], snapshot={"summary": "untrusted"} ) assert result["applied"] is True for item in reversed(deltas): await second.apply_extraction( "p1", "u1", items=[item], snapshot={"summary": "different"} ) first_state = await first.get_state("p1", "u1") second_state = await second.get_state("p1", "u1") assert first_state is not None and second_state is not None assert first_state["snapshot"] == second_state["snapshot"] assert first_state["snapshot"]["summary"] == "已沉淀 2 条项目记忆。" assert "untrusted" not in str(first_state["snapshot"]) @pytest.mark.asyncio async def test_sql_job_idempotency_lease_retry_and_new_success_recovers_ready(tmp_path) -> None: store, _sf, engine = await _sql_store(tmp_path) try: await store.mark_rebuild("project-1", "user-1") first, created = await store.enqueue_job( "project-1", "user-1", source_hash="hash-1", extractor_version="v1", payload={"source_message_ids": ["m1"]}, source_thread_id="thread-1", max_attempts=1, ) duplicate, duplicate_created = await store.enqueue_job( "project-1", "user-1", source_hash="hash-1", extractor_version="v1", payload={"source_message_ids": ["m1"]}, source_thread_id="thread-1", ) assert created is True assert duplicate_created is False assert duplicate["id"] == first["id"] assert (await store.get_state("project-1", "user-1"))["status"] == "rebuilding" lease_until = datetime.now(UTC) + timedelta(minutes=1) [claimed] = await store.claim_jobs( lease_owner="worker-1", lease_until=lease_until, limit=1 ) assert claimed["attempt"] == 1 failed = await store.fail_job( claimed["id"], lease_owner="worker-1", error="temporary failure" ) assert failed is not None and failed["status"] == "failed" assert (await store.get_state("project-1", "user-1"))["status"] == "error" await store.enqueue_job( "project-1", "user-1", source_hash="hash-2", extractor_version="v1", payload={"source_message_ids": ["m2"]}, source_thread_id="thread-2", ) [claimed] = await store.claim_jobs( lease_owner="worker-2", lease_until=lease_until, limit=1 ) applied = await store.apply_extraction( "project-1", "user-1", items=[{"logical_key": "new", "kind": "fact", "content": "New fact"}], source_thread_id="thread-2", job_id=claimed["id"], lease_owner="worker-2", ) assert applied["applied"] is True completed = await store.complete_job( claimed["id"], lease_owner="worker-2", result={"count": 1} ) assert completed is not None and completed["status"] == "completed" assert (await store.get_state("project-1", "user-1"))["status"] == "ready" finally: await engine.dispose() @pytest.mark.asyncio async def test_invalidate_and_clear_fence_running_jobs(tmp_path) -> None: store, _sf, engine = await _sql_store(tmp_path) try: await store.upsert_item( "project-1", "user-1", logical_key="old", kind="fact", content="Old fact", source_thread_id="thread-1", source_message_ids=["m1"], ) job, _ = await store.enqueue_job( "project-1", "user-1", source_hash="turn-1", extractor_version="v1", payload={"source_message_ids": ["m1"]}, source_thread_id="thread-1", ) [claimed] = await store.claim_jobs( lease_owner="worker", lease_until=datetime.now(UTC) + timedelta(minutes=1) ) assert claimed["id"] == job["id"] invalidated = await store.invalidate_source( "project-1", "user-1", source_thread_id="thread-1", source_message_ids=["m1"], ) assert invalidated["invalidated_items"] == 1 assert invalidated["invalidated_jobs"] == 1 assert invalidated["state"]["status"] == "rebuilding" fenced = await store.apply_extraction( "project-1", "user-1", items=[{"logical_key": "stale", "kind": "fact", "content": "Stale fact"}], source_thread_id="thread-1", job_id=job["id"], lease_owner="worker", ) assert fenced["applied"] is False assert await store.list_items("project-1", "user-1") == [] second, _ = await store.enqueue_job( "project-1", "user-1", source_hash="turn-2", extractor_version="v1", payload={"source_message_ids": ["m2"]}, source_thread_id="thread-2", ) [second_claimed] = await store.claim_jobs( lease_owner="worker", lease_until=datetime.now(UTC) + timedelta(minutes=1) ) assert second_claimed["id"] == second["id"] await store.clear("project-1", "user-1") fenced = await store.apply_extraction( "project-1", "user-1", items=[{"logical_key": "stale-2", "kind": "fact", "content": "Stale again"}], job_id=second["id"], lease_owner="worker", ) assert fenced["applied"] is False assert await store.list_items("project-1", "user-1") == [] finally: await engine.dispose() @pytest.mark.asyncio async def test_sql_project_delete_cascades_all_memory_rows(tmp_path) -> None: store, sf, engine = await _sql_store(tmp_path) try: await store.upsert_item( "project-1", "user-1", logical_key="fact", kind="fact", content="Remember me" ) await store.enqueue_job( "project-1", "user-1", source_hash="hash", extractor_version="v1", payload={}, ) async with sf() as session: await session.execute(delete(AnalysisProjectRow).where(AnalysisProjectRow.id == "project-1")) await session.commit() async with sf() as session: counts = [] for model in ( AnalysisProjectMemoryStateRow, AnalysisProjectMemoryItemRow, AnalysisProjectMemoryJobRow, ): counts.append((await session.execute(select(func.count()).select_from(model))).scalar_one()) assert counts == [0, 0, 0] finally: await engine.dispose()