From b5956d7e60564f7a0d11cfb565a0f27c1164003e Mon Sep 17 00:00:00 2001 From: Hongda Zhu Date: Sun, 6 Sep 2026 20:13:07 -0400 Subject: [PATCH] fix: preserve partial and turn_complete when after_model_callback returns a replacement A callback that returns a replacement LlmResponse (as the documented after_model_callback contract suggests) rarely sets the streaming-control fields. LlmResponse.partial defaults to None, which the streaming protocol treats as a final response, so during SSE streaming every replaced delta became a final event: clients rendered N final responses and the Runner persisted every fragment to the session as a separate complete model event. Treat None as unset and inherit partial/turn_complete from the response being replaced at both after-model-callback call sites (live and SSE). Explicit values set by a callback are still respected. Fixes #7035 --- .../adk/flows/llm_flows/base_llm_flow.py | 32 ++++- .../test_base_llm_flow_partial_handling.py | 131 ++++++++++++++++++ 2 files changed, 161 insertions(+), 2 deletions(-) diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index 975f0e528ca..d5946178e16 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -268,6 +268,32 @@ async def _handle_before_model_callback( return None +def _inherit_unset_streaming_fields( + original: LlmResponse, replacement: LlmResponse +) -> LlmResponse: + """Carries streaming-control fields over to a callback-built replacement. + + A callback that returns a replacement ``LlmResponse`` (as the + ``after_model_callback`` contract suggests) rarely sets ``partial`` or + ``turn_complete``. ``None`` means "unset": inherit the value from the + response being replaced so a replacement cannot silently turn a streaming + delta into a final response. An explicit value set by the callback is + respected. + + Args: + original: The response produced by the model. + replacement: The response returned by an after-model callback. + + Returns: + The replacement response, with unset streaming-control fields inherited. + """ + if replacement.partial is None and original.partial is not None: + replacement.partial = original.partial + if replacement.turn_complete is None and original.turn_complete is not None: + replacement.turn_complete = original.turn_complete + return replacement + + async def _handle_after_model_callback( invocation_context: InvocationContext, llm_response: LlmResponse, @@ -1049,7 +1075,7 @@ async def _call_llm_with_tracing() -> AsyncGenerator[LlmResponse, None]: model_response_event, ) ): - event = altered + event = _inherit_unset_streaming_fields(event, altered) # only yield partial response in SSE streaming mode if ( run_config.streaming_mode == StreamingMode.SSE @@ -1091,7 +1117,9 @@ async def _call_llm_with_tracing() -> AsyncGenerator[LlmResponse, None]: model_response_event, ) ): - llm_response = altered + llm_response = _inherit_unset_streaming_fields( + llm_response, altered + ) yield llm_response diff --git a/tests/unittests/flows/llm_flows/test_base_llm_flow_partial_handling.py b/tests/unittests/flows/llm_flows/test_base_llm_flow_partial_handling.py index c688a610735..57e51b66bc3 100644 --- a/tests/unittests/flows/llm_flows/test_base_llm_flow_partial_handling.py +++ b/tests/unittests/flows/llm_flows/test_base_llm_flow_partial_handling.py @@ -12,8 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import AsyncGenerator + from google.adk.agents.llm_agent import Agent +from google.adk.agents.run_config import RunConfig +from google.adk.agents.run_config import StreamingMode +from google.adk.flows.llm_flows.base_llm_flow import _inherit_unset_streaming_fields from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow +from google.adk.models.base_llm import BaseLlm +from google.adk.models.llm_request import LlmRequest from google.adk.models.llm_response import LlmResponse from google.genai import types import pytest @@ -162,3 +169,127 @@ async def test_run_async_breaks_on_first_partial_response(): assert len(events) == 1 assert events[0].partial is True assert events[0].content.parts[0].text == 'Partial response' + + +class _StreamingMockModel(BaseLlm): + """Streams queued responses as a single model turn.""" + + model: str = 'streaming-mock' + stream_responses: list[LlmResponse] = [] + + @classmethod + def supported_models(cls) -> list[str]: + return ['streaming-mock'] + + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + for response in self.stream_responses: + yield response + + +def _streamed_turn() -> list[LlmResponse]: + """Three partial deltas followed by the aggregated final response.""" + deltas = ['Hello ', 'brave ', 'world.'] + responses = [ + LlmResponse( + content=types.Content( + role='model', parts=[types.Part.from_text(text=delta)] + ), + partial=True, + ) + for delta in deltas + ] + responses.append( + LlmResponse( + content=types.Content( + role='model', + parts=[types.Part.from_text(text=''.join(deltas))], + ) + ) + ) + return responses + + +def _rebuilding_callback(callback_context, llm_response: LlmResponse): + """Returns a fresh LlmResponse, per the after_model_callback contract.""" + if not (llm_response.content and llm_response.content.parts): + return None + text = llm_response.content.parts[0].text or '' + return LlmResponse( + content=types.Content( + role='model', + parts=[types.Part.from_text(text=text.replace('brave', 'kind'))], + ) + ) + + +@pytest.mark.asyncio +async def test_after_model_callback_replacement_preserves_partial_in_sse(): + """A rebuilt replacement inherits `partial` from the streamed response.""" + agent = Agent( + name='test_agent', + model=_StreamingMockModel(stream_responses=_streamed_turn()), + after_model_callback=_rebuilding_callback, + ) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, + user_content='test message', + run_config=RunConfig(streaming_mode=StreamingMode.SSE), + ) + + flow = BaseLlmFlowForTesting() + events = [] + async for event in flow.run_async(invocation_context): + events.append(event) + + assert [event.partial for event in events] == [True, True, True, None] + assert events[-1].content.parts[0].text == 'Hello kind world.' + + +@pytest.mark.asyncio +async def test_after_model_callback_explicit_partial_false_is_respected(): + """A replacement that explicitly finalizes a delta is not overridden.""" + + def finalize_callback(callback_context, llm_response: LlmResponse): + if not llm_response.partial: + return None + return LlmResponse(content=llm_response.content, partial=False) + + agent = Agent( + name='test_agent', + model=_StreamingMockModel(stream_responses=_streamed_turn()), + after_model_callback=finalize_callback, + ) + invocation_context = await testing_utils.create_invocation_context( + agent=agent, + user_content='test message', + run_config=RunConfig(streaming_mode=StreamingMode.SSE), + ) + + flow = BaseLlmFlowForTesting() + events = [] + async for event in flow.run_async(invocation_context): + events.append(event) + + assert events[0].partial is False + + +def test_inherit_unset_streaming_fields_inherits_when_unset(): + original = LlmResponse(partial=True, turn_complete=True) + replacement = LlmResponse() + + result = _inherit_unset_streaming_fields(original, replacement) + + assert result.partial is True + assert result.turn_complete is True + + +def test_inherit_unset_streaming_fields_respects_explicit_values(): + original = LlmResponse(partial=True, turn_complete=True) + replacement = LlmResponse(partial=False, turn_complete=False) + + result = _inherit_unset_streaming_fields(original, replacement) + + assert result.partial is False + assert result.turn_complete is False