diff --git a/crates/buzz-acp/src/lib.rs b/crates/buzz-acp/src/lib.rs index 17d3030f635..b80e5ca6fc5 100644 --- a/crates/buzz-acp/src/lib.rs +++ b/crates/buzz-acp/src/lib.rs @@ -2976,6 +2976,7 @@ async fn tokio_main() -> Result<()> { Panic(tokio::task::JoinError), SteerAck(SteerAckEvent), Wake(u32, Result), + HoldDeadline, } loop { @@ -3117,6 +3118,8 @@ async fn tokio_main() -> Result<()> { } // Borrow result_rx and join_set simultaneously via split-borrow helper. + pool.retain_held_scopes(|scope| queue.has_pending_scope(scope)); + let hold_deadline = pool.next_hold_deadline(pool::HOLD_BUSY_OWNER_TIMEOUT); let pool_event: Option = { let (result_rx, join_set) = pool.rx_and_join_set(); tokio::select! { @@ -3159,6 +3162,9 @@ async fn tokio_main() -> Result<()> { _ => std::future::pending().await, } } => None, + _ = pool::AgentPool::wait_for_hold_deadline(hold_deadline), if pool_ready => { + Some(PoolEvent::HoldDeadline) + }, Some(Err(error)) = wake_tasks.join_next(), if !wake_tasks.is_empty() => { if let Some(attempt) = pool_lifecycle.waking_attempt() { let message = format!("pool wake task failed: {error}"); @@ -3992,6 +3998,21 @@ async fn tokio_main() -> Result<()> { } } } + Some(PoolEvent::HoldDeadline) => { + // A held thread must make progress even when every unrelated + // relay/timer source is quiet. The deadline is derived from + // the pool's first-held stamp, so this dispatch observes + // `ForkAfterHold` and claims an idle worker immediately. + for (scope, thread_tags) in dispatch_pending( + &mut pool, + &mut queue, + &ctx, + &mut last_activity, + observer.as_ref(), + ) { + typing_channels.insert(scope, thread_tags); + } + } None => {} // relay/heartbeat/shutdown branches handled inline above } } @@ -4379,7 +4400,7 @@ fn dispatch_pending( let mut held: Vec = Vec::new(); // One clock read for the whole cycle so every batch's bounded-hold window is // measured against the same instant. - let now = std::time::Instant::now(); + let now = tokio::time::Instant::now(); loop { let batch = match queue.flush_next() { Some(b) => b, @@ -4394,7 +4415,8 @@ fn dispatch_pending( // so an active channel cannot starve a sibling channel on a shared // worker. A held thread that outwaits the window forks a fresh session // rather than starve behind an unbounded turn. - match pool.hold_decision(&scope, now, pool::HOLD_BUSY_OWNER_TIMEOUT) { + let forked_after_hold = match pool.hold_decision(&scope, now, pool::HOLD_BUSY_OWNER_TIMEOUT) + { pool::HoldDecision::Hold { held_for, owner_index, @@ -4425,31 +4447,9 @@ fn dispatch_pending( pool::HoldDecision::ForkAfterHold { held_for, owner_index, - } => { - tracing::warn!( - channel = %channel_id, - scope = %scope.telemetry_label(), - owner_index, - held_for_secs = held_for.as_secs_f64(), - "busy-owner hold expired — forking fresh session on an idle worker" - ); - if let Some(observer) = observer { - observer.emit( - "busy_owner_hold_forked", - None, - &observer::context_for(Some(channel_id), None, None), - serde_json::json!({ - "scope": scope.telemetry_label(), - "ownerIndex": owner_index, - "heldForSecs": held_for.as_secs_f64(), - }), - ); - } - // Fall through to try_claim below (fork); record_scope_owner - // reassigns ownership to the new worker automatically. - } - pool::HoldDecision::Dispatch => {} - } + } => Some((held_for, owner_index)), + pool::HoldDecision::Dispatch => None, + }; let typing_scope = batch .events .last() @@ -4469,6 +4469,32 @@ fn dispatch_pending( break; } }; + // Consume a bounded hold only after a worker was actually claimed. + // If every slot is checked out, the expired stamp remains sticky and + // the worker-return event retries immediately instead of waiting for a + // fresh timeout window. + pool.clear_hold(&scope); + if let Some((held_for, owner_index)) = forked_after_hold { + tracing::warn!( + channel = %channel_id, + scope = %scope.telemetry_label(), + owner_index, + held_for_secs = held_for.as_secs_f64(), + "busy-owner hold expired — forking fresh session on an idle worker" + ); + if let Some(observer) = observer { + observer.emit( + "busy_owner_hold_forked", + None, + &observer::context_for(Some(channel_id), None, None), + serde_json::json!({ + "scope": scope.telemetry_label(), + "ownerIndex": owner_index, + "heldForSecs": held_for.as_secs_f64(), + }), + ); + } + } tracing::debug!(agent = agent.index, channel = %channel_id, scope = %scope.telemetry_label(), affinity_hit, "agent_claimed"); let recoverable_batch = match ctx.dedup_mode { @@ -4499,6 +4525,14 @@ fn dispatch_pending( let turn_id = Uuid::new_v4().to_string(); let task_turn_id = turn_id.clone(); + // Assign ownership before moving the worker into the task. If this is + // a bounded-hold fork, the new generation immediately invalidates the + // prior busy worker's copy when that worker eventually returns. + let owner_generation = pool.record_scope_owner(scope.clone(), agent.index); + agent + .state + .set_scope_owner_generation(scope.clone(), owner_generation); + let abort_handle = pool.join_set.spawn(async move { pool::run_prompt_task( agent, @@ -4525,9 +4559,6 @@ fn dispatch_pending( successful_steer_deliveries: HashSet::new(), }, ); - // Record this worker as the scope's session owner so a later dispatch - // while it is busy holds instead of forking a duplicate session. - pool.record_scope_owner(scope.clone(), agent_index); dispatched_channels.push((scope, typing_scope)); *last_activity = tokio::time::Instant::now(); } @@ -6199,7 +6230,7 @@ mod owner_control_command_tests { // The bounded hold decision stamps A's first-held time, then forks once // the window elapses rather than starving behind the busy owner. - let now = std::time::Instant::now(); + let now = tokio::time::Instant::now(); assert!( matches!( pool.hold_decision(&ta, now, pool::HOLD_BUSY_OWNER_TIMEOUT), @@ -6220,9 +6251,10 @@ mod owner_control_command_tests { "elapsed window => fork on an idle worker" ); assert!( - !pool.held_since_contains(&ta), - "fork clears the first-held stamp" + pool.held_since_contains(&ta), + "expired hold stays sticky until an idle worker is claimed" ); + pool.clear_hold(&ta); // A conversation scope never holds even with a busy recorded owner — // this is the cross-channel head-of-line-blocking regression guard. @@ -6254,6 +6286,55 @@ mod owner_control_command_tests { ); } + #[tokio::test] + async fn queue_cap_eviction_prunes_orphaned_hold_deadline() { + let mut pool = AgentPool::from_slots(vec![]); + let mut queue = EventQueue::new(DedupMode::Queue); + let channel_id = Uuid::new_v4(); + let held_scope = thread_scope(channel_id, &"a".repeat(64)); + let surviving_scope = thread_scope(channel_id, &"b".repeat(64)); + + pool.record_scope_owner(held_scope.clone(), 0); + let (tx, _rx) = tokio::sync::oneshot::channel(); + insert_task_meta(&mut pool, 0, surviving_scope.clone(), tx); + assert!(matches!( + pool.hold_decision( + &held_scope, + tokio::time::Instant::now(), + pool::HOLD_BUSY_OWNER_TIMEOUT + ), + pool::HoldDecision::Hold { .. } + )); + + let oldest = std::time::Instant::now() - Duration::from_secs(1); + queue.push(queue::QueuedEvent { + channel_id, + scope: held_scope.clone(), + event: make_event(KIND_STREAM_MESSAGE, "held", None), + received_at: oldest, + prompt_tag: "test".into(), + }); + for i in 0..500 { + queue.push(queue::QueuedEvent { + channel_id, + scope: surviving_scope.clone(), + event: make_event(KIND_STREAM_MESSAGE, &format!("new-{i}"), None), + received_at: std::time::Instant::now(), + prompt_tag: "test".into(), + }); + } + + assert!( + !queue.has_pending_scope(&held_scope), + "aggregate cap evicts the globally oldest scope" + ); + pool.retain_held_scopes(|scope| queue.has_pending_scope(scope)); + assert!( + !pool.held_since_contains(&held_scope), + "evicted scope cannot leave an immediately-ready deadline behind" + ); + } + #[test] fn project_owner_control_signs_only_addressable_project_events() { let keys = Keys::generate(); @@ -9334,6 +9415,15 @@ mod error_outcome_emission_tests { } } + fn bind_agent_scope_owner( + pool: &mut AgentPool, + agent: &mut OwnedAgent, + scope: scope::SessionScope, + ) { + let generation = pool.record_scope_owner(scope.clone(), agent.index); + agent.state.set_scope_owner_generation(scope, generation); + } + #[tokio::test] async fn successful_native_steer_is_transferred_to_live_session_delivery_state() { let channel_id = Uuid::new_v4(); @@ -9349,6 +9439,11 @@ mod error_outcome_emission_tests { ); let mut pool = AgentPool::from_slots(vec![None]); + bind_agent_scope_owner( + &mut pool, + &mut agent, + scope::SessionScope::Conversation { channel_id }, + ); let task_id = pool.join_set.spawn(async {}).id(); pool.task_map_mut().insert( task_id, @@ -9424,6 +9519,11 @@ mod error_outcome_emission_tests { ); let mut pool = AgentPool::from_slots(vec![None]); + bind_agent_scope_owner( + &mut pool, + &mut agent, + scope::SessionScope::Conversation { channel_id }, + ); let task_id = pool.join_set.spawn(async {}).id(); pool.task_map_mut().insert( task_id, @@ -10661,6 +10761,7 @@ mod error_outcome_emission_tests { .sessions .insert(session_scope.clone(), "healthy-session".into()); let mut pool = AgentPool::from_slots(vec![None]); + bind_agent_scope_owner(&mut pool, &mut agent, session_scope.clone()); let task_id = pool.join_set.spawn(async {}).id(); pool.task_map_mut().insert( task_id, diff --git a/crates/buzz-acp/src/pool.rs b/crates/buzz-acp/src/pool.rs index 06383d456d3..dbeafedda70 100644 --- a/crates/buzz-acp/src/pool.rs +++ b/crates/buzz-acp/src/pool.rs @@ -141,9 +141,18 @@ pub struct SessionState { /// Per-scope successful-delivery state. Created with the ACP session and /// cleared atomically with every invalidation path. pub deliveries: HashMap, + /// Pool-assigned ownership generation for each scope. A worker returning + /// after another worker forked the scope carries an older generation; the + /// pool uses this fence to discard that stale provider session before the + /// worker becomes claimable again. + scope_owner_generations: HashMap, } impl SessionState { + pub(crate) fn set_scope_owner_generation(&mut self, scope: SessionScope, generation: u64) { + self.scope_owner_generations.insert(scope, generation); + } + /// Invalidate the session (and turn counter) for a specific prompt source. pub fn invalidate(&mut self, source: &PromptSource) { match source { @@ -165,6 +174,7 @@ impl SessionState { self.core_sections.remove(scope); self.canvas_sections.remove(scope); self.deliveries.remove(scope); + self.scope_owner_generations.remove(scope); self.sessions.remove(scope).is_some() } @@ -179,6 +189,7 @@ impl SessionState { .chain(self.core_sections.keys()) .chain(self.canvas_sections.keys()) .chain(self.deliveries.keys()) + .chain(self.scope_owner_generations.keys()) .filter(|s| s.channel_id() == *channel_id) .cloned() .collect::>() @@ -203,6 +214,7 @@ impl SessionState { self.core_sections.clear(); self.canvas_sections.clear(); self.deliveries.clear(); + self.scope_owner_generations.clear(); } pub(crate) fn mark_scope_delivery_success( @@ -336,13 +348,23 @@ pub struct AgentPool { /// cause another worker to open a duplicate session for the same thread. /// Best-effort: stale entries (rotation, crash/respawn) self-heal on the /// next dispatch and are pruned on channel-wide session invalidation. - session_owners: HashMap, + session_owners: HashMap, + /// Monotonic validity fence assigned whenever a scope is dispatched. The + /// generation distinguishes a newly forked owner from every older copy of + /// that scope's provider session. + next_scope_owner_generation: u64, /// First time each scope was held for a busy owner, so the bounded hold can /// expire and fork rather than starve behind an unbounded turn. Derived /// state: cleared on every dispatch/invalidation path, and only ever holds /// `Thread` scopes (the sole variant [`hold_decision`](Self::hold_decision) /// stamps). - held_since: HashMap, + held_since: HashMap, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct SessionOwner { + agent_index: usize, + generation: u64, } /// Result returned by a completed prompt task. @@ -830,14 +852,33 @@ impl AgentPool { join_set: JoinSet::new(), task_map: HashMap::new(), session_owners: HashMap::new(), + next_scope_owner_generation: 1, held_since: HashMap::new(), } } - /// Record which worker is handling `scope` so a later dispatch can detect a - /// busy owner and avoid opening a duplicate session on another worker. - pub fn record_scope_owner(&mut self, scope: SessionScope, agent_index: usize) { - self.session_owners.insert(scope, agent_index); + /// Record `agent_index` as the newest owner of `scope`, returning the + /// generation that the caller must install on the checked-out worker. + /// Returning workers are accepted only while this exact + /// `(worker, generation)` pair remains authoritative. + pub fn record_scope_owner(&mut self, scope: SessionScope, agent_index: usize) -> u64 { + let generation = self.next_scope_owner_generation; + self.next_scope_owner_generation = self.next_scope_owner_generation.wrapping_add(1); + if self.next_scope_owner_generation == 0 { + // Preserve zero as an unassigned sentinel. Reaching this requires + // 2^64 dispatches in one process, but resetting safely is cheap: + // every previously tagged session becomes stale on return/claim. + self.next_scope_owner_generation = 1; + self.session_owners.clear(); + } + self.session_owners.insert( + scope, + SessionOwner { + agent_index, + generation, + }, + ); + generation } /// True when this scope should be **held** (left queued) rather than @@ -854,16 +895,21 @@ impl AgentPool { return false; } match self.session_owners.get(scope) { - Some(&owner_idx) => self.task_map.values().any(|m| m.agent_index == owner_idx), + Some(owner) => self + .task_map + .values() + .any(|m| m.agent_index == owner.agent_index), None => false, } } /// Decide whether to hold `scope`'s batch for its busy session owner, fork it - /// after a bounded hold, or dispatch immediately. Stamps and clears the - /// first-held time internally so the bounded window survives across dispatch - /// cycles without a dedicated timer; `now` and `timeout` are injected for - /// testability. + /// after a bounded hold, or dispatch immediately. Stamps the first-held time + /// so the bounded window survives across dispatch cycles; `now` and + /// `timeout` are injected for testability. An expired + /// stamp remains sticky until [`clear_hold`](Self::clear_hold) confirms a + /// worker was successfully claimed, so pool exhaustion cannot restart the + /// bounded window. /// /// Gated on the scope variant, not the session policy: `Conversation` scopes /// (channel-policy channels and all DMs) never hold — a busy owner there means @@ -873,18 +919,21 @@ impl AgentPool { pub fn hold_decision( &mut self, scope: &SessionScope, - now: std::time::Instant, + now: tokio::time::Instant, timeout: Duration, ) -> HoldDecision { if !scope.is_thread() || !self.should_hold_for_busy_owner(scope) { self.held_since.remove(scope); return HoldDecision::Dispatch; } - let owner_index = self.session_owners.get(scope).copied().unwrap_or_default(); + let owner_index = self + .session_owners + .get(scope) + .map(|owner| owner.agent_index) + .unwrap_or_default(); let first = *self.held_since.entry(scope.clone()).or_insert(now); let held_for = now.saturating_duration_since(first); if held_for >= timeout { - self.held_since.remove(scope); HoldDecision::ForkAfterHold { held_for, owner_index, @@ -911,7 +960,7 @@ impl AgentPool { if let Some(scope) = scope { let idx = self.agents.iter().position(|slot| { slot.as_ref() - .map(|a| a.state.sessions.contains_key(scope)) + .map(|a| self.agent_owns_scope(a, scope)) .unwrap_or(false) }); if let Some(i) = idx { @@ -925,7 +974,27 @@ impl AgentPool { } /// Return an agent to its slot after a task completes. - pub fn return_agent(&mut self, agent: OwnedAgent) { + pub fn return_agent(&mut self, mut agent: OwnedAgent) { + let stale_scopes: Vec = agent + .state + .sessions + .keys() + .filter(|scope| !self.agent_owns_scope(&agent, scope)) + .cloned() + .collect(); + for scope in stale_scopes { + tracing::info!( + agent = agent.index, + scope = %scope.telemetry_label(), + "discarding stale session after ownership changed" + ); + agent.state.invalidate_scope(&scope); + } + let live_scopes: HashSet = agent.state.sessions.keys().cloned().collect(); + agent + .state + .scope_owner_generations + .retain(|scope, _| live_scopes.contains(scope)); let idx = agent.index; if self.agents[idx].is_some() { // This is a bug: two tasks returned the same agent index. Log it @@ -945,16 +1014,64 @@ impl AgentPool { self.agents.iter().any(|slot| slot.is_some()) } + /// Confirm that pending work for `scope` successfully claimed a worker. + /// + /// In particular, an expired busy-owner hold must not be consumed until + /// this point: `try_claim` can fail while every worker remains checked out. + pub(crate) fn clear_hold(&mut self, scope: &SessionScope) { + self.held_since.remove(scope); + } + + /// Remove derived hold stamps for scopes that no longer have pending work. + pub(crate) fn retain_held_scopes( + &mut self, + mut has_pending_work: impl FnMut(&SessionScope) -> bool, + ) { + self.held_since.retain(|scope, _| has_pending_work(scope)); + } + /// Whether any idle agent already has a session for `scope`. /// Used to compute `affinity_hit` before calling `try_claim`. pub fn has_session_for(&self, scope: &SessionScope) -> bool { self.agents.iter().any(|slot| { slot.as_ref() - .map(|a| a.state.sessions.contains_key(scope)) + .map(|a| self.agent_owns_scope(a, scope)) .unwrap_or(false) }) } + fn agent_owns_scope(&self, agent: &OwnedAgent, scope: &SessionScope) -> bool { + let Some(owner) = self.session_owners.get(scope) else { + return false; + }; + owner.agent_index == agent.index + && agent.state.scope_owner_generations.get(scope) == Some(&owner.generation) + && agent.state.sessions.contains_key(scope) + } + + /// Earliest scheduled wake for a currently held scope that can claim a + /// worker. A worker return wakes the main loop independently, so arming an + /// already-expired timer while every slot is checked out would only spin. + pub(crate) fn next_hold_deadline(&self, timeout: Duration) -> Option { + if !self.any_idle() { + return None; + } + self.held_since + .values() + .map(|held_since| *held_since + timeout) + .min() + } + + /// Sleep until a held scope's scheduled wake, or remain pending when no + /// scope is held. This is the future polled directly by the main + /// `select!`, kept here so paused-time tests exercise the production seam. + pub(crate) async fn wait_for_hold_deadline(deadline: Option) { + match deadline { + Some(deadline) => tokio::time::sleep_until(deadline).await, + None => std::future::pending().await, + } + } + /// Count of agents that are alive: idle OR checked out (have a task_map entry). /// /// Used to detect when all agents have exited so the caller can respawn. @@ -7461,9 +7578,16 @@ printf '%s\n' '{{"jsonrpc":"2.0","id":0,"result":{{"stopReason":"end_turn"}}}}'" agent.state.sessions.insert(ta.clone(), "sess-a".into()); agent.state.sessions.insert(tb.clone(), "sess-b".into()); let mut pool = AgentPool::from_slots(vec![Some(agent)]); - pool.record_scope_owner(ta.clone(), 0); - pool.record_scope_owner(tb.clone(), 0); - let now = std::time::Instant::now(); + let ta_generation = pool.record_scope_owner(ta.clone(), 0); + let tb_generation = pool.record_scope_owner(tb.clone(), 0); + let agent = pool.agents[0].as_mut().expect("idle test agent"); + agent + .state + .set_scope_owner_generation(ta.clone(), ta_generation); + agent + .state + .set_scope_owner_generation(tb.clone(), tb_generation); + let now = tokio::time::Instant::now(); pool.held_since.insert(ta.clone(), now); pool.held_since.insert(tb.clone(), now); @@ -7515,12 +7639,12 @@ printf '%s\n' '{{"jsonrpc":"2.0","id":0,"result":{{"stopReason":"end_turn"}}}}'" /// An idle agent (slot 0) holding a provider session for `scope`, so /// `has_session_for(scope)` is true. - async fn idle_agent_with_session(scope: SessionScope) -> OwnedAgent { + async fn idle_agent_with_session(index: usize, scope: SessionScope) -> OwnedAgent { let acp = AcpClient::spawn("bash", &["-c".into(), "sleep 10".into()], &[], false) .await .expect("spawn dummy ACP"); let mut agent = OwnedAgent { - index: 0, + index, acp, state: SessionState::default(), model_capabilities: None, @@ -7608,7 +7732,7 @@ printf '%s\n' '{{"jsonrpc":"2.0","id":0,"result":{{"stopReason":"end_turn"}}}}'" }, ]; - let base = std::time::Instant::now(); + let base = tokio::time::Instant::now(); for row in rows { let ch = Uuid::new_v4(); let scope = if row.is_thread { @@ -7617,12 +7741,19 @@ printf '%s\n' '{{"jsonrpc":"2.0","id":0,"result":{{"stopReason":"end_turn"}}}}'" conv(ch) }; let slots = if row.has_session { - vec![Some(idle_agent_with_session(scope.clone()).await)] + vec![Some(idle_agent_with_session(0, scope.clone()).await)] } else { vec![] }; let mut pool = AgentPool::from_slots(slots); - if row.owner_busy { + if row.has_session { + let generation = pool.record_scope_owner(scope.clone(), 0); + pool.agents[0] + .as_mut() + .expect("idle test agent") + .state + .set_scope_owner_generation(scope.clone(), generation); + } else if row.owner_busy { pool.record_scope_owner(scope.clone(), 1); mark_agent_busy(&mut pool, 1, thread_scope(ch, &"b".repeat(64))); } @@ -7648,8 +7779,12 @@ printf '%s\n' '{{"jsonrpc":"2.0","id":0,"result":{{"stopReason":"end_turn"}}}}'" _ => panic!("{}: expected {:?}, got {decision:?}", row.name, row.expect), } - // held_since holds the scope only while a Hold is outstanding. - if matches!(decision, HoldDecision::Hold { .. }) { + // An expired hold remains sticky until a worker is successfully + // claimed; only immediate dispatch clears it here. + if matches!( + decision, + HoldDecision::Hold { .. } | HoldDecision::ForkAfterHold { .. } + ) { assert!( pool.held_since.contains_key(&scope), "{}: hold stamps held_since", @@ -7665,6 +7800,159 @@ printf '%s\n' '{{"jsonrpc":"2.0","id":0,"result":{{"stopReason":"end_turn"}}}}'" } } + #[tokio::test(start_paused = true)] + async fn held_scope_deadline_wakes_a_quiet_dispatch_loop() { + let channel_id = Uuid::new_v4(); + let scope = thread_scope(channel_id, &"a".repeat(64)); + let idle_scope = thread_scope(channel_id, &"c".repeat(64)); + let idle_agent = idle_agent_with_session(0, idle_scope).await; + let mut pool = AgentPool::from_slots(vec![Some(idle_agent)]); + pool.record_scope_owner(scope.clone(), 1); + mark_agent_busy(&mut pool, 1, thread_scope(channel_id, &"b".repeat(64))); + + let started = tokio::time::Instant::now(); + assert!(matches!( + pool.hold_decision(&scope, started, HOLD_BUSY_OWNER_TIMEOUT), + HoldDecision::Hold { .. } + )); + let deadline = pool + .next_hold_deadline(HOLD_BUSY_OWNER_TIMEOUT) + .expect("held scope schedules an independent wake"); + let wake = AgentPool::wait_for_hold_deadline(Some(deadline)); + tokio::pin!(wake); + + tokio::time::advance(HOLD_BUSY_OWNER_TIMEOUT - Duration::from_millis(1)).await; + assert!( + tokio::time::timeout(Duration::ZERO, &mut wake) + .await + .is_err(), + "quiet loop stays asleep before deadline" + ); + tokio::time::advance(Duration::from_millis(1)).await; + wake.await; + + assert!(matches!( + pool.hold_decision(&scope, tokio::time::Instant::now(), HOLD_BUSY_OWNER_TIMEOUT), + HoldDecision::ForkAfterHold { .. } + )); + } + + #[tokio::test] + async fn expired_hold_survives_pool_exhaustion_until_a_worker_is_claimable() { + let channel_id = Uuid::new_v4(); + let scope = thread_scope(channel_id, &"a".repeat(64)); + let mut pool = AgentPool::from_slots(vec![None]); + pool.record_scope_owner(scope.clone(), 1); + mark_agent_busy(&mut pool, 1, thread_scope(channel_id, &"b".repeat(64))); + + let started = tokio::time::Instant::now(); + assert!(matches!( + pool.hold_decision(&scope, started, HOLD_BUSY_OWNER_TIMEOUT), + HoldDecision::Hold { .. } + )); + assert!(matches!( + pool.hold_decision( + &scope, + started + HOLD_BUSY_OWNER_TIMEOUT, + HOLD_BUSY_OWNER_TIMEOUT + ), + HoldDecision::ForkAfterHold { .. } + )); + assert!( + pool.held_since.contains_key(&scope), + "failed claim must not restart the timeout" + ); + assert_eq!( + pool.next_hold_deadline(HOLD_BUSY_OWNER_TIMEOUT), + None, + "an expired hold cannot spin while all workers are checked out" + ); + + let idle_scope = thread_scope(channel_id, &"c".repeat(64)); + pool.agents[0] = Some(idle_agent_with_session(0, idle_scope).await); + assert_eq!( + pool.next_hold_deadline(HOLD_BUSY_OWNER_TIMEOUT), + Some(started + HOLD_BUSY_OWNER_TIMEOUT), + "worker availability immediately re-arms the expired deadline" + ); + assert!(matches!( + pool.hold_decision( + &scope, + started + HOLD_BUSY_OWNER_TIMEOUT + Duration::from_secs(1), + HOLD_BUSY_OWNER_TIMEOUT + ), + HoldDecision::ForkAfterHold { .. } + )); + pool.clear_hold(&scope); + assert!(!pool.held_since.contains_key(&scope)); + } + + #[tokio::test] + async fn forked_scope_discards_stale_session_when_busy_owner_returns() { + let channel_id = Uuid::new_v4(); + let scope = thread_scope(channel_id, &"a".repeat(64)); + let busy_scope = thread_scope(channel_id, &"b".repeat(64)); + let old_owner = idle_agent_with_session(0, scope.clone()).await; + let replacement = idle_agent_with_session(1, busy_scope.clone()).await; + let mut pool = AgentPool::from_slots(vec![Some(old_owner), Some(replacement)]); + + let mut old_owner = pool.try_claim(None).expect("claim worker 0"); + let old_generation = pool.record_scope_owner(scope.clone(), old_owner.index); + old_owner + .state + .set_scope_owner_generation(scope.clone(), old_generation); + mark_agent_busy(&mut pool, old_owner.index, busy_scope); + + let started = tokio::time::Instant::now(); + assert!(matches!( + pool.hold_decision(&scope, started, HOLD_BUSY_OWNER_TIMEOUT), + HoldDecision::Hold { .. } + )); + assert!(matches!( + pool.hold_decision( + &scope, + started + HOLD_BUSY_OWNER_TIMEOUT, + HOLD_BUSY_OWNER_TIMEOUT + ), + HoldDecision::ForkAfterHold { .. } + )); + + let mut replacement = pool + .try_claim(Some(&scope)) + .expect("idle worker receives forked scope"); + pool.clear_hold(&scope); + assert_eq!(replacement.index, 1); + replacement + .state + .sessions + .insert(scope.clone(), "fresh-session".into()); + let fresh_generation = pool.record_scope_owner(scope.clone(), replacement.index); + replacement + .state + .set_scope_owner_generation(scope.clone(), fresh_generation); + + // Both turns return. Slot order must not make worker 0's old provider + // context claimable after worker 1 became the authoritative owner. + pool.return_agent(replacement); + pool.task_map + .retain(|_, meta| meta.agent_index != old_owner.index); + pool.return_agent(old_owner); + assert!( + !pool.agents[0] + .as_ref() + .expect("worker 0 returned") + .state + .sessions + .contains_key(&scope), + "return cleanup removes the old provider session" + ); + + let claimed = pool + .try_claim(Some(&scope)) + .expect("authoritative owner remains claimable"); + assert_eq!(claimed.index, 1, "next turn resumes the forked session"); + } + #[test] fn test_rotate_after_natural_completion_invalidates_channel_state() { let (mut s, ch_a, ch_b) = make_state(); @@ -10440,7 +10728,7 @@ done"# pool.invalidate_scope_session(&scopes[1]); pool.record_scope_owner(scopes[0].clone(), 0); pool.held_since - .insert(scopes[0].clone(), std::time::Instant::now()); + .insert(scopes[0].clone(), tokio::time::Instant::now()); assert_eq!( pool.switch_idle_agent_model(channel_id, "model-b", Some("pick".into())), IdleSwitchResult::Switched, diff --git a/crates/buzz-acp/src/queue.rs b/crates/buzz-acp/src/queue.rs index b2203f62f9d..a419b78178d 100644 --- a/crates/buzz-acp/src/queue.rs +++ b/crates/buzz-acp/src/queue.rs @@ -782,6 +782,22 @@ impl EventQueue { self.queues.len() } + /// Whether `scope` still has work that can be reconstructed into a batch. + /// + /// Busy-owner hold timestamps are derived from pending queue state. Queue + /// cap eviction can retire a scope without going through a pool cleanup + /// path, so the dispatch loop uses this seam to prune orphaned holds before + /// scheduling their deadline wakeups. + pub(crate) fn has_pending_scope(&self, scope: &SessionScope) -> bool { + self.queues + .get(scope) + .is_some_and(|queue| !queue.is_empty()) + || self + .cancelled_batches + .get(scope) + .is_some_and(|events| !events.is_empty()) + } + /// Number of queued events for a specific scope (or channel, treated as its /// conversation scope). Test-only. #[cfg(test)]