From f442707444a03b92af8879173ca7a8facb422cfb Mon Sep 17 00:00:00 2001 From: npub1zurdm6fx3kksz8f8d8y4js3mvjyu3rshlpqejr8wa8hta5tsqqeshmes7w <1706dde9268dad011d2769c959423b6489c88e17f841990ceee9eebed1700033@sprout-oss.stage.blox.sqprod.co> Date: Sat, 18 Jul 2026 11:51:36 -0700 Subject: [PATCH] refactor(acp): consolidate top-level-sessions experiment maintenance surfaces Reduce the ongoing cost of the top-level-sessions experiment by collapsing its scattered maintenance surfaces without changing behavior: - pool: replace eager session_owners reconciliation (5 lifecycle call sites + retains in return_agent / invalidate_channel_sessions / model-switch paths) with a single lazy prune_session_owners() at the one read site in try_claim. Future slot-lifecycle paths can no longer forget the purge; the invariant converges at one choke point. - queue: extract pop_cancelled() / dispatch_cancelled() helpers so the three copy-pasted cancelled-batch blocks in flush_next become one-liners. - session key: drop the top_level_sessions bool param and dead thread-tag fallback from conversation_session_key; remove the redundant PromptContext.top_level_sessions field. The key is now purely data-driven from batch.conversation_root, which is only populated at queue-push when the experiment is enabled. Net -59 lines. cargo test -p buzz-acp: 530 passed; clippy and fmt clean. Disabled-mode paths degenerate identically (root=None everywhere). Co-authored-by: npub1zurdm6fx3kksz8f8d8y4js3mvjyu3rshlpqejr8wa8hta5tsqqeshmes7w <1706dde9268dad011d2769c959423b6489c88e17f841990ceee9eebed1700033@sprout-oss.stage.blox.sqprod.co> Signed-off-by: npub1zurdm6fx3kksz8f8d8y4js3mvjyu3rshlpqejr8wa8hta5tsqqeshmes7w <1706dde9268dad011d2769c959423b6489c88e17f841990ceee9eebed1700033@sprout-oss.stage.blox.sqprod.co> --- crates/buzz-acp/src/lib.rs | 16 +++-- crates/buzz-acp/src/pool.rs | 109 +++++++++++++---------------------- crates/buzz-acp/src/queue.rs | 108 +++++++++++++--------------------- 3 files changed, 87 insertions(+), 146 deletions(-) diff --git a/crates/buzz-acp/src/lib.rs b/crates/buzz-acp/src/lib.rs index 7f94ede7b..f291c0edd 100644 --- a/crates/buzz-acp/src/lib.rs +++ b/crates/buzz-acp/src/lib.rs @@ -1475,7 +1475,6 @@ async fn tokio_main() -> Result<()> { .as_deref() .and_then(|hex| nostr::PublicKey::from_hex(hex).ok()), memory_enabled: config.memory_enabled, - top_level_sessions: config.top_level_sessions, harness_name: crate::config::normalize_agent_command_identity(&config.agent_command), }); @@ -1633,7 +1632,6 @@ async fn tokio_main() -> Result<()> { if !slot.can_refill() { continue; } - pool.clear_slot_session_owners(idx); slot.respawn_in_flight = true; tracing::info!(agent = idx, "slot refill: spawning background respawn"); let cmd = config.agent_command.clone(); @@ -2670,7 +2668,7 @@ fn dispatch_pending( .last() .map(|event| queue::parse_thread_tags(&event.event)) .unwrap_or_default(); - let session_key = pool::conversation_session_key(&batch, ctx.top_level_sessions); + let session_key = pool::conversation_session_key(&batch); let affinity_hit = pool.has_session_for(&session_key); let mut agent = match pool.try_claim(Some(&session_key)) { Some(a) => a, @@ -2953,7 +2951,6 @@ fn handle_prompt_result( emit_turn_error(&death_message, None); let index = result.agent.index; - pool.clear_slot_session_owners(index); let slot_history = &mut crash_history[index]; if !spawn_respawn_task( result.agent, @@ -2994,7 +2991,6 @@ fn handle_prompt_result( emit_turn_error(&death_message, None); let index = result.agent.index; - pool.clear_slot_session_owners(index); let slot_history = &mut crash_history[index]; if !spawn_respawn_task( result.agent, @@ -3060,7 +3056,6 @@ fn handle_prompt_result( emit_turn_error(&e.to_string(), error_code); let index = result.agent.index; - pool.clear_slot_session_owners(index); let slot_history = &mut crash_history[index]; if !spawn_respawn_task( result.agent, @@ -3112,7 +3107,6 @@ fn recover_panicked_agent( return; }; let i = meta.agent_index; - pool.clear_slot_session_owners(i); // Requeue BEFORE mark_complete (same rationale as handle_prompt_result). if let Some(batch) = meta.recoverable_batch { @@ -4502,8 +4496,12 @@ mod error_outcome_emission_tests { pool.return_agent(owner); } pool.agents_mut()[0] = None; - assert_eq!(pool.clear_slot_session_owners(0), 2); - assert_eq!(pool.try_claim(Some(&keys[0])).unwrap().index, 1); + // Every reservation owned by the dead slot is pruned lazily on the + // next claim; both roots recreate on the surviving slot. + let first_claim = pool.try_claim(Some(&keys[0])).unwrap(); + assert_eq!(first_claim.index, 1); + pool.return_agent(first_claim); + assert_eq!(pool.try_claim(Some(&keys[1])).unwrap().index, 1); } #[tokio::test] diff --git a/crates/buzz-acp/src/pool.rs b/crates/buzz-acp/src/pool.rs index 734b8c58c..25c41a351 100644 --- a/crates/buzz-acp/src/pool.rs +++ b/crates/buzz-acp/src/pool.rs @@ -555,8 +555,6 @@ pub struct PromptContext { /// `[Agent Memory — core]` section. On by default; disabled via /// `--no-memory` / `BUZZ_ACP_NO_MEMORY`. pub memory_enabled: bool, - /// Whether ACP sessions are isolated by human conversation root. - pub top_level_sessions: bool, /// Harness identity string for NIP-AM `harness` field. Derived from the /// configured `agent_command` at startup (e.g. `"goose"`, `"buzz-agent"`). pub harness_name: String, @@ -595,20 +593,11 @@ impl AgentPool { // wait rather than creating a duplicate provider session in another slot. if let Some(key) = session_key { if key.root_event_id.is_some() { + self.prune_session_owners(); if let Some(&owner) = self.session_owners.get(key) { - let owner_busy = self.task_map.values().any(|meta| meta.agent_index == owner); - let owner_retains_session = self - .agents - .get(owner) - .and_then(Option::as_ref) - .is_some_and(|agent| agent.state.contains_session(key)); - if owner_busy || owner_retains_session { - return self.agents.get_mut(owner).and_then(Option::take); - } - // The slot died, was replaced, or evicted/invalidated this - // root while idle. Drop the stale reservation so a live - // slot can recreate the provider session. - self.session_owners.remove(key); + // Surviving the prune means the owner is busy (wait by + // returning None) or idle and retaining the session. + return self.agents.get_mut(owner).and_then(Option::take); } } let idx = self.agents.iter().position(|slot| { @@ -631,6 +620,23 @@ impl AgentPool { self.agents[idx].take() } + /// Drop root reservations that no longer bind: the owning slot is neither + /// busy (checked out on a task) nor idle-retaining the session. Runs before + /// every reservation read, so slot death, replacement, LRU eviction, and + /// channel invalidation all converge here instead of requiring explicit + /// cleanup calls at each lifecycle site. + fn prune_session_owners(&mut self) { + let agents = &self.agents; + let task_map = &self.task_map; + self.session_owners.retain(|key, owner| { + task_map.values().any(|meta| meta.agent_index == *owner) + || agents + .get(*owner) + .and_then(Option::as_ref) + .is_some_and(|agent| agent.state.contains_session(key)) + }); + } + /// Return an agent to its slot after a task completes. pub fn return_agent(&mut self, agent: OwnedAgent) { let idx = agent.index; @@ -644,11 +650,6 @@ impl AgentPool { "BUG: return_agent called for slot {idx} which is already occupied — overwriting" ); } - // Reconcile reservations whenever a slot returns. This clears roots - // evicted by the per-slot LRU, explicitly invalidated while checked - // out, or lost when a dead slot is replaced. - self.session_owners - .retain(|key, owner| *owner != idx || agent.state.contains_session(key)); self.agents[idx] = Some(agent); } @@ -759,14 +760,6 @@ impl AgentPool { &mut self.agents } - /// Remove every root reservation owned by a slot that died or is about to - /// be replaced. Returns the number removed for diagnostics and tests. - pub fn clear_slot_session_owners(&mut self, index: usize) -> usize { - let before = self.session_owners.len(); - self.session_owners.retain(|_, owner| *owner != index); - before - self.session_owners.len() - } - /// Remove the session for `channel_id` from all idle agents. /// /// Called when the agent is removed from a channel — stale sessions @@ -784,8 +777,8 @@ impl AgentPool { } } } - self.session_owners - .retain(|key, _| key.channel_id != channel_id); + // Freed root reservations are dropped lazily by prune_session_owners + // on the next root claim. count } @@ -841,33 +834,24 @@ impl AgentPool { agent.model_overridden = true; agent.state.invalidate_channel(&channel_id); } - self.session_owners - .retain(|key, _| key.channel_id != channel_id); + // Freed reservations are dropped lazily by prune_session_owners on the + // next root claim. IdleSwitchResult::Switched } /// Apply a busy channel model switch to every currently idle slot retaining /// another root for that channel. The checked-out owner receives the - /// control signal separately and invalidates its own channel state. + /// control signal separately and invalidates its own channel state; its + /// reservation survives pruning (busy), so a reply arriving during + /// cancellation still waits for it instead of falling through. pub fn switch_other_idle_channel_roots(&mut self, channel_id: Uuid, model_id: &str) { - let mut invalidated_slots = Vec::new(); - for (index, agent) in self.agents.iter_mut().enumerate() { - let Some(agent) = agent.as_mut() else { - continue; - }; + for agent in self.agents.iter_mut().flatten() { if agent.state.has_channel_state(&channel_id) { agent.desired_model = Some(model_id.to_string()); agent.model_overridden = true; agent.state.invalidate_channel(&channel_id); - invalidated_slots.push(index); } } - // Preserve the checked-out root owner's reservation until that task - // returns with its own channel state invalidated. Otherwise a reply - // arriving during cancellation could fall through to an idle slot. - self.session_owners.retain(|key, owner| { - key.channel_id != channel_id || !invalidated_slots.contains(owner) - }); } } @@ -1356,24 +1340,13 @@ fn send_prompt_result( }); } -pub(crate) fn conversation_session_key( - batch: &FlushBatch, - top_level_sessions: bool, -) -> ConversationSessionKey { - let root_event_id = top_level_sessions - .then(|| { - batch.conversation_root.clone().or_else(|| { - batch.events.last().map(|event| { - crate::queue::parse_thread_tags(&event.event) - .root_event_id - .unwrap_or_else(|| event.event.id.to_hex()) - }) - }) - }) - .flatten(); +/// Session identity for a batch. `FlushBatch::conversation_root` is populated +/// at queue-push time only when the top-level-sessions experiment is enabled, +/// so this is purely data-driven: no root means the legacy channel-scoped key. +pub(crate) fn conversation_session_key(batch: &FlushBatch) -> ConversationSessionKey { ConversationSessionKey { channel_id: batch.channel_id, - root_event_id, + root_event_id: batch.conversation_root.clone(), } } @@ -1400,7 +1373,7 @@ pub async fn run_prompt_task( ) { // Is this a channel prompt or a heartbeat? let source = match &batch { - Some(b) => PromptSource::Channel(conversation_session_key(b, ctx.top_level_sessions)), + Some(b) => PromptSource::Channel(conversation_session_key(b)), None => PromptSource::Heartbeat, }; let observer_channel_id = match &source { @@ -2679,7 +2652,8 @@ async fn fetch_conversation_context( // Experiment: seed only a fresh top-level conversation. Replies use their // thread context, and existing sessions already carry their own history. - if ctx.top_level_sessions && is_new_session { + // `conversation_root` is only ever set when the experiment is enabled. + if batch.conversation_root.is_some() && is_new_session { let triggering_ids: HashSet = batch .events .iter() @@ -4581,18 +4555,18 @@ mod tests { } #[test] - fn test_conversation_session_key_is_channel_scoped_when_disabled() { + fn test_conversation_session_key_is_channel_scoped_without_root() { let batch = one_event_batch(Uuid::new_v4()); - let key = conversation_session_key(&batch, false); + let key = conversation_session_key(&batch); assert_eq!(key.channel_id, batch.channel_id); assert_eq!(key.root_event_id, None); } #[test] - fn test_conversation_session_key_uses_explicit_root_when_enabled() { + fn test_conversation_session_key_uses_batch_root_when_present() { let mut batch = one_event_batch(Uuid::new_v4()); batch.conversation_root = Some("root-a".into()); - let key = conversation_session_key(&batch, true); + let key = conversation_session_key(&batch); assert_eq!(key.channel_id, batch.channel_id); assert_eq!(key.root_event_id.as_deref(), Some("root-a")); } @@ -5827,7 +5801,6 @@ mod tests { agent_owner_pubkey: owner_pubkey, memory_enabled: false, harness_name: "goose".to_string(), - top_level_sessions: false, } } diff --git a/crates/buzz-acp/src/queue.rs b/crates/buzz-acp/src/queue.rs index 8e71b2952..a07ce5278 100644 --- a/crates/buzz-acp/src/queue.rs +++ b/crates/buzz-acp/src/queue.rs @@ -293,40 +293,13 @@ impl EventQueue { .keys() .find(|id| !self.in_flight_channels.contains(id)) .copied(); - match cancelled_id { - Some(id) => { - // No new events to merge — re-dispatch the oldest cancelled - // batch unchanged under its original conversation root. - let cancelled = self - .cancelled_batches - .get_mut(&id) - .and_then(VecDeque::pop_front) - .expect("cancelled channel must have a batch"); - if self - .cancelled_batches - .get(&id) - .is_some_and(VecDeque::is_empty) - { - self.cancelled_batches.remove(&id); - } - self.in_flight_channels.insert(id); - self.in_flight_deadlines - .insert(id, now + self.in_flight_deadline); - self.in_flight_batch_sizes - .insert(id, cancelled.events.len()); - return Some(FlushBatch { - channel_id: id, - conversation_root: cancelled.conversation_root, - events: cancelled.events, - cancelled_events: vec![], - cancel_reason: Some(cancelled.reason), - }); - } - None => return None, - } + return cancelled_id.map(|id| self.dispatch_cancelled(id, now)); } }; + // Re-dispatch a cancelled batch whose conversation root differs from + // the queued head before starting the queued work — its root was + // interrupted first and must not merge into another root's turn. let queued_root = self .queues .get(&channel_id) @@ -336,32 +309,9 @@ impl EventQueue { .cancelled_batches .get(&channel_id) .and_then(|batches| batches.front()) - .and_then(|batch| batch.conversation_root.clone()); - if self.cancelled_batches.contains_key(&channel_id) && cancelled_root != queued_root { - let cancelled = self - .cancelled_batches - .get_mut(&channel_id) - .and_then(VecDeque::pop_front) - .expect("cancelled channel must have a batch"); - if self - .cancelled_batches - .get(&channel_id) - .is_some_and(VecDeque::is_empty) - { - self.cancelled_batches.remove(&channel_id); - } - self.in_flight_channels.insert(channel_id); - self.in_flight_deadlines - .insert(channel_id, now + self.in_flight_deadline); - self.in_flight_batch_sizes - .insert(channel_id, cancelled.events.len()); - return Some(FlushBatch { - channel_id, - conversation_root: cancelled.conversation_root, - events: cancelled.events, - cancelled_events: vec![], - cancel_reason: Some(cancelled.reason), - }); + .map(|batch| batch.conversation_root.clone()); + if cancelled_root.is_some_and(|root| root != queued_root) { + return Some(self.dispatch_cancelled(channel_id, now)); } // Drain up to MAX_BATCH_EVENTS; leave any remainder in the queue. @@ -401,18 +351,7 @@ impl EventQueue { self.in_flight_batch_sizes.insert(channel_id, events.len()); // Merge only a cancelled batch for the same conversation root. - let cancelled = self - .cancelled_batches - .get_mut(&channel_id) - .and_then(VecDeque::pop_front); - if self - .cancelled_batches - .get(&channel_id) - .is_some_and(VecDeque::is_empty) - { - self.cancelled_batches.remove(&channel_id); - } - let (cancelled_events, cancel_reason) = cancelled.map_or_else( + let (cancelled_events, cancel_reason) = self.pop_cancelled(channel_id).map_or_else( || (vec![], None), |batch| (batch.events, Some(batch.reason)), ); @@ -426,6 +365,37 @@ impl EventQueue { }) } + /// Pop the oldest cancelled batch for `channel_id`, dropping the map entry + /// once its queue is empty. `None` if the channel has no cancelled batches. + fn pop_cancelled(&mut self, channel_id: Uuid) -> Option { + let batches = self.cancelled_batches.get_mut(&channel_id)?; + let cancelled = batches.pop_front(); + if batches.is_empty() { + self.cancelled_batches.remove(&channel_id); + } + cancelled + } + + /// Re-dispatch the oldest cancelled batch for `channel_id` unchanged under + /// its original conversation root, marking the channel in-flight. + fn dispatch_cancelled(&mut self, channel_id: Uuid, now: Instant) -> FlushBatch { + let cancelled = self + .pop_cancelled(channel_id) + .expect("cancelled channel must have a batch"); + self.in_flight_channels.insert(channel_id); + self.in_flight_deadlines + .insert(channel_id, now + self.in_flight_deadline); + self.in_flight_batch_sizes + .insert(channel_id, cancelled.events.len()); + FlushBatch { + channel_id, + conversation_root: cancelled.conversation_root, + events: cancelled.events, + cancelled_events: vec![], + cancel_reason: Some(cancelled.reason), + } + } + /// Mark the prompt for `channel_id` as complete. /// /// Removes the channel from `in_flight_channels` and `in_flight_deadlines`.