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 ToolCapability, 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"


@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)