From ef482e886dcc6a451f7bbf82bff2d1f8519f38aa Mon Sep 17 00:00:00 2001 From: Jimmy <1fe240cd1a8cf775f6f3060f115e5a303181f3abf28ad4cb0c2515f4a02b36a8@meshllm.communities.buzz.xyz> Date: Tue, 1 Sep 2026 15:57:34 +1000 Subject: [PATCH] Reduce Goose embedding surface in buzz-agent Replace the monolithic Agent extension manager with a bounded direct RMCP registry while preserving lifecycle hooks, permission gating, skill loading, process-tree cleanup, and tool-result limits. Keep turn state in Buzz memory instead of a Goose sqlite session. Co-authored-by: Michael Neale Signed-off-by: Jimmy <1fe240cd1a8cf775f6f3060f115e5a303181f3abf28ad4cb0c2515f4a02b36a8@meshllm.communities.buzz.xyz> --- crates/buzz-agent/src/hooks.rs | 79 +- crates/buzz-agent/src/lib.rs | 232 +---- crates/buzz-agent/src/loop_drive.rs | 38 +- crates/buzz-agent/src/mcp.rs | 1155 ++++++++++++++++++++++++ crates/buzz-agent/src/ops.rs | 110 ++- crates/buzz-agent/src/session_store.rs | 107 --- crates/buzz-agent/src/skills.rs | 46 + crates/buzz-agent/src/tools.rs | 39 +- crates/buzz-agent/src/turn_state.rs | 153 ++-- 9 files changed, 1395 insertions(+), 564 deletions(-) create mode 100644 crates/buzz-agent/src/mcp.rs delete mode 100644 crates/buzz-agent/src/session_store.rs diff --git a/crates/buzz-agent/src/hooks.rs b/crates/buzz-agent/src/hooks.rs index 3f2e51a499..39f1fb33a8 100644 --- a/crates/buzz-agent/src/hooks.rs +++ b/crates/buzz-agent/src/hooks.rs @@ -34,9 +34,7 @@ use std::sync::Arc; -use goose::agents::Agent; -use goose::session::session_manager::Session; -use rmcp::model::CallToolRequestParams; +use crate::mcp::McpRegistry; /// Default cap on consecutive vetoes, mirroring goose's own /// `stop_hook_block_cap` (`agent.rs:108-119`). A tool that always objects must @@ -69,12 +67,8 @@ fn qualified(extension: &str, tool: &str) -> String { /// 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?; +pub async fn stop_objection(mcp: &Arc, extension: &str) -> Option { + let text = call_hook(mcp, extension, "_Stop").await?; let trimmed = text.trim(); // Empty response is the documented "no objection" signal // (`buzz-dev-mcp/src/todo.rs:100`). @@ -85,65 +79,40 @@ pub async fn stop_objection( /// /// 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?; +pub async fn post_compact_state(mcp: &Arc, extension: &str) -> Option { + let text = call_hook(mcp, 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, +async fn call_hook(mcp: &Arc, extension: &str, tool: &str) -> Option { + let cancel = tokio_util::sync::CancellationToken::new(); + let out = match mcp + .call_rmcp( + &qualified(extension, tool), + &format!("buzz_hook_{tool}"), + &serde_json::Value::Object(serde_json::Map::new()), + &cancel, ) - .await; - - let call = match result { - Ok(c) => c, - Err(e) => { - tracing::debug!(tool, error = %e, "hook tool unavailable"); + .await + { + Ok(out) => out, + Err(error) => { + tracing::debug!(tool, error = ?error, "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) + Some( + out.content + .iter() + .filter_map(|content| content.as_text().map(|text| text.text.as_str())) + .collect(), + ) } #[cfg(test)] diff --git a/crates/buzz-agent/src/lib.rs b/crates/buzz-agent/src/lib.rs index e4e9b3b971..8aff009334 100644 --- a/crates/buzz-agent/src/lib.rs +++ b/crates/buzz-agent/src/lib.rs @@ -24,12 +24,12 @@ pub mod agent; pub mod config; pub mod hooks; pub mod loop_drive; +pub mod mcp; pub mod model; pub mod ops; pub mod permission; pub mod prompt; pub mod provider; -pub mod session_store; pub mod skills; pub mod steer; pub mod tools; @@ -50,10 +50,6 @@ use serde_json::{json, Value}; use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; -use goose::agents::{Agent, AgentConfig, ExtensionConfig, GoosePlatform}; -use goose::config::PermissionManager; -use goose::session::session_manager::SessionType; - use config::{Config, MAX_PROMPT_BYTES, MAX_SYSTEM_PROMPT_BYTES, PROTOCOL_VERSION}; use types::McpServerStdio; use wire::{ @@ -64,7 +60,7 @@ use wire::{ /// One live ACP session, wrapping a Goose `Agent`. struct Session { - agent: Arc, + mcp: Arc, /// Session id, shared with goose so tool dispatch and provider request /// attribution agree on a name. No database sits behind it: the /// conversation lives in `history` and in the turn's own @@ -142,7 +138,7 @@ async fn build_agent( mcp_servers: &[McpServerStdio], ) -> Result< ( - Arc, + Arc, String, crate::model::SessionModel, Option, @@ -150,226 +146,38 @@ async fn build_agent( ), AgentError, > { - let session_manager = crate::session_store::session_manager(); - let permission_manager = PermissionManager::instance(); - - let agent = Agent::with_config(AgentConfig::new( - session_manager.clone(), - permission_manager, - None, - cfg.goose_mode, - // Session naming is a Goose UX affordance; the harness names sessions. - true, - GoosePlatform::GooseCli, - )); - - // One session row, and it is genuinely required -- I tried removing it. - // - // `Agent::add_extension_inner` (goose agent.rs:1465) resolves the - // extension's working directory by loading the session from the store, and - // fails setup outright if the row is absent: "Failed to get session ... - // Session not found". Every MCP server buzz mounts goes through that path, - // so without the row an agent has no tools at all. - // - // Nothing else needs it. The conversation stays in memory - // (`crate::turn_state`) and the provider config stays on buzz's own handle - // (`crate::model`), so a conversation of any length adds zero further - // rows. Removing this last one means an upstream change: goose would have - // to take the working directory as an argument rather than reading it back - // out of the session. - let session = session_manager - .create_session( - std::path::PathBuf::from(cwd), - "buzz-agent".to_string(), - SessionType::Acp, - cfg.goose_mode, - ) - .await - .map_err(|e| AgentError::Llm(format!("session create: {e}")))?; - let session_id = session.id.clone(); - - // Provider/model names come from the `GOOSE_*` variables `Config::from_env` - // projected out of the `BUZZ_AGENT_*` ones; goose's registry owns base-url - // and credential resolution from there. let provider_name = std::env::var("GOOSE_PROVIDER").unwrap_or_else(|_| "openai".to_string()); let model_name = cfg .model .clone() .or_else(|| std::env::var("GOOSE_MODEL").ok()) .ok_or_else(|| AgentError::Llm("no model configured".into()))?; - let provider = build_provider(&provider_name).await?; let model_config = goose::model_config::model_config_from_user_config(&provider_name, &model_name) - .map_err(|e| AgentError::LlmModelNotFound(e.to_string()))?; - - // Held here rather than pushed into the `Agent` with `update_provider`: - // that call writes the provider config to goose's sqlite session row, and - // buzz reads the provider and config back through `SessionModel` anyway. + .map_err(|error| AgentError::LlmModelNotFound(error.to_string()))?; let model = crate::model::SessionModel::new(provider, model_config, model_name); - // Persona. This is the seam that does NOT work over plain ACP: goose's own - // ACP server never reads `systemPrompt` (rg finds zero hits in - // 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() { - prompt.set_override(sp.to_string()).await; + if let Some(system_prompt) = system_prompt.or(cfg.system_prompt.as_deref()) { + if !system_prompt.trim().is_empty() { + prompt.set_override(system_prompt.to_string()).await; } } - for server in mcp_servers { - let envs: HashMap = server - .env - .iter() - .map(|e| (e.name.clone(), e.value.clone())) - .collect(); - - let ext = ExtensionConfig::Stdio { - name: server.name.clone(), - description: String::new(), - cmd: server.command.clone(), - args: server.args.clone(), - envs: goose::agents::extension::Envs::new(envs), - env_keys: Vec::new(), - timeout: Some(300), - cwd: Some(cwd.to_string()), - bundled: None, - available_tools: Vec::new(), - }; - - if let Err(e) = agent.add_extension(ext, &session_id).await { - // Match buzz-agent: a failed MCP server is a hard session error, - // not a silently tool-less agent. - return Err(AgentError::Mcp(format!("{}: {e}", server.name))); - } + let mcp = Arc::new(crate::mcp::McpRegistry::spawn_all(cfg, mcp_servers, cwd).await?); + let skill_index = crate::skills::skill_index(mcp.skills()); + if !skill_index.is_empty() { + prompt.add_extra("skills", skill_index).await; } - - // AGENTS.md hints are goose's job, not ours. - // - // `PromptManager::build_system_prompt` already calls `.with_hints()` - // (`prompt_manager.rs:255`), which runs the same git-root-anchored walk - // buzz-agent used to run itself. Adding our own copy as a prompt extra put - // every AGENTS.md into the system prompt *twice* -- verified against a - // scratch repo before this change, once under our `## Project Hints` and - // again under goose's `### Project Hints`. - // - // goose's loader is also a superset: `.goosehints` as well as `AGENTS.md`, - // the filename list configurable via `CONTEXT_FILE_NAMES`, gitignore-aware - // filtering, and `@file` reference expansion bounded at the git root. - // Skills are goose's too, for the same reason as hints. - // - // buzz-agent grew its own discovery, index and `load_skill` tool - // (`hints.rs` + `builtin.rs` + `builtin_client.rs`, ~1.2k lines) before - // goose had one. goose has since shipped `goose::skills` and registers a - // `SkillsClient` as a platform extension, so ours was a strictly smaller - // duplicate: same three project directories, but goose also reads the - // global ones, plugin-installed skills and its own builtins, and supports - // `args` templating we never had. - // - // Registered by handing over a *client* rather than by name, and that is - // the whole point: goose's platform factory - // (`platform_extensions/mod.rs:219-221`) calls plain `SkillsClient::new`, - // which leaves goose's own bundled skills switched on. Those are - // `goose-doc-guide` and `web-search` (`skills/builtins/`) — neither is a - // Buzz skill, and `web-search` advertises a capability Buzz does not - // provide, so every Buzz agent would offer the model a tool-shaped promise - // to shell out to `uvx ddgs`. `with_builtin_skills(false)` is only - // reachable off the constructor, so buzz builds the client itself and - // registers it with `add_client`. - // - // What that route costs: `add_client` does not consult - // `PLATFORM_EXTENSIONS`, so nothing here inherits the table's - // `unprefixed_tools: true`. It does not have to — `is_unprefixed_extension` - // (`extension_manager.rs:392-400`) keys off the *config*, and this passes - // the same `ExtensionConfig::Platform { name: "skills" }` the factory route - // does, so the lookup still hits the table entry and the model still sees a - // bare `load_skill`. Verified by driving both routes: same tool name, same - // filesystem skills, builtins gone. - // - // The prompt index is filtered to match. A skill listed in the index but - // absent from the client is a dead `load_skill` reference. - let skills: Vec<_> = goose::skills::discover_skills(Some(std::path::Path::new(cwd))) - .into_iter() - .filter(|s| s.source_type != goose::custom_requests::SourceType::BuiltinSkill) - .collect(); - if !skills.is_empty() { - let ext = ExtensionConfig::Platform { - name: goose::skills::EXTENSION_NAME.to_string(), - description: "Skills".to_string(), - display_name: Some("Skills".to_string()), - bundled: Some(true), - available_tools: Vec::new(), - }; - // The factory receives a context carrying the session, because - // `SkillsClient::new` reads `session.working_dir` for discovery and - // falls back to the *process* cwd without it — which for a - // desktop-spawned agent is not the nest. - let mut ctx = agent.extension_manager.get_context().clone(); - ctx.extension_manager = Some(Arc::downgrade(&agent.extension_manager)); - ctx.session = Some(Arc::new(session.clone())); - match goose::skills::SkillsClient::new(ctx) { - Ok(client) => { - let client = client.with_builtin_skills(false); - let info = goose::agents::mcp_client::McpClientTrait::get_info(&client).cloned(); - agent - .extension_manager - .add_client( - goose::skills::EXTENSION_NAME.to_string(), - ext, - Arc::new(client), - info, - ) - .await; - prompt - .add_extra("skills", crate::skills::skill_index(&skills)) - .await; - } - Err(e) => { - // Unlike an MCP server, a missing skills extension is not - // fatal: the agent still has every other tool. Losing it - // silently would be worse than losing it loudly. - tracing::warn!(error = %e, "skills extension unavailable"); - } - } - } - - 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 { + let hook_extension = mcp.hook_extension("_Stop"); + if let Some(extension) = &hook_extension { prompt .add_extra("buzz_hook_tools", hook_tool_guidance()) .await; - tracing::info!(extension = %ext, "lifecycle hooks available"); + tracing::info!(extension, "lifecycle hooks available"); } - - Ok((agent, session_id, model, hook_extension, prompt)) + Ok((mcp, uuid_like(), model, hook_extension, prompt)) } /// Keep the model's hands off the lifecycle hooks. @@ -597,7 +405,7 @@ async fn session_new(app: &Arc, id: Value, params: Value, wire_tx: &WireSen .await; } - let (agent, goose_session_id, model, hook_extension, prompt) = + let (mcp, goose_session_id, model, hook_extension, prompt) = match build_agent(&app.cfg, &p.cwd, p.system_prompt.as_deref(), &p.mcp_servers).await { Ok(v) => v, Err(e) => { @@ -631,7 +439,7 @@ async fn session_new(app: &Arc, id: Value, params: Value, wire_tx: &WireSen sessions.insert( sid.clone(), Session { - agent, + mcp, goose_session_id, history: Vec::new(), working_dir: std::path::PathBuf::from(&p.cwd), @@ -758,7 +566,7 @@ async fn session_prompt(app: &Arc, id: Value, params: Value, wire_tx: &Wire // Single-flight per session, and capture the agent handle. let ( - agent, + mcp, goose_session_id, pending_model, hook_extension, @@ -793,7 +601,7 @@ async fn session_prompt(app: &Arc, id: Value, params: Value, wire_tx: &Wire // Take the pending override so a `session/set_model` applies exactly // once, from the next prompt onward (buzz-agent `lib.rs:494-502`). ( - s.agent.clone(), + s.mcp.clone(), s.goose_session_id.clone(), s.pending_model.take(), s.hook_extension.clone(), @@ -839,7 +647,7 @@ async fn session_prompt(app: &Arc, id: Value, params: Value, wire_tx: &Wire let (result, tokens, conversation) = loop_drive::run_turn( loop_drive::TurnContext { - agent: &agent, + mcp: &mcp, session_id: &goose_session_id, wire_tx, cancel: &cancel, diff --git a/crates/buzz-agent/src/loop_drive.rs b/crates/buzz-agent/src/loop_drive.rs index b1524f3371..17f1fc6e4d 100644 --- a/crates/buzz-agent/src/loop_drive.rs +++ b/crates/buzz-agent/src/loop_drive.rs @@ -46,8 +46,8 @@ use std::sync::Arc; use futures::StreamExt; use tokio_util::sync::CancellationToken; -use goose::agents::Agent; -use goose::session::session_manager::Session; +use crate::mcp::McpRegistry; +use crate::turn_state::TurnSession; use goose_provider_types::conversation::message::{Message, MessageContent, ToolRequest}; use goose_provider_types::conversation::Conversation; @@ -63,7 +63,7 @@ use goose_agent::operation::{ConversationEffect, Emitter}; /// goose's own `GooseEffect`: the extra variants there (recipes, extension /// data, recorded usage) exist for operations buzz does not run, so naming the /// narrower type keeps unreachable cases out of `apply_effects` entirely. -type BuzzMachine<'a> = StateMachine<'a, Session, ConversationEffect>; +type BuzzMachine<'a> = StateMachine<'a, TurnSession, ConversationEffect>; use crate::wire::WireSender; @@ -103,7 +103,7 @@ enum Round { /// Everything a round needs that does not change between rounds. pub struct TurnContext<'a> { - pub agent: &'a Arc, + pub mcp: &'a Arc, pub session_id: &'a str, pub wire_tx: &'a WireSender, pub cancel: &'a CancellationToken, @@ -158,11 +158,8 @@ pub async fn run_turn( // The turn's conversation lives here, not in a database. See // `crate::turn_state` for why goose never needed one. - let model_config = ctx - .agent - .model_config_for_session(ctx.session_id) - .await - .ok(); + let (_provider, model_config, _model_id) = ctx.model.snapshot().await; + let model_config = Some(model_config); let mut state = crate::turn_state::TurnState::new( ctx.session_id.to_string(), ctx.working_dir.clone(), @@ -183,9 +180,8 @@ pub async fn run_turn( let outcome = crate::ops::Outcome::new(); let reply_guard_tools = if ctx.require_reply { Some( - ctx.agent - .list_tools(ctx.session_id, None) - .await + ctx.mcp + .rmcp_tools() .into_iter() .map(|tool| tool.name.to_string()) .collect(), @@ -198,14 +194,14 @@ pub async fn run_turn( outcome.clone(), ctx.cancel.clone(), ctx.hook_extension - .map(|extension| (Arc::clone(ctx.agent), extension.to_string())), + .map(|extension| (Arc::clone(ctx.mcp), extension.to_string())), reply_guard_tools, ); let start_machine = crate::ops::round_start( ctx.steers.clone(), crate::ops::BuzzCompactionOperation::new( ctx.model.clone(), - Arc::clone(ctx.agent), + Arc::clone(ctx.mcp), ctx.hook_extension.map(str::to_string), ctx.session_id.to_string(), ), @@ -321,7 +317,7 @@ enum Gate { /// Called at the top of every round, and again when a round wants to end -- /// that second call is what gives `_Stop` its veto, since its operation only /// applies to a conversation whose last message ends the turn. -async fn gate(machine: &BuzzMachine<'_>, session: &Session, emitter: &Emitter) -> Gate { +async fn gate(machine: &BuzzMachine<'_>, session: &TurnSession, emitter: &Emitter) -> Gate { match machine.step(session, emitter).await { Ok(Some(result)) if result.yield_to_client => Gate::Ended, Ok(Some(result)) => Gate::Applied(result.effects), @@ -462,8 +458,8 @@ async fn round( } let results = crate::tools::execute( - ctx.agent, - state.session(), + ctx.mcp, + &state.session().id, ctx.wire_tx, ctx.cancel, ctx.permissions, @@ -487,14 +483,14 @@ struct Inference { /// arrive and accumulating usage. async fn infer( ctx: &TurnContext<'_>, - session: &Session, + session: &TurnSession, conversation: &Conversation, tokens: &mut super::agent::TurnTokens, ) -> Result, AgentError> { let (provider, model_config, _model_id) = ctx.model.snapshot().await; let system_prompt = ctx.prompt.build(&session.working_dir).await; - let tools = ctx.agent.list_tools(ctx.session_id, None).await; + let tools = ctx.mcp.rmcp_tools(); let mut stream = provider .stream( @@ -825,7 +821,7 @@ mod tests { "summary" ); assert_eq!( - state.session().usage.total_tokens, + state.session().total_tokens, None, "a stale total would re-trigger compaction immediately" ); @@ -834,7 +830,7 @@ mod tests { // current occupancy, not a delta to add to the discarded conversation. state.set_total_tokens(Some(12_000)); assert_eq!( - state.session().usage.total_tokens, + state.session().total_tokens, Some(12_000), "the post-compaction round must not accumulate the old 190k total" ); diff --git a/crates/buzz-agent/src/mcp.rs b/crates/buzz-agent/src/mcp.rs new file mode 100644 index 0000000000..bd3ad0f6f8 --- /dev/null +++ b/crates/buzz-agent/src/mcp.rs @@ -0,0 +1,1155 @@ +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use arc_swap::ArcSwap; +use rmcp::model::CallToolRequestParams; +use rmcp::service::{RoleClient, RunningService}; +use rmcp::transport::TokioChildProcess; +use rmcp::ServiceError; +use rmcp::ServiceExt; +use serde_json::{Map, Value}; +use tokio::process::Command; +use tokio::sync::Mutex as AsyncMutex; +use tokio_util::sync::CancellationToken; + +use crate::config::Config; +use crate::types::{AgentError, McpServerStdio, ToolDef, ToolResult, ToolResultContent}; + +const SEP: &str = "__"; +const MAX_NAME_LEN: usize = 128; +const MAX_QNAME_LEN: usize = 64; +const MAX_TOOLS_PER_SESSION: usize = 128; +const MAX_DESCRIPTION_BYTES: usize = 1024; +const MAX_SCHEMA_BYTES: usize = 4096; +const MARKER_FIELD_MAX: usize = 256; +pub const MAX_MCP_SERVERS: usize = 16; + +/// Byte budgets for a single tool result. `total` bounds everything the +/// result may occupy in history (text + images); `text` bounds the text +/// portion alone, since text is where runaway outputs (build logs, file +/// dumps) live while images are legitimately large and self-limiting. +#[derive(Clone, Copy)] +pub struct ResultBudget { + pub total: usize, + pub text: usize, +} + +const PASSTHROUGH_ENV: &[&str] = &[ + // Core + "PATH", + "HOME", + "TERM", + "LANG", + "LC_ALL", + "TMPDIR", + "XDG_CONFIG_HOME", + // SSH — required for git clone/push over SSH (git@github.com:...) + "SSH_AUTH_SOCK", + "SSH_AGENT_PID", + // Git — operator-configured helpers and transport overrides + "GIT_ASKPASS", + "GIT_SSH_COMMAND", + "GIT_CONFIG_GLOBAL", + // Buzz identity — dev-mcp writes NOSTR_PRIVATE_KEY to a keyfile then + // removes it from its own env (children never see it). BUZZ_PRIVATE_KEY + // and BUZZ_RELAY_URL are kept for the buzz CLI. BUZZ_AUTH_TAG is a + // non-secret signed ownership attestation needed by portable owner-scoped + // CLI operations; MCP subprocesses are trusted like the agent runtime. + "NOSTR_PRIVATE_KEY", + "BUZZ_PRIVATE_KEY", + "BUZZ_RELAY_URL", + "BUZZ_AUTH_TAG", +]; + +// Windows has no $TMPDIR/$HOME. TMP/TEMP/USERPROFILE are what +// std::env::temp_dir() consults — without them it falls back to C:\Windows, +// which child processes can't write to (PermissionDenied). USERPROFILE is the +// always-set floor. APPDATA carries child-tool config (git, etc.). +#[cfg(windows)] +const PASSTHROUGH_ENV_WINDOWS: &[&str] = &["TMP", "TEMP", "USERPROFILE", "APPDATA"]; + +/// Environment retained by `spawn_one()` after `env_clear()` on Windows. +/// Shell resolver keys are shared with Doctor through the public contract. +#[cfg(windows)] +fn windows_child_passthrough_env() -> impl Iterator { + PASSTHROUGH_ENV_WINDOWS + .iter() + .copied() + .chain(crate::WINDOWS_SHELL_RESOLUTION_ENV.iter().copied()) +} + +type Client = RunningService; + +#[derive(Clone)] +struct ServerSpec { + name: String, + command: String, + args: Vec, + env: Vec<(String, String)>, + cwd: String, +} + +enum ClientState { + Healthy { + client: Arc, + pgid: Option, + tools: Arc>, + }, + Dead { + attempts: u32, + next_retry: Instant, + reason: String, + // Preserved from the last Healthy state so tools() filtering stays accurate while dead. + tools: Arc>, + }, +} + +struct Server { + name: String, + spec: ServerSpec, + client: ArcSwap, + restart_lock: AsyncMutex<()>, +} + +impl Drop for Server { + fn drop(&mut self) { + if let ClientState::Healthy { pgid: Some(p), .. } = &**self.client.load() { + killpg(*p, &self.name, "drop"); + } + } +} + +enum RestartCheck { + Healthy, + Ready { + attempt_n: u32, + prev_tools: Arc>, + }, +} + +fn check_restart_state(server: &Server, max_attempts: u32) -> Result { + match &**server.client.load() { + ClientState::Healthy { .. } => Ok(RestartCheck::Healthy), + ClientState::Dead { attempts, .. } if *attempts >= max_attempts => { + Err(AgentError::Mcp(format!( + "The MCP server '{}' is unavailable (exhausted). Its tools have been removed for this session.", + server.name + ))) + } + ClientState::Dead { next_retry, reason, .. } if Instant::now() < *next_retry => { + Err(AgentError::Mcp(format!( + "server '{}' is recovering (last error: {reason}). Try again later or use a different tool.", + server.name + ))) + } + ClientState::Dead { attempts, tools, .. } => Ok(RestartCheck::Ready { + attempt_n: attempts + 1, + prev_tools: tools.clone(), + }), + } +} + +struct Entry { + server_idx: usize, + bare: String, +} + +pub struct McpRegistry { + by_qname: HashMap, + defs: Vec, + servers: Vec>, + max_attempts: u32, + backoff_base: Duration, + backoff_max: Duration, + init_timeout: Duration, + skills: Vec, +} + +impl McpRegistry { + #[cfg(test)] + pub(crate) fn empty() -> Self { + Self { + by_qname: HashMap::new(), + defs: Vec::new(), + servers: Vec::new(), + max_attempts: 1, + backoff_base: Duration::from_millis(1), + backoff_max: Duration::from_millis(1), + init_timeout: Duration::from_secs(1), + skills: Vec::new(), + } + } + + pub async fn spawn_all( + _cfg: &Config, + servers: &[McpServerStdio], + cwd: &str, + ) -> Result { + if servers.len() > MAX_MCP_SERVERS { + return Err(AgentError::Mcp(format!( + "too many MCP servers: {} > {MAX_MCP_SERVERS}", + servers.len() + ))); + } + let mut reg = Self { + by_qname: HashMap::new(), + defs: Vec::new(), + servers: Vec::new(), + + max_attempts: env_u32("BUZZ_AGENT_MCP_RESTART_MAX_ATTEMPTS", 3).max(1), + backoff_base: Duration::from_millis( + env_u64("BUZZ_AGENT_MCP_RESTART_BASE_MS", 500).max(1), + ), + backoff_max: Duration::from_millis( + env_u64("BUZZ_AGENT_MCP_RESTART_MAX_MS", 30_000).max(1), + ), + init_timeout: Duration::from_secs( + env_u64("BUZZ_AGENT_MCP_INIT_TIMEOUT_SECS", 30).max(1), + ), + skills: goose::skills::discover_skills(Some(std::path::Path::new(cwd))) + .into_iter() + .filter(|skill| { + skill.source_type != goose::custom_requests::SourceType::BuiltinSkill + }) + .collect(), + }; + + let mut seen_names = HashSet::new(); + for s in servers { + if !valid_name(&s.name) || s.name.contains("__") { + return Err(AgentError::Mcp(format!("invalid server name: {}", s.name))); + } + if !seen_names.insert(s.name.clone()) { + return Err(AgentError::Mcp(format!( + "duplicate server name: {}", + s.name + ))); + } + let spec = ServerSpec { + name: s.name.clone(), + command: s.command.clone(), + args: s.args.clone(), + env: s + .env + .iter() + .map(|e| (e.name.clone(), e.value.clone())) + .collect(), + cwd: cwd.to_owned(), + }; + let (client, pgid, tool_names, raw_tools) = spawn_one(&spec, reg.init_timeout).await?; + let server_idx = reg.servers.len(); + let server = Arc::new(Server { + name: spec.name.clone(), + spec, + client: ArcSwap::from_pointee(ClientState::Healthy { + client: Arc::new(client), + pgid, + tools: Arc::new(tool_names), + }), + restart_lock: AsyncMutex::new(()), + }); + reg.servers.push(server); + + for t in raw_tools { + if reg.defs.len() >= MAX_TOOLS_PER_SESSION { + return Err(AgentError::Mcp(format!( + "too many tools (>{MAX_TOOLS_PER_SESSION})" + ))); + } + let bare = t.name.to_string(); + if !valid_name(&bare) || bare.contains("__") { + return Err(AgentError::Mcp(format!("invalid tool name: {bare}"))); + } + let qname = format!("{}{SEP}{}", s.name, bare); + if qname.len() > MAX_QNAME_LEN { + return Err(AgentError::Mcp(format!( + "qualified tool name too long: {} ({} > {MAX_QNAME_LEN})", + qname, + qname.len() + ))); + } + if reg.by_qname.contains_key(&qname) { + return Err(AgentError::Mcp(format!("duplicate tool: {qname}"))); + } + reg.defs.push(ToolDef { + name: qname.clone(), + description: clamp( + t.description.as_deref().unwrap_or("").to_owned(), + MAX_DESCRIPTION_BYTES, + ), + input_schema: cap_schema(&qname, Value::Object((*t.input_schema).clone())), + }); + reg.by_qname.insert(qname, Entry { server_idx, bare }); + } + } + Ok(reg) + } + + pub fn server_of(&self, qname: &str) -> Option<&str> { + self.by_qname + .get(qname) + .map(|e| self.servers[e.server_idx].name.as_str()) + } + + pub fn has(&self, qname: &str) -> bool { + self.by_qname.contains_key(qname) + } + + /// True if `qname` resolves to a hidden hook tool (bare name starts + /// with `_`). Used to reject hook calls coming from the LLM path — + /// hooks are only callable via `call_hooks`. + pub fn is_hook(&self, qname: &str) -> bool { + self.by_qname + .get(qname) + .map(|e| e.bare.starts_with('_')) + .unwrap_or(false) + } + + /// First extension exposing the requested lifecycle hook, in config order. + pub fn hook_extension(&self, hook: &str) -> Option { + self.servers + .iter() + .find(|server| { + self.by_qname + .contains_key(&format!("{}{SEP}{hook}", server.name)) + }) + .map(|server| server.name.clone()) + } + + pub fn tools(&self) -> Vec { + self.defs + .iter() + .filter(|d| { + let entry = match self.by_qname.get(&d.name) { + Some(e) => e, + None => return false, + }; + // Bare names starting with `_` are hooks — invisible to the LLM. + if entry.bare.starts_with('_') { + return false; + } + let server = &self.servers[entry.server_idx]; + match &**server.client.load() { + ClientState::Healthy { tools, .. } => tools.iter().any(|t| t == &entry.bare), + ClientState::Dead { + attempts, tools, .. + } => *attempts < self.max_attempts && tools.iter().any(|t| t == &entry.bare), + } + }) + .cloned() + .collect() + } + + /// Tool definitions in the GDK/provider-native RMCP shape. + pub fn rmcp_tools(&self) -> Vec { + let mut tools: Vec<_> = self + .defs + .iter() + .cloned() + .map(|tool| { + let schema = match tool.input_schema { + Value::Object(schema) => schema, + _ => Map::new(), + }; + rmcp::model::Tool::new(tool.name, tool.description, Arc::new(schema)) + }) + .collect(); + if !self.skills.is_empty() { + tools.push(crate::skills::load_skill_tool()); + } + tools + } + + pub fn skills(&self) -> &[goose::custom_requests::SourceEntry] { + &self.skills + } + + /// Call one qualified tool and return the provider-native RMCP result. + pub async fn call_rmcp( + &self, + qname: &str, + provider_id: &str, + arguments: &Value, + cancel: &CancellationToken, + ) -> Result { + if qname == "load_skill" { + return Ok(crate::skills::load_skill(&self.skills, arguments)); + } + let text = env_u64("BUZZ_AGENT_MAX_TOOL_RESULT_TEXT_BYTES", 50 * 1024) as usize; + let result = self + .call( + qname, + provider_id, + arguments, + ResultBudget { + total: 8 * 1024 * 1024, + text, + }, + cancel, + ) + .await + .map_err(|error| rmcp::model::ErrorData::internal_error(error.to_string(), None))?; + let content = result + .content + .into_iter() + .map(|block| match block { + ToolResultContent::Text(text) => rmcp::model::ContentBlock::text(text), + ToolResultContent::Image { data, mime_type } => { + rmcp::model::ContentBlock::image(data, mime_type) + } + }) + .collect(); + Ok(if result.is_error { + rmcp::model::CallToolResult::error(content) + } else { + rmcp::model::CallToolResult::success(content) + }) + } + + /// Kill the server's process group and mark it dead. Idempotent: + /// if the server is already Dead (or unknown), this is a no-op. + /// Counts as one attempt toward the restart budget so that a + /// pathological server (starts fine, deadlocks on every call) + /// eventually exhausts. + pub fn kill_server(&self, name: &str, reason: &str) { + let server = match self.servers.iter().find(|s| s.name == name) { + Some(s) => s, + None => return, + }; + let current = server.client.load_full(); + let (pgid, tools) = match &*current { + ClientState::Dead { .. } => return, + ClientState::Healthy { pgid, tools, .. } => (*pgid, tools.clone()), + }; + let dead = Arc::new(ClientState::Dead { + attempts: 1, + next_retry: Instant::now() + backoff(1, self.backoff_base, self.backoff_max), + reason: reason.to_owned(), + tools, + }); + // CAS so we don't clobber a concurrent restart that already + // transitioned the state. If the swap fails, the kill below is + // still safe — the pgid we read belonged to a process we observed + // as Healthy, and killpg on an already-reaped pgid is a no-op. + let prev = server.client.compare_and_swap(¤t, dead); + if Arc::ptr_eq(&prev, ¤t) { + if let Some(p) = pgid { + killpg(p, &server.name, "kill_server"); + } + tracing::error!( + "MCP server '{}' killed and marked dead (reason={reason})", + server.name + ); + } + } + + fn kill_and_mark_dead_if_current( + &self, + server: &Server, + failed_client: &Arc, + reason: &str, + ) { + let current = server.client.load_full(); + match &*current { + ClientState::Healthy { + client, + pgid, + tools, + } if Arc::ptr_eq(client, failed_client) => { + if let Some(p) = *pgid { + killpg(p, &server.name, "call_failed"); + } + let dead = Arc::new(ClientState::Dead { + attempts: 1, + next_retry: Instant::now() + backoff(1, self.backoff_base, self.backoff_max), + reason: reason.to_owned(), + tools: tools.clone(), + }); + let _ = server.client.compare_and_swap(¤t, dead); + tracing::error!( + "MCP server '{}' killed and marked dead (reason={reason})", + server.name + ); + } + _ => {} + } + } + + pub async fn call( + &self, + qname: &str, + provider_id: &str, + arguments: &Value, + budget: ResultBudget, + cancel: &CancellationToken, + ) -> Result { + let entry = self + .by_qname + .get(qname) + .ok_or_else(|| AgentError::Mcp(format!("unknown tool {qname}")))?; + let server = self.servers[entry.server_idx].clone(); + + let state = server.client.load(); + if let ClientState::Healthy { client, tools, .. } = &**state { + if !tools.iter().any(|t| t == &entry.bare) { + return Err(AgentError::Mcp(format!( + "tool '{qname}': no longer available; the MCP server restarted with a different tool set." + ))); + } + let client = client.clone(); + drop(state); + return self + .do_call( + &server, + &client, + &entry.bare, + qname, + provider_id, + arguments, + budget, + cancel, + ) + .await; + } + drop(state); + + self.maybe_restart(&server).await?; + let state = server.client.load(); + let client = match &**state { + ClientState::Healthy { client, tools, .. } => { + if !tools.iter().any(|t| t == &entry.bare) { + return Err(AgentError::Mcp(format!( + "tool '{qname}': no longer available; the MCP server restarted with a different tool set." + ))); + } + client.clone() + } + ClientState::Dead { reason, .. } => { + return Err(AgentError::Mcp(format!( + "tool '{qname}': server '{}' restart failed: {reason}", + server.name + ))); + } + }; + drop(state); + self.do_call( + &server, + &client, + &entry.bare, + qname, + provider_id, + arguments, + budget, + cancel, + ) + .await + } + + #[allow(clippy::too_many_arguments)] + async fn do_call( + &self, + server: &Server, + client: &Arc, + bare: &str, + qname: &str, + provider_id: &str, + arguments: &Value, + budget: ResultBudget, + cancel: &CancellationToken, + ) -> Result { + let arg_obj = match arguments { + Value::Object(m) => Some(m.clone()), + Value::Null => None, + _ => { + return Err(AgentError::Mcp(format!( + "tool {qname} arguments must be a JSON object" + ))) + } + }; + let mut params = CallToolRequestParams::default(); + params.name = bare.to_owned().into(); + params.arguments = arg_obj; + + use rmcp::model::{CallToolRequest, ClientRequest, ServerResult}; + use rmcp::service::PeerRequestOptions; + + let req = ClientRequest::CallToolRequest(CallToolRequest::new(params)); + let mut handle = client + .peer() + .send_cancellable_request(req, PeerRequestOptions::no_options()) + .await + .map_err(|e| AgentError::Mcp(format!("call {qname}: {e}")))?; + + // Early cancel check — watch::changed() only fires on NEW writes. + if cancel.is_cancelled() { + fire_and_forget_cancel(handle, qname); + return Err(AgentError::Cancelled); + } + + // Poll the inner oneshot directly so we can still own `handle` in + // the cancel branch (await_response would move it). + let raw: Result = tokio::select! { + biased; + _ = cancel.cancelled() => { + fire_and_forget_cancel(handle, qname); + return Err(AgentError::Cancelled); + } + r = &mut handle.rx => match r { + Ok(inner) => inner, + Err(_) => Err(ServiceError::TransportClosed), + }, + }; + + let res = match raw { + Ok(ServerResult::CallToolResult(r)) => r, + Ok(_) => { + return Err(AgentError::Mcp(format!( + "call {qname}: unexpected response type" + ))) + } + Err(e) => { + if is_transport_error(&e) { + self.kill_and_mark_dead_if_current( + server, + client, + &format!("call failed: {e}"), + ); + return Err(AgentError::Mcp(format!("call {qname}: {e}"))); + } + // Application-level JSON-RPC error (e.g. -32602 invalid params). + // Server is healthy — it correctly rejected bad input. Return to LLM. + return Ok(ToolResult { + provider_id: provider_id.to_owned(), + content: vec![ToolResultContent::Text(clamp( + format!("Tool call rejected: {e}"), + budget.text, + ))], + is_error: true, + }); + } + }; + let content = tool_result_content(&res.content, budget.total, budget.text); + Ok(ToolResult { + provider_id: provider_id.to_owned(), + content, + is_error: res.is_error.unwrap_or(false), + }) + } + + async fn maybe_restart(&self, server: &Server) -> Result<(), AgentError> { + match check_restart_state(server, self.max_attempts)? { + RestartCheck::Healthy => return Ok(()), + RestartCheck::Ready { .. } => {} + } + + let _guard = server.restart_lock.lock().await; + + let (attempt_n, prev_tools) = match check_restart_state(server, self.max_attempts)? { + RestartCheck::Healthy => return Ok(()), + RestartCheck::Ready { + attempt_n, + prev_tools, + } => (attempt_n, prev_tools), + }; + + let started = Instant::now(); + tracing::info!( + "MCP server '{}' restarting (attempt {attempt_n}/{})", + server.name, + self.max_attempts + ); + match spawn_one(&server.spec, self.init_timeout).await { + Ok((client, pgid, tool_names, _raw_tools)) => { + server.client.store(Arc::new(ClientState::Healthy { + client: Arc::new(client), + pgid, + tools: Arc::new(tool_names), + })); + + tracing::info!( + "MCP server '{}' restarted in {}ms (attempt {attempt_n})", + server.name, + started.elapsed().as_millis() + ); + Ok(()) + } + Err(e) => { + let reason = format!("restart failed: {e}"); + let permanent = attempt_n >= self.max_attempts; + let next_retry = if permanent { + Instant::now() + Duration::from_secs(86_400) + } else { + Instant::now() + backoff(attempt_n, self.backoff_base, self.backoff_max) + }; + server.client.store(Arc::new(ClientState::Dead { + attempts: attempt_n, + next_retry, + reason: reason.clone(), + tools: prev_tools, + })); + + tracing::error!( + "MCP server '{}' restart failed (attempt {attempt_n}/{}, permanent={permanent}): {reason}", + server.name, self.max_attempts + ); + Err(AgentError::Mcp(reason)) + } + } + } +} + +async fn spawn_one( + spec: &ServerSpec, + timeout: Duration, +) -> Result<(Client, Option, Vec, Vec), AgentError> { + let mut cmd = Command::new(&spec.command); + cmd.args(&spec.args); + cmd.env_clear(); + for k in PASSTHROUGH_ENV { + if let Ok(v) = std::env::var(k) { + cmd.env(k, v); + } + } + #[cfg(windows)] + for k in windows_child_passthrough_env() { + if let Ok(v) = std::env::var(k) { + cmd.env(k, v); + } + } + for (k, v) in &spec.env { + cmd.env(k, v); + } + cmd.current_dir(&spec.cwd); + cmd.stderr(std::process::Stdio::inherit()); + + #[cfg(unix)] + cmd.process_group(0); + + configure_no_window(&mut cmd); + + let transport = TokioChildProcess::new(cmd) + .map_err(|e| AgentError::Mcp(format!("spawn {}: {e}", spec.name)))?; + let pgid = transport.id(); + + struct PgidGuard { + pgid: Option, + name: String, + } + impl Drop for PgidGuard { + fn drop(&mut self) { + if let Some(p) = self.pgid.take() { + killpg(p, &self.name, "spawn_dropped"); + } + } + } + let mut guard = PgidGuard { + pgid, + name: spec.name.clone(), + }; + + let client: Client = match tokio::time::timeout(timeout, ().serve(transport)).await { + Ok(Ok(c)) => c, + Ok(Err(e)) => { + return Err(AgentError::Mcp(format!("init {}: {e}", spec.name))); + } + Err(_) => { + return Err(AgentError::Mcp(timeout_msg("init", &spec.name, timeout))); + } + }; + + let tools = match tokio::time::timeout(timeout, client.peer().list_all_tools()).await { + Ok(Ok(t)) => t, + Ok(Err(e)) => { + return Err(AgentError::Mcp(format!("list_tools {}: {e}", spec.name))); + } + Err(_) => { + return Err(AgentError::Mcp(timeout_msg( + "list_tools", + &spec.name, + timeout, + ))); + } + }; + let names: Vec = tools.iter().map(|t| t.name.to_string()).collect(); + guard.pgid = None; + Ok((client, pgid, names, tools)) +} + +/// Send `notifications/cancelled` to the MCP server, fire-and-forget. +/// Per MCP spec, cancellation notifications are best-effort; we never +/// block the agent on slow server stdio. +fn fire_and_forget_cancel( + handle: rmcp::service::RequestHandle, + qname: &str, +) { + let qname_owned = qname.to_owned(); + tokio::spawn(async move { + if let Err(e) = handle.cancel(Some("session cancelled".into())).await { + tracing::debug!("cancel notification failed for {qname_owned}: {e}"); + } + }); +} + +/// Returns `true` for errors indicating the MCP server process is dead or +/// unreachable. Returns `false` for application-level JSON-RPC errors where +/// the server is healthy but rejected the request (e.g. invalid params). +fn is_transport_error(e: &ServiceError) -> bool { + matches!( + e, + ServiceError::TransportSend(_) + | ServiceError::TransportClosed + | ServiceError::Timeout { .. } + | ServiceError::UnexpectedResponse + ) +} + +fn backoff(attempt: u32, base: Duration, max: Duration) -> Duration { + let shift = attempt.saturating_sub(1).min(20); + let scaled = base.saturating_mul(1u32 << shift); + let capped = scaled.min(max); + let ms = capped.as_millis() as u64; + let jitter_pct = jitter_percent(); + let jittered = (ms as i64) + ((ms as i64) * jitter_pct / 100); + Duration::from_millis(jittered.max(0) as u64) +} + +fn jitter_percent() -> i64 { + let mut buf = [0u8; 1]; + let _ = getrandom::fill(&mut buf); + ((buf[0] as i64) % 41) - 20 +} + +fn timeout_msg(stage: &str, name: &str, t: Duration) -> String { + format!("{stage} {name}: timeout after {}s", t.as_secs()) +} + +fn cap_schema(qname: &str, schema: Value) -> Value { + let size = serde_json::to_vec(&schema).map(|b| b.len()).unwrap_or(0); + if size <= MAX_SCHEMA_BYTES { + return schema; + } + tracing::warn!( + "tool {qname} schema is {size} bytes (>{MAX_SCHEMA_BYTES}); replacing with empty object" + ); + Value::Object(Map::new()) +} + +#[cfg(unix)] +fn killpg(pgid: u32, name: &str, stage: &str) { + use nix::sys::signal::{killpg as nix_killpg, Signal}; + use nix::unistd::Pid; + let result = nix_killpg(Pid::from_raw(pgid as i32), Signal::SIGKILL); + tracing::info!( + "killpg MCP {name} ({stage}) pgid={pgid} ok={}", + result.is_ok() + ); +} +#[cfg(not(unix))] +fn killpg(_pgid: u32, name: &str, stage: &str) { + tracing::info!("relying on Drop to kill MCP {name} ({stage})"); +} + +fn valid_name(s: &str) -> bool { + !s.is_empty() + && s.len() <= MAX_NAME_LEN + && s.bytes() + .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-') +} + +pub(crate) fn clamp(value: String, max: usize) -> String { + if value.len() <= max { + value + } else { + truncate_at_boundary(&value, max).to_string() + } +} + +fn truncate_at_boundary(s: &str, max: usize) -> &str { + if s.len() <= max { + return s; + } + let mut cut = max; + while cut > 0 && !s.is_char_boundary(cut) { + cut -= 1; + } + &s[..cut] +} + +/// Byte allowance reserved for the elision marker inside [`truncate_middle`]. +/// The marker is ~80 bytes; the slack keeps the arithmetic safely one-sided. +const ELISION_MARKER_ALLOWANCE: usize = 128; + +/// Truncate `s` to at most `max` bytes by eliding the *middle*, keeping the +/// head and tail. Tool output puts its conclusion at the end (test summaries, +/// error trailers) and its identity at the start; head-only truncation loses +/// the part the model needs most. The marker reports how much was elided so +/// the model knows to re-run a narrower command rather than trust the gap. +pub(crate) fn truncate_middle(s: &str, max: usize) -> String { + if s.len() <= max { + return s.to_owned(); + } + let keep = max.saturating_sub(ELISION_MARKER_ALLOWANCE); + if keep == 0 { + // Budget too small for head + marker + tail; degrade to a head cut. + return truncate_at_boundary(s, max).to_owned(); + } + let head = truncate_at_boundary(s, keep.div_ceil(2)); + let mut tail_start = s.len() - keep / 2; + while tail_start < s.len() && !s.is_char_boundary(tail_start) { + tail_start += 1; + } + let tail = &s[tail_start..]; + let elided = s.len() - head.len() - tail.len(); + format!( + "{head}\n[... {elided} of {} bytes elided from tool result ...]\n{tail}", + s.len() + ) +} + +/// Assemble tool-result content under two budgets: `max_bytes` bounds the +/// whole result (text + images), `max_text_bytes` bounds the text portion +/// alone. Images are large by nature and pass through whole or get elided +/// with a marker; text is middle-elided so the head (what ran) and tail +/// (how it ended) both survive. Every elision leaves an inline marker. +fn tool_result_content( + blocks: &[rmcp::model::ContentBlock], + max_bytes: usize, + max_text_bytes: usize, +) -> Vec { + let mut out = Vec::new(); + let mut text = String::new(); + let mut used = 0usize; // total bytes emitted (text + images) + let mut text_used = 0usize; // text bytes emitted + let short = |s: &str| truncate_at_boundary(s, MARKER_FIELD_MAX).to_owned(); + + // Flush accumulated text, middle-eliding to whatever budget remains. + let flush_text = |out: &mut Vec, + text: &mut String, + used: &mut usize, + text_used: &mut usize| { + if text.is_empty() { + return; + } + let budget = max_text_bytes + .saturating_sub(*text_used) + .min(max_bytes.saturating_sub(*used)); + let kept = truncate_middle(&std::mem::take(text), budget); + *used = used.saturating_add(kept.len()); + *text_used = text_used.saturating_add(kept.len()); + if !kept.is_empty() { + out.push(ToolResultContent::Text(kept)); + } + }; + + let append = |text: &mut String, s: &str| { + if !text.is_empty() { + text.push('\n'); + } + text.push_str(s); + }; + + for c in blocks { + match c { + rmcp::model::ContentBlock::Text(t) => append(&mut text, &t.text), + rmcp::model::ContentBlock::Image(i) => { + flush_text(&mut out, &mut text, &mut used, &mut text_used); + let image_bytes = i.data.len().saturating_add(i.mime_type.len()); + if used.saturating_add(image_bytes) <= max_bytes { + used = used.saturating_add(image_bytes); + out.push(ToolResultContent::Image { + data: i.data.clone(), + mime_type: i.mime_type.clone(), + }); + } else { + append( + &mut text, + &format!( + "[image elided: {}, {} base64 bytes exceeds remaining tool-result budget]", + short(&i.mime_type), + i.data.len() + ), + ); + } + } + rmcp::model::ContentBlock::Audio(a) => append( + &mut text, + &format!( + "[audio elided: {}, {} bytes]", + short(&a.mime_type), + a.data.len() + ), + ), + rmcp::model::ContentBlock::ResourceLink(r) => { + append(&mut text, &format!("[resource: {}]", short(&r.uri))); + } + rmcp::model::ContentBlock::Resource(_) => append(&mut text, "[resource elided]"), + _ => append(&mut text, "[unsupported content elided]"), + } + } + flush_text(&mut out, &mut text, &mut used, &mut text_used); + out +} + +/// Suppress the console window that Windows otherwise allocates for every +/// console-subsystem child process spawned from a GUI (non-console) parent. +/// No-op on non-Windows platforms. +fn configure_no_window(cmd: &mut Command) { + #[cfg(windows)] + { + const CREATE_NO_WINDOW: u32 = 0x0800_0000; + cmd.creation_flags(CREATE_NO_WINDOW); + } + #[cfg(not(windows))] + let _ = cmd; +} + +fn env_u64(key: &str, default: u64) -> u64 { + std::env::var(key) + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(default) +} + +fn env_u32(key: &str, default: u32) -> u32 { + std::env::var(key) + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(default) +} + +#[cfg(test)] +mod content_tests { + use super::*; + + #[test] + fn passthrough_includes_buzz_owner_attestation() { + assert!(PASSTHROUGH_ENV.contains(&"BUZZ_AUTH_TAG")); + } + use rmcp::model::ContentBlock as Content; + + #[cfg(windows)] + #[test] + fn windows_passthrough_includes_shell_resolution_vars() { + // Temp directories and every resolver key must survive `env_clear()`. + for var in ["TMP", "TEMP", "USERPROFILE"] { + assert!( + PASSTHROUGH_ENV_WINDOWS.contains(&var), + "{var} must pass through for Windows child processes" + ); + } + let child_env: Vec<_> = windows_child_passthrough_env().collect(); + for var in crate::WINDOWS_SHELL_RESOLUTION_ENV { + assert!( + child_env.contains(var), + "{var} must pass through so the MCP shell resolver matches Doctor" + ); + } + } + + #[test] + fn tool_result_content_preserves_images() { + let blocks = vec![ + Content::text("header"), + Content::image("aW1n", "image/png"), + Content::text("tail"), + ]; + let out = tool_result_content(&blocks, 1024, 1024); + assert_eq!(out.len(), 3); + assert!(matches!(&out[0], ToolResultContent::Text(t) if t == "header")); + assert!(matches!( + &out[1], + ToolResultContent::Image { data, mime_type } + if data == "aW1n" && mime_type == "image/png" + )); + assert!(matches!(&out[2], ToolResultContent::Text(t) if t == "tail")); + } + + #[test] + fn tool_result_content_elides_images_over_budget() { + let blocks = vec![Content::image("a".repeat(300), "image/png")]; + let out = tool_result_content(&blocks, 256, 256); + assert_eq!(out.len(), 1); + assert!(matches!(&out[0], ToolResultContent::Text(t) if t.contains("image elided"))); + } + + #[test] + fn oversized_text_is_middle_elided() { + let mut body = String::new(); + for i in 0..5000 { + body.push_str(&format!("line {i}\n")); + } + let blocks = vec![Content::text(body.clone())]; + let out = tool_result_content(&blocks, 1024 * 1024, 4096); + assert_eq!(out.len(), 1); + let ToolResultContent::Text(t) = &out[0] else { + panic!("expected text"); + }; + assert!(t.len() <= 4096, "text exceeds budget: {}", t.len()); + assert!(t.starts_with("line 0\n"), "head lost"); + assert!(t.ends_with("line 4999\n"), "tail lost"); + assert!( + t.contains("bytes elided from tool result"), + "missing elision marker" + ); + } + + #[test] + fn text_within_budget_is_untouched() { + let blocks = vec![Content::text("short output")]; + let out = tool_result_content(&blocks, 1024 * 1024, 4096); + assert_eq!(out.len(), 1); + assert!(matches!(&out[0], ToolResultContent::Text(t) if t == "short output")); + } + + #[test] + fn image_passes_whole_even_when_text_budget_is_small() { + let big_text = "x".repeat(10_000); + let img = "a".repeat(100_000); + let blocks = vec![ + Content::text(big_text), + Content::image(img.clone(), "image/png"), + ]; + let out = tool_result_content(&blocks, 8 * 1024 * 1024, 4096); + assert_eq!(out.len(), 2); + assert!(matches!(&out[0], ToolResultContent::Text(t) if t.len() <= 4096)); + assert!(matches!( + &out[1], + ToolResultContent::Image { data, .. } if data == &img + )); + } + + #[test] + fn truncate_middle_respects_max_and_boundaries() { + let s = "é".repeat(60_000); // 2-byte chars stress boundary handling + for max in [200usize, 1024, 50 * 1024] { + let out = super::truncate_middle(&s, max); + assert!(out.len() <= max, "max={max} got {}", out.len()); + assert!(std::str::from_utf8(out.as_bytes()).is_ok()); + } + assert_eq!(super::truncate_middle("ok", 1024), "ok"); + } + + #[test] + fn configure_no_window_is_a_noop_on_non_windows() { + // Cross-host: calling configure_no_window must not panic on any OS. + // On non-Windows the body is a cfg-gated no-op and the argument is + // consumed as `let _ = cmd`, so the only assertion is "didn't crash". + let mut cmd = Command::new("true"); + configure_no_window(&mut cmd); + } + + #[cfg(windows)] + #[test] + fn configure_no_window_compiles_and_applies_flag_on_windows() { + // On Windows, creation_flags(0x0800_0000) must be accepted without panicking. + // The call is a setter with no getter on tokio::process::Command, so the + // regression test confirms the flag is SET by checking the std inner command. + let mut cmd = Command::new("cmd.exe"); + configure_no_window(&mut cmd); + // std::process::Command on Windows does have as_inner / get_creation_flags via + // CommandExt — but tokio wraps it; we verify by ensuring the call compiles and + // the resulting spawn wouldn't OOM (build+flag-set is the full contract here). + // The real protection is the cfg-gated production path in spawn_one(). + } +} diff --git a/crates/buzz-agent/src/ops.rs b/crates/buzz-agent/src/ops.rs index 7dfcd1a764..ac0eb28aab 100644 --- a/crates/buzz-agent/src/ops.rs +++ b/crates/buzz-agent/src/ops.rs @@ -20,17 +20,17 @@ use std::sync::{Arc, Mutex}; +use crate::mcp::McpRegistry; +use crate::turn_state::TurnSession; use anyhow::Result; use async_trait::async_trait; -use goose::agents::Agent; -use goose::conversation::Conversation; -use goose::session::Session; use goose_agent::machine::{StateMachine, Step}; use goose_agent::operation::{ applied, assistant_turn_count, messages_since_kickoff, not_applicable, yielded, ConversationEffect, Emitter, Operation, OperationResult, }; use goose_provider_types::conversation::message::{Message, ToolRequest}; +use goose_provider_types::conversation::Conversation; use crate::types::StopReason; @@ -97,14 +97,14 @@ impl BuzzMaxRoundsOperation { } #[async_trait] -impl Operation for BuzzMaxRoundsOperation { +impl Operation for BuzzMaxRoundsOperation { fn name(&self) -> &'static str { "buzz_max_rounds" } async fn run( &self, - _session: &Session, + _session: &TurnSession, conversation: &Conversation, _emit: &Emitter, ) -> Result> { @@ -181,7 +181,7 @@ fn ends_turn(messages: &[Message]) -> bool { /// spans exactly that same turn. Three objections still end the turn; a new /// prompt still starts from zero. pub struct BuzzStopVetoOperation { - agent: Arc, + mcp: Arc, extension: String, block_cap: u32, } @@ -189,9 +189,9 @@ pub struct BuzzStopVetoOperation { impl BuzzStopVetoOperation { /// `extension` is the MCP extension serving `_Stop`; without one there is /// no hook to ask, and the operation never applies. - pub fn new(agent: Arc, extension: String, block_cap: u32) -> Self { + pub fn new(mcp: Arc, extension: String, block_cap: u32) -> Self { Self { - agent, + mcp, extension, block_cap, } @@ -203,14 +203,14 @@ impl BuzzStopVetoOperation { const OBJECTED: &str = "objected"; #[async_trait] -impl Operation for BuzzStopVetoOperation { +impl Operation for BuzzStopVetoOperation { fn name(&self) -> &'static str { "buzz_stop_veto" } async fn run( &self, - session: &Session, + _session: &TurnSession, conversation: &Conversation, _emit: &Emitter, ) -> Result> { @@ -233,9 +233,7 @@ impl Operation for BuzzStopVetoOperation { // A broken or absent hook must never trap a turn -- `stop_objection` // maps every failure to `None`, which is "no objection". - let Some(objection) = - crate::hooks::stop_objection(&self.agent, session, &self.extension).await - else { + let Some(objection) = crate::hooks::stop_objection(&self.mcp, &self.extension).await else { return not_applicable(); }; @@ -318,14 +316,14 @@ impl BuzzReplyGuardOperation { } #[async_trait] -impl Operation for BuzzReplyGuardOperation { +impl Operation for BuzzReplyGuardOperation { fn name(&self) -> &'static str { "buzz_reply_guard" } async fn run( &self, - _session: &Session, + _session: &TurnSession, conversation: &Conversation, _emit: &Emitter, ) -> Result> { @@ -392,14 +390,14 @@ impl BuzzSteerOperation { } #[async_trait] -impl Operation for BuzzSteerOperation { +impl Operation for BuzzSteerOperation { fn name(&self) -> &'static str { "buzz_steer" } async fn run( &self, - _session: &Session, + _session: &TurnSession, _conversation: &Conversation, _emit: &Emitter, ) -> Result> { @@ -428,7 +426,7 @@ impl Operation for BuzzSteerOperation { /// mid-work and announce its own housekeeping to the humans watching. pub struct BuzzCompactionOperation { model: crate::model::SessionModel, - agent: Arc, + mcp: Arc, hook_extension: Option, session_id: String, } @@ -436,13 +434,13 @@ pub struct BuzzCompactionOperation { impl BuzzCompactionOperation { pub fn new( model: crate::model::SessionModel, - agent: Arc, + mcp: Arc, hook_extension: Option, session_id: String, ) -> Self { Self { model, - agent, + mcp, hook_extension, session_id, } @@ -450,23 +448,37 @@ impl BuzzCompactionOperation { } #[async_trait] -impl Operation for BuzzCompactionOperation { +impl Operation for BuzzCompactionOperation { fn name(&self) -> &'static str { "buzz_compaction" } async fn run( &self, - session: &Session, + session: &TurnSession, conversation: &Conversation, _emit: &Emitter, ) -> Result> { let (provider, model_config, _model_id) = self.model.snapshot().await; + // Goose's legacy threshold facade still takes its application Session. + // Keep that coupling at this single adapter until the granular context + // crate exposes thresholding over model config + token occupancy. + let legacy_session = goose::session::Session { + id: session.id.clone(), + working_dir: session.working_dir.clone(), + model_config: session.model_config.clone(), + usage: goose_provider_types::conversation::token_usage::Usage { + total_tokens: session.total_tokens, + ..Default::default() + }, + conversation: Some(conversation.clone()), + ..Default::default() + }; if !goose::context_mgmt::check_if_compaction_needed( provider.as_ref(), conversation, None, - session, + &legacy_session, ) .await? { @@ -491,9 +503,7 @@ impl Operation for BuzzCompactionOperation { let mut effects = vec![ConversationEffect::ReplaceConversation(result.conversation)]; if let Some(extension) = &self.hook_extension { - if let Some(reported) = - crate::hooks::post_compact_state(&self.agent, session, extension).await - { + if let Some(reported) = crate::hooks::post_compact_state(&self.mcp, extension).await { // `[PostCompact]` prefix preserved from the inline version: // it is how the model tells re-injected state apart from a // user turn, and dropping it would change what it reads. @@ -518,18 +528,18 @@ pub fn round_gate( max_rounds: u32, outcome: Outcome, cancel: tokio_util::sync::CancellationToken, - stop_veto: Option<(Arc, String)>, + stop_veto: Option<(Arc, String)>, reply_guard_tools: Option>, -) -> StateMachine<'static, Session, ConversationEffect> { +) -> StateMachine<'static, TurnSession, ConversationEffect> { // Order is precedence. The round budget is checked first: once it is // spent the turn ends, and asking `_Stop` to veto a turn we are ending // anyway would dispatch a tool call whose answer cannot be honoured. - let mut steps: Vec> = vec![Step::Operation(Arc::new( + let mut steps: Vec> = vec![Step::Operation(Arc::new( BuzzMaxRoundsOperation::new(max_rounds, outcome), ))]; - if let Some((agent, extension)) = stop_veto { + if let Some((mcp, extension)) = stop_veto { steps.push(Step::Operation(Arc::new(BuzzStopVetoOperation::new( - agent, + mcp, extension, crate::hooks::stop_block_cap(), )))); @@ -557,7 +567,7 @@ pub fn round_start( steers: crate::steer::SteerQueue, compaction: BuzzCompactionOperation, cancel: tokio_util::sync::CancellationToken, -) -> StateMachine<'static, Session, ConversationEffect> { +) -> StateMachine<'static, TurnSession, ConversationEffect> { // Steer before compaction: a steer that arrives just as the window fills // should be part of what gets summarised, not appended to a conversation // that was compacted a moment earlier without it. @@ -578,8 +588,8 @@ mod tests { /// A bare agent with no extensions: enough to dispatch a tool call that /// will fail to resolve, which is the "hook unavailable" path. - async fn agent() -> Arc { - Arc::new(Agent::new()) + async fn agent() -> Arc { + Arc::new(McpRegistry::empty()) } fn emitter() -> Emitter { @@ -621,7 +631,7 @@ mod tests { let outcome = Outcome::new(); let op = BuzzMaxRoundsOperation::new(3, outcome.clone()); let result = op - .run(&Session::default(), &conversation(2), &emitter()) + .run(&TurnSession::default(), &conversation(2), &emitter()) .await .unwrap(); assert!(matches!(result, OperationResult::NotApplicable)); @@ -633,7 +643,7 @@ mod tests { let outcome = Outcome::new(); let op = BuzzMaxRoundsOperation::new(3, outcome.clone()); let result = op - .run(&Session::default(), &conversation(3), &emitter()) + .run(&TurnSession::default(), &conversation(3), &emitter()) .await .unwrap(); match result { @@ -662,7 +672,7 @@ mod tests { let op = BuzzStopVetoOperation::new(agent().await, "buzz-dev-mcp".to_string(), 3); let result = op .run( - &Session::default(), + &TurnSession::default(), &tool_calling_conversation(), &emitter(), ) @@ -677,7 +687,7 @@ mod tests { // A broken hook must not be able to hold a turn open. let op = BuzzStopVetoOperation::new(agent().await, "no-such-extension".to_string(), 3); let result = op - .run(&Session::default(), &conversation(1), &emitter()) + .run(&TurnSession::default(), &conversation(1), &emitter()) .await .unwrap(); assert!(matches!(result, OperationResult::NotApplicable)); @@ -713,7 +723,7 @@ mod tests { // At the cap the turn is allowed to end. let result = op - .run(&Session::default(), &conversation, &emitter()) + .run(&TurnSession::default(), &conversation, &emitter()) .await .unwrap(); assert!(matches!(result, OperationResult::NotApplicable)); @@ -742,7 +752,7 @@ mod tests { async fn the_reply_guard_reminds_when_nothing_was_published() { let op = BuzzReplyGuardOperation::new(MAX_REPLY_NAGS, ["developer__shell".to_string()]); let result = op - .run(&Session::default(), &conversation(1), &emitter()) + .run(&TurnSession::default(), &conversation(1), &emitter()) .await .unwrap(); match result { @@ -764,7 +774,7 @@ mod tests { ] { let result = op .run( - &Session::default(), + &TurnSession::default(), &turn_with(shell_request(command)), &emitter(), ) @@ -782,7 +792,7 @@ mod tests { let op = BuzzReplyGuardOperation::new(MAX_REPLY_NAGS, ["developer__shell".to_string()]); let result = op .run( - &Session::default(), + &TurnSession::default(), &turn_with(shell_request("ls -la")), &emitter(), ) @@ -796,7 +806,7 @@ mod tests { let op = BuzzReplyGuardOperation::new(MAX_REPLY_NAGS, ["real__shell".to_string()]); let result = op .run( - &Session::default(), + &TurnSession::default(), &turn_with(shell_request("buzz messages send --channel x --content hi")), &emitter(), ) @@ -822,7 +832,7 @@ mod tests { } let result = op .run( - &Session::default(), + &TurnSession::default(), &Conversation::new_unvalidated(messages), &emitter(), ) @@ -840,7 +850,7 @@ mod tests { let op = BuzzSteerOperation::new(steers.clone()); match op - .run(&Session::default(), &conversation(1), &emitter()) + .run(&TurnSession::default(), &conversation(1), &emitter()) .await .unwrap() { @@ -854,7 +864,7 @@ mod tests { // Drained, not peeked: a steer delivered twice would repeat the // instruction to the model on the next round. assert!(matches!( - op.run(&Session::default(), &conversation(1), &emitter()) + op.run(&TurnSession::default(), &conversation(1), &emitter()) .await .unwrap(), OperationResult::NotApplicable @@ -868,7 +878,7 @@ mod tests { // always-applying operation would spin it. let op = BuzzSteerOperation::new(crate::steer::SteerQueue::new()); assert!(matches!( - op.run(&Session::default(), &conversation(1), &emitter()) + op.run(&TurnSession::default(), &conversation(1), &emitter()) .await .unwrap(), OperationResult::NotApplicable @@ -881,7 +891,7 @@ mod tests { // `_Stop` off. An operator with a misbehaving hook depends on it. let op = BuzzStopVetoOperation::new(agent().await, "buzz-dev-mcp".to_string(), 0); let result = op - .run(&Session::default(), &conversation(1), &emitter()) + .run(&TurnSession::default(), &conversation(1), &emitter()) .await .unwrap(); assert!(matches!(result, OperationResult::NotApplicable)); @@ -903,7 +913,7 @@ mod tests { // actually saying it. let op = BuzzMaxRoundsOperation::new(1, Outcome::new()); let result = op - .run(&Session::default(), &conversation(1), &emitter()) + .run(&TurnSession::default(), &conversation(1), &emitter()) .await .unwrap(); match result { @@ -943,7 +953,7 @@ mod tests { let steered = Conversation::new_unvalidated(messages); let result = op - .run(&Session::default(), &steered, &emitter()) + .run(&TurnSession::default(), &steered, &emitter()) .await .unwrap(); assert!( diff --git a/crates/buzz-agent/src/session_store.rs b/crates/buzz-agent/src/session_store.rs deleted file mode 100644 index e9f72dd97d..0000000000 --- a/crates/buzz-agent/src/session_store.rs +++ /dev/null @@ -1,107 +0,0 @@ -//! Where buzz-agent's conversations are persisted. -//! -//! goose's [`SessionManager`] is the conversation store the loop reads and -//! writes each round. Its `instance()` singleton resolves to -//! `Paths::data_dir()` — for a normal install, the *user's own* goose data -//! directory. A buzz-agent using it would write its turns into the same -//! `sessions.db` as the person's `goose` CLI history: not a correctness bug -//! (SQLite handles concurrent writers), but wrong on two counts. Agent -//! conversations get interleaved with a human's personal sessions, and -//! buzz-agent grows an ever-larger file it does not own and cannot prune. -//! -//! So buzz-agent picks its own directory instead. `BUZZ_AGENT_SESSION_DIR` -//! (absolute) is the explicit knob; otherwise we derive a per-user default -//! under the OS data directory. Only if neither is available do we fall back -//! to goose's singleton, because a degraded shared store still beats refusing -//! to start. -//! -//! Note this is deliberately *not* `GOOSE_PATH_ROOT`: that would relocate -//! goose's config and keyring lookups too, and buzz-agent wants goose's -//! provider configuration to keep resolving exactly where goose puts it. - -use std::path::PathBuf; -use std::sync::Arc; - -use goose::session::session_manager::SessionManager; - -/// Absolute path to a directory buzz-agent should keep its session store in. -pub const SESSION_DIR_ENV: &str = "BUZZ_AGENT_SESSION_DIR"; - -/// Resolve the session-store directory, or `None` to use goose's default. -/// -/// Split out from [`session_manager`] so the precedence is testable without -/// touching the filesystem or building a SQLite pool. -fn resolve_dir(env_value: Option<&str>, data_dir: Option) -> Option { - // An explicit relative path is a misconfiguration we should not silently - // resolve against an arbitrary CWD -- the agent's CWD is the user's - // workspace, and a stray sessions.db there would be surprising. - if let Some(path) = env_value.map(PathBuf::from).filter(|p| p.is_absolute()) { - return Some(path); - } - data_dir.map(|dir| dir.join("buzz-agent")) -} - -/// The session store this process should use. -/// -/// Resolved once and shared: `SessionManager::new` builds its own SQLite pool, -/// so calling it per session or per round would open a new pool each time -/// against the same file. Every call site must get this one. -/// -/// Falls back to goose's shared singleton only when no private directory can -/// be determined at all. -pub fn session_manager() -> Arc { - static STORE: std::sync::OnceLock> = std::sync::OnceLock::new(); - STORE - .get_or_init(|| { - Arc::new( - match resolve_dir( - std::env::var(SESSION_DIR_ENV).ok().as_deref(), - dirs::data_dir(), - ) { - Some(dir) => { - tracing::debug!(dir = %dir.display(), "session store"); - SessionManager::new(dir) - } - None => { - tracing::warn!( - "no data dir available; falling back to goose's shared session store" - ); - SessionManager::instance() - } - }, - ) - }) - .clone() -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn an_absolute_env_path_wins() { - let dir = resolve_dir(Some("/tmp/buzz-sessions"), Some(PathBuf::from("/data"))); - assert_eq!(dir, Some(PathBuf::from("/tmp/buzz-sessions"))); - } - - #[test] - fn a_relative_env_path_is_ignored() { - // Resolving this against the agent's CWD would drop a sessions.db in - // whatever workspace the user happened to open. - let dir = resolve_dir(Some("relative/path"), Some(PathBuf::from("/data"))); - assert_eq!(dir, Some(PathBuf::from("/data/buzz-agent"))); - } - - #[test] - fn the_default_is_private_to_buzz_agent() { - // The point of the whole module: never the bare goose data dir, which - // is where the user's own CLI history lives. - let dir = resolve_dir(None, Some(PathBuf::from("/data"))); - assert_eq!(dir, Some(PathBuf::from("/data/buzz-agent"))); - } - - #[test] - fn without_a_data_dir_we_defer_to_goose() { - assert_eq!(resolve_dir(None, None), None); - } -} diff --git a/crates/buzz-agent/src/skills.rs b/crates/buzz-agent/src/skills.rs index 3330cd3ef4..b936d35f35 100644 --- a/crates/buzz-agent/src/skills.rs +++ b/crates/buzz-agent/src/skills.rs @@ -14,6 +14,52 @@ use goose::custom_requests::{SourceEntry, SourceType}; +/// Provider-native definition for the in-process skill loader. +pub fn load_skill_tool() -> rmcp::model::Tool { + let schema = serde_json::json!({ + "type": "object", + "required": ["name"], + "properties": { + "name": { "type": "string", "description": "Skill name, or skill-name/path for a supporting file." }, + "args": { "type": "string", "description": "Optional arguments for a parameterized skill." } + } + }); + rmcp::model::Tool::new( + "load_skill", + "Load a skill's full content into your context so you can follow its instructions.", + schema.as_object().cloned().unwrap_or_default(), + ) +} + +/// Execute the in-process skill loader using Goose's public skill rendering. +pub fn load_skill( + skills: &[SourceEntry], + arguments: &serde_json::Value, +) -> rmcp::model::CallToolResult { + let name = arguments + .get("name") + .and_then(serde_json::Value::as_str) + .unwrap_or(""); + let args = arguments.get("args").and_then(serde_json::Value::as_str); + if let Some(skill) = skills.iter().find(|skill| skill.name == name) { + return match goose::skills::loaded_skill_context_with_args(skill, args) { + Ok(rendered) => { + rmcp::model::CallToolResult::success(vec![rmcp::model::ContentBlock::text( + rendered, + )]) + } + Err(error) => { + rmcp::model::CallToolResult::error(vec![rmcp::model::ContentBlock::text(format!( + "Failed to parse skill arguments: {error}" + ))]) + } + }; + } + rmcp::model::CallToolResult::error(vec![rmcp::model::ContentBlock::text(format!( + "Skill '{name}' not found." + ))]) +} + /// Render the index of available skills for the system prompt. /// /// Empty for an empty slate, so the caller can skip the prompt extra entirely diff --git a/crates/buzz-agent/src/tools.rs b/crates/buzz-agent/src/tools.rs index 7873b422fe..2c9cb4d328 100644 --- a/crates/buzz-agent/src/tools.rs +++ b/crates/buzz-agent/src/tools.rs @@ -23,8 +23,7 @@ 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 crate::mcp::McpRegistry; use goose_provider_types::conversation::message::{ToolRequest, ToolResult}; use crate::permission::{AskSubject, PermissionBroker, PermissionDecision}; @@ -41,8 +40,8 @@ pub type Outcome = (String, ToolResult); /// produced them. #[allow(clippy::too_many_arguments)] pub async fn execute( - agent: &Arc, - session: &Session, + mcp: &Arc, + session_id: &str, wire_tx: &WireSender, cancel: &CancellationToken, permissions: &Arc, @@ -53,8 +52,8 @@ pub async fn execute( let mut pending = FuturesUnordered::new(); for (index, request) in requests.iter().enumerate() { pending.push(run_one( - agent, - session, + mcp, + session_id, wire_tx, cancel, permissions, @@ -73,7 +72,7 @@ pub async fn execute( 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; + emit_terminal(wire_tx, session_id, &id, &result).await; results.push((id, result)); } results @@ -100,8 +99,8 @@ pub fn cancelled_results(requests: &[ToolRequest]) -> Vec { #[allow(clippy::too_many_arguments)] async fn run_one( - agent: &Arc, - session: &Session, + mcp: &Arc, + session_id: &str, wire_tx: &WireSender, cancel: &CancellationToken, permissions: &Arc, @@ -117,7 +116,7 @@ async fn run_one( Err(e) => return (index, (id, Err(e.clone()))), }; - crate::agent::emit_tool_call(wire_tx, &session.id, &id, &call).await; + crate::agent::emit_tool_call(wire_tx, session_id, &id, &call).await; // Ask the client to authorize this call. Every non-authorizing outcome // fails closed: the model sees an ordinary tool failure and the turn @@ -131,7 +130,7 @@ async fn run_one( .request_permission( wire_tx, protocol_version, - &session.id, + session_id, AskSubject { tool_call_id: &id, tool_name: &call.name, @@ -174,15 +173,15 @@ async fn run_one( ); } - let (_request_id, dispatched) = agent - .dispatch_tool_call(call, id.clone(), Some(cancel.clone()), session) + let arguments = call + .arguments + .clone() + .map(serde_json::Value::Object) + .unwrap_or(serde_json::Value::Null); + let result = mcp + .call_rmcp(call.name.as_ref(), &id, &arguments, cancel) .await; - let result = match dispatched { - Ok(call_result) => call_result.result.await, - Err(e) => Err(e), - }; - (index, (id, result)) } @@ -232,7 +231,7 @@ fn reflect_on_failure( async fn emit_terminal( wire_tx: &WireSender, - session: &Session, + session_id: &str, id: &str, result: &ToolResult, ) { @@ -240,7 +239,7 @@ async fn emit_terminal( Ok(out) => out.is_error.unwrap_or(false), Err(_) => true, }; - crate::agent::emit_tool_call_update(wire_tx, &session.id, id, failed).await; + crate::agent::emit_tool_call_update(wire_tx, session_id, id, failed).await; } #[cfg(test)] diff --git a/crates/buzz-agent/src/turn_state.rs b/crates/buzz-agent/src/turn_state.rs index e6a2b495dc..cce3ed8ef4 100644 --- a/crates/buzz-agent/src/turn_state.rs +++ b/crates/buzz-agent/src/turn_state.rs @@ -1,95 +1,79 @@ -//! The conversation for one turn, held in memory. +//! The conversation and small amount of runtime metadata for one turn. //! -//! buzz-agent used to round-trip every message through goose's -//! `SessionManager` — a sqlite database — once per round. That was never -//! required: none of the goose APIs buzz calls reads the store. -//! -//! * `Provider::stream` takes `&[Message]`, not a session. -//! * `Agent::dispatch_tool_call` takes `&Session` but reads only -//! `working_dir` (for the subdirectory-hint tracker) and `id` (for the -//! `PreToolUse` hook context). -//! * `context_mgmt::{check_if_compaction_needed, compact_messages}` take a -//! `&Session` / `session_id`, but the id is only carried into a tokio -//! task-local that becomes an `agent-session-id` HTTP header for provider -//! request attribution (`goose/src/session_context.rs`), and the session is -//! read for `usage.total_tokens` and `model_config` — both plain fields. -//! -//! So the turn owns a [`Session`] value directly. goose still gets the -//! `&Session` its signatures ask for; it just isn't backed by a database. -//! -//! # Why this matters beyond dependency hygiene -//! -//! A `SessionManager` per agent process meant every buzz agent grew a sqlite -//! file it never read back. Worse, before the store was isolated they all -//! shared the user's own goose CLI history. An agent turn is ephemeral — the -//! durable record of a Buzz conversation is the relay, not a local database. +//! Buzz persists conversation history on the relay. The Goose GDK loop and +//! provider APIs only need an in-memory [`Conversation`], so carrying Goose's +//! application-level `Session` (and its sqlite-oriented fields) here couples +//! the embedder to state it does not use. -use goose::conversation::Conversation; -use goose::session::session_manager::Session; use goose_provider_types::conversation::message::Message; +use goose_provider_types::conversation::Conversation; -/// One turn's conversation and the session metadata goose needs alongside it. +/// Runtime metadata passed through Goose's generic state machine. +#[derive(Default)] +pub struct TurnSession { + /// ACP/Goose attribution id for this turn. + pub id: String, + /// Working directory used by prompt discovery and tool hints. + pub working_dir: std::path::PathBuf, + /// Model configuration used by the temporary legacy compaction adapter. + pub model_config: Option, + /// Latest provider-reported occupancy for the current conversation. + pub total_tokens: Option, + conversation: Conversation, +} + +impl goose_agent::machine::MachineSession for TurnSession { + fn id(&self) -> &str { + &self.id + } + + fn conversation(&self) -> Option<&Conversation> { + Some(&self.conversation) + } +} + +/// One turn's in-memory conversation and runtime metadata. pub struct TurnState { - session: Session, + session: TurnSession, } impl TurnState { - /// Start a turn from the session's identity and working directory. - /// - /// `model_config` is carried because `check_if_compaction_needed` reads it - /// from the session to resolve the context limit; without it goose falls - /// back to a default limit and compaction fires at the wrong point. pub fn new( id: String, working_dir: std::path::PathBuf, model_config: Option, ) -> Self { - let mut session = Session { - id, - working_dir, - model_config, - ..Session::default() - }; - session.conversation = Some(Conversation::new_unvalidated(Vec::new())); - Self { session } + Self { + session: TurnSession { + id, + working_dir, + model_config, + total_tokens: None, + conversation: Conversation::new_unvalidated(Vec::new()), + }, + } } - /// The session value goose's APIs take. - pub fn session(&self) -> &Session { + pub fn session(&self) -> &TurnSession { &self.session } - /// This turn's conversation. pub fn conversation(&self) -> Conversation { - self.session.conversation.clone().unwrap_or_default() + self.session.conversation.clone() } - /// Append a message to the turn. pub fn push(&mut self, message: Message) { - let mut messages = self - .session - .conversation - .take() - .map(|conversation| conversation.messages().to_vec()) - .unwrap_or_default(); + let mut messages = self.session.conversation.messages().to_vec(); messages.push(message); - self.session.message_count = messages.len(); - self.session.conversation = Some(Conversation::new_unvalidated(messages)); + self.session.conversation = Conversation::new_unvalidated(messages); } - /// Replace the conversation wholesale, as compaction does. pub fn replace(&mut self, conversation: Conversation) { - self.session.message_count = conversation.messages().len(); - self.session.conversation = Some(conversation); + self.session.conversation = conversation; } - /// Record the running token total. - /// - /// `check_if_compaction_needed` prefers `session.usage.total_tokens` over - /// re-counting the conversation, so keeping this current is what makes - /// compaction fire on the provider's numbers rather than an estimate. pub fn set_total_tokens(&mut self, total: Option) { - self.session.usage.total_tokens = total; + self.session.total_tokens = total; } } @@ -102,66 +86,37 @@ mod tests { } #[test] - fn a_new_turn_starts_empty_but_has_a_conversation() { - // `Some(empty)` rather than `None`: goose's compaction check bails on - // a session with no conversation at all. - let state = state(); - assert!(state.session().conversation.is_some()); - assert!(state.conversation().messages().is_empty()); + fn a_new_turn_starts_empty() { + assert!(state().conversation().messages().is_empty()); } #[test] - fn pushed_messages_are_visible_to_goose() { + fn pushed_messages_are_visible() { let mut state = state(); state.push(Message::user().with_text("hello")); state.push(Message::assistant().with_text("hi")); - - let conversation = state.conversation(); - assert_eq!(conversation.messages().len(), 2); - // goose reads the conversation through the session, so the two views - // must not drift. - assert_eq!( - state - .session() - .conversation - .as_ref() - .map(|c| c.messages().len()), - Some(2) - ); - assert_eq!(state.session().message_count, 2); + assert_eq!(state.conversation().messages().len(), 2); } #[test] fn replace_swaps_the_whole_conversation() { let mut state = state(); state.push(Message::user().with_text("one")); - state.push(Message::user().with_text("two")); - state.replace(Conversation::new_unvalidated(vec![ Message::user().with_text("summary") ])); - assert_eq!(state.conversation().messages().len(), 1); - assert_eq!(state.session().message_count, 1); } #[test] - fn total_tokens_reach_the_session_goose_reads() { - let mut state = state(); + fn runtime_metadata_is_preserved() { + let mut state = TurnState::new("abc".to_string(), std::path::PathBuf::from("/work"), None); state.set_total_tokens(Some(1234)); - assert_eq!(state.session().usage.total_tokens, Some(1234)); - } - - #[test] - fn the_working_dir_and_id_goose_reads_are_preserved() { - // These two fields are the entirety of what `dispatch_tool_call` uses - // the session for; getting them wrong breaks subdirectory hints and - // the PreToolUse hook context. - let state = TurnState::new("abc".to_string(), std::path::PathBuf::from("/work"), None); assert_eq!(state.session().id, "abc"); assert_eq!( state.session().working_dir, std::path::PathBuf::from("/work") ); + assert_eq!(state.session().total_tokens, Some(1234)); } }