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]