|
|
@@ -0,0 +1,311 @@
|
|
|
+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)
|
SamLee