From 32ee07df01f10dbee0e98ca9d412440a7fe9163d Mon Sep 17 00:00:00 2001
From: Benson Wang <91510760+wsa-2002@users.noreply.github.com>
Date: Mon, 9 Feb 2026 11:23:49 -0800
Subject: [PATCH] fix: prompt token may be None in streaming mode

Merge https://github.com/google/adk-python/pull/3462

**Please ensure you have read the [contribution guide](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) before creating a pull request.**

### Link to Issue or Description of Change

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

- Closes: #_issue_number_
- Related: #_issue_number_

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

**Problem:**
When using adk in streaming mode, `usage_metadata.prompt_token_count` may be `None` which will emit log
```Invalid type NoneType for attribute 'gen_ai.usage.input_tokens' value. Expected one of ['bool', 'str', 'bytes', 'int', 'float'] or a sequence of those types```

**Solution:**
Skip setting span attribute if prompt token count is None

**Unit Tests:**

- [x] All unit tests pass locally.

_Please include a summary of passed `pytest` results._

### Checklist

- [x] I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) document.
- [x] I have performed a self-review of my own code.
- [x] I have commented my code, particularly in hard-to-understand areas.
- [x] I have added tests that prove my fix is effective or that my feature works.
- [x] New and existing unit tests pass locally with my changes.
- [x] I have manually tested my changes end-to-end.
- [x] Any dependent changes have been merged and published in downstream modules.

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/3462 from wsa-2002:prompt-token-count-may-be-none-in-streaming-mode 94666862f70ed2577d5c55485e67f6da36a57bc6
PiperOrigin-RevId: 867693355
---
 src/google/adk/telemetry/tracing.py     |  9 ++---
 tests/unittests/telemetry/test_spans.py | 44 +++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/src/google/adk/telemetry/tracing.py b/src/google/adk/telemetry/tracing.py
index cd6b7071..fbb55ec9 100644
--- a/src/google/adk/telemetry/tracing.py
+++ b/src/google/adk/telemetry/tracing.py
@@ -327,10 +327,11 @@ def trace_call_llm(
     span.set_attribute('gcp.vertex.agent.llm_response', '{}')
 
   if llm_response.usage_metadata is not None:
-    span.set_attribute(
-        'gen_ai.usage.input_tokens',
-        llm_response.usage_metadata.prompt_token_count,
-    )
+    if llm_response.usage_metadata.prompt_token_count is not None:
+      span.set_attribute(
+          'gen_ai.usage.input_tokens',
+          llm_response.usage_metadata.prompt_token_count,
+      )
     if llm_response.usage_metadata.candidates_token_count is not None:
       span.set_attribute(
           'gen_ai.usage.output_tokens',
diff --git a/tests/unittests/telemetry/test_spans.py b/tests/unittests/telemetry/test_spans.py
index 2deebda1..bb084676 100644
--- a/tests/unittests/telemetry/test_spans.py
+++ b/tests/unittests/telemetry/test_spans.py
@@ -167,6 +167,50 @@ async def test_trace_call_llm(monkeypatch, mock_span_fixture):
   )
 
 
+@pytest.mark.asyncio
+async def test_trace_call_llm_with_no_usage_metadata(
+    monkeypatch, mock_span_fixture
+):
+  """Test trace_call_llm handles usage metadata with None token counts."""
+  monkeypatch.setattr(
+      'opentelemetry.trace.get_current_span', lambda: mock_span_fixture
+  )
+
+  agent = LlmAgent(name='test_agent')
+  invocation_context = await _create_invocation_context(agent)
+  llm_request = LlmRequest(
+      model='gemini-pro',
+      contents=[
+          types.Content(
+              role='user',
+              parts=[types.Part(text='Hello, how are you?')],
+          ),
+      ],
+      config=types.GenerateContentConfig(
+          top_p=0.95,
+          max_output_tokens=1024,
+      ),
+  )
+  llm_response = LlmResponse(
+      turn_complete=True,
+      finish_reason=types.FinishReason.STOP,
+      usage_metadata=types.GenerateContentResponseUsageMetadata(),
+  )
+  trace_call_llm(invocation_context, 'test_event_id', llm_request, llm_response)
+
+  expected_calls = [
+      mock.call('gen_ai.system', 'gcp.vertex.agent'),
+      mock.call('gen_ai.request.top_p', 0.95),
+      mock.call('gen_ai.request.max_tokens', 1024),
+      mock.call('gcp.vertex.agent.llm_response', mock.ANY),
+      mock.call('gen_ai.response.finish_reasons', ['stop']),
+  ]
+  assert mock_span_fixture.set_attribute.call_count == 10
+  mock_span_fixture.set_attribute.assert_has_calls(
+      expected_calls, any_order=True
+  )
+
+
 @pytest.mark.asyncio
 async def test_trace_call_llm_with_binary_content(
     monkeypatch, mock_span_fixture