From a7ed03e1753ac919425e19d98fc36d54b6099344 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Sat, 5 Sep 2026 11:44:23 -0400 Subject: [PATCH] Rebuild conclusion stage summaries from run events --- lib/components/fabro-store/src/run_state.rs | 189 ++++++++++++++- .../fabro-workflow/src/billing_rollup.rs | 129 +---------- .../fabro-workflow/src/pipeline/finalize.rs | 181 ++++----------- .../fabro-types/src/billing_rollup.rs | 217 ++++++++++++++++++ lib/foundation/fabro-types/src/lib.rs | 1 + 5 files changed, 442 insertions(+), 275 deletions(-) create mode 100644 lib/foundation/fabro-types/src/billing_rollup.rs diff --git a/lib/components/fabro-store/src/run_state.rs b/lib/components/fabro-store/src/run_state.rs index a6a4a12a7..b9a9fdacf 100644 --- a/lib/components/fabro-store/src/run_state.rs +++ b/lib/components/fabro-store/src/run_state.rs @@ -19,7 +19,7 @@ use fabro_types::{ RunStatus, RunTimestamps, SandboxProviderKind, StageCompletion, StageHandler, StageId, StageInferenceProjection, StageModelUsage, StageOutcome, StageProjection, StageState, StartRecord, SubAgentProjection, SubAgentStatus, TodoListKind, TodoListProjection, - TodoProjection, WorkflowRef, first_event_seq, timing, + TodoProjection, WorkflowRef, billing_rollup, first_event_seq, timing, }; use fabro_util::error::render_compact_with_causes; @@ -247,7 +247,7 @@ impl RunProjectionReducer for RunProjection { ts, )?; self.pending_control = None; - self.conclusion = Some(conclusion_from_completed(props, ts)?); + self.conclusion = Some(conclusion_from_completed(self, props, ts)?); self.pending_interviews.clear(); } EventBody::RunFailed(props) => { @@ -258,7 +258,7 @@ impl RunProjectionReducer for RunProjection { ts, )?; self.pending_control = None; - self.conclusion = Some(conclusion_from_failed(props, ts)); + self.conclusion = Some(conclusion_from_failed(self, props, ts)); self.pending_interviews.clear(); finalize_unfinished_stages_after_run_failed(self, props, ts); } @@ -1560,9 +1560,12 @@ fn diff_from_checkpoint_props(props: &CheckpointCompletedProps) -> RunDiff { } fn conclusion_from_completed( + projection: &RunProjection, props: &RunCompletedProps, timestamp: DateTime, ) -> Result { + let (stages, total_retries) = billing_rollup::billing_rollup_from_projection(projection, None) + .conclusion_stages(projection); Ok(Conclusion { timestamp, status: StageOutcome::from_str(&props.status) @@ -1570,9 +1573,9 @@ fn conclusion_from_completed( timing: props.timing, failure: None, final_git_commit_sha: props.final_git_commit_sha.clone(), - stages: Vec::new(), + stages, billing: props.billing.clone(), - total_retries: 0, + total_retries, diff: RunDiff { patch: props.final_patch.clone(), summary: props.diff_summary, @@ -1580,7 +1583,13 @@ fn conclusion_from_completed( }) } -fn conclusion_from_failed(props: &RunFailedProps, timestamp: DateTime) -> Conclusion { +fn conclusion_from_failed( + projection: &RunProjection, + props: &RunFailedProps, + timestamp: DateTime, +) -> Conclusion { + let (stages, total_retries) = billing_rollup::billing_rollup_from_projection(projection, None) + .conclusion_stages(projection); Conclusion { timestamp, status: StageOutcome::Failed { @@ -1589,9 +1598,9 @@ fn conclusion_from_failed(props: &RunFailedProps, timestamp: DateTime) -> C timing: props.timing, failure: Some(props.failure.clone()), final_git_commit_sha: props.final_git_commit_sha.clone(), - stages: Vec::new(), + stages, billing: props.billing.clone(), - total_retries: 0, + total_retries, diff: RunDiff { patch: props.final_patch.clone(), summary: props.diff_summary, @@ -4416,6 +4425,170 @@ mod tests { ); } + #[test] + fn terminal_conclusion_replays_stage_summaries_without_metadata() { + for terminal_name in ["run.completed", "run.failed"] { + let mut settings = WorkflowSettings::default(); + settings.run.meta_branch.enabled = false; + let mut events = vec![ + test_raw_event( + 1, + "run.created", + &json!({ + "settings": settings, + "graph": { "name": "test", "nodes": {}, "edges": [], "attrs": {} }, + "labels": {}, + "provenance": test_support::test_run_provenance() + }), + None, + ), + test_raw_event( + 2, + "run.runnable", + &json!({ "source": "start_requested" }), + None, + ), + test_raw_event(3, "run.starting", &json!({}), None), + test_raw_event(4, "run.running", &json!({}), None), + ]; + // Two executions of zebra share one conclusion row. First-event + // order differs from checkpoint order, and skipped has no completion. + for (seq, node, visit, millis, tokens) in [ + (5, "zebra", 1, 1200, 100), + (6, "apple", 1, 300, 20), + (7, "zebra", 2, 800, 200), + ] { + let mut props = completed_props(millis, StageOutcome::Succeeded); + props.billing = Some(test_usage("test-model", tokens, 10)); + events.push(test_stage_event( + seq, + EventBody::StageCompleted(props), + StageId::new(node, visit), + )); + } + events.push(test_raw_event(8, "checkpoint.completed", &json!({ + "status": "succeeded", + "current_node": "zebra", + "completed_nodes": ["apple", "zebra", "zebra"], + "node_retries": { "zebra": 3, "apple": 1 }, + "node_outcomes": { + "apple": Outcome::>::success(), + "zebra": Outcome::>::success(), + "skipped": Outcome::>::skipped("condition was false") + }, + "context_values": {}, + "node_visits": { "zebra": 2, "apple": 1, "skipped": 1 }, + "git_commit_sha": "checkpoint-sha" + }), Some("zebra"))); + let terminal_billing = usage_counts(&test_usage("test-model", 320, 30)); + let terminal_props = if terminal_name == "run.completed" { + json!({ + "status": "succeeded", "reason": "completed", + "timing": fabro_types::RunTiming::wall_only(9000), + "artifact_count": 0, "billing": terminal_billing, + "final_git_commit_sha": "final-sha", "final_patch": "final patch" + }) + } else { + let mut props = run_failed_props(FailureReason::WorkflowError); + props.timing = fabro_types::RunTiming::wall_only(9000); + props.billing = Some(terminal_billing.clone()); + props.final_git_commit_sha = Some("final-sha".to_string()); + props.final_patch = Some("final patch".to_string()); + serde_json::to_value(props).unwrap() + }; + events.push(test_raw_event(9, terminal_name, &terminal_props, None)); + for event in &mut events { + event.event.ts = test_dt("2026-04-07T12:00:00Z") + + chrono::Duration::seconds(i64::from(event.seq)); + } + + // Cross the persisted wire boundary before both incremental and full replay. + let events: Vec = + serde_json::from_slice(&serde_json::to_vec(&events).unwrap()).unwrap(); + let mut live = RunProjection::apply_events(&events[..1]).unwrap(); + for event in &events[1..] { + live.apply_event(event).unwrap(); + } + let replayed = RunProjection::apply_events(&events).unwrap(); + let conclusion = replayed.conclusion.as_ref().unwrap(); + assert_eq!( + serde_json::to_value(&live.conclusion).unwrap(), + serde_json::to_value(conclusion).unwrap(), + ); + assert_eq!(conclusion.timestamp, events.last().unwrap().event.ts); + assert_eq!(conclusion.timing.wall_time_ms, 9000); + assert_eq!(conclusion.billing, Some(terminal_billing)); + assert_eq!( + conclusion.final_git_commit_sha.as_deref(), + Some("final-sha") + ); + assert_eq!(conclusion.diff.patch.as_deref(), Some("final patch")); + assert_eq!(conclusion.failure.is_some(), terminal_name == "run.failed"); + insta::allow_duplicates! { + insta::assert_snapshot!(serde_json::to_string_pretty(&json!({ + "stages": conclusion.stages, + "total_retries": conclusion.total_retries, + })).unwrap(), @r###" + { + "stages": [ + { + "stage_id": "zebra", + "stage_label": "zebra", + "timing": { + "wall_time_ms": 2000, + "inference_time_ms": 0, + "tool_time_ms": 0, + "active_time_ms": 0 + }, + "billing_usd_micros": 320, + "retries": 2 + }, + { + "stage_id": "apple", + "stage_label": "apple", + "timing": { + "wall_time_ms": 300, + "inference_time_ms": 0, + "tool_time_ms": 0, + "active_time_ms": 0 + }, + "billing_usd_micros": 30, + "retries": 0 + }, + { + "stage_id": "skipped", + "stage_label": "skipped", + "timing": { + "wall_time_ms": 0, + "inference_time_ms": 0, + "tool_time_ms": 0, + "active_time_ms": 0 + }, + "retries": 0 + } + ], + "total_retries": 2 + } + "###); + } + } + } + + #[test] + fn terminal_conclusion_without_checkpoint_has_no_stage_summaries() { + let mut state = running_projection(); + let terminal = test_event( + 4, + EventBody::RunFailed(run_failed_props(FailureReason::WorkflowError)), + None, + ); + state.apply_event(&terminal).unwrap(); + let conclusion = state.conclusion.unwrap(); + assert!(conclusion.stages.is_empty()); + assert_eq!(conclusion.total_retries, 0); + assert_eq!(conclusion.timestamp, terminal.event.ts); + } + #[test] fn run_failed_with_final_patch_populates_projection() { let mut state = running_projection(); diff --git a/lib/components/fabro-workflow/src/billing_rollup.rs b/lib/components/fabro-workflow/src/billing_rollup.rs index 1b6487f36..46056eace 100644 --- a/lib/components/fabro-workflow/src/billing_rollup.rs +++ b/lib/components/fabro-workflow/src/billing_rollup.rs @@ -1,128 +1,7 @@ -use std::collections::HashMap; - -use fabro_model::Catalog; -use fabro_types::{BilledTokenCounts, ModelRef, RunProjection, RunTiming, StageTiming}; - -#[derive(Debug, Clone, PartialEq)] -pub struct ProjectionBillingStage { - pub node_id: String, - pub billing: BilledTokenCounts, - /// Per-node timing summed across every visit of that node within this - /// projection. `wall_time_ms`, `inference_time_ms`, `tool_time_ms`, and - /// `active_time_ms` are all summed in lockstep. - pub timing: StageTiming, - pub model: Option, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProjectionBillingByModel { - pub model: ModelRef, - pub stages: i64, - pub billing: BilledTokenCounts, -} - -#[derive(Debug, Clone, Default, PartialEq)] -pub struct ProjectionBillingRollup { - pub stages: Vec, - pub totals: BilledTokenCounts, - pub by_model: Vec, - /// Run-level timing summed across every stage visit. `wall_time_ms` is - /// the sum of stage visit wall times (not the run clock duration). - pub timing: RunTiming, - pub billed_visit_count: usize, -} - -impl ProjectionBillingRollup { - #[must_use] - pub fn billing_if_present(&self) -> Option { - (self.billed_visit_count > 0).then(|| self.totals.clone()) - } -} - -#[must_use] -pub fn billing_rollup_from_projection( - projection: &RunProjection, - catalog: Option<&Catalog>, -) -> ProjectionBillingRollup { - let mut stage_indices = HashMap::::new(); - let mut stages = Vec::::new(); - let mut by_model = HashMap::::new(); - let mut totals = BilledTokenCounts::default(); - let mut run_timing = RunTiming::default(); - let mut billed_visit_count = 0_usize; - - for (stage_id, stage) in projection.iter_stages() { - if projection.is_boundary_stage(stage_id.node_id()) { - continue; - } - let usage = stage.billed_usage(catalog); - let usage = usage.as_ref(); - if stage.completion.is_none() && stage.timing.is_none() && usage.is_zero() { - continue; - } - - let node_id = stage_id.node_id(); - let index = *stage_indices.entry(node_id.to_string()).or_insert_with(|| { - let index = stages.len(); - stages.push(ProjectionBillingStage { - node_id: node_id.to_string(), - billing: BilledTokenCounts::default(), - timing: StageTiming::default(), - model: None, - }); - index - }); - let row = &mut stages[index]; - - if let Some(timing) = stage.timing { - row.timing = row.timing.saturating_add(&timing); - run_timing = run_timing.saturating_add(&RunTiming::from(timing)); - } - - if !usage.is_zero() { - billed_visit_count += 1; - row.billing.add_counts(usage); - totals.add_counts(usage); - - if let Some(model) = &stage.model { - row.model = Some(model.clone()); - let model_entry = - by_model - .entry(model.clone()) - .or_insert_with(|| ProjectionBillingByModel { - model: model.clone(), - stages: 0, - billing: BilledTokenCounts::default(), - }); - model_entry.stages += 1; - model_entry.billing.add_counts(usage); - } - } - } - - let mut by_model = by_model.into_values().collect::>(); - by_model.sort_by(|left, right| { - let left_provider = left.model.provider.to_string(); - let right_provider = right.model.provider.to_string(); - left_provider - .cmp(&right_provider) - .then_with(|| left.model.model_id.cmp(&right.model.model_id)) - .then_with(|| { - left.model - .speed - .map(<&'static str>::from) - .cmp(&right.model.speed.map(<&'static str>::from)) - }) - }); - - ProjectionBillingRollup { - stages, - totals, - by_model, - timing: run_timing, - billed_visit_count, - } -} +pub use fabro_types::billing_rollup::{ + ProjectionBillingByModel, ProjectionBillingRollup, ProjectionBillingStage, + billing_rollup_from_projection, +}; #[cfg(test)] mod tests { diff --git a/lib/components/fabro-workflow/src/pipeline/finalize.rs b/lib/components/fabro-workflow/src/pipeline/finalize.rs index 58f1b6930..a91ebbd6e 100644 --- a/lib/components/fabro-workflow/src/pipeline/finalize.rs +++ b/lib/components/fabro-workflow/src/pipeline/finalize.rs @@ -1,4 +1,3 @@ -use std::collections::HashMap; use std::sync::Arc; use std::time::Instant; @@ -10,17 +9,17 @@ use fabro_util::error::collect_causes; use fabro_util::time::elapsed_ms; use super::types::{Concluded, Executed, FinalizeOptions, Finalized, PublishOutcome, Published}; +use crate::billing_rollup; use crate::error::{Error, run_failure_from_error, run_failure_from_outcome_failure}; use crate::event::{Event, RunNoticeCode, RunNoticeLevel}; use crate::outcome::{Outcome, StageOutcome}; -use crate::records::{Checkpoint, Conclusion, StageSummary}; +use crate::records::Conclusion; use crate::run_metadata::{MetadataSnapshot, metadata_push_failure_is_transient}; use crate::run_options::RunOptions; use crate::run_status::{FailureReason, RunStatus, SuccessReason}; use crate::runtime_store::RunStoreHandle; use crate::sandbox_git::{git_diff_with_timeout, list_diff_numstat, summarize_diff_numstat}; use crate::services::RunServices; -use crate::{ProjectionBillingRollup, billing_rollup_from_projection}; pub fn classify_engine_result( engine_result: &Result, @@ -65,22 +64,8 @@ pub(crate) async fn build_conclusion_from_store( final_git_commit_sha: Option, ) -> Conclusion { let projection = run_store.state().await.ok(); - let projection_order = projection - .as_ref() - .map(stage_projection_order) - .unwrap_or_default(); - let projection_billing = projection - .as_ref() - .map(|projection| billing_rollup_from_projection(projection, None)) - .unwrap_or_default(); - let checkpoint = projection - .as_ref() - .and_then(|state| state.current_checkpoint()); - - build_conclusion_from_parts( - checkpoint, - &projection_billing, - &projection_order, + build_conclusion_from_projection( + projection.as_ref(), status, failure, run_wall_time_ms, @@ -88,113 +73,32 @@ pub(crate) async fn build_conclusion_from_store( ) } -fn build_conclusion_from_parts( - checkpoint: Option<&Checkpoint>, - projection_billing: &ProjectionBillingRollup, - projection_order: &HashMap, +fn build_conclusion_from_projection( + projection: Option<&RunProjection>, status: StageOutcome, failure: Option, run_wall_time_ms: u64, final_git_commit_sha: Option, ) -> Conclusion { - // Looping workflows revisit nodes; `completed_nodes` accumulates duplicates - // while the other checkpoint maps are keyed by node_id. Dedupe to one row - // per node so the stages table matches the deduped billing total. - let (stages, total_retries) = if let Some(cp) = checkpoint { - let billing_by_node = projection_billing - .stages - .iter() - .map(|stage| (stage.node_id.as_str(), stage)) - .collect::>(); - let mut stage_rows = Vec::new(); - let mut seen = std::collections::HashSet::new(); - let mut retries_sum: u32 = 0; - let mut stage_order = Vec::new(); - - for (original_checkpoint_order, node_id) in cp.completed_nodes.iter().enumerate() { - if !seen.insert(node_id.as_str()) { - continue; - } - stage_order.push((original_checkpoint_order, node_id.as_str())); - } - let mut extra_node_outcomes = cp - .node_outcomes - .keys() - .filter(|node_id| !seen.contains(node_id.as_str())) - .map(String::as_str) - .collect::>(); - extra_node_outcomes.sort_unstable(); - let extra_offset = stage_order.len(); - for (extra_index, node_id) in extra_node_outcomes.into_iter().enumerate() { - seen.insert(node_id); - stage_order.push((extra_offset + extra_index, node_id)); - } - - for (original_checkpoint_order, node_id) in stage_order { - let retries = cp - .node_retries - .get(node_id) - .copied() - .unwrap_or(1) - .saturating_sub(1); - retries_sum += retries; - let billing = billing_by_node.get(node_id); - - let summary = StageSummary { - stage_id: node_id.to_string(), - stage_label: node_id.to_string(), - timing: billing - .map_or_else(fabro_types::StageTiming::default, |stage| stage.timing), - billing_usd_micros: billing.and_then(|stage| stage.billing.total_usd_micros), - retries, - }; - stage_rows.push(( - projection_order.get(node_id).copied().unwrap_or(u32::MAX), - original_checkpoint_order, - summary, - )); - } - stage_rows.sort_by(|left, right| { - left.0 - .cmp(&right.0) - .then_with(|| left.1.cmp(&right.1)) - .then_with(|| left.2.stage_id.cmp(&right.2.stage_id)) - }); - let stages = stage_rows - .into_iter() - .map(|(_, _, summary)| summary) - .collect(); - (stages, retries_sum) - } else { - (vec![], 0) - }; - + let billing = projection + .map(|projection| billing_rollup::billing_rollup_from_projection(projection, None)) + .unwrap_or_default(); + let (stages, total_retries) = projection + .map(|projection| billing.conclusion_stages(projection)) + .unwrap_or_default(); Conclusion { timestamp: chrono::Utc::now(), status, - timing: projection_billing.timing.with_wall_time(run_wall_time_ms), + timing: billing.timing.with_wall_time(run_wall_time_ms), failure, final_git_commit_sha, stages, - billing: projection_billing.billing_if_present(), + billing: billing.billing_if_present(), total_retries, diff: fabro_types::RunDiff::default(), } } -fn stage_projection_order(state: &RunProjection) -> HashMap { - let mut order = HashMap::new(); - for (stage_id, stage) in state.iter_stages() { - order - .entry(stage_id.node_id().to_string()) - .and_modify(|first_seq: &mut u32| { - *first_seq = (*first_seq).min(stage.first_event_seq.get()); - }) - .or_insert_with(|| stage.first_event_seq.get()); - } - order -} - /// `conclusion` is injected because the terminal event hasn't been emitted /// yet — the run store's `projection.conclusion` is still `None` at this point. pub async fn write_finalize_commit( @@ -439,7 +343,7 @@ async fn compute_final_patch( #[cfg(any(test, feature = "test-support"))] pub(crate) fn billing_from_projection(projection: &RunProjection) -> Option { - billing_rollup_from_projection(projection, None).billing_if_present() + billing_rollup::billing_rollup_from_projection(projection, None).billing_if_present() } pub(crate) fn build_terminal_event( @@ -543,21 +447,8 @@ pub async fn conclude(executed: Executed, options: &FinalizeOptions) -> Result, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProjectionBillingByModel { + pub model: ModelRef, + pub stages: i64, + pub billing: BilledTokenCounts, +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub struct ProjectionBillingRollup { + pub stages: Vec, + pub totals: BilledTokenCounts, + pub by_model: Vec, + /// Run-level timing summed across every stage visit. `wall_time_ms` is + /// the sum of stage visit wall times (not the run clock duration). + pub timing: RunTiming, + pub billed_visit_count: usize, +} + +impl ProjectionBillingRollup { + #[must_use] + pub fn billing_if_present(&self) -> Option { + (self.billed_visit_count > 0).then(|| self.totals.clone()) + } + + /// Reconstruct the conclusion's per-node summaries from checkpoint and + /// stage events. Repeated visits share one row, ordered by the node's + /// first stage event. + #[must_use] + pub fn conclusion_stages(&self, projection: &RunProjection) -> (Vec, u32) { + let projection_order = stage_projection_order(projection); + // Looping workflows revisit nodes; `completed_nodes` accumulates duplicates + // while the other checkpoint maps are keyed by node_id. Dedupe to one row + // per node so the stages table matches the deduped billing total. + if let Some(cp) = projection.current_checkpoint() { + let billing_by_node = self + .stages + .iter() + .map(|stage| (stage.node_id.as_str(), stage)) + .collect::>(); + let mut stage_rows = Vec::new(); + let mut seen = std::collections::HashSet::new(); + let mut retries_sum: u32 = 0; + let mut stage_order = Vec::new(); + + for (original_checkpoint_order, node_id) in cp.completed_nodes.iter().enumerate() { + if !seen.insert(node_id.as_str()) { + continue; + } + stage_order.push((original_checkpoint_order, node_id.as_str())); + } + let mut extra_node_outcomes = cp + .node_outcomes + .keys() + .filter(|node_id| !seen.contains(node_id.as_str())) + .map(String::as_str) + .collect::>(); + extra_node_outcomes.sort_unstable(); + let extra_offset = stage_order.len(); + for (extra_index, node_id) in extra_node_outcomes.into_iter().enumerate() { + seen.insert(node_id); + stage_order.push((extra_offset + extra_index, node_id)); + } + + for (original_checkpoint_order, node_id) in stage_order { + let retries = cp + .node_retries + .get(node_id) + .copied() + .unwrap_or(1) + .saturating_sub(1); + retries_sum += retries; + let billing = billing_by_node.get(node_id); + + let summary = StageSummary { + stage_id: node_id.to_string(), + stage_label: node_id.to_string(), + timing: billing.map_or_else(StageTiming::default, |stage| stage.timing), + billing_usd_micros: billing.and_then(|stage| stage.billing.total_usd_micros), + retries, + }; + stage_rows.push(( + projection_order.get(node_id).copied().unwrap_or(u32::MAX), + original_checkpoint_order, + summary, + )); + } + stage_rows.sort_by(|left, right| { + left.0 + .cmp(&right.0) + .then_with(|| left.1.cmp(&right.1)) + .then_with(|| left.2.stage_id.cmp(&right.2.stage_id)) + }); + let stages = stage_rows + .into_iter() + .map(|(_, _, summary)| summary) + .collect(); + (stages, retries_sum) + } else { + (vec![], 0) + } + } +} + +#[must_use] +pub fn billing_rollup_from_projection( + projection: &RunProjection, + catalog: Option<&Catalog>, +) -> ProjectionBillingRollup { + let mut stage_indices = HashMap::::new(); + let mut stages = Vec::::new(); + let mut by_model = HashMap::::new(); + let mut totals = BilledTokenCounts::default(); + let mut run_timing = RunTiming::default(); + let mut billed_visit_count = 0_usize; + + for (stage_id, stage) in projection.iter_stages() { + if projection.is_boundary_stage(stage_id.node_id()) { + continue; + } + let usage = stage.billed_usage(catalog); + let usage = usage.as_ref(); + if stage.completion.is_none() && stage.timing.is_none() && usage.is_zero() { + continue; + } + + let node_id = stage_id.node_id(); + let index = *stage_indices.entry(node_id.to_string()).or_insert_with(|| { + let index = stages.len(); + stages.push(ProjectionBillingStage { + node_id: node_id.to_string(), + billing: BilledTokenCounts::default(), + timing: StageTiming::default(), + model: None, + }); + index + }); + let row = &mut stages[index]; + + if let Some(timing) = stage.timing { + row.timing = row.timing.saturating_add(&timing); + run_timing = run_timing.saturating_add(&RunTiming::from(timing)); + } + + if !usage.is_zero() { + billed_visit_count += 1; + row.billing.add_counts(usage); + totals.add_counts(usage); + + if let Some(model) = &stage.model { + row.model = Some(model.clone()); + let model_entry = + by_model + .entry(model.clone()) + .or_insert_with(|| ProjectionBillingByModel { + model: model.clone(), + stages: 0, + billing: BilledTokenCounts::default(), + }); + model_entry.stages += 1; + model_entry.billing.add_counts(usage); + } + } + } + + let mut by_model = by_model.into_values().collect::>(); + by_model.sort_by(|left, right| { + let left_provider = left.model.provider.to_string(); + let right_provider = right.model.provider.to_string(); + left_provider + .cmp(&right_provider) + .then_with(|| left.model.model_id.cmp(&right.model.model_id)) + .then_with(|| { + left.model + .speed + .map(<&'static str>::from) + .cmp(&right.model.speed.map(<&'static str>::from)) + }) + }); + + ProjectionBillingRollup { + stages, + totals, + by_model, + timing: run_timing, + billed_visit_count, + } +} + +fn stage_projection_order(state: &RunProjection) -> HashMap { + let mut order = HashMap::new(); + for (stage_id, stage) in state.iter_stages() { + order + .entry(stage_id.node_id().to_string()) + .and_modify(|first_seq: &mut u32| { + *first_seq = (*first_seq).min(stage.first_event_seq.get()); + }) + .or_insert_with(|| stage.first_event_seq.get()); + } + order +} diff --git a/lib/foundation/fabro-types/src/lib.rs b/lib/foundation/fabro-types/src/lib.rs index 3e15fc39f..e5639e4ee 100644 --- a/lib/foundation/fabro-types/src/lib.rs +++ b/lib/foundation/fabro-types/src/lib.rs @@ -3,6 +3,7 @@ extern crate self as fabro_types; pub mod artifact; pub mod auth; pub mod billing; +pub mod billing_rollup; pub mod blob_hash; pub mod blob_ref; pub mod checkpoint;