Update typing

Author: Unshure

src/strands/plugins/decorator.py

@@ -1,111 +1,69 @@
 """Hook decorator for Plugin methods.
 
-This module provides the @hook decorator that marks methods as hook callbacks
-for automatic registration when the plugin is attached to an agent.
-
-The @hook decorator performs several functions:
-
-1. Marks methods as hook callbacks for automatic discovery by Plugin base class
-2. Infers event types from the callback's type hints (consistent with HookRegistry.add_callback)
-3. Supports both @hook and @hook() syntax
-4. Supports union types for multiple event types (e.g., BeforeModelCallEvent | AfterModelCallEvent)
-5. Stores hook metadata on the decorated method for later discovery
+Marks methods as hook callbacks for automatic registration when the plugin
+is attached to an agent. Infers event types from type hints and supports
+union types for multiple events.
 
 Example:
     ```python
-    from strands.plugins import Plugin, hook
-    from strands.hooks import BeforeModelCallEvent, AfterModelCallEvent
-
     class MyPlugin(Plugin):
-        name = "my-plugin"
-
         @hook
         def on_model_call(self, event: BeforeModelCallEvent):
             print(event)
-
-        @hook
-        def on_any_model_event(self, event: BeforeModelCallEvent | AfterModelCallEvent):
-            print(event)
     ```
 """
 
 from collections.abc import Callable
-from typing import TYPE_CHECKING, TypeVar, overload
+from typing import Generic, cast, overload
 
-if TYPE_CHECKING:
-    from ..hooks.registry import BaseHookEvent
+from ..hooks._type_inference import infer_event_types
+from ..hooks.registry import HookCallback, TEvent
 
-# Type for wrapped function
-T = TypeVar("T", bound=Callable[..., object])
+
+class _WrappedHookCallable(HookCallback, Generic[TEvent]):
+    """Wrapped version of HookCallback that includes a `_hook_event_types` argument."""
+
+    _hook_event_types: list[TEvent]
 
 
 # Handle @hook
 @overload
-def hook(__func: T) -> T: ...
+def hook(__func: HookCallback) -> _WrappedHookCallable: ...
 
 
 # Handle @hook()
 @overload
-def hook() -> Callable[[T], T]: ...
+def hook() -> Callable[[HookCallback], _WrappedHookCallable]: ...
 
 
 def hook(
-    func: T | None = None,
-) -> T | Callable[[T], T]:
-    """Decorator that marks a method as a hook callback for automatic registration.
-
-    This decorator enables declarative hook registration in Plugin classes. When a
-    Plugin is attached to an agent, methods marked with @hook are automatically
-    discovered and registered with the agent's hook registry.
+    func: HookCallback | None = None,
+) -> _WrappedHookCallable | Callable[[HookCallback], _WrappedHookCallable]:
+    """Mark a method as a hook callback for automatic registration.
 
-    The event type is inferred from the callback's type hint on the first parameter
-    (after 'self' for instance methods). Union types are supported for registering
-    a single callback for multiple event types.
-
-    The decorator can be used in two ways:
-    - As a simple decorator: `@hook`
-    - With parentheses: `@hook()`
+    Infers event type from the callback's type hint. Supports union types
+    for multiple events. Can be used as @hook or @hook().
 
     Args:
-        func: The function to decorate. When used as a simple decorator, this is
-            the function being decorated. When used with parentheses, this will be None.
+        func: The function to decorate.
 
     Returns:
-        The decorated function with hook metadata attached.
+        The decorated function with hook metadata.
 
     Raises:
-        ValueError: If the event type cannot be inferred from type hints, or if
-            the type hint is not a valid HookEvent subclass.
-
-    Example:
-        ```python
-        class MyPlugin(Plugin):
-            name = "my-plugin"
-
-            @hook
-            def on_model_call(self, event: BeforeModelCallEvent):
-                print(f"Model called: {event}")
-
-            @hook
-            def on_any_event(self, event: BeforeModelCallEvent | AfterModelCallEvent):
-                print(f"Event: {type(event).__name__}")
-        ```
+        ValueError: If event type cannot be inferred from type hints.
     """
 
-    def decorator(f: T) -> T:
-        # Import here to avoid circular dependency at runtime
-        from ..hooks._type_inference import infer_event_types
-
+    def decorator(f: HookCallback[TEvent]) -> _WrappedHookCallable[TEvent]:
         # Infer event types from type hints
-        event_types: list[type[BaseHookEvent]] = infer_event_types(f)  # type: ignore[arg-type]
+        event_types: list[type[TEvent]] = infer_event_types(f)
 
         # Store hook metadata on the function
-        f._hook_event_types = event_types  # type: ignore[attr-defined]
+        f_wrapped = cast(_WrappedHookCallable, f)
+        f_wrapped._hook_event_types = event_types
 
-        return f
+        return f_wrapped
 
-    # Handle both @hook and @hook() syntax
     if func is None:
         return decorator
-
     return decorator(func)

src/strands/plugins/plugin.py

@@ -9,7 +9,8 @@
 from collections.abc import Awaitable
 from typing import TYPE_CHECKING
 
-from strands.tools.decorator import DecoratedFunctionTool
+from ..tools.decorator import DecoratedFunctionTool
+from .decorator import _WrappedHookCallable
 
 if TYPE_CHECKING:
     from ..agent import Agent
@@ -74,17 +75,13 @@ def __init__(self) -> None:
         Scans the class for methods decorated with @hook and @tool and stores
         references for later registration when init_plugin is called.
         """
-        self._hooks: list[object] = []
+        self._hooks: list[_WrappedHookCallable] = []
         self._tools: list[DecoratedFunctionTool] = []
         self._discover_decorated_methods()
 
     def _discover_decorated_methods(self) -> None:
         """Scan class for @hook and @tool decorated methods."""
         for name in dir(self):
-            # Skip private and dunder methods
-            if name.startswith("_"):
-                continue
-
             try:
                 attr = getattr(self, name)
             except Exception:
@@ -118,7 +115,7 @@ def init_plugin(self, agent: "Agent") -> None | Awaitable[None]:
         for hook_callback in self._hooks:
             event_types = getattr(hook_callback, "_hook_event_types", [])
             for event_type in event_types:
-                agent.hooks.add_callback(event_type, hook_callback)
+                agent.add_hook(hook_callback, event_type)
                 logger.debug(
                     "plugin=<%s>, hook=<%s>, event_type=<%s> | registered hook",
                     self.name,

tests/strands/experimental/steering/core/test_handler.py

@@ -1,5 +1,6 @@
 """Unit tests for steering handler base class."""
 
+import inspect
 from unittest.mock import AsyncMock, Mock
 
 import pytest
@@ -8,6 +9,7 @@
 from strands.experimental.steering.core.context import SteeringContext, SteeringContextCallback, SteeringContextProvider
 from strands.experimental.steering.core.handler import SteeringHandler
 from strands.hooks.events import AfterModelCallEvent, BeforeToolCallEvent
+from strands.hooks.registry import HookRegistry
 from strands.plugins import Plugin
 
 
@@ -39,18 +41,14 @@ def test_steering_handler_is_plugin():
 
 def test_init_plugin():
     """Test init_plugin registers hooks on agent."""
-    from strands.hooks import HookRegistry
-
     handler = TestSteeringHandler()
     agent = Mock()
-    agent.hooks = HookRegistry()
-    agent.tool_registry = Mock()
 
     handler.init_plugin(agent)
 
-    # Verify hooks were auto-registered via @hook decorator
-    assert len(agent.hooks._registered_callbacks.get(BeforeToolCallEvent, [])) >= 1
-    assert len(agent.hooks._registered_callbacks.get(AfterModelCallEvent, [])) >= 1
+    # Verify hooks were registered (tool and model steering hooks)
+    assert agent.add_hook.call_count >= 2
+    agent.add_hook.assert_any_call(handler.provide_tool_steering_guidance, BeforeToolCallEvent)
 
 
 def test_steering_context_initialization():
@@ -174,7 +172,6 @@ async def steer_before_tool(self, *, agent, tool_use, **kwargs):
 
 def test_init_plugin_override():
     """Test that init_plugin can be overridden."""
-    from strands.hooks import HookRegistry
 
     class CustomHandler(SteeringHandler):
         async def steer_before_tool(self, *, agent, tool_use, **kwargs):
@@ -186,14 +183,11 @@ def init_plugin(self, agent):
 
     handler = CustomHandler()
     agent = Mock()
-    agent.hooks = HookRegistry()
-    agent.tool_registry = Mock()
 
     handler.init_plugin(agent)
 
-    # Should not register any hooks since parent init_plugin wasn't called
-    assert len(agent.hooks._registered_callbacks.get(BeforeToolCallEvent, [])) == 0
-    assert len(agent.hooks._registered_callbacks.get(AfterModelCallEvent, [])) == 0
+    # Should not register any hooks
+    assert agent.add_hook.call_count == 0
 
 
 # Integration tests with context providers
@@ -227,77 +221,68 @@ async def steer_before_tool(self, *, agent, tool_use, **kwargs):
 
 def test_handler_registers_context_provider_hooks():
     """Test that handler registers hooks from context callbacks."""
-    from strands.hooks import HookRegistry
-
     mock_callback = MockContextCallback()
     handler = TestSteeringHandlerWithProvider(context_callbacks=[mock_callback])
     agent = Mock()
-    agent.hooks = HookRegistry()
-    agent.tool_registry = Mock()
-    agent.add_hook = Mock()
 
     handler.init_plugin(agent)
 
-    # Should register 1 context callback via add_hook (steering hooks are auto-registered)
-    assert agent.add_hook.call_count >= 1
+    # Should register hooks for context callback and steering guidance
+    assert agent.add_hook.call_count >= 2
 
-    # Check that BeforeToolCallEvent was registered (either via add_hook or auto-registration)
+    # Check that BeforeToolCallEvent was registered
     call_args = [call[0] for call in agent.add_hook.call_args_list]
     event_types = [args[1] for args in call_args]
 
     # Context callback should be registered
-    assert (
-        BeforeToolCallEvent in event_types or len(agent.hooks._registered_callbacks.get(BeforeToolCallEvent, [])) >= 1
-    )
+    assert BeforeToolCallEvent in event_types
 
-
-def test_context_callbacks_receive_steering_context():
+@pytest.mark.asyncio
+async def test_context_callbacks_receive_steering_context():
     """Test that context callbacks receive the handler's steering context."""
     mock_callback = MockContextCallback()
     handler = TestSteeringHandlerWithProvider(context_callbacks=[mock_callback])
     agent = Mock()
-
+    agent.hooks = HookRegistry()
+    agent.tool_registry = Mock()
+    agent.add_hook = Mock(side_effect=lambda callback, event_type=None: agent.hooks.add_callback(event_type, callback))
     handler.init_plugin(agent)
 
-    # Get the registered callback for BeforeToolCallEvent
-    before_callback = None
-    for call in agent.add_hook.call_args_list:
-        if call[0][1] == BeforeToolCallEvent:
-            before_callback = call[0][0]
-            break
+    # Get the registered callbacks for BeforeToolCallEvent
+    callbacks = agent.hooks._registered_callbacks.get(BeforeToolCallEvent, [])
+    assert len(callbacks) > 0
 
-    assert before_callback is not None
-
-    # Create a mock event and call the callback
+    # The context callback is wrapped in a lambda, so we just call all callbacks
+    # and check if the steering context was updated
     event = Mock(spec=BeforeToolCallEvent)
     event.tool_use = {"name": "test_tool", "input": {}}
 
-    # The callback should execute without error and update the steering context
-    before_callback(event)
+    # Call all callbacks, handling both sync and async
+    for cb in callbacks:
+        try:
+            result = await cb(event)
+            if inspect.iscoroutine(result):
+                await result
+        except Exception:
+            pass  # Some callbacks might be async or have other requirements
 
-    # Verify the steering context was updated
+    # Verify the steering context was updated by at least one callback
     assert handler.steering_context.data.get("test_key") == "test_value"
 
 
 def test_multiple_context_callbacks_registered():
     """Test that multiple context callbacks are registered."""
-    from strands.hooks import HookRegistry
-
     callback1 = MockContextCallback()
     callback2 = MockContextCallback()
 
     handler = TestSteeringHandlerWithProvider(context_callbacks=[callback1, callback2])
     agent = Mock()
-    agent.hooks = HookRegistry()
-    agent.tool_registry = Mock()
-    agent.add_hook = Mock()
 
     handler.init_plugin(agent)
 
-    # Should register 2 context callbacks via add_hook, plus auto-registered @hook methods
-    assert agent.add_hook.call_count == 2  # Only context callbacks use add_hook
-    assert len(agent.hooks._registered_callbacks.get(BeforeToolCallEvent, [])) >= 1
-    assert len(agent.hooks._registered_callbacks.get(AfterModelCallEvent, [])) >= 1
+    # Should register one callback for each context provider plus tool and model steering guidance
+    expected_calls = 2 + 2  # 2 callbacks + 2 for steering guidance (tool and model)
+    assert agent.add_hook.call_count >= expected_calls
 
 
 def test_handler_initialization_with_callbacks():
@@ -509,14 +494,10 @@ async def test_default_steer_after_model_returns_proceed():
 
 def test_init_plugin_registers_model_steering():
     """Test that init_plugin registers model steering callback."""
-    from strands.hooks import HookRegistry
-
     handler = TestSteeringHandler()
     agent = Mock()
-    agent.hooks = HookRegistry()
-    agent.tool_registry = Mock()
 
     handler.init_plugin(agent)
 
-    # Verify model steering hook was auto-registered via @hook decorator
-    assert len(agent.hooks._registered_callbacks.get(AfterModelCallEvent, [])) >= 1
+    # Verify model steering hook was registered
+    agent.add_hook.assert_any_call(handler.provide_model_steering_guidance, AfterModelCallEvent)