fix: preserve partial and turn_complete when after_model_callback returns a replacement

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
 

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