Python: Include reasoning messages in MESSAGES_SNAPSHOT events (#4844)

* Include reasoning messages in MESSAGES_SNAPSHOT (#4843)

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>

* Python: Include reasoning messages in MESSAGES_SNAPSHOT events

Fixes #4843

* Fix PR review feedback for reasoning persistence (#4843)

- 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>

* Address review feedback for #4843: Python: agent-framework-ag-ui: include reasoning messages in MESSAGES_SNAPSHOT

---------

Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py

@@ -684,6 +684,10 @@ def _build_messages_snapshot(
             }
         )
 
+    # Add reasoning messages so frontends that reconcile state from
+    # MESSAGES_SNAPSHOT retain reasoning content after streaming ends.
+    all_messages.extend(flow.reasoning_messages)
+
     return MessagesSnapshotEvent(messages=all_messages)  # type: ignore[arg-type]
 
 
@@ -1061,7 +1065,9 @@ async def run_agent_stream(
 
     # Emit MessagesSnapshotEvent if we have tool calls or results
     # Feature #5: Suppress intermediate snapshots for predictive tools without confirmation
-    should_emit_snapshot = flow.pending_tool_calls or flow.tool_results or flow.accumulated_text
+    should_emit_snapshot = (
+        flow.pending_tool_calls or flow.tool_results or flow.accumulated_text or flow.reasoning_messages
+    )
     if should_emit_snapshot:
         # Check if we should suppress for predictive tool
         last_tool_name = None

python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py

@@ -604,6 +604,10 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Mes
         # Handle standard tool result messages early (role="tool") to preserve provider invariants
         # This path maps AG‑UI tool messages to function_result content with the correct tool_call_id
         role_str = normalize_agui_role(msg.get("role", "user"))
+        if role_str == "reasoning":
+            # Reasoning messages are UI-only state carried in MESSAGES_SNAPSHOT.
+            # They should not be forwarded to the LLM provider.
+            continue
         if role_str == "tool":
             # Prefer explicit tool_call_id fields; fall back to backend fields only if necessary
             tool_call_id = msg.get("tool_call_id") or msg.get("toolCallId")
@@ -1020,6 +1024,11 @@ def agui_messages_to_snapshot_format(messages: list[dict[str, Any]]) -> list[dic
             elif "toolCallId" not in normalized_msg:
                 normalized_msg["toolCallId"] = ""
 
+        # Normalize encrypted_value to encryptedValue for reasoning messages
+        if normalized_msg.get("role") == "reasoning" and "encrypted_value" in normalized_msg:
+            normalized_msg["encryptedValue"] = normalized_msg["encrypted_value"]
+            del normalized_msg["encrypted_value"]
+
         result.append(normalized_msg)
 
     return result

python/packages/ag-ui/agent_framework_ag_ui/_run_common.py

@@ -126,6 +126,8 @@ class FlowState:
     tool_results: list[dict[str, Any]] = field(default_factory=list)  # pyright: ignore[reportUnknownVariableType]
     tool_calls_ended: set[str] = field(default_factory=set)  # pyright: ignore[reportUnknownVariableType]
     interrupts: list[dict[str, Any]] = field(default_factory=list)  # pyright: ignore[reportUnknownVariableType]
+    reasoning_messages: list[dict[str, Any]] = field(default_factory=list)  # pyright: ignore[reportUnknownVariableType]
+    accumulated_reasoning: dict[str, str] = field(default_factory=dict)  # pyright: ignore[reportUnknownVariableType]
 
     def get_tool_name(self, call_id: str | None) -> str | None:
         """Get tool name by call ID."""
@@ -460,7 +462,7 @@ def _emit_mcp_tool_result(
     return _emit_tool_result_common(content.call_id, raw_output, flow, predictive_handler)
 
 
-def _emit_text_reasoning(content: Content) -> list[BaseEvent]:
+def _emit_text_reasoning(content: Content, flow: FlowState | None = None) -> list[BaseEvent]:
     """Emit AG-UI reasoning events for text_reasoning content.
 
     Uses the protocol-defined reasoning event types so that AG-UI consumers
@@ -470,6 +472,10 @@ def _emit_text_reasoning(content: Content) -> list[BaseEvent]:
     ``content.protected_data`` is present it is emitted as a
     ``ReasoningEncryptedValueEvent`` so that consumers can persist encrypted
     reasoning for state continuity without conflating it with display text.
+
+    When *flow* is provided the reasoning message is persisted into
+    ``flow.reasoning_messages`` so that ``_build_messages_snapshot`` can
+    include it in the final ``MESSAGES_SNAPSHOT``.
     """
     text = content.text or ""
     if not text and content.protected_data is None:
@@ -498,6 +504,36 @@ def _emit_text_reasoning(content: Content) -> list[BaseEvent]:
 
     events.append(ReasoningEndEvent(message_id=message_id))
 
+    # Persist reasoning into flow state for MESSAGES_SNAPSHOT.
+    # Accumulate reasoning text per message_id, similar to flow.accumulated_text,
+    # so that incremental deltas build the full reasoning string.
+    if flow is not None:
+        if text:
+            previous_text = flow.accumulated_reasoning.get(message_id, "")
+            flow.accumulated_reasoning[message_id] = previous_text + text
+        full_text = flow.accumulated_reasoning.get(message_id, text or "")
+
+        # Update existing reasoning entry for this message_id if present; otherwise append a new one.
+        existing_entry: dict[str, Any] | None = None
+        for entry in flow.reasoning_messages:
+            if isinstance(entry, dict) and entry.get("id") == message_id:
+                existing_entry = entry
+                break
+
+        if existing_entry is None:
+            reasoning_entry: dict[str, Any] = {
+                "id": message_id,
+                "role": "reasoning",
+                "content": full_text,
+            }
+            if content.protected_data is not None:
+                reasoning_entry["encryptedValue"] = content.protected_data
+            flow.reasoning_messages.append(reasoning_entry)
+        else:
+            existing_entry["content"] = full_text
+            if content.protected_data is not None:
+                existing_entry["encryptedValue"] = content.protected_data
+
     return events
 
 
@@ -527,6 +563,6 @@ def _emit_content(
     if content_type == "mcp_server_tool_result":
         return _emit_mcp_tool_result(content, flow, predictive_handler)
     if content_type == "text_reasoning":
-        return _emit_text_reasoning(content)
+        return _emit_text_reasoning(content, flow)
     logger.debug("Skipping unsupported content type in AG-UI emitter: %s", content_type)
     return []

python/packages/ag-ui/agent_framework_ag_ui/_utils.py

@@ -27,7 +27,7 @@ FRAMEWORK_TO_AGUI_ROLE: dict[str, str] = {
     "system": "system",
 }
 
-ALLOWED_AGUI_ROLES: set[str] = {"user", "assistant", "system", "tool"}
+ALLOWED_AGUI_ROLES: set[str] = {"user", "assistant", "system", "tool", "reasoning"}
 
 
 def generate_event_id() -> str:
@@ -82,7 +82,7 @@ def normalize_agui_role(raw_role: Any) -> str:
         raw_role: Raw role value from AG-UI message
 
     Returns:
-        Normalized role string (user, assistant, system, or tool)
+        Normalized role string (user, assistant, system, tool, or reasoning)
     """
     if not isinstance(raw_role, str):
         return "user"