fix(flows): inherit unset streaming fields on after_model_callback replacement

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -268,6 +268,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,
@@ -323,7 +354,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.
@@ -334,7 +368,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()
 
 

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',
+  ]