after_model_callback: a replacement LlmResponse silently loses `partial` during SSE streaming - every delta is persisted to the session as a final event

[zhuhongd]

Describe the Bug:

The documented contract for after_model_callback says a callback may return a replacement response ("When present, the actual model response will be ignored and the provided content will be returned to user" - agents/llm_agent.py). During SSE streaming, however, the callback is invoked on every streamedLlmResponse, including partial=True deltas (flows/llm_flows/base_llm_flow.py:1833-1840 on main), and the returned replacement is used wholesale:

if altered := (await self._handle_after_model_callback(...)):
    llm_response = altered

LlmResponse.partial defaults to None (models/llm_response.py), and per the streaming protocol described in LlmResponse's own docstring, a response with partial "false or unset" is a final response. So any callback that rebuilds an LlmResponse, exactly what the contract suggests, silently converts every streaming delta into a final response. Neither _handle_after_model_callback (module level, base_llm_flow.py:295) nor the call sites carry partial (or turn_complete) over from the response being replaced.

Downstream, Runner persists every non-partial event (runners.py: if not event.partial: await self.session_service.append_event(...)), so each delta fragment is appended to the session as a separate complete model event - the conversation history is corrupted for all subsequent turns.

To Reproduce:

Self-contained repro, no API key required (fake model follows the partial-deltas-then-aggregated-final protocol from LlmResponse's docstring):

# repro_partial_loss.py - run: pip install google-adk==2.8.0 && python repro_partial_loss.py
"""Repro: after_model_callback silently drops `partial` on streamed responses.

A fake streaming model yields 3 partial deltas + 1 final aggregated response (the protocol described in LlmResponse's own docstring). The
after_model_callback follows the documented contract ("When present, the actual model response will be ignored and the provided content will be returned to user") and returns a rebuilt LlmResponse.

Result: every streamed chunk loses partial=True, so (a) clients see N "final" responses instead of deltas, and (b) Runner persists every delta
into the session as a separate complete event (runners.py: `if not event.partial: append_event(...)`).
"""
import asyncio
from typing import AsyncGenerator

from google.adk.agents import LlmAgent
from google.adk.agents.run_config import RunConfig, StreamingMode
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.adk.runners import InMemoryRunner
from google.genai import types


def _text_response(text: str, partial: bool | None = None) -> LlmResponse:
  return LlmResponse(
      content=types.Content(role="model", parts=[types.Part(text=text)]),
      partial=partial,
  )


class FakeStreamingLlm(BaseLlm):
  """Emits: 'Hello ' / 'brave ' / 'world.' as partials, then the full text."""

  @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 ", "brave ", "world."]
    if stream:
      for d in deltas:
        yield _text_response(d, partial=True)
    yield _text_response("".join(deltas))  # aggregated final (partial unset)


def scrub_callback(callback_context, llm_response: LlmResponse):
  """A guardrail-style callback written exactly per the documented contract:
  build and return a replacement LlmResponse."""
  if not (llm_response.content and llm_response.content.parts):
    return None
  text = llm_response.content.parts[0].text or ""
  return _text_response(text.replace("brave", "[REDACTED]"))


async def run(label: str, callback):
  agent = LlmAgent(
      name="repro_agent",
      model=FakeStreamingLlm(model="fake-streaming"),
      after_model_callback=callback,
  )
  runner = InMemoryRunner(agent=agent, app_name="repro")
  session = await runner.session_service.create_session(
      app_name="repro", user_id="u"
  )
  print(f"\n=== {label} ===")
  n = 0
  async for event in runner.run_async(
      user_id="u",
      session_id=session.id,
      new_message=types.Content(role="user", parts=[types.Part(text="hi")]),
      run_config=RunConfig(streaming_mode=StreamingMode.SSE),
  ):
    n += 1
    text = ""
    if event.content and event.content.parts and event.content.parts[0].text:
      text = event.content.parts[0].text
    print(f"  event {n}: partial={event.partial!r:6} "
          f"final={event.is_final_response()!s:5} text={text!r}")
  stored = await runner.session_service.get_session(
      app_name="repro", user_id="u", session_id=session.id
  )
  model_events = [e for e in stored.events if e.author == "repro_agent"]
  print(f"  session now holds {len(model_events)} persisted model event(s):")
  for e in model_events:
    t = e.content.parts[0].text if (e.content and e.content.parts) else ""
    print(f"    - {t!r}")


async def main():
  await run("CONTROL: no after_model_callback", None)
  await run("BUG: doc-contract callback (rebuilds LlmResponse)", scrub_callback)

asyncio.run(main())

def scrub_callback_copy(callback_context, llm_response: LlmResponse):
  """Workaround: model_copy(update=...) preserves partial/turn_complete."""
  if not (llm_response.content and llm_response.content.parts):
    return None
  text = llm_response.content.parts[0].text or ""
  new_content = types.Content(
      role="model", parts=[types.Part(text=text.replace("brave", "[REDACTED]"))])
  return llm_response.model_copy(update={"content": new_content})

asyncio.run(run("WORKAROUND: model_copy(update=...)", scrub_callback_copy))

Actual output on google-adk 2.8.0:

=== CONTROL: no after_model_callback ===
  event 1: partial=True   final=False text='Hello '
  event 2: partial=True   final=False text='brave '
  event 3: partial=True   final=False text='world.'
  event 4: partial=None   final=True  text='Hello brave world.'
  session now holds 1 persisted model event(s):
    - 'Hello brave world.'

=== BUG: doc-contract callback (rebuilds LlmResponse) ===
  event 1: partial=None   final=True  text='Hello '
  event 2: partial=None   final=True  text='[REDACTED] '
  event 3: partial=None   final=True  text='world.'
  event 4: partial=None   final=True  text='Hello [REDACTED] world.'
  session now holds 4 persisted model event(s):
    - 'Hello '
    - '[REDACTED] '
    - 'world.'
    - 'Hello [REDACTED] world.'

=== WORKAROUND: model_copy(update=...) ===
  event 1: partial=True   final=False text='Hello '
  event 2: partial=True   final=False text='[REDACTED] '
  event 3: partial=True   final=False text='world.'
  event 4: partial=None   final=True  text='Hello [REDACTED] world.'
  session now holds 1 persisted model event(s):
    - 'Hello [REDACTED] world.'

Expected behavior:

A callback replacement that does not explicitly set streaming-control fields should not change the streaming semantics of the response it replaces: deltas stay partial=True, exactly one final event is persisted to the session (the WORKAROUND run in the output above shows the expected shape).

Impact:

1. Session corruption: every delta is appended to the session as a
   complete model event; later turns see N fragment messages as context.
2. Clients in SSE mode receive N "final" responses (is_final_response() == True for every delta) - duplicated/mangled rendering.
3. Same field-loss class affects turn_complete (live path checks
   event.turn_complete to close the live request queue - a rebuilt response
   loses it) and usage_metadata (token accounting silently dropped).

Workaround:

Rebuild via model_copy instead of the constructor:

return llm_response.model_copy(update={"content": new_content})

Verified (third run in the repro/output above): deltas stay partial, session gets exactly 1 event. But nothing in the docs points users to this, and the natural reading of the contract produces the broken version.

Suggested fix (happy to send a PR):

Treat None as "unset" and inherit streaming-control fields from the response being replaced, at both call sites (live: base_llm_flow.py:1791-1798, SSE: 1833-1840):

def _inherit_unset_streaming_fields(original, replacement):
    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

An explicit partial=False set by a callback is still respected.
Alternatively (or additionally): document the model_copy requirement in the after_model_callback contract.

Desktop:

• OS: macOS 15 (Darwin 25.6)
• Python: 3.13.12
• google-adk: 2.8.0 (also present in 1.18.0 — call sites at base_llm_flow.py:757-760 and 796-799 there; the pattern is unchanged across the 1.x→2.x line)

Additional context:

Found running a production multi-agent HR assistant where a guardrail-style after_model_callback (PII scrubbing) rebuilt responses per the documented contract; symptom was mid-stream output mangling plus duplicated assistant messages accumulating in session history.

[Bruce-Yii]

I reproduced this on current main (b018062) with the exact script from the report: control run persists 1 model event, the doc-contract callback run persists 4 (every delta stored as a final event). Root cause confirmed at both after_model_callback call sites in base_llm_flow.py (run_live SSE branch and the non-support_cfc branch): the replacement is used wholesale, so an unset partial/turn_complete silently flips every streamed delta to final, and Runner then persists each one.

Planned fix: in the module-level _handle_after_model_callback (the single choke point for plugin + agent callbacks), inherit partial/turn_complete from the replaced response when the replacement leaves them unset; an explicitly set value is still respected. Regression tests will cover inherit-partial, inherit-turn_complete, explicit-False-respected, and an SSE flow-level run asserting deltas stay partial with exactly one persisted final event.

If that direction looks right, I'll open the focused PR shortly.

[surajksharma07]

@zhuhongd The wholesale replacement at both after_model_callback call sites does flip every delta to final, confirmed with your script (4 persisted events instead of 1). The inherit-unset-fields approach checks out on my end too, ran it against both #7038 and #7036 and deltas stay partial with only one event landing in the session either way. Since you and Bruce-Yii arrived at basically the same fix independently might be worth syncing on which PR to carry forward so we're not reviewing duplicate work.

[Bruce-Yii]

#7036 and #7038 independently converged on the same inherit-unset-streaming-fields fix, and @surajksharma07 has confirmed both resolve the repro.

For coordination: #7036 was opened first and applies the inheritance at "_handle_after_model_callback", the shared choke point for plugin and agent callbacks. It also uses "model_copy()" only when inheritance is needed, so the callback-owned replacement object is not mutated.

Happy for maintainers to choose whichever implementation they prefer; I’ll avoid duplicating further work until that decision is made.