use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use agentplane::core::ProtectedField;
use agentplane::manifest::{Budgets, Manifest, ModelRef, ToolGrant};
use agentplane::model::fake::FakeProvider;
use agentplane::model::{ModelProvider, ToolDeclaration, Transcript};
use agentplane::prelude::*;
use agentplane::runtime::Agent;
use agentplane::tools::{Tool, ToolBox, ToolFailure};
use serde_json::{Value, json};
static FORMATTED: AtomicUsize = AtomicUsize::new(0);
#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
struct ReadFile {
path: String,
}
#[async_trait]
impl Tool for ReadFile {
const SERVER: &'static str = "fs";
const NAME: &'static str = "read";
fn mutates() -> bool {
false
}
async fn call(self) -> Result<Value, ToolFailure> {
Ok(json!({ "path": self.path, "text": "fn main(){println!(\"hi\")}" }))
}
}
#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
struct Format {}
#[async_trait]
impl Tool for Format {
const SERVER: &'static str = "repo";
const NAME: &'static str = "format";
async fn call(self) -> Result<Value, ToolFailure> {
FORMATTED.fetch_add(1, Ordering::SeqCst);
Ok(json!({ "changed": 1 }))
}
}
#[derive(Debug)]
struct Tidy {
offered: Vec<ToolDeclaration>,
}
#[async_trait]
impl Skill for Tidy {
fn descriptor(&self) -> SkillDescriptor {
SkillDescriptor::new("tidy").provides("code.tidy")
}
async fn invoke(
&self,
cx: &mut StepCtx<'_>,
input: Tainted<Value>,
) -> Result<Outcome, SkillError> {
let prompt = Tainted::object([
(
"system",
Tainted::trusted(json!("Tidy the working tree, then say what changed.")),
),
(
"messages",
Tainted::array([input.map(|i| json!({ "role": "user", "content": i }))]),
),
]);
let mut transcript = Transcript::new(&prompt);
for _turn in 0..4 {
let answer = cx
.complete_with(&prompt, |call| {
call.with_tools(self.offered.clone())
.continuing(&transcript)
})
.await?;
let calls = transcript.asked(&answer);
if calls.is_empty() {
return Ok(Outcome::done(answer.map(|c| json!({ "text": c.text }))));
}
for call in calls {
let arguments = Tainted::with_label(call.arguments.clone(), answer.label().clone());
match cx.call_tool_for(&call, &self.offered, arguments).await {
Ok(output) => transcript.answered(call, &output),
Err(error) => match cx.relayed(&error) {
Some(detail) => transcript.failed(call, detail),
None => return Err(error.into()),
},
}
}
}
Ok(Outcome::fail("the model did not finish within four turns"))
}
}
fn declaration() -> Result<Manifest, agentplane::manifest::ManifestError> {
Manifest::builder("tidier", "1.0.0")
.provides("code.tidy")
.privileged_model(ModelRef::new("fake", "coder-1"))
.egress_ceiling(Sensitivity::Internal)
.tool(ToolGrant {
description: Some("Read a file.".to_owned()),
mutates: false,
max_sensitivity: Some(Sensitivity::Internal),
..ToolGrant::new("tool://fs/read")
})
.tool(ToolGrant {
description: Some("Format the working tree.".to_owned()),
max_sensitivity: Some(Sensitivity::Internal),
protected_fields: vec![ProtectedField::one_of("", [json!({})])],
..ToolGrant::new("tool://repo/format")
})
.budgets(Budgets {
max_tokens: Some(100_000),
..Budgets::default()
})
.build()
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let store = Arc::new(RedbStore::open_in_memory()?);
let toolbox = ToolBox::new().with::<ReadFile>().with::<Format>();
let offered = [
(ToolId::new("fs", "read"), "read_file"),
(ToolId::new("repo", "format"), "format"),
]
.into_iter()
.map(|(id, name)| {
let (description, schema, _) = toolbox.declared(&id).expect("registered");
let (description, schema) = (description.to_owned(), schema.clone());
ToolDeclaration::offering(id, description, schema).named(name)
})
.collect();
let provider = FakeProvider::new();
provider
.will_call_tool("c1", "read_file", json!({ "path": "src/main.rs" }))
.will_call_tool("c2", "format", json!({}))
.will_say("Formatted src/main.rs.");
let driver: Arc<dyn ModelProvider> = provider.clone();
let manifest = declaration()?;
let rt = Runtime::builder_on(Arc::clone(&store))
.provider("fake", driver)
.agent(Agent::new(&manifest).skill(Tidy { offered }))
.toolbox(toolbox)
.build();
let run = rt
.start("code.tidy", Tainted::trusted(json!("tidy up")))
.run()
.await?;
println!("run → {}", run.status.as_str());
assert!(
matches!(run.status, RunStatus::Succeeded),
"{:?}",
run.status
);
assert_eq!(FORMATTED.load(Ordering::SeqCst), 1);
let label = run.output.as_ref().expect("an answer").label();
println!("answer → read {} source(s)", label.provenance.len());
assert!(
label
.provenance
.iter()
.any(|s| s.to_string() == "tool://fs/read"),
"{label:?}"
);
let turns = provider.asked();
println!("turns → {}", turns.len());
assert_eq!(turns.len(), 3);
assert_eq!(turns[2].exchanges[0].call.name, "format");
let replayed = rt
.replay(run.run_id, agentplane::runtime::Mode::Strict)
.await?;
println!("replay → {}", replayed.status.as_str());
assert_eq!(
FORMATTED.load(Ordering::SeqCst),
1,
"replay formatted nothing"
);
Ok(())
}