Makes every SDK tracing adapter emit real token usage for tracking: the OpenAI Agents SDK temporal models, the LangGraph handler, and litellm streaming/auto_send previously dropped usage and billed zero. Adds adk.tracing.turn_span() so agents record the per-turn aggregate (data["usage"] + cost_usd) instead of hand-rolling the contract that caused double-counting. Also fixes three bugs found along the way (uninstantiable tracing wrapper models, chunk concat wiping choices on empty-choices chunks, sys.modules leakage in the claude_agents tests); 30 new tests, tests/lib at 260 passing, ruff and pyright clean.

Greptile Summary

This PR fixes all three tracing adapters (litellm, OpenAI Agents SDK, LangGraph) so they emit real token usage onto spans for SGP billing, and introduces adk.tracing.turn_span() to record the per-turn rollup (data["usage"] + data["cost_usd"]) in the shape the backend bills from, preventing the double-counting bug that arose from hand-rolling that contract.

• TurnSpan / turn_span: new context manager that wraps the existing span() lifecycle and exposes record_usage(TurnUsage | dict, cost_usd?), writing the aggregate to span.data (not span.output) so child LLM spans can still carry output["usage"] without double-counting.
• LiteLLM streaming fix (litellm.py + completions.py): _stream_kwargs_with_usage appends stream_options={"include_usage": True} to every streaming call; concat_completion_chunks is switched from zip to zip_longest so a trailing usage-only chunk (empty choices) no longer wipes the assembled content choices.
• Cleanup: removes the uninstantiable TemporalTracingModelProvider/TemporalTracingResponsesModel/TemporalTracingChatCompletionsModel wrapper trio; fixes sys.modules stub leakage in claude_agents tests that was blocking collection of real submodule imports in later tests.

Confidence Score: 5/5

Safe to merge — all three adapters now emit real usage, the zip_longest fix is correct and well-tested, and the sys.modules cleanup is a targeted fix with no collateral effects.

The changes are well-contained billing instrumentation: new read paths (usage extraction) rather than modifications to request or message logic. The zip_longest fix is provably correct, the TurnSpan API cleanly prevents double-counting by routing aggregate usage to span.data instead of output, and 30 new tests cover every code path including edge cases.

No files require special attention.

Important Files Changed

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Agent
    participant TurnSpan
    participant TracingModule
    participant LiteLLMService
    participant Span

    Agent->>TracingModule: turn_span(trace_id, name, ...)
    TracingModule->>Span: start_span()
    TracingModule-->>Agent: yield TurnSpan(span)

    Agent->>LiteLLMService: chat_completion_stream_auto_send(...)
    Note over LiteLLMService: _stream_kwargs_with_usage() adds include_usage=True
    LiteLLMService->>Span: "span.output = {task_message, usage}"
    LiteLLMService-->>Agent: TaskMessage

    Agent->>TurnSpan: record_usage(result.usage, cost_usd)
    Note over TurnSpan: Aggregate to span.data["usage"] + cost_usd
    TurnSpan->>Span: "span.data = {usage, cost_usd}"

    Agent->>TracingModule: exit turn_span
    TracingModule->>Span: end_span()

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Agent
    participant TurnSpan
    participant TracingModule
    participant LiteLLMService
    participant Span

    Agent->>TracingModule: turn_span(trace_id, name, ...)
    TracingModule->>Span: start_span()
    TracingModule-->>Agent: yield TurnSpan(span)

    Agent->>LiteLLMService: chat_completion_stream_auto_send(...)
    Note over LiteLLMService: _stream_kwargs_with_usage() adds include_usage=True
    LiteLLMService->>Span: "span.output = {task_message, usage}"
    LiteLLMService-->>Agent: TaskMessage

    Agent->>TurnSpan: record_usage(result.usage, cost_usd)
    Note over TurnSpan: Aggregate to span.data["usage"] + cost_usd
    TurnSpan->>Span: "span.data = {usage, cost_usd}"

    Agent->>TracingModule: exit turn_span
    TracingModule->>Span: end_span()

Comments Outside Diff (1)

1. src/agentex/lib/adk/_modules/_langgraph_tracing.py, line 89-91 (link)

   Usage overwritten on multi-generation results

   _extract_usage is called inside the inner for generation in generation_list loop, so output["usage"] gets overwritten on each iteration. For n=1 calls (the common case) this is harmless, but if a caller ever requests n > 1 completions only the last generation's usage_metadata contributes to the billing blob — earlier generations are silently dropped. The llm_output fallback is safe because it reads the response-level aggregate and returns the same dict every iteration, but the usage_metadata path is per-generation and should accumulate.

   Prompt To Fix With AI

   This is a comment left during a code review.
   Path: src/agentex/lib/adk/_modules/_langgraph_tracing.py
   Line: 89-91
   
   Comment:
   **Usage overwritten on multi-generation results**
   
   `_extract_usage` is called inside the inner `for generation in generation_list` loop, so `output["usage"]` gets overwritten on each iteration. For `n=1` calls (the common case) this is harmless, but if a caller ever requests `n > 1` completions only the last generation's `usage_metadata` contributes to the billing blob — earlier generations are silently dropped. The `llm_output` fallback is safe because it reads the response-level aggregate and returns the same dict every iteration, but the `usage_metadata` path is per-generation and should accumulate.
   
   How can I resolve this? If you propose a fix, please make it concise.

   

Reviews (5): Last reviewed commit: "refactor(adk): record_usage accepts harn..." | Re-trigger Greptile