Python: Fix WorkflowAgent event handling and kwargs forwarding

python/packages/core/agent_framework/_workflows/_agent.py

@@ -141,7 +141,8 @@ async def run(
             checkpoint_storage: Runtime checkpoint storage. When provided with checkpoint_id,
                 used to load and restore the checkpoint. When provided without checkpoint_id,
                 enables checkpointing for this run.
-            **kwargs: Additional keyword arguments.
+            **kwargs: Additional keyword arguments passed through to underlying workflow
+                and ai_function tools.
 
         Returns:
             The final workflow response as an AgentRunResponse.
@@ -153,7 +154,7 @@ async def run(
         response_id = str(uuid.uuid4())
 
         async for update in self._run_stream_impl(
-            input_messages, response_id, thread, checkpoint_id, checkpoint_storage
+            input_messages, response_id, thread, checkpoint_id, checkpoint_storage, **kwargs
         ):
             response_updates.append(update)
 
@@ -187,7 +188,8 @@ async def run_stream(
             checkpoint_storage: Runtime checkpoint storage. When provided with checkpoint_id,
                 used to load and restore the checkpoint. When provided without checkpoint_id,
                 enables checkpointing for this run.
-            **kwargs: Additional keyword arguments.
+            **kwargs: Additional keyword arguments passed through to underlying workflow
+                and ai_function tools.
 
         Yields:
             AgentRunResponseUpdate objects representing the workflow execution progress.
@@ -198,7 +200,7 @@ async def run_stream(
         response_id = str(uuid.uuid4())
 
         async for update in self._run_stream_impl(
-            input_messages, response_id, thread, checkpoint_id, checkpoint_storage
+            input_messages, response_id, thread, checkpoint_id, checkpoint_storage, **kwargs
         ):
             response_updates.append(update)
             yield update
@@ -216,6 +218,7 @@ async def _run_stream_impl(
         thread: AgentThread,
         checkpoint_id: str | None = None,
         checkpoint_storage: CheckpointStorage | None = None,
+        **kwargs: Any,
     ) -> AsyncIterable[AgentRunResponseUpdate]:
         """Internal implementation of streaming execution.
 
@@ -225,6 +228,8 @@ async def _run_stream_impl(
             thread: The conversation thread containing message history.
             checkpoint_id: ID of checkpoint to restore from.
             checkpoint_storage: Runtime checkpoint storage.
+            **kwargs: Additional keyword arguments passed through to the underlying
+                workflow and ai_function tools.
 
         Yields:
             AgentRunResponseUpdate objects representing the workflow execution progress.
@@ -255,6 +260,7 @@ async def _run_stream_impl(
                 message=None,
                 checkpoint_id=checkpoint_id,
                 checkpoint_storage=checkpoint_storage,
+                **kwargs,
             )
         else:
             # Execute workflow with streaming (initial run or no function responses)
@@ -268,6 +274,7 @@ async def _run_stream_impl(
             event_stream = self.workflow.run_stream(
                 message=conversation_messages,
                 checkpoint_storage=checkpoint_storage,
+                **kwargs,
             )
 
         # Process events from the stream
@@ -295,13 +302,9 @@ def _convert_workflow_event_to_agent_update(
                 return None
 
             case WorkflowOutputEvent(data=data, source_executor_id=source_executor_id):
-                # Convert workflow output to an agent response update.
-                # Handle different data types appropriately.
                 if isinstance(data, AgentRunResponseUpdate):
-                    # Already an update, pass through
                     return data
                 if isinstance(data, ChatMessage):
-                    # Convert ChatMessage to update
                     return AgentRunResponseUpdate(
                         contents=list(data.contents),
                         role=data.role,
@@ -311,15 +314,9 @@ def _convert_workflow_event_to_agent_update(
                         created_at=datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
                         raw_representation=data,
                     )
-                # Determine contents based on data type
-                if isinstance(data, BaseContent):
-                    # Already a content type (TextContent, ImageContent, etc.)
-                    contents: list[Contents] = [cast(Contents, data)]
-                elif isinstance(data, str):
-                    contents = [TextContent(text=data)]
-                else:
-                    # Fallback: convert to string representation
-                    contents = [TextContent(text=str(data))]
+                contents = self._extract_contents(data)
+                if not contents:
+                    return None
                 return AgentRunResponseUpdate(
                     contents=contents,
                     role=Role.ASSISTANT,
@@ -405,6 +402,18 @@ def _extract_function_responses(self, input_messages: list[ChatMessage]) -> dict
                         raise AgentExecutionException("Unexpected content type while awaiting request info responses.")
         return function_responses
 
+    def _extract_contents(self, data: Any) -> list[Contents]:
+        """Recursively extract Contents from workflow output data."""
+        if isinstance(data, ChatMessage):
+            return list(data.contents)
+        if isinstance(data, list):
+            return [c for item in data for c in self._extract_contents(item)]
+        if isinstance(data, BaseContent):
+            return [cast(Contents, data)]
+        if isinstance(data, str):
+            return [TextContent(text=data)]
+        return [TextContent(text=str(data))]
+
     class _ResponseState(TypedDict):
         """State for grouping response updates by message_id."""
 

python/packages/core/tests/workflow/test_workflow_agent.py

@@ -422,6 +422,48 @@ async def raw_yielding_executor(messages: list[ChatMessage], ctx: WorkflowContex
         assert isinstance(updates[2].raw_representation, CustomData)
         assert updates[2].raw_representation.value == 42
 
+    async def test_workflow_as_agent_yield_output_with_list_of_chat_messages(self) -> None:
+        """Test that yield_output with list[ChatMessage] extracts contents from all messages.
+
+        Note: TextContent items are coalesced by _finalize_response, so multiple text contents
+        become a single merged TextContent in the final response.
+        """
+
+        @executor
+        async def list_yielding_executor(messages: list[ChatMessage], ctx: WorkflowContext) -> None:
+            # Yield a list of ChatMessages (as SequentialBuilder does)
+            msg_list = [
+                ChatMessage(role=Role.USER, contents=[TextContent(text="first message")]),
+                ChatMessage(role=Role.ASSISTANT, contents=[TextContent(text="second message")]),
+                ChatMessage(
+                    role=Role.ASSISTANT,
+                    contents=[TextContent(text="third"), TextContent(text="fourth")],
+                ),
+            ]
+            await ctx.yield_output(msg_list)
+
+        workflow = WorkflowBuilder().set_start_executor(list_yielding_executor).build()
+        agent = workflow.as_agent("list-msg-agent")
+
+        # Verify streaming returns the update with all 4 contents before coalescing
+        updates: list[AgentRunResponseUpdate] = []
+        async for update in agent.run_stream("test"):
+            updates.append(update)
+
+        assert len(updates) == 1
+        assert len(updates[0].contents) == 4
+        texts = [c.text for c in updates[0].contents if isinstance(c, TextContent)]
+        assert texts == ["first message", "second message", "third", "fourth"]
+
+        # Verify run() coalesces text contents (expected behavior)
+        result = await agent.run("test")
+
+        assert isinstance(result, AgentRunResponse)
+        assert len(result.messages) == 1
+        # TextContent items are coalesced into one
+        assert len(result.messages[0].contents) == 1
+        assert result.messages[0].text == "first messagesecond messagethirdfourth"
+
     async def test_thread_conversation_history_included_in_workflow_run(self) -> None:
         """Test that conversation history from thread is included when running WorkflowAgent.
 

python/packages/core/tests/workflow/test_workflow_kwargs.py

@@ -492,6 +492,117 @@ async def prepare_final_answer(self, context: MagenticContext) -> ChatMessage:
 # endregion
 
 
+# region WorkflowAgent (as_agent) kwargs Tests
+
+
+async def test_workflow_as_agent_run_propagates_kwargs_to_underlying_agent() -> None:
+    """Test that kwargs passed to workflow_agent.run() flow through to the underlying agents."""
+    agent = _KwargsCapturingAgent(name="inner_agent")
+    workflow = SequentialBuilder().participants([agent]).build()
+    workflow_agent = workflow.as_agent(name="TestWorkflowAgent")
+
+    custom_data = {"endpoint": "https://api.example.com", "version": "v1"}
+    user_token = {"user_name": "alice", "access_level": "admin"}
+
+    _ = await workflow_agent.run(
+        "test message",
+        custom_data=custom_data,
+        user_token=user_token,
+    )
+
+    # Verify inner agent received kwargs
+    assert len(agent.captured_kwargs) >= 1, "Inner agent should have been invoked at least once"
+    received = agent.captured_kwargs[0]
+    assert "custom_data" in received, "Inner agent should receive custom_data kwarg"
+    assert "user_token" in received, "Inner agent should receive user_token kwarg"
+    assert received["custom_data"] == custom_data
+    assert received["user_token"] == user_token
+
+
+async def test_workflow_as_agent_run_stream_propagates_kwargs_to_underlying_agent() -> None:
+    """Test that kwargs passed to workflow_agent.run_stream() flow through to the underlying agents."""
+    agent = _KwargsCapturingAgent(name="inner_agent")
+    workflow = SequentialBuilder().participants([agent]).build()
+    workflow_agent = workflow.as_agent(name="TestWorkflowAgent")
+
+    custom_data = {"session_id": "xyz123"}
+    api_token = "secret-token"
+
+    async for _ in workflow_agent.run_stream(
+        "test message",
+        custom_data=custom_data,
+        api_token=api_token,
+    ):
+        pass
+
+    # Verify inner agent received kwargs
+    assert len(agent.captured_kwargs) >= 1, "Inner agent should have been invoked at least once"
+    received = agent.captured_kwargs[0]
+    assert "custom_data" in received, "Inner agent should receive custom_data kwarg"
+    assert "api_token" in received, "Inner agent should receive api_token kwarg"
+    assert received["custom_data"] == custom_data
+    assert received["api_token"] == api_token
+
+
+async def test_workflow_as_agent_propagates_kwargs_to_multiple_agents() -> None:
+    """Test that kwargs flow to all agents when using workflow.as_agent()."""
+    agent1 = _KwargsCapturingAgent(name="agent1")
+    agent2 = _KwargsCapturingAgent(name="agent2")
+    workflow = SequentialBuilder().participants([agent1, agent2]).build()
+    workflow_agent = workflow.as_agent(name="MultiAgentWorkflow")
+
+    custom_data = {"batch_id": "batch-001"}
+
+    _ = await workflow_agent.run("test message", custom_data=custom_data)
+
+    # Both agents should have received kwargs
+    assert len(agent1.captured_kwargs) >= 1, "First agent should be invoked"
+    assert len(agent2.captured_kwargs) >= 1, "Second agent should be invoked"
+    assert agent1.captured_kwargs[0].get("custom_data") == custom_data
+    assert agent2.captured_kwargs[0].get("custom_data") == custom_data
+
+
+async def test_workflow_as_agent_kwargs_with_none_values() -> None:
+    """Test that kwargs with None values are passed through correctly via as_agent()."""
+    agent = _KwargsCapturingAgent(name="none_test_agent")
+    workflow = SequentialBuilder().participants([agent]).build()
+    workflow_agent = workflow.as_agent(name="NoneTestWorkflow")
+
+    _ = await workflow_agent.run("test", optional_param=None, other_param="value")
+
+    assert len(agent.captured_kwargs) >= 1
+    received = agent.captured_kwargs[0]
+    assert "optional_param" in received
+    assert received["optional_param"] is None
+    assert received["other_param"] == "value"
+
+
+async def test_workflow_as_agent_kwargs_with_complex_nested_data() -> None:
+    """Test that complex nested data structures flow through correctly via as_agent()."""
+    agent = _KwargsCapturingAgent(name="nested_agent")
+    workflow = SequentialBuilder().participants([agent]).build()
+    workflow_agent = workflow.as_agent(name="NestedDataWorkflow")
+
+    complex_data = {
+        "level1": {
+            "level2": {
+                "level3": ["a", "b", "c"],
+                "number": 42,
+            },
+            "list": [1, 2, {"nested": True}],
+        },
+    }
+
+    _ = await workflow_agent.run("test", complex_data=complex_data)
+
+    assert len(agent.captured_kwargs) >= 1
+    received = agent.captured_kwargs[0]
+    assert received.get("complex_data") == complex_data
+
+
+# endregion
+
+
 # region SubWorkflow (WorkflowExecutor) Tests
 
 

python/samples/getting_started/workflows/README.md

@@ -45,6 +45,7 @@ Once comfortable with these, explore the rest of the samples below.
 | Workflow as Agent (Reflection Pattern) | [agents/workflow_as_agent_reflection_pattern.py](./agents/workflow_as_agent_reflection_pattern.py) | Wrap a workflow so it can behave like an agent (reflection pattern) |
 | Workflow as Agent + HITL | [agents/workflow_as_agent_human_in_the_loop.py](./agents/workflow_as_agent_human_in_the_loop.py) | Extend workflow-as-agent with human-in-the-loop capability |
 | Workflow as Agent with Thread | [agents/workflow_as_agent_with_thread.py](./agents/workflow_as_agent_with_thread.py) | Use AgentThread to maintain conversation history across workflow-as-agent invocations |
+| Workflow as Agent kwargs | [agents/workflow_as_agent_kwargs.py](./agents/workflow_as_agent_kwargs.py) | Pass custom context (data, user tokens) via kwargs through workflow.as_agent() to @ai_function tools |
 | Handoff Workflow as Agent | [agents/handoff_workflow_as_agent.py](./agents/handoff_workflow_as_agent.py) | Use a HandoffBuilder workflow as an agent with HITL via FunctionCallContent/FunctionResultContent |
 
 ### checkpoint