jerry/buzz
1
0
Fork 0
mirror of https://github.com/block/buzz.git synced 2026-11-01 23:11:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(acp): fence tasks by membership generation

Browse Source 
Signed-off-by: Larry <8cf5a83f590ec0955b11647d1c88f796a98e088c30a492c58e0e46c3026ae7a4@buzz.block.builderlab.xyz>
This commit is contained in:
Larry
2026-08-12 17:57:24 -04:00
parent 9811b89186
commit 2e470fdbe0
2 changed files with 183 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

157
crates/buzz-acp/src/lib.rs
View File

@@ -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<PromptContext>,
membership_generations: &HashMap<Uuid, u64>,
last_activity: &mut tokio::time::Instant,
typing_channels: &mut HashMap<Uuid, ThreadTags>,
) {
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<PromptContext>,
membership_generations: &HashMap<Uuid, u64>,
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<Uuid>,
membership_generations: &HashMap<Uuid, u64>,
crash_history: &mut [SlotCircuit],
respawn_tx: &mpsc::Sender<RespawnResult>,
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<Uuid>,
membership_generations: &HashMap<Uuid, u64>,
typing_channels: &mut HashMap<Uuid, ThreadTags>,
crash_history: &mut [SlotCircuit],
respawn_tx: &mpsc::Sender<RespawnResult>,
@@ -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<Uuid>,
membership_generations: &HashMap<Uuid, u64>,
typing_channels: &mut HashMap<Uuid, ThreadTags>,
crash_history: &mut [SlotCircuit],
respawn_tx: &mpsc::Sender<RespawnResult>,
@@ -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,
);

62
crates/buzz-acp/src/pool.rs
View File

@@ -60,6 +60,9 @@ pub struct SuccessfulSteerDelivery {
pub struct TaskMeta {
pub agent_index: usize,
pub channel_id: Option<Uuid>,
/// Membership generation in which this channel task was dispatched.
/// Heartbeat tasks have no channel generation.
pub membership_generation: Option<u64>,
/// 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();