scriptoriumadmin/fabro
1
0
Fork 0
mirror of https://github.com/fabro-sh/fabro.git synced 2026-08-28 05:27:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

Merge remote-tracking branch 'origin/main'

Browse source
This commit is contained in:
Bryan Helmkamp 2026-05-09 15:40:18 -04:00
commit 237318f13f
No known key found for this signature in database
22 changed files with 783 additions and 176 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

37
apps/fabro-web/app/routes/run-billing.test.tsx
View file

@ -167,7 +167,7 @@ describe("RunBilling", () => {
expect(text).toContain("Stages will appear as soon as the run starts executing.");
});
test("renders an in-flight row with live runtime and includes its elapsed time in the footer", () => {
test("renders an in-flight row with live billing and includes its elapsed time in the footer", () => {
const originalNow = Date.now;
// Pin "now" to 30s after the in-flight row started.
const startedAt = "2026-04-29T12:00:00.000Z";
@ -180,19 +180,39 @@ describe("RunBilling", () => {
stages: [
{
stage: { id: "in-flight", name: "in-flight" },
model: null,
// Server reports 0 runtime / no billing; the row is still being executed.
billing: zeroBilling(),
model: { id: "claude-opus-4-6" },
billing: zeroBilling({
input_tokens: 1200,
output_tokens: 300,
total_tokens: 1500,
total_usd_micros: 240000,
}),
runtime_secs: 0,
started_at: startedAt,
state: "running",
},
],
// Server total is 0 because the in-flight row hasn't been finalized.
totals: {
runtime_secs: 0,
...zeroBilling(),
...zeroBilling({
input_tokens: 1200,
output_tokens: 300,
total_tokens: 1500,
total_usd_micros: 240000,
}),
},
by_model: [
{
model: { id: "claude-opus-4-6" },
stages: 1,
billing: zeroBilling({
input_tokens: 1200,
output_tokens: 300,
total_tokens: 1500,
total_usd_micros: 240000,
}),
},
],
}),
);
@ -201,6 +221,11 @@ describe("RunBilling", () => {
// first stage starts.
expect(text).not.toContain("No stages yet");
expect(text).toContain("in-flight");
expect(text).toContain("claude-opus-4-6");
expect(text).toContain("1.2k");
expect(text).toContain("0.3k");
expect(text).toContain("$0.24");
expect(text).toContain("By model");
// Both the row's runtime cell and the footer total should reflect
// ~30s elapsed since started_at.

30
docs/public/api-reference/fabro-api.yaml
View file

@ -6008,6 +6008,7 @@ components:
type: object
required:
- first_event_seq
- usage
properties:
first_event_seq:
type: integer
@ -6057,6 +6058,12 @@ components:
format: uint64
minimum: 0
description: Wall-clock duration of the stage's latest terminal attempt, if known.
usage:
$ref: "#/components/schemas/BilledTokenCounts"
model:
oneOf:
- $ref: "#/components/schemas/BillingModelRef"
- type: "null"
state:
oneOf:
- $ref: "#/components/schemas/StageState"
@ -6548,6 +6555,29 @@ components:
description: Billed USD amount in micros.
example: 720000
BillingModelRef:
description: Provider-qualified billing model identity used for cost estimates.
type: object
required:
- provider
- model_id
properties:
provider:
$ref: "#/components/schemas/Provider"
model_id:
type: string
speed:
oneOf:
- $ref: "#/components/schemas/BillingSpeed"
- type: "null"
BillingSpeed:
description: Optional provider-specific model speed tier used for cost estimates.
type: string
enum:
- standard
- fast
CodeLocation:
description: A file and line location in the codebase.
type: object

142
docs/superpowers/plans/2026-05-09-live-stage-token-projection.md
View file

@ -4,7 +4,7 @@
**Goal:** Make `StageProjection` the source of truth for current per-stage token usage so in-flight stages show running token counts on the billing page.
**Architecture:** Store zeroable token counters directly on each `StageProjection`, plus optional model identity once model usage is known. Projection application updates those counters from usage-bearing events as they arrive, while terminal stage events remain authoritative and replace live counts when final billing is present. Billing rollups and server responses read from the projection rather than re-scanning raw events.
**Architecture:** Store zeroable token counters directly on each `StageProjection`, plus optional `ModelRef` identity once model usage is known. Projection application updates those counters from usage-bearing events as they arrive, while terminal stage events replace live counts when they contain a later best-known billing snapshot. Billing rollups and server responses read from the projection rather than re-scanning raw events.
**Tech Stack:** Rust workspace crates `fabro-model`, `fabro-types`, `fabro-store`, `fabro-workflow`, `fabro-server`; OpenAPI-driven `fabro-api`; React route tests in `apps/fabro-web`.
@ -16,8 +16,8 @@
- Modify `lib/crates/fabro-types/src/run_projection.rs` for the new projection fields and attempt reset behavior.
- Modify `lib/crates/fabro-store/src/run_state.rs` so projection application updates live and terminal usage.
- Modify `lib/crates/fabro-workflow/src/billing_rollup.rs` and billing-related server code to read `StageProjection.usage`.
- Modify event conversion around `agent.message` to carry provider context and price live usage deltas when provider parsing succeeds.
- Modify `docs/public/api-reference/fabro-api.yaml` to expose `StageProjection.usage` and `StageProjection.model_id`.
- Modify event conversion around existing `agent.message` events to carry typed `ModelRef` billing identity instead of a string model id.
- Modify `docs/public/api-reference/fabro-api.yaml` to expose `StageProjection.usage` and `StageProjection.model`.
## Task 1: Add `BilledTokenCounts` Aggregation Methods
@ -89,20 +89,49 @@ Expected: all `fabro-model` billing tests pass.
#[serde(default)]
pub usage: BilledTokenCounts,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub model_id: Option<String>,
pub model: Option<ModelRef>,
```
- [ ] Remove the `#[serde(skip)] pub usage: Option<BilledModelUsage>` field.
- [ ] Initialize `usage` with `BilledTokenCounts::default()` and `model_id` with `None` in `StageProjection::new`.
- [ ] Initialize `usage` with `BilledTokenCounts::default()` and `model` with `None` in `StageProjection::new`.
- [ ] Confirm `StageProjection::begin_attempt` resets usage and model identity by relying on `*self = Self::new(self.first_event_seq)` before setting `started_at`, `handler`, and `state`.
- [ ] Update the OpenAPI `StageProjection` schema:
- add `usage: $ref: "#/components/schemas/BilledTokenCounts"`
- add `model_id: ["string", "null"]`
- add `model: oneOf [$ref: "#/components/schemas/BillingModelRef", null]`
- include `usage` in `required`
- [ ] Add an OpenAPI `BillingModelRef` schema for the existing Rust `fabro_types::ModelRef` shape:
```yaml
BillingModelRef:
description: Provider-qualified billing model identity used for cost estimates.
type: object
required:
- provider
- model_id
properties:
provider:
$ref: "#/components/schemas/Provider"
model_id:
type: string
speed:
oneOf:
- $ref: "#/components/schemas/BillingSpeed"
- type: "null"
BillingSpeed:
description: Optional provider-specific model speed tier used for cost estimates.
type: string
enum:
- standard
- fast
```
- [ ] Do not reuse the existing OpenAPI `ModelRef` schema name here; that schema is a string parser type for run settings, while this field is the billing `fabro_types::ModelRef`.
- [ ] Run:
```bash
@ -124,9 +153,7 @@ EventBody::AgentMessage(props) => {
return Ok(());
};
stage.usage.add_counts(&props.billing);
if !props.model.is_empty() {
stage.model_id = Some(props.model.clone());
}
stage.model = Some(props.model.clone());
}
```
@ -136,7 +163,7 @@ EventBody::AgentMessage(props) => {
stage.response = Some(props.response.clone());
if let Some(billing) = &props.billing {
stage.usage.replace_with_billed_usage(billing);
stage.model_id = Some(billing.model_id().to_string());
stage.model = Some(billing.model().clone());
}
```
@ -145,13 +172,13 @@ if let Some(billing) = &props.billing {
```rust
if let Some(billing) = &props.billing {
stage.usage.replace_with_billed_usage(billing);
stage.model_id = Some(billing.model_id().to_string());
stage.model = Some(billing.model().clone());
}
```
- [ ] Do not clear nonzero live usage on terminal events that have `billing: None`; this preserves best-known live usage for events that do not include final pricing.
- [ ] Update existing projection tests that assert `stage.usage.as_ref() == Some(...)` to assert flattened counters and `model_id`.
- [ ] Update existing projection tests that assert `stage.usage.as_ref() == Some(...)` to assert flattened counters and `model`.
- [ ] Add tests for live `agent.message` accumulation, terminal replacement, `stage.failed` replacement, and reset on a new `stage.started`.
@ -163,40 +190,70 @@ cargo nextest run -p fabro-store run_state
Expected: projection tests pass.
## Task 4: Keep Live Agent Usage Cost-Aware Where Possible
## Task 4: Carry Typed Model Identity on `agent.message`
**Files:**
- Modify: `lib/crates/fabro-agent/src/types.rs`
- Modify: `lib/crates/fabro-agent/src/session.rs`
- Modify: `lib/crates/fabro-workflow/src/event/convert.rs`
- Modify: `lib/crates/fabro-types/src/run_event/agent.rs`
- [ ] Add provider identity to `AgentEvent::AssistantMessage`:
- [ ] Change the existing `AgentEvent::AssistantMessage` payload to carry billing `ModelRef` instead of a string model id:
```rust
AssistantMessage {
text: String,
provider: String,
model: String,
model: ModelRef,
usage: TokenCounts,
tool_call_count: usize,
}
```
- [ ] Emit `provider: self.provider_profile.provider().to_string()` from `Session` when emitting `AgentEvent::AssistantMessage`.
- [ ] In workflow event conversion, price assistant message usage when provider parsing succeeds:
- [ ] Build the `ModelRef` at the source in `Session` when emitting `AgentEvent::AssistantMessage`:
```rust
let billing = provider.parse::<Provider>().ok().map_or_else(
|| billed_token_counts_from_llm(usage),
|provider| {
let billed = billed_model_usage_from_llm(model, provider, None, usage);
BilledTokenCounts::from_billed_usage(std::slice::from_ref(&billed))
let speed = self
.config
.speed
.as_deref()
.and_then(|value| value.parse::<Speed>().ok());
let model = ModelRef {
provider: self.provider_profile.provider(),
model_id: if response.model.is_empty() {
self.provider_profile.model().to_string()
} else {
response.model.clone()
},
);
speed,
};
```
- [ ] Keep the fallback to flattened token counts when provider parsing fails, so live tokens are never lost.
- [ ] Change `AgentMessageProps` to carry the same typed model identity:
```rust
pub struct AgentMessageProps {
pub text: String,
pub model: ModelRef,
pub billing: BilledTokenCounts,
pub tool_call_count: usize,
pub visit: u32,
}
```
- [ ] In workflow event conversion, price assistant message usage directly from the typed model:
```rust
let requested_speed = model.speed.map(<&'static str>::from);
let billed = billed_model_usage_from_llm(
&model.model_id,
model.provider,
requested_speed,
usage,
);
let billing = BilledTokenCounts::from_billed_usage(std::slice::from_ref(&billed));
```
- [ ] Set `AgentMessageProps.model` from event conversion with `model.clone()` for every `agent.message`.
- [ ] Update affected unit tests that pattern-match `AgentEvent::AssistantMessage`.
@ -219,6 +276,25 @@ Expected: agent event and event conversion tests pass.
- [ ] Replace `stage.usage.is_some()` checks with `!stage.usage.is_zero()`.
- [ ] Change rollup model fields from string ids to billing model refs:
```rust
pub struct ProjectionBillingStage {
pub node_id: String,
pub billing: BilledTokenCounts,
pub duration_ms: u64,
pub model: Option<ModelRef>,
}
pub struct ProjectionBillingByModel {
pub model: ModelRef,
pub stages: i64,
pub billing: BilledTokenCounts,
}
```
- [ ] Group by `HashMap<ModelRef, ProjectionBillingByModel>` and sort the final `by_model` vector by `(provider.to_string(), model_id, speed)` before returning it, so tests and API responses stay deterministic without reducing model identity to a string key.
- [ ] Replace `if let Some(usage) = stage.usage.as_ref()` rollup logic with direct `BilledTokenCounts` aggregation:
```rust
@ -226,11 +302,11 @@ if !stage.usage.is_zero() {
billed_visit_count += 1;
row.billing.add_counts(&stage.usage);
totals.add_counts(&stage.usage);
if let Some(model_id) = &stage.model_id {
row.model_id = Some(model_id.clone());
let model_entry = by_model.entry(model_id.clone()).or_insert_with(|| {
if let Some(model) = &stage.model {
row.model = Some(model.clone());
let model_entry = by_model.entry(model.clone()).or_insert_with(|| {
ProjectionBillingByModel {
model_id: model_id.clone(),
model: model.clone(),
stages: 0,
billing: BilledTokenCounts::default(),
}
@ -243,6 +319,8 @@ if !stage.usage.is_zero() {
- [ ] Replace open-coded token count accumulation in server aggregate billing with `add_counts`.
- [ ] Keep public `/runs/{id}/billing` response shape unchanged for this plan: server handlers convert `ModelRef` to the existing `ModelReference { id: model.model_id.clone() }` response object. Do not add a new public model identity response type in this task.
- [ ] Keep runtime behavior unchanged: live runtime still comes from `started_at` and terminal runtime still comes from `duration_ms`.
- [ ] Run:
@ -298,7 +376,7 @@ cd apps/fabro-web && bun test run-billing
## Acceptance Criteria
- In-flight stages included in `/runs/{id}/billing` can show nonzero token counts before `stage.completed`.
- Completed stages still use terminal billing as authoritative when terminal billing exists.
- Completed stages replace live billing with terminal billing when terminal billing exists.
- Retry/new visit behavior resets per-visit usage cleanly.
- Billing totals, by-model totals, and stage rows are derived from `StageProjection`.
- Old projections without `usage` deserialize with zero counters.
- Model identity uses the existing billing `ModelRef`; the plan does not introduce `StageModelIdentity` or store bare stage-level model-id strings.

22
lib/crates/fabro-agent/src/session.rs
View file

@ -13,7 +13,7 @@ use fabro_llm::types::{
use fabro_llm::{Error as LlmError, retry};
use fabro_mcp::config::{McpServerSettings, McpTransport};
use fabro_mcp::connection_manager::McpConnectionManager;
use fabro_model::Provider;
use fabro_model::{ModelRef, Provider, Speed};
use fabro_types::Principal;
use futures::StreamExt;
use tokio::sync::{Mutex as AsyncMutex, Notify, broadcast};
@ -1319,11 +1319,25 @@ impl Session {
});
// Emit AssistantMessage with enriched data from the response
let speed = self
.config
.speed
.as_deref()
.and_then(|value| value.parse::<Speed>().ok());
let model = ModelRef {
provider: self.provider_profile.provider(),
model_id: if response.model.is_empty() {
self.provider_profile.model().to_string()
} else {
response.model.clone()
},
speed,
};
self.event_emitter
.emit(self.id.clone(), AgentEvent::AssistantMessage {
text: text.clone(),
model: response.model.clone(),
usage: response.usage.clone(),
text: text.clone(),
model,
usage: response.usage.clone(),
tool_call_count: tool_calls.len(),
});

14
lib/crates/fabro-agent/src/types.rs
View file

@ -2,6 +2,7 @@ use std::time::SystemTime;
use fabro_llm::Error as LlmError;
use fabro_llm::types::{ContentPart, ThinkingData, TokenCounts, ToolCall, ToolResult};
use fabro_model::ModelRef;
use serde::{Deserialize, Serialize};
use crate::error::Error;
@ -117,7 +118,7 @@ pub enum AgentEvent {
},
AssistantMessage {
text: String,
model: String,
model: ModelRef,
usage: TokenCounts,
tool_call_count: usize,
},
@ -257,7 +258,8 @@ impl AgentEvent {
} => {
info!(
session_id,
model,
provider = %model.provider,
model = model.model_id.as_str(),
input_tokens = usage.input_tokens,
output_tokens = usage.output_tokens,
tool_call_count,
@ -427,6 +429,8 @@ pub struct SessionEvent {
#[cfg(test)]
mod tests {
use fabro_model::Provider;
use super::*;
#[test]
@ -666,7 +670,11 @@ mod tests {
};
let event = AgentEvent::AssistantMessage {
text: "Hello".into(),
model: "test-model".into(),
model: ModelRef {
provider: Provider::OpenAi,
model_id: "test-model".into(),
speed: None,
},
usage: usage.clone(),
tool_call_count: 2,
};

10
lib/crates/fabro-api/tests/run_projection_round_trip.rs
View file

@ -65,7 +65,15 @@ fn run_projection_round_trips_populated_projection() {
"script_invocation": null,
"script_timing": null,
"parallel_results": null,
"output": "done"
"output": "done",
"usage": {
"input_tokens": 0,
"output_tokens": 0,
"total_tokens": 0,
"reasoning_tokens": 0,
"cache_read_tokens": 0,
"cache_write_tokens": 0
}
}
}
});

8
lib/crates/fabro-api/tests/stage_projection_round_trip.rs
View file

@ -30,6 +30,14 @@ fn stage_projection_round_trips_representative_json() {
"termination": "exited",
"started_at": "2026-04-29T12:34:00Z",
"duration_ms": 56000,
"usage": {
"input_tokens": 0,
"output_tokens": 0,
"total_tokens": 0,
"reasoning_tokens": 0,
"cache_read_tokens": 0,
"cache_write_tokens": 0
},
"state": "succeeded"
});

2
lib/crates/fabro-cli/src/commands/run/run_progress/event.rs
View file

@ -382,7 +382,7 @@ pub(super) fn from_run_event(stored: &RunEvent) -> Option<ProgressEvent> {
EventBody::ParallelCompleted(_) => Some(ProgressEvent::ParallelCompleted),
EventBody::AgentMessage(props) => Some(ProgressEvent::AssistantMessage {
stage_node_id: node_id,
model: props.model.clone(),
model: props.model.model_id.clone(),
}),
EventBody::AgentToolStarted(props) => Some(ProgressEvent::ToolCallStarted {
stage_node_id: node_id,

8
lib/crates/fabro-cli/src/commands/run/run_progress/mod.rs
View file

@ -471,7 +471,7 @@ mod tests {
use chrono::{DateTime, Utc};
use fabro_agent::{AgentEvent, SandboxEvent};
use fabro_llm::types::TokenCounts;
use fabro_model::Provider;
use fabro_model::{ModelRef, Provider};
use fabro_types::{
MetadataSnapshotFailureKind, MetadataSnapshotPhase, ParallelBranchId, StageId, fixtures,
};
@ -550,7 +550,11 @@ mod tests {
fn assistant_message(stage: &str, model: &str) -> Event {
agent_event(stage, AgentEvent::AssistantMessage {
text: "done".into(),
model: model.into(),
model: ModelRef {
provider: Provider::OpenAi,
model_id: model.into(),
speed: None,
},
usage: TokenCounts::default(),
tool_call_count: 0,
})

180
lib/crates/fabro-model/src/billing.rs
View file

@ -362,6 +362,46 @@ impl BilledTokenCounts {
total_usd_micros: has_total.then_some(total_usd_micros),
}
}
pub fn add_counts(&mut self, source: &Self) {
self.input_tokens += source.input_tokens;
self.output_tokens += source.output_tokens;
self.total_tokens += source.total_tokens;
self.reasoning_tokens += source.reasoning_tokens;
self.cache_read_tokens += source.cache_read_tokens;
self.cache_write_tokens += source.cache_write_tokens;
if let Some(value) = source.total_usd_micros {
*self.total_usd_micros.get_or_insert(0) += value;
}
}
pub fn add_billed_usage(&mut self, usage: &BilledModelUsage) {
let tokens = usage.tokens();
self.input_tokens += tokens.input_tokens;
self.output_tokens += tokens.output_tokens;
self.reasoning_tokens += tokens.reasoning_tokens;
self.cache_read_tokens += tokens.cache_read_tokens;
self.cache_write_tokens += tokens.cache_write_tokens;
self.total_tokens += tokens.total_tokens();
if let Some(value) = usage.total_usd_micros {
*self.total_usd_micros.get_or_insert(0) += value;
}
}
pub fn replace_with_billed_usage(&mut self, usage: &BilledModelUsage) {
*self = Self::from_billed_usage(std::slice::from_ref(usage));
}
#[must_use]
pub fn is_zero(&self) -> bool {
self.input_tokens == 0
&& self.output_tokens == 0
&& self.total_tokens == 0
&& self.reasoning_tokens == 0
&& self.cache_read_tokens == 0
&& self.cache_write_tokens == 0
&& self.total_usd_micros.unwrap_or(0) == 0
}
}
impl Model {
@ -582,6 +622,146 @@ mod tests {
use super::*;
use crate::Catalog;
fn billed_usage(
input_tokens: i64,
output_tokens: i64,
total_usd_micros: Option<i64>,
) -> BilledModelUsage {
BilledModelUsage {
input: ModelBillingInput {
usage: ModelUsage {
model: ModelRef {
provider: Provider::OpenAi,
model_id: "gpt-5.4".to_string(),
speed: None,
},
tokens: TokenCounts {
input_tokens,
output_tokens,
reasoning_tokens: 3,
cache_read_tokens: 5,
cache_write_tokens: 7,
},
},
facts: ModelBillingFacts::OpenAi(OpenAiBillingFacts::default()),
},
total_usd_micros,
}
}
#[test]
fn billed_token_counts_add_counts_accumulates_cost_when_known() {
let mut counts = BilledTokenCounts {
input_tokens: 1,
output_tokens: 2,
total_tokens: 3,
reasoning_tokens: 4,
cache_read_tokens: 5,
cache_write_tokens: 6,
total_usd_micros: None,
};
counts.add_counts(&BilledTokenCounts {
input_tokens: 10,
output_tokens: 20,
total_tokens: 30,
reasoning_tokens: 40,
cache_read_tokens: 50,
cache_write_tokens: 60,
total_usd_micros: Some(70),
});
assert_eq!(counts, BilledTokenCounts {
input_tokens: 11,
output_tokens: 22,
total_tokens: 33,
reasoning_tokens: 44,
cache_read_tokens: 55,
cache_write_tokens: 66,
total_usd_micros: Some(70),
});
}
#[test]
fn billed_token_counts_add_billed_usage_preserves_unknown_cost() {
let mut counts = BilledTokenCounts::default();
counts.add_billed_usage(&billed_usage(10, 20, None));
assert_eq!(counts, BilledTokenCounts {
input_tokens: 10,
output_tokens: 20,
total_tokens: 45,
reasoning_tokens: 3,
cache_read_tokens: 5,
cache_write_tokens: 7,
total_usd_micros: None,
});
}
#[test]
fn billed_token_counts_add_billed_usage_accumulates_known_cost() {
let mut counts = BilledTokenCounts::default();
counts.add_billed_usage(&billed_usage(10, 20, Some(100)));
counts.add_billed_usage(&billed_usage(1, 2, Some(50)));
assert_eq!(counts.input_tokens, 11);
assert_eq!(counts.output_tokens, 22);
assert_eq!(counts.total_tokens, 63);
assert_eq!(counts.total_usd_micros, Some(150));
}
#[test]
fn billed_token_counts_replace_with_billed_usage_discards_previous_values() {
let mut counts = BilledTokenCounts {
input_tokens: 100,
output_tokens: 200,
total_tokens: 300,
reasoning_tokens: 400,
cache_read_tokens: 500,
cache_write_tokens: 600,
total_usd_micros: Some(700),
};
counts.replace_with_billed_usage(&billed_usage(1, 2, None));
assert_eq!(counts, BilledTokenCounts {
input_tokens: 1,
output_tokens: 2,
total_tokens: 18,
reasoning_tokens: 3,
cache_read_tokens: 5,
cache_write_tokens: 7,
total_usd_micros: None,
});
}
#[test]
fn billed_token_counts_is_zero_treats_missing_and_zero_cost_as_zero() {
assert!(BilledTokenCounts::default().is_zero());
assert!(
BilledTokenCounts {
total_usd_micros: Some(0),
..BilledTokenCounts::default()
}
.is_zero()
);
assert!(
!BilledTokenCounts {
input_tokens: 1,
..BilledTokenCounts::default()
}
.is_zero()
);
assert!(
!BilledTokenCounts {
total_usd_micros: Some(1),
..BilledTokenCounts::default()
}
.is_zero()
);
}
#[test]
fn openai_pricing_bills_cached_input_and_reasoning_output() {
let pricing = ModelPricing {

12
lib/crates/fabro-server/src/demo/mod.rs
View file

@ -1451,7 +1451,11 @@ mod runs {
"evt-detect-drift-2",
EventBody::AgentMessage(AgentMessageProps {
text: "I'll start by loading the environment configurations for both production and staging to compare them.".into(),
model: "Opus 4.6".into(),
model: fabro_model::ModelRef {
provider: fabro_model::Provider::Anthropic,
model_id: "claude-opus-4-6".into(),
speed: None,
},
billing: BilledTokenCounts::default(),
tool_call_count: 0,
visit: 1,
@ -1504,7 +1508,11 @@ mod runs {
"evt-detect-drift-7",
EventBody::AgentMessage(AgentMessageProps {
text: "I've detected drift in 3 resources between production and staging:\n\n1. **redis.max_connections** — production has 200, staging has 100\n2. **redis.tls** — enabled in production, disabled in staging\n3. **iam.session_duration** — production uses 3600s, staging uses 1800s".into(),
model: "Opus 4.6".into(),
model: fabro_model::ModelRef {
provider: fabro_model::Provider::Anthropic,
model_id: "claude-opus-4-6".into(),
speed: None,
},
billing: BilledTokenCounts::default(),
tool_call_count: 0,
visit: 1,

16
lib/crates/fabro-server/src/server.rs
View file

@ -602,18 +602,6 @@ pub(crate) struct ResolvedAppStateSettings {
pub(crate) manifest_run_settings: std::result::Result<RunNamespace, SharedError>,
}
fn accumulate_billed_token_counts(target: &mut BilledTokenCounts, source: &BilledTokenCounts) {
target.input_tokens += source.input_tokens;
target.output_tokens += source.output_tokens;
target.reasoning_tokens += source.reasoning_tokens;
target.cache_read_tokens += source.cache_read_tokens;
target.cache_write_tokens += source.cache_write_tokens;
target.total_tokens += source.total_tokens;
if let Some(value) = source.total_usd_micros {
*target.total_usd_micros.get_or_insert(0) += value;
}
}
fn accumulate_billing_rollup(
accumulator: &mut BillingAccumulator,
rollup: &fabro_workflow::ProjectionBillingRollup,
@ -623,10 +611,10 @@ fn accumulate_billing_rollup(
for model in &rollup.by_model {
let entry = accumulator
.by_model
.entry(model.model_id.clone())
.entry(model.model.model_id.clone())
.or_default();
entry.stages += model.stages;
accumulate_billed_token_counts(&mut entry.billing, &model.billing);
entry.billing.add_counts(&model.billing);
}
}

10
lib/crates/fabro-server/src/server/handler/billing.rs
View file

@ -86,7 +86,7 @@ async fn get_run_billing(
.map(|model| BillingByModel {
billing: model.billing.clone(),
model: ModelReference {
id: model.model_id.clone(),
id: model.model.model_id.clone(),
},
stages: model.stages,
})
@ -108,8 +108,10 @@ async fn get_run_billing(
.map(|stage| stage.billing.clone())
.unwrap_or_default(),
model: rollup_stage
.and_then(|stage| stage.model_id.as_ref())
.map(|id| ModelReference { id: id.clone() }),
.and_then(|stage| stage.model.as_ref())
.map(|model| ModelReference {
id: model.model_id.clone(),
}),
runtime_secs: row.runtime_secs,
stage: BillingStageRef {
id: row.node_id.clone(),
@ -191,7 +193,7 @@ fn billing_runtime_secs(stage: &StageProjection, now: DateTime<Utc>) -> Option<f
fn stage_has_billing_row(stage: &StageProjection) -> bool {
stage.completion.is_some()
|| stage.duration_ms.is_some()
|| stage.usage.is_some()
|| !stage.usage.is_zero()
|| stage.started_at.is_some()
|| stage.state.is_some()
}

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

@ -16,7 +16,7 @@ use fabro_interview::{
AnswerValue, ControlInterviewer, Interviewer, Question, WorkerControlMessage,
};
use fabro_llm::types::{Message as LlmMessage, Request as LlmRequest};
use fabro_model::Provider;
use fabro_model::{ModelRef, Provider};
use fabro_types::settings::ServerAuthMethod;
use fabro_types::{
AttrValue, AuthMethod, CommandTermination, FailureCategory, FailureDetail, Graph,
@ -7690,9 +7690,13 @@ fn aggregate_billing_counts_projection_rollup_usage_visits() {
},
by_model: vec![
fabro_workflow::ProjectionBillingByModel {
model_id: "gpt-old".to_string(),
stages: 1,
billing: BilledTokenCounts {
model: ModelRef {
provider: Provider::OpenAi,
model_id: "gpt-old".to_string(),
speed: None,
},
stages: 1,
billing: BilledTokenCounts {
input_tokens: 100,
output_tokens: 10,
total_tokens: 110,
@ -7703,9 +7707,13 @@ fn aggregate_billing_counts_projection_rollup_usage_visits() {
},
},
fabro_workflow::ProjectionBillingByModel {
model_id: "gpt-new".to_string(),
stages: 1,
billing: BilledTokenCounts {
model: ModelRef {
provider: Provider::OpenAi,
model_id: "gpt-new".to_string(),
speed: None,
},
stages: 1,
billing: BilledTokenCounts {
input_tokens: 200,
output_tokens: 20,
total_tokens: 220,

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

@ -319,6 +319,10 @@ impl RunProjectionReducer for RunProjection {
return Ok(());
};
stage.response = Some(props.response.clone());
if let Some(billing) = &props.billing {
stage.usage.replace_with_billed_usage(billing);
stage.model = Some(billing.model().clone());
}
}
EventBody::StageCompleted(props) => {
let response = props.response.clone();
@ -332,7 +336,10 @@ impl RunProjectionReducer for RunProjection {
stage.response = response;
stage.completion = Some(completion);
stage.duration_ms = Some(props.duration_ms);
stage.usage.clone_from(&props.billing);
if let Some(billing) = &props.billing {
stage.usage.replace_with_billed_usage(billing);
stage.model = Some(billing.model().clone());
}
stage.state = Some(StageState::from(outcome.status));
}
EventBody::StageFailed(props) => {
@ -350,9 +357,20 @@ impl RunProjectionReducer for RunProjection {
timestamp: ts,
});
stage.duration_ms = Some(props.duration_ms);
stage.usage.clone_from(&props.billing);
if let Some(billing) = &props.billing {
stage.usage.replace_with_billed_usage(billing);
stage.model = Some(billing.model().clone());
}
stage.state = Some(StageState::from(outcome));
}
EventBody::AgentMessage(props) => {
let Some(stage) = stage_at_stored_or_visit(self, stored, props.visit, event.seq)
else {
return Ok(());
};
stage.usage.add_counts(&props.billing);
stage.model = Some(props.model.clone());
}
EventBody::AgentSessionActivated(props) => {
let Some(stage) = stage_at_stored_or_visit(self, stored, props.visit, event.seq)
else {
@ -737,15 +755,15 @@ mod tests {
use chrono::Utc;
use fabro_types::run_event::run::RunFailedProps;
use fabro_types::run_event::{
AgentCliCancelledProps, AgentCliCompletedProps, AgentCliTimedOutProps,
AgentCliCancelledProps, AgentCliCompletedProps, AgentCliTimedOutProps, AgentMessageProps,
AgentSessionActivatedProps, AgentSessionEndedProps, AgentSessionStartedProps,
CheckpointCompletedProps, InterviewCompletedProps, InterviewOption, InterviewStartedProps,
RunControlEffectProps, StageCompletedProps, StageFailedProps, StagePromptProps,
StageRetryingProps, StageStartedProps,
};
use fabro_types::{
BilledModelUsage, BlockedReason, Checkpoint, CommandTermination, EventBody,
FailureCategory, FailureDetail, FailureReason, Outcome, QuestionType, RunBlobId,
BilledModelUsage, BilledTokenCounts, BlockedReason, Checkpoint, CommandTermination,
EventBody, FailureCategory, FailureDetail, FailureReason, Outcome, QuestionType, RunBlobId,
RunControlAction, RunEvent, RunStatus, StageOutcome, StageState, SuccessReason,
TerminalStatus, WorkflowSettings, first_event_seq, fixtures,
};
@ -806,6 +824,10 @@ mod tests {
serde_json::to_value(usage).unwrap()
}
fn usage_counts(usage: &BilledModelUsage) -> BilledTokenCounts {
BilledTokenCounts::from_billed_usage(std::slice::from_ref(usage))
}
fn test_raw_event(
seq: u32,
event: &str,
@ -1222,7 +1244,8 @@ mod tests {
let stage = state.stage(&StageId::new("build", 1)).unwrap();
assert_eq!(stage.duration_ms, Some(789));
assert_eq!(stage.usage.as_ref(), Some(&usage));
assert_eq!(stage.usage, usage_counts(&usage));
assert_eq!(stage.model.as_ref(), Some(usage.model()));
}
#[test]
@ -1263,7 +1286,8 @@ mod tests {
let stage = state.stage(&stage_id).unwrap();
assert_eq!(stage.duration_ms, Some(654));
assert_eq!(stage.usage.as_ref(), Some(&usage));
assert_eq!(stage.usage, usage_counts(&usage));
assert_eq!(stage.model.as_ref(), Some(usage.model()));
}
#[test]
@ -1307,9 +1331,11 @@ mod tests {
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!(first_stage.usage, usage_counts(&first_usage));
assert_eq!(first_stage.model.as_ref(), Some(first_usage.model()));
assert_eq!(second_stage.duration_ms, Some(222));
assert_eq!(second_stage.usage.as_ref(), Some(&second_usage));
assert_eq!(second_stage.usage, usage_counts(&second_usage));
assert_eq!(second_stage.model.as_ref(), Some(second_usage.model()));
}
#[test]
@ -1351,7 +1377,8 @@ mod tests {
);
let stage = state.stage(&scoped_stage_id).unwrap();
assert_eq!(stage.duration_ms, Some(333));
assert_eq!(stage.usage.as_ref(), Some(&usage));
assert_eq!(stage.usage, usage_counts(&usage));
assert_eq!(stage.model.as_ref(), Some(usage.model()));
assert_eq!(stage.response.as_deref(), Some("done"));
}
@ -1384,7 +1411,8 @@ mod tests {
);
let stage = state.stage(&scoped_stage_id).unwrap();
assert_eq!(stage.duration_ms, Some(444));
assert_eq!(stage.usage.as_ref(), Some(&usage));
assert_eq!(stage.usage, usage_counts(&usage));
assert_eq!(stage.model.as_ref(), Some(usage.model()));
let completion = stage.completion.as_ref().unwrap();
assert_eq!(completion.outcome, StageOutcome::Failed {
retry_requested: true,
@ -2329,6 +2357,28 @@ mod tests {
.expect("billing fixture should deserialize")
}
fn live_agent_message_props(billing: BilledTokenCounts) -> AgentMessageProps {
AgentMessageProps {
text: "assistant text".to_string(),
model: billed_usage().model().clone(),
billing,
tool_call_count: 0,
visit: 1,
}
}
fn live_counts(input_tokens: i64, output_tokens: i64) -> BilledTokenCounts {
BilledTokenCounts {
input_tokens,
output_tokens,
total_tokens: input_tokens + output_tokens,
reasoning_tokens: 0,
cache_read_tokens: 0,
cache_write_tokens: 0,
total_usd_micros: None,
}
}
#[test]
fn stage_started_records_started_at_and_running_state() {
let mut state = RunProjection::default();
@ -2348,6 +2398,175 @@ mod tests {
assert_eq!(stage.effective_state(), StageState::Running);
}
#[test]
fn agent_message_accumulates_live_usage_on_stage_projection() {
let mut state = RunProjection::default();
let stage_id = StageId::new("build", 1);
let model = billed_usage().model().clone();
state
.apply_event(&test_stage_event(
1,
EventBody::StageStarted(started_props()),
stage_id.clone(),
))
.unwrap();
state
.apply_event(&test_stage_event(
2,
EventBody::AgentMessage(live_agent_message_props(live_counts(10, 5))),
stage_id.clone(),
))
.unwrap();
state
.apply_event(&test_stage_event(
3,
EventBody::AgentMessage(live_agent_message_props(live_counts(20, 7))),
stage_id.clone(),
))
.unwrap();
let stage = state.stage(&stage_id).unwrap();
assert_eq!(stage.usage, live_counts(30, 12));
assert_eq!(stage.model, Some(model));
}
#[test]
fn stage_completed_replaces_live_usage_with_terminal_billing() {
let mut state = RunProjection::default();
let stage_id = StageId::new("build", 1);
let usage = billed_usage();
state
.apply_event(&test_stage_event(
1,
EventBody::StageStarted(started_props()),
stage_id.clone(),
))
.unwrap();
state
.apply_event(&test_stage_event(
2,
EventBody::AgentMessage(live_agent_message_props(live_counts(100, 50))),
stage_id.clone(),
))
.unwrap();
let mut props = completed_props(42, StageOutcome::Succeeded);
props.billing = Some(usage.clone());
state
.apply_event(&test_stage_event(
3,
EventBody::StageCompleted(props),
stage_id.clone(),
))
.unwrap();
let stage = state.stage(&stage_id).unwrap();
assert_eq!(stage.usage, usage_counts(&usage));
assert_eq!(stage.model.as_ref(), Some(usage.model()));
}
#[test]
fn stage_completed_without_billing_preserves_live_usage() {
let mut state = RunProjection::default();
let stage_id = StageId::new("build", 1);
let model = billed_usage().model().clone();
state
.apply_event(&test_stage_event(
1,
EventBody::StageStarted(started_props()),
stage_id.clone(),
))
.unwrap();
state
.apply_event(&test_stage_event(
2,
EventBody::AgentMessage(live_agent_message_props(live_counts(10, 5))),
stage_id.clone(),
))
.unwrap();
state
.apply_event(&test_stage_event(
3,
EventBody::StageCompleted(completed_props(42, StageOutcome::Succeeded)),
stage_id.clone(),
))
.unwrap();
let stage = state.stage(&stage_id).unwrap();
assert_eq!(stage.usage, live_counts(10, 5));
assert_eq!(stage.model, Some(model));
}
#[test]
fn stage_failed_replaces_live_usage_with_terminal_billing() {
let mut state = RunProjection::default();
let stage_id = StageId::new("build", 1);
let usage = billed_usage();
state
.apply_event(&test_stage_event(
1,
EventBody::StageStarted(started_props()),
stage_id.clone(),
))
.unwrap();
state
.apply_event(&test_stage_event(
2,
EventBody::AgentMessage(live_agent_message_props(live_counts(100, 50))),
stage_id.clone(),
))
.unwrap();
let mut props = failed_props(42, false);
props.billing = Some(usage.clone());
state
.apply_event(&test_stage_event(
3,
EventBody::StageFailed(props),
stage_id.clone(),
))
.unwrap();
let stage = state.stage(&stage_id).unwrap();
assert_eq!(stage.usage, usage_counts(&usage));
assert_eq!(stage.model.as_ref(), Some(usage.model()));
}
#[test]
fn stage_started_resets_live_usage_for_new_attempt() {
let mut state = RunProjection::default();
let stage_id = StageId::new("build", 1);
state
.apply_event(&test_stage_event(
1,
EventBody::StageStarted(started_props()),
stage_id.clone(),
))
.unwrap();
state
.apply_event(&test_stage_event(
2,
EventBody::AgentMessage(live_agent_message_props(live_counts(10, 5))),
stage_id.clone(),
))
.unwrap();
state
.apply_event(&test_stage_event(
3,
EventBody::StageStarted(started_props()),
stage_id.clone(),
))
.unwrap();
let stage = state.stage(&stage_id).unwrap();
assert!(stage.usage.is_zero());
assert_eq!(stage.model, None);
assert_eq!(stage.state, Some(StageState::Running));
}
#[test]
fn stage_completed_records_duration_usage_and_terminal_state() {
let mut state = RunProjection::default();
@ -2373,7 +2592,8 @@ mod tests {
let stage = state.stage(&stage_id).unwrap();
assert_eq!(stage.duration_ms, Some(42));
assert_eq!(stage.usage.as_ref(), Some(&usage));
assert_eq!(stage.usage, usage_counts(&usage));
assert_eq!(stage.model.as_ref(), Some(usage.model()));
assert_eq!(stage.state, Some(StageState::Succeeded));
assert_eq!(stage.effective_state(), StageState::Succeeded);
}

22
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::{
BilledModelUsage, Checkpoint, RunStatus, SandboxRecord, StageCompletion, StageOutcome,
StartRecord, TerminalStatus, WorkflowSettings, first_event_seq, fixtures,
BilledModelUsage, BilledTokenCounts, Checkpoint, RunStatus, SandboxRecord, StageCompletion,
StageOutcome, StartRecord, TerminalStatus, WorkflowSettings, first_event_seq, fixtures,
};
use serde_json::json;
@ -116,14 +116,21 @@ fn serializable_projection_round_trips_and_trims_bulky_node_fields() {
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());
let usage = sample_usage();
let usage_counts = BilledTokenCounts::from_billed_usage(std::slice::from_ref(&usage));
stage.usage = usage_counts.clone();
stage.model = Some(usage.model().clone());
stage.output = Some("output".to_string());
let serialized = serde_json::to_value(SerializableProjection(&projection))
.expect("projection should serialize");
assert!(
serialized["stages"]["build@2"].get("usage").is_none(),
"stage usage is server-internal and should not be serialized"
assert_eq!(
serialized["stages"]["build@2"]["usage"]["input_tokens"],
json!(123)
);
assert_eq!(
serialized["stages"]["build@2"]["model"]["model_id"],
json!("gpt-5.2")
);
let round_tripped: RunProjection =
serde_json::from_value(serialized).expect("serialized projection should deserialize");
@ -164,7 +171,8 @@ fn serializable_projection_round_trips_and_trims_bulky_node_fields() {
Some(json!([{ "stage": "fanout@1" }]))
);
assert_eq!(node.duration_ms, Some(1234));
assert_eq!(node.usage, None);
assert_eq!(node.usage, usage_counts);
assert_eq!(node.model.as_ref(), Some(usage.model()));
}
#[test]

3
lib/crates/fabro-types/src/run_event/agent.rs
View file

@ -2,6 +2,7 @@ use serde::{Deserialize, Serialize};
use serde_json::Value;
use super::BilledTokenCounts;
use crate::ModelRef;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct AgentSessionStartedProps {
@ -55,7 +56,7 @@ pub struct AgentInputProps {
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct AgentMessageProps {
pub text: String,
pub model: String,
pub model: ModelRef,
pub billing: BilledTokenCounts,
pub tool_call_count: usize,
pub visit: u32,

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

@ -5,8 +5,8 @@ use std::num::NonZeroU32;
use chrono::{DateTime, Utc};
use crate::{
BilledModelUsage, Checkpoint, Conclusion, DiffSummary, InterviewQuestionRecord,
InvalidTransition, PullRequestRecord, RunControlAction, RunId, RunSpec, RunStatus,
BilledTokenCounts, Checkpoint, Conclusion, DiffSummary, InterviewQuestionRecord,
InvalidTransition, ModelRef, PullRequestRecord, RunControlAction, RunId, RunSpec, RunStatus,
SandboxRecord, StageCompletion, StageHandler, StageId, StageState, StartRecord,
};
@ -65,11 +65,10 @@ pub struct StageProjection {
pub handler: Option<StageHandler>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub duration_ms: Option<u64>,
/// Server-internal billing usage for the latest attempt; not part of the
/// wire contract because `BilledModelUsage` is not modeled in OpenAPI.
/// Read only in-process by the billing handler.
#[serde(skip)]
pub usage: Option<BilledModelUsage>,
#[serde(default)]
pub usage: BilledTokenCounts,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub model: Option<ModelRef>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub state: Option<StageState>,
}
@ -90,7 +89,8 @@ impl StageProjection {
response: None,
completion: None,
duration_ms: None,
usage: None,
usage: BilledTokenCounts::default(),
model: None,
provider_used: None,
diff: None,
script_invocation: None,

103
lib/crates/fabro-workflow/src/billing_rollup.rs
View file

@ -1,20 +1,20 @@
use std::collections::{BTreeMap, HashMap};
use std::collections::HashMap;
use fabro_types::{BilledModelUsage, BilledTokenCounts, RunProjection};
use fabro_types::{BilledTokenCounts, ModelRef, RunProjection};
#[derive(Debug, Clone, PartialEq)]
pub struct ProjectionBillingStage {
pub node_id: String,
pub billing: BilledTokenCounts,
pub duration_ms: u64,
pub model_id: Option<String>,
pub model: Option<ModelRef>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProjectionBillingByModel {
pub model_id: String,
pub stages: i64,
pub billing: BilledTokenCounts,
pub model: ModelRef,
pub stages: i64,
pub billing: BilledTokenCounts,
}
#[derive(Debug, Clone, Default, PartialEq)]
@ -37,7 +37,7 @@ impl ProjectionBillingRollup {
pub fn billing_rollup_from_projection(projection: &RunProjection) -> ProjectionBillingRollup {
let mut stage_indices = HashMap::<String, usize>::new();
let mut stages = Vec::<ProjectionBillingStage>::new();
let mut by_model = BTreeMap::<String, ProjectionBillingByModel>::new();
let mut by_model = HashMap::<ModelRef, ProjectionBillingByModel>::new();
let mut totals = BilledTokenCounts::default();
let mut runtime_ms = 0_u64;
let mut billed_visit_count = 0_usize;
@ -46,7 +46,7 @@ pub fn billing_rollup_from_projection(projection: &RunProjection) -> ProjectionB
if is_boundary_stage(projection, stage_id.node_id()) {
continue;
}
if stage.completion.is_none() && stage.duration_ms.is_none() && stage.usage.is_none() {
if stage.completion.is_none() && stage.duration_ms.is_none() && stage.usage.is_zero() {
continue;
}
@ -57,7 +57,7 @@ pub fn billing_rollup_from_projection(projection: &RunProjection) -> ProjectionB
node_id: node_id.to_string(),
billing: BilledTokenCounts::default(),
duration_ms: 0,
model_id: None,
model: None,
});
index
});
@ -68,30 +68,46 @@ pub fn billing_rollup_from_projection(projection: &RunProjection) -> ProjectionB
runtime_ms = runtime_ms.saturating_add(duration_ms);
}
if let Some(usage) = stage.usage.as_ref() {
if !stage.usage.is_zero() {
billed_visit_count += 1;
row.model_id = Some(usage.model_id().to_string());
accumulate_usage(&mut row.billing, usage);
accumulate_usage(&mut totals, usage);
row.billing.add_counts(&stage.usage);
totals.add_counts(&stage.usage);
let model_id = usage.model_id().to_string();
let model_entry =
by_model
.entry(model_id.clone())
.or_insert_with(|| ProjectionBillingByModel {
model_id,
stages: 0,
billing: BilledTokenCounts::default(),
});
model_entry.stages += 1;
accumulate_usage(&mut model_entry.billing, 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(&stage.usage);
}
}
}
let mut by_model = by_model.into_values().collect::<Vec<_>>();
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: by_model.into_values().collect(),
by_model,
runtime_ms,
billed_visit_count,
}
@ -104,26 +120,13 @@ fn is_boundary_stage(projection: &RunProjection, node_id: &str) -> bool {
.is_some_and(|node| matches!(node.handler_type(), Some("start" | "exit")))
}
fn accumulate_usage(counts: &mut BilledTokenCounts, usage: &BilledModelUsage) {
let tokens = usage.tokens();
counts.input_tokens += tokens.input_tokens;
counts.output_tokens += tokens.output_tokens;
counts.reasoning_tokens += tokens.reasoning_tokens;
counts.cache_read_tokens += tokens.cache_read_tokens;
counts.cache_write_tokens += tokens.cache_write_tokens;
counts.total_tokens += tokens.total_tokens();
if let Some(value) = usage.total_usd_micros {
*counts.total_usd_micros.get_or_insert(0) += value;
}
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use fabro_types::{
AttrValue, BilledModelUsage, Graph, Node, RunProjection, RunSpec, StageCompletion,
StageOutcome, WorkflowSettings, first_event_seq, fixtures,
AttrValue, BilledModelUsage, BilledTokenCounts, Graph, Node, RunProjection, RunSpec,
StageCompletion, StageOutcome, WorkflowSettings, first_event_seq, fixtures,
};
use serde_json::json;
@ -158,7 +161,8 @@ mod tests {
let success_usage = test_usage("gpt-new", 200, 20);
let first = projection.stage_entry("verify", 1, first_event_seq(1));
first.duration_ms = Some(1200);
first.usage = Some(failed_usage);
first.usage = BilledTokenCounts::from_billed_usage(std::slice::from_ref(&failed_usage));
first.model = Some(failed_usage.model().clone());
first.completion = Some(StageCompletion {
outcome: StageOutcome::Failed {
retry_requested: true,
@ -169,7 +173,8 @@ mod tests {
});
let second = projection.stage_entry("verify", 2, first_event_seq(2));
second.duration_ms = Some(800);
second.usage = Some(success_usage);
second.usage = BilledTokenCounts::from_billed_usage(std::slice::from_ref(&success_usage));
second.model = Some(success_usage.model().clone());
second.completion = Some(StageCompletion {
outcome: StageOutcome::Succeeded,
notes: None,
@ -181,7 +186,13 @@ mod tests {
assert_eq!(rollup.stages.len(), 1);
assert_eq!(rollup.stages[0].node_id, "verify");
assert_eq!(rollup.stages[0].model_id.as_deref(), Some("gpt-new"));
assert_eq!(
rollup.stages[0]
.model
.as_ref()
.map(|model| model.model_id.as_str()),
Some("gpt-new")
);
assert_eq!(rollup.stages[0].duration_ms, 2000);
assert_eq!(rollup.stages[0].billing.input_tokens, 300);
assert_eq!(rollup.stages[0].billing.output_tokens, 30);
@ -194,10 +205,10 @@ mod tests {
assert_eq!(rollup.billed_visit_count, 2);
assert_eq!(rollup.by_model.len(), 2);
assert_eq!(rollup.by_model[0].model_id, "gpt-new");
assert_eq!(rollup.by_model[0].model.model_id, "gpt-new");
assert_eq!(rollup.by_model[0].stages, 1);
assert_eq!(rollup.by_model[0].billing.input_tokens, 200);
assert_eq!(rollup.by_model[1].model_id, "gpt-old");
assert_eq!(rollup.by_model[1].model.model_id, "gpt-old");
assert_eq!(rollup.by_model[1].stages, 1);
assert_eq!(rollup.by_model[1].billing.input_tokens, 100);
}
@ -219,7 +230,7 @@ mod tests {
assert_eq!(rollup.stages.len(), 1);
assert_eq!(rollup.stages[0].node_id, "build");
assert_eq!(rollup.stages[0].duration_ms, 25);
assert!(rollup.stages[0].model_id.is_none());
assert!(rollup.stages[0].model.is_none());
assert_eq!(rollup.stages[0].billing.input_tokens, 0);
assert_eq!(rollup.runtime_ms, 25);
assert!(rollup.by_model.is_empty());

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

@ -4,25 +4,13 @@ use ::fabro_types::{
};
use chrono::Utc;
use fabro_agent::{AgentEvent, SandboxEvent};
use fabro_llm::types::TokenCounts as LlmTokenCounts;
use uuid::Uuid;
use super::Event;
use super::stored_fields::stored_event_fields;
use crate::outcome::billed_model_usage_from_llm;
use crate::stage_scope::StageScope;
fn billed_token_counts_from_llm(usage: &LlmTokenCounts) -> BilledTokenCounts {
BilledTokenCounts {
input_tokens: usage.input_tokens,
output_tokens: usage.output_tokens,
total_tokens: usage.total_tokens(),
reasoning_tokens: usage.reasoning_tokens,
cache_read_tokens: usage.cache_read_tokens,
cache_write_tokens: usage.cache_write_tokens,
total_usd_micros: None,
}
}
fn stage_status_from_string(status: &str) -> StageOutcome {
status.parse().unwrap_or_else(|_| {
tracing::warn!(
@ -562,13 +550,23 @@ fn event_body_from_event(event: &Event) -> EventBody {
model,
usage,
tool_call_count,
} => EventBody::AgentMessage(fabro_types::AgentMessageProps {
text: text.clone(),
model: model.clone(),
billing: billed_token_counts_from_llm(usage),
tool_call_count: *tool_call_count,
visit: *visit,
}),
} => {
let requested_speed = model.speed.map(<&'static str>::from);
let billed = billed_model_usage_from_llm(
&model.model_id,
model.provider,
requested_speed,
usage,
);
let billing = BilledTokenCounts::from_billed_usage(std::slice::from_ref(&billed));
EventBody::AgentMessage(fabro_types::AgentMessageProps {
text: text.clone(),
model: model.clone(),
billing,
tool_call_count: *tool_call_count,
visit: *visit,
})
}
AgentEvent::ToolCallStarted {
tool_name,
tool_call_id,
@ -1276,6 +1274,7 @@ mod tests {
use chrono::Utc;
use fabro_agent::{AgentEvent, SandboxEvent};
use fabro_llm::types::TokenCounts as LlmTokenCounts;
use fabro_model::{ModelRef, Provider};
use super::*;
use crate::error::Error;
@ -1949,7 +1948,11 @@ mod tests {
visit: 1,
event: AgentEvent::AssistantMessage {
text: "ok".to_string(),
model: "claude-sonnet".to_string(),
model: ModelRef {
provider: Provider::Anthropic,
model_id: "claude-sonnet".to_string(),
speed: None,
},
usage: LlmTokenCounts::default(),
tool_call_count: 0,
},

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

@ -283,7 +283,7 @@ fn agent_actor_for_event(
AgentEvent::AssistantMessage { model, .. } => Some(Principal::Agent {
session_id: session_id.map(str::to_string),
parent_session_id: parent_session_id.map(str::to_string),
model: Some(model.clone()),
model: Some(model.model_id.clone()),
}),
AgentEvent::ToolCallStarted { .. }
| AgentEvent::ToolCallOutputDelta { .. }

11
lib/crates/fabro-workflow/src/pipeline/finalize.rs
View file

@ -647,8 +647,8 @@ mod tests {
use fabro_store::{Database, EventEnvelope, RunDatabase, RunProjection};
use fabro_types::run_event::{MetadataSnapshotFailureKind, MetadataSnapshotPhase};
use fabro_types::{
BilledModelUsage, EventBody, RunBlobId, RunEvent, RunId, StageCompletion, WorkflowSettings,
first_event_seq, fixtures,
BilledModelUsage, BilledTokenCounts, EventBody, RunBlobId, RunEvent, RunId,
StageCompletion, WorkflowSettings, first_event_seq, fixtures,
};
use object_store::memory::InMemory;
@ -934,7 +934,8 @@ mod tests {
let success_usage = test_usage("gpt-new", 200, 20);
let failed = projection.stage_entry("verify", 1, first_event_seq(1));
failed.duration_ms = Some(1200);
failed.usage = Some(failed_usage);
failed.usage = BilledTokenCounts::from_billed_usage(std::slice::from_ref(&failed_usage));
failed.model = Some(failed_usage.model().clone());
failed.completion = Some(StageCompletion {
outcome: StageOutcome::Failed {
retry_requested: true,
@ -945,7 +946,9 @@ mod tests {
});
let succeeded = projection.stage_entry("verify", 2, first_event_seq(2));
succeeded.duration_ms = Some(800);
succeeded.usage = Some(success_usage.clone());
succeeded.usage =
BilledTokenCounts::from_billed_usage(std::slice::from_ref(&success_usage));
succeeded.model = Some(success_usage.model().clone());
succeeded.completion = Some(StageCompletion {
outcome: StageOutcome::Succeeded,
notes: None,