Python: Await async callable aggregators in ConcurrentBuilder

python/packages/orchestrations/agent_framework_orchestrations/_concurrent.py

@@ -172,6 +172,10 @@ async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowCon
             else:
                 ret = await asyncio.to_thread(self._callback, results)
 
+        # Async callable objects and wrappers returning a coroutine are not coroutine functions
+        if inspect.isawaitable(ret):
+            ret = await ret
+
         # If the callback returned a value, finalize the workflow with it
         if ret is not None:
             await ctx.yield_output(ret)

python/packages/orchestrations/tests/test_concurrent.py

@@ -145,6 +145,34 @@ def summarize_sync(results: list[AgentExecutorResponse], _ctx: WorkflowContext[A
     assert output == "One | Two"
 
 
+class _AsyncCallableAggregator:
+    """Aggregator object whose __call__ is async, for example one that holds a client."""
+
+    async def __call__(self, results: list[AgentExecutorResponse]) -> str:
+        return " | ".join(sorted(r.agent_response.text for r in results))
+
+
+async def _join_texts(results: list[AgentExecutorResponse]) -> str:
+    return " | ".join(sorted(r.agent_response.text for r in results))
+
+
+@pytest.mark.parametrize(
+    "aggregator",
+    [
+        pytest.param(_AsyncCallableAggregator(), id="async_callable_object"),
+        pytest.param(lambda results: _join_texts(results), id="callable_returning_coroutine"),
+    ],
+)
+async def test_concurrent_custom_aggregator_awaits_awaitable_result(aggregator: Any) -> None:
+    e1 = _FakeAgentExec("agentA", "One")
+    e2 = _FakeAgentExec("agentB", "Two")
+
+    wf = ConcurrentBuilder(participants=[e1, e2]).with_aggregator(aggregator).build()
+    result = await wf.run("prompt: awaitable aggregator")
+
+    assert result.get_outputs() == ["One | Two"]
+
+
 def test_concurrent_custom_aggregator_uses_callback_name_for_id() -> None:
     e1 = _FakeAgentExec("agentA", "One")
     e2 = _FakeAgentExec("agentB", "Two")