fix: inherit usage_metadata when after_model_callback replaces response

Merge #7449

Fixes #7451

PiperOrigin-RevId: 996012278

Author: zhuhongd

src/google/adk/flows/llm_flows/core/_finalizer.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+from typing import Any
 from typing import AsyncGenerator
 from typing import Optional
 
@@ -169,24 +170,33 @@ async def handle_before_model_callback(
   return None
 
 
-def _inherit_unset_streaming_fields(
+def _inherit_unset_model_fields(
     original: LlmResponse, replacement: Optional[LlmResponse]
 ) -> Optional[LlmResponse]:
-  """Carries streaming-control fields from a replaced response.
+  """Carries model-owned 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.
+  the live path can close its request queue early.
+
+  ``usage_metadata`` is carried the same way. It measures the model call
+  (prompt, candidate and cached token counts as billed by the provider), not
+  the content of the response, so replacing the content does not change it.
+  A replacement that leaves it unset would otherwise erase the call from
+  token accounting: the finalized ``Event`` copies only non-None fields, so
+  the persisted session event, the BigQuery analytics plugin and A2A
+  converters all see no usage for that call.
+
+  An explicitly set value on the replacement is always respected.
 
   Args:
     original: The response being replaced.
     replacement: The callback-provided response, if there is one.
 
   Returns:
-    The replacement, with unset streaming-control fields filled in, or None
+    The replacement, with unset model-owned fields filled in, or None
     when there is no replacement. A copy is returned only when a field
     actually needs filling, so explicitly complete replacements keep their
     identity.
@@ -198,11 +208,13 @@ def _inherit_unset_streaming_fields(
       replacement, LlmResponse
   ):
     return replacement
-  updates = {}
+  updates: dict[str, Any] = {}
   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 replacement.usage_metadata is None and original.usage_metadata is not None:
+    updates['usage_metadata'] = original.usage_metadata
   if updates:
     return replacement.model_copy(update=updates)
   return replacement
@@ -266,7 +278,7 @@ async def _maybe_add_grounding_metadata(
       )
   )
   if callback_response:
-    return _inherit_unset_streaming_fields(
+    return _inherit_unset_model_fields(
         llm_response,
         await _maybe_add_grounding_metadata(callback_response),
     )
@@ -280,7 +292,7 @@ async def _maybe_add_grounding_metadata(
       llm_response=llm_response,
   )
   if callback_response:
-    return _inherit_unset_streaming_fields(
+    return _inherit_unset_model_fields(
         llm_response,
         await _maybe_add_grounding_metadata(callback_response),
     )

tests/unittests/flows/llm_flows/test_base_llm_flow_partial_handling.py

@@ -381,3 +381,87 @@ async def test_run_async_sse_rebuilt_responses_stay_partial():
       'scrubbed',
       'scrubbed',
   ]
+
+
+_USAGE = types.GenerateContentResponseUsageMetadata(
+    prompt_token_count=120,
+    candidates_token_count=8,
+    total_token_count=128,
+)
+
+
+@pytest.mark.asyncio
+async def test_after_model_callback_replacement_inherits_usage_metadata():
+  """A rebuilt replacement keeps the model call's token accounting."""
+  replacement = _text_response('scrubbed')
+  agent = Agent(
+      name='test_agent',
+      after_model_callback=[_rebuilding_callback(replacement)],
+  )
+  original = _text_response('Hello world.')
+  original.usage_metadata = _USAGE
+
+  result = await _run_handle_after_model_callback(agent, original)
+
+  assert result.usage_metadata == _USAGE
+  assert result.content.parts[0].text == 'scrubbed'
+  # The callback's object is not mutated; inheritance returns a copy.
+  assert replacement.usage_metadata is None
+
+
+@pytest.mark.asyncio
+async def test_after_model_callback_replacement_explicit_usage_respected():
+  """A replacement that reports its own usage is not overridden."""
+  own_usage = types.GenerateContentResponseUsageMetadata(
+      prompt_token_count=1, candidates_token_count=1, total_token_count=2
+  )
+  replacement = _text_response('scrubbed')
+  replacement.usage_metadata = own_usage
+  agent = Agent(
+      name='test_agent',
+      after_model_callback=[_rebuilding_callback(replacement)],
+  )
+  original = _text_response('Hello world.')
+  original.usage_metadata = _USAGE
+
+  result = await _run_handle_after_model_callback(agent, original)
+
+  assert result.usage_metadata == own_usage
+
+
+class _StreamingFakeModelWithUsage(_StreamingFakeModel):
+  """Like _StreamingFakeModel, but reports usage on the final response only."""
+
+  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)
+    final = _text_response(''.join(deltas))
+    final.usage_metadata = _USAGE
+    yield final
+
+
+@pytest.mark.asyncio
+async def test_run_async_sse_rebuilt_final_keeps_usage_metadata():
+  """End-to-end: the rebuilt final event still carries the call's usage."""
+  agent = Agent(
+      name='test_agent',
+      model=_StreamingFakeModelWithUsage(model='fake-streaming'),
+      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]
+  assert [event.usage_metadata for event in events] == [None, None, _USAGE]