From d6f1ab87c0d4b74975571b19ac5b9e4c5e8fec97 Mon Sep 17 00:00:00 2001 From: Scott Werner Date: Fri, 9 Oct 2026 16:57:59 -0400 Subject: [PATCH 1/2] Keep subagent usage when an agent stage completes A finished agent stage replaced its live usage, which included every subagent session, with the parent session's usage alone, so completed stage and run totals dropped descendant tokens and spend. The subagent accounts were also read as an array, while Petri reports them as an object keyed by session id, so no per-model rows were ever built. The completed stage now totals the parent's usage and each descendant account once, and builds per-model rows for the parent and every descendant. The parent's usage goes under the routes Petri reports in `pebble.usage_by_model` when present (lithoscomputer/petri#110), and under the stage's model until then. Fixes #933 Co-Authored-By: Claude Opus 5.5 --- .../fabro-petri/src/projection/engine.rs | 301 ++++++++++++++++-- 1 file changed, 276 insertions(+), 25 deletions(-) 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] From 81d63a0683f4881d642eab63be62a1ff99bf7cb8 Mon Sep 17 00:00:00 2001 From: Scott Werner Date: Mon, 12 Oct 2026 15:38:02 -0400 Subject: [PATCH 2/2] Track Petri's root usage breakdown by model Bump Petri to the commit that reports `pebble.usage_by_model`, the root session's usage by the route that spent it, and add an end-to-end test: an agent stage on a scripted model whose primary answers once and then fails mid-prompt keeps every token Petri reported, and files the primary's share and the fallback's share under their own models. Only Petri moves; Pebble and sandbox-driver stay at their locked revisions. Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 30 +-- .../fabro-petri/tests/agent_usage.rs | 218 ++++++++++++++++++ 2 files changed, 233 insertions(+), 15 deletions(-) create mode 100644 lib/components/fabro-petri/tests/agent_usage.rs diff --git a/Cargo.lock b/Cargo.lock index fc80f4d81..781078207 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/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), + ]); +}