diff --git a/Cargo.lock b/Cargo.lock index 56e38dbf6..a829bab70 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5166,7 +5166,7 @@ checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" [[package]] name = "petri-attractor-steps" version = "0.1.0" -source = "git+https://github.com/lithoscomputer/petri.git?branch=main#2548c104e1b0b5ac710abc59b25724d6b2c4d66d" +source = "git+https://github.com/lithoscomputer/petri.git?branch=main#e94a2fec1a4c88de43c46ac527968f67f1753b2b" dependencies = [ "async-trait", "globset", @@ -5197,7 +5197,7 @@ dependencies = [ [[package]] name = "petri-driver" version = "0.1.0" -source = "git+https://github.com/lithoscomputer/petri.git?branch=main#2548c104e1b0b5ac710abc59b25724d6b2c4d66d" +source = "git+https://github.com/lithoscomputer/petri.git?branch=main#e94a2fec1a4c88de43c46ac527968f67f1753b2b" dependencies = [ "async-trait", "getrandom 0.3.4", @@ -5217,7 +5217,7 @@ dependencies = [ [[package]] name = "petri-engine" version = "0.1.0" -source = "git+https://github.com/lithoscomputer/petri.git?branch=main#2548c104e1b0b5ac710abc59b25724d6b2c4d66d" +source = "git+https://github.com/lithoscomputer/petri.git?branch=main#e94a2fec1a4c88de43c46ac527968f67f1753b2b" dependencies = [ "petri-ir", "serde", @@ -5229,7 +5229,7 @@ dependencies = [ [[package]] name = "petri-execution" version = "0.1.0" -source = "git+https://github.com/lithoscomputer/petri.git?branch=main#2548c104e1b0b5ac710abc59b25724d6b2c4d66d" +source = "git+https://github.com/lithoscomputer/petri.git?branch=main#e94a2fec1a4c88de43c46ac527968f67f1753b2b" dependencies = [ "async-trait", "petri-driver", @@ -5253,7 +5253,7 @@ dependencies = [ [[package]] name = "petri-executor" version = "0.1.0" -source = "git+https://github.com/lithoscomputer/petri.git?branch=main#2548c104e1b0b5ac710abc59b25724d6b2c4d66d" +source = "git+https://github.com/lithoscomputer/petri.git?branch=main#e94a2fec1a4c88de43c46ac527968f67f1753b2b" dependencies = [ "async-trait", "libc", @@ -5268,7 +5268,7 @@ dependencies = [ [[package]] name = "petri-executor-sandbox" version = "0.1.0" -source = "git+https://github.com/lithoscomputer/petri.git?branch=main#2548c104e1b0b5ac710abc59b25724d6b2c4d66d" +source = "git+https://github.com/lithoscomputer/petri.git?branch=main#e94a2fec1a4c88de43c46ac527968f67f1753b2b" dependencies = [ "async-trait", "petri-executor", @@ -5290,7 +5290,7 @@ dependencies = [ [[package]] name = "petri-frontend" version = "0.1.0" -source = "git+https://github.com/lithoscomputer/petri.git?branch=main#2548c104e1b0b5ac710abc59b25724d6b2c4d66d" +source = "git+https://github.com/lithoscomputer/petri.git?branch=main#e94a2fec1a4c88de43c46ac527968f67f1753b2b" dependencies = [ "marked-yaml", "petri-ir", @@ -5304,7 +5304,7 @@ dependencies = [ [[package]] name = "petri-frontend-attractor" version = "0.1.0" -source = "git+https://github.com/lithoscomputer/petri.git?branch=main#2548c104e1b0b5ac710abc59b25724d6b2c4d66d" +source = "git+https://github.com/lithoscomputer/petri.git?branch=main#e94a2fec1a4c88de43c46ac527968f67f1753b2b" dependencies = [ "minijinja", "petri-frontend", @@ -5321,7 +5321,7 @@ dependencies = [ [[package]] name = "petri-frontend-fabro" version = "0.1.0" -source = "git+https://github.com/lithoscomputer/petri.git?branch=main#2548c104e1b0b5ac710abc59b25724d6b2c4d66d" +source = "git+https://github.com/lithoscomputer/petri.git?branch=main#e94a2fec1a4c88de43c46ac527968f67f1753b2b" dependencies = [ "petri-frontend", "petri-frontend-attractor", @@ -5337,7 +5337,7 @@ dependencies = [ [[package]] name = "petri-frontend-native" version = "0.1.0" -source = "git+https://github.com/lithoscomputer/petri.git?branch=main#2548c104e1b0b5ac710abc59b25724d6b2c4d66d" +source = "git+https://github.com/lithoscomputer/petri.git?branch=main#e94a2fec1a4c88de43c46ac527968f67f1753b2b" dependencies = [ "petri-frontend", "petri-ir", @@ -5348,7 +5348,7 @@ dependencies = [ [[package]] name = "petri-ir" version = "0.1.0" -source = "git+https://github.com/lithoscomputer/petri.git?branch=main#2548c104e1b0b5ac710abc59b25724d6b2c4d66d" +source = "git+https://github.com/lithoscomputer/petri.git?branch=main#e94a2fec1a4c88de43c46ac527968f67f1753b2b" dependencies = [ "regex", "serde", @@ -5361,7 +5361,7 @@ dependencies = [ [[package]] name = "petri-runtime" version = "0.1.0" -source = "git+https://github.com/lithoscomputer/petri.git?branch=main#2548c104e1b0b5ac710abc59b25724d6b2c4d66d" +source = "git+https://github.com/lithoscomputer/petri.git?branch=main#e94a2fec1a4c88de43c46ac527968f67f1753b2b" dependencies = [ "async-trait", "petri-driver", @@ -5382,7 +5382,7 @@ dependencies = [ [[package]] name = "petri-steps" version = "0.1.0" -source = "git+https://github.com/lithoscomputer/petri.git?branch=main#2548c104e1b0b5ac710abc59b25724d6b2c4d66d" +source = "git+https://github.com/lithoscomputer/petri.git?branch=main#e94a2fec1a4c88de43c46ac527968f67f1753b2b" dependencies = [ "async-trait", "petri-executor", @@ -5398,7 +5398,7 @@ dependencies = [ [[package]] name = "petri-store" version = "0.1.0" -source = "git+https://github.com/lithoscomputer/petri.git?branch=main#2548c104e1b0b5ac710abc59b25724d6b2c4d66d" +source = "git+https://github.com/lithoscomputer/petri.git?branch=main#e94a2fec1a4c88de43c46ac527968f67f1753b2b" dependencies = [ "async-trait", "getrandom 0.3.4", @@ -5413,7 +5413,7 @@ dependencies = [ [[package]] name = "petri-testkit" version = "0.1.0" -source = "git+https://github.com/lithoscomputer/petri.git?branch=main#2548c104e1b0b5ac710abc59b25724d6b2c4d66d" +source = "git+https://github.com/lithoscomputer/petri.git?branch=main#e94a2fec1a4c88de43c46ac527968f67f1753b2b" dependencies = [ "async-trait", "petri-driver", diff --git a/lib/components/fabro-petri/src/projection/engine.rs b/lib/components/fabro-petri/src/projection/engine.rs index 597737621..1c16ba328 100644 --- a/lib/components/fabro-petri/src/projection/engine.rs +++ b/lib/components/fabro-petri/src/projection/engine.rs @@ -6,10 +6,11 @@ use std::collections::BTreeMap; use chrono::{DateTime, Utc}; use fabro_types::{ - ModelUsage, ParallelBranchId, ParallelBranchResult, RunProjection, RunSandbox, + ModelRef, ModelUsage, ParallelBranchId, ParallelBranchResult, RunProjection, RunSandbox, RunSandboxFailure, StageCompletion, StageHandler, StageId, StageOutcome, StageProjection, StageState, StageTiming, first_event_seq, parse_blob_ref, timing, }; +use lithos_llm::types::Usage; use petri_execution::events::{Derived, RunEvent, Subject, ViewEvent, WaitState}; use petri_execution::{ExecutionId, InvocationId}; use petri_runtime::engine::{Admission, Event}; @@ -423,44 +424,294 @@ fn apply_metrics(stage: &mut StageProjection, metrics: &Metrics) { _ => (inference, tool), }; stage.set_authoritative_timing(StageTiming::new(wall, inference, tool).clamped_to_wall()); - if let Some(usage) = - usage_of(custom.get("pebble.usage")).or_else(|| usage_of(custom.get("prompt.usage"))) - { + if let Some(root) = usage_of(custom.get("pebble.usage")) { + apply_agent_usage(stage, metrics, root); + } else if let Some(usage) = usage_of(custom.get("prompt.usage")) { stage.usage = usage; } - if let Some(sessions) = custom +} + +/// A coding agent's usage onto its finished stage: the root session's own +/// and every descendant session's, each counted once. The root's usage goes +/// under the routes `pebble.usage_by_model` names, or under the stage's model +/// when Petri does not break it down; each descendant's goes under its own +/// model, or the stage's when it has none. +fn apply_agent_usage(stage: &mut StageProjection, metrics: &Metrics, root: Usage) { + let custom = &metrics.custom; + let root_rows = custom + .get("pebble.usage_by_model") + .and_then(Value::as_array) + .and_then(|routes| routes.iter().map(account).collect::>>()) + .unwrap_or_else(|| vec![(None, root)]); + let descendants: Vec<_> = custom .get("pebble.subagents") .and_then(|subagents| subagents.get("sessions")) - .and_then(Value::as_array) - { - let mut by_model: Vec = Vec::new(); - for session in sessions { - let provider = session.get("provider").and_then(Value::as_str); - let model = session.get("model").and_then(Value::as_str); - let Some(usage) = usage_of(session.get("usage")) else { - continue; - }; - let Some(model) = model.and_then(|model| model_ref(provider, model)) else { - continue; - }; - if let Some(entry) = by_model.iter_mut().find(|entry| entry.model == model) { - entry.usage = entry.usage.saturating_add(usage); - } else { - by_model.push(ModelUsage::new(model, usage)); - } - } - if !by_model.is_empty() { - stage.usage_by_model = by_model; + .and_then(Value::as_object) + .into_iter() + .flat_map(|sessions| sessions.values().filter_map(account)) + .collect(); + + stage.usage = descendants + .iter() + .fold(root, |total, (_, usage)| total.saturating_add(*usage)); + let mut by_model: Vec = Vec::new(); + for (model, usage) in root_rows.into_iter().chain(descendants) { + let Some(model) = model.or_else(|| stage.model.clone()) else { + continue; + }; + if let Some(entry) = by_model.iter_mut().find(|entry| entry.model == model) { + entry.usage = entry.usage.saturating_add(usage); + } else { + by_model.push(ModelUsage::new(model, usage)); } } + stage.usage_by_model = by_model; +} + +/// One usage account in Petri's metrics, `{ provider, model, usage }`: its +/// model when it names one, and its usage. +fn account(account: &Value) -> Option<(Option, Usage)> { + let usage = usage_of(account.get("usage"))?; + let provider = account.get("provider").and_then(Value::as_str); + let model = account + .get("model") + .and_then(Value::as_str) + .and_then(|model| model_ref(provider, model)); + Some((model, usage)) } #[cfg(test)] mod tests { + use lithos_llm::catalog::{ModelId, ProviderId}; + use lithos_llm::types::{Cost, CostSource, TokenCounts}; use petri_runtime::ir::FailureInfo; + use serde_json::json; use super::*; + /// Usage in the proportions of the scripted-provider family: `input` + /// tokens, half as many output, a fifth read from cache, a tenth written + /// to it, and the catalog cost. + fn usage(input: u64, usd_micros: u64) -> Usage { + Usage { + tokens: TokenCounts { + input, + output: input / 2, + cache_read: input / 5, + cache_write: input / 10, + ..TokenCounts::default() + }, + cost: Some(Cost { + usd_micros, + source: CostSource::Catalog, + }), + } + } + + fn model(provider: &str, model: &str) -> ModelRef { + ModelRef::new(ProviderId::new(provider), ModelId::new(model)) + } + + /// Petri's `pebble.subagents` metric: the descendants' summed usage and + /// their accounts, keyed by session id. + fn subagents(accounts: &[(&str, &str, &str, &str, Usage)]) -> Value { + let usage = accounts.iter().fold(Usage::default(), |sum, account| { + sum.saturating_add(account.4) + }); + let sessions: serde_json::Map = accounts + .iter() + .map(|(session, parent, provider, model, usage)| { + ( + (*session).to_owned(), + json!({ + "parent": parent, + "provider": provider, + "model": model, + "usage": usage, + "messages": 2, + "compactions": 0, + }), + ) + }) + .collect(); + json!({ + "spawned": accounts.len(), + "turns_started": accounts.len(), + "completed": accounts.len(), + "failed": 0, + "closed": accounts.len(), + "usage": usage, + "sessions": sessions, + }) + } + + fn agent_metrics(root: Usage, subagents: Value) -> Metrics { + let mut metrics = Metrics::default(); + metrics.custom.insert("pebble.usage".into(), json!(root)); + metrics.custom.insert("pebble.subagents".into(), subagents); + metrics + } + + fn agent_stage(model: ModelRef) -> StageProjection { + let mut stage = StageProjection::new(first_event_seq(1)); + stage.handler = Some(StageHandler::Agent); + stage.model = Some(model); + stage + } + + /// A parent, its child and its grandchild on one model: the finished + /// stage keeps all three sessions' usage, counted once, under that model. + #[test] + fn a_finished_agent_keeps_its_descendants_usage() { + let sonnet = model("anthropic", "claude-sonnet-5"); + let mut stage = agent_stage(sonnet.clone()); + let metrics = agent_metrics( + usage(20_000, 145_800), + subagents(&[ + ( + "child", + "root", + "anthropic", + "claude-sonnet-5", + usage(40_000, 291_600), + ), + ( + "grandchild", + "child", + "anthropic", + "claude-sonnet-5", + usage(60_000, 437_400), + ), + ]), + ); + + apply_metrics(&mut stage, &metrics); + + assert_eq!(stage.usage, usage(120_000, 874_800)); + assert_eq!(stage.usage_by_model, vec![ModelUsage::new( + sonnet, + usage(120_000, 874_800) + )]); + } + + /// Each session's usage goes under its own model; sessions sharing a + /// model share a row, and the rows sum to the stage's usage. + #[test] + fn a_finished_agent_splits_its_usage_by_each_sessions_model() { + let sonnet = model("anthropic", "claude-sonnet-5"); + let gpt = model("openai", "gpt-5.4"); + let mut stage = agent_stage(sonnet.clone()); + let metrics = agent_metrics( + usage(20_000, 145_800), + subagents(&[ + ("child", "root", "openai", "gpt-5.4", usage(40_000, 200_000)), + ( + "grandchild", + "child", + "anthropic", + "claude-sonnet-5", + usage(60_000, 437_400), + ), + ]), + ); + + apply_metrics(&mut stage, &metrics); + + assert_eq!(stage.usage, usage(120_000, 783_200)); + assert_eq!(stage.usage_by_model, vec![ + ModelUsage::new(sonnet, usage(80_000, 583_200)), + ModelUsage::new(gpt, usage(40_000, 200_000)), + ]); + } + + /// When Petri breaks the root's usage down by route, a fallback's + /// usage stays under the model that spent it, not the stage's model. + #[test] + fn the_roots_usage_goes_under_the_routes_petri_reports() { + let sonnet = model("anthropic", "claude-sonnet-5"); + let gpt = model("openai", "gpt-5.4"); + let mut stage = agent_stage(gpt.clone()); + let mut metrics = agent_metrics( + usage(30_000, 180_000), + subagents(&[("child", "root", "openai", "gpt-5.4", usage(40_000, 200_000))]), + ); + metrics.custom.insert( + "pebble.usage_by_model".into(), + json!([ + { "provider": "anthropic", "model": "claude-sonnet-5", "usage": usage(10_000, 72_900) }, + { "provider": "openai", "model": "gpt-5.4", "usage": usage(20_000, 107_100) }, + ]), + ); + + apply_metrics(&mut stage, &metrics); + + assert_eq!(stage.usage, usage(70_000, 380_000)); + assert_eq!(stage.usage_by_model, vec![ + ModelUsage::new(sonnet, usage(10_000, 72_900)), + ModelUsage::new(gpt, usage(60_000, 307_100)), + ]); + } + + /// An agent that started no subagents keeps its own usage, as before. + #[test] + fn a_finished_agent_without_subagents_keeps_its_own_usage() { + let sonnet = model("anthropic", "claude-sonnet-5"); + let mut stage = agent_stage(sonnet.clone()); + let metrics = agent_metrics(usage(20_000, 145_800), subagents(&[])); + + apply_metrics(&mut stage, &metrics); + + assert_eq!(stage.usage, usage(20_000, 145_800)); + assert_eq!(stage.usage_by_model, vec![ModelUsage::new( + sonnet, + usage(20_000, 145_800) + )]); + } + + /// The finished attempt's metrics replace the stage's usage rather than + /// add to it, so folding the same record again (a replay) or over the + /// live total changes nothing. + #[test] + fn applying_the_same_metrics_again_does_not_double_count() { + let sonnet = model("anthropic", "claude-sonnet-5"); + let mut stage = agent_stage(sonnet); + let metrics = agent_metrics( + usage(20_000, 145_800), + subagents(&[( + "child", + "root", + "anthropic", + "claude-sonnet-5", + usage(40_000, 291_600), + )]), + ); + stage.usage = usage(60_000, 437_400); + + apply_metrics(&mut stage, &metrics); + let once = (stage.usage, stage.usage_by_model.clone()); + apply_metrics(&mut stage, &metrics); + + assert_eq!(once.0, usage(60_000, 437_400)); + assert_eq!((stage.usage, stage.usage_by_model), once); + } + + /// A prompt stage's usage is unchanged: no model rows, so the rollup + /// puts its usage under its model. + #[test] + fn a_prompt_stage_keeps_its_usage_without_model_rows() { + let mut stage = StageProjection::new(first_event_seq(1)); + stage.handler = Some(StageHandler::Prompt); + let mut metrics = Metrics::default(); + metrics + .custom + .insert("prompt.usage".into(), json!(usage(10_000, 72_900))); + + apply_metrics(&mut stage, &metrics); + + assert_eq!(stage.usage, usage(10_000, 72_900)); + assert!(stage.usage_by_model.is_empty()); + } + /// A partial success reports the failure it was converted from, a /// timeout included, and its outcome stays a partial success. #[test] diff --git a/lib/components/fabro-petri/tests/agent_usage.rs b/lib/components/fabro-petri/tests/agent_usage.rs new file mode 100644 index 000000000..038e032f4 --- /dev/null +++ b/lib/components/fabro-petri/tests/agent_usage.rs @@ -0,0 +1,218 @@ +//! An agent stage's usage, from Petri's real step metrics to the projected +//! stage: one agent stage runs on the real step registry against a scripted +//! model whose primary route answers once and then fails mid-prompt, so the +//! prompt continues on the fallback route. The finished stage keeps the +//! usage Petri reported and files each route's share under the model that +//! spent it, as `pebble.usage_by_model` breaks it down. +//! +//! Built-in Host scopes run in process without a plugin executable. + +use std::collections::BTreeMap; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::Duration; + +use fabro_petri::projection::{Item, RunView}; +use fabro_petri::runtime::RuntimeSpec; +use fabro_store::platform_records::{PlatformRecord, RunCreatedRecord, StoredPlatformRecord}; +use fabro_types::{ModelRef, RunId, StageId, test_support as types_support}; +use lithos_llm::catalog::{ModelId, ProviderId}; +use lithos_llm::types::{ErrorKind, Usage}; +use pebble_coding_agent::test_support::{ + ScriptedCall, ScriptedFailure, scripted_client, text_response, tool_call_response, +}; +use petri_execution::events::{RunEvent, replay_run}; +use petri_execution::host::{self, HostRun}; +use petri_runtime::engine::Event; +use petri_runtime::executor::Retention; +use petri_runtime::frontend::CompileInputs; +use petri_runtime::ir::RunStatus; +use petri_runtime::{RunOptions, Runtime}; +use petri_store::{Access, MemoryRunStore, RunKey, RunStore}; +use serde_json::{Value, json}; +use smol_str::SmolStr; +use tokio::fs; + +/// One agent stage on the native backend, on the scripted catalog's +/// `test/model`. +const AGENT_WORKFLOW: &str = r#"digraph Agent { + graph [goal="Say hello", backend="api", default_max_retries=0] + start [shape=Mdiamond] + exit [shape=Msquare] + work [shape=box, prompt="Say hello", model="model", provider="test", max_retries=0] + start -> work -> exit +}"#; + +/// The scripted catalog's `small` is the stage model's one fallback. +const AGENT_SETTINGS: &str = "_version = 1\n\n[workflow]\ngraph = \"workflow.fabro\"\n\n\ + [run.model.fallbacks]\n\"model\" = [\"test:small\"]\n"; + +/// Write the agent bundle into `/.fabro/workflows/agent`; the +/// workflow file. +async fn install_bundle(root: &Path) -> PathBuf { + let bundle = root.join(".fabro").join("workflows").join("agent"); + fs::create_dir_all(&bundle) + .await + .expect("the bundle directory is creatable"); + fs::write(bundle.join("workflow.fabro"), AGENT_WORKFLOW) + .await + .expect("the workflow is writable"); + fs::write(bundle.join("workflow.toml"), AGENT_SETTINGS) + .await + .expect("the settings are writable"); + bundle.join("workflow.fabro") +} + +fn runtime( + run_dir: &Path, + run_id: &str, + model_client: lithos_llm::Client, + store: &Arc, +) -> Runtime { + let mut options = RunOptions::new(run_dir); + options.grace = Duration::from_secs(2); + options.retention = Retention::Never; + options.echo = false; + options.run_key = Some(RunKey::new(run_id)); + RuntimeSpec { + model_client: Some(model_client), + ..RuntimeSpec::default() + } + .runtime(true) + .store(Arc::clone(store) as Arc) + .options(options) +} + +/// Lower the bundle and run it to its end. +async fn run(rt: &Runtime, workflow: &Path) { + let lowered = rt + .check(workflow, None, None, &CompileInputs::new()) + .expect("the workflow file loads"); + let graph = lowered + .graph + .unwrap_or_else(|| panic!("the workflow lowers: {:?}", lowered.diagnostics)); + let report = host::run_configured(rt, HostRun::new(graph), |_, _| {}) + .await + .expect("the run completes"); + assert_eq!( + report.status, + RunStatus::Success, + "errors: {:?}; history: {:#?}", + report.state.errors(), + report.state.history() + ); +} + +/// The agent's finished step metrics, as Petri recorded them. +fn agent_metrics(events: &[RunEvent]) -> BTreeMap { + events + .iter() + .find_map(|event| match event.engine() { + Some(Event::StepFinished { outcome, .. }) + if outcome.metrics.custom.contains_key("pebble.usage") => + { + Some(outcome.metrics.custom.clone()) + } + _ => None, + }) + .expect("the agent's step finished with its usage") +} + +fn usage(value: &Value) -> Usage { + serde_json::from_value(value.clone()).expect("a usage") +} + +fn model(provider: &str, model: &str) -> ModelRef { + ModelRef::new(ProviderId::new(provider), ModelId::new(model)) +} + +#[tokio::test] +async fn a_fallback_mid_prompt_keeps_each_routes_usage_under_its_model() { + let root = tempfile::tempdir().expect("a temp dir"); + let workflow = install_bundle(root.path()).await; + let run_id = RunId::new(); + let (client, provider) = scripted_client(vec![ + ScriptedCall::response(tool_call_response( + "shell", + "work", + json!({"command": "true"}), + )), + ScriptedCall::Failure(ScriptedFailure::retryable( + ErrorKind::Server, + "primary down", + )), + ScriptedCall::response(text_response("Hello from small.")), + ]); + let store = Arc::new(MemoryRunStore::new()); + let rt = runtime( + &root.path().join("run"), + &run_id.to_string(), + client, + &store, + ); + + run(&rt, &workflow).await; + + let requested: Vec = provider + .requests() + .iter() + .map(|request| request.model().to_owned()) + .collect(); + assert_eq!(requested, ["test/model", "test/model", "test/small"]); + + // The run's events, folded as the projector folds them: the run's + // `run.created` record first, then Petri's events in record order. + let logs = store + .open(&RunKey::new(run_id.to_string()), Access::Read) + .await + .expect("the run opens"); + let events = replay_run(&*logs).await.expect("the record replays"); + let mut spec = types_support::test_run_spec(); + spec.run_id = run_id; + let created = StoredPlatformRecord { + seq: 1, + recorded_at: 0, + record: PlatformRecord::RunCreated(RunCreatedRecord { + spec, + title: Some("Say hello".to_string()), + parent_id: None, + retried_from: None, + web_url: None, + }), + position: None, + }; + let mut view = RunView::new(); + view.fold(&Item::Platform(&created), 1); + for (index, event) in events.iter().enumerate() { + view.fold(&Item::Petri(event), index as u64 + 2); + } + let projection = view.projection().expect("the run has a projection"); + let stage = projection + .stage(&StageId::new("work", 1)) + .expect("the agent stage is projected"); + + // Petri's breakdown, one entry per route in first-use order. + let metrics = agent_metrics(&events); + let routes = metrics["pebble.usage_by_model"] + .as_array() + .expect("Petri breaks the root's usage down by route"); + assert_eq!(routes.len(), 2, "{metrics:?}"); + let primary = usage(&routes[0]["usage"]); + let fallback = usage(&routes[1]["usage"]); + assert!(primary.total_tokens() > 0, "{metrics:?}"); + assert!(fallback.total_tokens() > 0, "{metrics:?}"); + + // The stage keeps every token Petri reported, and each route's share + // stays under the model that spent it. + assert_eq!(stage.usage, usage(&metrics["pebble.usage"])); + assert_eq!(stage.usage, primary.saturating_add(fallback)); + let rows: Vec<(ModelRef, Usage)> = stage + .usage_by_model + .iter() + .map(|row| (row.model.clone(), row.usage)) + .collect(); + assert_eq!(rows, [ + (model("test", "model"), primary), + (model("test", "small"), fallback), + ]); +}