From 1bb70adafd53c728856a2396975eb0b1368e72ae Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Sat, 1 Aug 2026 10:09:07 -0400 Subject: [PATCH] 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)]