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

from __future__ import annotations

import json

import pytest

from agent.tools.models import ToolCapability
from agent.tools.registry import ToolRegistry, get_tool_registry, tool


def _schema(name: str) -> dict:
    return {
        "type": "function",
        "function": {
            "name": name,
            "description": name,
            "parameters": {"type": "object", "properties": {}},
        },
    }


@pytest.mark.asyncio
async def test_registering_the_same_tool_twice_is_idempotent() -> None:
    registry = ToolRegistry()

    async def sample() -> str:
        return "ok"

    registry.register(
        sample,
        schema=_schema("sample"),
        groups=["test"],
        capabilities=[ToolCapability.READ],
    )
    assert await registry.execute("sample", {}) == "ok"

    registry.register(
        sample,
        schema=_schema("sample"),
        groups=["test"],
        capabilities=[ToolCapability.READ],
    )

    assert registry.get_tool_names() == ["sample"]
    assert registry.get_stats("sample")["sample"]["call_count"] == 1


@pytest.mark.asyncio
async def test_same_factory_definition_rebinds_a_fresh_host_adapter() -> None:
    registry = ToolRegistry()

    class Gateway:
        def __init__(self, value: str) -> None:
            self.value = value

    def make_tool(gateway: Gateway):
        async def host_tool() -> str:
            return gateway.value

        return host_tool

    registry.register(
        make_tool(Gateway("first")),
        schema=_schema("host_tool"),
        capabilities=[ToolCapability.READ],
    )
    assert await registry.execute("host_tool", {}) == "first"

    registry.register(
        make_tool(Gateway("second")),
        schema=_schema("host_tool"),
        capabilities=[ToolCapability.READ],
    )

    assert await registry.execute("host_tool", {}) == "second"
    assert registry.get_stats("host_tool")["host_tool"]["call_count"] == 2


def test_same_name_from_a_different_callable_is_rejected() -> None:
    registry = ToolRegistry()

    async def original() -> str:
        return "original"

    async def replacement() -> str:
        return "replacement"

    replacement.__name__ = original.__name__
    registry.register(original, schema=_schema("original"))

    with pytest.raises(
        ValueError,
        match=r"Duplicate tool registration for 'original'.*different callable",
    ) as exc_info:
        registry.register(replacement, schema=_schema("original"))

    message = str(exc_info.value)
    assert "original" in message
    assert "replacement" in message
    assert registry._tools["original"]["func"] is original


def test_reregistering_a_tool_with_different_metadata_is_rejected() -> None:
    registry = ToolRegistry()

    async def sample() -> str:
        return "ok"

    registry.register(sample, schema=_schema("sample"), groups=["read"])

    with pytest.raises(
        ValueError,
        match=r"Duplicate tool registration for 'sample'.*metadata differs: groups",
    ):
        registry.register(sample, schema=_schema("sample"), groups=["write"])


@pytest.mark.asyncio
async def test_legacy_execute_inject_values_are_defaults_not_overrides() -> None:
    registry = ToolRegistry()

    async def sample(value: str) -> str:
        return value

    registry.register(sample, schema=_schema("sample"))

    assert await registry.execute(
        "sample",
        {},
        inject_values={"value": "injected", "unknown": "ignored"},
    ) == "injected"
    assert await registry.execute(
        "sample",
        {"value": "explicit"},
        inject_values={"value": "injected"},
    ) == "explicit"


@pytest.mark.asyncio
async def test_host_can_replace_a_capability_matched_submission_protocol() -> None:
    registry = ToolRegistry()

    async def submit_attempt() -> str:
        return "framework"

    submit_attempt.__module__ = "agent.tools.builtin.example"

    registry.register(
        submit_attempt,
        schema=_schema("submit_attempt"),
        groups=["orchestration_worker"],
        capabilities=[ToolCapability.ATTEMPT_SUBMIT],
    )
    assert await registry.execute("submit_attempt", {}) == "framework"

    async def submit_attempt() -> str:  # noqa: F811
        return "host"

    submit_attempt.__module__ = "host.tools"

    registry.register(
        submit_attempt,
        schema={
            **_schema("submit_attempt"),
            "x-host-contract": True,
        },
        groups=["host"],
        capabilities=[ToolCapability.ATTEMPT_SUBMIT],
    )

    assert await registry.execute("submit_attempt", {}) == "host"
    assert registry.get_stats("submit_attempt")["submit_attempt"]["call_count"] == 2

    async def competing_host_tool() -> str:
        return "competing host"

    competing_host_tool.__name__ = "submit_attempt"
    competing_host_tool.__module__ = "another_host.tools"
    with pytest.raises(ValueError, match="different callable"):
        registry.register(
            competing_host_tool,
            schema=_schema("submit_attempt"),
            groups=["another_host"],
            capabilities=[ToolCapability.ATTEMPT_SUBMIT],
        )


def test_host_cannot_replace_a_builtin_with_different_capabilities() -> None:
    registry = ToolRegistry()

    async def read_contract() -> str:
        return "framework"

    read_contract.__module__ = "agent.tools.builtin.example"
    registry.register(
        read_contract,
        schema=_schema("read_contract"),
        capabilities=[ToolCapability.READ],
    )

    async def read_contract() -> str:  # noqa: F811
        return "host"

    read_contract.__module__ = "host.tools"
    with pytest.raises(ValueError, match="different callable"):
        registry.register(
            read_contract,
            schema=_schema("read_contract"),
            capabilities=[ToolCapability.WRITE],
        )


def test_builtin_tool_registration_contract() -> None:
    """Guard the security-relevant metadata without snapshotting optional tools."""
    registry = get_tool_registry()
    schemas = registry.get_schemas()
    schema_names = [schema["function"]["name"] for schema in schemas]

    assert len(schema_names) == len(set(schema_names))
    assert set(schema_names) == set(registry.get_tool_names())

    expected = {
        "read_file": (
            "agent.tools.builtin.file.read",
            {"core"},
            {ToolCapability.READ},
        ),
        "write_file": (
            "agent.tools.builtin.file.write",
            {"core"},
            {ToolCapability.WRITE},
        ),
        "glob_files": (
            "agent.tools.builtin.file.glob",
            {"core"},
            {ToolCapability.READ},
        ),
        "bash_command": (
            "agent.tools.builtin.bash",
            {"system"},
            {ToolCapability.WRITE},
        ),
        "agent": (
            "agent.tools.builtin.subagent",
            {"core"},
            {ToolCapability.AGENT_SPAWN},
        ),
        "task_plan": (
            "agent.tools.builtin.orchestration",
            {"orchestration_planner"},
            {ToolCapability.TASK_CONTROL},
        ),
        "submit_attempt": (
            "agent.tools.builtin.orchestration",
            {"orchestration_worker"},
            {ToolCapability.ATTEMPT_SUBMIT},
        ),
        "submit_validation": (
            "agent.tools.builtin.orchestration",
            {"orchestration_validator"},
            {ToolCapability.VALIDATION_SUBMIT},
        ),
    }

    actual = {
        name: (
            registry._tools[name]["func"].__module__,
            set(registry._tools[name]["groups"]),
            set(registry.get_capabilities(name)),
        )
        for name in expected
    }
    assert actual == expected, json.dumps(
        {
            name: {
                "module": module,
                "groups": sorted(groups),
                "capabilities": sorted(map(str, capabilities)),
            }
            for name, (module, groups, capabilities) in actual.items()
        },
        ensure_ascii=False,
        indent=2,
    )


def test_tool_decorator_applies_explicit_schema_descriptions() -> None:
    registry = get_tool_registry()

    try:
        @tool(
            description="Explicit tool description",
            param_descriptions={"value": "Explicit value description"},
        )
        async def governance_description_probe(value: str) -> str:
            """Docstring description.

            Args:
                value: Docstring value description.
            """
            return value

        schema = registry.get_schemas(["governance_description_probe"])[0]["function"]
        assert schema["description"] == "Explicit tool description"
        assert (
            schema["parameters"]["properties"]["value"]["description"]
            == "Explicit value description"
        )
    finally:
        registry._tools.pop("governance_description_probe", None)
        registry._stats.pop("governance_description_probe", None)