Thinking tokens not in traces #4829

# In trace_call_llm() (~line 329):
if llm_response.usage_metadata is not None:
    if llm_response.usage_metadata.prompt_token_count is not None:
        span.set_attribute('gen_ai.usage.input_tokens', ...)
    if llm_response.usage_metadata.candidates_token_count is not None:
        span.set_attribute('gen_ai.usage.output_tokens', ...)
    # thoughts_token_count is NOT exported here

    if llm_response.usage_metadata.thoughts_token_count is not None:
        span.set_attribute('gen_ai.usage.thinking_tokens', 
                           llm_response.usage_metadata.thoughts_token_count)

# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "google-adk>=1.26.0",
#     "google-genai>=1.65.0",
# ]
# ///
"""Minimal reproduction: thoughts_token_count present in Event but missing from traces.

Run with:
  export GOOGLE_GENAI_USE_VERTEXAI=true
  export GOOGLE_CLOUD_PROJECT=<your-project>
  uv run repro_thinking_trace.py
"""
import asyncio

from google.adk import Runner
from google.adk.agents import Agent
from google.adk.memory import InMemoryMemoryService
from google.adk.sessions import InMemorySessionService
from google.genai import types


def get_weather(city: str) -> dict:
    """Get current weather for a city."""
    return {"city": city, "temp_f": 72, "condition": "sunny"}


async def main():
    agent = Agent(
        model="gemini-2.5-flash",
        name="weather_agent",
        instruction="You are a weather assistant. Always use the get_weather tool.",
        tools=[get_weather],
        generate_content_config=types.GenerateContentConfig(
            thinking_config=types.ThinkingConfig(thinking_budget=2048),
        ),
        disallow_transfer_to_parent=True,
        disallow_transfer_to_peers=True,
    )

    session_service = InMemorySessionService()
    runner = Runner(
        agent=agent,
        app_name="repro",
        session_service=session_service,
        memory_service=InMemoryMemoryService(),
    )
    session = await session_service.create_session(
        user_id="test", app_name="repro"
    )

    message = types.Content(
        parts=[types.Part(text="What's the weather in San Francisco?")]
    )

    async for event in runner.run_async(
        new_message=message, user_id="test", session_id=session.id
    ):
        usage = getattr(event, "usage_metadata", None)
        if usage is not None:
            thoughts = getattr(usage, "thoughts_token_count", None)
            print(
                f"Event author={event.author}: "
                f"prompt={usage.prompt_token_count} "
                f"candidates={usage.candidates_token_count} "
                f"thoughts={thoughts} "
                f"total={usage.total_token_count}"
            )
            if thoughts is not None and thoughts > 0:
                print(
                    "  ^^^ thoughts_token_count is non-zero in Event, "
                    "but will NOT appear in Cloud Trace span attributes"
                )


if __name__ == "__main__":
    asyncio.run(main())