fabro/lib/components/fabro-workflow/src/lib.rs
Bryan Helmkamp ef86ad278a
Carry usage as lithos-llm's Usage and rename billing to usage
Re-pin lithos-llm to 55add4596b861a0623d00c3a54aa5c147c8d504b and
pebble to c91810fe51aece80359b9cd8efea971af0c46925, where token usage
and cost travel together as Usage { tokens: TokenCounts, cost:
Option<Cost> }. Fabro now carries that one type everywhere it used to
carry BilledTokenCounts, BilledModelUsage, UsdMicros, or a token count
beside a cost_usd_micros.

fabro-types: billing.rs is usage.rs with ModelRef, ModelUsage { model,
usage }, sum_usage, and usage_is_empty; billing_rollup.rs is
usage_rollup.rs with ProjectionUsageStage, ProjectionUsageByModel,
ProjectionUsageRollup, and usage_rollup_from_projection. Every usage
field is named usage: StageProjection.usage and usage_by_model,
Outcome<Option<ModelUsage>>, stage.completed and stage.failed usage and
usage_by_model, prompt.completed usage, run.completed and run.failed
usage (total_usd_micros is gone), Conclusion.usage, StageSummary.usage,
Run.usage. RunSize buckets by Cost.

fabro-workflow: model_usage_from_llm prices tokens from the catalog with
a Catalog cost source, with_reported_cost keeps a provider cost, and the
pebble handler's stage_usage groups pebble's accounts by model and sums
rows with Usage::saturating_add, so a total has a cost only when every
priced part was priced. The store fold's live usage is the agent's
usage plus its descendants'.

API: the OpenAPI spec deletes BilledTokenCounts, BilledModelUsage,
CompletionUsage, CompletionCost, TokenUsage, and RunBillingSummary,
adds TokenCounts, Cost, Usage, and ModelUsage, and renames every
billing schema, property, tag, path, and operation to usage. fabro-api
reuses lithos-llm's and fabro-types' types through with_replacement,
with a round-trip test per replacement.

Old stored runs get no migration: their pebble events in the old shape
read back with zero usage, and their rebuilt projections lose agent
usage.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
2026-09-14 12:31:34 -06:00

345 lines
12 KiB
Rust

#![cfg_attr(
test,
allow(
clippy::absolute_paths,
clippy::get_unwrap,
clippy::large_futures,
clippy::needless_borrows_for_generic_args,
clippy::option_option,
clippy::ptr_as_ptr,
clippy::ref_as_ptr,
clippy::cast_ptr_alignment,
clippy::uninlined_format_args,
clippy::unnecessary_literal_bound,
reason = "Test-only workflow helpers favor explicit fixtures over pedantic style lints."
)
)]
use std::collections::HashMap;
use std::sync::Arc;
use fabro_store::EventEnvelope;
use fabro_types::{EventBody, StageId};
/// Callback invoked when a workflow node starts executing.
pub type OnNodeCallback = Option<Arc<dyn Fn(&str) + Send + Sync>>;
/// Convert a Duration's milliseconds to u64, saturating on overflow.
pub(crate) fn millis_u64(d: std::time::Duration) -> u64 {
u64::try_from(d.as_millis()).unwrap_or(u64::MAX)
}
/// Extract the timing breakdown from a `stage.completed` / `stage.failed`
/// event body, or `None` for any other variant.
fn stage_completion_timing(body: &EventBody) -> Option<fabro_types::StageTiming> {
match body {
EventBody::StageCompleted(props) => Some(props.timing),
EventBody::StageFailed(props) => Some(props.timing),
_ => None,
}
}
/// Extract per-stage (node_id, visit) timing from `stage.completed` /
/// `stage.failed` events. Keys on the full [`StageId`] so multi-visit stages
/// (e.g. a looped `verify` node) keep distinct timings.
///
/// This is the canonical primitive; [`total_stage_timing_by_node`] and
/// [`latest_stage_timing_by_node`] are explicit rollups built on top of it.
pub fn extract_stage_timings_by_stage_id(
events: &[EventEnvelope],
) -> HashMap<StageId, fabro_types::StageTiming> {
let mut timings = HashMap::new();
for envelope in events {
let Some(timing) = stage_completion_timing(&envelope.event.body) else {
continue;
};
let Some(stage_id) = envelope.event.stage_id.as_ref() else {
continue;
};
timings.insert(stage_id.clone(), timing);
}
timings
}
/// Sum of timing in each node across every visit. Use for usage
/// where a retried node should count its full time. `wall_time_ms`,
/// `inference_time_ms`, `tool_time_ms`, and `active_time_ms` are all summed
/// per node.
pub fn total_stage_timing_by_node(
events: &[EventEnvelope],
) -> HashMap<String, fabro_types::StageTiming> {
let mut totals: HashMap<String, fabro_types::StageTiming> = HashMap::new();
for (stage_id, timing) in extract_stage_timings_by_stage_id(events) {
let entry = totals.entry(stage_id.node_id().to_string()).or_default();
*entry = entry.saturating_add(&timing);
}
totals
}
/// Timing of each node's most recent visit (the highest visit number). Use
/// for run summaries where the table shows one row per node and "the last
/// attempt" is the right representative.
pub fn latest_stage_timing_by_node(
events: &[EventEnvelope],
) -> HashMap<String, fabro_types::StageTiming> {
let mut entries: Vec<(StageId, fabro_types::StageTiming)> =
extract_stage_timings_by_stage_id(events)
.into_iter()
.collect();
entries.sort_by_key(|(stage_id, _)| stage_id.visit());
let mut latest = HashMap::new();
for (stage_id, timing) in entries {
latest.insert(stage_id.node_id().to_string(), timing);
}
latest
}
#[cfg(test)]
mod duration_tests {
use chrono::{TimeZone, Utc};
use fabro_store::EventEnvelope;
use fabro_types::run_event::{StageCompletedProps, StageFailedProps};
use fabro_types::{EventBody, RunEvent, StageId, StageOutcome, StageTiming, fixtures};
use super::{
extract_stage_timings_by_stage_id, latest_stage_timing_by_node, total_stage_timing_by_node,
};
fn completed_event(seq: u32, node: &str, visit: u32, wall_time_ms: u64) -> EventEnvelope {
let event = RunEvent {
id: format!("evt_{seq}"),
ts: Utc.with_ymd_and_hms(2026, 1, 1, 0, 0, 0).unwrap(),
run_id: fixtures::RUN_1,
node_id: Some(node.to_string()),
node_label: None,
stage_id: Some(StageId::new(node, visit)),
parallel_group_id: None,
parallel_branch_id: None,
session_id: None,
parent_session_id: None,
tool_call_id: None,
actor: None,
body: EventBody::StageCompleted(StageCompletedProps {
index: 0,
timing: StageTiming::wall_only(wall_time_ms),
status: StageOutcome::Succeeded,
preferred_label: None,
suggested_next_ids: vec![],
usage_by_model: Vec::new(),
usage: None,
failure: None,
notes: None,
files_touched: vec![],
context_updates: None,
jump_to_node: None,
context_values: None,
node_visits: None,
loop_failure_signatures: None,
restart_failure_signatures: None,
response: None,
attempt: 1,
max_attempts: 1,
}),
};
EventEnvelope { seq, event }
}
fn failed_event(seq: u32, node: &str, visit: u32, wall_time_ms: u64) -> EventEnvelope {
let event = RunEvent {
id: format!("evt_{seq}"),
ts: Utc.with_ymd_and_hms(2026, 1, 1, 0, 0, 0).unwrap(),
run_id: fixtures::RUN_1,
node_id: Some(node.to_string()),
node_label: None,
stage_id: Some(StageId::new(node, visit)),
parallel_group_id: None,
parallel_branch_id: None,
session_id: None,
parent_session_id: None,
tool_call_id: None,
actor: None,
body: EventBody::StageFailed(StageFailedProps {
index: 0,
failure: None,
will_retry: true,
timing: StageTiming::wall_only(wall_time_ms),
usage_by_model: Vec::new(),
usage: None,
}),
};
EventEnvelope { seq, event }
}
#[test]
fn extract_keys_timings_by_full_stage_id() {
let events = vec![
completed_event(1, "verify", 1, 100),
completed_event(2, "verify", 2, 200),
];
let timings = extract_stage_timings_by_stage_id(&events);
assert_eq!(
timings
.get(&StageId::new("verify", 1))
.map(|t| t.wall_time_ms),
Some(100)
);
assert_eq!(
timings
.get(&StageId::new("verify", 2))
.map(|t| t.wall_time_ms),
Some(200)
);
}
#[test]
fn total_sums_wall_time_across_visits_per_node() {
let events = vec![
completed_event(1, "verify", 1, 100),
completed_event(2, "verify", 2, 200),
completed_event(3, "build", 1, 50),
];
let totals = total_stage_timing_by_node(&events);
assert_eq!(totals.get("verify").map(|t| t.wall_time_ms), Some(300));
assert_eq!(totals.get("build").map(|t| t.wall_time_ms), Some(50));
}
#[test]
fn latest_picks_highest_visit_regardless_of_input_order() {
// Visit 2 appears in the events vector before visit 1; the result
// must still reflect visit 2's timing (the latest visit).
let events = vec![
completed_event(1, "verify", 2, 999),
completed_event(2, "verify", 1, 100),
];
let latest = latest_stage_timing_by_node(&events);
assert_eq!(latest.get("verify").map(|t| t.wall_time_ms), Some(999));
}
#[test]
fn stage_failed_timings_are_included() {
let events = vec![failed_event(1, "verify", 1, 75)];
let timings = extract_stage_timings_by_stage_id(&events);
assert_eq!(
timings
.get(&StageId::new("verify", 1))
.map(|t| t.wall_time_ms),
Some(75)
);
}
#[test]
fn total_sums_active_breakdown_across_visits() {
// Same node visited twice with different inference/tool breakdowns:
// the rollup must add inference, tool, and active fields, not just
// wall time. This guards against accidentally summing wall only.
fn timed_completed(seq: u32, visit: u32, timing: StageTiming) -> EventEnvelope {
let event = RunEvent {
id: format!("evt_{seq}"),
ts: Utc.with_ymd_and_hms(2026, 1, 1, 0, 0, 0).unwrap(),
run_id: fixtures::RUN_1,
node_id: Some("agent".to_string()),
node_label: None,
stage_id: Some(StageId::new("agent", visit)),
parallel_group_id: None,
parallel_branch_id: None,
session_id: None,
parent_session_id: None,
tool_call_id: None,
actor: None,
body: EventBody::StageCompleted(StageCompletedProps {
index: 0,
timing,
status: StageOutcome::Succeeded,
preferred_label: None,
suggested_next_ids: vec![],
usage_by_model: Vec::new(),
usage: None,
failure: None,
notes: None,
files_touched: vec![],
context_updates: None,
jump_to_node: None,
context_values: None,
node_visits: None,
loop_failure_signatures: None,
restart_failure_signatures: None,
response: None,
attempt: 1,
max_attempts: 1,
}),
};
EventEnvelope { seq, event }
}
let events = vec![
timed_completed(1, 1, StageTiming::new(1000, 600, 300)),
timed_completed(2, 2, StageTiming::new(700, 400, 200)),
];
let totals = total_stage_timing_by_node(&events);
let agent = totals.get("agent").copied().unwrap();
assert_eq!(agent.wall_time_ms, 1700);
assert_eq!(agent.inference_time_ms, 1000);
assert_eq!(agent.tool_time_ms, 500);
assert_eq!(agent.active_time_ms, 1500);
}
}
#[doc(hidden)]
pub mod agent_memory;
pub mod artifact;
pub mod artifact_snapshot;
pub mod artifact_upload;
pub mod command_log;
pub(crate) mod condition;
pub mod context;
pub mod error;
pub mod event;
pub mod file_resolver;
pub mod git;
pub(crate) mod git_bridge;
pub mod git_identity;
pub(crate) mod graph;
pub mod handler;
mod hook_context;
mod interview_runtime;
#[allow(
dead_code,
reason = "The lifecycle module remains crate-visible for tests and pending integrations."
)]
pub(crate) mod lifecycle;
pub mod model_fallback;
pub(crate) mod node_handler;
pub mod operations;
pub mod outcome;
pub mod pipeline;
pub mod pull_request;
pub mod records;
mod retry;
pub mod run_control;
pub(crate) mod run_dir;
pub mod run_lookup;
pub mod usage_rollup;
pub use error::{Error, FailureCategory, FailureSignature, FailureSignatureExt, Result};
pub use fabro_types::ManifestPath;
pub use steering_hub::{PairControlError, SteeringHub};
pub use usage_rollup::{
ProjectionUsageByModel, ProjectionUsageRollup, ProjectionUsageStage,
usage_rollup_from_projection,
};
pub mod run_materialization;
pub mod run_options;
pub mod run_status;
pub mod runtime_store;
pub mod sandbox_git;
pub(crate) mod sandbox_git_runtime;
pub mod services;
pub(crate) mod stage_execution;
mod stage_scope;
pub mod steering_hub;
#[cfg(any(test, feature = "test-support"))]
pub mod test_support;
#[doc(hidden)]
pub mod transforms;
pub mod web_search;
pub mod workflow_bundle;