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 ensures every SDK tracing adapter emits real token usage so SGP billing can read span.output["usage"] or the per-turn rollup span.data["usage"] + span.data["cost_usd"] instead of billing zero. It introduces TurnSpan / turn_span() to capture per-turn aggregate usage, and fixes three concrete bugs found along the way.

• concat_completion_chunks fix — replaces zip(…, strict=False) with zip_longest so the usage-only final chunk emitted by stream_options.include_usage: True no longer wipes the accumulated choices; tested with new tests/lib/utils/test_completions.py.
• LiteLLM streaming usage — adds _stream_kwargs_with_usage to default include_usage: True and propagates the resulting completion.usage onto span.output["usage"] for all three litellm paths; chunks are now always appended before the choices guard so the usage-only tail chunk reaches concat_completion_chunks.
• TurnSpan / TracingModule.turn_span() — new context manager that writes the turn's rollup usage to span.data["usage"] / span.data["cost_usd"]; accepts TurnUsage or a plain recognized-key dict; warns on non-dict span.data and unrecognized keys; deletes the uninstantiable TemporalTracingModelProvider wrappers; fixes sys.modules stub leakage in claude_agents tests.

Confidence Score: 5/5

Safe to merge — all three previously broken usage emission paths are fixed, the concat bug is correctly addressed, and 30 new tests cover the changed behavior end-to-end.

The changes are tightly scoped to tracing/billing paths with no risk of disrupting the main agent execution path. The zip_longest fix is the most load-bearing change, and it is both logically correct and verified by tests covering the usage-only tail chunk, summed usage, and the empty-choices edge case. The TurnSpan API is well-guarded (no-op when span is None, warning on unexpected types). Deletion of TemporalTracingModelProvider removes dead/broken code without changing any active execution path since TemporalStreamingModel already handles the tracing.

No files require special attention.

Important Files Changed

Filename	Overview
src/agentex/lib/utils/completions.py	Replaces zip with zip_longest so usage-only final chunks (choices=[]) no longer truncate the accumulated choices list; logic is correct and well-tested.
src/agentex/lib/adk/_modules/tracing.py	Adds TurnSpan class and turn_span() context manager; correctly extracts cost_usd from TurnUsage, warns on non-dict span.data, and preserves existing dict data.
src/agentex/lib/core/services/adk/providers/litellm.py	Adds _stream_kwargs_with_usage helper and emits completion.usage onto span.output for all three completion paths; moves chunks.append before the choices guard to capture usage-only tail chunks.
src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py	Deleted entirely — the abstract-class wrappers were uninstantiable and dropped usage; TemporalStreamingModel already handles tracing with real usage.
tests/lib/adk/test_tracing_module.py	Adds comprehensive TurnSpan tests covering TurnUsage, plain dict, cost override, unrecognized key warning, existing data preservation, and no-op when tracing is disabled.
tests/lib/adk/providers/test_litellm_usage.py	New test file covering all three litellm paths for usage emission, include_usage injection, and usage-only chunk handling.
tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py	New tests verifying TemporalStreamingModel captures ResponseUsage (including cached/reasoning token details) onto span.output and ModelResponse.usage.
tests/lib/test_claude_agents_activities.py	Fixes sys.modules stub leakage by tracking and removing placeholder package entries after module load.
tests/lib/test_claude_agents_hooks.py	Same sys.modules cleanup fix applied symmetrically to the hooks test file.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Agent
    participant turn_span as adk.tracing.turn_span()
    participant LiteLLM as LiteLLMService
    participant Gateway as LiteLLMGateway
    participant Backend as SGP Backend

    Agent->>turn_span: async with turn_span(trace_id, name, ...)
    turn_span-->>Agent: TurnSpan handle

    Agent->>LiteLLM: chat_completion_stream_auto_send(...)
    LiteLLM->>LiteLLM: _stream_kwargs_with_usage(llm_config)
    LiteLLM->>Gateway: "acompletion_stream(**completion_kwargs)"
    Gateway-->>LiteLLM: "delta chunks [choices=[...delta...]]"
    Gateway-->>LiteLLM: "usage-only chunk [choices=[], usage={...}]"
    Note over LiteLLM: chunks.append(every chunk), zip_longest fix keeps choices + usage
    LiteLLM->>LiteLLM: concat_completion_chunks(chunks)
    LiteLLM->>LiteLLM: "span.output = task_message + usage"

    Agent->>turn_span: turn.record_usage(result.usage)
    turn_span->>turn_span: "span.data = {usage: {...}, cost_usd: 0.xx}"

    turn_span->>Backend: end_span(span)
    Backend->>Backend: de-dup per-call output[usage] against aggregate data[usage]
    Backend-->>Agent: billed once per turn

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 turn_span as adk.tracing.turn_span()
    participant LiteLLM as LiteLLMService
    participant Gateway as LiteLLMGateway
    participant Backend as SGP Backend

    Agent->>turn_span: async with turn_span(trace_id, name, ...)
    turn_span-->>Agent: TurnSpan handle

    Agent->>LiteLLM: chat_completion_stream_auto_send(...)
    LiteLLM->>LiteLLM: _stream_kwargs_with_usage(llm_config)
    LiteLLM->>Gateway: "acompletion_stream(**completion_kwargs)"
    Gateway-->>LiteLLM: "delta chunks [choices=[...delta...]]"
    Gateway-->>LiteLLM: "usage-only chunk [choices=[], usage={...}]"
    Note over LiteLLM: chunks.append(every chunk), zip_longest fix keeps choices + usage
    LiteLLM->>LiteLLM: concat_completion_chunks(chunks)
    LiteLLM->>LiteLLM: "span.output = task_message + usage"

    Agent->>turn_span: turn.record_usage(result.usage)
    turn_span->>turn_span: "span.data = {usage: {...}, cost_usd: 0.xx}"

    turn_span->>Backend: end_span(span)
    Backend->>Backend: de-dup per-call output[usage] against aggregate data[usage]
    Backend-->>Agent: billed once per turn

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 (4): 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.