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), + ]); +}