From 33a1866603dcf9a1ea3841af516983426b3616a4 Mon Sep 17 00:00:00 2001
From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com>
Date: Mon, 3 Aug 2026 10:59:35 -0700
Subject: [PATCH] fix: StreamingResponseAggregator drops/duplicates function
 calls in streaming mode
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Fixes two issues in StreamingResponseAggregator that cause incorrect
agent behavior when using StreamingMode.SSE with multi-agent workflows
(transfer_to_agent + tools):

1. Progressive SSE path: close() could emit duplicate FunctionCall parts
   when streaming chunks overlap, causing the agent loop to re-execute
   the same tool call indefinitely. Fixed by deduplicating FC parts by
   their id in close().

2. Legacy SSE path: close() only emitted text parts — FunctionCall parts
   received during streaming were silently dropped. This caused the agent
   loop to exit prematurely without executing tools. Fixed by tracking
   FC parts in _function_call_parts and including them in close().

Fixes #6566
---
 src/google/adk/utils/streaming_utils.py       | 30 +++++-
 tests/unittests/utils/test_streaming_utils.py | 93 ++++++++++++++++++-
 2 files changed, 120 insertions(+), 3 deletions(-)

diff --git a/src/google/adk/utils/streaming_utils.py b/src/google/adk/utils/streaming_utils.py
index 1e8ac05a04e..571b3d05757 100644
--- a/src/google/adk/utils/streaming_utils.py
+++ b/src/google/adk/utils/streaming_utils.py
@@ -40,6 +40,9 @@ def __init__(self) -> None:
     self._citation_metadata: Optional[types.CitationMetadata] = None
     self._response = None
 
+    # For legacy streaming: track function call parts
+    self._function_call_parts: list[types.Part] = []
+
     # For progressive SSE streaming mode: accumulate parts in order
     self._parts_sequence: list[types.Part] = []
     self._current_text_buffer: list[str] = []
@@ -342,6 +345,16 @@ async def process_response(
       )
       self._thought_text = []
       self._text = []
+    
+    # Track function call parts for legacy aggregation
+    if (
+        llm_response.content
+        and llm_response.content.parts
+    ):
+      for part in llm_response.content.parts:
+        if part.function_call:
+          self._function_call_parts.append(part)
+
     yield llm_response
 
   def close(self) -> Optional[LlmResponse]:
@@ -377,8 +390,19 @@ def close(self) -> Optional[LlmResponse]:
       self._flush_text_buffer_to_sequence()
       self._flush_function_call_to_sequence()
 
-      final_parts = self._parts_sequence
-      content = types.ModelContent(parts=final_parts) if final_parts else None
+      # Deduplicate function call parts to prevent the agent loop from
+      # re-executing the same tool call. Streaming chunks can produce
+      # duplicate FunctionCall parts with the same id.
+      seen_fc_ids: set[str] = set()
+      deduped_parts: list[types.Part] = []
+      for part in self._parts_sequence:
+        if part.function_call and part.function_call.id:
+          if part.function_call.id in seen_fc_ids:
+            continue
+          seen_fc_ids.add(part.function_call.id)
+        deduped_parts.append(part)
+
+      content = types.ModelContent(parts=deduped_parts) if deduped_parts else None
 
       return LlmResponse(
           content=content,
@@ -398,6 +422,8 @@ def close(self) -> Optional[LlmResponse]:
       parts.append(types.Part(text=''.join(self._thought_text), thought=True))
     if self._text:
       parts.append(types.Part.from_text(text=''.join(self._text)))
+    # Include function call parts that were received during streaming
+    parts.extend(self._function_call_parts)
     content = types.ModelContent(parts=parts) if parts else None
 
     return LlmResponse(
diff --git a/tests/unittests/utils/test_streaming_utils.py b/tests/unittests/utils/test_streaming_utils.py
index ab03bc3c1d7..750b1164cea 100644
--- a/tests/unittests/utils/test_streaming_utils.py
+++ b/tests/unittests/utils/test_streaming_utils.py
@@ -288,7 +288,12 @@ async def test_pure_function_call_behavior_differs_by_mode(
         assert len(closed_response.content.parts) == 1
         assert closed_response.content.parts[0].function_call.name == "my_tool"
       else:
-        assert closed_response.content is None
+        # After the fix, legacy mode also preserves function call parts
+        assert closed_response.content is not None
+        assert any(
+            p.function_call and p.function_call.name == "my_tool"
+            for p in closed_response.content.parts
+        )
 
   @pytest.mark.asyncio
   @pytest.mark.parametrize(
@@ -489,6 +494,92 @@ async def test_non_progressive_merged_yield_propagates_model_version(self):
       assert merged_events, "expected a merged non-partial text event"
       assert merged_events[0].model_version == "gemini-test-2.0"
 
+  @pytest.mark.asyncio
+  async def test_progressive_close_deduplicates_function_calls(self):
+    with temporary_feature_override(FeatureName.PROGRESSIVE_SSE_STREAMING, True):
+      aggregator = streaming_utils.StreamingResponseAggregator()
+
+      part1 = types.Part(function_call=types.FunctionCall(name="test_func", args={"a": 1}, id="fc_123"))
+      part2 = types.Part(function_call=types.FunctionCall(name="test_func", args={"a": 1}, id="fc_123"))
+      part3 = types.Part(function_call=types.FunctionCall(name="test_func2", args={"b": 2}, id="fc_456"))
+
+      resp1 = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(parts=[part1])
+              )
+          ]
+      )
+      resp2 = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(parts=[part2])
+              )
+          ]
+      )
+      resp3 = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(parts=[part3])
+              )
+          ]
+      )
+
+      async for _ in aggregator.process_response(resp1):
+        pass
+      async for _ in aggregator.process_response(resp2):
+        pass
+      async for _ in aggregator.process_response(resp3):
+        pass
+
+      final_response = aggregator.close()
+
+      assert final_response is not None
+      assert final_response.content is not None
+      assert len(final_response.content.parts) == 2
+      assert final_response.content.parts[0].function_call.id == "fc_123"
+      assert final_response.content.parts[1].function_call.id == "fc_456"
+
+  @pytest.mark.asyncio
+  async def test_legacy_close_preserves_function_calls(self):
+    with temporary_feature_override(FeatureName.PROGRESSIVE_SSE_STREAMING, False):
+      aggregator = streaming_utils.StreamingResponseAggregator()
+
+      part1 = types.Part(text="Hello")
+      part2 = types.Part(function_call=types.FunctionCall(name="test_func", args={"a": 1}))
+
+      resp1 = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(parts=[part1])
+              )
+          ]
+      )
+      resp2 = types.GenerateContentResponse(
+          candidates=[
+              types.Candidate(
+                  content=types.Content(parts=[part2])
+              )
+          ]
+      )
+
+      async for _ in aggregator.process_response(resp1):
+        pass
+      async for _ in aggregator.process_response(resp2):
+        pass
+
+      final_response = aggregator.close()
+
+      # In legacy mode, text is flushed mid-stream (via intermediate yield)
+      # when a non-text chunk arrives. So close() only contains the FC part
+      # that was tracked separately via _function_call_parts.
+      assert final_response is not None
+      assert final_response.content is not None
+      assert any(
+          p.function_call and p.function_call.name == "test_func"
+          for p in final_response.content.parts
+      )
+
 
 class TestFunctionCallIdGeneration:
   """Tests for function call ID generation in streaming mode."""