deer-flow/backend/tests/test_thread_create_run_metadata_race.py

"""The detached run metadata writer participates in thread create admission."""

import asyncio
from contextlib import asynccontextmanager
from types import SimpleNamespace

import httpx
import pytest
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.store.memory import InMemoryStore
from sqlalchemy.exc import IntegrityError
from test_gateway_services import _stub_app_config as _base_config
from test_thread_create_idempotency import USER, app_for
from test_thread_create_idempotency import thread_store as _thread_store

from app.gateway import services
from app.gateway.routers import threads
from deerflow.runtime import RunManager, goal
from deerflow.runtime.events.store.memory import MemoryRunEventStore
from deerflow.runtime.runs.store.memory import MemoryRunStore
from deerflow.runtime.user_context import reset_current_user, set_current_user

thread_store = _thread_store
_stub_app_config = _base_config


@pytest.mark.anyio
@pytest.mark.parametrize("first_reader", ["run", "http"])
async def test_detached_run_metadata_and_http_create_share_admission(thread_store, monkeypatch, first_reader):
    thread_id = "cross-entry-create"
    first_missing = asyncio.Event()
    competing_attempt = asyncio.Event()
    release_first = asyncio.Event()
    original_get = thread_store.get
    original_create = thread_store.create
    real_lock = goal.goal_thread_lock
    creates = []
    first_task = None

    async def observed_get(identifier, **kwargs):
        row = await original_get(identifier, **kwargs)
        if identifier == thread_id and row is None:
            if asyncio.current_task() is first_task or (first_reader == "http" and first_task is None):
                first_missing.set()
                await release_first.wait()
            else:
                # Baseline reaches the second missing read; the fixed writer
                # instead reaches the lock observer before waiting for it.
                competing_attempt.set()
        return row

    @asynccontextmanager
    async def observed_lock(identifier):
        if identifier == thread_id and first_missing.is_set():
            competing_attempt.set()
        async with real_lock(identifier):
            yield

    saver = InMemorySaver()

    async def observed_create(identifier, **kwargs):
        row = await original_create(identifier, **kwargs)
        creates.append(row)
        if len(creates) == 1:
            await thread_store.update_status(identifier, "busy")
            await thread_store.update_display_name(identifier, "Existing title")
        return row

    monkeypatch.setattr(thread_store, "get", observed_get)
    monkeypatch.setattr(thread_store, "create", observed_create)
    monkeypatch.setattr(threads, "goal_thread_lock", observed_lock)
    monkeypatch.setattr(services, "goal_thread_lock", observed_lock)
    app = app_for(thread_store, saver)
    record = await RunManager().create(thread_id, metadata={"source": "run"}, kwargs={"input": {"title": "Run title"}})
    tasks = []

    async def ensure_from_run():
        token = set_current_user(USER)
        try:
            # Production run_after_metadata owns this writer as a detached
            # task too. No graph/model execution is simulated by this test.
            return await services._ensure_thread_metadata(SimpleNamespace(thread_store=thread_store), record, owner_user_id=None)
        finally:
            reset_current_user(token)

    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:

        async def post():
            return await client.post("/api/threads", json={"thread_id": thread_id, "metadata": {"source": "http"}})

        try:
            if first_reader == "run":
                first_task = asyncio.create_task(ensure_from_run())
                tasks.append(first_task)
                await asyncio.wait_for(first_missing.wait(), 5)
                tasks.append(asyncio.create_task(post()))
            else:
                # HTTP executes in the middleware's child task, so leave the
                # first-task marker absent until its first missing read parks.
                tasks.append(asyncio.create_task(post()))
                await asyncio.wait_for(first_missing.wait(), 5)
                first_task = tasks[0]
                tasks.append(asyncio.create_task(ensure_from_run()))
            await asyncio.wait_for(competing_attempt.wait(), 5)
            release_first.set()
            results = await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), 5)
            # The SQL loser may raise IntegrityError; run_after_metadata logs
            # that existing non-fatal failure. The winning row stays intact.
            assert all(isinstance(value, (dict, httpx.Response, IntegrityError)) for value in results)
            response = next(value for value in results if isinstance(value, httpx.Response))
            assert response.status_code == 200
            row = await original_get(thread_id, user_id=str(USER.id))
            assert len(creates) == 1
            assert response.json()["metadata"] == row["metadata"] == creates[0]["metadata"]
            assert row["status"] == "busy"
            assert row["display_name"] == "Existing title"
        finally:
            release_first.set()
            for task in tasks:
                if not task.done():
                    task.cancel()
            await asyncio.gather(*tasks, return_exceptions=True)
    assert not goal._goal_locks._entries_by_loop


@pytest.mark.anyio
async def test_start_run_detached_writer_cannot_replace_http_thread(thread_store, monkeypatch, _stub_app_config):
    """Use native admission/task attachment; only agent execution is stubbed."""
    from app.gateway.routers.thread_runs import RunCreateRequest

    missing = asyncio.Event()
    release = asyncio.Event()
    http_attempted = asyncio.Event()
    writer_locked = False
    writer_task = None
    original_get = thread_store.get
    original_create = thread_store.create
    original_ensure = services._ensure_thread_metadata
    real_lock = goal.goal_thread_lock
    creations = []

    async def tracked_ensure(*args, **kwargs):
        nonlocal writer_task
        writer_task = asyncio.current_task()
        return await original_ensure(*args, **kwargs)

    async def held_get(identifier, **kwargs):
        row = await original_get(identifier, **kwargs)
        if identifier == "native-admission" and row is None and asyncio.current_task() is writer_task:
            missing.set()
            await release.wait()
        return row

    async def counted_create(*args, **kwargs):
        row = await original_create(*args, **kwargs)
        creations.append(row)
        return row

    @asynccontextmanager
    async def service_lock(identifier):
        nonlocal writer_locked
        if asyncio.current_task() is writer_task:
            writer_locked = True
        async with real_lock(identifier):
            yield

    @asynccontextmanager
    async def http_lock(identifier):
        http_attempted.set()
        async with real_lock(identifier):
            yield

    async def no_agent(*args, **kwargs):
        pass

    monkeypatch.setattr(services, "_ensure_thread_metadata", tracked_ensure)
    monkeypatch.setattr(services, "goal_thread_lock", service_lock)
    monkeypatch.setattr(threads, "goal_thread_lock", http_lock)
    monkeypatch.setattr(thread_store, "get", held_get)
    monkeypatch.setattr(thread_store, "create", counted_create)
    monkeypatch.setattr(services, "resolve_agent_factory", lambda _name: object())
    monkeypatch.setattr(services, "run_agent", no_agent)
    app = app_for(thread_store, InMemorySaver())
    app.state.run_manager = RunManager(store=MemoryRunStore())
    app.state.stream_bridge = SimpleNamespace()
    app.state.store = InMemoryStore()
    app.state.run_event_store = MemoryRunEventStore()
    request = SimpleNamespace(headers={}, state=SimpleNamespace(user=USER), app=app)
    record = None
    http_task = None
    token = set_current_user(USER)
    try:
        record = await services.start_run(RunCreateRequest(input={"messages": [{"role": "user", "content": "hello"}]}, metadata={"source": "run"}), "native-admission", request)
        await asyncio.wait_for(missing.wait(), 5)
        assert writer_task is not record.task
        async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
            http_task = asyncio.create_task(client.post("/api/threads", json={"thread_id": "native-admission", "metadata": {"source": "http"}}))
            await asyncio.wait_for(http_attempted.wait(), 5)
            if not writer_locked:
                # Baseline permits this HTTP create to finish while the
                # detached writer still retains its earlier missing read.
                await asyncio.wait_for(asyncio.shield(http_task), 5)
            release.set()
            await asyncio.wait_for(asyncio.shield(record.task), 5)
            response = await asyncio.wait_for(http_task, 5)
            assert response.status_code == 200
            row = await original_get("native-admission", user_id=str(USER.id))
            assert len(creations) == 1
            assert response.json()["metadata"] == row["metadata"]
    finally:
        release.set()
        tasks = [task for task in [http_task, writer_task, record.task if record is not None else None] if task is not None]
        for task in tasks:
            if not task.done():
                task.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)
        reset_current_user(token)
    assert not goal._goal_locks._entries_by_loop


@pytest.mark.anyio
async def test_cancelled_metadata_waiter_releases_its_lock_entry(thread_store, monkeypatch):
    attempted = asyncio.Event()
    real_lock = goal.goal_thread_lock

    @asynccontextmanager
    async def observed_lock(identifier):
        attempted.set()
        async with real_lock(identifier):
            yield

    monkeypatch.setattr(services, "goal_thread_lock", observed_lock)
    record = await RunManager().create("cancelled-metadata")
    ctx = SimpleNamespace(thread_store=thread_store)
    task = None
    token = set_current_user(USER)
    try:
        async with real_lock(record.thread_id):
            task = asyncio.create_task(services._ensure_thread_metadata(ctx, record, owner_user_id=None))
            await asyncio.wait_for(attempted.wait(), 5)
            task.cancel()
            with pytest.raises(asyncio.CancelledError):
                await task
            assert await thread_store.get(record.thread_id) is None
        row = await asyncio.wait_for(services._ensure_thread_metadata(ctx, record, owner_user_id=None), 5)
        assert row["thread_id"] == record.thread_id
    finally:
        if task is not None:
            task.cancel()
            await asyncio.gather(task, return_exceptions=True)
        reset_current_user(token)
    assert not goal._goal_locks._entries_by_loop


@pytest.mark.anyio
@pytest.mark.parametrize("strict", [False, True], ids=["fallback", "strict"])
async def test_metadata_setup_budget_includes_create_admission(thread_store, monkeypatch, _stub_app_config, caplog, strict):
    from app.gateway.routers.thread_runs import RunCreateRequest
    from deerflow.runtime.runs.manager import RunStartOutcome
    from deerflow.runtime.runs.schemas import RunStatus

    thread_id = "metadata-admission-budget"
    waiting = asyncio.Event()
    metadata_created = asyncio.Event()
    admit_metadata = asyncio.Event()
    waiter_finished = asyncio.Event()
    original_ensure = services._ensure_thread_metadata
    real_lock = goal.goal_thread_lock
    metadata_task = None
    observed_budgets = []
    agent_calls = []

    async def tracked_ensure(*args, **kwargs):
        nonlocal metadata_task
        metadata_task = asyncio.current_task()
        metadata_created.set()
        await admit_metadata.wait()
        return await original_ensure(*args, **kwargs)

    @asynccontextmanager
    async def observed_lock(identifier):
        is_metadata_task = asyncio.current_task() is metadata_task
        if is_metadata_task:
            waiting.set()
        try:
            async with real_lock(identifier):
                yield
        finally:
            if is_metadata_task:
                waiter_finished.set()

    async def expire_metadata_budget(tasks, *, timeout, return_when):
        observed_budgets.append(timeout)
        # Expire only after the detached writer is waiting for admission.
        # No wall-clock sleep or five-second test delay is needed.
        await asyncio.wait_for(waiting.wait(), 5)
        done, pending = await asyncio.wait(tasks, timeout=0, return_when=return_when)
        assert not done
        return done, pending

    async def no_agent(_bridge, manager, record, **kwargs):
        agent_calls.append((await manager.try_start(record.run_id), kwargs))

    monkeypatch.setattr(services, "_ensure_thread_metadata", tracked_ensure)
    monkeypatch.setattr(services, "goal_thread_lock", observed_lock)
    monkeypatch.setattr(services, "asyncio", SimpleNamespace(**{**vars(asyncio), "wait": expire_metadata_budget}))
    monkeypatch.setattr(services, "resolve_agent_factory", lambda _name: object())
    monkeypatch.setattr(services, "run_agent", no_agent)
    app = app_for(thread_store, InMemorySaver())
    manager = RunManager(store=MemoryRunStore())
    app.state.run_manager = manager
    app.state.stream_bridge = SimpleNamespace()
    app.state.store = InMemoryStore()
    app.state.run_event_store = MemoryRunEventStore()
    request = SimpleNamespace(headers={}, state=SimpleNamespace(user=USER), app=app)
    record = None
    token = set_current_user(USER)
    try:
        row = await thread_store.create(thread_id, user_id=str(USER.id), metadata={"source": "http"})
        record = await services.start_run(RunCreateRequest(input={"messages": [{"role": "user", "content": "hello"}]}), thread_id, request, require_existing_thread=strict)
        await asyncio.wait_for(metadata_created.wait(), 5)
        # A same-thread operation wins admission after start_run releases its
        # caller-side lock and before the detached metadata writer enters it.
        async with real_lock(thread_id):
            admit_metadata.set()
            await asyncio.wait_for(asyncio.shield(record.task), 5)
            await asyncio.wait_for(waiter_finished.wait(), 5)
            assert metadata_task.cancelled()
            assert observed_budgets == [services._THREAD_METADATA_SETUP_TIMEOUT_SECONDS]
            assert len(agent_calls) == 1
            outcome, kwargs = agent_calls[0]
            assert "thread_incarnation" not in kwargs
            assert outcome is (RunStartOutcome.cancelled if strict else RunStartOutcome.started)
            assert record.status is (RunStatus.error if strict else RunStatus.running)
            assert (await manager.get(record.run_id)).status is record.status
            if strict:
                assert record.error == "Timed out verifying existing thread metadata"
            else:
                assert "MCP access will fail closed" in caplog.text
            assert "Timed out ensuring thread_meta" in caplog.text
            assert await thread_store.get(thread_id) == row
        assert not goal._goal_locks._entries_by_loop
        # The cancelled waiter neither writes late nor prevents a later setup.
        assert await services._ensure_thread_metadata(SimpleNamespace(thread_store=thread_store), record, owner_user_id=None, require_existing_thread=strict) == row
    finally:
        tasks = [task for task in [metadata_task, record.task if record is not None else None] if task is not None and task is not asyncio.current_task()]
        for task in tasks:
            if not task.done():
                task.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)
        reset_current_user(token)
    assert not goal._goal_locks._entries_by_loop