fix(telemetry): keep client token usage when the last response has none

[Somay-kousis]

Link to Issue or Description of Change

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

• Closes: gen_ai.client.token.usage records nothing when the last response of a model call has no usage #7560

2. Or, if no issue exists, describe the change:

Problem:

record_client_token_usage reads usage only from the last response of a model call. If that response has no usage_metadata, it records nothing for gen_ai.client.token.usage, even when an earlier response of the same call reported the usage.

LiteLlm streaming does this when the model hits max_tokens inside a tool call. It yields the text with the usage, clears the usage, then yields an error response without it (lite_llm.py:4313-4330). The call's tokens never reach the metric.

Solution:

Use TokenUsage.from_llm_responses, which takes the newest usage report. record_inference_telemetry already uses it for the per-agent token totals a few lines earlier, so the metric and the totals now read the same report. gen_ai.response.model still comes from the last response, as before.

When the last response does have usage, from_llm_responses picks that same report, so nothing changes for calls that already recorded. One edge moves: a last response whose usage has neither prompt nor candidate counts used to record a 0/0 point, and now it's skipped with the same warning, which matches the per-agent totals.

Testing Plan

Unit Tests:

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

New test test_record_client_token_usage_keeps_usage_before_trailing_response in tests/unittests/telemetry/test_metrics.py: a response with usage, then a MAX_TOKENS response without it.

Before the fix:

FAILED tests/unittests/telemetry/test_metrics.py::test_record_client_token_usage_keeps_usage_before_trailing_response
1 failed, 1 passed, 38 deselected in 1.77s

After the fix:

$ pytest tests/unittests/telemetry/test_metrics.py -k record_client_token_usage
2 passed, 38 deselected in 1.56s

$ pytest tests/unittests/telemetry tests/unittests/tools/model_consult
1052 passed, 7 xfailed, 44 warnings in 47.42s

pre-commit run --files on both files passes. mypy on _metrics.py reports one error at the "compaction" feature list, and main without this change reports the same error (from d01d5d8). This change adds none.

Manual End-to-End (E2E) Tests:

The script from the issue runs an LlmAgent on LiteLlm with a fake client that replays a stream cut off by max_tokens inside a tool call, in SSE mode, and reads the metric from an InMemoryMetricReader.

Before:

Skipping missing token usage metadata for agent weather and model openai/gpt-4o
adk module: src/google/adk/telemetry/_metrics.py
event: FinishReason.MAX_TOKENS usage: (120, 40)
event: FinishReason.MAX_TOKENS usage: None
gen_ai.client.token.usage: {}

After:

adk module: src/google/adk/telemetry/_metrics.py
event: FinishReason.MAX_TOKENS usage: (120, 40)
event: FinishReason.MAX_TOKENS usage: None
gen_ai.client.token.usage: {'input': 120, 'output': 40}

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.

Additional context

N/A

[gemini-code-assist]

Code Review

This pull request updates token usage tracking in streaming responses to use TokenUsage.from_llm_responses(responses) instead of relying solely on the last response chunk, which may lack metadata if the stream is cut off. A unit test was added to verify this behavior. The reviewer pointed out that responses[-1].model_version might also be missing in trailing chunks, and suggested searching backwards through the responses to find the first non-null model version.

[gemini-code-assist]

Since the last response in a stream can be a trailing chunk or error response that lacks metadata (which is the core issue this PR addresses), it is highly likely that responses[-1].model_version will also be None. To ensure we don't lose the resolved model version, we should retrieve the model version from the latest response chunk that actually populated it, falling back to llm_request.model only if none of the responses have it.

Suggested change

	response_model = responses[-1].model_version or llm_request.model
	response_model = next((r.model_version for r in reversed(responses) if r.model_version), None) or llm_request.model

[Somay-kousis]

Kept this as it was on purpose. Main already reads model_version from the last response, and in the LiteLlm stream this PR is about, the text and tool-call responses get the same stream-level model version, so the last one isn't missing it. I'd rather keep this PR to the usage fix and leave the model attribute unchanged.