Merge branch 'main' of https://github.com/strands-agents/sdk-python into interrupt_state

Author: pgrayy

src/strands/_async.py

@@ -1,6 +1,7 @@
 """Private async execution utilities."""
 
 import asyncio
+import contextvars
 from concurrent.futures import ThreadPoolExecutor
 from typing import Awaitable, Callable, TypeVar
 
@@ -27,5 +28,6 @@ def execute() -> T:
         return asyncio.run(execute_async())
 
     with ThreadPoolExecutor() as executor:
-        future = executor.submit(execute)
+        context = contextvars.copy_context()
+        future = executor.submit(context.run, execute)
         return future.result()

src/strands/agent/agent.py

@@ -170,22 +170,21 @@ async def acall() -> ToolResult:
                             self._agent._interrupt_state.deactivate()
                             raise RuntimeError("cannot raise interrupt in direct tool call")
 
-                    return tool_results[0]
+                    tool_result = tool_results[0]
 
-                tool_result = run_async(acall)
+                    if record_direct_tool_call is not None:
+                        should_record_direct_tool_call = record_direct_tool_call
+                    else:
+                        should_record_direct_tool_call = self._agent.record_direct_tool_call
 
-                if record_direct_tool_call is not None:
-                    should_record_direct_tool_call = record_direct_tool_call
-                else:
-                    should_record_direct_tool_call = self._agent.record_direct_tool_call
+                    if should_record_direct_tool_call:
+                        # Create a record of this tool execution in the message history
+                        await self._agent._record_tool_execution(tool_use, tool_result, user_message_override)
 
-                if should_record_direct_tool_call:
-                    # Create a record of this tool execution in the message history
-                    self._agent._record_tool_execution(tool_use, tool_result, user_message_override)
+                    return tool_result
 
-                # Apply window management
+                tool_result = run_async(acall)
                 self._agent.conversation_manager.apply_management(self._agent)
-
                 return tool_result
 
             return caller
@@ -533,15 +532,15 @@ async def structured_output_async(self, output_model: Type[T], prompt: AgentInpu
             category=DeprecationWarning,
             stacklevel=2,
         )
-        self.hooks.invoke_callbacks(BeforeInvocationEvent(agent=self))
+        await self.hooks.invoke_callbacks_async(BeforeInvocationEvent(agent=self))
         with self.tracer.tracer.start_as_current_span(
             "execute_structured_output", kind=trace_api.SpanKind.CLIENT
         ) as structured_output_span:
             try:
                 if not self.messages and not prompt:
                     raise ValueError("No conversation history or prompt provided")
 
-                temp_messages: Messages = self.messages + self._convert_prompt_to_messages(prompt)
+                temp_messages: Messages = self.messages + await self._convert_prompt_to_messages(prompt)
 
                 structured_output_span.set_attributes(
                     {
@@ -574,7 +573,7 @@ async def structured_output_async(self, output_model: Type[T], prompt: AgentInpu
                 return event["output"]
 
             finally:
-                self.hooks.invoke_callbacks(AfterInvocationEvent(agent=self))
+                await self.hooks.invoke_callbacks_async(AfterInvocationEvent(agent=self))
 
     def cleanup(self) -> None:
         """Clean up resources used by the agent.
@@ -657,7 +656,7 @@ async def stream_async(
             callback_handler = kwargs.get("callback_handler", self.callback_handler)
 
         # Process input and get message to add (if any)
-        messages = self._convert_prompt_to_messages(prompt)
+        messages = await self._convert_prompt_to_messages(prompt)
 
         self.trace_span = self._start_agent_trace_span(messages)
 
@@ -699,13 +698,13 @@ async def _run_loop(
         Yields:
             Events from the event loop cycle.
         """
-        self.hooks.invoke_callbacks(BeforeInvocationEvent(agent=self))
+        await self.hooks.invoke_callbacks_async(BeforeInvocationEvent(agent=self))
 
         try:
             yield InitEventLoopEvent()
 
             for message in messages:
-                self._append_message(message)
+                await self._append_message(message)
 
             structured_output_context = StructuredOutputContext(
                 structured_output_model or self._default_structured_output_model
@@ -731,7 +730,7 @@ async def _run_loop(
 
         finally:
             self.conversation_manager.apply_management(self)
-            self.hooks.invoke_callbacks(AfterInvocationEvent(agent=self))
+            await self.hooks.invoke_callbacks_async(AfterInvocationEvent(agent=self))
 
     async def _execute_event_loop_cycle(
         self, invocation_state: dict[str, Any], structured_output_context: StructuredOutputContext | None = None
@@ -780,7 +779,7 @@ async def _execute_event_loop_cycle(
             if structured_output_context:
                 structured_output_context.cleanup(self.tool_registry)
 
-    def _convert_prompt_to_messages(self, prompt: AgentInput) -> Messages:
+    async def _convert_prompt_to_messages(self, prompt: AgentInput) -> Messages:
         if self._interrupt_state.activated:
             return []
 
@@ -795,7 +794,7 @@ def _convert_prompt_to_messages(self, prompt: AgentInput) -> Messages:
                 tool_use_ids = [
                     content["toolUse"]["toolUseId"] for content in self.messages[-1]["content"] if "toolUse" in content
                 ]
-                self._append_message(
+                await self._append_message(
                     {
                         "role": "user",
                         "content": generate_missing_tool_result_content(tool_use_ids),
@@ -826,7 +825,7 @@ def _convert_prompt_to_messages(self, prompt: AgentInput) -> Messages:
             raise ValueError("Input prompt must be of type: `str | list[Contentblock] | Messages | None`.")
         return messages
 
-    def _record_tool_execution(
+    async def _record_tool_execution(
         self,
         tool: ToolUse,
         tool_result: ToolResult,
@@ -886,10 +885,10 @@ def _record_tool_execution(
         }
 
         # Add to message history
-        self._append_message(user_msg)
-        self._append_message(tool_use_msg)
-        self._append_message(tool_result_msg)
-        self._append_message(assistant_msg)
+        await self._append_message(user_msg)
+        await self._append_message(tool_use_msg)
+        await self._append_message(tool_result_msg)
+        await self._append_message(assistant_msg)
 
     def _start_agent_trace_span(self, messages: Messages) -> trace_api.Span:
         """Starts a trace span for the agent.
@@ -975,10 +974,10 @@ def _initialize_system_prompt(
         else:
             return None, None
 
-    def _append_message(self, message: Message) -> None:
+    async def _append_message(self, message: Message) -> None:
         """Appends a message to the agent's list of messages and invokes the callbacks for the MessageCreatedEvent."""
         self.messages.append(message)
-        self.hooks.invoke_callbacks(MessageAddedEvent(agent=self, message=message))
+        await self.hooks.invoke_callbacks_async(MessageAddedEvent(agent=self, message=message))
 
     def _redact_user_content(self, content: list[ContentBlock], redact_message: str) -> list[ContentBlock]:
         """Redact user content preserving toolResult blocks.

src/strands/event_loop/event_loop.py

@@ -227,7 +227,7 @@ async def event_loop_cycle(
             )
         structured_output_context.set_forced_mode()
         logger.debug("Forcing structured output tool")
-        agent._append_message(
+        await agent._append_message(
             {"role": "user", "content": [{"text": "You must format the previous response as structured output."}]}
         )
 
@@ -322,7 +322,7 @@ async def _handle_model_execution(
             model_id=model_id,
         )
         with trace_api.use_span(model_invoke_span):
-            agent.hooks.invoke_callbacks(
+            await agent.hooks.invoke_callbacks_async(
                 BeforeModelCallEvent(
                     agent=agent,
                 )
@@ -347,7 +347,7 @@ async def _handle_model_execution(
                 stop_reason, message, usage, metrics = event["stop"]
                 invocation_state.setdefault("request_state", {})
 
-                agent.hooks.invoke_callbacks(
+                await agent.hooks.invoke_callbacks_async(
                     AfterModelCallEvent(
                         agent=agent,
                         stop_response=AfterModelCallEvent.ModelStopResponse(
@@ -368,7 +368,7 @@ async def _handle_model_execution(
                 if model_invoke_span:
                     tracer.end_span_with_error(model_invoke_span, str(e), e)
 
-                agent.hooks.invoke_callbacks(
+                await agent.hooks.invoke_callbacks_async(
                     AfterModelCallEvent(
                         agent=agent,
                         exception=e,
@@ -402,7 +402,7 @@ async def _handle_model_execution(
 
         # Add the response message to the conversation
         agent.messages.append(message)
-        agent.hooks.invoke_callbacks(MessageAddedEvent(agent=agent, message=message))
+        await agent.hooks.invoke_callbacks_async(MessageAddedEvent(agent=agent, message=message))
 
         # Update metrics
         agent.event_loop_metrics.update_usage(usage)
@@ -507,7 +507,7 @@ async def _handle_tool_execution(
     }
 
     agent.messages.append(tool_result_message)
-    agent.hooks.invoke_callbacks(MessageAddedEvent(agent=agent, message=tool_result_message))
+    await agent.hooks.invoke_callbacks_async(MessageAddedEvent(agent=agent, message=tool_result_message))
 
     yield ToolResultMessageEvent(message=tool_result_message)
 

src/strands/hooks/registry.py

@@ -7,9 +7,10 @@
 via hook provider objects.
 """
 
+import inspect
 import logging
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Generator, Generic, Protocol, Type, TypeVar
+from typing import TYPE_CHECKING, Any, Awaitable, Generator, Generic, Protocol, Type, TypeVar
 
 from ..interrupt import Interrupt, InterruptException
 
@@ -122,10 +123,15 @@ class HookCallback(Protocol, Generic[TEvent]):
         ```python
         def my_callback(event: StartRequestEvent) -> None:
             print(f"Request started for agent: {event.agent.name}")
+
+        # Or
+
+        async def my_callback(event: StartRequestEvent) -> None:
+            # await an async operation
         ```
     """
 
-    def __call__(self, event: TEvent) -> None:
+    def __call__(self, event: TEvent) -> None | Awaitable[None]:
         """Handle a hook event.
 
         Args:
@@ -164,6 +170,10 @@ def my_handler(event: StartRequestEvent):
             registry.add_callback(StartRequestEvent, my_handler)
             ```
         """
+        # Related issue: https://github.com/strands-agents/sdk-python/issues/330
+        if event_type.__name__ == "AgentInitializedEvent" and inspect.iscoroutinefunction(callback):
+            raise ValueError("AgentInitializedEvent can only be registered with a synchronous callback")
+
         callbacks = self._registered_callbacks.setdefault(event_type, [])
         callbacks.append(callback)
 
@@ -189,6 +199,52 @@ def register_hooks(self, registry: HookRegistry):
         """
         hook.register_hooks(self)
 
+    async def invoke_callbacks_async(self, event: TInvokeEvent) -> tuple[TInvokeEvent, list[Interrupt]]:
+        """Invoke all registered callbacks for the given event.
+
+        This method finds all callbacks registered for the event's type and
+        invokes them in the appropriate order. For events with should_reverse_callbacks=True,
+        callbacks are invoked in reverse registration order. Any exceptions raised by callback
+        functions will propagate to the caller.
+
+        Additionally, this method aggregates interrupts raised by the user to instantiate human-in-the-loop workflows.
+
+        Args:
+            event: The event to dispatch to registered callbacks.
+
+        Returns:
+            The event dispatched to registered callbacks and any interrupts raised by the user.
+
+        Raises:
+            ValueError: If interrupt name is used more than once.
+
+        Example:
+            ```python
+            event = StartRequestEvent(agent=my_agent)
+            await registry.invoke_callbacks_async(event)
+            ```
+        """
+        interrupts: dict[str, Interrupt] = {}
+
+        for callback in self.get_callbacks_for(event):
+            try:
+                if inspect.iscoroutinefunction(callback):
+                    await callback(event)
+                else:
+                    callback(event)
+
+            except InterruptException as exception:
+                interrupt = exception.interrupt
+                if interrupt.name in interrupts:
+                    message = f"interrupt_name=<{interrupt.name}> | interrupt name used more than once"
+                    logger.error(message)
+                    raise ValueError(message) from exception
+
+                # Each callback is allowed to raise their own interrupt.
+                interrupts[interrupt.name] = interrupt
+
+        return event, list(interrupts.values())
+
     def invoke_callbacks(self, event: TInvokeEvent) -> tuple[TInvokeEvent, list[Interrupt]]:
         """Invoke all registered callbacks for the given event.
 
@@ -206,6 +262,7 @@ def invoke_callbacks(self, event: TInvokeEvent) -> tuple[TInvokeEvent, list[Inte
             The event dispatched to registered callbacks and any interrupts raised by the user.
 
         Raises:
+            RuntimeError: If at least one callback is async.
             ValueError: If interrupt name is used more than once.
 
         Example:
@@ -214,9 +271,13 @@ def invoke_callbacks(self, event: TInvokeEvent) -> tuple[TInvokeEvent, list[Inte
             registry.invoke_callbacks(event)
             ```
         """
+        callbacks = list(self.get_callbacks_for(event))
         interrupts: dict[str, Interrupt] = {}
 
-        for callback in self.get_callbacks_for(event):
+        if any(inspect.iscoroutinefunction(callback) for callback in callbacks):
+            raise RuntimeError(f"event=<{event}> | use invoke_callbacks_async to invoke async callback")
+
+        for callback in callbacks:
             try:
                 callback(event)
             except InterruptException as exception:

src/strands/multiagent/graph.py

@@ -453,7 +453,7 @@ def __init__(
         self._resume_from_session = False
         self.id = id
 
-        self.hooks.invoke_callbacks(MultiAgentInitializedEvent(self))
+        run_async(lambda: self.hooks.invoke_callbacks_async(MultiAgentInitializedEvent(self)))
 
     def __call__(
         self, task: str | list[ContentBlock], invocation_state: dict[str, Any] | None = None, **kwargs: Any
@@ -516,7 +516,7 @@ async def stream_async(
         if invocation_state is None:
             invocation_state = {}
 
-        self.hooks.invoke_callbacks(BeforeMultiAgentInvocationEvent(self, invocation_state))
+        await self.hooks.invoke_callbacks_async(BeforeMultiAgentInvocationEvent(self, invocation_state))
 
         logger.debug("task=<%s> | starting graph execution", task)
 
@@ -569,7 +569,7 @@ async def stream_async(
                 raise
             finally:
                 self.state.execution_time = round((time.time() - start_time) * 1000)
-                self.hooks.invoke_callbacks(AfterMultiAgentInvocationEvent(self))
+                await self.hooks.invoke_callbacks_async(AfterMultiAgentInvocationEvent(self))
                 self._resume_from_session = False
                 self._resume_next_nodes.clear()
 
@@ -776,7 +776,7 @@ def _is_node_ready_with_conditions(self, node: GraphNode, completed_batch: list[
 
     async def _execute_node(self, node: GraphNode, invocation_state: dict[str, Any]) -> AsyncIterator[Any]:
         """Execute a single node and yield TypedEvent objects."""
-        self.hooks.invoke_callbacks(BeforeNodeCallEvent(self, node.node_id, invocation_state))
+        await self.hooks.invoke_callbacks_async(BeforeNodeCallEvent(self, node.node_id, invocation_state))
 
         # Reset the node's state if reset_on_revisit is enabled, and it's being revisited
         if self.reset_on_revisit and node in self.state.completed_nodes:
@@ -920,7 +920,7 @@ async def _execute_node(self, node: GraphNode, invocation_state: dict[str, Any])
             raise
 
         finally:
-            self.hooks.invoke_callbacks(AfterNodeCallEvent(self, node.node_id, invocation_state))
+            await self.hooks.invoke_callbacks_async(AfterNodeCallEvent(self, node.node_id, invocation_state))
 
     def _accumulate_metrics(self, node_result: NodeResult) -> None:
         """Accumulate metrics from a node result."""