scriptoriumadmin/fabro
1
0
Fork 0
mirror of https://github.com/fabro-sh/fabro.git synced 2026-09-06 08:18:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Add token usage tracking to run conclusions and fix attach cleanup

Browse source 
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) <noreply@anthropic.com>
This commit is contained in:
Bryan Helmkamp 2026-03-22 14:24:09 -04:00
parent f28df7e536
commit 597601b314
7 changed files with 148 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
.ai/prompts/code-review-deep-2.md
View file

@ -17,7 +17,7 @@ Output a report with all the bugs using this format:
<bug>
<summary>up to 3 sentences</summary>
<severity>important OR nit</severity>
<pre_exiting>yes OR no</pre_existing>
<pre_existing>yes OR no</pre_existing>
<location>
<file>lib/crates/fabro-cli/src/commands/resume.rs</file>
<start_line>115</start_line>
@ -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.
</summary>
<severity>important</severity>
<pre_exiting>no</pre_existing>
<pre_existing>no</pre_existing>
<location>
<file>lib/crates/fabro-cli/src/commands/resume.rs</file>
<start_line>208</start_line>

13
lib/crates/fabro-cli/src/commands/attach.rs
View file

@ -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));
}
}

22
lib/crates/fabro-cli/src/commands/resume.rs
View file

@ -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<f64> = 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);

61
lib/crates/fabro-cli/src/commands/run.rs
View file

@ -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<f64> = 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),
)),
);
}
}

18
lib/crates/fabro-cli/src/commands/wait.rs
View file

@ -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);

30
lib/crates/fabro-workflows/src/conclusion.rs
View file

@ -31,6 +31,18 @@ pub struct Conclusion {
pub total_cost: Option<f64>,
#[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();

6
lib/crates/fabro-workflows/src/pull_request.rs
View file

@ -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,
}
}