From 4468385431c441dd8935ee8d14bb92f8e8a1167a Mon Sep 17 00:00:00 2001 From: suhasagg Date: Mon, 5 Oct 2026 19:09:08 +0530 Subject: [PATCH 1/2] fix(core): make text stream coalescing linear --- .../packages/core/agent_framework/_types.py | 93 ++++++++++++++++--- python/packages/core/tests/core/test_types.py | 30 ++++++ 2 files changed, 112 insertions(+), 11 deletions(-) diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index f98743794a4..ffaa28695eb 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -2304,35 +2304,106 @@ def _coalesce_text_content(contents: list[Content], type_str: Literal["text", "t """Take any subsequence Text or TextReasoningContent items and coalesce them into a single item.""" if not contents: return - coalesced_contents: list[Content] = [] + + if type_str == "text": + coalesced_contents: list[Content] = [] + run: list[Content] = [] + + def flush_run() -> None: + if not run: + return + + first = deepcopy(run[0]) + + if len(run) == 1: + coalesced_contents.append(first) + run.clear() + return + + first.text = "".join(content.text for content in run) + + # Earlier chunks take precedence for duplicate properties, matching + # _combine_additional_props and the existing repeated-add behavior. + additional_properties = first.additional_properties + for content in run[1:]: + for key, value in content.additional_properties.items(): + additional_properties.setdefault(key, value) + + annotation_groups = [content.annotations for content in run if content.annotations is not None] + if not annotation_groups: + first.annotations = None + elif len(annotation_groups) == 1: + first.annotations = annotation_groups[0] + else: + first.annotations = [annotation for annotations in annotation_groups for annotation in annotations] + + # deepcopy(run[0]) intentionally discards the first chunk's raw + # representation, matching the existing coalescing behavior. + raw_representations = [ + content.raw_representation for content in run[1:] if content.raw_representation is not None + ] + + if not raw_representations: + first.raw_representation = None + elif len(raw_representations) == 1: + first.raw_representation = raw_representations[0] + else: + first.raw_representation = [ + item for raw in raw_representations for item in (raw if isinstance(raw, list) else [raw]) + ] + + coalesced_contents.append(first) + run.clear() + + current_output_kind: Any = None + + for content in contents: + if content.type != "text": + flush_run() + coalesced_contents.append(content) + current_output_kind = None + continue + + output_kind = content.additional_properties.get(_MODEL_OUTPUT_KIND_KEY) + + if run and current_output_kind != output_kind: + flush_run() + + if not run: + current_output_kind = output_kind + + run.append(content) + + flush_run() + + contents.clear() + contents.extend(coalesced_contents) + return + + # Preserve the existing reasoning-content behavior. Reasoning chunks have + # additional merge boundaries involving ids, summaries and protected data. + coalesced_contents = [] first_new_content: Any | None = None + for content in contents: if content.type == type_str: if first_new_content is None: first_new_content = deepcopy(content) - elif type_str == "text" and first_new_content.additional_properties.get( - _MODEL_OUTPUT_KIND_KEY - ) != content.additional_properties.get(_MODEL_OUTPUT_KIND_KEY): - coalesced_contents.append(first_new_content) - first_new_content = deepcopy(content) else: try: first_new_content += content except AdditionItemMismatch: - # Different IDs means a new logical segment; flush the current one coalesced_contents.append(first_new_content) first_new_content = deepcopy(content) else: - # skip this content, it is not of the right type - # so write the existing one to the list and start a new one, - # once the right type is found again if first_new_content: coalesced_contents.append(first_new_content) first_new_content = None - # but keep the other content in the new list coalesced_contents.append(content) + if first_new_content: coalesced_contents.append(first_new_content) + contents.clear() contents.extend(coalesced_contents) diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 05d474687f4..15582cbbb86 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -2453,6 +2453,36 @@ def updates(content_type: Literal["text", "text_reasoning"]) -> list[ChatRespons assert [content.text for content in response.messages[0].contents] == ["Hello world"] +def test_text_coalescing_does_not_use_repeated_content_add(monkeypatch: pytest.MonkeyPatch) -> None: + """Coalescing many text chunks should aggregate them without repeated Content addition.""" + from agent_framework._types import _coalesce_text_content + + contents = [ + Content.from_text( + str(index), + additional_properties={"chunk": index}, + ) + for index in range(100) + ] + + add_calls = 0 + original_add = Content.__add__ + + def counting_add(self: Content, other: Content) -> Content: + nonlocal add_calls + add_calls += 1 + return original_add(self, other) + + monkeypatch.setattr(Content, "__add__", counting_add) + + _coalesce_text_content(contents, "text") + + assert len(contents) == 1 + assert contents[0].text == "".join(str(index) for index in range(100)) + assert contents[0].additional_properties == {"chunk": 0} + assert add_calls == 0 + + def test_text_reasoning_content_add_coverage(): """Test TextReasoningContent __add__ method for better coverage.""" From 3e8194ad9d16e7e1a82440da58d2c674cf1a2d5d Mon Sep 17 00:00:00 2001 From: suhasagg Date: Mon, 5 Oct 2026 19:32:08 +0530 Subject: [PATCH 2/2] fix(core): preserve copied annotations during coalescing --- .../packages/core/agent_framework/_types.py | 8 +++++-- python/packages/core/tests/core/test_types.py | 21 +++++++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index ffaa28695eb..a2035b14fdf 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -2329,14 +2329,18 @@ def flush_run() -> None: for key, value in content.additional_properties.items(): additional_properties.setdefault(key, value) - annotation_groups = [content.annotations for content in run if content.annotations is not None] + annotation_groups = [ + first.annotations, + *(content.annotations for content in run[1:] if content.annotations is not None), + ] + annotation_groups = [annotations for annotations in annotation_groups if annotations is not None] + if not annotation_groups: first.annotations = None elif len(annotation_groups) == 1: first.annotations = annotation_groups[0] else: first.annotations = [annotation for annotations in annotation_groups for annotation in annotations] - # deepcopy(run[0]) intentionally discards the first chunk's raw # representation, matching the existing coalescing behavior. raw_representations = [ diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 15582cbbb86..94366a69875 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -2483,6 +2483,27 @@ def counting_add(self: Content, other: Content) -> Content: assert add_calls == 0 +def test_text_coalescing_does_not_alias_first_chunk_annotations() -> None: + """Coalescing should preserve the deep-copied annotations from the first chunk.""" + from agent_framework._types import _coalesce_text_content + + annotation: Annotation = { + "type": "citation", + "url": "https://example.com", + } + first = Content.from_text("Hello", annotations=[annotation]) + second = Content.from_text(" world") + contents = [first, second] + + _coalesce_text_content(contents, "text") + + assert len(contents) == 1 + assert contents[0].text == "Hello world" + assert contents[0].annotations == first.annotations + assert contents[0].annotations is not first.annotations + assert contents[0].annotations[0] is not first.annotations[0] + + def test_text_reasoning_content_add_coverage(): """Test TextReasoningContent __add__ method for better coverage."""