From 2e470fdbe03bfffcefb5ce732ed0b26aa7b27b2d Mon Sep 17 00:00:00 2001 From: Larry <8cf5a83f590ec0955b11647d1c88f796a98e088c30a492c58e0e46c3026ae7a4@buzz.block.builderlab.xyz> Date: Wed, 12 Aug 2026 17:57:24 -0400 Subject: [PATCH] fix(acp): fence tasks by membership generation Signed-off-by: Larry <8cf5a83f590ec0955b11647d1c88f796a98e088c30a492c58e0e46c3026ae7a4@buzz.block.builderlab.xyz> --- crates/buzz-acp/src/lib.rs | 157 +++++++++++++++++++++++++++++------- crates/buzz-acp/src/pool.rs | 62 ++++++++++++-- 2 files changed, 183 insertions(+), 36 deletions(-) diff --git a/crates/buzz-acp/src/lib.rs b/crates/buzz-acp/src/lib.rs index a47ff405f0..93bad4116b 100644 --- a/crates/buzz-acp/src/lib.rs +++ b/crates/buzz-acp/src/lib.rs @@ -2378,9 +2378,13 @@ async fn tokio_main() -> Result<()> { // called on relay events or pool results, neither of which // arrive when the channel is silent. if queue.has_flushable_work() { - for (channel_id, thread_tags) in - dispatch_pending(&mut pool, &mut queue, &ctx, &mut last_activity) - { + for (channel_id, thread_tags) in dispatch_pending( + &mut pool, + &mut queue, + &ctx, + &membership_generations, + &mut last_activity, + ) { typing_channels.insert(channel_id, thread_tags); } } @@ -2427,9 +2431,13 @@ async fn tokio_main() -> Result<()> { // this, batches requeued during crash recovery sit idle until the // next relay event arrives — which can be minutes on quiet channels. if respawn_collected { - for (channel_id, thread_tags) in - dispatch_pending(&mut pool, &mut queue, &ctx, &mut last_activity) - { + for (channel_id, thread_tags) in dispatch_pending( + &mut pool, + &mut queue, + &ctx, + &membership_generations, + &mut last_activity, + ) { typing_channels.insert(channel_id, thread_tags); } } @@ -2890,7 +2898,13 @@ async fn tokio_main() -> Result<()> { // and drive fallback/dispatch as needed. sent_steer_pending } else if is_edit - && pool.native_steer_available(buzz_event.channel_id) + && pool.native_steer_available( + buzz_event.channel_id, + membership_generations + .get(&buzz_event.channel_id) + .copied() + .unwrap_or(0), + ) { let event_id = event_for_steer.id.to_hex(); let channel_id = buzz_event.channel_id; @@ -2969,7 +2983,7 @@ async fn tokio_main() -> Result<()> { } if pool_ready { for (channel_id, thread_tags) in - dispatch_pending(&mut pool, &mut queue, &ctx, &mut last_activity) + dispatch_pending(&mut pool, &mut queue, &ctx, &membership_generations, &mut last_activity) { typing_channels.insert(channel_id, thread_tags); } @@ -3069,7 +3083,7 @@ async fn tokio_main() -> Result<()> { } else if queue.has_flushable_work() { tracing::debug!("heartbeat_skipped_events"); for (channel_id, thread_tags) in - dispatch_pending(&mut pool, &mut queue, &ctx, &mut last_activity) + dispatch_pending(&mut pool, &mut queue, &ctx, &membership_generations, &mut last_activity) { typing_channels.insert(channel_id, thread_tags); } @@ -3143,6 +3157,7 @@ async fn tokio_main() -> Result<()> { *result, &mut heartbeat_in_flight, &removed_channels, + &membership_generations, &mut crash_history, &respawn_tx, &mut respawn_tasks, @@ -3158,6 +3173,7 @@ async fn tokio_main() -> Result<()> { &config, &mut heartbeat_in_flight, &removed_channels, + &membership_generations, &mut typing_channels, &mut crash_history, &respawn_tx, @@ -3167,9 +3183,13 @@ async fn tokio_main() -> Result<()> { { break; } - for (channel_id, thread_tags) in - dispatch_pending(&mut pool, &mut queue, &ctx, &mut last_activity) - { + for (channel_id, thread_tags) in dispatch_pending( + &mut pool, + &mut queue, + &ctx, + &membership_generations, + &mut last_activity, + ) { typing_channels.insert(channel_id, thread_tags); } } @@ -3182,6 +3202,7 @@ async fn tokio_main() -> Result<()> { join_error, &mut heartbeat_in_flight, &removed_channels, + &membership_generations, &mut typing_channels, &mut crash_history, &respawn_tx, @@ -3192,9 +3213,13 @@ async fn tokio_main() -> Result<()> { tracing::error!("all agents dead — exiting"); break; } - for (channel_id, thread_tags) in - dispatch_pending(&mut pool, &mut queue, &ctx, &mut last_activity) - { + for (channel_id, thread_tags) in dispatch_pending( + &mut pool, + &mut queue, + &ctx, + &membership_generations, + &mut last_activity, + ) { typing_channels.insert(channel_id, thread_tags); } } @@ -3235,6 +3260,7 @@ async fn tokio_main() -> Result<()> { channel_id, &event.id.to_hex(), &ctx, + &membership_generations, &mut last_activity, &mut typing_channels, ); @@ -3403,9 +3429,13 @@ async fn tokio_main() -> Result<()> { // tear down the in-flight task; on its completion the // queue drains. We still try here in case the in-flight // task has already returned. - for (channel_id, thread_tags) in - dispatch_pending(&mut pool, &mut queue, &ctx, &mut last_activity) - { + for (channel_id, thread_tags) in dispatch_pending( + &mut pool, + &mut queue, + &ctx, + &membership_generations, + &mut last_activity, + ) { typing_channels.insert(channel_id, thread_tags); } } @@ -3431,9 +3461,13 @@ async fn tokio_main() -> Result<()> { "ready", None, ); - for (channel_id, thread_tags) in - dispatch_pending(&mut pool, &mut queue, &ctx, &mut last_activity) - { + for (channel_id, thread_tags) in dispatch_pending( + &mut pool, + &mut queue, + &ctx, + &membership_generations, + &mut last_activity, + ) { typing_channels.insert(channel_id, thread_tags); } } @@ -3693,7 +3727,7 @@ fn try_native_steer( ack_tx, }; - match pool.send_steer(channel_id, request) { + match pool.send_steer(channel_id, membership_generation, request) { Ok(()) => { // Ordinary events are withheld after send. Edit preparation reserves // its event before leaving the main loop, so this is idempotent. @@ -3734,18 +3768,22 @@ fn try_native_steer( /// was withheld before asynchronous preparation, so release it, request the /// universal cancel+merge fallback, and immediately try dispatch in case the /// original turn already ended. +#[allow(clippy::too_many_arguments)] fn release_prepared_native_steer_fallback( pool: &mut AgentPool, queue: &mut EventQueue, channel_id: Uuid, event_id: &str, ctx: &Arc, + membership_generations: &HashMap, last_activity: &mut tokio::time::Instant, typing_channels: &mut HashMap, ) { queue.release_native_steer(channel_id, event_id); signal_in_flight_task(pool, channel_id, ControlSignal::Steer); - for (channel_id, thread_tags) in dispatch_pending(pool, queue, ctx, last_activity) { + for (channel_id, thread_tags) in + dispatch_pending(pool, queue, ctx, membership_generations, last_activity) + { typing_channels.insert(channel_id, thread_tags); } } @@ -3757,6 +3795,7 @@ fn dispatch_pending( pool: &mut AgentPool, queue: &mut EventQueue, ctx: &Arc, + membership_generations: &HashMap, last_activity: &mut tokio::time::Instant, ) -> Vec<(Uuid, ThreadTags)> { let mut dispatched_channels = Vec::new(); @@ -3830,6 +3869,12 @@ fn dispatch_pending( pool::TaskMeta { agent_index, channel_id: Some(channel_id), + membership_generation: Some( + membership_generations + .get(&channel_id) + .copied() + .unwrap_or(0), + ), turn_id, recoverable_batch, control_tx: Some(control_tx), @@ -3908,6 +3953,7 @@ fn handle_prompt_result( mut result: PromptResult, heartbeat_in_flight: &mut bool, removed_channels: &HashSet, + membership_generations: &HashMap, crash_history: &mut [SlotCircuit], respawn_tx: &mpsc::Sender, respawn_tasks: &mut tokio::task::JoinSet<()>, @@ -3916,12 +3962,22 @@ fn handle_prompt_result( ) -> LoopAction { let before = pool.task_map().len(); let agent_index = result.agent.index; - let successful_steer_deliveries = pool + let task_meta = pool .task_map() .values() - .find(|meta| meta.agent_index == agent_index) + .find(|meta| meta.agent_index == agent_index); + let successful_steer_deliveries = task_meta .map(|meta| meta.successful_steer_deliveries.clone()) .unwrap_or_default(); + let task_membership_is_current = task_meta + .and_then(|meta| meta.channel_id.zip(meta.membership_generation)) + .is_none_or(|(channel_id, generation)| { + membership_generations + .get(&channel_id) + .copied() + .unwrap_or(0) + == generation + }); pool.task_map_mut() .retain(|_, meta| meta.agent_index != agent_index); debug_assert_eq!(before, pool.task_map().len() + 1); @@ -3960,9 +4016,10 @@ fn handle_prompt_result( // every retry starts at attempt 1 — defeating exponential backoff and // dead-letter protection. if let Some(batch) = result.batch.take() { - // Don't requeue batches for channels the agent was removed from — - // those events are stale and should be silently dropped. - if !removed_channels.contains(&batch.channel_id) { + // Don't requeue batches from a removed membership generation. Re-add + // clears the current removal marker, but the task generation remains + // stale and its pre-removal work must not cross into the new epoch. + if !removed_channels.contains(&batch.channel_id) && task_membership_is_current { if matches!( result.outcome, PromptOutcome::Cancelled | PromptOutcome::CancelDrainTimeout(_) @@ -4056,9 +4113,10 @@ fn handle_prompt_result( tracing::debug!( channel_id = %batch.channel_id, events = batch.events.len(), - "dropping failed batch for removed channel" + task_membership_is_current, + "dropping failed batch from removed or stale membership generation" ); - hard_timeout_fate_suffix = Some(" — batch dropped (channel removed)"); + hard_timeout_fate_suffix = Some(" — batch dropped (channel membership changed)"); } } @@ -4301,6 +4359,7 @@ fn recover_panicked_agent( join_error: tokio::task::JoinError, heartbeat_in_flight: &mut bool, removed_channels: &HashSet, + membership_generations: &HashMap, typing_channels: &mut HashMap, crash_history: &mut [SlotCircuit], respawn_tx: &mpsc::Sender, @@ -4320,7 +4379,10 @@ fn recover_panicked_agent( // Requeue BEFORE mark_complete (same rationale as handle_prompt_result). if let Some(batch) = meta.recoverable_batch { if let Some(ch) = meta.channel_id { - if !removed_channels.contains(&ch) { + if !removed_channels.contains(&ch) + && meta.membership_generation + == Some(membership_generations.get(&ch).copied().unwrap_or(0)) + { // Dead-letter on exhaustion is logged inside requeue(); a // panic path has no outcome to report, so no notice here. let _ = queue.requeue(batch); @@ -4402,6 +4464,7 @@ fn drain_ready_join_results( config: &Config, heartbeat_in_flight: &mut bool, removed_channels: &HashSet, + membership_generations: &HashMap, typing_channels: &mut HashMap, crash_history: &mut [SlotCircuit], respawn_tx: &mpsc::Sender, @@ -4418,6 +4481,7 @@ fn drain_ready_join_results( join_error, heartbeat_in_flight, removed_channels, + membership_generations, typing_channels, crash_history, respawn_tx, @@ -4473,6 +4537,7 @@ fn dispatch_heartbeat( pool::TaskMeta { agent_index, channel_id: None, + membership_generation: None, turn_id, recoverable_batch: None, control_tx: None, @@ -5251,6 +5316,7 @@ mod owner_control_command_tests { pool::TaskMeta { agent_index: 0, channel_id: Some(channel_id), + membership_generation: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: Some(control_tx), @@ -7052,6 +7118,7 @@ mod error_outcome_emission_tests { crate::pool::TaskMeta { agent_index: 0, channel_id: Some(channel_id), + membership_generation: None, turn_id: "test-turn-id".into(), recoverable_batch: None, control_tx: None, @@ -7091,6 +7158,7 @@ mod error_outcome_emission_tests { result, &mut heartbeat_in_flight, &removed_channels, + &HashMap::new(), &mut crash_history, &respawn_tx, &mut respawn_tasks, @@ -7124,6 +7192,7 @@ mod error_outcome_emission_tests { crate::pool::TaskMeta { agent_index: 0, channel_id: Some(channel_id), + membership_generation: None, turn_id: "test-turn-id".into(), recoverable_batch: None, control_tx: None, @@ -7163,6 +7232,7 @@ mod error_outcome_emission_tests { result, &mut heartbeat_in_flight, &removed_channels, + &HashMap::new(), &mut crash_history, &respawn_tx, &mut respawn_tasks, @@ -7239,6 +7309,7 @@ mod error_outcome_emission_tests { crate::pool::TaskMeta { agent_index: 0, channel_id: Some(channel_id), + membership_generation: None, turn_id: "test-turn-id".into(), recoverable_batch: None, control_tx: None, @@ -7277,6 +7348,7 @@ mod error_outcome_emission_tests { result, &mut heartbeat_in_flight, &removed_channels, + &HashMap::new(), &mut crash_history, &respawn_tx, &mut respawn_tasks, @@ -7304,6 +7376,7 @@ mod error_outcome_emission_tests { crate::pool::TaskMeta { agent_index: 0, channel_id: None, + membership_generation: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -7340,6 +7413,7 @@ mod error_outcome_emission_tests { result, &mut heartbeat_in_flight, &removed_channels, + &HashMap::new(), &mut crash_history, &respawn_tx, &mut respawn_tasks, @@ -7381,6 +7455,7 @@ mod error_outcome_emission_tests { crate::pool::TaskMeta { agent_index: 0, channel_id: Some(channel_id), + membership_generation: None, turn_id: "panic-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -7413,6 +7488,7 @@ mod error_outcome_emission_tests { join_error, &mut heartbeat_in_flight, &removed_channels, + &HashMap::new(), &mut typing_channels, &mut crash_history, &respawn_tx, @@ -7474,6 +7550,7 @@ mod error_outcome_emission_tests { crate::pool::TaskMeta { agent_index: 0, channel_id: None, + membership_generation: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -7507,6 +7584,7 @@ mod error_outcome_emission_tests { result, &mut heartbeat_in_flight, &removed_channels, + &HashMap::new(), &mut crash_history, &respawn_tx, &mut respawn_tasks, @@ -7566,6 +7644,7 @@ mod error_outcome_emission_tests { crate::pool::TaskMeta { agent_index: 0, channel_id: None, + membership_generation: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -7598,6 +7677,7 @@ mod error_outcome_emission_tests { result, &mut heartbeat_in_flight, &removed_channels, + &HashMap::new(), &mut crash_history, &respawn_tx, &mut respawn_tasks, @@ -7672,6 +7752,7 @@ mod error_outcome_emission_tests { crate::pool::TaskMeta { agent_index: 0, channel_id: None, + membership_generation: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -7704,6 +7785,7 @@ mod error_outcome_emission_tests { result, &mut heartbeat_in_flight, &removed_channels, + &HashMap::new(), &mut crash_history, &respawn_tx, &mut respawn_tasks, @@ -7749,6 +7831,7 @@ mod error_outcome_emission_tests { crate::pool::TaskMeta { agent_index: 0, channel_id: None, + membership_generation: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -7796,6 +7879,7 @@ mod error_outcome_emission_tests { result, &mut heartbeat_in_flight, &removed_channels, + &HashMap::new(), &mut crash_history, &respawn_tx, &mut respawn_tasks, @@ -7844,6 +7928,7 @@ mod error_outcome_emission_tests { crate::pool::TaskMeta { agent_index: 0, channel_id: None, + membership_generation: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -7890,6 +7975,7 @@ mod error_outcome_emission_tests { result, &mut heartbeat_in_flight, &removed_channels, + &HashMap::new(), &mut crash_history, &respawn_tx, &mut respawn_tasks, @@ -7961,6 +8047,7 @@ mod error_outcome_emission_tests { crate::pool::TaskMeta { agent_index: 0, channel_id: None, + membership_generation: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -8006,6 +8093,7 @@ mod error_outcome_emission_tests { result, &mut heartbeat_in_flight, &removed_channels, + &HashMap::new(), &mut crash_history, &respawn_tx, &mut respawn_tasks, @@ -8101,6 +8189,7 @@ mod error_outcome_emission_tests { crate::pool::TaskMeta { agent_index: 0, channel_id: None, + membership_generation: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -8139,6 +8228,7 @@ mod error_outcome_emission_tests { result, &mut heartbeat_in_flight, &removed_channels, + &HashMap::new(), &mut crash_history, &respawn_tx, &mut respawn_tasks, @@ -8290,6 +8380,7 @@ mod error_outcome_emission_tests { crate::pool::TaskMeta { agent_index: 0, channel_id: None, + membership_generation: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -8322,6 +8413,7 @@ mod error_outcome_emission_tests { result, &mut heartbeat_in_flight, &removed_channels, + &HashMap::new(), &mut crash_history, &respawn_tx, &mut respawn_tasks, @@ -8376,6 +8468,7 @@ mod error_outcome_emission_tests { crate::pool::TaskMeta { agent_index: 0, channel_id: None, + membership_generation: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -8408,6 +8501,7 @@ mod error_outcome_emission_tests { result, &mut heartbeat_in_flight, &removed_channels, + &HashMap::new(), &mut crash_history, &respawn_tx, &mut respawn_tasks, @@ -8766,6 +8860,7 @@ mod native_edit_membership_lifecycle_tests { channel_id, &event_id, &prompt_context(), + &HashMap::new(), &mut last_activity, &mut typing_channels, ); diff --git a/crates/buzz-acp/src/pool.rs b/crates/buzz-acp/src/pool.rs index 9e561a2776..7be3a2bfd4 100644 --- a/crates/buzz-acp/src/pool.rs +++ b/crates/buzz-acp/src/pool.rs @@ -60,6 +60,9 @@ pub struct SuccessfulSteerDelivery { pub struct TaskMeta { pub agent_index: usize, pub channel_id: Option, + /// Membership generation in which this channel task was dispatched. + /// Heartbeat tasks have no channel generation. + pub membership_generation: Option, /// Identifies terminal events when the task panics before returning a result. pub turn_id: String, /// Clone of batch for Queue mode panic recovery. @@ -704,23 +707,31 @@ impl AgentPool { &mut self.task_map } - /// Return whether the current channel task can accept a native steer now. + /// Return whether the current-generation channel task can accept a native + /// steer now. + /// + /// Membership generation is part of task ownership: after remove/re-add, + /// a still-running task from the prior generation must not accept new work. /// /// This is a main-loop preflight for edit enrichment: avoid withholding an /// edit and performing REST work when the task has no steer transport or /// its capacity-one request slot is already occupied. Availability can /// still change while enrichment runs, so callers must retain the normal /// send-time fallback. - pub fn native_steer_available(&self, channel_id: Uuid) -> bool { + pub fn native_steer_available(&self, channel_id: Uuid, membership_generation: u64) -> bool { self.task_map .values() - .find(|meta| meta.channel_id == Some(channel_id)) + .find(|meta| { + meta.channel_id == Some(channel_id) + && meta.membership_generation == Some(membership_generation) + }) .and_then(|meta| meta.steer_tx.as_ref()) .is_some_and(|tx| !tx.is_closed() && tx.capacity() > 0) } /// Try to send a goose-native steer request to the in-flight task for - /// `channel_id`. + /// `channel_id` in `membership_generation`. A task from an earlier + /// remove/re-add generation is treated as completed for this request. /// /// Returns `Ok(())` if the request was accepted by the read loop's /// receiver (capacity-1 mpsc; one slot is the single in-flight steer @@ -744,12 +755,16 @@ impl AgentPool { pub fn send_steer( &mut self, channel_id: Uuid, + membership_generation: u64, request: SteerRequest, ) -> Result<(), SteerError> { let meta = self .task_map .values_mut() - .find(|m| m.channel_id == Some(channel_id)) + .find(|m| { + m.channel_id == Some(channel_id) + && m.membership_generation == Some(membership_generation) + }) .ok_or(SteerError::PromptCompleted)?; let tx = meta .steer_tx @@ -4495,6 +4510,43 @@ mod tests { } } + #[tokio::test] + async fn native_steer_rejects_task_from_prior_membership_generation() { + let channel_id = Uuid::new_v4(); + let old_generation = 4; + let current_generation = 5; + let mut pool = AgentPool::from_slots(vec![None]); + let task_id = pool.join_set.spawn(async {}).id(); + let (steer_tx, mut steer_rx) = tokio::sync::mpsc::channel(1); + pool.task_map_mut().insert( + task_id, + TaskMeta { + agent_index: 0, + channel_id: Some(channel_id), + membership_generation: Some(old_generation), + turn_id: "pre-removal-turn".into(), + recoverable_batch: None, + control_tx: None, + steer_tx: Some(steer_tx), + successful_steer_deliveries: HashSet::new(), + }, + ); + + assert!(pool.native_steer_available(channel_id, old_generation)); + assert!(!pool.native_steer_available(channel_id, current_generation)); + + let (ack_tx, _ack_rx) = tokio::sync::oneshot::channel(); + let request = SteerRequest { + prompt_blocks: vec!["post-readd work".into()], + ack_tx, + }; + assert!(matches!( + pool.send_steer(channel_id, current_generation, request), + Err(SteerError::PromptCompleted) + )); + assert!(steer_rx.try_recv().is_err()); + } + #[test] fn public_session_forwards_channel_origin_to_mcp() { let channel_id = Uuid::new_v4();