feat: add CancellationToken for graceful agent execution cancellation (#1772)

Author: jgoyani1

src/strands/agent/agent.py

@@ -240,6 +240,9 @@ def __init__(
         self.record_direct_tool_call = record_direct_tool_call
         self.load_tools_from_directory = load_tools_from_directory
 
+        # Create internal cancel signal for graceful cancellation using threading.Event
+        self._cancel_signal = threading.Event()
+
         self.tool_registry = ToolRegistry()
 
         # Process tool list if provided
@@ -327,6 +330,37 @@ def __init__(
 
         self.hooks.invoke_callbacks(AgentInitializedEvent(agent=self))
 
+    def cancel(self) -> None:
+        """Cancel the currently running agent invocation.
+
+        This method is thread-safe and can be called from any context
+        (e.g., another thread, web request handler, background task).
+
+        The agent will stop gracefully at the next checkpoint:
+        - During model response streaming
+        - Before tool execution
+
+        The agent will return a result with stop_reason="cancelled".
+
+        Example:
+            ```python
+            agent = Agent(model=model)
+
+            # Start agent in background
+            task = asyncio.create_task(agent.invoke_async("Hello"))
+
+            # Cancel from another context
+            agent.cancel()
+
+            result = await task
+            assert result.stop_reason == "cancelled"
+            ```
+
+        Note:
+            Multiple calls to cancel() are safe and idempotent.
+        """
+        self._cancel_signal.set()
+
     @property
     def system_prompt(self) -> str | None:
         """Get the system prompt as a string for backwards compatibility.
@@ -756,6 +790,9 @@ async def stream_async(
                     raise
 
         finally:
+            # Clear cancel signal to allow agent reuse after cancellation
+            self._cancel_signal.clear()
+
             if self._invocation_lock.locked():
                 self._invocation_lock.release()
 

src/strands/event_loop/event_loop.py

@@ -336,6 +336,7 @@ async def _handle_model_execution(
                     system_prompt_content=agent._system_prompt_content,
                     tool_choice=structured_output_context.tool_choice,
                     invocation_state=invocation_state,
+                    cancel_signal=agent._cancel_signal,
                 ):
                     yield event
 
@@ -465,6 +466,47 @@ async def _handle_tool_execution(
         tool_uses = [tool_use for tool_use in tool_uses if tool_use["toolUseId"] not in tool_use_ids]
 
     interrupts = []
+
+    # Check for cancellation before tool execution
+    # Add tool_result for each tool_use to maintain valid conversation state
+    if agent._cancel_signal.is_set():
+        logger.debug("tool_count=<%d> | cancellation detected before tool execution", len(tool_uses))
+
+        # Create cancellation tool_result for each tool_use to avoid invalid message state
+        # (tool_use without tool_result would be rejected on next invocation)
+        for tool_use in tool_uses:
+            cancel_result: ToolResult = {
+                "toolUseId": str(tool_use.get("toolUseId")),
+                "status": "error",
+                "content": [{"text": "Tool execution cancelled"}],
+            }
+            tool_results.append(cancel_result)
+
+        # Add tool results message to conversation if any tools were cancelled
+        cancelled_tool_result_message: Message | None = None
+        if tool_results:
+            _cancelled_msg: Message = {
+                "role": "user",
+                "content": [{"toolResult": result} for result in tool_results],
+            }
+            cancelled_tool_result_message = _cancelled_msg
+            agent.messages.append(_cancelled_msg)
+            await agent.hooks.invoke_callbacks_async(MessageAddedEvent(agent=agent, message=_cancelled_msg))
+            yield ToolResultMessageEvent(message=_cancelled_msg)
+
+        agent.event_loop_metrics.end_cycle(cycle_start_time, cycle_trace)
+        yield EventLoopStopEvent(
+            "cancelled",
+            message,
+            agent.event_loop_metrics,
+            invocation_state["request_state"],
+        )
+        if cycle_span:
+            tracer.end_event_loop_cycle_span(
+                span=cycle_span, message=message, tool_result_message=cancelled_tool_result_message
+            )
+        return
+
     tool_events = agent.tool_executor._execute(
         agent, tool_uses, tool_results, cycle_trace, cycle_span, invocation_state, structured_output_context
     )

src/strands/event_loop/streaming.py

@@ -2,6 +2,7 @@
 
 import json
 import logging
+import threading
 import time
 import warnings
 from collections.abc import AsyncGenerator, AsyncIterable
@@ -368,13 +369,16 @@ def extract_usage_metrics(event: MetadataEvent, time_to_first_byte_ms: int | Non
 
 
 async def process_stream(
-    chunks: AsyncIterable[StreamEvent], start_time: float | None = None
+    chunks: AsyncIterable[StreamEvent],
+    start_time: float | None = None,
+    cancel_signal: threading.Event | None = None,
 ) -> AsyncGenerator[TypedEvent, None]:
     """Processes the response stream from the API, constructing the final message and extracting usage metrics.
 
     Args:
         chunks: The chunks of the response stream from the model.
         start_time: Time when the model request is initiated
+        cancel_signal: Optional threading.Event to check for cancellation during streaming.
 
     Yields:
         The reason for stopping, the constructed message, and the usage metrics.
@@ -395,6 +399,19 @@ async def process_stream(
     metrics: Metrics = Metrics(latencyMs=0, timeToFirstByteMs=0)
 
     async for chunk in chunks:
+        # Check for cancellation during stream processing
+        if cancel_signal and cancel_signal.is_set():
+            logger.debug("cancellation detected during stream processing")
+            # Return cancelled stop reason with cancellation message
+            # The incomplete message in state["message"] is discarded and never added to agent.messages
+            yield ModelStopReason(
+                stop_reason="cancelled",
+                message={"role": "assistant", "content": [{"text": "Cancelled by user"}]},
+                usage=usage,
+                metrics=metrics,
+            )
+            return
+
         # Track first byte time when we get first content
         if first_byte_time is None and ("contentBlockDelta" in chunk or "contentBlockStart" in chunk):
             first_byte_time = time.time()
@@ -431,6 +448,7 @@ async def stream_messages(
     tool_choice: Any | None = None,
     system_prompt_content: list[SystemContentBlock] | None = None,
     invocation_state: dict[str, Any] | None = None,
+    cancel_signal: threading.Event | None = None,
     **kwargs: Any,
 ) -> AsyncGenerator[TypedEvent, None]:
     """Streams messages to the model and processes the response.
@@ -444,6 +462,7 @@ async def stream_messages(
         system_prompt_content: The authoritative system prompt content blocks that always contains the
             system prompt data.
         invocation_state: Caller-provided state/context that was passed to the agent when it was invoked.
+        cancel_signal: Optional threading.Event to check for cancellation during streaming.
         **kwargs: Additional keyword arguments for future extensibility.
 
     Yields:
@@ -463,5 +482,5 @@ async def stream_messages(
         invocation_state=invocation_state,
     )
 
-    async for event in process_stream(chunks, start_time):
+    async for event in process_stream(chunks, start_time, cancel_signal):
         yield event

src/strands/session/repository_session_manager.py

@@ -124,8 +124,8 @@ def sync_agent(self, agent: "Agent", **kwargs: Any) -> None:
         else:
             state_changed = current_state_version != last_synced.get("state_version")
             internal_state_changed = current_interrupt_state_version != last_synced.get("interrupt_state_version")
-            conversation_manager_state_changed = (
-                current_conversation_manager_state != last_synced.get("conversation_manager_state")
+            conversation_manager_state_changed = current_conversation_manager_state != last_synced.get(
+                "conversation_manager_state"
             )
 
         if not state_changed and not internal_state_changed and not conversation_manager_state_changed:

src/strands/types/event_loop.py

@@ -37,6 +37,7 @@ class Metrics(TypedDict, total=False):
 
 
 StopReason = Literal[
+    "cancelled",
     "content_filtered",
     "end_turn",
     "guardrail_intervened",
@@ -47,6 +48,7 @@ class Metrics(TypedDict, total=False):
 ]
 """Reason for the model ending its response generation.
 
+- "cancelled": Agent execution was cancelled via agent.cancel()
 - "content_filtered": Content was filtered due to policy violation
 - "end_turn": Normal completion of the response
 - "guardrail_intervened": Guardrail system intervened