hooks - before node call - cancel node

Author: pgrayy

src/strands/experimental/hooks/multiagent/events.py

@@ -35,11 +35,18 @@ class BeforeNodeCallEvent(BaseHookEvent):
         source: The multi-agent orchestrator instance
         node_id: ID of the node about to execute
         invocation_state: Configuration that user passes in
+        cancel_node: A user defined message that when set, will cancel the node execution with status FAILED.
+            The message will be emitted under a MultiAgentNodeCancel event. If set to `True`, Strands will cancel the
+            node using a default cancel message.
     """
 
     source: "MultiAgentBase"
     node_id: str
     invocation_state: dict[str, Any] | None = None
+    cancel_node: bool | str = False
+
+    def _can_write(self, name: str) -> bool:
+        return name in ["cancel_node"]
 
 
 @dataclass

src/strands/multiagent/graph.py

@@ -38,6 +38,7 @@
 from ..telemetry import get_tracer
 from ..types._events import (
     MultiAgentHandoffEvent,
+    MultiAgentNodeCancelEvent,
     MultiAgentNodeStartEvent,
     MultiAgentNodeStopEvent,
     MultiAgentNodeStreamEvent,
@@ -776,7 +777,6 @@ 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."""
-        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:
@@ -795,6 +795,18 @@ async def _execute_node(self, node: GraphNode, invocation_state: dict[str, Any])
 
         start_time = time.time()
         try:
+            before_event, _ = await self.hooks.invoke_callbacks_async(BeforeNodeCallEvent(self, node.node_id, invocation_state))
+
+            if before_event.cancel_node:
+                cancel_message = (
+                    before_event.cancel_node
+                    if isinstance(before_event.cancel_node, str)
+                    else "node cancelled by user"
+                )
+                logger.debug("reason=<%s> | cancelling execution", cancel_message)
+                yield MultiAgentNodeCancelEvent(node.node_id, cancel_message)
+                raise RuntimeError(cancel_message)
+
             # Build node input from satisfied dependencies
             node_input = self._build_node_input(node)
 

src/strands/multiagent/swarm.py

@@ -38,6 +38,7 @@
 from ..tools.decorator import tool
 from ..types._events import (
     MultiAgentHandoffEvent,
+    MultiAgentNodeCancelEvent,
     MultiAgentNodeStartEvent,
     MultiAgentNodeStopEvent,
     MultiAgentNodeStreamEvent,
@@ -680,9 +681,21 @@ async def _execute_swarm(self, invocation_state: dict[str, Any]) -> AsyncIterato
 
                 # TODO: Implement cancellation token to stop _execute_node from continuing
                 try:
-                    await self.hooks.invoke_callbacks_async(
+                    before_event, _ = await self.hooks.invoke_callbacks_async(
                         BeforeNodeCallEvent(self, current_node.node_id, invocation_state)
                     )
+
+                    if before_event.cancel_node:
+                        cancel_message = (
+                            before_event.cancel_node
+                            if isinstance(before_event.cancel_node, str)
+                            else "node cancelled by user"
+                        )
+                        logger.debug("reason=<%s> | cancelling execution", cancel_message)
+                        yield MultiAgentNodeCancelEvent(current_node.node_id, cancel_message)
+                        self.state.completion_status = Status.FAILED
+                        break
+
                     node_stream = self._stream_with_timeout(
                         self._execute_node(current_node, self.state.task, invocation_state),
                         self.node_timeout,
@@ -692,6 +705,13 @@ async def _execute_swarm(self, invocation_state: dict[str, Any]) -> AsyncIterato
                         yield event
 
                     self.state.node_history.append(current_node)
+
+                except Exception:
+                    logger.exception("node=<%s> | node execution failed", current_node.node_id)
+                    self.state.completion_status = Status.FAILED
+                    break
+
+                finally:
                     await self.hooks.invoke_callbacks_async(
                         AfterNodeCallEvent(self, current_node.node_id, invocation_state)
                     )
@@ -723,11 +743,6 @@ async def _execute_swarm(self, invocation_state: dict[str, Any]) -> AsyncIterato
                         self.state.completion_status = Status.COMPLETED
                         break
 
-                except Exception:
-                    logger.exception("node=<%s> | node execution failed", current_node.node_id)
-                    self.state.completion_status = Status.FAILED
-                    break
-
         except Exception:
             logger.exception("swarm execution failed")
             self.state.completion_status = Status.FAILED

src/strands/types/_events.py

@@ -524,3 +524,22 @@ def __init__(self, node_id: str, agent_event: dict[str, Any]) -> None:
                 "event": agent_event,  # Nest agent event to avoid field conflicts
             }
         )
+
+
+class MultiAgentNodeCancelEvent(TypedEvent):
+    """Event emitted when a user cancels node execution from their BeforeNodeCallEvent hook."""
+
+    def __init__(self, node_id: str, message: str) -> None:
+        """Initialize with cancel message.
+
+        Args:
+            node_id: Unique identifier for the node.
+            message: The node cancellation message.
+        """
+        super().__init__(
+            {
+                "type": "multiagent_node_cancel",
+                "node_id": node_id,
+                "message": message,
+            }
+        )