deer-flow/backend/tests/test_mcp_task_postgres.py

from __future__ import annotations

import asyncio
import os
import uuid
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit

import pytest
import pytest_asyncio
from sqlalchemy import event, select, text, update
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from deerflow.config.database_config import DatabaseConfig
from deerflow.persistence.engine import close_engine, get_engine, get_session_factory, init_engine_from_config
from deerflow.persistence.mcp_tasks import McpTaskRepository
from deerflow.persistence.mcp_tasks.model import McpTaskRow
from deerflow.persistence.thread_meta import ThreadMetaRepository
from deerflow.persistence.thread_meta.model import ThreadMetaRow

POSTGRES_URL = os.environ.get("TEST_POSTGRES_URI")

pytestmark = pytest.mark.skipif(
    not POSTGRES_URL,
    reason="requires TEST_POSTGRES_URI (real Postgres for row-lock interleaving)",
)


def _postgres_url(url: str) -> str:
    parts = urlsplit(url)
    query = urlencode([(key, value) for key, value in parse_qsl(parts.query, keep_blank_values=True) if key not in {"sslmode", "channel_binding"}])
    return urlunsplit(parts._replace(query=query))


@pytest_asyncio.fixture()
async def postgres_repositories():
    assert POSTGRES_URL is not None
    schema = f"mcp_incarnation_{uuid.uuid4().hex}"
    await init_engine_from_config(
        DatabaseConfig(
            backend="postgres",
            postgres_url=_postgres_url(POSTGRES_URL),
            postgres_schema=schema,
        )
    )
    session_factory = get_session_factory()
    assert session_factory is not None
    try:
        yield ThreadMetaRepository(session_factory), McpTaskRepository(session_factory), session_factory
    finally:
        engine = get_engine()
        assert engine is not None
        async with engine.begin() as connection:
            await connection.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
        await close_engine()


async def _create_task(repo: McpTaskRepository, task_id: str) -> None:
    await repo.create(
        task_id=task_id,
        user_id="user-1",
        thread_id="thread-1",
        run_id=None,
        tool_call_id=None,
        server_name="reports",
        driver_name="fake",
        remote_task_id=f"remote-{task_id}",
        task_name="Generate report",
        status="working",
        result=None,
        result_preview=None,
        result_truncated=False,
        result_artifact=None,
        error=None,
        input_required=None,
        next_poll_at=None,
    )


@pytest.mark.asyncio
@pytest.mark.parametrize("mutation", ["delete", "update_owner"])
async def test_postgres_task_create_serializes_with_thread_mutation(postgres_repositories, mutation: str) -> None:
    thread_repo, task_repo, session_factory = postgres_repositories
    created = await thread_repo.create("thread-1", user_id="user-1")

    async with session_factory() as blocker:
        locked = (await blocker.execute(select(ThreadMetaRow).where(ThreadMetaRow.thread_id == "thread-1").with_for_update())).scalar_one()
        assert locked.incarnation == created["incarnation"]

        if mutation == "delete":
            mutation_task = asyncio.create_task(thread_repo.delete("thread-1", user_id=None))
        else:
            mutation_task = asyncio.create_task(thread_repo.update_owner("thread-1", "user-2", user_id=None))
        await asyncio.sleep(0.1)
        assert not mutation_task.done()

        # PostgreSQL grants the mutation's earlier queued row-lock request
        # before this later FOR SHARE request. The task therefore observes the
        # committed delete/owner change rather than the pre-mutation row.
        create_task = asyncio.create_task(_create_task(task_repo, f"task-{mutation}"))
        await asyncio.sleep(0.1)
        assert not create_task.done()
        await blocker.commit()

    await asyncio.wait_for(mutation_task, timeout=5)
    await asyncio.wait_for(create_task, timeout=5)

    async with session_factory() as session:
        task = await session.get(McpTaskRow, f"task-{mutation}")
    assert task is not None
    assert task.thread_incarnation is None


@pytest.mark.asyncio
async def test_postgres_task_create_uses_share_lock(postgres_repositories) -> None:
    thread_repo, task_repo, _session_factory = postgres_repositories
    await thread_repo.create("thread-1", user_id="user-1")
    engine = get_engine()
    assert engine is not None
    statements: list[str] = []

    def capture_statement(_conn, _cursor, statement, _parameters, _context, _executemany):
        statements.append(" ".join(statement.upper().split()))

    event.listen(engine.sync_engine, "before_cursor_execute", capture_statement)
    try:
        await _create_task(task_repo, "task-share-lock")
    finally:
        event.remove(engine.sync_engine, "before_cursor_execute", capture_statement)

    assert any(statement.endswith("FOR SHARE") for statement in statements)
    assert not any(statement.endswith("FOR KEY SHARE") for statement in statements)


@pytest.mark.asyncio
async def test_postgres_repository_holds_share_lock_until_task_commit(postgres_repositories) -> None:
    thread_repo, _task_repo, session_factory = postgres_repositories
    created = await thread_repo.create("thread-1", user_id="user-1")
    engine = get_engine()
    assert engine is not None
    task_commit_entered = asyncio.Event()
    allow_task_commit = asyncio.Event()
    owner_update_started = asyncio.Event()
    task_backend_pid: int | None = None
    owner_backend_pid: int | None = None

    class PausingTaskCommitSession(AsyncSession):
        async def commit(self) -> None:
            nonlocal task_backend_pid
            contains_target_task = any(isinstance(instance, McpTaskRow) and instance.id == "task-lock-lifetime" for instance in self.new)
            if contains_target_task:
                task_backend_pid = await self.scalar(text("SELECT pg_backend_pid()"))
                task_commit_entered.set()
                await allow_task_commit.wait()
            await super().commit()

    task_session_factory = async_sessionmaker(
        engine,
        expire_on_commit=False,
        class_=PausingTaskCommitSession,
    )
    task_repo = McpTaskRepository(task_session_factory)

    def observe_owner_update(_conn, _cursor, statement, _parameters, _context, _executemany):
        normalized = " ".join(statement.upper().split())
        if normalized.startswith("UPDATE THREADS_META SET USER_ID"):
            owner_update_started.set()

    async def legacy_update_owner() -> None:
        nonlocal owner_backend_pid
        async with session_factory() as legacy_writer:
            owner_backend_pid = await legacy_writer.scalar(text("SELECT pg_backend_pid()"))
            await legacy_writer.execute(update(ThreadMetaRow).where(ThreadMetaRow.thread_id == "thread-1").values(user_id="user-2"))
            await legacy_writer.commit()

    event.listen(engine.sync_engine, "before_cursor_execute", observe_owner_update)
    task_create = asyncio.create_task(_create_task(task_repo, "task-lock-lifetime"))
    owner_update = None
    try:
        await asyncio.wait_for(task_commit_entered.wait(), timeout=5)
        owner_update = asyncio.create_task(legacy_update_owner())
        await asyncio.wait_for(owner_update_started.wait(), timeout=5)
        assert task_backend_pid is not None
        assert owner_backend_pid is not None
        async with session_factory() as observer:
            async with asyncio.timeout(5):
                while True:
                    if owner_update.done():
                        await owner_update
                        pytest.fail("owner update completed before the task transaction committed")
                    blockers = await observer.scalar(
                        text("SELECT pg_blocking_pids(:pid)"),
                        {"pid": owner_backend_pid},
                    )
                    if task_backend_pid in blockers:
                        break
                    await asyncio.sleep(0.01)
        allow_task_commit.set()
        await asyncio.wait_for(task_create, timeout=5)
        await asyncio.wait_for(owner_update, timeout=5)
    finally:
        allow_task_commit.set()
        event.remove(engine.sync_engine, "before_cursor_execute", observe_owner_update)
        if not task_create.done():
            task_create.cancel()
            await asyncio.gather(task_create, return_exceptions=True)
        if owner_update is not None and not owner_update.done():
            owner_update.cancel()
            await asyncio.gather(owner_update, return_exceptions=True)

    async with session_factory() as session:
        thread = await session.get(ThreadMetaRow, "thread-1")
        task = await session.get(McpTaskRow, "task-lock-lifetime")
    assert thread is not None
    assert thread.user_id == "user-2"
    assert task is not None
    assert task.thread_incarnation == created["incarnation"]