From 2220d885cda875144b52338b5becf6e5546f3f51 Mon Sep 17 00:00:00 2001 From: Xuan Yang Date: Tue, 3 Feb 2026 15:33:21 -0800 Subject: [PATCH] fix: Check `will_continue` for streaming function calls Related: https://github.com/google/adk-python/issues/4311 Co-authored-by: Xuan Yang PiperOrigin-RevId: 865074872 --- .../hello_world_stream_fc_args/agent.py | 10 ++ src/google/adk/utils/streaming_utils.py | 13 +- .../test_progressive_sse_streaming.py | 122 ++++++++++++++++++ 3 files changed, 140 insertions(+), 5 deletions(-) diff --git a/contributing/samples/hello_world_stream_fc_args/agent.py b/contributing/samples/hello_world_stream_fc_args/agent.py index 8b0593a0..0c34e199 100755 --- a/contributing/samples/hello_world_stream_fc_args/agent.py +++ b/contributing/samples/hello_world_stream_fc_args/agent.py @@ -29,6 +29,11 @@ def concat_number_and_string(num: int, s: str) -> str: return str(num) + ': ' + s +def write_document(document: str) -> dict[str, str]: + """Write a document.""" + return {'status': 'ok'} + + root_agent = Agent( model='gemini-3-pro-preview', name='hello_world_stream_fc_args', @@ -38,9 +43,14 @@ root_agent = Agent( You can use the `concat_number_and_string` tool to concatenate a number and a string. You should always call the concat_number_and_string tool to concatenate a number and a string. You should never concatenate on your own. + + You can use the `write_document` tool to write a document. + You should always call the write_document tool to write a document. + You should never write a document on your own. """, tools=[ concat_number_and_string, + write_document, ], generate_content_config=types.GenerateContentConfig( automatic_function_calling=types.AutomaticFunctionCallingConfig( diff --git a/src/google/adk/utils/streaming_utils.py b/src/google/adk/utils/streaming_utils.py index 3d3acfdb..d3751df9 100644 --- a/src/google/adk/utils/streaming_utils.py +++ b/src/google/adk/utils/streaming_utils.py @@ -188,7 +188,7 @@ class StreamingResponseAggregator: self._current_fc_id = fc.id # Process each partial argument - for partial_arg in getattr(fc, 'partial_args', []): + for partial_arg in fc.partial_args or []: json_path = partial_arg.json_path if not json_path: continue @@ -203,8 +203,7 @@ class StreamingResponseAggregator: self._set_value_by_json_path(json_path, value) # Check if function call is complete - fc_will_continue = getattr(fc, 'will_continue', False) - if not fc_will_continue: + if not fc.will_continue: # Function call complete, flush it self._flush_text_buffer_to_sequence() self._flush_function_call_to_sequence() @@ -216,9 +215,13 @@ class StreamingResponseAggregator: part: The part containing a function call """ fc = part.function_call + if not fc: + return - # Check if this is a streaming FC (has partialArgs) - if hasattr(fc, 'partial_args') and fc.partial_args: + # Check if this is a streaming FC (has partialArgs or will_continue=True) + # The first chunk of a streaming function call may have will_continue=True + # but no partial_args yet, so we need to check both conditions. + if fc.partial_args or fc.will_continue: # Streaming function call arguments # Save thought_signature from the part (first chunk should have it) diff --git a/tests/unittests/flows/llm_flows/test_progressive_sse_streaming.py b/tests/unittests/flows/llm_flows/test_progressive_sse_streaming.py index 37dfbb95..a5b984ed 100644 --- a/tests/unittests/flows/llm_flows/test_progressive_sse_streaming.py +++ b/tests/unittests/flows/llm_flows/test_progressive_sse_streaming.py @@ -14,6 +14,7 @@ """Tests for Progressive SSE Streaming Stage 1 implementation.""" +import asyncio from typing import Any from typing import AsyncGenerator @@ -26,6 +27,7 @@ from google.adk.models.llm_response import LlmResponse from google.adk.runners import InMemoryRunner from google.adk.utils.streaming_utils import StreamingResponseAggregator from google.genai import types +import pytest def get_weather(location: str) -> dict[str, Any]: @@ -633,6 +635,126 @@ def test_progressive_sse_handles_empty_function_call(): assert args["s"] == "ADK" +@pytest.mark.parametrize( + "first_chunk_partial_args", + [ + pytest.param(None, id="partial_args_none"), + pytest.param([], id="partial_args_empty_list"), + ], +) +def test_streaming_fc_chunk_with_will_continue_but_no_partial_args( + first_chunk_partial_args, +): + """Test streaming function call with will_continue=True but no partial_args.""" + + aggregator = StreamingResponseAggregator() + + # Chunk 1: FC name + will_continue=True, but NO partial_args (or empty list) + # This is the first chunk that Gemini 3 sends for streaming FC + chunk1_fc = types.FunctionCall( + name="my_tool", + id="fc_gemini3", + will_continue=True, + partial_args=first_chunk_partial_args, + ) + chunk1_part = types.Part( + function_call=chunk1_fc, + thought_signature=b"test_sig_123", + ) + chunk1 = types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content(role="model", parts=[chunk1_part]) + ) + ] + ) + + # Chunk 2: Middle chunk with partial_args, name is None + chunk2_fc = types.FunctionCall( + partial_args=[ + types.PartialArg(json_path="$.document", string_value="Once upon ") + ], + will_continue=True, + ) + chunk2 = types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content( + role="model", parts=[types.Part(function_call=chunk2_fc)] + ) + ) + ] + ) + + # Chunk 3: Another middle chunk continuing the string argument + chunk3_fc = types.FunctionCall( + partial_args=[ + types.PartialArg(json_path="$.document", string_value="a time...") + ], + will_continue=True, + ) + chunk3 = types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content( + role="model", parts=[types.Part(function_call=chunk3_fc)] + ) + ) + ] + ) + + # Chunk 4: Final chunk - no name, no partial_args, will_continue=False + # This signals the end of the streaming function call + chunk4_fc = types.FunctionCall( + will_continue=False, + ) + chunk4 = types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content( + role="model", parts=[types.Part(function_call=chunk4_fc)] + ), + finish_reason=types.FinishReason.STOP, + ) + ] + ) + + # Process all chunks through aggregator + async def process(): + results = [] + for chunk in [chunk1, chunk2, chunk3, chunk4]: + async for response in aggregator.process_response(chunk): + results.append(response) + return results + + processed_chunks = asyncio.run(process()) + + # All intermediate chunks should be marked as partial + assert all(chunk.partial for chunk in processed_chunks) + + # Get final aggregated response + final_response = aggregator.close() + + # Verify final aggregated response has the complete FC with accumulated args + assert final_response is not None + assert len(final_response.content.parts) == 1 + + fc_part = final_response.content.parts[0] + assert fc_part.function_call is not None + assert fc_part.function_call.name == "my_tool" + assert fc_part.function_call.id == "fc_gemini3" + + # Verify the document argument was correctly accumulated + args = fc_part.function_call.args + assert "document" in args + assert ( + args["document"] == "Once upon a time..." + ) # Concatenated from chunks 2 + 3 + + # Verify thought_signature was preserved from the first chunk + assert fc_part.thought_signature == b"test_sig_123" + + class PartialFunctionCallMockModel(BaseLlm): """A mock model that yields partial function call events followed by final."""