diff --git a/src/strands/tools/registry.py b/src/strands/tools/registry.py index 15150847d..fbb518e0c 100644 --- a/src/strands/tools/registry.py +++ b/src/strands/tools/registry.py @@ -279,6 +279,32 @@ def register_tool(self, tool: AgentTool) -> None: list(self.dynamic_tools.keys()), ) + def replace(self, 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: + new_tool: New tool implementation. Its name must match the tool being replaced. + + Raises: + ValueError: If the tool doesn't exist. + """ + tool_name = new_tool.tool_name + + if tool_name not in self.registry: + raise ValueError(f"Cannot replace tool '{tool_name}' - tool does not exist") + + # 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/tools/test_registry.py b/tests/strands/tools/test_registry.py index 9ae51dcfe..d44936f3e 100644 --- a/tests/strands/tools/test_registry.py +++ b/tests/strands/tools/test_registry.py @@ -511,3 +511,95 @@ 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(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(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(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(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(new_tool) + + assert registry.registry["my_tool"] == new_tool + assert registry.dynamic_tools["my_tool"] == new_tool