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

import asyncio
import json
from contextlib import suppress
from types import SimpleNamespace
from typing import Any

from beaver.engine import AgentLoop, AgentRunResult, EngineLoader
from beaver.engine import loop as loop_module
from beaver.engine.providers.base import LLMProvider, LLMResponse
from beaver.engine.providers.factory import ProviderBundle


def _run_result(run_id: str, output_text: str) -> AgentRunResult:
    return AgentRunResult(
        session_id="web:test",
        run_id=run_id,
        output_text=output_text,
        finish_reason="stop",
        tool_iterations=0,
    )


def test_running_loop_handles_reentrant_submit_direct(tmp_path) -> None:
    async def run_case() -> None:
        loop = AgentLoop(loader=EngineLoader(workspace=tmp_path))
        calls: list[str] = []

        async def fake_process_direct(task: str, **kwargs: Any) -> AgentRunResult:
            calls.append(task)
            if task == "outer":
                return await loop.submit_direct("inner", session_id="web:test")
            return _run_result(task, "inner completed")

        loop._process_direct_impl = fake_process_direct  # type: ignore[method-assign]

        loop_task = asyncio.create_task(loop.run())
        await asyncio.sleep(0)
        try:
            result = await asyncio.wait_for(loop.submit_direct("outer", session_id="web:test"), timeout=1)
        finally:
            await loop.stop()
            with suppress(asyncio.TimeoutError):
                await asyncio.wait_for(loop_task, timeout=1)
            if not loop_task.done():
                loop_task.cancel()
                with suppress(asyncio.CancelledError):
                    await loop_task

        assert result.output_text == "inner completed"
        assert calls == ["outer", "inner"]

    asyncio.run(run_case())


def test_web_search_loop_guard_keeps_successful_low_quality_results_available() -> None:
    guard = loop_module._WebSearchLoopGuard()
    low_quality = json.dumps(
        {
            "success": True,
            "query": "weather beijing",
            "quality": "low",
            "results": [{"title": "Example", "url": "https://example.com", "snippet": ""}],
        }
    )

    assert guard.observe_result("web_search", low_quality) is None
    assert guard.observe_result("web_search", low_quality) is None

    assert guard.observe_result("web_search", low_quality) is None


def test_web_search_loop_guard_resets_after_useful_result() -> None:
    guard = loop_module._WebSearchLoopGuard()
    failed_search = json.dumps({"success": False, "query": "weather", "error": "timeout"})
    useful = json.dumps({"success": True, "query": "weather", "quality": "high", "results": []})

    assert guard.observe_result("web_search", failed_search) is None
    assert guard.observe_result("web_search", useful) is None
    assert guard.observe_result("web_search", failed_search) is None
    assert guard.observe_result("web_search", failed_search) is None

    assert guard.observe_result("web_search", failed_search) is not None


class RecordingProvider(LLMProvider):
    def __init__(self) -> None:
        super().__init__()
        self.tool_names_by_call: list[list[str]] = []

    async def chat(
        self,
        messages: list[dict],
        tools: list[dict] | None = None,
        model: str | None = None,
        max_tokens: int | None = None,
        temperature: float = 0.7,
        thinking_enabled: bool | None = None,
    ) -> LLMResponse:
        self.tool_names_by_call.append(
            [
                str(tool.get("function", {}).get("name") or tool.get("name"))
                for tool in tools or []
            ]
        )
        return LLMResponse("done", provider_name="stub", model="stub-model")

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


def _bundle(provider: RecordingProvider) -> ProviderBundle:
    return ProviderBundle(
        main_runtime=SimpleNamespace(model="stub-model", provider_name="stub"),
        main_provider=provider,
    )


def test_task_mode_root_does_not_expose_agent_team_tool(tmp_path) -> None:
    provider = RecordingProvider()
    loop = AgentLoop(loader=EngineLoader(workspace=tmp_path))

    asyncio.run(
        loop.process_direct(
            "compare financial reports",
            session_id="session",
            task_id="task-1",
            task_mode=True,
            include_skill_assembly=False,
            provider_bundle=_bundle(provider),
        )
    )

    assert "run_agent_team" not in provider.tool_names_by_call[0]


def test_default_engine_registry_does_not_register_agent_team_tool(tmp_path) -> None:
    loaded = AgentLoop(loader=EngineLoader(workspace=tmp_path)).boot()

    assert loaded.tool_registry is not None
    assert loaded.tool_registry.get("run_agent_team") is None
    assert "run_agent_team" not in loaded.tools


def test_non_task_and_team_node_do_not_expose_agent_team_tool(tmp_path) -> None:
    non_task_provider = RecordingProvider()
    loop = AgentLoop(loader=EngineLoader(workspace=tmp_path))
    asyncio.run(
        loop.process_direct(
            "ordinary chat",
            session_id="session",
            include_skill_assembly=False,
            provider_bundle=_bundle(non_task_provider),
        )
    )

    team_node_provider = RecordingProvider()
    asyncio.run(
        loop.process_direct(
            "team child",
            session_id="session:team:child",
            parent_session_id="session",
            source="team:child",
            task_id="task-1",
            task_mode=True,
            include_skill_assembly=False,
            provider_bundle=_bundle(team_node_provider),
        )
    )

    assert "run_agent_team" not in non_task_provider.tool_names_by_call[0]
    assert "run_agent_team" not in team_node_provider.tool_names_by_call[0]