deer-flow/backend/tests/test_history_batch_queries.py

"""Cross-store contracts used by thread-global history pagination."""

from __future__ import annotations

import pytest

from deerflow.runtime import RunManager, RunStatus
from deerflow.runtime.events.store.memory import MemoryRunEventStore
from deerflow.runtime.runs.manager import EditReplayVisibility
from deerflow.runtime.runs.store.memory import MemoryRunStore


async def _seed_ai_messages(store):
    await store.put(
        thread_id="t1",
        run_id="r1",
        event_type="llm.ai.response",
        category="message",
        content={"type": "ai", "content": "first"},
        metadata={"caller": "lead_agent"},
    )
    await store.put(
        thread_id="t1",
        run_id="r1",
        event_type="llm.ai.response",
        category="message",
        content={"type": "ai", "content": "middleware"},
        metadata={"caller": "middleware:title"},
    )
    last = await store.put(
        thread_id="t1",
        run_id="r1",
        event_type="llm.ai.response",
        category="message",
        content={"type": "ai", "content": "last"},
        metadata={"caller": "lead_agent"},
    )
    other = await store.put(
        thread_id="t1",
        run_id="r2",
        event_type="llm.ai.response",
        category="message",
        content={"type": "ai", "content": "other"},
        metadata={"caller": "lead_agent"},
    )
    await store.put(
        thread_id="t1",
        run_id="r_mw",
        event_type="llm.ai.response",
        category="message",
        content={"type": "ai", "content": "middleware only"},
        metadata={"caller": "middleware:title"},
    )
    return {"r1": last["seq"], "r2": other["seq"]}


@pytest.mark.anyio
async def test_memory_event_store_returns_global_last_non_middleware_ai_seq():
    store = MemoryRunEventStore()
    expected = await _seed_ai_messages(store)
    result = await store.get_last_visible_ai_seq_by_run("t1", {"r1", "r2", "r_mw", "missing"})
    assert result == expected
    assert "r_mw" not in result


@pytest.mark.anyio
async def test_memory_event_store_defensively_rechecks_message_category():
    store = MemoryRunEventStore()
    expected = await store.put(
        thread_id="t1",
        run_id="r1",
        event_type="llm.ai.response",
        category="message",
        content={"type": "ai", "content": "visible"},
        metadata={"caller": "lead_agent"},
    )
    mutated = await store.put(
        thread_id="t1",
        run_id="r1",
        event_type="llm.ai.response",
        category="message",
        content={"type": "ai", "content": "no longer a message"},
        metadata={"caller": "lead_agent"},
    )
    # Memory projections intentionally share their row dictionaries. Recheck
    # category at read time so an accidental mutation cannot violate the same
    # contract that the DB and JSONL stores enforce explicitly.
    mutated["category"] = "trace"

    assert await store.get_last_visible_ai_seq_by_run("t1", {"r1"}) == {"r1": expected["seq"]}


@pytest.mark.anyio
async def test_jsonl_event_store_returns_global_last_non_middleware_ai_seq(tmp_path):
    from deerflow.runtime.events.store.jsonl import JsonlRunEventStore

    store = JsonlRunEventStore(base_dir=tmp_path)
    expected = await _seed_ai_messages(store)
    result = await store.get_last_visible_ai_seq_by_run("t1", {"r1", "r2", "r_mw", "missing"})
    assert result == expected
    assert "r_mw" not in result


@pytest.mark.anyio
async def test_db_event_store_returns_global_last_non_middleware_ai_seq(tmp_path):
    from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
    from deerflow.runtime.events.store.db import DbRunEventStore

    await init_engine("sqlite", url=f"sqlite+aiosqlite:///{tmp_path / 'events.db'}", sqlite_dir=str(tmp_path))
    try:
        store = DbRunEventStore(get_session_factory())
        expected = await _seed_ai_messages(store)
        result = await store.get_last_visible_ai_seq_by_run("t1", {"r1", "r2", "r_mw", "missing"})
        assert result == expected
        assert "r_mw" not in result
    finally:
        await close_engine()


@pytest.mark.anyio
async def test_memory_run_store_supersession_is_unbounded_and_owner_scoped():
    store = MemoryRunStore()
    for index in range(105):
        await store.put(f"normal-{index}", thread_id="t1", user_id="alice", status="success")
    await store.put(
        "regen-success",
        thread_id="t1",
        user_id="alice",
        status="success",
        metadata={"regenerate_from_run_id": "source-success"},
    )
    await store.put(
        "regen-failed",
        thread_id="t1",
        user_id="alice",
        status="error",
        metadata={"regenerate_from_run_id": "source-failed"},
    )
    await store.put(
        "regen-bob",
        thread_id="t1",
        user_id="bob",
        status="success",
        metadata={"regenerate_from_run_id": "source-bob"},
    )

    assert await store.list_successful_regenerate_sources("t1", user_id="alice") == {"source-success"}


@pytest.mark.anyio
async def test_run_repository_batch_queries_are_unbounded_and_owner_scoped(tmp_path):
    from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
    from deerflow.persistence.run import RunRepository

    await init_engine("sqlite", url=f"sqlite+aiosqlite:///{tmp_path / 'runs.db'}", sqlite_dir=str(tmp_path))
    try:
        repo = RunRepository(get_session_factory())
        for index in range(105):
            await repo.put(f"normal-{index}", thread_id="t1", user_id="alice", status="success")
        await repo.put(
            "regen-a",
            thread_id="t1",
            user_id="alice",
            status="success",
            metadata={"regenerate_from_run_id": "source-a"},
        )
        await repo.put(
            "regen-b",
            thread_id="t1",
            user_id="bob",
            status="success",
            metadata={"regenerate_from_run_id": "source-b"},
        )

        assert await repo.list_successful_regenerate_sources("t1", user_id="alice") == {"source-a"}
        rows = await repo.get_many_by_thread("t1", {"normal-0", "regen-a", "regen-b"}, user_id="alice")
        assert set(rows) == {"normal-0", "regen-a"}
    finally:
        await close_engine()


@pytest.mark.anyio
async def test_run_manager_prefers_latest_in_memory_regenerate_status():
    store = MemoryRunStore()
    await store.put(
        "regen",
        thread_id="t1",
        status="success",
        metadata={"regenerate_from_run_id": "source"},
    )
    manager = RunManager(store=store)
    # Simulate the same logical run being newer in memory than its persisted
    # successful snapshot.
    persisted = await manager.get("regen")
    assert persisted is not None
    manager._runs["regen"] = persisted
    manager._index_run_locked(persisted)
    persisted.status = RunStatus.error

    assert await manager.list_successful_regenerate_sources("t1", user_id=None) == set()


@pytest.mark.anyio
async def test_run_manager_uses_latest_attempt_for_shared_regenerate_source():
    manager = RunManager()
    older = await manager.create(
        "t1",
        metadata={"regenerate_from_run_id": "source"},
    )
    older.status = RunStatus.success
    newer = await manager.create(
        "t1",
        metadata={"regenerate_from_run_id": "source"},
    )
    newer.status = RunStatus.error

    assert await manager.list_successful_regenerate_sources("t1", user_id=None) == set()


@pytest.mark.anyio
async def test_run_manager_batch_history_methods_default_to_current_user():
    from types import SimpleNamespace

    from deerflow.runtime.user_context import reset_current_user, set_current_user

    store = MemoryRunStore()
    await store.put(
        "regen-alice",
        thread_id="shared-thread",
        user_id="alice",
        status="success",
        metadata={"regenerate_from_run_id": "source-alice"},
    )
    await store.put(
        "regen-bob",
        thread_id="shared-thread",
        user_id="bob",
        status="success",
        metadata={"regenerate_from_run_id": "source-bob"},
    )
    manager = RunManager(store=store)
    token = set_current_user(SimpleNamespace(id="alice"))
    try:
        sources = await manager.list_successful_regenerate_sources("shared-thread")
        records = await manager.get_many_by_thread("shared-thread", {"regen-alice", "regen-bob"})
    finally:
        reset_current_user(token)

    assert sources == {"source-alice"}
    assert set(records) == {"regen-alice"}


@pytest.mark.anyio
async def test_run_manager_batch_history_methods_fail_closed_without_user_context():
    from deerflow.runtime import user_context

    manager = RunManager(store=MemoryRunStore())
    token = user_context._current_user.set(None)
    try:
        with pytest.raises(RuntimeError, match="user_id=AUTO"):
            await manager.list_successful_regenerate_sources("t1")
        with pytest.raises(RuntimeError, match="user_id=AUTO"):
            await manager.get_many_by_thread("t1", {"run-1"})
    finally:
        user_context._current_user.reset(token)


@pytest.mark.anyio
async def test_run_manager_batch_history_methods_allow_explicit_unscoped_access():
    store = MemoryRunStore()
    await store.put(
        "regen-alice",
        thread_id="shared-thread",
        user_id="alice",
        status="success",
        metadata={"regenerate_from_run_id": "source-alice"},
    )
    await store.put(
        "regen-bob",
        thread_id="shared-thread",
        user_id="bob",
        status="success",
        metadata={"regenerate_from_run_id": "source-bob"},
    )
    manager = RunManager(store=store)

    sources = await manager.list_successful_regenerate_sources("shared-thread", user_id=None)
    records = await manager.get_many_by_thread("shared-thread", {"regen-alice", "regen-bob"}, user_id=None)

    assert sources == {"source-alice", "source-bob"}
    assert set(records) == {"regen-alice", "regen-bob"}


@pytest.mark.anyio
async def test_memory_run_store_lists_edit_replay_runs_unbounded_and_owner_scoped():
    store = MemoryRunStore()
    for index in range(105):
        await store.put(f"normal-{index}", thread_id="t1", user_id="alice", status="success")
    await store.put(
        "edit-success",
        thread_id="t1",
        user_id="alice",
        status="success",
        metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-success"},
    )
    await store.put(
        "edit-error",
        thread_id="t1",
        user_id="alice",
        status="error",
        metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-error"},
    )
    await store.put(
        "edit-bob",
        thread_id="t1",
        user_id="bob",
        status="success",
        metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-bob"},
    )

    rows = await store.list_edit_regenerate_runs("t1", user_id="alice")

    assert [row["run_id"] for row in rows] == ["edit-success", "edit-error"]


@pytest.mark.anyio
async def test_run_repository_lists_edit_replay_runs_unbounded_and_owner_scoped(tmp_path):
    from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
    from deerflow.persistence.run import RunRepository

    await init_engine("sqlite", url=f"sqlite+aiosqlite:///{tmp_path / 'runs.db'}", sqlite_dir=str(tmp_path))
    try:
        repo = RunRepository(get_session_factory())
        for index in range(105):
            await repo.put(f"normal-{index}", thread_id="t1", user_id="alice", status="success")
        await repo.put(
            "edit-success",
            thread_id="t1",
            user_id="alice",
            status="success",
            metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-success"},
        )
        await repo.put(
            "edit-error",
            thread_id="t1",
            user_id="alice",
            status="error",
            metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-error"},
        )
        await repo.put(
            "edit-bob",
            thread_id="t1",
            user_id="bob",
            status="success",
            metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-bob"},
        )

        rows = await repo.list_edit_regenerate_runs("t1", user_id="alice")

        assert [row["run_id"] for row in rows] == ["edit-success", "edit-error"]
    finally:
        await close_engine()


@pytest.mark.anyio
async def test_run_manager_computes_edit_replay_visibility_by_latest_attempt():
    manager = RunManager()
    running = await manager.create(
        "t1",
        metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-running"},
    )
    running.status = RunStatus.running
    success = await manager.create(
        "t1",
        metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-success"},
    )
    success.status = RunStatus.success
    failed = await manager.create(
        "t1",
        metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-failed"},
    )
    failed.status = RunStatus.error
    older_success = await manager.create(
        "t1",
        metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-retried"},
    )
    older_success.status = RunStatus.success
    newer_failed = await manager.create(
        "t1",
        metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-retried"},
    )
    newer_failed.status = RunStatus.interrupted
    older_failed_then_success = await manager.create(
        "t1",
        metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-retry-success"},
    )
    older_failed_then_success.status = RunStatus.error
    newer_success_after_failure = await manager.create(
        "t1",
        metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-retry-success"},
    )
    newer_success_after_failure.status = RunStatus.success

    visibility = await manager.list_edit_replay_visibility("t1", user_id=None)

    assert visibility == EditReplayVisibility(
        hidden_source_run_ids={"source-running", "source-success", "source-retry-success"},
        hidden_attempt_run_ids={failed.run_id, newer_failed.run_id, older_failed_then_success.run_id},
    )