2020 TracingActivityName ,
2121)
2222from agentex .lib .core .tracing .tracer import AsyncTracer
23+ from agentex .lib .core .harness .types import TurnUsage
2324from agentex .types .span import Span
2425from agentex .lib .utils .logging import make_logger
2526from agentex .lib .utils .model_utils import BaseModel
3031DEFAULT_RETRY_POLICY = RetryPolicy (maximum_attempts = 1 )
3132TEMPORAL_SPAN_ACTIVITY_DROPPED_METRIC = "agentex.tracing.temporal_span_activity.dropped"
3233
34+ # Token key spellings the backend accepts when billing usage from spans.
35+ RECOGNIZED_USAGE_KEYS = frozenset (
36+ {
37+ "input_tokens" ,
38+ "prompt_tokens" ,
39+ "output_tokens" ,
40+ "completion_tokens" ,
41+ "cached_input_tokens" ,
42+ "cached_tokens" ,
43+ "reasoning_tokens" ,
44+ "total_tokens" ,
45+ "cost_usd" ,
46+ }
47+ )
48+
3349
3450def _record_temporal_span_activity_dropped (event_type : str ) -> None :
3551 try :
@@ -42,6 +58,80 @@ def _record_temporal_span_activity_dropped(event_type: str) -> None:
4258 pass
4359
4460
61+ class TurnSpan :
62+ """Handle for a turn-level (rollup) span, yielded by ``TracingModule.turn_span``.
63+
64+ Encapsulates the billing contract so agents cannot double-count usage:
65+ the turn's aggregate usage goes to ``span.data["usage"]`` (+
66+ ``span.data["cost_usd"]``) via :meth:`record_usage`. The backend keeps the
67+ aggregate and de-dups any per-call ``output["usage"]`` children against it.
68+ Never hand-write usage into ``output`` on a rollup span — that is the
69+ double-count bug this helper exists to prevent.
70+
71+ All methods no-op when tracing is disabled (``span`` is None), so agent
72+ code needs no ``if span:`` guards.
73+ """
74+
75+ def __init__ (self , span : Span | None ):
76+ self .span = span
77+
78+ def record_usage (
79+ self ,
80+ usage : TurnUsage | dict [str , Any ] | None = None ,
81+ cost_usd : float | None = None ,
82+ ) -> None :
83+ """Record the turn's aggregate usage on the span's ``data``.
84+
85+ Pass the harness ``TurnUsage`` (e.g. ``LangGraphTurn.usage()`` or
86+ ``run_turn(...).usage``) — its ``cost_usd`` is stamped automatically —
87+ or a plain dict with backend-recognized token spellings
88+ (``prompt_tokens``/``completion_tokens`` also work). An explicit
89+ ``cost_usd`` argument overrides any cost carried by ``usage``. The
90+ usage must be this turn's own tokens, not a session-cumulative total.
91+ """
92+ if self .span is None :
93+ return
94+
95+ blob : dict [str , Any ]
96+ if isinstance (usage , TurnUsage ):
97+ blob = usage .model_dump (exclude_none = True )
98+ # cost lives beside the blob as data["cost_usd"], not inside it
99+ blob_cost = blob .pop ("cost_usd" , None )
100+ if cost_usd is None :
101+ cost_usd = blob_cost
102+ elif usage is not None :
103+ blob = dict (usage )
104+ if not any (key in RECOGNIZED_USAGE_KEYS for key in blob ):
105+ logger .warning (
106+ "TurnSpan.record_usage: usage has no recognized token keys and will "
107+ f"not be billed. Got keys { sorted (blob )}
108+ f"{ sorted (RECOGNIZED_USAGE_KEYS )}
109+ )
110+ else :
111+ blob = {}
112+
113+ if self .span .data is not None and not isinstance (self .span .data , dict ):
114+ logger .warning (
115+ f"TurnSpan.record_usage: span.data is { type (self .span .data ).__name__ }
116+ "(expected dict or None); existing data will be replaced."
117+ )
118+ data = self .span .data if isinstance (self .span .data , dict ) else {}
119+ if blob :
120+ data ["usage" ] = blob
121+ if cost_usd is not None :
122+ data ["cost_usd" ] = cost_usd
123+ self .span .data = data
124+
125+ @property
126+ def output (self ) -> Any :
127+ return self .span .output if self .span is not None else None
128+
129+ @output .setter
130+ def output (self , value : Any ) -> None :
131+ if self .span is not None :
132+ self .span .output = value
133+
134+
45135class TracingModule :
46136 """
47137 Module for managing tracing and span operations in Agentex.
@@ -156,6 +246,47 @@ async def span(
156246 retry_policy = retry_policy ,
157247 )
158248
249+ @asynccontextmanager
250+ async def turn_span (
251+ self ,
252+ trace_id : str ,
253+ name : str ,
254+ input : list [Any ] | dict [str , Any ] | BaseModel | None = None ,
255+ data : list [Any ] | dict [str , Any ] | BaseModel | None = None ,
256+ parent_id : str | None = None ,
257+ task_id : str | None = None ,
258+ start_to_close_timeout : timedelta = timedelta (seconds = 5 ),
259+ heartbeat_timeout : timedelta = timedelta (seconds = 5 ),
260+ retry_policy : RetryPolicy = DEFAULT_RETRY_POLICY ,
261+ ) -> AsyncGenerator [TurnSpan , None ]:
262+ """Span for one agent turn, with usage recorded as the billable aggregate.
263+
264+ Same lifecycle as :meth:`span`, but yields a :class:`TurnSpan` whose
265+ ``record_usage(usage=..., cost_usd=...)`` writes the turn's rollup
266+ usage to ``span.data`` — the shape the backend bills once per turn.
267+ Per-call child spans (LLM adapters) may still carry
268+ ``output["usage"]``; the backend de-dups them against this aggregate.
269+
270+ Example (with a harness turn, e.g. ``LangGraphTurn`` / ``run_turn``)::
271+
272+ async with adk.tracing.turn_span(trace_id=task.id, name="turn", input={...}, task_id=task.id) as turn:
273+ result = await run_turn(...)
274+ turn.output = {"response": result.final_output}
275+ turn.record_usage(result.usage) # TurnUsage, cost_usd included
276+ """
277+ async with self .span (
278+ trace_id = trace_id ,
279+ name = name ,
280+ input = input ,
281+ data = data ,
282+ parent_id = parent_id ,
283+ task_id = task_id ,
284+ start_to_close_timeout = start_to_close_timeout ,
285+ heartbeat_timeout = heartbeat_timeout ,
286+ retry_policy = retry_policy ,
287+ ) as span :
288+ yield TurnSpan (span )
289+
159290 async def start_span (
160291 self ,
161292 trace_id : str ,
0 commit comments