Support setting the current_tool to None to support method interception/limits

Author: zastrowm

src/strands/event_loop/event_loop.py

@@ -14,6 +14,7 @@
 from typing import TYPE_CHECKING, Any, AsyncGenerator
 
 from ..experimental.hooks import AfterToolInvocationEvent, BeforeToolInvocationEvent
+from ..experimental.hooks.registry import get_registry
 from ..telemetry.metrics import Trace
 from ..telemetry.tracer import get_tracer
 from ..tools.executor import run_tools, validate_and_prepare_tools
@@ -288,40 +289,61 @@ def run_tool(agent: "Agent", tool_use: ToolUse, kwargs: dict[str, Any]) -> ToolG
         }
     )
 
-    before_event = BeforeToolInvocationEvent(
-        agent=agent,
-        selected_tool=tool_func,
-        tool_use=tool_use,
-        kwargs=kwargs,
+    before_event = get_registry(agent).invoke_callbacks(
+        BeforeToolInvocationEvent(
+            agent=agent,
+            selected_tool=tool_func,
+            tool_use=tool_use,
+            kwargs=kwargs,
+        )
     )
-    agent._hooks.invoke_callbacks(before_event)
 
     try:
         selected_tool = before_event.selected_tool
         tool_use = before_event.tool_use
 
         # Check if tool exists
         if not selected_tool:
-            logger.error(
-                "tool_name=<%s>, available_tools=<%s> | tool not found in registry",
-                tool_name,
-                list(agent.tool_registry.registry.keys()),
-            )
-            return {
+            if tool_func == selected_tool:
+                logger.error(
+                    "tool_name=<%s>, available_tools=<%s> | tool not found in registry",
+                    tool_name,
+                    list(agent.tool_registry.registry.keys()),
+                )
+            else:
+                logger.debug(
+                    "tool_name=<%s>, tool_use_id=<%s> | a hook resulted in a non-existing tool call",
+                    tool_name,
+                    tool_use,
+                )
+
+            result: ToolResult = {
                 "toolUseId": str(tool_use.get("toolUseId")),
                 "status": "error",
                 "content": [{"text": f"Unknown tool: {tool_name}"}],
             }
+            # for every Before event call, we need to have an AfterEvent call
+            after_event = get_registry(agent).invoke_callbacks(
+                AfterToolInvocationEvent(
+                    agent=agent,
+                    selected_tool=selected_tool,
+                    tool_use=tool_use,
+                    kwargs=kwargs,
+                    result=result,
+                )
+            )
+            return after_event.result
 
         result = yield from selected_tool.stream(tool_use, **kwargs)
-        after_event = AfterToolInvocationEvent(
-            agent=agent,
-            selected_tool=selected_tool,
-            tool_use=tool_use,
-            kwargs=kwargs,
-            result=result,
+        after_event = get_registry(agent).invoke_callbacks(
+            AfterToolInvocationEvent(
+                agent=agent,
+                selected_tool=selected_tool,
+                tool_use=tool_use,
+                kwargs=kwargs,
+                result=result,
+            )
         )
-        agent._hooks.invoke_callbacks(after_event)
         return after_event.result
 
     except Exception as e:
@@ -331,15 +353,16 @@ def run_tool(agent: "Agent", tool_use: ToolUse, kwargs: dict[str, Any]) -> ToolG
             "status": "error",
             "content": [{"text": f"Error: {str(e)}"}],
         }
-        after_event = AfterToolInvocationEvent(
-            agent=agent,
-            selected_tool=selected_tool,
-            tool_use=tool_use,
-            kwargs=kwargs,
-            result=error_result,
-            exception=e,
+        after_event = get_registry(agent).invoke_callbacks(
+            AfterToolInvocationEvent(
+                agent=agent,
+                selected_tool=selected_tool,
+                tool_use=tool_use,
+                kwargs=kwargs,
+                result=error_result,
+                exception=e,
+            )
         )
-        agent._hooks.invoke_callbacks(after_event)
         return after_event.result
 
 

src/strands/experimental/hooks/registry.py

@@ -72,6 +72,9 @@ def __setattr__(self, name: str, value: Any) -> None:
 T = TypeVar("T", bound=Callable)
 TEvent = TypeVar("TEvent", bound=HookEvent, contravariant=True)
 
+# Non-contravariant generic for when invoking events
+TInvokeEvent = TypeVar("TInvokeEvent", bound=HookEvent)
+
 
 class HookProvider(Protocol):
     """Protocol for objects that provide hook callbacks to an agent.
@@ -178,7 +181,7 @@ def register_hooks(self, registry: HookRegistry):
         """
         hook.register_hooks(self)
 
-    def invoke_callbacks(self, event: TEvent) -> None:
+    def invoke_callbacks(self, event: TInvokeEvent) -> TInvokeEvent:
         """Invoke all registered callbacks for the given event.
 
         This method finds all callbacks registered for the event's type and
@@ -191,6 +194,9 @@ def invoke_callbacks(self, event: TEvent) -> None:
         Raises:
             Any exceptions raised by callback functions will propagate to the caller.
 
+        Returns:
+            The event dispatched to registered callbacks.
+
         Example:
             ```python
             event = StartRequestEvent(agent=my_agent)
@@ -200,6 +206,8 @@ def invoke_callbacks(self, event: TEvent) -> None:
         for callback in self.get_callbacks_for(event):
             callback(event)
 
+        return event
+
     def get_callbacks_for(self, event: TEvent) -> Generator[HookCallback[TEvent], None, None]:
         """Get callbacks registered for the given event in the appropriate order.
 
@@ -227,3 +235,18 @@ def get_callbacks_for(self, event: TEvent) -> Generator[HookCallback[TEvent], No
             yield from reversed(callbacks)
         else:
             yield from callbacks
+
+
+def get_registry(agent: "Agent") -> HookRegistry:
+    """*Experimental*: Get the hooks registry for the provided agent.
+
+    This function is available while hooks are in experimental preview.
+
+    Args:
+        agent: The agent whose hook registry should be returned.
+
+    Returns:
+        The HookRegistry for the given agent.
+
+    """
+    return agent._hooks

tests/strands/event_loop/test_event_loop.py

@@ -7,7 +7,7 @@
 import strands
 import strands.telemetry
 from strands.event_loop.event_loop import run_tool
-from strands.experimental.hooks import AfterToolInvocationEvent, BeforeToolInvocationEvent, HookRegistry
+from strands.experimental.hooks import AfterToolInvocationEvent, BeforeToolInvocationEvent, HookProvider, HookRegistry
 from strands.telemetry.metrics import EventLoopMetrics
 from strands.tools.registry import ToolRegistry
 from strands.types.exceptions import ContextWindowOverflowException, EventLoopException, ModelThrottledException
@@ -824,6 +824,35 @@ def test_run_tool_hooks(agent, generate, hook_provider, tool_times_2):
     )
 
 
+def test_run_tool_hooks_on_missing_tool(agent, tool_registry, generate, hook_provider):
+    """Test that AfterToolInvocation hook is invoked even when tool throws exception."""
+    process = run_tool(
+        agent=agent,
+        tool_use={"toolUseId": "test", "name": "missing_tool", "input": {"x": 5}},
+        kwargs={},
+    )
+
+    _, result = generate(process)
+
+    assert len(hook_provider.events_received) == 2
+
+    assert hook_provider.events_received[0] == BeforeToolInvocationEvent(
+        agent=agent,
+        selected_tool=None,
+        tool_use={"input": {"x": 5}, "name": "missing_tool", "toolUseId": "test"},
+        kwargs=ANY,
+    )
+
+    assert hook_provider.events_received[1] == AfterToolInvocationEvent(
+        agent=agent,
+        selected_tool=None,
+        tool_use={"input": {"x": 5}, "name": "missing_tool", "toolUseId": "test"},
+        kwargs=ANY,
+        result={"content": [{"text": "Unknown tool: missing_tool"}], "status": "error", "toolUseId": "test"},
+        exception=None,
+    )
+
+
 def test_run_tool_hook_after_tool_invocation_on_exception(agent, tool_registry, generate, hook_provider):
     """Test that AfterToolInvocation hook is invoked even when tool throws exception."""
     error = ValueError("Tool failed")
@@ -906,5 +935,79 @@ def modify_hook(event: AfterToolInvocationEvent):
 
     _, result = generate(process)
 
-    # Should return modified result instead of original (5 * 2 = 10)
     assert result == updated_result
+
+
+def test_run_tool_hook_after_tool_invocation_updates_with_missing_tool(agent, tool_times_2, generate, hook_registry):
+    """Test that modifying properties on AfterToolInvocation takes effect."""
+
+    updated_result = {"toolUseId": "modified", "status": "success", "content": [{"text": "modified_result"}]}
+
+    def modify_hook(event: AfterToolInvocationEvent):
+        # Modify result to change the output
+        event.result = updated_result
+
+    hook_registry.add_callback(AfterToolInvocationEvent, modify_hook)
+
+    process = run_tool(
+        agent=agent,
+        tool_use={"toolUseId": "test", "name": "missing_tool", "input": {"x": 5}},
+        kwargs={},
+    )
+
+    _, result = generate(process)
+
+    assert result == updated_result
+
+
+def test_run_tool_hook_update_result_with_missing_tool(agent, generate, tool_registry, hook_registry):
+    """Test that modifying properties on AfterToolInvocation takes effect."""
+
+    @strands.tool
+    def test_quota():
+        return "9"
+
+    tool_registry.register_tool(test_quota)
+
+    class ExampleProvider(HookProvider):
+        def register_hooks(self, registry: "HookRegistry") -> None:
+            registry.add_callback(BeforeToolInvocationEvent, self.before_tool_call)
+            registry.add_callback(AfterToolInvocationEvent, self.after_tool_call)
+
+        def before_tool_call(self, event: BeforeToolInvocationEvent):
+            if event.tool_use.get("name") == "test_quota":
+                event.selected_tool = None
+
+        def after_tool_call(self, event: AfterToolInvocationEvent):
+            if event.tool_use.get("name") == "test_quota":
+                event.result = {
+                    "status": "error",
+                    "toolUseId": "test",
+                    "content": [{"text": "This tool has been used too many times!"}],
+                }
+
+    hook_registry.add_hook(ExampleProvider())
+
+    with patch.object(strands.event_loop.event_loop, "logger") as mock_logger:
+        process = run_tool(
+            agent=agent,
+            tool_use={"toolUseId": "test", "name": "test_quota", "input": {"x": 5}},
+            kwargs={},
+        )
+
+        _, result = generate(process)
+
+    assert result == {
+        "status": "error",
+        "toolUseId": "test",
+        "content": [{"text": "This tool has been used too many times!"}],
+    }
+
+    assert mock_logger.debug.call_args_list == [
+        call("tool_use=<%s> | streaming", {"toolUseId": "test", "name": "test_quota", "input": {"x": 5}}),
+        call(
+            "tool_name=<%s>, tool_use_id=<%s> | a hook resulted in a non-existing tool call",
+            "test_quota",
+            {"toolUseId": "test", "name": "test_quota", "input": {"x": 5}},
+        ),
+    ]