From fa37b1f7c2fe7532002bdd48a24366afda4b1b11 Mon Sep 17 00:00:00 2001 From: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> Date: Tue, 13 Oct 2026 19:44:16 +0530 Subject: [PATCH 1/3] Python: fix(ag-ui): keep repeated single-chunk text deltas A duplicate delta equal to the whole accumulated text was always treated as a full-message replay and dropped. When the message only streamed one chunk so far, that same shape is a genuine repeated token (like the 2s in 22) and the text was silently lost. Only treat the duplicate as a replay when it carries cumulative annotations or the text already streamed in more than one chunk, and track the emitted delta count on FlowState. Fixes #9336 Signed-off-by: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> --- .../agent_framework_ag_ui/_run_common.py | 15 ++++++++++++- .../agent_framework_ag_ui/_workflow_run.py | 1 + python/packages/ag-ui/tests/ag_ui/test_run.py | 22 ++++++++++++++++--- 3 files changed, 34 insertions(+), 4 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index 0f931ff88e2..f8bba08b4c2 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -537,6 +537,9 @@ class FlowState: waiting_for_approval: bool = False current_state: dict[str, Any] = field(default_factory=dict) # pyright: ignore[reportUnknownVariableType] accumulated_text: str = "" + # Number of text deltas emitted for the currently open message. Used to tell a + # genuine repeated token apart from a full-message replay chunk (#9336). + text_delta_count: int = 0 pending_tool_calls: list[dict[str, Any]] = field(default_factory=list) # pyright: ignore[reportUnknownVariableType] tool_calls_by_id: dict[str, dict[str, Any]] = field(default_factory=dict) # pyright: ignore[reportUnknownVariableType] tool_results: list[dict[str, Any]] = field(default_factory=list) # pyright: ignore[reportUnknownVariableType] @@ -641,10 +644,17 @@ def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> li if not flow.message_id: flow.message_id = generate_event_id() flow.accumulated_text = "" + flow.text_delta_count = 0 _open_text_segment(flow, flow.message_id) events.append(TextMessageStartEvent(message_id=flow.message_id, role="assistant")) - elif flow.accumulated_text and content.text == flow.accumulated_text: + elif flow.accumulated_text and content.text == flow.accumulated_text and ( + content.annotations or flow.text_delta_count > 1 + ): # Guard against full-message replay chunks that can appear after streaming deltas. + # A duplicate equal to the whole accumulated text is only a safe replay signal when + # it carries cumulative annotations or the text already streamed in more than one + # chunk. With a single prior chunk a genuine repeated token is just as likely, and + # dropping it would lose text (issue #9336). logger.debug("Skipping duplicate full-text delta for message_id=%s", flow.message_id) if not content.annotations: return [] @@ -660,6 +670,7 @@ def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> li ): flow.message_id = generate_event_id() flow.accumulated_text = "" + flow.text_delta_count = 0 events.append(TextMessageEndEvent(message_id=current_message_id)) events.append(TextMessageStartEvent(message_id=flow.message_id, role="assistant")) @@ -674,6 +685,7 @@ def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> li if content.text and not duplicate_text: events.append(TextMessageContentEvent(message_id=flow.message_id, delta=content.text)) flow.accumulated_text += content.text + flow.text_delta_count += 1 segment["text"] = flow.accumulated_text if content.annotations: annotations = cast( @@ -921,6 +933,7 @@ def _emit_tool_result_common( events.append(TextMessageEndEvent(message_id=flow.message_id)) flow.message_id = None flow.accumulated_text = "" + flow.text_delta_count = 0 return events diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 7017931580b..4656cadb8f4 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -1316,6 +1316,7 @@ def _drain_open_message() -> list[TextMessageEndEvent]: current_message_id = flow.message_id flow.message_id = None flow.accumulated_text = "" + flow.text_delta_count = 0 return [TextMessageEndEvent(message_id=current_message_id)] def _drain_open_tool_calls() -> list[ToolCallEndEvent]: diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py index 8961d331f76..db8831a7a4e 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -651,10 +651,10 @@ def test_emit_text_continues_existing_message(): def test_emit_text_skips_duplicate_full_message_delta(): - """Test _emit_text skips replayed full-message chunks on an open message.""" + """Test _emit_text skips a replay of the full text after multi-chunk streaming.""" flow = FlowState() - flow.message_id = "existing-id" - flow.accumulated_text = "Case complete." + _emit_text(Content.from_text("Case"), flow) + _emit_text(Content.from_text(" complete."), flow) content = Content.from_text("Case complete.") events = _emit_text(content, flow) @@ -663,6 +663,22 @@ def test_emit_text_skips_duplicate_full_message_delta(): assert flow.accumulated_text == "Case complete." +def test_emit_text_keeps_repeated_single_chunk_delta(): + """A repeated token equal to the whole short text streams as a real delta (#9336).""" + flow = FlowState() + first = _emit_text(Content.from_text("2"), flow) + second = _emit_text(Content.from_text("2"), flow) + assert len(first) == 2 # message start + first delta + assert len(second) == 1 + assert flow.accumulated_text == "22" + + fresh_flow = FlowState() + _emit_text(Content.from_text("ha"), fresh_flow) + repeated = _emit_text(Content.from_text("ha"), fresh_flow) + assert len(repeated) == 1 + assert fresh_flow.accumulated_text == "haha" + + def test_emit_text_skips_when_waiting_for_approval(): """Test _emit_text skips when waiting for approval.""" flow = FlowState() From d72d5b97d48880d652d4df37066b84027601d742 Mon Sep 17 00:00:00 2001 From: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> Date: Wed, 14 Oct 2026 17:43:02 +0530 Subject: [PATCH 2/3] Python: fix(ag-ui): require annotation evidence before suppressing a duplicate The replay guard still treated any annotated duplicate equal to the accumulated text as a cumulative replay, dropping a genuine repeated token that happened to carry annotations. A replay re-sends the annotations already recorded on the message segment, so the guard now suppresses only when the duplicate's annotation set covers the recorded occurrences. A repeated delta with new annotations streams normally. Follow-up to the review on #9339. Signed-off-by: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> --- .../agent_framework_ag_ui/_run_common.py | 36 +++++++++++++-- python/packages/ag-ui/tests/ag_ui/test_run.py | 44 +++++++++++++++++++ 2 files changed, 76 insertions(+), 4 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index f8bba08b4c2..23cab5e5678 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -580,6 +580,33 @@ def _text_segment_for(flow: FlowState, message_id: str) -> dict[str, Any] | None return None +def _annotations_replay_recorded_occurrences(flow: FlowState, annotations: Any) -> bool: + """Whether an equal-text duplicate's annotations re-send every recorded occurrence. + + A cumulative replay re-sends the annotations the message already carries; a genuine + repeated delta brings new annotations. Comparing against the segment's recorded + annotations separates the two without guessing from the delta alone. + """ + if not annotations: + return False + segment = _text_segment_for(flow, flow.message_id) if flow.message_id else None + if segment is None: + return False + recorded = list(segment.get("annotations", [])) + if not recorded: + return False + incoming = cast( + "list[dict[str, Any]]", + make_json_safe( + [ + {key: value for key, value in annotation.items() if key != "raw_representation"} + for annotation in annotations + ] + ), + ) + return all(annotation in incoming for annotation in recorded) + + def _new_tool_call_segment_id(flow: FlowState) -> str: """Allocate a snapshot ID, reusing a tool-only opening ID at most once.""" segment_ids = {segment.get("id") for segment in flow.snapshot_segments if segment.get("id")} @@ -648,13 +675,14 @@ def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> li _open_text_segment(flow, flow.message_id) events.append(TextMessageStartEvent(message_id=flow.message_id, role="assistant")) elif flow.accumulated_text and content.text == flow.accumulated_text and ( - content.annotations or flow.text_delta_count > 1 + flow.text_delta_count > 1 or _annotations_replay_recorded_occurrences(flow, content.annotations) ): # Guard against full-message replay chunks that can appear after streaming deltas. # A duplicate equal to the whole accumulated text is only a safe replay signal when - # it carries cumulative annotations or the text already streamed in more than one - # chunk. With a single prior chunk a genuine repeated token is just as likely, and - # dropping it would lose text (issue #9336). + # the text already streamed in more than one chunk, or when its annotations are + # cumulative: a replay re-sends the annotations already recorded for the message, + # while a genuine repeated delta carries new ones. Otherwise a single prior chunk + # leaves the two indistinguishable and dropping it would lose text (issue #9336). logger.debug("Skipping duplicate full-text delta for message_id=%s", flow.message_id) if not content.annotations: return [] diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py index db8831a7a4e..8764ecc0cb9 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -3817,3 +3817,47 @@ def capture(*args: Any, **kwargs: Any) -> Any: ] assert observed == [None], f"client state selected a conversation (use_service_session={use_service_session})" + + +def test_emit_text_keeps_repeated_annotated_single_chunk_delta() -> None: + """A genuine repeated token with NEW annotations is kept (#9336 review round). + + An equal delta after one chunk is only a replay when its annotations re-send the + recorded ones. A second identical token carrying a fresh citation streams normally. + """ + flow = FlowState() + first = _emit_text( + Content.from_text("ha", annotations=[Annotation(type="citation", url="https://example.com/a")]), flow + ) + second = _emit_text( + Content.from_text("ha", annotations=[Annotation(type="citation", url="https://example.com/b")]), flow + ) + + assert flow.accumulated_text == "haha" + text_events = [event for event in [*first, *second] if isinstance(event, TextMessageContentEvent)] + assert [event.delta for event in text_events] == ["ha", "ha"] + citation_events = [event for event in [*first, *second] if isinstance(event, CustomEvent)] + assert len(citation_events) == 2 + + +def test_emit_text_suppresses_cumulative_annotated_replay_after_one_chunk() -> None: + """An equal delta whose annotations re-send the recorded set is still a replay.""" + flow = FlowState() + first_citation = Annotation(type="citation", url="https://example.com/first") + _emit_text(Content.from_text("Answer", annotations=[first_citation]), flow) + + replay = _emit_text( + Content.from_text( + "Answer", + annotations=[ + first_citation, + Annotation(type="citation", url="https://example.com/second"), + ], + ), + flow, + ) + + assert flow.accumulated_text == "Answer" + citation_events = [event for event in replay if isinstance(event, CustomEvent)] + assert len(citation_events) == 1 + assert citation_events[0].value["annotations"][0]["url"] == "https://example.com/second" From c33d7e4d2ef4d78e85b9b04cddeadbe1e8e9f510 Mon Sep 17 00:00:00 2001 From: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:01:34 +0530 Subject: [PATCH 3/3] Python: fix(ag-ui): keep one-chunk annotated duplicates like unannotated ones A duplicate equal to the whole accumulated text was still suppressed after a single prior chunk whenever its annotations re-sent the recorded occurrences, which dropped a genuine repeated token carrying the same citation. Suppression now needs more than one prior chunk for annotated duplicates too, matching the unannotated policy from the issue. The genuine path emits the delta with its annotations, spans rebased to the new chunk position, and the existing occurrence reconciliation still applies on the suppression path. Follow-up to the review on #9339. Signed-off-by: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> --- .../agent_framework_ag_ui/_run_common.py | 37 ++--------- python/packages/ag-ui/tests/ag_ui/test_run.py | 61 ++++++++++++++++--- 2 files changed, 56 insertions(+), 42 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index f41c79148ca..edf6680ea2d 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -581,33 +581,6 @@ def _text_segment_for(flow: FlowState, message_id: str) -> dict[str, Any] | None return None -def _annotations_replay_recorded_occurrences(flow: FlowState, annotations: Any) -> bool: - """Whether an equal-text duplicate's annotations re-send every recorded occurrence. - - A cumulative replay re-sends the annotations the message already carries; a genuine - repeated delta brings new annotations. Comparing against the segment's recorded - annotations separates the two without guessing from the delta alone. - """ - if not annotations: - return False - segment = _text_segment_for(flow, flow.message_id) if flow.message_id else None - if segment is None: - return False - recorded = list(segment.get("annotations", [])) - if not recorded: - return False - incoming = cast( - "list[dict[str, Any]]", - make_json_safe( - [ - {key: value for key, value in annotation.items() if key != "raw_representation"} - for annotation in annotations - ] - ), - ) - return all(annotation in incoming for annotation in recorded) - - def _new_tool_call_segment_id(flow: FlowState) -> str: """Allocate a snapshot ID, reusing a tool-only opening ID at most once.""" segment_ids = {segment.get("id") for segment in flow.snapshot_segments if segment.get("id")} @@ -676,14 +649,14 @@ def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> li _open_text_segment(flow, flow.message_id) events.append(TextMessageStartEvent(message_id=flow.message_id, role="assistant")) elif flow.accumulated_text and content.text == flow.accumulated_text and ( - flow.text_delta_count > 1 or _annotations_replay_recorded_occurrences(flow, content.annotations) + flow.text_delta_count > 1 ): # Guard against full-message replay chunks that can appear after streaming deltas. # A duplicate equal to the whole accumulated text is only a safe replay signal when - # the text already streamed in more than one chunk, or when its annotations are - # cumulative: a replay re-sends the annotations already recorded for the message, - # while a genuine repeated delta carries new ones. Otherwise a single prior chunk - # leaves the two indistinguishable and dropping it would lose text (issue #9336). + # the text already streamed in more than one chunk. With a single prior chunk the + # duplicate is indistinguishable from a genuine repeated token, annotated or not, + # so the one-chunk policy stays the same as the unannotated path: keep the text + # (issue #9336) and let the annotation emission reconcile occurrences. logger.debug("Skipping duplicate full-text delta for message_id=%s", flow.message_id) if not content.annotations: return [] diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py index 8764ecc0cb9..78b8ddececd 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -788,7 +788,12 @@ def test_emit_text_annotation_only_update_starts_a_message_without_an_empty_delt def test_emit_text_replay_emits_only_new_annotations() -> None: - """Suppress repeated text and citations without dropping final-only citations.""" + """A one-chunk duplicate keeps its text and emits only new citation occurrences. + + After a single streamed chunk a duplicate equal to the whole text is treated as a + genuine delta, matching the unannotated policy from issue #9336. The annotation + emission still reconciles occurrences, so a re-sent citation is not duplicated. + """ flow = FlowState() first = Annotation(type="citation", url="https://example.com/first") second = Annotation(type="citation", url="https://example.com/second") @@ -797,11 +802,15 @@ def test_emit_text_replay_emits_only_new_annotations() -> None: events = _emit_text(Content.from_text("Answer", annotations=[first, second]), flow) replay_events = _emit_text(Content.from_text("Answer", annotations=[first, second]), flow) - assert len(events) == 1 - assert isinstance(events[0], CustomEvent) - assert events[0].value == {"messageId": flow.message_id, "annotations": [second]} - assert replay_events == [] - assert flow.accumulated_text == "Answer" + text_events = [event for event in events if isinstance(event, TextMessageContentEvent)] + citation_events = [event for event in events if isinstance(event, CustomEvent)] + assert [event.delta for event in text_events] == ["Answer"] + assert [event.value["annotations"] for event in citation_events] == [[first, second]] + replay_text_events = [event for event in replay_events if isinstance(event, TextMessageContentEvent)] + assert [event.delta for event in replay_text_events] == ["Answer"] + replay_citation_events = [event for event in replay_events if isinstance(event, CustomEvent)] + assert [event.value["annotations"] for event in replay_citation_events] == [[first, second]] + assert flow.accumulated_text == "AnswerAnswerAnswer" def test_emit_text_preserves_identical_annotation_occurrences() -> None: @@ -816,7 +825,11 @@ def test_emit_text_preserves_identical_annotation_occurrences() -> None: assert len(first_events) == len(second_events) == 1 assert all(isinstance(event, CustomEvent) for event in [*first_events, *second_events]) - assert replay_events == [] + # One prior text chunk: the duplicate streams as genuine text with its annotations. + replay_text_events = [event for event in replay_events if isinstance(event, TextMessageContentEvent)] + assert [event.delta for event in replay_text_events] == ["Answer"] + replay_citation_events = [event for event in replay_events if isinstance(event, CustomEvent)] + assert [event.value["annotations"] for event in replay_citation_events] == [[annotation, annotation]] @pytest.mark.parametrize("suppression", ["skip_text", "waiting_for_approval"]) @@ -3841,7 +3854,11 @@ def test_emit_text_keeps_repeated_annotated_single_chunk_delta() -> None: def test_emit_text_suppresses_cumulative_annotated_replay_after_one_chunk() -> None: - """An equal delta whose annotations re-send the recorded set is still a replay.""" + """A one-chunk cumulative duplicate keeps its text and emits only the new citation. + + Suppression of an equal duplicate needs more than one prior chunk; with one chunk + the policy matches the unannotated path and the text streams as a genuine delta. + """ flow = FlowState() first_citation = Annotation(type="citation", url="https://example.com/first") _emit_text(Content.from_text("Answer", annotations=[first_citation]), flow) @@ -3857,7 +3874,31 @@ def test_emit_text_suppresses_cumulative_annotated_replay_after_one_chunk() -> N flow, ) - assert flow.accumulated_text == "Answer" + assert flow.accumulated_text == "AnswerAnswer" + text_events = [event for event in replay if isinstance(event, TextMessageContentEvent)] + assert [event.delta for event in text_events] == ["Answer"] citation_events = [event for event in replay if isinstance(event, CustomEvent)] assert len(citation_events) == 1 - assert citation_events[0].value["annotations"][0]["url"] == "https://example.com/second" + assert [a["url"] for a in citation_events[0].value["annotations"]] == [ + "https://example.com/first", + "https://example.com/second", + ] + + # A later multi-chunk message still suppresses full-text replays. + multi_flow = FlowState() + _emit_text(Content.from_text("Pre "), multi_flow) + _emit_text(Content.from_text("fix"), multi_flow) + replay_events = _emit_text(Content.from_text("Pre fix"), multi_flow) + assert replay_events == [] + assert multi_flow.accumulated_text == "Pre fix" + + # The exact case from the review: two genuine identical annotated tokens. + same_flow = FlowState() + same = Annotation(type="citation", url="https://example.com/same") + _emit_text(Content.from_text("ha", annotations=[same]), same_flow) + repeat = _emit_text(Content.from_text("ha", annotations=[same]), same_flow) + assert same_flow.accumulated_text == "haha" + text_events = [event for event in repeat if isinstance(event, TextMessageContentEvent)] + assert [event.delta for event in text_events] == ["ha"] + repeat_citations = [event for event in repeat if isinstance(event, CustomEvent)] + assert [event.value["annotations"] for event in repeat_citations] == [[same]]