|
|
@@ -3,7 +3,7 @@ import pytest
|
|
|
from agent.core.presets import AgentPreset, register_preset
|
|
|
from agent.core.runner import AgentRunner, RunConfig
|
|
|
from agent.orchestration.models import AgentRole, CompletionPolicy
|
|
|
-from agent.tools.models import ToolResult
|
|
|
+from agent.tools.models import ToolCapability, ToolResult
|
|
|
from agent.tools.registry import ToolRegistry
|
|
|
from agent.trace.models import Trace
|
|
|
|
|
|
@@ -104,6 +104,138 @@ async def test_framework_context_overrides_host_spoofing():
|
|
|
assert context["coordinator"] == "real"
|
|
|
|
|
|
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_validator_capabilities_hide_and_reject_arbitrary_write_tool():
|
|
|
+ registry = ToolRegistry()
|
|
|
+ executed = False
|
|
|
+
|
|
|
+ async def db_modify():
|
|
|
+ nonlocal executed
|
|
|
+ executed = True
|
|
|
+ return "modified"
|
|
|
+
|
|
|
+ async def db_lookup():
|
|
|
+ return "row"
|
|
|
+
|
|
|
+ registry.register(
|
|
|
+ db_modify,
|
|
|
+ schema=_schema("db_modify"),
|
|
|
+ capabilities=[ToolCapability.WRITE],
|
|
|
+ )
|
|
|
+ registry.register(
|
|
|
+ db_lookup,
|
|
|
+ schema=_schema("db_lookup"),
|
|
|
+ capabilities=[ToolCapability.READ],
|
|
|
+ )
|
|
|
+ register_preset(
|
|
|
+ "test_validator_capabilities",
|
|
|
+ AgentPreset(
|
|
|
+ role=AgentRole.VALIDATOR,
|
|
|
+ allowed_tools=["db_modify", "db_lookup"],
|
|
|
+ max_iterations=5,
|
|
|
+ skills=[],
|
|
|
+ ),
|
|
|
+ )
|
|
|
+ runner = AgentRunner(tool_registry=registry, llm_call=lambda **_: None, task_coordinator=object())
|
|
|
+ config = RunConfig(
|
|
|
+ agent_type="test_validator_capabilities",
|
|
|
+ completion_policy=CompletionPolicy.EXPLICIT_VALIDATION,
|
|
|
+ tools=["db_modify", "db_lookup"],
|
|
|
+ )
|
|
|
+ names = {
|
|
|
+ item["function"]["name"] for item in runner._get_run_tool_schemas(config)
|
|
|
+ }
|
|
|
+ assert names == {"db_lookup"}
|
|
|
+
|
|
|
+ trace = Trace(trace_id="validator", mode="agent", agent_role="validator", context={})
|
|
|
+ result = await runner._execute_authorized_tool(
|
|
|
+ "db_modify", {}, "forged", config, trace, None, 1
|
|
|
+ )
|
|
|
+ assert result["error"] == "unauthorized_tool"
|
|
|
+ assert executed is False
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_worker_rejects_agent_spawn_capability_with_unrelated_name():
|
|
|
+ registry = ToolRegistry()
|
|
|
+ executed = False
|
|
|
+
|
|
|
+ async def spawn_helper():
|
|
|
+ nonlocal executed
|
|
|
+ executed = True
|
|
|
+ return "spawned"
|
|
|
+
|
|
|
+ registry.register(
|
|
|
+ spawn_helper,
|
|
|
+ schema=_schema("spawn_helper"),
|
|
|
+ capabilities=[ToolCapability.AGENT_SPAWN],
|
|
|
+ )
|
|
|
+ register_preset(
|
|
|
+ "test_worker_spawn_capability",
|
|
|
+ AgentPreset(
|
|
|
+ role=AgentRole.WORKER,
|
|
|
+ allowed_tools=["spawn_helper"],
|
|
|
+ max_iterations=5,
|
|
|
+ skills=[],
|
|
|
+ ),
|
|
|
+ )
|
|
|
+ runner = AgentRunner(tool_registry=registry, llm_call=lambda **_: None, task_coordinator=object())
|
|
|
+ config = RunConfig(
|
|
|
+ agent_type="test_worker_spawn_capability",
|
|
|
+ completion_policy=CompletionPolicy.EXPLICIT_VALIDATION,
|
|
|
+ tools=["spawn_helper"],
|
|
|
+ )
|
|
|
+ assert runner._get_run_tool_schemas(config) == []
|
|
|
+ trace = Trace(trace_id="worker", mode="agent", agent_role="worker", context={})
|
|
|
+ result = await runner._execute_authorized_tool(
|
|
|
+ "spawn_helper", {}, "forged", config, trace, None, 1
|
|
|
+ )
|
|
|
+ assert result["error"] == "unauthorized_tool"
|
|
|
+ assert executed is False
|
|
|
+
|
|
|
+
|
|
|
+@pytest.mark.asyncio
|
|
|
+async def test_unclassified_tool_fails_closed_only_in_explicit_mode():
|
|
|
+ registry = ToolRegistry()
|
|
|
+ calls = 0
|
|
|
+
|
|
|
+ async def mystery_tool():
|
|
|
+ nonlocal calls
|
|
|
+ calls += 1
|
|
|
+ return "ok"
|
|
|
+
|
|
|
+ registry.register(mystery_tool, schema=_schema("mystery_tool"))
|
|
|
+ register_preset(
|
|
|
+ "test_unclassified_worker",
|
|
|
+ AgentPreset(
|
|
|
+ role=AgentRole.WORKER,
|
|
|
+ allowed_tools=["mystery_tool"],
|
|
|
+ max_iterations=5,
|
|
|
+ skills=[],
|
|
|
+ ),
|
|
|
+ )
|
|
|
+ runner = AgentRunner(tool_registry=registry, llm_call=lambda **_: None, task_coordinator=object())
|
|
|
+ explicit = RunConfig(
|
|
|
+ agent_type="test_unclassified_worker",
|
|
|
+ completion_policy=CompletionPolicy.EXPLICIT_VALIDATION,
|
|
|
+ tools=["mystery_tool"],
|
|
|
+ )
|
|
|
+ worker_trace = Trace(trace_id="worker", mode="agent", agent_role="worker", context={})
|
|
|
+ rejected = await runner._execute_authorized_tool(
|
|
|
+ "mystery_tool", {}, "explicit", explicit, worker_trace, None, 1
|
|
|
+ )
|
|
|
+ assert rejected["error"] == "unauthorized_tool"
|
|
|
+ assert calls == 0
|
|
|
+
|
|
|
+ legacy = RunConfig(tools=["mystery_tool"], tool_groups=[])
|
|
|
+ legacy_trace = Trace(trace_id="legacy", mode="agent", agent_role="legacy", context={})
|
|
|
+ assert {x["function"]["name"] for x in runner._get_run_tool_schemas(legacy)} == {"mystery_tool"}
|
|
|
+ assert await runner._execute_authorized_tool(
|
|
|
+ "mystery_tool", {}, "legacy", legacy, legacy_trace, None, 1
|
|
|
+ ) == "ok"
|
|
|
+ assert calls == 1
|
|
|
+
|
|
|
+
|
|
|
def _disabled_knowledge():
|
|
|
from agent.tools.builtin.knowledge import KnowledgeConfig
|
|
|
return KnowledgeConfig(enable_extraction=False, enable_completion_extraction=False, enable_injection=False)
|