use super::{
activity::emit_activity,
config::SubagentRunConfig,
cwd::resolve_child_cwd,
dto::{SubagentStatus, SubagentTask, SubagentTaskResult, SubagentsOutput, SubagentsSummary},
worker::{SubagentRunInput, failed_result_with_session, format_duration, run_one_subagent},
};
use crate::{
agent::cancellation::{AgentCancellation, AgentCancellationHandle},
output::{ActivityEvent, ActivityId, ActivityStatus},
};
use std::{
path::PathBuf,
sync::{
Arc, Mutex,
atomic::{AtomicBool, Ordering as AtomicOrdering},
mpsc,
},
thread,
time::{Duration, Instant},
};
const DEFAULT_SCHEDULER_POLL_INTERVAL: Duration = Duration::from_millis(250);
const DEFAULT_SCHEDULER_STALL_AFTER: Duration = Duration::from_secs(300);
const SCHEDULER_WORKER_SHUTDOWN_TIMEOUT: Duration = Duration::from_millis(500);
pub(super) const SCHEDULER_EVENT_CHANNEL_BOUND: usize = 512;
#[derive(Debug, Clone)]
pub(super) struct PreparedSubagentTask {
pub(super) index: usize,
pub(super) id: String,
pub(super) task: SubagentTask,
pub(super) cwd: PathBuf,
pub(super) cancellation: AgentCancellation,
pub(super) cancel_handle: AgentCancellationHandle,
}
pub(super) type SharedTaskQueue = Arc<Mutex<std::vec::IntoIter<PreparedSubagentTask>>>;
#[cfg(test)]
pub(super) type SharedTaskResults = Arc<Mutex<Vec<Option<SubagentTaskResult>>>>;
#[derive(Clone, Copy)]
pub(super) struct SchedulerWaitConfig {
pub(super) poll_interval: Duration,
pub(super) stall_after: Duration,
}
impl Default for SchedulerWaitConfig {
fn default() -> Self {
Self {
poll_interval: DEFAULT_SCHEDULER_POLL_INTERVAL,
stall_after: DEFAULT_SCHEDULER_STALL_AFTER,
}
}
}
#[derive(Clone)]
pub(super) struct TaskActivityFinisher {
finished: Arc<AtomicBool>,
}
impl TaskActivityFinisher {
pub(super) fn new() -> Self {
Self {
finished: Arc::new(AtomicBool::new(false)),
}
}
pub(super) fn finish(
&self,
config: &SubagentRunConfig,
id: ActivityId,
status: ActivityStatus,
) {
if !self.finished.swap(true, AtomicOrdering::SeqCst) {
emit_activity(
config,
ActivityEvent::Finished {
id,
status,
metadata: None,
},
);
}
}
}
pub(super) struct SchedulerTaskState {
prepared: PreparedSubagentTask,
session_id: Option<String>,
session_path: Option<PathBuf>,
last_progress: Option<Instant>,
completed: bool,
finisher: TaskActivityFinisher,
cancel_handle: AgentCancellationHandle,
}
pub(super) enum WorkerEvent {
TaskStarted {
index: usize,
},
TaskSession {
index: usize,
session_id: String,
session_path: PathBuf,
},
TaskProgress {
index: usize,
at: Instant,
},
TaskFinished {
index: usize,
result: Box<SubagentTaskResult>,
},
WorkerFailed(String),
}
pub(super) struct SchedulerTaskReporter {
pub(super) index: usize,
pub(super) sender: mpsc::SyncSender<WorkerEvent>,
}
impl SchedulerTaskReporter {
pub(super) fn progress(&self) {
send_worker_event(
&self.sender,
WorkerEvent::TaskProgress {
index: self.index,
at: Instant::now(),
},
);
}
pub(super) fn session(&self, session_id: String, session_path: PathBuf) {
send_terminal_worker_event(
&self.sender,
WorkerEvent::TaskSession {
index: self.index,
session_id,
session_path,
},
);
}
}
pub(super) struct SubagentScheduler {
queue: SharedTaskQueue,
states: Vec<SchedulerTaskState>,
concurrency: usize,
config: SubagentRunConfig,
batch_id: ActivityId,
pub(super) wait: SchedulerWaitConfig,
}
impl SubagentScheduler {
pub(super) fn new(
tasks: Vec<SubagentTask>,
concurrency: usize,
config: SubagentRunConfig,
batch_id: ActivityId,
) -> anyhow::Result<Self> {
let prepared = tasks
.into_iter()
.enumerate()
.map(|(index, task)| {
let cwd = resolve_child_cwd(
&config.parent_cwd,
task.cwd.as_deref(),
config.parent_tools.subagents_absolute_paths(),
)?;
let (cancellation, cancel_handle) = config.cancellation.child_token();
Ok(PreparedSubagentTask {
index,
id: format!("g{}", index + 1),
task,
cwd,
cancellation,
cancel_handle,
})
})
.collect::<anyhow::Result<Vec<_>>>()?;
let states = prepared
.iter()
.cloned()
.map(|prepared| {
let cancel_handle = prepared.cancel_handle.clone();
SchedulerTaskState {
prepared,
session_id: None,
session_path: None,
last_progress: None,
completed: false,
finisher: TaskActivityFinisher::new(),
cancel_handle,
}
})
.collect();
Ok(Self {
concurrency,
states,
queue: Arc::new(Mutex::new(prepared.into_iter())),
config,
batch_id,
wait: SchedulerWaitConfig::default(),
})
}
pub(super) fn run(mut self) -> anyhow::Result<SubagentsOutput> {
let (sender, receiver) = mpsc::sync_channel(SCHEDULER_EVENT_CHANNEL_BOUND);
let mut workers = Vec::new();
for _ in 0..self.concurrency {
workers.push(Self::spawn_worker(
Arc::clone(&self.queue),
self.config.clone(),
self.batch_id.clone(),
sender.clone(),
self.states
.iter()
.map(|state| state.finisher.clone())
.collect(),
));
}
drop(sender);
let mut results = vec![None; self.states.len()];
let mut remaining = self.states.len();
let mut worker_error = None;
let mut observed_stall = false;
while remaining > 0 {
if self.config.cancellation.is_canceled() {
let cleanup_warning =
self.cancel_and_drain_workers(&mut workers, &receiver, &mut results)?;
let error: anyhow::Error = self.config.cancellation.check().unwrap_err();
return Err(match cleanup_warning {
Some(warning) => error.context(warning),
None => error,
});
}
match receiver.recv_timeout(self.wait.poll_interval) {
Ok(event) => {
if let Err(error) = self.apply_worker_event(event, &mut results, &mut remaining)
&& worker_error.is_none()
{
let cleanup_warning =
self.cancel_and_drain_workers(&mut workers, &receiver, &mut results)?;
worker_error = Some(match cleanup_warning {
Some(warning) => error.context(warning),
None => error,
});
break;
}
}
Err(mpsc::RecvTimeoutError::Timeout) => {}
Err(mpsc::RecvTimeoutError::Disconnected) => break,
}
observed_stall |= self.fail_stalled_tasks(&mut results, &mut remaining)?;
if observed_stall && self.no_active_unfinished_tasks() {
self.fail_unstarted_tasks(&mut results, &mut remaining)?;
}
}
self.config.cancellation.check()?;
if !observed_stall && remaining == 0 {
self.wait_for_worker_completion(&mut workers, &receiver, &mut results, &mut remaining)?;
}
if observed_stall {
self.wait_for_worker_shutdown_blocking(
&mut workers,
&receiver,
&mut results,
&mut remaining,
)?;
}
if self.config.cancellation.is_canceled() {
self.config.cancellation.check()?;
}
if let Some(error) = worker_error {
return Err(error);
}
let results = results
.into_iter()
.enumerate()
.map(|(index, result)| {
result.ok_or_else(|| anyhow::anyhow!("subagent result missing at index {index}"))
})
.collect::<anyhow::Result<Vec<_>>>()?;
Ok(output_from_results(results))
}
fn spawn_worker(
queue: SharedTaskQueue,
config: SubagentRunConfig,
batch_id: ActivityId,
sender: mpsc::SyncSender<WorkerEvent>,
finishers: Vec<TaskActivityFinisher>,
) -> thread::JoinHandle<()> {
thread::spawn(move || {
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
Self::run_worker(queue, config, batch_id, sender.clone(), finishers)
}));
match result {
Ok(Ok(())) => {}
Ok(Err(error)) => {
send_terminal_worker_event(
&sender,
WorkerEvent::WorkerFailed(error.to_string()),
);
}
Err(_) => {
send_terminal_worker_event(
&sender,
WorkerEvent::WorkerFailed("subagent worker panicked".to_string()),
);
}
}
})
}
fn wait_for_worker_completion(
&mut self,
workers: &mut Vec<thread::JoinHandle<()>>,
receiver: &mpsc::Receiver<WorkerEvent>,
results: &mut [Option<SubagentTaskResult>],
remaining: &mut usize,
) -> anyhow::Result<()> {
self.wait_for_worker_shutdown_until(workers, receiver, results, remaining, None)
.map(|_| ())
}
fn wait_for_worker_shutdown_blocking(
&mut self,
workers: &mut Vec<thread::JoinHandle<()>>,
receiver: &mpsc::Receiver<WorkerEvent>,
results: &mut [Option<SubagentTaskResult>],
remaining: &mut usize,
) -> anyhow::Result<()> {
let deadline = Instant::now() + SCHEDULER_WORKER_SHUTDOWN_TIMEOUT;
self.wait_for_worker_shutdown_until(workers, receiver, results, remaining, Some(deadline))
.map(|_| ())
}
fn wait_for_worker_shutdown_until(
&mut self,
workers: &mut Vec<thread::JoinHandle<()>>,
receiver: &mpsc::Receiver<WorkerEvent>,
results: &mut [Option<SubagentTaskResult>],
remaining: &mut usize,
deadline: Option<Instant>,
) -> anyhow::Result<Option<String>> {
while !workers.is_empty() {
if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
break;
}
let wait = deadline
.map(|deadline| deadline.saturating_duration_since(Instant::now()))
.unwrap_or(self.wait.poll_interval)
.min(self.wait.poll_interval);
match receiver.recv_timeout(wait) {
Ok(event) => self.apply_worker_event(event, results, remaining)?,
Err(mpsc::RecvTimeoutError::Timeout | mpsc::RecvTimeoutError::Disconnected) => {}
}
join_finished_workers(workers);
}
while let Ok(event) = receiver.try_recv() {
self.apply_worker_event(event, results, remaining)?;
}
if workers.is_empty() {
Ok(None)
} else {
for index in 0..self.states.len() {
self.fail_task(
index,
"subagent worker shutdown exceeded 500ms; detached worker cleanup incomplete and still occupies capacity until it exits"
.to_string(),
results,
remaining,
)?;
}
workers.clear();
Ok(Some(
"subagent worker cleanup incomplete after 500ms grace period; detached worker still occupies capacity until it exits".to_string(),
))
}
}
pub(super) fn cancel_and_drain_workers(
&mut self,
workers: &mut Vec<thread::JoinHandle<()>>,
receiver: &mpsc::Receiver<WorkerEvent>,
results: &mut [Option<SubagentTaskResult>],
) -> anyhow::Result<Option<String>> {
self.cancel_unfinished_tasks();
let deadline = Instant::now() + SCHEDULER_WORKER_SHUTDOWN_TIMEOUT;
self.wait_for_worker_shutdown_until(workers, receiver, results, &mut 0, Some(deadline))
}
fn run_worker(
queue: SharedTaskQueue,
config: SubagentRunConfig,
batch_id: ActivityId,
sender: mpsc::SyncSender<WorkerEvent>,
finishers: Vec<TaskActivityFinisher>,
) -> anyhow::Result<()> {
loop {
config.cancellation.check()?;
let Some(prepared) = Self::next_task(&queue)? else {
break;
};
if prepared.cancellation.is_canceled() {
continue;
}
let index = prepared.index;
send_worker_event(&sender, WorkerEvent::TaskStarted { index });
let reporter = SchedulerTaskReporter {
index,
sender: sender.clone(),
};
let result = run_one_subagent(SubagentRunInput {
id: prepared.id,
task: prepared.task,
cwd: prepared.cwd,
config: &config,
cancellation: prepared.cancellation,
batch_id: &batch_id,
reporter: Some(reporter),
finisher: finishers[index].clone(),
});
send_terminal_worker_event(
&sender,
WorkerEvent::TaskFinished {
index,
result: Box::new(result),
},
);
config.cancellation.check()?;
}
Ok(())
}
fn apply_worker_event(
&mut self,
event: WorkerEvent,
results: &mut [Option<SubagentTaskResult>],
remaining: &mut usize,
) -> anyhow::Result<()> {
match event {
WorkerEvent::TaskStarted { index } => {
let state = self.state_mut(index)?;
state.last_progress = Some(Instant::now());
}
WorkerEvent::TaskSession {
index,
session_id,
session_path,
} => {
let state = self.state_mut(index)?;
state.session_id = Some(session_id);
state.session_path = Some(session_path);
state.last_progress = Some(Instant::now());
}
WorkerEvent::TaskProgress { index, at } => {
self.state_mut(index)?.last_progress = Some(at);
}
WorkerEvent::TaskFinished { index, result } => {
let state = self.state_mut(index)?;
if !state.completed {
state.completed = true;
results[index] = Some(*result);
*remaining = remaining.saturating_sub(1);
}
}
WorkerEvent::WorkerFailed(error) => return Err(anyhow::anyhow!(error)),
}
Ok(())
}
fn fail_stalled_tasks(
&mut self,
results: &mut [Option<SubagentTaskResult>],
remaining: &mut usize,
) -> anyhow::Result<bool> {
let mut failed_any = false;
let now = Instant::now();
for index in 0..self.states.len() {
let Some(last_progress) = self.states[index].last_progress else {
continue;
};
if self.states[index].completed
|| now.duration_since(last_progress) < self.wait.stall_after
{
continue;
}
let elapsed = now.duration_since(last_progress);
let error = format!(
"subagent stalled with no activity for {}",
format_duration(elapsed)
);
self.fail_task(index, error, results, remaining)?;
failed_any = true;
}
Ok(failed_any)
}
fn fail_unstarted_tasks(
&mut self,
results: &mut [Option<SubagentTaskResult>],
remaining: &mut usize,
) -> anyhow::Result<()> {
for index in 0..self.states.len() {
if self.states[index].completed || self.states[index].last_progress.is_some() {
continue;
}
self.fail_task(
index,
"subagent did not start before scheduler workers stalled".to_string(),
results,
remaining,
)?;
}
Ok(())
}
fn no_active_unfinished_tasks(&self) -> bool {
self.states
.iter()
.filter(|state| !state.completed)
.all(|state| state.last_progress.is_none())
}
fn cancel_unfinished_tasks(&mut self) {
for state in &mut self.states {
if !state.completed {
state.cancel_handle.cancel();
}
}
}
fn fail_task(
&mut self,
index: usize,
error: String,
results: &mut [Option<SubagentTaskResult>],
remaining: &mut usize,
) -> anyhow::Result<()> {
let config = self.config.clone();
let batch_id = self.batch_id.clone();
let state = self.state_mut(index)?;
if state.completed {
return Ok(());
}
state.completed = true;
state.cancel_handle.cancel();
let task_activity_id = ActivityId::new(format!(
"{}/{}",
batch_id.as_str(),
state.prepared.id.as_str()
));
state
.finisher
.finish(&config, task_activity_id, ActivityStatus::Failed);
results[index] = Some(failed_result_with_session(
state.prepared.id.clone(),
state.prepared.task.clone(),
state.prepared.cwd.clone(),
state.session_id.clone(),
state.session_path.clone(),
error,
));
*remaining = remaining.saturating_sub(1);
Ok(())
}
fn state_mut(&mut self, index: usize) -> anyhow::Result<&mut SchedulerTaskState> {
self.states
.get_mut(index)
.ok_or_else(|| anyhow::anyhow!("subagent task index out of bounds: {index}"))
}
pub(super) fn next_task(
queue: &SharedTaskQueue,
) -> anyhow::Result<Option<PreparedSubagentTask>> {
queue
.lock()
.map_err(|_| anyhow::anyhow!("subagent queue mutex poisoned"))
.map(|mut queue| queue.next())
}
#[cfg(test)]
pub(super) fn apply_worker_failure_for_test(&mut self, error: String) -> (String, Vec<bool>) {
let mut results = vec![None; self.states.len()];
let mut remaining = self.states.len();
let error = self
.apply_worker_event(
WorkerEvent::WorkerFailed(error),
&mut results,
&mut remaining,
)
.unwrap_err()
.to_string();
self.cancel_unfinished_tasks();
let cancellations = self
.states
.iter()
.map(|state| state.prepared.cancellation.is_canceled())
.collect();
(error, cancellations)
}
#[cfg(test)]
pub(super) fn record_result(
results: &SharedTaskResults,
index: usize,
result: SubagentTaskResult,
) -> anyhow::Result<()> {
let mut results = results
.lock()
.map_err(|_| anyhow::anyhow!("subagent results mutex poisoned"))?;
let slot = results
.get_mut(index)
.ok_or_else(|| anyhow::anyhow!("subagent result index out of bounds: {index}"))?;
*slot = Some(result);
Ok(())
}
#[cfg(test)]
pub(super) fn finish_results(
results: SharedTaskResults,
) -> anyhow::Result<Vec<SubagentTaskResult>> {
Arc::try_unwrap(results)
.map_err(|_| anyhow::anyhow!("subagent results still shared"))?
.into_inner()
.map_err(|_| anyhow::anyhow!("subagent results mutex poisoned"))?
.into_iter()
.enumerate()
.map(|(index, result)| {
result.ok_or_else(|| anyhow::anyhow!("subagent result missing at index {index}"))
})
.collect()
}
}
pub(super) fn join_finished_workers(workers: &mut Vec<thread::JoinHandle<()>>) {
let mut index = 0;
while index < workers.len() {
if workers[index].is_finished() {
let handle = workers.swap_remove(index);
let _ = handle.join();
} else {
index += 1;
}
}
}
pub(super) fn send_worker_event(sender: &mpsc::SyncSender<WorkerEvent>, event: WorkerEvent) {
let _ = sender.try_send(event);
}
pub(super) fn send_terminal_worker_event(
sender: &mpsc::SyncSender<WorkerEvent>,
event: WorkerEvent,
) {
let _ = sender.send(event);
}
pub(super) fn output_from_results(results: Vec<SubagentTaskResult>) -> SubagentsOutput {
let completed = results
.iter()
.filter(|result| result.status == SubagentStatus::Completed)
.count();
let failed = results
.iter()
.filter(|result| result.status == SubagentStatus::Failed)
.count();
let total_tokens = results
.iter()
.map(|result| result.total_tokens)
.try_fold(0_u64, |total, tokens| {
tokens.map(|tokens| total.saturating_add(tokens))
});
SubagentsOutput {
summary: SubagentsSummary {
total: results.len(),
completed,
failed,
total_tokens,
},
results,
}
}