diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 14c3050952..6d6bc58068 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -1380,6 +1380,13 @@ def _add_text_content(self, other: Content) -> Content: def _add_text_reasoning_content(self, other: Content) -> Content: """Add two TextReasoningContent instances.""" + # Ensure we do not silently merge contents with conflicting ids + if self.id and other.id and self.id != other.id: + raise AdditionItemMismatch( + f"Cannot add text_reasoning content with different ids: {self.id!r} != {other.id!r}" + ) + combined_id = self.id or other.id + # Concatenate text, handling None values self_text = self.text or "" # type: ignore[attr-defined] other_text = other.text or "" # type: ignore[attr-defined] @@ -1390,6 +1397,7 @@ def _add_text_reasoning_content(self, other: Content) -> Content: return Content( "text_reasoning", + id=combined_id, text=combined_text, protected_data=protected_data, annotations=_combine_annotations(self.annotations, other.annotations), @@ -1880,7 +1888,12 @@ def _coalesce_text_content(contents: list[Content], type_str: Literal["text", "t if first_new_content is None: first_new_content = deepcopy(content) else: - first_new_content += content + 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, diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 8e5a0dc7d8..9d1b39f29e 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -43,7 +43,7 @@ add_usage_details, validate_tool_mode, ) -from agent_framework.exceptions import ContentError +from agent_framework.exceptions import AdditionItemMismatch, ContentError @fixture @@ -1526,6 +1526,88 @@ def test_text_reasoning_content_iadd_coverage(): assert t1.text == "Thinking 1 Thinking 2" +def test_text_reasoning_content_add_preserves_id(): + """Test that coalescing text_reasoning Content preserves the id field.""" + + t1 = Content.from_text_reasoning(id="rs_abc123", text="Thinking part 1") + t2 = Content.from_text_reasoning(id="rs_abc123", text=" part 2") + + result = t1 + t2 + assert result.text == "Thinking part 1 part 2" + assert result.id == "rs_abc123" + + +def test_text_reasoning_content_add_id_fallback_to_other(): + """Test that coalescing falls back to other's id when self has no id.""" + + t1 = Content.from_text_reasoning(text="Thinking part 1") + t2 = Content.from_text_reasoning(id="rs_abc123", text=" part 2") + + result = t1 + t2 + assert result.id == "rs_abc123" + + +def test_text_reasoning_content_add_preserves_id_with_encrypted_content(): + """Test that id and encrypted_content both survive coalescing for round-trip.""" + + t1 = Content.from_text_reasoning( + id="rs_abc123", + text="Thinking", + additional_properties={"encrypted_content": "enc_blob_data"}, + ) + t2 = Content.from_text_reasoning(id="rs_abc123", text=" more") + + result = t1 + t2 + assert result.text == "Thinking more" + assert result.id == "rs_abc123" + assert result.additional_properties.get("encrypted_content") == "enc_blob_data" + + +def test_text_reasoning_content_add_conflicting_ids_raises(): + """Test that coalescing text_reasoning Content with different ids raises AdditionItemMismatch.""" + + t1 = Content.from_text_reasoning(id="rs_abc123", text="Thinking part 1") + t2 = Content.from_text_reasoning(id="rs_xyz789", text=" part 2") + + with pytest.raises(AdditionItemMismatch, match="different ids"): + t1 + t2 + + +def test_text_reasoning_content_add_neither_has_id(): + """Test that coalescing text_reasoning Content when neither has an id results in None id.""" + + t1 = Content.from_text_reasoning(text="Thinking part 1") + t2 = Content.from_text_reasoning(text=" part 2") + + result = t1 + t2 + assert result.text == "Thinking part 1 part 2" + assert result.id is None + + +def test_coalesce_text_reasoning_with_different_ids(): + """Test that _coalesce_text_content keeps separate text_reasoning items when IDs differ. + + Regression test: streaming responses can produce multiple text_reasoning + segments with distinct IDs. These must not be merged into one. + """ + from agent_framework._types import _coalesce_text_content + + contents = [ + Content.from_text_reasoning(id="rs_aaa", text="Thinking A1"), + Content.from_text_reasoning(id="rs_aaa", text=" A2"), + Content.from_text_reasoning(id="rs_bbb", text="Thinking B1"), + Content.from_text_reasoning(id="rs_bbb", text=" B2"), + ] + + _coalesce_text_content(contents, "text_reasoning") + + assert len(contents) == 2 + assert contents[0].id == "rs_aaa" + assert contents[0].text == "Thinking A1 A2" + assert contents[1].id == "rs_bbb" + assert contents[1].text == "Thinking B1 B2" + + def test_comprehensive_to_dict_exclude_options(): """Test to_dict methods with various exclude options for better coverage."""