mirror of
https://github.com/block/buzz.git
synced 2026-10-27 12:40:30 +00:00
feat(acp): add --multiple-event-handling flag with cancel+re-prompt (#216)
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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<String>,
|
||||
pub subscribe_mode: SubscribeMode,
|
||||
pub dedup_mode: DedupMode,
|
||||
pub multiple_event_handling: MultipleEventHandling,
|
||||
pub ignore_self: bool,
|
||||
pub kinds_override: Option<Vec<u32>>,
|
||||
pub channels_override: Option<Vec<String>>,
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -51,6 +51,9 @@ pub struct TaskMeta {
|
||||
pub channel_id: Option<Uuid>,
|
||||
/// Clone of batch for Queue mode panic recovery.
|
||||
pub recoverable_batch: Option<FlushBatch>,
|
||||
/// Cancel signal for the in-flight prompt task.
|
||||
/// `None` for heartbeat tasks (not cancellable) and after signal is consumed.
|
||||
pub cancel_tx: Option<tokio::sync::oneshot::Sender<()>>,
|
||||
}
|
||||
|
||||
/// 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<String>,
|
||||
ctx: Arc<PromptContext>,
|
||||
result_tx: mpsc::UnboundedSender<PromptResult>,
|
||||
cancel_rx: Option<tokio::sync::oneshot::Receiver<()>>,
|
||||
) {
|
||||
// ── 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 {
|
||||
|
||||
@@ -65,6 +65,10 @@ pub struct BatchEvent {
|
||||
pub struct FlushBatch {
|
||||
pub channel_id: Uuid,
|
||||
pub events: Vec<BatchEvent>,
|
||||
/// 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<BatchEvent>,
|
||||
}
|
||||
|
||||
// ── EventQueue ────────────────────────────────────────────────────────────────
|
||||
@@ -125,6 +129,10 @@ pub struct EventQueue {
|
||||
/// Per-channel retry attempt counter for exponential backoff / dead-lettering.
|
||||
retry_counts: HashMap<Uuid, u32>,
|
||||
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<Uuid, Vec<BatchEvent>>,
|
||||
}
|
||||
|
||||
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"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user