fix: preserve partial and turn_complete when after_model_callback returns a replacement

[zhuhongd]

Link to Issue or Description of Change

1. Link to an existing issue:

• 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:

The documented after_model_callback contract invites callbacks to return a replacement response, but during SSE streaming the callback fires on every streamed LlmResponse and the replacement is used wholesale. LlmResponse.partial defaults to None, which the streaming protocol treats as final, so any rebuilt replacement silently converts every delta into a final response: clients render N "final" responses, and the Runner persists every fragment to the session as a separate complete model event (session history corruption). Self-contained repro in #7035.

Solution:

Treat None as "unset" and inherit partial / turn_complete from the response being replaced, at both after-model-callback call sites (live and SSE paths), via a small helper _inherit_unset_streaming_fields. Explicit values set by a callback (e.g. partial=False to deliberately finalize) are still respected. This makes the documented callback pattern safe without removing any behavior a callback could previously express intentionally.

Testing Plan

Unit Tests:

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

Added to tests/unittests/flows/llm_flows/test_base_llm_flow_partial_handling.py:

• test_after_model_callback_replacement_preserves_partial_in_sse — a
  doc-contract callback (rebuilds LlmResponse) on a streamed turn of 3
  partial deltas + aggregated final now yields [True, True, True, None]
  partial flags (was [None, None, None, None] before the fix).
• test_after_model_callback_explicit_partial_false_is_respected — a
  callback that deliberately finalizes a delta is not overridden.
• test_inherit_unset_streaming_fields_inherits_when_unset and
  test_inherit_unset_streaming_fields_respects_explicit_values — direct
  unit tests of the helper, including turn_complete.

tests/unittests/flows/llm_flows/test_base_llm_flow_partial_handling.py
  8 passed (4 pre-existing + 4 new)

tests/unittests/flows/llm_flows/  (full directory regression)
  716 passed

Manual End-to-End (E2E) Tests:

Ran the self-contained repro script from #7035 (fake streaming model, no API
key required) against this branch: the BUG scenario now matches CONTROL —
deltas stay partial=True and the session persists exactly one final model
event instead of one event per delta. The model_copy workaround scenario is
unchanged.

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 dependent changes)

Additional context

Formatted with isort + pyink per pyproject.toml. Found while running a
production multi-agent assistant where a guardrail-style
after_model_callback rebuilt responses per the documented contract.