feat(tools): add replace method to ToolRegistry

Author: Ratish1

src/strands/tools/registry.py

@@ -279,6 +279,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.
 

tests/strands/tools/test_registry.py

@@ -511,3 +511,112 @@ def test_validate_tool_spec_with_ref_property():
     assert props["ref_field"] == {"$ref": "#/$defs/SomeType"}
     assert "type" not in props["ref_field"]
     assert "description" not in props["ref_field"]
+
+
+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