jerry/buzz
1
0
Fork 0
mirror of https://github.com/block/buzz.git synced 2026-11-01 15:01:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(acp): follow-up hardening from PR #200 independent review (#201)

Browse Source
This commit is contained in:
tlongwell-block
2026-03-28 08:11:49 -04:00
committed by GitHub
parent e33c6bc0f2
commit 195fdfff2f
5 changed files with 254 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
Cargo.lock generated
View File

@@ -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",

3
crates/sprout-acp/Cargo.toml
View File

@@ -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 }

135
crates/sprout-acp/src/acp.rs
View File

@@ -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));
}
}

106
crates/sprout-acp/src/queue.rs
View File

@@ -119,6 +119,8 @@ pub struct EventQueue {
in_flight_channels: HashSet<Uuid>,
/// Per-channel deadline for auto-expiring stuck in-flight entries.
in_flight_deadlines: HashMap<Uuid, Instant>,
/// Number of events in each in-flight batch (for expiry logging).
in_flight_batch_sizes: HashMap<Uuid, usize>,
retry_after: HashMap<Uuid, Instant>,
/// Per-channel retry attempt counter for exponential backoff / dead-lettering.
retry_counts: HashMap<Uuid, u32>,
@@ -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]

52
crates/sprout-acp/src/relay.rs
View File

@@ -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<u64> {
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");