Python: (ag-ui): fix Workflow.as_agent() streaming regression (#3875)

* fix Workflow.as_agent() streaming regression in ag-ui

* Address PR feedback

* PR feedback

python/packages/ag-ui/tests/ag_ui/test_endpoint.py

@@ -6,6 +6,7 @@ import json
 
 import pytest
 from agent_framework import Agent, ChatResponseUpdate, Content
+from agent_framework.orchestrations import SequentialBuilder
 from fastapi import FastAPI, Header, HTTPException
 from fastapi.params import Depends
 from fastapi.testclient import TestClient
@@ -165,6 +166,28 @@ async def test_endpoint_event_streaming(build_chat_client):
     assert found_run_finished
 
 
+async def test_endpoint_with_workflow_as_agent_stream_output(build_chat_client):
+    """Test endpoint handles workflow-as-agent stream outputs."""
+    app = FastAPI()
+    brainstorm_agent = Agent(name="brainstorm", instructions="Brainstorm ideas", client=build_chat_client("Idea"))
+    reviewer_agent = Agent(name="reviewer", instructions="Review ideas", client=build_chat_client("Review"))
+    agent = SequentialBuilder(participants=[brainstorm_agent, reviewer_agent]).build().as_agent()
+
+    add_agent_framework_fastapi_endpoint(app, agent, path="/workflow-like")
+
+    client = TestClient(app)
+    response = client.post("/workflow-like", json={"messages": [{"role": "user", "content": "Hello"}]})
+
+    assert response.status_code == 200
+    content = response.content.decode("utf-8")
+    lines = [line for line in content.split("\n") if line.startswith("data: ")]
+    event_types = [json.loads(line[6:]).get("type") for line in lines]
+
+    assert "RUN_STARTED" in event_types
+    assert "TEXT_MESSAGE_CONTENT" in event_types
+    assert "RUN_FINISHED" in event_types
+
+
 async def test_endpoint_error_handling(build_chat_client):
     """Test endpoint error handling during request parsing."""
     app = FastAPI()

python/packages/ag-ui/tests/ag_ui/test_run.py

@@ -2,11 +2,13 @@
 
 """Tests for _run.py helper functions and FlowState."""
 
+import pytest
 from ag_ui.core import (
     TextMessageEndEvent,
     TextMessageStartEvent,
 )
-from agent_framework import Content, Message
+from agent_framework import AgentResponseUpdate, Content, Message, ResponseStream
+from agent_framework.exceptions import AgentExecutionException
 
 from agent_framework_ag_ui._run import (
     FlowState,
@@ -16,6 +18,7 @@ from agent_framework_ag_ui._run import (
     _emit_tool_result,
     _has_only_tool_calls,
     _inject_state_context,
+    _normalize_response_stream,
     _should_suppress_intermediate_snapshot,
 )
 
@@ -179,6 +182,54 @@ class TestFlowState:
         assert result[0]["id"] == "call_2"
 
 
+class TestNormalizeResponseStream:
+    """Tests for _normalize_response_stream helper."""
+
+    async def test_accepts_response_stream(self):
+        """Accept standard ResponseStream values."""
+
+        async def _stream():
+            yield AgentResponseUpdate(contents=[Content.from_text("hello")], role="assistant")
+
+        stream = await _normalize_response_stream(ResponseStream(_stream()))
+        updates = [update async for update in stream]
+
+        assert len(updates) == 1
+        assert updates[0].contents[0].text == "hello"
+
+    async def test_accepts_async_iterable(self):
+        """Accept workflow-style async generator streams."""
+
+        async def _stream():
+            yield AgentResponseUpdate(contents=[Content.from_text("hello")], role="assistant")
+
+        stream = await _normalize_response_stream(_stream())
+        updates = [update async for update in stream]
+
+        assert len(updates) == 1
+        assert updates[0].contents[0].text == "hello"
+
+    async def test_accepts_awaitable_resolving_to_async_iterable(self):
+        """Accept awaitables that resolve to async iterable streams."""
+
+        async def _stream():
+            yield AgentResponseUpdate(contents=[Content.from_text("hello")], role="assistant")
+
+        async def _resolve():
+            return _stream()
+
+        stream = await _normalize_response_stream(_resolve())
+        updates = [update async for update in stream]
+
+        assert len(updates) == 1
+        assert updates[0].contents[0].text == "hello"
+
+    async def test_rejects_non_stream_values(self):
+        """Reject unsupported stream return values."""
+        with pytest.raises(AgentExecutionException):
+            await _normalize_response_stream("not-a-stream")
+
+
 class TestCreateStateContextMessage:
     """Tests for _create_state_context_message function."""