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

[Bruce-Yii]

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

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

Problem:

after_model_callback may return a replacement LlmResponse ("the actual model response will be ignored and the provided content will be returned to user"). During streaming, the callback runs on every response including partial=True deltas, and the replacement was used wholesale. Since LlmResponse.partial defaults to None — and unset means final per the streaming protocol — any callback that rebuilds the response silently converted every delta into a final response. Runner persists every non-partial event, so each delta fragment was appended to the session as a separate complete model event (4 persisted events instead of 1 in the issue's repro), and SSE clients saw N "final" responses. The same field-loss class affects turn_complete (live path closes its request queue on it).

Solution:

In the module-level _handle_after_model_callback (the single choke point for plugin + agent callbacks, covering both the live and SSE call sites), inherit partial/turn_complete from the replaced response when the replacement leaves them unset. An explicitly set value is always respected. Inheritance returns a model_copy only when a field actually needs filling, so explicitly complete replacements keep their identity and callback-owned objects are never mutated. usage_metadata is deliberately left out of this change to keep the fix minimal (noted as residual risk below).

Testing Plan

Unit Tests:

• I have added or updated unit tests for my change.
• All unit tests pass locally.

New tests in tests/unittests/flows/llm_flows/test_base_llm_flow_partial_handling.py (5 added, file now 9 tests):

• test_after_model_callback_replacement_inherits_partial — rebuilt replacement keeps the delta's partial=True; the callback's object is not mutated.
• test_after_model_callback_replacement_inherits_turn_complete.
• test_after_model_callback_replacement_explicit_partial_respected — explicit partial=False wins (guards against over-correction).
• test_after_model_callback_replacement_without_streaming_keeps_identity — non-streaming pass-through untouched.
• test_run_async_sse_rebuilt_responses_stay_partial — flow-level SSE run: deltas stay partial, final stays final.

Fail-before verified by stashing the src/ fix: the 3 behavior tests fail on pristine main (b018062), the 2 guard tests pass both before and after.

pytest results (Python 3.12, Windows):

• tests/unittests/flows/llm_flows/test_base_llm_flow_partial_handling.py + test_base_llm_flow.py: 109 passed
• tests/unittests/flows/llm_flows/ (minus 2 files needing the uninstalled optional google.cloud speech dep — pre-existing collection errors, also present without this change): 676 passed, 1 skipped
• test_llm_agent_callbacks.py + test_model_callback_chain.py + tests/unittests/runners: 95 passed, 1 skipped, 3 xfailed
• ruff check clean; pyink==25.12 (CI pin) and isort==8.0.1 clean. (ruff format disagrees with the repo's pyink style on pristine files too, so it was intentionally not applied.)

Manual End-to-End (E2E) Tests:

Ran the issue's verbatim repro script (InMemoryRunner, fake streaming model, no API key) before/after:

• Before: BUG run events all partial=None/final=True, 4 persisted model events.
• After: events partial=True/True/True/None, 1 persisted model event (Hello [REDACTED] world.), matching the issue's documented model_copy workaround shape.

Checklist

• I have read the CONTRIBUTING.md document.
• I have performed a self-review of my own code.
• I have commented my code, particularly in hard-to-understand areas.
• I have added tests that prove my fix is effective or that my feature works.
• New and existing unit tests pass locally with my changes.
• I have manually tested my changes end-to-end.
• Any dependent changes have been merged and published in downstream modules. (N/A — no dependents.)

Additional context

• Residual risk: usage_metadata has the same loss shape (a rebuilt response drops token accounting). Left out deliberately for minimal scope; happy to follow up if maintainers want it covered too.
• Note: I have not yet signed the Google CLA — will complete it if/when requested by the cla bot.
• AI assistance: an AI coding agent performed the reproduction, implementation, and tests under my direction; I reviewed the diff and take responsibility for it.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.