feat: Add hooks for before/after tool calls + allow hooks to update values

Add the ability to intercept/modify tool calls by implementing support for BeforeToolInvocation & AfterToolInvocation hooks

Author: zastrowm

src/strands/agent/agent.py

@@ -283,7 +283,7 @@ def __init__(
         self.load_tools_from_directory = load_tools_from_directory
 
         self.tool_registry = ToolRegistry()
-        self.tool_handler = AgentToolHandler(tool_registry=self.tool_registry)
+        self.tool_handler = AgentToolHandler(agent=self, tool_registry=self.tool_registry)
 
         # Process tool list if provided
         if tools is not None:

src/strands/experimental/hooks/__init__.py

@@ -29,13 +29,15 @@ def log_end(self, event: EndRequestEvent) -> None:
 type-safe system that supports multiple subscribers per event type.
 """
 
-from .events import AgentInitializedEvent, EndRequestEvent, StartRequestEvent
+from .events import AfterToolInvocation, AgentInitializedEvent, BeforeToolInvocation, EndRequestEvent, StartRequestEvent
 from .registry import HookCallback, HookEvent, HookProvider, HookRegistry
 
 __all__ = [
     "AgentInitializedEvent",
     "StartRequestEvent",
     "EndRequestEvent",
+    "BeforeToolInvocation",
+    "AfterToolInvocation",
     "HookEvent",
     "HookProvider",
     "HookCallback",

src/strands/experimental/hooks/events.py

@@ -4,7 +4,9 @@
 """
 
 from dataclasses import dataclass
+from typing import Any, Optional
 
+from ...types.tools import AgentTool, ToolResult, ToolUse
 from .registry import HookEvent
 
 
@@ -56,9 +58,63 @@ class EndRequestEvent(HookEvent):
 
     @property
     def should_reverse_callbacks(self) -> bool:
-        """Return True to invoke callbacks in reverse order for proper cleanup.
+        """True to invoke callbacks in reverse order."""
+        return True
+
+
+@dataclass
+class BeforeToolInvocation(HookEvent):
+    """Event triggered before a tool is invoked.
+
+    This event is fired just before the agent executes a tool, allowing hook
+    providers to inspect, modify, or replace the tool that will be executed.
+    The selected_tool can be modified by hook callbacks to change which tool
+    gets executed.
+
+    Attributes:
+        selected_tool: The tool that will be invoked. Can be modified by hooks
+            to change which tool gets executed. This may be None if tool lookup failed.
+        tool_use: The tool parameters that will be passed to selected_tool.
+        kwargs: Keyword arguments that will be passed to the tool.
+    """
+
+    selected_tool: Optional[AgentTool]
+    tool_use: ToolUse
+    kwargs: dict[str, Any]
+
+    def _can_write(self, name: str) -> bool:
+        return name in ["selected_tool", "tool_use"]
+
+
+@dataclass
+class AfterToolInvocation(HookEvent):
+    """Event triggered after a tool invocation completes.
 
-        Returns:
-            True, indicating callbacks should be invoked in reverse order.
-        """
+    This event is fired after the agent has finished executing a tool,
+    regardless of whether the execution was successful or resulted in an error.
+    Hook providers can use this event for cleanup, logging, or post-processing.
+
+    Note: This event uses reverse callback ordering, meaning callbacks registered
+    later will be invoked first during cleanup.
+
+    Attributes:
+        selected_tool: The tool that was invoked. It may be None if tool lookup failed.
+        tool_use: The tool parameters that were passed to the tool invoked.
+        kwargs: Keyword arguments that were passed to the tool
+        result: The result of the tool invocation. Either a ToolResult on success
+            or an Exception if the tool execution failed.
+    """
+
+    selected_tool: Optional[AgentTool]
+    tool_use: ToolUse
+    kwargs: dict[str, Any]
+    result: ToolResult
+    exception: Optional[Exception] = None
+
+    def _can_write(self, name: str) -> bool:
+        return name == "result"
+
+    @property
+    def should_reverse_callbacks(self) -> bool:
+        """True to invoke callbacks in reverse order."""
         return True

src/strands/experimental/hooks/registry.py

@@ -8,7 +8,7 @@
 """
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Callable, Generator, Generic, Protocol, Type, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, Generator, Generic, Protocol, Type, TypeVar
 
 if TYPE_CHECKING:
     from ...agent import Agent
@@ -34,6 +34,40 @@ def should_reverse_callbacks(self) -> bool:
         """
         return False
 
+    ### Code below is infrastructure for disallowing updates to properties ###
+    ### that aren't expected to be updated from hook callbacks.            ###
+
+    def _can_write(self, name: str) -> bool:
+        """Check if the given property can be written to.
+
+        Args:
+            name: The name of the property to check.
+
+        Returns:
+            True if the property can be written to, False otherwise.
+        """
+        return False
+
+    def __post_init__(self) -> None:
+        """Disallow writes to non-approved properties."""
+        # This is needed as otherwise the class can't be initialized at all, so we trigger
+        # this after class initialization
+        super().__setattr__("_disallow_writes", True)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        """Prevent setting attributes on hook events.
+
+        Raises:
+            AttributeError: Always raised to prevent setting attributes on hook events.
+        """
+        #  Allow setting attributes:
+        #    - during init (when __dict__) doesn't exist
+        #    - if the subclass specifically said the property is writable
+        if not hasattr(self, "_disallow_writes") or self._can_write(name):
+            return super().__setattr__(name, value)
+
+        raise AttributeError(f"Property {name} is not writable")
+
 
 T = TypeVar("T", bound=Callable)
 TEvent = TypeVar("TEvent", bound=HookEvent, contravariant=True)

src/strands/handlers/tool_handler.py

@@ -1,15 +1,19 @@
 """This module provides handlers for managing tool invocations."""
 
 import logging
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
+from ..experimental.hooks import AfterToolInvocation, BeforeToolInvocation
 from ..tools.registry import ToolRegistry
 from ..types.content import Messages
 from ..types.models import Model
-from ..types.tools import ToolConfig, ToolGenerator, ToolHandler, ToolUse
+from ..types.tools import ToolConfig, ToolGenerator, ToolHandler, ToolResult, ToolUse
 
 logger = logging.getLogger(__name__)
 
+if TYPE_CHECKING:
+    from ..agent import Agent
+
 
 class AgentToolHandler(ToolHandler):
     """Handler for processing tool invocations in agent.
@@ -18,13 +22,15 @@ class AgentToolHandler(ToolHandler):
     invoking them with the appropriate parameters.
     """
 
-    def __init__(self, tool_registry: ToolRegistry) -> None:
+    def __init__(self, agent: "Agent", tool_registry: ToolRegistry) -> None:
         """Initialize handler.
 
         Args:
+            agent: The agent associated with this tool handler.
             tool_registry: Registry of available tools.
         """
         self.tool_registry = tool_registry
+        self.agent = agent
 
     def process(
         self,
@@ -55,44 +61,75 @@ def process(
             The final tool result or an error response if the tool fails or is not found.
         """
         logger.debug("tool=<%s> | invoking", tool)
-        tool_use_id = tool["toolUseId"]
         tool_name = tool["name"]
 
         # Get the tool info
         tool_info = self.tool_registry.dynamic_tools.get(tool_name)
         tool_func = tool_info if tool_info is not None else self.tool_registry.registry.get(tool_name)
 
+        # Add standard arguments to kwargs for Python tools
+        kwargs.update(
+            {
+                "model": model,
+                "system_prompt": system_prompt,
+                "messages": messages,
+                "tool_config": tool_config,
+            }
+        )
+
+        before_event = BeforeToolInvocation(
+            agent=self.agent,
+            selected_tool=tool_func,
+            tool_use=tool,
+            kwargs=kwargs,
+        )
+        self.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 tool_func:
+            if not selected_tool:
                 logger.error(
                     "tool_name=<%s>, available_tools=<%s> | tool not found in registry",
                     tool_name,
                     list(self.tool_registry.registry.keys()),
                 )
                 return {
-                    "toolUseId": tool_use_id,
+                    "toolUseId": str(tool_use.get("toolUseId")),
                     "status": "error",
                     "content": [{"text": f"Unknown tool: {tool_name}"}],
                 }
-            # Add standard arguments to kwargs for Python tools
-            kwargs.update(
-                {
-                    "model": model,
-                    "system_prompt": system_prompt,
-                    "messages": messages,
-                    "tool_config": tool_config,
-                }
-            )
 
-            result = tool_func.invoke(tool, **kwargs)
-            yield {"result": result}  # Placeholder until tool_func becomes a generator from which we can yield from
-            return result
+            result = selected_tool.invoke(tool_use, **kwargs)
+            after_event = AfterToolInvocation(
+                agent=self.agent,
+                selected_tool=selected_tool,
+                tool_use=tool_use,
+                kwargs=kwargs,
+                result=result,
+            )
+            self.agent._hooks.invoke_callbacks(after_event)
+            yield {
+                "result": after_event.result
+            }  # Placeholder until tool_func becomes a generator from which we can yield from
+            return after_event.result
 
         except Exception as e:
             logger.exception("tool_name=<%s> | failed to process tool", tool_name)
-            return {
-                "toolUseId": tool_use_id,
+            error_result: ToolResult = {
+                "toolUseId": str(tool_use.get("toolUseId")),
                 "status": "error",
                 "content": [{"text": f"Error: {str(e)}"}],
             }
+            after_event = AfterToolInvocation(
+                agent=self.agent,
+                selected_tool=selected_tool,
+                tool_use=tool_use,
+                kwargs=kwargs,
+                result=error_result,
+                exception=e,
+            )
+            self.agent._hooks.invoke_callbacks(after_event)
+            return after_event.result

tests/fixtures/mock_hook_provider.py

@@ -14,7 +14,7 @@ def get_events(self) -> deque[HookEvent]:
 
     def register_hooks(self, registry: HookRegistry) -> None:
         for event_type in self.events_types:
-            registry.add_callback(event_type, self._add_event)
+            registry.add_callback(event_type, self.add_event)
 
-    def _add_event(self, event: HookEvent) -> None:
+    def add_event(self, event: HookEvent) -> None:
         self.events_received.append(event)

tests/strands/agent/test_agent_hooks.py

@@ -1,20 +1,27 @@
-import unittest.mock
-from unittest.mock import call
+from unittest.mock import ANY, Mock, call, patch
 
 import pytest
 from pydantic import BaseModel
 
 import strands
 from strands import Agent
-from strands.experimental.hooks import AgentInitializedEvent, EndRequestEvent, StartRequestEvent
+from strands.experimental.hooks import (
+    AfterToolInvocation,
+    AgentInitializedEvent,
+    BeforeToolInvocation,
+    EndRequestEvent,
+    StartRequestEvent,
+)
 from strands.types.content import Messages
 from tests.fixtures.mock_hook_provider import MockHookProvider
 from tests.fixtures.mocked_model_provider import MockedModelProvider
 
 
 @pytest.fixture
 def hook_provider():
-    return MockHookProvider([AgentInitializedEvent, StartRequestEvent, EndRequestEvent])
+    return MockHookProvider(
+        [AgentInitializedEvent, StartRequestEvent, EndRequestEvent, AfterToolInvocation, BeforeToolInvocation]
+    )
 
 
 @pytest.fixture
@@ -71,7 +78,7 @@ class User(BaseModel):
     return User(name="Jane Doe", age=30)
 
 
-@unittest.mock.patch("strands.experimental.hooks.registry.HookRegistry.invoke_callbacks")
+@patch("strands.experimental.hooks.registry.HookRegistry.invoke_callbacks")
 def test_agent__init__hooks(mock_invoke_callbacks):
     """Verify that the AgentInitializedEvent is emitted on Agent construction."""
     agent = Agent()
@@ -87,9 +94,19 @@ def test_agent__call__hooks(agent, hook_provider, agent_tool, tool_use):
     agent("test message")
 
     events = hook_provider.get_events()
-    assert len(events) == 2
 
+    assert len(events) == 4
     assert events.popleft() == StartRequestEvent(agent=agent)
+    assert events.popleft() == BeforeToolInvocation(
+        agent=agent, selected_tool=agent_tool, tool_use=tool_use, kwargs=ANY
+    )
+    assert events.popleft() == AfterToolInvocation(
+        agent=agent,
+        selected_tool=agent_tool,
+        tool_use=tool_use,
+        kwargs=ANY,
+        result={"content": [{"text": "!loot a dekovni I"}], "status": "success", "toolUseId": "123"},
+    )
     assert events.popleft() == EndRequestEvent(agent=agent)
 
 
@@ -105,16 +122,26 @@ async def test_agent_stream_async_hooks(agent, hook_provider, agent_tool, tool_u
         pass
 
     events = hook_provider.get_events()
-    assert len(events) == 2
 
+    assert len(events) == 4
     assert events.popleft() == StartRequestEvent(agent=agent)
+    assert events.popleft() == BeforeToolInvocation(
+        agent=agent, selected_tool=agent_tool, tool_use=tool_use, kwargs=ANY
+    )
+    assert events.popleft() == AfterToolInvocation(
+        agent=agent,
+        selected_tool=agent_tool,
+        tool_use=tool_use,
+        kwargs=ANY,
+        result={"content": [{"text": "!loot a dekovni I"}], "status": "success", "toolUseId": "123"},
+    )
     assert events.popleft() == EndRequestEvent(agent=agent)
 
 
 def test_agent_structured_output_hooks(agent, hook_provider, user, agenerator):
     """Verify that the correct hook events are emitted as part of structured_output."""
 
-    agent.model.structured_output = unittest.mock.Mock(return_value=agenerator([{"output": user}]))
+    agent.model.structured_output = Mock(return_value=agenerator([{"output": user}]))
     agent.structured_output(type(user), "example prompt")
 
     assert hook_provider.events_received == [StartRequestEvent(agent=agent), EndRequestEvent(agent=agent)]
@@ -124,7 +151,7 @@ def test_agent_structured_output_hooks(agent, hook_provider, user, agenerator):
 async def test_agent_structured_async_output_hooks(agent, hook_provider, user, agenerator):
     """Verify that the correct hook events are emitted as part of structured_output_async."""
 
-    agent.model.structured_output = unittest.mock.Mock(return_value=agenerator([{"output": user}]))
+    agent.model.structured_output = Mock(return_value=agenerator([{"output": user}]))
     await agent.structured_output_async(type(user), "example prompt")
 
     assert hook_provider.events_received == [StartRequestEvent(agent=agent), EndRequestEvent(agent=agent)]