diff --git a/Cargo.lock b/Cargo.lock index d7c2c59c87..e8681fa7d4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2034,6 +2034,18 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "43a88da9dd148bbcdce323dd6ac47d369b4769d4a3b78c6c52389b9269f77932" +[[package]] +name = "nix" +version = "0.31.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d6d0705320c1e6ba1d912b5e37cf18071b6c2e9b7fa8215a1e8a7651966f5d3" +dependencies = [ + "bitflags", + "cfg-if", + "cfg_aliases", + "libc", +] + [[package]] name = "nom" version = "7.1.3" @@ -3293,6 +3305,7 @@ dependencies = [ "clap", "evalexpr", "futures-util", + "nix", "nostr", "reqwest", "rustls", diff --git a/crates/sprout-acp/Cargo.toml b/crates/sprout-acp/Cargo.toml index 8c78dcd4cd..eee76a2bff 100644 --- a/crates/sprout-acp/Cargo.toml +++ b/crates/sprout-acp/Cargo.toml @@ -48,6 +48,9 @@ url = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } +# Process-group kill (safe wrapper around killpg) +nix = { version = "0.31", default-features = false, features = ["signal"] } + # Error handling thiserror = { workspace = true } anyhow = { workspace = true } diff --git a/crates/sprout-acp/src/acp.rs b/crates/sprout-acp/src/acp.rs index 8b6f1c25d4..6d70d43141 100644 --- a/crates/sprout-acp/src/acp.rs +++ b/crates/sprout-acp/src/acp.rs @@ -148,13 +148,19 @@ impl AcpClient { /// Call this when you need guaranteed cleanup — e.g., in `run_models` /// before process exit. pub async fn shutdown(&mut self) { - // Send SIGKILL to the direct child. On Unix, process_group(0) was set - // at spawn time, so the child runs in its own process group. Ideally - // we'd kill the entire group via kill(-pgid, SIGKILL), but that - // requires `unsafe` (libc::kill) which is denied by #![deny(unsafe_code)]. - // When the direct child dies, its subprocesses become orphans adopted - // by init — acceptable for a harness that respawns or exits. - let _ = self.child.start_kill(); + // Kill the entire process group when possible. The child was spawned + // with process_group(0), so its PID == its PGID. Killing the group + // ensures subprocesses (MCP servers, tool processes) are cleaned up + // rather than orphaned to init. + // + // Falls back to start_kill() (direct child only) on non-Unix or if + // the child has been polled to completion (id() returns None). + match self.child.id() { + Some(pid) if kill_process_group(pid) => {} + _ => { + let _ = self.child.start_kill(); + } + } // Bounded wait: if the child doesn't exit within 5s after SIGKILL, // give up and let Drop/OS handle it. An unbounded wait here would // wedge the harness during respawn or shutdown if a child is stuck. @@ -611,17 +617,16 @@ impl AcpClient { } }; - // Check if this is a response to our expected request (has matching id). - // Compare as serde_json::Value so both numeric and string IDs work. + // Check if this is a response to our expected request (has matching id + // AND no `method` field — a `method` field means it's an agent-initiated + // request, not a response, even if the id happens to match). if let Some(id) = msg.get("id") { - if *id == serde_json::json!(expected_id) { + if *id == serde_json::json!(expected_id) && msg.get("method").is_none() { if let Some(error) = msg.get("error") { return Err(AcpError::Protocol(error.to_string())); } return Ok(msg["result"].clone()); } - // Has an id but not ours — fall through to method dispatch - // (e.g., session/request_permission has both id and method). } // Dispatch by method name (notifications and agent-initiated requests). @@ -719,9 +724,10 @@ impl AcpClient { // Malformed lines (skipped above) don't count as real agent activity. idle_deadline = Instant::now() + idle_timeout; - // Check for matching response. + // Check for matching response (has matching id AND no `method` + // field — a `method` field means agent-initiated request, not response). if let Some(id) = msg.get("id") { - if *id == serde_json::json!(expected_id) { + if *id == serde_json::json!(expected_id) && msg.get("method").is_none() { if let Some(error) = msg.get("error") { return Err(AcpError::Protocol(error.to_string())); } @@ -1032,21 +1038,45 @@ pub fn resolve_model_switch_method( impl Drop for AcpClient { fn drop(&mut self) { - // Best-effort SIGKILL + reap. We cannot `await` in Drop (sync context), - // so we use `start_kill()` followed by `try_wait()`. The `try_wait()` - // reaps the child if it has already exited (likely after SIGKILL on most - // platforms). If the process hasn't exited yet, it will be reaped by the - // OS when the harness process itself exits (init adopts orphans). - // - // Callers SHOULD still call `shutdown().await` whenever possible for - // guaranteed reaping. Drop is the safety net for panics and early exits. - let _ = self.child.start_kill(); + // Best-effort SIGKILL + reap. We cannot `await` in Drop (sync context). + // Kill the process group when possible so subprocesses don't leak. + // Callers SHOULD still call `shutdown().await` for guaranteed reaping. + match self.child.id() { + Some(pid) if kill_process_group(pid) => {} + _ => { + let _ = self.child.start_kill(); + } + } // Non-blocking reap attempt — prevents zombie accumulation in the // common case where SIGKILL takes effect before Drop returns. let _ = self.child.try_wait(); } } +/// Send SIGKILL to an entire process group. Returns `true` if the signal was sent. +/// +/// The child is spawned with `process_group(0)`, so its PID equals its PGID. +/// Killing the group ensures subprocesses (MCP servers, tool processes) are +/// cleaned up rather than orphaned to init on repeated crash-recovery cycles. +/// +/// Uses `nix::sys::signal::killpg` — a safe wrapper around the POSIX `killpg` +/// syscall — so the crate's `#![deny(unsafe_code)]` policy is preserved. +#[cfg(unix)] +fn kill_process_group(pid: u32) -> bool { + use nix::sys::signal::{killpg, Signal}; + use nix::unistd::Pid; + + // pid == pgid because the child was spawned with process_group(0). + killpg(Pid::from_raw(pid as i32), Signal::SIGKILL).is_ok() +} + +/// Fallback for non-Unix: process-group kill not available. +/// Returns `false` so the caller falls back to `child.start_kill()`. +#[cfg(not(unix))] +fn kill_process_group(_pid: u32) -> bool { + false +} + // ─── Tests ─────────────────────────────────────────────────────────────────── #[cfg(test)] @@ -1675,6 +1705,34 @@ mod tests { assert!(matches!(result, Err(AcpError::AgentExited))); } + /// A message with both `id` and `method` is an agent-initiated request, + /// not a response. The response matcher must not consume it even if the + /// id happens to match the expected value. + #[tokio::test] + async fn agent_request_with_matching_id_not_consumed_as_response() { + // The script sends an agent-initiated request (has both id and method) + // whose id matches what we're waiting for (0), then sends the real + // response. The request should be dispatched (triggering -32601 since + // "test/method" is unknown), and the real response should be returned. + let script = r#" + echo '{"jsonrpc":"2.0","id":0,"method":"test/method","params":{}}' + read -t 2 _reply + echo '{"jsonrpc":"2.0","id":0,"result":{"ok":true}}' + sleep 1 + "#; + let mut client = spawn_script(script).await; + let hard_deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5); + let result = client + .read_until_response_with_idle_timeout( + 0, + std::time::Duration::from_secs(3), + hard_deadline, + ) + .await; + assert!(result.is_ok(), "expected Ok response, got {result:?}"); + assert_eq!(result.unwrap()["ok"], serde_json::json!(true)); + } + #[tokio::test] async fn idle_fires_before_hard_when_idle_is_shorter() { let mut client = spawn_script("sleep 10").await; @@ -1688,4 +1746,35 @@ mod tests { "idle should fire before hard when idle << hard, got {result:?}" ); } + + /// Same as `agent_request_with_matching_id_not_consumed_as_response` but + /// exercises the non-idle `read_until_response` path (via `send_request`). + #[tokio::test] + async fn agent_request_not_consumed_via_send_request() { + // Script: wait for the initialize request, reply, then send an + // agent-initiated request with id=1 (matching the next send_request id), + // wait for the -32601 error reply, then send the real response. + let script = r#" + read -t 2 _init + echo '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":1,"agentCapabilities":{}}}' + read -t 2 _req + echo '{"jsonrpc":"2.0","id":1,"method":"test/unknown","params":{}}' + read -t 2 _err_reply + echo '{"jsonrpc":"2.0","id":1,"result":{"worked":true}}' + sleep 1 + "#; + let mut client = spawn_script(script).await; + // initialize consumes id=0 + let _init = client + .initialize() + .await + .expect("initialize should succeed"); + // send_request uses id=1 — the agent's request with id=1 and method + // must not be consumed as the response. + let result = client + .send_request("test/echo", serde_json::json!({})) + .await; + assert!(result.is_ok(), "expected Ok, got {result:?}"); + assert_eq!(result.unwrap()["worked"], serde_json::json!(true)); + } } diff --git a/crates/sprout-acp/src/queue.rs b/crates/sprout-acp/src/queue.rs index 0805278339..9f07c2c40b 100644 --- a/crates/sprout-acp/src/queue.rs +++ b/crates/sprout-acp/src/queue.rs @@ -119,6 +119,8 @@ pub struct EventQueue { in_flight_channels: HashSet, /// Per-channel 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, @@ -132,6 +134,7 @@ impl EventQueue { queues: HashMap::new(), in_flight_channels: HashSet::new(), in_flight_deadlines: HashMap::new(), + in_flight_batch_sizes: HashMap::new(), retry_after: HashMap::new(), retry_counts: HashMap::new(), dedup_mode, @@ -185,9 +188,13 @@ impl EventQueue { .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, - "BUG: in-flight channel expired without mark_complete — auto-releasing" + lost_events, + deadline_secs = IN_FLIGHT_DEADLINE_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); @@ -228,6 +235,7 @@ impl EventQueue { channel_id, now + Duration::from_secs(IN_FLIGHT_DEADLINE_SECS), ); + self.in_flight_batch_sizes.insert(channel_id, events.len()); Some(FlushBatch { channel_id, events }) } @@ -245,6 +253,7 @@ impl EventQueue { 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); let now = Instant::now(); match self.retry_after.get(&channel_id) { // Active throttle → channel was requeued; keep retry_counts intact. @@ -356,16 +365,26 @@ impl EventQueue { /// Does NOT set `retry_after`. Does NOT remove from `in_flight_channels` — /// caller must call `mark_complete` separately. pub fn requeue_preserve_timestamps(&mut self, batch: FlushBatch) { - let queue = self.queues.entry(batch.channel_id).or_default(); + let channel_id = batch.channel_id; + let queue = self.queues.entry(channel_id).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 { - channel_id: batch.channel_id, + channel_id, event: be.event, prompt_tag: be.prompt_tag, received_at: be.received_at, }); } + // Enforce per-channel cap: trim newest (back) events if over limit. + while queue.len() > MAX_PENDING_PER_CHANNEL { + queue.pop_back(); + tracing::warn!( + channel_id = %channel_id, + limit = MAX_PENDING_PER_CHANNEL, + "requeue_preserve overflow — dropped newest event to enforce cap" + ); + } } /// Returns `true` if any channel has pending events that are not in-flight @@ -385,9 +404,13 @@ impl EventQueue { .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, - "BUG: in-flight channel expired without mark_complete — auto-releasing" + lost_events, + deadline_secs = IN_FLIGHT_DEADLINE_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); @@ -1495,6 +1518,81 @@ mod tests { assert!(q.flush_next().is_some()); } + // ── Test 20b: requeue_preserve_timestamps enforces per-channel cap ──────── + + #[test] + fn test_requeue_preserve_timestamps_enforces_cap() { + 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 { + q.push(make_queued(ch, &format!("fill-{i}"))); + } + assert_eq!(q.pending_count(), MAX_PENDING_PER_CHANNEL); + + // 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; + assert_eq!(q.pending_count(), 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!(q.pending_count(), MAX_PENDING_PER_CHANNEL); + + // Requeue the original batch — without cap enforcement this would + // push the queue to MAX_PENDING_PER_CHANNEL + batch_size. + q.requeue_preserve_timestamps(batch); + + // Cap must be enforced: queue should not exceed MAX_PENDING_PER_CHANNEL. + assert!( + q.pending_count() <= MAX_PENDING_PER_CHANNEL, + "queue exceeded cap: {} > {}", + q.pending_count(), + MAX_PENDING_PER_CHANNEL, + ); + } + + // ── Test 20c: requeue_preserve overflow trims newest, keeps requeued ───── + + #[test] + fn test_requeue_preserve_timestamps_overflow_keeps_requeued_events() { + 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 { + q.push(make_queued(ch, &format!("original-{i}"))); + } + + // Flush a batch — these are the "requeued" events we want to survive. + let batch = q.flush_next().expect("should flush"); + let batch_size = batch.events.len(); + + // Push new events to fill back to cap. + for i in 0..batch_size { + q.push(make_queued(ch, &format!("new-{i}"))); + } + + // Capture the content of the first requeued event for verification. + let requeued_first_content = batch.events[0].event.content.to_string(); + + // Requeue — older events go to front, overflow trims from back (newest). + q.requeue_preserve_timestamps(batch); + q.mark_complete(ch); + + // The requeued events should be at the front of the queue. + let batch2 = q.flush_next().expect("should flush after requeue"); + assert_eq!( + batch2.events[0].event.content.to_string(), + requeued_first_content, + "requeued events should be at the front (oldest), not trimmed" + ); + } + // ── Test 21: has_flushable_work returns correct results ─────────────────── #[test] diff --git a/crates/sprout-acp/src/relay.rs b/crates/sprout-acp/src/relay.rs index e368641eb0..ec42f55c1d 100644 --- a/crates/sprout-acp/src/relay.rs +++ b/crates/sprout-acp/src/relay.rs @@ -756,6 +756,27 @@ impl BgState { true } + /// Compute the `since` timestamp for a channel (re)subscribe. + /// + /// Picks the earliest of `last_seen` and `channel_dropped_since` so + /// the replay window covers both successfully processed events and any + /// that were dropped due to queue pressure. Falls back to the per-channel + /// `subscribe_since` (set at first subscribe) or `startup_watermark`. + fn channel_since(&self, channel_id: &Uuid) -> Option { + let last_seen = self.last_seen.get(channel_id).copied(); + let dropped = self.channel_dropped_since.get(channel_id).copied(); + match (last_seen, dropped) { + (Some(l), Some(d)) => Some(l.min(d)), + (Some(l), None) => Some(l), + (None, Some(d)) => Some(d), + (None, None) => self + .subscribe_since + .get(channel_id) + .copied() + .or(self.startup_watermark), + } + } + /// Clear all per-channel state for a channel that is being unsubscribed. /// Prevents stale replay on re-subscribe and avoids unbounded state growth /// for channels that are removed and never re-added. @@ -1522,15 +1543,7 @@ async fn handle_ws_message( if !state.active_subscriptions.contains_key(&channel_id) { debug!("ignoring CLOSED for already-unsubscribed channel {channel_id}"); } else { - let last_seen = state.last_seen.get(&channel_id).copied(); - let dropped = state.channel_dropped_since.get(&channel_id).copied(); - let subscribe_ts = state.subscribe_since.get(&channel_id).copied(); - let since = match (last_seen, dropped) { - (Some(l), Some(d)) => Some(l.min(d)), - (Some(l), None) => Some(l), - (None, Some(d)) => Some(d), - (None, None) => subscribe_ts.or(state.startup_watermark), - }; + let since = state.channel_since(&channel_id); let filter = match state.active_filters.get(&channel_id).cloned() { Some(f) => f, None => { @@ -1697,17 +1710,7 @@ async fn resubscribe_after_reconnect( channels.len() ); for channel_id in channels { - let last_seen = state.last_seen.get(&channel_id).copied(); - let dropped = state.channel_dropped_since.get(&channel_id).copied(); - // Fall back to per-channel subscribe_since (not startup_watermark) - // so channels joined after startup don't replay stale history. - let subscribe_ts = state.subscribe_since.get(&channel_id).copied(); - let since = match (last_seen, dropped) { - (Some(l), Some(d)) => Some(l.min(d)), - (Some(l), None) => Some(l), - (None, Some(d)) => Some(d), - (None, None) => subscribe_ts.or(state.startup_watermark), - }; + let since = state.channel_since(&channel_id); let filter = match state.active_filters.get(&channel_id).cloned() { Some(f) => f, None => { @@ -3057,14 +3060,7 @@ mod tests { state.channel_dropped_since.insert(channel_id, 900); // Compute the since value the reconnect path would use. - let last_seen = state.last_seen.get(&channel_id).copied(); - let dropped = state.channel_dropped_since.get(&channel_id).copied(); - let since = match (last_seen, dropped) { - (Some(l), Some(d)) => Some(l.min(d)), - (Some(l), None) => Some(l), - (None, Some(d)) => Some(d), - (None, None) => None, - }; + let since = state.channel_since(&channel_id); // The since passed to send_subscribe (which subtracts SINCE_SKEW_SECS internally). assert_eq!(since, Some(900), "since should be min(1000, 900) = 900");