Python: fix(ag-ui): keep repeated single-chunk text deltas

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

@@ -538,6 +538,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]
@@ -642,10 +645,18 @@ 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 (
+        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. 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 []
@@ -661,6 +672,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"))
 
@@ -675,6 +687,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(
@@ -922,6 +935,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
 

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

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()
@@ -772,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")
@@ -781,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:
@@ -800,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"])
@@ -3801,3 +3830,75 @@ 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:
+    """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)
+
+    replay = _emit_text(
+        Content.from_text(
+            "Answer",
+            annotations=[
+                first_citation,
+                Annotation(type="citation", url="https://example.com/second"),
+            ],
+        ),
+        flow,
+    )
+
+    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 [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]]