From 597601b31401f0f0922c9338bdfd1d57409ad843 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Sun, 22 Mar 2026 14:24:09 -0400 Subject: [PATCH] Add token usage tracking to run conclusions and fix attach cleanup Track input, output, cache, and reasoning tokens in the Conclusion struct so the run summary can display token usage even when cost pricing is unavailable. The summary now shows cache read/write stats and reasoning tokens when present. Also fix attach_run to kill the engine child process on timeout or cancellation instead of orphaning it, and return exit code 1 on cancel. Co-Authored-By: Claude Opus 4.6 (1M context) --- .ai/prompts/code-review-deep-2.md | 4 +- lib/crates/fabro-cli/src/commands/attach.rs | 13 +++- lib/crates/fabro-cli/src/commands/resume.rs | 22 +++++++ lib/crates/fabro-cli/src/commands/run.rs | 61 ++++++++++++++++++- lib/crates/fabro-cli/src/commands/wait.rs | 18 ++++++ lib/crates/fabro-workflows/src/conclusion.rs | 30 +++++++++ .../fabro-workflows/src/pull_request.rs | 6 ++ 7 files changed, 148 insertions(+), 6 deletions(-) diff --git a/.ai/prompts/code-review-deep-2.md b/.ai/prompts/code-review-deep-2.md index f904ed5cb..6bc81586b 100644 --- a/.ai/prompts/code-review-deep-2.md +++ b/.ai/prompts/code-review-deep-2.md @@ -17,7 +17,7 @@ Output a report with all the bugs using this format: up to 3 sentences important OR nit - yes OR no + yes OR no lib/crates/fabro-cli/src/commands/resume.rs 115 @@ -49,7 +49,7 @@ Here is a real-world example: `prepare_from_checkpoint` unconditionally creates a `LocalSandbox` via `local_sandbox_with_callback`, completely ignoring the `--sandbox` flag and TOML config. A user running `fabro resume --checkpoint logs/checkpoint.json --workflow w.fabro --sandbox docker` will silently get a local sandbox instead of Docker; to fix this, call `resolve_sandbox_provider(args.sandbox.map(Into::into), None, run_defaults)` just as `prepare_from_branch` does. important - no + no lib/crates/fabro-cli/src/commands/resume.rs 208 diff --git a/lib/crates/fabro-cli/src/commands/attach.rs b/lib/crates/fabro-cli/src/commands/attach.rs index 6f3bf7a91..6a04c753a 100644 --- a/lib/crates/fabro-cli/src/commands/attach.rs +++ b/lib/crates/fabro-cli/src/commands/attach.rs @@ -48,13 +48,24 @@ pub async fn attach_run( tokio::time::sleep(std::time::Duration::from_millis(100)).await; wait_count += 1; if wait_count > 100 { + // Kill the engine to avoid orphaning it + if let Some(mut child) = engine_child { + let _ = child.kill(); + let _ = child.wait(); + } bail!( "Timed out waiting for progress.jsonl to appear in {}", run_dir.display() ); } if cancelled.load(Ordering::Relaxed) { - return Ok(ExitCode::from(0)); + if kill_on_detach { + if let Some(mut child) = engine_child { + let _ = child.kill(); + let _ = child.wait(); + } + } + return Ok(ExitCode::from(1)); } } diff --git a/lib/crates/fabro-cli/src/commands/resume.rs b/lib/crates/fabro-cli/src/commands/resume.rs index f1698e58a..919ea9d84 100644 --- a/lib/crates/fabro-cli/src/commands/resume.rs +++ b/lib/crates/fabro-cli/src/commands/resume.rs @@ -1301,6 +1301,13 @@ async fn run_resumed( let checkpoint_loaded = Checkpoint::load(&run_dir.join("checkpoint.json")).ok(); let stage_durations = fabro_retro::retro::extract_stage_durations(&run_dir); + let mut total_input_tokens: i64 = 0; + let mut total_output_tokens: i64 = 0; + let mut total_cache_read_tokens: i64 = 0; + let mut total_cache_write_tokens: i64 = 0; + let mut total_reasoning_tokens: i64 = 0; + let mut has_pricing = false; + let (stages, total_cost, total_retries) = if let Some(ref cp) = checkpoint_loaded { let mut stages = Vec::new(); let mut cost_sum: Option = None; @@ -1319,6 +1326,15 @@ async fn run_resumed( let cost = outcome.and_then(|o| o.usage.as_ref()).and_then(|u| u.cost); if let Some(c) = cost { *cost_sum.get_or_insert(0.0) += c; + has_pricing = true; + } + + if let Some(usage) = outcome.and_then(|o| o.usage.as_ref()) { + total_input_tokens += usage.input_tokens; + total_output_tokens += usage.output_tokens; + total_cache_read_tokens += usage.cache_read_tokens.unwrap_or(0); + total_cache_write_tokens += usage.cache_write_tokens.unwrap_or(0); + total_reasoning_tokens += usage.reasoning_tokens.unwrap_or(0); } stages.push(fabro_workflows::conclusion::StageSummary { @@ -1343,6 +1359,12 @@ async fn run_resumed( stages, total_cost, total_retries, + total_input_tokens, + total_output_tokens, + total_cache_read_tokens, + total_cache_write_tokens, + total_reasoning_tokens, + has_pricing, }; let _ = conclusion.save(&run_dir.join("conclusion.json")); fabro_workflows::run_status::write_run_status(&run_dir, run_status, status_reason); diff --git a/lib/crates/fabro-cli/src/commands/run.rs b/lib/crates/fabro-cli/src/commands/run.rs index a71e1e64d..80c2233f2 100644 --- a/lib/crates/fabro-cli/src/commands/run.rs +++ b/lib/crates/fabro-cli/src/commands/run.rs @@ -1567,6 +1567,13 @@ pub async fn run_command( let checkpoint = Checkpoint::load(&run_dir.join("checkpoint.json")).ok(); let stage_durations = fabro_retro::retro::extract_stage_durations(&run_dir); + let mut total_input_tokens: i64 = 0; + let mut total_output_tokens: i64 = 0; + let mut total_cache_read_tokens: i64 = 0; + let mut total_cache_write_tokens: i64 = 0; + let mut total_reasoning_tokens: i64 = 0; + let mut has_pricing = false; + let (stages, total_cost, total_retries) = if let Some(ref cp) = checkpoint { let mut stages = Vec::new(); let mut cost_sum: Option = None; @@ -1585,6 +1592,15 @@ pub async fn run_command( let cost = outcome.and_then(|o| o.usage.as_ref()).and_then(|u| u.cost); if let Some(c) = cost { *cost_sum.get_or_insert(0.0) += c; + has_pricing = true; + } + + if let Some(usage) = outcome.and_then(|o| o.usage.as_ref()) { + total_input_tokens += usage.input_tokens; + total_output_tokens += usage.output_tokens; + total_cache_read_tokens += usage.cache_read_tokens.unwrap_or(0); + total_cache_write_tokens += usage.cache_write_tokens.unwrap_or(0); + total_reasoning_tokens += usage.reasoning_tokens.unwrap_or(0); } stages.push(fabro_workflows::conclusion::StageSummary { @@ -1609,6 +1625,12 @@ pub async fn run_command( stages, total_cost, total_retries, + total_input_tokens, + total_output_tokens, + total_cache_read_tokens, + total_cache_write_tokens, + total_reasoning_tokens, + has_pricing, }; let _ = conclusion.save(&run_dir.join("conclusion.json")); fabro_workflows::run_status::write_run_status(&run_dir, run_status, status_reason); @@ -1876,13 +1898,46 @@ pub fn print_run_summary(run_dir: &Path, run_id: &str, styles: &Styles) { HumanDuration(Duration::from_millis(conclusion.duration_ms)) ); - if let Some(cost) = conclusion.total_cost { - if cost > 0.0 { + let total_tokens = conclusion.total_input_tokens + conclusion.total_output_tokens; + if total_tokens > 0 { + if conclusion.has_pricing { + if let Some(cost) = conclusion.total_cost { + if cost > 0.0 { + eprintln!( + "{}", + styles.dim.apply_to(format!( + "Cost: {} ({} toks)", + format_cost(cost), + format_tokens_human(total_tokens) + )) + ); + } + } + } else { eprintln!( "{}", styles .dim - .apply_to(format!("Cost: {}", format_cost(cost))) + .apply_to(format!("Toks: {}", format_tokens_human(total_tokens))) + ); + } + if conclusion.total_cache_read_tokens > 0 { + eprintln!( + "{}", + styles.dim.apply_to(format!( + "Cache: {} read, {} write", + format_tokens_human(conclusion.total_cache_read_tokens), + format_tokens_human(conclusion.total_cache_write_tokens), + )), + ); + } + if conclusion.total_reasoning_tokens > 0 { + eprintln!( + "{}", + styles.dim.apply_to(format!( + "Reasoning: {} tokens", + format_tokens_human(conclusion.total_reasoning_tokens), + )), ); } } diff --git a/lib/crates/fabro-cli/src/commands/wait.rs b/lib/crates/fabro-cli/src/commands/wait.rs index a6054f1b2..f1f617d68 100644 --- a/lib/crates/fabro-cli/src/commands/wait.rs +++ b/lib/crates/fabro-cli/src/commands/wait.rs @@ -155,6 +155,12 @@ mod tests { stages: vec![], total_cost: Some(0.42), total_retries: 0, + total_input_tokens: 0, + total_output_tokens: 0, + total_cache_read_tokens: 0, + total_cache_write_tokens: 0, + total_reasoning_tokens: 0, + has_pricing: false, }; let json = build_json_output(RunStatus::Succeeded, "ABC123", Some(&conclusion)); assert_eq!(json["run_id"], "ABC123"); @@ -189,6 +195,12 @@ mod tests { stages: vec![], total_cost: None, total_retries: 0, + total_input_tokens: 0, + total_output_tokens: 0, + total_cache_read_tokens: 0, + total_cache_write_tokens: 0, + total_reasoning_tokens: 0, + has_pricing: false, }; let json = build_json_output(RunStatus::Failed, "JKL012", Some(&conclusion)); assert!(json.get("total_cost").is_none()); @@ -207,6 +219,12 @@ mod tests { stages: vec![], total_cost: Some(0.15), total_retries: 0, + total_input_tokens: 0, + total_output_tokens: 0, + total_cache_read_tokens: 0, + total_cache_write_tokens: 0, + total_reasoning_tokens: 0, + has_pricing: false, }; // Just verify no panic; actual stderr output is hard to capture print_human_output(RunStatus::Succeeded, "ABC123", Some(&conclusion), &styles); diff --git a/lib/crates/fabro-workflows/src/conclusion.rs b/lib/crates/fabro-workflows/src/conclusion.rs index 830380d7b..4abe21317 100644 --- a/lib/crates/fabro-workflows/src/conclusion.rs +++ b/lib/crates/fabro-workflows/src/conclusion.rs @@ -31,6 +31,18 @@ pub struct Conclusion { pub total_cost: Option, #[serde(default)] pub total_retries: u32, + #[serde(default)] + pub total_input_tokens: i64, + #[serde(default)] + pub total_output_tokens: i64, + #[serde(default)] + pub total_cache_read_tokens: i64, + #[serde(default)] + pub total_cache_write_tokens: i64, + #[serde(default)] + pub total_reasoning_tokens: i64, + #[serde(default)] + pub has_pricing: bool, } impl Conclusion { @@ -72,6 +84,12 @@ mod tests { ], total_cost: Some(0.15), total_retries: 1, + total_input_tokens: 5000, + total_output_tokens: 1500, + total_cache_read_tokens: 2000, + total_cache_write_tokens: 500, + total_reasoning_tokens: 300, + has_pricing: true, } } @@ -127,6 +145,12 @@ mod tests { stages: vec![], total_cost: None, total_retries: 0, + total_input_tokens: 0, + total_output_tokens: 0, + total_cache_read_tokens: 0, + total_cache_write_tokens: 0, + total_reasoning_tokens: 0, + has_pricing: false, }; conclusion.save(&path).unwrap(); @@ -152,6 +176,12 @@ mod tests { stages: vec![], total_cost: None, total_retries: 0, + total_input_tokens: 0, + total_output_tokens: 0, + total_cache_read_tokens: 0, + total_cache_write_tokens: 0, + total_reasoning_tokens: 0, + has_pricing: false, }; conclusion.save(&path).unwrap(); let loaded = Conclusion::load(&path).unwrap(); diff --git a/lib/crates/fabro-workflows/src/pull_request.rs b/lib/crates/fabro-workflows/src/pull_request.rs index dcd9666c6..d700afe2e 100644 --- a/lib/crates/fabro-workflows/src/pull_request.rs +++ b/lib/crates/fabro-workflows/src/pull_request.rs @@ -458,6 +458,12 @@ mod tests { ], total_cost: Some(0.42), total_retries: 0, + total_input_tokens: 0, + total_output_tokens: 0, + total_cache_read_tokens: 0, + total_cache_write_tokens: 0, + total_reasoning_tokens: 0, + has_pricing: true, } }