From a6f23619c4b1938083005571e791efe9fb8b0117 Mon Sep 17 00:00:00 2001 From: ashrafchowdury Date: Thu, 16 Jul 2026 19:07:21 +0600 Subject: [PATCH 1/3] feat: enhance agent token usage tracking with cache metrics --- .../sdk/agents/adapters/vercel/stream.py | 2 +- sdks/python/agenta/sdk/agents/tracing.py | 8 ++++ sdks/python/agenta/sdk/agents/wire_models.py | 8 +++- .../unit/agents/golden/run_result.ok.json | 4 +- .../pytest/unit/agents/test_wire_contract.py | 11 ++++- .../runner/src/engines/sandbox_agent/usage.ts | 16 ++++++- services/runner/src/protocol.ts | 10 ++++- services/runner/src/tracing/otel.ts | 44 +++++++++++++++++-- .../runner/tests/unit/wire-contract.test.ts | 6 ++- .../components/AgentChatSlice/assets/trace.ts | 18 ++++++-- .../metrics/ExecutionMetricsDisplay.tsx | 18 ++++++++ 11 files changed, 129 insertions(+), 16 deletions(-) diff --git a/sdks/python/agenta/sdk/agents/adapters/vercel/stream.py b/sdks/python/agenta/sdk/agents/adapters/vercel/stream.py index 0adf7b18b0..cc50eec87a 100644 --- a/sdks/python/agenta/sdk/agents/adapters/vercel/stream.py +++ b/sdks/python/agenta/sdk/agents/adapters/vercel/stream.py @@ -781,7 +781,7 @@ def _approval_tool_call_id(payload: Dict[str, Any]) -> Optional[Any]: def _usage_metadata(data: Dict[str, Any]) -> Dict[str, Any]: return { key: data[key] - for key in ("input", "output", "total", "cost") + for key in ("input", "output", "total", "cost", "cacheRead", "cacheWrite") if data.get(key) is not None } diff --git a/sdks/python/agenta/sdk/agents/tracing.py b/sdks/python/agenta/sdk/agents/tracing.py index 66a9b8e8c7..6d9133e1c4 100644 --- a/sdks/python/agenta/sdk/agents/tracing.py +++ b/sdks/python/agenta/sdk/agents/tracing.py @@ -229,6 +229,14 @@ def record_usage(usage: Optional[Dict[str, Any]]) -> None: span.set_attribute("gen_ai.usage.prompt_tokens", input_tokens) span.set_attribute("gen_ai.usage.completion_tokens", output_tokens) span.set_attribute("gen_ai.usage.total_tokens", int(usage.get("total") or 0)) + cache_read = usage.get("cacheRead") + if cache_read: + span.set_attribute("gen_ai.usage.cache_read.input_tokens", int(cache_read)) + cache_write = usage.get("cacheWrite") + if cache_write: + span.set_attribute( + "gen_ai.usage.cache_creation.input_tokens", int(cache_write) + ) cost = usage.get("cost") if cost: span.set_attribute("gen_ai.usage.cost", float(cost)) diff --git a/sdks/python/agenta/sdk/agents/wire_models.py b/sdks/python/agenta/sdk/agents/wire_models.py index eb5366cc1d..887bdac1cd 100644 --- a/sdks/python/agenta/sdk/agents/wire_models.py +++ b/sdks/python/agenta/sdk/agents/wire_models.py @@ -354,12 +354,18 @@ class WireHarnessCapabilities(_WireModel): class WireAgentUsage(_WireModel): - """Token / cost usage rolled onto a workflow span.""" + """Token / cost usage rolled onto a workflow span. + + ``input`` counts non-cached prompt tokens only; the cached portion is reported + separately so consumers can show a breakdown that sums to ``total``. + """ input: Optional[int] = None output: Optional[int] = None total: Optional[int] = None cost: Optional[float] = None + cache_read: Optional[int] = Field(default=None, alias="cacheRead") + cache_write: Optional[int] = Field(default=None, alias="cacheWrite") # --------------------------------------------------------------------------- diff --git a/sdks/python/oss/tests/pytest/unit/agents/golden/run_result.ok.json b/sdks/python/oss/tests/pytest/unit/agents/golden/run_result.ok.json index 0943d2d047..9b42f9788e 100644 --- a/sdks/python/oss/tests/pytest/unit/agents/golden/run_result.ok.json +++ b/sdks/python/oss/tests/pytest/unit/agents/golden/run_result.ok.json @@ -6,11 +6,11 @@ ], "events": [ {"type": "message", "text": "Hello!"}, - {"type": "usage", "input": 10, "output": 5, "total": 15, "cost": 0.001}, + {"type": "usage", "input": 10, "output": 5, "total": 135, "cost": 0.001, "cacheRead": 100, "cacheWrite": 20}, {"type": "done", "stopReason": "end_turn"}, {"text": "an event with no type, dropped on parse"} ], - "usage": {"input": 10, "output": 5, "total": 15, "cost": 0.001}, + "usage": {"input": 10, "output": 5, "total": 135, "cost": 0.001, "cacheRead": 100, "cacheWrite": 20}, "stopReason": "end_turn", "capabilities": { "textMessages": true, diff --git a/sdks/python/oss/tests/pytest/unit/agents/test_wire_contract.py b/sdks/python/oss/tests/pytest/unit/agents/test_wire_contract.py index 1bd083a4bb..462c6d534a 100644 --- a/sdks/python/oss/tests/pytest/unit/agents/test_wire_contract.py +++ b/sdks/python/oss/tests/pytest/unit/agents/test_wire_contract.py @@ -566,7 +566,16 @@ def test_result_from_wire_parses_ok(golden): # The event with no `type` is dropped on parse; the other three survive. assert [e.type for e in result.events] == ["message", "usage", "done"] assert result.events[0].data == {"type": "message", "text": "Hello!"} - assert result.usage == {"input": 10, "output": 5, "total": 15, "cost": 0.001} + # Cache fields ride alongside input/output so consumers can show a breakdown + # that sums to total (input counts non-cached prompt tokens only). + assert result.usage == { + "input": 10, + "output": 5, + "total": 135, + "cost": 0.001, + "cacheRead": 100, + "cacheWrite": 20, + } assert result.stop_reason == "end_turn" assert result.session_id == "sess-42" assert result.model == "gpt-5.5" diff --git a/services/runner/src/engines/sandbox_agent/usage.ts b/services/runner/src/engines/sandbox_agent/usage.ts index 1d08b3447b..a5dbcce034 100644 --- a/services/runner/src/engines/sandbox_agent/usage.ts +++ b/services/runner/src/engines/sandbox_agent/usage.ts @@ -33,10 +33,22 @@ export function mergePromptAndStreamUsage( const promptUsage = promptResult?.usage; const inputTokens = promptUsage?.inputTokens ?? streamUsage?.input ?? 0; const outputTokens = promptUsage?.outputTokens ?? streamUsage?.output ?? 0; - const total = inputTokens + outputTokens || streamUsage?.total || 0; + const cacheRead = promptUsage?.cachedReadTokens ?? streamUsage?.cacheRead; + const cacheWrite = promptUsage?.cachedWriteTokens ?? streamUsage?.cacheWrite; + const total = + inputTokens + outputTokens + (cacheRead ?? 0) + (cacheWrite ?? 0) || + streamUsage?.total || + 0; const cost = streamUsage?.cost ?? 0; return total > 0 || cost > 0 - ? { input: inputTokens, output: outputTokens, total, cost } + ? { + input: inputTokens, + output: outputTokens, + total, + cost, + ...(cacheRead != null ? { cacheRead } : {}), + ...(cacheWrite != null ? { cacheWrite } : {}), + } : undefined; } diff --git a/services/runner/src/protocol.ts b/services/runner/src/protocol.ts index 760ea7e58e..31ae95c552 100644 --- a/services/runner/src/protocol.ts +++ b/services/runner/src/protocol.ts @@ -367,6 +367,8 @@ export type AgentEvent = output?: number; total?: number; cost?: number; + cacheRead?: number; + cacheWrite?: number; } | { type: "error"; message: string } | { type: "done"; stopReason?: string }; @@ -374,12 +376,18 @@ export type AgentEvent = /** A live event sink the engines call as each event is built. */ export type EmitEvent = (event: AgentEvent) => void; -/** Run token/cost totals, rolled up onto the caller's workflow span. */ +/** Run token/cost totals, rolled up onto the caller's workflow span. + * `input` counts non-cached prompt tokens only (Pi semantics); the cached portion is + * reported separately so consumers can show a breakdown that sums to `total`. */ export interface AgentUsage { input: number; output: number; total: number; cost: number; + /** Prompt tokens served from the provider's cache. Absent when the harness doesn't report it. */ + cacheRead?: number; + /** Prompt tokens written to the provider's cache. Absent when the harness doesn't report it. */ + cacheWrite?: number; } export interface AgentRunRequest { diff --git a/services/runner/src/tracing/otel.ts b/services/runner/src/tracing/otel.ts index 921b128511..aad864f9ec 100644 --- a/services/runner/src/tracing/otel.ts +++ b/services/runner/src/tracing/otel.ts @@ -614,7 +614,14 @@ export interface AgentaOtel { /** Flush this run's trace to Agenta. Await before the process/response ends. */ flush: () => Promise; /** Run totals (tokens + cost) summed across turns, for roll-up onto the parent. */ - usage: () => { input: number; output: number; total: number; cost: number }; + usage: () => { + input: number; + output: number; + total: number; + cost: number; + cacheRead: number; + cacheWrite: number; + }; } /** @@ -654,16 +661,28 @@ export function createAgentaOtel( // returned so the caller can roll them up onto the workflow span in its own process // (the agent and workflow spans are exported in separate OTLP batches, so Agenta's // per-batch cumulative roll-up cannot bridge them on its own). - const runUsage = { input: 0, output: 0, total: 0, cost: 0 }; + const runUsage = { + input: 0, + output: 0, + total: 0, + cost: 0, + cacheRead: 0, + cacheWrite: 0, + }; function accumulateUsage(msg: any): void { const u = msg?.usage; if (!u) return; const input = u.input ?? 0; const output = u.output ?? 0; + const cacheRead = u.cacheRead ?? 0; + const cacheWrite = u.cacheWrite ?? 0; runUsage.input += input; runUsage.output += output; - runUsage.total += u.totalTokens ?? input + output; + runUsage.cacheRead += cacheRead; + runUsage.cacheWrite += cacheWrite; + // Pi's totalTokens includes cache reads/writes; mirror that in the fallback sum. + runUsage.total += u.totalTokens ?? input + output + cacheRead + cacheWrite; if (u.cost?.total != null) runUsage.cost += u.cost.total; } @@ -829,6 +848,16 @@ export function createAgentaOtel( runUsage.output, ); agentSpan.setAttribute("gen_ai.usage.total_tokens", runUsage.total); + if (runUsage.cacheRead > 0) + agentSpan.setAttribute( + "gen_ai.usage.cache_read.input_tokens", + runUsage.cacheRead, + ); + if (runUsage.cacheWrite > 0) + agentSpan.setAttribute( + "gen_ai.usage.cache_creation.input_tokens", + runUsage.cacheWrite, + ); if (runUsage.cost > 0) agentSpan.setAttribute("gen_ai.usage.cost", runUsage.cost); } @@ -1146,6 +1175,13 @@ export function createSandboxAgentOtel( span.setAttribute("gen_ai.usage.prompt_tokens", u.input); span.setAttribute("gen_ai.usage.completion_tokens", u.output); span.setAttribute("gen_ai.usage.total_tokens", u.total); + if (u.cacheRead != null && u.cacheRead > 0) + span.setAttribute("gen_ai.usage.cache_read.input_tokens", u.cacheRead); + if (u.cacheWrite != null && u.cacheWrite > 0) + span.setAttribute( + "gen_ai.usage.cache_creation.input_tokens", + u.cacheWrite, + ); if (u.cost > 0) span.setAttribute("gen_ai.usage.cost", u.cost); } @@ -1427,6 +1463,8 @@ export function createSandboxAgentOtel( output: usage?.output ?? 0, total: typeof total === "number" ? total : usage?.total ?? 0, cost: typeof cost === "number" ? cost : usage?.cost ?? 0, + ...(usage?.cacheRead != null ? { cacheRead: usage.cacheRead } : {}), + ...(usage?.cacheWrite != null ? { cacheWrite: usage.cacheWrite } : {}), }; record({ type: "usage", ...usage }); } diff --git a/services/runner/tests/unit/wire-contract.test.ts b/services/runner/tests/unit/wire-contract.test.ts index c9700b45e2..94ef363b76 100644 --- a/services/runner/tests/unit/wire-contract.test.ts +++ b/services/runner/tests/unit/wire-contract.test.ts @@ -262,11 +262,15 @@ describe("wire contract: results (vs Python golden)", () => { typed.map((e) => e.type), ["message", "usage", "done"], ); + // Cache fields ride alongside input/output so consumers can show a breakdown + // that sums to total (input counts non-cached prompt tokens only). assert.deepEqual(res.usage, { input: 10, output: 5, - total: 15, + total: 135, cost: 0.001, + cacheRead: 100, + cacheWrite: 20, }); assert.equal(res.stopReason, "end_turn"); assert.equal(res.sessionId, "sess-42"); diff --git a/web/oss/src/components/AgentChatSlice/assets/trace.ts b/web/oss/src/components/AgentChatSlice/assets/trace.ts index a8811f4573..c78e708172 100644 --- a/web/oss/src/components/AgentChatSlice/assets/trace.ts +++ b/web/oss/src/components/AgentChatSlice/assets/trace.ts @@ -51,13 +51,19 @@ export interface MessageUsageMetrics { completionTokens?: number totalTokens?: number totalCost?: number + cacheReadTokens?: number + cacheWriteTokens?: number } /** * Usage (tokens + cost) the service stamps onto `message.metadata.usage` via the - * `finish` part's messageMetadata (`{input, output, total, cost}`), mapped to the - * metrics-display field names. The trace supplies latency; this supplies tokens/cost - * (the agent-run trace summary doesn't surface them on the Pi/local path). + * `finish` part's messageMetadata (`{input, output, total, cost, cacheRead, cacheWrite}`), + * mapped to the metrics-display field names. The trace supplies latency; this supplies + * tokens/cost (the agent-run trace summary doesn't surface them on the Pi/local path). + * + * The wire `input` counts non-cached prompt tokens only (Pi semantics); the displayed + * prompt count folds the cached portion back in so Prompt + Completion = Total, with + * the cache split surfaced separately. */ export const getMessageUsage = (message: UIMessage): MessageUsageMetrics | undefined => { const usage = (message.metadata as {usage?: Record} | undefined)?.usage @@ -68,9 +74,13 @@ export const getMessageUsage = (message: UIMessage): MessageUsageMetrics | undef const output = num(usage.output) const total = num(usage.total) const cost = num(usage.cost) - if (input !== undefined) out.promptTokens = input + const cacheRead = num(usage.cacheRead) + const cacheWrite = num(usage.cacheWrite) + if (input !== undefined) out.promptTokens = input + (cacheRead ?? 0) + (cacheWrite ?? 0) if (output !== undefined) out.completionTokens = output if (total !== undefined) out.totalTokens = total if (cost !== undefined) out.totalCost = cost + if (cacheRead !== undefined) out.cacheReadTokens = cacheRead + if (cacheWrite !== undefined) out.cacheWriteTokens = cacheWrite return Object.keys(out).length > 0 ? out : undefined } diff --git a/web/packages/agenta-ui/src/components/presentational/metrics/ExecutionMetricsDisplay.tsx b/web/packages/agenta-ui/src/components/presentational/metrics/ExecutionMetricsDisplay.tsx index e176c1e27b..26f2a97799 100644 --- a/web/packages/agenta-ui/src/components/presentational/metrics/ExecutionMetricsDisplay.tsx +++ b/web/packages/agenta-ui/src/components/presentational/metrics/ExecutionMetricsDisplay.tsx @@ -44,6 +44,10 @@ export interface ExecutionMetricsData { promptTokens?: number /** Completion/output tokens */ completionTokens?: number + /** Prompt tokens served from the provider's cache (subset of promptTokens) */ + cacheReadTokens?: number + /** Prompt tokens written to the provider's cache (subset of promptTokens) */ + cacheWriteTokens?: number /** Total cost in dollars */ totalCost?: number } @@ -129,6 +133,20 @@ export const ExecutionMetricsDisplay = memo(function ExecutionMetricsDisplay({ Prompt {formatTokens(metrics.promptTokens)} + {metrics.cacheReadTokens !== undefined && + metrics.cacheReadTokens > 0 && ( +
+ Cached (read) + {formatTokens(metrics.cacheReadTokens)} +
+ )} + {metrics.cacheWriteTokens !== undefined && + metrics.cacheWriteTokens > 0 && ( +
+ Cached (write) + {formatTokens(metrics.cacheWriteTokens)} +
+ )}
Completion {formatTokens(metrics.completionTokens)} From 161a7f86d7c8cf67365e69f9e268e4392c8e4301 Mon Sep 17 00:00:00 2001 From: ashrafchowdury Date: Thu, 16 Jul 2026 19:24:48 +0600 Subject: [PATCH 2/3] feat: update agent usage tracking and metrics display for improved clarity and accuracy --- .../sdk/agents/adapters/vercel/stream.py | 13 +-- .../pytest/unit/agents/test_wire_contract.py | 3 +- services/runner/src/tracing/otel.ts | 86 +++++++------------ .../runner/tests/unit/wire-contract.test.ts | 3 +- .../components/AgentChatSlice/assets/trace.ts | 5 +- .../metrics/ExecutionMetricsDisplay.tsx | 31 ++++--- 6 files changed, 61 insertions(+), 80 deletions(-) diff --git a/sdks/python/agenta/sdk/agents/adapters/vercel/stream.py b/sdks/python/agenta/sdk/agents/adapters/vercel/stream.py index cc50eec87a..3ce5fef740 100644 --- a/sdks/python/agenta/sdk/agents/adapters/vercel/stream.py +++ b/sdks/python/agenta/sdk/agents/adapters/vercel/stream.py @@ -11,6 +11,7 @@ from ...dtos import AgentResult from ...streaming import AgentStream from ...utils.wire import sanitize_runner_error +from ...wire_models import WireAgentUsage from .messages import TOOL_APPROVAL_REQUEST log = get_module_logger(__name__) @@ -778,12 +779,14 @@ def _approval_tool_call_id(payload: Dict[str, Any]) -> Optional[Any]: return None +# The wire model is the single source of truth for the usage field set (camelCase on the wire). +_USAGE_KEYS = tuple( + field.alias or name for name, field in WireAgentUsage.model_fields.items() +) + + def _usage_metadata(data: Dict[str, Any]) -> Dict[str, Any]: - return { - key: data[key] - for key in ("input", "output", "total", "cost", "cacheRead", "cacheWrite") - if data.get(key) is not None - } + return {key: data[key] for key in _USAGE_KEYS if data.get(key) is not None} def _committed_revision_data(tool_name: Any, output: Any) -> Optional[Dict[str, Any]]: diff --git a/sdks/python/oss/tests/pytest/unit/agents/test_wire_contract.py b/sdks/python/oss/tests/pytest/unit/agents/test_wire_contract.py index 462c6d534a..a5de7def4a 100644 --- a/sdks/python/oss/tests/pytest/unit/agents/test_wire_contract.py +++ b/sdks/python/oss/tests/pytest/unit/agents/test_wire_contract.py @@ -566,8 +566,7 @@ def test_result_from_wire_parses_ok(golden): # The event with no `type` is dropped on parse; the other three survive. assert [e.type for e in result.events] == ["message", "usage", "done"] assert result.events[0].data == {"type": "message", "text": "Hello!"} - # Cache fields ride alongside input/output so consumers can show a breakdown - # that sums to total (input counts non-cached prompt tokens only). + # input excludes cached tokens — see WireAgentUsage in wire_models.py. assert result.usage == { "input": 10, "output": 5, diff --git a/services/runner/src/tracing/otel.ts b/services/runner/src/tracing/otel.ts index aad864f9ec..ed40381ad8 100644 --- a/services/runner/src/tracing/otel.ts +++ b/services/runner/src/tracing/otel.ts @@ -172,12 +172,16 @@ function redactSpan(span: ReadableSpan, redactors: Iterable): void { } const status = span.status as { message?: string }; if (typeof status.message === "string") { - status.message = redactor.redactString(status.message, "spans") ?? status.message; + status.message = + redactor.redactString(status.message, "spans") ?? status.message; } } } -function redactAttributes(attrs: Record, redactor: Redactor): void { +function redactAttributes( + attrs: Record, + redactor: Redactor, +): void { for (const [key, value] of Object.entries(attrs)) { if (typeof value === "string") { attrs[key] = redactor.redactString(value, "spans"); @@ -296,7 +300,8 @@ class TraceBatchProcessor implements SpanProcessor { // Fall back to the env default only for a span whose OWN run's target is unknown // (untagged span, or the run already released) — never to another run's target, or a // batch could still land on an unintended endpoint/auth. - const target = (runId ? byRun?.get(runId) : undefined) ?? defaultTarget(); + const target = + (runId ? byRun?.get(runId) : undefined) ?? defaultTarget(); return new Promise((resolve) => { try { getExporter(target).export(orderParentFirst(group), (result) => { @@ -602,6 +607,24 @@ function applyAssistant( return undefined; } +/** Stamp run-total usage (tokens, cache split, cost) as gen_ai.usage.* span attributes. */ +function stampUsage(span: Span, u: AgentUsage | undefined): void { + if (!u) return; + span.setAttribute("gen_ai.usage.input_tokens", u.input); + span.setAttribute("gen_ai.usage.output_tokens", u.output); + span.setAttribute("gen_ai.usage.prompt_tokens", u.input); + span.setAttribute("gen_ai.usage.completion_tokens", u.output); + span.setAttribute("gen_ai.usage.total_tokens", u.total); + if ((u.cacheRead ?? 0) > 0) + span.setAttribute("gen_ai.usage.cache_read.input_tokens", u.cacheRead!); + if ((u.cacheWrite ?? 0) > 0) + span.setAttribute( + "gen_ai.usage.cache_creation.input_tokens", + u.cacheWrite!, + ); + if (u.cost > 0) span.setAttribute("gen_ai.usage.cost", u.cost); +} + // --------------------------------------------------------------------------- // Extension factory (one per run; state is closure-scoped) // --------------------------------------------------------------------------- @@ -614,14 +637,7 @@ export interface AgentaOtel { /** Flush this run's trace to Agenta. Await before the process/response ends. */ flush: () => Promise; /** Run totals (tokens + cost) summed across turns, for roll-up onto the parent. */ - usage: () => { - input: number; - output: number; - total: number; - cost: number; - cacheRead: number; - cacheWrite: number; - }; + usage: () => Required; } /** @@ -839,28 +855,7 @@ export function createAgentaOtel( ); // Stamp the run total on the agent span so it shows the agent's tokens/cost even // though Agenta cannot roll the per-turn LLM spans up across batches. - if (runUsage.total > 0) { - agentSpan.setAttribute("gen_ai.usage.input_tokens", runUsage.input); - agentSpan.setAttribute("gen_ai.usage.output_tokens", runUsage.output); - agentSpan.setAttribute("gen_ai.usage.prompt_tokens", runUsage.input); - agentSpan.setAttribute( - "gen_ai.usage.completion_tokens", - runUsage.output, - ); - agentSpan.setAttribute("gen_ai.usage.total_tokens", runUsage.total); - if (runUsage.cacheRead > 0) - agentSpan.setAttribute( - "gen_ai.usage.cache_read.input_tokens", - runUsage.cacheRead, - ); - if (runUsage.cacheWrite > 0) - agentSpan.setAttribute( - "gen_ai.usage.cache_creation.input_tokens", - runUsage.cacheWrite, - ); - if (runUsage.cost > 0) - agentSpan.setAttribute("gen_ai.usage.cost", runUsage.cost); - } + if (runUsage.total > 0) stampUsage(agentSpan, runUsage); agentSpan.end(); agentSpan = undefined; agentCtx = undefined; @@ -1168,23 +1163,6 @@ export function createSandboxAgentOtel( } } - function stampUsage(span: Span, u: AgentUsage | undefined): void { - if (!u) return; - span.setAttribute("gen_ai.usage.input_tokens", u.input); - span.setAttribute("gen_ai.usage.output_tokens", u.output); - span.setAttribute("gen_ai.usage.prompt_tokens", u.input); - span.setAttribute("gen_ai.usage.completion_tokens", u.output); - span.setAttribute("gen_ai.usage.total_tokens", u.total); - if (u.cacheRead != null && u.cacheRead > 0) - span.setAttribute("gen_ai.usage.cache_read.input_tokens", u.cacheRead); - if (u.cacheWrite != null && u.cacheWrite > 0) - span.setAttribute( - "gen_ai.usage.cache_creation.input_tokens", - u.cacheWrite, - ); - if (u.cost > 0) span.setAttribute("gen_ai.usage.cost", u.cost); - } - function setUsage(finalUsage: AgentUsage | undefined): void { if (!finalUsage) return; usage = finalUsage; @@ -1458,13 +1436,13 @@ export function createSandboxAgentOtel( // which the sandbox-agent engine reads. Keep total + cost here and leave the split to the caller. const cost = update.cost?.amount; const total = update.used; + // Spread first so optional fields (cache split) carry forward untouched. usage = { + ...usage, input: usage?.input ?? 0, output: usage?.output ?? 0, - total: typeof total === "number" ? total : usage?.total ?? 0, - cost: typeof cost === "number" ? cost : usage?.cost ?? 0, - ...(usage?.cacheRead != null ? { cacheRead: usage.cacheRead } : {}), - ...(usage?.cacheWrite != null ? { cacheWrite: usage.cacheWrite } : {}), + total: typeof total === "number" ? total : (usage?.total ?? 0), + cost: typeof cost === "number" ? cost : (usage?.cost ?? 0), }; record({ type: "usage", ...usage }); } diff --git a/services/runner/tests/unit/wire-contract.test.ts b/services/runner/tests/unit/wire-contract.test.ts index 94ef363b76..a406dedc83 100644 --- a/services/runner/tests/unit/wire-contract.test.ts +++ b/services/runner/tests/unit/wire-contract.test.ts @@ -262,8 +262,7 @@ describe("wire contract: results (vs Python golden)", () => { typed.map((e) => e.type), ["message", "usage", "done"], ); - // Cache fields ride alongside input/output so consumers can show a breakdown - // that sums to total (input counts non-cached prompt tokens only). + // input excludes cached tokens — see AgentUsage in protocol.ts. assert.deepEqual(res.usage, { input: 10, output: 5, diff --git a/web/oss/src/components/AgentChatSlice/assets/trace.ts b/web/oss/src/components/AgentChatSlice/assets/trace.ts index c78e708172..11c3355c45 100644 --- a/web/oss/src/components/AgentChatSlice/assets/trace.ts +++ b/web/oss/src/components/AgentChatSlice/assets/trace.ts @@ -61,9 +61,8 @@ export interface MessageUsageMetrics { * mapped to the metrics-display field names. The trace supplies latency; this supplies * tokens/cost (the agent-run trace summary doesn't surface them on the Pi/local path). * - * The wire `input` counts non-cached prompt tokens only (Pi semantics); the displayed - * prompt count folds the cached portion back in so Prompt + Completion = Total, with - * the cache split surfaced separately. + * The wire `input` excludes cached tokens, so the displayed prompt count folds the + * cached portion back in (Prompt + Completion = Total), with the split kept separately. */ export const getMessageUsage = (message: UIMessage): MessageUsageMetrics | undefined => { const usage = (message.metadata as {usage?: Record} | undefined)?.usage diff --git a/web/packages/agenta-ui/src/components/presentational/metrics/ExecutionMetricsDisplay.tsx b/web/packages/agenta-ui/src/components/presentational/metrics/ExecutionMetricsDisplay.tsx index 26f2a97799..45fbd146b7 100644 --- a/web/packages/agenta-ui/src/components/presentational/metrics/ExecutionMetricsDisplay.tsx +++ b/web/packages/agenta-ui/src/components/presentational/metrics/ExecutionMetricsDisplay.tsx @@ -133,20 +133,23 @@ export const ExecutionMetricsDisplay = memo(function ExecutionMetricsDisplay({ Prompt {formatTokens(metrics.promptTokens)}
- {metrics.cacheReadTokens !== undefined && - metrics.cacheReadTokens > 0 && ( -
- Cached (read) - {formatTokens(metrics.cacheReadTokens)} -
- )} - {metrics.cacheWriteTokens !== undefined && - metrics.cacheWriteTokens > 0 && ( -
- Cached (write) - {formatTokens(metrics.cacheWriteTokens)} -
- )} + {( + [ + ["Cached (read)", metrics.cacheReadTokens], + ["Cached (write)", metrics.cacheWriteTokens], + ] as const + ).map( + ([label, value]) => + !!value && ( +
+ {label} + {formatTokens(value)} +
+ ), + )}
Completion {formatTokens(metrics.completionTokens)} From c2b6c4ede102a301d0561204bd91bfc8a5a4c5f9 Mon Sep 17 00:00:00 2001 From: ashrafchowdury Date: Thu, 16 Jul 2026 19:26:45 +0600 Subject: [PATCH 3/3] feat: enhance token metrics tracking by adding cache read and write metrics --- api/oss/src/core/tracing/utils/trees.py | 97 ++++--------------- .../pytest/unit/tracing/utils/test_trees.py | 36 +++++++ .../TraceSidePanel/TraceDetails/index.tsx | 25 +++++ .../newObservability/selectors/tracing.ts | 31 +++++- 4 files changed, 112 insertions(+), 77 deletions(-) diff --git a/api/oss/src/core/tracing/utils/trees.py b/api/oss/src/core/tracing/utils/trees.py index eca466d41b..f115387baa 100644 --- a/api/oss/src/core/tracing/utils/trees.py +++ b/api/oss/src/core/tracing/utils/trees.py @@ -318,98 +318,43 @@ def _set_cumulative(span: OTelFlatSpan, costs: dict): ) +TOKEN_KEYS = ("prompt", "completion", "total", "cache_read", "cache_creation") + + def cumulate_tokens( spans_id_tree: OrderedDict, spans_idx: Dict[str, OTelFlatSpan], ) -> None: - def _get_incremental(span: OTelFlatSpan): - _tokens = { - "prompt": 0.0, - "completion": 0.0, - "total": 0.0, - } + def _get(span: OTelFlatSpan, scope: str): + attr: dict = span.attributes or {} - if span.attributes is None: - return _tokens + tokens = attr.get("ag", {}).get("metrics", {}).get("tokens", {}).get(scope, {}) + if not isinstance(tokens, dict): + tokens = {} - attr: dict = span.attributes + return {key: tokens.get(key, 0.0) for key in TOKEN_KEYS} - return { - "prompt": ( - attr.get("ag", {}) - .get("metrics", {}) - .get("tokens", {}) - .get("incremental", {}) - .get("prompt", 0.0) - ), - "completion": ( - attr.get("ag", {}) - .get("metrics", {}) - .get("tokens", {}) - .get("incremental", {}) - .get("completion", 0.0) - ), - "total": ( - attr.get("ag", {}) - .get("metrics", {}) - .get("tokens", {}) - .get("incremental", {}) - .get("total", 0.0) - ), - } + def _get_incremental(span: OTelFlatSpan): + return _get(span, "incremental") def _get_cumulative(span: OTelFlatSpan): - _tokens = { - "prompt": 0.0, - "completion": 0.0, - "total": 0.0, - } - - if span.attributes is None: - return _tokens - - attr: dict = span.attributes - - return { - "prompt": ( - attr.get("ag", {}) - .get("metrics", {}) - .get("tokens", {}) - .get("cumulative", {}) - .get("prompt", 0.0) - ), - "completion": ( - attr.get("ag", {}) - .get("metrics", {}) - .get("tokens", {}) - .get("cumulative", {}) - .get("completion", 0.0) - ), - "total": ( - attr.get("ag", {}) - .get("metrics", {}) - .get("tokens", {}) - .get("cumulative", {}) - .get("total", 0.0) - ), - } + return _get(span, "cumulative") def _accumulate(a: dict, b: dict): - return { - "prompt": a.get("prompt", 0.0) + b.get("prompt", 0.0), - "completion": a.get("completion", 0.0) + b.get("completion", 0.0), - "total": a.get("total", 0.0) + b.get("total", 0.0), - } + return {key: a.get(key, 0.0) + b.get(key, 0.0) for key in TOKEN_KEYS} def _set_cumulative(span: OTelFlatSpan, tokens: dict): if span.attributes is None: span.attributes = {} - if ( - tokens.get("prompt", 0.0) != 0.0 - or tokens.get("completion", 0.0) != 0.0 - or tokens.get("total", 0.0) != 0.0 - ): + # Keep the stored shape stable for non-cache traces: cache keys only when non-zero. + tokens = { + key: value + for key, value in tokens.items() + if value != 0.0 or key in ("prompt", "completion", "total") + } + + if any(value != 0.0 for value in tokens.values()): if "ag" not in span.attributes or not isinstance( span.attributes["ag"], dict, diff --git a/api/oss/tests/pytest/unit/tracing/utils/test_trees.py b/api/oss/tests/pytest/unit/tracing/utils/test_trees.py index 5414693054..e72104232f 100644 --- a/api/oss/tests/pytest/unit/tracing/utils/test_trees.py +++ b/api/oss/tests/pytest/unit/tracing/utils/test_trees.py @@ -141,6 +141,42 @@ def test_cumulate_tokens_and_costs_propagate_from_children_to_parent(): assert root_costs["total"] == pytest.approx(1.2) +def test_cumulate_tokens_propagates_cache_split_when_present(): + root = _span( + span_id=ROOT_UUID, + span_name="root", + prompt_tokens=1, + completion_tokens=2, + ) + child = _span( + span_id=CHILD_A_UUID, + parent_id=ROOT_UUID, + span_name="child", + prompt_tokens=4, + completion_tokens=5, + start_offset_s=1, + ) + child.attributes["ag"]["metrics"]["tokens"]["incremental"].update( + {"cache_read": 100.0, "cache_creation": 20.0} + ) + + span_idx = parse_span_dtos_to_span_idx([root, child]) + tree = parse_span_idx_to_span_id_tree(span_idx) + + cumulate_tokens(tree, span_idx) + + root_tokens = span_idx[ROOT_UUID].attributes["ag"]["metrics"]["tokens"][ + "cumulative" + ] + assert root_tokens == { + "prompt": 5.0, + "completion": 7.0, + "total": 12.0, + "cache_read": 100.0, + "cache_creation": 20.0, + } + + def test_cumulate_errors_propagates_scalar_counts_from_children_to_parent(): root = _span( span_id=ROOT_UUID, diff --git a/web/oss/src/components/SharedDrawers/TraceDrawer/components/TraceSidePanel/TraceDetails/index.tsx b/web/oss/src/components/SharedDrawers/TraceDrawer/components/TraceSidePanel/TraceDetails/index.tsx index ce1201b169..00dd7fc1ec 100644 --- a/web/oss/src/components/SharedDrawers/TraceDrawer/components/TraceSidePanel/TraceDetails/index.tsx +++ b/web/oss/src/components/SharedDrawers/TraceDrawer/components/TraceSidePanel/TraceDetails/index.tsx @@ -7,6 +7,8 @@ import StatusRenderer from "@/oss/components/pages/observability/components/Stat import ResultTag from "@/oss/components/ResultTag/ResultTag" import {TraceSpanNode} from "@/oss/services/tracing/types" import { + formattedSpanCacheReadTokensAtomFamily, + formattedSpanCacheWriteTokensAtomFamily, formattedSpanCompletionTokensAtomFamily, formattedSpanCostAtomFamily, formattedSpanLatencyAtomFamily, @@ -30,6 +32,12 @@ const TraceDetails = ({activeTrace}: {activeTrace: TraceSpanNode}) => { const formattedCompletionTokens = useAtomValue( formattedSpanCompletionTokensAtomFamily(activeTrace), ) + const formattedCacheReadTokens = useAtomValue( + formattedSpanCacheReadTokensAtomFamily(activeTrace), + ) + const formattedCacheWriteTokens = useAtomValue( + formattedSpanCacheWriteTokensAtomFamily(activeTrace), + ) const traceStartTime = useAtomValue(spanStartTimeAtomFamily(activeTrace)) const traceEndTime = useAtomValue(spanEndTimeAtomFamily(activeTrace)) return ( @@ -110,6 +118,23 @@ const TraceDetails = ({activeTrace}: {activeTrace: TraceSpanNode}) => {
{formattedPromptTokens}
Prompt tokens
+ {( + [ + ["Cached (read)", formattedCacheReadTokens], + ["Cached (write)", formattedCacheWriteTokens], + ] as const + ).map( + ([label, value]) => + value && ( + +
{value}
+
{label}
+ + ), + )}
{formattedCompletionTokens}
Completion tokens
diff --git a/web/oss/src/state/newObservability/selectors/tracing.ts b/web/oss/src/state/newObservability/selectors/tracing.ts index 0596e498e0..b9d7fc9f72 100644 --- a/web/oss/src/state/newObservability/selectors/tracing.ts +++ b/web/oss/src/state/newObservability/selectors/tracing.ts @@ -14,9 +14,23 @@ export const getTokens = (span?: TraceSpanNode) => { return tokens?.cumulative?.total ?? tokens?.incremental?.total ?? null } +export const getCacheReadTokens = (span?: TraceSpanNode) => { + const tokens = getTokenMetrics(span) + return tokens?.cumulative?.cache_read ?? tokens?.incremental?.cache_read ?? null +} + +export const getCacheWriteTokens = (span?: TraceSpanNode) => { + const tokens = getTokenMetrics(span) + return tokens?.cumulative?.cache_creation ?? tokens?.incremental?.cache_creation ?? null +} + export const getPromptTokens = (span?: TraceSpanNode) => { const tokens = getTokenMetrics(span) - return tokens?.cumulative?.prompt ?? tokens?.incremental?.prompt ?? null + const prompt = tokens?.cumulative?.prompt ?? tokens?.incremental?.prompt ?? null + if (prompt === null) return null + // `prompt` excludes cached tokens (cache_read/cache_creation are reported separately), + // so fold them back in — displayed Prompt + Completion = Total, cache shown as sub-rows. + return prompt + (getCacheReadTokens(span) ?? 0) + (getCacheWriteTokens(span) ?? 0) } export const getCompletionTokens = (span?: TraceSpanNode) => { @@ -152,6 +166,21 @@ export const formattedSpanCompletionTokensAtomFamily = atomFamily((span?: TraceS atom(() => formatTokenUsage(getCompletionTokens(span))), ) +// Null when the span has no cache metrics, so the drawer can omit the rows entirely. +export const formattedSpanCacheReadTokensAtomFamily = atomFamily((span?: TraceSpanNode) => + atom(() => { + const tokens = getCacheReadTokens(span) + return tokens ? formatTokenUsage(tokens) : null + }), +) + +export const formattedSpanCacheWriteTokensAtomFamily = atomFamily((span?: TraceSpanNode) => + atom(() => { + const tokens = getCacheWriteTokens(span) + return tokens ? formatTokenUsage(tokens) : null + }), +) + export const formattedSpanCostAtomFamily = atomFamily((span?: TraceSpanNode) => atom(() => formatCurrency(getCost(span))), )