feat(tracing): emit token usage on spans for SGP billing

[levilentz]

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

Filename	Overview
src/agentex/lib/adk/_modules/tracing.py	Adds TurnSpan class and turn_span context manager; handles non-dict span.data by warning and replacing; clean guard when span is None
src/agentex/lib/core/services/adk/providers/litellm.py	Adds _stream_kwargs_with_usage to force include_usage on streaming calls; attaches usage to span output in all three completion paths
src/agentex/lib/utils/completions.py	Fixes zip() truncation bug: zip_longest preserves choices from whichever side is longer, so the final usage-only chunk no longer wipes existing choices
src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py	Deleted — uninstantiable wrapper classes replaced by usage tracking directly in TemporalStreamingModel
tests/lib/test_claude_agents_activities.py	Fixes sys.modules leakage: tracks and removes placeholder package stubs after module load
tests/lib/adk/providers/test_litellm_usage.py	New tests covering usage propagation for all three LiteLLMService completion paths
tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py	New tests asserting TemporalStreamingModel copies ResponseUsage to both ModelResponse.usage and span.output["usage"]

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()

Loading

%%{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()

Loading

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}

[declan-scale]

Can we rebase and base this pr off of next

[declan-scale]

Checked this against the harness work now on next (#412–#438) since it's been rebased — looks well-integrated, low regression risk:

• ⚠️ Deleting temporal_tracing_model.py removes TemporalTracingModelProvider from the plugin's public __all__. Clean within this repo (no dangling refs in src/tests/examples), but if any external agent repo imports it directly that's a breaking change — worth a quick confirm.

All harness conformance suites (langgraph/codex/claude_code + conformance) pass. LGTM from a compatibility standpoint.