From 9e799806ad932bb553a5d0bd96b6d250868e14c9 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Fri, 2 Oct 2026 04:08:45 +0800 Subject: [PATCH 1/3] Python: fix(core): fold streamed text runs in one pass instead of repeated += _coalesce_text_content folded chunks with first_new_content += content, and every Content.__add__ rebuilds the text string, re-merges additional_properties, and re-concatenates raw_representation lists, so aggregating n chunks cost O(n^2) (#8908). Collect each mergeable run and fold it in one pass in _merge_content_run: one join per segment, a single props merge (the earliest chunk still wins per key), one flattened raw-representation list. The fold reproduces the old semantics exactly, including the run head's raw representation being dropped (the old path deep-copied the head, and deepcopy discards raw_representation via _SHALLOW_COPY_FIELDS) and the text_reasoning id / reasoning-text split boundaries. 16000 chunks: 0.302s -> 0.011s on this machine. Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> --- .../packages/core/agent_framework/_types.py | 114 ++++++++++++---- python/packages/core/tests/core/test_types.py | 127 ++++++++++++++++++ 2 files changed, 214 insertions(+), 27 deletions(-) diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index ba1703c0caf..c84cd1b6800 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -2280,39 +2280,99 @@ def _merge_function_call_content(message: Message, content: Content) -> None: message.contents.append(content) +def _merge_content_run(run: list[Content], type_str: Literal["text", "text_reasoning"]) -> Content: + """Fold a collected run of same-type contents in one pass. + + The result matches folding the run with repeated ``+=``: the earliest chunk + wins per additional-property key, annotations concatenate in order, and raw + representations flatten into one list (left untouched when a single chunk + carries the only value). The run head's own raw representation is dropped, + matching the old path whose initial ``deepcopy`` discards it + (``_SHALLOW_COPY_FIELDS``). + """ + merged_props: dict[str, Any] = {} + for content in reversed(run): + merged_props.update(content.additional_properties) + annotation_seqs = [c.annotations for c in run if c.annotations is not None] + annotations = [a for seq in annotation_seqs for a in seq] if annotation_seqs else None + non_null_raws = [c.raw_representation for c in run[1:] if c.raw_representation is not None] + raw_representation: Any + if not non_null_raws: + raw_representation = None + elif len(non_null_raws) == 1: + raw_representation = non_null_raws[0] + else: + raw_representation = [] + for raw in non_null_raws: + raw_representation.extend(raw if isinstance(raw, list) else [raw]) + if type_str == "text": + return Content( + "text", + text="".join(c.text or "" for c in run), + annotations=annotations, + additional_properties=merged_props, + raw_representation=raw_representation, + ) + return Content( + "text_reasoning", + id=next((c.id for c in run if c.id), None), + text=None if all(c.text is None for c in run) else "".join(c.text or "" for c in run), + protected_data=next((c.protected_data for c in reversed(run) if c.protected_data is not None), None), + annotations=annotations, + additional_properties=merged_props, + raw_representation=raw_representation, + ) + + def _coalesce_text_content(contents: list[Content], type_str: Literal["text", "text_reasoning"]) -> None: - """Take any subsequence Text or TextReasoningContent items and coalesce them into a single item.""" + """Take any subsequence Text or TextReasoningContent items and coalesce them into a single item. + + Mergeable runs are collected and folded in a single pass, so aggregating n + chunks costs O(n) instead of the O(n^2) of repeated ``+=``. + """ if not contents: return coalesced_contents: list[Content] = [] - first_new_content: Any | None = None + run: list[Content] = [] + run_id: str | None = None + run_has_text = False + run_has_reasoning_text = False + + def flush() -> None: + nonlocal run_id, run_has_text, run_has_reasoning_text + if run: + coalesced_contents.append(deepcopy(run[0]) if len(run) == 1 else _merge_content_run(run, type_str)) + run.clear() + run_id = None + run_has_text = run_has_reasoning_text = False + 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 + if content.type != type_str: + flush() coalesced_contents.append(content) - if first_new_content: - coalesced_contents.append(first_new_content) + continue + if run: + if type_str == "text": + # A folded run keeps the first chunk's value for this key, so the + # run head alone decides whether the next chunk joins or splits. + if run[0].additional_properties.get(_MODEL_OUTPUT_KIND_KEY) != content.additional_properties.get( + _MODEL_OUTPUT_KIND_KEY + ): + flush() + else: + other_id: str | None = getattr(content, "id", None) + if (run_id and other_id and run_id != other_id) or ( + run_has_text + and getattr(content, "text", None) + and (run_has_reasoning_text != ("reasoning_text" in content.additional_properties)) + ): + flush() + run.append(content) + if type_str == "text_reasoning": + run_id = run_id or getattr(content, "id", None) + run_has_text = run_has_text or bool(getattr(content, "text", None)) + run_has_reasoning_text = run_has_reasoning_text or "reasoning_text" in content.additional_properties + flush() 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 9360e1a7609..51ed2647653 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -6,6 +6,7 @@ import json import warnings from collections.abc import AsyncIterable, Awaitable, Callable, Sequence +from copy import deepcopy from dataclasses import dataclass from datetime import datetime, timezone from functools import partial @@ -41,6 +42,7 @@ GROUP_TOKEN_COUNT_KEY, ) from agent_framework._types import ( + _MODEL_OUTPUT_KIND_KEY, _append_instructions, _get_data_bytes, _get_data_bytes_as_str, @@ -2494,6 +2496,131 @@ def test_coalesce_text_reasoning_with_different_ids(): assert contents[1].text == "Thinking B1 B2" +def _reference_coalesce(contents: list[Content], type_str: Literal["text", "text_reasoning"]) -> None: + """The pre-fix fold, kept as the equivalence oracle: deepcopy the run head, then += the rest.""" + if not contents: + return + coalesced: list[Content] = [] + acc: Content | None = None + for content in contents: + if content.type == type_str: + if acc is None: + acc = deepcopy(content) + elif type_str == "text" and acc.additional_properties.get( + _MODEL_OUTPUT_KIND_KEY + ) != content.additional_properties.get(_MODEL_OUTPUT_KIND_KEY): + coalesced.append(acc) + acc = deepcopy(content) + else: + try: + acc += content + except AdditionItemMismatch: + coalesced.append(acc) + acc = deepcopy(content) + else: + if acc: + coalesced.append(acc) + acc = None + coalesced.append(content) + if acc: + coalesced.append(acc) + contents.clear() + contents.extend(coalesced) + + +@pytest.mark.parametrize("type_str", ["text", "text_reasoning"]) +def test_coalesce_matches_repeated_add(type_str: Literal["text", "text_reasoning"]) -> None: + """The one-pass fold must reproduce the repeated-+= result on a mixed stream.""" + from agent_framework._types import _coalesce_text_content + + def make(i: int, **kwargs: Any) -> Content: + if type_str == "text": + return Content.from_text(f"chunk{i} ", **kwargs) + return Content.from_text_reasoning(text=f"chunk{i} ", **kwargs) + + streams = [ + # plain run + [make(i) for i in range(5)], + # run with props, annotations and raw representations of mixed shapes + [ + make(0, additional_properties={"a": 1}, raw_representation={"n": 0}), + make(1, additional_properties={"a": 2, "b": 1}, raw_representation=[{"n": 1}]), + make(2), + make(3, raw_representation={"n": 3}), + ], + # split by an unrelated content type, then resume + [make(0), Content.from_data(b"\x00", media_type="application/octet-stream"), make(1), make(2)], + # single chunk stays a plain copy + [make(0, raw_representation={"n": 0})], + ] + if type_str == "text": + marker = {_MODEL_OUTPUT_KIND_KEY: "refusal"} + streams.append([make(0), make(1, additional_properties=marker), make(2, additional_properties=marker), make(3)]) + else: + streams.append([make(0, id="rs_a"), make(1, id="rs_a"), make(2, id="rs_b"), make(3)]) + streams.append([ + make(0, additional_properties={"reasoning_text": True}), + make(1), + make(2, additional_properties={"reasoning_text": True}), + ]) + + for stream in streams: + expected = [deepcopy(c) for c in stream] + _reference_coalesce(expected, type_str) + actual = list(stream) + _coalesce_text_content(actual, type_str) + assert actual == expected, type_str + + +def test_coalesce_fold_does_not_readd_chunks() -> None: + """Aggregating n chunks must not invoke Content.__add__ per chunk (the old O(n^2) path).""" + from agent_framework._types import _coalesce_text_content + + calls = 0 + original_add = Content.__add__ + + def counting_add(self: Content, other: Content) -> Content: + nonlocal calls + calls += 1 + return original_add(self, other) + + contents = [Content.from_text(f"c{i} ", raw_representation={"i": i}) for i in range(2000)] + try: + Content.__add__ = counting_add # type: ignore[method-assign] + _coalesce_text_content(contents, "text") + finally: + Content.__add__ = original_add # type: ignore[method-assign] + + assert calls == 0 + assert len(contents) == 1 + assert contents[0].text == "".join(f"c{i} " for i in range(2000)) + # The run head's raw representation is dropped by the fold's deepcopy semantics; + # the remaining 1999 entries flatten in order. + assert contents[0].raw_representation == [{"i": i} for i in range(1, 2000)] + + +def test_coalesce_text_reasoning_one_pass_semantics() -> None: + """text_reasoning specifics: id from the first tagged chunk, protected_data from the last non-null.""" + from agent_framework._types import _coalesce_text_content + + contents = [ + Content.from_text_reasoning(id="rs_a", text="t1", protected_data="sig1"), + Content.from_text_reasoning(text="t2"), + Content.from_text_reasoning(text=None, protected_data="sig2"), + ] + _coalesce_text_content(contents, "text_reasoning") + assert len(contents) == 1 + assert contents[0].id == "rs_a" + assert contents[0].text == "t1t2" + assert contents[0].protected_data == "sig2" + + # a fully text-less run keeps text=None rather than "" + none_run = [Content.from_text_reasoning(text=None), Content.from_text_reasoning(text=None)] + _coalesce_text_content(none_run, "text_reasoning") + assert len(none_run) == 1 + assert none_run[0].text is None + + def test_agent_response_from_updates_preserves_refusal_marker() -> None: marker = {"model_output_kind": "refusal"} response = AgentResponse.from_updates([ From 376748d4e98330c01cd276b172451ceceef859d8 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Fri, 2 Oct 2026 12:52:30 +0800 Subject: [PATCH 2/3] Python: fix(core): match the repeated-+= id fallback and detach the run head in the one-pass fold Two exactness gaps against the old fold, both found by fuzzing old vs new: - id for a text_reasoning run fell back to None when no chunk carried a truthy id, but repeated `self.id or other.id` yields the last chunk's value there (so an all-falsy run ending in id="" kept ""). Fall back to run[-1].id. - the old path deep-copied the run head before folding, so the aggregate never aliased the head's nested additional_properties values or annotation dicts. The one-pass merge reused them. Deep-copy the head's contributions the same way. Also harden the equivalence tests: compare raw_representation explicitly (Content.__eq__ skips it, and the oracle's expected side now uses shallow copies so deepcopy does not strip raw up front), cover annotations on head/later chunks, and add a seeded fuzz test folding random streams through the live __add__ so future drift between the add-path and the one-pass fold fails CI. Cross-reference _merge_content_run from _add_text_content/_add_text_reasoning_content. --- .../packages/core/agent_framework/_types.py | 27 ++++-- python/packages/core/tests/core/test_types.py | 89 ++++++++++++++++++- 2 files changed, 108 insertions(+), 8 deletions(-) diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index c84cd1b6800..96a86ba280e 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -1693,7 +1693,11 @@ def __add__(self, other: Content) -> Content: raise ContentError(f"Addition not supported for content type: {self.type}") def _add_text_content(self, other: Content) -> Content: - """Add two TextContent instances.""" + """Add two TextContent instances. + + Keep the merge rules here in sync with ``_merge_content_run``, which + folds a run of chunks in one pass without calling ``__add__``. + """ return Content( "text", text=self.text + other.text, # type: ignore[attr-defined, operator] @@ -1703,7 +1707,12 @@ def _add_text_content(self, other: Content) -> Content: ) def _add_text_reasoning_content(self, other: Content) -> Content: - """Add two TextReasoningContent instances.""" + """Add two TextReasoningContent instances. + + Keep the mismatch checks and merge rules here in sync with + ``_coalesce_text_content`` / ``_merge_content_run``, which fold a run of + chunks in one pass without calling ``__add__``. + """ # Ensure we do not silently merge contents with conflicting ids if self.id and other.id and self.id != other.id: raise AdditionItemMismatch( @@ -2288,12 +2297,18 @@ def _merge_content_run(run: list[Content], type_str: Literal["text", "text_reaso representations flatten into one list (left untouched when a single chunk carries the only value). The run head's own raw representation is dropped, matching the old path whose initial ``deepcopy`` discards it - (``_SHALLOW_COPY_FIELDS``). + (``_SHALLOW_COPY_FIELDS``). The head's nested values are deep-copied for the + same reason, so mutating the source head after the fold cannot leak into + the aggregate. """ + head = run[0] merged_props: dict[str, Any] = {} - for content in reversed(run): + for content in reversed(run[1:]): merged_props.update(content.additional_properties) - annotation_seqs = [c.annotations for c in run if c.annotations is not None] + merged_props.update(deepcopy(head.additional_properties)) + annotation_seqs = [ + deepcopy(c.annotations) if c is head else c.annotations for c in run if c.annotations is not None + ] annotations = [a for seq in annotation_seqs for a in seq] if annotation_seqs else None non_null_raws = [c.raw_representation for c in run[1:] if c.raw_representation is not None] raw_representation: Any @@ -2315,7 +2330,7 @@ def _merge_content_run(run: list[Content], type_str: Literal["text", "text_reaso ) return Content( "text_reasoning", - id=next((c.id for c in run if c.id), None), + id=next((c.id for c in run if c.id), run[-1].id), text=None if all(c.text is None for c in run) else "".join(c.text or "" for c in run), protected_data=next((c.protected_data for c in reversed(run) if c.protected_data is not None), None), annotations=annotations, diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 51ed2647653..878b685f123 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -6,7 +6,7 @@ import json import warnings from collections.abc import AsyncIterable, Awaitable, Callable, Sequence -from copy import deepcopy +from copy import copy, deepcopy from dataclasses import dataclass from datetime import datetime, timezone from functools import partial @@ -2548,6 +2548,13 @@ def make(i: int, **kwargs: Any) -> Content: make(2), make(3, raw_representation={"n": 3}), ], + # annotations on the head, on a later chunk, and on both + [ + make(0, annotations=[{"type": "citation", "url": "https://a"}]), + make(1), + make(2, annotations=[{"type": "citation", "url": "https://b"}]), + ], + [make(0), make(1, annotations=[{"type": "citation", "url": "https://b"}]), make(2)], # split by an unrelated content type, then resume [make(0), Content.from_data(b"\x00", media_type="application/octet-stream"), make(1), make(2)], # single chunk stays a plain copy @@ -2565,11 +2572,15 @@ def make(i: int, **kwargs: Any) -> Content: ]) for stream in streams: - expected = [deepcopy(c) for c in stream] + # Shallow copies: deepcopy would strip raw_representation up front and make + # the explicit raw comparison below vacuous. + expected = [copy(c) for c in stream] _reference_coalesce(expected, type_str) actual = list(stream) _coalesce_text_content(actual, type_str) assert actual == expected, type_str + # Content.__eq__ excludes raw_representation, so compare it explicitly. + assert [c.raw_representation for c in actual] == [c.raw_representation for c in expected], type_str def test_coalesce_fold_does_not_readd_chunks() -> None: @@ -2621,6 +2632,80 @@ def test_coalesce_text_reasoning_one_pass_semantics() -> None: assert none_run[0].text is None +def test_coalesce_text_reasoning_empty_id_matches_repeated_add() -> None: + """An empty-string id survives the fold exactly like repeated += (``None or ""`` yields ``""``).""" + from agent_framework._types import _coalesce_text_content + + contents = [Content.from_text_reasoning(text="a"), Content.from_text_reasoning(id="", text="b")] + _coalesce_text_content(contents, "text_reasoning") + assert len(contents) == 1 + assert contents[0].id == "" + + +def test_coalesce_detaches_run_head_nested_values() -> None: + """Mutating the source head's nested values after the fold must not leak into the aggregate.""" + from agent_framework._types import _coalesce_text_content + + head = Content.from_text( + "h ", + additional_properties={"k": {"nested": 1}}, + annotations=[{"type": "citation", "url": "https://a"}], + ) + contents = [head, Content.from_text("t")] + _coalesce_text_content(contents, "text") + head.additional_properties["k"]["nested"] = 99 + head.annotations[0]["url"] = "https://mutated" + assert contents[0].additional_properties["k"]["nested"] == 1 + assert contents[0].annotations[0]["url"] == "https://a" + + +@pytest.mark.parametrize("type_str", ["text", "text_reasoning"]) +def test_coalesce_matches_repeated_add_fuzz(type_str: Literal["text", "text_reasoning"]) -> None: + """Seeded random streams: the one-pass fold must match repeated += exactly. + + The oracle folds through the live ``__add__``, so any future drift between + the add-path and the one-pass fold fails here. + """ + import random + + from agent_framework._types import _coalesce_text_content + + rng = random.Random(20261002) + + def rand_chunk() -> Content: + kwargs: dict[str, Any] = {} + if rng.random() < 0.5: + props: dict[str, Any] = {rng.choice(["a", "b"]): rng.randint(0, 3)} + if type_str == "text" and rng.random() < 0.2: + props[_MODEL_OUTPUT_KIND_KEY] = rng.choice(["refusal", "regular"]) + if type_str == "text_reasoning" and rng.random() < 0.3: + props["reasoning_text"] = True + kwargs["additional_properties"] = props + if rng.random() < 0.3: + annotation: Annotation = {"type": "citation", "url": f"https://x/{rng.randint(0, 9)}"} + kwargs["annotations"] = [annotation] + if rng.random() < 0.4: + kwargs["raw_representation"] = ( + {"n": rng.randint(0, 5)} if rng.random() < 0.5 else [{"n": rng.randint(0, 5)}] + ) + if type_str == "text": + return Content.from_text(rng.choice(["", "x", "y "]), **kwargs) + return Content.from_text_reasoning( + id=rng.choice([None, "", "rs_a", "rs_b"]), text=rng.choice([None, "", "t "]), **kwargs + ) + + for _ in range(300): + stream = [rand_chunk() for _ in range(rng.randint(0, 8))] + if stream and rng.random() < 0.3: + stream.insert(rng.randrange(len(stream)), Content.from_data(b"\x00", media_type="application/octet-stream")) + expected = [copy(c) for c in stream] + _reference_coalesce(expected, type_str) + actual = list(stream) + _coalesce_text_content(actual, type_str) + assert actual == expected + assert [c.raw_representation for c in actual] == [c.raw_representation for c in expected] + + def test_agent_response_from_updates_preserves_refusal_marker() -> None: marker = {"model_output_kind": "refusal"} response = AgentResponse.from_updates([ From 3dd8732caced5c3d2b03e5f887a12771f40e9544 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Fri, 2 Oct 2026 20:49:05 +0800 Subject: [PATCH 3/3] Python: test(core): narrow the optional props before subscripting in the detach test The four test type checkers all rejected subscripting Content's optional additional_properties/annotations directly. Bind and assert them first. Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> --- python/packages/core/tests/core/test_types.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 878b685f123..c67a89533ec 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -2653,10 +2653,16 @@ def test_coalesce_detaches_run_head_nested_values() -> None: ) contents = [head, Content.from_text("t")] _coalesce_text_content(contents, "text") - head.additional_properties["k"]["nested"] = 99 - head.annotations[0]["url"] = "https://mutated" - assert contents[0].additional_properties["k"]["nested"] == 1 - assert contents[0].annotations[0]["url"] == "https://a" + props = head.additional_properties + annotations = head.annotations + assert props is not None and annotations is not None + props["k"]["nested"] = 99 + annotations[0]["url"] = "https://mutated" + folded_props = contents[0].additional_properties + folded_annotations = contents[0].annotations + assert folded_props is not None and folded_annotations is not None + assert folded_props["k"]["nested"] == 1 + assert folded_annotations[0]["url"] == "https://a" @pytest.mark.parametrize("type_str", ["text", "text_reasoning"])