use std::time::SystemTime;
use mentra::tool::{ParallelToolContext, ToolContext, ToolResult};
use super::{describe::listed, files, no_remote_agents};
use crate::agents::{
Attachment, ExchangeBounds, Failure, Message, Outcome, PeerError, Request, Resolved, Route,
Started, StillWorking, Unresolved,
remote::{Conversation, Outstanding, Remembered, RemoteAgents, Waiting},
};
pub(super) fn authorize(agents: &RemoteAgents, name: &str) -> Result<(), String> {
match agents.route(name) {
Route::Known | Route::Maybe => Ok(()),
Route::Backoff(wait) => Err(backing_off(name, wait)),
Route::Refused(reason) => Err(refused(name, &reason)),
Route::Unknown => Err(unknown(agents, name)),
}
}
fn refused(name: &str, reason: &str) -> String {
format!(
"{reason}, so there is nothing to hand `@{name}` to; write the task with no `@` to hand \
it to a subagent here"
)
}
pub(super) fn unknown(agents: &RemoteAgents, name: &str) -> String {
let configured = agents.names();
if configured.is_empty() {
return no_remote_agents(name);
}
format!(
"there is no remote agent named `{name}`; the ones this workspace configures are {}, \
and `@@` begins a task with a literal `@`",
listed(&configured)
)
}
fn backing_off(name: &str, wait: std::time::Duration) -> String {
format!(
"the remote agent `{name}` could not be reached when it was last tried; `@{name}` tries \
it again in {}s",
seconds(wait)
)
}
fn seconds(wait: std::time::Duration) -> u64 {
wait.as_secs_f64().ceil().max(1.0) as u64
}
fn unresolved(agents: &RemoteAgents, name: &str, why: Unresolved) -> String {
match why {
Unresolved::Backoff(wait) => backing_off(name, wait),
Unresolved::Unreachable {
error,
retry_in: Some(retry_in),
} => format!(
"the remote agent `{name}` could not be reached: {error}. `@{name}` tries it again \
in {}s",
seconds(retry_in)
),
Unresolved::Unreachable {
error,
retry_in: None,
} => format!("the remote agent `{name}` could not be reached: {error}"),
Unresolved::Refused(reason) => refused(name, &reason),
Unresolved::Unknown => unknown(agents, name),
}
}
pub(super) struct Task<'a> {
pub(super) body: &'a str,
pub(super) metadata: Option<&'a serde_json::Map<String, serde_json::Value>>,
pub(super) files: Vec<Attachment>,
}
pub(super) async fn delegate(
ctx: ToolContext<'_>,
agents: &RemoteAgents,
name: &str,
task: Task<'_>,
remembered: &Remembered,
) -> ToolResult {
let Task {
body,
metadata,
files,
} = task;
let Resolved {
peer, id, endpoint, ..
} = agents
.resolve(name)
.await
.map_err(|why| unresolved(agents, name, why))?;
let id = id.as_str();
let caller = ctx.agent_id.clone();
let options = ctx.child_run_options();
let hop = Hop {
hop: agents.hop().saturating_add(1),
deadline: options.deadline,
};
let before = agents.conversation(&caller, id, remembered);
let request = request(name, body, &before, hop)?;
let request = match request {
Request::Send(message) => {
Request::Send(message.with_metadata(metadata.cloned()).with_files(files))
}
wait => wait,
};
let ctx: ParallelToolContext = ctx.into();
let progress = |text: String| ctx.emit_progress(text);
let save = |next: Conversation| {
if let Err(error) = agents.record(&caller, id, next, remembered) {
ctx.emit_progress(format!(
"the conversation with `{name}` was not saved: {error}"
));
}
};
let started = |started: Started| {
save(Conversation {
context_id: started.context_id.or_else(|| before.context_id.clone()),
outstanding: Some(Outstanding {
task_id: started.task_id,
waiting_for: Waiting::Agent,
endpoint: Some(endpoint.clone()),
}),
});
};
let outcome = peer
.exchange(
request,
&ExchangeBounds {
options: &options,
progress: &progress,
started: &started,
},
)
.await;
let (outcome, returned) = match outcome {
Ok(mut outcome) => {
let returned = take_files(&mut outcome);
(Ok(outcome), returned)
}
Err(error) => (Err(error), None),
};
let (next, result) = settle(name, before, outcome, &endpoint);
save(next);
let Some((folder, files)) = returned else {
return result;
};
let root = ctx
.resolve_working_directory(None)
.unwrap_or_else(|_| ctx.working_directory().to_path_buf());
let saved = files::save(&root, &folder, files, |path, bytes| {
ctx.write_file_bytes(path, bytes)
});
match result {
Ok(text) => Ok(said(text, saved)),
Err(text) => Err(said(text, saved)),
}
}
fn take_files(outcome: &mut Outcome) -> Option<(String, Vec<Attachment>)> {
let (task_id, files) = match outcome {
Outcome::Replied { files, .. } => (None, files),
Outcome::Finished { task_id, files, .. }
| Outcome::Working { task_id, files, .. }
| Outcome::NeedsInput { task_id, files, .. }
| Outcome::NeedsAuthorization { task_id, files, .. }
| Outcome::Failed { task_id, files, .. } => (Some(task_id.clone()), files),
#[allow(unreachable_patterns)]
_ => return None,
};
if files.is_empty() {
return None;
}
let folder = task_id.unwrap_or_else(|| {
let now = SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default();
format!("reply-{}", now.as_millis())
});
Some((folder, std::mem::take(files)))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct Hop {
hop: u32,
deadline: Option<SystemTime>,
}
fn request(
name: &str,
body: &str,
conversation: &Conversation,
hop: Hop,
) -> Result<Request, String> {
match (&conversation.outstanding, body.is_empty()) {
(Some(outstanding), true) => Ok(Request::wait(&outstanding.task_id)),
(None, true) => Err(format!(
"`{name}` has nothing outstanding to wait on; write `@{name} <task>` to hand it one"
)),
(
Some(Outstanding {
waiting_for: Waiting::Agent,
..
}),
false,
) => Err(format!(
"`{name}` is still working on the last task you handed it, and takes one task at a \
time; write `@{name}` alone to keep waiting for that one first"
)),
(Some(outstanding), false) => Ok(Request::Send(
Message::new(body)
.in_context(conversation.context_id.clone())
.answering(Some(outstanding.task_id.clone()))
.at_hop(hop.hop, hop.deadline),
)),
(None, false) => Ok(Request::Send(
Message::new(body)
.in_context(conversation.context_id.clone())
.at_hop(hop.hop, hop.deadline),
)),
}
}
fn settle(
name: &str,
before: Conversation,
outcome: Result<Outcome, PeerError>,
endpoint: &str,
) -> (Conversation, ToolResult) {
let outstanding = |task_id: String, waiting_for: Waiting| {
Some(Outstanding {
task_id,
waiting_for,
endpoint: Some(endpoint.to_string()),
})
};
let settled = |context_id: Option<String>| Conversation {
context_id: context_id.or_else(|| before.context_id.clone()),
outstanding: None,
};
match outcome {
Ok(Outcome::Replied {
context_id, text, ..
}) => (
settled(context_id),
Ok(said(format!("Remote agent `{name}` replied:"), text)),
),
Ok(Outcome::Finished {
context_id, text, ..
}) => (
settled(context_id),
Ok(said(format!("Remote agent `{name}` replied:"), text)),
),
Ok(Outcome::NeedsInput {
task_id,
context_id,
text,
..
}) => (
Conversation {
context_id: context_id.or_else(|| before.context_id.clone()),
outstanding: outstanding(task_id, Waiting::Input),
},
Ok(said(
format!(
"Remote agent `{name}` needs more before it can finish. Answer with \
`@{name} <your answer>`."
),
text,
)),
),
Ok(Outcome::NeedsAuthorization {
context_id, text, ..
}) => (
settled(context_id),
Err(said(
format!(
"Remote agent `{name}` needs authorization this conversation cannot give \
it, so the task stopped there."
),
text,
)),
),
Ok(Outcome::Failed {
context_id,
text,
how,
..
}) => (
settled(context_id),
Err(said(
format!(
"Remote agent `{name}` {}.",
match how {
Failure::Rejected => "rejected the task",
Failure::Canceled => "reports the task was canceled",
_ => "reports the task failed",
}
),
text,
)),
),
Ok(Outcome::Working {
task_id,
context_id,
text,
why,
..
}) => {
let why = match why {
StillWorking::WaitBound(waited) => {
format!("did not finish within {}s", waited.as_secs().max(1))
}
_ => "was still working when this run began to stop".to_string(),
};
(
Conversation {
context_id: context_id.or_else(|| before.context_id.clone()),
outstanding: outstanding(task_id, Waiting::Agent),
},
Ok(said(
format!(
"Remote agent `{name}` {why}; it keeps working remotely. Write `@{name}` \
alone to keep waiting for it, and do not hand it the task again."
),
text,
)),
)
}
Err(error) if moved(&before, endpoint) && cause(&error).is_some() => (
Conversation::default(),
Err(format!(
"Remote agent `{name}` has moved to a new address since it was handed the task it \
was working on, and the new address could not pick that task up ({}), so it is \
forgotten and the conversation with it starts over. Hand it the task again with \
`@{name} <task>` if it still needs doing.",
cause(&error).expect("checked by the guard")
)),
),
Err(PeerError::UnknownTask) if before.outstanding.is_some() => (
Conversation::default(),
Err(format!(
"Remote agent `{name}` no longer knows the task it was working on — it may have \
moved to a new address — so that task is forgotten and the conversation with \
it starts over. Hand it the task again with `@{name} <task>` if it still needs \
doing."
)),
),
Err(PeerError::UnknownTask) => (
before,
Err(format!(
"Handing the task to remote agent `{name}` failed: it does not know the task \
this answers"
)),
),
Err(PeerError::Send(error)) => (
before,
Err(format!(
"Handing the task to remote agent `{name}` failed: {error}. No task came back, \
so there is nothing to wait on. If the task is safe to repeat, such as a \
question, hand it over again with `@{name} <task>`; if it is not, say that it \
may have been received."
)),
),
Err(PeerError::Check { task_id, error }) => (
Conversation {
outstanding: outstanding(task_id, Waiting::Agent),
..before
},
Err(format!(
"Checking on remote agent `{name}` failed: {error}. The task may still be \
running; write `@{name}` alone to check on it again."
)),
),
Err(PeerError::Abandoned { canceled, .. }) => (
Conversation {
outstanding: None,
..before
},
Err(match canceled {
Ok(()) => format!(
"The run was cancelled while remote agent `{name}` was working, and the task \
was asked to cancel."
),
Err(error) => format!(
"The run was cancelled while remote agent `{name}` was working, and asking \
the task to cancel failed: {error}"
),
}),
),
Err(error) => (
before,
Err(format!(
"Remote agent `{name}` could not be reached: {error}"
)),
),
}
}
fn moved(conversation: &Conversation, endpoint: &str) -> bool {
conversation
.outstanding
.as_ref()
.and_then(|outstanding| outstanding.endpoint.as_deref())
.is_some_and(|recorded| recorded != endpoint)
}
fn cause(error: &PeerError) -> Option<String> {
match error {
PeerError::Send(error) | PeerError::Check { error, .. } => Some(error.clone()),
PeerError::UnknownTask => Some(error.to_string()),
_ => None,
}
}
fn said(headline: String, output: String) -> String {
if output.is_empty() {
return headline;
}
format!("{headline}\n\n{output}")
}
#[cfg(test)]
mod tests;