task group

Author: pgrayy

src/strands/experimental/bidi/_async/__init__.py

@@ -2,9 +2,10 @@
 
 from typing import Awaitable, Callable
 
+from ._task_group import _TaskGroup
 from ._task_pool import _TaskPool
 
-__all__ = ["_TaskPool"]
+__all__ = ["_TaskGroup", "_TaskPool"]
 
 
 async def stop_all(*funcs: Callable[..., Awaitable[None]]) -> None:
@@ -28,6 +29,6 @@ async def stop_all(*funcs: Callable[..., Awaitable[None]]) -> None:
     if exceptions:
         exceptions.append(RuntimeError("failed stop sequence"))
         for i in range(1, len(exceptions)):
-            exceptions[i].__cause__ = exceptions[i - 1]
+            exceptions[i].__context__ = exceptions[i - 1]
 
         raise exceptions[-1]

src/strands/experimental/bidi/_async/_task_group.py

@@ -0,0 +1,61 @@
+"""Manage a group of async tasks.
+
+This is intended to mimic the behaviors of asyncio.TaskGroup released in Python 3.11.
+
+- Docs: https://docs.python.org/3/library/asyncio-task.html#task-groups
+"""
+
+import asyncio
+from typing import Any, Coroutine
+
+
+class _TaskGroup:
+    """Implementation of asyncio.TaskGroup for use in Python 3.10.
+
+    Attributes:
+        _tasks: List of tasks in group.
+    """
+
+    _tasks: list[asyncio.Task]
+
+    def create_task(self, coro: Coroutine[Any, Any, Any]) -> asyncio.Task:
+        """Create an async task and add to group.
+
+        Returns:
+            The created task.
+        """
+        task = asyncio.create_task(coro)
+        self._tasks.append(task)
+        return task
+
+    async def __aenter__(self) -> "_TaskGroup":
+        """Setup self managed task group context."""
+        self._tasks = []
+        return self
+
+    async def __aexit__(self, *_: Any) -> None:
+        """Execute tasks in group.
+
+        The following execution rules are enforced:
+        - The context stops executing all tasks if at least one task raises an Exception or the context is cancelled.
+        - The context re-raises Exceptions to the caller.
+        - The context re-raises CancelledErrors to the caller only if the context itself was cancelled.
+        """
+        try:
+            await asyncio.gather(*self._tasks)
+
+        except (Exception, asyncio.CancelledError) as error:
+            for task in self._tasks:
+                task.cancel()
+
+            await asyncio.gather(*self._tasks, return_exceptions=True)
+
+            if not isinstance(error, asyncio.CancelledError):
+                raise
+
+            context_task = asyncio.current_task()
+            if context_task and context_task.cancelling() > 0:  # context itself was cancelled
+                raise
+
+        finally:
+            self._tasks = []

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

@@ -30,7 +30,7 @@
 from ....types.tools import AgentTool
 from ...hooks.events import BidiAgentInitializedEvent, BidiMessageAddedEvent
 from ...tools import ToolProvider
-from .._async import stop_all
+from .._async import _TaskGroup, stop_all
 from ..models.model import BidiModel
 from ..models.nova_sonic import BidiNovaSonicModel
 from ..types.agent import BidiAgentInput
@@ -390,22 +390,9 @@ async def run_outputs(inputs_task: asyncio.Task) -> None:
             for start in [*input_starts, *output_starts]:
                 await start(self)
 
-            inputs_task = asyncio.create_task(run_inputs())
-            outputs_task = asyncio.create_task(run_outputs(inputs_task))
-
-            try:
-                await asyncio.gather(inputs_task, outputs_task)
-            except (Exception, asyncio.CancelledError) as error:
-                inputs_task.cancel()
-                outputs_task.cancel()
-                await asyncio.gather(inputs_task, outputs_task, return_exceptions=True)
-
-                if not isinstance(error, asyncio.CancelledError):
-                    raise
-
-                run_task = asyncio.current_task()
-                if run_task and run_task.cancelling() > 0:  # externally cancelled
-                    raise
+            async with _TaskGroup() as task_group:
+                inputs_task = task_group.create_task(run_inputs())
+                task_group.create_task(run_outputs(inputs_task))
 
         finally:
             input_stops = [input_.stop for input_ in inputs if isinstance(input_, BidiInput)]

tests/strands/experimental/bidi/_async/test_task_group.py

@@ -0,0 +1,59 @@
+import asyncio
+import unittest.mock
+
+import pytest
+
+from strands.experimental.bidi._async._task_group import _TaskGroup
+
+
+@pytest.mark.asyncio
+async def test_task_group__aexit__():
+    coro = unittest.mock.AsyncMock()
+
+    async with _TaskGroup() as task_group:
+        task_group.create_task(coro())
+
+    coro.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_task_group__aexit__exception():
+    wait_event = asyncio.Event()
+    async def wait():
+        await wait_event.wait()
+
+    async def fail():
+        raise ValueError("test error")
+
+    with pytest.raises(ValueError, match="test error"):
+        async with _TaskGroup() as task_group:
+            wait_task = task_group.create_task(wait())
+            fail_task = task_group.create_task(fail())
+
+    assert wait_task.cancelled()
+    assert not fail_task.cancelled()
+
+
+@pytest.mark.asyncio
+async def test_task_group__aexit__cancelled():
+    wait_event = asyncio.Event()
+    async def wait():
+        await wait_event.wait()
+
+    tasks = []
+
+    run_event = asyncio.Event()
+    async def run():
+        async with _TaskGroup() as task_group:
+            tasks.append(task_group.create_task(wait()))
+            run_event.set()
+
+    run_task = asyncio.create_task(run())
+    await run_event.wait()
+    run_task.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await run_task
+
+    wait_task = tasks[0]
+    assert wait_task.cancelled()