diff --git a/crates/buzz-acp/src/lib.rs b/crates/buzz-acp/src/lib.rs index 44309138d..8254eb8af 100644 --- a/crates/buzz-acp/src/lib.rs +++ b/crates/buzz-acp/src/lib.rs @@ -36,8 +36,8 @@ use filter::SubscriptionRule; use futures_util::FutureExt; use nostr::{PublicKey, ToBech32}; use pool::{ - AgentPool, ControlSignal, IdleSwitchResult, OwnedAgent, PromptContext, PromptOutcome, - PromptResult, PromptSource, SessionState, TimeoutKind, + AgentPool, ClaimOutcome, ControlSignal, ConversationSessionKey, IdleSwitchResult, OwnedAgent, + PromptContext, PromptOutcome, PromptResult, PromptSource, SessionState, TimeoutKind, }; use queue::{CancelReason, EventQueue, FlushBatch, QueuedEvent, ThreadTags}; use relay::{HarnessRelay, RelayEventPublisher}; @@ -735,6 +735,11 @@ fn handle_relay_observer_control_event( } /// Handle a `cancel_turn` control frame: signal the in-flight task to cancel. +/// +/// An optional `rootEventId` targets one thread scope exactly. Without it, +/// the frame acts only when exactly one scope is in flight for the channel; +/// with several, the target is ambiguous and nothing is touched +/// (`status: "ambiguous_target"`). Control never fans out across roots. fn handle_cancel_turn_control( payload: &serde_json::Value, pool: &mut AgentPool, @@ -748,9 +753,29 @@ fn handle_cancel_turn_control( tracing::warn!("observer cancel_turn control frame missing valid channelId"); return; }; + let root_event_id = payload + .get("rootEventId") + .and_then(|value| value.as_str()) + .map(str::to_owned); - let fired = signal_in_flight_task(pool, channel_id, ControlSignal::Cancel); - let status = if fired { "sent" } else { "no_active_turn" }; + let status = match resolve_control_frame_scope(pool, channel_id, root_event_id) { + Ok(Some(scope)) => { + if signal_in_flight_task(pool, &scope, ControlSignal::Cancel) { + "sent" + } else { + "no_active_turn" + } + } + Ok(None) => "no_active_turn", + Err(in_flight) => { + tracing::warn!( + channel_id = %channel_id, + in_flight, + "cancel_turn without rootEventId while multiple scopes in flight — refusing" + ); + "ambiguous_target" + } + }; if let Some(observer) = observer { observer.emit( "control_result", @@ -769,6 +794,26 @@ fn handle_cancel_turn_control( } } +/// Resolve a control frame's target scope from its optional `rootEventId`. +/// +/// `Some(root)` targets that thread scope exactly (in-flight or not — the +/// signal send reports `no_active_turn` if it isn't). `None` falls back to +/// the channel-level exactly-one rule; `Err(n)` reports `n` in-flight scopes +/// (ambiguous — touch nothing). +fn resolve_control_frame_scope( + pool: &AgentPool, + channel_id: Uuid, + root_event_id: Option, +) -> Result, usize> { + match root_event_id { + Some(root) => Ok(Some(ConversationSessionKey { + channel_id, + root_event_id: Some(root), + })), + None => resolve_channel_control_scope(pool, channel_id), + } +} + /// Handle a `switch_model` control frame (Phase 3a, Option ii). /// /// Busy path: deliver `SwitchModel` over the in-flight task's oneshot — the @@ -780,6 +825,11 @@ fn handle_cancel_turn_control( /// Idle path: validate against the cached catalog *before* invalidating /// (pre-cancel guard), then set `desired_model` + invalidate. The override /// takes visible effect on the agent's next turn. +/// +/// An optional `rootEventId` targets one thread scope exactly. Without it, +/// a busy channel is only switchable when exactly one scope is in flight; +/// with several, the target is ambiguous and nothing is touched +/// (`status: "ambiguous_target"`). Control never fans out across roots. fn handle_switch_model_control( payload: &serde_json::Value, pool: &mut AgentPool, @@ -797,35 +847,50 @@ fn handle_switch_model_control( tracing::warn!("observer switch_model control frame missing modelId"); return; }; + let root_event_id = payload + .get("rootEventId") + .and_then(|value| value.as_str()) + .map(str::to_owned); - // A turn is in flight for this channel iff a task_map entry exists. The - // agent is moved out of the pool during a turn, so the control oneshot is - // the only reachable lever; an idle channel has no such entry. - let turn_in_flight = pool - .task_map() - .values() - .any(|m| m.channel_id == Some(channel_id)); - - let status = if turn_in_flight { - // Busy path: deliver over the oneshot. `false` means the oneshot was - // already consumed this turn (a prior cancel/interrupt) — the turn is - // already ending, so the switch cannot land on it. - if signal_in_flight_task( - pool, - channel_id, - ControlSignal::SwitchModel(model_id.to_string()), - ) { - pool.switch_other_idle_channel_roots(channel_id, model_id); - "sent" - } else { - "turn_ending" + // A turn is in flight for a scope iff a task_map entry exists. The agent + // is moved out of the pool during a turn, so the control oneshot is the + // only reachable lever; an idle scope has no such entry. + let status = match resolve_control_frame_scope(pool, channel_id, root_event_id) { + Ok(Some(scope)) + if pool + .task_map() + .values() + .any(|m| m.scope.as_ref() == Some(&scope)) => + { + // Busy path: deliver over the oneshot. `false` means the oneshot + // was already consumed this turn (a prior cancel/interrupt) — the + // turn is already ending, so the switch cannot land on it. + if signal_in_flight_task( + pool, + &scope, + ControlSignal::SwitchModel(model_id.to_string()), + ) { + pool.switch_other_idle_channel_roots(channel_id, model_id); + "sent" + } else { + "turn_ending" + } } - } else { - // Idle path: validate against the cached catalog before invalidating. - match pool.switch_idle_agent_model(channel_id, model_id) { - IdleSwitchResult::Switched => "switched", - IdleSwitchResult::UnsupportedModel => "unsupported_model", - IdleSwitchResult::NoIdleAgent => "no_active_turn", + Ok(_) => { + // Idle path: validate against the cached catalog before invalidating. + match pool.switch_idle_agent_model(channel_id, model_id) { + IdleSwitchResult::Switched => "switched", + IdleSwitchResult::UnsupportedModel => "unsupported_model", + IdleSwitchResult::NoIdleAgent => "no_active_turn", + } + } + Err(in_flight) => { + tracing::warn!( + channel_id = %channel_id, + in_flight, + "switch_model without rootEventId while multiple scopes in flight — refusing" + ); + "ambiguous_target" } }; @@ -997,10 +1062,10 @@ struct RespawnResult { /// stream. /// /// Carries enough identity to operate on the right withheld event in -/// `EventQueue::withheld_native_steer`: `channel_id` is the routing key, +/// `EventQueue::withheld_native_steer`: `scope` is the routing key, /// `event_id` is the hex id of the single event the steer carried. struct SteerAckEvent { - channel_id: Uuid, + scope: ConversationSessionKey, event_id: String, /// `Ok` if the read loop sent any of the locked `SteerAck` variants. /// `Err` if the oneshot was dropped without a send — should not happen @@ -1515,7 +1580,7 @@ async fn tokio_main() -> Result<()> { } else { None }; - let mut typing_channels: HashMap = HashMap::new(); + let mut typing_scopes: HashMap = HashMap::new(); let mut presence_task: Option> = None; // Runs at the TOP of every loop iteration via Instant check — cannot be @@ -1652,8 +1717,8 @@ 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) { - typing_channels.insert(channel_id, thread_tags); + for (scope, thread_tags) in dispatch_pending(&mut pool, &mut queue, &ctx) { + typing_scopes.insert(scope, thread_tags); } } } @@ -1688,8 +1753,8 @@ 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) { - typing_channels.insert(channel_id, thread_tags); + for (scope, thread_tags) in dispatch_pending(&mut pool, &mut queue, &ctx) { + typing_scopes.insert(scope, thread_tags); } } @@ -1834,7 +1899,7 @@ async fn tokio_main() -> Result<()> { // Track removed channels so checked-out agents get // their sessions stripped when they return to the pool. removed_channels.insert(ch); - typing_channels.remove(&ch); + typing_scopes.retain(|scope, _| scope.channel_id != ch); // Best-effort: clean up 👀 on drained events. // Note: the relay revokes membership before // emitting the notification, so this DELETE may @@ -1908,16 +1973,44 @@ async fn tokio_main() -> Result<()> { if is_cancel { if let Some(owner) = owner_cache.get() { if buzz_event.event.pubkey.to_hex() == *owner { - let fired = signal_in_flight_task( - &mut pool, + // Resolve the target scope from the + // event's thread tags: a threaded + // !cancel hits its root exactly; a + // top-level one acts only when a + // single scope is in flight (never + // fan-out across roots). + match resolve_event_control_scope( + &pool, + &config, + &buzz_event.event, buzz_event.channel_id, - ControlSignal::Cancel, - ); - if !fired { - tracing::warn!( - channel_id = %buzz_event.channel_id, - "!cancel received but no in-flight task — no-op" - ); + ) { + Ok(Some(scope)) => { + if !signal_in_flight_task( + &mut pool, + &scope, + ControlSignal::Cancel, + ) { + tracing::warn!( + channel_id = %buzz_event.channel_id, + "!cancel received but no in-flight task — no-op" + ); + } + } + Ok(None) => { + tracing::warn!( + channel_id = %buzz_event.channel_id, + "!cancel received but no in-flight task — no-op" + ); + } + Err(in_flight) => { + tracing::warn!( + channel_id = %buzz_event.channel_id, + in_flight, + "top-level !cancel while multiple threads in flight — \ + ambiguous target, touching nothing (reply in the thread to cancel it)" + ); + } } continue; // consume event — do NOT push to queue } @@ -1946,23 +2039,46 @@ async fn tokio_main() -> Result<()> { if is_rotate { if let Some(owner) = owner_cache.get() { if buzz_event.event.pubkey.to_hex() == *owner { - let fired = signal_in_flight_task( - &mut pool, + // Same scope resolution as !cancel: a + // threaded !rotate hits its root; a + // top-level one needs an unambiguous + // (single or zero) in-flight target. + match resolve_event_control_scope( + &pool, + &config, + &buzz_event.event, buzz_event.channel_id, - ControlSignal::Rotate, - ); - if fired { - tracing::info!( - channel_id = %buzz_event.channel_id, - "!rotate received — cancelling in-flight turn and rotating session" - ); - } else { - let invalidated = pool.invalidate_channel_sessions(buzz_event.channel_id); - tracing::info!( - channel_id = %buzz_event.channel_id, - invalidated, - "!rotate received — invalidated idle channel session(s)" - ); + ) { + Ok(scope) => { + let fired = scope.is_some_and(|scope| { + signal_in_flight_task( + &mut pool, + &scope, + ControlSignal::Rotate, + ) + }); + if fired { + tracing::info!( + channel_id = %buzz_event.channel_id, + "!rotate received — cancelling in-flight turn and rotating session" + ); + } else { + let invalidated = pool.invalidate_channel_sessions(buzz_event.channel_id); + tracing::info!( + channel_id = %buzz_event.channel_id, + invalidated, + "!rotate received — invalidated idle channel session(s)" + ); + } + } + Err(in_flight) => { + tracing::warn!( + channel_id = %buzz_event.channel_id, + in_flight, + "top-level !rotate while multiple threads in flight — \ + ambiguous target, touching nothing (reply in the thread to rotate it)" + ); + } } continue; // consume event — do NOT push to queue } @@ -2026,17 +2142,22 @@ async fn tokio_main() -> Result<()> { // backed payload) so the cost is negligible. let event_for_steer = buzz_event.event.clone(); let prompt_tag_for_steer = prompt_tag.clone(); + let conversation_root = if matches!(config.session_scope, config::SessionScope::Thread) { + let tags = queue::parse_thread_tags(&event_for_steer); + Some(tags.root_event_id.unwrap_or_else(|| event_id_hex.clone())) + } else { + None + }; + let scope = ConversationSessionKey { + channel_id: buzz_event.channel_id, + root_event_id: conversation_root.clone(), + }; let accepted = queue.push(QueuedEvent { channel_id: buzz_event.channel_id, event: buzz_event.event, received_at: std::time::Instant::now(), prompt_tag, - conversation_root: if matches!(config.session_scope, config::SessionScope::Thread) { - let tags = queue::parse_thread_tags(&event_for_steer); - Some(tags.root_event_id.unwrap_or_else(|| event_id_hex.clone())) - } else { - None - }, + conversation_root, }); // 👀 — immediate "seen" reaction, only if the event // was actually queued (not dropped by DedupMode::Drop). @@ -2051,9 +2172,11 @@ async fn tokio_main() -> Result<()> { }); } // Event is already queued. If mode requires it AND - // the channel has an in-flight task, fire cancel — + // the scope has an in-flight task, fire cancel — // OR take the non-cancelling (ACP steer) fork for Steer signals. - if accepted && queue.is_channel_in_flight(buzz_event.channel_id) { + // Scope-exact: a busy sibling thread in the same + // channel never triggers a signal for this one. + if accepted && queue.is_scope_in_flight(&scope) { // Author eligibility (owner ∪ allowlist ∪ siblings) // is already enforced by the inbound author gate // above, so the mid-turn signal fires for every @@ -2080,7 +2203,7 @@ async fn tokio_main() -> Result<()> { && try_native_steer( &mut pool, &mut queue, - buzz_event.channel_id, + &scope, event_for_steer, prompt_tag_for_steer, &steer_ack_tx, @@ -2088,16 +2211,16 @@ async fn tokio_main() -> Result<()> { if !native_attempted { signal_in_flight_task( &mut pool, - buzz_event.channel_id, + &scope, signal, ); } } } - for (channel_id, thread_tags) in + for (scope, thread_tags) in dispatch_pending(&mut pool, &mut queue, &ctx) { - typing_channels.insert(channel_id, thread_tags); + typing_scopes.insert(scope, thread_tags); } } None => { @@ -2120,10 +2243,10 @@ async fn tokio_main() -> Result<()> { let _ = result_rx; if queue.has_flushable_work() { tracing::debug!("heartbeat_skipped_events"); - for (channel_id, thread_tags) in + for (scope, thread_tags) in dispatch_pending(&mut pool, &mut queue, &ctx) { - typing_channels.insert(channel_id, thread_tags); + typing_scopes.insert(scope, thread_tags); } } else if pool.any_idle() { dispatch_heartbeat(&mut pool, &ctx, &mut heartbeat_in_flight); @@ -2162,14 +2285,17 @@ async fn tokio_main() -> Result<()> { // Use try_publish (non-blocking) for typing indicators — // they're ephemeral and must not block the main loop during // relay reconnection (#35). - for (&ch, thread_tags) in &typing_channels { + for (scope, thread_tags) in &typing_scopes { if let Ok(event) = relay.build_typing_event( - ch, + scope.channel_id, thread_tags.root_event_id.as_deref(), thread_tags.parent_event_id.as_deref(), ) { if let Err(e) = relay.try_publish_event(event) { - tracing::debug!("typing indicator dropped for {ch}: {e}"); + tracing::debug!( + "typing indicator dropped for {}: {e}", + scope.channel_id + ); } } } @@ -2184,9 +2310,9 @@ async fn tokio_main() -> Result<()> { match pool_event { Some(PoolEvent::Result(result)) => { - // Stop typing indicator for the completed channel. + // Stop typing indicator for the completed scope. if let PromptSource::Channel(key) = &result.source { - typing_channels.remove(&key.channel_id); + typing_scopes.remove(key); } if handle_prompt_result( &mut pool, @@ -2210,7 +2336,7 @@ async fn tokio_main() -> Result<()> { &config, &mut heartbeat_in_flight, &removed_channels, - &mut typing_channels, + &mut typing_scopes, &mut crash_history, &respawn_tx, &mut respawn_tasks, @@ -2219,8 +2345,8 @@ async fn tokio_main() -> Result<()> { { break; } - for (channel_id, thread_tags) in dispatch_pending(&mut pool, &mut queue, &ctx) { - typing_channels.insert(channel_id, thread_tags); + for (scope, thread_tags) in dispatch_pending(&mut pool, &mut queue, &ctx) { + typing_scopes.insert(scope, thread_tags); } } Some(PoolEvent::Panic(join_error)) => { @@ -2232,7 +2358,7 @@ async fn tokio_main() -> Result<()> { join_error, &mut heartbeat_in_flight, &removed_channels, - &mut typing_channels, + &mut typing_scopes, &mut crash_history, &respawn_tx, &mut respawn_tasks, @@ -2242,12 +2368,12 @@ 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) { - typing_channels.insert(channel_id, thread_tags); + for (scope, thread_tags) in dispatch_pending(&mut pool, &mut queue, &ctx) { + typing_scopes.insert(scope, thread_tags); } } Some(PoolEvent::SteerAck(SteerAckEvent { - channel_id, + scope, event_id, ack, })) => { @@ -2330,7 +2456,8 @@ async fn tokio_main() -> Result<()> { Err(_recv_err) => (true, false, false), }; tracing::info!( - channel = %channel_id, + channel = %scope.channel_id, + root = scope.root_event_id.as_deref().unwrap_or(""), event_id = %event_id, ?ack, release_withheld, @@ -2339,28 +2466,28 @@ async fn tokio_main() -> Result<()> { "non-cancelling steer ack received" ); if drop_withheld { - queue.remove_event(channel_id, &event_id); + queue.remove_event(&scope, &event_id); } if release_withheld { - queue.release_native_steer(channel_id, &event_id); + queue.release_native_steer(&scope, &event_id); } if signal_fallback { // Universal cancel+merge fallback. Note: the // queued event has already been released to the - // front of `queues[channel_id]`, so the cancel + // front of `queues[scope]`, so the cancel // will pick it up as part of the merged batch and // re-prompt the agent. - signal_in_flight_task(&mut pool, channel_id, ControlSignal::Steer); + signal_in_flight_task(&mut pool, &scope, ControlSignal::Steer); } // After releasing a withheld event, give dispatch a chance // to re-flush. If the prompt is still in flight, the - // channel stays `in_flight_channels` and `flush_next` + // scope stays `in_flight_scopes` and `flush_next` // skips it — but a Steer fallback signal sent above will // 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) { - typing_channels.insert(channel_id, thread_tags); + for (scope, thread_tags) in dispatch_pending(&mut pool, &mut queue, &ctx) { + typing_scopes.insert(scope, thread_tags); } } None => {} // relay/heartbeat/shutdown branches handled inline above @@ -2519,21 +2646,26 @@ fn mode_gate_signal( } } -/// Send a control signal to the in-flight task for `channel_id`. +/// Send a control signal to the in-flight task for `scope`. /// Returns `true` if a signal was sent, `false` if no in-flight task was found. fn signal_in_flight_task( pool: &mut AgentPool, - channel_id: uuid::Uuid, + scope: &ConversationSessionKey, mode: ControlSignal, ) -> bool { let entry = pool .task_map_mut() .values_mut() - .find(|m| m.channel_id == Some(channel_id)); + .find(|m| m.scope.as_ref() == Some(scope)); if let Some(meta) = entry { if let Some(tx) = meta.control_tx.take() { - tracing::info!(channel = %channel_id, ?mode, "control signal sent to in-flight task"); + tracing::info!( + channel = %scope.channel_id, + root = scope.root_event_id.as_deref().unwrap_or(""), + ?mode, + "control signal sent to in-flight task" + ); let _ = tx.send(mode); return true; } @@ -2541,6 +2673,57 @@ fn signal_in_flight_task( false } +/// Resolve a channel-level control action (no explicit root) to a single +/// in-flight scope. +/// +/// Control never fans out across roots: with more than one scope in flight +/// for the channel the target is ambiguous and the caller must touch +/// nothing. `Ok(None)` means no turn is in flight at all. +fn resolve_channel_control_scope( + pool: &AgentPool, + channel_id: uuid::Uuid, +) -> Result, usize> { + let mut scopes = pool + .task_map() + .values() + .filter_map(|m| m.scope.as_ref()) + .filter(|scope| scope.channel_id == channel_id); + let Some(first) = scopes.next() else { + return Ok(None); + }; + let extra = scopes.count(); + if extra == 0 { + Ok(Some(first.clone())) + } else { + Err(extra + 1) + } +} + +/// Resolve the scope an inbound control event (owner `!cancel` / `!rotate`) +/// targets. +/// +/// Thread mode: a threaded event targets its outermost root exactly; a +/// top-level event falls back to the channel-level exactly-one rule via +/// [`resolve_channel_control_scope`]. Channel mode: always the channel scope. +fn resolve_event_control_scope( + pool: &AgentPool, + config: &Config, + event: &nostr::Event, + channel_id: uuid::Uuid, +) -> Result, usize> { + if matches!(config.session_scope, config::SessionScope::Thread) { + match queue::parse_thread_tags(event).root_event_id { + Some(root) => Ok(Some(ConversationSessionKey { + channel_id, + root_event_id: Some(root), + })), + None => resolve_channel_control_scope(pool, channel_id), + } + } else { + Ok(Some(ConversationSessionKey::channel(channel_id))) + } +} + /// Attempt the non-cancelling (ACP) steer for a freshly-queued event. /// /// Caller invariants: @@ -2568,11 +2751,12 @@ fn signal_in_flight_task( fn try_native_steer( pool: &mut AgentPool, queue: &mut EventQueue, - channel_id: uuid::Uuid, + scope: &ConversationSessionKey, event: nostr::Event, prompt_tag: String, steer_ack_tx: &mpsc::UnboundedSender, ) -> bool { + let channel_id = scope.channel_id; // Build the steer body: framing strings come from // `queue::native_steer_framing()` (Eva's drift-proof requirement — // native and cancel+merge fallback share these so the agent gets the @@ -2602,14 +2786,14 @@ fn try_native_steer( ack_tx, }; - match pool.send_steer(channel_id, request) { + match pool.send_steer(scope, request) { Ok(()) => { // Withhold the queued event synchronously BEFORE spawning // the watcher: this closes the race where `mark_complete` - // clears `in_flight_channels` and a stray `flush_next` could + // clears `in_flight_scopes` and a stray `flush_next` could // re-deliver the event via normal dispatch. See - // `EventQueue::mark_native_steer_pending` docs at queue.rs:606. - let withheld = queue.mark_native_steer_pending(channel_id, &event_id_hex); + // `EventQueue::mark_native_steer_pending` docs. + let withheld = queue.mark_native_steer_pending(scope, &event_id_hex); if !withheld { // Race: the event was already drained out of the queue // before we got here (e.g. a concurrent flush picked it @@ -2627,10 +2811,11 @@ fn try_native_steer( } let ack_tx_clone = steer_ack_tx.clone(); let event_id_for_watcher = event_id_hex.clone(); + let scope_for_watcher = scope.clone(); tokio::spawn(async move { let ack = ack_rx.await; let _ = ack_tx_clone.send(SteerAckEvent { - channel_id, + scope: scope_for_watcher, event_id: event_id_for_watcher, ack, }); @@ -2651,12 +2836,23 @@ fn try_native_steer( // ── dispatch_pending ────────────────────────────────────────────────────────── /// Flush queued work to available agents. +/// +/// Busy scopes don't block the loop: a batch whose retained root is owned by +/// a busy slot is requeued (timestamps preserved) and dispatch continues with +/// other scopes — head-of-line blocking across roots is exactly what +/// thread-scoped sessions exist to remove. Only pool exhaustion stops the +/// loop. fn dispatch_pending( pool: &mut AgentPool, queue: &mut EventQueue, ctx: &Arc, -) -> Vec<(Uuid, ThreadTags)> { - let mut dispatched_channels = Vec::new(); +) -> Vec<(ConversationSessionKey, ThreadTags)> { + let mut dispatched_scopes = Vec::new(); + // Batches whose retained root's owner is busy this pass. Held (scope + // still marked in-flight) so `flush_next` cannot re-select the scope — + // restoring it mid-loop would livelock on flush → busy → restore. + // Restored after the loop. + let mut busy_batches: Vec = Vec::new(); loop { let batch = match queue.flush_next() { Some(b) => b, @@ -2668,15 +2864,23 @@ fn dispatch_pending( .last() .map(|event| queue::parse_thread_tags(&event.event)) .unwrap_or_default(); - let session_key = pool::conversation_session_key(&batch); + let session_key = batch.scope_key(); let affinity_hit = pool.has_session_for(&session_key); let mut agent = match pool.try_claim(Some(&session_key)) { - Some(a) => a, - None => { + ClaimOutcome::Claimed(a) => a, + ClaimOutcome::BusyOwner => { + tracing::debug!( + channel = %channel_id, + root = session_key.root_event_id.as_deref().unwrap_or(""), + "busy_owner — holding scope, continuing dispatch" + ); + busy_batches.push(batch); + continue; + } + ClaimOutcome::Exhausted => { let pending = queue.pending_channels(); - tracing::debug!(pending_channels = pending, "pool_exhausted"); - queue.requeue_preserve_timestamps(batch); - queue.mark_complete(channel_id); + tracing::debug!(pending_scopes = pending, "pool_exhausted"); + busy_batches.push(batch); break; } }; @@ -2727,21 +2931,27 @@ fn dispatch_pending( abort_handle.id(), pool::TaskMeta { agent_index, - channel_id: Some(channel_id), + scope: Some(session_key.clone()), turn_id, recoverable_batch, control_tx: Some(control_tx), steer_tx, }, ); - dispatched_channels.push((channel_id, typing_scope)); + dispatched_scopes.push((session_key, typing_scope)); + } + // Return held busy/exhausted batches to the queue exactly as flushed. + for batch in busy_batches { + let scope = batch.scope_key(); + queue.restore_unclaimed(batch); + queue.mark_complete(&scope); } tracing::debug!( - dispatched = dispatched_channels.len(), + dispatched = dispatched_scopes.len(), queue_depth = queue.pending_channels(), "dispatch_pending" ); - dispatched_channels + dispatched_scopes } /// Spawn a task that posts a user-visible failure notice to the relay. @@ -2863,7 +3073,7 @@ fn handle_prompt_result( } match &result.source { - PromptSource::Channel(key) => queue.mark_complete(key.channel_id), + PromptSource::Channel(key) => queue.mark_complete(key), PromptSource::Heartbeat => *heartbeat_in_flight = false, } @@ -3095,7 +3305,7 @@ fn recover_panicked_agent( join_error: tokio::task::JoinError, heartbeat_in_flight: &mut bool, removed_channels: &HashSet, - typing_channels: &mut HashMap, + typing_scopes: &mut HashMap, crash_history: &mut [SlotCircuit], respawn_tx: &mpsc::Sender, respawn_tasks: &mut tokio::task::JoinSet<()>, @@ -3110,25 +3320,25 @@ 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 let Some(ref scope) = meta.scope { + if !removed_channels.contains(&scope.channel_id) { // Dead-letter on exhaustion is logged inside requeue(); a // panic path has no outcome to report, so no notice here. let _ = queue.requeue(batch); tracing::warn!("requeued batch for panicked agent {i}"); } else { tracing::debug!( - channel_id = %ch, + channel_id = %scope.channel_id, "dropping panicked batch for removed channel" ); } } } - if let Some(ch) = meta.channel_id { - queue.mark_complete(ch); - typing_channels.remove(&ch); - tracing::warn!("cleared wedged in-flight channel {ch} from panicked agent {i}"); + if let Some(ref scope) = meta.scope { + queue.mark_complete(scope); + typing_scopes.remove(scope); + tracing::warn!("cleared wedged in-flight scope {scope:?} from panicked agent {i}"); } else { *heartbeat_in_flight = false; tracing::warn!("cleared wedged heartbeat_in_flight from panicked agent {i}"); @@ -3138,7 +3348,11 @@ fn recover_panicked_agent( observer.emit( "agent_panic", Some(i), - &observer::context_for(meta.channel_id, None, Some(meta.turn_id)), + &observer::context_for( + meta.scope.as_ref().map(|s| s.channel_id), + None, + Some(meta.turn_id), + ), serde_json::json!({ "outcome": "panic", "error": format!("Agent task panicked: {join_error}"), @@ -3193,7 +3407,7 @@ fn drain_ready_join_results( config: &Config, heartbeat_in_flight: &mut bool, removed_channels: &HashSet, - typing_channels: &mut HashMap, + typing_scopes: &mut HashMap, crash_history: &mut [SlotCircuit], respawn_tx: &mpsc::Sender, respawn_tasks: &mut tokio::task::JoinSet<()>, @@ -3209,7 +3423,7 @@ fn drain_ready_join_results( join_error, heartbeat_in_flight, removed_channels, - typing_channels, + typing_scopes, crash_history, respawn_tx, respawn_tasks, @@ -3232,8 +3446,8 @@ fn dispatch_heartbeat( return; } let agent = match pool.try_claim(None) { - Some(a) => a, - None => return, + pool::ClaimOutcome::Claimed(a) => a, + pool::ClaimOutcome::BusyOwner | pool::ClaimOutcome::Exhausted => return, }; let prompt_text = ctx @@ -3263,7 +3477,7 @@ fn dispatch_heartbeat( abort_handle.id(), pool::TaskMeta { agent_index, - channel_id: None, + scope: None, turn_id, recoverable_batch: None, control_tx: None, @@ -3842,8 +4056,8 @@ mod owner_control_command_tests { #[tokio::test] async fn signal_in_flight_task_sends_rotate_once() { let mut pool = AgentPool::from_slots(vec![]); - let channel_id = Uuid::new_v4(); - let other_channel_id = Uuid::new_v4(); + let scope = ConversationSessionKey::channel(Uuid::new_v4()); + let other_scope = ConversationSessionKey::channel(Uuid::new_v4()); let (control_tx, control_rx) = tokio::sync::oneshot::channel(); let abort_handle = pool.join_set.spawn(async {}); @@ -3851,7 +4065,7 @@ mod owner_control_command_tests { abort_handle.id(), pool::TaskMeta { agent_index: 0, - channel_id: Some(channel_id), + scope: Some(scope.clone()), turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: Some(control_tx), @@ -3861,18 +4075,18 @@ mod owner_control_command_tests { assert!(!signal_in_flight_task( &mut pool, - other_channel_id, + &other_scope, ControlSignal::Rotate )); assert!(signal_in_flight_task( &mut pool, - channel_id, + &scope, ControlSignal::Rotate )); assert_eq!(control_rx.await.unwrap(), ControlSignal::Rotate); assert!(!signal_in_flight_task( &mut pool, - channel_id, + &scope, ControlSignal::Rotate )); } @@ -4143,7 +4357,9 @@ mod build_mcp_servers_tests { /// Env-var-touching tests must run serially — env vars are process-global. static ENV_LOCK: Mutex<()> = Mutex::new(()); - fn test_config() -> Config { + /// `pub(super)` so sibling test modules (e.g. `dispatch_scope_tests`) + /// can borrow a baseline Config instead of duplicating the literal. + pub(super) fn test_config() -> Config { Config { keys: nostr::Keys::generate(), relay_url: "ws://localhost:3000".into(), @@ -4402,6 +4618,7 @@ mod error_outcome_emission_tests { let mut owner = pool .try_claim(Some(&key)) + .claimed() .expect("root should reserve slot 0"); assert_eq!(owner.index, 0); let task_id = pool.join_set.spawn(async {}).id(); @@ -4409,7 +4626,7 @@ mod error_outcome_emission_tests { task_id, crate::pool::TaskMeta { agent_index: 0, - channel_id: Some(key.channel_id), + scope: Some(key.clone()), turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -4418,13 +4635,13 @@ mod error_outcome_emission_tests { ); assert!(pool.any_idle(), "slot 1 remains idle"); assert!( - pool.try_claim(Some(&key)).is_none(), + matches!(pool.try_claim(Some(&key)), pool::ClaimOutcome::BusyOwner), "reply must wait for the reserved root owner" ); pool.task_map_mut().remove(&task_id); owner.state.insert_session(key.clone(), "session-a".into()); pool.return_agent(owner); - assert_eq!(pool.try_claim(Some(&key)).unwrap().index, 0); + assert_eq!(pool.try_claim(Some(&key)).claimed().unwrap().index, 0); } async fn pool_with_retained_root() -> (AgentPool, pool::ConversationSessionKey) { @@ -4435,7 +4652,7 @@ mod error_outcome_emission_tests { channel_id: Uuid::new_v4(), root_event_id: Some("root-a".into()), }; - let mut owner = pool.try_claim(Some(&key)).unwrap(); + let mut owner = pool.try_claim(Some(&key)).claimed().unwrap(); owner.state.insert_session(key.clone(), "session-a".into()); pool.return_agent(owner); (pool, key) @@ -4444,7 +4661,7 @@ mod error_outcome_emission_tests { #[tokio::test] async fn root_owner_is_cleared_after_lru_eviction() { let (mut pool, key) = pool_with_retained_root().await; - let mut owner = pool.try_claim(Some(&key)).unwrap(); + let mut owner = pool.try_claim(Some(&key)).claimed().unwrap(); for index in 0..=pool::SessionState::MAX_CHANNEL_SESSIONS { owner.state.insert_session( pool::ConversationSessionKey { @@ -4457,7 +4674,7 @@ mod error_outcome_emission_tests { assert!(!owner.state.contains_session(&key)); pool.return_agent(owner); pool.agents_mut()[0] = None; - assert_eq!(pool.try_claim(Some(&key)).unwrap().index, 1); + assert_eq!(pool.try_claim(Some(&key)).claimed().unwrap().index, 1); } #[tokio::test] @@ -4465,14 +4682,14 @@ mod error_outcome_emission_tests { let (mut pool, key) = pool_with_retained_root().await; assert_eq!(pool.invalidate_channel_sessions(key.channel_id), 1); pool.agents_mut()[0] = None; - assert_eq!(pool.try_claim(Some(&key)).unwrap().index, 1); + assert_eq!(pool.try_claim(Some(&key)).claimed().unwrap().index, 1); } #[tokio::test] async fn stale_dead_root_owner_falls_back_to_live_slot() { let (mut pool, key) = pool_with_retained_root().await; pool.agents_mut()[0] = None; - assert_eq!(pool.try_claim(Some(&key)).unwrap().index, 1); + assert_eq!(pool.try_claim(Some(&key)).claimed().unwrap().index, 1); } #[tokio::test] @@ -4489,7 +4706,7 @@ mod error_outcome_emission_tests { }) .collect(); for key in &keys { - let mut owner = pool.try_claim(Some(key)).unwrap(); + let mut owner = pool.try_claim(Some(key)).claimed().unwrap(); owner .state .insert_session(key.clone(), format!("session-{key:?}")); @@ -4498,10 +4715,10 @@ mod error_outcome_emission_tests { pool.agents_mut()[0] = None; // 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(); + let first_claim = pool.try_claim(Some(&keys[0])).claimed().unwrap(); assert_eq!(first_claim.index, 1); pool.return_agent(first_claim); - assert_eq!(pool.try_claim(Some(&keys[1])).unwrap().index, 1); + assert_eq!(pool.try_claim(Some(&keys[1])).claimed().unwrap().index, 1); } #[tokio::test] @@ -4552,7 +4769,7 @@ mod error_outcome_emission_tests { let mut second = dummy_agent(1).await; second.state.insert_session(idle_key, "session-idle".into()); let mut pool = AgentPool::from_slots(vec![Some(first), Some(second)]); - let busy = pool.try_claim(Some(&busy_key)).unwrap(); + let busy = pool.try_claim(Some(&busy_key)).claimed().unwrap(); pool.switch_other_idle_channel_roots(channel_id, "new-model"); let idle = pool.agents_mut()[1].as_ref().unwrap(); assert!(!idle.state.has_channel_state(&channel_id)); @@ -4575,7 +4792,7 @@ mod error_outcome_emission_tests { task_id, crate::pool::TaskMeta { agent_index: 0, - channel_id: None, + scope: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -4651,7 +4868,7 @@ mod error_outcome_emission_tests { task_id, crate::pool::TaskMeta { agent_index: 0, - channel_id: Some(channel_id), + scope: Some(ConversationSessionKey::channel(channel_id)), turn_id: "panic-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -4740,7 +4957,7 @@ mod error_outcome_emission_tests { task_id, crate::pool::TaskMeta { agent_index: 0, - channel_id: None, + scope: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -4828,7 +5045,7 @@ mod error_outcome_emission_tests { task_id, crate::pool::TaskMeta { agent_index: 0, - channel_id: None, + scope: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -4868,7 +5085,7 @@ mod error_outcome_emission_tests { ); ( queue.pending_channels(), - queue.queued_event_count(&channel_id), + queue.queued_event_count(&ConversationSessionKey::channel(channel_id)), ) }; @@ -4941,7 +5158,7 @@ mod error_outcome_emission_tests { task_id, crate::pool::TaskMeta { agent_index: 0, - channel_id: None, + scope: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -5081,7 +5298,7 @@ mod error_outcome_emission_tests { task_id, crate::pool::TaskMeta { agent_index: 0, - channel_id: None, + scope: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, control_tx: None, @@ -5181,6 +5398,294 @@ mod error_outcome_emission_tests { } } +#[cfg(test)] +mod dispatch_scope_tests { + //! Pins the concurrency contract of thread-scoped dispatch: + //! + //! - Distinct roots in one channel dispatch to distinct workers in one + //! `dispatch_pending` pass — no cross-root head-of-line blocking. + //! - A retained root whose owning slot is busy is skipped (batch restored + //! losslessly), while other scopes keep dispatching. The root never + //! migrates to another slot. + //! - Pool exhaustion restores the undispatched batch losslessly. + + use super::*; + use crate::acp::AcpClient; + use crate::pool::{AgentPool, OwnedAgent, PromptContext}; + use crate::queue::QueuedEvent; + use nostr::{EventBuilder, Keys, Kind}; + + async fn dummy_agent(index: usize) -> OwnedAgent { + OwnedAgent { + index, + acp: AcpClient::spawn("cat", &[], &[], false) + .await + .expect("spawn cat as inert agent"), + state: Default::default(), + model_capabilities: None, + desired_model: None, + model_overridden: false, + agent_name: "unknown".into(), + goose_system_prompt_supported: None, + protocol_version: 1, + } + } + + fn test_ctx() -> Arc { + let keys = nostr::Keys::generate(); + Arc::new(PromptContext { + mcp_servers: vec![], + initial_message: None, + idle_timeout: Duration::from_secs(60), + max_turn_duration: Duration::from_secs(120), + turn_liveness_interval: Duration::ZERO, + dedup_mode: DedupMode::Queue, + system_prompt: None, + team_instructions: None, + heartbeat_prompt: None, + base_prompt: None, + cwd: ".".to_string(), + rest_client: relay::RestClient { + http: reqwest::Client::new(), + base_url: "http://127.0.0.1:0".to_string(), + keys: keys.clone(), + auth_tag_json: None, + }, + channel_info: std::collections::HashMap::new(), + context_message_limit: 0, + max_turns_per_session: 0, + permission_mode: config::PermissionMode::Default, + agent_keys: keys, + agent_owner_pubkey: None, + memory_enabled: false, + harness_name: "goose".to_string(), + }) + } + + fn push_rooted(queue: &mut EventQueue, ch: Uuid, root: &str, content: &str) { + let keys = Keys::generate(); + let event = EventBuilder::new(Kind::Custom(9), content) + .tags([]) + .sign_with_keys(&keys) + .unwrap(); + queue.push(QueuedEvent { + channel_id: ch, + conversation_root: Some(root.into()), + event, + received_at: std::time::Instant::now(), + prompt_tag: "test".into(), + }); + } + + fn key(ch: Uuid, root: &str) -> ConversationSessionKey { + ConversationSessionKey { + channel_id: ch, + root_event_id: Some(root.into()), + } + } + + #[tokio::test] + async fn two_roots_in_one_channel_dispatch_concurrently_to_distinct_workers() { + let mut pool = + AgentPool::from_slots(vec![Some(dummy_agent(0).await), Some(dummy_agent(1).await)]); + let mut queue = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + push_rooted(&mut queue, ch, "root-a", "for a"); + push_rooted(&mut queue, ch, "root-b", "for b"); + + let dispatched = dispatch_pending(&mut pool, &mut queue, &test_ctx()); + + let scopes: Vec<_> = dispatched.iter().map(|(s, _)| s.clone()).collect(); + assert!(scopes.contains(&key(ch, "root-a"))); + assert!(scopes.contains(&key(ch, "root-b"))); + assert!(queue.is_scope_in_flight(&key(ch, "root-a"))); + assert!(queue.is_scope_in_flight(&key(ch, "root-b"))); + // Distinct workers: both slots checked out, distinct task scopes. + assert!(!pool.any_idle()); + let task_scopes: HashSet<_> = pool + .task_map() + .values() + .filter_map(|m| m.scope.clone()) + .collect(); + assert_eq!(task_scopes.len(), 2); + } + + #[tokio::test] + async fn busy_root_owner_is_skipped_without_blocking_other_scopes() { + let mut pool = + AgentPool::from_slots(vec![Some(dummy_agent(0).await), Some(dummy_agent(1).await)]); + let mut queue = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + + // Slot 0 retains root-a's session… + let mut owner = pool + .try_claim(Some(&key(ch, "root-a"))) + .claimed() + .expect("claim slot 0 for root-a"); + owner + .state + .insert_session(key(ch, "root-a"), "session-a".into()); + pool.return_agent(owner); + // …and is then checked out busy on an unrelated scope. + let busy = pool + .try_claim(Some(&key(ch, "root-busy"))) + .claimed() + .expect("slot 0 is first idle"); + assert_eq!(busy.index, 0); + let task_id = pool.join_set.spawn(std::future::pending()).id(); + pool.task_map_mut().insert( + task_id, + pool::TaskMeta { + agent_index: 0, + scope: Some(key(ch, "root-busy")), + turn_id: "busy-turn".into(), + recoverable_batch: None, + control_tx: None, + steer_tx: None, + }, + ); + drop(busy); // stays checked out — slot 0 empty + + push_rooted(&mut queue, ch, "root-a", "must wait for owner"); + push_rooted(&mut queue, ch, "root-b", "must not be blocked"); + + let dispatched = dispatch_pending(&mut pool, &mut queue, &test_ctx()); + + // root-b dispatched (on slot 1) despite root-a being older and busy. + let scopes: Vec<_> = dispatched.iter().map(|(s, _)| s.clone()).collect(); + assert_eq!(scopes, vec![key(ch, "root-b")]); + // root-a restored losslessly: queued again, not in flight, and NOT + // migrated to slot 1. + assert!(!queue.is_scope_in_flight(&key(ch, "root-a"))); + assert_eq!(queue.queued_event_count(&key(ch, "root-a")), 1); + assert!( + pool.task_map() + .values() + .all(|m| m.scope != Some(key(ch, "root-a"))), + "busy-owned root must never run on another slot" + ); + } + + #[tokio::test] + async fn pool_exhaustion_restores_undispatched_batch_losslessly() { + let mut pool = AgentPool::from_slots(vec![Some(dummy_agent(0).await)]); + let mut queue = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + push_rooted(&mut queue, ch, "root-a", "gets the only agent"); + push_rooted(&mut queue, ch, "root-b", "waits for capacity"); + + let dispatched = dispatch_pending(&mut pool, &mut queue, &test_ctx()); + + let scopes: Vec<_> = dispatched.iter().map(|(s, _)| s.clone()).collect(); + assert_eq!(scopes, vec![key(ch, "root-a")]); + assert!(!queue.is_scope_in_flight(&key(ch, "root-b"))); + assert_eq!(queue.queued_event_count(&key(ch, "root-b")), 1); + // Next pass with a freed agent picks root-b up. + pool.agents_mut()[0] = Some(dummy_agent(0).await); + let next = dispatch_pending(&mut pool, &mut queue, &test_ctx()); + // Note: root-a is still in flight, so only root-b dispatches. + let next_scopes: Vec<_> = next.iter().map(|(s, _)| s.clone()).collect(); + assert_eq!(next_scopes, vec![key(ch, "root-b")]); + } + + // ── control-scope resolution ───────────────────────────────────────── + + fn insert_in_flight_scope(pool: &mut AgentPool, scope: ConversationSessionKey) { + let task_id = pool.join_set.spawn(std::future::pending()).id(); + let agent_index = pool.task_map().len(); + pool.task_map_mut().insert( + task_id, + pool::TaskMeta { + agent_index, + scope: Some(scope), + turn_id: Uuid::new_v4().to_string(), + recoverable_batch: None, + control_tx: None, + steer_tx: None, + }, + ); + } + + #[tokio::test] + async fn channel_control_resolves_only_when_exactly_one_scope_in_flight() { + let mut pool = AgentPool::from_slots(vec![]); + let ch = Uuid::new_v4(); + + // No turns: nothing to act on. + assert_eq!(resolve_channel_control_scope(&pool, ch), Ok(None)); + + // Exactly one: unambiguous. + insert_in_flight_scope(&mut pool, key(ch, "root-a")); + assert_eq!( + resolve_channel_control_scope(&pool, ch), + Ok(Some(key(ch, "root-a"))) + ); + + // Two scopes in the channel: ambiguous, touch nothing. + insert_in_flight_scope(&mut pool, key(ch, "root-b")); + assert_eq!(resolve_channel_control_scope(&pool, ch), Err(2)); + + // A busy scope in a DIFFERENT channel never bleeds in. + let other = Uuid::new_v4(); + assert_eq!(resolve_channel_control_scope(&pool, other), Ok(None)); + } + + #[tokio::test] + async fn control_frame_root_event_id_targets_one_scope_exactly() { + let mut pool = AgentPool::from_slots(vec![]); + let ch = Uuid::new_v4(); + insert_in_flight_scope(&mut pool, key(ch, "root-a")); + insert_in_flight_scope(&mut pool, key(ch, "root-b")); + + // Explicit root bypasses the ambiguity rule. + assert_eq!( + resolve_control_frame_scope(&pool, ch, Some("root-b".into())), + Ok(Some(key(ch, "root-b"))) + ); + // Without it, two in-flight scopes are ambiguous. + assert_eq!(resolve_control_frame_scope(&pool, ch, None), Err(2)); + } + + #[tokio::test] + async fn threaded_event_control_targets_its_own_root() { + let mut pool = AgentPool::from_slots(vec![]); + let ch = Uuid::new_v4(); + insert_in_flight_scope(&mut pool, key(ch, "aaaa")); + insert_in_flight_scope(&mut pool, key(ch, "bbbb")); + + let mut config = build_mcp_servers_tests::test_config(); + config.session_scope = config::SessionScope::Thread; + + // A reply inside thread bbbb targets bbbb exactly — never fans out. + let keys = Keys::generate(); + let threaded = EventBuilder::new(Kind::Custom(9), "!cancel") + .tags([nostr::Tag::parse(["e", "bbbb", "", "root"]).unwrap()]) + .sign_with_keys(&keys) + .unwrap(); + assert_eq!( + resolve_event_control_scope(&pool, &config, &threaded, ch), + Ok(Some(key(ch, "bbbb"))) + ); + + // A top-level !cancel with two threads in flight is ambiguous. + let top_level = EventBuilder::new(Kind::Custom(9), "!cancel") + .tags([]) + .sign_with_keys(&keys) + .unwrap(); + assert_eq!( + resolve_event_control_scope(&pool, &config, &top_level, ch), + Err(2) + ); + + // Channel mode: always the channel scope, regardless of threads. + config.session_scope = config::SessionScope::Channel; + assert_eq!( + resolve_event_control_scope(&pool, &config, &threaded, ch), + Ok(Some(ConversationSessionKey::channel(ch))) + ); + } +} + #[cfg(test)] mod observer_payload_trim_tests { use super::*; diff --git a/crates/buzz-acp/src/pool.rs b/crates/buzz-acp/src/pool.rs index 25c41a351..7a8cf5e7f 100644 --- a/crates/buzz-acp/src/pool.rs +++ b/crates/buzz-acp/src/pool.rs @@ -40,13 +40,18 @@ use crate::queue::{ }; use crate::relay::{ChannelInfo, RestClient}; +// Scheduling/session identity lives in queue.rs next to the scope-keyed +// queue state machine; re-exported here so pool consumers keep one name. +pub use crate::queue::ConversationSessionKey; + // FlushBatch and BatchEvent derive Clone (added in queue.rs) so we can store // a recoverable copy in TaskMeta for panic recovery in Queue mode. /// Metadata stored per in-flight task for panic recovery. pub struct TaskMeta { pub agent_index: usize, - pub channel_id: Option, + /// Conversation scope of the in-flight prompt; `None` for heartbeats. + pub scope: Option, /// Identifies terminal events when the task panics before returning a result. pub turn_id: String, /// Clone of batch for Queue mode panic recovery. @@ -74,22 +79,6 @@ pub struct AgentModelCapabilities { pub available_models_raw: Option, } -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct ConversationSessionKey { - pub channel_id: Uuid, - pub root_event_id: Option, -} - -impl ConversationSessionKey { - #[cfg(test)] - pub fn channel(channel_id: Uuid) -> Self { - Self { - channel_id, - root_event_id: None, - } - } -} - /// Per-channel session IDs and turn counters. /// /// Separated from `OwnedAgent` so the state machine is testable without @@ -579,25 +568,32 @@ impl AgentPool { } } - /// Try to claim an idle agent for the given channel (or heartbeat if `None`). + /// Try to claim an idle agent for the given scope (or heartbeat if `None`). /// - /// Pass 1: prefer an agent that already has a session for `channel_id`. + /// Pass 1: prefer an agent that already has a session for the scope. /// Pass 2: any idle agent. /// - /// Returns `None` if all agents are checked out. - pub fn try_claim( - &mut self, - session_key: Option<&ConversationSessionKey>, - ) -> Option { + /// Retained roots have strict slot affinity: if the root's owning slot is + /// busy, the claim returns [`ClaimOutcome::BusyOwner`] so the dispatcher + /// can skip this scope and keep dispatching other scopes — a busy owner + /// must not stall unrelated roots, and the root must never migrate to + /// another slot (that would fork the provider session). + /// [`ClaimOutcome::Exhausted`] means no idle agent exists at all, so the + /// dispatch loop should stop. + pub fn try_claim(&mut self, session_key: Option<&ConversationSessionKey>) -> ClaimOutcome { // Pass 1: retained roots have strict ownership. If their slot is busy, // 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) { - // 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); + // Surviving the prune means the owner is busy (skip this + // scope, keep dispatching others) or idle and retaining + // the session. + return match self.agents.get_mut(owner).and_then(Option::take) { + Some(agent) => ClaimOutcome::Claimed(agent), + None => ClaimOutcome::BusyOwner, + }; } } let idx = self.agents.iter().position(|slot| { @@ -607,17 +603,21 @@ impl AgentPool { }); if let Some(i) = idx { self.session_owners.insert(key.clone(), i); - return self.agents[i].take(); + return ClaimOutcome::Claimed( + self.agents[i].take().expect("position matched idle slot"), + ); } } // Pass 2: first idle agent. Reserve root affinity immediately so another // turn for the same root cannot fall through while this slot is checked out. - let idx = self.agents.iter().position(|slot| slot.is_some())?; + let Some(idx) = self.agents.iter().position(|slot| slot.is_some()) else { + return ClaimOutcome::Exhausted; + }; if let Some(key) = session_key.filter(|key| key.root_event_id.is_some()) { self.session_owners.insert(key.clone(), idx); } - self.agents[idx].take() + ClaimOutcome::Claimed(self.agents[idx].take().expect("position matched idle slot")) } /// Drop root reservations that no longer bind: the owning slot is neither @@ -703,19 +703,19 @@ impl AgentPool { /// watcher, to close the result-vs-ack race. /// /// Returns `Err(SteerError::PromptCompleted)` if no task is in flight - /// for `channel_id` (the prompt completed between the mode-gate check - /// and this call, or the channel was never in flight). This is + /// for `scope` (the prompt completed between the mode-gate check + /// and this call, or the scope was never in flight). This is /// semantically a soft no-op — the caller should release any withheld /// event and let normal dispatch handle delivery. pub fn send_steer( &mut self, - channel_id: Uuid, + scope: &ConversationSessionKey, request: SteerRequest, ) -> Result<(), SteerError> { let meta = self .task_map .values_mut() - .find(|m| m.channel_id == Some(channel_id)) + .find(|m| m.scope.as_ref() == Some(scope)) .ok_or(SteerError::PromptCompleted)?; let tx = meta .steer_tx @@ -855,6 +855,33 @@ impl AgentPool { } } +/// Outcome of [`AgentPool::try_claim`]. +// The variant size skew is fine: a ClaimOutcome is matched and consumed +// immediately at the claim site, never stored — boxing would add churn +// on the hot dispatch path for no benefit. +#[allow(clippy::large_enum_variant)] +pub enum ClaimOutcome { + /// An idle agent was checked out for the scope. + Claimed(OwnedAgent), + /// The scope is a retained root whose owning slot is busy on another + /// turn. Skip this scope — do NOT stop the dispatch loop, and do NOT + /// run the root on a different slot (that would fork its session). + BusyOwner, + /// No idle agent in the pool. Stop the dispatch loop. + Exhausted, +} + +impl ClaimOutcome { + /// The claimed agent, if any. Test-only convenience. + #[cfg(test)] + pub fn claimed(self) -> Option { + match self { + ClaimOutcome::Claimed(agent) => Some(agent), + ClaimOutcome::BusyOwner | ClaimOutcome::Exhausted => None, + } + } +} + /// Outcome of [`AgentPool::switch_idle_agent_model`]. #[derive(Debug, PartialEq, Eq)] pub enum IdleSwitchResult { @@ -1340,16 +1367,6 @@ fn send_prompt_result( }); } -/// 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: batch.conversation_root.clone(), - } -} - /// Core async function spawned for each prompt. /// /// Lifecycle: @@ -1373,7 +1390,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)), + Some(b) => PromptSource::Channel(b.scope_key()), None => PromptSource::Heartbeat, }; let observer_channel_id = match &source { @@ -4557,7 +4574,7 @@ mod tests { #[test] fn test_conversation_session_key_is_channel_scoped_without_root() { let batch = one_event_batch(Uuid::new_v4()); - let key = conversation_session_key(&batch); + let key = batch.scope_key(); assert_eq!(key.channel_id, batch.channel_id); assert_eq!(key.root_event_id, None); } @@ -4566,7 +4583,7 @@ mod tests { 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); + let key = batch.scope_key(); assert_eq!(key.channel_id, batch.channel_id); assert_eq!(key.root_event_id.as_deref(), Some("root-a")); } diff --git a/crates/buzz-acp/src/queue.rs b/crates/buzz-acp/src/queue.rs index a07ce5278..7f12197ff 100644 --- a/crates/buzz-acp/src/queue.rs +++ b/crates/buzz-acp/src/queue.rs @@ -20,8 +20,8 @@ use uuid::Uuid; use crate::config::DedupMode; -/// Maximum events queued per channel before oldest events are dropped. -const MAX_PENDING_PER_CHANNEL: usize = 500; +/// Maximum events queued per conversation scope before oldest events are dropped. +const MAX_PENDING_PER_SCOPE: usize = 500; /// Maximum events drained into a single batch. const MAX_BATCH_EVENTS: usize = 50; @@ -41,6 +41,29 @@ const IN_FLIGHT_DEADLINE_BUFFER_SECS: u64 = 100; /// Default in-flight deadline: default max_turn (7200s) + 100s buffer. const DEFAULT_IN_FLIGHT_DEADLINE_SECS: u64 = 7300; +/// Scheduling and session identity for one conversation scope. +/// +/// `root_event_id = None` → channel scope (legacy): the whole channel is one +/// scope. `Some(root)` → thread scope: one scope per outermost thread root. +/// All queue bookkeeping, in-flight tracking, retries, session affinity, and +/// control routing key off this — channel-scoped mode is simply the +/// degenerate single-scope-per-channel case. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct ConversationSessionKey { + pub channel_id: Uuid, + pub root_event_id: Option, +} + +impl ConversationSessionKey { + /// Channel-scoped key (legacy mode / heartbeat-adjacent paths). + pub fn channel(channel_id: Uuid) -> Self { + Self { + channel_id, + root_event_id: None, + } + } +} + /// An event waiting in the queue. #[derive(Debug, Clone)] pub struct QueuedEvent { @@ -49,10 +72,20 @@ pub struct QueuedEvent { pub received_at: Instant, /// Tag identifying which rule (or mode) matched this event. pub prompt_tag: String, - /// Conversation root used only by the top-level-session experiment. + /// Outermost thread root when thread-scoped sessions are enabled; + /// `None` in channel-scoped mode. pub conversation_root: Option, } +impl QueuedEvent { + fn scope_key(&self) -> ConversationSessionKey { + ConversationSessionKey { + channel_id: self.channel_id, + root_event_id: self.conversation_root.clone(), + } + } +} + /// A single event inside a [`FlushBatch`]. #[derive(Debug, Clone)] pub struct BatchEvent { @@ -92,84 +125,100 @@ pub struct FlushBatch { pub cancel_reason: Option, } +impl FlushBatch { + /// The conversation scope this batch belongs to. + pub fn scope_key(&self) -> ConversationSessionKey { + ConversationSessionKey { + channel_id: self.channel_id, + root_event_id: self.conversation_root.clone(), + } + } +} + #[derive(Debug, Clone)] struct CancelledBatch { - conversation_root: Option, events: Vec, reason: CancelReason, } -/// Per-channel event queue with per-channel in-flight enforcement. +/// Per-scope event queue with per-scope in-flight enforcement. +/// +/// A "scope" is a [`ConversationSessionKey`]: the whole channel in channel +/// mode (`root_event_id = None`), or one outermost thread root in thread +/// mode. Every piece of bookkeeping below is scope-keyed, so channel mode is +/// the degenerate one-scope-per-channel case and needs no special paths. /// /// # State Machine /// /// ```text /// State: -/// queues: Map> (capped at MAX_PENDING_PER_CHANNEL) -/// in_flight_channels: HashSet -/// in_flight_deadlines: Map (auto-expire after in_flight_deadline) -/// retry_after: Map -/// retry_counts: Map (dead-letter after MAX_RETRIES) +/// queues: Map> (capped at MAX_PENDING_PER_SCOPE) +/// in_flight_scopes: HashSet +/// in_flight_deadlines: Map (auto-expire after in_flight_deadline) +/// retry_after: Map +/// retry_counts: Map (dead-letter after MAX_RETRIES) /// dedup_mode: DedupMode /// /// Transitions: /// push(event): -/// if dedup_mode == Drop AND in_flight_channels.contains(event.channel_id): +/// if dedup_mode == Drop AND in_flight_scopes.contains(event.scope): /// debug log + discard -/// else if queues[channel].len() >= MAX_PENDING_PER_CHANNEL: +/// else if queues[scope].len() >= MAX_PENDING_PER_SCOPE: /// drop oldest (pop_front), warn, push_back new event /// else: -/// queues[event.channel_id].push_back(event) +/// queues[event.scope].push_back(event) /// /// flush_next() → Option: /// expire any stuck in-flight entries past their deadline -/// candidates = channels where queue non-empty -/// AND NOT in in_flight_channels -/// AND (no retry_after OR retry_after[c] <= now) +/// candidates = scopes where queue non-empty +/// AND NOT in in_flight_scopes +/// AND (no retry_after OR retry_after[s] <= now) /// if candidates empty: return None -/// channel = pick candidate with oldest head event (min received_at) -/// events = drain up to MAX_BATCH_EVENTS from queues[channel] -/// in_flight_channels.insert(channel) -/// in_flight_deadlines.insert(channel, now + in_flight_deadline) -/// return Some(FlushBatch { channel, events }) +/// scope = pick candidate with oldest head event (min received_at) +/// events = drain up to MAX_BATCH_EVENTS from queues[scope] +/// in_flight_scopes.insert(scope) +/// in_flight_deadlines.insert(scope, now + in_flight_deadline) +/// return Some(FlushBatch { scope, events }) /// -/// mark_complete(channel_id): -/// in_flight_channels.remove(channel_id) -/// in_flight_deadlines.remove(channel_id) -/// retry_counts.remove(channel_id) +/// mark_complete(scope): +/// in_flight_scopes.remove(scope) +/// in_flight_deadlines.remove(scope) +/// retry_counts.remove(scope) /// clean up expired retry_after entry if present /// /// requeue(batch): -/// increment retry_counts[channel] -/// if retry_counts[channel] > MAX_RETRIES: dead-letter (log ERROR, return batch to caller) +/// increment retry_counts[scope] +/// if retry_counts[scope] > MAX_RETRIES: dead-letter (log ERROR, return batch to caller) /// else: push_front with original received_at, set exponential backoff retry_after with jitter /// ``` pub struct EventQueue { - queues: HashMap>, - in_flight_channels: HashSet, - /// Per-channel deadline for auto-expiring stuck in-flight entries. - in_flight_deadlines: HashMap, + queues: HashMap>, + in_flight_scopes: HashSet, + /// Per-scope deadline for auto-expiring stuck in-flight entries. + in_flight_deadlines: HashMap, /// Number of events in each in-flight batch (for expiry logging). - in_flight_batch_sizes: HashMap, - retry_after: HashMap, - /// Per-channel retry attempt counter for exponential backoff / dead-lettering. - retry_counts: HashMap, + in_flight_batch_sizes: HashMap, + retry_after: HashMap, + /// Per-scope retry attempt counter for exponential backoff / dead-lettering. + retry_counts: HashMap, dedup_mode: DedupMode, - /// Cancelled batches awaiting redispatch, preserving conversation identity. - /// Multiple roots may wait behind the same channel, but are never merged. - cancelled_batches: HashMap>, + /// Cancelled batches awaiting redispatch. Scope-keyed: batches from + /// different roots live under different keys and are never merged; + /// repeat cancels for the SAME scope merge into one batch (most recent + /// reason wins), so each scope holds at most one entry. + cancelled_batches: HashMap, /// Events withheld from `queues` while a goose-native steer is in flight /// for that event. Invisible to `flush_next` / `has_flushable_work` / /// `drain` (the events have been moved out of `queues`), so the queue's /// no-double-deliver invariant holds without any change to the hot drain /// path. Populated by [`mark_native_steer_pending`]; drained back to the /// queue front by [`release_native_steer`] (preserving original - /// `received_at` fairness, same discipline as `requeue_preserve_timestamps` - /// at line 453). Bulk recovery on in-flight deadline expiry is performed - /// by `flush_next` / `has_flushable_work` (recover, not log-and-drop — + /// `received_at` fairness, same discipline as `requeue_preserve_timestamps`). + /// Bulk recovery on in-flight deadline expiry is performed by + /// `flush_next` / `has_flushable_work` (recover, not log-and-drop — /// the events were never delivered to the agent). - withheld_native_steer: HashMap>, - /// Duration after which an in-flight channel is auto-expired as orphaned. + withheld_native_steer: HashMap>, + /// Duration after which an in-flight scope is auto-expired as orphaned. /// Must be strictly greater than `max_turn_duration` so a turn running to /// the hard cap returns via `mark_complete` before the backstop fires. in_flight_deadline: Duration, @@ -184,7 +233,7 @@ impl EventQueue { pub fn new(dedup_mode: DedupMode) -> Self { Self { queues: HashMap::new(), - in_flight_channels: HashSet::new(), + in_flight_scopes: HashSet::new(), in_flight_deadlines: HashMap::new(), in_flight_batch_sizes: HashMap::new(), retry_after: HashMap::new(), @@ -204,29 +253,29 @@ impl EventQueue { self } - /// Push an event into the queue for its channel. + /// Push an event into the queue for its conversation scope. /// - /// In [`DedupMode::Drop`], events for any currently in-flight channel are + /// In [`DedupMode::Drop`], events for any currently in-flight scope are /// silently discarded (debug-logged). /// /// Returns `true` if the event was accepted, `false` if dropped. pub fn push(&mut self, event: QueuedEvent) -> bool { - if matches!(self.dedup_mode, DedupMode::Drop) - && self.in_flight_channels.contains(&event.channel_id) - { + let scope = event.scope_key(); + if matches!(self.dedup_mode, DedupMode::Drop) && self.in_flight_scopes.contains(&scope) { tracing::debug!( channel_id = %event.channel_id, - "dropping event for in-flight channel (drop mode)" + root = scope.root_event_id.as_deref().unwrap_or(""), + "dropping event for in-flight scope (drop mode)" ); return false; } - let queue = self.queues.entry(event.channel_id).or_default(); - // Enforce per-channel depth cap: drop oldest to make room. - if queue.len() >= MAX_PENDING_PER_CHANNEL { + let queue = self.queues.entry(scope).or_default(); + // Enforce per-scope depth cap: drop oldest to make room. + if queue.len() >= MAX_PENDING_PER_SCOPE { queue.pop_front(); tracing::warn!( channel_id = %event.channel_id, - limit = MAX_PENDING_PER_CHANNEL, + limit = MAX_PENDING_PER_SCOPE, "queue depth cap reached — dropped oldest event" ); } @@ -237,95 +286,46 @@ impl EventQueue { /// Try to flush the next batch. /// /// Returns `None` if all non-in-flight, non-throttled queues are empty. - /// Otherwise picks the channel with the oldest pending event (FIFO fairness - /// across channels), drains ALL events for that channel into a single batch, - /// inserts into `in_flight_channels`, and returns the batch. + /// Otherwise picks the scope with the oldest pending event (FIFO fairness + /// across scopes), drains up to [`MAX_BATCH_EVENTS`] for that scope into a + /// single batch, inserts into `in_flight_scopes`, and returns the batch. + /// Every event in a scope's queue shares the same conversation root by + /// construction, so a batch can never mix roots. pub fn flush_next(&mut self) -> Option { let now = Instant::now(); + self.expire_stuck_in_flight(now); - // Auto-expire any stuck in-flight entries that missed mark_complete. - let expired: Vec = self - .in_flight_deadlines - .iter() - .filter(|(_, deadline)| now >= **deadline) - .map(|(id, _)| *id) - .collect(); - for id in expired { - let lost_events = self.in_flight_batch_sizes.remove(&id).unwrap_or(0); - tracing::error!( - channel_id = %id, - lost_events, - deadline_secs = self.in_flight_deadline.as_secs(), - "BUG: in-flight channel expired without mark_complete — \ - auto-releasing; {lost_events} dispatched event(s) orphaned" - ); - self.in_flight_channels.remove(&id); - self.in_flight_deadlines.remove(&id); - // Recover any withheld goose-native steer events for the expired - // channel back to the queue front so normal dispatch delivers - // them. Unlike the in-flight batch above (already delivered to a - // now-hung prompt — nothing to recover), these events were never - // delivered to the agent. - self.recover_withheld_for_expired_channel(id); - } - - // Find the channel whose head event has the oldest received_at, - // excluding in-flight channels and throttled channels. - let channel_id = self + // Find the scope whose head event has the oldest received_at, + // excluding in-flight scopes and throttled scopes. + let scope = self .queues .iter() - .filter(|(id, q)| { + .filter(|(key, q)| { !q.is_empty() - && !self.in_flight_channels.contains(id) - && self.retry_after.get(id).is_none_or(|&t| t <= now) + && !self.in_flight_scopes.contains(key) + && self.retry_after.get(key).is_none_or(|&t| t <= now) }) .min_by_key(|(_, q)| q.front().unwrap().received_at) - .map(|(id, _)| *id); + .map(|(key, _)| key.clone()); - // Fallback: if no queued events are ready but a channel has cancelled + // Fallback: if no queued events are ready but a scope has cancelled // events waiting (e.g., explicit !cancel with no new @mention), flush // those as a regular batch (re-dispatch unchanged). - let channel_id = match channel_id { - Some(id) => id, + let scope = match scope { + Some(key) => key, None => { - let cancelled_id = self + let cancelled_key = self .cancelled_batches .keys() - .find(|id| !self.in_flight_channels.contains(id)) - .copied(); - return cancelled_id.map(|id| self.dispatch_cancelled(id, now)); + .find(|key| !self.in_flight_scopes.contains(key)) + .cloned(); + return cancelled_key.map(|key| self.dispatch_cancelled(key, 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) - .and_then(|queue| queue.front()) - .and_then(|event| event.conversation_root.clone()); - let cancelled_root = self - .cancelled_batches - .get(&channel_id) - .and_then(|batches| batches.front()) - .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. - let queue = self.queues.entry(channel_id).or_default(); - // Never merge independent experiment roots into one ACP turn. Drain only - // the contiguous head conversation while retaining channel serialization. - let head_root = queue - .front() - .and_then(|event| event.conversation_root.clone()); - let same_root_count = queue - .iter() - .take_while(|event| event.conversation_root == head_root) - .count(); - let drain_count = MAX_BATCH_EVENTS.min(same_root_count); + let queue = self.queues.entry(scope.clone()).or_default(); + let drain_count = MAX_BATCH_EVENTS.min(queue.len()); let mut events: Vec = queue .drain(..drain_count) .map(|qe| BatchEvent { @@ -341,87 +341,112 @@ impl EventQueue { events.sort_by_key(|be| be.event.created_at); // Remove the queue entry if now empty. - if self.queues.get(&channel_id).is_some_and(|q| q.is_empty()) { - self.queues.remove(&channel_id); + if self.queues.get(&scope).is_some_and(|q| q.is_empty()) { + self.queues.remove(&scope); } - self.in_flight_channels.insert(channel_id); + self.in_flight_scopes.insert(scope.clone()); self.in_flight_deadlines - .insert(channel_id, now + self.in_flight_deadline); - self.in_flight_batch_sizes.insert(channel_id, events.len()); + .insert(scope.clone(), now + self.in_flight_deadline); + self.in_flight_batch_sizes + .insert(scope.clone(), events.len()); - // Merge only a cancelled batch for the same conversation root. - let (cancelled_events, cancel_reason) = self.pop_cancelled(channel_id).map_or_else( + // Merge any cancelled batch waiting under this same scope. + let (cancelled_events, cancel_reason) = self.pop_cancelled(&scope).map_or_else( || (vec![], None), |batch| (batch.events, Some(batch.reason)), ); Some(FlushBatch { - channel_id, - conversation_root: head_root, + channel_id: scope.channel_id, + conversation_root: scope.root_event_id, events, cancelled_events, cancel_reason, }) } - /// 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); + /// Auto-expire any stuck in-flight entries that missed `mark_complete`. + /// Shared by `flush_next` and `has_flushable_work`. + fn expire_stuck_in_flight(&mut self, now: Instant) { + let expired: Vec = self + .in_flight_deadlines + .iter() + .filter(|(_, deadline)| now >= **deadline) + .map(|(key, _)| key.clone()) + .collect(); + for key in expired { + let lost_events = self.in_flight_batch_sizes.remove(&key).unwrap_or(0); + tracing::error!( + channel_id = %key.channel_id, + root = key.root_event_id.as_deref().unwrap_or(""), + lost_events, + deadline_secs = self.in_flight_deadline.as_secs(), + "BUG: in-flight scope expired without mark_complete — \ + auto-releasing; {lost_events} dispatched event(s) orphaned" + ); + self.in_flight_scopes.remove(&key); + self.in_flight_deadlines.remove(&key); + // Recover any withheld goose-native steer events for the expired + // scope back to the queue front so normal dispatch delivers + // them. Unlike the in-flight batch above (already delivered to a + // now-hung prompt — nothing to recover), these events were never + // delivered to the agent. + self.recover_withheld_for_expired_scope(&key); } - 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 { + /// Take the cancelled batch for `scope`, if any. + fn pop_cancelled(&mut self, scope: &ConversationSessionKey) -> Option { + self.cancelled_batches.remove(scope) + } + + /// Re-dispatch the oldest cancelled batch for `scope` unchanged, + /// marking the scope in-flight. + fn dispatch_cancelled(&mut self, scope: ConversationSessionKey, now: Instant) -> FlushBatch { let cancelled = self - .pop_cancelled(channel_id) - .expect("cancelled channel must have a batch"); - self.in_flight_channels.insert(channel_id); + .pop_cancelled(&scope) + .expect("cancelled scope must have a batch"); + self.in_flight_scopes.insert(scope.clone()); self.in_flight_deadlines - .insert(channel_id, now + self.in_flight_deadline); + .insert(scope.clone(), now + self.in_flight_deadline); self.in_flight_batch_sizes - .insert(channel_id, cancelled.events.len()); + .insert(scope.clone(), cancelled.events.len()); FlushBatch { - channel_id, - conversation_root: cancelled.conversation_root, + channel_id: scope.channel_id, + conversation_root: scope.root_event_id, events: cancelled.events, cancelled_events: vec![], cancel_reason: Some(cancelled.reason), } } - /// Mark the prompt for `channel_id` as complete. + /// Mark the prompt for `scope` as complete. /// - /// Removes the channel from `in_flight_channels` and `in_flight_deadlines`. + /// Removes the scope from `in_flight_scopes` and `in_flight_deadlines`. /// - /// If the channel was NOT requeued (no active `retry_after` throttle), the - /// retry counter is reset — the channel is healthy and the next failure - /// starts fresh. If the channel WAS requeued, `retry_counts` is left intact + /// If the scope was NOT requeued (no active `retry_after` throttle), the + /// retry counter is reset — the scope is healthy and the next failure + /// starts fresh. If the scope WAS requeued, `retry_counts` is left intact /// so the backoff sequence continues on the next attempt. /// /// Also cleans up any already-expired `retry_after` entry. - pub fn mark_complete(&mut self, channel_id: Uuid) { - self.in_flight_channels.remove(&channel_id); - self.in_flight_deadlines.remove(&channel_id); - self.in_flight_batch_sizes.remove(&channel_id); + pub fn mark_complete(&mut self, scope: &ConversationSessionKey) { + self.in_flight_scopes.remove(scope); + self.in_flight_deadlines.remove(scope); + self.in_flight_batch_sizes.remove(scope); let now = Instant::now(); - match self.retry_after.get(&channel_id) { - // Active throttle → channel was requeued; keep retry_counts intact. + match self.retry_after.get(scope) { + // Active throttle → scope was requeued; keep retry_counts intact. Some(&deadline) if deadline > now => {} // Expired or absent throttle → successful completion; reset counter // and clean up the stale retry_after entry. Some(_) => { - self.retry_after.remove(&channel_id); - self.retry_counts.remove(&channel_id); + self.retry_after.remove(scope); + self.retry_counts.remove(scope); } None => { - self.retry_counts.remove(&channel_id); + self.retry_counts.remove(scope); } } } @@ -441,13 +466,14 @@ impl EventQueue { /// failure notice can be posted to the channel. Returns `None` when the /// batch was requeued for another attempt. /// - /// Note: does NOT remove from `in_flight_channels` — caller must call + /// Note: does NOT remove from `in_flight_scopes` — caller must call /// `mark_complete` separately. pub fn requeue(&mut self, batch: FlushBatch) -> Option { + let scope = batch.scope_key(); let channel_id = batch.channel_id; let conversation_root = batch.conversation_root.clone(); let attempt = { - let count = self.retry_counts.entry(channel_id).or_insert(0); + let count = self.retry_counts.entry(scope.clone()).or_insert(0); *count += 1; *count }; @@ -455,16 +481,17 @@ impl EventQueue { if attempt > MAX_RETRIES { tracing::error!( channel_id = %channel_id, + root = scope.root_event_id.as_deref().unwrap_or(""), attempt, events = batch.events.len(), "dead-lettering batch after {} retries — discarding {} events", MAX_RETRIES, batch.events.len(), ); - self.retry_counts.remove(&channel_id); - // Also clear retry_after so fresh traffic on this channel isn't + self.retry_counts.remove(&scope); + // Also clear retry_after so fresh traffic on this scope isn't // throttled by stale backoff from the discarded poison batch. - self.retry_after.remove(&channel_id); + self.retry_after.remove(&scope); return Some(batch); } @@ -490,7 +517,7 @@ impl EventQueue { "requeueing failed batch with backoff" ); - let queue = self.queues.entry(channel_id).or_default(); + let queue = self.queues.entry(scope.clone()).or_default(); // Push to front in reverse order so original order is preserved. for be in batch.events.into_iter().rev() { queue.push_front(QueuedEvent { @@ -501,18 +528,18 @@ impl EventQueue { conversation_root: conversation_root.clone(), }); } - // Enforce per-channel cap: trim oldest (back) events if requeue pushed + // Enforce per-scope cap: trim oldest (back) events if requeue pushed // the queue over the limit. Without this, repeated requeue+push cycles // can grow the queue unboundedly. - while queue.len() > MAX_PENDING_PER_CHANNEL { + while queue.len() > MAX_PENDING_PER_SCOPE { queue.pop_back(); tracing::warn!( channel_id = %channel_id, - limit = MAX_PENDING_PER_CHANNEL, + limit = MAX_PENDING_PER_SCOPE, "requeue overflow — dropped oldest event to enforce cap" ); } - self.retry_after.insert(channel_id, Instant::now() + delay); + self.retry_after.insert(scope, Instant::now() + delay); None } @@ -522,12 +549,13 @@ impl EventQueue { /// retry without penalizing the channel's position in the fairness queue /// and without imposing a retry throttle. /// - /// Does NOT set `retry_after`. Does NOT remove from `in_flight_channels` — + /// Does NOT set `retry_after`. Does NOT remove from `in_flight_scopes` — /// caller must call `mark_complete` separately. pub fn requeue_preserve_timestamps(&mut self, batch: FlushBatch) { + let scope = batch.scope_key(); let channel_id = batch.channel_id; let conversation_root = batch.conversation_root.clone(); - let queue = self.queues.entry(channel_id).or_default(); + let queue = self.queues.entry(scope).or_default(); // Push to front in reverse order so original order is preserved. for be in batch.events.into_iter().rev() { queue.push_front(QueuedEvent { @@ -538,52 +566,86 @@ impl EventQueue { conversation_root: conversation_root.clone(), }); } - // Enforce per-channel cap: trim newest (back) events if over limit. - while queue.len() > MAX_PENDING_PER_CHANNEL { + // Enforce per-scope cap: trim newest (back) events if over limit. + while queue.len() > MAX_PENDING_PER_SCOPE { queue.pop_back(); tracing::warn!( channel_id = %channel_id, - limit = MAX_PENDING_PER_CHANNEL, + limit = MAX_PENDING_PER_SCOPE, "requeue_preserve overflow — dropped newest event to enforce cap" ); } } + /// Restore a flushed batch that could not be dispatched (no claimable + /// agent) to its exact pre-flush state: merged cancelled events return + /// to the cancelled store under their original reason, fresh events + /// return to the queue front with original timestamps. No retry + /// accounting. Caller must still call `mark_complete` to release the + /// in-flight entry. + /// + /// This is the undo of `flush_next` — `requeue_preserve_timestamps` + /// alone would silently drop `cancelled_events` (or strip the cancel + /// framing from a re-dispatched cancelled batch). + pub fn restore_unclaimed(&mut self, mut batch: FlushBatch) { + match batch.cancel_reason { + // Re-dispatch of a cancelled batch (`dispatch_cancelled` puts + // the cancelled events in `events`): return it whole to the + // cancelled store. + Some(reason) if batch.cancelled_events.is_empty() => { + self.requeue_as_cancelled(batch, reason); + } + // Merged batch: split back — cancelled part to the cancelled + // store, fresh part to the queue. + Some(reason) => { + let cancelled = std::mem::take(&mut batch.cancelled_events); + self.requeue_as_cancelled( + FlushBatch { + channel_id: batch.channel_id, + conversation_root: batch.conversation_root.clone(), + events: cancelled, + cancelled_events: vec![], + cancel_reason: Some(reason), + }, + reason, + ); + self.requeue_preserve_timestamps(batch); + } + None => self.requeue_preserve_timestamps(batch), + } + } + /// Requeue a cancelled batch so its events appear as `cancelled_events` - /// in the next `FlushBatch` for this channel (enabling the annotated + /// in the next `FlushBatch` for its scope (enabling the annotated /// merged-prompt format in `format_prompt()`). /// /// `reason` records why the turn was cancelled (steer vs interrupt) so the - /// merged prompt is framed correctly. On a double-cancel, the most recent - /// reason wins. + /// merged prompt is framed correctly. On a double-cancel for the same + /// scope, events accumulate and the most recent reason wins. /// /// Unlike `requeue_preserve_timestamps`, events are NOT pushed back into /// the generic queue — they are stored separately and merged by /// `flush_next()`. No retry throttle, no backoff. pub fn requeue_as_cancelled(&mut self, batch: FlushBatch, reason: CancelReason) { - let channel_id = batch.channel_id; - let conversation_root = batch.conversation_root; + let scope = batch.scope_key(); let mut events = batch.cancelled_events; events.extend(batch.events); - let batches = self.cancelled_batches.entry(channel_id).or_default(); - if let Some(existing) = batches - .iter_mut() - .find(|existing| existing.conversation_root == conversation_root) - { - // Preserve any already-cancelled events for this same root. The most - // recent cancellation reason wins, matching prior double-cancel behavior. - existing.events.extend(events); - existing.reason = reason; - } else { - batches.push_back(CancelledBatch { - conversation_root, - events, - reason, - }); + match self.cancelled_batches.entry(scope) { + std::collections::hash_map::Entry::Occupied(mut existing) => { + // Preserve any already-cancelled events for this same scope. + // The most recent cancellation reason wins, matching prior + // double-cancel behavior. + let existing = existing.get_mut(); + existing.events.extend(events); + existing.reason = reason; + } + std::collections::hash_map::Entry::Vacant(slot) => { + slot.insert(CancelledBatch { events, reason }); + } } } - /// Returns `true` if any channel has pending events that are not in-flight + /// Returns `true` if any scope has pending events that are not in-flight /// and not throttled by `retry_after`. /// /// Also auto-expires any stuck in-flight entries whose deadline has passed. @@ -591,112 +653,101 @@ impl EventQueue { /// full `flush_next` call. pub fn has_flushable_work(&mut self) -> bool { let now = Instant::now(); + self.expire_stuck_in_flight(now); - // Auto-expire stuck in-flight entries (same logic as flush_next). - let expired: Vec = self - .in_flight_deadlines - .iter() - .filter(|(_, deadline)| now >= **deadline) - .map(|(id, _)| *id) - .collect(); - for id in expired { - let lost_events = self.in_flight_batch_sizes.remove(&id).unwrap_or(0); - tracing::error!( - channel_id = %id, - lost_events, - deadline_secs = self.in_flight_deadline.as_secs(), - "BUG: in-flight channel expired without mark_complete — \ - auto-releasing; {lost_events} dispatched event(s) orphaned" - ); - self.in_flight_channels.remove(&id); - self.in_flight_deadlines.remove(&id); - // Symmetric with the flush_next expiry block: recover withheld - // goose-native steer events for the expired channel so they are - // not permanently orphaned in the side table. - self.recover_withheld_for_expired_channel(id); - } - - self.queues.iter().any(|(id, q)| { + self.queues.iter().any(|(key, q)| { !q.is_empty() - && !self.in_flight_channels.contains(id) - && self.retry_after.get(id).is_none_or(|&t| t <= now) + && !self.in_flight_scopes.contains(key) + && self.retry_after.get(key).is_none_or(|&t| t <= now) }) || self .cancelled_batches .keys() - .any(|id| !self.in_flight_channels.contains(id)) + .any(|key| !self.in_flight_scopes.contains(key)) } - /// Number of channels with pending events. + /// Number of scopes with pending events. pub fn pending_channels(&self) -> usize { self.queues.len() } - /// Number of queued events for a specific channel. Test-only. + /// Number of queued events for a specific scope. Test-only. #[cfg(test)] - pub fn queued_event_count(&self, channel_id: &Uuid) -> usize { - self.queues.get(channel_id).map_or(0, |q| q.len()) + pub fn queued_event_count(&self, scope: &ConversationSessionKey) -> usize { + self.queues.get(scope).map_or(0, |q| q.len()) } - /// Drop all queued (non-in-flight) events for a channel. + /// Drop all queued (non-in-flight) events for a channel, across ALL of + /// its scopes. /// /// Used when the agent is removed from a channel — any pending events /// for that channel are stale and should not be prompted. Does NOT /// affect in-flight prompts (those will complete normally; the agent /// may fail to act if it lost access, but that's handled by the relay). /// - /// Also clears any `retry_after` throttle for the channel. + /// Also clears any `retry_after` throttles for the channel's scopes. /// /// Returns the event IDs of dropped events so the caller can clean up /// any reactions (👀) that were added at queue-push time. pub fn drain_channel(&mut self, channel_id: Uuid) -> Vec { - let ids = self - .queues - .remove(&channel_id) - .map(|q| q.into_iter().map(|e| e.event.id.to_hex()).collect()) - .unwrap_or_default(); - self.retry_after.remove(&channel_id); - self.retry_counts.remove(&channel_id); - self.cancelled_batches.remove(&channel_id); - self.withheld_native_steer.remove(&channel_id); - // Preserve in_flight_channels AND in_flight_deadlines: the in-flight - // task will eventually complete (calling mark_complete) or the deadline - // will expire (auto-cleaning the channel). Removing deadlines without - // removing in_flight_channels would disable auto-expiry and leave a - // wedged task permanently blocking the channel. + let mut ids = Vec::new(); + self.queues.retain(|key, q| { + if key.channel_id != channel_id { + return true; + } + ids.extend(q.iter().map(|e| e.event.id.to_hex())); + false + }); + self.retry_after + .retain(|key, _| key.channel_id != channel_id); + self.retry_counts + .retain(|key, _| key.channel_id != channel_id); + self.cancelled_batches + .retain(|key, _| key.channel_id != channel_id); + self.withheld_native_steer + .retain(|key, _| key.channel_id != channel_id); + // Preserve in_flight_scopes AND in_flight_deadlines: each in-flight + // task will eventually complete (calling mark_complete) or its deadline + // will expire (auto-cleaning the scope). Removing deadlines without + // removing in_flight_scopes would disable auto-expiry and leave a + // wedged task permanently blocking the scope. ids } - /// Whether a prompt is currently in-flight for the given channel. - pub fn is_channel_in_flight(&self, channel_id: Uuid) -> bool { - self.in_flight_channels.contains(&channel_id) + /// Whether a prompt is currently in-flight for the given scope. + pub fn is_scope_in_flight(&self, scope: &ConversationSessionKey) -> bool { + self.in_flight_scopes.contains(scope) } // ── Goose-native steer withhold (side table) ────────────────────────── // // While a goose-native `_goose/unstable/session/steer` write is in flight // for a specific queued event, that event is moved out of `queues` into - // `withheld_native_steer` so `flush_next` / `has_flushable_work` / the - // contiguous drain at line 285 cannot see it — closing the race window - // between `mark_complete` (which clears `in_flight_channels`) and the - // ack arriving on the main loop. On `Success` the event is consumed - // (`remove_event`); on `Err` / `PromptCompletedNeutral` it is released - // back to the queue front (`release_native_steer`), preserving its - // original `received_at` for FIFO fairness. + // `withheld_native_steer` so `flush_next` / `has_flushable_work` / + // `drain` cannot see it — closing the race window between + // `mark_complete` (which clears `in_flight_scopes`) and the ack arriving + // on the main loop. On `Success` the event is consumed (`remove_event`); + // on `Err` / `PromptCompletedNeutral` it is released back to the queue + // front (`release_native_steer`), preserving its original `received_at` + // for FIFO fairness. - /// Move a queued event out of `queues[channel_id]` into the side table + /// Move a queued event out of `queues[scope]` into the side table /// to withhold it from `flush_next` while a goose-native steer is in /// flight. /// /// Returns `true` if the event was found and withheld, `false` if the - /// event id was not present in `queues[channel_id]` (race-safe no-op: + /// event id was not present in `queues[scope]` (race-safe no-op: /// the event may have already been drained, removed, or never queued). /// /// Must be called synchronously from the mode-gate fork immediately /// after `pool.send_steer` returns `Ok(())` and before any watcher task /// is spawned, so the withhold is established before `mark_complete` / /// any subsequent `flush_next` tick can run. - pub fn mark_native_steer_pending(&mut self, channel_id: Uuid, event_id: &str) -> bool { - let Some(q) = self.queues.get_mut(&channel_id) else { + pub fn mark_native_steer_pending( + &mut self, + scope: &ConversationSessionKey, + event_id: &str, + ) -> bool { + let Some(q) = self.queues.get_mut(scope) else { return false; }; let Some(pos) = q.iter().position(|qe| qe.event.id.to_hex() == event_id) else { @@ -706,27 +757,27 @@ impl EventQueue { .remove(pos) .expect("position came from iter so remove must succeed"); if q.is_empty() { - self.queues.remove(&channel_id); + self.queues.remove(scope); } self.withheld_native_steer - .entry(channel_id) + .entry(scope.clone()) .or_default() .push(qe); true } /// Release a single withheld event back to the front of - /// `queues[channel_id]`, preserving its original `received_at`. + /// `queues[scope]`, preserving its original `received_at`. /// /// Called on `SteerAck::Err(_)` and `SteerAck::PromptCompletedNeutral` /// (delivery unknown after prompt completion; restoring queued event /// for normal dispatch). Idempotent: a no-op if the event was already /// removed or never withheld. /// - /// Push-to-front matches the discipline of `requeue_preserve_timestamps` - /// at line 453, preserving fairness across channels. - pub fn release_native_steer(&mut self, channel_id: Uuid, event_id: &str) { - let Some(entries) = self.withheld_native_steer.get_mut(&channel_id) else { + /// Push-to-front matches the discipline of `requeue_preserve_timestamps`, + /// preserving fairness across scopes. + pub fn release_native_steer(&mut self, scope: &ConversationSessionKey, event_id: &str) { + let Some(entries) = self.withheld_native_steer.get_mut(scope) else { return; }; let Some(pos) = entries @@ -737,18 +788,18 @@ impl EventQueue { }; let qe = entries.remove(pos); if entries.is_empty() { - self.withheld_native_steer.remove(&channel_id); + self.withheld_native_steer.remove(scope); } // Push to FRONT so original `received_at` keeps the event at the head - // of the channel's queue. Per-channel cap is enforced below in case + // of the scope's queue. Per-scope cap is enforced below in case // a flood of events arrived during the ack window. - let queue = self.queues.entry(channel_id).or_default(); + let queue = self.queues.entry(scope.clone()).or_default(); queue.push_front(qe); - while queue.len() > MAX_PENDING_PER_CHANNEL { + while queue.len() > MAX_PENDING_PER_SCOPE { queue.pop_back(); tracing::warn!( - channel_id = %channel_id, - limit = MAX_PENDING_PER_CHANNEL, + channel_id = %scope.channel_id, + limit = MAX_PENDING_PER_SCOPE, "release_native_steer overflow — dropped newest event to enforce cap" ); } @@ -760,53 +811,54 @@ impl EventQueue { /// Called on `SteerAck::Success` — the agent received the steer, so the /// event has been "delivered" via the non-cancelling path and must not /// be redelivered via normal dispatch. Idempotent across both stores. - pub fn remove_event(&mut self, channel_id: Uuid, event_id: &str) { - if let Some(entries) = self.withheld_native_steer.get_mut(&channel_id) { + pub fn remove_event(&mut self, scope: &ConversationSessionKey, event_id: &str) { + if let Some(entries) = self.withheld_native_steer.get_mut(scope) { entries.retain(|qe| qe.event.id.to_hex() != event_id); if entries.is_empty() { - self.withheld_native_steer.remove(&channel_id); + self.withheld_native_steer.remove(scope); } } - if let Some(q) = self.queues.get_mut(&channel_id) { + if let Some(q) = self.queues.get_mut(scope) { q.retain(|qe| qe.event.id.to_hex() != event_id); if q.is_empty() { - self.queues.remove(&channel_id); + self.queues.remove(scope); } } } - /// Bulk-release every withheld event for `channel_id` back to the queue + /// Bulk-release every withheld event for `scope` back to the queue /// front, preserving relative FIFO order. /// - /// Called from the `in_flight_deadline` expiry blocks in - /// `flush_next` and `has_flushable_work` — if a steer ack never arrives - /// (read loop hung, watcher never posted), the withheld events would - /// otherwise be permanently orphaned. Recover, do not log-and-drop: the - /// events were never delivered to the agent, so normal dispatch must - /// have a chance to deliver them. + /// Called from the `in_flight_deadline` expiry path in + /// `expire_stuck_in_flight` — if a steer ack never arrives (read loop + /// hung, watcher never posted), the withheld events would otherwise be + /// permanently orphaned. Recover, do not log-and-drop: the events were + /// never delivered to the agent, so normal dispatch must have a chance + /// to deliver them. /// /// Iterates the stored entries in reverse so per-entry `push_front` /// composes to original-FIFO order at the queue front (same discipline - /// as `requeue_preserve_timestamps` at line 453). - fn recover_withheld_for_expired_channel(&mut self, channel_id: Uuid) { - let Some(entries) = self.withheld_native_steer.remove(&channel_id) else { + /// as `requeue_preserve_timestamps`). + fn recover_withheld_for_expired_scope(&mut self, scope: &ConversationSessionKey) { + let Some(entries) = self.withheld_native_steer.remove(scope) else { return; }; let n = entries.len(); - let queue = self.queues.entry(channel_id).or_default(); + let queue = self.queues.entry(scope.clone()).or_default(); for qe in entries.into_iter().rev() { queue.push_front(qe); } - while queue.len() > MAX_PENDING_PER_CHANNEL { + while queue.len() > MAX_PENDING_PER_SCOPE { queue.pop_back(); tracing::warn!( - channel_id = %channel_id, - limit = MAX_PENDING_PER_CHANNEL, + channel_id = %scope.channel_id, + limit = MAX_PENDING_PER_SCOPE, "withheld-steer recovery overflow — dropped newest event to enforce cap" ); } tracing::warn!( - channel_id = %channel_id, + channel_id = %scope.channel_id, + root = scope.root_event_id.as_deref().unwrap_or(""), recovered = n, "in-flight expiry recovered withheld steer event(s) — \ steer ack never arrived; normal dispatch will deliver" @@ -816,12 +868,12 @@ impl EventQueue { /// Compact expired metadata entries to prevent unbounded map growth. /// /// Removes `retry_after` entries whose deadline has already passed, and - /// cleans up orphaned `retry_counts` entries for channels that have no + /// cleans up orphaned `retry_counts` entries for scopes that have no /// queued events, no active throttle, and no in-flight prompt. Without - /// this, channels that completed their retry cycle but never received + /// this, scopes that completed their retry cycle but never received /// fresh traffic would leak a `u32` entry in `retry_counts` indefinitely. /// - /// The in-flight guard is critical: a channel whose throttle expired and + /// The in-flight guard is critical: a scope whose throttle expired and /// whose queue is empty because it was flushed may still have a retry /// attempt in flight. Removing its `retry_counts` would reset the /// backoff sequence if that attempt fails and requeues. @@ -832,13 +884,13 @@ impl EventQueue { pub fn compact_expired_state(&mut self) { let now = Instant::now(); self.retry_after.retain(|_, deadline| *deadline > now); - // Remove retry_counts for channels with no active throttle, no + // Remove retry_counts for scopes with no active throttle, no // queued events, AND no in-flight prompt — they completed their // retry cycle and are truly idle. - self.retry_counts.retain(|ch, _| { - self.retry_after.contains_key(ch) - || self.queues.get(ch).is_some_and(|q| !q.is_empty()) - || self.in_flight_channels.contains(ch) + self.retry_counts.retain(|key, _| { + self.retry_after.contains_key(key) + || self.queues.get(key).is_some_and(|q| !q.is_empty()) + || self.in_flight_scopes.contains(key) }); } } @@ -1727,8 +1779,13 @@ mod tests { q.queues.values().map(|q| q.len()).sum() } + /// Channel-scoped key shorthand for tests exercising legacy scoping. + fn scope(ch: Uuid) -> ConversationSessionKey { + ConversationSessionKey::channel(ch) + } + fn any_in_flight(q: &EventQueue) -> bool { - !q.in_flight_channels.is_empty() + !q.in_flight_scopes.is_empty() } #[test] @@ -1742,7 +1799,7 @@ mod tests { } #[test] - fn test_distinct_conversation_roots_are_flushed_in_separate_channel_batches() { + fn test_distinct_conversation_roots_flush_concurrently_in_separate_batches() { let mut q = EventQueue::new(DedupMode::Queue); let ch = Uuid::new_v4(); let mut first = make_queued(ch, "first root"); @@ -1757,41 +1814,142 @@ mod tests { assert_eq!(first_batch.events.len(), 1); assert_eq!(pending_count(&q), 1); - // Scheduling remains channel-serialized even though session identity is root-scoped. - assert!(q.flush_next().is_none()); - q.mark_complete(ch); + // Scheduling is scope-keyed: a different root in the same channel + // flushes while the first is still in flight. let second_batch = q .flush_next() - .expect("second root should flush after completion"); + .expect("second root should flush concurrently"); assert_eq!(second_batch.conversation_root.as_deref(), Some("root-b")); assert_eq!(second_batch.events.len(), 1); + assert!(q.is_scope_in_flight(&first_batch.scope_key())); + assert!(q.is_scope_in_flight(&second_batch.scope_key())); + + // Each root stays serialized within itself: nothing left to flush. + assert!(q.flush_next().is_none()); } #[test] - fn test_cancelled_root_is_redispatched_before_different_queued_root() { + fn test_cancelled_root_redispatches_alongside_other_queued_roots() { let mut q = EventQueue::new(DedupMode::Queue); let ch = Uuid::new_v4(); let mut root_a = make_queued(ch, "cancelled A"); root_a.conversation_root = Some("root-a".into()); q.push(root_a); let batch_a = q.flush_next().expect("root A should flush"); + let scope_a = batch_a.scope_key(); q.requeue_as_cancelled(batch_a, CancelReason::Steer); - q.mark_complete(ch); + q.mark_complete(&scope_a); let mut root_b = make_queued(ch, "queued B"); root_b.conversation_root = Some("root-b".into()); q.push(root_b); - let redispatched_a = q.flush_next().expect("cancelled A should redispatch first"); - assert_eq!(redispatched_a.conversation_root.as_deref(), Some("root-a")); - assert_eq!(redispatched_a.events[0].event.content, "cancelled A"); - assert!(redispatched_a.cancelled_events.is_empty()); - q.mark_complete(ch); - - let dispatched_b = q.flush_next().expect("root B should remain queued"); + // Fresh root B dispatches without waiting on A's cancelled work… + let dispatched_b = q.flush_next().expect("root B should dispatch"); assert_eq!(dispatched_b.conversation_root.as_deref(), Some("root-b")); assert_eq!(dispatched_b.events[0].event.content, "queued B"); assert!(dispatched_b.cancelled_events.is_empty()); + + // …and A's cancelled batch redispatches concurrently under its own + // scope, with B still in flight. + let redispatched_a = q.flush_next().expect("cancelled A should redispatch"); + assert_eq!(redispatched_a.conversation_root.as_deref(), Some("root-a")); + assert_eq!(redispatched_a.events[0].event.content, "cancelled A"); + assert!(redispatched_a.cancelled_events.is_empty()); + assert!(q.is_scope_in_flight(&redispatched_a.scope_key())); + assert!(q.is_scope_in_flight(&dispatched_b.scope_key())); + } + + /// Helper: a root-scoped queued event. + fn make_queued_rooted(ch: Uuid, root: &str, content: &str) -> QueuedEvent { + let mut qe = make_queued(ch, content); + qe.conversation_root = Some(root.into()); + qe + } + + #[test] + fn test_restore_unclaimed_fresh_batch_is_exact_flush_undo() { + let mut q = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + q.push(make_queued_rooted(ch, "root-a", "one")); + q.push(make_queued_rooted(ch, "root-a", "two")); + + let batch = q.flush_next().expect("flush"); + let scope_a = batch.scope_key(); + let original_received: Vec<_> = batch.events.iter().map(|e| e.received_at).collect(); + + q.restore_unclaimed(batch); + q.mark_complete(&scope_a); + + // Exact undo: same events, same order, same timestamps, no retry + // accounting, nothing left in the cancelled store. + let refetched = q.flush_next().expect("restored batch reflushes"); + assert_eq!(refetched.conversation_root.as_deref(), Some("root-a")); + assert_eq!(refetched.events.len(), 2); + assert_eq!(refetched.events[0].event.content, "one"); + assert_eq!(refetched.events[1].event.content, "two"); + let restored_received: Vec<_> = refetched.events.iter().map(|e| e.received_at).collect(); + assert_eq!(restored_received, original_received); + assert!(refetched.cancelled_events.is_empty()); + assert!(q.retry_counts.is_empty()); + assert!(q.cancelled_batches.is_empty()); + } + + #[test] + fn test_restore_unclaimed_merged_batch_preserves_cancel_framing() { + let mut q = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + q.push(make_queued_rooted(ch, "root-a", "old")); + let batch = q.flush_next().expect("flush old"); + let scope_a = batch.scope_key(); + q.push(make_queued_rooted(ch, "root-a", "new")); + q.requeue_as_cancelled(batch, CancelReason::Steer); + q.mark_complete(&scope_a); + + // Merged flush: events=[new], cancelled_events=[old]. + let merged = q.flush_next().expect("merged flush"); + assert_eq!(merged.events.len(), 1); + assert_eq!(merged.cancelled_events.len(), 1); + assert_eq!(merged.cancel_reason, Some(CancelReason::Steer)); + + // No agent available — restore. The cancelled portion must go back + // to the cancelled store (framing intact), the fresh event back to + // the queue. + q.restore_unclaimed(merged); + q.mark_complete(&scope_a); + + let remerged = q.flush_next().expect("re-merged flush"); + assert_eq!(remerged.events.len(), 1); + assert_eq!(remerged.events[0].event.content, "new"); + assert_eq!(remerged.cancelled_events.len(), 1); + assert_eq!(remerged.cancelled_events[0].event.content, "old"); + assert_eq!(remerged.cancel_reason, Some(CancelReason::Steer)); + } + + #[test] + fn test_restore_unclaimed_cancelled_redispatch_returns_to_cancelled_store() { + let mut q = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + q.push(make_queued_rooted(ch, "root-a", "cancelled work")); + let batch = q.flush_next().expect("flush"); + let scope_a = batch.scope_key(); + q.requeue_as_cancelled(batch, CancelReason::Interrupt); + q.mark_complete(&scope_a); + + // Cancelled-only redispatch (dispatch_cancelled path). + let redispatch = q.flush_next().expect("cancelled redispatch"); + assert_eq!(redispatch.cancel_reason, Some(CancelReason::Interrupt)); + assert!(redispatch.cancelled_events.is_empty()); + + // Restore: must return whole to the cancelled store under the + // ORIGINAL reason, not become a plain queued event. + q.restore_unclaimed(redispatch); + q.mark_complete(&scope_a); + assert!(q.cancelled_batches.contains_key(&scope_a)); + + let again = q.flush_next().expect("redispatch again"); + assert_eq!(again.cancel_reason, Some(CancelReason::Interrupt)); + assert_eq!(again.events[0].event.content, "cancelled work"); } #[test] @@ -1853,7 +2011,7 @@ mod tests { assert!(q.flush_next().is_none()); // Complete the in-flight prompt. - q.mark_complete(ch); + q.mark_complete(&scope(ch)); assert!(!any_in_flight(&q)); // Now flush should succeed. @@ -1949,7 +2107,7 @@ mod tests { assert_eq!(pending_count(&q), 1); assert_eq!(q.queues.len(), 1); - q.mark_complete(ch_a); + q.mark_complete(&scope(ch_a)); // Second flush picks B. let batch_b = q.flush_next().expect("second flush"); @@ -2102,7 +2260,7 @@ mod tests { // The mode gate fires Steer → cancel → requeue as cancelled, carrying // the steer reason (exactly the lib.rs requeue path). q.requeue_as_cancelled(batch, CancelReason::Steer); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // The re-prompt the agent actually receives. let merged = q.flush_next().unwrap(); @@ -2246,10 +2404,10 @@ mod tests { // Simulate failure — requeue the batch. queue.requeue(batch); - queue.mark_complete(ch); + queue.mark_complete(&scope(ch)); // retry_after is set, so manually clear it for this test. - queue.retry_after.remove(&ch); + queue.retry_after.remove(&scope(ch)); // Should be able to flush again and get the same events in order. let batch2 = queue.flush_next().unwrap(); @@ -2274,7 +2432,7 @@ mod tests { // Requeue ch_a (simulating failure) and complete. queue.requeue(batch_a); - queue.mark_complete(ch_a); + queue.mark_complete(&scope(ch_a)); // After requeue, ch_a has retry_after set (5s), so ch_b goes first. let next_batch = queue.flush_next().unwrap(); @@ -2635,7 +2793,7 @@ mod tests { q.push(make_queued(ch, "dropped")); assert_eq!(pending_count(&q), 0, "event should be dropped"); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // Nothing to flush. assert!(q.flush_next().is_none()); } @@ -2654,7 +2812,7 @@ mod tests { q.push(make_queued(ch_b, "B-event")); assert_eq!(pending_count(&q), 1); - q.mark_complete(ch_a); + q.mark_complete(&scope(ch_a)); let batch_b = q.flush_next().expect("flush B"); assert_eq!(batch_b.channel_id, ch_b); } @@ -2678,14 +2836,14 @@ mod tests { assert_eq!(batch_b.channel_id, ch_b); // Both in-flight. - assert_eq!(q.in_flight_channels.len(), 2); + assert_eq!(q.in_flight_scopes.len(), 2); // Complete A only. - q.mark_complete(ch_a); + q.mark_complete(&scope(ch_a)); assert!(any_in_flight(&q)); // B still in-flight. // Complete B. - q.mark_complete(ch_b); + q.mark_complete(&scope(ch_b)); assert!(!any_in_flight(&q)); } @@ -2729,8 +2887,8 @@ mod tests { q.push(make_queued(ch_b, "B-dropped")); assert_eq!(pending_count(&q), 0); - q.mark_complete(ch_a); - q.mark_complete(ch_b); + q.mark_complete(&scope(ch_a)); + q.mark_complete(&scope(ch_b)); } #[test] @@ -2760,9 +2918,9 @@ mod tests { // All in-flight. assert!(q.flush_next().is_none()); - q.mark_complete(ch_a); - q.mark_complete(ch_b); - q.mark_complete(ch_c); + q.mark_complete(&scope(ch_a)); + q.mark_complete(&scope(ch_b)); + q.mark_complete(&scope(ch_c)); } #[test] @@ -2777,18 +2935,18 @@ mod tests { let _batch_a = q.flush_next().expect("flush A"); let _batch_b = q.flush_next().expect("flush B"); - assert_eq!(q.in_flight_channels.len(), 2); + assert_eq!(q.in_flight_scopes.len(), 2); // Complete only A. - q.mark_complete(ch_a); - assert_eq!(q.in_flight_channels.len(), 1); - assert!(q.in_flight_channels.contains(&ch_b)); - assert!(!q.in_flight_channels.contains(&ch_a)); + q.mark_complete(&scope(ch_a)); + assert_eq!(q.in_flight_scopes.len(), 1); + assert!(q.in_flight_scopes.contains(&scope(ch_b))); + assert!(!q.in_flight_scopes.contains(&scope(ch_a))); // B still in-flight. assert!(any_in_flight(&q)); - q.mark_complete(ch_b); + q.mark_complete(&scope(ch_b)); assert!(!any_in_flight(&q)); } @@ -2811,7 +2969,7 @@ mod tests { // requeue_preserve_timestamps should keep the original timestamp. q.requeue_preserve_timestamps(batch); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // No retry_after set — should be immediately flushable. let batch2 = q.flush_next().expect("flush after requeue_preserve"); @@ -2827,10 +2985,10 @@ mod tests { let batch = q.flush_next().expect("flush"); q.requeue_preserve_timestamps(batch); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // No retry_after — channel should be immediately flushable. - assert!(!q.retry_after.contains_key(&ch)); + assert!(!q.retry_after.contains_key(&scope(ch))); assert!(q.flush_next().is_some()); } @@ -2839,34 +2997,34 @@ mod tests { let mut q = EventQueue::new(DedupMode::Queue); let ch = Uuid::new_v4(); - // Fill the channel to MAX_PENDING_PER_CHANNEL. - for i in 0..MAX_PENDING_PER_CHANNEL { + // Fill the channel to MAX_PENDING_PER_SCOPE. + for i in 0..MAX_PENDING_PER_SCOPE { q.push(make_queued(ch, &format!("fill-{i}"))); } - assert_eq!(pending_count(&q), MAX_PENDING_PER_CHANNEL); + assert_eq!(pending_count(&q), MAX_PENDING_PER_SCOPE); // Flush a batch (removes some events from the queue). let batch = q.flush_next().expect("should flush"); let batch_size = batch.events.len(); - let remaining = MAX_PENDING_PER_CHANNEL - batch_size; + let remaining = MAX_PENDING_PER_SCOPE - batch_size; assert_eq!(pending_count(&q), remaining); // Push more events while the batch is "in-flight" — fill back to cap. for i in 0..batch_size { q.push(make_queued(ch, &format!("new-{i}"))); } - assert_eq!(pending_count(&q), MAX_PENDING_PER_CHANNEL); + assert_eq!(pending_count(&q), MAX_PENDING_PER_SCOPE); // Requeue the original batch — without cap enforcement this would - // push the queue to MAX_PENDING_PER_CHANNEL + batch_size. + // push the queue to MAX_PENDING_PER_SCOPE + batch_size. q.requeue_preserve_timestamps(batch); - // Cap must be enforced: queue should not exceed MAX_PENDING_PER_CHANNEL. + // Cap must be enforced: queue should not exceed MAX_PENDING_PER_SCOPE. assert!( - pending_count(&q) <= MAX_PENDING_PER_CHANNEL, + pending_count(&q) <= MAX_PENDING_PER_SCOPE, "queue exceeded cap: {} > {}", pending_count(&q), - MAX_PENDING_PER_CHANNEL, + MAX_PENDING_PER_SCOPE, ); } @@ -2875,8 +3033,8 @@ mod tests { let mut q = EventQueue::new(DedupMode::Queue); let ch = Uuid::new_v4(); - // Push exactly MAX_PENDING_PER_CHANNEL events with identifiable content. - for i in 0..MAX_PENDING_PER_CHANNEL { + // Push exactly MAX_PENDING_PER_SCOPE events with identifiable content. + for i in 0..MAX_PENDING_PER_SCOPE { q.push(make_queued(ch, &format!("original-{i}"))); } @@ -2894,7 +3052,7 @@ mod tests { // Requeue — older events go to front, overflow trims from back (newest). q.requeue_preserve_timestamps(batch); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // The requeued events should be at the front of the queue. let batch2 = q.flush_next().expect("should flush after requeue"); @@ -2921,14 +3079,14 @@ mod tests { assert!(!q.has_flushable_work()); // Complete — no pending events, no flushable work. - q.mark_complete(ch); + q.mark_complete(&scope(ch)); assert!(!q.has_flushable_work()); // Requeue with retry_after — throttled, no flushable work. q.push(make_queued(ch, "msg2")); let batch2 = q.flush_next().expect("flush2"); q.requeue(batch2); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); assert!( !q.has_flushable_work(), "throttled channel should not be flushable" @@ -2936,7 +3094,7 @@ mod tests { // Manually expire the retry_after to simulate time passing. q.retry_after - .insert(ch, Instant::now() - Duration::from_secs(1)); + .insert(scope(ch), Instant::now() - Duration::from_secs(1)); assert!( q.has_flushable_work(), "expired throttle should be flushable" @@ -2951,26 +3109,26 @@ mod tests { q.push(make_queued(ch, "poison")); for attempt in 1..=MAX_RETRIES { q.retry_after - .insert(ch, Instant::now() - Duration::from_secs(1)); + .insert(scope(ch), Instant::now() - Duration::from_secs(1)); let batch = q.flush_next().expect("flush"); assert!( q.requeue(batch).is_none(), "attempt {attempt} should requeue, not dead-letter" ); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); } // The MAX_RETRIES+1'th failure dead-letters: batch is returned. q.retry_after - .insert(ch, Instant::now() - Duration::from_secs(1)); + .insert(scope(ch), Instant::now() - Duration::from_secs(1)); let batch = q.flush_next().expect("flush"); let dead = q.requeue(batch).expect("should dead-letter"); assert_eq!(dead.channel_id, ch); assert_eq!(dead.events.len(), 1); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // Retry state is cleared so fresh traffic isn't throttled. - assert!(!q.retry_counts.contains_key(&ch)); - assert!(!q.retry_after.contains_key(&ch)); + assert!(!q.retry_counts.contains_key(&scope(ch))); + assert!(!q.retry_after.contains_key(&scope(ch))); } #[test] @@ -2984,7 +3142,7 @@ mod tests { // Requeue sets retry_after. q.requeue(batch); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // Channel is throttled — flush_next should return None (no other channels). assert!(q.flush_next().is_none()); @@ -2996,8 +3154,8 @@ mod tests { // After retry_after expires, ch should be flushable again. q.retry_after - .insert(ch, Instant::now() - Duration::from_secs(1)); - q.mark_complete(ch2); + .insert(scope(ch), Instant::now() - Duration::from_secs(1)); + q.mark_complete(&scope(ch2)); let batch3 = q .flush_next() .expect("ch should be flushable after throttle expires"); @@ -3697,7 +3855,7 @@ mod tests { q.push(make_queued(ch, "msg")); let batch = q.flush_next().unwrap(); q.requeue(batch); // sets retry_after - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // Channel is throttled — verify drain clears it. assert!(!q.has_flushable_work()); @@ -3741,26 +3899,26 @@ mod tests { q.push(make_queued(ch, "msg1")); let batch = q.flush_next().unwrap(); q.requeue(batch); - q.mark_complete(ch); - assert!(q.retry_after.contains_key(&ch)); - assert!(q.retry_counts.contains_key(&ch)); + q.mark_complete(&scope(ch)); + assert!(q.retry_after.contains_key(&scope(ch))); + assert!(q.retry_counts.contains_key(&scope(ch))); // The requeued event is back in the queue. Flush it again so the // queue is empty (simulating a successful retry dispatch). // We need to wait for retry_after to expire first. q.retry_after - .insert(ch, Instant::now() - Duration::from_secs(1)); + .insert(scope(ch), Instant::now() - Duration::from_secs(1)); let _batch2 = q.flush_next().unwrap(); // Now mark_complete with no active throttle — clears retry_counts. - q.mark_complete(ch); - assert!(!q.retry_counts.contains_key(&ch)); + q.mark_complete(&scope(ch)); + assert!(!q.retry_counts.contains_key(&scope(ch))); // Re-create the orphan scenario: manually insert stale retry_counts // with no queue, no throttle, and no in-flight. - q.retry_counts.insert(ch, 3); + q.retry_counts.insert(scope(ch), 3); q.compact_expired_state(); assert!( - !q.retry_counts.contains_key(&ch), + !q.retry_counts.contains_key(&scope(ch)), "orphaned retry_counts should be removed" ); } @@ -3774,21 +3932,21 @@ mod tests { q.push(make_queued(ch, "msg1")); let batch = q.flush_next().unwrap(); q.requeue(batch); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // Expire the throttle so the requeued event can be flushed. q.retry_after - .insert(ch, Instant::now() - Duration::from_secs(1)); + .insert(scope(ch), Instant::now() - Duration::from_secs(1)); let _batch2 = q.flush_next().unwrap(); // Channel is now in-flight with empty queue and expired throttle. - assert!(q.in_flight_channels.contains(&ch)); - assert!(q.queues.get(&ch).is_none_or(|q| q.is_empty())); + assert!(q.in_flight_scopes.contains(&scope(ch))); + assert!(q.queues.get(&scope(ch)).is_none_or(|q| q.is_empty())); // compact must NOT remove retry_counts — the in-flight attempt // may fail and requeue, which needs the existing count. q.compact_expired_state(); assert!( - q.retry_counts.contains_key(&ch), + q.retry_counts.contains_key(&scope(ch)), "retry_counts must survive while channel is in-flight" ); } @@ -3800,11 +3958,11 @@ mod tests { // Manually set up: retry_counts exists, queue is non-empty, no throttle. q.push(make_queued(ch, "msg1")); - q.retry_counts.insert(ch, 2); + q.retry_counts.insert(scope(ch), 2); q.compact_expired_state(); assert!( - q.retry_counts.contains_key(&ch), + q.retry_counts.contains_key(&scope(ch)), "retry_counts should survive when queue is non-empty" ); } @@ -3825,7 +3983,7 @@ mod tests { // Cancel the original batch and release the channel. q.requeue_as_cancelled(batch, CancelReason::Interrupt); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // flush_next should merge: events=[new-1], cancelled_events=[old-1, old-2]. let next = q.flush_next().unwrap(); @@ -3847,27 +4005,27 @@ mod tests { let batch = q.flush_next().unwrap(); q.push(make_queued(ch, "new")); q.requeue_as_cancelled(batch, CancelReason::Steer); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); let merged = q.flush_next().unwrap(); assert_eq!( merged.cancel_reason, Some(CancelReason::Steer), "steer reason should reach the merged batch" ); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // Fallback path (no new event): reason still rides through. q.push(make_queued(ch, "only")); let batch = q.flush_next().unwrap(); q.requeue_as_cancelled(batch, CancelReason::Interrupt); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); let fallback = q.flush_next().unwrap(); assert_eq!( fallback.cancel_reason, Some(CancelReason::Interrupt), "interrupt reason should reach the re-dispatched batch" ); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // A normal (non-cancel) flush carries no reason. q.push(make_queued(ch, "plain")); @@ -3883,12 +4041,12 @@ mod tests { let batch1 = q.flush_next().unwrap(); q.push(make_queued(ch, "new-1")); q.requeue_as_cancelled(batch1, CancelReason::Interrupt); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); let batch2 = q.flush_next().unwrap(); // Second cancel with a different reason — the latest reason wins. q.requeue_as_cancelled(batch2, CancelReason::Steer); q.push(make_queued(ch, "new-2")); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); let batch3 = q.flush_next().unwrap(); assert_eq!(batch3.cancel_reason, Some(CancelReason::Steer)); } @@ -3906,7 +4064,7 @@ mod tests { // Cancel the batch (no new events pushed) and release the channel. q.requeue_as_cancelled(batch, CancelReason::Interrupt); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // Fallback path: cancelled events become regular events, cancelled_events is empty. let next = q.flush_next().unwrap(); @@ -3930,7 +4088,7 @@ mod tests { q.push(make_queued(ch, "msg")); let batch = q.flush_next().unwrap(); q.requeue_as_cancelled(batch, CancelReason::Interrupt); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // Channel has only cancelled events — should still be considered flushable. assert!( @@ -3948,7 +4106,7 @@ mod tests { q.push(make_queued(ch, "msg")); let batch = q.flush_next().unwrap(); q.requeue_as_cancelled(batch, CancelReason::Interrupt); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // drain_channel should clear cancelled_batches for the channel. q.drain_channel(ch); @@ -3976,7 +4134,7 @@ mod tests { // First cancel: store 2 cancelled events. q.requeue_as_cancelled(batch1, CancelReason::Interrupt); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // Second flush: events=[new-1], cancelled_events=[orig-1, orig-2]. let batch2 = q.flush_next().unwrap(); @@ -3989,7 +4147,7 @@ mod tests { // Push 1 more new event and release channel. q.push(make_queued(ch, "new-2")); - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // Third flush: events=[new-2], cancelled_events=[orig-1, orig-2, new-1]. let batch3 = q.flush_next().unwrap(); @@ -4431,7 +4589,7 @@ mod tests { let event_id = qe.event.id.to_hex(); q.push(qe); - assert!(q.mark_native_steer_pending(ch, &event_id)); + assert!(q.mark_native_steer_pending(&scope(ch), &event_id)); assert!( q.flush_next().is_none(), @@ -4442,7 +4600,10 @@ mod tests { "withheld-only channel must not register as flushable work" ); assert_eq!(pending_count(&q), 0); - assert_eq!(q.withheld_native_steer.get(&ch).map(|v| v.len()), Some(1)); + assert_eq!( + q.withheld_native_steer.get(&scope(ch)).map(|v| v.len()), + Some(1) + ); } /// Earlier events on the same channel must flush normally during the @@ -4468,7 +4629,7 @@ mod tests { q.push(e3); // Steer in flight for e3 — withhold it from normal dispatch. - assert!(q.mark_native_steer_pending(ch, &e3_id)); + assert!(q.mark_native_steer_pending(&scope(ch), &e3_id)); // Earlier events flush as a normal batch; e3 is invisible. let batch = q @@ -4480,10 +4641,10 @@ mod tests { assert_eq!(batch.events[1].event.id.to_hex(), e2_id); // Earlier batch completes; channel is no longer in flight. - q.mark_complete(ch); + q.mark_complete(&scope(ch)); // Ack arrives as Err or PromptCompletedNeutral → release e3. - q.release_native_steer(ch, &e3_id); + q.release_native_steer(&scope(ch), &e3_id); let next = q.flush_next().expect("released e3 should now flush"); assert_eq!(next.channel_id, ch); @@ -4510,17 +4671,17 @@ mod tests { // Simulate a prompt in flight for `ch`, then withhold the queued // event for an in-flight goose-native steer. - q.in_flight_channels.insert(ch); - q.in_flight_deadlines.insert(ch, Instant::now()); - q.in_flight_batch_sizes.insert(ch, 1); - assert!(q.mark_native_steer_pending(ch, &event_id)); + q.in_flight_scopes.insert(scope(ch)); + q.in_flight_deadlines.insert(scope(ch), Instant::now()); + q.in_flight_batch_sizes.insert(scope(ch), 1); + assert!(q.mark_native_steer_pending(&scope(ch), &event_id)); // Force the in-flight deadline to be in the past, simulating the // steer ack never arriving and the read loop hanging long enough // for `in_flight_deadline` to elapse. Same expiry-simulation // trick used by `test_retry_throttle_blocks_requeue_channel`. q.in_flight_deadlines - .insert(ch, Instant::now() - Duration::from_secs(1)); + .insert(scope(ch), Instant::now() - Duration::from_secs(1)); // `has_flushable_work` runs the expiry block first; it must recover // the withheld event so the channel registers as flushable. @@ -4568,24 +4729,27 @@ mod tests { // tests. What matters here is that the bulk-recovery path // (reverse iter + push_front) composes to original FIFO at the // queue front. - assert!(q.mark_native_steer_pending(ch, &e1_id)); - assert!(q.mark_native_steer_pending(ch, &e2_id)); - assert!(q.mark_native_steer_pending(ch, &e3_id)); + assert!(q.mark_native_steer_pending(&scope(ch), &e1_id)); + assert!(q.mark_native_steer_pending(&scope(ch), &e2_id)); + assert!(q.mark_native_steer_pending(&scope(ch), &e3_id)); assert_eq!(pending_count(&q), 0); - assert_eq!(q.withheld_native_steer.get(&ch).map(|v| v.len()), Some(3)); + assert_eq!( + q.withheld_native_steer.get(&scope(ch)).map(|v| v.len()), + Some(3) + ); // Trigger expiry → bulk-release path. - q.in_flight_channels.insert(ch); + q.in_flight_scopes.insert(scope(ch)); q.in_flight_deadlines - .insert(ch, Instant::now() - Duration::from_secs(1)); - q.in_flight_batch_sizes.insert(ch, 3); + .insert(scope(ch), Instant::now() - Duration::from_secs(1)); + q.in_flight_batch_sizes.insert(scope(ch), 3); assert!(q.has_flushable_work()); // After recovery, the queue front-to-back order must match the // original FIFO: e1, e2, e3. let recovered: Vec = q .queues - .get(&ch) + .get(&scope(ch)) .expect("queue restored") .iter() .map(|qe| qe.event.id.to_hex())