mod depth;
mod execute;
mod parse;
use std::collections::BTreeSet;
use async_trait::async_trait;
use mentra::tool::{
ParallelToolContext, RuntimeToolDescriptor, ToolApprovalCategory, ToolAuthorizationPreview,
ToolCapability, ToolContext, ToolDefinition, ToolDurability, ToolExecutionCategory,
ToolExecutor, ToolResult, ToolSideEffectLevel,
};
use serde_json::{Value, json};
use parse::{INPUT_FIELD, Mode, Spawn, parse};
pub(crate) use parse::{LOCAL_TARGET, Mode as SpawnMode, is_target_name, parse as parse_spawn};
pub use depth::MAX_DEPTH;
pub const SPAWN: &str = "spawn";
const DESCRIPTION: &str = "\
Hand work to something else and read the result back.
Takes one string. If it starts with `!`, the rest runs as a shell command in \
this workspace and you get its output — `!cargo test -q`. Anything else is a \
task handed to a subagent, which works on it and returns a final answer — \
`find every TODO under src/ and summarise them`.
To delegate a task whose own text starts with `!`, double it: `!!important, …`.
Commands are put to the operator before they run, so ask for one command that \
does the job rather than several that each need answering.";
fn targets_paragraph(targets: &[String]) -> String {
format!(
"A command can also say where it runs: `!@<target> <command>` runs it on that target \
rather than here — `!@{first} <command>`. Registered targets: {names}. A command with \
no `@` runs where basis itself is running, which is what you want unless the work \
needs one of those targets.",
first = targets[0],
names = listed(targets),
)
}
fn listed(targets: &[String]) -> String {
targets
.iter()
.map(|name| format!("`{name}`"))
.collect::<Vec<_>>()
.join(", ")
}
#[derive(Debug, Default)]
pub struct SpawnTool {
depth: depth::Depth,
targets: Vec<String>,
}
impl SpawnTool {
pub fn new() -> Self {
Self::default()
}
pub fn with_targets(targets: impl IntoIterator<Item = String>) -> Self {
Self {
depth: depth::Depth::default(),
targets: targets
.into_iter()
.collect::<BTreeSet<_>>()
.into_iter()
.collect(),
}
}
fn authorize_target(&self, spawn: &Spawn) -> Result<(), String> {
let Some(target) = spawn.target() else {
return Ok(());
};
if self.targets.iter().any(|name| name == target) {
return Ok(());
}
Err(if self.targets.is_empty() {
format!(
"spawn has no command targets registered, so `!@{target}` has nowhere to run; \
write the command with no `@` to run it where basis is running"
)
} else {
format!(
"spawn has no command target named `{target}`; the registered targets are {}, \
and a command with no `@` runs where basis is running",
listed(&self.targets)
)
})
}
}
impl ToolDefinition for SpawnTool {
fn descriptor(&self) -> RuntimeToolDescriptor {
RuntimeToolDescriptor::builder(SPAWN)
.description(description(&self.targets))
.input_schema(input_schema(&self.targets))
.capabilities(vec![
ToolCapability::ProcessExec,
ToolCapability::FilesystemWrite,
ToolCapability::Delegation,
])
.side_effect_level(ToolSideEffectLevel::Process)
.durability(ToolDurability::Ephemeral)
.execution_category(ToolExecutionCategory::ExclusiveLocalMutation)
.approval_category(ToolApprovalCategory::Process)
.build()
}
}
fn description(targets: &[String]) -> String {
if targets.is_empty() {
return DESCRIPTION.to_string();
}
format!("{DESCRIPTION}\n\n{}", targets_paragraph(targets))
}
fn input_schema(targets: &[String]) -> Value {
let mut description = "A shell command when it starts with `!`, otherwise a task to \
delegate. Write `!!` to begin a task with a literal `!`."
.to_string();
if !targets.is_empty() {
description.push_str(&format!(
" Write `!@<target> <command>` to run a command on one of this runtime's targets \
({}).",
listed(targets)
));
}
json!({
"type": "object",
"properties": {
INPUT_FIELD: {
"type": "string",
"description": description,
}
},
"required": [INPUT_FIELD],
})
}
#[async_trait]
impl ToolExecutor for SpawnTool {
fn authorization_preview(
&self,
ctx: &ParallelToolContext,
input: &Value,
) -> Result<ToolAuthorizationPreview, String> {
let spawn = parse(input)?;
self.authorize_target(&spawn)?;
if spawn.mode() == Mode::Agent {
self.depth.authorize_delegation(&ctx.agent_id)?;
}
let cwd = ctx.resolve_working_directory(None)?;
Ok(preview(&spawn, cwd, &self.descriptor(), input))
}
fn execution_category(&self, input: &Value) -> ToolExecutionCategory {
parse(input).map_or(ToolExecutionCategory::ExclusiveLocalMutation, |spawn| {
execution_category(spawn.mode())
})
}
async fn execute_mut(&self, mut ctx: ToolContext<'_>, input: Value) -> ToolResult {
let spawn = parse(&input)?;
match spawn.mode() {
Mode::Command => {
self.authorize_target(&spawn)?;
execute::command(&ctx, spawn.body(), spawn.target()).await
}
Mode::Agent => {
let depth = self.depth.authorize_delegation(&ctx.agent_id)?;
execute::delegate(&self.depth, &mut ctx, spawn.body(), depth).await
}
}
}
}
fn preview(
spawn: &Spawn,
cwd: std::path::PathBuf,
descriptor: &RuntimeToolDescriptor,
raw_input: &Value,
) -> ToolAuthorizationPreview {
let structured_input = json!({
"mode": spawn.mode().as_str(),
"body": spawn.body(),
"cwd": cwd,
"target": spawn.target().unwrap_or(LOCAL_TARGET),
});
ToolAuthorizationPreview {
working_directory: cwd,
capabilities: capabilities(spawn.mode()),
side_effect_level: side_effect_level(spawn.mode()),
durability: descriptor.durability,
execution_category: execution_category(spawn.mode()),
approval_category: approval_category(spawn.mode()),
raw_input: raw_input.clone(),
structured_input,
}
}
const fn side_effect_level(mode: Mode) -> ToolSideEffectLevel {
match mode {
Mode::Command => ToolSideEffectLevel::Process,
Mode::Agent => ToolSideEffectLevel::LocalState,
}
}
fn capabilities(mode: Mode) -> Vec<ToolCapability> {
match mode {
Mode::Command => vec![ToolCapability::ProcessExec, ToolCapability::FilesystemWrite],
Mode::Agent => vec![ToolCapability::Delegation],
}
}
const fn approval_category(mode: Mode) -> ToolApprovalCategory {
match mode {
Mode::Command => ToolApprovalCategory::Process,
Mode::Agent => ToolApprovalCategory::Delegation,
}
}
const fn execution_category(mode: Mode) -> ToolExecutionCategory {
match mode {
Mode::Command => ToolExecutionCategory::ExclusiveLocalMutation,
Mode::Agent => ToolExecutionCategory::Delegation,
}
}
#[cfg(test)]
mod tests;