deer-flow/backend/tests/test_batch_lease_lifecycle.py

"""Real database regressions for durable batch lease generations and setup."""

import asyncio
import threading
from datetime import UTC, datetime, timedelta
from enum import Enum
from types import SimpleNamespace

import pytest
import pytest_asyncio
from sqlalchemy.ext.asyncio import AsyncSession

from deerflow.config.database_config import DatabaseConfig
from deerflow.config.subagent_batches_config import SubagentBatchesConfig
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config
from deerflow.persistence.subagent_batches import SubagentBatchRepository
from deerflow.persistence.subagent_batches.model import SubagentBatchItemRow, SubagentBatchRow
from deerflow.subagents import batch_service as service_module
from deerflow.subagents.batch_service import SubagentBatchService


class SubagentStatus(Enum):
    RUNNING = "running"
    COMPLETED = "completed"

    @property
    def is_terminal(self):
        return self is SubagentStatus.COMPLETED


@pytest_asyncio.fixture
async def repository(tmp_path):
    await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
    try:
        yield SubagentBatchRepository(get_session_factory())
    finally:
        await close_engine()


async def create_batch(repository, *, count=1, max_attempts=3):
    return await repository.create_batch(
        batch_id="batch-lease-test",
        user_id="user-1",
        thread_id="thread-1",
        run_id="run-1",
        tool_call_id="call-1",
        submission_key="run-1:call-1",
        title="Lease recovery",
        subagent_type="general-purpose",
        items=[{"key": f"record-{i}", "prompt": f"Process {i}"} for i in range(count)],
        max_live_items=count,
        max_running_items=1,
        max_attempts=max_attempts,
        execution_spec={"subagent_config": {"name": "general-purpose", "description": "test", "system_prompt": "Work carefully."}},
    )


async def row_snapshot(repository, item_id):
    async with repository._sf() as session:
        item = await session.get(SubagentBatchItemRow, item_id)
        return {c.name: getattr(item, c.name) for c in item.__table__.columns}


@pytest.mark.asyncio
@pytest.mark.parametrize("operation", ["renew", "running", "finalize", "requeue"])
@pytest.mark.parametrize("peer_action", ["reclaim", "cancel"])
async def test_in_flight_old_acknowledgement_cannot_change_reclaimed_row(repository, monkeypatch, operation, peer_action):
    await create_batch(repository)
    now = datetime.now(UTC)
    first = (await repository.claim_items(now=now, lease_owner="old-claim", lease_seconds=10, limit=1))[0]
    item_id = first["id"]
    paused, release = asyncio.Event(), asyncio.Event()
    original = AsyncSession.execute
    old_task = None

    async def execute(session, statement, *args, **kwargs):
        result = await original(session, statement, *args, **kwargs)
        if asyncio.current_task() is old_task and not paused.is_set() and getattr(statement, "is_select", False) and "subagent_batch_items" in str(statement) and len(statement.selected_columns) > 1:
            # Pause after the old operation's actual database read, not before
            # it starts. SQLite ignores SELECT FOR UPDATE; a peer can reclaim.
            paused.set()
            await release.wait()
        return result

    monkeypatch.setattr(AsyncSession, "execute", execute)

    async def acknowledge():
        stamp = now + timedelta(seconds=1)  # Within the OLD lease; expiry must not mask a missing fence.
        if operation == "renew":
            return await repository.renew_item_lease(item_id, lease_owner="old-claim", lease_seconds=10, now=stamp)
        if operation == "running":
            return await repository.mark_item_running(item_id, lease_owner="old-claim", now=stamp)
        if operation == "requeue":
            return await repository.requeue_item_after_admission_failure(item_id, lease_owner="old-claim", error="stale rejection", now=stamp)
        return await repository.finalize_item(
            item_id,
            lease_owner="old-claim",
            succeeded=True,
            result="stale result",
            result_preview="stale result",
            result_truncated=False,
            error=None,
            stop_reason=None,
            token_usage={"total_tokens": 99},
            model_name="old-model",
            completed_at=stamp,
        )

    old_task = asyncio.create_task(acknowledge())
    try:
        await asyncio.wait_for(paused.wait(), 3)
        if peer_action == "reclaim":
            replacement = await repository.claim_items(now=now + timedelta(seconds=11), lease_owner="new-claim", lease_seconds=10, limit=1)
            assert replacement[0]["id"] == item_id and replacement[0]["attempt"] == 2
        else:
            await repository.cancel_batch("batch-lease-test", user_id="user-1")
        expected = await row_snapshot(repository, item_id)
    finally:
        release.set()
        outcome = await asyncio.wait_for(old_task, 3)
    assert outcome == ({"valid": False, "cancel_requested": peer_action == "cancel"} if operation == "renew" else False)
    assert await row_snapshot(repository, item_id) == expected


@pytest.mark.asyncio
@pytest.mark.parametrize("phase", ["pending", "queued"])
@pytest.mark.parametrize("peer_action", ["claim", "cancel", "pause"])
async def test_stale_claim_pass_preserves_peer_claim_or_control(repository, monkeypatch, phase, peer_action):
    await create_batch(repository)
    now = datetime.now(UTC)
    item_id = (await repository.list_items("batch-lease-test", user_id="user-1"))[0]["id"]
    if phase == "queued":
        # Promote without claiming so the paused pass has not acquired SQLite's
        # writer lock before its runnable-item read.
        assert await repository.claim_items(now=now, lease_owner="seed", lease_seconds=10, limit=1, exclude_item_ids=(item_id,)) == []
    paused, release = asyncio.Event(), asyncio.Event()
    original = AsyncSession.execute
    stale = None

    async def execute(session, statement, *args, **kwargs):
        result = await original(session, statement, *args, **kwargs)
        if asyncio.current_task() is stale and not paused.is_set() and getattr(statement, "is_select", False) and "subagent_batch_items.status =" in str(statement) and len(statement.selected_columns) > 1:
            if phase in statement.compile().params.values():
                paused.set()
                await release.wait()
        return result

    monkeypatch.setattr(AsyncSession, "execute", execute)
    stale = asyncio.create_task(repository.claim_items(now=now, lease_owner="stale-claim", lease_seconds=10, limit=1))
    try:
        await asyncio.wait_for(paused.wait(), 3)
        if peer_action == "claim":
            peer = await repository.claim_items(now=now + timedelta(seconds=1), lease_owner="peer-claim", lease_seconds=10, limit=1)
            assert len(peer) == 1 and peer[0]["attempt"] == 1
        elif peer_action == "cancel":
            await repository.cancel_batch("batch-lease-test", user_id="user-1")
        else:
            await repository.pause_batch("batch-lease-test", user_id="user-1")
        expected = await row_snapshot(repository, item_id)
        expected_batch = await repository.get_batch("batch-lease-test", user_id="user-1")
    finally:
        release.set()
        claimed = await asyncio.wait_for(stale, 3)
    assert claimed == [], "a stale read must not grant ownership over the peer's claim or control"
    assert await row_snapshot(repository, item_id) == expected
    assert await repository.get_batch("batch-lease-test", user_id="user-1") == expected_batch


@pytest.mark.asyncio
@pytest.mark.parametrize("case", ["missing", "wrong_owner", "item_cancelled", "batch_cancelled", "owned"])
async def test_renewal_reports_only_explicit_cancellation(repository, case):
    await create_batch(repository)
    now = datetime.now(UTC)
    item = (await repository.claim_items(now=now, lease_owner="claim", lease_seconds=10, limit=1))[0]
    if case in ("item_cancelled", "batch_cancelled"):
        async with repository._sf() as session:
            if case == "item_cancelled":
                (await session.get(SubagentBatchItemRow, item["id"])).cancel_requested_at = now
            else:
                (await session.get(SubagentBatchRow, "batch-lease-test")).status = "cancelled"
            await session.commit()
    outcome = await repository.renew_item_lease("missing" if case == "missing" else item["id"], lease_owner="peer" if case == "wrong_owner" else "claim", lease_seconds=10, now=now)
    assert outcome == {"valid": case == "owned", "cancel_requested": case in ("item_cancelled", "batch_cancelled")}


@pytest.mark.asyncio
async def test_execution_requires_an_explicit_claim_owner(repository):
    service = make_service(repository)
    with pytest.raises(TypeError, match="lease_owner"):
        await service._execute_item({"id": "missing"})


def setup_executor(monkeypatch, started, release):
    launched = []
    monkeypatch.setattr(service_module, "SubagentStatus", SubagentStatus)

    def assemble(**_kwargs):
        started.set()
        assert release.wait(10), "test must release the real assembly worker"
        return []

    class Executor:
        def __init__(self, **_kwargs):
            pass

        def execute_async(self, _prompt, task_id=None):
            launched.append(task_id)
            return "execution-1"

    result = SimpleNamespace(status=SubagentStatus.COMPLETED, result="correct source result", error=None, stop_reason=None, token_usage_records=None, ai_messages=[])
    monkeypatch.setattr("deerflow.tools.get_available_tools", assemble)
    monkeypatch.setattr(service_module, "resolve_subagent_model_name", lambda *_a, **_kw: "test-model")
    monkeypatch.setattr(service_module, "SubagentExecutor", Executor)
    monkeypatch.setattr(service_module, "get_background_task_result", lambda _id: result)
    monkeypatch.setattr(service_module, "cleanup_background_task", lambda _id: None)
    return launched


def make_service(repository):
    return SubagentBatchService(repository=repository, config=SubagentBatchesConfig(lease_seconds=10), runtime_config=SubagentRuntimeConfig(max_running=3), app_config=SimpleNamespace())


@pytest.mark.asyncio
async def test_scheduler_does_not_reclaim_its_still_initializing_item(repository, monkeypatch):
    await create_batch(repository, max_attempts=1)
    started, release = threading.Event(), threading.Event()
    launched = setup_executor(monkeypatch, started, release)
    service = make_service(repository)
    now = datetime.now(UTC)
    await service.run_once(now=now)
    tasks = list(service._executions.values())
    try:
        assert await asyncio.to_thread(started.wait, 3)
        # Advance the scheduler's logical claim clock past its lease, while a
        # real assembly-pool worker still owns initialization. No private lock
        # is held or manufactured by the test.
        await service.run_once(now=now + timedelta(seconds=11))
        item = (await repository.list_items("batch-lease-test", user_id="user-1"))[0]
        assert item["status"] == "leased", "live initialization must not exhaust its retry budget"
        assert item["attempt"] == 1
    finally:
        release.set()
        await asyncio.wait_for(asyncio.gather(*tasks), 3)
    assert len(launched) == 1
    item = (await repository.list_items("batch-lease-test", user_id="user-1", include_result=True))[0]
    assert item["status"] == "succeeded" and item["result"] == "correct source result"


@pytest.mark.asyncio
async def test_initialization_keeps_renewing_for_other_workers(repository, monkeypatch):
    await create_batch(repository)
    started, release = threading.Event(), threading.Event()
    setup_executor(monkeypatch, started, release)
    service = make_service(repository)
    renewals, twice = [], asyncio.Event()
    original = repository.renew_item_lease

    async def renew(*args, **kwargs):
        result = await original(*args, **kwargs)
        renewals.append(kwargs["now"])
        if len(renewals) >= 2:
            twice.set()
        return result

    monkeypatch.setattr(repository, "renew_item_lease", renew)
    now = datetime.now(UTC)
    await service.run_once(now=now)
    tasks = list(service._executions.values())
    try:
        assert await asyncio.to_thread(started.wait, 3)
        try:
            await asyncio.wait_for(twice.wait(), 5)
        except TimeoutError:
            pytest.fail("a blocked tool assembly received no periodic lease heartbeat")
        # Just beyond the first renewal's expiry, but before the second one's.
        peer_time = renewals[0] + timedelta(seconds=10, milliseconds=1)
        assert await repository.claim_items(now=peer_time, lease_owner="peer", lease_seconds=10, limit=1) == []
    finally:
        release.set()
        await asyncio.wait_for(asyncio.gather(*tasks), 3)


@pytest.mark.asyncio
async def test_each_scheduler_claim_has_a_fresh_ownership_generation(repository, monkeypatch):
    await create_batch(repository, count=2)
    started, release = threading.Event(), threading.Event()
    release.set()
    setup_executor(monkeypatch, started, release)
    service = make_service(repository)
    owners = []
    original = repository.claim_items

    async def claim(**kwargs):
        owners.append(kwargs["lease_owner"])
        return await original(**kwargs)

    monkeypatch.setattr(repository, "claim_items", claim)
    for _ in range(2):
        await service.run_once(now=datetime.now(UTC))
        await asyncio.gather(*list(service._executions.values()))
        await asyncio.sleep(0)  # Deliver the normal task-registry done callbacks.
    assert len(owners) == len(set(owners)) == 2, "worker identity must not be reused as a claim-generation fence"
    assert all(len(owner) <= 128 for owner in owners)
    items = await repository.list_items("batch-lease-test", user_id="user-1", include_result=True)
    assert [item["status"] for item in items] == ["succeeded", "succeeded"]
    assert all("lease_owner" not in item for item in items)


@pytest.mark.asyncio
async def test_recovery_does_not_overwrite_a_renewal_after_its_expiry_read(repository, monkeypatch):
    await create_batch(repository)
    now = datetime.now(UTC)
    item = (await repository.claim_items(now=now, lease_owner="live-worker", lease_seconds=10, limit=1))[0]
    paused, release = asyncio.Event(), asyncio.Event()
    original = AsyncSession.execute
    recovery = None

    async def execute(session, statement, *args, **kwargs):
        result = await original(session, statement, *args, **kwargs)
        if asyncio.current_task() is recovery and not paused.is_set() and getattr(statement, "is_select", False) and "lease_expires_at <" in str(statement):
            paused.set()
            await release.wait()
        return result

    monkeypatch.setattr(AsyncSession, "execute", execute)
    recovery = asyncio.create_task(repository.claim_items(now=now + timedelta(seconds=11), lease_owner="recoverer", lease_seconds=10, limit=1))
    try:
        await asyncio.wait_for(paused.wait(), 3)
        assert (await repository.renew_item_lease(item["id"], lease_owner="live-worker", lease_seconds=10, now=now + timedelta(seconds=12)))["valid"]
        expected = await row_snapshot(repository, item["id"])
    finally:
        release.set()
        claimed = await asyncio.wait_for(recovery, 3)
    assert claimed == []
    assert await row_snapshot(repository, item["id"]) == expected


@pytest.mark.asyncio
async def test_lost_heartbeat_cancels_execution_before_a_long_poll_interval(repository, monkeypatch):
    await create_batch(repository)
    started, release = threading.Event(), threading.Event()
    release.set()
    setup_executor(monkeypatch, started, release)
    service = make_service(repository)
    service._config = SubagentBatchesConfig(lease_seconds=10, poll_interval_seconds=60)
    running, cancelled = asyncio.Event(), asyncio.Event()
    mark_running = repository.mark_item_running

    async def mark(*args, **kwargs):
        valid = await mark_running(*args, **kwargs)
        if valid:
            running.set()
        return valid

    monkeypatch.setattr(repository, "mark_item_running", mark)
    result = SimpleNamespace(status=SubagentStatus.RUNNING, result=None, error=None, stop_reason=None, token_usage_records=None, ai_messages=[])

    def read_result(_id):
        return result

    def cancel(_id):
        result.status = SubagentStatus.COMPLETED
        cancelled.set()

    monkeypatch.setattr(service_module, "get_background_task_result", read_result)
    monkeypatch.setattr(service_module, "request_cancel_background_task", cancel)
    await service.run_once(now=datetime.now(UTC))
    task = next(iter(service._executions.values()))
    try:
        await asyncio.wait_for(running.wait(), 3)
        # Another worker handles the control request; no local cancellation shortcut.
        await repository.cancel_batch("batch-lease-test", user_id="user-1")
        await asyncio.wait_for(cancelled.wait(), 8)
        await asyncio.wait_for(task, 4)
    finally:
        if not task.done():
            task.cancel()
        await asyncio.gather(task, return_exceptions=True)
    item = (await repository.list_items("batch-lease-test", user_id="user-1", include_result=True))[0]
    assert item["status"] == "cancelled" and item["result"] is None


@pytest.mark.asyncio
async def test_heartbeat_exit_drains_an_in_flight_renewal_across_repeated_cancellation():
    from deerflow.subagents.batch_lease import BatchItemLease

    renewing, finish_renewal, exit_context = asyncio.Event(), asyncio.Event(), asyncio.Event()
    calls = 0

    class Repository:
        async def renew_item_lease(self, *_args, **_kwargs):
            nonlocal calls
            calls += 1
            if calls > 1:
                renewing.set()
                await finish_renewal.wait()
            return {"valid": True, "cancel_requested": False}

    async def supervise():
        async with BatchItemLease(Repository(), item_id="item", owner="generation", seconds=3):
            await exit_context.wait()

    task = asyncio.create_task(supervise())
    try:
        await asyncio.wait_for(renewing.wait(), 3)
        task.cancel()
        await asyncio.sleep(0)
        task.cancel()
        await asyncio.sleep(0)
        assert not task.done(), "caller must retain its in-flight database renewal"
        finish_renewal.set()
        with pytest.raises(asyncio.CancelledError):
            await asyncio.wait_for(task, 3)
        assert calls == 2
        assert not any(t.get_name() == "batch-item-lease-item" for t in asyncio.all_tasks())
    finally:
        finish_renewal.set()
        exit_context.set()
        if not task.done():
            task.cancel()
        await asyncio.gather(task, return_exceptions=True)