From 463d6d26d7a7d6a28a7f043f7dfc2f084f005bd0 Mon Sep 17 00:00:00 2001 From: tlongwell-block <109685178+tlongwell-block@users.noreply.github.com> Date: Sat, 4 Apr 2026 15:34:32 -0400 Subject: [PATCH] feat(acp): add --multiple-event-handling flag with cancel+re-prompt (#216) --- crates/sprout-acp/src/acp.rs | 5 + crates/sprout-acp/src/config.rs | 48 +++++- crates/sprout-acp/src/main.rs | 120 ++++++++++++- crates/sprout-acp/src/pool.rs | 123 +++++++++++++- crates/sprout-acp/src/queue.rs | 287 ++++++++++++++++++++++++++++++-- 5 files changed, 559 insertions(+), 24 deletions(-) diff --git a/crates/sprout-acp/src/acp.rs b/crates/sprout-acp/src/acp.rs index 6d70d43141..198228c743 100644 --- a/crates/sprout-acp/src/acp.rs +++ b/crates/sprout-acp/src/acp.rs @@ -376,6 +376,11 @@ impl AcpClient { self.send_notification("session/cancel", params).await } + /// Returns `true` if a `session/prompt` request is currently in flight. + pub fn has_in_flight_prompt(&self) -> bool { + self.last_prompt_id.is_some() + } + /// Cancel a turn cleanly, handling any pending permission request first. /// /// Steps: diff --git a/crates/sprout-acp/src/config.rs b/crates/sprout-acp/src/config.rs index 17d28e3c3a..613edca9fb 100644 --- a/crates/sprout-acp/src/config.rs +++ b/crates/sprout-acp/src/config.rs @@ -42,6 +42,23 @@ pub enum DedupMode { Queue, } +/// How to handle new @mentions while a turn is already in-flight for that channel. +#[derive(Debug, Clone, Copy, PartialEq, clap::ValueEnum)] +pub enum MultipleEventHandling { + /// Queue new events while a turn is in-flight. Deliver after current turn + /// completes. Existing behavior — zero code change in this path. + Queue, + /// Cancel the in-flight turn and re-dispatch a merged prompt combining + /// the original events with the new ones, for ANY new @mention. + /// Requires DedupMode::Queue. + Interrupt, + /// Cancel the in-flight turn only when the new @mention is from the agent + /// owner (resolved via owner_cache). All other authors queue normally. + /// Requires DedupMode::Queue. + #[value(name = "owner-interrupt")] + OwnerInterrupt, +} + /// Inbound author gate: which authors' events the harness forwards to the agent. /// /// - `owner-only` — only the agent's registered owner (default). @@ -265,6 +282,17 @@ pub struct CliArgs { #[arg(long, env = "SPROUT_ACP_DEDUP", default_value = "queue", value_enum)] pub dedup: DedupMode, + /// How to handle new @mentions while a turn is already in-flight. + /// queue: events wait (default). interrupt: cancel+re-prompt on any mention. + /// owner-interrupt: cancel only for agent owner's mentions. + #[arg( + long, + env = "SPROUT_ACP_MULTIPLE_EVENT_HANDLING", + default_value = "queue", + value_enum + )] + pub multiple_event_handling: MultipleEventHandling, + #[arg(long, env = "SPROUT_ACP_NO_IGNORE_SELF")] pub no_ignore_self: bool, @@ -353,6 +381,7 @@ pub struct Config { pub initial_message: Option, pub subscribe_mode: SubscribeMode, pub dedup_mode: DedupMode, + pub multiple_event_handling: MultipleEventHandling, pub ignore_self: bool, pub kinds_override: Option>, pub channels_override: Option>, @@ -613,6 +642,20 @@ impl Config { HashSet::new() }; + // ── Multiple-event-handling validation ────────────────────────────── + if matches!( + args.multiple_event_handling, + MultipleEventHandling::Interrupt | MultipleEventHandling::OwnerInterrupt + ) && matches!(args.dedup, DedupMode::Drop) + { + return Err(ConfigError::ConfigFile( + "--multiple-event-handling=interrupt (or owner-interrupt) requires --dedup=queue. \ + DedupMode::Drop discards events during the cancel drain window, \ + producing incomplete merged prompts." + .into(), + )); + } + let config = Config { keys, api_token: args.api_token, @@ -629,6 +672,7 @@ impl Config { initial_message: args.initial_message, subscribe_mode: args.subscribe, dedup_mode: args.dedup, + multiple_event_handling: args.multiple_event_handling, ignore_self: !args.no_ignore_self, kinds_override: args.kinds, channels_override: args.channels, @@ -656,7 +700,7 @@ impl Config { other => format!("respond_to={other}"), }; format!( - "relay={} pubkey={} agent_cmd={} {} mcp_cmd={} idle_timeout={}s max_turn={}s agents={} heartbeat={}s subscribe={:?} dedup={:?} ignore_self={} context_limit={} max_turns_per_session={} presence={} typing={} model={} permission_mode={} {}", + "relay={} pubkey={} agent_cmd={} {} mcp_cmd={} idle_timeout={}s max_turn={}s agents={} heartbeat={}s subscribe={:?} dedup={:?} meh={:?} ignore_self={} context_limit={} max_turns_per_session={} presence={} typing={} model={} permission_mode={} {}", self.relay_url, self.keys.public_key().to_hex(), self.agent_command, @@ -668,6 +712,7 @@ impl Config { self.heartbeat_interval_secs, self.subscribe_mode, self.dedup_mode, + self.multiple_event_handling, self.ignore_self, self.context_message_limit, self.max_turns_per_session, @@ -986,6 +1031,7 @@ mod tests { initial_message: None, subscribe_mode: mode, dedup_mode: DedupMode::Queue, + multiple_event_handling: MultipleEventHandling::Queue, ignore_self: true, kinds_override: None, channels_override: None, diff --git a/crates/sprout-acp/src/main.rs b/crates/sprout-acp/src/main.rs index 636bd1b4d8..d7811719e4 100644 --- a/crates/sprout-acp/src/main.rs +++ b/crates/sprout-acp/src/main.rs @@ -14,7 +14,7 @@ use std::time::Duration; use acp::{AcpClient, EnvVar, McpServer}; use anyhow::Result; use clap::Parser; -use config::{Config, DedupMode, ModelsArgs, RespondTo, SubscribeMode}; +use config::{Config, DedupMode, ModelsArgs, MultipleEventHandling, RespondTo, SubscribeMode}; use filter::SubscriptionRule; use futures_util::FutureExt; use nostr::ToBech32; @@ -1132,6 +1132,40 @@ async fn tokio_main() -> Result<()> { } // ── End shutdown command handling ────────────────── + // ── Cancel command handling ────────────────────── + // Mirrors !shutdown: kind:9, content "!cancel", from + // owner, mentions THIS agent. Must be BEFORE + // queue.push() — the event content is moved by push. + // + // Mode-independent: !cancel fires regardless of + // --multiple-event-handling. It is explicit user + // intent, not an automatic policy decision. + let is_cancel = kind_u32 == KIND_STREAM_MESSAGE + && sprout_event.event.content.trim() == "!cancel" + && sprout_event.event.tags.iter().any(|t| { + t.as_slice().first().map(|s| s.as_str()) == Some("p") + && t.as_slice().get(1).map(|s| s.as_str()) == Some(pubkey_hex.as_str()) + }); + if is_cancel { + let owner = owner_cache + .get_or_resolve(&rest_client_for_presence, &pubkey_hex) + .await; + if let Some(owner) = owner { + if sprout_event.event.pubkey.to_hex() == *owner { + let fired = cancel_in_flight_task(&mut pool, sprout_event.channel_id); + if !fired { + tracing::warn!( + channel_id = %sprout_event.channel_id, + "!cancel received but no in-flight task — no-op" + ); + } + continue; // consume event — do NOT push to queue + } + } + // Not from owner — fall through to normal prompt handling. + } + // ── End cancel command handling ─────────────────── + // ── Inbound author gate ────────────────────────── // Coarse security policy: drop events from disallowed // authors before they reach subscription rules or the @@ -1189,6 +1223,9 @@ async fn tokio_main() -> Result<()> { continue; } }; + // Capture author pubkey before queue.push() moves + // sprout_event.event (needed for mode gate below). + let author_hex = sprout_event.event.pubkey.to_hex(); let event_id_hex = sprout_event.event.id.to_hex(); let accepted = queue.push(QueuedEvent { channel_id: sprout_event.channel_id, @@ -1207,6 +1244,28 @@ async fn tokio_main() -> Result<()> { pool::reaction_add(&rc, &event_id_hex, "👀").await; }); } + // ── Multiple-event-handling mode gate ───────────── + // Event is already queued. If mode requires it AND + // the channel has an in-flight task, fire cancel. + if accepted && queue.is_channel_in_flight(sprout_event.channel_id) { + let should_cancel = match config.multiple_event_handling { + MultipleEventHandling::Queue => false, + MultipleEventHandling::Interrupt => true, + MultipleEventHandling::OwnerInterrupt => { + let owner = owner_cache + .get_or_resolve(&rest_client_for_presence, &pubkey_hex) + .await; + match owner { + Some(o) => author_hex == *o, + None => false, + } + } + }; + if should_cancel { + cancel_in_flight_task(&mut pool, sprout_event.channel_id); + } + } + // ── End mode gate ──────────────────────────────── typing_channels.extend(dispatch_pending(&mut pool, &mut queue, &ctx)); } None => { @@ -1448,6 +1507,26 @@ enum LoopAction { Exit, } +// ── cancel_in_flight_task ───────────────────────────────────────────────────── + +/// Send a cancel signal to the in-flight task for `channel_id`. +/// Returns `true` if a signal was sent, `false` if no in-flight task was found. +fn cancel_in_flight_task(pool: &mut AgentPool, channel_id: uuid::Uuid) -> bool { + let entry = pool + .task_map_mut() + .values_mut() + .find(|m| m.channel_id == Some(channel_id)); + + if let Some(meta) = entry { + if let Some(tx) = meta.cancel_tx.take() { + let _ = tx.send(()); + tracing::info!(channel = %channel_id, "cancel signal sent to in-flight task"); + return true; + } + } + false +} + // ── dispatch_pending ────────────────────────────────────────────────────────── /// Flush queued work to available agents. @@ -1487,8 +1566,18 @@ fn dispatch_pending( // Prompt text is now built inside run_prompt_task (needs async for // context fetching). Pass None for prompt_text; batch carries the data. + let (cancel_tx, cancel_rx) = tokio::sync::oneshot::channel::<()>(); + let abort_handle = pool.join_set.spawn(async move { - pool::run_prompt_task(agent, Some(batch), None, ctx_clone, result_tx).await; + pool::run_prompt_task( + agent, + Some(batch), + None, + ctx_clone, + result_tx, + Some(cancel_rx), + ) + .await; }); pool.task_map_mut().insert( @@ -1497,6 +1586,7 @@ fn dispatch_pending( agent_index, channel_id: Some(channel_id), recoverable_batch, + cancel_tx: Some(cancel_tx), }, ); dispatched_channels.push(channel_id); @@ -1538,7 +1628,14 @@ fn handle_prompt_result( // Don't requeue batches for channels the agent was removed from — // those events are stale and should be silently dropped. if !removed_channels.contains(&batch.channel_id) { - queue.requeue(batch); + if matches!(result.outcome, PromptOutcome::Cancelled) { + // Cancel re-prompt: store as cancelled events so flush_next() + // merges them into the next FlushBatch.cancelled_events, + // enabling the annotated merged-prompt format. + queue.requeue_as_cancelled(batch); + } else { + queue.requeue(batch); + } } else { tracing::debug!( channel_id = %batch.channel_id, @@ -1565,6 +1662,7 @@ fn handle_prompt_result( PromptOutcome::Error(_) => "error", PromptOutcome::Timeout => "timeout", PromptOutcome::AgentExited => "exited", + PromptOutcome::Cancelled => "cancelled", }; let agent_index = result.agent.index; @@ -1611,6 +1709,19 @@ fn handle_prompt_result( // 2. Application-class (IdleTimeout, HardTimeout, Json): the pipe is // intact but the prompt failed. Return the agent to the pool so it // can be reused for the next event. + + // Intentional cancel — agent is healthy, return it to the pool. + // No respawn, no retry penalty. The cancelled batch was already stored + // via requeue_as_cancelled() above and will be merged into the next + // FlushBatch by flush_next(). + PromptOutcome::Cancelled => { + tracing::debug!( + agent = agent_index, + outcome = outcome_label, + "agent_returned (cancelled)" + ); + pool.return_agent(result.agent); + } PromptOutcome::Error(ref e) => { let is_transport_error = matches!( e, @@ -1797,7 +1908,7 @@ fn dispatch_heartbeat( let agent_index = agent.index; let abort_handle = pool.join_set.spawn(async move { - pool::run_prompt_task(agent, None, Some(prompt_text), ctx_clone, result_tx).await; + pool::run_prompt_task(agent, None, Some(prompt_text), ctx_clone, result_tx, None).await; }); pool.task_map_mut().insert( @@ -1806,6 +1917,7 @@ fn dispatch_heartbeat( agent_index, channel_id: None, recoverable_batch: None, + cancel_tx: None, }, ); *heartbeat_in_flight = true; diff --git a/crates/sprout-acp/src/pool.rs b/crates/sprout-acp/src/pool.rs index be45dc4499..da8e251114 100644 --- a/crates/sprout-acp/src/pool.rs +++ b/crates/sprout-acp/src/pool.rs @@ -51,6 +51,9 @@ pub struct TaskMeta { pub channel_id: Option, /// Clone of batch for Queue mode panic recovery. pub recoverable_batch: Option, + /// Cancel signal for the in-flight prompt task. + /// `None` for heartbeat tasks (not cancellable) and after signal is consumed. + pub cancel_tx: Option>, } /// Agent-level model capabilities. Populated on first session creation. @@ -157,6 +160,9 @@ pub enum PromptOutcome { Error(AcpError), AgentExited, Timeout, + /// Intentional cancel via `!cancel` command or interrupt mode. + /// Agent is healthy — no respawn, no retry penalty. + Cancelled, } /// Immutable config subset shared (via `Arc`) by all spawned prompt tasks. @@ -590,6 +596,7 @@ pub async fn run_prompt_task( prompt_text: Option, ctx: Arc, result_tx: mpsc::UnboundedSender, + cancel_rx: Option>, ) { // ── Determine source and resolve/create session ─────────────────────── @@ -851,15 +858,112 @@ pub async fn run_prompt_task( // ── Send the actual prompt ──────────────────────────────────────────── - let prompt_result = agent - .acp - .session_prompt_with_idle_timeout( - &session_id, - &prompt_text, - ctx.idle_timeout, - ctx.max_turn_duration, - ) - .await; + // ── Cancel-aware prompt dispatch ────────────────────────────────────── + // When cancel_rx is Some (channel tasks), wrap the prompt in select! so + // the main loop can interrupt it. Heartbeats (cancel_rx=None) take the + // simple await path — they are not cancellable. + let prompt_result = match cancel_rx { + None => { + // Heartbeat / non-cancellable path. + agent + .acp + .session_prompt_with_idle_timeout( + &session_id, + &prompt_text, + ctx.idle_timeout, + ctx.max_turn_duration, + ) + .await + } + Some(rx) => { + tokio::select! { + biased; + result = agent.acp.session_prompt_with_idle_timeout( + &session_id, + &prompt_text, + ctx.idle_timeout, + ctx.max_turn_duration, + ) => result, + _ = rx => { + // Cancel signal received. Guard against Race 1: the turn may + // have completed naturally just as cancel fired. + if agent.acp.has_in_flight_prompt() { + // Prompt is genuinely in-flight — cancel it. + match agent.acp.cancel_with_cleanup(&session_id, ctx.idle_timeout).await { + Ok(stop_reason) => { + log_stop_reason(&source, &stop_reason); + agent.state.invalidate(&source); + let _ = result_tx.send(PromptResult { + agent, + source, + outcome: PromptOutcome::Cancelled, + batch: requeue_batch_if_queue(&ctx, batch), + }); + return; + } + Err(AcpError::AgentExited) => { + agent.state.invalidate_all(); + let _ = result_tx.send(PromptResult { + agent, + source, + outcome: PromptOutcome::AgentExited, + batch: requeue_batch_if_queue(&ctx, batch), + }); + return; + } + Err(AcpError::IdleTimeout(_) | AcpError::HardTimeout) => { + // Cancel drain timed out — agent state uncertain. + agent.state.invalidate(&source); + let _ = result_tx.send(PromptResult { + agent, + source, + outcome: PromptOutcome::Timeout, + batch: requeue_batch_if_queue(&ctx, batch), + }); + return; + } + Err(e) => { + agent.state.invalidate(&source); + let _ = result_tx.send(PromptResult { + agent, + source, + outcome: PromptOutcome::Error(e), + batch: requeue_batch_if_queue(&ctx, batch), + }); + return; + } + } + } else { + // Race 1 resolution: turn completed naturally before cancel + // could fire. last_prompt_id is None — cleared by + // session_prompt_with_idle_timeout() on success. The prompt + // future was dropped by select! — its Ok result is gone. + // + // Note: this `else` branch (last_prompt_id is None) cannot + // fire during the pre-prompt phase because `biased` select! + // polls the prompt arm first. That arm sets last_prompt_id + // synchronously before its first yield point, so by the time + // the cancel arm can win, last_prompt_id is already Some. + // This branch only fires when the turn genuinely completed + // and last_prompt_id was cleared by the success path. + // + // MUST send a PromptResult or the main loop deadlocks. + tracing::debug!( + target: "pool::prompt", + "cancel signal arrived but turn already completed — treating as success" + ); + let _ = result_tx.send(PromptResult { + agent, + source, + outcome: PromptOutcome::Ok(StopReason::EndTurn), + batch: None, // turn succeeded — batch was processed, no requeue + }); + return; + } + } + } + } + }; match prompt_result { Ok(stop_reason) => { @@ -1936,6 +2040,7 @@ mod tests { prompt_tag: "@mention".into(), received_at: std::time::Instant::now(), }], + cancelled_events: vec![], }; let context = ConversationContext::Thread { messages: vec![ContextMessage { diff --git a/crates/sprout-acp/src/queue.rs b/crates/sprout-acp/src/queue.rs index 9f07c2c40b..4abbaeeeac 100644 --- a/crates/sprout-acp/src/queue.rs +++ b/crates/sprout-acp/src/queue.rs @@ -65,6 +65,10 @@ pub struct BatchEvent { pub struct FlushBatch { pub channel_id: Uuid, pub events: Vec, + /// Events from a cancelled batch that triggered this re-prompt. + /// Empty for normal (non-cancel) batches. When non-empty, `format_prompt()` + /// produces a merged prompt with annotated sections. + pub cancelled_events: Vec, } // ── EventQueue ──────────────────────────────────────────────────────────────── @@ -125,6 +129,10 @@ pub struct EventQueue { /// Per-channel retry attempt counter for exponential backoff / dead-lettering. retry_counts: HashMap, dedup_mode: DedupMode, + /// Events from cancelled batches, keyed by channel. Merged into the next + /// `FlushBatch` for that channel as `cancelled_events` so `format_prompt()` + /// can produce annotated "[Previous request — interrupted]" sections. + cancelled_batches: HashMap>, } impl EventQueue { @@ -138,6 +146,7 @@ impl EventQueue { retry_after: HashMap::new(), retry_counts: HashMap::new(), dedup_mode, + cancelled_batches: HashMap::new(), } } @@ -211,7 +220,38 @@ impl EventQueue { && self.retry_after.get(id).is_none_or(|&t| t <= now) }) .min_by_key(|(_, q)| q.front().unwrap().received_at) - .map(|(id, _)| *id)?; + .map(|(id, _)| *id); + + // Fallback: if no queued events are ready but a channel 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, + None => { + let cancelled_id = self + .cancelled_batches + .keys() + .find(|id| !self.in_flight_channels.contains(id)) + .copied(); + match cancelled_id { + Some(id) => { + // Move cancelled events into the regular events slot. + // No new events to merge — re-dispatch the original batch. + let cancelled = self.cancelled_batches.remove(&id).unwrap_or_default(); + self.in_flight_channels.insert(id); + self.in_flight_deadlines + .insert(id, now + Duration::from_secs(IN_FLIGHT_DEADLINE_SECS)); + self.in_flight_batch_sizes.insert(id, cancelled.len()); + return Some(FlushBatch { + channel_id: id, + events: cancelled, + cancelled_events: vec![], + }); + } + None => return None, + } + } + }; // Drain up to MAX_BATCH_EVENTS; leave any remainder in the queue. let queue = self.queues.entry(channel_id).or_default(); @@ -237,7 +277,17 @@ impl EventQueue { ); self.in_flight_batch_sizes.insert(channel_id, events.len()); - Some(FlushBatch { channel_id, events }) + // Merge any cancelled events stored by requeue_as_cancelled(). + let cancelled_events = self + .cancelled_batches + .remove(&channel_id) + .unwrap_or_default(); + + Some(FlushBatch { + channel_id, + events, + cancelled_events, + }) } /// Mark the prompt for `channel_id` as complete. @@ -387,6 +437,20 @@ impl EventQueue { } } + /// Requeue a cancelled batch so its events appear as `cancelled_events` + /// in the next `FlushBatch` for this channel (enabling the annotated + /// merged-prompt format in `format_prompt()`). + /// + /// 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) { + let entry = self.cancelled_batches.entry(batch.channel_id).or_default(); + // Preserve any already-cancelled events from a prior cancel (double-cancel). + entry.extend(batch.cancelled_events); + entry.extend(batch.events); + } + /// Returns `true` if any channel has pending events that are not in-flight /// and not throttled by `retry_after`. /// @@ -420,7 +484,10 @@ impl EventQueue { !q.is_empty() && !self.in_flight_channels.contains(id) && self.retry_after.get(id).is_none_or(|&t| t <= now) - }) + }) || self + .cancelled_batches + .keys() + .any(|id| !self.in_flight_channels.contains(id)) } /// Whether any prompt is currently in flight. @@ -463,6 +530,7 @@ impl EventQueue { .unwrap_or_default(); self.retry_after.remove(&channel_id); self.retry_counts.remove(&channel_id); + self.cancelled_batches.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 @@ -905,16 +973,47 @@ pub fn format_prompt( sections.push(format_conversation_context(ctx, profile_lookup)); } - // 4. Event block(s). + // 4a. Cancelled events section (cancel + re-prompt). + if !batch.cancelled_events.is_empty() { + let mut s = "[Previous request — interrupted before completion]".to_string(); + for (i, be) in batch.cancelled_events.iter().enumerate() { + s.push_str(&format!( + "\n\n--- Event {} ({}) ---\n{}", + i + 1, + be.prompt_tag, + format_event_block(batch.channel_id, channel_info, be, profile_lookup) + )); + } + sections.push(s); + } + + // 4b. Event block(s). + let has_cancelled = !batch.cancelled_events.is_empty(); let event_section = if batch.events.len() == 1 { let be = &batch.events[0]; - format!( - "[Sprout event: {}]\n{}", - be.prompt_tag, - format_event_block(batch.channel_id, channel_info, be, profile_lookup) - ) + if has_cancelled { + format!( + "[New request — supersedes previous]\n\n--- Event 1 ({}) ---\n{}", + be.prompt_tag, + format_event_block(batch.channel_id, channel_info, be, profile_lookup) + ) + } else { + format!( + "[Sprout event: {}]\n{}", + be.prompt_tag, + format_event_block(batch.channel_id, channel_info, be, profile_lookup) + ) + } } else { - let mut s = format!("[Sprout events — {} events]", batch.events.len()); + let header = if has_cancelled { + format!( + "[New request — supersedes previous — {} events]", + batch.events.len() + ) + } else { + format!("[Sprout events — {} events]", batch.events.len()) + }; + let mut s = header; for (i, be) in batch.events.iter().enumerate() { s.push_str(&format!( "\n\n--- Event {} ({}) ---\n{}", @@ -927,6 +1026,16 @@ pub fn format_prompt( }; sections.push(event_section); + // Closing note for cancel + re-prompt. + if has_cancelled { + sections.push( + "Note: The previous request was interrupted. Please address the new request.\n\ + If the new request is unrelated to the previous one, you may briefly acknowledge\n\ + the interruption." + .to_string(), + ); + } + sections.join("\n\n") } @@ -1160,6 +1269,7 @@ mod tests { prompt_tag: "@mention".into(), received_at: Instant::now(), }], + cancelled_events: vec![], }; let prompt = format_prompt(&batch, None, None, None, None); @@ -1255,6 +1365,7 @@ mod tests { received_at: Instant::now(), }, ], + cancelled_events: vec![], }; let prompt = format_prompt(&batch, None, None, None, None); @@ -1283,6 +1394,7 @@ mod tests { prompt_tag: "test".into(), received_at: Instant::now(), }], + cancelled_events: vec![], }; let prompt = format_prompt(&batch, Some("You are a triage bot."), None, None, None); @@ -1758,6 +1870,7 @@ mod tests { prompt_tag: "test".into(), received_at: Instant::now(), }], + cancelled_events: vec![], }; let ci = PromptChannelInfo { name: "engineering".into(), @@ -1780,6 +1893,7 @@ mod tests { prompt_tag: "dm".into(), received_at: Instant::now(), }], + cancelled_events: vec![], }; let ci = PromptChannelInfo { name: "DM".into(), @@ -1809,6 +1923,7 @@ mod tests { prompt_tag: "@mention".into(), received_at: Instant::now(), }], + cancelled_events: vec![], }; let prompt = format_prompt(&batch, None, None, None, None); @@ -1835,6 +1950,7 @@ mod tests { prompt_tag: "@mention".into(), received_at: Instant::now(), }], + cancelled_events: vec![], }; let ctx = ConversationContext::Thread { messages: vec![ @@ -1870,6 +1986,7 @@ mod tests { prompt_tag: "dm".into(), received_at: Instant::now(), }], + cancelled_events: vec![], }; let ci = PromptChannelInfo { name: "DM".into(), @@ -1909,6 +2026,7 @@ mod tests { prompt_tag: "@mention".into(), received_at: Instant::now(), }], + cancelled_events: vec![], }; let ctx = ConversationContext::Thread { messages: vec![ContextMessage { @@ -2014,6 +2132,7 @@ mod tests { prompt_tag: "dm".into(), received_at: Instant::now(), }], + cancelled_events: vec![], }; let ci = PromptChannelInfo { name: "DM".into(), @@ -2061,6 +2180,7 @@ mod tests { prompt_tag: "dm".into(), received_at: Instant::now(), }], + cancelled_events: vec![], }; let ci = PromptChannelInfo { name: "DM".into(), @@ -2092,6 +2212,7 @@ mod tests { prompt_tag: "test".into(), received_at: Instant::now(), }], + cancelled_events: vec![], }; let prompt = format_prompt(&batch, None, None, None, None); @@ -2114,6 +2235,7 @@ mod tests { prompt_tag: "test".into(), received_at: Instant::now(), }], + cancelled_events: vec![], }; let prompt = format_prompt(&batch, None, None, None, None); @@ -2135,6 +2257,7 @@ mod tests { prompt_tag: "test".into(), received_at: Instant::now(), }], + cancelled_events: vec![], }; let prompt = format_prompt(&batch, None, None, None, None); @@ -2295,4 +2418,148 @@ mod tests { "retry_counts should survive when queue is non-empty" ); } + + // ── Test: requeue_as_cancelled merges into flush_next ──────────────────── + + #[test] + fn test_requeue_as_cancelled_merges_in_flush_next() { + let mut q = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + + // Push 2 events, flush into a batch. + q.push(make_queued(ch, "old-1")); + q.push(make_queued(ch, "old-2")); + let batch = q.flush_next().unwrap(); + assert_eq!(batch.events.len(), 2); + + // Push 1 new event while channel is in-flight. + q.push(make_queued(ch, "new-1")); + + // Cancel the original batch and release the channel. + q.requeue_as_cancelled(batch); + q.mark_complete(ch); + + // flush_next should merge: events=[new-1], cancelled_events=[old-1, old-2]. + let next = q.flush_next().unwrap(); + assert_eq!(next.events.len(), 1, "should have 1 new event"); + assert_eq!( + next.cancelled_events.len(), + 2, + "should have 2 cancelled events" + ); + } + + // ── Test: requeue_as_cancelled fallback (no new events) ────────────────── + + #[test] + fn test_requeue_as_cancelled_no_new_events_fallback() { + let mut q = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + + // Push 1 event, flush into a batch. + q.push(make_queued(ch, "only-event")); + let batch = q.flush_next().unwrap(); + + // Cancel the batch (no new events pushed) and release the channel. + q.requeue_as_cancelled(batch); + q.mark_complete(ch); + + // Fallback path: cancelled events become regular events, cancelled_events is empty. + let next = q.flush_next().unwrap(); + assert_eq!( + next.events.len(), + 1, + "cancelled event re-dispatched as regular event" + ); + assert!( + next.cancelled_events.is_empty(), + "no merge needed — cancelled_events should be empty" + ); + } + + // ── Test: has_flushable_work accounts for cancelled_batches ────────────── + + #[test] + fn test_has_flushable_work_with_cancelled_only() { + let mut q = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + + // Push, flush, cancel — no new events queued. + q.push(make_queued(ch, "msg")); + let batch = q.flush_next().unwrap(); + q.requeue_as_cancelled(batch); + q.mark_complete(ch); + + // Channel has only cancelled events — should still be considered flushable. + assert!( + q.has_flushable_work(), + "cancelled-only channel should be flushable" + ); + } + + // ── Test: drain_channel clears cancelled_batches ────────────────────────── + + #[test] + fn test_drain_channel_clears_cancelled_batches() { + let mut q = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + + // Push, flush, cancel. + q.push(make_queued(ch, "msg")); + let batch = q.flush_next().unwrap(); + q.requeue_as_cancelled(batch); + q.mark_complete(ch); + + // drain_channel should clear cancelled_batches for the channel. + q.drain_channel(ch); + + assert!(!q.has_flushable_work(), "nothing left after drain"); + assert!( + q.flush_next().is_none(), + "flush_next should return None after drain" + ); + } + + // ── Test: double-cancel accumulates all events ──────────────────────────── + + #[test] + fn test_double_cancel_preserves_all_events() { + let mut q = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + + // First flush: 2 events. + q.push(make_queued(ch, "orig-1")); + q.push(make_queued(ch, "orig-2")); + let batch1 = q.flush_next().unwrap(); + assert_eq!(batch1.events.len(), 2); + + // Push 1 new event while in-flight. + q.push(make_queued(ch, "new-1")); + + // First cancel: store 2 cancelled events. + q.requeue_as_cancelled(batch1); + q.mark_complete(ch); + + // Second flush: events=[new-1], cancelled_events=[orig-1, orig-2]. + let batch2 = q.flush_next().unwrap(); + assert_eq!(batch2.events.len(), 1); + assert_eq!(batch2.cancelled_events.len(), 2); + + // Second cancel: requeue_as_cancelled should accumulate all 3 events + // (2 from cancelled_events + 1 from events). + q.requeue_as_cancelled(batch2); + + // Push 1 more new event and release channel. + q.push(make_queued(ch, "new-2")); + q.mark_complete(ch); + + // Third flush: events=[new-2], cancelled_events=[orig-1, orig-2, new-1]. + let batch3 = q.flush_next().unwrap(); + assert_eq!(batch3.events.len(), 1, "should have 1 newest event"); + assert_eq!( + batch3.cancelled_events.len(), + 3, + "should accumulate all 3 cancelled events" + ); + } }