Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Manjunath Janardhan (manjunathshiva) wants to merge 4 commits into
base: main
Choose a base branch
from
+808 −4
Draft

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

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ec0a686
Python: fix: keep AG-UI workflow reasoning in thread snapshots
manjunathshiva Sep 4, 2026
55be3e9
Python: fix: keep encrypted-only AG-UI workflow reasoning on message end
manjunathshiva Sep 4, 2026
b7c83fb
Python: replay AG-UI workflow reasoning in the order it streamed
manjunathshiva Sep 8, 2026
6338c32
Merge remote-tracking branch 'origin/main' into python-agui-workflow-…
manjunathshiva Sep 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 68 additions & 0 deletions python/packages/ag-ui/agent_framework_ag_ui/_workflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,12 @@
from ag_ui.core import (
BaseEvent,
MessagesSnapshotEvent,
ReasoningEncryptedValueEvent,
ReasoningEndEvent,
ReasoningMessageContentEvent,
ReasoningMessageEndEvent,
ReasoningMessageStartEvent,
ReasoningStartEvent,
RunErrorEvent,
RunFinishedEvent,
RunStartedEvent,
Expand Down Expand Up @@ -122,6 +128,7 @@ def __init__(self, raw_messages: list[dict[str, Any]]) -> None:
self._synthesized_messages = agui_messages_to_snapshot_format(raw_messages)
self._emitted_messages: list[dict[str, Any]] | None = None
self._open_text_message: dict[str, Any] | None = None
self._open_reasoning_message: dict[str, Any] | None = None
self._tool_call_message: dict[str, Any] | None = None
self._tool_calls_by_id: dict[str, dict[str, Any]] = {}
self.state: dict[str, Any] | None = None
Expand Down Expand Up @@ -165,14 +172,27 @@ def observe(self, event: BaseEvent) -> None:
self._observe_tool_call_args(event)
elif isinstance(event, ToolCallResultEvent):
self._observe_tool_call_result(event)
elif isinstance(event, ReasoningStartEvent):
# A new reasoning block supersedes anything still open from the last one.
self._flush_open_reasoning_message()
elif isinstance(event, ReasoningMessageStartEvent):
self._observe_reasoning_start(event)
elif isinstance(event, ReasoningMessageContentEvent):
Comment on lines +187 to +192

@moonbox3 Evan Mattson (moonbox3) Sep 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

What happens when a workflow emits an output event before a later intermediate or deprecated data event? The assistant text remains in _open_text_message, then build() flushes the newer reasoning first, so hydration reverses the order that streamed. Could either reasoning-start branch flush _open_text_message before opening the reasoning block?

@manjunathshiva Manjunath Janardhan (manjunathshiva) Sep 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Confirmed, and reproduced as a failing test before changing anything: text streams, no
TextMessageEndEvent arrives because the workflow keeps running, reasoning opens, and build()
flushes reasoning first so hydration replays it before the text that streamed earlier.

Rather than reorder build(), I made the two slots peers so they can never both be open, which
removes the dependence on that order entirely. _observe_text_start and _observe_tool_call_start
already flushed open reasoning; reasoning-open paths now flush open text.

Working through it turned up three more instances of the same gap, all fixed here:

  • _observe_text_content opens a message when text resumes without a start event, and did not flush
    open reasoning. So in text -> reasoning -> text the reasoning still landed last.
  • _observe_tool_call_result did not flush open reasoning either, so reasoning that streamed before
    a tool result replayed after it. _observe_tool_call_start does flush, which is why this only
    showed on the result path.
  • Same method, worse symptom: _observe_text_content also opened a message over one already open
    under a different id, discarding its content outright, where _observe_text_start flushes first.
    Two content events with different ids and no start events keep only the second today. Concurrent
    executors reach this by interleaving content events, which is exactly what the live run does, so I
    fixed it in the same place rather than leave a silent text loss next to the flush I was adding.

I checked that the _observe_tool_call_result change is safe against the constraint the comment
there records. A reasoning block can now land between a tool call and its result, but agui_messages_to_agent_framework drops
role == "reasoning" before provider conversion, so the call and result come back adjacent. Nothing
asserted that across the two modules, so there is now a test that converts the snapshot and checks
the adjacency rather than trusting the comment.

Worth knowing that this ordering is not hypothetical: on a live fan-out workflow the reasoning, text
and tool events interleave across concurrently scheduled executors, and the alternation count varied
between 4 and 10 across three runs of the same prompt.

One consequence worth your call. Splitting an open text message means a message that resumes under
the same message_id would replay twice under that id. I re-identify the later fragment the way
_observe_tool_call_start already re-identifies a split message, and only when the id genuinely
collides, so a message that never resumes keeps the id it streamed under. That preserves streamed
order at the cost of one synthetic id.

The alternative is to merge the resumed fragment back into the earlier message, which keeps ids
stable but puts the reasoning after text that streamed later -- the defect this comment is about, in
a smaller form. I chose faithful order; say the word if you would rather have stable ids.

I checked the consumer side before settling on that, since splitting a message and changing an id
could plausibly break thread continuation. Two things make it safe, and both look deliberate rather
than accidental:

  • _snapshot_messages_match only requires equal ids for non-assistant roles. For an assistant
    message with mismatched ids it falls through to _canonical_snapshot_message, which pops id
    before comparing -- so assistant ids are already non-load-bearing for identity, which is what
    makes the existing _observe_tool_call_start re-identification safe too.
  • A conforming client accumulates content per message_id and so holds one merged message where
    the snapshot now holds two fragments. Running _reconstruct_messages_from_thread_snapshot with
    that mismatch appends only the new user turn and does not duplicate the assistant history,
    because reconstruction is backend-authoritative and client-supplied assistant messages are
    filtered out of the incoming suffix.

So the split is contained to the stored shape and does not leak into the reconstructed transcript.

self._observe_reasoning_content(event)
elif isinstance(event, (ReasoningMessageEndEvent, ReasoningEndEvent)):
self._observe_reasoning_end(event)
elif isinstance(event, ReasoningEncryptedValueEvent):
self._observe_reasoning_encrypted_value(event)

def build(self) -> AGUIThreadSnapshot:
"""Return the replayable thread snapshot."""
self._flush_open_reasoning_message()
self._flush_open_text_message()
messages = self._emitted_messages if self._emitted_messages is not None else self._synthesized_messages
return AGUIThreadSnapshot(messages=messages, state=self.state, interrupt=self.interrupt)

def _observe_text_start(self, event: TextMessageStartEvent) -> None:
self._flush_open_reasoning_message()
if self._open_text_message is not None and self._open_text_message.get("id") != event.message_id:
Comment on lines 213 to 215

@moonbox3 Evan Mattson (moonbox3) Sep 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Following up on the earlier encrypted-only case: could we keep an empty ended reasoning message addressable when intervening text starts? This flush drops the empty shell, so a later ReasoningEncryptedValueEvent finds neither _open_reasoning_message nor a synthesized message and the protected value still disappears on hydration. The new late-value path handles intervening output only when the reasoning message also had visible text.

@manjunathshiva Manjunath Janardhan (manjunathshiva) Sep 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

You're right that the value is still lost, and this is a real hole in the earlier fix rather than an
edge case. Keeping the message open past REASONING_MESSAGE_END only helps when the value arrives
before anything else flushes it. A block carrying only protected data has no content event at all,
so intervening output flushed an empty message, and the late ReasoningEncryptedValueEvent then
found neither _open_reasoning_message nor a synthesized message to attach to. My earlier test only
covered the case where the reasoning also had visible text, which is exactly the gap you identified.

I went to a live fan-out workflow on gpt-5-mini to check the shape of this against a real provider,
and it refines your description in one way worth passing on. With a flow, the encrypted value is
emitted between REASONING_MESSAGE_START and REASONING_MESSAGE_END, not after the message end -- the
order in your comment is what the no-flow branch of _emit_text_reasoning produces. So on the flow
path the vulnerable window is narrower than stated, and it is a concurrency race rather than a
deterministic emission order.

It is still a real loss. The live run confirms both preconditions: reasoning arrives with an
encrypted value and no visible text at all, and concurrent executors interleave their events.
Replaying the captured event order with a concurrent text message starting in the gap before the
value arrives loses the encrypted reasoning on the current branch and keeps it with this change.
That is now a regression test built from the captured ordering rather than a synthesized one.

Two other things from that run: the provider emits ReasoningEncryptedValueEvent twice for one
reasoning item with the same entity_id, which is idempotent here but pinned by the test; and the
saved snapshot ends up with the reasoning row positioned where it streamed and dropped before
provider conversion, with no duplicate ids.

The fix itself: the flush now keeps the empty message addressable in the position it streamed, and
build() filters it out of the snapshot only while nothing has claimed it. Genuinely empty reasoning
is still absent, so the existing test that drops reasoning with neither text nor an encrypted value
still holds. I deliberately did not attach the value on arrival instead, which would have been the
smaller change: the message would then be appended after the output that flushed it, so the value
would survive but the reasoning would replay in the wrong place, trading this bug for the one in your
other comment. Position has to be decided at flush time.

The filter compares by object identity, not by id, so a caller-supplied message from raw_messages
that happens to share an id with a synthesized one cannot be removed with it. To be straight about
the strength of that choice: build() runs once per run at the single call site, so filtering rather
than deleting is not observable through the current path. I went that way because it keeps build()
a projection of accumulated state instead of mutating it -- design hygiene, not a second defect.

One consequential cleanup while in here: _observe_reasoning_start had an id check that my change
made dead, and following it through showed a repeated start for the block already open used to
discard the deltas already folded into it. It is now a no-op, with a test.

self._flush_open_text_message()
self._open_text_message = {"id": event.message_id, "role": event.role, "content": ""}
Expand All @@ -188,6 +208,7 @@ def _observe_text_end(self, event: TextMessageEndEvent) -> None:
self._flush_open_text_message()

def _observe_tool_call_start(self, event: ToolCallStartEvent) -> None:
self._flush_open_reasoning_message()
parent_message_id = event.parent_message_id
if (
self._open_text_message is not None
Expand Down Expand Up @@ -245,6 +266,53 @@ def _flush_open_text_message(self) -> None:
self._tool_call_message = None
self._open_text_message = None

def _observe_reasoning_start(self, event: ReasoningMessageStartEvent) -> None:
if self._open_reasoning_message is not None and self._open_reasoning_message.get("id") != event.message_id:
self._flush_open_reasoning_message()
self._open_reasoning_message = {"id": event.message_id, "role": "reasoning", "content": ""}

def _observe_reasoning_content(self, event: ReasoningMessageContentEvent) -> None:
if self._open_reasoning_message is None or self._open_reasoning_message.get("id") != event.message_id:
self._flush_open_reasoning_message()
self._open_reasoning_message = {"id": event.message_id, "role": "reasoning", "content": ""}
self._open_reasoning_message["content"] = f"{self._open_reasoning_message.get('content', '')}{event.delta}"

def _observe_reasoning_end(self, event: ReasoningMessageEndEvent | ReasoningEndEvent) -> None:
if self._open_reasoning_message is None or self._open_reasoning_message.get("id") != event.message_id:
return
if isinstance(event, ReasoningMessageEndEvent):
# REASONING_MESSAGE_END closes the message, not the block, and an
# encrypted value is block-scoped so it legitimately trails it -- that is
# the order `_emit_text_reasoning` produces without a flow. Keep the
# message open so protected-data-only reasoning still has somewhere to
# land; REASONING_END, the next block, or build() finalizes it.
return
self._flush_open_reasoning_message()
Comment thread
manjunathshiva marked this conversation as resolved.

def _observe_reasoning_encrypted_value(self, event: ReasoningEncryptedValueEvent) -> None:
# Only message-scoped encrypted values belong on a reasoning message; the
# protocol also uses this event for other subtypes.
if event.subtype != "message":
return
if self._open_reasoning_message is not None and self._open_reasoning_message.get("id") == event.entity_id:
self._open_reasoning_message["encryptedValue"] = event.encrypted_value
return
# Intervening text or tool output can flush the message before its encrypted
# value arrives; attach it to the message we already synthesized.
for message in reversed(self._synthesized_messages):
if message.get("role") == "reasoning" and message.get("id") == event.entity_id:
message["encryptedValue"] = event.encrypted_value
return

def _flush_open_reasoning_message(self) -> None:
if self._open_reasoning_message is None:
return
# An encrypted-value-only block carries no display text but still has to
# survive hydration, so it counts as content worth keeping.
if self._open_reasoning_message.get("content") or self._open_reasoning_message.get("encryptedValue"):
self._synthesized_messages.append(self._open_reasoning_message)
self._open_reasoning_message = None


class AgentFrameworkWorkflow:
"""Base AG-UI workflow wrapper.
Expand Down
Loading
Loading