scriptoriumadmin/fabro
1
0
Fork 0
mirror of https://github.com/fabro-sh/fabro.git synced 2026-09-14 23:22:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

refactor(projection): capture per-visit stage billing

Browse source 
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.
This commit is contained in:
Bryan Helmkamp 2026-05-04 20:01:25 -04:00
parent 9b15a9cd3c
commit 55b6fff829
No known key found for this signature in database
10 changed files with 233 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
lib/crates/fabro-server/src/server/tests.rs
View file

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

162
lib/crates/fabro-store/src/run_state.rs
View file

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

30
lib/crates/fabro-store/tests/serializable_projection.rs
View file

@ -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]

2
lib/crates/fabro-types/src/run_event/stage.rs
View file

@ -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<BilledModelUsage>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]

12
lib/crates/fabro-types/src/run_projection.rs
View file

@ -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<String>,
pub response: Option<String>,
pub completion: Option<StageCompletion>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub duration_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub usage: Option<BilledModelUsage>,
pub provider_used: Option<serde_json::Value>,
pub diff: Option<String>,
pub script_invocation: Option<serde_json::Value>,
@ -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,

1
lib/crates/fabro-workflow/src/error.rs
View file

@ -1838,6 +1838,7 @@ mod tests {
failure: failure.clone(),
will_retry: false,
duration_ms: 0,
billing: None,
actor: None,
};

29
lib/crates/fabro-workflow/src/event/convert.rs
View file

@ -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]

1
lib/crates/fabro-workflow/src/event/events.rs
View file

@ -206,6 +206,7 @@ pub enum Event {
failure: FailureDetail,
will_retry: bool,
duration_ms: u64,
billing: Option<BilledModelUsage>,
#[serde(default, skip_serializing_if = "Option::is_none")]
actor: Option<Principal>,
},

1
lib/crates/fabro-workflow/src/lib.rs
View file

@ -211,6 +211,7 @@ mod duration_tests {
failure: None,
will_retry: true,
duration_ms,
billing: None,
}),
};
EventEnvelope { seq, event }

2
lib/crates/fabro-workflow/src/lifecycle/event.rs
View file

@ -224,6 +224,7 @@ impl RunLifecycle<WorkflowGraph> for EventLifecycle {
failure,
will_retry: true,
duration_ms,
billing: outcome.usage.clone(),
actor,
},
&scope,
@ -275,6 +276,7 @@ impl RunLifecycle<WorkflowGraph> for EventLifecycle {
failure,
will_retry: false,
duration_ms,
billing: outcome.usage.clone(),
actor,
},
&scope,