Python: Include reasoning messages in MESSAGES_SNAPSHOT events#4844
Python: Include reasoning messages in MESSAGES_SNAPSHOT events#4844moonbox3 wants to merge 4 commits intomicrosoft:mainfrom
Conversation
FlowState now tracks reasoning messages emitted during a run. _emit_text_reasoning() persists reasoning (including encrypted_value) into flow.reasoning_messages, and _build_messages_snapshot() appends them to the final MESSAGES_SNAPSHOT event. Changes: - Add reasoning_messages field to FlowState - Update _emit_text_reasoning() to accept optional flow parameter - Include reasoning_messages in _build_messages_snapshot() - Add 'reasoning' to ALLOWED_AGUI_ROLES so normalize_agui_role() preserves the role through snapshot round-trips - Skip reasoning messages in agui_messages_to_agent_framework() since they are UI-only state and should not be forwarded to LLM providers - Add regression tests for snapshot emission, encrypted value preservation, and multi-turn round-trip with reasoning Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
moonbox3
left a comment
There was a problem hiding this comment.
Automated Code Review
Reviewers: 4 | Confidence: 97% | Result: All clear
Reviewed: Correctness, Security Reliability, Test Coverage, Design Approach
Automated review by moonbox3's agents
Python Test Coverage Report •
Python Unit Test Overview
|
|||||||||||||||||||||||||||||||||||||||||||||
There was a problem hiding this comment.
Pull request overview
This PR updates the Python AG-UI runner so that streamed reasoning (text_reasoning) is persisted and included in the final MESSAGES_SNAPSHOT, preventing frontends that reconcile from snapshots from losing reasoning after streaming ends (fixes #4843).
Changes:
- Add
reasoning_messagestoFlowStateand persist reasoning into it when reasoning content is emitted. - Append persisted reasoning messages into
_build_messages_snapshotand emit snapshots when reasoning exists. - Extend AG-UI role normalization to allow
role="reasoning"and filter reasoning messages out of inbound provider/LLM history.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| python/packages/ag-ui/agent_framework_ag_ui/_run_common.py | Adds FlowState storage for reasoning and persists reasoning messages during streaming. |
| python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py | Includes reasoning messages in MESSAGES_SNAPSHOT and snapshot emission conditions. |
| python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py | Filters inbound role="reasoning" messages so they aren’t forwarded to the LLM provider. |
| python/packages/ag-ui/agent_framework_ag_ui/_utils.py | Allows reasoning as an AG-UI role. |
| python/packages/ag-ui/tests/ag_ui/test_run.py | Adds tests for reasoning persistence + snapshot inclusion. |
| python/packages/ag-ui/tests/ag_ui/test_message_adapters.py | Adds tests for multi-turn round-trip behavior with reasoning present. |
| python/packages/ag-ui/tests/ag_ui/test_utils.py | Adds role normalization test for reasoning. |
- Accumulate reasoning text per message_id (append deltas) instead of storing only the current chunk, matching flow.accumulated_text pattern - Use camelCase encryptedValue in snapshot JSON to match AG-UI protocol conventions (toolCallId, encryptedValue) - Normalize snake_case encrypted_value to encryptedValue in agui_messages_to_snapshot_format for input compatibility - Update normalize_agui_role docstring to include reasoning role - Add tests for incremental reasoning accumulation and key normalization Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…g-ui: include reasoning messages in MESSAGES_SNAPSHOT
moonbox3
left a comment
There was a problem hiding this comment.
Automated Code Review
Reviewers: 4 | Confidence: 97% | Result: All clear
Reviewed: Correctness, Security Reliability, Test Coverage, Design Approach
Automated review by moonbox3's agents
|
@robinsk fyi |
Motivation and Context
When an LLM produces reasoning (chain-of-thought) content during streaming, it was emitted as real-time reasoning events but never included in the final
MESSAGES_SNAPSHOT. Frontends that reconcile state from snapshots would lose all reasoning content after streaming ended.Fixes #4843
Description
The root cause was that
_emit_text_reasoningdid not persist reasoning intoFlowState, and_build_messages_snapshothad no awareness of reasoning messages. The fix adds areasoning_messageslist toFlowState, populates it when reasoning content is emitted, appends those messages to the snapshot in_build_messages_snapshot, and adds"reasoning"to the allowed AG-UI roles. The inbound adapter is updated to skip reasoning messages so they are not forwarded to the LLM provider on subsequent turns.Contribution Checklist