From 4aaf0fbd6436d7c48a4bbaf10ec5843dd8629925 Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Wed, 2 Sep 2026 06:54:25 +0530 Subject: [PATCH 1/2] Python: fix: surface AG-UI workflow intermediate events as reasoning The AG-UI workflow runner only handled "output" and "data" workflow events, so "intermediate" events (the modern designation produced by executors listed under intermediate_output_from) fell through to the generic CustomEvent fallback instead of being converted to AG-UI reasoning events. Add "intermediate" to the handled set and route intermediate text content to text_reasoning so consumers render it as a collapsible "thinking" block rather than a final assistant message. The deprecated "data" alias is treated the same as "intermediate" (per its definition as the intermediate compatibility alias in _events.py). "output" keeps the terminal-message behavior, and non-text content (tool calls and results) still emits as its native AG-UI events. Fixes #8000 --- .../agent_framework_ag_ui/_workflow_run.py | 54 ++++++++++++--- .../ag-ui/tests/ag_ui/test_workflow_run.py | 66 +++++++++++++++++++ 2 files changed, 110 insertions(+), 10 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index a01851d1dc..58d9905bd5 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -954,6 +954,25 @@ def _workflow_payload_to_contents(payload: Any) -> list[Content] | None: return None +def _as_reasoning_content(content: Content) -> Content: + """Re-tag plain text content as ``text_reasoning``. + + Intermediate workflow output should surface as AG-UI reasoning (a collapsible + "thinking" block) rather than a final assistant message. Only ``text`` content + is converted; tool calls, results, and other content types pass through + unchanged so they still emit as their native AG-UI events. + """ + if content.type != "text": + return content + return Content.from_text_reasoning( + id=content.id, + text=content.text, + annotations=content.annotations, + additional_properties=content.additional_properties or None, + raw_representation=content.raw_representation, + ) + + def _event_name(event: Any) -> str: event_type = getattr(event, "type", None) if isinstance(event_type, str) and event_type: @@ -1274,7 +1293,12 @@ def _drain_open_message() -> list[TextMessageEndEvent]: yield CustomEvent(name=_INTERRUPT_CARD_EVENT_NAME, value=interrupt_event_value) continue - if event_type in {"output", "data"}: + if event_type in {"output", "intermediate", "data"}: + # "intermediate" (and its deprecated alias "data") carry non-terminal + # output. Their text is surfaced as AG-UI reasoning so consumers render + # it as a collapsible "thinking" block instead of a final assistant + # message. "output" keeps the terminal-message behavior. + is_intermediate = event_type in {"intermediate", "data"} output_payload = getattr(event, "data", None) if isinstance(output_payload, BaseEvent): yield output_payload @@ -1293,15 +1317,25 @@ def _drain_open_message() -> list[TextMessageEndEvent]: yield out_event contents = _workflow_payload_to_contents(output_payload) if contents: - output_text = _text_from_contents(contents) - skip_text = bool(output_text and output_text == last_assistant_text) - for content in contents: - for out_event in _emit_content(content, flow, predictive_handler=None, skip_text=skip_text): - yield out_event - if flow.message_id and flow.accumulated_text: - last_assistant_text = flow.accumulated_text.strip() or last_assistant_text - elif output_text: - last_assistant_text = output_text + if is_intermediate: + # Reasoning is a separate channel from the final assistant + # message, so the last_assistant_text dedup does not apply. + for content in contents: + reasoning_content = _as_reasoning_content(content) + for out_event in _emit_content( + reasoning_content, flow, predictive_handler=None, skip_text=False + ): + yield out_event + else: + output_text = _text_from_contents(contents) + skip_text = bool(output_text and output_text == last_assistant_text) + for content in contents: + for out_event in _emit_content(content, flow, predictive_handler=None, skip_text=skip_text): + yield out_event + if flow.message_id and flow.accumulated_text: + last_assistant_text = flow.accumulated_text.strip() or last_assistant_text + elif output_text: + last_assistant_text = output_text else: yield CustomEvent(name="workflow_output", value=make_json_safe(output_payload)) continue diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py index 7ee5d835c5..eaf21b6eec 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py @@ -120,6 +120,72 @@ async def start(message: Any, ctx: WorkflowContext[Any, str]) -> None: assert custom_events[0].value == {"progress": 10} # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] +async def test_workflow_run_maps_intermediate_text_to_reasoning_events(): + """Intermediate workflow output is surfaced as AG-UI reasoning, not a generic CustomEvent.""" + + @executor(id="thinker") + async def thinker(message: Any, ctx: WorkflowContext[str, str]) -> None: + # Intermediate-designated executor: text should render as reasoning. + await ctx.yield_output("Analyzing the problem...") + await ctx.send_message("go") + + @executor(id="finalizer") + async def finalizer(message: str, ctx: WorkflowContext[None, str]) -> None: + # Output-designated executor: final assistant text. + await ctx.yield_output("Here's my answer!") + + workflow = ( + WorkflowBuilder( + start_executor=thinker, + output_from=[finalizer], + intermediate_output_from=[thinker], + ) + .add_edge(thinker, finalizer) + .build() + ) + input_data = {"messages": [{"role": "user", "content": "solve it"}]} + + events = [event async for event in run_workflow_stream(input_data, workflow)] + event_types = [event.type for event in events] + + # The intermediate text renders as reasoning ... + assert "REASONING_MESSAGE_CONTENT" in event_types + reasoning_deltas = [event.delta for event in events if event.type == "REASONING_MESSAGE_CONTENT"] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert "Analyzing the problem..." in reasoning_deltas + + # ... and is not swallowed by the generic custom-event fallback. + assert not [event for event in events if event.type == "CUSTOM" and event.name == "intermediate"] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + + # The final output still renders as an assistant text message. + assert "TEXT_MESSAGE_CONTENT" in event_types + text_deltas = [event.delta for event in events if event.type == "TEXT_MESSAGE_CONTENT"] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert "Here's my answer!" in text_deltas + + +async def test_workflow_run_maps_data_alias_text_to_reasoning_events(): + """The deprecated ``type='data'`` alias is treated like ``intermediate`` and renders as reasoning.""" + + @executor(id="emitter") + async def emitter(message: Any, ctx: WorkflowContext[Any, str]) -> None: + # Deprecated compatibility alias for an intermediate emission. + await ctx.add_event(WorkflowEvent.emit("emitter", "legacy reasoning")) + await ctx.yield_output("final answer") + + workflow = WorkflowBuilder(start_executor=emitter).build() + input_data = {"messages": [{"role": "user", "content": "go"}]} + + with pytest.warns(DeprecationWarning): + events = [event async for event in run_workflow_stream(input_data, workflow)] + + reasoning_deltas = [event.delta for event in events if event.type == "REASONING_MESSAGE_CONTENT"] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert "legacy reasoning" in reasoning_deltas + assert not [event for event in events if event.type == "CUSTOM" and event.name == "data"] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + + # The genuine output is still emitted as an assistant text message. + text_deltas = [event.delta for event in events if event.type == "TEXT_MESSAGE_CONTENT"] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert "final answer" in text_deltas + + async def test_workflow_and_agent_spans_use_supplied_agui_thread_id(monkeypatch: pytest.MonkeyPatch) -> None: """Workflow spans use supplied AG-UI threads without replacing provider fallback behavior.""" import agent_framework.observability as observability From 13a8b1a2d89a24db5e963f7a3513be636ca5ae86 Mon Sep 17 00:00:00 2001 From: Manjunath Janardhan Date: Wed, 2 Sep 2026 14:53:15 +0530 Subject: [PATCH 2/2] Python: address AG-UI intermediate-reasoning review feedback Follow-up to the intermediate-reasoning change on this branch, addressing maintainer review: - Preserve protected_data when re-tagging intermediate text as reasoning so the ReasoningEncryptedValueEvent and the snapshot encryptedValue survive. - Close any open reasoning block (and assistant text message) before every terminal event (RUN_FINISHED / RUN_ERROR) and before a request_info tool call, via an idempotent _drain_open_blocks() helper. Previously the post-loop cleanup emitted the reasoning end events after the terminal event, so a client that stopped at RUN_FINISHED never saw them; a request_info tool call could likewise sit inside an unclosed reasoning block. - Keep role-less AgentResponseUpdate text (and tool content) instead of dropping it to a CUSTOM workflow_output. Explicit non-assistant roles and approval requests stay excluded. --- .../agent_framework_ag_ui/_workflow_run.py | 50 ++++++- .../ag-ui/tests/ag_ui/test_workflow_run.py | 124 +++++++++++++++++- 2 files changed, 166 insertions(+), 8 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 58d9905bd5..a0e4a3dfdd 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -931,6 +931,17 @@ def _workflow_payload_to_contents(payload: Any) -> list[Content] | None: if isinstance(payload, AgentResponseUpdate): contents = list(payload.contents or []) role_field = payload.role + if role_field is None: + # ``role`` is optional and streamed continuation chunks routinely omit it. + # Keep their text -- previously dropped, so role-less text surfaced as a + # CUSTOM workflow_output instead of reasoning/assistant text -- alongside tool + # content. Approval requests stay excluded (see _TOOL_CONTENT_TYPES): a + # role-less approval interrupt from streamed content has no pending request to + # resume against. + role_less_contents = [ + content for content in contents if content.type == "text" or content.type in _TOOL_CONTENT_TYPES + ] + return role_less_contents or None if isinstance(role_field, str): role = role_field else: @@ -967,6 +978,11 @@ def _as_reasoning_content(content: Content) -> Content: return Content.from_text_reasoning( id=content.id, text=content.text, + # Carry encrypted reasoning metadata through unchanged: _emit_text_reasoning + # turns protected_data into a ReasoningEncryptedValueEvent and an + # ``encryptedValue`` on the snapshot entry, so dropping it here would break + # reasoning state continuity for intermediate content that carries it. + protected_data=content.protected_data, annotations=content.annotations, additional_properties=content.additional_properties or None, raw_representation=content.raw_representation, @@ -1150,6 +1166,22 @@ def _drain_open_message() -> list[TextMessageEndEvent]: flow.accumulated_text = "" return [TextMessageEndEvent(message_id=current_message_id)] + def _drain_open_blocks() -> list[BaseEvent]: + """Close any open reasoning block and assistant text message. + + Emitted before content that must not sit inside an open block: a terminal event + (RUN_FINISHED / RUN_ERROR, which must be the final events in the stream) or a + request_info tool call (non-reasoning message content). Otherwise the block's + REASONING_* / TEXT_MESSAGE_* end events would be flushed only by the post-loop + cleanup -- after the terminal event, or after the tool call. Both inner helpers + are no-ops when their block is not open, so this is always safe to call (a later + cleanup pass then simply does nothing). + """ + events: list[BaseEvent] = [] + events.extend(_close_reasoning_block(flow)) + events.extend(_drain_open_message()) + return events + fwd_kwargs: dict[str, Any] = {} if "forwarded_props" in input_data: forwarded_props = input_data["forwarded_props"] @@ -1205,6 +1237,10 @@ def _drain_open_message() -> list[TextMessageEndEvent]: run_started_emitted = True if event_type == "failed": + # Close any open reasoning block / text message so RUN_ERROR stays the + # last event a client receives for this run. + for end_event in _drain_open_blocks(): + yield end_event details = getattr(event, "details", None) yield RunErrorEvent(message=_details_message(details), code=_details_code(details)) run_error_emitted = True @@ -1218,9 +1254,9 @@ def _drain_open_message() -> list[TextMessageEndEvent]: else: state_value = str(getattr(state, "value", state)) if state_value in _TERMINAL_STATES and not terminal_emitted: - # Close any open assistant text message before the terminal event so - # RUN_FINISHED is always the last emitted event. - for end_event in _drain_open_message(): + # Close any open reasoning block and assistant text message before the + # terminal event so RUN_FINISHED is always the last emitted event. + for end_event in _drain_open_blocks(): yield end_event if not interrupts: interrupts.extend(_interrupts_from_pending_requests(await _pending_request_events(workflow))) @@ -1273,7 +1309,10 @@ def _drain_open_message() -> list[TextMessageEndEvent]: continue if event_type == "request_info": - for end_event in _drain_open_message(): + # A request_info emits a tool call (non-reasoning message content), so any + # open reasoning block / text message must be closed first -- otherwise the + # tool call would sit inside an unclosed reasoning block. + for end_event in _drain_open_blocks(): yield end_event request_payload = _request_payload_from_request_event(event) if request_payload is None: @@ -1348,6 +1387,9 @@ def _drain_open_message() -> list[TextMessageEndEvent]: if not run_started_emitted: yield RunStartedEvent(run_id=run_id, thread_id=thread_id) run_started_emitted = True + # Close any open reasoning block / text message so RUN_ERROR stays the final event. + for end_event in _drain_open_blocks(): + yield end_event if not run_error_emitted: yield RunErrorEvent(message=str(exc), code=type(exc).__name__) run_error_emitted = True diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py index eaf21b6eec..f1cd16c959 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py @@ -34,6 +34,7 @@ from pydantic import BaseModel from agent_framework_ag_ui._workflow_run import ( + _as_reasoning_content, _coerce_content, _coerce_json_value, _coerce_message, @@ -186,6 +187,112 @@ async def emitter(message: Any, ctx: WorkflowContext[Any, str]) -> None: assert "final answer" in text_deltas +def test_as_reasoning_content_preserves_protected_data(): + """Re-tagging text as reasoning keeps encrypted protected_data and passes non-text through.""" + text = Content("text", text="thinking", protected_data="enc-blob") + + reasoning = _as_reasoning_content(text) + + assert reasoning.type == "text_reasoning" + assert reasoning.text == "thinking" + # Without this the ReasoningEncryptedValueEvent and snapshot encryptedValue are lost. + assert reasoning.protected_data == "enc-blob" + + # Non-text content (e.g. a tool call) is returned unchanged. + call = Content.from_function_call(call_id="c1", name="tool", arguments="{}") + assert _as_reasoning_content(call) is call + + +async def test_workflow_run_closes_reasoning_before_run_finished(): + """Intermediate reasoning with no terminal text still closes before the terminal event.""" + + @executor(id="thinker") + async def thinker(message: Any, ctx: WorkflowContext[Any, str]) -> None: + # Intermediate output only -- no terminal assistant text follows. + await ctx.yield_output("Thinking, but no final answer...") + + workflow = WorkflowBuilder( + start_executor=thinker, + output_from=[], + intermediate_output_from=[thinker], + ).build() + input_data = {"messages": [{"role": "user", "content": "go"}]} + + events = [event async for event in run_workflow_stream(input_data, workflow)] + event_types = [event.type for event in events] + + assert "REASONING_END" in event_types + assert "RUN_FINISHED" in event_types + run_finished_idx = next(i for i, event in enumerate(events) if event.type == "RUN_FINISHED") + # Every reasoning event -- including the closing REASONING_MESSAGE_END / REASONING_END -- + # must precede RUN_FINISHED so clients that stop at the terminal event get a complete stream. + reasoning_idxs = [i for i, event in enumerate(events) if "REASONING" in str(event.type)] + assert reasoning_idxs + assert max(reasoning_idxs) < run_finished_idx + + +async def test_workflow_run_closes_reasoning_before_request_info(): + """An open reasoning block is closed before a request_info tool call (not left spanning it).""" + + @executor(id="asker") + async def asker(message: Any, ctx: WorkflowContext[Any, str]) -> None: + # Intermediate reasoning, then a human-in-the-loop request in the same run. + await ctx.yield_output("Thinking before I ask...") + await ctx.request_info("Need approval", str, request_id="approval-1") + + workflow = WorkflowBuilder( + start_executor=asker, + output_from=[], + intermediate_output_from=[asker], + ).build() + input_data = {"messages": [{"role": "user", "content": "go"}]} + + events = [event async for event in run_workflow_stream(input_data, workflow)] + event_types = [event.type for event in events] + + assert "REASONING_END" in event_types + reasoning_end_idx = max(i for i, event in enumerate(events) if "REASONING" in str(event.type)) + request_start_idx = next( + i + for i, event in enumerate(events) + if event.type == "TOOL_CALL_START" and getattr(event, "tool_call_id", None) == "approval-1" + ) + # The reasoning block must be fully closed before the request_info tool call. + assert reasoning_end_idx < request_start_idx + + +async def test_workflow_run_roleless_intermediate_update_becomes_reasoning(): + """A role-less AgentResponseUpdate on the intermediate path surfaces its text as reasoning.""" + + @executor(id="thinker") + async def thinker(message: Any, ctx: WorkflowContext[str, Any]) -> None: + # role=None is the common shape for streamed continuation chunks. + await ctx.yield_output(AgentResponseUpdate(contents=[Content.from_text("Role-less thought")], role=None)) + await ctx.send_message("go") + + @executor(id="finalizer") + async def finalizer(message: str, ctx: WorkflowContext[None, str]) -> None: + await ctx.yield_output("Done.") + + workflow = ( + WorkflowBuilder( + start_executor=thinker, + output_from=[finalizer], + intermediate_output_from=[thinker], + ) + .add_edge(thinker, finalizer) + .build() + ) + input_data = {"messages": [{"role": "user", "content": "go"}]} + + events = [event async for event in run_workflow_stream(input_data, workflow)] + + reasoning_deltas = [event.delta for event in events if event.type == "REASONING_MESSAGE_CONTENT"] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert "Role-less thought" in reasoning_deltas + # The role-less text must not be dropped into a generic custom event. + assert not [event for event in events if event.type == "CUSTOM" and event.name == "workflow_output"] # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + + async def test_workflow_and_agent_spans_use_supplied_agui_thread_id(monkeypatch: pytest.MonkeyPatch) -> None: """Workflow spans use supplied AG-UI threads without replacing provider fallback behavior.""" import agent_framework.observability as observability @@ -1904,9 +2011,10 @@ def test_agent_response_update_non_assistant(self): assert _workflow_payload_to_contents(update) is None def test_agent_response_update_none_role(self): - """AgentResponseUpdate with None role returns None.""" - update = AgentResponseUpdate(contents=[Content.from_text(text="hi")], role=None) - assert _workflow_payload_to_contents(update) is None + """AgentResponseUpdate with None role keeps text (role-less continuation chunks).""" + text = Content.from_text(text="hi") + update = AgentResponseUpdate(contents=[text], role=None) + assert _workflow_payload_to_contents(update) == [text] def test_agent_response_update_function_call_without_role(self) -> None: """Function call content passes through without role metadata.""" @@ -1949,11 +2057,19 @@ def test_agent_response_update_mcp_tool_result_without_role(self) -> None: assert _workflow_payload_to_contents(update) == [mcp_result] def test_agent_response_update_mixed_content_without_role(self) -> None: - """Non-assistant updates keep tool content and drop text content.""" + """Role-less updates keep both text and tool content in order.""" text = Content.from_text(text="calling the tool") function_call = Content.from_function_call(call_id="call-1", name="search", arguments={"query": "weather"}) update = AgentResponseUpdate(contents=[text, function_call], role=None) + assert _workflow_payload_to_contents(update) == [text, function_call] + + def test_agent_response_update_explicit_non_assistant_role_drops_text(self) -> None: + """An explicit non-assistant role still keeps only tool content and drops text.""" + text = Content.from_text(text="calling the tool") + function_call = Content.from_function_call(call_id="call-1", name="search", arguments={"query": "weather"}) + update = AgentResponseUpdate(contents=[text, function_call], role="tool") + assert _workflow_payload_to_contents(update) == [function_call] def test_agent_response_update_assistant_text(self) -> None: