From 8b5106902c88c160fa77037d982532b3a4244ac5 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Fri, 31 Jul 2026 08:19:10 -0400 Subject: [PATCH 1/4] feat: reuse completed subagent sessions --- .../src/commands/run/run_progress/event.rs | 12 + .../src/commands/run/run_progress/mod.rs | 51 +- .../run/run_progress/stage_display.rs | 26 + lib/components/fabro-agent/src/cli.rs | 36 +- .../fabro-agent/src/profiles/claude5_tools.rs | 18 +- lib/components/fabro-agent/src/subagent.rs | 1036 +++++++++++++---- lib/components/fabro-agent/src/types.rs | 132 ++- lib/components/fabro-store/src/run_state.rs | 94 +- .../fabro-workflow/src/event/convert.rs | 47 +- .../fabro-workflow/src/event/names.rs | 24 +- .../fabro-workflow/src/handler/llm/api.rs | 5 +- .../fabro-types/src/run_event/agent.rs | 43 +- .../fabro-types/src/run_event/mod.rs | 35 + 13 files changed, 1206 insertions(+), 353 deletions(-) diff --git a/lib/apps/fabro-cli/src/commands/run/run_progress/event.rs b/lib/apps/fabro-cli/src/commands/run/run_progress/event.rs index eca9e8d67..9b8509c0e 100644 --- a/lib/apps/fabro-cli/src/commands/run/run_progress/event.rs +++ b/lib/apps/fabro-cli/src/commands/run/run_progress/event.rs @@ -194,6 +194,12 @@ pub(super) enum ProgressEvent { agent_id: String, task: String, }, + SubagentTurnStarted { + stage_node_id: String, + agent_id: String, + task: String, + generation: u64, + }, SubagentCompleted { stage_node_id: String, agent_id: String, @@ -431,6 +437,12 @@ pub(super) fn from_run_event(stored: &RunEvent) -> Option { agent_id: props.agent_id.clone(), task: props.task.clone(), }), + EventBody::AgentSubTurnStarted(props) => Some(ProgressEvent::SubagentTurnStarted { + stage_node_id: node_id, + agent_id: props.agent_id.clone(), + task: props.task.clone(), + generation: props.generation, + }), EventBody::AgentSubCompleted(props) => Some(ProgressEvent::SubagentCompleted { stage_node_id: node_id, agent_id: props.agent_id.clone(), diff --git a/lib/apps/fabro-cli/src/commands/run/run_progress/mod.rs b/lib/apps/fabro-cli/src/commands/run/run_progress/mod.rs index d12fbdd88..44fbb484b 100644 --- a/lib/apps/fabro-cli/src/commands/run/run_progress/mod.rs +++ b/lib/apps/fabro-cli/src/commands/run/run_progress/mod.rs @@ -371,6 +371,20 @@ impl ProgressUI { self.stage .on_subagent_spawned(renderer, &stage_node_id, &agent_id, &task); } + ProgressEvent::SubagentTurnStarted { + stage_node_id, + agent_id, + task, + generation, + } => { + self.stage.on_subagent_turn_started( + renderer, + &stage_node_id, + &agent_id, + &task, + generation, + ); + } ProgressEvent::SubagentCompleted { stage_node_id, agent_id, @@ -970,13 +984,15 @@ mod tests { }, }), agent_event("code", AgentEvent::SubAgentSpawned { - agent_id: "a1".into(), - depth: 1, - task: "review recent changes".into(), + agent_id: "a1".into(), + depth: 1, + task: "review recent changes".into(), + generation: 1, }), agent_event("code", AgentEvent::SubAgentCompleted { agent_id: "a1".into(), depth: 1, + generation: 1, success: true, turns_used: 3, }), @@ -1332,9 +1348,10 @@ mod tests { emit( &mut ui, agent_event("code", AgentEvent::SubAgentSpawned { - agent_id: "a1".into(), - depth: 1, - task: "review recent changes".into(), + agent_id: "a1".into(), + depth: 1, + task: "review recent changes".into(), + generation: 1, }), ); emit( @@ -1342,10 +1359,30 @@ mod tests { agent_event("code", AgentEvent::SubAgentCompleted { agent_id: "a1".into(), depth: 1, + generation: 1, success: true, turns_used: 3, }), ); + emit( + &mut ui, + agent_event("code", AgentEvent::SubAgentTurnStarted { + agent_id: "a1".into(), + depth: 1, + task: "fix the review findings".into(), + generation: 2, + }), + ); + emit( + &mut ui, + agent_event("code", AgentEvent::SubAgentCompleted { + agent_id: "a1".into(), + depth: 1, + generation: 2, + success: true, + turns_used: 2, + }), + ); emit(&mut ui, Event::SetupStarted { command_count: 1 }); emit(&mut ui, Event::SetupCommandCompleted { command: "bun install".into(), @@ -1363,6 +1400,8 @@ mod tests { ⚠ retry: gpt-5-mini attempt 2 (busy, delay 1s) ▸ subagent[a1] "review recent changes" ✓ subagent[a1] (3 turns) + ↻ subagent[a1] turn 2 "fix the review findings" + ✓ subagent[a1] (2 turns) ✓ [1/1] bun install 2s Setup: 1 command (2s) ✓ Code 5s (1 turns, 0 tools, 1.5k toks) diff --git a/lib/apps/fabro-cli/src/commands/run/run_progress/stage_display.rs b/lib/apps/fabro-cli/src/commands/run/run_progress/stage_display.rs index e71072fb2..235f70918 100644 --- a/lib/apps/fabro-cli/src/commands/run/run_progress/stage_display.rs +++ b/lib/apps/fabro-cli/src/commands/run/run_progress/stage_display.rs @@ -655,6 +655,32 @@ impl StageDisplay { ); } + pub(super) fn on_subagent_turn_started( + &mut self, + renderer: &ProgressRenderer, + stage_node_id: &str, + agent_id: &str, + task: &str, + generation: u64, + ) { + if !self.verbose { + return; + } + + self.insert_subagent_line_for_stage( + renderer, + stage_node_id, + &renderer + .styles() + .dim + .apply_to(format!( + "\u{21bb} subagent[{agent_id}] turn {generation} \"{}\"", + styles::truncate(task, 50) + )) + .to_string(), + ); + } + fn finish_stage( &mut self, renderer: &ProgressRenderer, diff --git a/lib/components/fabro-agent/src/cli.rs b/lib/components/fabro-agent/src/cli.rs index 666235737..d6f31d23a 100644 --- a/lib/components/fabro-agent/src/cli.rs +++ b/lib/components/fabro-agent/src/cli.rs @@ -692,7 +692,7 @@ pub async fn run_with_args_and_client_and_catalog( agent_id, depth, task, - .. + generation, } => { let task_preview = if task.len() > 60 { &task[..task.floor_char_boundary(60)] @@ -702,40 +702,64 @@ pub async fn run_with_args_and_client_and_catalog( eprintln!( " {}", s.dim.apply_to(format!( - "{child_prefix}\u{25b6} subagent {agent_id} spawned (depth={depth}) task={task_preview:?}" + "{child_prefix}\u{25b6} subagent {agent_id} spawned (depth={depth}, generation={generation}) task={task_preview:?}" + )), + ); + } + AgentEvent::SubAgentTurnStarted { + agent_id, + depth, + task, + generation, + } => { + let task_preview = if task.len() > 60 { + &task[..task.floor_char_boundary(60)] + } else { + task + }; + eprintln!( + " {}", + s.dim.apply_to(format!( + "{child_prefix}\u{25b6} subagent {agent_id} turn started (depth={depth}, generation={generation}) task={task_preview:?}" )), ); } AgentEvent::SubAgentCompleted { agent_id, depth, + generation, success, turns_used, } => { eprintln!( " {}", s.dim.apply_to(format!( - "{child_prefix}\u{25a0} subagent {agent_id} completed (depth={depth}, success={success}, turns={turns_used})" + "{child_prefix}\u{25a0} subagent {agent_id} completed (depth={depth}, generation={generation}, success={success}, turns={turns_used})" )), ); } AgentEvent::SubAgentFailed { agent_id, depth, + generation, error, } => { eprintln!( " {}", s.red.apply_to(format!( - "{child_prefix}\u{2717} subagent {agent_id} failed (depth={depth}): {error}" + "{child_prefix}\u{2717} subagent {agent_id} failed (depth={depth}, generation={generation}): {error}" )), ); } - AgentEvent::SubAgentClosed { agent_id, depth } => { + AgentEvent::SubAgentClosed { + agent_id, + depth, + generation, + } => { eprintln!( " {}", s.dim.apply_to(format!( - "{child_prefix}\u{25a0} subagent {agent_id} closed (depth={depth})" + "{child_prefix}\u{25a0} subagent {agent_id} closed (depth={depth}, generation={generation})" )), ); } diff --git a/lib/components/fabro-agent/src/profiles/claude5_tools.rs b/lib/components/fabro-agent/src/profiles/claude5_tools.rs index ab319b965..d7c2b3e80 100644 --- a/lib/components/fabro-agent/src/profiles/claude5_tools.rs +++ b/lib/components/fabro-agent/src/profiles/claude5_tools.rs @@ -383,7 +383,7 @@ pub(crate) fn make_task_stop_tool(supervisor: SubAgentSupervisor) -> RegisteredT RegisteredTool { definition: definition( NativeTool::StopAgent, - "Stop a running background agent by task ID.", + "Stop a running or completed background agent by task ID.", serde_json::json!({ "type": "object", "properties": { @@ -416,7 +416,7 @@ pub(crate) fn make_send_message_tool(supervisor: SubAgentSupervisor) -> Register RegisteredTool { definition: definition( NativeTool::MessageAgent, - "Send additional instructions to a running background agent by its task ID.", + "Send additional instructions to a background agent by its task ID. A running agent receives them at a safe turn boundary. A completed agent starts another turn in the same session with its existing history.", serde_json::json!({ "type": "object", "properties": { @@ -590,11 +590,17 @@ mod tests { assert_schema(&make_task_stop_tool(supervisor.clone()), &["task_id"], &[ "task_id", ]); - assert_schema( - &make_send_message_tool(supervisor), - &["message", "summary", "to"], - &["message", "to"], + let send_message = make_send_message_tool(supervisor); + assert_schema(&send_message, &["message", "summary", "to"], &[ + "message", "to", + ]); + assert!( + send_message + .definition + .description + .contains("completed agent") ); + assert!(send_message.definition.description.contains("same session")); } #[tokio::test] diff --git a/lib/components/fabro-agent/src/subagent.rs b/lib/components/fabro-agent/src/subagent.rs index 476094a94..2025cbd04 100644 --- a/lib/components/fabro-agent/src/subagent.rs +++ b/lib/components/fabro-agent/src/subagent.rs @@ -1,11 +1,11 @@ use std::collections::{HashMap, VecDeque}; -use std::sync::{Arc, Mutex, RwLock}; +use std::sync::{Arc, Mutex, RwLock, Weak}; use std::time::Duration; use fabro_llm::types::ToolDefinition; use fabro_util::error as util_error; use futures::future; -use tokio::sync::{oneshot, watch}; +use tokio::sync::{mpsc, oneshot, watch}; use tokio::task::{AbortHandle, JoinHandle}; use tokio::time::{Instant, timeout_at}; use tokio_util::sync::CancellationToken; @@ -14,7 +14,7 @@ use crate::error::{Error, InterruptReason}; use crate::session::{Session, SessionShutdownReason}; use crate::tool_registry::{RegisteredTool, ToolSource}; use crate::tools::required_str; -use crate::types::{AgentEvent, SessionEvent}; +use crate::types::{AgentEvent, SessionEvent, SessionState}; pub type SessionFactory = Arc Session + Send + Sync>; @@ -87,9 +87,26 @@ pub enum SubAgentStatus { } const SUBAGENT_SHUTDOWN_GRACE: Duration = Duration::from_secs(5); +const INITIAL_SUBAGENT_GENERATION: u64 = 1; +const SUBAGENT_COMMAND_CAPACITY: usize = 1; + +#[derive(Debug)] +enum SubAgentCommand { + Start { generation: u64, prompt: String }, +} + +struct ParentNotificationState { + description: String, + pending_generations: VecDeque, +} struct SubAgent { status: watch::Sender, + generation: u64, + results: HashMap>, + reusable: bool, + command_tx: mpsc::Sender, + runner_stop: CancellationToken, cleanup_done: watch::Sender, cleanup_started: bool, monitor_task: Option>, @@ -99,14 +116,14 @@ struct SubAgent { followup_queue: Arc>>, cancel_token: CancellationToken, depth: usize, - /// Task description, set when the parent should receive this child's - /// terminal result automatically. Cleared once the result is delivered, - /// the parent retrieves it explicitly, or the agent is shut down. + /// Registration for generations whose results should be delivered to the + /// parent automatically. The description remains available so a later + /// turn in the same child session can register its own result. /// - /// Keeping this beside the status it is delivered with means a - /// notification cannot be registered before -- or suppressed after -- the - /// state it describes: there is only one lock and one ordering. - parent_notification: Option, + /// Keeping this beside the generation results means a notification cannot + /// be registered before -- or suppressed after -- the state it describes: + /// there is only one lock and one ordering. + parent_notification: Option, /// Spawn order, so a batch is delivered oldest-first rather than in /// whatever order the map happens to iterate. spawn_seq: u64, @@ -114,6 +131,7 @@ struct SubAgent { impl Drop for SubAgent { fn drop(&mut self) { + self.runner_stop.cancel(); self.cancel_token.cancel(); self.child_abort_handle.abort(); if let Some(task) = self.monitor_task.take() { @@ -130,13 +148,16 @@ impl Drop for SubAgent { #[derive(Default)] struct SupervisorState { - agents: HashMap, - next_spawn_seq: u64, + agents: HashMap, + next_spawn_seq: u64, + lifecycle_events: VecDeque, + lifecycle_draining: bool, } struct ShutdownWork { agent_id: String, depth: usize, + generation: u64, close_running_agent: bool, status: watch::Sender, cleanup_done: watch::Sender, @@ -144,10 +165,13 @@ struct ShutdownWork { event_forwarder: Option>, child_abort_handle: AbortHandle, cancel_token: CancellationToken, + runner_stop: CancellationToken, + state: Weak>, } impl Drop for ShutdownWork { fn drop(&mut self) { + self.runner_stop.cancel(); self.cancel_token.cancel(); self.child_abort_handle.abort(); if let Some(task) = self.monitor_task.take() { @@ -182,47 +206,33 @@ fn signal_notifications(changed: &watch::Sender) { }); } -fn spawn_result_monitor( - child_task: JoinHandle>, - status: watch::Sender, - event_callback: Arc>>, - notifications_changed: Arc>, - agent_id: String, - depth: usize, -) -> JoinHandle<()> { - tokio::spawn(async move { - let task_result = match child_task.await { - Ok(result) => result, - Err(err) => Err(Error::InvalidState(format!( - "Agent task failed to join: {err}" - ))), - }; - let committed = status.send_if_modified(|current| { - if matches!(current, SubAgentStatus::Running) { - *current = SubAgentStatus::Finished(task_result.clone()); - true - } else { - false - } - }); - if !committed { +fn queue_lifecycle_event(state: &mut SupervisorState, event: AgentEvent) { + state.lifecycle_events.push_back(event); +} + +/// Deliver lifecycle callbacks in the same order as the state transitions +/// that queued them. Callbacks run without the supervisor lock held and may +/// safely call back into the supervisor. +fn drain_lifecycle_events( + state: &Arc>, + event_callback: &Arc>>, +) { + { + let mut state = state.lock().expect("subagent state lock poisoned"); + if state.lifecycle_draining { return; } - // The status this agent will be delivered with is now committed. - signal_notifications(¬ifications_changed); + state.lifecycle_draining = true; + } - let event = match &task_result { - Ok(result) => AgentEvent::SubAgentCompleted { - agent_id: agent_id.clone(), - depth, - success: result.success, - turns_used: result.turns_used, - }, - Err(error) => AgentEvent::SubAgentFailed { - agent_id: agent_id.clone(), - depth, - error: error.clone(), - }, + loop { + let event = { + let mut state = state.lock().expect("subagent state lock poisoned"); + let Some(event) = state.lifecycle_events.pop_front() else { + state.lifecycle_draining = false; + return; + }; + event }; let callback = event_callback .read() @@ -231,6 +241,237 @@ fn spawn_result_monitor( if let Some(callback) = callback { callback(SubAgentCallbackEvent::Lifecycle(event)); } + } +} + +enum TurnCommit { + Continue(String), + Finished, + Stopping, +} + +fn completion_event( + agent_id: &str, + depth: usize, + generation: u64, + result: &Result, +) -> AgentEvent { + match result { + Ok(result) => AgentEvent::SubAgentCompleted { + agent_id: agent_id.to_string(), + depth, + generation, + success: result.success, + turns_used: result.turns_used, + }, + Err(error) => AgentEvent::SubAgentFailed { + agent_id: agent_id.to_string(), + depth, + generation, + error: error.clone(), + }, + } +} + +/// Commit one generation result or claim a follow-up that raced its final +/// boundary. The supervisor state lock is acquired before the follow-up queue +/// lock, which is also the ordering used by `send_input`. +fn commit_turn_result( + state: &Arc>, + event_callback: &Arc>>, + notifications_changed: &watch::Sender, + agent_id: &str, + depth: usize, + generation: u64, + result: &Result, + reusable: bool, +) -> TurnCommit { + let outcome = { + let mut state = state.lock().expect("subagent state lock poisoned"); + let Some(agent) = state.agents.get_mut(agent_id) else { + return TurnCommit::Stopping; + }; + if agent.generation != generation + || !matches!(*agent.status.borrow(), SubAgentStatus::Running) + { + return TurnCommit::Stopping; + } + + if reusable { + let next_prompt = agent + .followup_queue + .lock() + .expect("followup queue lock poisoned") + .pop_front(); + if let Some(next_prompt) = next_prompt { + return TurnCommit::Continue(next_prompt); + } + } + + agent.results.insert(generation, result.clone()); + agent.reusable = reusable; + agent + .status + .send_replace(SubAgentStatus::Finished(result.clone())); + queue_lifecycle_event( + &mut state, + completion_event(agent_id, depth, generation, result), + ); + TurnCommit::Finished + }; + + signal_notifications(notifications_changed); + drain_lifecycle_events(state, event_callback); + outcome +} + +async fn run_subagent_session( + mut session: Session, + state: Weak>, + event_callback: Arc>>, + notifications_changed: Arc>, + agent_id: String, + depth: usize, + initial_prompt: String, + mut command_rx: mpsc::Receiver, + runner_stop: CancellationToken, + start_rx: oneshot::Receiver<()>, +) { + if start_rx.await.is_err() { + return; + } + + if let Err(error) = session.initialize().await { + if let Some(state) = state.upgrade() { + let result = Err(error); + commit_turn_result( + &state, + &event_callback, + ¬ifications_changed, + &agent_id, + depth, + INITIAL_SUBAGENT_GENERATION, + &result, + false, + ); + } + runner_stop.cancelled().await; + let reason = if session.cancel_token().is_cancelled() { + SessionShutdownReason::Cancelled + } else { + SessionShutdownReason::Error + }; + session.shutdown(reason).await; + return; + } + + let mut command = SubAgentCommand::Start { + generation: INITIAL_SUBAGENT_GENERATION, + prompt: initial_prompt, + }; + 'commands: loop { + let SubAgentCommand::Start { + generation, + mut prompt, + } = command; + let generation_start_turns = session.history().turns().len(); + + loop { + let result = session + .process_input_with_output(&prompt) + .await + .and_then(|output| { + output.ok_or_else(|| { + Error::InvalidState( + "Subagent completed without a non-empty final response".to_string(), + ) + }) + }) + .map(|output| SubAgentResult { + output, + success: true, + turns_used: session + .history() + .turns() + .len() + .saturating_sub(generation_start_turns), + }); + let reusable = + session.state() == SessionState::Idle && !session.cancel_token().is_cancelled(); + let Some(state) = state.upgrade() else { + break 'commands; + }; + match commit_turn_result( + &state, + &event_callback, + ¬ifications_changed, + &agent_id, + depth, + generation, + &result, + reusable, + ) { + TurnCommit::Continue(next_prompt) => prompt = next_prompt, + TurnCommit::Finished => break, + TurnCommit::Stopping => break 'commands, + } + } + + command = tokio::select! { + biased; + () = runner_stop.cancelled() => break, + command = command_rx.recv() => { + let Some(command) = command else { + break; + }; + command + } + }; + } + + let reason = if session.cancel_token().is_cancelled() { + SessionShutdownReason::Cancelled + } else { + SessionShutdownReason::Completed + }; + session.shutdown(reason).await; +} + +fn spawn_runner_monitor( + runner_task: JoinHandle<()>, + state: Weak>, + event_callback: Arc>>, + notifications_changed: Arc>, + agent_id: String, + depth: usize, +) -> JoinHandle<()> { + tokio::spawn(async move { + let Err(error) = runner_task.await else { + return; + }; + let Some(state) = state.upgrade() else { + return; + }; + let task_result = Err(Error::InvalidState(format!( + "Agent task failed to join: {error}" + ))); + let generation = { + let state = state.lock().expect("subagent state lock poisoned"); + let Some(agent) = state.agents.get(&agent_id) else { + return; + }; + agent.generation + }; + commit_turn_result( + &state, + &event_callback, + ¬ifications_changed, + &agent_id, + depth, + generation, + &task_result, + false, + ); }) } @@ -265,6 +506,7 @@ impl SubAgentSupervisor { .expect("subagent callback lock poisoned") = Some(cb); } + #[cfg(test)] fn emit_event(&self, event: AgentEvent) { let callback = self .event_callback @@ -299,7 +541,7 @@ impl SubAgentSupervisor { fn spawn_inner( &self, - mut session: Session, + session: Session, task_prompt: String, depth: usize, parent_notification_description: Option, @@ -351,55 +593,50 @@ impl SubAgentSupervisor { None }; - let task_prompt_for_spawn = task_prompt.clone(); let (start_tx, start_rx) = oneshot::channel(); - let child_task = tokio::spawn(async move { - let _ = start_rx.await; - let result = async { - session.initialize().await?; - let output = session - .process_input_with_output(&task_prompt_for_spawn) - .await? - .ok_or_else(|| { - Error::InvalidState( - "Subagent completed without a non-empty final response".to_string(), - ) - })?; - let turns = session.history().turns(); - Ok(SubAgentResult { - output, - success: true, - turns_used: turns.len(), - }) - } - .await; - let reason = match &result { - Ok(_) => SessionShutdownReason::Completed, - Err(Error::Interrupted(_)) => SessionShutdownReason::Cancelled, - Err(_) => SessionShutdownReason::Error, - }; - session.shutdown(reason).await; - result - }); - let child_abort_handle = child_task.abort_handle(); - let (status, _) = watch::channel(SubAgentStatus::Running); - let (cleanup_done, _) = watch::channel(false); + let (command_tx, command_rx) = mpsc::channel(SUBAGENT_COMMAND_CAPACITY); + let runner_stop = CancellationToken::new(); let child_depth = depth + 1; - let monitor_task = spawn_result_monitor( - child_task, - status.clone(), + let runner_task = tokio::spawn(run_subagent_session( + session, + Arc::downgrade(&self.state), + Arc::clone(&self.event_callback), + Arc::clone(&self.notifications_changed), + agent_id.clone(), + child_depth, + task_prompt.clone(), + command_rx, + runner_stop.clone(), + start_rx, + )); + let child_abort_handle = runner_task.abort_handle(); + let monitor_task = spawn_runner_monitor( + runner_task, + Arc::downgrade(&self.state), Arc::clone(&self.event_callback), Arc::clone(&self.notifications_changed), agent_id.clone(), child_depth, ); + let (status, _) = watch::channel(SubAgentStatus::Running); + let (cleanup_done, _) = watch::channel(false); { let mut state = self.state.lock().expect("subagent state lock poisoned"); let spawn_seq = state.next_spawn_seq; state.next_spawn_seq = state.next_spawn_seq.saturating_add(1); + let parent_notification = + parent_notification_description.map(|description| ParentNotificationState { + description, + pending_generations: VecDeque::from([INITIAL_SUBAGENT_GENERATION]), + }); state.agents.insert(agent_id.clone(), SubAgent { status, + generation: INITIAL_SUBAGENT_GENERATION, + results: HashMap::new(), + reusable: false, + command_tx, + runner_stop, cleanup_done, cleanup_started: false, monitor_task: Some(monitor_task), @@ -409,42 +646,92 @@ impl SubAgentSupervisor { followup_queue, cancel_token, depth: child_depth, - parent_notification: parent_notification_description, + parent_notification, spawn_seq, }); + queue_lifecycle_event(&mut state, AgentEvent::SubAgentSpawned { + agent_id: agent_id.clone(), + depth: child_depth, + task: task_prompt, + generation: INITIAL_SUBAGENT_GENERATION, + }); } signal_notifications(&self.notifications_changed); - - self.emit_event(AgentEvent::SubAgentSpawned { - agent_id: agent_id.clone(), - depth: child_depth, - task: task_prompt, - }); + drain_lifecycle_events(&self.state, &self.event_callback); let _ = start_tx.send(()); Ok(agent_id) } pub fn send_input(&self, agent_id: &str, message: &str) -> Result<(), Error> { - let followup_queue = { - let state = self.state.lock().expect("subagent state lock poisoned"); - let agent = state.agents.get(agent_id).ok_or_else(|| { + let resumed = { + let mut state = self.state.lock().expect("subagent state lock poisoned"); + let agent = state.agents.get_mut(agent_id).ok_or_else(|| { Error::InvalidState(format!( "No agent found with id: {agent_id} (it was never spawned)" )) })?; - if !matches!(*agent.status.borrow(), SubAgentStatus::Running) { - return Err(Error::InvalidState(format!( - "Agent {agent_id} is not running" - ))); + let status = agent.status.borrow().clone(); + match status { + SubAgentStatus::Running => { + agent + .followup_queue + .lock() + .expect("followup queue lock poisoned") + .push_back(message.to_string()); + None + } + SubAgentStatus::Finished(_) => { + if !agent.reusable { + return Err(Error::InvalidState(format!( + "Agent {agent_id} cannot accept more input because its session ended" + ))); + } + let permit = agent + .command_tx + .clone() + .try_reserve_owned() + .map_err(|error| { + Error::InvalidState(format!( + "Agent {agent_id} could not start another turn: {error}" + )) + })?; + let generation = agent.generation.checked_add(1).ok_or_else(|| { + Error::InvalidState(format!( + "Agent {agent_id} exhausted its turn generation" + )) + })?; + agent.generation = generation; + agent.reusable = false; + agent.status.send_replace(SubAgentStatus::Running); + if let Some(notification) = &mut agent.parent_notification { + notification.pending_generations.push_back(generation); + } + let depth = agent.depth; + queue_lifecycle_event(&mut state, AgentEvent::SubAgentTurnStarted { + agent_id: agent_id.to_string(), + depth, + task: message.to_string(), + generation, + }); + Some((permit, generation)) + } + SubAgentStatus::Closing | SubAgentStatus::Closed => { + return Err(Error::InvalidState(format!( + "Agent {agent_id} has been closed" + ))); + } } - Arc::clone(&agent.followup_queue) }; - followup_queue - .lock() - .expect("followup queue lock poisoned") - .push_back(message.to_string()); + if let Some((permit, generation)) = resumed { + signal_notifications(&self.notifications_changed); + drain_lifecycle_events(&self.state, &self.event_callback); + permit.send(SubAgentCommand::Start { + generation, + prompt: message.to_string(), + }); + } Ok(()) } @@ -454,45 +741,52 @@ impl SubAgentSupervisor { agent_id: &str, cancel: &CancellationToken, ) -> Result { - let mut status = { + let (generation, mut status) = { let state = self.state.lock().expect("subagent state lock poisoned"); - state - .agents - .get(agent_id) - .ok_or_else(|| { - Error::InvalidState(format!( - "No agent found with id: {agent_id} (it was never spawned)" - )) - })? - .status - .subscribe() + let agent = state.agents.get(agent_id).ok_or_else(|| { + Error::InvalidState(format!( + "No agent found with id: {agent_id} (it was never spawned)" + )) + })?; + (agent.generation, agent.status.subscribe()) }; loop { - let current = status.borrow().clone(); - match current { - SubAgentStatus::Running => { - tokio::select! { - biased; - () = cancel.cancelled() => { - self.ensure_closed(agent_id).await?; - return Err(Error::Interrupted(InterruptReason::Cancelled)); - } - changed = status.changed() => { - changed.map_err(|_| { - Error::InvalidState(format!( - "Agent {agent_id} result observer closed unexpectedly" - )) - })?; - } - } + let current = { + let state = self.state.lock().expect("subagent state lock poisoned"); + let agent = state.agents.get(agent_id).ok_or_else(|| { + Error::InvalidState(format!( + "No agent found with id: {agent_id} (it was never spawned)" + )) + })?; + if let Some(result) = agent.results.get(&generation) { + return result.clone(); } - SubAgentStatus::Finished(result) => return result, + let current = agent.status.borrow().clone(); + current + }; + match current { SubAgentStatus::Closing | SubAgentStatus::Closed => { return Err(Error::InvalidState(format!( "Agent {agent_id} has been closed" ))); } + SubAgentStatus::Running | SubAgentStatus::Finished(_) => {} + } + + tokio::select! { + biased; + () = cancel.cancelled() => { + self.ensure_closed(agent_id).await?; + return Err(Error::Interrupted(InterruptReason::Cancelled)); + } + changed = status.changed() => { + changed.map_err(|_| { + Error::InvalidState(format!( + "Agent {agent_id} result observer closed unexpectedly" + )) + })?; + } } } } @@ -505,8 +799,12 @@ impl SubAgentSupervisor { state .agents .get_mut(agent_id) - .and_then(|agent| agent.parent_notification.take()) - .is_some() + .and_then(|agent| agent.parent_notification.as_mut()) + .is_some_and(|notification| { + let cleared = !notification.pending_generations.is_empty(); + notification.pending_generations.clear(); + cleared + }) }; if cleared { signal_notifications(&self.notifications_changed); @@ -542,38 +840,49 @@ impl SubAgentSupervisor { let mut ready = Vec::new(); let mut awaiting_result = false; for (agent_id, agent) in &state.agents { - let Some(description) = agent.parent_notification.as_ref() else { + let Some(notification) = agent.parent_notification.as_ref() else { continue; }; - let finished = match &*agent.status.borrow() { - SubAgentStatus::Finished(result) => Some(result.clone()), - SubAgentStatus::Running => { + for generation in ¬ification.pending_generations { + if let Some(result) = agent.results.get(generation) { + ready.push(( + agent.spawn_seq, + *generation, + SubAgentParentNotification { + agent_id: agent_id.clone(), + description: notification.description.clone(), + result: result.clone(), + }, + )); + } else if agent.generation == *generation + && matches!(*agent.status.borrow(), SubAgentStatus::Running) + { awaiting_result = true; - None } - // Being torn down, so no result is coming. Ignoring - // these is what keeps a shutdown that races delivery - // from parking the parent forever. - SubAgentStatus::Closing | SubAgentStatus::Closed => None, - }; - if let Some(result) = finished { - ready.push((agent.spawn_seq, SubAgentParentNotification { - agent_id: agent_id.clone(), - description: description.clone(), - result, - })); } } if !ready.is_empty() { - ready.sort_by_key(|(spawn_seq, _)| *spawn_seq); + ready.sort_by_key(|(spawn_seq, generation, _)| (*spawn_seq, *generation)); + let delivered = ready + .iter() + .map(|(_, generation, notification)| { + (notification.agent_id.clone(), *generation) + }) + .collect::>(); let batch: Vec<_> = ready .into_iter() - .map(|(_, notification)| notification) + .map(|(_, _, notification)| notification) .collect(); - for notification in &batch { - if let Some(agent) = state.agents.get_mut(¬ification.agent_id) { - agent.parent_notification = None; + for (agent_id, generation) in delivered { + if let Some(notification) = state + .agents + .get_mut(&agent_id) + .and_then(|agent| agent.parent_notification.as_mut()) + { + notification + .pending_generations + .retain(|pending| *pending != generation); } } return Ok(Some(batch)); @@ -613,52 +922,39 @@ impl SubAgentSupervisor { )) })?; - let close_running_agent = loop { - let current = agent.status.borrow().clone(); - match current { - SubAgentStatus::Running => { - if agent.status.send_if_modified(|status| { - if matches!(status, SubAgentStatus::Running) { - *status = SubAgentStatus::Closing; - true - } else { - false - } - }) { - break true; - } - } - SubAgentStatus::Finished(_) if strict => { - return Err(Error::InvalidState(format!( - "Agent {agent_id} is not running" - ))); - } - SubAgentStatus::Finished(_) => break false, - SubAgentStatus::Closing | SubAgentStatus::Closed if strict => { - return Err(Error::InvalidState(format!( - "Agent {agent_id} is already closed" - ))); - } - SubAgentStatus::Closing => { - return Ok(ShutdownDisposition::Follow(agent.cleanup_done.subscribe())); - } - SubAgentStatus::Closed => return Ok(ShutdownDisposition::Done), + let close_running_agent = match agent.status.borrow().clone() { + SubAgentStatus::Running => true, + SubAgentStatus::Finished(_) => false, + SubAgentStatus::Closing | SubAgentStatus::Closed if strict => { + return Err(Error::InvalidState(format!( + "Agent {agent_id} is already closed" + ))); } + SubAgentStatus::Closing => { + return Ok(ShutdownDisposition::Follow(agent.cleanup_done.subscribe())); + } + SubAgentStatus::Closed => return Ok(ShutdownDisposition::Done), }; + agent.status.send_replace(SubAgentStatus::Closing); - // Shutdown is committed, so this child's result will never reach the - // parent. The early returns above leave the notification intact, so a - // rejected shutdown cannot discard a result the parent is owed. + // Shutdown is committed, so no pending result will reach the parent. agent.parent_notification = None; if agent.cleanup_started { - return Ok(ShutdownDisposition::Follow(agent.cleanup_done.subscribe())); + return if strict { + Err(Error::InvalidState(format!( + "Agent {agent_id} is already closed" + ))) + } else { + Ok(ShutdownDisposition::Follow(agent.cleanup_done.subscribe())) + }; } agent.cleanup_started = true; Ok(ShutdownDisposition::Lead(ShutdownWork { agent_id: agent_id.to_string(), depth: agent.depth, + generation: agent.generation, close_running_agent, status: agent.status.clone(), cleanup_done: agent.cleanup_done.clone(), @@ -666,6 +962,8 @@ impl SubAgentSupervisor { event_forwarder: agent.event_forwarder.take(), child_abort_handle: agent.child_abort_handle.clone(), cancel_token: agent.cancel_token.clone(), + runner_stop: agent.runner_stop.clone(), + state: Arc::downgrade(&self.state), })) } @@ -675,6 +973,7 @@ impl SubAgentSupervisor { ) { let _cleanup_done = CleanupDoneGuard(work.cleanup_done.clone()); let deadline = Instant::now() + SUBAGENT_SHUTDOWN_GRACE; + work.runner_stop.cancel(); if work.close_running_agent { work.cancel_token.cancel(); } @@ -693,27 +992,25 @@ impl SubAgentSupervisor { } } - let emit_closed = work.close_running_agent - && work.status.send_if_modified(|status| { - if matches!(status, SubAgentStatus::Closing) { - *status = SubAgentStatus::Closed; - true - } else { - false - } - }); + let emit_closed = work.status.send_if_modified(|status| { + if matches!(status, SubAgentStatus::Closing) { + *status = SubAgentStatus::Closed; + true + } else { + false + } + }); if emit_closed { - let callback = event_callback - .read() - .expect("subagent callback lock poisoned") - .clone(); - if let Some(callback) = callback { - callback(SubAgentCallbackEvent::Lifecycle( - AgentEvent::SubAgentClosed { - agent_id: work.agent_id.clone(), - depth: work.depth, - }, - )); + if let Some(state) = work.state.upgrade() { + { + let mut state = state.lock().expect("subagent state lock poisoned"); + queue_lifecycle_event(&mut state, AgentEvent::SubAgentClosed { + agent_id: work.agent_id.clone(), + depth: work.depth, + generation: work.generation, + }); + } + drain_lifecycle_events(&state, &event_callback); } } } @@ -767,7 +1064,7 @@ impl SubAgentSupervisor { Ok(()) } - /// Strict user-facing close: only a currently running child may be closed. + /// Close a running or idle child that is no longer needed. pub async fn close_agent(&self, agent_id: &str) -> Result<(), Error> { let disposition = self.begin_shutdown(agent_id, true)?; signal_notifications(&self.notifications_changed); @@ -783,9 +1080,8 @@ impl SubAgentSupervisor { Ok(()) } - /// Cooperatively shut down active children and join every owned child and - /// event-forwarding task. Finished children are reaped without rewriting - /// their terminal result. + /// Cooperatively shut down all children and join every owned runner and + /// event-forwarding task. pub async fn shutdown_all(&self) { let ids = { let state = self.state.lock().expect("subagent state lock poisoned"); @@ -834,21 +1130,46 @@ impl SubAgentSupervisor { let child_abort_handle = child_task.abort_handle(); let (status, _) = watch::channel(SubAgentStatus::Running); let (cleanup_done, _) = watch::channel(false); + let (command_tx, command_rx) = mpsc::channel(SUBAGENT_COMMAND_CAPACITY); + drop(command_rx); + let runner_stop = CancellationToken::new(); let depth = 1; - let monitor_task = spawn_result_monitor( - child_task, - status.clone(), - Arc::clone(&self.event_callback), - Arc::clone(&self.notifications_changed), - agent_id.clone(), - depth, - ); - self.state - .lock() - .expect("subagent state lock poisoned") - .agents - .insert(agent_id, SubAgent { + let (monitor_start_tx, monitor_start_rx) = oneshot::channel(); + let state = Arc::downgrade(&self.state); + let event_callback = Arc::clone(&self.event_callback); + let notifications_changed = Arc::clone(&self.notifications_changed); + let monitored_agent_id = agent_id.clone(); + let monitor_task = tokio::spawn(async move { + let _ = monitor_start_rx.await; + let task_result = match child_task.await { + Ok(result) => result, + Err(error) => Err(Error::InvalidState(format!( + "Agent task failed to join: {error}" + ))), + }; + let Some(state) = state.upgrade() else { + return; + }; + commit_turn_result( + &state, + &event_callback, + ¬ifications_changed, + &monitored_agent_id, + depth, + INITIAL_SUBAGENT_GENERATION, + &task_result, + false, + ); + }); + { + let mut state = self.state.lock().expect("subagent state lock poisoned"); + state.agents.insert(agent_id, SubAgent { status, + generation: INITIAL_SUBAGENT_GENERATION, + results: HashMap::new(), + reusable: false, + command_tx, + runner_stop, cleanup_done, cleanup_started: false, monitor_task: Some(monitor_task), @@ -861,6 +1182,8 @@ impl SubAgentSupervisor { cancel_token, depth, }); + } + let _ = monitor_start_tx.send(()); } } @@ -910,7 +1233,7 @@ pub fn make_send_input_tool(supervisor: SubAgentSupervisor) -> RegisteredTool { RegisteredTool { definition: ToolDefinition { name: "send_input".into(), - description: "Send a follow-up message to a running subagent when new information or corrected instructions are needed.".into(), + description: "Send a follow-up message to a subagent. A running agent receives it at a safe turn boundary. A completed agent starts another turn in the same session with its existing history.".into(), parameters: serde_json::json!({ "type": "object", "properties": { @@ -983,7 +1306,7 @@ pub fn make_close_agent_tool(supervisor: SubAgentSupervisor) -> RegisteredTool { RegisteredTool { definition: ToolDefinition { name: "close_agent".into(), - description: "Close a running subagent that is no longer needed.".into(), + description: "Close a running or completed subagent that is no longer needed.".into(), parameters: serde_json::json!({ "type": "object", "properties": { @@ -1039,6 +1362,8 @@ mod tests { assert!(spawn.definition.description.contains("independent work")); assert!(spawn.definition.description.contains("context isolation")); assert!(send.definition.description.contains("follow-up")); + assert!(send.definition.description.contains("completed agent")); + assert!(send.definition.description.contains("same session")); assert!(wait.definition.description.contains("synthesize")); assert!(close.definition.description.contains("no longer needed")); @@ -1115,6 +1440,56 @@ mod tests { supervisor.shutdown_all().await; } + #[tokio::test] + async fn a_reused_agent_delivers_each_generation_to_the_parent() { + let supervisor = SubAgentSupervisor::new(3); + let child = make_session(vec![ + text_response("first result"), + text_response("remediation result"), + ]) + .await; + let agent_id = supervisor + .spawn_with_parent_notification( + child, + "implement".to_string(), + "Implement the work".to_string(), + 0, + ) + .unwrap(); + + assert_eq!( + supervisor.wait(&agent_id).await.unwrap().output, + "first result" + ); + let first = supervisor + .next_parent_notification_batch(&CancellationToken::new()) + .await + .unwrap() + .expect("generation one should be delivered"); + assert_eq!(first[0].result.as_ref().unwrap().output, "first result"); + + supervisor + .send_input(&agent_id, "Fix the review findings") + .unwrap(); + assert_eq!( + supervisor.wait(&agent_id).await.unwrap().output, + "remediation result" + ); + let second = supervisor + .next_parent_notification_batch(&CancellationToken::new()) + .await + .unwrap() + .expect("generation two should be delivered"); + assert_eq!(second.len(), 1); + assert_eq!( + second[0].result.as_ref().unwrap().output, + "remediation result" + ); + assert_eq!(second[0].description, "Implement the work"); + + supervisor.shutdown_all().await; + } + #[tokio::test] async fn batches_are_delivered_in_spawn_order() { let supervisor = SubAgentSupervisor::new(3); @@ -1208,7 +1583,7 @@ mod tests { } #[tokio::test] - async fn rejected_stop_of_a_finished_agent_keeps_its_notification() { + async fn closing_a_finished_agent_discards_its_pending_notification() { let supervisor = SubAgentSupervisor::new(3); let child = make_session(vec![text_response("child result")]).await; let agent_id = supervisor @@ -1226,18 +1601,19 @@ mod tests { .await .unwrap(); - // Stopping a finished agent is rejected... - let error = supervisor.close_agent(&agent_id).await.unwrap_err(); - assert!(matches!(error, Error::InvalidState(_)), "{error:?}"); + supervisor.close_agent(&agent_id).await.unwrap(); - // ...so it must not have discarded the result the parent is owed. - let batch = supervisor - .next_parent_notification_batch(&CancellationToken::new()) - .await - .unwrap() - .expect("a rejected stop must leave the pending result deliverable"); - assert_eq!(batch.len(), 1); - assert_eq!(batch[0].agent_id, agent_id); + assert!( + supervisor + .next_parent_notification_batch(&CancellationToken::new()) + .await + .unwrap() + .is_none() + ); + assert!(matches!( + supervisor.status(&agent_id), + Some(SubAgentStatus::Closed) + )); supervisor.shutdown_all().await; } @@ -1577,8 +1953,9 @@ mod tests { // Manager without callback should not panic on emit let manager = SubAgentSupervisor::new(3); manager.emit_event(AgentEvent::SubAgentClosed { - agent_id: "x".into(), - depth: 0, + agent_id: "x".into(), + depth: 0, + generation: 1, }); } @@ -1638,15 +2015,112 @@ mod tests { } #[tokio::test] - async fn send_input_to_completed_agent_errors() { + async fn send_input_to_running_agent_joins_the_current_generation() { + let (callback, events) = captured_events(); let manager = SubAgentSupervisor::new(3); - let session = make_session(vec![text_response("done")]).await; + manager.set_event_callback(callback); + let session = make_session(vec![ + text_response("initial"), + text_response("after follow-up"), + ]) + .await; let agent_id = manager.spawn(session, "Do something".into(), 0).unwrap(); - let _ = manager.wait(&agent_id).await.unwrap(); - let result = manager.send_input(&agent_id, "hello"); - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("is not running")); + manager + .send_input(&agent_id, "Use this additional information") + .unwrap(); + let result = manager.wait(&agent_id).await.unwrap(); + + assert_eq!(result.output, "after follow-up"); + { + let events = events.lock().unwrap(); + assert!(!events.iter().any(|event| matches!( + event, + SubAgentCallbackEvent::Lifecycle(AgentEvent::SubAgentTurnStarted { .. }) + ))); + assert!(events.iter().any(|event| matches!( + event, + SubAgentCallbackEvent::Lifecycle(AgentEvent::SubAgentCompleted { + generation: 1, + .. + }) + ))); + } + + manager.shutdown_all().await; + } + + #[tokio::test] + async fn send_input_to_completed_agent_reuses_its_session_and_history() { + let (callback, events) = captured_events(); + let manager = SubAgentSupervisor::new(3); + manager.set_event_callback(callback); + let provider = Arc::new(CapturingLlmProvider::new()); + let provider_ref = provider.clone(); + let client = make_client(provider as Arc).await; + let profile = Arc::new(TestProfile::new()); + let env = Arc::new(MockSandbox::default()); + let session = Session::new(client, profile, env, SessionOptions::default(), None); + let agent_id = manager.spawn(session, "Do something".into(), 0).unwrap(); + let first = manager.wait(&agent_id).await.unwrap(); + assert_eq!(first.output, "captured"); + + manager + .send_input(&agent_id, "Fix the review findings") + .unwrap(); + let second = manager.wait(&agent_id).await.unwrap(); + assert_eq!(second.output, "captured"); + assert_eq!(second.turns_used, first.turns_used); + + { + let captured = provider_ref.captured_request.lock().unwrap(); + let request = captured + .as_ref() + .expect("second request should be captured"); + assert!(request.messages.iter().any(|message| { + message.role == Role::User && message.text().contains("Do something") + })); + assert!(request.messages.iter().any(|message| { + message.role == Role::Assistant && message.text().contains("captured") + })); + assert!(request.messages.iter().any(|message| { + message.role == Role::User && message.text().contains("Fix the review findings") + })); + } + + { + let events = events.lock().unwrap(); + let spawn_count = events + .iter() + .filter(|event| { + matches!( + event, + SubAgentCallbackEvent::Lifecycle(AgentEvent::SubAgentSpawned { .. }) + ) + }) + .count(); + assert_eq!(spawn_count, 1); + assert!(events.iter().any(|event| matches!( + event, + SubAgentCallbackEvent::Lifecycle(AgentEvent::SubAgentTurnStarted { + generation: 2, + .. + }) + ))); + let completed_generations = events + .iter() + .filter_map(|event| match event { + SubAgentCallbackEvent::Lifecycle(AgentEvent::SubAgentCompleted { + generation, + .. + }) => Some(*generation), + _ => None, + }) + .collect::>(); + assert_eq!(completed_generations, vec![1, 2]); + } + + manager.shutdown_all().await; } #[tokio::test] @@ -1658,7 +2132,7 @@ mod tests { let result = manager.send_input(&agent_id, "hello"); assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("is not running")); + assert!(result.unwrap_err().to_string().contains("has been closed")); } #[tokio::test] @@ -1674,7 +2148,7 @@ mod tests { } #[tokio::test] - async fn close_completed_agent_preserves_finished_result() { + async fn close_completed_agent_closes_its_idle_session() { let manager = SubAgentSupervisor::new(3); let session = make_session(vec![text_response("done")]).await; let agent_id = manager.spawn(session, "Do something".into(), 0).unwrap(); @@ -1684,11 +2158,10 @@ mod tests { Some(SubAgentStatus::Finished(Ok(_))) )); - let result = manager.close_agent(&agent_id).await; - assert!(result.is_err()); + manager.close_agent(&agent_id).await.unwrap(); assert!(matches!( manager.status(&agent_id), - Some(SubAgentStatus::Finished(Ok(_))) + Some(SubAgentStatus::Closed) )); } @@ -1779,6 +2252,59 @@ mod tests { assert_eq!(completion_count, 1); } + #[tokio::test] + async fn wait_returns_its_target_generation_after_a_later_turn_starts() { + let supervisor = SubAgentSupervisor::new(3); + let child_cancel = CancellationToken::new(); + let task_cancel = child_cancel.clone(); + let child = tokio::spawn(async move { + task_cancel.cancelled().await; + Ok(SubAgentResult { + output: "unused".to_string(), + success: true, + turns_used: 1, + }) + }); + let agent_id = "generation-aware-wait".to_string(); + supervisor.supervise_test_task(agent_id.clone(), child, child_cancel, None); + + let wait_cancel = CancellationToken::new(); + let mut wait = Box::pin(supervisor.wait_with_cancel(&agent_id, &wait_cancel)); + assert!( + futures::poll!(wait.as_mut()).is_pending(), + "generation one should still be running" + ); + + { + let mut state = supervisor + .state + .lock() + .expect("subagent state lock poisoned"); + let agent = state.agents.get_mut(&agent_id).unwrap(); + let first_result = Ok(SubAgentResult { + output: "generation one".to_string(), + success: true, + turns_used: 1, + }); + agent + .results + .insert(INITIAL_SUBAGENT_GENERATION, first_result.clone()); + agent + .status + .send_replace(SubAgentStatus::Finished(first_result)); + agent.generation = INITIAL_SUBAGENT_GENERATION + 1; + agent.status.send_replace(SubAgentStatus::Running); + } + + let result = time::timeout(Duration::from_secs(1), wait) + .await + .expect("the coalesced status updates must not hide generation one") + .unwrap(); + assert_eq!(result.output, "generation one"); + + supervisor.close_agent(&agent_id).await.unwrap(); + } + #[tokio::test] async fn uncooperative_child_and_forwarder_are_aborted_after_grace() { struct DropProbe(Arc); diff --git a/lib/components/fabro-agent/src/types.rs b/lib/components/fabro-agent/src/types.rs index c5bd99deb..38559b180 100644 --- a/lib/components/fabro-agent/src/types.rs +++ b/lib/components/fabro-agent/src/types.rs @@ -224,6 +224,10 @@ pub struct McpToolSummary { pub original_name: String, } +const fn initial_subagent_generation() -> u64 { + 1 +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub enum AgentEvent { SessionStarted { @@ -352,24 +356,38 @@ pub enum AgentEvent { phase: LlmRetryPhase, }, SubAgentSpawned { - agent_id: String, - depth: usize, - task: String, + agent_id: String, + depth: usize, + task: String, + #[serde(default = "initial_subagent_generation")] + generation: u64, + }, + SubAgentTurnStarted { + agent_id: String, + depth: usize, + task: String, + generation: u64, }, SubAgentCompleted { agent_id: String, depth: usize, + #[serde(default = "initial_subagent_generation")] + generation: u64, success: bool, turns_used: usize, }, SubAgentFailed { - agent_id: String, - depth: usize, - error: Error, + agent_id: String, + depth: usize, + #[serde(default = "initial_subagent_generation")] + generation: u64, + error: Error, }, SubAgentClosed { - agent_id: String, - depth: usize, + agent_id: String, + depth: usize, + #[serde(default = "initial_subagent_generation")] + generation: u64, }, McpServerReady { server_name: String, @@ -586,35 +604,57 @@ impl AgentEvent { agent_id, depth, task, + generation, } => { - debug!(session_id, agent_id, depth, task, "Sub-agent spawned"); + debug!( + session_id, + agent_id, depth, generation, task, "Sub-agent spawned" + ); + } + Self::SubAgentTurnStarted { + agent_id, + depth, + task, + generation, + } => { + debug!( + session_id, + agent_id, depth, generation, task, "Sub-agent turn started" + ); } Self::SubAgentCompleted { agent_id, depth, + generation, success, turns_used, } => { debug!( session_id, - agent_id, depth, success, turns_used, "Sub-agent completed" + agent_id, depth, generation, success, turns_used, "Sub-agent completed" ); } Self::SubAgentFailed { agent_id, depth, + generation, error, } => { warn!( session_id, agent_id, depth, + generation, error = %error, "Sub-agent failed" ); } - Self::SubAgentClosed { agent_id, depth } => { - debug!(session_id, agent_id, depth, "Sub-agent closed"); + Self::SubAgentClosed { + agent_id, + depth, + generation, + } => { + debug!(session_id, agent_id, depth, generation, "Sub-agent closed"); } Self::McpServerReady { server_name, @@ -765,9 +805,10 @@ mod tests { #[test] fn subagent_spawned_constructible() { let event = AgentEvent::SubAgentSpawned { - agent_id: "sa-1".into(), - depth: 1, - task: "list files".into(), + agent_id: "sa-1".into(), + depth: 1, + task: "list files".into(), + generation: 1, }; assert!(matches!(event, AgentEvent::SubAgentSpawned { depth: 1, @@ -780,6 +821,7 @@ mod tests { let event = AgentEvent::SubAgentCompleted { agent_id: "sa-1".into(), depth: 1, + generation: 1, success: true, turns_used: 5, }; @@ -793,9 +835,10 @@ mod tests { #[test] fn subagent_failed_constructible() { let event = AgentEvent::SubAgentFailed { - agent_id: "sa-1".into(), - depth: 0, - error: Error::ToolExecution("timeout".into()), + agent_id: "sa-1".into(), + depth: 0, + generation: 1, + error: Error::ToolExecution("timeout".into()), }; assert!(matches!(event, AgentEvent::SubAgentFailed { depth: 0, .. })); } @@ -803,8 +846,9 @@ mod tests { #[test] fn subagent_closed_constructible() { let event = AgentEvent::SubAgentClosed { - agent_id: "sa-1".into(), - depth: 2, + agent_id: "sa-1".into(), + depth: 2, + generation: 1, }; assert!(matches!(event, AgentEvent::SubAgentClosed { depth: 2, .. })); } @@ -813,29 +857,52 @@ mod tests { fn subagent_events_serde_round_trip() { let events = vec![ AgentEvent::SubAgentSpawned { - agent_id: "sa-1".into(), - depth: 0, - task: "test".into(), + agent_id: "sa-1".into(), + depth: 0, + task: "test".into(), + generation: 1, + }, + AgentEvent::SubAgentTurnStarted { + agent_id: "sa-1".into(), + depth: 0, + task: "fix it".into(), + generation: 2, }, AgentEvent::SubAgentCompleted { agent_id: "sa-1".into(), depth: 0, + generation: 2, success: true, turns_used: 3, }, AgentEvent::SubAgentFailed { - agent_id: "sa-1".into(), - depth: 0, - error: Error::ToolExecution("oops".into()), + agent_id: "sa-1".into(), + depth: 0, + generation: 2, + error: Error::ToolExecution("oops".into()), }, AgentEvent::SubAgentClosed { - agent_id: "sa-1".into(), - depth: 0, + agent_id: "sa-1".into(), + depth: 0, + generation: 2, }, ]; let json = serde_json::to_string(&events).unwrap(); let deserialized: Vec = serde_json::from_str(&json).unwrap(); - assert_eq!(deserialized.len(), 4); + assert_eq!(deserialized.len(), 5); + } + + #[test] + fn legacy_subagent_event_defaults_to_the_initial_generation() { + let event: AgentEvent = serde_json::from_str( + r#"{"SubAgentSpawned":{"agent_id":"sa-1","depth":0,"task":"test"}}"#, + ) + .unwrap(); + + assert!(matches!(event, AgentEvent::SubAgentSpawned { + generation: 1, + .. + })); } #[test] @@ -1054,9 +1121,10 @@ mod tests { #[test] fn subagent_failed_carries_agent_error() { let event = AgentEvent::SubAgentFailed { - agent_id: "sa-1".into(), - depth: 0, - error: Error::ToolExecution("cmd failed".into()), + agent_id: "sa-1".into(), + depth: 0, + generation: 1, + error: Error::ToolExecution("cmd failed".into()), }; let json = serde_json::to_string(&event).unwrap(); let deserialized: AgentEvent = serde_json::from_str(&json).unwrap(); diff --git a/lib/components/fabro-store/src/run_state.rs b/lib/components/fabro-store/src/run_state.rs index 05945e33f..d3a8d5aae 100644 --- a/lib/components/fabro-store/src/run_state.rs +++ b/lib/components/fabro-store/src/run_state.rs @@ -689,6 +689,15 @@ impl RunProjectionReducer for RunProjection { status: SubAgentStatus::Running, }); } + EventBody::AgentSubTurnStarted(props) => { + let Some(stage) = stage_at_stored_or_visit(self, stored, props.visit, event.seq) + else { + return Ok(()); + }; + if let Some(subagent) = subagent_mut(stage, &props.agent_id) { + subagent.status = SubAgentStatus::Running; + } + } EventBody::AgentSubCompleted(props) => { let Some(stage) = stage_at_stored_or_visit(self, stored, props.visit, event.seq) else { @@ -1611,9 +1620,10 @@ mod tests { AgentSessionDeactivatedProps, AgentSessionEndedProps, AgentSessionStartedProps, AgentSkillActivatedProps, AgentSkillActivationSource, AgentSkillSummary, AgentSkillsDiscoveredProps, AgentSteeringInjectedProps, AgentSubClosedProps, - AgentSubCompletedProps, AgentSubFailedProps, AgentSubSpawnedProps, AgentToolCategory, - AgentToolSource, AgentToolStartedProps, AgentToolSummary, AgentToolsAvailableProps, - CheckpointCompletedProps, InterviewCompletedProps, InterviewOption, InterviewStartedProps, + AgentSubCompletedProps, AgentSubFailedProps, AgentSubSpawnedProps, + AgentSubTurnStartedProps, AgentToolCategory, AgentToolSource, AgentToolStartedProps, + AgentToolSummary, AgentToolsAvailableProps, CheckpointCompletedProps, + InterviewCompletedProps, InterviewOption, InterviewStartedProps, ParallelBranchCompletedProps, ParallelBranchStartedProps, RunCompletedProps, RunControlEffectProps, StageCompletedProps, StageFailedProps, StagePromptProps, StageRetryingProps, StageStartedProps, @@ -6434,10 +6444,11 @@ mod tests { .apply_event(&test_stage_event( 1, EventBody::AgentSubSpawned(AgentSubSpawnedProps { - agent_id: "sub-1".to_string(), - depth: 1, - task: "write tests".to_string(), - visit: 1, + agent_id: "sub-1".to_string(), + depth: 1, + task: "write tests".to_string(), + generation: 1, + visit: 1, }), stage_id.clone(), )) @@ -6455,6 +6466,7 @@ mod tests { EventBody::AgentSubCompleted(AgentSubCompletedProps { agent_id: "sub-1".to_string(), depth: 1, + generation: 1, success: true, turns_used: 3, visit: 1, @@ -6471,23 +6483,64 @@ mod tests { state .apply_event(&test_stage_event( 3, + EventBody::AgentSubTurnStarted(AgentSubTurnStartedProps { + agent_id: "sub-1".to_string(), + depth: 1, + task: "fix the review findings".to_string(), + generation: 2, + visit: 1, + }), + stage_id.clone(), + )) + .unwrap(); + let stage = state.stage(&stage_id).unwrap(); + assert_eq!(stage.subagents.len(), 1); + assert_eq!(stage.subagents[0].task, "write tests"); + assert_eq!(stage.subagents[0].status, SubAgentStatus::Running); + + state + .apply_event(&test_stage_event( + 4, + EventBody::AgentSubCompleted(AgentSubCompletedProps { + agent_id: "sub-1".to_string(), + depth: 1, + generation: 2, + success: true, + turns_used: 5, + visit: 1, + }), + stage_id.clone(), + )) + .unwrap(); + let stage = state.stage(&stage_id).unwrap(); + assert_eq!(stage.subagents.len(), 1); + assert_eq!(stage.subagents[0].status, SubAgentStatus::Completed { + success: true, + turns_used: 5, + }); + + state + .apply_event(&test_stage_event( + 5, EventBody::AgentSubSpawned(AgentSubSpawnedProps { - agent_id: "sub-2".to_string(), - depth: 2, - task: "debug failure".to_string(), - visit: 1, + agent_id: "sub-2".to_string(), + depth: 2, + task: "debug failure".to_string(), + generation: 1, + visit: 1, }), stage_id.clone(), )) .unwrap(); state .apply_event(&test_stage_event( - 4, + 6, EventBody::AgentSubFailed(AgentSubFailedProps { - agent_id: "sub-2".to_string(), - depth: 2, - error: json!({ "message": "boom" }), - visit: 1, + agent_id: "sub-2".to_string(), + depth: 2, + generation: 1, + error: json!({ "message": "boom" }), + visit: 1, }), stage_id.clone(), )) @@ -6499,11 +6552,12 @@ mod tests { state .apply_event(&test_stage_event( - 5, + 7, EventBody::AgentSubClosed(AgentSubClosedProps { - agent_id: "sub-2".to_string(), - depth: 2, - visit: 1, + agent_id: "sub-2".to_string(), + depth: 2, + generation: 1, + visit: 1, }), stage_id.clone(), )) diff --git a/lib/components/fabro-workflow/src/event/convert.rs b/lib/components/fabro-workflow/src/event/convert.rs index a6275bfad..3a99f6b64 100644 --- a/lib/components/fabro-workflow/src/event/convert.rs +++ b/lib/components/fabro-workflow/src/event/convert.rs @@ -758,20 +758,36 @@ fn event_body_from_event(event: &Event) -> EventBody { agent_id, depth, task, + generation, } => EventBody::AgentSubSpawned(fabro_types::AgentSubSpawnedProps { - agent_id: agent_id.clone(), - depth: *depth, - task: task.clone(), - visit: *visit, + agent_id: agent_id.clone(), + depth: *depth, + task: task.clone(), + generation: *generation, + visit: *visit, + }), + AgentEvent::SubAgentTurnStarted { + agent_id, + depth, + task, + generation, + } => EventBody::AgentSubTurnStarted(fabro_types::AgentSubTurnStartedProps { + agent_id: agent_id.clone(), + depth: *depth, + task: task.clone(), + generation: *generation, + visit: *visit, }), AgentEvent::SubAgentCompleted { agent_id, depth, + generation, success, turns_used, } => EventBody::AgentSubCompleted(fabro_types::AgentSubCompletedProps { agent_id: agent_id.clone(), depth: *depth, + generation: *generation, success: *success, turns_used: *turns_used, visit: *visit, @@ -779,18 +795,25 @@ fn event_body_from_event(event: &Event) -> EventBody { AgentEvent::SubAgentFailed { agent_id, depth, + generation, error, } => EventBody::AgentSubFailed(fabro_types::AgentSubFailedProps { - agent_id: agent_id.clone(), - depth: *depth, - error: serde_json::to_value(error).expect("agent Error derives Serialize with no custom logic that can fail"), - visit: *visit, + agent_id: agent_id.clone(), + depth: *depth, + generation: *generation, + error: serde_json::to_value(error).expect("agent Error derives Serialize with no custom logic that can fail"), + visit: *visit, }), - AgentEvent::SubAgentClosed { agent_id, depth } => { + AgentEvent::SubAgentClosed { + agent_id, + depth, + generation, + } => { EventBody::AgentSubClosed(fabro_types::AgentSubClosedProps { - agent_id: agent_id.clone(), - depth: *depth, - visit: *visit, + agent_id: agent_id.clone(), + depth: *depth, + generation: *generation, + visit: *visit, }) } AgentEvent::McpServerReady { diff --git a/lib/components/fabro-workflow/src/event/names.rs b/lib/components/fabro-workflow/src/event/names.rs index d76b50e9a..ee051c602 100644 --- a/lib/components/fabro-workflow/src/event/names.rs +++ b/lib/components/fabro-workflow/src/event/names.rs @@ -86,6 +86,7 @@ pub fn event_name(event: &Event) -> &'static str { AgentEvent::CompactionCompleted { .. } => "agent.compaction.completed", AgentEvent::LlmRetry { .. } => "agent.llm.retry", AgentEvent::SubAgentSpawned { .. } => "agent.sub.spawned", + AgentEvent::SubAgentTurnStarted { .. } => "agent.sub.turn.started", AgentEvent::SubAgentCompleted { .. } => "agent.sub.completed", AgentEvent::SubAgentFailed { .. } => "agent.sub.failed", AgentEvent::SubAgentClosed { .. } => "agent.sub.closed", @@ -184,9 +185,10 @@ mod tests { stage: "code".to_string(), visit: 1, event: AgentEvent::SubAgentSpawned { - agent_id: "a1".to_string(), - depth: 1, - task: "do it".to_string(), + agent_id: "a1".to_string(), + depth: 1, + task: "do it".to_string(), + generation: 1, }, session_id: None, parent_session_id: None, @@ -194,6 +196,22 @@ mod tests { }), "agent.sub.spawned" ); + assert_eq!( + event_name(&Event::Agent { + stage: "code".to_string(), + visit: 1, + event: AgentEvent::SubAgentTurnStarted { + agent_id: "a1".to_string(), + depth: 1, + task: "fix it".to_string(), + generation: 2, + }, + session_id: None, + parent_session_id: None, + tool_call_id: None, + }), + "agent.sub.turn.started" + ); assert_eq!( event_name(&Event::Agent { stage: "code".to_string(), diff --git a/lib/components/fabro-workflow/src/handler/llm/api.rs b/lib/components/fabro-workflow/src/handler/llm/api.rs index 5dad09d34..519f5d263 100644 --- a/lib/components/fabro-workflow/src/handler/llm/api.rs +++ b/lib/components/fabro-workflow/src/handler/llm/api.rs @@ -3990,8 +3990,9 @@ enabled = true session.sub_agent_event_callback()( fabro_agent::subagent::SubAgentCallbackEvent::Lifecycle(AgentEvent::SubAgentClosed { - agent_id: "child-1".to_string(), - depth: 1, + agent_id: "child-1".to_string(), + depth: 1, + generation: 1, }), ); let session_id = session.id().to_string(); diff --git a/lib/foundation/fabro-types/src/run_event/agent.rs b/lib/foundation/fabro-types/src/run_event/agent.rs index 8b27a3be1..dd70b8561 100644 --- a/lib/foundation/fabro-types/src/run_event/agent.rs +++ b/lib/foundation/fabro-types/src/run_event/agent.rs @@ -378,16 +378,29 @@ pub struct AgentLlmFirstOutputProps { #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct AgentSubSpawnedProps { - pub agent_id: String, - pub depth: usize, - pub task: String, - pub visit: u32, + pub agent_id: String, + pub depth: usize, + pub task: String, + #[serde(default = "initial_subagent_generation")] + pub generation: u64, + pub visit: u32, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct AgentSubTurnStartedProps { + pub agent_id: String, + pub depth: usize, + pub task: String, + pub generation: u64, + pub visit: u32, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct AgentSubCompletedProps { pub agent_id: String, pub depth: usize, + #[serde(default = "initial_subagent_generation")] + pub generation: u64, pub success: bool, pub turns_used: usize, pub visit: u32, @@ -395,17 +408,25 @@ pub struct AgentSubCompletedProps { #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct AgentSubFailedProps { - pub agent_id: String, - pub depth: usize, - pub error: Value, - pub visit: u32, + pub agent_id: String, + pub depth: usize, + #[serde(default = "initial_subagent_generation")] + pub generation: u64, + pub error: Value, + pub visit: u32, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct AgentSubClosedProps { - pub agent_id: String, - pub depth: usize, - pub visit: u32, + pub agent_id: String, + pub depth: usize, + #[serde(default = "initial_subagent_generation")] + pub generation: u64, + pub visit: u32, +} + +const fn initial_subagent_generation() -> u64 { + 1 } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] diff --git a/lib/foundation/fabro-types/src/run_event/mod.rs b/lib/foundation/fabro-types/src/run_event/mod.rs index 27c3b87e9..59892030f 100644 --- a/lib/foundation/fabro-types/src/run_event/mod.rs +++ b/lib/foundation/fabro-types/src/run_event/mod.rs @@ -242,6 +242,8 @@ pub enum EventBody { AgentLlmRetry(AgentLlmRetryProps), #[serde(rename = "agent.sub.spawned")] AgentSubSpawned(AgentSubSpawnedProps), + #[serde(rename = "agent.sub.turn.started")] + AgentSubTurnStarted(AgentSubTurnStartedProps), #[serde(rename = "agent.sub.completed")] AgentSubCompleted(AgentSubCompletedProps), #[serde(rename = "agent.sub.failed")] @@ -510,6 +512,7 @@ impl EventBody { Self::AgentLlmFirstOutput(_) => "agent.llm.first_output", Self::AgentLlmRetry(_) => "agent.llm.retry", Self::AgentSubSpawned(_) => "agent.sub.spawned", + Self::AgentSubTurnStarted(_) => "agent.sub.turn.started", Self::AgentSubCompleted(_) => "agent.sub.completed", Self::AgentSubFailed(_) => "agent.sub.failed", Self::AgentSubClosed(_) => "agent.sub.closed", @@ -682,6 +685,7 @@ fn is_known_event_name(event: &str) -> bool { | "agent.llm.first_output" | "agent.llm.retry" | "agent.sub.spawned" + | "agent.sub.turn.started" | "agent.sub.completed" | "agent.sub.failed" | "agent.sub.closed" @@ -2497,6 +2501,37 @@ mod tests { assert_eq!(parsed, body); } + #[test] + fn subagent_generations_are_typed_and_legacy_events_default_to_one() { + let started = EventBody::AgentSubTurnStarted(AgentSubTurnStartedProps { + agent_id: "sub-1".to_string(), + depth: 1, + task: "fix the review findings".to_string(), + generation: 2, + visit: 1, + }); + let value = serde_json::to_value(&started).unwrap(); + assert_eq!(value["event"], "agent.sub.turn.started"); + assert_eq!(value["properties"]["generation"], 2); + assert_eq!(serde_json::from_value::(value).unwrap(), started); + + let legacy: EventBody = serde_json::from_value(json!({ + "event": "agent.sub.completed", + "properties": { + "agent_id": "sub-1", + "depth": 1, + "success": true, + "turns_used": 3, + "visit": 1 + } + })) + .unwrap(); + let EventBody::AgentSubCompleted(props) = legacy else { + panic!("expected subagent completion"); + }; + assert_eq!(props.generation, 1); + } + #[test] fn agent_tool_source_and_category_use_public_json_shape() { assert_eq!( From 1bb70adafd53c728856a2396975eb0b1368e72ae Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Sat, 1 Aug 2026 10:09:07 -0400 Subject: [PATCH 2/4] refactor: simplify subagent session reuse Review pass over the reuse change. No intended behavior changes. - share one definition of the initial generation from fabro-types instead of three copies across fabro-types, fabro-agent, and the supervisor - give each child one SubAgentHandle instead of threading the supervisor's state, callback, and notification sender through five functions, and collapse the repeated signal-then-drain pairs into publish() - move `reusable` inside SubAgentStatus::Finished so a closed agent can no longer be marked reusable - clear the lifecycle draining flag with an RAII guard, so one panicking callback cannot silence every later lifecycle event - tear down a session that failed to initialize right away rather than holding it and its sandbox until the parent closes the agent - look agents up through SupervisorState::agent/agent_mut instead of five copies of the same not-found error - drop the unreachable cleanup_started branch and the test-only emit_event whose only caller was its own test - render subagent starts from one ProgressEvent and one display method, deriving the spawn/turn distinction from the generation - set projected subagent status through one helper instead of four identical reducer arms - drive the generation-pinned wait test through spawn/send_input rather than hand-writing private supervisor state Co-Authored-By: Claude Opus 5 (1M context) --- .../src/commands/run/run_progress/event.rs | 14 +- .../src/commands/run/run_progress/mod.rs | 12 +- .../run/run_progress/stage_display.rs | 41 +- lib/components/fabro-agent/src/cli.rs | 22 +- .../fabro-agent/src/profiles/claude5_tools.rs | 2 +- lib/components/fabro-agent/src/subagent.rs | 552 ++++++++---------- lib/components/fabro-agent/src/types.rs | 12 +- lib/components/fabro-store/src/run_state.rs | 87 ++- lib/foundation/fabro-types/src/lib.rs | 8 +- .../fabro-types/src/run_event/agent.rs | 11 +- 10 files changed, 360 insertions(+), 401 deletions(-) diff --git a/lib/apps/fabro-cli/src/commands/run/run_progress/event.rs b/lib/apps/fabro-cli/src/commands/run/run_progress/event.rs index 9b8509c0e..a8596e6f1 100644 --- a/lib/apps/fabro-cli/src/commands/run/run_progress/event.rs +++ b/lib/apps/fabro-cli/src/commands/run/run_progress/event.rs @@ -189,12 +189,9 @@ pub(super) enum ProgressEvent { LlmRequestFinished { stage_node_id: String, }, - SubagentSpawned { - stage_node_id: String, - agent_id: String, - task: String, - }, - SubagentTurnStarted { + /// A subagent started work. Generation 1 is the spawn; later generations + /// are further turns in the same child session. + SubagentStarted { stage_node_id: String, agent_id: String, task: String, @@ -432,12 +429,13 @@ pub(super) fn from_run_event(stored: &RunEvent) -> Option { stage_node_id: node_id, }) } - EventBody::AgentSubSpawned(props) => Some(ProgressEvent::SubagentSpawned { + EventBody::AgentSubSpawned(props) => Some(ProgressEvent::SubagentStarted { stage_node_id: node_id, agent_id: props.agent_id.clone(), task: props.task.clone(), + generation: props.generation, }), - EventBody::AgentSubTurnStarted(props) => Some(ProgressEvent::SubagentTurnStarted { + EventBody::AgentSubTurnStarted(props) => Some(ProgressEvent::SubagentStarted { stage_node_id: node_id, agent_id: props.agent_id.clone(), task: props.task.clone(), diff --git a/lib/apps/fabro-cli/src/commands/run/run_progress/mod.rs b/lib/apps/fabro-cli/src/commands/run/run_progress/mod.rs index 44fbb484b..9e7e85030 100644 --- a/lib/apps/fabro-cli/src/commands/run/run_progress/mod.rs +++ b/lib/apps/fabro-cli/src/commands/run/run_progress/mod.rs @@ -363,21 +363,13 @@ impl ProgressUI { ProgressEvent::LlmRequestFinished { stage_node_id } => { self.stage.on_llm_request_finished(&stage_node_id); } - ProgressEvent::SubagentSpawned { - stage_node_id, - agent_id, - task, - } => { - self.stage - .on_subagent_spawned(renderer, &stage_node_id, &agent_id, &task); - } - ProgressEvent::SubagentTurnStarted { + ProgressEvent::SubagentStarted { stage_node_id, agent_id, task, generation, } => { - self.stage.on_subagent_turn_started( + self.stage.on_subagent_started( renderer, &stage_node_id, &agent_id, diff --git a/lib/apps/fabro-cli/src/commands/run/run_progress/stage_display.rs b/lib/apps/fabro-cli/src/commands/run/run_progress/stage_display.rs index 235f70918..f2b7dbc79 100644 --- a/lib/apps/fabro-cli/src/commands/run/run_progress/stage_display.rs +++ b/lib/apps/fabro-cli/src/commands/run/run_progress/stage_display.rs @@ -3,7 +3,7 @@ use std::convert::TryFrom; use std::time::Duration; use chrono::{DateTime, Utc}; -use fabro_types::LlmOutputKind; +use fabro_types::{INITIAL_SUBAGENT_GENERATION, LlmOutputKind}; use fabro_workflow::outcome::{StageOutcome, format_cost}; use indicatif::ProgressBar; @@ -606,17 +606,26 @@ impl StageDisplay { ); } - pub(super) fn on_subagent_spawned( + /// Show a subagent starting work. Generation 1 is the spawn; a later + /// generation is another turn in the same child session, so it reads as a + /// return to work rather than a new agent. + pub(super) fn on_subagent_started( &mut self, renderer: &ProgressRenderer, stage_node_id: &str, agent_id: &str, task: &str, + generation: u64, ) { if !self.verbose { return; } + let (glyph, turn) = if generation > INITIAL_SUBAGENT_GENERATION { + ("\u{21bb}", format!("turn {generation} ")) + } else { + ("\u{25b8}", String::new()) + }; self.insert_subagent_line_for_stage( renderer, stage_node_id, @@ -624,7 +633,7 @@ impl StageDisplay { .styles() .dim .apply_to(format!( - "\u{25b8} subagent[{agent_id}] \"{}\"", + "{glyph} subagent[{agent_id}] {turn}\"{}\"", styles::truncate(task, 50) )) .to_string(), @@ -655,32 +664,6 @@ impl StageDisplay { ); } - pub(super) fn on_subagent_turn_started( - &mut self, - renderer: &ProgressRenderer, - stage_node_id: &str, - agent_id: &str, - task: &str, - generation: u64, - ) { - if !self.verbose { - return; - } - - self.insert_subagent_line_for_stage( - renderer, - stage_node_id, - &renderer - .styles() - .dim - .apply_to(format!( - "\u{21bb} subagent[{agent_id}] turn {generation} \"{}\"", - styles::truncate(task, 50) - )) - .to_string(), - ); - } - fn finish_stage( &mut self, renderer: &ProgressRenderer, diff --git a/lib/components/fabro-agent/src/cli.rs b/lib/components/fabro-agent/src/cli.rs index d6f31d23a..720b8b0af 100644 --- a/lib/components/fabro-agent/src/cli.rs +++ b/lib/components/fabro-agent/src/cli.rs @@ -693,25 +693,19 @@ pub async fn run_with_args_and_client_and_catalog( depth, task, generation, - } => { - let task_preview = if task.len() > 60 { - &task[..task.floor_char_boundary(60)] - } else { - task - }; - eprintln!( - " {}", - s.dim.apply_to(format!( - "{child_prefix}\u{25b6} subagent {agent_id} spawned (depth={depth}, generation={generation}) task={task_preview:?}" - )), - ); } - AgentEvent::SubAgentTurnStarted { + | AgentEvent::SubAgentTurnStarted { agent_id, depth, task, generation, } => { + let started = + if matches!(event.event, AgentEvent::SubAgentSpawned { .. }) { + "spawned" + } else { + "turn started" + }; let task_preview = if task.len() > 60 { &task[..task.floor_char_boundary(60)] } else { @@ -720,7 +714,7 @@ pub async fn run_with_args_and_client_and_catalog( eprintln!( " {}", s.dim.apply_to(format!( - "{child_prefix}\u{25b6} subagent {agent_id} turn started (depth={depth}, generation={generation}) task={task_preview:?}" + "{child_prefix}\u{25b6} subagent {agent_id} {started} (depth={depth}, generation={generation}) task={task_preview:?}" )), ); } diff --git a/lib/components/fabro-agent/src/profiles/claude5_tools.rs b/lib/components/fabro-agent/src/profiles/claude5_tools.rs index d7c2b3e80..d9dd53aa6 100644 --- a/lib/components/fabro-agent/src/profiles/claude5_tools.rs +++ b/lib/components/fabro-agent/src/profiles/claude5_tools.rs @@ -333,7 +333,7 @@ pub(crate) fn make_task_output_tool(supervisor: SubAgentSupervisor) -> Registere } match supervisor.status(task_id) { - Some(SubAgentStatus::Finished(result)) => { + Some(SubAgentStatus::Finished { result, .. }) => { return finished_output(&supervisor, task_id, result); } Some(SubAgentStatus::Running) if !block => { diff --git a/lib/components/fabro-agent/src/subagent.rs b/lib/components/fabro-agent/src/subagent.rs index 2025cbd04..d6f29f4de 100644 --- a/lib/components/fabro-agent/src/subagent.rs +++ b/lib/components/fabro-agent/src/subagent.rs @@ -3,6 +3,7 @@ use std::sync::{Arc, Mutex, RwLock, Weak}; use std::time::Duration; use fabro_llm::types::ToolDefinition; +use fabro_types::INITIAL_SUBAGENT_GENERATION; use fabro_util::error as util_error; use futures::future; use tokio::sync::{mpsc, oneshot, watch}; @@ -81,18 +82,29 @@ fn escape_notification_xml(value: &str) -> String { #[derive(Debug, Clone)] pub enum SubAgentStatus { Running, - Finished(Result), + /// The turn ended. `reusable` reports whether the child session survived it + /// and can start another turn, so a finished-but-spent agent and a + /// finished-and-ready one cannot be confused. + Finished { + result: Result, + reusable: bool, + }, Closing, Closed, } const SUBAGENT_SHUTDOWN_GRACE: Duration = Duration::from_secs(5); -const INITIAL_SUBAGENT_GENERATION: u64 = 1; +/// One idle child accepts one next turn. `send_input` reserves this single slot +/// before it makes the agent running, so an agent can never be running with no +/// turn on its way; input for a running agent goes to the follow-up queue +/// instead. const SUBAGENT_COMMAND_CAPACITY: usize = 1; +/// Start the next turn of an existing child session. #[derive(Debug)] -enum SubAgentCommand { - Start { generation: u64, prompt: String }, +struct StartTurn { + generation: u64, + prompt: String, } struct ParentNotificationState { @@ -104,11 +116,9 @@ struct SubAgent { status: watch::Sender, generation: u64, results: HashMap>, - reusable: bool, - command_tx: mpsc::Sender, + command_tx: mpsc::Sender, runner_stop: CancellationToken, cleanup_done: watch::Sender, - cleanup_started: bool, monitor_task: Option>, event_forwarder: Option>, cleanup_task: Option>, @@ -154,9 +164,32 @@ struct SupervisorState { lifecycle_draining: bool, } +impl SupervisorState { + fn agent(&self, agent_id: &str) -> Result<&SubAgent, Error> { + self.agents + .get(agent_id) + .ok_or_else(|| unknown_agent(agent_id)) + } + + fn agent_mut(&mut self, agent_id: &str) -> Result<&mut SubAgent, Error> { + self.agents + .get_mut(agent_id) + .ok_or_else(|| unknown_agent(agent_id)) + } + + fn queue_lifecycle_event(&mut self, event: AgentEvent) { + self.lifecycle_events.push_back(event); + } +} + +fn unknown_agent(agent_id: &str) -> Error { + Error::InvalidState(format!( + "No agent found with id: {agent_id} (it was never spawned)" + )) +} + struct ShutdownWork { - agent_id: String, - depth: usize, + handle: SubAgentHandle, generation: u64, close_running_agent: bool, status: watch::Sender, @@ -166,7 +199,6 @@ struct ShutdownWork { child_abort_handle: AbortHandle, cancel_token: CancellationToken, runner_stop: CancellationToken, - state: Weak>, } impl Drop for ShutdownWork { @@ -206,30 +238,44 @@ fn signal_notifications(changed: &watch::Sender) { }); } -fn queue_lifecycle_event(state: &mut SupervisorState, event: AgentEvent) { - state.lifecycle_events.push_back(event); +/// Clear the draining flag however the drain ends, so one panicking callback +/// cannot silence every later lifecycle event. +struct DrainingGuard<'a>(&'a Arc>); + +impl Drop for DrainingGuard<'_> { + fn drop(&mut self) { + self.0 + .lock() + .expect("subagent state lock poisoned") + .lifecycle_draining = false; + } } -/// Deliver lifecycle callbacks in the same order as the state transitions -/// that queued them. Callbacks run without the supervisor lock held and may -/// safely call back into the supervisor. +/// Deliver lifecycle callbacks in the same order as the state transitions that +/// queued them. +/// +/// The queue exists for cross-thread ordering: a runner thread that releases +/// the lock after committing one generation would otherwise race a `send_input` +/// thread emitting the next generation's start, and consumers would see the +/// turns out of order. Callbacks run with no lock held, so one may also call +/// back into the supervisor without deadlocking. fn drain_lifecycle_events( state: &Arc>, event_callback: &Arc>>, ) { { - let mut state = state.lock().expect("subagent state lock poisoned"); - if state.lifecycle_draining { + let mut locked = state.lock().expect("subagent state lock poisoned"); + if locked.lifecycle_draining { return; } - state.lifecycle_draining = true; + locked.lifecycle_draining = true; } + let _draining = DrainingGuard(state); loop { let event = { - let mut state = state.lock().expect("subagent state lock poisoned"); - let Some(event) = state.lifecycle_events.pop_front() else { - state.lifecycle_draining = false; + let mut locked = state.lock().expect("subagent state lock poisoned"); + let Some(event) = locked.lifecycle_events.pop_front() else { return; }; event @@ -273,67 +319,109 @@ fn completion_event( } } -/// Commit one generation result or claim a follow-up that raced its final -/// boundary. The supervisor state lock is acquired before the follow-up queue -/// lock, which is also the ordering used by `send_input`. -fn commit_turn_result( - state: &Arc>, - event_callback: &Arc>>, - notifications_changed: &watch::Sender, - agent_id: &str, - depth: usize, - generation: u64, - result: &Result, - reusable: bool, -) -> TurnCommit { - let outcome = { - let mut state = state.lock().expect("subagent state lock poisoned"); - let Some(agent) = state.agents.get_mut(agent_id) else { +/// One child's view of its supervisor: the shared state plus the identity every +/// lifecycle transition needs. +/// +/// The state reference is weak because a child task reaches its supervisor +/// through this handle, and a strong reference would close the cycle +/// state -> `SubAgent` -> runner task -> handle. +#[derive(Clone)] +struct SubAgentHandle { + state: Weak>, + event_callback: Arc>>, + notifications_changed: Arc>, + agent_id: String, + depth: usize, +} + +impl SubAgentHandle { + /// Commit one generation result, or claim a follow-up that raced its final + /// boundary. The supervisor state lock is acquired before the follow-up + /// queue lock, which is also the ordering used by `send_input`. + fn commit_turn_result( + &self, + generation: u64, + result: &Result, + reusable: bool, + ) -> TurnCommit { + let Some(state) = self.state.upgrade() else { return TurnCommit::Stopping; }; - if agent.generation != generation - || !matches!(*agent.status.borrow(), SubAgentStatus::Running) - { - return TurnCommit::Stopping; - } - - if reusable { - let next_prompt = agent - .followup_queue - .lock() - .expect("followup queue lock poisoned") - .pop_front(); - if let Some(next_prompt) = next_prompt { - return TurnCommit::Continue(next_prompt); + let outcome = { + let mut locked = state.lock().expect("subagent state lock poisoned"); + let Ok(agent) = locked.agent_mut(&self.agent_id) else { + return TurnCommit::Stopping; + }; + if agent.generation != generation + || !matches!(*agent.status.borrow(), SubAgentStatus::Running) + { + return TurnCommit::Stopping; } - } - agent.results.insert(generation, result.clone()); - agent.reusable = reusable; - agent - .status - .send_replace(SubAgentStatus::Finished(result.clone())); - queue_lifecycle_event( - &mut state, - completion_event(agent_id, depth, generation, result), - ); - TurnCommit::Finished - }; + if reusable { + let next_prompt = agent + .followup_queue + .lock() + .expect("followup queue lock poisoned") + .pop_front(); + if let Some(next_prompt) = next_prompt { + return TurnCommit::Continue(next_prompt); + } + } - signal_notifications(notifications_changed); - drain_lifecycle_events(state, event_callback); - outcome + agent.results.insert(generation, result.clone()); + agent.status.send_replace(SubAgentStatus::Finished { + result: result.clone(), + reusable, + }); + locked.queue_lifecycle_event(completion_event( + &self.agent_id, + self.depth, + generation, + result, + )); + TurnCommit::Finished + }; + + self.publish(&state); + outcome + } + + /// The generation this agent is on now, or `None` once the supervisor or + /// the agent itself is gone. + fn current_generation(&self) -> Option { + let state = self.state.upgrade()?; + let locked = state.lock().expect("subagent state lock poisoned"); + locked + .agent(&self.agent_id) + .ok() + .map(|agent| agent.generation) + } + + fn queue_and_publish(&self, event: AgentEvent) { + let Some(state) = self.state.upgrade() else { + return; + }; + state + .lock() + .expect("subagent state lock poisoned") + .queue_lifecycle_event(event); + self.publish(&state); + } + + /// Wake notification waiters and deliver queued lifecycle callbacks. Always + /// called with no supervisor lock held. + fn publish(&self, state: &Arc>) { + signal_notifications(&self.notifications_changed); + drain_lifecycle_events(state, &self.event_callback); + } } async fn run_subagent_session( mut session: Session, - state: Weak>, - event_callback: Arc>>, - notifications_changed: Arc>, - agent_id: String, - depth: usize, + handle: SubAgentHandle, initial_prompt: String, - mut command_rx: mpsc::Receiver, + mut command_rx: mpsc::Receiver, runner_stop: CancellationToken, start_rx: oneshot::Receiver<()>, ) { @@ -342,35 +430,20 @@ async fn run_subagent_session( } if let Err(error) = session.initialize().await { - if let Some(state) = state.upgrade() { - let result = Err(error); - commit_turn_result( - &state, - &event_callback, - ¬ifications_changed, - &agent_id, - depth, - INITIAL_SUBAGENT_GENERATION, - &result, - false, - ); - } - runner_stop.cancelled().await; - let reason = if session.cancel_token().is_cancelled() { - SessionShutdownReason::Cancelled - } else { - SessionShutdownReason::Error - }; - session.shutdown(reason).await; + handle.commit_turn_result(INITIAL_SUBAGENT_GENERATION, &Err(error), false); + // A session that never initialized has no history worth reusing, so + // release it and its sandbox now rather than holding both until the + // parent closes the agent. + session.shutdown(shutdown_reason(&session, true)).await; return; } - let mut command = SubAgentCommand::Start { + let mut command = StartTurn { generation: INITIAL_SUBAGENT_GENERATION, prompt: initial_prompt, }; 'commands: loop { - let SubAgentCommand::Start { + let StartTurn { generation, mut prompt, } = command; @@ -398,19 +471,7 @@ async fn run_subagent_session( }); let reusable = session.state() == SessionState::Idle && !session.cancel_token().is_cancelled(); - let Some(state) = state.upgrade() else { - break 'commands; - }; - match commit_turn_result( - &state, - &event_callback, - ¬ifications_changed, - &agent_id, - depth, - generation, - &result, - reusable, - ) { + match handle.commit_turn_result(generation, &result, reusable) { TurnCommit::Continue(next_prompt) => prompt = next_prompt, TurnCommit::Finished => break, TurnCommit::Stopping => break 'commands, @@ -429,49 +490,35 @@ async fn run_subagent_session( }; } - let reason = if session.cancel_token().is_cancelled() { - SessionShutdownReason::Cancelled - } else { - SessionShutdownReason::Completed - }; - session.shutdown(reason).await; + session.shutdown(shutdown_reason(&session, false)).await; } -fn spawn_runner_monitor( - runner_task: JoinHandle<()>, - state: Weak>, - event_callback: Arc>>, - notifications_changed: Arc>, - agent_id: String, - depth: usize, -) -> JoinHandle<()> { +/// Cancellation always wins as the reported reason; otherwise a session that +/// failed to start reports an error and one that ran reports completion. +fn shutdown_reason(session: &Session, failed_to_start: bool) -> SessionShutdownReason { + if session.cancel_token().is_cancelled() { + SessionShutdownReason::Cancelled + } else if failed_to_start { + SessionShutdownReason::Error + } else { + SessionShutdownReason::Completed + } +} + +/// Report a runner that died without committing its own result, so the agent +/// never sits in `Running` with nothing left to run. +fn spawn_runner_monitor(runner_task: JoinHandle<()>, handle: SubAgentHandle) -> JoinHandle<()> { tokio::spawn(async move { let Err(error) = runner_task.await else { return; }; - let Some(state) = state.upgrade() else { + let Some(generation) = handle.current_generation() else { return; }; let task_result = Err(Error::InvalidState(format!( "Agent task failed to join: {error}" ))); - let generation = { - let state = state.lock().expect("subagent state lock poisoned"); - let Some(agent) = state.agents.get(&agent_id) else { - return; - }; - agent.generation - }; - commit_turn_result( - &state, - &event_callback, - ¬ifications_changed, - &agent_id, - depth, - generation, - &task_result, - false, - ); + handle.commit_turn_result(generation, &task_result, false); }) } @@ -499,6 +546,24 @@ impl SubAgentSupervisor { } } + /// A child's view of this supervisor, for the tasks that run that child. + fn handle(&self, agent_id: String, depth: usize) -> SubAgentHandle { + SubAgentHandle { + state: Arc::downgrade(&self.state), + event_callback: Arc::clone(&self.event_callback), + notifications_changed: Arc::clone(&self.notifications_changed), + agent_id, + depth, + } + } + + /// Wake notification waiters and deliver queued lifecycle callbacks, after + /// the state lock has been released. + fn publish(&self) { + signal_notifications(&self.notifications_changed); + drain_lifecycle_events(&self.state, &self.event_callback); + } + pub fn set_event_callback(&self, cb: SubAgentEventCallback) { *self .event_callback @@ -506,18 +571,6 @@ impl SubAgentSupervisor { .expect("subagent callback lock poisoned") = Some(cb); } - #[cfg(test)] - fn emit_event(&self, event: AgentEvent) { - let callback = self - .event_callback - .read() - .expect("subagent callback lock poisoned") - .clone(); - if let Some(cb) = callback { - cb(SubAgentCallbackEvent::Lifecycle(event)); - } - } - pub fn spawn( &self, session: Session, @@ -597,27 +650,17 @@ impl SubAgentSupervisor { let (command_tx, command_rx) = mpsc::channel(SUBAGENT_COMMAND_CAPACITY); let runner_stop = CancellationToken::new(); let child_depth = depth + 1; + let handle = self.handle(agent_id.clone(), child_depth); let runner_task = tokio::spawn(run_subagent_session( session, - Arc::downgrade(&self.state), - Arc::clone(&self.event_callback), - Arc::clone(&self.notifications_changed), - agent_id.clone(), - child_depth, + handle.clone(), task_prompt.clone(), command_rx, runner_stop.clone(), start_rx, )); let child_abort_handle = runner_task.abort_handle(); - let monitor_task = spawn_runner_monitor( - runner_task, - Arc::downgrade(&self.state), - Arc::clone(&self.event_callback), - Arc::clone(&self.notifications_changed), - agent_id.clone(), - child_depth, - ); + let monitor_task = spawn_runner_monitor(runner_task, handle); let (status, _) = watch::channel(SubAgentStatus::Running); let (cleanup_done, _) = watch::channel(false); @@ -634,11 +677,9 @@ impl SubAgentSupervisor { status, generation: INITIAL_SUBAGENT_GENERATION, results: HashMap::new(), - reusable: false, command_tx, runner_stop, cleanup_done, - cleanup_started: false, monitor_task: Some(monitor_task), event_forwarder, cleanup_task: None, @@ -649,15 +690,14 @@ impl SubAgentSupervisor { parent_notification, spawn_seq, }); - queue_lifecycle_event(&mut state, AgentEvent::SubAgentSpawned { + state.queue_lifecycle_event(AgentEvent::SubAgentSpawned { agent_id: agent_id.clone(), depth: child_depth, task: task_prompt, generation: INITIAL_SUBAGENT_GENERATION, }); } - signal_notifications(&self.notifications_changed); - drain_lifecycle_events(&self.state, &self.event_callback); + self.publish(); let _ = start_tx.send(()); Ok(agent_id) @@ -666,11 +706,7 @@ impl SubAgentSupervisor { pub fn send_input(&self, agent_id: &str, message: &str) -> Result<(), Error> { let resumed = { let mut state = self.state.lock().expect("subagent state lock poisoned"); - let agent = state.agents.get_mut(agent_id).ok_or_else(|| { - Error::InvalidState(format!( - "No agent found with id: {agent_id} (it was never spawned)" - )) - })?; + let agent = state.agent_mut(agent_id)?; let status = agent.status.borrow().clone(); match status { SubAgentStatus::Running => { @@ -681,8 +717,8 @@ impl SubAgentSupervisor { .push_back(message.to_string()); None } - SubAgentStatus::Finished(_) => { - if !agent.reusable { + SubAgentStatus::Finished { reusable, .. } => { + if !reusable { return Err(Error::InvalidState(format!( "Agent {agent_id} cannot accept more input because its session ended" ))); @@ -702,13 +738,12 @@ impl SubAgentSupervisor { )) })?; agent.generation = generation; - agent.reusable = false; agent.status.send_replace(SubAgentStatus::Running); if let Some(notification) = &mut agent.parent_notification { notification.pending_generations.push_back(generation); } let depth = agent.depth; - queue_lifecycle_event(&mut state, AgentEvent::SubAgentTurnStarted { + state.queue_lifecycle_event(AgentEvent::SubAgentTurnStarted { agent_id: agent_id.to_string(), depth, task: message.to_string(), @@ -725,9 +760,8 @@ impl SubAgentSupervisor { }; if let Some((permit, generation)) = resumed { - signal_notifications(&self.notifications_changed); - drain_lifecycle_events(&self.state, &self.event_callback); - permit.send(SubAgentCommand::Start { + self.publish(); + permit.send(StartTurn { generation, prompt: message.to_string(), }); @@ -743,22 +777,14 @@ impl SubAgentSupervisor { ) -> Result { let (generation, mut status) = { let state = self.state.lock().expect("subagent state lock poisoned"); - let agent = state.agents.get(agent_id).ok_or_else(|| { - Error::InvalidState(format!( - "No agent found with id: {agent_id} (it was never spawned)" - )) - })?; + let agent = state.agent(agent_id)?; (agent.generation, agent.status.subscribe()) }; loop { let current = { let state = self.state.lock().expect("subagent state lock poisoned"); - let agent = state.agents.get(agent_id).ok_or_else(|| { - Error::InvalidState(format!( - "No agent found with id: {agent_id} (it was never spawned)" - )) - })?; + let agent = state.agent(agent_id)?; if let Some(result) = agent.results.get(&generation) { return result.clone(); } @@ -771,7 +797,7 @@ impl SubAgentSupervisor { "Agent {agent_id} has been closed" ))); } - SubAgentStatus::Running | SubAgentStatus::Finished(_) => {} + SubAgentStatus::Running | SubAgentStatus::Finished { .. } => {} } tokio::select! { @@ -916,15 +942,11 @@ impl SubAgentSupervisor { fn begin_shutdown(&self, agent_id: &str, strict: bool) -> Result { let mut state = self.state.lock().expect("subagent state lock poisoned"); - let agent = state.agents.get_mut(agent_id).ok_or_else(|| { - Error::InvalidState(format!( - "No agent found with id: {agent_id} (it was never spawned)" - )) - })?; + let agent = state.agent_mut(agent_id)?; let close_running_agent = match agent.status.borrow().clone() { SubAgentStatus::Running => true, - SubAgentStatus::Finished(_) => false, + SubAgentStatus::Finished { .. } => false, SubAgentStatus::Closing | SubAgentStatus::Closed if strict => { return Err(Error::InvalidState(format!( "Agent {agent_id} is already closed" @@ -935,25 +957,18 @@ impl SubAgentSupervisor { } SubAgentStatus::Closed => return Ok(ShutdownDisposition::Done), }; + // Reaching here means the status was Running or Finished, so this call + // is the one that commits shutdown: the arms above return for a status + // already Closing or Closed, and the only write out of Closing is + // `run_shutdown`'s move to Closed. + debug_assert!(agent.cleanup_task.is_none()); agent.status.send_replace(SubAgentStatus::Closing); // Shutdown is committed, so no pending result will reach the parent. agent.parent_notification = None; - if agent.cleanup_started { - return if strict { - Err(Error::InvalidState(format!( - "Agent {agent_id} is already closed" - ))) - } else { - Ok(ShutdownDisposition::Follow(agent.cleanup_done.subscribe())) - }; - } - agent.cleanup_started = true; - Ok(ShutdownDisposition::Lead(ShutdownWork { - agent_id: agent_id.to_string(), - depth: agent.depth, + handle: self.handle(agent_id.to_string(), agent.depth), generation: agent.generation, close_running_agent, status: agent.status.clone(), @@ -963,14 +978,10 @@ impl SubAgentSupervisor { child_abort_handle: agent.child_abort_handle.clone(), cancel_token: agent.cancel_token.clone(), runner_stop: agent.runner_stop.clone(), - state: Arc::downgrade(&self.state), })) } - async fn run_shutdown( - mut work: ShutdownWork, - event_callback: Arc>>, - ) { + async fn run_shutdown(mut work: ShutdownWork) { let _cleanup_done = CleanupDoneGuard(work.cleanup_done.clone()); let deadline = Instant::now() + SUBAGENT_SHUTDOWN_GRACE; work.runner_stop.cancel(); @@ -1001,25 +1012,18 @@ impl SubAgentSupervisor { } }); if emit_closed { - if let Some(state) = work.state.upgrade() { - { - let mut state = state.lock().expect("subagent state lock poisoned"); - queue_lifecycle_event(&mut state, AgentEvent::SubAgentClosed { - agent_id: work.agent_id.clone(), - depth: work.depth, - generation: work.generation, - }); - } - drain_lifecycle_events(&state, &event_callback); - } + work.handle.queue_and_publish(AgentEvent::SubAgentClosed { + agent_id: work.handle.agent_id.clone(), + depth: work.handle.depth, + generation: work.generation, + }); } } fn spawn_shutdown(&self, work: ShutdownWork) -> watch::Receiver { let cleanup_done = work.cleanup_done.subscribe(); - let agent_id = work.agent_id.clone(); - let event_callback = Arc::clone(&self.event_callback); - let cleanup_task = tokio::spawn(Self::run_shutdown(work, event_callback)); + let agent_id = work.handle.agent_id.clone(); + let cleanup_task = tokio::spawn(Self::run_shutdown(work)); let mut state = self.state.lock().expect("subagent state lock poisoned"); let agent = state .agents @@ -1135,10 +1139,7 @@ impl SubAgentSupervisor { let runner_stop = CancellationToken::new(); let depth = 1; let (monitor_start_tx, monitor_start_rx) = oneshot::channel(); - let state = Arc::downgrade(&self.state); - let event_callback = Arc::clone(&self.event_callback); - let notifications_changed = Arc::clone(&self.notifications_changed); - let monitored_agent_id = agent_id.clone(); + let handle = self.handle(agent_id.clone(), depth); let monitor_task = tokio::spawn(async move { let _ = monitor_start_rx.await; let task_result = match child_task.await { @@ -1147,19 +1148,7 @@ impl SubAgentSupervisor { "Agent task failed to join: {error}" ))), }; - let Some(state) = state.upgrade() else { - return; - }; - commit_turn_result( - &state, - &event_callback, - ¬ifications_changed, - &monitored_agent_id, - depth, - INITIAL_SUBAGENT_GENERATION, - &task_result, - false, - ); + handle.commit_turn_result(INITIAL_SUBAGENT_GENERATION, &task_result, false); }); { let mut state = self.state.lock().expect("subagent state lock poisoned"); @@ -1167,11 +1156,9 @@ impl SubAgentSupervisor { status, generation: INITIAL_SUBAGENT_GENERATION, results: HashMap::new(), - reusable: false, command_tx, runner_stop, cleanup_done, - cleanup_started: false, monitor_task: Some(monitor_task), event_forwarder, cleanup_task: None, @@ -1718,7 +1705,7 @@ mod tests { assert!(agent_result.turns_used > 0); assert!(matches!( manager.status(&agent_id), - Some(SubAgentStatus::Finished(Ok(_))) + Some(SubAgentStatus::Finished { result: Ok(_), .. }) )); } @@ -1948,17 +1935,6 @@ mod tests { assert_eq!(grandchild.parent_session_id.as_deref(), Some("child")); } - #[test] - fn no_callback_does_not_panic() { - // Manager without callback should not panic on emit - let manager = SubAgentSupervisor::new(3); - manager.emit_event(AgentEvent::SubAgentClosed { - agent_id: "x".into(), - depth: 0, - generation: 1, - }); - } - #[tokio::test] async fn close_all_closes_all_agents() { let manager = SubAgentSupervisor::new(3); @@ -1995,7 +1971,7 @@ mod tests { assert_eq!(result2.output, "cached output"); assert!(matches!( manager.status(&agent_id), - Some(SubAgentStatus::Finished(Ok(_))) + Some(SubAgentStatus::Finished { result: Ok(_), .. }) )); } @@ -2155,7 +2131,7 @@ mod tests { let _ = manager.wait(&agent_id).await.unwrap(); assert!(matches!( manager.status(&agent_id), - Some(SubAgentStatus::Finished(Ok(_))) + Some(SubAgentStatus::Finished { result: Ok(_), .. }) )); manager.close_agent(&agent_id).await.unwrap(); @@ -2199,7 +2175,7 @@ mod tests { time::timeout(Duration::from_secs(1), async { while !matches!( supervisor.status(&agent_id), - Some(SubAgentStatus::Finished(Ok(_))) + Some(SubAgentStatus::Finished { result: Ok(_), .. }) ) { yield_now().await; } @@ -2255,19 +2231,16 @@ mod tests { #[tokio::test] async fn wait_returns_its_target_generation_after_a_later_turn_starts() { let supervisor = SubAgentSupervisor::new(3); - let child_cancel = CancellationToken::new(); - let task_cancel = child_cancel.clone(); - let child = tokio::spawn(async move { - task_cancel.cancelled().await; - Ok(SubAgentResult { - output: "unused".to_string(), - success: true, - turns_used: 1, - }) - }); - let agent_id = "generation-aware-wait".to_string(); - supervisor.supervise_test_task(agent_id.clone(), child, child_cancel, None); + let child = make_session(vec![ + text_response("generation one"), + text_response("generation two"), + ]) + .await; + let agent_id = supervisor.spawn(child, "implement".to_string(), 0).unwrap(); + // Registering the wait pins it to generation one. It stays unpolled + // from here, so generation one's completion and generation two's start + // reach it as a single coalesced watch update. let wait_cancel = CancellationToken::new(); let mut wait = Box::pin(supervisor.wait_with_cancel(&agent_id, &wait_cancel)); assert!( @@ -2275,26 +2248,15 @@ mod tests { "generation one should still be running" ); - { - let mut state = supervisor - .state - .lock() - .expect("subagent state lock poisoned"); - let agent = state.agents.get_mut(&agent_id).unwrap(); - let first_result = Ok(SubAgentResult { - output: "generation one".to_string(), - success: true, - turns_used: 1, - }); - agent - .results - .insert(INITIAL_SUBAGENT_GENERATION, first_result.clone()); - agent - .status - .send_replace(SubAgentStatus::Finished(first_result)); - agent.generation = INITIAL_SUBAGENT_GENERATION + 1; - agent.status.send_replace(SubAgentStatus::Running); + while !matches!( + supervisor.status(&agent_id), + Some(SubAgentStatus::Finished { .. }) + ) { + yield_now().await; } + supervisor + .send_input(&agent_id, "Fix the review findings") + .unwrap(); let result = time::timeout(Duration::from_secs(1), wait) .await diff --git a/lib/components/fabro-agent/src/types.rs b/lib/components/fabro-agent/src/types.rs index 38559b180..4cbd85736 100644 --- a/lib/components/fabro-agent/src/types.rs +++ b/lib/components/fabro-agent/src/types.rs @@ -224,10 +224,6 @@ pub struct McpToolSummary { pub original_name: String, } -const fn initial_subagent_generation() -> u64 { - 1 -} - #[derive(Debug, Clone, Serialize, Deserialize)] pub enum AgentEvent { SessionStarted { @@ -359,7 +355,7 @@ pub enum AgentEvent { agent_id: String, depth: usize, task: String, - #[serde(default = "initial_subagent_generation")] + #[serde(default = "fabro_types::initial_subagent_generation")] generation: u64, }, SubAgentTurnStarted { @@ -371,7 +367,7 @@ pub enum AgentEvent { SubAgentCompleted { agent_id: String, depth: usize, - #[serde(default = "initial_subagent_generation")] + #[serde(default = "fabro_types::initial_subagent_generation")] generation: u64, success: bool, turns_used: usize, @@ -379,14 +375,14 @@ pub enum AgentEvent { SubAgentFailed { agent_id: String, depth: usize, - #[serde(default = "initial_subagent_generation")] + #[serde(default = "fabro_types::initial_subagent_generation")] generation: u64, error: Error, }, SubAgentClosed { agent_id: String, depth: usize, - #[serde(default = "initial_subagent_generation")] + #[serde(default = "fabro_types::initial_subagent_generation")] generation: u64, }, McpServerReady { diff --git a/lib/components/fabro-store/src/run_state.rs b/lib/components/fabro-store/src/run_state.rs index d3a8d5aae..ed520dcd6 100644 --- a/lib/components/fabro-store/src/run_state.rs +++ b/lib/components/fabro-store/src/run_state.rs @@ -689,46 +689,54 @@ impl RunProjectionReducer for RunProjection { status: SubAgentStatus::Running, }); } + // A reused subagent stays one projected row: the spawn task and + // generation 1 identify it, and every later generation only moves + // its status. The per-turn task and generation stay in the event + // log for consumers that need each turn. EventBody::AgentSubTurnStarted(props) => { - let Some(stage) = stage_at_stored_or_visit(self, stored, props.visit, event.seq) - else { - return Ok(()); - }; - if let Some(subagent) = subagent_mut(stage, &props.agent_id) { - subagent.status = SubAgentStatus::Running; - } + set_subagent_status( + self, + stored, + props.visit, + event.seq, + &props.agent_id, + SubAgentStatus::Running, + ); } EventBody::AgentSubCompleted(props) => { - let Some(stage) = stage_at_stored_or_visit(self, stored, props.visit, event.seq) - else { - return Ok(()); - }; - if let Some(subagent) = subagent_mut(stage, &props.agent_id) { - subagent.status = SubAgentStatus::Completed { + set_subagent_status( + self, + stored, + props.visit, + event.seq, + &props.agent_id, + SubAgentStatus::Completed { success: props.success, turns_used: props.turns_used, - }; - } + }, + ); } EventBody::AgentSubFailed(props) => { - let Some(stage) = stage_at_stored_or_visit(self, stored, props.visit, event.seq) - else { - return Ok(()); - }; - if let Some(subagent) = subagent_mut(stage, &props.agent_id) { - subagent.status = SubAgentStatus::Failed { + set_subagent_status( + self, + stored, + props.visit, + event.seq, + &props.agent_id, + SubAgentStatus::Failed { error: props.error.clone(), - }; - } + }, + ); } EventBody::AgentSubClosed(props) => { - let Some(stage) = stage_at_stored_or_visit(self, stored, props.visit, event.seq) - else { - return Ok(()); - }; - if let Some(subagent) = subagent_mut(stage, &props.agent_id) { - subagent.status = SubAgentStatus::Closed; - } + set_subagent_status( + self, + stored, + props.visit, + event.seq, + &props.agent_id, + SubAgentStatus::Closed, + ); } EventBody::AgentSkillsDiscovered(props) => { let Some(stage) = stage_at_stored_or_visit(self, stored, props.visit, event.seq) @@ -898,6 +906,25 @@ fn apply_todo_deleted(stage: &mut StageProjection, props: &TodoDeletedProps) { } } +/// Move an already-projected subagent to a new lifecycle status. Every +/// subagent event after the spawn updates the same row, so reuse shows one +/// agent returning to running rather than a second agent appearing. +fn set_subagent_status( + state: &mut RunProjection, + stored: &RunEvent, + visit: u32, + seq: u32, + agent_id: &str, + status: SubAgentStatus, +) { + let Some(stage) = stage_at_stored_or_visit(state, stored, visit, seq) else { + return; + }; + if let Some(subagent) = subagent_mut(stage, agent_id) { + subagent.status = status; + } +} + fn subagent_mut<'a>( stage: &'a mut StageProjection, agent_id: &str, diff --git a/lib/foundation/fabro-types/src/lib.rs b/lib/foundation/fabro-types/src/lib.rs index 7bbde04ab..309505c4a 100644 --- a/lib/foundation/fabro-types/src/lib.rs +++ b/lib/foundation/fabro-types/src/lib.rs @@ -114,10 +114,10 @@ pub use run_blob_id::RunBlobId; pub use run_event::{ AgentMcpToolSummary, AgentMemoryFileProps, AgentSkillActivationSource, AgentSkillSummary, AgentToolCategory, AgentToolSource, AgentToolSummary, AgentToolsAvailableProps, EventBody, - ExecOutputTail, FailoverProps, InterviewOption, LlmOutputKind, LlmRetryPhase, - MetadataSnapshotFailureKind, MetadataSnapshotPhase, RunEvent, RunNoticeCode, RunNoticeLevel, - RunPairEndedReason, RunPairFailedReason, RunRunnableSource, SessionCapability, - TodoCreatedProps, TodoDeletedProps, TodoUpdatedProps, + ExecOutputTail, FailoverProps, INITIAL_SUBAGENT_GENERATION, InterviewOption, LlmOutputKind, + LlmRetryPhase, MetadataSnapshotFailureKind, MetadataSnapshotPhase, RunEvent, RunNoticeCode, + RunNoticeLevel, RunPairEndedReason, RunPairFailedReason, RunRunnableSource, SessionCapability, + TodoCreatedProps, TodoDeletedProps, TodoUpdatedProps, initial_subagent_generation, }; pub use run_failure::RunFailure; pub use run_id::{RunId, fixtures}; diff --git a/lib/foundation/fabro-types/src/run_event/agent.rs b/lib/foundation/fabro-types/src/run_event/agent.rs index dd70b8561..a20ee4aaf 100644 --- a/lib/foundation/fabro-types/src/run_event/agent.rs +++ b/lib/foundation/fabro-types/src/run_event/agent.rs @@ -425,8 +425,15 @@ pub struct AgentSubClosedProps { pub visit: u32, } -const fn initial_subagent_generation() -> u64 { - 1 +/// The generation of a subagent's first turn. Events stored before subagent +/// session reuse existed carry no generation, so they read back as this. +pub const INITIAL_SUBAGENT_GENERATION: u64 = 1; + +/// Serde default for the generation of a stored subagent event. Public so +/// crates with their own subagent event types share this one definition. +#[must_use] +pub const fn initial_subagent_generation() -> u64 { + INITIAL_SUBAGENT_GENERATION } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] From 8c19b1e94a3e83f47cb2e51ab5c47352505d0a37 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Sat, 1 Aug 2026 10:13:27 -0400 Subject: [PATCH 3/4] fix: keep forwarding a reused child's events after a broadcast lag The subagent event forwarder left its loop on any `recv` error, including `Lagged`. A lagged broadcast receiver stays usable, so one transient lag silenced the child for the rest of its life while the task completed normally and shutdown joined it without noticing. Session reuse widens that window from a single turn to the whole parent session. Also borrow each result's output when rendering a parent notification instead of cloning it. Co-Authored-By: Claude Opus 5 (1M context) --- lib/components/fabro-agent/src/subagent.rs | 24 +++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/lib/components/fabro-agent/src/subagent.rs b/lib/components/fabro-agent/src/subagent.rs index d6f29f4de..2a90384fc 100644 --- a/lib/components/fabro-agent/src/subagent.rs +++ b/lib/components/fabro-agent/src/subagent.rs @@ -1,3 +1,4 @@ +use std::borrow::Cow; use std::collections::{HashMap, VecDeque}; use std::sync::{Arc, Mutex, RwLock, Weak}; use std::time::Duration; @@ -6,7 +7,7 @@ use fabro_llm::types::ToolDefinition; use fabro_types::INITIAL_SUBAGENT_GENERATION; use fabro_util::error as util_error; use futures::future; -use tokio::sync::{mpsc, oneshot, watch}; +use tokio::sync::{broadcast, mpsc, oneshot, watch}; use tokio::task::{AbortHandle, JoinHandle}; use tokio::time::{Instant, timeout_at}; use tokio_util::sync::CancellationToken; @@ -48,9 +49,14 @@ fn format_parent_notification_batch(notifications: &[SubAgentParentNotification] .iter() .map(|notification| { let (status, result) = match ¬ification.result { - Ok(result) if result.success => ("completed", result.output.clone()), - Ok(result) => ("failed", result.output.clone()), - Err(error) => ("failed", util_error::collect_chain(error).join(": ")), + Ok(result) if result.success => { + ("completed", Cow::Borrowed(result.output.as_str())) + } + Ok(result) => ("failed", Cow::Borrowed(result.output.as_str())), + Err(error) => ( + "failed", + Cow::Owned(util_error::collect_chain(error).join(": ")), + ), }; format!( "\n {}\n {status}\n \ @@ -621,7 +627,15 @@ impl SubAgentSupervisor { let mut rx = session.subscribe(); let callback = Arc::clone(&self.event_callback); Some(tokio::spawn(async move { - while let Ok(event) = rx.recv().await { + loop { + let event = match rx.recv().await { + Ok(event) => event, + // A lagged receiver stays usable, and a reused child + // forwards for the whole parent session. Giving up here + // would silence the child for the rest of its life. + Err(broadcast::error::RecvError::Lagged(_)) => continue, + Err(broadcast::error::RecvError::Closed) => break, + }; // Skip streaming / noise events if event.event.is_streaming_noise() || matches!( From a617eaa5eb39bd0dea285d2e621f53b0da80372f Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Sat, 1 Aug 2026 10:17:10 -0400 Subject: [PATCH 4/4] docs: record why a resumed turn's send cannot hang an agent Review asked twice whether `permit.send` can leave an agent Running with no turn on its way. It cannot, and the reasoning is not local to the call, so state it there. Co-Authored-By: Claude Opus 5 (1M context) --- lib/components/fabro-agent/src/subagent.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/lib/components/fabro-agent/src/subagent.rs b/lib/components/fabro-agent/src/subagent.rs index 2a90384fc..697aa8893 100644 --- a/lib/components/fabro-agent/src/subagent.rs +++ b/lib/components/fabro-agent/src/subagent.rs @@ -775,6 +775,14 @@ impl SubAgentSupervisor { if let Some((permit, generation)) = resumed { self.publish(); + // This send cannot fail: `OwnedPermit::send` returns `()`, and the + // capacity it needs was reserved above while the state lock was + // held. Should a concurrent close drop the receiver first, the + // command is discarded and the agent still cannot hang, because + // every path that stops the runner -- `run_shutdown` and the two + // drop impls -- is reached only after `begin_shutdown` has set + // `Closing` under this same lock. A waiter then observes the close + // and stops instead of waiting on a turn that will never run. permit.send(StartTurn { generation, prompt: message.to_string(),