fix: throw exceptions from ConcurrentToolExecutor (#1797)

Co-authored-by: Patrick Gray <pgrayy@amazon.com>

Author: charles-dyfis-net

src/strands/tools/executors/concurrent.py

@@ -48,34 +48,43 @@ async def _execute(
         task_events = [asyncio.Event() for _ in tool_uses]
         stop_event = object()
 
-        tasks = [
-            asyncio.create_task(
-                self._task(
-                    agent,
-                    tool_use,
-                    tool_results,
-                    cycle_trace,
-                    cycle_span,
-                    invocation_state,
-                    task_id,
-                    task_queue,
-                    task_events[task_id],
-                    stop_event,
-                    structured_output_context,
+        tasks = []
+        try:
+            for task_id, tool_use in enumerate(tool_uses):
+                tasks.append(
+                    asyncio.create_task(
+                        self._task(
+                            agent,
+                            tool_use,
+                            tool_results,
+                            cycle_trace,
+                            cycle_span,
+                            invocation_state,
+                            task_id,
+                            task_queue,
+                            task_events[task_id],
+                            stop_event,
+                            structured_output_context,
+                        )
+                    )
                 )
-            )
-            for task_id, tool_use in enumerate(tool_uses)
-        ]
 
-        task_count = len(tasks)
-        while task_count:
-            task_id, event = await task_queue.get()
-            if event is stop_event:
-                task_count -= 1
-                continue
+            task_count = len(tasks)
+            while task_count:
+                task_id, event = await task_queue.get()
+                if event is stop_event:
+                    task_count -= 1
+                    continue
 
-            yield event
-            task_events[task_id].set()
+                if isinstance(event, Exception):
+                    raise event
+
+                yield event
+                task_events[task_id].set()
+        finally:
+            for task in tasks:
+                task.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
 
     async def _task(
         self,
@@ -115,5 +124,8 @@ async def _task(
                 await task_event.wait()
                 task_event.clear()
 
+        except Exception as e:
+            task_queue.put_nowait((task_id, e))
+
         finally:
             task_queue.put_nowait((task_id, stop_event))

tests/strands/tools/executors/conftest.py

@@ -1,3 +1,4 @@
+import asyncio
 import threading
 import unittest.mock
 
@@ -90,13 +91,24 @@ def func(tool_context: ToolContext) -> str:
 
 
 @pytest.fixture
-def tool_registry(weather_tool, temperature_tool, exception_tool, thread_tool, interrupt_tool):
+def slow_tool():
+    @strands.tool(name="slow_tool")
+    async def func():
+        """A tool that blocks until cancelled."""
+        await asyncio.sleep(3)
+
+    return func
+
+
+@pytest.fixture
+def tool_registry(weather_tool, temperature_tool, exception_tool, thread_tool, interrupt_tool, slow_tool):
     registry = ToolRegistry()
     registry.register_tool(weather_tool)
     registry.register_tool(temperature_tool)
     registry.register_tool(exception_tool)
     registry.register_tool(thread_tool)
     registry.register_tool(interrupt_tool)
+    registry.register_tool(slow_tool)
     return registry
 
 

tests/strands/tools/executors/test_concurrent.py

@@ -1,6 +1,6 @@
 import pytest
 
-from strands.hooks import BeforeToolCallEvent
+from strands.hooks import AfterToolCallEvent, BeforeToolCallEvent
 from strands.interrupt import Interrupt
 from strands.tools.executors import ConcurrentToolExecutor
 from strands.tools.structured_output._structured_output_context import StructuredOutputContext
@@ -76,3 +76,30 @@ def interrupt_callback(event):
     tru_results = tool_results
     exp_results = [exp_events[1].tool_result]
     assert tru_results == exp_results
+
+
+@pytest.mark.asyncio
+async def test_concurrent_executor_reraises_exceptions(
+    executor, agent, tool_results, cycle_trace, cycle_span, invocation_state, structured_output_context, alist
+):
+    """Test that hook re-raised exceptions propagate and cancel remaining tasks."""
+
+    def reraise_callback(event):
+        if event.exception is not None:
+            raise event.exception
+
+    agent.hooks.add_callback(AfterToolCallEvent, reraise_callback)
+
+    tool_uses = [
+        {"name": "exception_tool", "toolUseId": "1", "input": {}},
+        {"name": "slow_tool", "toolUseId": "2", "input": {}},
+    ]
+
+    stream = executor._execute(
+        agent, tool_uses, tool_results, cycle_trace, cycle_span, invocation_state, structured_output_context
+    )
+
+    with pytest.raises(RuntimeError, match="Tool error"):
+        await alist(stream)
+
+    assert tool_results == []