cancellation - agent loop (strands-agents#63)

Author: pgrayy

src/strands/experimental/bidi/agent/agent.py

@@ -20,6 +20,7 @@
 from .... import _identifier
 from ....agent.state import AgentState
 from ....hooks import HookProvider, HookRegistry
+from ....interrupt import _InterruptState
 from ....tools.caller import _ToolCaller
 from ....tools.executors import ConcurrentToolExecutor
 from ....tools.executors._executor import ToolExecutor
@@ -146,6 +147,9 @@ def __init__(
         # Emit initialization event
         self.hooks.invoke_callbacks(BidiAgentInitializedEvent(agent=self))
 
+        # TODO: Determine if full support is required
+        self._interrupt_state = _InterruptState()
+
         self._started = False
 
     @property
@@ -270,6 +274,10 @@ async def start(self, invocation_state: dict[str, Any] | None = None) -> None:
                 This allows passing custom data (user_id, session_id, database connections, etc.)
                 that tools can access via their invocation_state parameter.
 
+        Raises:
+            RuntimeError:
+                If agent already started.
+
         Example:
             ```python
             await agent.start(invocation_state={

src/strands/experimental/bidi/agent/loop.py

@@ -5,7 +5,7 @@
 
 import asyncio
 import logging
-from typing import TYPE_CHECKING, Any, AsyncIterable, Awaitable
+from typing import TYPE_CHECKING, Any, AsyncIterable
 
 from ....types._events import ToolInterruptEvent, ToolResultEvent, ToolResultMessageEvent, ToolUseStreamEvent
 from ....types.content import Message
@@ -18,6 +18,7 @@
 from ...hooks.events import (
     BidiInterruptionEvent as BidiInterruptionHookEvent,
 )
+from .._async import _TaskPool, stop_all
 from ..types.events import BidiInterruptionEvent, BidiOutputEvent, BidiTranscriptStreamEvent
 
 if TYPE_CHECKING:
@@ -31,21 +32,14 @@ class _BidiAgentLoop:
 
     Attributes:
         _agent: BidiAgent instance to loop.
+        _started: Flag if agent loop has started.
+        _task_pool: Track active async tasks created in loop.
         _event_queue: Queue model and tool call events for receiver.
-        _stop_event: Sentinel to mark end of loop.
-        _tasks: Track active async tasks created in loop.
-        _active: Flag if agent loop is started.
         _invocation_state: Optional context to pass to tools during execution.
             This allows passing custom data (user_id, session_id, database connections, etc.)
             that tools can access via their invocation_state parameter.
     """
 
-    _event_queue: asyncio.Queue
-    _stop_event: object
-    _tasks: set
-    _active: bool
-    _invocation_state: dict[str, Any]
-
     def __init__(self, agent: "BidiAgent") -> None:
         """Initialize members of the agent loop.
 
@@ -55,8 +49,10 @@ def __init__(self, agent: "BidiAgent") -> None:
             agent: Bidirectional agent to loop over.
         """
         self._agent = agent
-        self._active = False
-        self._invocation_state = {}
+        self._started = False
+        self._task_pool = _TaskPool()
+        self._event_queue: asyncio.Queue
+        self._invocation_state: dict[str, Any]
 
     async def start(self, invocation_state: dict[str, Any] | None = None) -> None:
         """Start the agent loop.
@@ -67,19 +63,15 @@ async def start(self, invocation_state: dict[str, Any] | None = None) -> None:
             invocation_state: Optional context to pass to tools during execution.
                 This allows passing custom data (user_id, session_id, database connections, etc.)
                 that tools can access via their invocation_state parameter.
+
+        Raises:
+            RuntimeError:
+                If loop already started.
         """
-        if self.active:
-            return
+        if self._started:
+            raise RuntimeError("loop already started | call stop before starting again")
 
         logger.debug("agent loop starting")
-
-        self._invocation_state = invocation_state or {}
-
-        self._event_queue = asyncio.Queue(maxsize=1)
-        self._stop_event = object()
-        self._tasks = set()
-
-        # Emit before invocation event
         await self._agent.hooks.invoke_callbacks_async(BidiBeforeInvocationEvent(agent=self._agent))
 
         await self._agent.model.start(
@@ -88,99 +80,88 @@ async def start(self, invocation_state: dict[str, Any] | None = None) -> None:
             messages=self._agent.messages,
         )
 
-        self._create_task(self._run_model())
+        self._event_queue = asyncio.Queue(maxsize=1)
+
+        self._task_pool = _TaskPool()
+        self._task_pool.create(self._run_model())
 
-        self._active = True
+        self._invocation_state = invocation_state or {}
+        self._started = True
 
     async def stop(self) -> None:
         """Stop the agent loop."""
-        if not self.active:
-            return
-
         logger.debug("agent loop stopping")
 
+        self._started = False
         self._invocation_state = {}
 
-        try:
-            # Cancel all tasks
-            for task in self._tasks:
-                task.cancel()
-
-            # Wait briefly for tasks to finish their current operations
-            await asyncio.gather(*self._tasks, return_exceptions=True)
+        async def stop_tasks() -> None:
+            await self._task_pool.cancel()
 
-            # Stop the model
+        async def stop_model() -> None:
             await self._agent.model.stop()
 
-            # Clean up the event queue
-            if not self._event_queue.empty():
-                self._event_queue.get_nowait()
-            self._event_queue.put_nowait(self._stop_event)
-
-            self._active = False
-
+        try:
+            await stop_all(stop_tasks, stop_model)
         finally:
-            # Emit after invocation event (reverse order for cleanup)
             await self._agent.hooks.invoke_callbacks_async(BidiAfterInvocationEvent(agent=self._agent))
 
     async def receive(self) -> AsyncIterable[BidiOutputEvent]:
-        """Receive model and tool call events."""
+        """Receive model and tool call events.
+
+        Raises:
+            RuntimeError: If start has not been called.
+        """
+        if not self._started:
+            raise RuntimeError("loop not started | call start before receiving")
+
         while True:
             event = await self._event_queue.get()
-            if event is self._stop_event:
-                break
+            if isinstance(event, Exception):
+                raise event
 
             yield event
 
-    @property
-    def active(self) -> bool:
-        """True if agent loop started, False otherwise."""
-        return self._active
-
-    def _create_task(self, coro: Awaitable[None]) -> None:
-        """Utilitly to create async task.
-
-        Adds a clean up callback to run after task completes.
-        """
-        task: asyncio.Task[None] = asyncio.create_task(coro)  # type: ignore
-        task.add_done_callback(lambda task: self._tasks.remove(task))
-
-        self._tasks.add(task)
-
     async def _run_model(self) -> None:
         """Task for running the model.
 
         Events are streamed through the event queue.
         """
         logger.debug("model task starting")
 
-        async for event in self._agent.model.receive():  # type: ignore
-            await self._event_queue.put(event)
-
-            if isinstance(event, BidiTranscriptStreamEvent):
-                if event["is_final"]:
-                    message: Message = {"role": event["role"], "content": [{"text": event["text"]}]}
-                    self._agent.messages.append(message)
+        try:
+            async for event in self._agent.model.receive():  # type: ignore
+                await self._event_queue.put(event)
+
+                if isinstance(event, BidiTranscriptStreamEvent):
+                    if event["is_final"]:
+                        message: Message = {"role": event["role"], "content": [{"text": event["text"]}]}
+                        self._agent.messages.append(message)
+                        await self._agent.hooks.invoke_callbacks_async(
+                            BidiMessageAddedEvent(agent=self._agent, message=message)
+                        )
+
+                elif isinstance(event, ToolUseStreamEvent):
+                    tool_use = event["current_tool_use"]
+                    self._task_pool.create(self._run_tool(tool_use))
+
+                    tool_message: Message = {"role": "assistant", "content": [{"toolUse": tool_use}]}
+                    self._agent.messages.append(tool_message)
                     await self._agent.hooks.invoke_callbacks_async(
-                        BidiMessageAddedEvent(agent=self._agent, message=message)
+                        BidiMessageAddedEvent(agent=self._agent, message=tool_message)
                     )
 
-            elif isinstance(event, ToolUseStreamEvent):
-                tool_use = event["current_tool_use"]
-                self._create_task(self._run_tool(tool_use))
-
-                tool_message: Message = {"role": "assistant", "content": [{"toolUse": tool_use}]}
-                self._agent.messages.append(tool_message)
-
-            elif isinstance(event, BidiInterruptionEvent):
-                # Emit interruption hook event
-                await self._agent.hooks.invoke_callbacks_async(
-                    BidiInterruptionHookEvent(
-                        agent=self._agent,
-                        reason=event["reason"],
-                        interrupted_response_id=event.get("interrupted_response_id"),
+                elif isinstance(event, BidiInterruptionEvent):
+                    await self._agent.hooks.invoke_callbacks_async(
+                        BidiInterruptionHookEvent(
+                            agent=self._agent,
+                            reason=event["reason"],
+                            interrupted_response_id=event.get("interrupted_response_id"),
+                        )
                     )
-                )
+
+        except Exception as error:
+            await self._event_queue.put(error)
 
     async def _run_tool(self, tool_use: ToolUse) -> None:
         """Task for running tool requested by the model using the tool executor."""
@@ -196,30 +177,34 @@ async def _run_tool(self, tool_use: ToolUse) -> None:
             "system_prompt": self._agent.system_prompt,
         }
 
-        tool_events = self._agent.tool_executor._stream(
-            self._agent,
-            tool_use,
-            tool_results,
-            invocation_state,
-            structured_output_context=None,
-        )
-
-        async for event in tool_events:
-            if isinstance(event, ToolInterruptEvent):
-                raise RuntimeError(
-                    "Tool interruption is not yet supported in BidiAgent. "
-                    "ToolInterruptEvent received but cannot be handled in bidirectional streaming context."
-                )
-            await self._event_queue.put(event)
-            if isinstance(event, ToolResultEvent):
-                result = event.tool_result
-
-        await self._agent.model.send(ToolResultEvent(result))
-
-        message: Message = {
-            "role": "user",
-            "content": [{"toolResult": result}],
-        }
-        self._agent.messages.append(message)
-        await self._agent.hooks.invoke_callbacks_async(BidiMessageAddedEvent(agent=self._agent, message=message))
-        await self._event_queue.put(ToolResultMessageEvent(message))
+        try:
+            tool_events = self._agent.tool_executor._stream(
+                self._agent,
+                tool_use,
+                tool_results,
+                invocation_state,
+                structured_output_context=None,
+            )
+
+            async for event in tool_events:
+                if isinstance(event, ToolInterruptEvent):
+                    self._agent._interrupt_state.deactivate()
+                    interrupt_names = [interrupt.name for interrupt in event.interrupts]
+                    raise RuntimeError(f"interrupts={interrupt_names} | tool interrupts are not supported in bidi")
+
+                await self._event_queue.put(event)
+                if isinstance(event, ToolResultEvent):
+                    result = event.tool_result
+
+            await self._agent.model.send(ToolResultEvent(result))
+
+            message: Message = {
+                "role": "user",
+                "content": [{"toolResult": result}],
+            }
+            self._agent.messages.append(message)
+            await self._agent.hooks.invoke_callbacks_async(BidiMessageAddedEvent(agent=self._agent, message=message))
+            await self._event_queue.put(ToolResultMessageEvent(message))
+
+        except Exception as error:
+            await self._event_queue.put(error)