after_model_callback: a replacement LlmResponse drops usage_metadata, erasing the model call from token accounting #7451

# repro_usage_loss.py
"""Repro: after_model_callback replacement drops usage_metadata on current main."""
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

USAGE = types.GenerateContentResponseUsageMetadata(
    prompt_token_count=120, candidates_token_count=8, total_token_count=128)


def _resp(text, partial=None, usage=None):
  return LlmResponse(
      content=types.Content(role="model", parts=[types.Part(text=text)]),
      partial=partial, usage_metadata=usage)


class FakeLlm(BaseLlm):
  @classmethod
  def supported_models(cls): return [".*"]
  async def generate_content_async(self, llm_request: LlmRequest, stream=False
                                   ) -> AsyncGenerator[LlmResponse, None]:
    if stream:
      for d in ["Hello ", "world."]:
        yield _resp(d, partial=True)
    yield _resp("Hello world.", usage=USAGE)  # provider reports usage on final


def rebuild_cb(callback_context, llm_response):
  if not (llm_response.content and llm_response.content.parts): return None
  t = llm_response.content.parts[0].text or ""
  return _resp(t.replace("world", "[REDACTED]"))


async def run(label, cb, stream):
  agent = LlmAgent(name="a", model=FakeLlm(model="fake"), after_model_callback=cb)
  runner = InMemoryRunner(agent=agent, app_name="r")
  s = await runner.session_service.create_session(app_name="r", user_id="u")
  cfg = RunConfig(streaming_mode=StreamingMode.SSE if stream else StreamingMode.NONE)
  finals = []
  async for ev in runner.run_async(user_id="u", session_id=s.id,
      new_message=types.Content(role="user", parts=[types.Part(text="hi")]), run_config=cfg):
    if not ev.partial: finals.append(ev)
  stored = await runner.session_service.get_session(app_name="r", user_id="u", session_id=s.id)
  persisted = [e for e in stored.events if e.author == "a"]
  tok = lambda e: e.usage_metadata.total_token_count if e.usage_metadata else None
  print(f"{label:22} final_event.usage={tok(finals[-1])!s:5} persisted.usage={tok(persisted[-1])!s:5} partial_flags={[e.partial for e in persisted]}")

async def main():
  import google.adk; print("google-adk", google.adk.__version__)
  await run("CONTROL non-stream", None, False)
  await run("REBUILD non-stream", rebuild_cb, False)
  await run("CONTROL SSE", None, True)
  await run("REBUILD SSE", rebuild_cb, True)
asyncio.run(main())