From 237fb4a12f6ca646a16a3a3f046a470b1df5feb8 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Mon, 4 May 2026 18:54:31 -0400 Subject: [PATCH] refactor(workflow): finish StageId migration in projection and durations Make iter_stages() and iter_stages_mut() yield in first_event_seq order so callers don't need to re-sort. Drops boilerplate sorts in the billing handler and the run-dump builder. Replace extract_stage_durations_from_events with two intent-explicit rollups built on extract_stage_durations_by_stage_id: - total_stage_duration_by_node (sum across visits) for billing/usage - latest_stage_duration_by_node (highest visit) for finalize/retro The old function silently picked an arbitrary visit's duration per node; the new helpers preserve all visit data and document the rollup choice at the call site. Co-Authored-By: Claude Opus 4.7 (1M context) --- lib/crates/fabro-dump/src/lib.rs | 7 +- lib/crates/fabro-server/src/server.rs | 4 +- .../src/server/handler/billing.rs | 11 +- lib/crates/fabro-types/src/run_projection.rs | 69 ++++++- lib/crates/fabro-workflow/src/lib.rs | 171 ++++++++++++++++-- .../fabro-workflow/src/pipeline/finalize.rs | 4 +- .../fabro-workflow/src/pipeline/retro.rs | 2 +- 7 files changed, 230 insertions(+), 38 deletions(-) diff --git a/lib/crates/fabro-dump/src/lib.rs b/lib/crates/fabro-dump/src/lib.rs index 69a670565..5b2da9403 100644 --- a/lib/crates/fabro-dump/src/lib.rs +++ b/lib/crates/fabro-dump/src/lib.rs @@ -66,18 +66,13 @@ impl RunDump { entries.push(RunDumpEntry::text("graph.fabro", graph_source.clone())); } - let mut stages: Vec<_> = state.iter_stages().collect(); + let stages: Vec<_> = state.iter_stages().collect(); if stages.len() > MAX_STAGES_IN_DUMP { bail!( "run dump supports at most {MAX_STAGES_IN_DUMP} stages with the current path prefix width (got {})", stages.len() ); } - stages.sort_by(|(left_id, left), (right_id, right)| { - left.first_event_seq - .cmp(&right.first_event_seq) - .then_with(|| left_id.cmp(right_id)) - }); let mut stage_ranks = HashMap::new(); for (index, (stage_id, _)) in stages.iter().enumerate() { diff --git a/lib/crates/fabro-server/src/server.rs b/lib/crates/fabro-server/src/server.rs index eb35ef8b0..9fdee7e5f 100644 --- a/lib/crates/fabro-server/src/server.rs +++ b/lib/crates/fabro-server/src/server.rs @@ -2788,7 +2788,7 @@ async fn execute_run_in_process(state: Arc, run_id: RunId) { // Accumulate aggregate usage after execution completes. if let Some(ref cp) = checkpoint { let stage_durations = match run_store.list_events().await { - Ok(events) => fabro_workflow::extract_stage_durations_from_events(&events), + Ok(events) => fabro_workflow::total_stage_duration_by_node(&events), Err(err) => { tracing::warn!(run_id = %run_id, error = %err, "Failed to load run events from store"); HashMap::default() @@ -3105,7 +3105,7 @@ async fn execute_run_subprocess(state: Arc, run_id: RunId) { if let Some(ref checkpoint) = final_state.checkpoint { let stage_durations = match run_store.list_events().await { - Ok(events) => fabro_workflow::extract_stage_durations_from_events(&events), + Ok(events) => fabro_workflow::total_stage_duration_by_node(&events), Err(err) => { tracing::warn!(run_id = %run_id, error = %err, "Failed to load run events from store"); HashMap::default() diff --git a/lib/crates/fabro-server/src/server/handler/billing.rs b/lib/crates/fabro-server/src/server/handler/billing.rs index d3d56b3a2..7861ba824 100644 --- a/lib/crates/fabro-server/src/server/handler/billing.rs +++ b/lib/crates/fabro-server/src/server/handler/billing.rs @@ -2,7 +2,7 @@ use std::num::NonZeroU32; use std::sync::Arc; use fabro_store::RunProjectionReducer; -use fabro_types::{EventBody, RunProjection, StageId, StageProjection}; +use fabro_types::{EventBody, RunProjection, StageId}; use super::super::{ ApiError, AppState, BilledTokenCounts, BillingByModel, BillingStageRef, EventEnvelope, HashMap, @@ -82,11 +82,8 @@ async fn list_run_stages( let stage_durations = fabro_workflow::extract_stage_durations_by_stage_id(&events); let lifecycle_states = latest_stage_states(&events); - let mut entries: Vec<(&StageId, &StageProjection)> = projection.iter_stages().collect(); - entries.sort_by_key(|(_, stage)| stage.first_event_seq); - - let mut stages = Vec::with_capacity(entries.len()); - for (stage_id, stage_projection) in entries { + let mut stages = Vec::new(); + for (stage_id, stage_projection) in projection.iter_stages() { let node_id = stage_id.node_id().to_string(); let Some(visit) = NonZeroU32::new(stage_id.visit()) else { tracing::warn!( @@ -156,7 +153,7 @@ async fn get_run_billing( }; let stage_durations = match run_store.list_events().await { - Ok(events) => fabro_workflow::extract_stage_durations_from_events(&events), + Ok(events) => fabro_workflow::total_stage_duration_by_node(&events), Err(err) => { return ApiError::new(StatusCode::INTERNAL_SERVER_ERROR, err.to_string()) .into_response(); diff --git a/lib/crates/fabro-types/src/run_projection.rs b/lib/crates/fabro-types/src/run_projection.rs index de2e4a1a5..fa1aa74bf 100644 --- a/lib/crates/fabro-types/src/run_projection.rs +++ b/lib/crates/fabro-types/src/run_projection.rs @@ -99,12 +99,22 @@ impl RunProjection { self.stages.get(stage) } + /// Iterate stages in `first_event_seq` order (the chronological order in + /// which each stage's first lifecycle event was recorded). Internal + /// storage is a `HashMap`, so iteration would otherwise be + /// non-deterministic; every caller wants chronological order, so we sort + /// here once instead of asking each caller to remember. pub fn iter_stages(&self) -> impl Iterator { - self.stages.iter() + let mut entries: Vec<(&StageId, &StageProjection)> = self.stages.iter().collect(); + entries.sort_by_key(|(_, stage)| stage.first_event_seq); + entries.into_iter() } + /// Mutable counterpart of [`iter_stages`]. Same chronological ordering. pub fn iter_stages_mut(&mut self) -> impl Iterator { - self.stages.iter_mut() + let mut entries: Vec<(&StageId, &mut StageProjection)> = self.stages.iter_mut().collect(); + entries.sort_by_key(|(_, stage)| stage.first_event_seq); + entries.into_iter() } pub fn is_empty(&self) -> bool { @@ -186,3 +196,58 @@ impl RunProjection { } } } + +#[cfg(test)] +mod iter_stages_tests { + use std::num::NonZeroU32; + + use super::RunProjection; + + fn seq(n: u32) -> NonZeroU32 { + NonZeroU32::new(n).unwrap() + } + + #[test] + fn iter_stages_yields_chronological_order_across_nodes() { + let mut p = RunProjection::default(); + // Insert in non-monotonic seq order to exercise the sort. + p.stage_entry("c", 1, seq(30)); + p.stage_entry("a", 1, seq(10)); + p.stage_entry("b", 1, seq(20)); + + let order: Vec<&str> = p + .iter_stages() + .map(|(stage_id, _)| stage_id.node_id()) + .collect(); + assert_eq!(order, vec!["a", "b", "c"]); + } + + #[test] + fn iter_stages_orders_visits_within_a_node() { + let mut p = RunProjection::default(); + // Visit 2 inserted first; visit 1's earlier first_event_seq must still + // win the chronological ordering. + p.stage_entry("verify", 2, seq(50)); + p.stage_entry("verify", 1, seq(20)); + + let visits: Vec = p + .iter_stages() + .map(|(stage_id, _)| stage_id.visit()) + .collect(); + assert_eq!(visits, vec![1, 2]); + } + + #[test] + fn iter_stages_mut_yields_chronological_order() { + let mut p = RunProjection::default(); + p.stage_entry("c", 1, seq(30)); + p.stage_entry("a", 1, seq(10)); + p.stage_entry("b", 1, seq(20)); + + let order: Vec = p + .iter_stages_mut() + .map(|(stage_id, _)| stage_id.node_id().to_string()) + .collect(); + assert_eq!(order, vec!["a", "b", "c"]); + } +} diff --git a/lib/crates/fabro-workflow/src/lib.rs b/lib/crates/fabro-workflow/src/lib.rs index 61289ec85..32de269e1 100644 --- a/lib/crates/fabro-workflow/src/lib.rs +++ b/lib/crates/fabro-workflow/src/lib.rs @@ -97,25 +97,12 @@ fn stage_completion_duration_ms(body: &EventBody) -> Option { } } -pub fn extract_stage_durations_from_events(events: &[EventEnvelope]) -> HashMap { - let mut durations = HashMap::new(); - for envelope in events { - let Some(duration_ms) = stage_completion_duration_ms(&envelope.event.body) else { - continue; - }; - let Some(node_id) = envelope.event.node_id.as_deref() else { - continue; - }; - durations.insert(node_id.to_string(), duration_ms); - } - durations -} - /// Extract per-stage (node_id, visit) durations from `stage.completed` / -/// `stage.failed` events. Differs from -/// [`extract_stage_durations_from_events`] by keying on the full -/// [`StageId`] instead of just `node_id`, so multi-visit -/// stages (e.g. a looped `verify` node) keep distinct durations. +/// `stage.failed` events. Keys on the full [`StageId`] so multi-visit stages +/// (e.g. a looped `verify` node) keep distinct durations. +/// +/// This is the canonical primitive; [`total_stage_duration_by_node`] and +/// [`latest_stage_duration_by_node`] are explicit rollups built on top of it. pub fn extract_stage_durations_by_stage_id(events: &[EventEnvelope]) -> HashMap { let mut durations = HashMap::new(); for envelope in events { @@ -130,6 +117,154 @@ pub fn extract_stage_durations_by_stage_id(events: &[EventEnvelope]) -> HashMap< durations } +/// Total duration spent in each node, summed across every visit. Use for +/// billing/usage where a retried node should count its full time. +pub fn total_stage_duration_by_node(events: &[EventEnvelope]) -> HashMap { + let mut totals: HashMap = HashMap::new(); + for (stage_id, duration_ms) in extract_stage_durations_by_stage_id(events) { + *totals.entry(stage_id.node_id().to_string()).or_default() += duration_ms; + } + totals +} + +/// Duration of each node's most recent visit (the highest visit number). Use +/// for run summaries and retros where the table shows one row per node and +/// "the last attempt" is the right representative. +pub fn latest_stage_duration_by_node(events: &[EventEnvelope]) -> HashMap { + let mut entries: Vec<(StageId, u64)> = extract_stage_durations_by_stage_id(events) + .into_iter() + .collect(); + entries.sort_by_key(|(stage_id, _)| stage_id.visit()); + let mut latest = HashMap::new(); + for (stage_id, duration_ms) in entries { + latest.insert(stage_id.node_id().to_string(), duration_ms); + } + 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, fixtures}; + + use super::{ + extract_stage_durations_by_stage_id, latest_stage_duration_by_node, + total_stage_duration_by_node, + }; + + fn completed_event(seq: u32, node: &str, visit: u32, duration_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, + duration_ms, + status: StageOutcome::Succeeded, + preferred_label: None, + suggested_next_ids: vec![], + billing: 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, duration_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, + duration_ms, + }), + }; + EventEnvelope { seq, event } + } + + #[test] + fn extract_keys_durations_by_full_stage_id() { + let events = vec![ + completed_event(1, "verify", 1, 100), + completed_event(2, "verify", 2, 200), + ]; + let durations = extract_stage_durations_by_stage_id(&events); + assert_eq!( + durations.get(&StageId::new("verify", 1)).copied(), + Some(100) + ); + assert_eq!( + durations.get(&StageId::new("verify", 2)).copied(), + Some(200) + ); + } + + #[test] + fn total_sums_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_duration_by_node(&events); + assert_eq!(totals.get("verify").copied(), Some(300)); + assert_eq!(totals.get("build").copied(), 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 duration (the latest visit). + let events = vec![ + completed_event(1, "verify", 2, 999), + completed_event(2, "verify", 1, 100), + ]; + let latest = latest_stage_duration_by_node(&events); + assert_eq!(latest.get("verify").copied(), Some(999)); + } + + #[test] + fn stage_failed_durations_are_included() { + let events = vec![failed_event(1, "verify", 1, 75)]; + let durations = extract_stage_durations_by_stage_id(&events); + assert_eq!(durations.get(&StageId::new("verify", 1)).copied(), Some(75)); + } +} + #[doc(hidden)] pub mod artifact; pub mod artifact_snapshot; diff --git a/lib/crates/fabro-workflow/src/pipeline/finalize.rs b/lib/crates/fabro-workflow/src/pipeline/finalize.rs index 921697369..997de0f4e 100644 --- a/lib/crates/fabro-workflow/src/pipeline/finalize.rs +++ b/lib/crates/fabro-workflow/src/pipeline/finalize.rs @@ -78,7 +78,7 @@ pub(crate) async fn build_conclusion_from_store( .as_ref() .and_then(|state| state.checkpoint.as_ref()); let stage_durations = events_result - .map(|events| crate::extract_stage_durations_from_events(&events)) + .map(|events| crate::latest_stage_duration_by_node(&events)) .unwrap_or_default(); build_conclusion_from_parts( @@ -503,7 +503,7 @@ pub async fn finalize(retroed: Retroed, options: &FinalizeOptions) -> Result Option { return None; } }; - let stage_durations = crate::extract_stage_durations_from_events(&events); + let stage_durations = crate::latest_stage_duration_by_node(&events); let mut retro = derive_retro( options.run_id, &options.workflow_name,