mirror of
https://github.com/block/buzz.git
synced 2026-11-01 06:51:17 +00:00
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 <michael.neale@gmail.com>
This commit is contained in:
Michael Neale
@@ -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.
|
||||
|
||||
@@ -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`.
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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<App>, params: Value) {
|
||||
/// registry) and installs it with the new `ModelConfig`. `SharedProvider` is an
|
||||
/// `Arc<Mutex<Option<..>>>` 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<Arc<dyn goose::providers::base::Provider>, 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<String> {
|
||||
keys.iter()
|
||||
.filter_map(|k| std::env::var(k).ok())
|
||||
.find(|v| !v.trim().is_empty())
|
||||
}
|
||||
|
||||
async fn apply_model(
|
||||
agent: &Arc<Agent>,
|
||||
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
|
||||
|
||||
478
crates/buzz-agent/src/mesh.rs
Normal file
478
crates/buzz-agent/src/mesh.rs
Normal file
@@ -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<Instant>,
|
||||
consecutive_available: u8,
|
||||
collective_enabled: bool,
|
||||
cooldown_until: Option<Instant>,
|
||||
}
|
||||
|
||||
/// 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<bool> {
|
||||
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::<Value>(&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<dyn Provider>,
|
||||
/// Router base URL, for the `/models` probe.
|
||||
base_url: String,
|
||||
api_key: String,
|
||||
http: reqwest::Client,
|
||||
state: Mutex<State>,
|
||||
}
|
||||
|
||||
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<dyn Provider>, 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::<Value>().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<MessageStream, 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.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<usize, ProviderError> {
|
||||
self.inner.get_context_limit(model_config).await
|
||||
}
|
||||
|
||||
fn retry_config(&self) -> RetryConfig {
|
||||
self.inner.retry_config()
|
||||
}
|
||||
|
||||
async fn fetch_supported_models(&self) -> Result<Vec<String>, ProviderError> {
|
||||
self.inner.fetch_supported_models().await
|
||||
}
|
||||
|
||||
async fn fetch_supported_model_info(&self) -> Result<Vec<ModelInfo>, 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::<Vec<_>>() })
|
||||
}
|
||||
|
||||
#[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()
|
||||
)));
|
||||
}
|
||||
}
|
||||
366
crates/buzz-agent/tests/mesh_auto.rs
Normal file
366
crates/buzz-agent/tests/mesh_auto.rs
Normal file
@@ -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<String>, Arc<AtomicUsize>) {
|
||||
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<Value> = 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::<Value>(&body)
|
||||
.ok()
|
||||
.and_then(|v| v["model"].as_str().map(str::to_owned))
|
||||
.unwrap_or_else(|| "<none>".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<std::process::ChildStdout>,
|
||||
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::<Value>(&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<String>) -> Vec<String> {
|
||||
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"
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user