diff --git a/buzz-agent-core/Cargo.lock b/buzz-agent-core/Cargo.lock index eeb3a8f8ba..a4d07d7b4f 100644 --- a/buzz-agent-core/Cargo.lock +++ b/buzz-agent-core/Cargo.lock @@ -750,6 +750,7 @@ dependencies = [ "futures", "goose", "goose-provider-types", + "rmcp", "serde", "serde_json", "tempfile", diff --git a/buzz-agent-core/Cargo.toml b/buzz-agent-core/Cargo.toml index fdfc62bc92..8e32304b42 100644 --- a/buzz-agent-core/Cargo.toml +++ b/buzz-agent-core/Cargo.toml @@ -35,6 +35,12 @@ path = "src/lib.rs" name = "buzz-agent-core" path = "src/main.rs" +# Test-only fake MCP server, copied from crates/buzz-agent. Built +# unconditionally because cargo can't gate bins on `cfg(test)`. +[[bin]] +name = "fake-mcp" +path = "tests/bin/fake_mcp.rs" + [dependencies] # Pinned to a release commit deliberately: goose's public `Agent` surface is # not semver-stable (many fields are `pub(super)`) and its MSRV moves quickly. @@ -57,6 +63,7 @@ serde_json = "1" anyhow = "1" tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } +rmcp = { version = "1.8", default-features = false, features = ["schemars"] } [dev-dependencies] tempfile = "3" diff --git a/buzz-agent-core/src/agent.rs b/buzz-agent-core/src/agent.rs index a1732aad23..8a29a5eced 100644 --- a/buzz-agent-core/src/agent.rs +++ b/buzz-agent-core/src/agent.rs @@ -86,6 +86,10 @@ pub fn prompt_to_text(blocks: &[ContentBlock]) -> String { /// responsible for emitting `usage_update` *before* the `session/prompt` /// response — that ordering is load-bearing for kind-44200 metrics /// (`buzz-agent/src/lib.rs:701-706`). +/// +/// `hook_extension` names the MCP extension carrying `_Stop`/`_PostCompact` +/// (buzz-dev-mcp). When set, the turn is not allowed to end while `_Stop` +/// objects — see [`crate::hooks`]. #[allow(clippy::too_many_arguments)] pub async fn run_turn( agent: Arc, @@ -94,10 +98,92 @@ pub async fn run_turn( max_rounds: Option, wire_tx: &WireSender, cancel: CancellationToken, + hook_extension: Option<&str>, ) -> (Result, TurnTokens) { - let text = prompt_to_text(&prompt); let mut tokens = TurnTokens::default(); + let mut next_message = Message::user().with_text(prompt_to_text(&prompt)); + let mut stop_blocks: u32 = 0; + // Outer loop exists solely for the `_Stop` veto: goose's `reply()` stream + // ends when the model stops, so continuing means re-entering `reply()` + // with the objection appended. Goose's own Stop hook does the equivalent + // internally (`agent.rs:2891-2917`); we do it here because its + // `hook_manager` is private. + loop { + let stop = match drive_stream( + &agent, + session_id, + next_message, + max_rounds, + wire_tx, + &cancel, + &mut tokens, + hook_extension, + ) + .await + { + Ok(s) => s, + Err(e) => return (Err(e), tokens), + }; + + // Only a clean end-of-turn is vetoable. Cancellation and refusals pass + // through untouched. + if !matches!(stop, StopReason::EndTurn) || cancel.is_cancelled() { + return (Ok(stop), tokens); + } + + let Some(extension) = hook_extension else { + return (Ok(stop), tokens); + }; + + if stop_blocks >= crate::hooks::MAX_STOP_BLOCKS { + tracing::warn!( + blocks = stop_blocks, + "_Stop veto cap reached; ending turn anyway" + ); + return (Ok(stop), tokens); + } + + let Some(session) = current_session(session_id).await else { + return (Ok(stop), tokens); + }; + let Some(objection) = crate::hooks::stop_objection(&agent, &session, extension).await + else { + return (Ok(stop), tokens); + }; + + stop_blocks += 1; + tracing::info!(blocks = stop_blocks, "_Stop hook vetoed end of turn"); + + // Agent-visible, user-invisible — the objection steers the model + // without appearing in the channel, matching both buzz-agent and + // goose's own Deny handling. + next_message = Message::user() + .with_text(format!("[Stop] {objection}")) + .with_visibility(false, true); + } +} + +/// Look up the goose `Session` record needed to dispatch a hook tool. +async fn current_session(session_id: &str) -> Option { + goose::session::session_manager::SessionManager::instance() + .get_session(session_id, false) + .await + .ok() +} + +/// Drive a single `reply()` stream to completion. +#[allow(clippy::too_many_arguments)] +async fn drive_stream( + agent: &Arc, + session_id: &str, + message: Message, + max_rounds: Option, + wire_tx: &WireSender, + cancel: &CancellationToken, + tokens: &mut TurnTokens, + hook_extension: Option<&str>, +) -> Result { let session_config = SessionConfig { id: session_id.to_string(), schedule_id: None, @@ -106,15 +192,11 @@ pub async fn run_turn( }; let mut stream = match agent - .reply( - Message::user().with_text(text), - session_config, - Some(cancel.clone()), - ) + .reply(message, session_config, Some(cancel.clone())) .await { Ok(s) => s, - Err(e) => return (Err(AgentError::Llm(e.to_string())), tokens), + Err(e) => return Err(AgentError::Llm(e.to_string())), }; let mut keepalive = tokio::time::interval(KEEPALIVE_INTERVAL); @@ -122,6 +204,7 @@ pub async fn run_turn( keepalive.tick().await; // first tick is immediate; discard it let mut stop = StopReason::EndTurn; + let mut compacted = false; loop { tokio::select! { @@ -143,7 +226,9 @@ pub async fn run_turn( let Some(event) = next else { break }; match event { Ok(ev) => { - if let Some(reason) = handle_event(ev, session_id, wire_tx, &mut tokens).await { + if let Some(reason) = + handle_event(ev, session_id, wire_tx, tokens, &mut compacted).await + { stop = reason; break; } @@ -159,18 +244,41 @@ pub async fn run_turn( } else { AgentError::Llm(msg) }; - return (Err(err), tokens); + return Err(err); } } } } } + // Goose compacted history mid-turn. Re-inject the todo list so it survives + // the truncation — this is what buzz-agent's `_PostCompact` existed for + // (`handoff.rs:73-81`). Steering is the right channel: goose drains it at + // the next round boundary (`agent.rs:1951-1974`). + if compacted && !cancel.is_cancelled() { + if let Some(extension) = hook_extension { + if let Some(session) = current_session(session_id).await { + if let Some(state) = + crate::hooks::post_compact_state(agent, &session, extension).await + { + agent + .steer( + session_id, + Message::user() + .with_text(format!("[PostCompact] {state}")) + .with_visibility(false, true), + ) + .await; + } + } + } + } + if cancel.is_cancelled() { stop = StopReason::Cancelled; } - (Ok(stop), tokens) + Ok(stop) } /// Translate one `AgentEvent` into ACP notifications. @@ -181,6 +289,7 @@ async fn handle_event( session_id: &str, wire_tx: &WireSender, tokens: &mut TurnTokens, + compacted: &mut bool, ) -> Option { match event { AgentEvent::Message(msg) => { @@ -206,10 +315,12 @@ async fn handle_event( } // Compaction happened. buzz-agent surfaced this as a `[Context - // Handoff]` history rewrite; here Goose has already done it and we - // only need to note it. + // Handoff]` history rewrite; Goose has already done the rewrite, so we + // only flag it — the caller then asks `_PostCompact` for state to + // re-inject once the stream settles. AgentEvent::HistoryReplaced(_) => { tracing::info!(target: "buzz_agent::compaction", "history compacted"); + *compacted = true; None } diff --git a/buzz-agent-core/src/hooks.rs b/buzz-agent-core/src/hooks.rs new file mode 100644 index 0000000000..6fdb8db5be --- /dev/null +++ b/buzz-agent-core/src/hooks.rs @@ -0,0 +1,144 @@ +//! buzz-agent's `_Stop` / `_PostCompact` lifecycle hooks, driven from our own +//! turn loop rather than from goose's. +//! +//! # Why not goose's hook system +//! +//! Goose has a blocking Stop hook with exactly buzz-agent's semantics +//! (`agent.rs:2891-2917`): on `HookDecision::Deny` the objection is pushed +//! into the conversation as an agent-visible / user-invisible message and the +//! loop continues, bounded by a block cap. But it is unreachable for us: +//! +//! * `Agent::hook_manager` is private, and the only setter is +//! `#[cfg(test)]`-gated (`agent.rs:430`). +//! * Hooks are otherwise discovered as *subprocesses* from +//! `/.goose/plugins/*/hooks/hooks.json` (`hooks/mod.rs:238-253`). +//! +//! The subprocess route looked viable until you notice where the answer lives: +//! `buzz-dev-mcp`'s todo list is **in-process state** (`todo.rs:49`, a +//! `Mutex>`), and that process's stdio is owned by goose. A separate +//! hook binary spawned by goose cannot see it. Writing a `hooks.json` would +//! give us a hook that always answers "no objection" — worse than no hook, +//! because it would look like it worked. +//! +//! # What we do instead +//! +//! We own the outer loop, so we ask the tool ourselves. +//! `Agent::dispatch_tool_call` is public (`agent.rs:1059`), so between rounds +//! we call `_Stop` on the same extension the model uses, and if it objects we +//! feed the objection back via `Agent::steer` — which goose drains at the +//! round boundary (`agent.rs:1951-1974`) and which also prevents the turn from +//! ending while a steer is pending (`agent.rs:2876-2878`). +//! +//! Net effect is buzz-agent's behaviour, with no goose changes: the agent does +//! not get to stop while its todo list has open items. + +use std::sync::Arc; + +use goose::agents::Agent; +use goose::session::session_manager::Session; +use rmcp::model::CallToolRequestParams; + +/// Hard cap on consecutive vetoes, mirroring goose's own `stop_hook_block_cap` +/// (`agent.rs:108-119`). A tool that always objects must not trap the turn. +pub const MAX_STOP_BLOCKS: u32 = 3; + +/// Tools are advertised to the model prefixed with their extension name +/// (`extension_manager.rs:1425`), and dispatch resolves the same prefixed +/// form. `_Stop` is hidden from the model but still callable by us. +fn qualified(extension: &str, tool: &str) -> String { + format!("{extension}__{tool}") +} + +/// Ask `_Stop` whether the turn may end. +/// +/// Returns `Some(objection)` when the agent should keep working. Any failure — +/// no such extension, tool error, malformed result — is `None`. A broken hook +/// must never trap a turn; that is also goose's rule for its own hooks +/// ("a misbehaving hook MUST NOT block", `hooks/mod.rs:361`). +pub async fn stop_objection( + agent: &Arc, + session: &Session, + extension: &str, +) -> Option { + let text = call_hook(agent, session, extension, "_Stop").await?; + let trimmed = text.trim(); + // Empty response is the documented "no objection" signal + // (`buzz-dev-mcp/src/todo.rs:100`). + (!trimmed.is_empty()).then(|| trimmed.to_string()) +} + +/// Ask `_PostCompact` for state to re-inject after goose compacted history. +/// +/// buzz-agent called this so the todo list survives truncation; goose emits +/// `AgentEvent::HistoryReplaced` when it compacts, which is our trigger. +pub async fn post_compact_state( + agent: &Arc, + session: &Session, + extension: &str, +) -> Option { + let text = call_hook(agent, session, extension, "_PostCompact").await?; + let trimmed = text.trim(); + (!trimmed.is_empty()).then(|| trimmed.to_string()) +} + +/// Dispatch a zero-argument hook tool and flatten its result to text. +async fn call_hook( + agent: &Arc, + session: &Session, + extension: &str, + tool: &str, +) -> Option { + let params = CallToolRequestParams::new(qualified(extension, tool)) + .with_arguments(serde_json::Map::new()); + + let (_id, result) = agent + .dispatch_tool_call( + params, + format!("buzz_hook_{tool}"), + // Deliberately not cancellable: this runs at a turn boundary and + // must resolve either way. + None, + session, + ) + .await; + + let call = match result { + Ok(c) => c, + Err(e) => { + tracing::debug!(tool, error = %e, "hook tool unavailable"); + return None; + } + }; + + let out = match call.result.await { + Ok(o) => o, + Err(e) => { + tracing::debug!(tool, error = ?e, "hook tool failed"); + return None; + } + }; + + if out.is_error.unwrap_or(false) { + tracing::debug!(tool, "hook tool returned is_error"); + return None; + } + + let mut text = String::new(); + for content in out.content.iter() { + if let Some(raw) = content.as_text() { + text.push_str(&raw.text); + } + } + Some(text) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn qualifies_tool_names_with_the_extension_prefix() { + // Must match extension_manager.rs:1425 or dispatch cannot resolve it. + assert_eq!(qualified("developer", "_Stop"), "developer___Stop"); + } +} diff --git a/buzz-agent-core/src/lib.rs b/buzz-agent-core/src/lib.rs index bb8b94c2d9..fbad500847 100644 --- a/buzz-agent-core/src/lib.rs +++ b/buzz-agent-core/src/lib.rs @@ -21,6 +21,7 @@ pub mod agent; pub mod config; +pub mod hooks; pub mod types; pub mod wire; @@ -61,6 +62,9 @@ struct Session { /// the override takes effect "from the next prompt" and never mutates a /// turn already in flight (`buzz-agent/src/lib.rs:494-502`). pending_model: Option, + /// Name of the MCP extension carrying `_Stop`/`_PostCompact`, if any. + /// See [`crate::hooks`] for why we dispatch these ourselves. + hook_extension: Option, accumulated_input_tokens: u64, accumulated_output_tokens: u64, } @@ -82,7 +86,7 @@ async fn build_agent( cwd: &str, system_prompt: Option<&str>, mcp_servers: &[McpServerStdio], -) -> Result<(Arc, String), AgentError> { +) -> Result<(Arc, String, Option), AgentError> { let session_manager = Arc::new(SessionManager::instance()); let permission_manager = PermissionManager::instance(); @@ -174,7 +178,49 @@ async fn build_agent( } } - Ok((Arc::new(agent), session.id)) + let agent = Arc::new(agent); + + // Find the extension carrying `_Stop`, by asking rather than assuming the + // name — `buzz-acp` derives it from the MCP binary's file stem + // (`buzz-acp/src/lib.rs:4145-4149`), so it is not a fixed string. + // + // KNOWN DEVIATION: buzz-agent hid `_`-prefixed tools from the model + // (`agent.rs:328-336`) while still calling them itself. Goose's + // `available_tools` allowlist gates advertising *and* dispatch through the + // same cache (`extension_manager.rs:1421`, `:1698`), so hiding them would + // also make them undispatchable — which would break the veto. They stay + // visible; `hook_tool_guidance()` tells the model to leave them alone. + let mut hook_extension = None; + for tool in agent.list_tools(&session.id, None).await { + if tool.name.ends_with("___Stop") { + hook_extension = tool + .name + .strip_suffix("___Stop") + .map(|prefix| prefix.to_string()); + break; + } + } + + if let Some(ext) = &hook_extension { + agent + .extend_system_prompt("buzz_hook_tools".to_string(), hook_tool_guidance()) + .await; + tracing::info!(extension = %ext, "lifecycle hooks available"); + } + + Ok((agent, session.id, hook_extension)) +} + +/// Keep the model's hands off the lifecycle hooks. +/// +/// They are ordinary MCP tools on the wire, so a generic harness advertises +/// them. buzz-agent solved this by hiding them; we cannot (see `build_agent`), +/// so we ask instead. +fn hook_tool_guidance() -> String { + "Tools whose names begin with an underscore (`_Stop`, `_PostCompact`) are \ + lifecycle hooks invoked automatically by the runtime. Never call them \ + yourself. Use the `todo` tool to manage your task list." + .to_string() } pub async fn serve(cfg: Config) -> anyhow::Result<()> { @@ -350,7 +396,7 @@ async fn session_new(app: &Arc, id: Value, params: Value, wire_tx: &WireSen .await; } - let (agent, goose_session_id) = + let (agent, goose_session_id, hook_extension) = match build_agent(&app.cfg, &p.cwd, p.system_prompt.as_deref(), &p.mcp_servers).await { Ok(v) => v, Err(e) => { @@ -379,6 +425,7 @@ async fn session_new(app: &Arc, id: Value, params: Value, wire_tx: &WireSen cancel: None, model_override: None, pending_model: None, + hook_extension, accumulated_input_tokens: 0, accumulated_output_tokens: 0, }, @@ -472,7 +519,7 @@ async fn session_prompt(app: &Arc, id: Value, params: Value, wire_tx: &Wire let cancel = CancellationToken::new(); // Single-flight per session, and capture the agent handle. - let (agent, goose_session_id, pending_model) = { + let (agent, goose_session_id, pending_model, hook_extension) = { let mut sessions = app.sessions.lock().await; let Some(s) = sessions.get_mut(&p.session_id) else { return wire::send( @@ -501,6 +548,7 @@ async fn session_prompt(app: &Arc, id: Value, params: Value, wire_tx: &Wire s.agent.clone(), s.goose_session_id.clone(), s.pending_model.take(), + s.hook_extension.clone(), ) }; @@ -538,6 +586,7 @@ async fn session_prompt(app: &Arc, id: Value, params: Value, wire_tx: &Wire app.cfg.max_rounds, wire_tx, cancel, + hook_extension.as_deref(), ) .await; diff --git a/buzz-agent-core/tests/bin/fake_mcp.rs b/buzz-agent-core/tests/bin/fake_mcp.rs new file mode 100644 index 0000000000..0bbd1d3478 --- /dev/null +++ b/buzz-agent-core/tests/bin/fake_mcp.rs @@ -0,0 +1,301 @@ +//! Tiny fake MCP server for integration tests. +//! +//! Reads JSON-RPC line frames on stdin and replies on stdout. Driven by +//! environment variables so tests can simulate misbehavior: +//! +//! FAKE_MCP_HANG_INIT=1 — never reply to `initialize` (init timeout) +//! FAKE_MCP_HANG_TOOLS=1 — never reply to `tools/list` (list timeout) +//! FAKE_MCP_TOOL_COUNT=N — return N tools (default: 1) +//! FAKE_MCP_HUGE_DESC=1 — every tool description is 100 KB +//! FAKE_MCP_DESC_SIZE=N — every tool description is N bytes (overrides HUGE_DESC) +//! FAKE_MCP_TOOL_DELAY=N — `tools/call` sleeps N seconds before replying +//! (use a large value, e.g. 999, to simulate hang) +//! FAKE_MCP_RESULT_SIZE=N — `tools/call` returns an N-byte text result +//! (default: the literal "ok"); grows history +//! FAKE_MCP_PID_FILE=path — write the child PID to `path` on startup +//! (for tests that want to verify the child died) +//! FAKE_MCP_SPAWN_GRANDCHILD=1 +//! — on `tools/call`, spawn a `sleep 999` +//! grandchild before hanging. Its PID is +//! written to FAKE_MCP_GRANDCHILD_PID_FILE +//! so a test can verify the entire process +//! tree dies on timeout. +//! FAKE_MCP_GRANDCHILD_PID_FILE=path +//! — path to write the grandchild PID to. +//! 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 +//! (use a large value to simulate hang/timeout) +//! FAKE_MCP_STOP_COUNT=N — `_Stop` returns STOP_TEXT for the first N +//! invocations; empty string thereafter. If +//! unset, every call returns STOP_TEXT. +//! FAKE_MCP_POSTCOMPACT_HOOK=1 +//! — expose a `_PostCompact` hook tool +//! FAKE_MCP_POSTCOMPACT_TEXT=text +//! — `_PostCompact` returns this (default: "") + +use std::io::{BufRead, Write}; + +use serde_json::{json, Value}; + +fn env_flag(k: &str) -> bool { + std::env::var(k).map(|v| v != "0").unwrap_or(false) +} + +fn env_usize(k: &str, default: usize) -> usize { + std::env::var(k) + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(default) +} + +fn env_u64(k: &str, default: u64) -> u64 { + std::env::var(k) + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(default) +} + +fn write_response(id: Value, result: Value) { + let msg = json!({ "jsonrpc": "2.0", "id": id, "result": result }); + let mut s = serde_json::to_string(&msg).expect("serialize"); + s.push('\n'); + let mut out = std::io::stdout().lock(); + out.write_all(s.as_bytes()).expect("write"); + out.flush().expect("flush"); +} + +fn hang_forever() -> ! { + loop { + std::thread::sleep(std::time::Duration::from_secs(60)); + } +} + +fn make_tools( + count: usize, + desc: &str, + include_stop_hook: bool, + include_post_compact_hook: bool, +) -> Vec { + let mut tools: Vec = (0..count) + .map(|i| { + json!({ + "name": format!("tool_{i}"), + "description": desc, + "inputSchema": { "type": "object", "properties": {} }, + }) + }) + .collect(); + if include_stop_hook { + tools.push(json!({ + "name": "_Stop", + "description": "stop hook", + "inputSchema": { "type": "object", "properties": {} }, + })); + } + if include_post_compact_hook { + tools.push(json!({ + "name": "_PostCompact", + "description": "post compact hook", + "inputSchema": { "type": "object", "properties": {} }, + })); + } + tools +} + +fn main() { + // Optional: write our own PID so a test can later check the process is gone. + if let Ok(path) = std::env::var("FAKE_MCP_PID_FILE") { + let pid = std::process::id().to_string(); + let _ = std::fs::write(&path, pid); + } + + let hang_init = env_flag("FAKE_MCP_HANG_INIT"); + let hang_tools = env_flag("FAKE_MCP_HANG_TOOLS"); + let tool_count = env_usize("FAKE_MCP_TOOL_COUNT", 1); + // FAKE_MCP_DESC_SIZE wins over FAKE_MCP_HUGE_DESC when set. + let desc: String = if let Some(n) = std::env::var("FAKE_MCP_DESC_SIZE") + .ok() + .and_then(|v| v.parse::().ok()) + { + "x".repeat(n) + } else if env_flag("FAKE_MCP_HUGE_DESC") { + "x".repeat(100_000) + } else { + "fake tool".to_owned() + }; + let tool_delay_secs = env_u64("FAKE_MCP_TOOL_DELAY", 0); + // Tool-call result text size in bytes (default: the literal "ok"). Lets a + // test grow session history by a controlled amount via a tool result. + let result_size = env_u64("FAKE_MCP_RESULT_SIZE", 0) as usize; + let stop_hook = env_flag("FAKE_MCP_STOP_HOOK"); + let stop_text = std::env::var("FAKE_MCP_STOP_TEXT").unwrap_or_else(|_| "keep going".to_owned()); + let stop_delay_secs = env_u64("FAKE_MCP_STOP_DELAY", 0); + // 0 means "unset" → unlimited; any positive value caps the number of + // calls that return STOP_TEXT before flipping to empty string. + let stop_count_limit: usize = env_usize("FAKE_MCP_STOP_COUNT", usize::MAX); + let mut stop_calls_seen: usize = 0; + let post_compact_hook = env_flag("FAKE_MCP_POSTCOMPACT_HOOK"); + let post_compact_text = std::env::var("FAKE_MCP_POSTCOMPACT_TEXT").unwrap_or_default(); + + // 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. + let cancel_log_path = std::env::var("FAKE_MCP_CANCEL_LOG").ok(); + let (tx, rx) = std::sync::mpsc::channel::<(Value, Option)>(); + let cancel_log_for_thread = cancel_log_path.clone(); + std::thread::spawn(move || { + let stdin = std::io::stdin(); + let lines = stdin.lock().lines(); + for line in lines { + let line = match line { + Ok(l) => l, + Err(_) => return, + }; + if line.trim().is_empty() { + continue; + } + let msg: Value = match serde_json::from_str(&line) { + Ok(v) => v, + Err(_) => continue, + }; + let method = msg.get("method").and_then(Value::as_str).unwrap_or(""); + let id = msg.get("id").cloned(); + // Notifications carry no id. Log cancellations if configured. + if id.is_none() || id == Some(Value::Null) { + if method == "notifications/cancelled" { + if let Some(ref path) = cancel_log_for_thread { + use std::io::Write as _; + if let Ok(mut f) = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(path) + { + let _ = writeln!(f, "{}", line.trim()); + } + } + } + continue; + } + // Send requests (with id) to the main processing loop. + let _ = tx.send((msg, id)); + } + }); + + while let Ok((msg, id_opt)) = rx.recv() { + let method = msg.get("method").and_then(Value::as_str).unwrap_or(""); + let id = id_opt.unwrap_or(Value::Null); + + match method { + "initialize" => { + if hang_init { + hang_forever(); + } + write_response( + id, + json!({ + "protocolVersion": "2025-06-18", + "capabilities": { "tools": {} }, + "serverInfo": { "name": "fake-mcp", "version": "0.0.0" }, + }), + ); + } + "tools/list" => { + if hang_tools { + hang_forever(); + } + write_response( + id, + json!({ + "tools": make_tools(tool_count, &desc, stop_hook, post_compact_hook) + }), + ); + } + "tools/call" => { + // Signal that the request was received (for tests that + // need to wait until the call is in-flight before cancelling). + // Write the request id so tests can correlate with cancel. + if let Ok(path) = std::env::var("FAKE_MCP_CALL_RECEIVED") { + let id_str = serde_json::to_string(&id).unwrap_or_else(|_| "?".into()); + let _ = std::fs::write(&path, id_str); + } + let called_name = msg + .get("params") + .and_then(|p| p.get("name")) + .and_then(Value::as_str) + .unwrap_or(""); + // 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") { + let child = std::process::Command::new("sleep") + .arg("999") + .spawn() + .expect("spawn grandchild"); + if let Ok(path) = std::env::var("FAKE_MCP_GRANDCHILD_PID_FILE") { + let _ = std::fs::write(&path, child.id().to_string()); + } + std::mem::forget(child); + } + if called_name == "_Stop" { + if stop_delay_secs > 0 { + std::thread::sleep(std::time::Duration::from_secs(stop_delay_secs)); + } + // Once we exceed the configured count, return empty + // text so the agent treats it as no objection. This + // lets a test exercise the "objected then cleared" + // path without relying on the rejection budget. + let payload = if stop_calls_seen < stop_count_limit { + stop_text.clone() + } else { + String::new() + }; + stop_calls_seen = stop_calls_seen.saturating_add(1); + write_response( + id, + json!({ + "content": [{ "type": "text", "text": payload }], + "isError": false, + }), + ); + continue; + } + if called_name == "_PostCompact" { + write_response( + id, + json!({ + "content": [{ "type": "text", "text": post_compact_text }], + "isError": false, + }), + ); + continue; + } + if tool_delay_secs > 0 { + std::thread::sleep(std::time::Duration::from_secs(tool_delay_secs)); + } + let result_text = if result_size > 0 { + "x".repeat(result_size) + } else { + "ok".to_owned() + }; + write_response( + id, + json!({ + "content": [{ "type": "text", "text": result_text }], + "isError": false, + }), + ); + } + _ => { + // Unknown method: respond with an error so rmcp doesn't hang. + let err = json!({ + "jsonrpc": "2.0", "id": id, + "error": { "code": -32601, "message": format!("method not found: {method}") }, + }); + let mut s = serde_json::to_string(&err).unwrap(); + s.push('\n'); + let mut out = std::io::stdout().lock(); + let _ = out.write_all(s.as_bytes()); + let _ = out.flush(); + } + } + } +} diff --git a/buzz-agent-core/tests/stop_hook.rs b/buzz-agent-core/tests/stop_hook.rs new file mode 100644 index 0000000000..42b56b7d9f --- /dev/null +++ b/buzz-agent-core/tests/stop_hook.rs @@ -0,0 +1,337 @@ +//! Proves the `_Stop` veto: the turn must not end while the hook objects. +//! +//! This is buzz-agent's most load-bearing non-standard behaviour — it is why +//! the agent finishes its todo list instead of stopping halfway. Goose has the +//! same mechanism internally, but its `hook_manager` is private and its hooks +//! are subprocesses that cannot see buzz-dev-mcp's in-process todo state, so +//! we drive it from our own loop (see `src/hooks.rs`). +//! +//! Setup: a fake MCP server exposing `_Stop`, plus a fake provider that counts +//! how many times it is asked to generate. One user prompt with a hook that +//! objects twice must produce three provider calls, not one. + +use std::io::{BufRead, BufReader, Write}; +use std::net::TcpListener; +use std::process::{Child, Command, Stdio}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +use serde_json::{json, Value}; + +/// Fake OpenAI-compatible provider that always answers "done" and counts calls. +fn spawn_counting_provider() -> (String, Arc) { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind"); + let addr = listener.local_addr().expect("addr"); + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + + std::thread::spawn(move || { + for stream in listener.incoming() { + let Ok(mut stream) = stream else { continue }; + let mut reader = BufReader::new(stream.try_clone().expect("clone")); + + let mut request_line = String::new(); + if reader.read_line(&mut request_line).is_err() { + continue; + } + let mut content_length = 0usize; + loop { + let mut line = String::new(); + if reader.read_line(&mut line).is_err() { + break; + } + let t = line.trim_end(); + if t.is_empty() { + break; + } + if let Some(v) = t + .strip_prefix("content-length: ") + .or_else(|| t.strip_prefix("Content-Length: ")) + { + content_length = v.parse().unwrap_or(0); + } + } + let mut body = vec![0u8; content_length]; + use std::io::Read; + let _ = reader.read_exact(&mut body); + + if request_line.contains("/models") { + let payload = json!({ "object": "list", "data": [] }).to_string(); + let _ = stream.write_all( + format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + payload.len(), payload + ).as_bytes(), + ); + let _ = stream.flush(); + continue; + } + + counter.fetch_add(1, Ordering::SeqCst); + + let chunk = |delta: Value, finish: Value, usage: Value| { + let mut c = json!({ + "id": "c", "object": "chat.completion.chunk", "created": 1, + "model": "fake-model", + "choices": [{ "index": 0, "delta": delta, "finish_reason": finish }], + }); + if !usage.is_null() { + c["usage"] = usage; + } + format!("data: {c}\n\n") + }; + let sse = format!( + "{}{}data: [DONE]\n\n", + chunk( + json!({ "role": "assistant", "content": "done" }), + Value::Null, + Value::Null + ), + chunk( + json!({}), + json!("stop"), + json!({ "prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7 }) + ), + ); + let _ = stream.write_all( + format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + sse.len(), sse + ).as_bytes(), + ); + let _ = stream.flush(); + } + }); + + (format!("http://{addr}"), calls) +} + +struct Harness { + child: Child, + stdin: std::process::ChildStdin, + stdout: BufReader, + next_id: i64, +} + +impl Harness { + fn start(base_url: &str, home: &std::path::Path) -> Self { + let mut child = Command::new(env!("CARGO_BIN_EXE_buzz-agent-core")) + .env("BUZZ_AGENT_PROVIDER", "openai-compat") + .env("BUZZ_AGENT_MODEL", "fake-model") + .env("OPENAI_COMPAT_API_KEY", "k") + .env("OPENAI_COMPAT_BASE_URL", base_url) + .env("HOME", home) + .env("XDG_CONFIG_HOME", home.join("config")) + .env("XDG_DATA_HOME", home.join("data")) + .env("GOOSE_DISABLE_KEYRING", "1") + .env("RUST_LOG", "warn") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::inherit()) + .spawn() + .expect("spawn"); + let stdin = child.stdin.take().expect("stdin"); + let stdout = BufReader::new(child.stdout.take().expect("stdout")); + Self { + child, + stdin, + stdout, + next_id: 0, + } + } + + fn request(&mut self, method: &str, params: Value) -> i64 { + self.next_id += 1; + let id = self.next_id; + writeln!( + self.stdin, + "{}", + json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }) + ) + .expect("write"); + self.stdin.flush().expect("flush"); + id + } + + fn await_response(&mut self, id: i64) -> Value { + loop { + let mut line = String::new(); + let n = self.stdout.read_line(&mut line).expect("read"); + assert_ne!(n, 0, "agent closed stdout awaiting id={id}"); + let Ok(msg) = serde_json::from_str::(&line) else { + continue; + }; + if msg.get("id").and_then(Value::as_i64) == Some(id) { + return msg; + } + } + } +} + +impl Drop for Harness { + fn drop(&mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} + +/// Declare the fake MCP server, with `_Stop` objecting `objections` times. +fn mcp_servers(objections: usize) -> Value { + json!([{ + "name": "devtools", + "command": env!("CARGO_BIN_EXE_fake-mcp"), + "args": [], + "env": [ + { "name": "FAKE_MCP_STOP_HOOK", "value": "1" }, + { "name": "FAKE_MCP_STOP_TEXT", "value": "1 item still open" }, + { "name": "FAKE_MCP_STOP_COUNT", "value": objections.to_string() }, + ], + }]) +} + +fn new_session(h: &mut Harness, cwd: &std::path::Path, objections: usize) -> String { + let id = h.request("initialize", json!({ "protocolVersion": 2 })); + let _ = h.await_response(id); + + let id = h.request( + "session/new", + json!({ + "cwd": cwd.to_str().unwrap(), + "mcpServers": mcp_servers(objections), + }), + ); + let resp = h.await_response(id); + resp["result"]["sessionId"] + .as_str() + .unwrap_or_else(|| panic!("session/new failed: {resp}")) + .to_string() +} + +#[test] +fn stop_hook_veto_extends_the_turn() { + let (base_url, calls) = spawn_counting_provider(); + let home = tempfile::tempdir().expect("home"); + let cwd = tempfile::tempdir().expect("cwd"); + let mut h = Harness::start(&base_url, home.path()); + + // Object twice, then allow. + let session_id = new_session(&mut h, cwd.path(), 2); + + let id = h.request( + "session/prompt", + json!({ + "sessionId": session_id, + "prompt": [{ "type": "text", "text": "do the work" }], + }), + ); + let resp = h.await_response(id); + + assert_eq!( + resp["result"]["stopReason"], "end_turn", + "turn should still end cleanly once the hook stops objecting: {resp}" + ); + + // The model says "done" immediately, so without the veto this is 1. + // Two objections must force two more generations. + let n = calls.load(Ordering::SeqCst); + assert_eq!( + n, 3, + "expected 3 provider calls (initial + 2 vetoes), got {n} — the _Stop veto did not fire" + ); +} + +#[test] +fn stop_hook_veto_is_capped() { + let (base_url, calls) = spawn_counting_provider(); + let home = tempfile::tempdir().expect("home"); + let cwd = tempfile::tempdir().expect("cwd"); + let mut h = Harness::start(&base_url, home.path()); + + // Always objects. Must not trap the turn forever. + let session_id = new_session(&mut h, cwd.path(), usize::MAX); + + let id = h.request( + "session/prompt", + json!({ + "sessionId": session_id, + "prompt": [{ "type": "text", "text": "do the work" }], + }), + ); + let resp = h.await_response(id); + + assert_eq!( + resp["result"]["stopReason"], "end_turn", + "a permanently objecting hook must not hang the turn: {resp}" + ); + + // MAX_STOP_BLOCKS = 3, so: initial + 3 vetoes = 4, then give up. + let n = calls.load(Ordering::SeqCst); + assert_eq!( + n, 4, + "expected the veto to stop at MAX_STOP_BLOCKS (1 + 3 = 4 calls), got {n}" + ); +} + +#[test] +fn no_stop_hook_means_no_extra_rounds() { + let (base_url, calls) = spawn_counting_provider(); + let home = tempfile::tempdir().expect("home"); + let cwd = tempfile::tempdir().expect("cwd"); + let mut h = Harness::start(&base_url, home.path()); + + let id = h.request("initialize", json!({ "protocolVersion": 2 })); + let _ = h.await_response(id); + + // No MCP servers at all: nothing to veto with. + let id = h.request( + "session/new", + json!({ "cwd": cwd.path().to_str().unwrap(), "mcpServers": [] }), + ); + let resp = h.await_response(id); + let session_id = resp["result"]["sessionId"].as_str().unwrap().to_string(); + + let id = h.request( + "session/prompt", + json!({ + "sessionId": session_id, + "prompt": [{ "type": "text", "text": "hi" }], + }), + ); + let resp = h.await_response(id); + assert_eq!(resp["result"]["stopReason"], "end_turn"); + + let n = calls.load(Ordering::SeqCst); + assert_eq!(n, 1, "no hook should mean exactly one generation, got {n}"); +} + +#[test] +fn stop_hook_tools_are_discovered_by_suffix() { + // The extension name comes from buzz-acp's MCP binary file stem + // (`buzz-acp/src/lib.rs:4145`), so discovery must not hardcode a name. + // Here the server is called "devtools", not "developer" or "buzz-dev-mcp". + let (base_url, calls) = spawn_counting_provider(); + let home = tempfile::tempdir().expect("home"); + let cwd = tempfile::tempdir().expect("cwd"); + let mut h = Harness::start(&base_url, home.path()); + + let session_id = new_session(&mut h, cwd.path(), 1); + + let id = h.request( + "session/prompt", + json!({ + "sessionId": session_id, + "prompt": [{ "type": "text", "text": "go" }], + }), + ); + let resp = h.await_response(id); + assert_eq!(resp["result"]["stopReason"], "end_turn"); + + let n = calls.load(Ordering::SeqCst); + assert_eq!( + n, 2, + "one objection under a non-obvious extension name should yield 2 calls, got {n}" + ); + + let _ = Duration::from_secs(0); +}