use std::collections::{BTreeSet, HashMap};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use agent_base::llm_trait::LlmProvider;
use agent_base::{
AgentBuilder, AgentError, AgentResult, AgentRuntime, AllowAllApprovalHandler, ApprovalHandler,
DenyAllApprovalHandler, DenyAllToolPolicy, Language, ReasoningEffort, RuntimeEvent, SessionId,
Tool, ToolPolicy,
};
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;
use super::budget::{SpawnTicket, usage_total};
use super::child_builder::ChildBuilder;
use super::child_config::ChildConfig;
use super::config::{AgentAutonomy, ChildPermissionMode, MultiAgentConfig};
use super::control::AgentControl;
use super::limiter::{AgentExecutionLimiter, ExecutionSlot};
use super::mailbox::{ChildMailbox, MailboxHub, MailboxResult, MailboxStatus};
use super::path::AgentPath;
use super::registry::{AgentLifecycleEvent, AgentRegistry, RegistrySnapshot};
pub use super::runtime::watcher::{ChildReport, ChildResultEvent};
mod build;
mod fork;
mod outcome;
mod spawn;
pub mod watcher;
pub struct MultiAgentRuntime {
registry: Arc<Mutex<AgentRegistry>>,
mailbox: Arc<MailboxHub>,
control: Arc<AgentControl>,
limiter: Arc<AgentExecutionLimiter>,
client: Arc<dyn LlmProvider>,
business_tools: Vec<Arc<dyn Tool>>,
child_excluded_tools: Vec<String>,
child_reasoning_effort: Option<ReasoningEffort>,
child_read_only: bool,
child_fork_history: Option<String>,
autonomy: AgentAutonomy,
write_tools: Vec<String>,
task_timeout: Option<Duration>,
event_tx: Mutex<Option<tokio::sync::mpsc::UnboundedSender<RuntimeEvent>>>,
root_cancel: CancellationToken,
join_set: Mutex<JoinSet<()>>,
child_cancels: Arc<Mutex<HashMap<AgentPath, CancellationToken>>>,
error_recovery: Option<Arc<dyn agent_base::ToolErrorRecovery>>,
language: Language,
child_permission_mode: ChildPermissionMode,
tool_policy: Option<Arc<dyn ToolPolicy>>,
approval_handler: Option<Arc<dyn ApprovalHandler>>,
session_manager: Mutex<Option<Arc<agent_base::engine::SessionManager>>>,
}
impl MultiAgentRuntime {
#[allow(clippy::too_many_arguments)] pub fn new(
config: MultiAgentConfig,
client: Arc<dyn LlmProvider>,
business_tools: Vec<Arc<dyn Tool>>,
root_cancel: CancellationToken,
error_recovery: Option<Arc<dyn agent_base::ToolErrorRecovery>>,
language: Language,
tool_policy: Option<Arc<dyn ToolPolicy>>,
approval_handler: Option<Arc<dyn ApprovalHandler>>,
) -> Self {
let child_permission_mode = config.child_permission_mode;
let child_excluded_tools = config.child_excluded_tools.clone();
let child_reasoning_effort = config.child_reasoning_effort.clone();
let child_read_only = config.child_read_only;
let child_fork_history = config.child_fork_history.clone();
let control = Arc::new(AgentControl::new(&config.control));
let limiter = Arc::clone(control.limiter());
let autonomy = config.control.autonomy;
let write_tools = config.control.write_tools.clone();
let task_timeout = config.control.task_timeout;
let registry = Arc::new(Mutex::new(AgentRegistry::new(config)));
Self {
registry,
mailbox: Arc::new(MailboxHub::new()),
control,
limiter,
client,
business_tools,
child_excluded_tools,
child_reasoning_effort,
child_read_only,
child_fork_history,
autonomy,
write_tools,
task_timeout,
event_tx: Mutex::new(None),
root_cancel,
join_set: Mutex::new(JoinSet::new()),
child_cancels: Arc::new(Mutex::new(HashMap::new())),
error_recovery,
language,
child_permission_mode,
tool_policy,
approval_handler,
session_manager: Mutex::new(None),
}
}
pub fn client(&self) -> &Arc<dyn LlmProvider> {
&self.client
}
pub fn child_fork_history(&self) -> Option<&str> {
self.child_fork_history.as_deref()
}
pub fn set_event_sender(&self, tx: tokio::sync::mpsc::UnboundedSender<RuntimeEvent>) {
*self.event_tx.lock().unwrap() = Some(tx);
}
pub fn set_session_manager(&self, session_manager: Arc<agent_base::engine::SessionManager>) {
*self.session_manager.lock().unwrap() = Some(session_manager);
}
pub fn send_message(&self, agent_path: &str, message: String) -> Result<bool, String> {
let path = self.parse_path(agent_path)?;
Ok(self.mailbox.send_message(&path, message))
}
pub fn send_task(
&self,
agent_path: &str,
task: String,
interrupt: bool,
) -> Result<bool, String> {
let path = self.parse_path(agent_path)?;
if !self.mailbox.contains(&path) {
return Err("agent not found".to_string());
}
let mut registry = self.registry.lock().unwrap();
registry.set_task(&path, task.clone());
registry.touch(&path);
let tracked = registry.note_enqueued(&path);
drop(registry);
let sent = self.mailbox.send_task(&path, task, interrupt);
if !sent && tracked {
self.registry.lock().unwrap().note_send_failed(&path);
}
Ok(sent)
}
pub async fn wait_for_result(&self, agent_path: Option<&str>, timeout_ms: u64) -> WaitResult {
let filter_path = match agent_path {
Some(s) => match AgentPath::parse(s) {
Some(p) => Some(p),
None => {
return WaitResult {
status: "error".to_string(),
result: Some(format!("invalid agent path: {}", s)),
agent_path: None,
has_more: false,
denied_tools: vec![],
};
}
},
None => None,
};
let mut seq = self.mailbox.subscribe_seq();
let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_millis(timeout_ms);
loop {
let result = match &filter_path {
Some(path) => self.mailbox.try_recv_result(path),
None => self.mailbox.try_recv_any(),
};
if let Some(r) = result {
let has_more = self.mailbox.total_pending_results() > 0;
let (status_str, result_text) = match r.status {
MailboxStatus::Ok => ("ok".to_string(), r.result),
MailboxStatus::Error => ("error".to_string(), r.result),
MailboxStatus::Closed => ("closed".to_string(), r.result),
};
return WaitResult {
status: status_str,
result: result_text,
agent_path: Some(r.agent_path.to_string()),
has_more,
denied_tools: r.denied_tools,
};
}
if let Some(path) = &filter_path {
let gone =
!self.mailbox.contains(path) && !self.registry.lock().unwrap().contains(path);
if gone {
return WaitResult {
status: "closed".to_string(),
result: None,
agent_path: Some(path.to_string()),
has_more: false,
denied_tools: vec![],
};
}
}
let now = tokio::time::Instant::now();
if now >= deadline {
return WaitResult {
status: "timeout".to_string(),
result: None,
agent_path: None,
has_more: false,
denied_tools: vec![],
};
}
let remaining = deadline - now;
tokio::select! {
_ = seq.changed() => {
continue;
}
_ = tokio::time::sleep(remaining) => {
return WaitResult {
status: "timeout".to_string(),
result: None,
agent_path: None,
has_more: false,
denied_tools: vec![],
};
}
}
}
}
pub fn try_wait(&self, agent_path: Option<&str>) -> WaitResult {
let filter_path = match agent_path {
Some(s) => match AgentPath::parse(s) {
Some(p) => Some(p),
None => {
return WaitResult {
status: "error".to_string(),
result: Some(format!("invalid agent path: {}", s)),
agent_path: None,
has_more: false,
denied_tools: vec![],
};
}
},
None => None,
};
let result = match &filter_path {
Some(path) => self.mailbox.try_recv_result(path),
None => self.mailbox.try_recv_any(),
};
if let Some(r) = result {
let has_more = self.mailbox.total_pending_results() > 0;
let (status_str, result_text) = match r.status {
MailboxStatus::Ok => ("ok".to_string(), r.result),
MailboxStatus::Error => ("error".to_string(), r.result),
MailboxStatus::Closed => ("closed".to_string(), r.result),
};
return WaitResult {
status: status_str,
result: result_text,
agent_path: Some(r.agent_path.to_string()),
has_more,
denied_tools: r.denied_tools,
};
}
if let Some(path) = &filter_path {
let gone =
!self.mailbox.contains(path) && !self.registry.lock().unwrap().contains(path);
if gone {
return WaitResult {
status: "closed".to_string(),
result: None,
agent_path: Some(path.to_string()),
has_more: false,
denied_tools: vec![],
};
}
}
WaitResult {
status: "pending".to_string(),
result: None,
agent_path: filter_path.map(|p| p.to_string()),
has_more: false,
denied_tools: vec![],
}
}
pub fn start_watcher(
self: &Arc<Self>,
) -> (
tokio::task::JoinHandle<()>,
mpsc::UnboundedReceiver<ChildResultEvent>,
) {
let (child_result_tx, child_result_rx) = mpsc::unbounded_channel();
let summarizer = Arc::new(crate::focus::ProgressSummarizer::new(
Arc::clone(&self.client),
crate::focus::DEFAULT_SUMMARY_TIMEOUT,
));
let handle = watcher::spawn_watcher_with_watchdog(
Arc::clone(&self.mailbox),
Arc::clone(&self.registry),
Some(summarizer),
Some(child_result_tx),
self.root_cancel.clone(),
);
(handle, child_result_rx)
}
pub fn close_agent(&self, agent_path: &str) -> Result<CloseResult, String> {
let path = self.parse_path(agent_path)?;
let previous_status = {
let registry = self.registry.lock().unwrap();
registry
.get(&path)
.map(|e| e.status().name().to_string())
.unwrap_or_else(|| "unknown".to_string())
};
let closed = {
let cancels = self.child_cancels.lock().unwrap();
match cancels.get(&path) {
Some(token) => {
if token.is_cancelled() {
false
} else {
token.cancel();
true
}
}
None => false,
}
};
Ok(CloseResult {
closed,
previous_status,
message: if closed {
"agent closed".to_string()
} else {
"agent not found".to_string()
},
})
}
pub fn list_agents(&self) -> Vec<AgentInfo> {
self.registry
.lock()
.unwrap()
.snapshot()
.agents
.into_iter()
.map(|a| AgentInfo {
agent_path: a.path,
status: a.status,
tool_calls: a.tool_calls,
running_secs: a.running_secs,
last_activity_secs: a.last_activity_secs,
task: a.task,
pending_results: a.pending_results,
})
.collect()
}
pub fn busy_count(&self) -> usize {
self.registry.lock().unwrap().busy_count()
}
pub fn subscribe_lifecycle(&self) -> watch::Receiver<Arc<RegistrySnapshot>> {
self.registry.lock().unwrap().subscribe()
}
pub fn recent_lifecycle_events(&self, max: usize) -> Vec<AgentLifecycleEvent> {
self.registry.lock().unwrap().recent_events(max)
}
pub fn mailbox(&self) -> &Arc<MailboxHub> {
&self.mailbox
}
pub fn registry(&self) -> &Mutex<AgentRegistry> {
&self.registry
}
pub fn control(&self) -> &Arc<AgentControl> {
&self.control
}
pub fn cancel_all(&self) {
let mut cancels = self.child_cancels.lock().unwrap();
for (_, token) in cancels.drain() {
token.cancel();
}
}
}
impl Drop for MultiAgentRuntime {
fn drop(&mut self) {
self.cancel_all();
let mut js = self.join_set.lock().unwrap();
while let Some(result) = js.try_join_next() {
if let Err(e) = result
&& e.is_panic()
{
tracing::error!(
error = %e,
"child agent task panicked"
);
}
}
}
}
impl MultiAgentRuntime {
fn parse_path(&self, s: &str) -> Result<AgentPath, String> {
AgentPath::parse(s).ok_or_else(|| format!("invalid agent path: '{}'", s))
}
}
#[derive(Clone, Debug)]
pub struct WaitResult {
pub status: String,
pub result: Option<String>,
pub agent_path: Option<String>,
pub has_more: bool,
pub denied_tools: Vec<String>,
}
#[derive(Clone, Debug)]
pub struct CloseResult {
pub closed: bool,
pub previous_status: String,
pub message: String,
}
#[derive(Clone, Debug, serde::Serialize)]
pub struct AgentInfo {
pub agent_path: String,
pub status: String,
pub tool_calls: usize,
#[serde(skip_serializing_if = "Option::is_none")]
pub running_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub last_activity_secs: Option<u64>,
pub task: Option<String>,
#[serde(skip_serializing_if = "is_zero")]
pub pending_results: usize,
}
fn is_zero(n: &usize) -> bool {
*n == 0
}
#[cfg(test)]
mod tests;