deer-flow/backend/tests/test_subagent_runtime.py

import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from deerflow.config.subagent_batches_config import SubagentBatchesConfig
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
from deerflow.subagents import SubagentRuntime


def test_runtime_rejects_batch_repository_without_enabled_batch_config() -> None:
    with pytest.raises(ValueError, match="batch_config.enabled"):
        SubagentRuntime(
            SubagentRuntimeConfig(),
            batch_repository=MagicMock(),
            batch_config=SubagentBatchesConfig(enabled=False),
        )


def test_runtime_rejects_batch_repository_without_app_config_snapshot() -> None:
    with pytest.raises(ValueError, match="explicit app_config snapshot"):
        SubagentRuntime(
            SubagentRuntimeConfig(),
            batch_repository=MagicMock(),
            batch_config=SubagentBatchesConfig(enabled=True),
        )


def test_runtime_uses_one_caller_owned_app_config_snapshot() -> None:
    app_config = SimpleNamespace(
        subagent_runtime=SubagentRuntimeConfig(max_running=11),
        subagents=SimpleNamespace(max_total_per_run=14),
        subagent_batches=SubagentBatchesConfig(enabled=False),
    )

    runtime = SubagentRuntime.from_app_config(app_config)

    assert runtime.config.max_running == 11
    assert runtime.max_total_per_run == 14
    assert runtime.app_config is app_config
    assert runtime.batch_submitter is None


@pytest.mark.asyncio
async def test_runtime_owns_batch_worker_lifecycle_and_shared_capacity() -> None:
    service = MagicMock()
    service.start = AsyncMock()
    service.stop = AsyncMock()
    repository = MagicMock()
    app_config = MagicMock()

    with patch(
        "deerflow.subagents.batch_service.SubagentBatchService",
        return_value=service,
    ) as service_type:
        runtime = SubagentRuntime(
            SubagentRuntimeConfig(max_running=9),
            batch_repository=repository,
            batch_config=SubagentBatchesConfig(enabled=True),
            app_config=app_config,
        )
        assert runtime.batch_submitter is None

        async with runtime:
            assert runtime.batch_submitter is service

        assert runtime.batch_submitter is None

    service_type.assert_called_once_with(
        repository=repository,
        config=runtime.batch_config,
        runtime_config=runtime.config,
        app_config=app_config,
        execution_capacity=runtime.execution_capacity,
    )
    service.start.assert_awaited_once_with()
    service.stop.assert_awaited_once_with()


@pytest.mark.asyncio
async def test_runtime_stop_drains_owned_batch_worker_across_repeated_cancellation() -> None:
    stop_started = asyncio.Event()
    allow_stop = asyncio.Event()

    async def blocking_stop() -> None:
        stop_started.set()
        await allow_stop.wait()

    service = MagicMock()
    service.start = AsyncMock()
    service.stop = AsyncMock(side_effect=blocking_stop)
    repository = MagicMock()
    app_config = MagicMock()

    with patch(
        "deerflow.subagents.batch_service.SubagentBatchService",
        return_value=service,
    ):
        runtime = SubagentRuntime(
            SubagentRuntimeConfig(max_running=1),
            batch_repository=repository,
            batch_config=SubagentBatchesConfig(enabled=True),
            app_config=app_config,
        )
        await runtime.start()

        stop_task = asyncio.create_task(runtime.stop())
        await asyncio.wait_for(stop_started.wait(), timeout=1)

        stop_task.cancel()
        for _ in range(10):
            await asyncio.sleep(0)

        assert runtime.batch_submitter is None
        assert not stop_task.done(), "runtime stop released ownership after the first cancellation"

        stop_task.cancel()
        for _ in range(10):
            await asyncio.sleep(0)
        assert not stop_task.done(), "runtime stop released ownership after repeated cancellation"

        allow_stop.set()
        with pytest.raises(asyncio.CancelledError):
            await stop_task

        service.stop.assert_awaited_once_with()
        assert runtime.batch_submitter is None


@pytest.mark.asyncio
@pytest.mark.parametrize("cancel_caller", [False, True])
@pytest.mark.parametrize("stop_error_type", [RuntimeError, asyncio.CancelledError])
async def test_runtime_stop_preserves_caller_cancellation_when_service_fails(cancel_caller: bool, stop_error_type: type[BaseException]) -> None:
    stop_error = stop_error_type("service shutdown failed")
    stop_started = asyncio.Event()
    allow_stop = asyncio.Event()

    async def failing_stop() -> None:
        stop_started.set()
        await allow_stop.wait()
        raise stop_error

    service = MagicMock()
    service.start = AsyncMock()
    service.stop = AsyncMock(side_effect=failing_stop)
    with patch("deerflow.subagents.batch_service.SubagentBatchService", return_value=service):
        runtime = SubagentRuntime(
            batch_repository=MagicMock(),
            batch_config=SubagentBatchesConfig(enabled=True),
            app_config=MagicMock(),
        )
        await runtime.start()
        stop_task = asyncio.create_task(runtime.stop())
        try:
            await asyncio.wait_for(stop_started.wait(), timeout=1)
            if cancel_caller:
                stop_task.cancel("first caller cancellation")
                await asyncio.sleep(0)
                stop_task.cancel("second caller cancellation")
                await asyncio.sleep(0)
                assert not stop_task.done()
            allow_stop.set()

            expected_error = asyncio.CancelledError if cancel_caller else type(stop_error)
            with pytest.raises(expected_error) as raised:
                await stop_task
            if cancel_caller:
                assert raised.value.args == ("first caller cancellation",)
                assert raised.value.__cause__ is stop_error
            else:
                assert raised.value is stop_error
            service.stop.assert_awaited_once_with()
            assert runtime.batch_submitter is None
        finally:
            allow_stop.set()
            await asyncio.gather(stop_task, return_exceptions=True)