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

src/google/adk/telemetry/_metrics.py

@@ -580,11 +580,11 @@ def record_client_token_usage(
   if not responses:
     return
 
-  # The assumption is that token usage in streaming responses is cumulative.
-  # The last response chunk contains the total usage for the entire request.
-  # Summing them up across all response chunks would result in overcounting.
-  last_response = responses[-1]
-  if not last_response.usage_metadata:
+  # Usage in streaming responses is cumulative, so summing it would overcount.
+  # The newest report is the whole figure. It is not always on the last
+  # response: a stream cut off inside a tool call ends with one that has none.
+  token_usage = TokenUsage.from_llm_responses(responses)
+  if token_usage is None:
     logger.warning(
         "Skipping missing token usage metadata for agent %s and model %s",
         agent_name,
@@ -598,10 +598,9 @@ def record_client_token_usage(
   # thoughts tokens for "output".
   # `cached_content_token_count` is omitted as it's already included in prompt tokens.
   # `total_token_count` is omitted as SemConv expects input/output breakdown.
-  token_usage = TokenUsage.from_usage_metadata(last_response.usage_metadata)
   input_token_count = token_usage.input_tokens or 0
   output_token_count = token_usage.output_tokens or 0
-  response_model = last_response.model_version or llm_request.model
+  response_model = responses[-1].model_version or llm_request.model
   base_attrs = {
       gen_ai_attributes.GEN_AI_AGENT_NAME: agent_name,
       gen_ai_attributes.GEN_AI_OPERATION_NAME: "generate_content",

tests/unittests/telemetry/test_metrics.py

@@ -17,6 +17,7 @@
 import os
 from unittest import mock
 
+from google.adk.models.llm_response import LlmResponse
 from google.adk.telemetry import _hallucination
 from google.adk.telemetry import _metrics
 from google.adk.telemetry import _token_usage
@@ -458,6 +459,45 @@ def test_record_client_token_usage(mock_meter_setup):
   }
 
 
+def test_record_client_token_usage_keeps_usage_before_trailing_response(
+    mock_meter_setup,
+):
+  """Usage reported before the last response of a call is still recorded.
+
+  A LiteLlm stream cut off by max_tokens inside a tool call yields the text
+  with the usage, then an error response that carries none.
+  """
+  llm_request = mock.MagicMock(model="test-model")
+  responses = [
+      LlmResponse(
+          content=types.ModelContent("Let me check the weather."),
+          finish_reason=types.FinishReason.MAX_TOKENS,
+          model_version="test-model-v1",
+          usage_metadata=types.GenerateContentResponseUsageMetadata(
+              prompt_token_count=120, candidates_token_count=40
+          ),
+      ),
+      LlmResponse(
+          finish_reason=types.FinishReason.MAX_TOKENS,
+          model_version="test-model-v1",
+      ),
+  ]
+
+  _metrics.record_client_token_usage(
+      agent_name="test_agent",
+      llm_request=llm_request,
+      responses=responses,
+  )
+
+  recorded = {
+      kwargs["attributes"]["gen_ai.token.type"]: args[0]
+      for args, kwargs in mock_meter_setup[
+          "client_token_usage"
+      ].record.call_args_list
+  }
+  assert recorded == {"input": 120, "output": 40}
+
+
 @pytest.fixture(name="call_count_histograms")
 def _call_count_histograms(monkeypatch):
   """Redirects the two per-invocation call-count histograms."""