diff --git a/apps/fabro-web/app/routes/run-billing.test.tsx b/apps/fabro-web/app/routes/run-billing.test.tsx index cf8f595e4..ab733b53b 100644 --- a/apps/fabro-web/app/routes/run-billing.test.tsx +++ b/apps/fabro-web/app/routes/run-billing.test.tsx @@ -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. diff --git a/docs/public/api-reference/fabro-api.yaml b/docs/public/api-reference/fabro-api.yaml index a63439f18..f76b53473 100644 --- a/docs/public/api-reference/fabro-api.yaml +++ b/docs/public/api-reference/fabro-api.yaml @@ -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 diff --git a/docs/superpowers/plans/2026-05-09-live-stage-token-projection.md b/docs/superpowers/plans/2026-05-09-live-stage-token-projection.md index 5a2cb8844..d19aa2a03 100644 --- a/docs/superpowers/plans/2026-05-09-live-stage-token-projection.md +++ b/docs/superpowers/plans/2026-05-09-live-stage-token-projection.md @@ -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, +pub model: Option, ``` - [ ] Remove the `#[serde(skip)] pub usage: Option` 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::().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::().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, +} + +pub struct ProjectionBillingByModel { + pub model: ModelRef, + pub stages: i64, + pub billing: BilledTokenCounts, +} +``` + +- [ ] Group by `HashMap` 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. diff --git a/lib/crates/fabro-agent/src/session.rs b/lib/crates/fabro-agent/src/session.rs index b6171c8a2..045038a63 100644 --- a/lib/crates/fabro-agent/src/session.rs +++ b/lib/crates/fabro-agent/src/session.rs @@ -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::().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(), }); diff --git a/lib/crates/fabro-agent/src/types.rs b/lib/crates/fabro-agent/src/types.rs index 4e4cb952e..ad9103060 100644 --- a/lib/crates/fabro-agent/src/types.rs +++ b/lib/crates/fabro-agent/src/types.rs @@ -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, }; diff --git a/lib/crates/fabro-api/tests/run_projection_round_trip.rs b/lib/crates/fabro-api/tests/run_projection_round_trip.rs index 412d04d4d..1bbcf878c 100644 --- a/lib/crates/fabro-api/tests/run_projection_round_trip.rs +++ b/lib/crates/fabro-api/tests/run_projection_round_trip.rs @@ -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 + } } } }); diff --git a/lib/crates/fabro-api/tests/stage_projection_round_trip.rs b/lib/crates/fabro-api/tests/stage_projection_round_trip.rs index a7b397be3..7d740316c 100644 --- a/lib/crates/fabro-api/tests/stage_projection_round_trip.rs +++ b/lib/crates/fabro-api/tests/stage_projection_round_trip.rs @@ -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" }); diff --git a/lib/crates/fabro-cli/src/commands/run/run_progress/event.rs b/lib/crates/fabro-cli/src/commands/run/run_progress/event.rs index 431ad12d4..77d4596e2 100644 --- a/lib/crates/fabro-cli/src/commands/run/run_progress/event.rs +++ b/lib/crates/fabro-cli/src/commands/run/run_progress/event.rs @@ -382,7 +382,7 @@ pub(super) fn from_run_event(stored: &RunEvent) -> Option { 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, diff --git a/lib/crates/fabro-cli/src/commands/run/run_progress/mod.rs b/lib/crates/fabro-cli/src/commands/run/run_progress/mod.rs index 58398efb0..01c3534e7 100644 --- a/lib/crates/fabro-cli/src/commands/run/run_progress/mod.rs +++ b/lib/crates/fabro-cli/src/commands/run/run_progress/mod.rs @@ -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, }) diff --git a/lib/crates/fabro-model/src/billing.rs b/lib/crates/fabro-model/src/billing.rs index b07df50cb..55aa3e9c3 100644 --- a/lib/crates/fabro-model/src/billing.rs +++ b/lib/crates/fabro-model/src/billing.rs @@ -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, + ) -> 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 { diff --git a/lib/crates/fabro-server/src/demo/mod.rs b/lib/crates/fabro-server/src/demo/mod.rs index 0b78bd61e..216d7392a 100644 --- a/lib/crates/fabro-server/src/demo/mod.rs +++ b/lib/crates/fabro-server/src/demo/mod.rs @@ -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, diff --git a/lib/crates/fabro-server/src/server.rs b/lib/crates/fabro-server/src/server.rs index 15ca01ebe..448ff14ee 100644 --- a/lib/crates/fabro-server/src/server.rs +++ b/lib/crates/fabro-server/src/server.rs @@ -602,18 +602,6 @@ pub(crate) struct ResolvedAppStateSettings { pub(crate) manifest_run_settings: std::result::Result, } -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); } } diff --git a/lib/crates/fabro-server/src/server/handler/billing.rs b/lib/crates/fabro-server/src/server/handler/billing.rs index 73ff31d1b..c8321e2c4 100644 --- a/lib/crates/fabro-server/src/server/handler/billing.rs +++ b/lib/crates/fabro-server/src/server/handler/billing.rs @@ -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) -> Option 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() } diff --git a/lib/crates/fabro-server/src/server/tests.rs b/lib/crates/fabro-server/src/server/tests.rs index 5b6002bb4..3cf2f21ea 100644 --- a/lib/crates/fabro-server/src/server/tests.rs +++ b/lib/crates/fabro-server/src/server/tests.rs @@ -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, diff --git a/lib/crates/fabro-store/src/run_state.rs b/lib/crates/fabro-store/src/run_state.rs index e184bba6b..92fafa906 100644 --- a/lib/crates/fabro-store/src/run_state.rs +++ b/lib/crates/fabro-store/src/run_state.rs @@ -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); } diff --git a/lib/crates/fabro-store/tests/serializable_projection.rs b/lib/crates/fabro-store/tests/serializable_projection.rs index 5fed2cf4a..2dfd8269e 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::{ - 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] diff --git a/lib/crates/fabro-types/src/run_event/agent.rs b/lib/crates/fabro-types/src/run_event/agent.rs index 73f5b85f9..fad5215cd 100644 --- a/lib/crates/fabro-types/src/run_event/agent.rs +++ b/lib/crates/fabro-types/src/run_event/agent.rs @@ -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, diff --git a/lib/crates/fabro-types/src/run_projection.rs b/lib/crates/fabro-types/src/run_projection.rs index efed303c0..eab52571e 100644 --- a/lib/crates/fabro-types/src/run_projection.rs +++ b/lib/crates/fabro-types/src/run_projection.rs @@ -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, #[serde(default, skip_serializing_if = "Option::is_none")] pub duration_ms: Option, - /// 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, + #[serde(default)] + pub usage: BilledTokenCounts, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub state: Option, } @@ -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, diff --git a/lib/crates/fabro-workflow/src/billing_rollup.rs b/lib/crates/fabro-workflow/src/billing_rollup.rs index 1f116c34f..3538c7661 100644 --- a/lib/crates/fabro-workflow/src/billing_rollup.rs +++ b/lib/crates/fabro-workflow/src/billing_rollup.rs @@ -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, + pub model: Option, } #[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::::new(); let mut stages = Vec::::new(); - let mut by_model = BTreeMap::::new(); + let mut by_model = HashMap::::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::>(); + 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()); diff --git a/lib/crates/fabro-workflow/src/event/convert.rs b/lib/crates/fabro-workflow/src/event/convert.rs index 1dacbb917..ff2d1ef25 100644 --- a/lib/crates/fabro-workflow/src/event/convert.rs +++ b/lib/crates/fabro-workflow/src/event/convert.rs @@ -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, }, diff --git a/lib/crates/fabro-workflow/src/event/stored_fields.rs b/lib/crates/fabro-workflow/src/event/stored_fields.rs index 80c610b53..afa202c1f 100644 --- a/lib/crates/fabro-workflow/src/event/stored_fields.rs +++ b/lib/crates/fabro-workflow/src/event/stored_fields.rs @@ -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 { .. } diff --git a/lib/crates/fabro-workflow/src/pipeline/finalize.rs b/lib/crates/fabro-workflow/src/pipeline/finalize.rs index be68d208c..1eb2af3e7 100644 --- a/lib/crates/fabro-workflow/src/pipeline/finalize.rs +++ b/lib/crates/fabro-workflow/src/pipeline/finalize.rs @@ -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,