From 5446548f1ede5ab62754431dd494836ef772b612 Mon Sep 17 00:00:00 2001 From: Ratish1 Date: Fri, 14 Nov 2025 18:18:48 +0400 Subject: [PATCH] feat(tools): add replace method to ToolRegistry --- src/strands/tools/registry.py | 28 +++++++ tests/strands/agent/test_agent.py | 4 +- tests/strands/tools/test_registry.py | 109 +++++++++++++++++++++++++++ 3 files changed, 139 insertions(+), 2 deletions(-) diff --git a/src/strands/tools/registry.py b/src/strands/tools/registry.py index c80b80f64..290987661 100644 --- a/src/strands/tools/registry.py +++ b/src/strands/tools/registry.py @@ -277,6 +277,34 @@ def register_tool(self, tool: AgentTool) -> None: list(self.dynamic_tools.keys()), ) + def replace(self, tool_name: str, new_tool: AgentTool) -> None: + """Replace an existing tool with a new implementation. + + This performs an atomic swap of the tool implementation in the registry. + The replacement takes effect on the next agent invocation. + + Args: + tool_name: Name of the tool to replace. + new_tool: New tool implementation. + + Raises: + ValueError: If the tool doesn't exist or if names don't match. + """ + if tool_name not in self.registry: + raise ValueError(f"Cannot replace tool '{tool_name}' - tool does not exist") + + if new_tool.tool_name != tool_name: + raise ValueError(f"Tool names must match - expected '{tool_name}', got '{new_tool.tool_name}'") + + # Atomic replacement in main registry + self.registry[tool_name] = new_tool + + # Update dynamic_tools to match new tool's dynamic status + if new_tool.is_dynamic: + self.dynamic_tools[tool_name] = new_tool + elif tool_name in self.dynamic_tools: + del self.dynamic_tools[tool_name] + def get_tools_dirs(self) -> List[Path]: """Get all tool directory paths. diff --git a/tests/strands/agent/test_agent.py b/tests/strands/agent/test_agent.py index 3a0bc2dfb..550422cfe 100644 --- a/tests/strands/agent/test_agent.py +++ b/tests/strands/agent/test_agent.py @@ -2240,8 +2240,8 @@ def test_agent_backwards_compatibility_single_text_block(): # Should extract text for backwards compatibility assert agent.system_prompt == text - - + + @pytest.mark.parametrize( "content, expected", [ diff --git a/tests/strands/tools/test_registry.py b/tests/strands/tools/test_registry.py index c700016f6..2056f68af 100644 --- a/tests/strands/tools/test_registry.py +++ b/tests/strands/tools/test_registry.py @@ -389,3 +389,112 @@ async def track_load_tools(*args, **kwargs): # Verify add_consumer was called with the registry ID mock_provider.add_consumer.assert_called_once_with(registry._registry_id) + + +def test_tool_registry_replace_existing_tool(): + """Test replacing an existing tool.""" + old_tool = MagicMock() + old_tool.tool_name = "my_tool" + old_tool.is_dynamic = False + old_tool.supports_hot_reload = False + + new_tool = MagicMock() + new_tool.tool_name = "my_tool" + new_tool.is_dynamic = False + + registry = ToolRegistry() + registry.register_tool(old_tool) + registry.replace("my_tool", new_tool) + + assert registry.registry["my_tool"] == new_tool + + +def test_tool_registry_replace_nonexistent_tool(): + """Test replacing a tool that doesn't exist raises ValueError.""" + new_tool = MagicMock() + new_tool.tool_name = "my_tool" + + registry = ToolRegistry() + + with pytest.raises(ValueError, match="Cannot replace tool 'my_tool' - tool does not exist"): + registry.replace("my_tool", new_tool) + + +def test_tool_registry_replace_with_different_name(): + """Test replacing with different name raises ValueError.""" + old_tool = MagicMock() + old_tool.tool_name = "old_tool" + old_tool.is_dynamic = False + old_tool.supports_hot_reload = False + + new_tool = MagicMock() + new_tool.tool_name = "new_tool" + + registry = ToolRegistry() + registry.register_tool(old_tool) + + with pytest.raises(ValueError, match="Tool names must match"): + registry.replace("old_tool", new_tool) + + +def test_tool_registry_replace_dynamic_tool(): + """Test replacing a dynamic tool updates both registries.""" + old_tool = MagicMock() + old_tool.tool_name = "dynamic_tool" + old_tool.is_dynamic = True + old_tool.supports_hot_reload = True + + new_tool = MagicMock() + new_tool.tool_name = "dynamic_tool" + new_tool.is_dynamic = True + + registry = ToolRegistry() + registry.register_tool(old_tool) + registry.replace("dynamic_tool", new_tool) + + assert registry.registry["dynamic_tool"] == new_tool + assert registry.dynamic_tools["dynamic_tool"] == new_tool + + +def test_tool_registry_replace_dynamic_with_non_dynamic(): + """Test replacing a dynamic tool with non-dynamic tool removes from dynamic_tools.""" + old_tool = MagicMock() + old_tool.tool_name = "my_tool" + old_tool.is_dynamic = True + old_tool.supports_hot_reload = True + + new_tool = MagicMock() + new_tool.tool_name = "my_tool" + new_tool.is_dynamic = False + + registry = ToolRegistry() + registry.register_tool(old_tool) + + assert "my_tool" in registry.dynamic_tools + + registry.replace("my_tool", new_tool) + + assert registry.registry["my_tool"] == new_tool + assert "my_tool" not in registry.dynamic_tools + + +def test_tool_registry_replace_non_dynamic_with_dynamic(): + """Test replacing a non-dynamic tool with dynamic tool adds to dynamic_tools.""" + old_tool = MagicMock() + old_tool.tool_name = "my_tool" + old_tool.is_dynamic = False + old_tool.supports_hot_reload = False + + new_tool = MagicMock() + new_tool.tool_name = "my_tool" + new_tool.is_dynamic = True + + registry = ToolRegistry() + registry.register_tool(old_tool) + + assert "my_tool" not in registry.dynamic_tools + + registry.replace("my_tool", new_tool) + + assert registry.registry["my_tool"] == new_tool + assert registry.dynamic_tools["my_tool"] == new_tool