diff --git a/nanobot/agent/tools/registry.py b/nanobot/agent/tools/registry.py index 137038c0..3d185d57 100644 --- a/nanobot/agent/tools/registry.py +++ b/nanobot/agent/tools/registry.py @@ -14,14 +14,17 @@ class ToolRegistry: def __init__(self): self._tools: dict[str, Tool] = {} + self._cached_definitions: list[dict[str, Any]] | None = None def register(self, tool: Tool) -> None: """Register a tool.""" self._tools[tool.name] = tool + self._cached_definitions = None def unregister(self, name: str) -> None: """Unregister a tool by name.""" self._tools.pop(name, None) + self._cached_definitions = None def get(self, name: str) -> Tool | None: """Get a tool by name.""" @@ -46,8 +49,12 @@ class ToolRegistry: """Get tool definitions with stable ordering for cache-friendly prompts. Built-in tools are sorted first as a stable prefix, then MCP tools are - sorted and appended. + sorted and appended. The result is cached until the next + register/unregister call. """ + if self._cached_definitions is not None: + return self._cached_definitions + definitions = [tool.to_schema() for tool in self._tools.values()] builtins: list[dict[str, Any]] = [] mcp_tools: list[dict[str, Any]] = [] @@ -60,7 +67,8 @@ class ToolRegistry: builtins.sort(key=self._schema_name) mcp_tools.sort(key=self._schema_name) - return builtins + mcp_tools + self._cached_definitions = builtins + mcp_tools + return self._cached_definitions def prepare_call( self, diff --git a/tests/tools/test_tool_registry.py b/tests/tools/test_tool_registry.py index f9e8ce5e..ca60f30e 100644 --- a/tests/tools/test_tool_registry.py +++ b/tests/tools/test_tool_registry.py @@ -71,3 +71,33 @@ def test_prepare_call_other_tools_keep_generic_object_validation() -> None: assert tool is not None assert params == ["TODO"] assert error == "Error: Invalid parameters for tool 'grep': parameters must be an object, got list" + + +def test_get_definitions_returns_cached_result() -> None: + registry = ToolRegistry() + registry.register(_FakeTool("read_file")) + first = registry.get_definitions() + assert registry._cached_definitions is not None + second = registry.get_definitions() + assert first == second + + +def test_register_invalidates_cache() -> None: + registry = ToolRegistry() + registry.register(_FakeTool("read_file")) + first = registry.get_definitions() + registry.register(_FakeTool("write_file")) + second = registry.get_definitions() + assert first is not second + assert len(second) == 2 + + +def test_unregister_invalidates_cache() -> None: + registry = ToolRegistry() + registry.register(_FakeTool("read_file")) + registry.register(_FakeTool("write_file")) + first = registry.get_definitions() + registry.unregister("write_file") + second = registry.get_definitions() + assert first is not second + assert len(second) == 1