Python: fix: keep AG-UI workflow reasoning in thread snapshots - #8058
Conversation
`_emit_text_reasoning` persists each workflow reasoning message into
`flow.reasoning_messages`, and `_WorkflowSnapshotBuilder.observe` folded
`TextMessage*` and `ToolCall*` events into the synthesized snapshot but had
no `Reasoning*` branch. Reasoning events fell through silently, so
intermediate output rendered live and then vanished when the thread was
hydrated from a snapshot -- while the same run through the agent path kept
it.
Fold reasoning into the builder: accumulate deltas per message_id and emit
entries in the shape the agent path already 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 replays in the position it streamed in.
Reasoning deliberately does not close an open tool-call group: it is UI-only
state that `agui_messages_to_agent_framework` drops, so it cannot break the
tool_calls/result adjacency providers require.
Fixes microsoft#8054
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
agent-framework-automation
Bot
added
the
python
GitHub Actions
Inactive
There was a problem hiding this comment.
🟡 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.
Address review: an encrypted-value-only reasoning message was dropped for the event order `_emit_text_reasoning` produces without a flow, where REASONING_MESSAGE_END precedes REASONING_ENCRYPTED_VALUE. The message carries no display text, so closing it at REASONING_MESSAGE_END discarded it and left the encrypted value with nothing to attach to. The underlying mistake was conflating two protocol levels. REASONING_MESSAGE_END closes the message; REASONING_END closes the block; and an encrypted value is block-scoped, so it legitimately trails the message end. Keep the message open past REASONING_MESSAGE_END and let REASONING_END, a new REASONING_START, intervening text/tool output, or build() finalize it. Handle REASONING_START so a new block also closes anything the previous one left open. Adds regression coverage for the reported order, for an empty block with neither text nor an encrypted value still being dropped, for a block closed by the next REASONING_START, for an encrypted value arriving after intervening text, and for end events naming a message that was never opened. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
GitHub Actions
Inactive
Evan Mattson (moonbox3)
added
the
ag-ui
Review on microsoft#8058 found two ordering defects in `_WorkflowSnapshotBuilder`. An `output` event followed by a later `intermediate` event left both the open text message and the open reasoning message unflushed, so `build()`'s fixed flush order decided the sequence and hydration reversed what streamed. Nothing flushed open text when reasoning opened, even though `_observe_text_start` and `_observe_tool_call_start` already flush open reasoning. Reasoning is now flushed wherever other output is appended and text is flushed wherever reasoning opens, so the two slots can never both be open and `build()`'s order cannot matter. That covers the `_observe_text_content` resume path and `_observe_tool_call_result`, which had the same gap. Reasoning is UI-only and dropped before provider conversion, so a block that now lands between a tool call and its result still converts back with the two adjacent, which provider APIs require; a test pins that across the two modules. A block carrying only protected data has no content event at all, so output arriving before its encrypted value flushed an empty message that the late `ReasoningEncryptedValueEvent` could no longer find, and the protected value was lost. The empty message now stays addressable in the position it streamed and `build()` filters it out only while nothing has claimed it, so genuinely empty reasoning is still absent from the snapshot. Filtering rather than deleting is not observable through the one call path today, which builds once per run; it keeps `build()` a projection of accumulated state so the choice does not have to be revisited if the builder is ever driven incrementally. A live fan-out workflow on gpt-5-mini shows both preconditions for that loss. Reasoning arrives with an encrypted value and no visible text, the value is emitted twice for one item, and with a flow it lands between REASONING_MESSAGE_START and REASONING_MESSAGE_END rather than after it. Concurrent executors interleave their events, so a text message can start in the gap before the value arrives; replaying the captured order with that interleaving loses the encrypted reasoning before this change and keeps it after. That same interleaving exposed a second loss on the no-start text path: `_observe_text_content` opened a message over one already open under a different id, discarding its content, where `_observe_text_start` flushes first. It now flushes as well. Splitting an open text message on reasoning meant a resumed message reusing its id replayed twice under that id; the later fragment is now re-identified the way `_observe_tool_call_start` already re-identifies a split message, and only when a collision is real. A repeated start for the block already open is now a no-op rather than silently discarding the deltas folded into it. Detecting that collision by scanning the accumulated messages on every flush made snapshot building quadratic -- flat at about 8us per message before, rising to 124us per message by 4000 messages. Appends now go through one method that maintains an id index, which restores linear scaling, and a test pins the index against a future append that bypasses it.
GitHub Actions
Inactive
GitHub Actions
Inactive
|
Two things that belong on the PR rather than in either review thread. A performance regression I introduced and fixed. Detecting the message-id collision by scanning Live validation on real Foundry infrastructure. Two of the questions in review are about what the
Replaying the captured event order with a concurrent text message in the gap before the value arrives Validation: |
|
Please address the file conflicts when you can. Thanks. |
…reasoning-snapshot-8054 # Conflicts: # python/packages/ag-ui/agent_framework_ag_ui/_workflow.py # python/packages/ag-ui/tests/ag_ui/test_snapshots.py
GitHub Actions
Inactive
|
Done — merged Worth a note on how it was resolved, because it was not a pick-one-side conflict. #8129 changed the
One question while you are here: this is still a draft because #8054 has not had a pick between the |
|
Thanks for the update. Before this is ready, could you please:
Once those are addressed, please re-request review. Thanks! |
|
Both review discussions are resolved — each had my reply from 2026-09-08 and was waiting only on the On re-requesting review: this is still a draft for one reason, and it is a decision rather than
No pick came back, so this PR implements (a). If you would rather have (b) or (c), part of this So unless you tell me otherwise, I will take (a) as accepted, rebase on current Two judgment calls in the description are also yours rather than mechanics, and neither blocks a Current state: |
…reasoning-snapshot-8054
GitHub Actions
Active
GitHub Actions
Active
|
Please fix failing tests. |
…reasoning-snapshot-8054
GitHub Actions
Active
|
Green now, in The run was 1 failed / 12568 passed, and the single failure was Your own #8380, "stabilize Hyperlight output cleanup test", fixes it — it landed on Current state: ag-ui 1414 passing, Separately, the open question from my previous comment still stands whenever you have a moment: this The workflows on this head are at |
Motivation & Context
With an AG-UI snapshot store enabled, a
Workflowrun'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_contentdelegates to fortext_reasoningcontent, persists each reasoning message intoflow.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.observefoldsTextMessage*andToolCall*events into the synthesized snapshot but had noReasoning*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?
_WorkflowSnapshotBuildernow folds reasoning events, mirroring the existing text handling:observegains branches forReasoningMessageStartEvent,ReasoningMessageContentEvent,ReasoningMessageEndEvent/ReasoningEndEvent, andReasoningEncryptedValueEvent._observe_reasoning_*/_flush_open_reasoning_messagehelpers accumulate deltas permessage_idand emit entries in the same shape the agent path produces:{"id", "role": "reasoning", "content", ["encryptedValue"]}.Review found that the first version of this only ordered correctly for some event streams, so the
flush points are now symmetric rather than partial:
reasoning opens. The two slots can therefore never both be open, which is what makes
build()'sflush order stop deciding the sequence. Previously an
outputevent followed by a laterintermediateevent left both open and hydration replayed them in the wrong order._observe_text_content(text resuming with nostart event),
_observe_tool_call_result, and reasoning opened from a content event with nostart event.
its encrypted value used to flush an empty message that the late
ReasoningEncryptedValueEventcould no longer find, silently dropping the protected value. The empty message now stays
addressable in the position it streamed, and
build()leaves it out of the snapshot only whilenothing has claimed it, so genuinely empty reasoning is still absent.
message_idwould replaytwice under that id. The later fragment is re-identified the way
_observe_tool_call_startalready re-identifies a split message, and only when a collision is real.
_observe_text_contentopened a message over onealready open under a different id and discarded its content, where
_observe_text_startflushesfirst. It now flushes as well.
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, and it replays in the position it streamed rather than wherever
build()happened to putit.
Reasoning still does not close an open tool-call group -- verified, two tool calls with a reasoning
block between them stay in one assistant message. A reasoning row can now sit between a tool call
and its result, which is safe because
_message_adapters.py:691-694dropsrole == "reasoning"when converting back to provider messages, so the call and result come back adjacent. Nothing
asserted that across the two modules before; a test now converts the snapshot and checks it.
One behaviour change worth calling out: a text message that resumes after an interleaved reasoning
block replays as two messages, and the later fragment carries a generated id rather than the one it
streamed under. A message that never resumes keeps its original id.
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 executorpassthrough events are yielded directly and never touch
flow, soobservesees strictly morethan
_build_messages_snapshotwould. The second shape -- emitting a terminalMessagesSnapshotEventfrom the workflow runner -- is left to the issue, since an emitted snapshotis treated as authoritative and would need to reproduce those bypassing events too.
What do you want reviewers to focus on?
Two judgment calls rather than the mechanics.
First, the re-identification above. It preserves streamed order at the cost of one generated id.
The alternative is to merge the resumed fragment back into the earlier message, which keeps ids
stable but replays the reasoning after text that streamed later -- a smaller version of the defect
this PR fixes. Say if you would rather have stable ids.
Second, the empty-reasoning message being filtered at
build()rather than dropped at flush.Filtering is not observable through the single call path today, which builds once per run; it keeps
build()a projection of accumulated state rather than a mutation of it. Attaching the value onarrival instead would have been smaller but would replay the reasoning after the output that
flushed it.
A note on measurement, since two of these depend on what the provider and scheduler actually do
rather than on the folding logic: a live fan-out workflow on
gpt-5-minishows reasoning arrivingwith an encrypted value and no visible text, the same value emitted twice for one item, and
concurrent executors interleaving their events. Details are in a PR comment.
Related Issue
Fixes #8054
Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.