mod background;
mod child;
mod depth;
mod describe;
mod execute;
mod files;
mod parse;
mod reach;
mod remote;
pub(crate) mod teammate;
use std::{collections::BTreeSet, sync::Arc};
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 describe::{Names, Remote, Taught, description, input_schema, listed};
use parse::{Mode, Spawn, parse};
use teammate::Teammate;
pub use parse::{
ANY_FIELD as SPAWN_ANY_FIELD, BACKGROUND_FIELD as SPAWN_BACKGROUND_FIELD,
INPUT_FIELD as SPAWN_INPUT_FIELD, Mode as SpawnMode, REACH_FIELD as SPAWN_REACH_FIELD,
Spawn as SpawnInput, WAIT_FIELD as SPAWN_WAIT_FIELD, classify_spawn_input,
parse as parse_spawn_input,
};
pub(crate) use parse::{LOCAL_TARGET, is_namespaced, is_target_name};
pub use reach::{Reach, ReachPolicy};
pub(crate) use background::{Children, JoinAtEnd};
pub use background::{DEFAULT_DELEGATION_WIDTH, TaskFuture, TaskScope};
pub use child::{ChildContext, ChildSpec};
pub use depth::DEFAULT_DELEGATION_DEPTH;
pub(crate) use child::ChildPolicy;
pub const SPAWN: &str = "spawn";
#[derive(Clone)]
pub struct SpawnTool {
depth: depth::Depth,
width: usize,
targets: Vec<String>,
policy: Option<ChildPolicy>,
agents: Option<Arc<crate::runtime::agents::AgentRegistry>>,
remote: Option<Arc<crate::agents::remote::RemoteAgents>>,
remembers: crate::agents::remote::Remembered,
owner: Option<String>,
refresh: Option<crate::runtime::Refresh>,
reach: ReachPolicy,
task_scope: Option<TaskScope>,
}
impl Default for SpawnTool {
fn default() -> Self {
Self::with_targets_and_depth(Vec::<String>::new(), depth::DEFAULT_DELEGATION_DEPTH)
}
}
#[derive(Debug)]
enum SpawnDecision {
Command,
Agent {
depth: usize,
spec: Box<ChildSpec>,
workspace: Option<std::path::PathBuf>,
},
Teammate {
depth: usize,
name: String,
existing: Option<Arc<Teammate>>,
spec: Box<ChildSpec>,
workspace: Option<std::path::PathBuf>,
},
Remote {
agents: Arc<crate::agents::remote::RemoteAgents>,
name: String,
},
Wait {
ids: Vec<String>,
any: bool,
},
}
struct Named {
reach: Reach,
agent_id: Option<String>,
}
impl std::fmt::Debug for SpawnTool {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SpawnTool")
.field("depth", &self.depth)
.field("width", &self.width)
.field("targets", &self.targets)
.field("policy", &self.policy.as_ref().map(|_| "<child policy>"))
.field("agents", &self.agents.as_ref().map(|_| "<agent ledger>"))
.field("owner", &self.owner)
.field("reach", &self.reach)
.field("task_scope", &self.task_scope.as_ref().map(|_| "<scope>"))
.field("remote", &self.taught().remote)
.finish()
}
}
impl SpawnTool {
pub fn new() -> Self {
Self::default()
}
pub fn with_targets(targets: impl IntoIterator<Item = String>) -> Self {
Self::with_targets_and_depth(targets, depth::DEFAULT_DELEGATION_DEPTH)
}
pub fn with_targets_and_depth(
targets: impl IntoIterator<Item = String>,
max_depth: usize,
) -> Self {
Self {
depth: depth::Depth::new(max_depth),
width: background::DEFAULT_DELEGATION_WIDTH,
targets: targets
.into_iter()
.collect::<BTreeSet<_>>()
.into_iter()
.collect(),
policy: None,
agents: None,
remote: None,
remembers: crate::agents::remote::Remembered::Not,
owner: None,
refresh: None,
reach: ReachPolicy::default(),
task_scope: None,
}
}
#[must_use]
pub fn with_task_scope<F>(self, scope: F) -> Self
where
F: Fn(TaskFuture) -> TaskFuture + Send + Sync + 'static,
{
self.with_task_scope_arc(Arc::new(scope))
}
#[must_use]
pub(crate) fn with_task_scope_arc(self, scope: TaskScope) -> Self {
Self {
task_scope: Some(scope),
..self
}
}
#[must_use]
pub fn with_delegation_width(self, width: usize) -> Self {
Self { width, ..self }
}
#[must_use]
pub fn with_reach_policy(self, reach: ReachPolicy) -> Self {
Self { reach, ..self }
}
#[must_use]
pub fn with_child_policy<F>(self, policy: F) -> Self
where
F: Fn(&ChildContext<'_>) -> ChildSpec + Send + Sync + 'static,
{
self.with_child_policy_arc(Arc::new(policy))
}
#[must_use]
pub(crate) fn with_child_policy_arc(self, policy: ChildPolicy) -> Self {
Self {
policy: Some(policy),
..self
}
}
#[must_use]
pub(crate) fn with_agents(self, agents: Arc<crate::runtime::agents::AgentRegistry>) -> Self {
Self {
agents: Some(agents),
..self
}
}
pub(crate) fn for_agent(&self, owner: &str, refresh: crate::runtime::Refresh) -> Self {
Self {
owner: Some(owner.to_string()),
refresh: Some(refresh),
..self.clone()
}
}
pub(crate) fn for_remote_agents(
&self,
remote: &Arc<crate::agents::remote::RemoteAgents>,
) -> Self {
Self {
remote: Some(Arc::clone(remote)),
remembers: crate::agents::remote::Remembered::Not,
..self.clone()
}
}
fn taught(&self) -> Names {
let teammates = self.owner.as_ref().and_then(|owner| {
Some(self.agents.as_ref()?.of(owner)?.teammates.listed())
});
Names {
remote: self.taught_remote(),
teammates,
reach: if self.reach.offers_a_choice() {
self.reach.allowed().collect()
} else {
Vec::new()
},
width: self.agents.as_ref().map(|_| self.width),
}
}
fn taught_remote(&self) -> Remote {
let Some(remote) = &self.remote else {
return Remote::default();
};
let taught = remote.taught();
Remote {
agents: taught
.agents
.into_iter()
.map(|agent| Taught {
name: agent.name,
summary: agent.summary,
})
.collect(),
note: taught.note,
}
}
#[must_use]
pub(crate) fn remembering(self) -> Self {
Self {
remembers: crate::agents::remote::Remembered::ByAgent,
..self
}
}
#[must_use]
pub(crate) fn remembering_under(self, key: String) -> Self {
Self {
remembers: crate::agents::remote::Remembered::ByKey(key),
..self
}
}
fn remote_agents_of(&self, agent_id: &str) -> Option<Arc<crate::agents::remote::RemoteAgents>> {
self.agents.as_ref()?.of(agent_id)?.remote.clone()
}
fn children_of(&self, agent_id: &str) -> Option<Children> {
self.agents
.as_ref()?
.of(agent_id)
.map(|tools| tools.background.clone())
}
fn remote_route(&self, agent_id: &str, name: &str) -> Result<Option<SpawnDecision>, String> {
let Some(agents) = self.remote_agents_of(agent_id) else {
return Ok(None);
};
if agents.route(name) == crate::agents::Route::Unknown {
return Ok(None);
}
remote::authorize(&agents, name)?;
Ok(Some(SpawnDecision::Remote {
agents,
name: name.to_string(),
}))
}
fn no_remote_agent(&self, agent_id: &str, name: &str) -> String {
if let Some(agents) = self.remote_agents_of(agent_id) {
return remote::unknown(&agents, name);
}
no_remote_agents(name)
}
fn decide_named(
&self,
spawn: &Spawn,
agent_id: &str,
name: &str,
resolve_workspace: impl FnOnce() -> Result<std::path::PathBuf, String>,
) -> Result<SpawnDecision, String> {
let reach = self.reach.decide(spawn.reach())?;
let row = self.agents.as_ref().and_then(|agents| agents.of(agent_id));
let mine = row.as_ref().and_then(|row| row.teammates.find(name));
if reach != Reach::Remote
&& let Some(existing) = mine
{
return self.teammate_turn(spawn, agent_id, name, Some(existing), resolve_workspace);
}
if reach != Reach::Local
&& let Some(remote) = self.remote_route(agent_id, name)?
{
return Ok(remote);
}
if reach == Reach::Remote {
return Err(self.no_remote_agent(agent_id, name));
}
if is_namespaced(name) {
return Err(format!(
"`{name}` is a namespaced name, which reaches a remote agent, and no resolver \
here knows it; a teammate's name has no `/`, as `@{}`",
name.rsplit('/').next().unwrap_or(name)
));
}
if row.is_none() {
return Err(format!(
"spawn keeps no teammates on this runtime, so there is nobody to hand `@{name}` \
to; write the task with no `@` to hand it to a subagent here, or `@@` to begin \
it with a literal `@`"
));
}
self.teammate_turn(spawn, agent_id, name, None, resolve_workspace)
}
fn teammate_turn(
&self,
spawn: &Spawn,
agent_id: &str,
name: &str,
existing: Option<Arc<Teammate>>,
resolve_workspace: impl FnOnce() -> Result<std::path::PathBuf, String>,
) -> Result<SpawnDecision, String> {
if spawn.body().is_empty() {
return Err(match &existing {
Some(_) => format!(
"`@{name}` is a teammate here and has nothing outstanding to wait on; write \
`@{name} <task>` to hand it one"
),
None => format!(
"`@{name}` names nobody yet; write `@{name} <task>` to make a teammate called \
`{name}` and hand it that task"
),
});
}
let depth = self.depth.authorize_delegation(agent_id)?;
let (spec, workspace) = match (&existing, self.policy.as_ref()) {
(None, Some(policy)) => {
let workspace = resolve_workspace()?;
let spec =
policy(&ChildContext::new(spawn.body(), agent_id, &workspace).named(name));
(spec, Some(workspace))
}
_ => (ChildSpec::inherit(), None),
};
Ok(SpawnDecision::Teammate {
depth,
name: name.to_string(),
existing,
spec: Box::new(spec),
workspace,
})
}
fn child_spec(
&self,
prompt: &str,
parent_agent_id: &str,
workspace_dir: &std::path::Path,
) -> ChildSpec {
match &self.policy {
Some(policy) => policy(&ChildContext::new(prompt, parent_agent_id, workspace_dir)),
None => ChildSpec::inherit(),
}
}
fn authorize_target(&self, spawn: &Spawn) -> Result<(), String> {
let (Mode::Command, Some(target)) = (spawn.mode(), 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)
)
})
}
fn decide_wait(
&self,
agent_id: &str,
keys: &[String],
any: bool,
) -> Result<SpawnDecision, String> {
let Some(children) = self.children_of(agent_id) else {
return Err(
"spawn keeps no background work on this runtime, so there is nothing to wait for"
.to_string(),
);
};
let mut ids = Vec::with_capacity(keys.len());
for key in keys {
let Some(id) = children.find(key) else {
return Err(nothing_running(&children, key));
};
if !ids.contains(&id) {
ids.push(id);
}
}
Ok(SpawnDecision::Wait { ids, any })
}
fn authorize_background(&self, agent_id: &str) -> Result<(), String> {
let Some(children) = self.children_of(agent_id) else {
return Err(
"spawn keeps no background work on this runtime; leave `background` out and \
wait for the answer"
.to_string(),
);
};
let running = children.running();
if running >= self.width {
return Err(background::too_wide(running, self.width));
}
Ok(())
}
fn decide(
&self,
spawn: &Spawn,
agent_id: &str,
resolve_workspace: impl FnOnce() -> Result<std::path::PathBuf, String>,
) -> Result<SpawnDecision, String> {
self.authorize_target(spawn)?;
if spawn.background() {
self.authorize_background(agent_id)?;
}
match spawn.mode() {
Mode::Command => Ok(SpawnDecision::Command),
Mode::Wait => self.decide_wait(agent_id, spawn.wait(), spawn.any()),
Mode::Agent => {
if let Some(name) = spawn.target() {
if spawn.body().is_empty()
&& let Some(children) = self.children_of(agent_id)
&& let Some(id) = children.find(name)
{
return Ok(SpawnDecision::Wait {
ids: vec![id],
any: false,
});
}
return self.decide_named(spawn, agent_id, name, resolve_workspace);
}
let depth = self.depth.authorize_delegation(agent_id)?;
let (spec, workspace) = match self.policy.as_ref() {
Some(_) => {
let workspace = resolve_workspace()?;
let spec = self.child_spec(spawn.body(), agent_id, &workspace);
(spec, Some(workspace))
}
None => (ChildSpec::inherit(), None),
};
Ok(SpawnDecision::Agent {
depth,
spec: Box::new(spec),
workspace,
})
}
}
}
}
fn no_remote_agents(name: &str) -> String {
format!(
"spawn has no remote agents to hand `@{name}` to; write the task with no `@` to hand it \
to a subagent here, or `@@` to begin it with a literal `@`"
)
}
fn nothing_running(children: &Children, key: &str) -> String {
let running: Vec<String> = children
.listed()
.into_iter()
.map(|(id, name, _)| match name {
Some(name) => format!("{name} (id {id})"),
None => id,
})
.collect();
if running.is_empty() {
return format!(
"nothing is running in the background under `{key}`, or at all; start a task with \
`background: true` before waiting for it"
);
}
format!(
"nothing is running in the background under `{key}`; what is: {}",
listed(&running)
)
}
impl ToolDefinition for SpawnTool {
fn descriptor(&self) -> RuntimeToolDescriptor {
let names = self.taught();
RuntimeToolDescriptor::builder(SPAWN)
.description(description(&self.targets, &names))
.input_schema(input_schema(&self.targets, &names))
.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()
}
}
#[async_trait]
impl ToolExecutor for SpawnTool {
fn authorization_preview(
&self,
ctx: &ParallelToolContext,
input: &Value,
) -> Result<ToolAuthorizationPreview, String> {
let spawn = parse(input)?;
let decision = self.decide(&spawn, &ctx.agent_id, || {
ctx.resolve_working_directory(None)
})?;
parameters_reach_a_remote_agent(&spawn, &decision)?;
let (cwd, child, named) = match decision {
SpawnDecision::Command | SpawnDecision::Wait { .. } => {
(ctx.resolve_working_directory(None)?, None, None)
}
SpawnDecision::Remote { .. } => (
ctx.resolve_working_directory(None)?,
None,
Some(Named {
reach: Reach::Remote,
agent_id: None,
}),
),
SpawnDecision::Agent {
spec, workspace, ..
} => {
let cwd = match workspace {
Some(workspace) => workspace,
None => ctx.resolve_working_directory(None)?,
};
(cwd, spec.preview_value(), None)
}
SpawnDecision::Teammate {
spec,
workspace,
existing,
..
} => {
let cwd = match workspace {
Some(workspace) => workspace,
None => ctx.resolve_working_directory(None)?,
};
let child = existing.is_none().then(|| spec.preview_value()).flatten();
(
cwd,
child,
Some(Named {
reach: Reach::Local,
agent_id: existing.map(|teammate| teammate.id().to_string()),
}),
)
}
};
Ok(preview(
&spawn,
cwd,
&self.descriptor(),
input,
child,
named,
))
}
fn execution_category(&self, input: &Value) -> ToolExecutionCategory {
parse(input).map_or(ToolExecutionCategory::ExclusiveLocalMutation, |spawn| {
execution_category(&spawn)
})
}
async fn execute_mut(&self, mut ctx: ToolContext<'_>, input: Value) -> ToolResult {
let spawn = parse(&input)?;
let decision = self.decide(&spawn, &ctx.agent_id, || {
ctx.resolve_working_directory(None)
})?;
parameters_reach_a_remote_agent(&spawn, &decision)?;
let sent = files::read(spawn.files(), |path| ctx.read_file_bytes(path))?;
match decision {
SpawnDecision::Command => execute::command(&ctx, spawn.body(), spawn.target()).await,
SpawnDecision::Wait { ids, any } => execute::wait(self, &mut ctx, &ids, any).await,
SpawnDecision::Remote { agents, name } => {
let task = remote::Task {
body: spawn.body().to_string(),
metadata: spawn.metadata().cloned(),
files: sent.into_iter().map(|sent| sent.file).collect(),
};
if spawn.background() {
let children = execute::children_of(self, &ctx.agent_id)?;
return remote::delegate_in_background(
&ctx,
self.task_scope.as_ref(),
children,
agents,
&name,
task,
self.remembers.clone(),
);
}
remote::delegate(ctx, &agents, &name, task, &self.remembers).await
}
SpawnDecision::Agent { depth, spec, .. } => {
let body = files::named(spawn.body(), &sent);
if spawn.background() {
return execute::delegate_in_background(self, &mut ctx, &body, depth, *spec)
.await;
}
execute::delegate(self, &mut ctx, &body, depth, *spec).await
}
SpawnDecision::Teammate {
depth,
name,
existing,
spec,
..
} => {
let body = files::named(spawn.body(), &sent);
execute::teammate(
self,
&mut ctx,
&body,
execute::TeammateCall {
name,
depth,
existing,
spec: *spec,
background: spawn.background(),
},
)
.await
}
}
}
}
fn parameters_reach_a_remote_agent(spawn: &Spawn, decision: &SpawnDecision) -> Result<(), String> {
match (spawn.metadata(), decision) {
(Some(_), SpawnDecision::Remote { .. }) | (None, _) => Ok(()),
(Some(_), _) => Err(format!(
"`{}` goes to a remote agent, and `@{}` is not one here; leave the parameters out, \
or name a remote agent",
parse::METADATA_FIELD,
spawn.target().unwrap_or_default()
)),
}
}
fn preview(
spawn: &Spawn,
cwd: std::path::PathBuf,
descriptor: &RuntimeToolDescriptor,
raw_input: &Value,
child: Option<Value>,
named: Option<Named>,
) -> ToolAuthorizationPreview {
let mut structured_input = json!({
"mode": spawn.mode().as_str(),
"body": spawn.body(),
"cwd": cwd,
"target": spawn.target().unwrap_or(LOCAL_TARGET),
});
let object = structured_input
.as_object_mut()
.expect("built as an object two lines up");
if let Some(child) = child {
object.insert("child".to_string(), child);
}
if let Some(named) = &named {
object.insert("reach".to_string(), json!(named.reach.as_str()));
if let Some(agent_id) = &named.agent_id {
object.insert("agent_id".to_string(), json!(agent_id));
}
}
if let Some(metadata) = spawn.metadata() {
object.insert("metadata".to_string(), Value::Object(metadata.clone()));
}
if !spawn.files().is_empty() {
object.insert(parse::FILES_FIELD.to_string(), json!(spawn.files()));
}
if spawn.background() {
object.insert(parse::BACKGROUND_FIELD.to_string(), json!(true));
}
if spawn.mode() == Mode::Wait {
object.insert(parse::WAIT_FIELD.to_string(), json!(spawn.wait()));
if spawn.any() {
object.insert(parse::ANY_FIELD.to_string(), json!(true));
}
}
ToolAuthorizationPreview {
working_directory: cwd,
capabilities: capabilities(spawn.mode()),
side_effect_level: side_effect_level(spawn, named.as_ref()),
durability: descriptor.durability,
execution_category: execution_category(spawn),
approval_category: approval_category(spawn),
raw_input: raw_input.clone(),
structured_input,
}
}
fn side_effect_level(spawn: &Spawn, named: Option<&Named>) -> ToolSideEffectLevel {
match (spawn.mode(), named.map(|named| named.reach)) {
(Mode::Command, _) => ToolSideEffectLevel::Process,
(Mode::Wait, _) => ToolSideEffectLevel::None,
(Mode::Agent, Some(Reach::Remote | Reach::Auto)) => ToolSideEffectLevel::External,
(Mode::Agent, Some(Reach::Local) | None) => ToolSideEffectLevel::LocalState,
}
}
fn capabilities(mode: Mode) -> Vec<ToolCapability> {
match mode {
Mode::Command => vec![ToolCapability::ProcessExec, ToolCapability::FilesystemWrite],
Mode::Agent | Mode::Wait => vec![ToolCapability::Delegation],
}
}
const fn approval_category(spawn: &Spawn) -> ToolApprovalCategory {
match spawn.mode() {
Mode::Command => ToolApprovalCategory::Process,
Mode::Agent if spawn.background() => ToolApprovalCategory::Background,
Mode::Agent | Mode::Wait => ToolApprovalCategory::Delegation,
}
}
const fn execution_category(spawn: &Spawn) -> ToolExecutionCategory {
match spawn.mode() {
Mode::Command => ToolExecutionCategory::ExclusiveLocalMutation,
Mode::Agent if spawn.background() => ToolExecutionCategory::BackgroundJob,
Mode::Agent | Mode::Wait => ToolExecutionCategory::Delegation,
}
}
#[cfg(test)]
mod tests;