mirror of
https://github.com/fabro-sh/fabro.git
synced 2026-10-13 04:01:43 +00:00
Merge pull request #996 from fabro-sh/preserve-subagent-usage
Some checks are pending
Rust / Format (push) Waiting to run
Rust / Test (macOS) (push) Waiting to run
Rust / Clippy (push) Waiting to run
Rust / Rustdoc (push) Waiting to run
Rust / Generated Docs (push) Waiting to run
Rust / Test (Linux) (push) Waiting to run
Rust / Sandbox providers (Docker) (push) Waiting to run
Rust / Process titles (musl) (push) Waiting to run
TypeScript / Typecheck (push) Waiting to run
TypeScript / Test (push) Waiting to run
TypeScript / Build (push) Waiting to run
Some checks are pending
Rust / Format (push) Waiting to run
Rust / Test (macOS) (push) Waiting to run
Rust / Clippy (push) Waiting to run
Rust / Rustdoc (push) Waiting to run
Rust / Generated Docs (push) Waiting to run
Rust / Test (Linux) (push) Waiting to run
Rust / Sandbox providers (Docker) (push) Waiting to run
Rust / Process titles (musl) (push) Waiting to run
TypeScript / Typecheck (push) Waiting to run
TypeScript / Test (push) Waiting to run
TypeScript / Build (push) Waiting to run
Keep subagent usage when an agent stage completes
This commit is contained in:
Scott Werner
GitHub
commit
d5b5435c11
3 changed files with 509 additions and 40 deletions
30
Cargo.lock
generated
30
Cargo.lock
generated
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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::<Option<Vec<_>>>())
|
||||
.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<ModelUsage> = 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<ModelUsage> = 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<ModelRef>, 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<String, Value> = 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]
|
||||
|
|
|
|||
218
lib/components/fabro-petri/tests/agent_usage.rs
Normal file
218
lib/components/fabro-petri/tests/agent_usage.rs
Normal file
|
|
@ -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 `<root>/.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<MemoryRunStore>,
|
||||
) -> 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<dyn RunStore>)
|
||||
.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<SmolStr, Value> {
|
||||
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<String> = 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),
|
||||
]);
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue