script-build-and-framework-tmp/agent/tests/test_runner_policy.py

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)