135 lines
4.4 KiB
Python
135 lines
4.4 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
from beaver.coordinator.models import AgentDescriptor, ExecutionGraph, ExecutionNode
|
|
from beaver.coordinator.registry import AgentRegistry, RegisteredAgent, TargetResolver
|
|
from beaver.tasks import TaskRecord
|
|
|
|
|
|
def _task() -> TaskRecord:
|
|
return TaskRecord(
|
|
task_id="task-1",
|
|
session_id="session-1",
|
|
description="implement tests",
|
|
goal="implement tests",
|
|
constraints=[],
|
|
priority=0,
|
|
status="open",
|
|
creator="test",
|
|
created_at="now",
|
|
updated_at="now",
|
|
)
|
|
|
|
|
|
def test_registry_starts_empty_and_filters_disabled(tmp_path) -> None:
|
|
registry = AgentRegistry(tmp_path)
|
|
|
|
assert registry.list_agents() == []
|
|
|
|
registry.upsert_agent(
|
|
RegisteredAgent(
|
|
agent_id="tester",
|
|
name="tester",
|
|
display_name="Tester",
|
|
role="testing",
|
|
description="Runs checks.",
|
|
system_prompt="test",
|
|
)
|
|
)
|
|
registry.disable_agent("tester")
|
|
|
|
assert "tester" not in {agent.agent_id for agent in registry.list_active_agents()}
|
|
|
|
|
|
def test_registry_drops_legacy_builtin_agents(tmp_path) -> None:
|
|
registry_path = tmp_path / "agents" / "registry.json"
|
|
registry_path.parent.mkdir(parents=True)
|
|
registry_path.write_text(
|
|
json.dumps(
|
|
{
|
|
"version": 1,
|
|
"agents": [
|
|
{
|
|
"agent_id": "researcher",
|
|
"name": "researcher",
|
|
"display_name": "Researcher",
|
|
"role": "research",
|
|
"description": "legacy builtin",
|
|
"system_prompt": "research",
|
|
"source": "builtin",
|
|
},
|
|
{
|
|
"agent_id": "workspace-agent",
|
|
"name": "workspace-agent",
|
|
"display_name": "Workspace Agent",
|
|
"role": "workspace",
|
|
"description": "user configured",
|
|
"system_prompt": "work",
|
|
"source": "workspace",
|
|
},
|
|
],
|
|
}
|
|
)
|
|
+ "\n",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
registry = AgentRegistry(tmp_path)
|
|
|
|
assert [agent.agent_id for agent in registry.list_agents()] == ["workspace-agent"]
|
|
|
|
|
|
def test_resolver_selects_registered_agent_by_role_and_capabilities(tmp_path) -> None:
|
|
registry = AgentRegistry(tmp_path)
|
|
registry.upsert_agent(
|
|
RegisteredAgent(
|
|
agent_id="security-reviewer",
|
|
name="security-reviewer",
|
|
display_name="Security Reviewer",
|
|
role="security review",
|
|
description="Reviews auth, permissions, and data exposure risk.",
|
|
system_prompt="review security",
|
|
capabilities=["security", "review", "auth"],
|
|
priority=90,
|
|
)
|
|
)
|
|
resolver = TargetResolver(registry)
|
|
graph = ExecutionGraph(
|
|
strategy="sequence",
|
|
nodes=[
|
|
ExecutionNode(
|
|
node_id="review",
|
|
task="review auth handling",
|
|
agent=AgentDescriptor(
|
|
name="reviewer",
|
|
role="security review",
|
|
metadata={"requested_capabilities": ["security"]},
|
|
),
|
|
)
|
|
],
|
|
)
|
|
|
|
resolved, reports = resolver.resolve_graph(graph, task=_task(), user_message="review auth", attempt_index=1)
|
|
|
|
assert resolved.nodes[0].agent.metadata["agent_id"] == "security-reviewer"
|
|
assert reports[0].fallback_used is False
|
|
assert reports[0].selected_agent_id == "security-reviewer"
|
|
|
|
|
|
def test_resolver_falls_back_to_ephemeral_agent_when_no_match(tmp_path) -> None:
|
|
registry = AgentRegistry(tmp_path)
|
|
for agent in registry.list_agents():
|
|
registry.disable_agent(agent.agent_id)
|
|
resolver = TargetResolver(registry)
|
|
graph = ExecutionGraph(
|
|
strategy="sequence",
|
|
nodes=[ExecutionNode("rare", "rare work", AgentDescriptor(name="rare", role="rare"))],
|
|
)
|
|
|
|
resolved, reports = resolver.resolve_graph(graph, task=_task(), user_message="rare work", attempt_index=1)
|
|
|
|
assert resolved.nodes[0].agent.name == "rare"
|
|
assert resolved.nodes[0].agent.metadata["resolution"] == "fallback_ephemeral"
|
|
assert reports[0].fallback_used is True
|