From 55b6fff829d616e777a93384797cbdd6fdfc3d39 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Mon, 4 May 2026 20:01:25 -0400 Subject: [PATCH] refactor(projection): capture per-visit stage billing Preserve per-visit `duration_ms` and model usage on `StageProjection` for completed and failed stages while leaving checkpoint payloads unchanged. Thread failed-stage billing through the workflow event conversion path so retry and terminal failure visits can retain token usage in the projection. --- lib/crates/fabro-server/src/server/tests.rs | 3 + lib/crates/fabro-store/src/run_state.rs | 162 +++++++++++++++++- .../tests/serializable_projection.rs | 30 +++- lib/crates/fabro-types/src/run_event/stage.rs | 2 + lib/crates/fabro-types/src/run_projection.rs | 12 +- lib/crates/fabro-workflow/src/error.rs | 1 + .../fabro-workflow/src/event/convert.rs | 29 +++- lib/crates/fabro-workflow/src/event/events.rs | 1 + lib/crates/fabro-workflow/src/lib.rs | 1 + .../fabro-workflow/src/lifecycle/event.rs | 2 + 10 files changed, 233 insertions(+), 10 deletions(-) diff --git a/lib/crates/fabro-server/src/server/tests.rs b/lib/crates/fabro-server/src/server/tests.rs index 093e0f858..d984d30bd 100644 --- a/lib/crates/fabro-server/src/server/tests.rs +++ b/lib/crates/fabro-server/src/server/tests.rs @@ -2231,6 +2231,7 @@ async fn list_run_stages_projects_retrying_until_completion() { failure: FailureDetail::new("try again", FailureCategory::TransientInfra), will_retry: true, duration_ms: 10, + billing: None, actor: None, }, ) @@ -2614,6 +2615,7 @@ async fn list_run_stages_shows_retrying_after_failed_event() { failure: FailureDetail::new("flake", FailureCategory::TransientInfra), will_retry: true, duration_ms: 5, + billing: None, actor: None, }, ) @@ -2693,6 +2695,7 @@ async fn list_run_stages_shows_retrying_when_failed_will_retry() { failure: FailureDetail::new("flake", FailureCategory::TransientInfra), will_retry: true, duration_ms: 5, + billing: None, actor: None, }, ) diff --git a/lib/crates/fabro-store/src/run_state.rs b/lib/crates/fabro-store/src/run_state.rs index 77b36796e..cd19243d4 100644 --- a/lib/crates/fabro-store/src/run_state.rs +++ b/lib/crates/fabro-store/src/run_state.rs @@ -320,6 +320,8 @@ impl RunProjectionReducer for RunProjection { let stage = self.stage_entry(node_id, visit, first_event_seq(event.seq)); stage.response = response; stage.completion = Some(completion); + stage.duration_ms = Some(props.duration_ms); + stage.usage.clone_from(&props.billing); } EventBody::StageFailed(props) => { let failure_reason = props.failure.as_ref().map(|detail| detail.message.clone()); @@ -334,6 +336,8 @@ impl RunProjectionReducer for RunProjection { failure_reason, timestamp: ts, }); + stage.duration_ms = Some(props.duration_ms); + stage.usage.clone_from(&props.billing); } EventBody::AgentSessionStarted(props) => { let Some(stage) = stage_at_visit(self, stored, props.visit, event.seq) else { @@ -613,12 +617,12 @@ mod tests { use fabro_types::run_event::run::RunFailedProps; use fabro_types::run_event::{ CheckpointCompletedProps, InterviewCompletedProps, InterviewOption, InterviewStartedProps, - RunControlEffectProps, StagePromptProps, StageStartedProps, + RunControlEffectProps, StageCompletedProps, StagePromptProps, StageStartedProps, }; use fabro_types::{ - BlockedReason, Checkpoint, EventBody, FailureReason, Outcome, QuestionType, RunBlobId, - RunControlAction, RunEvent, RunStatus, StageOutcome, SuccessReason, TerminalStatus, - WorkflowSettings, first_event_seq, fixtures, + BilledModelUsage, BlockedReason, Checkpoint, EventBody, FailureReason, Outcome, + QuestionType, RunBlobId, RunControlAction, RunEvent, RunStatus, StageOutcome, + SuccessReason, TerminalStatus, WorkflowSettings, first_event_seq, fixtures, }; use serde_json::json; @@ -651,6 +655,32 @@ mod tests { event } + fn test_usage(model_id: &str, input_tokens: i64, output_tokens: i64) -> BilledModelUsage { + serde_json::from_value(json!({ + "input": { + "usage": { + "model": { + "provider": "openai", + "model_id": model_id + }, + "tokens": { + "input_tokens": input_tokens, + "output_tokens": output_tokens + } + }, + "facts": { + "provider": "open_ai" + } + }, + "total_usd_micros": input_tokens + output_tokens + })) + .unwrap() + } + + fn usage_json(usage: &BilledModelUsage) -> serde_json::Value { + serde_json::to_value(usage).unwrap() + } + fn test_raw_event( seq: u32, event: &str, @@ -861,6 +891,130 @@ mod tests { assert_eq!(stage.prompt.as_deref(), Some("prompt")); } + #[test] + fn stage_completed_event_captures_duration_and_usage_per_visit() { + let mut state = RunProjection::default(); + let usage = test_usage("gpt-5.2", 123, 45); + + state + .apply_event(&test_event( + 3, + EventBody::StageCompleted(StageCompletedProps { + index: 0, + duration_ms: 789, + status: StageOutcome::Succeeded, + preferred_label: None, + suggested_next_ids: Vec::new(), + billing: Some(usage.clone()), + failure: None, + notes: None, + files_touched: Vec::new(), + context_updates: None, + jump_to_node: None, + context_values: None, + node_visits: None, + loop_failure_signatures: None, + restart_failure_signatures: None, + response: Some("done".to_string()), + attempt: 1, + max_attempts: 1, + }), + Some("build"), + )) + .unwrap(); + + let stage = state.stage(&StageId::new("build", 1)).unwrap(); + assert_eq!(stage.duration_ms, Some(789)); + assert_eq!(stage.usage.as_ref(), Some(&usage)); + } + + #[test] + fn stage_failed_event_captures_duration_and_usage_per_visit() { + let mut state = RunProjection::default(); + let stage_id = StageId::new("build", 1); + let usage = test_usage("gpt-5.2", 321, 54); + + state + .apply_event(&test_stage_event( + 2, + EventBody::StageStarted(StageStartedProps { + index: 0, + handler_type: "agent".to_string(), + attempt: 1, + max_attempts: 1, + }), + stage_id.clone(), + )) + .unwrap(); + state + .apply_event(&test_raw_event( + 3, + "stage.failed", + &json!({ + "index": 0, + "failure": { + "message": "provider failed", + "failure_class": "transient_infra" + }, + "will_retry": false, + "duration_ms": 654, + "billing": usage_json(&usage) + }), + Some("build"), + )) + .unwrap(); + + let stage = state.stage(&stage_id).unwrap(); + assert_eq!(stage.duration_ms, Some(654)); + assert_eq!(stage.usage.as_ref(), Some(&usage)); + } + + #[test] + fn two_visits_of_one_node_retain_distinct_usage() { + let mut state = RunProjection::default(); + let first_usage = test_usage("gpt-5.2", 100, 10); + let second_usage = test_usage("gpt-5.2", 200, 20); + + for (seq, visit, duration_ms, usage) in [ + (3, 1usize, 111, first_usage.clone()), + (4, 2usize, 222, second_usage.clone()), + ] { + state + .apply_event(&test_event( + seq, + EventBody::StageCompleted(StageCompletedProps { + index: 0, + duration_ms, + status: StageOutcome::Succeeded, + preferred_label: None, + suggested_next_ids: Vec::new(), + billing: Some(usage), + failure: None, + notes: None, + files_touched: Vec::new(), + context_updates: None, + jump_to_node: None, + context_values: None, + node_visits: Some(BTreeMap::from([("build".to_string(), visit)])), + loop_failure_signatures: None, + restart_failure_signatures: None, + response: None, + attempt: 1, + max_attempts: 1, + }), + Some("build"), + )) + .unwrap(); + } + + let first_stage = state.stage(&StageId::new("build", 1)).unwrap(); + let second_stage = state.stage(&StageId::new("build", 2)).unwrap(); + assert_eq!(first_stage.duration_ms, Some(111)); + assert_eq!(first_stage.usage.as_ref(), Some(&first_usage)); + assert_eq!(second_stage.duration_ms, Some(222)); + assert_eq!(second_stage.usage.as_ref(), Some(&second_usage)); + } + #[test] fn checkpoint_completed_creates_projection_entry_for_skipped_stage() { let mut state = RunProjection::default(); diff --git a/lib/crates/fabro-store/tests/serializable_projection.rs b/lib/crates/fabro-store/tests/serializable_projection.rs index 6d51065ba..2b24cd54e 100644 --- a/lib/crates/fabro-store/tests/serializable_projection.rs +++ b/lib/crates/fabro-store/tests/serializable_projection.rs @@ -5,8 +5,8 @@ use fabro_store::{RunProjection, SerializableProjection, StageId}; use fabro_types::graph::Graph; use fabro_types::run::RunSpec; use fabro_types::{ - Checkpoint, RunStatus, SandboxRecord, StageCompletion, StageOutcome, StartRecord, - TerminalStatus, WorkflowSettings, first_event_seq, fixtures, + BilledModelUsage, Checkpoint, RunStatus, SandboxRecord, StageCompletion, StageOutcome, + StartRecord, TerminalStatus, WorkflowSettings, first_event_seq, fixtures, }; use serde_json::json; @@ -52,6 +52,28 @@ fn sample_checkpoint() -> Checkpoint { } } +fn sample_usage() -> BilledModelUsage { + serde_json::from_value(json!({ + "input": { + "usage": { + "model": { + "provider": "openai", + "model_id": "gpt-5.2" + }, + "tokens": { + "input_tokens": 123, + "output_tokens": 45 + } + }, + "facts": { + "provider": "open_ai" + } + }, + "total_usd_micros": 168 + })) + .expect("sample usage should deserialize") +} + #[test] fn serializable_projection_round_trips_and_trims_bulky_node_fields() { let stage_id = StageId::new("build", 2); @@ -94,6 +116,8 @@ fn serializable_projection_round_trips_and_trims_bulky_node_fields() { stage.script_invocation = Some(json!({ "command": "cargo test" })); stage.script_timing = Some(json!({ "duration_ms": 10 })); stage.parallel_results = Some(json!([{ "stage": "fanout@1" }])); + stage.duration_ms = Some(1234); + stage.usage = Some(sample_usage()); stage.stdout = Some("stdout".to_string()); stage.stderr = Some("stderr".to_string()); @@ -138,6 +162,8 @@ fn serializable_projection_round_trips_and_trims_bulky_node_fields() { node.parallel_results, Some(json!([{ "stage": "fanout@1" }])) ); + assert_eq!(node.duration_ms, Some(1234)); + assert_eq!(node.usage, Some(sample_usage())); } #[test] diff --git a/lib/crates/fabro-types/src/run_event/stage.rs b/lib/crates/fabro-types/src/run_event/stage.rs index 9f1781d82..1b6609884 100644 --- a/lib/crates/fabro-types/src/run_event/stage.rs +++ b/lib/crates/fabro-types/src/run_event/stage.rs @@ -57,6 +57,8 @@ pub struct StageFailedProps { pub will_retry: bool, #[serde(default)] pub duration_ms: u64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub billing: Option, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] diff --git a/lib/crates/fabro-types/src/run_projection.rs b/lib/crates/fabro-types/src/run_projection.rs index fa1aa74bf..712ceaa7f 100644 --- a/lib/crates/fabro-types/src/run_projection.rs +++ b/lib/crates/fabro-types/src/run_projection.rs @@ -4,9 +4,9 @@ use std::num::NonZeroU32; use chrono::{DateTime, Utc}; use crate::{ - Checkpoint, Conclusion, InterviewQuestionRecord, InvalidTransition, PullRequestRecord, Retro, - RunControlAction, RunId, RunSpec, RunStatus, SandboxRecord, StageCompletion, StageId, - StartRecord, + BilledModelUsage, Checkpoint, Conclusion, InterviewQuestionRecord, InvalidTransition, + PullRequestRecord, Retro, RunControlAction, RunId, RunSpec, RunStatus, SandboxRecord, + StageCompletion, StageId, StartRecord, }; #[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)] @@ -44,6 +44,10 @@ pub struct StageProjection { pub prompt: Option, pub response: Option, pub completion: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub duration_ms: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub usage: Option, pub provider_used: Option, pub diff: Option, pub script_invocation: Option, @@ -78,6 +82,8 @@ impl StageProjection { prompt: None, response: None, completion: None, + duration_ms: None, + usage: None, provider_used: None, diff: None, script_invocation: None, diff --git a/lib/crates/fabro-workflow/src/error.rs b/lib/crates/fabro-workflow/src/error.rs index 1d08b2777..6a3ff0d06 100644 --- a/lib/crates/fabro-workflow/src/error.rs +++ b/lib/crates/fabro-workflow/src/error.rs @@ -1838,6 +1838,7 @@ mod tests { failure: failure.clone(), will_retry: false, duration_ms: 0, + billing: None, actor: None, }; diff --git a/lib/crates/fabro-workflow/src/event/convert.rs b/lib/crates/fabro-workflow/src/event/convert.rs index 815aaec11..4aa9df5b1 100644 --- a/lib/crates/fabro-workflow/src/event/convert.rs +++ b/lib/crates/fabro-workflow/src/event/convert.rs @@ -290,12 +290,14 @@ fn event_body_from_event(event: &Event) -> EventBody { failure, will_retry, duration_ms, + billing, .. } => EventBody::StageFailed(fabro_types::StageFailedProps { index: *index, failure: Some(failure.clone()), will_retry: *will_retry, duration_ms: *duration_ms, + billing: billing.clone(), }), Event::StageRetrying { index, @@ -1178,7 +1180,7 @@ mod tests { use crate::error::Error; use crate::event::test_support::user_principal; use crate::event::{Event, StageScope}; - use crate::outcome::FailureDetail; + use crate::outcome::{BilledModelUsage, FailureDetail}; #[derive(Debug)] struct EventTestCause; @@ -1200,6 +1202,28 @@ mod tests { } } + fn test_usage(model_id: &str, input_tokens: i64, output_tokens: i64) -> BilledModelUsage { + serde_json::from_value(serde_json::json!({ + "input": { + "usage": { + "model": { + "provider": "openai", + "model_id": model_id + }, + "tokens": { + "input_tokens": input_tokens, + "output_tokens": output_tokens + } + }, + "facts": { + "provider": "open_ai" + } + }, + "total_usd_micros": input_tokens + output_tokens + })) + .unwrap() + } + #[test] fn run_event_stage_completed_places_node_fields_in_header() { let stored = to_run_event_at( @@ -1279,6 +1303,7 @@ mod tests { #[test] fn run_event_stage_failure_keeps_failure_detail() { + let usage = test_usage("gpt-5.2", 321, 54); let stored = to_run_event(&fixtures::RUN_3, &Event::StageFailed { node_id: "code".to_string(), name: "Code".to_string(), @@ -1289,6 +1314,7 @@ mod tests { ), will_retry: true, duration_ms: 5000, + billing: Some(usage.clone()), actor: None, }); @@ -1297,6 +1323,7 @@ mod tests { assert_eq!(properties["failure"]["message"], "lint failed"); assert_eq!(properties["failure"]["failure_class"], "deterministic"); assert_eq!(properties["will_retry"], true); + assert_eq!(properties["billing"], serde_json::to_value(&usage).unwrap()); } #[test] diff --git a/lib/crates/fabro-workflow/src/event/events.rs b/lib/crates/fabro-workflow/src/event/events.rs index a3b2e2b5b..c753e20c6 100644 --- a/lib/crates/fabro-workflow/src/event/events.rs +++ b/lib/crates/fabro-workflow/src/event/events.rs @@ -206,6 +206,7 @@ pub enum Event { failure: FailureDetail, will_retry: bool, duration_ms: u64, + billing: Option, #[serde(default, skip_serializing_if = "Option::is_none")] actor: Option, }, diff --git a/lib/crates/fabro-workflow/src/lib.rs b/lib/crates/fabro-workflow/src/lib.rs index 32de269e1..fe814118d 100644 --- a/lib/crates/fabro-workflow/src/lib.rs +++ b/lib/crates/fabro-workflow/src/lib.rs @@ -211,6 +211,7 @@ mod duration_tests { failure: None, will_retry: true, duration_ms, + billing: None, }), }; EventEnvelope { seq, event } diff --git a/lib/crates/fabro-workflow/src/lifecycle/event.rs b/lib/crates/fabro-workflow/src/lifecycle/event.rs index 482f0311a..82dbd0501 100644 --- a/lib/crates/fabro-workflow/src/lifecycle/event.rs +++ b/lib/crates/fabro-workflow/src/lifecycle/event.rs @@ -224,6 +224,7 @@ impl RunLifecycle for EventLifecycle { failure, will_retry: true, duration_ms, + billing: outcome.usage.clone(), actor, }, &scope, @@ -275,6 +276,7 @@ impl RunLifecycle for EventLifecycle { failure, will_retry: false, duration_ms, + billing: outcome.usage.clone(), actor, }, &scope,