From 0687c0bc94407786f09f90343c5a6da700ededf0 Mon Sep 17 00:00:00 2001 From: Aditya kumar singh <143548997+Adityakk9031@users.noreply.github.com> Date: Thu, 17 Sep 2026 18:06:35 +0530 Subject: [PATCH 1/2] feat(usage): aggregate token usage and expose in CLI view and across harnesses - Implement Default, total_tokens(), is_zero(), Add, AddAssign, and Sum for Usage in common.rs - Implement Transcript::total_usage() in transcript.rs - Preserve token usage in fx harness (events, session.json, and usage-v2.json) - Support expanded token usage variants (promptTokens, completionTokens, cache tokens) in Cursor Desktop - Parse and serialize token usage in Hermes harness message rows - Display session token accounting and assistant turn token counts in CLI view - Document token usage display in docs/usage.md - Add comprehensive integration tests in tests/integration/usage.rs --- cli/src/view.rs | 40 +++++- docs/usage.md | 2 +- src/common.rs | 95 +++++++++++++- src/harness/cursor_desktop.rs | 31 ++++- src/harness/fx.rs | 238 ++++++++++++++++++++++++++++------ src/harness/hermes.rs | 73 ++++++++++- src/transcript.rs | 19 ++- tests/integration/main.rs | 1 + tests/integration/usage.rs | 223 +++++++++++++++++++++++++++++++ 9 files changed, 664 insertions(+), 58 deletions(-) create mode 100644 tests/integration/usage.rs diff --git a/cli/src/view.rs b/cli/src/view.rs index 9055772..2da1bdf 100644 --- a/cli/src/view.rs +++ b/cli/src/view.rs @@ -515,6 +515,30 @@ fn human_header(out: &mut String, common: &Transcript, span: &Span, colo &format!("{shown} of {}", common.body.len()), color, ); + if let Some(usage) = common.total_usage() { + if !usage.is_zero() { + let mut parts = vec![ + format!("{} in", usage.input_tokens), + format!("{} out", usage.output_tokens), + ]; + if let Some(cached) = usage.cache_read_input_tokens { + if cached > 0 { + parts.push(format!("{cached} cached")); + } + } + if let Some(created) = usage.cache_creation_input_tokens { + if created > 0 { + parts.push(format!("{created} cache write")); + } + } + human_field( + out, + "Tokens", + &format!("{} ({})", usage.total_tokens(), parts.join(", ")), + color, + ); + } + } } /// Render `messages` under `filters`, returning the line index of each @@ -548,15 +572,23 @@ fn human_messages( continue; } let ordinal = start + offset + 1; - let role = match message.role { - Role::User => "User", - Role::Assistant => "Assistant", + let role_label = match (message.role, &message.usage) { + (Role::Assistant, Some(u)) if !u.is_zero() => { + format!("Assistant · {} tokens", u.total_tokens()) + } + (Role::Assistant, _) => "Assistant".to_string(), + (Role::User, _) => "User".to_string(), }; lines += out[counted..].bytes().filter(|byte| *byte == b'\n').count(); counted = out.len(); // The rule follows the blank line `human_rule` opens with. message_starts.push(lines + 1); - human_rule(out, &format!("Message #{ordinal} · {role}"), width, color); + human_rule( + out, + &format!("Message #{ordinal} · {role_label}"), + width, + color, + ); for (index, block) in message.content.iter().enumerate() { if filters.shows_block(block) { blocks.render(out, (offset, index), block); diff --git a/docs/usage.md b/docs/usage.md index 8d52555..42a0868 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -39,7 +39,7 @@ A session id is any unambiguous prefix of the full id, or the session's exact ti - A [Simple](formats/simple.md) document instead of an id — `txcript continue ./run.json --with claude_code`, or `my-agent | txcript continue - --with claude_code` — brings any agent's transcript in the same way; `--with` is required since a document has no harness of its own. - The launch command is per-harness and overridable: set `TRANSCRIPT__RESUME_CMD` to a `{id}` template, e.g. `TRANSCRIPT_CODEX_RESUME_CMD="codex resume {id}"`. -`view` in a terminal opens a built-in pager: `u`, `a`, `t`, and `r` hide or show user messages, assistant messages, tool calls, and reasoning; `]` and `[` jump between messages; `/` searches what is shown. Images are drawn inline on terminals that can show them (Ghostty, kitty, WezTerm, Konsole). Set `TXCRIPT_PAGER` to use an external pager instead, or pass `--no-pager` to print the view directly. Piped or redirected, `view` prints the same compact text the MCP server serves. Either way each message is numbered by a `── #N ──` rule, and `#range` selects messages by those printed ordinals, 1-based and inclusive: +`view` in a terminal opens a built-in pager: `u`, `a`, `t`, and `r` hide or show user messages, assistant messages, tool calls, and reasoning; `]` and `[` jump between messages; `/` searches what is shown. Images are drawn inline on terminals that can show them (Ghostty, kitty, WezTerm, Konsole). Session metadata in the header includes total token accounting (`Tokens: ( in, out, cached)`), and assistant turn rules display individual turn token counts (`Message #N · Assistant · tokens`) when reported by the source harness. Set `TXCRIPT_PAGER` to use an external pager instead, or pass `--no-pager` to print the view directly. Piped or redirected, `view` prints the same compact text the MCP server serves. Either way each message is numbered by a `── #N ──` rule, and `#range` selects messages by those printed ordinals, 1-based and inclusive: - `abc#7`: message 7 only - `abc#5-12`: messages 5 through 12 diff --git a/src/common.rs b/src/common.rs index 8574209..16eea72 100644 --- a/src/common.rs +++ b/src/common.rs @@ -125,7 +125,7 @@ pub enum StopReason { } /// Token accounting for one assistant turn. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)] pub struct Usage { pub input_tokens: u64, pub output_tokens: u64, @@ -135,6 +135,63 @@ pub struct Usage { pub cache_creation_input_tokens: Option, } +impl Usage { + /// Compute total tokens represented by this usage record, including cached and creation tokens. + #[must_use] + pub fn total_tokens(&self) -> u64 { + self.input_tokens + .saturating_add(self.output_tokens) + .saturating_add(self.cache_read_input_tokens.unwrap_or(0)) + .saturating_add(self.cache_creation_input_tokens.unwrap_or(0)) + } + + /// Whether all token counts in this record are zero or unset. + #[must_use] + pub fn is_zero(&self) -> bool { + self.input_tokens == 0 + && self.output_tokens == 0 + && self.cache_read_input_tokens.unwrap_or(0) == 0 + && self.cache_creation_input_tokens.unwrap_or(0) == 0 + } +} + +impl std::ops::Add for Usage { + type Output = Self; + + fn add(self, rhs: Self) -> Self::Output { + let combine_opt = |a: Option, b: Option| match (a, b) { + (Some(x), Some(y)) => Some(x.saturating_add(y)), + (Some(x), None) | (None, Some(x)) => Some(x), + (None, None) => None, + }; + + Self { + input_tokens: self.input_tokens.saturating_add(rhs.input_tokens), + output_tokens: self.output_tokens.saturating_add(rhs.output_tokens), + cache_read_input_tokens: combine_opt( + self.cache_read_input_tokens, + rhs.cache_read_input_tokens, + ), + cache_creation_input_tokens: combine_opt( + self.cache_creation_input_tokens, + rhs.cache_creation_input_tokens, + ), + } + } +} + +impl std::ops::AddAssign for Usage { + fn add_assign(&mut self, rhs: Self) { + *self = *self + rhs; + } +} + +impl std::iter::Sum for Usage { + fn sum>(iter: I) -> Self { + iter.fold(Self::default(), |acc, u| acc + u) + } +} + /// A base64-encoded inline image. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct ImageSource { @@ -574,6 +631,42 @@ mod tests { assert!(matches!(tool, Tool::Bash { .. })); assert_eq!(tool.to_canonical().1, input); } + + #[test] + fn usage_arithmetic_and_aggregation() { + let u1 = Usage { + input_tokens: 100, + output_tokens: 50, + cache_read_input_tokens: Some(25), + cache_creation_input_tokens: None, + }; + let u2 = Usage { + input_tokens: 200, + output_tokens: 150, + cache_read_input_tokens: Some(30), + cache_creation_input_tokens: Some(10), + }; + + assert_eq!(u1.total_tokens(), 175); + assert_eq!(u2.total_tokens(), 390); + assert!(!u1.is_zero()); + assert!(Usage::default().is_zero()); + + let sum = u1 + u2; + assert_eq!(sum.input_tokens, 300); + assert_eq!(sum.output_tokens, 200); + assert_eq!(sum.cache_read_input_tokens, Some(55)); + assert_eq!(sum.cache_creation_input_tokens, Some(10)); + assert_eq!(sum.total_tokens(), 565); + + let mut acc = u1; + acc += u2; + assert_eq!(acc, sum); + + let list = vec![u1, u2, Usage::default()]; + let iter_sum: Usage = list.into_iter().sum(); + assert_eq!(iter_sum, sum); + } } /// One find/replace within a [`Tool::MultiEdit`]. diff --git a/src/harness/cursor_desktop.rs b/src/harness/cursor_desktop.rs index 895c851..f5a647a 100644 --- a/src/harness/cursor_desktop.rs +++ b/src/harness/cursor_desktop.rs @@ -332,19 +332,36 @@ fn flush(messages: &mut Vec, assistant: &mut Option) { fn bubble_usage(bubble: &Value) -> Option { let input = bubble .pointer("/tokenCount/inputTokens") + .or_else(|| bubble.pointer("/tokenCount/promptTokens")) + .or_else(|| bubble.pointer("/tokenUsage/promptTokens")) + .or_else(|| bubble.pointer("/tokenUsage/inputTokens")) .and_then(Value::as_u64) .unwrap_or(0); let output = bubble .pointer("/tokenCount/outputTokens") + .or_else(|| bubble.pointer("/tokenCount/completionTokens")) + .or_else(|| bubble.pointer("/tokenUsage/completionTokens")) + .or_else(|| bubble.pointer("/tokenUsage/outputTokens")) .and_then(Value::as_u64) .unwrap_or(0); - // All-zero counts are the serializer default, not an observation. - (input > 0 || output > 0).then_some(Usage { - input_tokens: input, - output_tokens: output, - cache_read_input_tokens: None, - cache_creation_input_tokens: None, - }) + let cache_read = bubble + .pointer("/tokenCount/cacheReadTokens") + .or_else(|| bubble.pointer("/tokenCount/cachedTokens")) + .or_else(|| bubble.pointer("/tokenUsage/cacheReadTokens")) + .or_else(|| bubble.pointer("/tokenUsage/cachedTokens")) + .and_then(Value::as_u64); + let cache_write = bubble + .pointer("/tokenCount/cacheCreationTokens") + .or_else(|| bubble.pointer("/tokenUsage/cacheCreationTokens")) + .and_then(Value::as_u64); + + (input > 0 || output > 0 || cache_read.unwrap_or(0) > 0 || cache_write.unwrap_or(0) > 0) + .then_some(Usage { + input_tokens: input, + output_tokens: output, + cache_read_input_tokens: cache_read, + cache_creation_input_tokens: cache_write, + }) } /// Emit the `ToolUse` on the assistant message and its paired result on the diff --git a/src/harness/fx.rs b/src/harness/fx.rs index e0401fa..260bd1d 100644 --- a/src/harness/fx.rs +++ b/src/harness/fx.rs @@ -33,12 +33,12 @@ //! `txcript-meta.json` sidecar that fx ignores: same-harness round trips keep //! reasoning, while fx resume simply renders the conversation without it. //! -//! Known representational losses through `Common`: per-turn token usage and -//! per-message model (fx stores one session model), `replace_all` on edits -//! (fx's `edit_file` has no such flag), `terminal` action/profile and other -//! non-Bash tool argument shapes when a session leaves fx, and the recovery -//! checkpoint state. `execution.files` (fx's changed-file panel) is emitted -//! empty; the conversation itself is intact. +//! Known representational losses through `Common`: per-message model (fx +//! stores one session model), `replace_all` on edits (fx's `edit_file` has no +//! such flag), `terminal` action/profile and other non-Bash tool argument +//! shapes when a session leaves fx, and the recovery checkpoint state. +//! `execution.files` (fx's changed-file panel) is emitted empty; the +//! conversation itself is intact. use std::collections::HashMap; use std::fs; @@ -50,7 +50,7 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value, json}; use uuid::Uuid; -use crate::common::{Block, ImageSource, Message, Meta, Role, StopReason, Tool, ToolOutput}; +use crate::common::{Block, ImageSource, Message, Meta, Role, StopReason, Tool, ToolOutput, Usage}; use crate::error::{Error, Result}; use crate::harness::jsonl; use crate::transcript::{Codec, Common, Discovered, Harness, Saved, Store, TextCodec, Transcript}; @@ -148,6 +148,7 @@ fn body_to_messages(body: &FxSession, fallback_ts: DateTime) -> Vec) -> Vec) -> Option { + let payload = payload?; + let input = payload + .get("total_input_tokens") + .or_else(|| payload.get("input_tokens")) + .and_then(Value::as_u64) + .unwrap_or(0); + let output = payload + .get("total_output_tokens") + .or_else(|| payload.get("output_tokens")) + .and_then(Value::as_u64) + .unwrap_or(0); + let cache_read = payload + .get("total_cache_read_tokens") + .or_else(|| payload.get("cache_read_tokens")) + .and_then(Value::as_u64); + let cache_write = payload + .get("total_cache_write_tokens") + .or_else(|| payload.get("cache_write_tokens")) + .and_then(Value::as_u64); + + (input > 0 || output > 0 || cache_read.unwrap_or(0) > 0 || cache_write.unwrap_or(0) > 0) + .then_some(Usage { + input_tokens: input, + output_tokens: output, + cache_read_input_tokens: cache_read, + cache_creation_input_tokens: cache_write, + }) +} + +fn parse_session_usage(body: &FxSession) -> Option { + if let Some(usage_val) = &body.usage { + let input = usage_val + .get("input_tokens") + .and_then(Value::as_u64) + .unwrap_or(0); + let output = usage_val + .get("output_tokens") + .and_then(Value::as_u64) + .unwrap_or(0); + let cache_read = usage_val.get("cache_read_tokens").and_then(Value::as_u64); + let cache_write = usage_val.get("cache_write_tokens").and_then(Value::as_u64); + if input > 0 || output > 0 || cache_read.unwrap_or(0) > 0 || cache_write.unwrap_or(0) > 0 { + return Some(Usage { + input_tokens: input, + output_tokens: output, + cache_read_input_tokens: cache_read, + cache_creation_input_tokens: cache_write, + }); + } + } + if let Some(session_val) = &body.session { + let input = session_val + .get("total_input_tokens") + .and_then(Value::as_u64) + .unwrap_or(0); + let output = session_val + .get("total_output_tokens") + .and_then(Value::as_u64) + .unwrap_or(0); + if input > 0 || output > 0 { + return Some(Usage { + input_tokens: input, + output_tokens: output, + cache_read_input_tokens: None, + cache_creation_input_tokens: None, + }); + } + } + None +} + /// Emit one committed (or interrupted) turn as `Common` messages. Reasoning /// blocks for the turn's `n`-th assistant message are keyed `(turn_idx, n)` in /// the private sidecar and spliced back in the same order `from_common` wrote /// them. +#[allow(clippy::too_many_arguments)] fn emit_turn( turn: &Value, ts: DateTime, @@ -178,8 +261,10 @@ fn emit_turn( model: Option<&str>, reasoning: &HashMap<(u64, u64), Vec>, images: &HashMap<&str, &FxImage>, + turn_usage: Option, out: &mut Vec, ) { + let start_idx = out.len(); // The user prompt opens the turn. let user_content = user_blocks(turn.get("user"), images); if !user_content.is_empty() { @@ -219,39 +304,46 @@ fn emit_turn( content.push(block); } push_assistant(content, Some(StopReason::Aborted), out); - return; - } - - // Committed turn: intermediate steps, then the concluding text. - if let Some(steps) = turn - .pointer("/execution/tool_steps") - .and_then(Value::as_array) - { - for step in steps { - let mut content = Vec::new(); - if let Some(text) = nonempty_str(step.get("assistant")) { - content.push(Block::Text { text }); - } - if let Some(calls) = step.get("tool_calls").and_then(Value::as_array) { - content.extend(calls.iter().filter_map(tool_use_block)); - } - push_assistant(content, None, out); + } else { + // Committed turn: intermediate steps, then the concluding text. + if let Some(steps) = turn + .pointer("/execution/tool_steps") + .and_then(Value::as_array) + { + for step in steps { + let mut content = Vec::new(); + if let Some(text) = nonempty_str(step.get("assistant")) { + content.push(Block::Text { text }); + } + if let Some(calls) = step.get("tool_calls").and_then(Value::as_array) { + content.extend(calls.iter().filter_map(tool_use_block)); + } + push_assistant(content, None, out); - if let Some(results) = step.get("tool_results").and_then(Value::as_array) { - let blocks: Vec = results.iter().filter_map(tool_result_block).collect(); - if !blocks.is_empty() { - out.push(plain(Role::User, blocks, ts)); + if let Some(results) = step.get("tool_results").and_then(Value::as_array) { + let blocks: Vec = results.iter().filter_map(tool_result_block).collect(); + if !blocks.is_empty() { + out.push(plain(Role::User, blocks, ts)); + } } } } + + let final_text = nonempty_str(turn.get("assistant")); + let final_content: Vec = final_text + .map(|text| Block::Text { text }) + .into_iter() + .collect(); + push_assistant(final_content, Some(StopReason::EndTurn), out); } - let final_text = nonempty_str(turn.get("assistant")); - let final_content: Vec = final_text - .map(|text| Block::Text { text }) - .into_iter() - .collect(); - push_assistant(final_content, Some(StopReason::EndTurn), out); + if let Some(usage) = turn_usage + && let Some(last_asst) = out[start_idx..] + .iter_mut() + .rfind(|m| m.role == Role::Assistant) + { + last_asst.usage = Some(usage); + } } fn user_blocks(user: Option<&Value>, images: &HashMap<&str, &FxImage>) -> Vec { @@ -522,7 +614,7 @@ fn body_from_messages(meta: &Meta, messages: &[Message]) -> FxSession { reasoning: Vec::new(), tool_names: HashMap::new(), }; - let mut turn_payloads: Vec<(i64, Value)> = Vec::new(); + let mut turn_payloads: Vec<(i64, Value, u64, u64)> = Vec::new(); let mut i = 0; let mut turn_idx: u64 = 0; while i < messages.len() { @@ -540,11 +632,19 @@ fn body_from_messages(meta: &Meta, messages: &[Message]) -> FxSession { i += 1; } let body = &messages[body_start..i]; + let mut turn_input = 0u64; + let mut turn_output = 0u64; + for msg in body { + if let Some(u) = msg.usage { + turn_input = turn_input.saturating_add(u.input_tokens); + turn_output = turn_output.saturating_add(u.output_tokens); + } + } let ts = prompt .or_else(|| body.first()) .map_or(meta.timestamp, |m| m.timestamp); let payload = builder.build_turn(turn_idx, prompt, body); - turn_payloads.push((ts.timestamp_millis(), payload)); + turn_payloads.push((ts.timestamp_millis(), payload, turn_input, turn_output)); turn_idx += 1; } @@ -553,7 +653,7 @@ fn body_from_messages(meta: &Meta, messages: &[Message]) -> FxSession { let mut events: Vec = Vec::new(); events.push(session_started(&session_id, &generation, created_ms, meta)); let mut last_ts = created_ms; - for (idx, (ts_ms, turn)) in turn_payloads.into_iter().enumerate() { + for (idx, (ts_ms, turn, turn_input, turn_output)) in turn_payloads.into_iter().enumerate() { last_ts = ts_ms; let seq = u64::try_from(idx).unwrap_or(0) + 2; events.push(json!({ @@ -565,13 +665,30 @@ fn body_from_messages(meta: &Meta, messages: &[Message]) -> FxSession { "kind": "history_turn_committed", "payload": { "conversation_language": "und", - "total_input_tokens": 0, - "total_output_tokens": 0, + "total_input_tokens": turn_input, + "total_output_tokens": turn_output, "turn": turn, }, })); } + let mut total_input = 0u64; + let mut total_output = 0u64; + let mut total_cache_read: Option = None; + let mut total_cache_write: Option = None; + for m in messages { + if let Some(u) = m.usage { + total_input = total_input.saturating_add(u.input_tokens); + total_output = total_output.saturating_add(u.output_tokens); + if let Some(r) = u.cache_read_input_tokens { + total_cache_read = Some(total_cache_read.unwrap_or(0).saturating_add(r)); + } + if let Some(w) = u.cache_creation_input_tokens { + total_cache_write = Some(total_cache_write.unwrap_or(0).saturating_add(w)); + } + } + } + assemble_body( meta, &session_id, @@ -579,6 +696,10 @@ fn body_from_messages(meta: &Meta, messages: &[Message]) -> FxSession { &authority_id, created_ms, last_ts, + total_input, + total_output, + total_cache_read, + total_cache_write, events, builder, ) @@ -594,6 +715,10 @@ fn assemble_body( authority_id: &str, created_ms: i64, last_ts: i64, + total_input: u64, + total_output: u64, + total_cache_read: Option, + total_cache_write: Option, events: Vec, builder: TurnBuilder, ) -> FxSession { @@ -619,8 +744,8 @@ fn assemble_body( "workspace_root": meta.cwd.clone().unwrap_or_default(), "conversation_language": "und", "history_len": events.len().saturating_sub(1), - "total_input_tokens": 0, - "total_output_tokens": 0, + "total_input_tokens": total_input, + "total_output_tokens": total_output, "last_event_seq": last_seq, "event_log_bytes": event_log_bytes, "generation_base_seq": 1, @@ -652,13 +777,42 @@ fn assemble_body( let reasoning = (!builder.reasoning.is_empty()).then(|| json!({ "entries": builder.reasoning })); + let usage = if total_input > 0 + || total_output > 0 + || total_cache_read.unwrap_or(0) > 0 + || total_cache_write.unwrap_or(0) > 0 + { + Some(json!({ + "billing": "complete", + "api_duration_complete": true, + "wall_duration_complete": true, + "code_complete": true, + "next_sequence": 1, + "settled_through_sequence": 0, + "api_duration_ms": 0, + "wall_duration_ms": 0, + "total_cost": 0, + "input_tokens": total_input, + "output_tokens": total_output, + "cache_read_tokens": total_cache_read.unwrap_or(0), + "cache_write_tokens": total_cache_write.unwrap_or(0), + "billable_web_search_calls": 0, + "lines_added": 0, + "lines_removed": 0, + "models": [], + "pending": [], + })) + } else { + None + }; + FxSession { events, session: Some(session), authority: Some(authority), commit: Some(commit), display: Some(display), - usage: None, + usage, checkpoint: None, images: builder.images, reasoning, diff --git a/src/harness/hermes.rs b/src/harness/hermes.rs index 6073375..d7be211 100644 --- a/src/harness/hermes.rs +++ b/src/harness/hermes.rs @@ -25,7 +25,7 @@ use std::collections::HashMap; use std::path::Path; use std::path::PathBuf; -use crate::common::{Block, ImageSource, Message, Meta, Role, StopReason, Tool, ToolOutput}; +use crate::common::{Block, ImageSource, Message, Meta, Role, StopReason, Tool, ToolOutput, Usage}; use crate::error::{Error, Result}; use crate::transcript::{Codec, Common, Discovered, Harness, Saved, Store, TextCodec, Transcript}; @@ -395,10 +395,62 @@ fn assistant_message(row: &Value, meta: &Meta, timestamp: DateTime) -> Opti .get("finish_reason") .and_then(Value::as_str) .map(parse_finish_reason), - usage: None, + usage: parse_row_usage(row), }) } +fn parse_row_usage(row: &Value) -> Option { + if let Some(u) = row.get("usage") { + let input = u + .get("prompt_tokens") + .or_else(|| u.get("input_tokens")) + .and_then(Value::as_u64) + .unwrap_or(0); + let output = u + .get("completion_tokens") + .or_else(|| u.get("output_tokens")) + .and_then(Value::as_u64) + .unwrap_or(0); + let cache_read = u + .get("cached_tokens") + .or_else(|| u.get("cache_read_tokens")) + .or_else(|| u.get("cache_read_input_tokens")) + .and_then(Value::as_u64); + let cache_write = u + .get("cache_creation_tokens") + .or_else(|| u.get("cache_write_tokens")) + .or_else(|| u.get("cache_creation_input_tokens")) + .and_then(Value::as_u64); + if input > 0 || output > 0 || cache_read.unwrap_or(0) > 0 || cache_write.unwrap_or(0) > 0 { + return Some(Usage { + input_tokens: input, + output_tokens: output, + cache_read_input_tokens: cache_read, + cache_creation_input_tokens: cache_write, + }); + } + } + let input = row + .get("prompt_tokens") + .or_else(|| row.get("input_tokens")) + .and_then(Value::as_u64) + .unwrap_or(0); + let output = row + .get("completion_tokens") + .or_else(|| row.get("output_tokens")) + .and_then(Value::as_u64) + .unwrap_or(0); + if input > 0 || output > 0 { + return Some(Usage { + input_tokens: input, + output_tokens: output, + cache_read_input_tokens: None, + cache_creation_input_tokens: None, + }); + } + None +} + fn content_blocks(value: Option<&Value>) -> Vec { match value { None | Some(Value::Null) => Vec::new(), @@ -652,6 +704,7 @@ fn export_from_messages(meta: &Meta, messages: &[Message]) -> Value { }) } +#[allow(clippy::too_many_lines)] fn rows_from_messages(session_id: &str, messages: &[Message]) -> Vec { let mut rows = Vec::new(); let mut next_id = 1_u64; @@ -731,6 +784,22 @@ fn rows_from_messages(session_id: &str, messages: &[Message]) -> Vec { "finish_reason".into(), json!(finish_reason(message.stop_reason.as_ref(), has_tools)), ); + if let Some(usage) = &message.usage + && !usage.is_zero() + { + let mut usage_obj = json!({ + "prompt_tokens": usage.input_tokens, + "completion_tokens": usage.output_tokens, + "total_tokens": usage.total_tokens(), + }); + if let Some(cache_read) = usage.cache_read_input_tokens { + usage_obj["cached_tokens"] = json!(cache_read); + } + if let Some(cache_write) = usage.cache_creation_input_tokens { + usage_obj["cache_creation_tokens"] = json!(cache_write); + } + object.insert("usage".into(), usage_obj); + } } rows.push(row); next_id += 1; diff --git a/src/transcript.rs b/src/transcript.rs index eaf14fc..c635383 100644 --- a/src/transcript.rs +++ b/src/transcript.rs @@ -10,7 +10,7 @@ use std::str::FromStr; use serde::{Deserialize, Serialize}; -use crate::common::{Message, Meta}; +use crate::common::{Message, Meta, Usage}; use crate::error::Result; /// A transcript in some representation `H`. @@ -134,6 +134,23 @@ impl CropError { } impl Transcript { + /// Aggregate token usage reported across all turns in this transcript. + /// + /// Returns `None` when no message in the body records usage. + #[must_use] + pub fn total_usage(&self) -> Option { + let mut total: Option = None; + for msg in &self.body { + if let Some(turn_usage) = msg.usage { + total = match total { + Some(acc) => Some(acc + turn_usage), + None => Some(turn_usage), + }; + } + } + total + } + /// Resolve a [`Span`] to its messages, borrowing from this transcript. /// `None` when the span reaches past the end of the session. #[must_use] diff --git a/tests/integration/main.rs b/tests/integration/main.rs index 4eb9cce..4da926a 100644 --- a/tests/integration/main.rs +++ b/tests/integration/main.rs @@ -26,6 +26,7 @@ mod pi; mod properties; mod simple; mod store_delete; +mod usage; #[cfg(feature = "search")] mod search; diff --git a/tests/integration/usage.rs b/tests/integration/usage.rs new file mode 100644 index 0000000..4d07791 --- /dev/null +++ b/tests/integration/usage.rs @@ -0,0 +1,223 @@ +#![allow(clippy::expect_used, clippy::panic, clippy::unwrap_used)] + +//! Integration tests for token usage aggregation, arithmetic, and cross-harness preservation. + +use chrono::{DateTime, Utc}; +use txcript::common::{Block, Message, Meta, Role, StopReason, Usage}; +use txcript::harness::fx::Fx; +use txcript::harness::hermes::Hermes; +use txcript::{Codec, Transcript}; + +fn sample_meta() -> Meta { + Meta { + id: "usage-session-test".into(), + timestamp: DateTime::::UNIX_EPOCH, + cwd: Some("/work/test".into()), + git_branch: Some("feat/usage".into()), + title: Some("Token Usage Test".into()), + cli_version: Some("0.14.4".into()), + model: Some("claude-sonnet".into()), + } +} + +fn user_msg(text: &str) -> Message { + Message { + role: Role::User, + content: vec![Block::Text { text: text.into() }], + timestamp: DateTime::::UNIX_EPOCH, + model: None, + stop_reason: None, + usage: None, + } +} + +fn asst_msg(text: &str, usage: Option) -> Message { + Message { + role: Role::Assistant, + content: vec![Block::Text { text: text.into() }], + timestamp: DateTime::::UNIX_EPOCH, + model: Some("claude-sonnet".into()), + stop_reason: Some(StopReason::EndTurn), + usage, + } +} + +#[test] +fn usage_arithmetic_methods_and_sums() { + let empty = Usage::default(); + assert!(empty.is_zero()); + assert_eq!(empty.total_tokens(), 0); + + let u1 = Usage { + input_tokens: 150, + output_tokens: 45, + cache_read_input_tokens: Some(30), + cache_creation_input_tokens: None, + }; + assert!(!u1.is_zero()); + assert_eq!(u1.total_tokens(), 225); + + let u2 = Usage { + input_tokens: 250, + output_tokens: 80, + cache_read_input_tokens: Some(50), + cache_creation_input_tokens: Some(20), + }; + assert_eq!(u2.total_tokens(), 400); + + let combined = u1 + u2; + assert_eq!(combined.input_tokens, 400); + assert_eq!(combined.output_tokens, 125); + assert_eq!(combined.cache_read_input_tokens, Some(80)); + assert_eq!(combined.cache_creation_input_tokens, Some(20)); + assert_eq!(combined.total_tokens(), 625); + + let mut accum = u1; + accum += u2; + assert_eq!(accum, combined); + + let list = vec![u1, u2, Usage::default()]; + let summed: Usage = list.into_iter().sum(); + assert_eq!(summed, combined); +} + +#[test] +fn transcript_total_usage_aggregation() { + // 1. Transcript with no usage on any turn + let empty_t = Transcript::new(sample_meta(), vec![user_msg("hi"), asst_msg("hello", None)]); + assert_eq!(empty_t.total_usage(), None); + + // 2. Transcript with multiple assistant turns reporting usage + let u1 = Usage { + input_tokens: 100, + output_tokens: 20, + cache_read_input_tokens: Some(10), + cache_creation_input_tokens: None, + }; + let u2 = Usage { + input_tokens: 200, + output_tokens: 50, + cache_read_input_tokens: None, + cache_creation_input_tokens: Some(15), + }; + let populated_t = Transcript::new( + sample_meta(), + vec![ + user_msg("step 1"), + asst_msg("res 1", Some(u1)), + user_msg("step 2"), + asst_msg("res 2", Some(u2)), + ], + ); + + let total = populated_t.total_usage().expect("should aggregate usage"); + assert_eq!(total.input_tokens, 300); + assert_eq!(total.output_tokens, 70); + assert_eq!(total.cache_read_input_tokens, Some(10)); + assert_eq!(total.cache_creation_input_tokens, Some(15)); + assert_eq!(total.total_tokens(), 395); +} + +#[test] +fn fx_token_usage_preservation_roundtrip() { + let u1 = Usage { + input_tokens: 120, + output_tokens: 30, + cache_read_input_tokens: Some(10), + cache_creation_input_tokens: None, + }; + let u2 = Usage { + input_tokens: 180, + output_tokens: 70, + cache_read_input_tokens: Some(20), + cache_creation_input_tokens: Some(5), + }; + + let original = Transcript::new( + sample_meta(), + vec![ + user_msg("first prompt"), + asst_msg("first response", Some(u1)), + user_msg("second prompt"), + asst_msg("second response", Some(u2)), + ], + ); + + // Convert to Fx native representation + let fx_transcript = Fx::from_common(&original).expect("from_common should succeed"); + + // Verify session.json token totals + let session = fx_transcript + .body + .session + .as_ref() + .expect("session.json exists"); + assert_eq!(session["total_input_tokens"], 300); + assert_eq!(session["total_output_tokens"], 100); + + // Verify usage-v2.json sidecar was created and has aggregated counts + let usage_v2 = fx_transcript + .body + .usage + .as_ref() + .expect("usage-v2.json exists"); + assert_eq!(usage_v2["input_tokens"], 300); + assert_eq!(usage_v2["output_tokens"], 100); + assert_eq!(usage_v2["cache_read_tokens"], 30); + assert_eq!(usage_v2["cache_write_tokens"], 5); + + // Convert back to Common + let roundtrip = Fx::to_common(&fx_transcript).expect("to_common should succeed"); + + // Check that total_usage matches + let rt_total = roundtrip + .total_usage() + .expect("roundtrip should have usage"); + assert_eq!(rt_total.input_tokens, 300); + assert_eq!(rt_total.output_tokens, 100); +} + +#[test] +fn hermes_token_usage_preservation_roundtrip() { + let u = Usage { + input_tokens: 450, + output_tokens: 120, + cache_read_input_tokens: Some(80), + cache_creation_input_tokens: Some(25), + }; + + let original = Transcript::new( + sample_meta(), + vec![ + user_msg("solve this problem"), + asst_msg("solved it", Some(u)), + ], + ); + + // Convert to Hermes native + let hermes_transcript = Hermes::from_common(&original).expect("Hermes from_common"); + + // Verify the assistant row contains the serialized usage object + let rows = hermes_transcript.body["messages"] + .as_array() + .expect("messages array"); + let asst_row = rows + .iter() + .find(|r| r["role"] == "assistant") + .expect("assistant row"); + let row_usage = &asst_row["usage"]; + assert_eq!(row_usage["prompt_tokens"], 450); + assert_eq!(row_usage["completion_tokens"], 120); + assert_eq!(row_usage["cached_tokens"], 80); + assert_eq!(row_usage["cache_creation_tokens"], 25); + assert_eq!(row_usage["total_tokens"], 675); + + // Convert back to Common + let roundtrip = Hermes::to_common(&hermes_transcript).expect("Hermes to_common"); + let rt_total = roundtrip.total_usage().expect("roundtrip has total usage"); + assert_eq!(rt_total.input_tokens, 450); + assert_eq!(rt_total.output_tokens, 120); + assert_eq!(rt_total.cache_read_input_tokens, Some(80)); + assert_eq!(rt_total.cache_creation_input_tokens, Some(25)); + assert_eq!(rt_total.total_tokens(), 675); +} From 6bbbbbb4a379c4a90d07b054dd3e426c69d20e1c Mon Sep 17 00:00:00 2001 From: Aditya kumar singh <143548997+Adityakk9031@users.noreply.github.com> Date: Thu, 17 Sep 2026 20:01:09 +0530 Subject: [PATCH 2/2] fix(cli): collapse nested if conditions in human_header --- cli/src/view.rs | 44 ++++++++++++++++++++++---------------------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/cli/src/view.rs b/cli/src/view.rs index 2da1bdf..f9847b5 100644 --- a/cli/src/view.rs +++ b/cli/src/view.rs @@ -515,29 +515,29 @@ fn human_header(out: &mut String, common: &Transcript, span: &Span, colo &format!("{shown} of {}", common.body.len()), color, ); - if let Some(usage) = common.total_usage() { - if !usage.is_zero() { - let mut parts = vec![ - format!("{} in", usage.input_tokens), - format!("{} out", usage.output_tokens), - ]; - if let Some(cached) = usage.cache_read_input_tokens { - if cached > 0 { - parts.push(format!("{cached} cached")); - } - } - if let Some(created) = usage.cache_creation_input_tokens { - if created > 0 { - parts.push(format!("{created} cache write")); - } - } - human_field( - out, - "Tokens", - &format!("{} ({})", usage.total_tokens(), parts.join(", ")), - color, - ); + if let Some(usage) = common.total_usage() + && !usage.is_zero() + { + let mut parts = vec![ + format!("{} in", usage.input_tokens), + format!("{} out", usage.output_tokens), + ]; + if let Some(cached) = usage.cache_read_input_tokens + && cached > 0 + { + parts.push(format!("{cached} cached")); + } + if let Some(created) = usage.cache_creation_input_tokens + && created > 0 + { + parts.push(format!("{created} cache write")); } + human_field( + out, + "Tokens", + &format!("{} ({})", usage.total_tokens(), parts.join(", ")), + color, + ); } }