|
@@ -0,0 +1,109 @@
|
|
|
|
|
+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.registry import ToolRegistry
|
|
|
|
|
+from agent.trace.models import Trace
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def _schema(name):
|
|
|
|
|
+ return {
|
|
|
|
|
+ "type": "function",
|
|
|
|
|
+ "function": {"name": name, "description": name, "parameters": {"type": "object", "properties": {}}},
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+@pytest.mark.asyncio
|
|
|
|
|
+async def test_terminal_tool_completes_without_final_assistant_text(tmp_path):
|
|
|
|
|
+ registry = ToolRegistry()
|
|
|
|
|
+
|
|
|
|
|
+ async def finish():
|
|
|
|
|
+ return ToolResult(
|
|
|
|
|
+ title="done", output="submitted", terminate_run=True, result_summary="structured-result"
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+ registry.register(finish, schema=_schema("finish"), groups=["test"])
|
|
|
|
|
+ calls = 0
|
|
|
|
|
+
|
|
|
|
|
+ async def llm_call(**kwargs):
|
|
|
|
|
+ nonlocal calls
|
|
|
|
|
+ calls += 1
|
|
|
|
|
+ return {
|
|
|
|
|
+ "content": "",
|
|
|
|
|
+ "tool_calls": [{"id": "call-1", "type": "function", "function": {"name": "finish", "arguments": "{}"}}],
|
|
|
|
|
+ "finish_reason": "tool_calls",
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
|
|
+ from agent.trace.store import FileSystemTraceStore
|
|
|
|
|
+ runner = AgentRunner(
|
|
|
|
|
+ trace_store=FileSystemTraceStore(str(tmp_path)),
|
|
|
|
|
+ tool_registry=registry,
|
|
|
|
|
+ llm_call=llm_call,
|
|
|
|
|
+ )
|
|
|
|
|
+ result = await runner.run_result(
|
|
|
|
|
+ [{"role": "user", "content": "finish"}],
|
|
|
|
|
+ RunConfig(tools=["finish"], tool_groups=[], knowledge=_disabled_knowledge()),
|
|
|
|
|
+ )
|
|
|
|
|
+ assert calls == 1
|
|
|
|
|
+ assert result["status"] == "completed"
|
|
|
|
|
+ assert result["summary"] == "structured-result"
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+@pytest.mark.asyncio
|
|
|
|
|
+async def test_execution_layer_rejects_forged_worker_tool_call():
|
|
|
|
|
+ registry = ToolRegistry()
|
|
|
|
|
+ executed = False
|
|
|
|
|
+
|
|
|
|
|
+ async def agent():
|
|
|
|
|
+ nonlocal executed
|
|
|
|
|
+ executed = True
|
|
|
|
|
+ return "bad"
|
|
|
|
|
+
|
|
|
|
|
+ registry.register(agent, schema=_schema("agent"))
|
|
|
|
|
+ register_preset(
|
|
|
|
|
+ "test_worker_policy",
|
|
|
|
|
+ AgentPreset(role=AgentRole.WORKER, allowed_tools=["agent"], max_iterations=5, skills=[]),
|
|
|
|
|
+ )
|
|
|
|
|
+ runner = AgentRunner(tool_registry=registry, llm_call=lambda **_: None, task_coordinator=object())
|
|
|
|
|
+ config = RunConfig(
|
|
|
|
|
+ agent_type="test_worker_policy",
|
|
|
|
|
+ completion_policy=CompletionPolicy.EXPLICIT_VALIDATION,
|
|
|
|
|
+ tools=["agent"],
|
|
|
|
|
+ tool_groups=None,
|
|
|
|
|
+ )
|
|
|
|
|
+ trace = Trace(trace_id="worker-trace", mode="agent", agent_role="worker", context={})
|
|
|
|
|
+ result = await runner._execute_authorized_tool(
|
|
|
|
|
+ "agent", {}, "call", config, trace, None, 1
|
|
|
|
|
+ )
|
|
|
|
|
+ assert result["error"] == "unauthorized_tool"
|
|
|
|
|
+ assert executed is False
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+@pytest.mark.asyncio
|
|
|
|
|
+async def test_framework_context_overrides_host_spoofing():
|
|
|
|
|
+ registry = ToolRegistry()
|
|
|
|
|
+ register_preset(
|
|
|
|
|
+ "test_worker_context",
|
|
|
|
|
+ AgentPreset(role=AgentRole.WORKER, allowed_tools=[], max_iterations=5, skills=[]),
|
|
|
|
|
+ )
|
|
|
|
|
+ runner = AgentRunner(tool_registry=registry, llm_call=lambda **_: None, task_coordinator="real")
|
|
|
|
|
+ config = RunConfig(
|
|
|
|
|
+ agent_type="test_worker_context",
|
|
|
|
|
+ completion_policy=CompletionPolicy.EXPLICIT_VALIDATION,
|
|
|
|
|
+ context={"role": "planner", "task_id": "spoofed", "coordinator": "fake"},
|
|
|
|
|
+ )
|
|
|
|
|
+ trace = Trace(
|
|
|
|
|
+ trace_id="trace", mode="agent", agent_role="worker",
|
|
|
|
|
+ context={"root_trace_id": "root", "task_id": "real-task", "attempt_id": "attempt"},
|
|
|
|
|
+ )
|
|
|
|
|
+ context = runner._build_protected_tool_context(config, trace, None, 1)
|
|
|
|
|
+ assert context["role"] == "worker"
|
|
|
|
|
+ assert context["task_id"] == "real-task"
|
|
|
|
|
+ assert context["coordinator"] == "real"
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def _disabled_knowledge():
|
|
|
|
|
+ from agent.tools.builtin.knowledge import KnowledgeConfig
|
|
|
|
|
+ return KnowledgeConfig(enable_extraction=False, enable_completion_extraction=False, enable_injection=False)
|