Skip to content

fix(integrations): merge AG-UI text message chunk continuations instead of dropping them - #3023

Open
BlueX888 wants to merge 2 commits into
vibrantlabsai:mainfrom
BlueX888:fix/prep-agui-text-chunk-continuation-dropped
Open

fix(integrations): merge AG-UI text message chunk continuations instead of dropping them#3023
BlueX888 wants to merge 2 commits into
vibrantlabsai:mainfrom
BlueX888:fix/prep-agui-text-chunk-continuation-dropped

Conversation

@BlueX888

@BlueX888 BlueX888 commented Sep 20, 2026 •
edited
Loading

Copy link
Copy Markdown

Issue Link / Problem Description

No tracking issue — reporting the defect directly, per the contributing guide.

TextMessageChunkEvent is the AG-UI shorthand that "stands in for a start, content and end sequence, for producers that cannot know in advance where a message begins" (its own model docstring). Every field on it is optional: role is documented as "Who the message is from, on the chunk that opens it", message_id as "The message this chunk belongs to. Absent continues the message already open", and delta carries the fragment to append.

AGUIEventCollector._handle_text_message_chunk treated each chunk as a complete, self-contained message. It read the role with getattr(event, "role", "assistant") — a default that never applies, because the field exists and is None on continuation chunks — so any chunk after the first fell through to the unexpected-role branch, logged a warning, and its delta was discarded. A streamed response was therefore truncated to its first chunk.

Reproduction on the current main:

from ag_ui.core import TextMessageChunkEvent
from ragas.integrations.ag_ui import convert_to_ragas_messages, extract_response

messages = convert_to_ragas_messages([
    TextMessageChunkEvent(message_id="msg-1", role="assistant", delta="The weather "),
    TextMessageChunkEvent(message_id="msg-1", delta="in SF is "),
    TextMessageChunkEvent(message_id="msg-1", delta="sunny, 72F."),
])
# AIMessage(content='The weather ')   <- trailing chunks dropped
# WARNING ragas.integrations.ag_ui:ag_ui.py:553 Unexpected message role in chunk event: None

Changes Made

src/ragas/integrations/ag_ui.py

  • _handle_text_message_chunk now keeps the position of the chunk message currently open (_active_text_chunk, reset to None in clear()). A chunk that repeats that message_id, or omits message_id entirely, appends its delta to that open message and returns, so the message keeps the role it was opened with.
  • A chunk with a new message_id, or the first chunk of a stream, still opens a message exactly as before. Defaulting role to "assistant" moved after the continuation check, so it only applies when a message is actually opened — the fallback the code always intended.
  • The now-dead role value at append time made the AIMessage/HumanMessage construction identical in both branches; the messages are appended once at the end and the open-chunk position recorded. No behaviour change for single-chunk events (the documented example) or for metadata=True.

Not changed: the TextMessageStart/Content/End triad path, ToolCallChunk, MessagesSnapshot, and extract_response.

Testing

How to Test

  • Automated tests added/updated

Three regression tests in tests/unit/integrations/test_ag_ui.py, next to the existing test_text_message_chunk:

  • test_text_message_chunk_continuation_is_merged_into_open_message — three chunks sharing one message_id, role only on the first; asserts one AIMessage whose content is the full concatenation and that extract_response returns it.
  • test_text_message_chunk_without_message_id_continues_open_message — a producer that omits message_id on later chunks.
  • test_text_message_chunk_continuation_keeps_opening_role — a role="user" message keeps HumanMessage type through its continuation chunk.

Results, all run with ./.venv/bin/python -m pytest -p no:randomly -q:

Command Unpatched tree Patched tree
pytest tests/unit/integrations/test_ag_ui.py 3 failed, 43 passed in 0.13s 46 passed in 0.35s
pytest tests/unit/integrations — 73 passed in 0.08s
pytest tests/unit (4 environment-broken modules excluded) 789 passed, 62 skipped 789 passed, 62 skipped
ruff format --check / ruff check on the two changed files — 2 files already formatted / All checks passed!

Unpatched failure for the first test:

>       assert messages[0].content == "The weather in SF is sunny, 72F."
E       AssertionError: assert 'The weather ' == 'The weather ...s sunny, 72F.'
E         - The weather in SF is sunny, 72F.
E         + The weather
tests/unit/integrations/test_ag_ui.py:549: AssertionError
WARNING  ragas.integrations.ag_ui:ag_ui.py:553 Unexpected message role in chunk event: None

Local sandbox notes (not attributable to this change): test_chrf_score.py, test_import.py and test_prechunked_generation.py need sacrebleu/haystack that are not installed here, and TestVersionExperiment needs a git remote layout absent in this worktree; all of them fail identically on the unpatched tree. pyright reports 6 errors in integrations/langsmith.py, tracing/langfuse.py and tracing/mlflow.py on both trees, none in ag_ui.py.

References

Screenshots/Examples (if applicable)

Patch this before / after:

before: AIMessage(content='The weather ')                                  # 2 of 3 chunks lost
after:  AIMessage(content='The weather in SF is sunny, 72F.')

Known limitation (not addressed here)

Two malformed-stream cases from the AG-UI streaming spec are tolerated rather than rejected, and are outside this fix:

  • A TEXT_MESSAGE_CHUNK with no message_id that arrives after a closing event opens a new assistant message. The spec wants an opening chunk to carry messageId; the text is kept and the closed message stays intact.
  • A repeat of the open message_id that declares a different role keeps the role the message was opened with and ignores the conflicting value.

Both need explicit stream validation, which changes behavior for streams that currently produce a response, so they belong in a separate change with their own tests. contextablemark, who wrote the original integration, reviewed e27853f and asked for this to be recorded rather than folded into the fix.

@BlueX888
fix: merge AG-UI text message chunk continuations into the open message
f696667 
TextMessageChunkEvent.role is None on every chunk except the one that
opens a message, so continuation chunks fell through to the unexpected
role branch and their delta was discarded, truncating the collected
response. Chunks that repeat the open message_id, or omit it, now extend
that message and inherit its role instead of being dropped.
@contextablemark

contextablemark commented Oct 11, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Thanks for this. I wrote the original AG-UI integration (#2395), reproduced this defect on main, and validated this patch locally. It correctly fixes the reported truncation. Before merge I'd ask for one small follow-up, described below.

The bug (confirmed): TextMessageChunkEvent carries role only on the chunk that opens a message. Every field on the model is optional, and role/delta default to None. But _handle_text_message_chunk treated every chunk as a complete message on its own. Follow-on chunks fell through to the unexpected-role branch and their deltas were discarded. I got exactly the reported truncation: 3 chunks → AIMessage(content='The weather ') plus Unexpected message role in chunk event: None warnings.

Validation (exact patch at f696667):

Edge cases I tested beyond the PR's tests, all handled correctly:

  • delta=None on an opening chunk now produces an empty message. Unpatched, it raised a pydantic ValidationError (AIMessage.content must be str). On a follow-on chunk, delta=None appends nothing.
  • A follow-on chunk claiming a different role is ignored. The open message keeps the role it was opened with.
  • metadata=True runs: follow-on chunks inherit the opening chunk's metadata, and single-chunk events behave exactly as before.

Requested follow-up: close the open chunk message at stream boundaries. Tracking only one open chunk message is correct. Per the AG-UI spec (Closing a chunk stream), a chunk that opens a different message closes the previous one. The spec also says the open chunk stream closes when:

  • a message, tool-call or step event arrives (RAW, ACTIVITY_SNAPSHOT, ACTIVITY_DELTA and REASONING_ENCRYPTED_VALUE are the exceptions), or
  • a run-level event arrives (RUN_STARTED, RUN_FINISHED, RUN_ERROR, MESSAGES_SNAPSHOT).

_active_text_chunk is only reset in clear(), so a later chunk without message_id attaches to a stale message. On this branch:

Event sequence Result
chunk(m1, "Checking.") → tool call → tool result → chunk(no id, "It is sunny.") AIMessage('Checking.It is sunny.') placed before the ToolMessage
chunk(m1) in run 1 → RUN_FINISHED → RUN_STARTED → chunk(no id) merged into run 1's message
chunk(m1) → Start/Content/End(m2) → chunk(no id) appended to m1
chunk(m1, assistant) → chunk(s1, system) → chunk(s1, no role) system chunk rejected with a warning but the slot isn't updated, so the follow-on text becomes an AIMessage

Strictly, a chunk with no message_id after a closing event is a producer error. Still, since this handler accepts follow-on chunks without an id, it shouldn't let them attach to a closed message. A minimal fix:

  • in process_event, set self._active_text_chunk = None for any event other than TEXT_MESSAGE_CHUNK (excluding the spec's no-op events above);
  • also reset it in the unexpected-role branch before return.

A regression test for the tool-call row would cover it.

@anistark, you reviewed and merged the original integration. Could you take a look when you have a chance, or point us to the current maintainer for src/ragas/integrations/?

— Mark

Edit: corrected my earlier note on delta=None (unpatched code raised a validation error and did not append "None"). Also replaced my earlier design note about tracking only one open chunk message, which had the spec semantics backwards, with the boundary issue above.

@BlueX888

BlueX888 commented Oct 12, 2026

Copy link
Copy Markdown
Author

Thanks @contextablemark — this is an exceptionally thorough validation, and the boundary table made the gap easy to see.

Pushed e27853f with the requested follow-up:

  • process_event now closes the open chunk message (_active_text_chunk = None) on any event other than TEXT_MESSAGE_CHUNK, except the spec's no-op events (RAW, ACTIVITY_SNAPSHOT, ACTIVITY_DELTA, REASONING_ENCRYPTED_VALUE). The no-op set is resolved via getattr since older releases (e.g. 0.1.9) don't define all four members.
  • The unexpected-role branch also clears the slot before returning, per your suggestion.
  • Added regression tests for the tool-call boundary row and the run-boundary row from your table, plus a guard test that a RAW event between chunks does not close the open message.

One behavior note: after a boundary, a chunk with no message_id still isn't dropped outright — the pre-existing role or "assistant" defaulting makes it open a new AIMessage (so "It is sunny." lands after the ToolMessage instead of merging into "Checking."). The closed message is no longer corrupted, which is what your fix targets; if you'd rather have roleless chunks with no open message dropped entirely, that's a small follow-up I'm happy to make.

Verified: 49 passed on both ag-ui-protocol==1.0.0 and ==0.1.9; ruff and pyright clean. Before the fix the two boundary tests fail exactly per your table ('Checking.It is sunny.' != 'Checking.').

Also noting for @anistark: the check-docs failure on these PRs happens at the checkout step — the pull_request_target workflow refuses to check out fork PR heads — so it's a repo-side workflow guard affecting fork PRs, not something in the patch.

@contextablemark

contextablemark commented Oct 12, 2026

Copy link
Copy Markdown
Contributor

Thanks for addressing the stream-boundary request in e27853f. I reviewed the update and reran the AG-UI test file against the updated PR head: 49 passed on both ag-ui-protocol==0.1.9 and ==1.0.0. The tool/run boundary regressions and the RAW preservation test pass. The stale-message corruption I called out is fixed.

On your question about chunks arriving without an ID after a boundary: the current AG-UI 1.0 streaming spec requires an opening text chunk to carry messageId, and requires conflicting repeated opening fields to be treated as protocol violations. I confirmed that the current patch still:

  • opens a new assistant message for an ID-less chunk after a boundary;
  • silently retains the original role when a same-ID continuation declares a conflicting role.

Those behaviors mean I would not describe the handler as fully protocol-compliant. However, strict rejection may reasonably belong in a separate follow-up, rather than expanding this truncation/boundary fix into a broader validation change. From my perspective, I would not hold this fix up solely for that work, provided the remaining behavior is documented and tracked. A follow-up should surface malformed streams explicitly rather than silently discard text. The final scope decision belongs to the maintainers.

I also inspected the failing check-docs log: checkout is refused because this is fork code in a privileged pull_request_target workflow. That corroborates your explanation; the failure occurs before patch tests run. It needs a safe maintainer-side workflow resolution, not simply bypassing the checkout security guard.

One important clarification: I’m the author of the original AG-UI integration, not a Ragas maintainer. My comments are contributor review and validation, not merge authority. Unfortunately, the project appears to have stalled on maintenance: the latest merge I could verify was February 24, 2026, with none in the last 90 days despite ongoing contributor activity. I don’t know whether that is a pause or a handover, so I can’t promise when this will receive maintainer attention.

I support the updated fix for the reported truncation and boundary defects. Thanks for responding so thoroughly.

— Mark

@BlueX888

BlueX888 commented Oct 25, 2026

Copy link
Copy Markdown
Author

Recorded both cases in the PR description, under "Known limitation (not addressed here)": a TEXT_MESSAGE_CHUNK without message_id after a closing event opens a new assistant message instead of being rejected, and a same-ID continuation with a conflicting role keeps the role the message was opened with. The section also states why strict rejection is separate work: it changes behavior for streams that currently produce a response and needs its own tests.

Head is still e27853f. No issue was filed for the follow-up, since whether the handler should reject or tolerate those streams is a scope call for the maintainers.

check-docs on this head is the fork-checkout refusal in claude-docs-check.yml, failing before any analysis of the diff; #3038 is open against the same workflow.

This branch has not been deployed

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

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@BlueX888 @contextablemark