Python: fix: keep AG-UI workflow reasoning in thread snapshots

[Manjunath Janardhan (manjunathshiva)]

Motivation & Context

With an AG-UI snapshot store enabled, a Workflow run's intermediate reasoning renders live and then disappears when the thread is hydrated. The same run through the agent path keeps it, so live and replayed output disagree.

The reasoning is already recorded and simply never read back. _emit_text_reasoning (_run_common.py:1051), which _emit_content delegates to for text_reasoning content, persists each reasoning message into flow.reasoning_messages — its docstring says so outright. The agent runner honours that contract (_agent_run.py:2100-2102, emitted terminally at :3202-3212); the workflow runner calls the same _emit_content, filling the same list, and never consumes it.

_WorkflowSnapshotBuilder.observe folds TextMessage* and ToolCall* events into the synthesized snapshot but had no Reasoning* branch, so reasoning events fell through silently and never reached the store.

Evan Mattson (@moonbox3) raised this in review on #8003 and confirmed a follow-up PR was the right home for it ("Follow up is fine, thanks."). This is that follow-up, built on top of #8003 now that it has merged.

Description & Review Guide

• What are the major changes?

  _WorkflowSnapshotBuilder now folds reasoning events, mirroring the existing text handling:

  • observe gains branches for ReasoningMessageStartEvent, ReasoningMessageContentEvent, ReasoningMessageEndEvent / ReasoningEndEvent, and ReasoningEncryptedValueEvent.
  • New _observe_reasoning_* / _flush_open_reasoning_message helpers accumulate deltas per message_id and emit entries in the same shape the agent path produces: {"id", "role": "reasoning", "content", ["encryptedValue"]}.
  • Reasoning is flushed at build(), and at text-message and tool-call starts, so a block that streamed before other output is replayed in the position it streamed in.

• What is the impact of these changes?

  Reasoning that streamed during a workflow run is still present after the thread is hydrated from a snapshot. Nothing else about the snapshot changes: reasoning is deliberately not allowed to close an open tool-call group, because _message_adapters.py:691-694 drops role == "reasoning" when converting back to provider messages — it is UI-only state, so it cannot break the tool_calls/result adjacency that providers require.

  This is the smaller of the two shapes discussed on the issue. It reads the real event stream rather than flow, which matters for workflows: request_info/interrupt tool calls and executor passthrough events are yielded directly and never touch flow, so observe sees strictly more than _build_messages_snapshot would. The second shape — emitting a terminal MessagesSnapshotEvent from the workflow runner — is left to the issue, since an emitted snapshot is treated as authoritative and would need to reproduce those bypassing events too.

• What do you want reviewers to focus on?

  The flush points. _flush_open_reasoning_message is called from build(), _observe_text_start and _observe_tool_call_start; the runner normally closes a reasoning block via _close_reasoning_block before any non-reasoning content, but raw ToolCall* events that bypass _emit_content can arrive with a block still open, which is what the tool-call-start flush covers.

Related Issue

Fixes #8054

Contribution Checklist

• The code builds clean without any errors or warnings
• All unit tests pass, and I have added new tests where possible
• The PR follows the Contribution Guidelines
• This PR is linked to an issue and there is no other open PR for this issue (see Related Issue above).
• This is not a breaking change. If it is a breaking change, add the breaking change label (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.

[Copilot]

🟡 Changes recommended

The moderate encrypted-only reasoning loss must be fixed before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Preserves AG-UI workflow reasoning in thread snapshots so hydrated output matches live streaming.

Changes:

• Folds reasoning events and encrypted values into workflow snapshots.
• Preserves reasoning output order.
• Adds unit and end-to-end regression coverage.

File summaries

File	Summary
python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py	Tests reasoning persistence through snapshot hydration.
python/packages/ag-ui/tests/ag_ui/test_snapshots.py	Tests reasoning synthesis, ordering, and encryption.
python/packages/ag-ui/agent_framework_ag_ui/_workflow.py	Adds reasoning snapshot synthesis; encrypted-only reasoning can be dropped when the message-end event precedes its encrypted value.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 1
• Review effort level: Balanced

---

💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.