Skip to content

Python: Include reasoning messages in MESSAGES_SNAPSHOT events#4844

Open
moonbox3 wants to merge 4 commits intomicrosoft:mainfrom
moonbox3:agent/fix-4843-1
Open

Python: Include reasoning messages in MESSAGES_SNAPSHOT events#4844
moonbox3 wants to merge 4 commits intomicrosoft:mainfrom
moonbox3:agent/fix-4843-1

Conversation

@moonbox3
Copy link
Contributor

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_reasoning did not persist reasoning into FlowState, and _build_messages_snapshot had no awareness of reasoning messages. The fix adds a reasoning_messages list to FlowState, 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

  • The code builds clean without any errors or warnings
  • The PR follows the Contribution Guidelines
  • All unit tests pass, and I have added new tests where possible
  • Is this a breaking change? If yes, add "[BREAKING]" prefix to the title of the PR.

Note: PR autogenerated by moonbox3's agent

Copilot and others added 2 commits March 23, 2026 05:19
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>
Copilot AI review requested due to automatic review settings March 23, 2026 05:26
@moonbox3 moonbox3 self-assigned this Mar 23, 2026
Copy link
Contributor Author

@moonbox3 moonbox3 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Automated Code Review

Reviewers: 4 | Confidence: 97% | Result: All clear

Reviewed: Correctness, Security Reliability, Test Coverage, Design Approach


Automated review by moonbox3's agents

@markwallace-microsoft
Copy link
Member

markwallace-microsoft commented Mar 23, 2026

Python Test Coverage

Python Test Coverage Report •
FileStmtsMissCoverMissing
packages/ag-ui/agent_framework_ag_ui
   _agent_run.py4855089%149–156, 195–196, 203, 300, 312, 316, 318–319, 335, 362–363, 399–403, 517–519, 531–533, 631, 639, 752, 754–755, 808, 825–826, 833, 901, 924, 932, 934, 937, 943, 996, 999, 1009–1010, 1017
   _message_adapters.py5975590%102–103, 112–115, 118–122, 124–129, 132, 141–147, 187, 318, 408–411, 413–415, 462–464, 518, 521, 523, 526, 529, 545, 562, 584, 684, 700–701, 772, 794, 864, 899–900, 968, 1011
   _run_common.py2641096%311–312, 318–323, 567–568
   _utils.py100199%74
TOTAL27321321988% 

Python Unit Test Overview

Tests Skipped Failures Errors Time
5362 20 💤 0 ❌ 0 🔥 1m 22s ⏱️

Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_messages to FlowState and persist reasoning into it when reasoning content is emitted.
  • Append persisted reasoning messages into _build_messages_snapshot and 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.

Copilot and others added 2 commits March 23, 2026 05:35
- 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
Copy link
Contributor Author

@moonbox3 moonbox3 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Automated Code Review

Reviewers: 4 | Confidence: 97% | Result: All clear

Reviewed: Correctness, Security Reliability, Test Coverage, Design Approach


Automated review by moonbox3's agents

@moonbox3 moonbox3 added the ag-ui label Mar 23, 2026
@moonbox3
Copy link
Contributor Author

@robinsk fyi

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Python: agent-framework-ag-ui: include reasoning messages in MESSAGES_SNAPSHOT

3 participants