From 386e63f2732391a884698aa4e7398af0939c8608 Mon Sep 17 00:00:00 2001 From: Hayt <41ea58f1e64c243627e8acde7c89be667052ee6e17d8f021c1195be4324ebf04@buzz.block.builderlab.xyz> Date: Wed, 12 Aug 2026 18:39:33 -0400 Subject: [PATCH 1/3] feat(buzz-agent): gate LLM tool calls on session/request_permission MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every LLM-issued MCP tool call now asks the client to authorize it before executing. buzz-agent always asks; the client applies BUZZ_ACP_PERMISSION_POLICY. The agent never reads policy, matching the layering of the other ACP harnesses. A crate-local PermissionBroker owns the full correlation lifecycle: a process-wide admission semaphore (BUZZ_AGENT_MAX_PENDING_PERMISSIONS, default 32) acquired before any correlation entry is inserted, a monotonic id allocator, an abort-safe PendingPermission lease whose Drop synchronously removes the entry and releases the slot, claim-before-wake delivery for at-most-once resolution, and a single absolute deadline (BUZZ_AGENT_PERMISSION_TIMEOUT_SECS, default 330s) shared by admission and response wait. Cancellation races inside the wait, never depending on the outer abort drain. The request builder is version-aware, keyed on the protocol version negotiated at initialize and stored for the connection lifetime: v2 nests the tool call under subject, v1 uses the legacy top-level shape. Authorization is fail-closed: execute IFF outcome is "selected" and the selected optionId equals the offered allow option; every other shape denies with a synthetic tool error and the turn continues. Argument-shape validation is hoisted ahead of the ask so a malformed call is rejected locally without prompting. load_skill and _Stop/ _PostCompact lifecycle hooks are exempt — they are not model-issued. Tests: a subprocess + fake-MCP boundary suite proves no call reaches MCP before approval, exact-allow reaches it once, and reject/cancelled/ error/malformed/unknown-outcome/wrong-option/stale-id all fail closed, plus crossed-parallel isolation and the two exemptions; broker unit tests prove timeout, abort, and multi-session admission invariants with an injectable deadline. Co-authored-by: Will Pfleger Signed-off-by: Will Pfleger --- crates/buzz-agent/src/agent.rs | 44 +- crates/buzz-agent/src/config.rs | 25 + crates/buzz-agent/src/lib.rs | 36 +- crates/buzz-agent/src/llm.rs | 2 + crates/buzz-agent/src/mcp.rs | 33 +- crates/buzz-agent/src/permission.rs | 633 ++++++++++++++ crates/buzz-agent/src/wire.rs | 198 ++++- crates/buzz-agent/tests/bin/fake_mcp.rs | 43 + crates/buzz-agent/tests/fake_llm.rs | 24 +- .../buzz-agent/tests/permission_boundary.rs | 820 ++++++++++++++++++ crates/buzz-agent/tests/regressions.rs | 79 +- 11 files changed, 1881 insertions(+), 56 deletions(-) create mode 100644 crates/buzz-agent/src/permission.rs create mode 100644 crates/buzz-agent/tests/permission_boundary.rs diff --git a/crates/buzz-agent/src/agent.rs b/crates/buzz-agent/src/agent.rs index 9258ce449f3..0743a75667b 100644 --- a/crates/buzz-agent/src/agent.rs +++ b/crates/buzz-agent/src/agent.rs @@ -14,6 +14,7 @@ use crate::hints::SkillEntry; use crate::llm::Llm; use crate::mcp::McpRegistry; use crate::mcp::ResultBudget; +use crate::permission::PermissionDecision; use crate::types::{ AgentError, CacheTotalState, ContentBlock, HistoryItem, PricingIdentity, ProviderStop, @@ -142,6 +143,14 @@ pub struct RunCtx<'a> { pub system_prompt: &'a str, pub llm: &'a Llm, pub mcp: &'a Arc, + /// Process-wide permission broker (owned by `App`). Every LLM-issued MCP + /// tool call asks the client to authorize it through this broker before + /// executing. Shared across all sessions so the global admission cap bounds + /// simultaneously-outstanding asks process-wide. + pub permissions: &'a Arc, + /// ACP protocol version negotiated at `initialize`, fixed for the + /// connection. Selects the `session/request_permission` wire shape. + pub protocol_version: u32, /// Skills discovered at session creation; used by the built-in `load_skill` tool. pub skills: &'a [SkillEntry], pub wire: &'a WireSender, @@ -882,8 +891,10 @@ impl RunCtx<'_> { total: MAX_TOOL_RESULT_BYTES, text: self.cfg.max_tool_result_text_bytes, }; - let cancel = self.cancel.clone(); + let mut cancel = self.cancel.clone(); let sem = Arc::clone(&sem); + let permissions = Arc::clone(self.permissions); + let protocol_version = self.protocol_version; set.spawn(async move { // Acquire a permit; if the semaphore is closed (cancel), // emit a terminal wire update and skip the call. @@ -894,6 +905,37 @@ impl RunCtx<'_> { return (i, InvokeOutcome::Failed("cancelled".into())); } }; + // Argument-shape validation BEFORE the ask: a malformed + // non-object argument can never execute, so reject it locally + // without prompting the user to approve a doomed call. + if let Err(e) = crate::mcp::validate_arg_shape(&call.name, &call.arguments) { + let msg = e.to_string(); + emit_failed(&wire, &session_id, &call, &msg).await; + return (i, InvokeOutcome::Failed(msg)); + } + // Ask the client to authorize this call. The broker owns the + // full correlation lifecycle and races cancellation internally; + // every non-authorizing outcome fails closed. + match permissions + .request_permission(&wire, protocol_version, &session_id, &call, &mut cancel) + .await + { + PermissionDecision::Allowed => {} + PermissionDecision::Denied(msg) => { + emit_failed(&wire, &session_id, &call, msg).await; + return (i, InvokeOutcome::Failed(msg.into())); + } + PermissionDecision::Cancelled => { + emit_failed(&wire, &session_id, &call, "cancelled").await; + return (i, InvokeOutcome::Failed("cancelled".into())); + } + } + // Cancellation recheck: a cancel may have landed while we + // waited for approval. Do not start the call in that case. + if *cancel.borrow() { + emit_failed(&wire, &session_id, &call, "cancelled").await; + return (i, InvokeOutcome::Failed("cancelled".into())); + } emit_in_progress(&wire, &session_id, &call).await; let outcome = invoke_tool_inner(&mcp, &call, timeout, budget, cancel).await; match &outcome { diff --git a/crates/buzz-agent/src/config.rs b/crates/buzz-agent/src/config.rs index b29bf3d2fb8..ad61f4211fb 100644 --- a/crates/buzz-agent/src/config.rs +++ b/crates/buzz-agent/src/config.rs @@ -813,6 +813,18 @@ pub struct Config { /// Set via `BUZZ_AGENT_MAX_HANDOFFS`. Default 10. pub max_handoffs: usize, pub max_parallel_tools: usize, + /// Process-wide cap on simultaneously-outstanding `session/request_permission` + /// asks. Bounds the [`PermissionBroker`](crate::permission::PermissionBroker) + /// correlation map independently of the per-turn tool semaphore (which is + /// fresh per turn) and of `max_sessions` (unbounded by default). Default 32. + /// Set via `BUZZ_AGENT_MAX_PENDING_PERMISSIONS`; validated `>= 1`. + pub max_pending_permissions: usize, + /// Single absolute deadline for a permission ask — shared by broker + /// admission and the response wait, so a saturated call cannot live for two + /// full timeout windows. Default 330s, chosen to outlast the client's 300s + /// auto-deny so the answer (or auto-deny) lands first. Set via + /// `BUZZ_AGENT_PERMISSION_TIMEOUT_SECS`; validated `>= 1`. + pub permission_timeout: Duration, pub hook_timeout: Duration, /// Maximum `_Stop` rejections per prompt. Default 3. Set to 0 to /// disable `_Stop` hooks entirely (agent always honors end_turn). @@ -956,6 +968,11 @@ impl Config { max_context_tokens: parse_env("BUZZ_AGENT_MAX_CONTEXT_TOKENS", 200_000u64)?, max_handoffs: parse_env("BUZZ_AGENT_MAX_HANDOFFS", 10)?, max_parallel_tools: parse_env("BUZZ_AGENT_MAX_PARALLEL_TOOLS", 8usize)?, + max_pending_permissions: parse_env("BUZZ_AGENT_MAX_PENDING_PERMISSIONS", 32usize)?, + permission_timeout: Duration::from_secs(parse_env( + "BUZZ_AGENT_PERMISSION_TIMEOUT_SECS", + 330u64, + )?), hook_timeout: Duration::from_millis(parse_env("BUZZ_AGENT_HOOK_TIMEOUT_MS", 2500u64)?), stop_max_rejections: parse_env("BUZZ_AGENT_STOP_MAX_REJECTIONS", 3u32)?, require_reply: parse_env("BUZZ_AGENT_REQUIRE_REPLY", 0u8)? != 0, @@ -1002,6 +1019,8 @@ impl Config { max_context_tokens: 200_001, max_handoffs: 0, max_parallel_tools: 1, + max_pending_permissions: 32, + permission_timeout: Duration::from_secs(330), hook_timeout: Duration::from_secs(1), stop_max_rejections: 0, require_reply: false, @@ -1063,6 +1082,12 @@ impl Config { if self.max_parallel_tools < 1 { return Err("config: BUZZ_AGENT_MAX_PARALLEL_TOOLS must be >= 1".into()); } + if self.max_pending_permissions < 1 { + return Err("config: BUZZ_AGENT_MAX_PENDING_PERMISSIONS must be >= 1".into()); + } + if self.permission_timeout < MIN_TIMEOUT { + return Err("config: BUZZ_AGENT_PERMISSION_TIMEOUT_SECS must be >= 1".into()); + } if self.mcp_max_restart_attempts < 1 { return Err("config: BUZZ_AGENT_MCP_RESTART_MAX_ATTEMPTS must be >= 1".into()); } diff --git a/crates/buzz-agent/src/lib.rs b/crates/buzz-agent/src/lib.rs index 3e4ee3cd527..f94d60a2aa2 100644 --- a/crates/buzz-agent/src/lib.rs +++ b/crates/buzz-agent/src/lib.rs @@ -8,6 +8,7 @@ mod handoff; mod hints; mod llm; mod mcp; +mod permission; pub mod types; mod wire; @@ -31,6 +32,7 @@ pub const WINDOWS_SHELL_RESOLUTION_ENV: &[&str] = &[ use std::collections::HashMap; use std::path::Path; +use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::Arc; use serde_json::{json, Value}; @@ -53,6 +55,17 @@ struct App { cfg: Config, llm: Arc, sessions: Mutex>, + /// ACP protocol version negotiated at `initialize`, stored for the whole + /// connection lifetime. The `session/request_permission` wire shape derives + /// from this value — never from a later mutable session field — so a strict + /// client always receives exactly the shape it negotiated. Defaults to + /// [`PROTOCOL_VERSION`] before `initialize`; no prompt (and thus no + /// permission ask) can run before then. + negotiated_version: AtomicU32, + /// Owns the entire `session/request_permission` correlation lifecycle: + /// process-wide admission, id allocation, response delivery, and abort-safe + /// cleanup. See [`permission::PermissionBroker`]. + permissions: Arc, /// Cached model catalog for Databricks providers. Populated lazily on the /// first successful `session/new` discovery call. Failed discovery is never /// cached: static-token authentication errors reject session creation, while @@ -180,10 +193,16 @@ async fn async_main() { let cfg = Config::from_env().unwrap_or_else(|e| die(e)); let llm = Arc::new(Llm::new(&cfg).unwrap_or_else(|e| die(e.to_string()))); let max_line = cfg.max_line_bytes; + let permissions = Arc::new(permission::PermissionBroker::new( + cfg.max_pending_permissions, + cfg.permission_timeout, + )); let app = Arc::new(App { cfg, llm, sessions: Mutex::new(HashMap::new()), + negotiated_version: AtomicU32::new(PROTOCOL_VERSION), + permissions, models_cache: tokio::sync::OnceCell::new(), }); let (wire_tx, wire_rx) = mpsc::channel::(64); @@ -234,7 +253,10 @@ async fn dispatch(app: &Arc, msg: Value, wire_tx: &WireSender) { handle_request(app, id, method, params, wire_tx).await } Inbound::Notification { method, params } => handle_notification(app, &method, params).await, - Inbound::Ignored => {} + // Client's answer to a `session/request_permission` we issued. The + // broker matches it to a live correlation id (waking that waiter) or + // ignores an unknown/late id. + Inbound::Response { id, result } => app.permissions.deliver(&id, result), Inbound::Invalid { id, code, message } => { wire::send(wire_tx, wire::err(id, code, &message)).await } @@ -249,7 +271,7 @@ async fn handle_request( wire_tx: &WireSender, ) { match method.as_str() { - "initialize" => initialize(id, params, wire_tx).await, + "initialize" => initialize(app, id, params, wire_tx).await, "session/new" => { let app = app.clone(); let wire_tx = wire_tx.clone(); @@ -290,7 +312,7 @@ async fn handle_notification(app: &Arc, method: &str, params: Value) { } } -async fn initialize(id: Value, params: Value, wire_tx: &WireSender) { +async fn initialize(app: &Arc, id: Value, params: Value, wire_tx: &WireSender) { let p: InitializeParams = match decode(params, "initialize") { Ok(p) => p, Err(m) => return reject(wire_tx, id, INVALID_PARAMS, &m).await, @@ -302,6 +324,12 @@ async fn initialize(id: Value, params: Value, wire_tx: &WireSender) { // RFD. Revisit when that RFD merges; otherwise a genuine upstream-v2 agent // would silently lose `[Base]`. let negotiated_version = p.protocol_version.min(PROTOCOL_VERSION); + // Store the negotiated version for the connection lifetime: the + // `session/request_permission` wire shape derives from this value, never + // from a later mutable session field, so a strict client always receives + // exactly the shape it negotiated at `initialize`. + app.negotiated_version + .store(negotiated_version, Ordering::Relaxed); wire::send( wire_tx, wire::ok( @@ -730,6 +758,8 @@ async fn run_prompt(app: Arc, id: Value, params: Value, wire_tx: WireSender system_prompt: &effective_system_prompt, llm: &app.llm, mcp: &mcp, + permissions: &app.permissions, + protocol_version: app.negotiated_version.load(Ordering::Relaxed), skills: &skills, wire: &wire_tx, cancel: &mut cancel_rx, diff --git a/crates/buzz-agent/src/llm.rs b/crates/buzz-agent/src/llm.rs index a963de1e7c1..e66f230d615 100644 --- a/crates/buzz-agent/src/llm.rs +++ b/crates/buzz-agent/src/llm.rs @@ -2603,6 +2603,8 @@ mod tests { max_context_tokens: 200_000, max_handoffs: 1, max_parallel_tools: 1, + max_pending_permissions: 32, + permission_timeout: Duration::from_secs(330), hook_timeout: Duration::from_secs(1), stop_max_rejections: 0, require_reply: false, diff --git a/crates/buzz-agent/src/mcp.rs b/crates/buzz-agent/src/mcp.rs index 9ae125a0b76..42c9cc48780 100644 --- a/crates/buzz-agent/src/mcp.rs +++ b/crates/buzz-agent/src/mcp.rs @@ -594,15 +594,7 @@ impl McpRegistry { budget: ResultBudget, cancel: &mut watch::Receiver, ) -> Result { - let arg_obj = match arguments { - Value::Object(m) => Some(m.clone()), - Value::Null => None, - _ => { - return Err(AgentError::Mcp(format!( - "tool {qname} arguments must be a JSON object" - ))) - } - }; + let arg_obj = validate_arg_shape(qname, arguments)?; let mut params = CallToolRequestParams::default(); params.name = bare.to_owned().into(); params.arguments = arg_obj; @@ -812,6 +804,29 @@ async fn spawn_one( Ok((client, pgid, names, tools)) } +/// Validate that tool-call arguments are a shape the MCP transport can carry: +/// a JSON object (`Some(map)`) or absent (`None`). Any other JSON type is a +/// malformed call that the transport would reject. +/// +/// Hoisted out of `do_call` so the permission gate can run it *before* asking +/// the user: a malformed non-object argument is rejected locally without +/// prompting for approval of a call that could never execute. `do_call` runs +/// it again as the single authoritative shape check — the duplicate is a cheap +/// idempotent match, and keeping it here means no code path can reach the +/// transport with an unvalidated shape. +pub fn validate_arg_shape( + qname: &str, + arguments: &Value, +) -> Result>, AgentError> { + match arguments { + Value::Object(m) => Ok(Some(m.clone())), + Value::Null => Ok(None), + _ => Err(AgentError::Mcp(format!( + "tool {qname} arguments must be a JSON object" + ))), + } +} + /// Send `notifications/cancelled` to the MCP server, fire-and-forget. /// Per MCP spec, cancellation notifications are best-effort; we never /// block the agent on slow server stdio. diff --git a/crates/buzz-agent/src/permission.rs b/crates/buzz-agent/src/permission.rs new file mode 100644 index 00000000000..0baef5137c1 --- /dev/null +++ b/crates/buzz-agent/src/permission.rs @@ -0,0 +1,633 @@ +//! `session/request_permission` broker. +//! +//! buzz-agent asks the client to authorize every LLM-issued MCP tool call +//! *before* executing it; the client applies `BUZZ_ACP_PERMISSION_POLICY` and +//! answers. The agent never reads the policy — it always asks, matching the +//! layering of every other ACP harness. This module owns the whole request +//! correlation lifecycle so the rest of the agent only sees a single +//! `Allowed`/`Denied`/`Cancelled` decision. +//! +//! ## Invariants +//! +//! - **Process-wide admission.** The broker owns a global [`Semaphore`] +//! (`BUZZ_AGENT_MAX_PENDING_PERMISSIONS`) acquired *before* any correlation +//! entry is inserted. The per-turn `execute_parallel` semaphore is fresh per +//! turn and sessions are unbounded by default, so only this global cap bounds +//! simultaneously outstanding asks process-wide. +//! - **Abort-safe cleanup.** A successful admission returns a +//! [`PendingPermission`] lease that owns the admission permit and the +//! correlation id. Its `Drop` synchronously removes the still-pending entry +//! and releases the slot, covering task abort/panic that bypasses the normal +//! `run_prompt` tail. +//! - **Claim-before-wake / at-most-once.** [`PermissionBroker::deliver`] removes +//! the entry *before* waking the waiter, so each id resolves at most once and +//! a later lease `Drop` is a harmless no-op. +//! - **Unknown/late ids ignored.** A response whose id is not a live entry is +//! logged and dropped. +//! - **Single absolute deadline.** Admission wait and response wait share one +//! absolute deadline computed at gate entry, so a saturated call cannot live +//! for two full timeout windows. +//! - **Cancellation races inside the wait.** The waiter selects on the turn's +//! cancel receiver directly; resolution never depends on the outer abort +//! drain. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use serde_json::Value; +use tokio::sync::{oneshot, watch, OwnedSemaphorePermit, Semaphore}; +use tokio::time::Instant; + +use crate::types::ToolCall; +use crate::wire::{self, WireSender, ALLOW_OPTION_ID}; + +/// Model-visible tool error when a call is not authorized. Rides the normal +/// tool-failure path so the turn continues, matching every other tool error. +pub const PERMISSION_DENIED_MSG: &str = "permission denied: the tool call was not authorized"; + +/// Model-visible tool error when the client never answers within the deadline. +pub const PERMISSION_TIMEOUT_MSG: &str = + "permission request timed out: the tool call was not authorized"; + +/// Outcome of asking the client to authorize one tool call. +#[derive(Debug, PartialEq, Eq)] +pub enum PermissionDecision { + /// The client selected the offered allow option — execute the tool. + Allowed, + /// Every non-authorizing shape (reject, cancelled outcome, JSON-RPC error, + /// malformed response, unknown outcome, wrong/unknown optionId, timeout, + /// wire-channel closure). Fails closed with the given model-visible reason; + /// the turn continues. + Denied(&'static str), + /// The turn was cancelled while admitting or waiting. No tool runs and the + /// caller propagates cancellation exactly as the existing cancel path does. + Cancelled, +} + +/// Broker owned by `App` for the connection lifetime. +pub struct PermissionBroker { + /// Global admission cap. Acquired before any entry is inserted. + sem: Arc, + /// Live correlation entries: outbound request id -> response sender. + pending: Arc>>>, + /// Monotonic id allocator. Never reused within a process lifetime, so a + /// late response for a removed id can never collide with a fresh request. + next_id: AtomicU64, + /// Absolute deadline budget shared by admission + response wait. + timeout: Duration, +} + +impl PermissionBroker { + /// `max_pending` is validated `>= 1` by config; `timeout` is injectable so + /// broker unit tests exercise the timeout/abort paths without a 330s wait. + pub fn new(max_pending: usize, timeout: Duration) -> Self { + Self { + sem: Arc::new(Semaphore::new(max_pending.max(1))), + pending: Arc::new(Mutex::new(HashMap::new())), + next_id: AtomicU64::new(0), + timeout, + } + } + + /// Number of live (unresolved, un-dropped) correlation entries. Test-only + /// observability for the drop-guard and delivery invariants. + #[cfg(test)] + pub fn pending_count(&self) -> usize { + self.pending.lock().unwrap().len() + } + + /// Free admission slots. Test-only, so a test can prove a terminal path + /// (delivery, timeout, cancel, drop) actually released the capacity it + /// held rather than leaking it. + #[cfg(test)] + pub fn available_permits(&self) -> usize { + self.sem.available_permits() + } + + /// Deliver a client response to its waiter. Claims (removes) the entry + /// before waking so the id resolves at most once; unknown/late ids are + /// logged and ignored. `result` is the JSON-RPC `result` field (or + /// `Value::Null` for an error/malformed response — every such shape fails + /// the authorization predicate and denies). + pub fn deliver(&self, id: &Value, result: Value) { + let Some(key) = parse_id(id) else { + tracing::debug!(target: "permission", "ignoring response with unrecognized id {id}"); + return; + }; + // Claim before wake: remove first, then send into the removed sender. + let sender = self.pending.lock().unwrap().remove(&key); + match sender { + Some(tx) => { + // The receiver may already be gone (waiter cancelled/timed out + // and dropped the lease); a failed send is a harmless no-op. + let _ = tx.send(result); + } + None => { + tracing::debug!(target: "permission", "ignoring unknown/late permission id {id}"); + } + } + } + + /// Ask the client to authorize `call`, returning the decision. + /// + /// Sequence: acquire global admission (racing cancel + deadline) → insert + /// correlation entry (held by an abort-safe lease) → send the version-aware + /// request → wait for the response (racing cancel + deadline). One absolute + /// deadline bounds both waits. + pub async fn request_permission( + &self, + wire: &WireSender, + version: u32, + session_id: &str, + call: &ToolCall, + cancel: &mut watch::Receiver, + ) -> PermissionDecision { + let deadline = Instant::now() + self.timeout; + + // ── Admission ────────────────────────────────────────────────────── + // Early cancel check: watch::changed() only fires on NEW writes. + if *cancel.borrow() { + return PermissionDecision::Cancelled; + } + let permit = tokio::select! { + biased; + _ = cancel.changed() => return PermissionDecision::Cancelled, + _ = tokio::time::sleep_until(deadline) => { + return PermissionDecision::Denied(PERMISSION_TIMEOUT_MSG); + } + p = Arc::clone(&self.sem).acquire_owned() => match p { + Ok(p) => p, + // Semaphore is never closed in production; treat as fail-closed. + Err(_) => return PermissionDecision::Denied(PERMISSION_DENIED_MSG), + }, + }; + + // Insert the correlation entry under the owned permit. The lease's Drop + // removes the entry + releases the slot on every exit path below, + // including task abort. + let mut lease = self.register(permit); + + // ── Send the version-aware request ───────────────────────────────── + let params = wire::request_permission_params( + version, + session_id, + &call.provider_id, + &call.name, + &call.arguments, + ); + wire::send( + wire, + wire::request_permission(lease.id_value.clone(), params), + ) + .await; + + // ── Response wait ────────────────────────────────────────────────── + if *cancel.borrow() { + return PermissionDecision::Cancelled; + } + tokio::select! { + biased; + _ = cancel.changed() => PermissionDecision::Cancelled, + r = &mut lease.rx => match r { + Ok(result) => evaluate(&result), + // Sender dropped without sending — should not happen (delivery + // always sends before drop); fail closed. + Err(_) => PermissionDecision::Denied(PERMISSION_DENIED_MSG), + }, + _ = tokio::time::sleep_until(deadline) => { + PermissionDecision::Denied(PERMISSION_TIMEOUT_MSG) + } + } + // `lease` drops here: entry removed (no-op if delivered) + slot released. + } + + /// Allocate an id, insert its response sender, and return the abort-safe + /// lease holding the receiver + owned permit. + fn register(&self, permit: OwnedSemaphorePermit) -> PendingPermission { + let id = self.next_id.fetch_add(1, Ordering::Relaxed); + let (tx, rx) = oneshot::channel(); + self.pending.lock().unwrap().insert(id, tx); + PendingPermission { + id, + id_value: Value::String(format!("perm-{id}")), + rx, + pending: Arc::clone(&self.pending), + _permit: permit, + } + } +} + +/// Abort-safe correlation lease. Owns the admission permit and the correlation +/// id; its `Drop` synchronously removes the still-pending entry and releases +/// the slot. Delivery removes the entry first, so a later drop is a no-op. +struct PendingPermission { + id: u64, + id_value: Value, + rx: oneshot::Receiver, + pending: Arc>>>, + _permit: OwnedSemaphorePermit, +} + +impl Drop for PendingPermission { + fn drop(&mut self) { + // Synchronous, non-async removal — safe from a Drop and required for + // abort/panic paths. No-op if delivery already claimed the entry. + self.pending.lock().unwrap().remove(&self.id); + // `_permit` drops → global admission slot released. + } +} + +/// The authorization predicate, stated once: execute IFF the client selected an +/// option AND the selected `optionId` equals exactly this request's offered +/// allow-option id. Every other shape fails closed. +fn evaluate(result: &Value) -> PermissionDecision { + let outcome = &result["outcome"]; + if outcome["outcome"] == "selected" && outcome["optionId"].as_str() == Some(ALLOW_OPTION_ID) { + PermissionDecision::Allowed + } else { + PermissionDecision::Denied(PERMISSION_DENIED_MSG) + } +} + +/// Recover the correlation key from an outbound request id echoed by the +/// client. Only ids we minted (`perm-`) are ours; anything else is a +/// foreign/stale id and is ignored. +fn parse_id(id: &Value) -> Option { + id.as_str()?.strip_prefix("perm-")?.parse().ok() +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + use tokio::sync::mpsc; + + const LONG: Duration = Duration::from_secs(30); + const SHORT: Duration = Duration::from_millis(60); + + fn tool_call() -> ToolCall { + ToolCall { + provider_id: "fake".into(), + name: "fake__shell".into(), + arguments: json!({ "command": "ls" }), + provider_extra: serde_json::Map::new(), + } + } + + fn selected(option_id: &str) -> Value { + json!({ "outcome": { "outcome": "selected", "optionId": option_id } }) + } + + /// Pull the next outbound frame off the wire and return its JSON-RPC `id`. + /// Reading it also proves the request was registered and sent (delivery + /// only happens after `register`). + async fn next_request_id(rx: &mut mpsc::Receiver) -> Value { + let wire::WireMsg::Notify(v) = tokio::time::timeout(Duration::from_secs(1), rx.recv()) + .await + .expect("a request frame") + .expect("wire open"); + assert_eq!(v["method"], "session/request_permission"); + v["id"].clone() + } + + // ── Authorization predicate (fail-closed) ──────────────────────────────── + + #[test] + fn test_selected_allow_option_authorizes() { + assert_eq!( + evaluate(&selected(ALLOW_OPTION_ID)), + PermissionDecision::Allowed + ); + } + + #[test] + fn test_every_non_allow_shape_denies() { + // reject, wrong/unknown option id, unknown outcome, missing fields, + // empty object — the full adversarial set the predicate must reject. + let denied = [ + selected("reject_once"), + selected("some_unknown_option"), + json!({ "outcome": { "outcome": "cancelled" } }), + json!({ "outcome": { "outcome": "selected" } }), // no optionId + json!({ "outcome": { "outcome": "banana", "optionId": ALLOW_OPTION_ID } }), + json!({ "outcome": {} }), + json!({}), + Value::Null, + ]; + for shape in denied { + assert_eq!( + evaluate(&shape), + PermissionDecision::Denied(PERMISSION_DENIED_MSG), + "shape must fail closed: {shape}" + ); + } + } + + // ── Id correlation ─────────────────────────────────────────────────────── + + #[test] + fn test_parse_id_accepts_only_minted_ids() { + assert_eq!(parse_id(&json!("perm-0")), Some(0)); + assert_eq!(parse_id(&json!("perm-42")), Some(42)); + assert_eq!(parse_id(&json!("perm-x")), None); + assert_eq!(parse_id(&json!("42")), None); // foreign numeric-string id + assert_eq!(parse_id(&json!(42)), None); // foreign numeric id + assert_eq!(parse_id(&Value::Null), None); + } + + // ── Delivery: exact allow / deny ────────────────────────────────────────── + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_deliver_allow_authorizes_and_frees_slot() { + let broker = Arc::new(PermissionBroker::new(4, LONG)); + let (tx, mut rx) = mpsc::channel(8); + let (_cancel_tx, mut cancel_rx) = watch::channel(false); + + let b = Arc::clone(&broker); + let call = tool_call(); + let task = tokio::spawn(async move { + b.request_permission(&tx, 2, "ses_a", &call, &mut cancel_rx) + .await + }); + + let id = next_request_id(&mut rx).await; + assert_eq!(broker.pending_count(), 1); + assert_eq!(broker.available_permits(), 3); + + broker.deliver(&id, selected(ALLOW_OPTION_ID)); + assert_eq!(task.await.unwrap(), PermissionDecision::Allowed); + assert_eq!(broker.pending_count(), 0, "entry claimed on delivery"); + assert_eq!( + broker.available_permits(), + 4, + "slot released after decision" + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_deliver_reject_denies_and_frees_slot() { + let broker = Arc::new(PermissionBroker::new(4, LONG)); + let (tx, mut rx) = mpsc::channel(8); + let (_cancel_tx, mut cancel_rx) = watch::channel(false); + + let b = Arc::clone(&broker); + let call = tool_call(); + let task = tokio::spawn(async move { + b.request_permission(&tx, 2, "ses_a", &call, &mut cancel_rx) + .await + }); + + let id = next_request_id(&mut rx).await; + broker.deliver(&id, selected("reject_once")); + assert_eq!( + task.await.unwrap(), + PermissionDecision::Denied(PERMISSION_DENIED_MSG) + ); + assert_eq!(broker.pending_count(), 0); + assert_eq!(broker.available_permits(), 4); + } + + // ── Stale / unknown id ignored ──────────────────────────────────────────── + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_unknown_id_does_not_unblock_waiter() { + let broker = Arc::new(PermissionBroker::new(4, LONG)); + let (tx, mut rx) = mpsc::channel(8); + let (_cancel_tx, mut cancel_rx) = watch::channel(false); + + let b = Arc::clone(&broker); + let call = tool_call(); + let task = tokio::spawn(async move { + b.request_permission(&tx, 2, "ses_a", &call, &mut cancel_rx) + .await + }); + + let real_id = next_request_id(&mut rx).await; + // A stale/foreign id is dropped; the live entry survives. + broker.deliver(&json!("perm-999"), selected(ALLOW_OPTION_ID)); + broker.deliver(&json!(1), selected(ALLOW_OPTION_ID)); + broker.deliver(&Value::Null, selected(ALLOW_OPTION_ID)); + assert_eq!(broker.pending_count(), 1, "waiter still pending"); + + // The correct id resolves it exactly once. + broker.deliver(&real_id, selected(ALLOW_OPTION_ID)); + assert_eq!(task.await.unwrap(), PermissionDecision::Allowed); + assert_eq!(broker.pending_count(), 0); + } + + // ── Timeout ─────────────────────────────────────────────────────────────── + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_timeout_denies_and_removes_state() { + let broker = Arc::new(PermissionBroker::new(4, SHORT)); + let (tx, _rx) = mpsc::channel(8); + let (_cancel_tx, mut cancel_rx) = watch::channel(false); + let call = tool_call(); + + // No delivery ever arrives: the shared deadline denies. + let decision = broker + .request_permission(&tx, 2, "ses_a", &call, &mut cancel_rx) + .await; + assert_eq!(decision, PermissionDecision::Denied(PERMISSION_TIMEOUT_MSG)); + assert_eq!( + broker.pending_count(), + 0, + "timeout removes correlation state" + ); + assert_eq!(broker.available_permits(), 4, "timeout releases the slot"); + } + + // ── Cancellation while waiting ──────────────────────────────────────────── + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_cancel_while_waiting_returns_cancelled_and_removes_state() { + let broker = Arc::new(PermissionBroker::new(4, LONG)); + let (tx, mut rx) = mpsc::channel(8); + let (cancel_tx, mut cancel_rx) = watch::channel(false); + + let b = Arc::clone(&broker); + let call = tool_call(); + let task = tokio::spawn(async move { + b.request_permission(&tx, 2, "ses_a", &call, &mut cancel_rx) + .await + }); + + let _id = next_request_id(&mut rx).await; + assert_eq!(broker.pending_count(), 1); + cancel_tx.send(true).unwrap(); + assert_eq!(task.await.unwrap(), PermissionDecision::Cancelled); + assert_eq!(broker.pending_count(), 0); + assert_eq!(broker.available_permits(), 4); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_precancelled_turn_never_sends_request() { + let broker = Arc::new(PermissionBroker::new(4, LONG)); + let (tx, mut rx) = mpsc::channel(8); + let (_cancel_tx, mut cancel_rx) = watch::channel(true); // already cancelled + let call = tool_call(); + + let decision = broker + .request_permission(&tx, 2, "ses_a", &call, &mut cancel_rx) + .await; + assert_eq!(decision, PermissionDecision::Cancelled); + assert_eq!(broker.pending_count(), 0, "no entry inserted"); + assert_eq!(broker.available_permits(), 4); + assert!(rx.try_recv().is_err(), "no request frame emitted"); + } + + // ── Abort-safe drop guard ───────────────────────────────────────────────── + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_abort_while_waiting_leaves_zero_pending_and_reusable_slot() { + let broker = Arc::new(PermissionBroker::new(1, LONG)); + let (tx, mut rx) = mpsc::channel(8); + let (_cancel_tx, mut cancel_rx) = watch::channel(false); + + let b = Arc::clone(&broker); + let call = tool_call(); + let task = tokio::spawn(async move { + b.request_permission(&tx, 2, "ses_a", &call, &mut cancel_rx) + .await + }); + + // Registered + sent → then hard-abort the task (bypasses every normal + // exit path). The lease's Drop must still run. + let _id = next_request_id(&mut rx).await; + assert_eq!(broker.pending_count(), 1); + assert_eq!(broker.available_permits(), 0); + task.abort(); + let _ = task.await; + assert_eq!( + broker.pending_count(), + 0, + "drop guard removed the entry on abort" + ); + assert_eq!( + broker.available_permits(), + 1, + "drop guard released the slot on abort" + ); + } + + // ── Process-wide admission cap across multiple sessions ─────────────────── + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_admission_cap_bounds_entries_across_sessions() { + // Capacity 2, shared by all sessions. Three distinct sessions ask at + // once; only two can register/send while the cap is saturated. The + // third is admitted only after a slot frees. Frame ids arrive in + // nondeterministic order across tasks, so the test never maps a task + // handle to a specific id — it proves the bound structurally (frame + // count + pending_count) and that every task ultimately resolves. + let broker = Arc::new(PermissionBroker::new(2, LONG)); + let (tx, mut rx) = mpsc::channel(8); + let (_c, cancel_rx) = watch::channel(false); + + let spawn_req = |session: &'static str| { + let b = Arc::clone(&broker); + let tx = tx.clone(); + let mut cancel = cancel_rx.clone(); + let call = tool_call(); + tokio::spawn(async move { + b.request_permission(&tx, 2, session, &call, &mut cancel) + .await + }) + }; + + let tasks = [spawn_req("ses_a"), spawn_req("ses_b"), spawn_req("ses_c")]; + + // Only two frames appear while the cap is 2; the third is blocked in + // admission with no entry and no frame. + let id1 = next_request_id(&mut rx).await; + let id2 = next_request_id(&mut rx).await; + assert_eq!(broker.pending_count(), 2); + assert_eq!(broker.available_permits(), 0); + assert!( + tokio::time::timeout(SHORT, rx.recv()).await.is_err(), + "third session must not send a request while the cap is saturated" + ); + assert_eq!( + broker.pending_count(), + 2, + "cap holds: no third entry inserted" + ); + + // Free one slot → the third session is admitted and sends its frame. + broker.deliver(&id1, selected(ALLOW_OPTION_ID)); + let id3 = next_request_id(&mut rx).await; + assert_eq!(broker.pending_count(), 2, "still bounded after churn"); + + // Resolve the two remaining live entries. + broker.deliver(&id2, selected(ALLOW_OPTION_ID)); + broker.deliver(&id3, selected(ALLOW_OPTION_ID)); + + // Every session resolved to Allowed — none stranded or timed out. + for t in tasks { + assert_eq!(t.await.unwrap(), PermissionDecision::Allowed); + } + assert_eq!(broker.pending_count(), 0); + assert_eq!(broker.available_permits(), 2); + } + + // ── Admission-phase cancel: fail-closed, zero entries ───────────────────── + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_cancel_during_admission_inserts_no_entry() { + // Saturate the single slot directly (test-module access to `sem`) so the + // request under test blocks in the admission phase, before any insert. + let broker = Arc::new(PermissionBroker::new(1, LONG)); + let held = Arc::clone(&broker.sem).acquire_owned().await.unwrap(); + assert_eq!(broker.available_permits(), 0); + + let (tx, mut rx) = mpsc::channel(8); + let (cancel_tx, mut cancel_rx) = watch::channel(false); + let b = Arc::clone(&broker); + let call = tool_call(); + let task = tokio::spawn(async move { + b.request_permission(&tx, 2, "ses_a", &call, &mut cancel_rx) + .await + }); + + // It cannot proceed past admission: no entry, no frame. + assert!(tokio::time::timeout(SHORT, rx.recv()).await.is_err()); + assert_eq!(broker.pending_count(), 0, "blocked before insert"); + + cancel_tx.send(true).unwrap(); + assert_eq!(task.await.unwrap(), PermissionDecision::Cancelled); + assert_eq!( + broker.pending_count(), + 0, + "cancel during admission inserts nothing" + ); + drop(held); + assert_eq!(broker.available_permits(), 1); + } + + // ── Admission-phase deadline: fail-closed deny, zero entries ────────────── + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_deadline_during_admission_denies_with_no_entry() { + let broker = Arc::new(PermissionBroker::new(1, SHORT)); + let held = Arc::clone(&broker.sem).acquire_owned().await.unwrap(); + + let (tx, mut rx) = mpsc::channel(8); + let (_cancel_tx, mut cancel_rx) = watch::channel(false); + let call = tool_call(); + + // Slot never frees within the deadline → admission times out. + let decision = broker + .request_permission(&tx, 2, "ses_a", &call, &mut cancel_rx) + .await; + assert_eq!(decision, PermissionDecision::Denied(PERMISSION_TIMEOUT_MSG)); + assert_eq!( + broker.pending_count(), + 0, + "no entry inserted on admission timeout" + ); + assert!(rx.try_recv().is_err(), "no request frame emitted"); + drop(held); + } +} diff --git a/crates/buzz-agent/src/wire.rs b/crates/buzz-agent/src/wire.rs index b4c876e0fe3..2015e22378c 100644 --- a/crates/buzz-agent/src/wire.rs +++ b/crates/buzz-agent/src/wire.rs @@ -27,7 +27,15 @@ pub enum Inbound { method: String, params: Value, }, - Ignored, + /// A bare JSON-RPC response (id present, no method) — the client's answer + /// to a request buzz-agent issued. Today the only such request is + /// `session/request_permission`. `result` carries the JSON-RPC `result` + /// field, or `Null` for an `error`/malformed response; every non-`selected` + /// shape fails the broker's authorization predicate and denies. + Response { + id: Value, + result: Value, + }, Invalid { id: Value, code: i32, @@ -109,9 +117,14 @@ pub fn classify(msg: &Value) -> Inbound { params, }, (Some(m), None) => Inbound::Notification { method: m, params }, - // Bare responses (id present, no method) are unexpected — buzz-agent - // does not issue requests to the client. Ignore silently. - (None, Some(_)) => Inbound::Ignored, + // Bare responses (id present, no method) answer a request buzz-agent + // issued — today only `session/request_permission`. Route the `result` + // (or `Null` on an `error`/absent result) to the permission broker, + // which matches it to a live correlation id or ignores it if unknown. + (None, Some(id)) => Inbound::Response { + id, + result: msg.get("result").cloned().unwrap_or(Value::Null), + }, (None, None) => Inbound::Invalid { id: Value::Null, code: INVALID_REQUEST, @@ -120,6 +133,79 @@ pub fn classify(msg: &Value) -> Inbound { } } +/// `optionId`/`kind` of the single allow option offered on every +/// `session/request_permission`. buzz-acp's answering side selects the option +/// whose `kind == "allow_once"` (never by hardcoded `optionId`), and the +/// authorization predicate on this side requires the returned `optionId` to +/// equal exactly this value. Keeping option id and kind identical means both +/// sides agree without a separate lookup table. +pub const ALLOW_OPTION_ID: &str = "allow_once"; + +/// The two options offered on every permission request: allow-once and +/// reject-once. First cut ships only these (no session-scoped grant), so every +/// offered option is already in the desktop card's exact actionable allowlist. +fn permission_options() -> Value { + json!([ + { "optionId": ALLOW_OPTION_ID, "name": "Allow", "kind": ALLOW_OPTION_ID }, + { "optionId": "reject_once", "name": "Deny", "kind": "reject_once" }, + ]) +} + +/// Build `session/request_permission` params for the negotiated protocol +/// version. No hybrid shapes — the request must match exactly what the client +/// negotiated at `initialize`, or a strict client can reject it before policy +/// is applied. +/// +/// - **v2** (what buzz-agent negotiates with current buzz-acp): tool context +/// lives under `subject: {type: "tool_call", toolCall}` with top-level +/// `title` and `options`. +/// - **v1** (still negotiated when a client requests it): the legacy shape with +/// `toolCall` (carrying `kind`) directly at the params level. +pub fn request_permission_params( + version: u32, + session_id: &str, + tool_call_id: &str, + title: &str, + raw_input: &Value, +) -> Value { + if version >= 2 { + json!({ + "sessionId": session_id, + "title": title, + "subject": { + "type": "tool_call", + "toolCall": { + "toolCallId": tool_call_id, + "title": title, + "rawInput": raw_input, + }, + }, + "options": permission_options(), + }) + } else { + json!({ + "sessionId": session_id, + "toolCall": { + "toolCallId": tool_call_id, + "title": title, + "kind": "other", + "rawInput": raw_input, + }, + "options": permission_options(), + }) + } +} + +/// Build an outbound JSON-RPC request `session/request_permission` frame. +pub fn request_permission(id: Value, params: Value) -> Value { + json!({ + "jsonrpc": "2.0", + "id": id, + "method": "session/request_permission", + "params": params, + }) +} + pub fn ok(id: Value, result: Value) -> Value { json!({ "jsonrpc": "2.0", "id": id, "result": result }) } @@ -534,4 +620,108 @@ mod tests { assert_eq!(payload["accumulatedInputTokens"], serde_json::json!(1000)); assert_eq!(payload["accumulatedOutputTokens"], serde_json::json!(200)); } + + // ── request_permission_params: version-aware wire shape ────────────────── + + /// v2 (what buzz-agent negotiates with current buzz-acp): tool context is + /// nested under `subject: {type: "tool_call", toolCall}` with top-level + /// `title` and `options`, matching the ACP v2 `RequestPermissionRequest`. + #[test] + fn request_permission_params_v2_nests_tool_call_under_subject() { + let raw = json!({ "command": "ls" }); + let p = request_permission_params(2, "ses_1", "fake__shell", "fake__shell", &raw); + + assert_eq!(p["sessionId"], "ses_1"); + assert_eq!(p["title"], "fake__shell"); + assert_eq!(p["subject"]["type"], "tool_call"); + assert_eq!(p["subject"]["toolCall"]["toolCallId"], "fake__shell"); + assert_eq!(p["subject"]["toolCall"]["title"], "fake__shell"); + assert_eq!(p["subject"]["toolCall"]["rawInput"], raw); + // No hybrid: v2 must NOT carry a top-level `toolCall`. + assert!(p.get("toolCall").is_none(), "v2 must not use the v1 shape"); + assert_options(&p["options"]); + } + + /// v1 (still negotiated when a client requests it): the legacy shape with + /// `toolCall` (carrying `kind`) directly at the params level, no `subject`. + #[test] + fn request_permission_params_v1_uses_legacy_top_level_tool_call() { + let raw = json!({ "command": "ls" }); + let p = request_permission_params(1, "ses_1", "fake__shell", "fake__shell", &raw); + + assert_eq!(p["sessionId"], "ses_1"); + assert_eq!(p["toolCall"]["toolCallId"], "fake__shell"); + assert_eq!(p["toolCall"]["title"], "fake__shell"); + assert_eq!(p["toolCall"]["kind"], "other"); + assert_eq!(p["toolCall"]["rawInput"], raw); + // No hybrid: v1 must NOT carry the v2 `subject` or top-level `title`. + assert!(p.get("subject").is_none(), "v1 must not use the v2 shape"); + assert!(p.get("title").is_none(), "v1 has no top-level title"); + assert_options(&p["options"]); + } + + /// Both offered options are exactly allow-once and reject-once, with + /// `optionId == kind` so buzz-acp's `kind`-based selector and this side's + /// `optionId`-based predicate agree without a lookup table. + fn assert_options(options: &Value) { + let opts = options.as_array().expect("options is an array"); + assert_eq!(opts.len(), 2, "first cut offers exactly two options"); + assert_eq!(opts[0]["optionId"], ALLOW_OPTION_ID); + assert_eq!(opts[0]["kind"], ALLOW_OPTION_ID); + assert_eq!(opts[0]["name"], "Allow"); + assert_eq!(opts[1]["optionId"], "reject_once"); + assert_eq!(opts[1]["kind"], "reject_once"); + assert_eq!(opts[1]["name"], "Deny"); + } + + /// The outbound frame wraps params in a JSON-RPC request whose id echoes + /// back verbatim so the broker can correlate the response. + #[test] + fn request_permission_frame_is_a_correlatable_jsonrpc_request() { + let params = request_permission_params(2, "ses_1", "t", "t", &json!({})); + let frame = request_permission(json!("perm-7"), params); + assert_eq!(frame["jsonrpc"], "2.0"); + assert_eq!(frame["id"], "perm-7"); + assert_eq!(frame["method"], "session/request_permission"); + assert_eq!(frame["params"]["sessionId"], "ses_1"); + } + + // ── classify: bare responses route to the broker ───────────────────────── + + /// A bare JSON-RPC response (id, no method) is the client's answer to a + /// request buzz-agent issued; it routes to the broker with its `result`. + #[test] + fn classify_bare_response_routes_to_broker() { + let msg = json!({ + "jsonrpc": "2.0", + "id": "perm-3", + "result": { "outcome": { "outcome": "selected", "optionId": ALLOW_OPTION_ID } }, + }); + match classify(&msg) { + Inbound::Response { id, result } => { + assert_eq!(id, json!("perm-3")); + assert_eq!(result["outcome"]["outcome"], "selected"); + } + other => panic!("expected Response, got {other:?}"), + } + } + + /// A JSON-RPC error response (id, `error`, no `result`) still routes to the + /// broker but with `result == Null`, which the authorization predicate + /// fails closed. buzz-agent never leaves the waiter hanging on an error. + #[test] + fn classify_error_response_routes_with_null_result() { + let msg = json!({ + "jsonrpc": "2.0", + "id": "perm-3", + "error": { "code": -32601, "message": "method not found" }, + }); + match classify(&msg) { + Inbound::Response { id, result } => { + assert_eq!(id, json!("perm-3")); + assert_eq!(result, Value::Null, "error/absent result → Null → deny"); + } + other => panic!("expected Response, got {other:?}"), + } + } } diff --git a/crates/buzz-agent/tests/bin/fake_mcp.rs b/crates/buzz-agent/tests/bin/fake_mcp.rs index 1b7f3461624..8d96779bbac 100644 --- a/crates/buzz-agent/tests/bin/fake_mcp.rs +++ b/crates/buzz-agent/tests/bin/fake_mcp.rs @@ -23,6 +23,12 @@ //! tree dies on timeout. //! FAKE_MCP_GRANDCHILD_PID_FILE=path //! — path to write the grandchild PID to. +//! FAKE_MCP_CANCEL_LOG=path — append each `notifications/cancelled` frame to +//! `path` (one JSON line per notification). +//! FAKE_MCP_CALL_LOG=path — append the tool name of each `tools/call` to +//! `path` (one name per line). Lets a test assert +//! a tool was invoked exactly once, or never — the +//! permission gate's core proof. //! FAKE_MCP_STOP_HOOK=1 — expose a `_Stop` hook tool //! FAKE_MCP_STOP_TEXT=text — `_Stop` returns this text (default: "keep going") //! FAKE_MCP_STOP_DELAY=N — `_Stop` sleeps N seconds before replying @@ -39,6 +45,11 @@ //! `command` string. Lets a test drive the //! reply guard's recognition of a real, //! registered shell tool. +//! FAKE_MCP_NAMED_TOOLS=a,b — expose one no-arg tool per comma-separated bare +//! name (each registered as `__`), in +//! addition to any `FAKE_MCP_TOOL_COUNT` tools. Lets +//! a test issue parallel calls to distinctly named +//! tools and tell them apart in `FAKE_MCP_CALL_LOG`. use std::io::{BufRead, Write}; @@ -83,6 +94,7 @@ fn make_tools( include_stop_hook: bool, include_post_compact_hook: bool, include_shell_tool: bool, + named_tools: &[String], ) -> Vec { let mut tools: Vec = (0..count) .map(|i| { @@ -93,6 +105,13 @@ fn make_tools( }) }) .collect(); + for name in named_tools { + tools.push(json!({ + "name": name, + "description": "named test tool", + "inputSchema": { "type": "object", "properties": {} }, + })); + } if include_stop_hook { tools.push(json!({ "name": "_Stop", @@ -156,6 +175,14 @@ fn main() { let post_compact_hook = env_flag("FAKE_MCP_POSTCOMPACT_HOOK"); let shell_tool = env_flag("FAKE_MCP_SHELL_TOOL"); let post_compact_text = std::env::var("FAKE_MCP_POSTCOMPACT_TEXT").unwrap_or_default(); + // One extra no-arg tool per comma-separated bare name. + let named_tools: Vec = std::env::var("FAKE_MCP_NAMED_TOOLS") + .unwrap_or_default() + .split(',') + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(str::to_owned) + .collect(); // Use a channel-based stdin reader so notifications (which carry no id) // are captured even while the main thread is sleeping during a tool call. @@ -231,6 +258,7 @@ fn main() { stop_hook, post_compact_hook, shell_tool, + &named_tools, ) }), ); @@ -248,6 +276,21 @@ fn main() { .and_then(|p| p.get("name")) .and_then(Value::as_str) .unwrap_or(""); + // Append every invoked tool name so a test can prove a call + // reached the server exactly once (or never). This fires for + // ALL tools/call, including `_Stop`/`_PostCompact` hooks, so a + // test can also prove hooks are NOT permission-gated by + // observing they still reach the server without an ask. + if let Ok(path) = std::env::var("FAKE_MCP_CALL_LOG") { + use std::io::Write as _; + if let Ok(mut f) = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&path) + { + let _ = writeln!(f, "{called_name}"); + } + } // Optionally spawn a long-sleeping grandchild so the test // can verify process-group killing reaches the whole tree. if env_flag("FAKE_MCP_SPAWN_GRANDCHILD") { diff --git a/crates/buzz-agent/tests/fake_llm.rs b/crates/buzz-agent/tests/fake_llm.rs index 4253ef329c1..304f1883f26 100644 --- a/crates/buzz-agent/tests/fake_llm.rs +++ b/crates/buzz-agent/tests/fake_llm.rs @@ -267,6 +267,23 @@ fn openai_tool_call(id: &str, name: &str, args: Value) -> Value { }) } +/// Select the offered option whose `kind == "allow_once"` and return the +/// `session/request_permission` response. Mirrors buzz-acp's answering side, +/// which selects by `kind`, never by a hardcoded `optionId`. Centralizing this +/// means a future option-id rename can't silently turn allow into a denial. +fn approve_permission(request: &Value) -> Value { + let option_id = request["params"]["options"] + .as_array() + .and_then(|opts| opts.iter().find(|o| o["kind"] == "allow_once")) + .and_then(|o| o["optionId"].as_str()) + .expect("request must offer an allow_once option"); + json!({ + "jsonrpc": "2.0", + "id": request["id"], + "result": { "outcome": { "outcome": "selected", "optionId": option_id } }, + }) +} + async fn init_session(h: &mut Harness) -> String { h.send( "initialize", @@ -396,12 +413,7 @@ async fn unsupported_image_response_recovers_without_replaying_image() { loop { let message = h.recv().await; if message.get("method") == Some(&json!("session/request_permission")) { - h.write(json!({ - "jsonrpc": "2.0", - "id": message["id"], - "result": { "outcome": { "outcome": "selected", "optionId": "allow" } }, - })) - .await; + h.write(approve_permission(&message)).await; } else if message["id"] == json!(prompt_id) { assert_eq!(message["result"]["stopReason"], "end_turn"); break; diff --git a/crates/buzz-agent/tests/permission_boundary.rs b/crates/buzz-agent/tests/permission_boundary.rs new file mode 100644 index 00000000000..f6651f0efbc --- /dev/null +++ b/crates/buzz-agent/tests/permission_boundary.rs @@ -0,0 +1,820 @@ +//! Production authorization-boundary tests for the `session/request_permission` +//! surface. +//! +//! These drive a real `buzz-agent` subprocess against a fake MCP server and a +//! capturing LLM, and prove the security invariant end to end: an LLM-issued +//! MCP tool call reaches the server IFF the client selected the offered +//! allow-once option, and every other outcome fails closed without invoking the +//! tool. `fake_mcp` appends each invoked *bare* tool name to `FAKE_MCP_CALL_LOG` +//! (fired for `_Stop`/`_PostCompact` too), so "reached the tool exactly once" +//! and "never reached the tool" are both directly observable from disk. +//! +//! Timeout/abort/multi-session-cap state invariants live in the broker-seam unit +//! tests (`src/permission.rs`), which use an injectable deadline and inspect the +//! private correlation map — neither of which a subprocess can do. +//! +//! The `Harness` is copied from `regressions.rs`: each integration test file is +//! its own binary, so a self-contained harness is the established convention. + +use std::collections::VecDeque; +use std::process::Stdio; +use std::sync::{Arc, Mutex as StdMutex}; +use std::time::Duration; + +use serde_json::{json, Value}; +use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; +use tokio::net::TcpListener; +use tokio::sync::Mutex; + +// ───────────────────────────────────────────────────────────────────────────── +// Capturing LLM +// ───────────────────────────────────────────────────────────────────────────── + +struct CapturingLlm { + url: String, + #[allow(dead_code)] + captured: Arc>>, +} + +async fn spawn_capturing_llm(responses: Vec) -> CapturingLlm { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let queue = Arc::new(Mutex::new(VecDeque::from(responses))); + let captured: Arc>> = Arc::new(Mutex::new(Vec::new())); + let cap2 = captured.clone(); + tokio::spawn(async move { + loop { + let (mut sock, _) = match listener.accept().await { + Ok(p) => p, + Err(_) => return, + }; + let queue = queue.clone(); + let captured = cap2.clone(); + tokio::spawn(async move { + let mut buf = Vec::new(); + let mut tmp = [0u8; 8192]; + while !buf.windows(4).any(|w| w == b"\r\n\r\n") { + match sock.read(&mut tmp).await { + Ok(0) | Err(_) => return, + Ok(n) => buf.extend_from_slice(&tmp[..n]), + } + if buf.len() > 4_000_000 { + return; + } + } + let header_end = buf.windows(4).position(|w| w == b"\r\n\r\n").unwrap() + 4; + let headers = &buf[..header_end]; + let mut body_len = 0usize; + for line in headers.split(|b| *b == b'\n') { + let line = std::str::from_utf8(line).unwrap_or(""); + if let Some(rest) = line.to_ascii_lowercase().strip_prefix("content-length:") { + body_len = rest.trim().trim_end_matches('\r').parse().unwrap_or(0); + } + } + while buf.len() < header_end + body_len { + match sock.read(&mut tmp).await { + Ok(0) | Err(_) => return, + Ok(n) => buf.extend_from_slice(&tmp[..n]), + } + } + if let Ok(req) = serde_json::from_slice::(&buf[header_end..]) { + captured.lock().await.push(req); + } + let body = queue + .lock() + .await + .pop_front() + .unwrap_or_else(|| json!({ "error": "no canned response" })); + let body_s = serde_json::to_string(&body).unwrap(); + let resp = format!( + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + body_s.len(), + body_s, + ); + let _ = sock.write_all(resp.as_bytes()).await; + let _ = sock.shutdown().await; + }); + } + }); + CapturingLlm { url, captured } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Subprocess harness +// ───────────────────────────────────────────────────────────────────────────── + +struct Harness { + child: tokio::process::Child, + stdin: tokio::process::ChildStdin, + stdout: BufReader, + stderr: Arc>, + next_id: i64, +} + +impl Harness { + async fn spawn_with_env(base_url: &str, extra: &[(&str, &str)]) -> Self { + let bin = env!("CARGO_BIN_EXE_buzz-agent"); + let mut cmd = tokio::process::Command::new(bin); + cmd.env("BUZZ_AGENT_PROVIDER", "openai") + .env("OPENAI_COMPAT_API_KEY", "test") + .env("OPENAI_COMPAT_MODEL", "fake-model") + .env("OPENAI_COMPAT_BASE_URL", base_url) + .env("BUZZ_AGENT_LLM_TIMEOUT_SECS", "5") + .env("BUZZ_AGENT_TOOL_TIMEOUT_SECS", "5") + .env("BUZZ_AGENT_MAX_ROUNDS", "8") + .env("BUZZ_AGENT_MCP_INIT_TIMEOUT_SECS", "2"); + for (k, v) in extra { + cmd.env(k, v); + } + cmd.stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true); + let mut child = cmd.spawn().expect("spawn buzz-agent"); + let stdin = child.stdin.take().unwrap(); + let stdout = BufReader::new(child.stdout.take().unwrap()); + let stderr = child.stderr.take().unwrap(); + let stderr_buf = Arc::new(StdMutex::new(String::new())); + let stderr_out = Arc::clone(&stderr_buf); + tokio::spawn(async move { + let mut reader = BufReader::new(stderr); + let mut line = String::new(); + loop { + line.clear(); + let n = match reader.read_line(&mut line).await { + Ok(n) => n, + Err(_) => break, + }; + if n == 0 { + break; + } + if let Ok(mut out) = stderr_out.lock() { + out.push_str(&line); + } + } + }); + Self { + child, + stdin, + stdout, + stderr: stderr_buf, + next_id: 1, + } + } + + async fn spawn(base_url: &str) -> Self { + Self::spawn_with_env(base_url, &[]).await + } + + async fn send(&mut self, method: &str, params: Value) -> i64 { + let id = self.next_id; + self.next_id += 1; + self.write(json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params })) + .await; + id + } + + async fn notify(&mut self, method: &str, params: Value) { + self.write(json!({ "jsonrpc": "2.0", "method": method, "params": params })) + .await; + } + + async fn write(&mut self, msg: Value) { + let mut s = serde_json::to_string(&msg).unwrap(); + s.push('\n'); + self.stdin.write_all(s.as_bytes()).await.unwrap(); + self.stdin.flush().await.unwrap(); + } + + async fn recv(&mut self) -> Value { + let mut line = String::new(); + let n = tokio::time::timeout(Duration::from_secs(15), self.stdout.read_line(&mut line)) + .await + .expect("recv timeout") + .expect("read line"); + assert!(n > 0, "agent EOF; stderr={}", self.stderr_text()); + serde_json::from_str(&line).expect("non-JSON line") + } + + async fn recv_until bool>(&mut self, mut pred: F) -> Value { + loop { + let v = self.recv().await; + if pred(&v) { + return v; + } + } + } + + async fn shutdown(mut self) { + drop(self.stdin); + let _ = tokio::time::timeout(Duration::from_secs(2), self.child.wait()).await; + let _ = self.child.start_kill(); + } + + fn stderr_text(&self) -> String { + self.stderr.lock().map(|s| s.clone()).unwrap_or_default() + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// LLM response builders +// ───────────────────────────────────────────────────────────────────────────── + +fn openai_text(content: &str) -> Value { + json!({ + "id": "cc-1", "object": "chat.completion", "model": "fake-model", + "choices": [{ + "index": 0, + "message": { "role": "assistant", "content": content }, + "finish_reason": "stop", + }], + }) +} + +/// One assistant turn issuing `calls`, each `(id, qualified_name, arguments)`. +fn openai_tool_calls(calls: &[(&str, &str, Value)]) -> Value { + let tool_calls: Vec = calls + .iter() + .map(|(id, name, args)| { + json!({ + "id": id, "type": "function", + "function": { "name": name, "arguments": args.to_string() }, + }) + }) + .collect(); + json!({ + "id": "cc-tc", "object": "chat.completion", "model": "fake-model", + "choices": [{ + "index": 0, + "message": { "role": "assistant", "content": null, "tool_calls": tool_calls }, + "finish_reason": "tool_calls", + }], + }) +} + +fn shell_call(id: &str) -> Value { + openai_tool_calls(&[(id, "fake__shell", json!({ "command": "ls" }))]) +} + +// ───────────────────────────────────────────────────────────────────────────── +// Permission-response builders + drivers +// ───────────────────────────────────────────────────────────────────────────── + +/// Answer a permission request by selecting the offered option whose +/// `kind == "allow_once"` — mirroring buzz-acp's answering side, never a +/// hardcoded `optionId`. A future option-id rename can't silently deny. +fn approve(req: &Value) -> Value { + let option_id = req["params"]["options"] + .as_array() + .and_then(|opts| opts.iter().find(|o| o["kind"] == "allow_once")) + .and_then(|o| o["optionId"].as_str()) + .expect("request must offer an allow_once option"); + resp_selected(&req["id"], option_id) +} + +/// Bare JSON-RPC response selecting `option_id`. +fn resp_selected(id: &Value, option_id: &str) -> Value { + json!({ + "jsonrpc": "2.0", + "id": id.clone(), + "result": { "outcome": { "outcome": "selected", "optionId": option_id } }, + }) +} + +/// Bare JSON-RPC response carrying an arbitrary `result` shape. +fn resp_result(id: &Value, result: Value) -> Value { + json!({ "jsonrpc": "2.0", "id": id.clone(), "result": result }) +} + +/// Bare JSON-RPC *error* response (id present, `error`, no `result`). +fn resp_error(id: &Value) -> Value { + json!({ + "jsonrpc": "2.0", + "id": id.clone(), + "error": { "code": -32601, "message": "method not found" }, + }) +} + +/// The tool title carried by a permission request, across both wire shapes. +fn perm_title(req: &Value) -> String { + let p = &req["params"]; + p["title"] + .as_str() + .or_else(|| p["toolCall"]["title"].as_str()) + .or_else(|| p["subject"]["toolCall"]["title"].as_str()) + .unwrap_or("") + .to_owned() +} + +/// Drive one prompt to completion. For each `session/request_permission`, +/// `decide(&req)` returns `Some(response)` to answer or `None` to leave it +/// unanswered. Returns the final prompt response and every request seen. +async fn drive( + h: &mut Harness, + sid: &str, + prompt: &str, + mut decide: impl FnMut(&Value) -> Option, +) -> (Value, Vec) { + let p = h + .send( + "session/prompt", + json!({ "sessionId": sid, "prompt": [{ "type": "text", "text": prompt }] }), + ) + .await; + let mut requests: Vec = Vec::new(); + loop { + let v = h.recv().await; + if v.get("method") == Some(&json!("session/request_permission")) { + requests.push(v.clone()); + if let Some(resp) = decide(&v) { + h.write(resp).await; + } + continue; + } + if v["id"] == json!(p) { + return (v, requests); + } + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// Session init + call-log helpers +// ───────────────────────────────────────────────────────────────────────────── + +fn call_log_path(tag: &str) -> String { + let p = std::env::temp_dir().join(format!( + "buzz_perm_calllog_{tag}_{}_{:x}.log", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_nanos() + )); + let s = p.to_string_lossy().to_string(); + let _ = std::fs::remove_file(&s); + s +} + +/// Invoked tool names recorded by fake_mcp (bare names, one per `tools/call`). +/// A missing file means zero invocations. +fn call_log_lines(path: &str) -> Vec { + std::fs::read_to_string(path) + .unwrap_or_default() + .lines() + .map(str::trim) + .filter(|l| !l.is_empty()) + .map(str::to_owned) + .collect() +} + +/// Initialize + create a session with a single fake MCP server named `fake`, +/// negotiating `protocol_version` and passing `mcp_env` to the server. `cwd` +/// controls skill discovery (`.agents/skills`). +async fn init( + h: &mut Harness, + protocol_version: u32, + cwd: &str, + mcp_env: &[(&str, &str)], +) -> String { + let fake_mcp = env!("CARGO_BIN_EXE_fake-mcp"); + let env: Vec = mcp_env + .iter() + .map(|(k, v)| json!({ "name": k, "value": v })) + .collect(); + h.send( + "initialize", + json!({ "protocolVersion": protocol_version, "clientCapabilities": {} }), + ) + .await; + let _ = h.recv().await; + let servers = if mcp_env.is_empty() { + json!([]) + } else { + json!([{ "name": "fake", "command": fake_mcp, "args": [], "env": env }]) + }; + h.send("session/new", json!({ "cwd": cwd, "mcpServers": servers })) + .await; + let r = h + .recv_until(|v| v.get("result").is_some() || v.get("error").is_some()) + .await; + r["result"]["sessionId"] + .as_str() + .unwrap_or_else(|| panic!("session/new failed: {r}, stderr={}", h.stderr_text())) + .to_owned() +} + +fn stop_reason(resp: &Value) -> String { + resp["result"]["stopReason"] + .as_str() + .unwrap_or("") + .to_owned() +} + +// ═════════════════════════════════════════════════════════════════════════════ +// The authorization boundary +// ═════════════════════════════════════════════════════════════════════════════ + +/// Exact allow reaches the tool exactly once — and never *before* approval. +/// The pre-approval check proves the gate precedes the MCP call, not just that +/// the tally ends at one. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_allow_reaches_tool_exactly_once_and_not_before_approval() { + let log = call_log_path("allow"); + let llm = spawn_capturing_llm(vec![shell_call("tc1"), openai_text("done")]).await; + let mut h = Harness::spawn(&llm.url).await; + let sid = init( + &mut h, + 1, + "/tmp", + &[("FAKE_MCP_SHELL_TOOL", "1"), ("FAKE_MCP_CALL_LOG", &log)], + ) + .await; + + let log_for_check = log.clone(); + let (resp, requests) = drive(&mut h, &sid, "go", |req| { + // Before answering, the tool must not have run. + assert!( + call_log_lines(&log_for_check).is_empty(), + "tool invoked BEFORE approval" + ); + Some(approve(req)) + }) + .await; + + assert_eq!(stop_reason(&resp), "end_turn"); + assert_eq!(requests.len(), 1, "exactly one call → exactly one ask"); + assert_eq!( + call_log_lines(&log), + vec!["shell"], + "approved tool reached MCP exactly once" + ); + h.shutdown().await; +} + +/// Selecting the offered reject option never invokes the tool, and the model +/// sees a permission-denied tool error (the turn continues to `end_turn`). +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_reject_option_never_invokes_tool() { + let log = call_log_path("reject"); + let llm = spawn_capturing_llm(vec![shell_call("tc1"), openai_text("understood")]).await; + let mut h = Harness::spawn(&llm.url).await; + let sid = init( + &mut h, + 1, + "/tmp", + &[("FAKE_MCP_SHELL_TOOL", "1"), ("FAKE_MCP_CALL_LOG", &log)], + ) + .await; + + let (resp, requests) = drive(&mut h, &sid, "go", |req| { + Some(resp_selected(&req["id"], "reject_once")) + }) + .await; + + assert_eq!(stop_reason(&resp), "end_turn"); + assert_eq!(requests.len(), 1); + assert!( + call_log_lines(&log).is_empty(), + "rejected tool must never reach MCP" + ); + h.shutdown().await; +} + +/// Every non-authorizing response shape fails closed: the tool never runs. +/// One subprocess per shape keeps the failure attributable. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_adversarial_outcomes_never_invoke_tool() { + // (tag, response-for-request builder) — the full fail-closed matrix. + type Shape = (&'static str, fn(&Value) -> Value); + let shapes: Vec = vec![ + ("cancelled", |req| { + resp_result(&req["id"], json!({ "outcome": { "outcome": "cancelled" } })) + }), + ("jsonrpc_error", |req| resp_error(&req["id"])), + ("missing_outcome", |req| resp_result(&req["id"], json!({}))), + ("empty_outcome", |req| { + resp_result(&req["id"], json!({ "outcome": {} })) + }), + ("selected_no_option", |req| { + resp_result(&req["id"], json!({ "outcome": { "outcome": "selected" } })) + }), + ("unknown_outcome", |req| { + resp_result( + &req["id"], + json!({ "outcome": { "outcome": "banana", "optionId": "allow_once" } }), + ) + }), + ("wrong_option_id", |req| { + resp_selected(&req["id"], "not_an_offered_option") + }), + ]; + + for (tag, build) in shapes { + let log = call_log_path(tag); + let llm = spawn_capturing_llm(vec![shell_call("tc1"), openai_text("ok")]).await; + let mut h = Harness::spawn(&llm.url).await; + let sid = init( + &mut h, + 1, + "/tmp", + &[("FAKE_MCP_SHELL_TOOL", "1"), ("FAKE_MCP_CALL_LOG", &log)], + ) + .await; + + let (resp, requests) = drive(&mut h, &sid, "go", |req| Some(build(req))).await; + + assert_eq!( + stop_reason(&resp), + "end_turn", + "shape {tag}: turn should continue" + ); + assert_eq!(requests.len(), 1, "shape {tag}: exactly one ask"); + assert!( + call_log_lines(&log).is_empty(), + "shape {tag}: non-authorizing outcome must never reach MCP" + ); + h.shutdown().await; + } +} + +/// A stale/unknown/foreign response id is ignored and does not unblock the live +/// waiter; the correct id then authorizes exactly once. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_stale_id_ignored_then_real_id_authorizes() { + let log = call_log_path("staleid"); + let llm = spawn_capturing_llm(vec![shell_call("tc1"), openai_text("done")]).await; + let mut h = Harness::spawn(&llm.url).await; + let sid = init( + &mut h, + 1, + "/tmp", + &[("FAKE_MCP_SHELL_TOOL", "1"), ("FAKE_MCP_CALL_LOG", &log)], + ) + .await; + + let p = h + .send( + "session/prompt", + json!({ "sessionId": sid, "prompt": [{ "type": "text", "text": "go" }] }), + ) + .await; + + // Wait for the ask. + let req = h + .recv_until(|v| v.get("method") == Some(&json!("session/request_permission"))) + .await; + + // Feed several ignorable responses first: a minted-but-never-issued id, a + // foreign numeric id, and a null id. None may unblock the real waiter. + h.write(resp_selected(&json!("perm-9999"), "allow_once")) + .await; + h.write(resp_selected(&json!(7), "allow_once")).await; + h.write(resp_selected(&Value::Null, "allow_once")).await; + // Give the agent a beat to (wrongly) act on any of them. + tokio::time::sleep(Duration::from_millis(150)).await; + assert!( + call_log_lines(&log).is_empty(), + "stale/foreign ids must not authorize the pending call" + ); + + // The real id resolves it exactly once. + h.write(approve(&req)).await; + let resp = h.recv_until(|v| v["id"] == json!(p)).await; + assert_eq!(stop_reason(&resp), "end_turn"); + assert_eq!( + call_log_lines(&log), + vec!["shell"], + "only the correct id authorizes, exactly once" + ); + h.shutdown().await; +} + +/// `session/cancel` while a permission ask is outstanding terminates the turn +/// promptly, executes nothing, and needs no `cancelled` permission response — +/// a Buzz client always answers, but a non-Buzz ACP client may violate the spec +/// by staying silent, and cancellation must not depend on that answer. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_cancel_while_waiting_executes_nothing_without_client_answer() { + let log = call_log_path("cancelwait"); + let llm = spawn_capturing_llm(vec![shell_call("tc1"), openai_text("done")]).await; + let mut h = Harness::spawn(&llm.url).await; + let sid = init( + &mut h, + 1, + "/tmp", + &[("FAKE_MCP_SHELL_TOOL", "1"), ("FAKE_MCP_CALL_LOG", &log)], + ) + .await; + + let p = h + .send( + "session/prompt", + json!({ "sessionId": sid, "prompt": [{ "type": "text", "text": "go" }] }), + ) + .await; + + // Wait for the ask, then cancel WITHOUT ever answering the permission. + let _req = h + .recv_until(|v| v.get("method") == Some(&json!("session/request_permission"))) + .await; + h.notify("session/cancel", json!({ "sessionId": sid })) + .await; + + let resp = h.recv_until(|v| v["id"] == json!(p)).await; + assert_eq!( + stop_reason(&resp), + "cancelled", + "cancel resolves the turn without a client permission answer" + ); + assert!( + call_log_lines(&log).is_empty(), + "a cancelled ask must never reach the tool" + ); + h.shutdown().await; +} + +/// Two parallel calls each get their own ask (distinct ids); crossed decisions — +/// deny the first-asked, allow the second-asked — authorize only the allowed +/// call. Serial admission (`max_parallel_tools=1`) serializes the asks: the +/// second ask fires only after the first resolves. Which tool is admitted first +/// is a tokio scheduling detail, so the test denies whichever is asked first and +/// proves only the allowed (second) call reached MCP — order-agnostic. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_crossed_parallel_decisions_authorize_only_matching_call() { + let log = call_log_path("crossed"); + // Two distinct registered tools so the call log distinguishes them by name. + let llm = spawn_capturing_llm(vec![ + openai_tool_calls(&[ + ("tc-alpha", "fake__alpha", json!({})), + ("tc-bravo", "fake__bravo", json!({})), + ]), + openai_text("done"), + ]) + .await; + let mut h = Harness::spawn_with_env(&llm.url, &[("BUZZ_AGENT_MAX_PARALLEL_TOOLS", "1")]).await; + let sid = init( + &mut h, + 1, + "/tmp", + &[ + ("FAKE_MCP_NAMED_TOOLS", "alpha,bravo"), + ("FAKE_MCP_CALL_LOG", &log), + ], + ) + .await; + + // Deny whichever tool is asked first; allow the second. Ids are distinct. + let mut seen_ids: Vec = Vec::new(); + let mut allowed_title = String::new(); + let (resp, requests) = drive(&mut h, &sid, "go", |req| { + assert!( + !seen_ids.contains(&req["id"]), + "each parallel call must get a distinct request id" + ); + seen_ids.push(req["id"].clone()); + if seen_ids.len() == 1 { + Some(resp_selected(&req["id"], "reject_once")) // deny the first-asked + } else { + allowed_title = perm_title(req); + Some(approve(req)) // allow the second-asked + } + }) + .await; + + assert_eq!(stop_reason(&resp), "end_turn"); + assert_eq!(requests.len(), 2, "one ask per parallel call"); + // The call log records bare tool names; the title carries the qualified + // `__`. Only the allowed (second-asked) call reached MCP. + let allowed_bare = allowed_title + .strip_prefix("fake__") + .expect("qualified title") + .to_owned(); + assert_eq!( + call_log_lines(&log), + vec![allowed_bare], + "only the allowed (second-asked) call reached MCP; the denied one did not" + ); + h.shutdown().await; +} + +/// The built-in `load_skill` tool is not an MCP call and is exempt from the +/// permission boundary: it executes with no `session/request_permission`. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_load_skill_emits_no_permission_request() { + let tmp = tempfile::TempDir::new().unwrap(); + let cwd = tmp.path(); + let skill_dir = cwd.join(".agents/skills/my-skill"); + std::fs::create_dir_all(&skill_dir).unwrap(); + std::fs::write( + skill_dir.join("SKILL.md"), + "---\nname: my-skill\ndescription: A skill\n---\nSKILL_BODY_77\n", + ) + .unwrap(); + + let llm = spawn_capturing_llm(vec![ + openai_tool_calls(&[("tc-ls", "load_skill", json!({ "name": "my-skill" }))]), + openai_text("done"), + ]) + .await; + let mut h = Harness::spawn(&llm.url).await; + // No MCP server: `load_skill` is a built-in, and skills come from `cwd`. + let sid = init(&mut h, 1, cwd.to_str().unwrap(), &[]).await; + + let (resp, requests) = drive(&mut h, &sid, "use my-skill", |_| { + panic!("load_skill must not trigger a permission request") + }) + .await; + + assert_eq!(stop_reason(&resp), "end_turn"); + assert!(requests.is_empty(), "built-in load_skill is exempt"); + h.shutdown().await; +} + +/// Lifecycle hooks (`_Stop`, `_PostCompact`) invoke MCP through `call_hooks`, +/// not the model-issued tool path, so they are exempt: the hook reaches the +/// server (call log records it) with no permission ask. Here the `_Stop` hook +/// objects once, forcing a hook invocation the test can observe. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_stop_hook_reaches_mcp_without_permission_ask() { + let log = call_log_path("stophook"); + // Text turn triggers the _Stop gate → hook objects once → agent loops → + // second text turn, hook silent → end_turn. No model-issued tool call. + let llm = spawn_capturing_llm(vec![ + openai_text("premature"), + openai_text("really done"), + openai_text("unexpected"), + ]) + .await; + let mut h = Harness::spawn_with_env( + &llm.url, + &[ + ("MCP_HOOK_SERVERS", "fake"), + ("BUZZ_AGENT_STOP_MAX_REJECTIONS", "10"), + ], + ) + .await; + let sid = init( + &mut h, + 1, + "/tmp", + &[ + ("FAKE_MCP_TOOL_COUNT", "1"), + ("FAKE_MCP_STOP_HOOK", "1"), + ("FAKE_MCP_STOP_TEXT", "you have open work"), + ("FAKE_MCP_STOP_COUNT", "1"), + ("FAKE_MCP_CALL_LOG", &log), + ], + ) + .await; + + let (resp, requests) = drive(&mut h, &sid, "go", |_| { + panic!("a lifecycle hook must not trigger a permission request") + }) + .await; + + assert_eq!(stop_reason(&resp), "end_turn"); + assert!(requests.is_empty(), "_Stop hook is exempt from the ask"); + assert!( + call_log_lines(&log).contains(&"_Stop".to_owned()), + "the _Stop hook still reached MCP without an ask; log={:?}", + call_log_lines(&log) + ); + h.shutdown().await; +} + +/// Under a v2-negotiated connection, the emitted `session/request_permission` +/// carries the v2 shape (`subject.toolCall`, top-level `title`), never the v1 +/// legacy top-level `toolCall`. Complements the pure v1/v2 builder unit tests +/// with an end-to-end proof that the negotiated version reaches the wire. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_v2_negotiation_emits_v2_request_shape() { + let log = call_log_path("v2shape"); + let llm = spawn_capturing_llm(vec![shell_call("tc1"), openai_text("done")]).await; + let mut h = Harness::spawn(&llm.url).await; + let sid = init( + &mut h, + 2, // negotiate v2 + "/tmp", + &[("FAKE_MCP_SHELL_TOOL", "1"), ("FAKE_MCP_CALL_LOG", &log)], + ) + .await; + + let (resp, requests) = drive(&mut h, &sid, "go", |req| { + let params = &req["params"]; + // v2: tool context under `subject.toolCall`, with top-level `title`. + assert_eq!(params["subject"]["type"], "tool_call", "v2 uses subject"); + assert_eq!(params["subject"]["toolCall"]["title"], "fake__shell"); + assert_eq!(params["title"], "fake__shell", "v2 has top-level title"); + assert!( + params.get("toolCall").is_none(), + "v2 must not carry the v1 top-level toolCall" + ); + Some(approve(req)) + }) + .await; + + assert_eq!(stop_reason(&resp), "end_turn"); + assert_eq!(requests.len(), 1); + assert_eq!(call_log_lines(&log), vec!["shell"]); + h.shutdown().await; +} diff --git a/crates/buzz-agent/tests/regressions.rs b/crates/buzz-agent/tests/regressions.rs index 6a4f347f6bb..513b9887a6a 100644 --- a/crates/buzz-agent/tests/regressions.rs +++ b/crates/buzz-agent/tests/regressions.rs @@ -198,6 +198,24 @@ impl Harness { } } + /// Like `recv_until`, but auto-approves any `session/request_permission` + /// seen while waiting. Tests here exercise tool execution, not the + /// permission boundary (that lives in `permission_boundary.rs`), so a + /// model-issued tool call must be approved to reach the server. + async fn recv_until_approving bool>(&mut self, mut pred: F) -> Value { + loop { + let v = self.recv().await; + if v.get("method") == Some(&json!("session/request_permission")) { + let resp = approve_permission(&v); + self.write(resp).await; + continue; + } + if pred(&v) { + return v; + } + } + } + async fn shutdown(mut self) { drop(self.stdin); let _ = tokio::time::timeout(Duration::from_secs(2), self.child.wait()).await; @@ -270,6 +288,23 @@ fn openai_tool_call(id: &str, name: &str, args: Value) -> Value { }) } +/// Select the offered option whose `kind == "allow_once"` and return the +/// `session/request_permission` response. Mirrors buzz-acp's answering side, +/// which selects by `kind`, never by a hardcoded `optionId`. Centralizing this +/// means a future option-id rename can't silently turn allow into a denial. +fn approve_permission(request: &Value) -> Value { + let option_id = request["params"]["options"] + .as_array() + .and_then(|opts| opts.iter().find(|o| o["kind"] == "allow_once")) + .and_then(|o| o["optionId"].as_str()) + .expect("request must offer an allow_once option"); + json!({ + "jsonrpc": "2.0", + "id": request["id"], + "result": { "outcome": { "outcome": "selected", "optionId": option_id } }, + }) +} + async fn init_session(h: &mut Harness, mcp_servers: Value) -> String { h.send( "initialize", @@ -676,13 +711,7 @@ async fn per_turn_tool_call_cap_enforced() { loop { let v = h.recv().await; if v.get("method") == Some(&json!("session/request_permission")) { - let id = v["id"].clone(); - h.write(json!({ - "jsonrpc": "2.0", - "id": id, - "result": { "outcome": { "outcome": "selected", "optionId": "allow" } }, - })) - .await; + h.write(approve_permission(&v)).await; continue; } if v.get("method") == Some(&json!("session/update")) @@ -858,7 +887,7 @@ async fn hook_stop_blocks_premature_end() { json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}), ) .await; - let r = h.recv_until(|v| v["id"] == json!(p)).await; + let r = h.recv_until_approving(|v| v["id"] == json!(p)).await; assert!(r.get("result").is_some(), "errored: {r}"); assert_eq!(r["result"]["stopReason"], "end_turn"); @@ -936,7 +965,7 @@ async fn hook_stop_budget_exhausted() { json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}), ) .await; - let r = h.recv_until(|v| v["id"] == json!(p)).await; + let r = h.recv_until_approving(|v| v["id"] == json!(p)).await; assert!(r.get("result").is_some(), "errored: {r}"); assert_eq!(r["result"]["stopReason"], "end_turn"); @@ -1517,7 +1546,7 @@ async fn stale_usage_plus_history_growth_triggers_handoff() { json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}), ) .await; - let _ = h.recv_until(|v| v["id"] == json!(p)).await; + let _ = h.recv_until_approving(|v| v["id"] == json!(p)).await; // req1 (tool_call) + summarize (handoff) + req2 (done) = 3. Without the // growth estimate we'd see only 2 (stale 8500 < 9000, no handoff). let captured = llm.captured.lock().await.len(); @@ -1788,7 +1817,7 @@ async fn cancel_sends_notifications_cancelled_to_any_mcp_server() { .await; // Wait for tool call to be in-progress. - h.recv_until(|v| { + h.recv_until_approving(|v| { v.get("params") .and_then(|p| p.get("update")) .and_then(|u| u.get("status")) @@ -1903,13 +1932,7 @@ async fn prompt_to_completion(h: &mut Harness, sid: &str) -> Value { loop { let v = h.recv().await; if v.get("method") == Some(&json!("session/request_permission")) { - let id = v["id"].clone(); - h.write(json!({ - "jsonrpc": "2.0", - "id": id, - "result": { "outcome": { "outcome": "selected", "optionId": "allow" } }, - })) - .await; + h.write(approve_permission(&v)).await; continue; } if v["id"] == json!(p) { @@ -2664,7 +2687,9 @@ async fn max_tokens_recovery_can_proceed_to_tool_call() { json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}), ) .await; - let reply = h.recv_until(|v| v["id"] == json!(prompt_id)).await; + let reply = h + .recv_until_approving(|v| v["id"] == json!(prompt_id)) + .await; assert_eq!(reply["result"]["stopReason"], "end_turn", "{reply}"); let requests = llm.captured.lock().await; assert_eq!(requests.len(), 3); @@ -3642,13 +3667,7 @@ async fn handoff_cap_binds_within_a_single_turn() { } if v.get("method") == Some(&json!("session/request_permission")) { - let id = v["id"].clone(); - h.write(json!({ - "jsonrpc": "2.0", - "id": id, - "result": { "outcome": { "outcome": "selected", "optionId": "allow" } }, - })) - .await; + h.write(approve_permission(&v)).await; continue; } if v["id"] == json!(p2) { @@ -3795,13 +3814,7 @@ async fn failed_summarize_burns_handoff_attempt_budget() { loop { let v = h.recv().await; if v.get("method") == Some(&json!("session/request_permission")) { - let id = v["id"].clone(); - h.write(json!({ - "jsonrpc": "2.0", - "id": id, - "result": { "outcome": { "outcome": "selected", "optionId": "allow" } }, - })) - .await; + h.write(approve_permission(&v)).await; continue; } if v["id"] == json!(p2) { From d8cf983c21284d23eca4e5d67c4a99df802e8fb0 Mon Sep 17 00:00:00 2001 From: Hayt <41ea58f1e64c243627e8acde7c89be667052ee6e17d8f021c1195be4324ebf04@buzz.block.builderlab.xyz> Date: Wed, 12 Aug 2026 20:58:28 -0400 Subject: [PATCH 2/3] fix(buzz-agent): make undeliverable permission asks terminal, prove claim-before-wake ordering, reject noncanonical ids MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address three review findings on the permission broker's ask surface. An undeliverable ask is now terminal. wire::send_checked surfaces the mpsc send failure that occurs exactly when the writer task has exited on a closed/broken stdout; request_permission fails closed immediately on that error (dropping the lease removes the entry and releases the permit synchronously) instead of leaving a resident waiter to time out. async_main now selects on both the reader and the writer JoinHandle: writer death cancels every session and closes the reader lifecycle rather than reading input while asks wait out their deadline for a reply that can never be written. Claim-before-wake ordering is now mutation-sensitive. Production enforces it structurally — the oneshot sender is consumed by remove, so a send-before-remove mutant cannot compile without swapping the channel type. A test-only wake observer fires synchronously in the waiter's response arm and asserts the entry is already absent from pending at the wake; a faithful wake-before-claim mutant makes it observe false and the test goes red while the behavioral delivery tests stay green. parse_id now requires an exact canonical round-trip, so noncanonical aliases (perm-01, perm-+0, perm-00) that u64::parse would accept are rejected as foreign ids rather than resolving live asks. Co-authored-by: Will Pfleger Signed-off-by: Will Pfleger --- crates/buzz-agent/src/lib.rs | 41 ++++-- crates/buzz-agent/src/permission.rs | 190 +++++++++++++++++++++++++++- crates/buzz-agent/src/wire.rs | 34 ++++- 3 files changed, 246 insertions(+), 19 deletions(-) diff --git a/crates/buzz-agent/src/lib.rs b/crates/buzz-agent/src/lib.rs index f94d60a2aa2..a879bd621b5 100644 --- a/crates/buzz-agent/src/lib.rs +++ b/crates/buzz-agent/src/lib.rs @@ -206,21 +206,40 @@ async fn async_main() { models_cache: tokio::sync::OnceCell::new(), }); let (wire_tx, wire_rx) = mpsc::channel::(64); - let writer = tokio::spawn(wire::writer_task(wire_rx)); - if let Err(e) = read_loop( - BufReader::new(tokio::io::stdin()), - app.clone(), - wire_tx, - max_line, - ) - .await - { - tracing::error!("io: reader: {e}"); + let mut writer = tokio::spawn(wire::writer_task(wire_rx)); + // Whichever ends first drives shutdown. The reader ending is the normal + // path (stdin EOF/error). The writer ending while the reader still runs + // means stdout is closed/broken: no reply can ever be written, so we must + // stop reading and cancel every session rather than leave the process + // reading input while outstanding permission asks wait out their full + // deadline for a response that can never arrive. + tokio::select! { + r = read_loop( + BufReader::new(tokio::io::stdin()), + app.clone(), + wire_tx, + max_line, + ) => { + if let Err(e) = r { + tracing::error!("io: reader: {e}"); + } + cancel_all_sessions(&app).await; + let _ = writer.await; + } + _ = &mut writer => { + tracing::error!("io: writer exited (stdout closed); shutting down connection"); + cancel_all_sessions(&app).await; + } } +} + +/// Signal every live session to cancel. Run on connection teardown so in-flight +/// prompts — including any waiting on a `session/request_permission` response — +/// resolve promptly instead of waiting out their deadline. +async fn cancel_all_sessions(app: &Arc) { for session in app.sessions.lock().await.values() { let _ = session.cancel_tx.send(true); } - let _ = writer.await; } async fn read_loop( diff --git a/crates/buzz-agent/src/permission.rs b/crates/buzz-agent/src/permission.rs index 0baef5137c1..09d7421cfc2 100644 --- a/crates/buzz-agent/src/permission.rs +++ b/crates/buzz-agent/src/permission.rs @@ -24,6 +24,11 @@ //! a later lease `Drop` is a harmless no-op. //! - **Unknown/late ids ignored.** A response whose id is not a live entry is //! logged and dropped. +//! - **Undeliverable asks are terminal.** If the output wire is closed when the +//! request is enqueued, [`PermissionBroker::request_permission`] fails closed +//! immediately (dropping the lease removes the entry and releases the permit) +//! rather than leaving a resident waiter to time out — a closed wire can never +//! carry the reply. //! - **Single absolute deadline.** Admission wait and response wait share one //! absolute deadline computed at gate entry, so a saturated call cannot live //! for two full timeout windows. @@ -51,6 +56,12 @@ pub const PERMISSION_DENIED_MSG: &str = "permission denied: the tool call was no pub const PERMISSION_TIMEOUT_MSG: &str = "permission request timed out: the tool call was not authorized"; +/// Model-visible tool error when the permission request cannot be delivered +/// because the output wire is closed. Terminal and immediate — no waiter is +/// left resident, since a closed wire can never carry a reply. +pub const PERMISSION_WIRE_CLOSED_MSG: &str = + "permission request undeliverable: the tool call was not authorized"; + /// Outcome of asking the client to authorize one tool call. #[derive(Debug, PartialEq, Eq)] pub enum PermissionDecision { @@ -77,8 +88,20 @@ pub struct PermissionBroker { next_id: AtomicU64, /// Absolute deadline budget shared by admission + response wait. timeout: Duration, + /// Test-only: invoked by the waiter the instant it observes a delivered + /// response, with `true` iff the entry was already claimed (removed) from + /// `pending` before the wake. Makes claim-before-wake ordering + /// mutation-sensitive — a wake-before-claim mutant reports `false`, which a + /// purely behavioral test cannot detect (the waiter reads the oneshot once + /// either way). + #[cfg(test)] + wake_observer: Mutex>, } +/// Test-only wake-boundary observer; see [`PermissionBroker::wake_observer`]. +#[cfg(test)] +type WakeObserver = Arc; + impl PermissionBroker { /// `max_pending` is validated `>= 1` by config; `timeout` is injectable so /// broker unit tests exercise the timeout/abort paths without a 330s wait. @@ -88,6 +111,31 @@ impl PermissionBroker { pending: Arc::new(Mutex::new(HashMap::new())), next_id: AtomicU64::new(0), timeout, + #[cfg(test)] + wake_observer: Mutex::new(None), + } + } + + /// Test-only: register a callback the waiter fires the instant it observes a + /// delivered response, with `true` iff the correlation entry was already + /// claimed (removed) before the wake. Used to prove claim-before-wake + /// ordering in a way a wake-before-claim mutant cannot satisfy. + #[cfg(test)] + pub fn set_wake_observer(&self, observer: WakeObserver) { + *self.wake_observer.lock().unwrap() = Some(observer); + } + + /// Test-only: fire the wake observer (if any) with the claimed-before-wake + /// status of `id`. Called synchronously by the waiter the moment it receives + /// its response, so the observed `pending` state is exactly the state at the + /// wake — deterministic in production (removal happens-before the send) and + /// violated by a wake-before-claim mutant. + #[cfg(test)] + fn observe_wake(&self, id: u64) { + let claimed = !self.pending.lock().unwrap().contains_key(&id); + let observer = self.wake_observer.lock().unwrap().clone(); + if let Some(observer) = observer { + observer(claimed); } } @@ -177,21 +225,39 @@ impl PermissionBroker { &call.name, &call.arguments, ); - wire::send( + if wire::send_checked( wire, wire::request_permission(lease.id_value.clone(), params), ) - .await; + .await + .is_err() + { + // The output wire is closed: this ask will never be written and no + // reply can ever arrive. Fail closed now — dropping `lease` here + // removes the entry and releases the permit synchronously — instead + // of leaving the entry resident until the deadline expires. + return PermissionDecision::Denied(PERMISSION_WIRE_CLOSED_MSG); + } // ── Response wait ────────────────────────────────────────────────── if *cancel.borrow() { return PermissionDecision::Cancelled; } + #[cfg(test)] + let id = lease.id; tokio::select! { biased; _ = cancel.changed() => PermissionDecision::Cancelled, r = &mut lease.rx => match r { - Ok(result) => evaluate(&result), + Ok(result) => { + // The waiter observes delivery here. At this instant the + // entry must already be claimed (removed) — delivery removes + // before it sends. The observer is test-only and a no-op in + // production. + #[cfg(test)] + self.observe_wake(id); + evaluate(&result) + } // Sender dropped without sending — should not happen (delivery // always sends before drop); fail closed. Err(_) => PermissionDecision::Denied(PERMISSION_DENIED_MSG), @@ -252,10 +318,14 @@ fn evaluate(result: &Value) -> PermissionDecision { } /// Recover the correlation key from an outbound request id echoed by the -/// client. Only ids we minted (`perm-`) are ours; anything else is a -/// foreign/stale id and is ignored. +/// client. Only ids we minted (`perm-`, canonical decimal) are ours; a +/// noncanonical alias (`perm-01`, `perm-+0`, `perm-00`) or any other string is +/// a foreign/stale id and is ignored. Requiring an exact round-trip means only +/// the string the broker actually minted correlates — no alias is ever live. fn parse_id(id: &Value) -> Option { - id.as_str()?.strip_prefix("perm-")?.parse().ok() + let s = id.as_str()?; + let n: u64 = s.strip_prefix("perm-")?.parse().ok()?; + (format!("perm-{n}") == s).then_some(n) } #[cfg(test)] @@ -337,6 +407,31 @@ mod tests { assert_eq!(parse_id(&Value::Null), None); } + /// Noncanonical strings that `u64::parse` would otherwise accept as aliases + /// of a minted id must NOT correlate. Only the exact string the broker + /// minted (`format!("perm-{n}")`) is live; leading zeros, a sign, or + /// whitespace make the id foreign and it is ignored. Without the exact + /// round-trip check these would resolve live asks under ids the broker + /// never issued. + #[test] + fn test_parse_id_rejects_noncanonical_aliases() { + for alias in [ + "perm-00", // extra leading zero + "perm-01", // leading zero + "perm-+0", // explicit sign + "perm-0x1", // hex + "perm- 1", // leading space + "perm-1 ", // trailing space + "perm-1_000", // digit separator + ] { + assert_eq!( + parse_id(&json!(alias)), + None, + "alias must be foreign: {alias}" + ); + } + } + // ── Delivery: exact allow / deny ────────────────────────────────────────── #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -439,7 +534,88 @@ mod tests { assert_eq!(broker.available_permits(), 4, "timeout releases the slot"); } - // ── Cancellation while waiting ──────────────────────────────────────────── + // ── Undeliverable ask (closed wire) is terminal ─────────────────────────── + + /// When the output wire is closed, the ask can never be written and no + /// reply can ever arrive. `request_permission` must fail closed + /// *immediately* — denying with the wire-closed reason and leaving zero + /// pending entries and zero held permits — rather than registering an entry + /// that waits out the full deadline. Uses a LONG timeout so a wrong + /// implementation that waits the deadline would visibly hang the test far + /// past its own assertions. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_closed_wire_denies_immediately_without_leaking_state() { + let broker = Arc::new(PermissionBroker::new(4, LONG)); + // Drop the receiver so every send fails: the writer is gone. + let (tx, rx) = mpsc::channel(8); + drop(rx); + let (_cancel_tx, mut cancel_rx) = watch::channel(false); + let call = tool_call(); + + // Bound the whole call: correct behavior returns at once; a regression + // that waits the deadline blows this timeout instead of hanging LONG. + let decision = tokio::time::timeout( + Duration::from_secs(2), + broker.request_permission(&tx, 2, "ses_a", &call, &mut cancel_rx), + ) + .await + .expect("closed wire must deny immediately, not wait the deadline"); + + assert_eq!( + decision, + PermissionDecision::Denied(PERMISSION_WIRE_CLOSED_MSG) + ); + assert_eq!( + broker.pending_count(), + 0, + "undeliverable ask leaves no resident entry" + ); + assert_eq!( + broker.available_permits(), + 4, + "undeliverable ask releases its admission slot" + ); + } + + // ── Claim-before-wake ordering (mutation-sensitive) ─────────────────────── + + /// The waiter must observe the correlation entry already *claimed* (removed + /// from `pending`) at the instant it wakes with the delivered response — + /// `deliver` removes before it sends. The wake observer fires synchronously + /// inside the waiter's response arm, so it captures the exact `pending` + /// state at the wake. A wake-before-claim mutant (send first, remove after) + /// makes the observed state `false` and fails this assertion; the behavioral + /// delivery tests cannot detect that mutant because the waiter reads the + /// oneshot exactly once either way. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_waiter_observes_entry_claimed_before_wake() { + let broker = Arc::new(PermissionBroker::new(4, LONG)); + let claimed_at_wake = Arc::new(Mutex::new(None::)); + let sink = Arc::clone(&claimed_at_wake); + broker.set_wake_observer(Arc::new(move |claimed| { + *sink.lock().unwrap() = Some(claimed); + })); + + let (tx, mut rx) = mpsc::channel(8); + let (_cancel_tx, mut cancel_rx) = watch::channel(false); + let b = Arc::clone(&broker); + let call = tool_call(); + let task = tokio::spawn(async move { + b.request_permission(&tx, 2, "ses_a", &call, &mut cancel_rx) + .await + }); + + let id = next_request_id(&mut rx).await; + broker.deliver(&id, selected(ALLOW_OPTION_ID)); + assert_eq!(task.await.unwrap(), PermissionDecision::Allowed); + assert_eq!( + *claimed_at_wake.lock().unwrap(), + Some(true), + "entry must be claimed (removed) before the waiter is woken" + ); + } + + // ── Cancellation while waiting ─────────────────────────────────────────── #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn test_cancel_while_waiting_returns_cancelled_and_removes_state() { diff --git a/crates/buzz-agent/src/wire.rs b/crates/buzz-agent/src/wire.rs index 2015e22378c..2d1f0d6c575 100644 --- a/crates/buzz-agent/src/wire.rs +++ b/crates/buzz-agent/src/wire.rs @@ -353,7 +353,18 @@ pub fn session_update_with_goose_meta(sid: &str, update: Value, goose_meta: Valu } pub async fn send(wire: &WireSender, msg: Value) { - let _ = wire.send(WireMsg::Notify(msg)).await; + let _ = send_checked(wire, msg).await; +} + +/// Enqueue a frame, reporting whether the writer accepted it. Unlike mpsc's +/// non-blocking `try_send`, this awaits channel capacity; it fails only when +/// the writer task has dropped its receiver, which happens exactly when the +/// writer has exited because stdout is closed/broken. A frame that fails here +/// will never be written, so callers that correlate a response — the +/// permission broker — must fail closed immediately rather than wait out a +/// deadline for a reply that can never arrive. +pub async fn send_checked(wire: &WireSender, msg: Value) -> Result<(), ()> { + wire.send(WireMsg::Notify(msg)).await.map_err(|_| ()) } pub async fn read_bounded_line( @@ -724,4 +735,25 @@ mod tests { other => panic!("expected Response, got {other:?}"), } } + + // ── send_checked: observable wire closure ──────────────────────────────── + + /// `send_checked` reports `Ok` while the writer's receiver is alive and + /// `Err` once it is gone (writer task exited on closed/broken stdout). This + /// is the contract the permission broker relies on to fail an undeliverable + /// ask closed immediately instead of waiting out its deadline for a reply + /// that can never be written. + #[tokio::test] + async fn send_checked_reports_closure_when_writer_gone() { + let (tx, rx) = mpsc::channel::(4); + assert!( + send_checked(&tx, json!({ "ok": 1 })).await.is_ok(), + "send succeeds while the writer receiver is alive" + ); + drop(rx); // writer exited → receiver dropped + assert!( + send_checked(&tx, json!({ "ok": 2 })).await.is_err(), + "send reports failure once the writer is gone" + ); + } } From a3256c068b17b21c9c0ea3e0129f952665363855 Mon Sep 17 00:00:00 2001 From: Hayt <41ea58f1e64c243627e8acde7c89be667052ee6e17d8f021c1195be4324ebf04@buzz.block.builderlab.xyz> Date: Wed, 12 Aug 2026 22:00:37 -0400 Subject: [PATCH 3/3] fix(buzz-agent): make writer flush fatal and bound permission enqueue by the deadline Close two wire-closure terminality gaps in the permission broker's ask surface. Writer flush failure is now connection-fatal. A blocking stdout can report Ok from write_all when it only schedules the underlying write and surface the real error at flush; the writer previously discarded that flush error and kept waiting on its receiver, so a genuinely dead stdout left the writer alive, the async_main writer-death arm never fired, and an accepted ask stayed resident until its deadline. write_frames (extracted, generic over its AsyncWrite sink for tests) now returns on either write_all or flush error, dropping the receiver so the connection supervisor cancels every session. Permission enqueue is now the third phase governed by the single absolute deadline. request_permission previously bare-awaited send_checked after registering the entry; a full-but-live channel makes that await wait for capacity, racing neither cancellation nor the deadline, so a stalled writer could hold the ask and its global permit past the advertised deadline and session/cancel could not resolve it. The send now runs in the same biased select as admission and response: cancel wins Cancelled, the shared deadline wins the timeout deny, send error stays the wire-closed deny. Co-authored-by: Will Pfleger Signed-off-by: Will Pfleger --- crates/buzz-agent/src/permission.rs | 194 ++++++++++++++++++++++++++-- crates/buzz-agent/src/wire.rs | 26 +++- 2 files changed, 203 insertions(+), 17 deletions(-) diff --git a/crates/buzz-agent/src/permission.rs b/crates/buzz-agent/src/permission.rs index 09d7421cfc2..1dde42a27c0 100644 --- a/crates/buzz-agent/src/permission.rs +++ b/crates/buzz-agent/src/permission.rs @@ -225,18 +225,26 @@ impl PermissionBroker { &call.name, &call.arguments, ); - if wire::send_checked( - wire, - wire::request_permission(lease.id_value.clone(), params), - ) - .await - .is_err() - { - // The output wire is closed: this ask will never be written and no - // reply can ever arrive. Fail closed now — dropping `lease` here - // removes the entry and releases the permit synchronously — instead - // of leaving the entry resident until the deadline expires. - return PermissionDecision::Denied(PERMISSION_WIRE_CLOSED_MSG); + // ── Send (deadline- and cancel-governed) ─────────────────────────── + // Enqueue is the third phase under the single absolute deadline. A + // full-but-live channel makes `send_checked` wait for capacity; racing + // it against cancel + the deadline means a stalled writer cannot hold + // the ask (and its global permit) past the advertised deadline, and + // `session/cancel` resolves it promptly. On send error the wire is + // closed: fail closed at once — dropping `lease` removes the entry and + // releases the permit synchronously. + let request = wire::request_permission(lease.id_value.clone(), params); + tokio::select! { + biased; + _ = cancel.changed() => return PermissionDecision::Cancelled, + _ = tokio::time::sleep_until(deadline) => { + return PermissionDecision::Denied(PERMISSION_TIMEOUT_MSG); + } + r = wire::send_checked(wire, request) => { + if r.is_err() { + return PermissionDecision::Denied(PERMISSION_WIRE_CLOSED_MSG); + } + } } // ── Response wait ────────────────────────────────────────────────── @@ -332,6 +340,10 @@ fn parse_id(id: &Value) -> Option { mod tests { use super::*; use serde_json::json; + use std::io; + use std::pin::Pin; + use std::task::{Context, Poll}; + use tokio::io::AsyncWrite; use tokio::sync::mpsc; const LONG: Duration = Duration::from_secs(30); @@ -362,6 +374,29 @@ mod tests { v["id"].clone() } + /// An `AsyncWrite` that accepts every write but fails on `flush`, modelling + /// Tokio's blocking stdout when the pipe has broken: `write_all` reports + /// `Ok` (the underlying blocking write is only scheduled) and the real + /// error surfaces at `flush`. Used to prove the writer treats flush failure + /// as connection-fatal. + struct FlushFailSink; + + impl AsyncWrite for FlushFailSink { + fn poll_write( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Poll::Ready(Ok(buf.len())) + } + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Err(io::Error::from(io::ErrorKind::BrokenPipe))) + } + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + // ── Authorization predicate (fail-closed) ──────────────────────────────── #[test] @@ -577,6 +612,141 @@ mod tests { ); } + // ── Writer flush failure is connection-fatal ────────────────────────────── + + /// A blocking stdout can report `Ok` from `write_all` (the underlying + /// blocking write is only scheduled) and surface the real error at `flush`. + /// The writer must therefore treat flush failure exactly like write failure + /// — return, dropping its receiver — so the connection supervisor observes + /// writer death and cancels every session (which resolves any waiting ask). + /// Modelled here: `write_frames` fed one frame over a sink that accepts the + /// write but fails flush must terminate promptly; a cancellation wired to + /// that termination (as `async_main`'s writer arm does) then unblocks a + /// registered permission waiter, leaving zero pending entries and all + /// permits free — not after the injected deadline. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_flush_failure_kills_writer_and_resolves_waiter() { + let broker = Arc::new(PermissionBroker::new(4, LONG)); + // A live output channel; its frames are drained by `write_frames` into + // the flush-failing sink, modelling the real writer over a broken pipe. + let (wire_tx, wire_rx) = mpsc::channel(8); + let (cancel_tx, mut cancel_rx) = watch::channel(false); + + // Supervisor: run the writer over the failing sink; when it returns + // (flush error → connection-fatal), propagate cancellation exactly like + // `async_main`'s writer-death arm. + let writer = tokio::spawn(async move { + wire::write_frames(wire_rx, FlushFailSink).await; + let _ = cancel_tx.send(true); + }); + + let b = Arc::clone(&broker); + let call = tool_call(); + let task = tokio::spawn(async move { + b.request_permission(&wire_tx, 2, "ses_a", &call, &mut cancel_rx) + .await + }); + + // The ask registers and enqueues its frame; the writer accepts the + // write, fails the flush, returns, and the supervisor cancels. The + // waiter must resolve via that cancellation, not the LONG deadline. + let decision = tokio::time::timeout(Duration::from_secs(2), task) + .await + .expect("flush failure must cancel the waiter, not wait the deadline") + .unwrap(); + + writer.await.unwrap(); + assert_eq!(decision, PermissionDecision::Cancelled); + assert_eq!( + broker.pending_count(), + 0, + "writer death resolves the waiter and leaves no resident entry" + ); + assert_eq!( + broker.available_permits(), + 4, + "writer death releases the held admission slot" + ); + } + + // ── Enqueue backpressure is deadline- and cancel-governed ───────────────── + + /// A full-but-live output channel makes `send_checked` wait for capacity. + /// The send phase races the single absolute deadline, so a stalled writer + /// cannot hold the ask (and its global permit) past the advertised + /// deadline: `request_permission` must return the timeout deny within the + /// deadline and leave zero pending entries and zero held permits. Uses a + /// SHORT deadline bounded by a longer outer timeout, so a regression that + /// waits forever on capacity blows the outer bound. This is a distinct seam + /// from the dropped-receiver test: here the receiver is alive but never + /// drains. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_full_channel_send_is_bounded_by_the_deadline() { + let broker = Arc::new(PermissionBroker::new(4, SHORT)); + // Capacity-1 channel, prefilled and never drained: the next send blocks + // on capacity while the receiver stays alive (writer present, stalled). + let (tx, _rx) = mpsc::channel(1); + tx.send(wire::WireMsg::Notify(json!({ "fill": 1 }))) + .await + .unwrap(); + let (_cancel_tx, mut cancel_rx) = watch::channel(false); + let call = tool_call(); + + let decision = tokio::time::timeout( + Duration::from_secs(2), + broker.request_permission(&tx, 2, "ses_a", &call, &mut cancel_rx), + ) + .await + .expect("a stalled writer must not hold the send past the deadline"); + + assert_eq!( + decision, + PermissionDecision::Denied(PERMISSION_TIMEOUT_MSG), + "a full-but-live channel resolves via the deadline, not the wire-closed path" + ); + assert_eq!( + broker.pending_count(), + 0, + "a timed-out enqueue leaves no resident entry" + ); + assert_eq!( + broker.available_permits(), + 4, + "a timed-out enqueue releases its admission slot" + ); + } + + /// Cancellation must also resolve an ask stuck enqueueing on a stalled + /// writer: with a full-but-live channel and a LONG deadline, a + /// `session/cancel` returns `Cancelled` promptly rather than waiting the + /// deadline out. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_cancel_resolves_a_blocked_enqueue() { + let broker = Arc::new(PermissionBroker::new(4, LONG)); + let (tx, _rx) = mpsc::channel(1); + tx.send(wire::WireMsg::Notify(json!({ "fill": 1 }))) + .await + .unwrap(); + let (cancel_tx, mut cancel_rx) = watch::channel(false); + let b = Arc::clone(&broker); + let call = tool_call(); + let task = tokio::spawn(async move { + b.request_permission(&tx, 2, "ses_a", &call, &mut cancel_rx) + .await + }); + + // Give the task time to admit + block on the full channel, then cancel. + tokio::time::sleep(Duration::from_millis(50)).await; + cancel_tx.send(true).unwrap(); + let decision = tokio::time::timeout(Duration::from_secs(2), task) + .await + .expect("cancel must resolve a blocked enqueue promptly") + .unwrap(); + assert_eq!(decision, PermissionDecision::Cancelled); + assert_eq!(broker.pending_count(), 0); + assert_eq!(broker.available_permits(), 4); + } + // ── Claim-before-wake ordering (mutation-sensitive) ─────────────────────── /// The waiter must observe the correlation entry already *claimed* (removed diff --git a/crates/buzz-agent/src/wire.rs b/crates/buzz-agent/src/wire.rs index 2d1f0d6c575..cb959ea5648 100644 --- a/crates/buzz-agent/src/wire.rs +++ b/crates/buzz-agent/src/wire.rs @@ -1,6 +1,6 @@ use serde::Deserialize; use serde_json::{json, Value}; -use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWriteExt}; +use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWrite, AsyncWriteExt}; use tokio::sync::mpsc; use crate::types::{ContentBlock, McpServerStdio}; @@ -413,8 +413,25 @@ pub async fn read_bounded_line( } } -pub async fn writer_task(mut rx: mpsc::Receiver) { - let mut stdout = tokio::io::stdout(); +pub async fn writer_task(rx: mpsc::Receiver) { + write_frames(rx, tokio::io::stdout()).await; +} + +/// Drain `rx`, writing each frame to `out` as a newline-terminated JSON line. +/// Generic over the sink so tests can inject an `AsyncWrite` that fails on +/// flush; production passes stdout. +/// +/// Both `write_all` and `flush` failure are connection-fatal: they return, +/// dropping `rx` so `async_main`'s writer-death arm cancels every session. +/// Flush must be fatal too — a blocking stdout can report `Ok` from +/// `write_all` when it only schedules the underlying write and surface the +/// real error at `flush`, so ignoring flush failure would leave a dead stdout +/// undetected and strand any correlated ask waiting for a reply that can never +/// be written. +pub(crate) async fn write_frames( + mut rx: mpsc::Receiver, + mut out: W, +) { while let Some(msg) = rx.recv().await { let WireMsg::Notify(v) = msg; let mut s = match serde_json::to_string(&v) { @@ -425,10 +442,9 @@ pub async fn writer_task(mut rx: mpsc::Receiver) { } }; s.push('\n'); - if stdout.write_all(s.as_bytes()).await.is_err() { + if out.write_all(s.as_bytes()).await.is_err() || out.flush().await.is_err() { return; } - let _ = stdout.flush().await; } }