use std::{
path::{Path, PathBuf},
sync::Arc,
time::Duration,
};
use async_trait::async_trait;
use mentra::tool::{
ParallelToolContext, RuntimeToolDescriptor, ToolApprovalCategory, ToolAuthorizationPreview,
ToolCapability, ToolDefinition, ToolDurability, ToolExecutionCategory, ToolExecutor,
ToolResult,
};
use serde_json::{Value, json};
use crate::subprocess::{self, Completion};
use super::manifest::DeclaredToolSpec;
const TIMEOUT_BACKSTOP: Duration = Duration::from_secs(5);
#[derive(Debug, Clone)]
pub struct DeclaredTool {
spec: Arc<DeclaredToolSpec>,
workspace: PathBuf,
}
impl DeclaredTool {
pub fn new(spec: DeclaredToolSpec, workspace: impl Into<PathBuf>) -> Self {
Self {
spec: Arc::new(spec),
workspace: workspace.into(),
}
}
pub fn name(&self) -> &str {
&self.spec.name
}
pub fn spec(&self) -> &DeclaredToolSpec {
&self.spec
}
}
impl ToolDefinition for DeclaredTool {
fn descriptor(&self) -> RuntimeToolDescriptor {
RuntimeToolDescriptor::builder(&self.spec.name)
.description(&self.spec.description)
.input_schema(self.spec.input_schema.clone())
.capabilities(vec![ToolCapability::ProcessExec])
.side_effect_level(self.spec.side_effect.level())
.durability(ToolDurability::Ephemeral)
.execution_category(ToolExecutionCategory::ExclusiveLocalMutation)
.approval_category(ToolApprovalCategory::Process)
.execution_timeout(self.spec.timeout() + TIMEOUT_BACKSTOP)
.build()
}
}
#[async_trait]
impl ToolExecutor for DeclaredTool {
fn authorization_preview(
&self,
_ctx: &ParallelToolContext,
input: &Value,
) -> Result<ToolAuthorizationPreview, String> {
preview(&self.spec, &self.workspace, &self.descriptor(), input)
}
async fn execute(&self, _ctx: ParallelToolContext, input: Value) -> ToolResult {
run(Arc::clone(&self.spec), self.workspace.clone(), input).await
}
}
fn preview(
spec: &DeclaredToolSpec,
workspace: &Path,
descriptor: &RuntimeToolDescriptor,
input: &Value,
) -> Result<ToolAuthorizationPreview, String> {
check_input(spec, input)?;
let cwd = spec.working_directory(workspace);
Ok(ToolAuthorizationPreview {
capabilities: descriptor.capabilities.clone(),
side_effect_level: descriptor.side_effect_level,
durability: descriptor.durability,
execution_category: descriptor.execution_category,
approval_category: descriptor.approval_category,
raw_input: input.clone(),
structured_input: json!({
"tool": spec.name,
"command": spec.command,
"cwd": cwd,
"input": input,
}),
working_directory: cwd,
})
}
async fn run(spec: Arc<DeclaredToolSpec>, workspace: PathBuf, input: Value) -> ToolResult {
check_input(&spec, &input)?;
let payload = serde_json::to_string(&input).map_err(|error| {
format!(
"{} was called with input basis could not serialize: {error}",
spec.name
)
})?;
let running = Arc::clone(&spec);
let completion = tokio::task::spawn_blocking(move || {
subprocess::execute(
&running.command,
&running.working_directory(&workspace),
&running.env,
&payload,
running.timeout(),
)
})
.await
.map_err(|error| format!("{} could not be run: {error}", spec.name))?;
answer(&spec, completion)
}
fn answer(spec: &DeclaredToolSpec, completion: std::io::Result<Completion>) -> ToolResult {
let completion = completion.map_err(|error| {
format!("{} could not be started: {error}", spec.name)
})?;
let (code, stdout, stderr) = match completion {
Completion::TimedOut => {
return Err(format!(
"{} did not finish within {} seconds and was stopped",
spec.name,
spec.timeout().as_secs()
));
}
Completion::Exited {
code,
stdout,
stderr,
} => (code, stdout, stderr),
};
match code {
Some(0) => Ok(succeeded(spec, stdout)),
Some(code) => Err(failed(spec, code, &stdout, &stderr)),
None => Err(format!(
"{} was killed by a signal before it answered",
spec.name
)),
}
}
fn succeeded(spec: &DeclaredToolSpec, stdout: String) -> String {
if stdout.trim().is_empty() {
return format!("{} finished and printed nothing", spec.name);
}
stdout.trim_end_matches('\n').to_string()
}
fn failed(spec: &DeclaredToolSpec, code: i32, stdout: &str, stderr: &str) -> String {
let explanation = if stderr.trim().is_empty() {
subprocess::truncated_output(stdout)
} else {
stderr.trim().to_string()
};
if explanation.is_empty() {
return format!("{} exited {code} and said nothing", spec.name);
}
format!("{} exited {code}: {explanation}", spec.name)
}
fn check_input(spec: &DeclaredToolSpec, input: &Value) -> Result<(), String> {
if !input.is_object() {
return Err(format!(
"{} takes a JSON object matching its input schema",
spec.name
));
}
let missing: Vec<&str> = spec
.input_schema
.get("required")
.and_then(Value::as_array)
.map(|required| {
required
.iter()
.filter_map(Value::as_str)
.filter(|field| input.get(*field).is_none())
.collect()
})
.unwrap_or_default();
if missing.is_empty() {
return Ok(());
}
Err(format!(
"{} was called without {}, which its input schema requires",
spec.name,
missing
.iter()
.map(|field| format!("`{field}`"))
.collect::<Vec<_>>()
.join(", ")
))
}
#[cfg(test)]
mod tests;