beaver_project/app-instance/backend/tests/unit/test_task_execution_planner.py

from __future__ import annotations

import asyncio
from types import SimpleNamespace

from beaver.engine.providers.base import LLMProvider, LLMResponse
from beaver.engine.providers.factory import ProviderBundle
from beaver.tasks import TaskExecutionPlanner, TaskRecord


class PlannerProvider(LLMProvider):
    def __init__(self, response: str) -> None:
        super().__init__()
        self.response = response

    async def chat(
        self,
        messages: list[dict],
        tools: list[dict] | None = None,
        model: str | None = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
    ) -> LLMResponse:
        return LLMResponse(content=self.response, finish_reason="stop", provider_name="stub", model="stub-model")

    def get_default_model(self) -> str:
        return "stub-model"


class HangingPlannerProvider(LLMProvider):
    async def chat(
        self,
        messages: list[dict],
        tools: list[dict] | None = None,
        model: str | None = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
    ) -> LLMResponse:
        await asyncio.sleep(10)
        return LLMResponse(content='{"mode":"team"}', finish_reason="stop", provider_name="stub", model="stub-model")

    def get_default_model(self) -> str:
        return "stub-model"


def _task() -> TaskRecord:
    return TaskRecord(
        task_id="task-1",
        session_id="session-1",
        description="implement workflow",
        goal="implement workflow",
        constraints=[],
        priority=0,
        status="open",
        creator="test",
        created_at="now",
        updated_at="now",
    )


def _bundle(response: str) -> ProviderBundle:
    return ProviderBundle(
        main_runtime=SimpleNamespace(model="stub-model", provider_name="stub"),
        main_provider=PlannerProvider(response),
    )


def _hanging_bundle() -> ProviderBundle:
    return ProviderBundle(
        main_runtime=SimpleNamespace(model="stub-model", provider_name="stub"),
        main_provider=HangingPlannerProvider(),
    )


def test_planner_selects_single_mode() -> None:
    plan = asyncio.run(
        TaskExecutionPlanner().plan(
            task=_task(),
            user_message="implement workflow",
            attempt_index=1,
            provider_bundle=_bundle('{"mode":"single","reason":"main agent is enough"}'),
        )
    )

    assert plan.mode == "single"
    assert plan.graph is None
    assert plan.reason == "main agent is enough"


def test_planner_builds_team_graph() -> None:
    plan = asyncio.run(
        TaskExecutionPlanner().plan(
            task=_task(),
            user_message="implement workflow",
            attempt_index=1,
            provider_bundle=_bundle(
                """
                {
                  "mode": "team",
                  "reason": "needs parallel review",
                  "strategy": "dag",
                  "nodes": [
                    {"node_id": "research", "task": "research options", "agent": {"name": "researcher"}},
                    {"node_id": "review", "task": "review result", "agent": {"name": "reviewer"}, "depends_on": ["research"]}
                  ],
                  "final_synthesis_instruction": "merge the findings"
                }
                """
            ),
        )
    )

    assert plan.is_team
    assert plan.graph is not None
    assert plan.graph.strategy == "dag"
    assert [node.node_id for node in plan.graph.nodes] == ["research", "review"]
    assert plan.graph.nodes[1].depends_on == ["research"]
    assert plan.final_synthesis_instruction == "merge the findings"


def test_planner_timeout_falls_back_to_single() -> None:
    plan = asyncio.run(
        TaskExecutionPlanner().plan(
            task=_task(),
            user_message="implement workflow",
            attempt_index=1,
            provider_bundle=_hanging_bundle(),
            timeout_seconds=0.01,
        )
    )

    assert plan.mode == "single"
    assert plan.reason == "planner_failed"
    assert "TimeoutError" in (plan.fallback_error or "")


def test_planner_team_nodes_can_target_skills_without_agent_roles() -> None:
    plan = TaskExecutionPlanner().from_json(
        """
        {
          "mode": "team",
          "reason": "needs skill-guided review",
          "strategy": "sequence",
          "nodes": [
            {
              "node_id": "api_review",
              "task": "review API compatibility",
              "skill_query": "API contract compatibility review",
              "required_capabilities": ["schema compatibility"]
            }
          ]
        }
        """
    )

    assert plan.is_team
    assert plan.graph is not None
    node = plan.graph.nodes[0]
    assert node.agent.name == "api_review"
    assert node.agent.role == ""
    assert node.agent.metadata["skill_query"] == "API contract compatibility review"
    assert node.agent.metadata["required_capabilities"] == ["schema compatibility"]


def test_planner_invalid_outputs_fallback_to_single() -> None:
    planner = TaskExecutionPlanner()
    invalid_json = planner.from_json("not json")
    unknown_strategy = planner.from_json(
        '{"mode":"team","strategy":"moa","nodes":[{"node_id":"a","task":"a","agent":{"name":"a"}}]}'
    )
    too_many_nodes = planner.from_json(
        '{"mode":"team","strategy":"parallel","nodes":['
        + ",".join(
            '{"node_id":"n%s","task":"work","agent":{"name":"n%s"}}' % (index, index)
            for index in range(7)
        )
        + "]}"
    )
    cyclic = planner.from_json(
        """
        {
          "mode": "team",
          "strategy": "dag",
          "nodes": [
            {"node_id": "a", "task": "a", "agent": {"name": "a"}, "depends_on": ["b"]},
            {"node_id": "b", "task": "b", "agent": {"name": "b"}, "depends_on": ["a"]}
          ]
        }
        """
    )

    assert invalid_json.mode == "single"
    assert unknown_strategy.mode == "single"
    assert too_many_nodes.mode == "single"
    assert cyclic.mode == "single"