diff --git a/src/google/adk/flows/llm_flows/_model_response_finalizer.py b/src/google/adk/flows/llm_flows/_model_response_finalizer.py index 595407d090..943978f0ad 100644 --- a/src/google/adk/flows/llm_flows/_model_response_finalizer.py +++ b/src/google/adk/flows/llm_flows/_model_response_finalizer.py @@ -123,6 +123,37 @@ async def handle_before_model_callback( return None +def _inherit_unset_streaming_fields( + original: LlmResponse, replacement: LlmResponse +) -> LlmResponse: + """Carries streaming-control fields from a replaced response. + + A callback replacement that leaves ``partial``/``turn_complete`` unset must + not change the streaming semantics of the response it replaces. Otherwise + every streamed delta looks final downstream: SSE clients render N final + responses, ``Runner`` persists each delta as a separate session event, and + the live path can close its request queue early. An explicitly set value on + the replacement is always respected. + + Args: + original: The response being replaced. + replacement: The callback-provided response. + + Returns: + The replacement, with unset streaming-control fields filled in. A copy is + returned only when a field actually needs filling, so explicitly complete + replacements keep their identity. + """ + updates = {} + if replacement.partial is None and original.partial is not None: + updates['partial'] = original.partial + if replacement.turn_complete is None and original.turn_complete is not None: + updates['turn_complete'] = original.turn_complete + if updates: + return replacement.model_copy(update=updates) + return replacement + + async def handle_after_model_callback( invocation_context: InvocationContext, llm_response: LlmResponse, @@ -178,7 +209,10 @@ async def _maybe_add_grounding_metadata( ) ) if callback_response: - return await _maybe_add_grounding_metadata(callback_response) + return _inherit_unset_streaming_fields( + llm_response, + await _maybe_add_grounding_metadata(callback_response), + ) # If no overrides are provided from the plugins, further run the canonical # callbacks. @@ -189,7 +223,10 @@ async def _maybe_add_grounding_metadata( llm_response=llm_response, ) if callback_response: - return await _maybe_add_grounding_metadata(callback_response) + return _inherit_unset_streaming_fields( + llm_response, + await _maybe_add_grounding_metadata(callback_response), + ) return await _maybe_add_grounding_metadata() 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 c688a61073..fb1389adc2 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,16 @@ # 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.events.event import Event +from google.adk.flows.llm_flows.base_llm_flow import _handle_after_model_callback 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 +170,157 @@ 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' + + +def _text_response(text, partial=None, turn_complete=None): + return LlmResponse( + content=types.Content( + role='model', parts=[types.Part.from_text(text=text)] + ), + partial=partial, + turn_complete=turn_complete, + ) + + +def _rebuilding_callback(replacement): + """Returns a callback that rebuilds the response, dropping control fields.""" + + async def _callback(callback_context, llm_response): + del callback_context, llm_response + return replacement + + return _callback + + +async def _run_handle_after_model_callback(agent, llm_response): + invocation_context = await testing_utils.create_invocation_context( + agent=agent + ) + event = Event( + id=Event.new_id(), + invocation_id=invocation_context.invocation_id, + author=agent.name, + ) + return await _handle_after_model_callback( + invocation_context, llm_response, event + ) + + +@pytest.mark.asyncio +async def test_after_model_callback_replacement_inherits_partial(): + """A rebuilt replacement keeps the delta's partial flag (issue #7035).""" + replacement = _text_response('Hello ') + agent = Agent( + name='test_agent', + after_model_callback=[_rebuilding_callback(replacement)], + ) + + result = await _run_handle_after_model_callback( + agent, _text_response('Hello ', partial=True) + ) + + assert result.partial is True + assert result.content.parts[0].text == 'Hello ' + # The callback's object is not mutated; inheritance returns a copy. + assert replacement.partial is None + + +@pytest.mark.asyncio +async def test_after_model_callback_replacement_inherits_turn_complete(): + """A rebuilt replacement keeps turn_complete from the replaced response.""" + replacement = _text_response('done') + agent = Agent( + name='test_agent', + after_model_callback=[_rebuilding_callback(replacement)], + ) + + result = await _run_handle_after_model_callback( + agent, _text_response('done', turn_complete=True) + ) + + assert result.turn_complete is True + + +@pytest.mark.asyncio +async def test_after_model_callback_replacement_explicit_partial_respected(): + """An explicitly set partial=False on the replacement is not overridden.""" + agent = Agent( + name='test_agent', + after_model_callback=[ + _rebuilding_callback(_text_response('final', partial=False)) + ], + ) + + result = await _run_handle_after_model_callback( + agent, _text_response('final', partial=True) + ) + + assert result.partial is False + + +@pytest.mark.asyncio +async def test_after_model_callback_replacement_without_streaming_keeps_identity(): + """Non-streaming replacements pass through untouched (no copy).""" + replacement = _text_response('final') + agent = Agent( + name='test_agent', + after_model_callback=[_rebuilding_callback(replacement)], + ) + + result = await _run_handle_after_model_callback( + agent, _text_response('final') + ) + + assert result is replacement + + +class _StreamingFakeModel(BaseLlm): + """Yields two partial deltas then the aggregated final response.""" + + model: str = 'fake-streaming' + + @classmethod + def create(cls): + return cls(model='fake-streaming') + + @classmethod + def supported_models(cls) -> list[str]: + return ['.*'] + + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + deltas = ['Hello ', 'world.'] + if stream: + for delta in deltas: + yield _text_response(delta, partial=True) + yield _text_response(''.join(deltas)) + + +@pytest.mark.asyncio +async def test_run_async_sse_rebuilt_responses_stay_partial(): + """End-to-end: rebuilding callbacks must not flip SSE deltas to final.""" + agent = Agent( + name='test_agent', + model=_StreamingFakeModel.create(), + after_model_callback=[_rebuilding_callback(_text_response('scrubbed'))], + ) + 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, None] + # The callback ran on every response (content replaced), but the streaming + # semantics of the originals survived. + assert [event.content.parts[0].text for event in events] == [ + 'scrubbed', + 'scrubbed', + 'scrubbed', + ]