From 9869f24319cff0be8f6c8ce31eea5559a566d08d Mon Sep 17 00:00:00 2001 From: Michael Neale Date: Tue, 11 Aug 2026 14:50:05 +1000 Subject: [PATCH] refactor(buzz-agent): buzz owns the agent loop; goose supplies the parts Inverts the previous design. buzz-agent no longer calls goose::agents::Agent::reply -- it drives its own round loop and calls goose for the four heavy components: model call Provider::stream (via Agent::provider) tool surface Agent::list_tools tool execution Agent::dispatch_tool_call system prompt PromptManager (our own instance) compaction context_mgmt::{check_if_compaction_needed, compact_messages} Under Agent::reply every buzz-specific behaviour had to be smuggled in around goose's turn policy: the _Stop veto needed an outer loop purely to re-enter reply(), and [Reflect] had to be delivered as a *steer* because the tool result itself was out of reach. Owning the loop removes the smuggling -- the veto is a branch, and [Reflect] goes back on the tool result where buzz-agent originally put it and where the model reads it in context. New modules: loop_drive.rs the round loop: inference, tools, compaction, _Stop veto, steer drain, cancellation, max-rounds bound tools.rs parallel dispatch, announce->terminal wire invariant, [Reflect] on failure prompt.rs our PromptManager (goose's is pub(super) to Agent::reply) steer.rs our steer queue (goose's drain is pub(crate) to reply) agent.rs shrinks to what goose knows nothing about: ACP session/update emission and the keepalive ticker. Tool lifecycle now emits from tools.rs, which is the only place that knows when a call starts and ends -- emitting from streamed content as well would double-announce every tool. Behaviour preserved and covered by the existing suite: Fizz persona, _Stop veto (cap 3), [Reflect] (cap 8), _PostCompact re-injection, steering without cancellation, cancel leaving no unresolved tool call, usage accounting, and the model catalog. 81 unit + 21 integration tests pass; clippy clean (the one hints.rs warning predates this branch). Co-authored-by: Michael Neale Signed-off-by: Michael Neale --- crates/buzz-agent/src/agent.rs | 550 ++++------------------------ crates/buzz-agent/src/lib.rs | 76 ++-- crates/buzz-agent/src/loop_drive.rs | 537 +++++++++++++++++++++++++++ crates/buzz-agent/src/prompt.rs | 108 ++++++ crates/buzz-agent/src/steer.rs | 86 +++++ crates/buzz-agent/src/tools.rs | 209 +++++++++++ 6 files changed, 1061 insertions(+), 505 deletions(-) create mode 100644 crates/buzz-agent/src/loop_drive.rs create mode 100644 crates/buzz-agent/src/prompt.rs create mode 100644 crates/buzz-agent/src/steer.rs create mode 100644 crates/buzz-agent/src/tools.rs diff --git a/crates/buzz-agent/src/agent.rs b/crates/buzz-agent/src/agent.rs index 43984f7cc9..305f9c4f07 100644 --- a/crates/buzz-agent/src/agent.rs +++ b/crates/buzz-agent/src/agent.rs @@ -1,49 +1,18 @@ -//! The turn driver: Goose's agent loop, translated onto buzz-agent's ACP wire. +//! ACP wire emission for buzz-agent's turn loop. //! -//! This replaces roughly 8,000 lines of `crates/buzz-agent`: -//! -//! | buzz-agent file | lines | who owns it now | -//! |---|---:|---| -//! | `llm.rs` | 3846 | `goose::providers` (superset of the 4 providers) | -//! | `mcp.rs` | 1139 | `goose::agents::extension_manager` | -//! | `auth.rs` | 845 | `goose::providers` (incl. Databricks OAuth) | -//! | `agent.rs` | 746 | `goose::agents::Agent::reply` | -//! | `hints.rs` | 726 | kept -- goose's loader can't replace it (below) | -//! | `catalog.rs` | 631 | `goose::providers::init` model discovery | -//! | `builtin.rs` | 575 | kept -- served as a goose *frontend* tool (below) | -//! | `handoff.rs` | 430 | `goose::context_mgmt` (auto-compaction) | -//! -//! Two entries above are NOT goose substitutions; the old buzz-agent modules -//! are kept and wired into goose instead, because goose's own equivalents -//! don't cover them: -//! -//! * **`builtin.rs` / `load_skill`.** `Agent::with_config` loads zero -//! extensions and `build_agent` only adds the `mcpServers` the harness -//! declares, so goose's `skills` platform extension is never loaded. Instead -//! `build_agent` registers `load_skill` as a goose *platform* extension -//! backed by an `McpClientTrait` (see [`crate::builtin_client`]), so goose -//! dispatches it on its ordinary tool path. -//! * **`AGENTS.md` hints / skill index.** Goose's loader keys off -//! `GOOSE_HINTS_FILENAME` (`.goosehints`); the old code walked the directory -//! chain for `AGENTS.md` plus `~/AGENTS.md`. Every repo here ships the -//! latter shape, so `hints.rs` is kept and its output injected via -//! `extend_system_prompt` at session build. -//! -//! What this file keeps is the part Goose does *not* know about: the mapping -//! from `AgentEvent` onto the exact `session/update` notifications `buzz-acp` +//! The loop itself lives in [`crate::loop_drive`] — buzz-agent owns it and +//! calls goose for the model, the tools, the prompt and compaction. What stays +//! here is the part goose knows nothing about: the mapping from goose +//! `MessageContent` onto the exact `session/update` notifications `buzz-acp` //! consumes (`acp.rs:1528-1627`), and the `keepalive` ticker that resets the //! harness idle clock. -use std::sync::Arc; - -use futures::StreamExt; use serde_json::{json, Value}; -use tokio_util::sync::CancellationToken; -use goose::agents::{Agent, AgentEvent, SessionConfig}; -use goose_provider_types::conversation::message::{Message, MessageContent}; +use goose_provider_types::conversation::message::MessageContent; +use rmcp::model::CallToolRequestParams; -use crate::types::{AgentError, ContentBlock, StopReason}; +use crate::types::ContentBlock; use crate::wire::{self, WireSender}; /// How often to emit `keepalive` while waiting on the provider. @@ -58,32 +27,7 @@ use crate::wire::{self, WireSender}; /// alive — but there is no traffic during a long *pre-first-token* wait (big /// prompt processing, provider queueing, reasoning models). So the ticker /// stays. -const KEEPALIVE_INTERVAL: std::time::Duration = std::time::Duration::from_secs(30); - -/// Appended by buzz-agent to every failed tool result (`agent.rs:21-22`) so -/// the model diagnoses a failure instead of blindly retrying it. -/// -/// buzz-agent mutated the tool result itself. Goose gives no interception -/// point for that: its `PostToolUseFailure` hook is fire-and-forget and its -/// output is discarded (`agent.rs:589-620`). So we deliver the same text as a -/// steer instead — goose drains pending steers at the round boundary -/// (`agent.rs:1951-1974`), which is exactly when the model would next act on -/// the failed result. Agent-visible, user-invisible. -const ERROR_REFLECTION: &str = - "[Reflect] Before retrying, identify the cause and change your approach."; - -/// Cap on reflections per turn, so a tool failing in a loop cannot flood the -/// conversation. -const MAX_REFLECTIONS: usize = 8; - -/// How long to let goose unwind after `session/cancel` before giving up. -/// -/// Cancellation is cooperative: goose has to notice the token, send -/// `notifications/cancelled` to each in-flight MCP request -/// (`mcp_client.rs:687-690`), emit the resulting tool responses, and end the -/// stream. Dropping the stream instead skips all of that. This bounds the -/// wait so a wedged MCP child cannot hold the turn open indefinitely. -const CANCEL_DRAIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); +pub(crate) const KEEPALIVE_INTERVAL: std::time::Duration = std::time::Duration::from_secs(30); /// Per-turn token accounting, mirroring buzz-agent's contract. #[derive(Debug, Default, Clone, Copy)] @@ -129,351 +73,70 @@ pub fn prompt_to_text(blocks: &[ContentBlock]) -> String { out } -/// Drive one `session/prompt` turn to completion. -/// -/// Returns the ACP stop reason plus the turn's token counts. The caller is -/// 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, - session_id: &str, - prompt: Vec, - max_rounds: Option, - wire_tx: &WireSender, - cancel: CancellationToken, - hook_extension: Option<&str>, -) -> (Result, TurnTokens) { - let mut tokens = TurnTokens::default(); - let mut next_message = Message::user().with_text(prompt_to_text(&prompt)); - let mut stop_blocks: u32 = 0; +/// Reset the harness idle clock during provider silence. +pub(crate) async fn emit_keepalive(session_id: &str, wire_tx: &WireSender) { + wire::send( + wire_tx, + wire::session_update(session_id, json!({ "sessionUpdate": "keepalive" })), + ) + .await; +} - // 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, +/// Announce a tool call as in-progress. +/// +/// Every announcement must be matched by exactly one +/// [`emit_tool_call_update`], or the desktop renders a spinner that never +/// resolves. buzz-agent's loop guarantees the pairing including on cancel. +pub(crate) async fn emit_tool_call( + wire_tx: &WireSender, + session_id: &str, + tool_call_id: &str, + call: &CallToolRequestParams, +) { + let raw = serde_json::to_value(&call.arguments).unwrap_or(Value::Null); + wire::send( + wire_tx, + wire::session_update( 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); - } + json!({ + "sessionUpdate": "tool_call", + "toolCallId": tool_call_id, + "title": call.name.to_string(), + "kind": "other", + "status": "in_progress", + "rawInput": raw, + }), + ), + ) + .await; } -/// 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, +/// Move an announced tool call to its terminal state. +pub(crate) async fn emit_tool_call_update( 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, - max_turns: max_rounds, - retry_config: None, - }; - - let mut stream = match agent - .reply(message, session_config, Some(cancel.clone())) - .await - { - Ok(s) => s, - Err(e) => return Err(AgentError::Llm(e.to_string())), - }; - - let mut keepalive = tokio::time::interval(KEEPALIVE_INTERVAL); - keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - keepalive.tick().await; // first tick is immediate; discard it - - let mut stop = StopReason::EndTurn; - let mut compacted = false; - let mut reflections = 0usize; - - // Set once cancellation is observed; bounds how long we let goose unwind. - let mut drain: Option>> = None; - - // Tool calls announced to the harness, minus those that reached a terminal - // state. Anything still here when the stream ends gets a synthetic - // terminal update — see the cancel arm below. - let mut open_tool_calls: Vec = Vec::new(); - - loop { - tokio::select! { - // Keep the harness idle clock alive during provider silence. - _ = keepalive.tick() => { - wire::send( - wire_tx, - wire::session_update(session_id, json!({ "sessionUpdate": "keepalive" })), - ) - .await; - } - - // Cancellation is a cooperative DRAIN, not an abort. - // - // Breaking here drops `stream`, which drops the futures goose is - // awaiting — so `mcp_client.rs:688` never reaches its - // `cancel_token.cancelled()` arm and never sends - // `notifications/cancelled`. The MCP child keeps running its tool - // after the turn is over, and any announced `tool_call` never - // reaches a terminal state, leaving the desktop UI spinning - // (the invariant buzz-agent held at `agent.rs:470-477`). - // - // So we keep polling the stream and let goose unwind: it emits the - // tool responses, sends the MCP cancellations, and ends the stream - // itself. `drain` only bounds how long we are willing to wait. - _ = cancel.cancelled(), if drain.is_none() => { - stop = StopReason::Cancelled; - drain = Some(Box::pin(tokio::time::sleep(CANCEL_DRAIN_TIMEOUT))); - } - - // Goose did not unwind in time. Give up and synthesise terminal - // states, because a stuck spinner is worse than a wrong status. - _ = async { drain.as_mut().unwrap().await }, if drain.is_some() => { - tracing::warn!("cancel drain timed out; forcing terminal tool states"); - break; - } - - next = stream.next() => { - let Some(event) = next else { break }; - match event { - Ok(ev) => { - if let Some(reason) = handle_event( - ev, - agent, - session_id, - wire_tx, - tokens, - &mut compacted, - &mut reflections, - &mut open_tool_calls, - ) - .await - { - stop = reason; - break; - } - } - Err(e) => { - let msg = e.to_string(); - // Preserve buzz-agent's error taxonomy so the harness's - // JSON-RPC code mapping stays meaningful. - let err = if msg.contains("auth") || msg.contains("401") { - AgentError::LlmAuth(msg) - } else if msg.contains("model") && msg.contains("not found") { - AgentError::LlmModelNotFound(msg) - } else { - AgentError::Llm(msg) - }; - return Err(err); - } - } - } - } - } - - // Anything announced but never resolved gets a synthetic terminal update. - // buzz-agent guaranteed this invariant (`agent.rs:470-477`) because the - // desktop renders an unresolved `tool_call` as a spinner forever. Normally - // the drain above means this list is already empty; it only fires when - // goose failed to unwind inside CANCEL_DRAIN_TIMEOUT. - for tool_call_id in open_tool_calls.drain(..) { - wire::send( - wire_tx, - wire::session_update( - session_id, - json!({ - "sessionUpdate": "tool_call_update", - "toolCallId": tool_call_id, - "status": "failed", - }), - ), - ) - .await; - } - - // 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) + session_id: &str, + tool_call_id: &str, + failed: bool, +) { + wire::send( + wire_tx, + wire::session_update( + session_id, + json!({ + "sessionUpdate": "tool_call_update", + "toolCallId": tool_call_id, + "status": if failed { "failed" } else { "completed" }, + }), + ), + ) + .await; } -/// Translate one `AgentEvent` into ACP notifications. +/// Emit streamed assistant content onto the ACP wire. /// -/// Returns `Some(StopReason)` if the event terminates the turn. -#[allow(clippy::too_many_arguments)] -async fn handle_event( - event: AgentEvent, - agent: &Arc, - session_id: &str, - wire_tx: &WireSender, - tokens: &mut TurnTokens, - compacted: &mut bool, - reflections: &mut usize, - open_tool_calls: &mut Vec, -) -> Option { - match event { - AgentEvent::Message(msg) => { - for content in &msg.content { - emit_content(content, session_id, wire_tx).await; - - if let MessageContent::ToolRequest(req) = content { - open_tool_calls.push(req.id.clone()); - } - - if let MessageContent::ToolResponse(resp) = content { - open_tool_calls.retain(|id| id != &resp.id); - - let failed = match &resp.tool_result { - Ok(r) => r.is_error.unwrap_or(false), - Err(_) => true, - }; - if failed && *reflections < MAX_REFLECTIONS { - *reflections += 1; - agent - .steer( - session_id, - Message::user() - .with_text(ERROR_REFLECTION) - .with_visibility(false, true), - ) - .await; - } - } - } - None - } - - // Usage arrives per provider chunk, already cost-enriched. Accumulate - // rather than overwrite: `MessageUsage` is suppressed when the last - // assistant message has no user-visible content (goose - // `agent.rs:299`), so it is not a reliable sole source. - AgentEvent::Usage(usage) => { - let u = &usage.usage; - if let Some(i) = u.input_tokens { - tokens.input = Some(tokens.input.unwrap_or(0) + i.max(0) as u64); - } - if let Some(o) = u.output_tokens { - tokens.output = Some(tokens.output.unwrap_or(0) + o.max(0) as u64); - } - let cached = u.cache_read_input_tokens.unwrap_or(0).max(0) as u64 - + u.cache_write_input_tokens.unwrap_or(0).max(0) as u64; - if u.cache_read_input_tokens.is_some() || u.cache_write_input_tokens.is_some() { - tokens.cached_input = Some(tokens.cached_input.unwrap_or(0) + cached); - } - if let Some(t) = u.total_tokens { - tokens.total = Some(tokens.total.unwrap_or(0) + t.max(0) as u64); - } - None - } - - // Compaction happened. buzz-agent surfaced this as a `[Context - // 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 - } - - _ => None, - } -} - -/// Map message content onto the nine `sessionUpdate` variants `buzz-acp` -/// recognises (`acp.rs:1528-1627`). Anything else it debug-logs and drops. -async fn emit_content(content: &MessageContent, session_id: &str, wire_tx: &WireSender) { +/// Text and reasoning only — tool lifecycle is emitted by [`crate::tools`], +/// which is the only place that knows when a call actually starts and ends. +pub(crate) async fn emit_content(content: &MessageContent, session_id: &str, wire_tx: &WireSender) { match content { MessageContent::Text(t) if !t.text.is_empty() => { wire::send( @@ -503,84 +166,11 @@ async fn emit_content(content: &MessageContent, session_id: &str, wire_tx: &Wire .await; } - // Tool lifecycle. buzz-agent guaranteed every announced tool reached a - // terminal state, or the desktop UI shows a stuck spinner - // (`agent.rs:470-477`). Goose emits request and response as separate - // messages, so the pairing is by tool-call id. - MessageContent::ToolRequest(req) => { - let (name, raw) = match &req.tool_call { - Ok(call) => ( - call.name.to_string(), - serde_json::to_value(&call.arguments).unwrap_or(Value::Null), - ), - Err(e) => (String::from("unknown"), json!({ "error": e.to_string() })), - }; - wire::send( - wire_tx, - wire::session_update( - session_id, - json!({ - "sessionUpdate": "tool_call", - "toolCallId": req.id, - "title": name, - "kind": "other", - "status": "in_progress", - "rawInput": raw, - }), - ), - ) - .await; - } - - // Frontend tool calls (`load_skill`) are stripped out of the normal - // `ToolRequest` flow by goose and arrive as this variant instead. The - // desktop still needs the announcement: its response comes back as a - // plain `ToolResponse`, and an update for a never-announced id would - // break the announce→terminal pairing. - MessageContent::FrontendToolRequest(req) => { - let (name, raw) = match &req.tool_call { - Ok(call) => ( - call.name.to_string(), - serde_json::to_value(&call.arguments).unwrap_or(Value::Null), - ), - Err(e) => (String::from("unknown"), json!({ "error": e.to_string() })), - }; - wire::send( - wire_tx, - wire::session_update( - session_id, - json!({ - "sessionUpdate": "tool_call", - "toolCallId": req.id, - "title": name, - "kind": "other", - "status": "in_progress", - "rawInput": raw, - }), - ), - ) - .await; - } - - MessageContent::ToolResponse(resp) => { - let failed = match &resp.tool_result { - Ok(r) => r.is_error.unwrap_or(false), - Err(_) => true, - }; - wire::send( - wire_tx, - wire::session_update( - session_id, - json!({ - "sessionUpdate": "tool_call_update", - "toolCallId": resp.id, - "status": if failed { "failed" } else { "completed" }, - }), - ), - ) - .await; - } - + // Tool lifecycle is NOT emitted from here. buzz-agent's loop owns + // it: `crate::tools` announces each call immediately before dispatch + // and emits the terminal update immediately after, including for calls + // a cancelled turn never ran. Emitting from the streamed content too + // would double-announce every tool. _ => {} } } diff --git a/crates/buzz-agent/src/lib.rs b/crates/buzz-agent/src/lib.rs index 07f3941975..aacdf392e6 100644 --- a/crates/buzz-agent/src/lib.rs +++ b/crates/buzz-agent/src/lib.rs @@ -26,7 +26,11 @@ pub mod builtin_client; pub mod config; pub mod hints; pub mod hooks; +pub mod loop_drive; pub mod mesh; +pub mod prompt; +pub mod steer; +pub mod tools; pub mod types; pub mod wire; @@ -75,6 +79,12 @@ struct Session { /// Name of the MCP extension carrying `_Stop`/`_PostCompact`, if any. /// See [`crate::hooks`] for why we dispatch these ourselves. hook_extension: Option, + /// This session's system prompt. Owned here rather than inside goose's + /// `Agent`, because buzz-agent builds the prompt itself each round. + prompt: crate::prompt::SessionPrompt, + /// Mid-turn steer queue. goose's own is `pub(crate)` to `Agent::reply`, + /// which buzz-agent no longer calls. See [`crate::steer`]. + steers: crate::steer::SteerQueue, accumulated_input_tokens: u64, accumulated_output_tokens: u64, /// Subset of `accumulated_input_tokens`, published as @@ -102,7 +112,15 @@ async fn build_agent( cwd: &str, system_prompt: Option<&str>, mcp_servers: &[McpServerStdio], -) -> Result<(Arc, String, Option), AgentError> { +) -> Result< + ( + Arc, + String, + Option, + crate::prompt::SessionPrompt, + ), + AgentError, +> { let session_manager = Arc::new(SessionManager::instance()); let permission_manager = PermissionManager::instance(); @@ -155,9 +173,13 @@ async fn build_agent( // goose/crates/goose/src), and both PRs that would have wired it — // buzz#1290 and goose#9971 — are closed unmerged. Driving the library // directly is what makes Fizz (and `[Base]`) actually arrive. + // buzz-agent owns the loop, so it owns the prompt: goose's own + // `PromptManager` is only reachable from `Agent::reply`, which we no longer + // call. See `crate::prompt`. + let prompt = crate::prompt::SessionPrompt::new(cfg.goose_mode); if let Some(sp) = system_prompt.or(cfg.system_prompt.as_deref()) { if !sp.trim().is_empty() { - agent.override_system_prompt(sp.to_string()).await; + prompt.set_override(sp.to_string()).await; } } @@ -193,9 +215,7 @@ async fn build_agent( // (`prompt_manager.rs:170-198`), so this survives `override_system_prompt`. let (hints, skills) = crate::hints::build_hints_section(std::path::Path::new(cwd)); if !hints.trim().is_empty() { - agent - .extend_system_prompt("buzz_hints".to_string(), hints) - .await; + prompt.add_extra("buzz_hints", hints).await; } // `load_skill` runs in-process, registered the way Maple does it: a @@ -242,13 +262,13 @@ async fn build_agent( } if let Some(ext) = &hook_extension { - agent - .extend_system_prompt("buzz_hook_tools".to_string(), hook_tool_guidance()) + prompt + .add_extra("buzz_hook_tools", hook_tool_guidance()) .await; tracing::info!(extension = %ext, "lifecycle hooks available"); } - Ok((agent, session.id, hook_extension)) + Ok((agent, session.id, hook_extension, prompt)) } /// Keep the model's hands off the lifecycle hooks. @@ -461,7 +481,7 @@ async fn session_new(app: &Arc, id: Value, params: Value, wire_tx: &WireSen .await; } - let (agent, goose_session_id, hook_extension) = + let (agent, goose_session_id, hook_extension, prompt) = match build_agent(&app.cfg, &p.cwd, p.system_prompt.as_deref(), &p.mcp_servers).await { Ok(v) => v, Err(e) => { @@ -503,6 +523,8 @@ async fn session_new(app: &Arc, id: Value, params: Value, wire_tx: &WireSen model_override: None, pending_model: None, hook_extension, + prompt, + steers: crate::steer::SteerQueue::new(), accumulated_input_tokens: 0, accumulated_output_tokens: 0, accumulated_cached_input_tokens: 0, @@ -599,7 +621,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, hook_extension) = { + let (agent, goose_session_id, pending_model, hook_extension, prompt, steers) = { let mut sessions = app.sessions.lock().await; let Some(s) = sessions.get_mut(&p.session_id) else { return wire::send( @@ -629,6 +651,8 @@ async fn session_prompt(app: &Arc, id: Value, params: Value, wire_tx: &Wire s.goose_session_id.clone(), s.pending_model.take(), s.hook_extension.clone(), + s.prompt.clone(), + s.steers.clone(), ) }; @@ -659,14 +683,19 @@ async fn session_prompt(app: &Arc, id: Value, params: Value, wire_tx: &Wire ) .await; - let (result, tokens) = agent::run_turn( - agent, - &goose_session_id, - p.prompt, - app.cfg.max_rounds, - wire_tx, - cancel, - hook_extension.as_deref(), + let (result, tokens) = loop_drive::run_turn( + loop_drive::TurnContext { + agent: &agent, + session_id: &goose_session_id, + wire_tx, + cancel: &cancel, + hook_extension: hook_extension.as_deref(), + max_rounds: app.cfg.max_rounds, + prompt: &prompt, + steers: &steers, + }, + goose_provider_types::conversation::message::Message::user() + .with_text(agent::prompt_to_text(&p.prompt)), ) .await; @@ -835,7 +864,7 @@ async fn steer(app: &Arc, id: Value, params: Value, wire_tx: &WireSender) { .await; } - let (agent, goose_session_id) = { + let steers = { let sessions = app.sessions.lock().await; let Some(s) = sessions.get(&p.session_id) else { return wire::send( @@ -861,15 +890,12 @@ async fn steer(app: &Arc, id: Value, params: Value, wire_tx: &WireSender) { ) .await; } - (s.agent.clone(), s.goose_session_id.clone()) + s.steers.clone() }; let message_id = format!("steer_{}", uuid_like()); - agent - .steer( - &goose_session_id, - goose_provider_types::conversation::message::Message::user().with_text(text), - ) + steers + .push(goose_provider_types::conversation::message::Message::user().with_text(text)) .await; wire::send( diff --git a/crates/buzz-agent/src/loop_drive.rs b/crates/buzz-agent/src/loop_drive.rs new file mode 100644 index 0000000000..9e066fd6b2 --- /dev/null +++ b/crates/buzz-agent/src/loop_drive.rs @@ -0,0 +1,537 @@ +//! buzz-agent's own agent loop, built from goose components. +//! +//! # Who owns what +//! +//! This is the deliberate inversion of the earlier design, where the turn was +//! `goose::agents::Agent::reply` and buzz-agent only translated the resulting +//! event stream onto the ACP wire. **buzz-agent owns the loop.** goose is a +//! parts bin: +//! +//! | concern | goose API used | +//! |---|---| +//! | model call | `Provider::stream` (via `Agent::provider`) | +//! | tool surface | `Agent::list_tools` | +//! | tool execution | `Agent::dispatch_tool_call` | +//! | system prompt | `Agent::build_turn_system_prompt` → `PromptManager` | +//! | compaction | `goose::context_mgmt::{check_if_compaction_needed, compact_messages}` | +//! | conversation store | `SessionManager` (sqlite), read/written by us | +//! +//! Everything that decides *turn shape* stays here: round structure, the +//! `_Stop` veto, `[Reflect]` on tool failure, `[PostCompact]` re-injection, +//! steer draining, cancellation semantics, and the ACP `session/update` +//! emission. +//! +//! # Why own the loop +//! +//! `Agent::reply` bakes in goose's own turn policy — its compaction trigger, +//! its stop-hook handling, its steer drain points, its max-turns rule. Every +//! buzz-specific behaviour then has to be smuggled in around the edges +//! (see the old `run_turn`: an outer loop purely to re-enter `reply()` for the +//! `_Stop` veto, and `[Reflect]` delivered as a *steer* because the tool +//! result itself was out of reach). Driving the components directly removes +//! the smuggling: the veto is a branch, `[Reflect]` is appended to the tool +//! result where buzz-agent originally put it. +//! +//! # Round structure +//! +//! One round is: build prompt+tools → stream from the provider → persist the +//! assistant message → if it asked for tools, run them, append results, and +//! loop; otherwise the turn wants to end, at which point `_Stop` may veto it. +//! +//! Compaction is checked at the top of each round rather than mid-stream, so a +//! rewrite of history can never race a partially-consumed response. + +use std::sync::Arc; + +use futures::StreamExt; +use tokio_util::sync::CancellationToken; + +use goose::agents::Agent; +use goose::session::session_manager::{Session, SessionManager}; +use goose_provider_types::conversation::message::{Message, MessageContent, ToolRequest}; +use goose_provider_types::conversation::Conversation; + +use crate::types::{AgentError, StopReason}; +use crate::wire::WireSender; + +/// Appended to every failed tool result so the model diagnoses the failure +/// instead of blindly retrying it. +/// +/// This is buzz-agent's original behaviour, restored. Under `Agent::reply` the +/// tool result was owned by goose and unreachable, so the text had to be +/// delivered as a steer instead; owning the loop means we can put it back on +/// the result where the model reads it in context. +const ERROR_REFLECTION: &str = + "[Reflect] Before retrying, identify the cause and change your approach."; + +/// Cap on reflections per turn, so a tool failing in a loop cannot flood the +/// conversation. +const MAX_REFLECTIONS: usize = 8; + +/// Hard cap on consecutive `_Stop` vetoes, mirroring goose's own +/// `stop_hook_block_cap`. A tool that always objects must not trap the turn. +pub const MAX_STOP_BLOCKS: u32 = 3; + +/// Rounds allowed when the caller sets no limit. +/// +/// buzz-agent's loop was previously bounded by goose's `max_turns`. Owning the +/// loop means owning the bound, and an unbounded loop against a model that +/// tool-calls forever is a runaway spend. +const DEFAULT_MAX_ROUNDS: u32 = 1000; + +/// Outcome of a single round, as decided by this loop. +enum Round { + /// The model asked for tools; they ran and their results are in history. + Continued, + /// The model produced a final answer. + WantsToEnd, + /// A terminal condition the loop must surrender to. + Stopped(StopReason), +} + +/// Everything a round needs that does not change between rounds. +pub struct TurnContext<'a> { + pub agent: &'a Arc, + pub session_id: &'a str, + pub wire_tx: &'a WireSender, + pub cancel: &'a CancellationToken, + pub hook_extension: Option<&'a str>, + pub max_rounds: Option, + /// This session's system prompt. buzz-agent owns it because goose's own + /// `PromptManager` is only reachable from `Agent::reply`. + pub prompt: &'a crate::prompt::SessionPrompt, + /// Messages injected mid-turn by `session/steer`. Drained at the round + /// boundary; a pending steer also blocks the end of the turn. + pub steers: &'a crate::steer::SteerQueue, +} + +/// Drive one `session/prompt` turn to completion. +/// +/// Returns the ACP stop reason plus the turn's token counts. The caller emits +/// `usage_update` *before* the `session/prompt` response — that ordering is +/// load-bearing for kind-44200 metrics. +pub async fn run_turn( + ctx: TurnContext<'_>, + prompt: Message, +) -> (Result, super::agent::TurnTokens) { + let mut tokens = super::agent::TurnTokens::default(); + let sessions = SessionManager::instance(); + + // The user's prompt is persisted before the first inference so a resumed + // or inspected session sees the same history the model saw. + if let Err(e) = sessions.add_message(ctx.session_id, &prompt).await { + return (Err(AgentError::Llm(format!("persist prompt: {e}"))), tokens); + } + + let max_rounds = ctx.max_rounds.unwrap_or(DEFAULT_MAX_ROUNDS); + let mut rounds = 0u32; + let mut stop_blocks = 0u32; + let mut reflections = 0usize; + + loop { + if ctx.cancel.is_cancelled() { + return (Ok(StopReason::Cancelled), tokens); + } + if rounds >= max_rounds { + tracing::warn!(rounds, "max rounds reached; ending turn"); + return (Ok(StopReason::MaxTurnRequests), tokens); + } + rounds += 1; + + // Steers land here, at the round boundary — never mid-inference, + // where they would race a partially streamed response. + for message in ctx.steers.drain().await { + if let Err(e) = sessions.add_message(ctx.session_id, &message).await { + return (Err(AgentError::Llm(format!("persist steer: {e}"))), tokens); + } + } + + // Compaction at the round boundary, never mid-stream. + if let Err(e) = maybe_compact(&ctx, &sessions).await { + tracing::warn!(error = %e, "compaction failed; continuing uncompacted"); + } + + match round(&ctx, &sessions, &mut tokens, &mut reflections).await { + Err(e) => return (Err(e), tokens), + Ok(Round::Continued) => continue, + Ok(Round::Stopped(reason)) => return (Ok(reason), tokens), + Ok(Round::WantsToEnd) => { + // A steer that arrived while the model was finishing must be + // answered, not dropped — so it also blocks the end of turn. + if ctx.steers.is_pending().await && !ctx.cancel.is_cancelled() { + continue; + } + + // The turn wants to end. `_Stop` gets a veto. + match stop_veto(&ctx, &sessions, &mut stop_blocks).await { + Some(objection) => { + let message = Message::user() + .with_text(format!("[Stop] {objection}")) + .with_visibility(false, true); + if let Err(e) = sessions.add_message(ctx.session_id, &message).await { + return (Err(AgentError::Llm(format!("persist stop: {e}"))), tokens); + } + continue; + } + None => { + // A steer arriving after this point belongs to no run. + ctx.steers.clear().await; + return (Ok(StopReason::EndTurn), tokens); + } + } + } + } + } +} + +/// One inference plus, if the model asked for them, one batch of tool calls. +async fn round( + ctx: &TurnContext<'_>, + sessions: &SessionManager, + tokens: &mut super::agent::TurnTokens, + reflections: &mut usize, +) -> Result { + let session = load_session(sessions, ctx.session_id).await?; + let conversation = session.conversation.clone().unwrap_or_default(); + + let assistant = match infer(ctx, &session, &conversation, tokens).await? { + Some(message) => message, + // A provider that returns nothing is not a turn we can continue. + None => return Ok(Round::Stopped(StopReason::EndTurn)), + }; + + sessions + .add_message(ctx.session_id, &assistant) + .await + .map_err(|e| AgentError::Llm(format!("persist assistant: {e}")))?; + + let requests: Vec = assistant + .content + .iter() + .filter_map(|content| match content { + MessageContent::ToolRequest(req) => Some(req.clone()), + _ => None, + }) + .collect(); + + if requests.is_empty() { + return Ok(Round::WantsToEnd); + } + + if ctx.cancel.is_cancelled() { + // Announced tool calls must still reach a terminal state or the + // desktop renders a spinner forever. + let results = crate::tools::cancelled_results(&requests); + for request in &requests { + crate::agent::emit_tool_call_update(ctx.wire_tx, ctx.session_id, &request.id, true) + .await; + } + persist_tool_results(ctx, sessions, results).await?; + return Ok(Round::Stopped(StopReason::Cancelled)); + } + + // Tool arguments feed goose's subdirectory-hint tracker, so the next + // round's prompt picks up an AGENTS.md in a directory we just touched. + for request in &requests { + if let Ok(call) = &request.tool_call { + ctx.prompt + .record_tool_arguments(&call.arguments, &session.working_dir) + .await; + } + } + + let results = crate::tools::execute( + ctx.agent, + &session, + ctx.wire_tx, + ctx.cancel, + &requests, + reflections, + ) + .await; + persist_tool_results(ctx, sessions, results).await?; + + Ok(Round::Continued) +} + +/// Stream one assistant response, emitting chunks onto the ACP wire as they +/// arrive and accumulating usage. +async fn infer( + ctx: &TurnContext<'_>, + session: &Session, + conversation: &Conversation, + tokens: &mut super::agent::TurnTokens, +) -> Result, AgentError> { + let provider = ctx + .agent + .provider() + .await + .map_err(|e| AgentError::Llm(e.to_string()))?; + let model_config = ctx + .agent + .model_config_for_session(ctx.session_id) + .await + .map_err(|e| AgentError::LlmModelNotFound(e.to_string()))?; + + let system_prompt = ctx.prompt.build(&session.working_dir).await; + let tools = ctx.agent.list_tools(ctx.session_id, None).await; + + let mut stream = provider + .stream( + &model_config, + &system_prompt, + conversation.messages(), + &tools, + ) + .await + .map_err(|e| classify_provider_error(&e.to_string()))?; + + let mut accumulated: Option = None; + let mut keepalive = tokio::time::interval(crate::agent::KEEPALIVE_INTERVAL); + keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + keepalive.tick().await; // first tick is immediate; discard it + + loop { + tokio::select! { + // Keep the harness idle clock alive during provider silence. + _ = keepalive.tick() => { + crate::agent::emit_keepalive(ctx.session_id, ctx.wire_tx).await; + } + + // Cancelling the model call is safe to do abruptly: unlike tool + // dispatch there is no child process to notify, and the partial + // response is discarded rather than persisted. + _ = ctx.cancel.cancelled() => { + return Ok(accumulated); + } + + next = stream.next() => { + let Some(item) = next else { break }; + let (message, usage) = item.map_err(|e| classify_provider_error(&e.to_string()))?; + + if let Some(usage) = usage { + accumulate_usage(tokens, &usage); + } + + let Some(chunk) = message else { continue }; + for content in &chunk.content { + crate::agent::emit_content(content, ctx.session_id, ctx.wire_tx).await; + } + accumulated = Some(match accumulated { + None => chunk, + Some(mut prev) => { + merge_chunk(&mut prev, chunk); + prev + } + }); + } + } + } + + Ok(accumulated) +} + +/// Fold a streamed chunk into the message being accumulated. +/// +/// Text arrives fragmented and must be coalesced or the persisted history +/// contains one message per token; tool calls arrive whole and are appended. +fn merge_chunk(target: &mut Message, chunk: Message) { + for content in chunk.content { + match (target.content.last_mut(), &content) { + (Some(MessageContent::Text(last)), MessageContent::Text(new)) => { + last.text.push_str(&new.text); + } + (Some(MessageContent::Thinking(last)), MessageContent::Thinking(new)) => { + last.thinking.push_str(&new.thinking); + } + _ => target.content.push(content), + } + } +} + +fn accumulate_usage( + tokens: &mut super::agent::TurnTokens, + usage: &goose::providers::base::ProviderUsage, +) { + let u = &usage.usage; + if let Some(i) = u.input_tokens { + tokens.input = Some(tokens.input.unwrap_or(0) + i.max(0) as u64); + } + if let Some(o) = u.output_tokens { + tokens.output = Some(tokens.output.unwrap_or(0) + o.max(0) as u64); + } + if u.cache_read_input_tokens.is_some() || u.cache_write_input_tokens.is_some() { + let cached = u.cache_read_input_tokens.unwrap_or(0).max(0) as u64 + + u.cache_write_input_tokens.unwrap_or(0).max(0) as u64; + tokens.cached_input = Some(tokens.cached_input.unwrap_or(0) + cached); + } + if let Some(t) = u.total_tokens { + tokens.total = Some(tokens.total.unwrap_or(0) + t.max(0) as u64); + } +} + +/// Persist tool results as a single user message, matching the shape every +/// provider wire format expects (one result per outstanding request). +async fn persist_tool_results( + ctx: &TurnContext<'_>, + sessions: &SessionManager, + results: Vec<( + String, + goose_provider_types::conversation::message::ToolResult, + )>, +) -> Result<(), AgentError> { + if results.is_empty() { + return Ok(()); + } + let mut message = Message::user(); + for (id, result) in results { + message = message.with_tool_response(id, result); + } + sessions + .add_message(ctx.session_id, &message) + .await + .map_err(|e| AgentError::Llm(format!("persist tool results: {e}"))) +} + +/// Compact history when it is close to the context limit, then re-inject the +/// state `_PostCompact` reports so a todo list survives the truncation. +async fn maybe_compact(ctx: &TurnContext<'_>, sessions: &SessionManager) -> anyhow::Result<()> { + let session = sessions.get_session(ctx.session_id, true).await?; + let Some(conversation) = session.conversation.clone() else { + return Ok(()); + }; + let provider = ctx.agent.provider().await?; + + if !goose::context_mgmt::check_if_compaction_needed( + provider.as_ref(), + &conversation, + None, + &session, + ) + .await? + { + return Ok(()); + } + + let model_config = ctx.agent.model_config_for_session(ctx.session_id).await?; + let result = goose::context_mgmt::compact_messages( + provider.as_ref(), + &model_config, + ctx.session_id, + &conversation, + false, + ) + .await?; + + sessions + .replace_conversation(ctx.session_id, &result.conversation) + .await?; + tracing::info!(target: "buzz_agent::compaction", "history compacted"); + + // `_PostCompact` state is what buzz-agent's handoff existed to preserve. + if let Some(extension) = ctx.hook_extension { + let session = sessions.get_session(ctx.session_id, true).await?; + if let Some(state) = crate::hooks::post_compact_state(ctx.agent, &session, extension).await + { + let message = Message::user() + .with_text(format!("[PostCompact] {state}")) + .with_visibility(false, true); + sessions.add_message(ctx.session_id, &message).await?; + } + } + + Ok(()) +} + +/// Ask `_Stop` whether the turn may end. `Some(objection)` means keep working. +async fn stop_veto( + ctx: &TurnContext<'_>, + sessions: &SessionManager, + stop_blocks: &mut u32, +) -> Option { + if ctx.cancel.is_cancelled() { + return None; + } + let extension = ctx.hook_extension?; + + if *stop_blocks >= MAX_STOP_BLOCKS { + tracing::warn!(blocks = *stop_blocks, "_Stop veto cap reached; ending turn"); + return None; + } + + let session = sessions.get_session(ctx.session_id, false).await.ok()?; + let objection = crate::hooks::stop_objection(ctx.agent, &session, extension).await?; + + *stop_blocks += 1; + tracing::info!(blocks = *stop_blocks, "_Stop hook vetoed end of turn"); + Some(objection) +} + +async fn load_session(sessions: &SessionManager, id: &str) -> Result { + sessions + .get_session(id, true) + .await + .map_err(|e| AgentError::Llm(format!("load session: {e}"))) +} + +/// Preserve buzz-agent's error taxonomy so the harness's JSON-RPC code mapping +/// stays meaningful. +fn classify_provider_error(msg: &str) -> AgentError { + let lower = msg.to_lowercase(); + if lower.contains("auth") || lower.contains("401") { + AgentError::LlmAuth(msg.to_string()) + } else if lower.contains("model") && lower.contains("not found") { + AgentError::LlmModelNotFound(msg.to_string()) + } else { + AgentError::Llm(msg.to_string()) + } +} + +/// Text appended to a failed tool result, exposed for the tool module. +pub(crate) fn reflection_text() -> &'static str { + ERROR_REFLECTION +} + +pub(crate) fn max_reflections() -> usize { + MAX_REFLECTIONS +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn merge_coalesces_consecutive_text() { + let mut target = Message::assistant().with_text("hel"); + merge_chunk(&mut target, Message::assistant().with_text("lo")); + assert_eq!(target.as_concat_text(), "hello"); + assert_eq!(target.content.len(), 1, "text must coalesce, not append"); + } + + #[test] + fn merge_appends_distinct_content_kinds() { + let mut target = Message::assistant().with_text("thinking about it"); + merge_chunk( + &mut target, + Message::assistant() + .with_tool_response("id", Ok(rmcp::model::CallToolResult::success(vec![]))), + ); + assert_eq!(target.content.len(), 2); + } + + #[test] + fn provider_errors_keep_buzz_agents_taxonomy() { + assert!(matches!( + classify_provider_error("401 Unauthorized"), + AgentError::LlmAuth(_) + )); + assert!(matches!( + classify_provider_error("model gpt-9 not found"), + AgentError::LlmModelNotFound(_) + )); + assert!(matches!( + classify_provider_error("connection reset"), + AgentError::Llm(_) + )); + } +} diff --git a/crates/buzz-agent/src/prompt.rs b/crates/buzz-agent/src/prompt.rs new file mode 100644 index 0000000000..514be7d41f --- /dev/null +++ b/crates/buzz-agent/src/prompt.rs @@ -0,0 +1,108 @@ +//! buzz-agent's system prompt, assembled with goose's `PromptManager`. +//! +//! `Agent::prompt_manager` is `pub(super)` in goose, and `Agent::reply` is the +//! only thing that reads it. Since buzz-agent owns the loop it also owns the +//! prompt: we keep our own `PromptManager` and call it at the top of each +//! round, which is exactly what goose's inference step does internally. +//! +//! Owning it is not just a workaround for the visibility — it is the seam that +//! makes the persona work at all. goose's ACP server never reads +//! `systemPrompt`, so an agent driven over plain ACP silently loses Fizz and +//! the `[Base]` operating manual. Building the prompt here puts them in front +//! of the model. + +use std::path::Path; +use std::sync::Arc; + +use goose::agents::PromptManager; +use goose::config::GooseMode; +use tokio::sync::Mutex; + +/// The system prompt for one session. +/// +/// Cloneable so a turn can hold it without borrowing the session map. +#[derive(Clone)] +pub struct SessionPrompt { + manager: Arc>, + mode: GooseMode, +} + +impl SessionPrompt { + pub fn new(mode: GooseMode) -> Self { + Self { + manager: Arc::new(Mutex::new(PromptManager::new())), + mode, + } + } + + /// Replace goose's base template with the harness-supplied persona. + pub async fn set_override(&self, template: String) { + self.manager + .lock() + .await + .set_system_prompt_override(template); + } + + /// Append a keyed section (AGENTS.md hints, the skill index, hook + /// guidance). Extras survive an override — goose appends them to whichever + /// base template is in force. + pub async fn add_extra(&self, key: &str, instruction: String) { + self.manager + .lock() + .await + .add_system_prompt_extra(key.to_string(), instruction); + } + + /// Render the prompt for one inference. + /// + /// Called per round rather than once per session because + /// `build_system_prompt` also folds in subdirectory hints discovered from + /// the tool arguments seen so far — a prompt built once at session start + /// would never pick them up. + pub async fn build(&self, working_dir: &Path) -> String { + self.manager + .lock() + .await + .build_system_prompt(working_dir, Vec::new(), self.mode) + } + + /// Feed tool arguments to the subdirectory-hint tracker, so entering a new + /// part of the tree pulls in that directory's AGENTS.md on the next round. + pub async fn record_tool_arguments( + &self, + arguments: &Option>, + working_dir: &Path, + ) { + self.manager + .lock() + .await + .record_tool_arguments(arguments, working_dir); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn override_replaces_the_base_template() { + let prompt = SessionPrompt::new(GooseMode::Auto); + prompt.set_override("You are Fizz.".to_string()).await; + let built = prompt.build(Path::new(".")).await; + assert!(built.contains("You are Fizz.")); + } + + #[tokio::test] + async fn extras_survive_an_override() { + // The persona arrives as an override and the hints as an extra; losing + // either one is a silent failure the harness cannot detect. + let prompt = SessionPrompt::new(GooseMode::Auto); + prompt.set_override("You are Fizz.".to_string()).await; + prompt + .add_extra("buzz_hints", "Always ship the sprocket first.".to_string()) + .await; + let built = prompt.build(Path::new(".")).await; + assert!(built.contains("You are Fizz.")); + assert!(built.contains("Always ship the sprocket first.")); + } +} diff --git a/crates/buzz-agent/src/steer.rs b/crates/buzz-agent/src/steer.rs new file mode 100644 index 0000000000..a9a09922ea --- /dev/null +++ b/crates/buzz-agent/src/steer.rs @@ -0,0 +1,86 @@ +//! Mid-turn steering, owned by buzz-agent. +//! +//! Steering is a message injected into a turn already in flight, without +//! cancelling it. goose has an equivalent queue on its `Agent`, but +//! `drain_pending_steers` / `has_pending_steers` are `pub(crate)` — only +//! `Agent::reply` can consume them. Since buzz-agent owns the loop, it owns +//! the queue too. +//! +//! The semantics match what buzz-acp expects and what goose does internally: +//! +//! * the message is queued, never applied mid-inference — it lands at the +//! **round boundary**, so it cannot corrupt a partially streamed response; +//! * a pending steer **prevents the turn from ending**, so a user who steers +//! just as the model finishes is not silently ignored; +//! * steering never restarts or cancels the turn. + +use std::collections::VecDeque; +use std::sync::Arc; + +use goose_provider_types::conversation::message::Message; +use tokio::sync::Mutex; + +/// A session's pending steer messages. +/// +/// Cloneable and shared: `session/steer` pushes from the JSON-RPC dispatch +/// task while the turn drains from the loop task. +#[derive(Clone, Default)] +pub struct SteerQueue { + inner: Arc>>, +} + +impl SteerQueue { + pub fn new() -> Self { + Self::default() + } + + pub async fn push(&self, message: Message) { + self.inner.lock().await.push_back(message); + } + + pub async fn is_pending(&self) -> bool { + !self.inner.lock().await.is_empty() + } + + /// Take everything queued so far. + /// + /// Draining and appending is a single logical step for the loop: a steer + /// that arrives *during* the drain stays queued for the next boundary + /// rather than being lost. + pub async fn drain(&self) -> Vec { + self.inner.lock().await.drain(..).collect() + } + + /// Drop anything queued. Used when a turn ends so a late steer cannot leak + /// into the next one. + pub async fn clear(&self) { + self.inner.lock().await.clear(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn drain_takes_everything_in_order() { + let queue = SteerQueue::new(); + queue.push(Message::user().with_text("first")).await; + queue.push(Message::user().with_text("second")).await; + assert!(queue.is_pending().await); + + let drained = queue.drain().await; + assert_eq!(drained.len(), 2); + assert_eq!(drained[0].as_concat_text(), "first"); + assert_eq!(drained[1].as_concat_text(), "second"); + assert!(!queue.is_pending().await, "drain must empty the queue"); + } + + #[tokio::test] + async fn clear_discards_a_late_steer() { + let queue = SteerQueue::new(); + queue.push(Message::user().with_text("too late")).await; + queue.clear().await; + assert!(!queue.is_pending().await); + } +} diff --git a/crates/buzz-agent/src/tools.rs b/crates/buzz-agent/src/tools.rs new file mode 100644 index 0000000000..59aaa7a2bb --- /dev/null +++ b/crates/buzz-agent/src/tools.rs @@ -0,0 +1,209 @@ +//! Tool execution for buzz-agent's own loop. +//! +//! goose's `Agent::dispatch_tool_call` resolves the prefixed tool name, runs +//! the MCP call, and applies goose's own pre/post hooks — so we use it rather +//! than reaching into `ExtensionManager` directly. What we keep is the policy +//! around it: parallel dispatch, the announce→terminal wire invariant, and +//! `[Reflect]` appended to failed results. + +use std::sync::Arc; + +use futures::stream::{FuturesUnordered, StreamExt}; +use rmcp::model::{CallToolResult, ContentBlock, ErrorData}; +use tokio_util::sync::CancellationToken; + +use goose::agents::Agent; +use goose::session::session_manager::Session; +use goose_provider_types::conversation::message::{ToolRequest, ToolResult}; + +use crate::wire::WireSender; + +/// One tool's outcome, keyed by the request id the model used. +pub type Outcome = (String, ToolResult); + +/// Run every tool the model asked for. +/// +/// Calls are dispatched concurrently — a turn that reads three files should +/// not pay for them serially — but results are returned in **request order**, +/// because providers require tool results to line up with the requests that +/// produced them. +pub async fn execute( + agent: &Arc, + session: &Session, + wire_tx: &WireSender, + cancel: &CancellationToken, + requests: &[ToolRequest], + reflections: &mut usize, +) -> Vec { + let mut pending = FuturesUnordered::new(); + for (index, request) in requests.iter().enumerate() { + pending.push(run_one(agent, session, wire_tx, cancel, index, request)); + } + + let mut collected: Vec> = (0..requests.len()).map(|_| None).collect(); + while let Some((index, outcome)) = pending.next().await { + collected[index] = Some(outcome); + } + + let mut results = Vec::with_capacity(collected.len()); + for outcome in collected.into_iter().flatten() { + let (id, result) = outcome; + let result = reflect_on_failure(result, reflections); + emit_terminal(wire_tx, session, &id, &result).await; + results.push((id, result)); + } + results +} + +/// Synthesise results for tools that were never run because the turn was +/// cancelled first. +/// +/// Every announced tool call must reach a terminal state; the desktop renders +/// an unresolved one as a spinner that never stops. +pub fn cancelled_results(requests: &[ToolRequest]) -> Vec { + requests + .iter() + .map(|request| { + ( + request.id.clone(), + Ok(CallToolResult::error(vec![ContentBlock::text( + "Tool call was cancelled before execution", + )])), + ) + }) + .collect() +} + +async fn run_one( + agent: &Arc, + session: &Session, + wire_tx: &WireSender, + cancel: &CancellationToken, + index: usize, + request: &ToolRequest, +) -> (usize, Outcome) { + let id = request.id.clone(); + + // A tool call the provider itself failed to parse never reaches MCP. + let call = match &request.tool_call { + Ok(call) => call.clone(), + Err(e) => return (index, (id, Err(e.clone()))), + }; + + crate::agent::emit_tool_call(wire_tx, &session.id, &id, &call).await; + + let (_request_id, dispatched) = agent + .dispatch_tool_call(call, id.clone(), Some(cancel.clone()), session) + .await; + + let result = match dispatched { + Ok(call_result) => call_result.result.await, + Err(e) => Err(e), + }; + + (index, (id, result)) +} + +/// Append `[Reflect]` to a failed tool result. +/// +/// buzz-agent's original behaviour: the model reads the reflection in the same +/// content block as the failure, rather than as a separate steer message +/// arriving at an unrelated point in the conversation. Capped so a tool +/// failing in a loop cannot flood the context. +fn reflect_on_failure( + result: ToolResult, + reflections: &mut usize, +) -> ToolResult { + let failed = match &result { + Ok(out) => out.is_error.unwrap_or(false), + Err(_) => true, + }; + if !failed || *reflections >= crate::loop_drive::max_reflections() { + return result; + } + *reflections += 1; + + match result { + Ok(mut out) => { + out.content + .push(ContentBlock::text(crate::loop_drive::reflection_text())); + Ok(out) + } + // An `ErrorData` has no content to append to, so the reflection is + // folded into its message instead of being dropped. + Err(e) => Err(ErrorData::new( + e.code, + format!("{} {}", e.message, crate::loop_drive::reflection_text()), + e.data, + )), + } +} + +async fn emit_terminal( + wire_tx: &WireSender, + session: &Session, + id: &str, + result: &ToolResult, +) { + let failed = match result { + Ok(out) => out.is_error.unwrap_or(false), + Err(_) => true, + }; + crate::agent::emit_tool_call_update(wire_tx, &session.id, id, failed).await; +} + +#[cfg(test)] +mod tests { + use super::*; + + fn request(id: &str) -> ToolRequest { + ToolRequest { + id: id.to_string(), + tool_call: Ok(rmcp::model::CallToolRequestParams::new("x")), + metadata: None, + tool_meta: None, + } + } + + #[test] + fn cancellation_still_terminates_every_announced_call() { + // The desktop spins forever on an unresolved tool call, so a cancelled + // turn must still produce one result per request. + let requests = vec![request("a"), request("b")]; + let results = cancelled_results(&requests); + assert_eq!(results.len(), 2); + assert_eq!(results[0].0, "a"); + assert_eq!(results[1].0, "b"); + } + + #[test] + fn reflection_is_appended_to_a_failed_result() { + let mut reflections = 0; + let failed = Ok(CallToolResult::error(vec![ContentBlock::text("boom")])); + let out = reflect_on_failure(failed, &mut reflections).unwrap(); + let text: String = out + .content + .iter() + .filter_map(|c| c.as_text().map(|t| t.text.clone())) + .collect(); + assert!(text.contains("[Reflect]")); + assert_eq!(reflections, 1); + } + + #[test] + fn success_is_left_alone() { + let mut reflections = 0; + let ok = Ok(CallToolResult::success(vec![ContentBlock::text("fine")])); + let out = reflect_on_failure(ok, &mut reflections).unwrap(); + assert_eq!(out.content.len(), 1); + assert_eq!(reflections, 0); + } + + #[test] + fn reflection_stops_at_the_cap() { + let mut reflections = crate::loop_drive::max_reflections(); + let failed = Ok(CallToolResult::error(vec![ContentBlock::text("boom")])); + let out = reflect_on_failure(failed, &mut reflections).unwrap(); + assert_eq!(out.content.len(), 1, "capped: no reflection appended"); + } +}