From 43d36985a366bf29b97e3ee43c0f8facc8d64aef Mon Sep 17 00:00:00 2001 From: Michael Neale Date: Mon, 27 Jul 2026 20:23:35 +1000 Subject: [PATCH] feat(buzz-agent): restore relay-mesh adaptive MoA routing Buzz owns the meaning of relay-mesh `auto`, and the swap had silently dropped it. The desktop sets BUZZ_AGENT_PREFER_MESH_FOR_AUTO=1 on every relay-mesh agent (relay_mesh.rs:41-44); the old loop honoured it per request (old llm.rs:406-470) by polling the router's /models catalog and sending mesh-llm's virtual Mixture-of-Agents model instead of `auto` whenever the mesh could support it. I previously described this as "pins whatever auto resolved to at startup". That was wrong: `auto` is a router-side id, so nothing resolves it -- the agent just sent `auto` forever and MoA never engaged at all. For mesh-llm lab work that is the entire feature missing, not a degraded version of it. mesh::MeshAutoProvider wraps goose's provider and rewrites ModelConfig.model_name per call. Provider requires only get_name + stream, so wrapping is cheap -- and this is precisely the kind of interception that is only possible with goose as a library; an out-of-process ACP agent has no seam for it. Hysteresis is identical to the old implementation, deliberately: 5s catalog TTL, two consecutive positive observations to enable, immediate disable plus a 30s cooldown on a negative one, and an unreachable/malformed catalog preserves the last confirmed route rather than treating a failed probe as evidence the mesh vanished. A mid-request contraction (503 "MoA requires >=2 models" or error.type=moa_failure) cools down and retries once on `auto`, so the turn still completes. Other 5xx must NOT be treated as contractions -- that would mask real outages behind a silent retry -- and there is a test pinning that. 4 end-to-end tests against a fake mesh-llm router assert what actually goes on the wire: two-turn confirmation before MoA engages, single-model mesh never routes to MoA, contraction produces a mesh->auto retry pair without failing the turn, and the policy is inert (no extra /models polls) when the flag is absent. That last test initially failed on an absolute catalog-hit count -- my assertion was wrong, not the code: session/new polls /models for the desktop model picker and goose does its own lazy capability lookup. Rewritten as a differential across the TTL boundary, which isolates the policy's own poll. 48 tests green, fmt + clippy clean workspace-wide. Signed-off-by: Michael Neale --- crates/buzz-agent/HANDTEST.md | 5 + crates/buzz-agent/src/config.rs | 21 +- crates/buzz-agent/src/lib.rs | 56 +++- crates/buzz-agent/src/mesh.rs | 478 +++++++++++++++++++++++++++ crates/buzz-agent/tests/mesh_auto.rs | 366 ++++++++++++++++++++ 5 files changed, 902 insertions(+), 24 deletions(-) create mode 100644 crates/buzz-agent/src/mesh.rs create mode 100644 crates/buzz-agent/tests/mesh_auto.rs diff --git a/crates/buzz-agent/HANDTEST.md b/crates/buzz-agent/HANDTEST.md index 8661798b20..20b4c5d44f 100644 --- a/crates/buzz-agent/HANDTEST.md +++ b/crates/buzz-agent/HANDTEST.md @@ -86,6 +86,11 @@ running that relay side by side. — though Databricks *model discovery* still uses our own code, moved to `buzz-model-catalog` (the desktop cannot link goose: native `sqlite3` collision with its own rusqlite). +* **Relay-mesh MoA is restored but only tested against a fake router.** With + `mesh-llm` running, a relay-mesh agent on `auto` should switch to the virtual + `mesh` model within ~2 turns once ≥2 models are live, and fall back to `auto` + within 30s of the mesh shrinking. Grep the agent log for + `relay-mesh auto:`. * Binary is roughly +22.7 MiB raw / +6.0 MiB gzip. * Nothing changed in packaging or the harness catalog — same sidecar name, so nothing needed to. diff --git a/crates/buzz-agent/src/config.rs b/crates/buzz-agent/src/config.rs index f0cce6262d..683698c550 100644 --- a/crates/buzz-agent/src/config.rs +++ b/crates/buzz-agent/src/config.rs @@ -204,24 +204,9 @@ impl Config { set_if_absent("GOOSE_CONTEXT_LIMIT", &ctx); } - // `BUZZ_AGENT_PREFER_MESH_FOR_AUTO` is still injected by the desktop - // (`managed_agents/relay_mesh.rs:42`) but is NOT honoured any more. - // - // It used to make the old loop re-resolve the relay-mesh `auto` model - // against the `/models` catalog mid-run, so a long-lived agent could - // join or leave MoA without restarting (`llm.rs:410-440`). Goose - // resolves the model once at session start and has no equivalent hook. - // - // The agent still works — it just pins whatever `auto` resolved to at - // startup. Warn rather than fail: silently ignoring a config the user - // set is how "why is MoA not kicking in" becomes a day of debugging. - if env_str("BUZZ_AGENT_PREFER_MESH_FOR_AUTO").is_some_and(|v| v != "0") { - tracing::warn!( - "BUZZ_AGENT_PREFER_MESH_FOR_AUTO is set but no longer supported: \ - the relay-mesh `auto` model is resolved once at session start \ - and will not follow catalog changes mid-run" - ); - } + // `BUZZ_AGENT_PREFER_MESH_FOR_AUTO` is consumed in `build_provider` + // (lib.rs), not here: it selects a provider *wrapper* rather than an + // env var goose reads. See `mesh.rs`. } } diff --git a/crates/buzz-agent/src/lib.rs b/crates/buzz-agent/src/lib.rs index 1be7aa126d..d46d5c4025 100644 --- a/crates/buzz-agent/src/lib.rs +++ b/crates/buzz-agent/src/lib.rs @@ -22,6 +22,7 @@ pub mod agent; pub mod config; pub mod hooks; +pub mod mesh; pub mod types; pub mod wire; @@ -129,9 +130,7 @@ async fn build_agent( .or_else(|| std::env::var("GOOSE_MODEL").ok()) .ok_or_else(|| AgentError::Llm("no model configured".into()))?; - let provider = goose::providers::create(&provider_name, Vec::new()) - .await - .map_err(|e| map_provider_error(&e.to_string()))?; + 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()))?; @@ -793,15 +792,60 @@ async fn cancel_session(app: &Arc, params: Value) { /// registry) and installs it with the new `ModelConfig`. `SharedProvider` is an /// `Arc>>` precisely so this is hot-swappable /// (goose `agents/types.rs:11-12`). +/// Construct the goose provider, wrapping it in Buzz's relay-mesh `auto` +/// policy when the desktop asked for it. +/// +/// `BUZZ_AGENT_PREFER_MESH_FOR_AUTO=1` is set on every relay-mesh agent +/// (`desktop/src-tauri/src/managed_agents/relay_mesh.rs:41-44`). It means "when +/// the configured model is `auto`, dynamically use mesh-llm's virtual +/// Mixture-of-Agents model whenever the live catalog can support it". Goose +/// resolves a model once per session and has no hook for that, so +/// [`mesh::MeshAutoProvider`] re-resolves per request instead — see that module +/// for why this is only possible with goose-as-a-library. +async fn build_provider( + provider_name: &str, +) -> Result, AgentError> { + let provider = goose::providers::create(provider_name, Vec::new()) + .await + .map_err(|e| map_provider_error(&e.to_string()))?; + + let prefer_mesh = std::env::var("BUZZ_AGENT_PREFER_MESH_FOR_AUTO") + .is_ok_and(|v| !v.trim().is_empty() && v != "0"); + if !prefer_mesh { + return Ok(provider); + } + + // The policy needs to poll the router's own `/models`, so it needs the + // base URL and key. Without them there is nothing to probe, and silently + // pinning `auto` would look like MoA is broken. + let Some(base_url) = env_first(&["OPENAI_BASE_URL", "OPENAI_COMPAT_BASE_URL"]) else { + tracing::warn!( + "BUZZ_AGENT_PREFER_MESH_FOR_AUTO is set but no OpenAI base URL is \ + configured; relay-mesh auto policy disabled" + ); + return Ok(provider); + }; + let api_key = env_first(&["OPENAI_API_KEY", "OPENAI_COMPAT_API_KEY"]).unwrap_or_default(); + + tracing::info!(%base_url, "relay-mesh auto policy enabled"); + Ok(Arc::new(mesh::MeshAutoProvider::new( + provider, base_url, api_key, + ))) +} + +fn env_first(keys: &[&str]) -> Option { + keys.iter() + .filter_map(|k| std::env::var(k).ok()) + .find(|v| !v.trim().is_empty()) +} + async fn apply_model( agent: &Arc, session_id: &str, model_id: &str, ) -> Result<(), AgentError> { let provider_name = std::env::var("GOOSE_PROVIDER").unwrap_or_else(|_| "openai".to_string()); - let provider = goose::providers::create(&provider_name, Vec::new()) - .await - .map_err(|e| map_provider_error(&e.to_string()))?; + let provider = build_provider(&provider_name).await?; let model_config = goose::model_config::model_config_from_user_config(&provider_name, model_id) .map_err(|e| AgentError::LlmModelNotFound(e.to_string()))?; agent diff --git a/crates/buzz-agent/src/mesh.rs b/crates/buzz-agent/src/mesh.rs new file mode 100644 index 0000000000..56d276bdeb --- /dev/null +++ b/crates/buzz-agent/src/mesh.rs @@ -0,0 +1,478 @@ +//! Relay-mesh `auto` policy: adaptive Mixture-of-Agents routing. +//! +//! # What this restores +//! +//! Buzz owns the meaning of relay-mesh's `auto` model. The desktop sets +//! `BUZZ_AGENT_PREFER_MESH_FOR_AUTO=1` on every relay-mesh agent +//! (`desktop/src-tauri/src/managed_agents/relay_mesh.rs:41-44`), and the +//! pre-goose agent loop honoured it per request (old `llm.rs:406-470`): +//! +//! * Poll the router's `/models` catalog, cached for 5s. +//! * If the virtual `mesh` model is advertised **and** ≥2 distinct physical +//! models are present, send `mesh` (Mixture-of-Agents) instead of `auto`. +//! * If the mesh contracts mid-request, cool down for 30s and retry once with +//! plain `auto`. +//! +//! The point is that a long-running agent adopts or leaves MoA as lab nodes +//! come and go, without a restart. Goose resolves a model once at session +//! start, so without this the agent would send `auto` forever and MoA would +//! never engage at all. +//! +//! # How +//! +//! [`MeshAutoProvider`] wraps goose's real provider and rewrites +//! `ModelConfig.model_name` per call. `Provider` only requires `get_name` and +//! `stream`, so wrapping is cheap — and this is the kind of thing that is only +//! possible because we drive goose as a library: an out-of-process ACP agent +//! has no seam to intercept model resolution. +//! +//! Hysteresis is deliberately identical to the old implementation: two +//! consecutive positive observations to enable, immediate disable on a negative +//! one, and an unknown (unreachable/malformed) catalog preserves the last +//! confirmed state rather than treating a failed probe as evidence the mesh +//! vanished. + +use std::collections::BTreeSet; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use async_trait::async_trait; +use goose::providers::base::Provider; +use goose_provider_types::base::{MessageStream, ModelInfo}; +use goose_provider_types::conversation::message::Message; +use goose_provider_types::conversation::token_usage::ProviderUsage; +use goose_provider_types::errors::ProviderError; +use goose_provider_types::model::ModelConfig; +use goose_provider_types::retry::RetryConfig; +use rmcp::model::Tool; +use serde_json::Value; +use tokio::sync::Mutex; + +/// mesh-llm's virtual Mixture-of-Agents model. +const VIRTUAL_MODEL: &str = "mesh"; +/// The router's ordinary single-model routing. +const AUTO_MODEL: &str = "auto"; +/// How long one catalog observation is trusted. +const CATALOG_TTL: Duration = Duration::from_secs(5); +/// Budget for the catalog probe itself — it sits in the inference path. +const CATALOG_TIMEOUT: Duration = Duration::from_secs(2); +/// Backoff after MoA is withdrawn or fails, so a flapping mesh cannot thrash. +const COOLDOWN: Duration = Duration::from_secs(30); +/// Consecutive positive observations required before enabling MoA. +const ENABLE_OBSERVATIONS: u8 = 2; + +/// The 503 body mesh-llm returns when MoA is configured but under-provisioned. +const MOA_UNAVAILABLE_MESSAGE: &str = "MoA requires ≥2 models available in the mesh"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Observation { + Available, + Unavailable, + /// Probe failed or was unparseable. Explicitly *not* evidence of absence. + Unknown, +} + +#[derive(Debug, Default)] +struct State { + last_checked: Option, + consecutive_available: u8, + collective_enabled: bool, + cooldown_until: Option, +} + +/// Does this catalog justify Mixture-of-Agents? +/// +/// Requires the virtual `mesh` model to be advertised *and* at least two +/// distinct physical models behind it. `@main` suffixes are stripped so the +/// same model pinned to a branch does not count twice — MoA over one model is +/// pointless. +/// +/// `None` means the response had no `data` array, i.e. not a catalog we +/// understand; the caller maps that to [`Observation::Unknown`]. +fn catalog_supports_collective(catalog: &Value) -> Option { + let models = catalog.get("data").and_then(Value::as_array)?; + let mut has_virtual_mesh = false; + let mut physical = BTreeSet::new(); + for id in models + .iter() + .filter_map(|model| model.get("id").and_then(Value::as_str)) + .map(str::trim) + .filter(|id| !id.is_empty()) + { + if id == VIRTUAL_MODEL { + has_virtual_mesh = true; + } else if id != AUTO_MODEL { + physical.insert(id.replace("@main", "")); + } + } + Some(has_virtual_mesh && physical.len() >= 2) +} + +/// Did this error mean "the mesh shrank", as opposed to a generic failure? +/// +/// Two shapes, both from the old implementation (`llm.rs:1362-1388`): a 503 +/// whose `error.message` is exactly [`MOA_UNAVAILABLE_MESSAGE`], or any server +/// error whose `error.type` is `moa_failure`. Anything else is a real error and +/// must surface — silently retrying every 5xx as `auto` would mask outages. +fn is_mesh_contraction(err: &ProviderError) -> bool { + let text = err.to_string(); + if text.contains(MOA_UNAVAILABLE_MESSAGE) { + return true; + } + // Provider errors carry the body inline; find the JSON and inspect it. + let Some(start) = text.find('{') else { + return false; + }; + serde_json::from_str::(&text[start..]) + .ok() + .is_some_and(|v| { + let message = v.pointer("/error/message").and_then(Value::as_str); + let kind = v.pointer("/error/type").and_then(Value::as_str); + message == Some(MOA_UNAVAILABLE_MESSAGE) || kind == Some("moa_failure") + }) +} + +/// Wraps a provider and applies Buzz's relay-mesh `auto` policy. +pub struct MeshAutoProvider { + inner: Arc, + /// Router base URL, for the `/models` probe. + base_url: String, + api_key: String, + http: reqwest::Client, + state: Mutex, +} + +impl std::fmt::Debug for MeshAutoProvider { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("MeshAutoProvider") + .field("base_url", &self.base_url) + .finish_non_exhaustive() + } +} + +impl MeshAutoProvider { + /// Wrap `inner`, probing `base_url/models` to decide when to use MoA. + pub fn new(inner: Arc, base_url: String, api_key: String) -> Self { + Self { + inner, + base_url: base_url.trim_end_matches('/').to_string(), + api_key, + http: reqwest::Client::new(), + state: Mutex::new(State::default()), + } + } + + /// Should this call be routed to MoA? + /// + /// Only rewrites the configured `auto`; an explicitly chosen model is never + /// second-guessed. + async fn resolve_model(&self, configured: &str) -> String { + if configured != AUTO_MODEL { + return configured.to_string(); + } + + let mut state = self.state.lock().await; + let now = Instant::now(); + + if let Some(until) = state.cooldown_until { + if now < until { + return AUTO_MODEL.to_string(); + } + state.cooldown_until = None; + } + + if state + .last_checked + .is_some_and(|at| at.elapsed() < CATALOG_TTL) + { + return self.model_for(state.collective_enabled); + } + + let observation = self.observe().await; + let checked_at = Instant::now(); + state.last_checked = Some(checked_at); + match observation { + Observation::Available => { + state.consecutive_available = state.consecutive_available.saturating_add(1); + if state.consecutive_available >= ENABLE_OBSERVATIONS { + state.collective_enabled = true; + } + } + Observation::Unavailable => { + if state.collective_enabled { + state.cooldown_until = Some(checked_at + COOLDOWN); + } + state.consecutive_available = 0; + state.collective_enabled = false; + } + // A failed probe is not evidence a stable mesh disappeared. + // Inference stays authoritative via the contraction fallback. + Observation::Unknown => {} + } + + let model = self.model_for(state.collective_enabled); + tracing::debug!( + configured_model = configured, + request_model = %model, + "relay-mesh auto: resolved request model from live catalog" + ); + model + } + + fn model_for(&self, collective: bool) -> String { + if collective { + VIRTUAL_MODEL + } else { + AUTO_MODEL + } + .to_string() + } + + /// Note that inference just told us the mesh is too small, and back off. + async fn cool_down(&self) { + let now = Instant::now(); + let mut state = self.state.lock().await; + state.last_checked = Some(now); + state.consecutive_available = 0; + state.collective_enabled = false; + state.cooldown_until = Some(now + COOLDOWN); + } + + async fn observe(&self) -> Observation { + let url = format!("{}/models", self.base_url); + let response = match self + .http + .get(&url) + .bearer_auth(&self.api_key) + .timeout(CATALOG_TIMEOUT) + .send() + .await + { + Ok(response) => response, + Err(error) => { + tracing::debug!( + %error, + "relay-mesh auto: catalog probe failed; preserving last confirmed route" + ); + return Observation::Unknown; + } + }; + if !response.status().is_success() { + tracing::debug!( + status = %response.status(), + "relay-mesh auto: catalog probe unsuccessful; preserving last confirmed route" + ); + return Observation::Unknown; + } + match response.json::().await { + Ok(catalog) => match catalog_supports_collective(&catalog) { + Some(true) => Observation::Available, + Some(false) => Observation::Unavailable, + None => { + tracing::debug!( + "relay-mesh auto: catalog has no data array; preserving last confirmed route" + ); + Observation::Unknown + } + }, + Err(error) => { + tracing::debug!( + %error, + "relay-mesh auto: invalid catalog response; preserving last confirmed route" + ); + Observation::Unknown + } + } + } + + fn config_for(model_config: &ModelConfig, model: &str) -> ModelConfig { + let mut resolved = model_config.clone(); + resolved.model_name = model.to_string(); + resolved + } +} + +#[async_trait] +impl Provider for MeshAutoProvider { + fn get_name(&self) -> &str { + self.inner.get_name() + } + + async fn stream( + &self, + model_config: &ModelConfig, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result { + let model = self.resolve_model(&model_config.model_name).await; + let attempting_moa = model == VIRTUAL_MODEL; + + let resolved = Self::config_for(model_config, &model); + match self.inner.stream(&resolved, system, messages, tools).await { + Err(error) if attempting_moa && is_mesh_contraction(&error) => { + // The mesh shrank between discovery and inference. Back off and + // retry once with ordinary routing so the turn still completes. + self.cool_down().await; + tracing::warn!( + attempted_model = VIRTUAL_MODEL, + fallback_model = AUTO_MODEL, + provider_message = %error, + "relay-mesh auto: collective request failed; retrying once with auto" + ); + let fallback = Self::config_for(model_config, AUTO_MODEL); + self.inner.stream(&fallback, system, messages, tools).await + } + other => other, + } + } + + async fn complete( + &self, + model_config: &ModelConfig, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result<(Message, ProviderUsage), ProviderError> { + let model = self.resolve_model(&model_config.model_name).await; + let attempting_moa = model == VIRTUAL_MODEL; + + let resolved = Self::config_for(model_config, &model); + match self + .inner + .complete(&resolved, system, messages, tools) + .await + { + Err(error) if attempting_moa && is_mesh_contraction(&error) => { + self.cool_down().await; + tracing::warn!( + attempted_model = VIRTUAL_MODEL, + fallback_model = AUTO_MODEL, + provider_message = %error, + "relay-mesh auto: collective request failed; retrying once with auto" + ); + let fallback = Self::config_for(model_config, AUTO_MODEL); + self.inner + .complete(&fallback, system, messages, tools) + .await + } + other => other, + } + } + + // Everything below must delegate: the wrapper exists only to pick a model. + + async fn get_context_limit(&self, model_config: &ModelConfig) -> Result { + self.inner.get_context_limit(model_config).await + } + + fn retry_config(&self) -> RetryConfig { + self.inner.retry_config() + } + + async fn fetch_supported_models(&self) -> Result, ProviderError> { + self.inner.fetch_supported_models().await + } + + async fn fetch_supported_model_info(&self) -> Result, ProviderError> { + self.inner.fetch_supported_model_info().await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn catalog(ids: &[&str]) -> Value { + json!({ "data": ids.iter().map(|id| json!({ "id": id })).collect::>() }) + } + + #[test] + fn collective_needs_virtual_mesh_and_two_physical_models() { + assert_eq!( + catalog_supports_collective(&catalog(&["mesh", "qwen3", "llama3"])), + Some(true) + ); + } + + #[test] + fn one_physical_model_is_not_a_collective() { + // MoA over a single model is pointless, and mesh-llm 503s on it. + assert_eq!( + catalog_supports_collective(&catalog(&["mesh", "qwen3"])), + Some(false) + ); + } + + #[test] + fn virtual_mesh_must_be_advertised() { + assert_eq!( + catalog_supports_collective(&catalog(&["qwen3", "llama3"])), + Some(false) + ); + } + + #[test] + fn branch_pins_do_not_count_as_distinct_models() { + // `qwen3` and `qwen3@main` are the same model. + assert_eq!( + catalog_supports_collective(&catalog(&["mesh", "qwen3", "qwen3@main"])), + Some(false) + ); + } + + #[test] + fn auto_is_not_counted_as_a_physical_model() { + assert_eq!( + catalog_supports_collective(&catalog(&["mesh", "auto", "qwen3"])), + Some(false) + ); + } + + #[test] + fn missing_data_array_is_unknown_not_false() { + // Must be distinguishable: unknown preserves the last confirmed route, + // false actively disables MoA. + assert_eq!(catalog_supports_collective(&json!({})), None); + assert_eq!( + catalog_supports_collective(&json!({"data": []})), + Some(false) + ); + } + + #[test] + fn blank_model_ids_are_ignored() { + assert_eq!( + catalog_supports_collective(&catalog(&["mesh", "qwen3", " ", ""])), + Some(false) + ); + } + + #[test] + fn recognises_the_moa_unavailable_503() { + let err = ProviderError::ServerError(format!( + "503: {{\"error\":{{\"message\":\"{MOA_UNAVAILABLE_MESSAGE}\"}}}}" + )); + assert!(is_mesh_contraction(&err)); + } + + #[test] + fn recognises_moa_failure_type() { + let err = + ProviderError::ServerError("500: {\"error\":{\"type\":\"moa_failure\"}}".to_string()); + assert!(is_mesh_contraction(&err)); + } + + #[test] + fn ordinary_server_errors_are_not_contractions() { + // Critical: treating every 5xx as a contraction would silently mask + // real outages behind an `auto` retry. + let err = + ProviderError::ServerError("500: {\"error\":{\"message\":\"internal\"}}".to_string()); + assert!(!is_mesh_contraction(&err)); + assert!(!is_mesh_contraction(&ProviderError::RequestFailed( + "timeout".to_string() + ))); + } +} diff --git a/crates/buzz-agent/tests/mesh_auto.rs b/crates/buzz-agent/tests/mesh_auto.rs new file mode 100644 index 0000000000..b13b287e78 --- /dev/null +++ b/crates/buzz-agent/tests/mesh_auto.rs @@ -0,0 +1,366 @@ +//! Proves the relay-mesh `auto` policy end to end, over the real stdio wire. +//! +//! The unit tests in `src/mesh.rs` cover the pure decisions (catalog parsing, +//! error classification). This covers the part that actually matters: does the +//! request the provider sends carry `"model":"mesh"` instead of `"auto"`? +//! +//! Hysteresis makes this a two-turn test, deliberately. `ENABLE_OBSERVATIONS` +//! is 2 and one observation is cached for `CATALOG_TTL` (5s), so: +//! +//! turn 1 → first observation → still `auto` +//! (wait out the TTL) +//! turn 2 → second observation → `mesh` +//! +//! That two-strike rule is the whole point — it stops a briefly-flapping peer +//! from bouncing a long-running agent in and out of MoA. + +use std::io::{BufRead, BufReader, Write}; +use std::net::TcpListener; +use std::process::{Child, Command, Stdio}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{mpsc, Arc}; +use std::time::Duration; + +use serde_json::{json, Value}; + +/// Fake mesh-llm router. +/// +/// `/models` reports the virtual `mesh` model plus `physical_models` real ones. +/// `/chat/completions` records the requested model and answers "ok" — unless +/// `moa_503` is set, in which case a `mesh` request gets mesh-llm's real +/// under-provisioned 503 so the contraction fallback can be exercised. +fn spawn_router( + physical_models: usize, + moa_503: bool, +) -> (String, mpsc::Receiver, Arc) { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind"); + let addr = listener.local_addr().expect("addr"); + let (tx, rx) = mpsc::channel(); + let catalog_hits = Arc::new(AtomicUsize::new(0)); + let hits = catalog_hits.clone(); + + std::thread::spawn(move || { + for stream in listener.incoming() { + let Ok(mut stream) = stream else { continue }; + let mut reader = BufReader::new(stream.try_clone().expect("clone")); + + let mut request_line = String::new(); + if reader.read_line(&mut request_line).is_err() { + continue; + } + let mut content_length = 0usize; + loop { + let mut line = String::new(); + if reader.read_line(&mut line).is_err() { + break; + } + let t = line.trim_end(); + if t.is_empty() { + break; + } + if let Some(v) = t + .strip_prefix("content-length: ") + .or_else(|| t.strip_prefix("Content-Length: ")) + { + content_length = v.parse().unwrap_or(0); + } + } + let mut body = vec![0u8; content_length]; + use std::io::Read; + let _ = reader.read_exact(&mut body); + + let respond = |stream: &mut std::net::TcpStream, + status: &str, + ct: &str, + payload: String| { + let _ = stream.write_all( + format!( + "HTTP/1.1 {status}\r\ncontent-type: {ct}\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + payload.len(), + payload + ) + .as_bytes(), + ); + let _ = stream.flush(); + }; + + if request_line.contains("/models") { + hits.fetch_add(1, Ordering::SeqCst); + let mut ids: Vec = vec![json!({"id": "mesh"}), json!({"id": "auto"})]; + for i in 0..physical_models { + ids.push(json!({ "id": format!("model-{i}") })); + } + respond( + &mut stream, + "200 OK", + "application/json", + json!({ "object": "list", "data": ids }).to_string(), + ); + continue; + } + + // Record which model the agent actually asked for. + let requested = serde_json::from_slice::(&body) + .ok() + .and_then(|v| v["model"].as_str().map(str::to_owned)) + .unwrap_or_else(|| "".into()); + let _ = tx.send(requested.clone()); + + if moa_503 && requested == "mesh" { + // mesh-llm's real under-provisioned response. + respond( + &mut stream, + "503 Service Unavailable", + "application/json", + json!({"error": {"message": "MoA requires ≥2 models available in the mesh"}}) + .to_string(), + ); + continue; + } + + let usage = json!({"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}); + let text = json!({ + "id": "c", "object": "chat.completion.chunk", "created": 1, "model": requested, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, + "finish_reason": null}], + }); + let done = json!({ + "id": "c", "object": "chat.completion.chunk", "created": 1, "model": "m", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": usage, + }); + respond( + &mut stream, + "200 OK", + "text/event-stream", + format!("data: {text}\n\ndata: {done}\n\ndata: [DONE]\n\n"), + ); + } + }); + + (format!("http://{addr}"), rx, catalog_hits) +} + +struct Harness { + child: Child, + stdin: std::process::ChildStdin, + stdout: BufReader, + next_id: i64, +} + +impl Harness { + fn start(base_url: &str, home: &std::path::Path, prefer_mesh: bool) -> Self { + let mut cmd = Command::new(env!("CARGO_BIN_EXE_buzz-agent")); + cmd.env("BUZZ_AGENT_PROVIDER", "relay-mesh") + .env("BUZZ_AGENT_MODEL", "auto") + .env("OPENAI_COMPAT_API_KEY", "k") + .env("OPENAI_COMPAT_BASE_URL", base_url) + .env("HOME", home) + .env("XDG_CONFIG_HOME", home.join("config")) + .env("XDG_DATA_HOME", home.join("data")) + .env("GOOSE_DISABLE_KEYRING", "1") + .env("RUST_LOG", "warn") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::inherit()); + if prefer_mesh { + cmd.env("BUZZ_AGENT_PREFER_MESH_FOR_AUTO", "1"); + } + let mut child = cmd.spawn().expect("spawn"); + let stdin = child.stdin.take().expect("stdin"); + let stdout = BufReader::new(child.stdout.take().expect("stdout")); + Self { + child, + stdin, + stdout, + next_id: 0, + } + } + + fn request(&mut self, method: &str, params: Value) -> i64 { + self.next_id += 1; + let id = self.next_id; + writeln!( + self.stdin, + "{}", + json!({"jsonrpc": "2.0", "id": id, "method": method, "params": params}) + ) + .expect("write"); + self.stdin.flush().expect("flush"); + id + } + + fn await_response(&mut self, id: i64) -> Value { + loop { + let mut line = String::new(); + let n = self.stdout.read_line(&mut line).expect("read"); + assert_ne!(n, 0, "agent closed stdout awaiting id={id}"); + let Ok(msg) = serde_json::from_str::(&line) else { + continue; + }; + if msg.get("id").and_then(Value::as_i64) == Some(id) { + return msg; + } + } + } + + fn new_session(&mut self, cwd: &std::path::Path) -> String { + let id = self.request("initialize", json!({"protocolVersion": 2})); + let _ = self.await_response(id); + let id = self.request( + "session/new", + json!({"cwd": cwd.to_str().unwrap(), "mcpServers": []}), + ); + let resp = self.await_response(id); + resp["result"]["sessionId"] + .as_str() + .unwrap_or_else(|| panic!("session/new failed: {resp}")) + .to_string() + } + + fn prompt(&mut self, session_id: &str, text: &str) -> Value { + let id = self.request( + "session/prompt", + json!({"sessionId": session_id, "prompt": [{"type": "text", "text": text}]}), + ); + self.await_response(id) + } +} + +impl Drop for Harness { + fn drop(&mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} + +/// Collect every model the router was asked for. +fn drain(rx: &mpsc::Receiver) -> Vec { + let mut seen = Vec::new(); + while let Ok(model) = rx.recv_timeout(Duration::from_millis(500)) { + seen.push(model); + } + seen +} + +#[test] +fn moa_engages_after_two_confirmations() { + // 2 physical models + virtual mesh = MoA is viable. + let (base_url, requests, catalog_hits) = spawn_router(2, false); + let home = tempfile::tempdir().expect("home"); + let cwd = tempfile::tempdir().expect("cwd"); + let mut h = Harness::start(&base_url, home.path(), true); + let session = h.new_session(cwd.path()); + + let resp = h.prompt(&session, "one"); + assert!(resp["result"]["stopReason"].is_string(), "turn 1: {resp}"); + let first = drain(&requests); + assert_eq!( + first, + vec!["auto"], + "first observation must NOT enable MoA — two strikes are required so a \ + flapping peer cannot bounce a long-running agent" + ); + + // Wait out CATALOG_TTL so the next turn re-observes. + std::thread::sleep(Duration::from_millis(5_200)); + + let resp = h.prompt(&session, "two"); + assert!(resp["result"]["stopReason"].is_string(), "turn 2: {resp}"); + let second = drain(&requests); + assert_eq!( + second, + vec!["mesh"], + "second confirmation must switch to the virtual MoA model" + ); + + assert!( + catalog_hits.load(Ordering::SeqCst) >= 2, + "expected the catalog to be re-polled after the TTL — that re-poll is \ + what lets a long-running agent adopt MoA without a restart" + ); +} + +#[test] +fn moa_stays_off_with_only_one_physical_model() { + // MoA over a single model is pointless, and mesh-llm 503s on it. + let (base_url, requests, _) = spawn_router(1, false); + let home = tempfile::tempdir().expect("home"); + let cwd = tempfile::tempdir().expect("cwd"); + let mut h = Harness::start(&base_url, home.path(), true); + let session = h.new_session(cwd.path()); + + h.prompt(&session, "one"); + std::thread::sleep(Duration::from_millis(5_200)); + h.prompt(&session, "two"); + + let seen = drain(&requests); + assert!( + seen.iter().all(|m| m == "auto"), + "a single-model mesh must never route to MoA, got {seen:?}" + ); +} + +#[test] +fn contraction_mid_request_falls_back_to_auto() { + // Catalog says MoA is fine, but inference 503s: the mesh shrank in between. + let (base_url, requests, _) = spawn_router(2, true); + let home = tempfile::tempdir().expect("home"); + let cwd = tempfile::tempdir().expect("cwd"); + let mut h = Harness::start(&base_url, home.path(), true); + let session = h.new_session(cwd.path()); + + h.prompt(&session, "one"); // observation 1 → auto + std::thread::sleep(Duration::from_millis(5_200)); + + // Observation 2 enables MoA, the mesh request 503s, and the turn must still + // complete via a single retry on `auto`. + let resp = h.prompt(&session, "two"); + assert!( + resp["result"]["stopReason"].is_string(), + "a contraction must not fail the turn: {resp}" + ); + + let seen = drain(&requests); + assert!( + seen.windows(2).any(|w| w[0] == "mesh" && w[1] == "auto"), + "expected a mesh attempt immediately followed by an auto retry, got {seen:?}" + ); +} + +#[test] +fn policy_is_inert_when_the_desktop_did_not_ask_for_it() { + // No BUZZ_AGENT_PREFER_MESH_FOR_AUTO: `auto` must be sent untouched, and + // the catalog must not be probed at all. + let (base_url, requests, catalog_hits) = spawn_router(2, false); + let home = tempfile::tempdir().expect("home"); + let cwd = tempfile::tempdir().expect("cwd"); + let mut h = Harness::start(&base_url, home.path(), false); + let session = h.new_session(cwd.path()); + + // /models is polled for reasons that have nothing to do with this policy: + // `session/new` populates the desktop model picker via `discover_models`, + // and goose's own provider does a lazy capability lookup on its first + // request. So an absolute count proves nothing. + // + // The differential does. The policy re-observes once per turn whose cached + // observation has expired, so measuring across the TTL boundary isolates + // it: enabled adds a poll, disabled adds none. + h.prompt(&session, "one"); + let after_first = catalog_hits.load(Ordering::SeqCst); + + std::thread::sleep(Duration::from_millis(5_200)); + + h.prompt(&session, "two"); + let after_second = catalog_hits.load(Ordering::SeqCst); + + let seen = drain(&requests); + assert!( + seen.iter().all(|m| m == "auto"), + "without the flag the model must pass through untouched, got {seen:?}" + ); + assert_eq!( + after_second, after_first, + "a disabled policy must not poll /models across the TTL boundary" + ); +}