From 28cdefec6744ae659cc8b86b4c830e3960fc1c36 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Mon, 23 Feb 2026 22:01:17 -0500 Subject: [PATCH] Add token usage and cost tracking to pipeline runs Surface LLM token consumption and dollar cost at every verbosity level: default mode appends per-stage tokens/cost, verbose modes include it in event summary/detail, and the Pipeline Result section shows a total. Cost is computed from the catalog pricing for Anthropic models; providers without pricing (OpenAI, Gemini) show token counts only. Dry runs with zero tokens omit the cost line entirely. Co-Authored-By: Claude Opus 4.6 --- attractor-web/src/api.ts | 2 +- crates/attractor/src/cli/backend.rs | 36 +++++++++----- crates/attractor/src/cli/mod.rs | 50 ++++++++++++++++++++ crates/attractor/src/cli/run.rs | 60 +++++++++++++++++++++--- crates/attractor/src/condition.rs | 1 + crates/attractor/src/engine.rs | 1 + crates/attractor/src/event.rs | 3 ++ crates/attractor/src/handler/codergen.rs | 18 ++++--- crates/attractor/src/handler/fan_in.rs | 4 +- crates/attractor/src/handler/parallel.rs | 1 + crates/attractor/src/outcome.rs | 16 ++++++- crates/attractor/tests/integration.rs | 15 +++--- 12 files changed, 172 insertions(+), 35 deletions(-) diff --git a/attractor-web/src/api.ts b/attractor-web/src/api.ts index 718c53bdd..ecabf1bc4 100644 --- a/attractor-web/src/api.ts +++ b/attractor-web/src/api.ts @@ -54,7 +54,7 @@ export type PipelineEvent = | { PipelineCompleted: { duration_ms: number; artifact_count: number } } | { PipelineFailed: { error: string; duration_ms: number } } | { StageStarted: { name: string; index: number } } - | { StageCompleted: { name: string; index: number; duration_ms: number; status: string; preferred_label?: string; suggested_next_ids: string[] } } + | { StageCompleted: { name: string; index: number; duration_ms: number; status: string; preferred_label?: string; suggested_next_ids: string[]; usage?: { model: string; input_tokens: number; output_tokens: number } } } | { StageFailed: { name: string; index: number; error: string; will_retry: boolean } } | { StageRetrying: { name: string; index: number; attempt: number; delay_ms: number } } | { ParallelStarted: { branch_count: number } } diff --git a/crates/attractor/src/cli/backend.rs b/crates/attractor/src/cli/backend.rs index d787d8e86..73b5356f1 100644 --- a/crates/attractor/src/cli/backend.rs +++ b/crates/attractor/src/cli/backend.rs @@ -15,6 +15,7 @@ use crate::context::Context; use crate::error::AttractorError; use crate::graph::Node; use crate::handler::codergen::{CodergenBackend, CodergenResult}; +use crate::outcome::StageUsage; /// LLM backend that delegates to an `agent` Session per invocation. pub struct AgentBackend { @@ -152,19 +153,30 @@ impl CodergenBackend for AgentBackend { AttractorError::Handler(format!("Agent session failed: {e}")) })?; + // Aggregate token usage from all assistant turns. + let (mut turn_count, mut tool_call_count, mut input_tokens, mut output_tokens) = + (0usize, 0usize, 0i64, 0i64); + for turn in session.history().turns() { + if let Turn::Assistant { + tool_calls, usage, .. + } = turn + { + turn_count += 1; + tool_call_count += tool_calls.len(); + input_tokens += usage.input_tokens; + output_tokens += usage.output_tokens; + } + } + + let stage_usage = StageUsage { + model: self.model.clone(), + input_tokens, + output_tokens, + }; + // Print session summary to stderr. if self.verbose >= 1 { - let (mut turn_count, mut tool_call_count, mut total_tokens) = (0usize, 0usize, 0i64); - for turn in session.history().turns() { - if let Turn::Assistant { - tool_calls, usage, .. - } = turn - { - turn_count += 1; - tool_call_count += tool_calls.len(); - total_tokens += usage.total_tokens; - } - } + let total_tokens = input_tokens + output_tokens; let token_str = if total_tokens >= 1000 { format!("{}k tokens", total_tokens / 1000) } else { @@ -194,7 +206,7 @@ impl CodergenBackend for AgentBackend { }) .unwrap_or_default(); - Ok(CodergenResult::Text(response)) + Ok(CodergenResult::Text { text: response, usage: Some(stage_usage) }) } } diff --git a/crates/attractor/src/cli/mod.rs b/crates/attractor/src/cli/mod.rs index 51894249e..ce0ff2671 100644 --- a/crates/attractor/src/cli/mod.rs +++ b/crates/attractor/src/cli/mod.rs @@ -11,6 +11,7 @@ use std::path::PathBuf; use terminal::Styles; use crate::event::PipelineEvent; +use crate::outcome::StageUsage; use crate::validation::{Diagnostic, Severity}; #[derive(Parser)] @@ -187,6 +188,7 @@ pub fn format_event_summary(event: &PipelineEvent, styles: &Styles) -> String { status, preferred_label, suggested_next_ids, + usage, } => { let mut s = format!("[STAGE_COMPLETED] name={name} index={index} duration={duration_ms}ms status={status}"); if let Some(label) = preferred_label { @@ -195,6 +197,15 @@ pub fn format_event_summary(event: &PipelineEvent, styles: &Styles) -> String { if !suggested_next_ids.is_empty() { s.push_str(&format!(" suggested_next_ids={}", suggested_next_ids.join(","))); } + if let Some(u) = usage { + let total = u.input_tokens + u.output_tokens; + let tokens_str = format_tokens_human(total); + if let Some(cost) = compute_stage_cost(u) { + s.push_str(&format!(" tokens={tokens_str} cost={}", format_cost(cost))); + } else { + s.push_str(&format!(" tokens={tokens_str}")); + } + } s } PipelineEvent::StageFailed { @@ -296,6 +307,7 @@ pub fn format_event_detail(event: &PipelineEvent, styles: &Styles) -> String { status, preferred_label, suggested_next_ids, + usage, } => { let mut s = format!("{d}── STAGE_COMPLETED ──────────────────────────{r}\n {d}name:{r} {name}\n {d}index:{r} {index}\n {d}duration_ms:{r} {duration_ms}\n {d}status:{r} {status}\n"); if let Some(label) = preferred_label { @@ -304,6 +316,18 @@ pub fn format_event_detail(event: &PipelineEvent, styles: &Styles) -> String { if !suggested_next_ids.is_empty() { s.push_str(&format!(" {d}suggested_next_ids:{r} {}\n", suggested_next_ids.join(", "))); } + if let Some(u) = usage { + let total = u.input_tokens + u.output_tokens; + s.push_str(&format!(" {d}model:{r} {}\n", u.model)); + s.push_str(&format!(" {d}tokens:{r} {} ({} in / {} out)\n", + format_tokens_human(total), + format_tokens_human(u.input_tokens), + format_tokens_human(u.output_tokens), + )); + if let Some(cost) = compute_stage_cost(u) { + s.push_str(&format!(" {d}cost:{r} {}\n", format_cost(cost))); + } + } s } PipelineEvent::StageFailed { @@ -367,3 +391,29 @@ pub fn format_event_detail(event: &PipelineEvent, styles: &Styles) -> String { } } } + +/// Compute the dollar cost for a stage's token usage, if pricing is available. +#[must_use] +pub fn compute_stage_cost(usage: &StageUsage) -> Option { + let info = llm::catalog::get_model_info(&usage.model)?; + let input_rate = info.input_cost_per_million?; + let output_rate = info.output_cost_per_million?; + Some(usage.input_tokens as f64 * input_rate / 1_000_000.0 + + usage.output_tokens as f64 * output_rate / 1_000_000.0) +} + +/// Format a dollar cost for display (e.g. `"$1.23"`). +#[must_use] +pub fn format_cost(cost: f64) -> String { + format!("${cost:.2}") +} + +/// Format a token count for human display (e.g. `"15.2k"` or `"850"`). +#[must_use] +pub fn format_tokens_human(tokens: i64) -> String { + if tokens >= 1000 { + format!("{:.1}k", tokens as f64 / 1000.0) + } else { + tokens.to_string() + } +} diff --git a/crates/attractor/src/cli/run.rs b/crates/attractor/src/cli/run.rs index c80a47ed2..1d7977481 100644 --- a/crates/attractor/src/cli/run.rs +++ b/crates/attractor/src/cli/run.rs @@ -1,5 +1,5 @@ use std::path::PathBuf; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::time::Instant; use anyhow::bail; @@ -18,7 +18,16 @@ use crate::pipeline::PipelineBuilder; use crate::validation::Severity; use super::backend::AgentBackend; -use super::{format_duration_human, format_event_detail, format_event_summary, print_diagnostics, read_dot_file, RunArgs}; +use super::{compute_stage_cost, format_cost, format_duration_human, format_event_detail, format_event_summary, format_tokens_human, print_diagnostics, read_dot_file, RunArgs}; + +/// Accumulates token usage and cost across all pipeline stages. +#[derive(Default)] +struct CostAccumulator { + total_input_tokens: i64, + total_output_tokens: i64, + total_cost: f64, + has_pricing: bool, +} /// Execute a full pipeline run. /// @@ -69,6 +78,24 @@ pub async fn run_command(args: RunArgs, styles: &'static Styles) -> anyhow::Resu // 3. Build event emitter let mut emitter = EventEmitter::new(); + + // Cost accumulator — shared across all verbosity levels + let accumulator = Arc::new(Mutex::new(CostAccumulator::default())); + let acc_clone = Arc::clone(&accumulator); + emitter.on_event(move |event| { + if let crate::event::PipelineEvent::StageCompleted { usage, .. } = event { + if let Some(u) = usage { + let mut acc = acc_clone.lock().unwrap(); + acc.total_input_tokens += u.input_tokens; + acc.total_output_tokens += u.output_tokens; + if let Some(cost) = compute_stage_cost(u) { + acc.total_cost += cost; + acc.has_pricing = true; + } + } + } + }); + if args.verbose >= 2 { emitter.on_event(move |event| { eprint!("{}", format_event_detail(event, styles)); @@ -80,12 +107,22 @@ pub async fn run_command(args: RunArgs, styles: &'static Styles) -> anyhow::Resu } else { emitter.on_event(move |event| { match event { - crate::event::PipelineEvent::StageCompleted { name, duration_ms, status, .. } => { - eprintln!( - "{dim}Stage \"{name}\" completed ({status}) in {duration}{reset}", + crate::event::PipelineEvent::StageCompleted { name, duration_ms, status, usage, .. } => { + let mut line = format!( + "{dim}Stage \"{name}\" completed ({status}) in {duration}", duration = format_duration_human(*duration_ms), - dim = styles.dim, reset = styles.reset, + dim = styles.dim, ); + if let Some(u) = usage { + let total = u.input_tokens + u.output_tokens; + let tokens_str = format_tokens_human(total); + if let Some(cost) = compute_stage_cost(u) { + line.push_str(&format!(" \u{2014} {tokens_str} tokens ({})", format_cost(cost))); + } else { + line.push_str(&format!(" \u{2014} {tokens_str} tokens")); + } + } + eprintln!("{line}{reset}", reset = styles.reset); } crate::event::PipelineEvent::StageFailed { name, .. } => { eprintln!( @@ -198,6 +235,17 @@ pub async fn run_command(args: RunArgs, styles: &'static Styles) -> anyhow::Resu eprintln!("Status: {status_color}{status_str}{reset}", reset = styles.reset); eprintln!("Duration: {}", format_duration_human(run_duration_ms)); + let acc = accumulator.lock().unwrap(); + let total_tokens = acc.total_input_tokens + acc.total_output_tokens; + if total_tokens > 0 { + if acc.has_pricing { + eprintln!("Cost: {} ({} tokens)", format_cost(acc.total_cost), format_tokens_human(total_tokens)); + } else { + eprintln!("Tokens: {}", format_tokens_human(total_tokens)); + } + } + drop(acc); + if let Some(notes) = &outcome.notes { eprintln!("Notes: {notes}"); } diff --git a/crates/attractor/src/condition.rs b/crates/attractor/src/condition.rs index 7705bddc3..88a6391fa 100644 --- a/crates/attractor/src/condition.rs +++ b/crates/attractor/src/condition.rs @@ -154,6 +154,7 @@ mod tests { context_updates: std::collections::HashMap::new(), notes: None, failure_reason: None, + usage: None, } } diff --git a/crates/attractor/src/engine.rs b/crates/attractor/src/engine.rs index f424a38c3..405cdc2a4 100644 --- a/crates/attractor/src/engine.rs +++ b/crates/attractor/src/engine.rs @@ -863,6 +863,7 @@ impl PipelineEngine { status: outcome.status.to_string(), preferred_label: outcome.preferred_label.clone(), suggested_next_ids: outcome.suggested_next_ids.clone(), + usage: outcome.usage.clone(), }); self.inform( &format!("Stage completed: {}", node.label()), diff --git a/crates/attractor/src/event.rs b/crates/attractor/src/event.rs index 8d0b8ddd2..1d0a55871 100644 --- a/crates/attractor/src/event.rs +++ b/crates/attractor/src/event.rs @@ -1,5 +1,7 @@ use serde::{Deserialize, Serialize}; +use crate::outcome::StageUsage; + /// Events emitted during pipeline execution for observability. #[derive(Debug, Clone, Serialize, Deserialize)] pub enum PipelineEvent { @@ -26,6 +28,7 @@ pub enum PipelineEvent { status: String, preferred_label: Option, suggested_next_ids: Vec, + usage: Option, }, StageFailed { name: String, diff --git a/crates/attractor/src/handler/codergen.rs b/crates/attractor/src/handler/codergen.rs index b2440a212..c94959a53 100644 --- a/crates/attractor/src/handler/codergen.rs +++ b/crates/attractor/src/handler/codergen.rs @@ -5,13 +5,16 @@ use async_trait::async_trait; use crate::context::Context; use crate::error::AttractorError; use crate::graph::{Graph, Node}; -use crate::outcome::Outcome; +use crate::outcome::{Outcome, StageUsage}; use super::{EngineServices, Handler}; /// Result from a `CodergenBackend` invocation. pub enum CodergenResult { - Text(String), + Text { + text: String, + usage: Option, + }, Full(Outcome), } @@ -199,7 +202,7 @@ impl Handler for CodergenHandler { let thread_id = context .get("internal.thread_id") .and_then(|v| v.as_str().map(String::from)); - let response_text = if let Some(backend) = &self.backend { + let (response_text, stage_usage) = if let Some(backend) = &self.backend { match backend.run(node, &prompt, context, thread_id.as_deref()).await { Ok(CodergenResult::Full(outcome)) => { let status_json = serde_json::to_string_pretty(&outcome) @@ -207,7 +210,7 @@ impl Handler for CodergenHandler { tokio::fs::write(stage_dir.join("status.json"), &status_json).await?; return Ok(outcome); } - Ok(CodergenResult::Text(text)) => text, + Ok(CodergenResult::Text { text, usage }) => (text, usage), Err(e) if e.is_retryable() => { return Err(e); } @@ -216,7 +219,7 @@ impl Handler for CodergenHandler { } } } else { - format!("[Simulated] Response for stage: {}", node.id) + (format!("[Simulated] Response for stage: {}", node.id), None) }; // 5. Execute post-hook (spec 9.7) @@ -246,6 +249,7 @@ impl Handler for CodergenHandler { // 7b. Parse routing directives from response text extract_status_fields(&response_text, &mut outcome); + outcome.usage = stage_usage; let status_json = serde_json::to_string_pretty(&outcome) .unwrap_or_else(|_| "{}".to_string()); @@ -512,7 +516,7 @@ mod tests { ) -> Result { *self.captured_thread_id.lock().unwrap() = Some(thread_id.map(String::from)); - Ok(CodergenResult::Text("ok".to_string())) + Ok(CodergenResult::Text { text: "ok".to_string(), usage: None }) } } @@ -557,7 +561,7 @@ mod tests { ) -> Result { *self.captured_thread_id.lock().unwrap() = Some(thread_id.map(String::from)); - Ok(CodergenResult::Text("ok".to_string())) + Ok(CodergenResult::Text { text: "ok".to_string(), usage: None }) } } diff --git a/crates/attractor/src/handler/fan_in.rs b/crates/attractor/src/handler/fan_in.rs index 5a6075bbb..e454b6fa9 100644 --- a/crates/attractor/src/handler/fan_in.rs +++ b/crates/attractor/src/handler/fan_in.rs @@ -190,7 +190,7 @@ async fn llm_evaluate( score: 0.0, }) } - Ok(CodergenResult::Text(text)) => { + Ok(CodergenResult::Text { text, .. }) => { // Write response to logs tokio::fs::write(stage_dir.join("response.md"), &text).await?; @@ -369,7 +369,7 @@ mod tests { _thread_id: Option<&str>, ) -> Result { // Return text that contains the ID "branch_b" - Ok(CodergenResult::Text("The best candidate is branch_b".to_string())) + Ok(CodergenResult::Text { text: "The best candidate is branch_b".to_string(), usage: None }) } } diff --git a/crates/attractor/src/handler/parallel.rs b/crates/attractor/src/handler/parallel.rs index b6fba9985..6d5c3f298 100644 --- a/crates/attractor/src/handler/parallel.rs +++ b/crates/attractor/src/handler/parallel.rs @@ -297,6 +297,7 @@ impl Handler for ParallelHandler { } else { None }, + usage: None, }; if is_fail { diff --git a/crates/attractor/src/outcome.rs b/crates/attractor/src/outcome.rs index ddf201b8c..c7679235a 100644 --- a/crates/attractor/src/outcome.rs +++ b/crates/attractor/src/outcome.rs @@ -43,6 +43,14 @@ impl FromStr for StageStatus { } } +/// Token usage from a single pipeline stage. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StageUsage { + pub model: String, + pub input_tokens: i64, + pub output_tokens: i64, +} + /// The result of executing a node handler. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Outcome { @@ -57,6 +65,8 @@ pub struct Outcome { pub notes: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub failure_reason: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub usage: Option, } impl Outcome { @@ -69,6 +79,7 @@ impl Outcome { context_updates: HashMap::new(), notes: None, failure_reason: None, + usage: None, } } @@ -80,6 +91,7 @@ impl Outcome { context_updates: HashMap::new(), notes: None, failure_reason: Some(reason.into()), + usage: None, } } @@ -91,10 +103,11 @@ impl Outcome { context_updates: HashMap::new(), notes: None, failure_reason: Some(reason.into()), + usage: None, } } - #[must_use] + #[must_use] pub fn skipped() -> Self { Self { status: StageStatus::Skipped, @@ -103,6 +116,7 @@ impl Outcome { context_updates: HashMap::new(), notes: None, failure_reason: None, + usage: None, } } } diff --git a/crates/attractor/tests/integration.rs b/crates/attractor/tests/integration.rs index b0734c99f..fb439b4d5 100644 --- a/crates/attractor/tests/integration.rs +++ b/crates/attractor/tests/integration.rs @@ -1057,11 +1057,14 @@ impl CodergenBackend for MockCodergenBackend { _context: &Context, _thread_id: Option<&str>, ) -> Result { - Ok(CodergenResult::Text(format!( - "Response for {}: processed prompt '{}'", - node.id, - &prompt[..prompt.len().min(50)] - ))) + Ok(CodergenResult::Text { + text: format!( + "Response for {}: processed prompt '{}'", + node.id, + &prompt[..prompt.len().min(50)] + ), + usage: None, + }) } } @@ -5098,7 +5101,7 @@ mod real_llm { .complete(&request) .await .map_err(|e| AttractorError::Handler(e.to_string()))?; - Ok(CodergenResult::Text(response.text())) + Ok(CodergenResult::Text { text: response.text(), usage: None }) } }