External/buzz
1
0
Fork 0
mirror of https://github.com/block/buzz.git synced 2026-08-18 06:50:31 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat(buzz-agent): config parity — thinking effort, model switching, normalized token limits (#1470)

Browse Source 
Signed-off-by: Will Pfleger <pfleger.will@gmail.com>
Co-authored-by: npub1mn7jgtj4w2pd0g0zeuhxsa6jy6p0rewxz4kujt98my82ahfmp72sxjexk7 <dcfd242e557282d7a1e2cf2e6877522682f1e5c6156dc92ca7d90eaedd3b0f95@sprout-oss.stage.blox.sqprod.co>
This commit is contained in:
Will Pfleger
2026-07-02 22:52:38 +00:00
committed by GitHub
co-authored by npub1mn7jgtj4w2pd0g0zeuhxsa6jy6p0rewxz4kujt98my82ahfmp72sxjexk7
parent f9d06ae21a
commit 3e282a2418
13 changed files with 3186 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/buzz-agent/src/agent.rs
+4 -1
View File
@@ -23,6 +23,9 @@ const ERROR_REFLECTION_SUFFIX: &str =
pub struct RunCtx<'a> {
pub cfg: &'a Config,
/// Effective model for this session. Usually equals `cfg.model`; overridden
/// per-session by `session/set_model`. All LLM calls use this value.
pub effective_model: &'a str,
pub session_id: &'a str,
pub system_prompt: &'a str,
pub llm: &'a Llm,
@@ -113,7 +116,7 @@ impl RunCtx<'_> {
let response = tokio::select! {
biased;
_ = self.cancel.changed() => return Ok(StopReason::Cancelled),
r = self.llm.complete(self.cfg, self.system_prompt, self.history, &tools) => r?,
r = self.llm.complete(self.cfg, self.system_prompt, self.history, &tools, self.effective_model) => r?,
_ = async {
// Keepalive ticker: emit a lightweight session update every 30s
// while waiting on the LLM provider. This resets the ACP harness
crates/buzz-agent/src/catalog.rs
+33
View File
@@ -32,6 +32,39 @@ pub struct ModelEntry {
pub const DATABRICKS_V2_KNOWN_MODELS: &[&str] =
&["databricks-gpt-5-5", "databricks-claude-opus-4-7"];
/// Returns the discovery-failure fallback catalog for a Databricks provider.
///
/// This is the list of models advertised by `session/new` when
/// `discover_databricks_models` returns an error (e.g., no token available).
///
/// - `DatabricksV2` falls back to [`DATABRICKS_V2_KNOWN_MODELS`] so the
/// model-picker is always populated for AI Gateway v2 users.
/// - Legacy `Databricks` falls back to only the configured model — the
/// `DATABRICKS_V2_KNOWN_MODELS` IDs are AI Gateway v2 endpoints that the
/// `/serving-endpoints/{model}/invocations` API may not serve.
///
/// Extracting this as a pure function makes the split testable without
/// spawning an async runtime or making network calls.
pub fn discovery_failure_fallback(provider: Provider, configured_model: &str) -> Vec<ModelEntry> {
match provider {
Provider::DatabricksV2 => DATABRICKS_V2_KNOWN_MODELS
.iter()
.map(|id| ModelEntry {
id: id.to_string(),
name: id.to_string(),
})
.collect(),
Provider::Databricks => vec![ModelEntry {
id: configured_model.to_string(),
name: configured_model.to_string(),
}],
_ => vec![ModelEntry {
id: configured_model.to_string(),
name: configured_model.to_string(),
}],
}
}
/// Discover available models for a Databricks provider.
///
/// Returns a non-empty `Vec<ModelEntry>` on success. Returns
crates/buzz-agent/src/config.rs
+1724
View File
File diff suppressed because it is too large Load Diff
crates/buzz-agent/src/handoff.rs
+1
View File
@@ -36,6 +36,7 @@ impl RunCtx<'_> {
HANDOFF_SYSTEM_PROMPT,
&prompt,
HANDOFF_MAX_OUTPUT_TOKENS,
self.effective_model,
) => match r {
Ok(s) if !s.trim().is_empty() => s,
Ok(_) => {
crates/buzz-agent/src/lib.rs
+275 -3
View File
@@ -31,14 +31,20 @@ use crate::mcp::McpRegistry;
use crate::types::{ContentBlock, HistoryItem};
use crate::wire::{
classify, Inbound, InitializeParams, SessionCancelParams, SessionNewParams,
SessionPromptParams, SessionSteerParams, WireMsg, WireSender, INVALID_PARAMS, METHOD_NOT_FOUND,
PARSE_ERROR,
SessionPromptParams, SessionSetModelParams, SessionSteerParams, WireMsg, WireSender,
INVALID_PARAMS, METHOD_NOT_FOUND, PARSE_ERROR,
};
struct App {
cfg: Config,
llm: Arc<Llm>,
sessions: Mutex<HashMap<String, Session>>,
/// Cached model catalog for Databricks providers. Populated lazily on the
/// first successful `session/new` discovery call. When discovery fails (e.g.
/// auth missing or a transient network error) the cell is intentionally left
/// empty so the next `session/new` call retries — a transient failure never
/// pins the degraded fallback catalog for the process lifetime.
models_cache: tokio::sync::OnceCell<Vec<ModelEntry>>,
}
struct Session {
@@ -71,6 +77,10 @@ struct Session {
/// with it so the gate can account for history appended since.
last_request_history_bytes: Option<usize>,
effective_system_prompt: Arc<str>,
/// Per-session model override set by `session/set_model`. When `Some`,
/// overrides `App::cfg.model` for all LLM calls on this session. Persists
/// across `session/prompt` calls until changed.
effective_model: Option<String>,
}
fn die(msg: String) -> ! {
@@ -135,6 +145,7 @@ async fn async_main() {
cfg,
llm,
sessions: Mutex::new(HashMap::new()),
models_cache: tokio::sync::OnceCell::new(),
});
let (wire_tx, wire_rx) = mpsc::channel::<WireMsg>(64);
let writer = tokio::spawn(wire::writer_task(wire_rx));
@@ -206,6 +217,9 @@ async fn handle_request(
tokio::spawn(async move { session_new(&app, id, params, &wire_tx).await });
}
"session/prompt" => spawn_prompt(app.clone(), id, params, wire_tx.clone()),
"session/set_model" => {
set_model_session(app, id, params, wire_tx).await;
}
"session/cancel" => {
cancel_session(app, params).await;
wire::send(wire_tx, wire::ok(id, Value::Null)).await;
@@ -267,6 +281,32 @@ async fn initialize(id: Value, params: Value, wire_tx: &WireSender) {
.await;
}
/// Resolve the Databricks model catalog for one `session/new` call.
///
/// Tries to use a previously-cached successful discovery result. If the cache is empty,
/// runs `discover` and — on success — populates the cache for future calls. On failure
/// the cell is intentionally left empty so the next session retries; the provider-aware
/// fallback is returned for the immediate response only.
///
/// Extracted from `session_new` so that tests can drive this path with an injected
/// discovery future without requiring a full `App` / transport stack.
async fn resolve_models_catalog(
cache: &tokio::sync::OnceCell<Vec<ModelEntry>>,
provider: crate::config::Provider,
model: &str,
discover: impl std::future::Future<Output = Result<Vec<ModelEntry>, AgentError>>,
) -> Vec<ModelEntry> {
match cache.get_or_try_init(|| discover).await {
Ok(cached) => cached.clone(),
Err(e) => {
tracing::warn!(
"model catalog discovery failed: {e}; using fallback (will retry next session)"
);
crate::catalog::discovery_failure_fallback(provider, model)
}
}
}
async fn session_new(app: &Arc<App>, id: Value, params: Value, wire_tx: &WireSender) {
let p: SessionNewParams = match decode(params, "session/new") {
Ok(p) => p,
@@ -365,10 +405,55 @@ async fn session_new(app: &Arc<App>, id: Value, params: Value, wire_tx: &WireSen
last_request_input_tokens: None,
last_request_history_bytes: None,
effective_system_prompt,
effective_model: None,
},
);
drop(sessions);
wire::send(wire_tx, wire::ok(id, json!({ "sessionId": session_id }))).await;
// Build a models catalog for the `session/new` response. For Databricks
// providers this advertises available models so the desktop ModelPicker and
// pool can resolve `session/set_model` switches. For Anthropic/OpenAI we
// report only the configured model — live switching on those providers
// effectively requires respawn.
//
// `models_cache` caches only a successful discovery result (`get_or_try_init`
// leaves the cell empty on error so the next `session/new` call retries). On
// discovery failure the fallback is used for the immediate response without
// being written to the cell.
let available_models: Vec<Value> = {
use crate::config::Provider;
match app.cfg.provider {
Provider::Databricks | Provider::DatabricksV2 => {
let models = resolve_models_catalog(
&app.models_cache,
app.cfg.provider,
&app.cfg.model,
discover_databricks_models(&app.cfg),
)
.await;
models
.iter()
.map(|m| json!({ "modelId": m.id, "name": m.name }))
.collect()
}
_ => vec![json!({ "modelId": app.cfg.model, "name": app.cfg.model })],
}
};
wire::send(
wire_tx,
wire::ok(
id,
json!({
"sessionId": session_id,
"models": {
"currentModelId": app.cfg.model,
"availableModels": available_models,
},
}),
),
)
.await;
}
fn decode<T: serde::de::DeserializeOwned>(params: Value, stage: &str) -> Result<T, String> {
@@ -387,6 +472,55 @@ async fn cancel_session(app: &Arc<App>, params: Value) {
}
}
/// Handle `session/set_model`: apply a per-session model override immediately.
///
/// Validation:
/// - Unknown `sessionId` → `invalid_params`.
/// - Empty `modelId` → `invalid_params`.
///
/// On success: stores `model_id` on the session and responds `{ sessionId, modelId }`.
/// The override is picked up by the next `session/prompt` call on this session.
async fn set_model_session(app: &Arc<App>, id: Value, params: Value, wire_tx: &WireSender) {
let p: SessionSetModelParams = match decode(params, "session/set_model") {
Ok(p) => p,
Err(m) => return reject(wire_tx, id, INVALID_PARAMS, &m).await,
};
if p.model_id.trim().is_empty() {
return reject(
wire_tx,
id,
INVALID_PARAMS,
"session/set_model: modelId must not be empty",
)
.await;
}
let mut sessions = app.sessions.lock().await;
let Some(s) = sessions.get_mut(&p.session_id) else {
return reject(
wire_tx,
id,
INVALID_PARAMS,
"session/set_model: unknown session",
)
.await;
};
s.effective_model = Some(p.model_id.clone());
tracing::info!(
session_id = %p.session_id,
model_id = %p.model_id,
"session/set_model: model overridden"
);
drop(sessions);
wire::send(
wire_tx,
wire::ok(
id,
json!({ "sessionId": p.session_id, "modelId": p.model_id }),
),
)
.await;
}
/// Handle `_goose/unstable/session/steer`: queue user input into the in-flight
/// prompt. Validation mirrors goose's `on_steer_session`:
/// - empty prompt → `invalid_params`
@@ -487,6 +621,7 @@ async fn run_prompt(app: Arc<App>, id: Value, params: Value, wire_tx: WireSender
mut last_request_history_bytes,
mut cancel_rx,
effective_system_prompt,
effective_model_override,
run_id,
mut steer_rx,
) = match acquire_session(&app, &p.session_id).await {
@@ -512,8 +647,13 @@ async fn run_prompt(app: Arc<App>, id: Value, params: Value, wire_tx: WireSender
),
)
.await;
// Resolve effective model: session override wins over config default.
let effective_model_str = effective_model_override
.as_deref()
.unwrap_or(&app.cfg.model);
let mut ctx = RunCtx {
cfg: &app.cfg,
effective_model: effective_model_str,
session_id: &sid,
system_prompt: &effective_system_prompt,
llm: &app.llm,
@@ -570,6 +710,7 @@ async fn acquire_session(
Option<usize>,
watch::Receiver<bool>,
Arc<str>,
Option<String>,
String,
mpsc::UnboundedReceiver<Vec<ContentBlock>>,
),
@@ -593,6 +734,7 @@ async fn acquire_session(
s.active_run_id = Some(run_id.clone());
let (steer_tx, steer_rx) = mpsc::unbounded_channel();
s.steer_tx = Some(steer_tx);
let effective_model = s.effective_model.clone();
Ok((
s.id.clone(),
s.mcp.clone(),
@@ -605,6 +747,7 @@ async fn acquire_session(
s.last_request_history_bytes,
rx,
Arc::clone(&s.effective_system_prompt),
effective_model,
run_id,
steer_rx,
))
@@ -615,3 +758,132 @@ fn session_token() -> Result<String, String> {
getrandom::fill(&mut b).map_err(|e| format!("rng: getrandom failed: {e}"))?;
Ok(b.iter().map(|x| format!("{x:02x}")).collect())
}
#[cfg(test)]
mod tests {
use crate::catalog::{discovery_failure_fallback, ModelEntry, DATABRICKS_V2_KNOWN_MODELS};
use crate::config::Provider;
use crate::types::AgentError;
/// Regression: a discovery error must not pin the models_cache for the process lifetime.
///
/// `resolve_models_catalog` uses `get_or_try_init` so an `Err` leaves the `OnceCell`
/// empty and the next `session/new` retries discovery. This test calls
/// `resolve_models_catalog` directly — the same function `session_new` calls — so
/// reverting `session_new` to `get_or_init` (or any other cache-on-error variant) would
/// break this test, not just the standalone `OnceCell` semantics.
#[tokio::test]
async fn models_cache_does_not_pin_on_discovery_error() {
let cache: tokio::sync::OnceCell<Vec<ModelEntry>> = tokio::sync::OnceCell::new();
let provider = Provider::DatabricksV2;
let model = "my-configured-model";
// First call — discovery fails. Cell must remain empty; fallback returned.
let first = crate::resolve_models_catalog(&cache, provider, model, async {
Err::<Vec<ModelEntry>, AgentError>(AgentError::LlmAuth("transient failure".into()))
})
.await;
assert!(
cache.get().is_none(),
"cell must be empty after a discovery error — next session must retry"
);
let expected_fallback = discovery_failure_fallback(provider, model);
assert_eq!(
first, expected_fallback,
"error path must return the provider-aware fallback"
);
// Second call — discovery succeeds. Cell is now populated and returned.
let discovered = vec![ModelEntry {
id: "databricks-meta-llama-3-1-70b-instruct".into(),
name: "databricks-meta-llama-3-1-70b-instruct".into(),
}];
let discovered_clone = discovered.clone();
let second = crate::resolve_models_catalog(&cache, provider, model, async move {
Ok::<Vec<ModelEntry>, AgentError>(discovered_clone)
})
.await;
assert_eq!(
second, discovered,
"second call must return the discovered catalog"
);
assert!(
cache.get().is_some(),
"cell must be populated after successful discovery"
);
assert_eq!(
cache.get().unwrap(),
&discovered,
"cache must hold the successful discovery result"
);
}
/// Regression: legacy `Provider::Databricks` must not advertise v2 AI Gateway model IDs
/// on discovery failure (Wes W1). This test calls `discovery_failure_fallback` directly —
/// the same helper used by `session_new` — and verifies the split behavior. It FAILS if
/// the arm is un-split (i.e., if both providers return the v2 catalog on failure).
#[test]
fn databricks_discovery_failure_fallback_legacy_returns_configured_model_only() {
let configured = "my-serving-endpoint";
let result = discovery_failure_fallback(Provider::Databricks, configured);
// Legacy Databricks must advertise exactly the configured model — nothing more.
assert_eq!(
result.len(),
1,
"legacy Databricks fallback must contain exactly one entry, got: {result:?}"
);
assert_eq!(
result[0].id, configured,
"legacy Databricks fallback must be the configured model"
);
// Crucially: must NOT contain any DATABRICKS_V2_KNOWN_MODELS entry.
let v2_ids: Vec<&str> = DATABRICKS_V2_KNOWN_MODELS.to_vec();
for id in &result {
assert!(
!v2_ids.contains(&id.id.as_str()),
"legacy Databricks fallback must not include v2 ID '{}' — that endpoint \
may not be served by /serving-endpoints/{{model}}/invocations",
id.id
);
}
}
#[test]
fn databricks_discovery_failure_fallback_v2_returns_known_models_catalog() {
let configured = "my-configured-model";
let result = discovery_failure_fallback(Provider::DatabricksV2, configured);
// DatabricksV2 must return the full DATABRICKS_V2_KNOWN_MODELS list.
assert_eq!(
result.len(),
DATABRICKS_V2_KNOWN_MODELS.len(),
"DatabricksV2 fallback must return all known models"
);
let result_ids: Vec<&str> = result.iter().map(|m| m.id.as_str()).collect();
for known_id in DATABRICKS_V2_KNOWN_MODELS {
assert!(
result_ids.contains(known_id),
"DatabricksV2 fallback must include known model '{known_id}'"
);
}
}
#[test]
fn databricks_discovery_failure_fallback_split_verified() {
// This test FAILS if the v1/v2 arms are merged back into one — it directly verifies
// that the two providers' error-path behavior diverges (Wes W1 protection).
let v1 = discovery_failure_fallback(Provider::Databricks, "my-endpoint");
let v2 = discovery_failure_fallback(Provider::DatabricksV2, "my-endpoint");
let v1_ids: Vec<&str> = v1.iter().map(|m| m.id.as_str()).collect();
let v2_ids: Vec<&str> = v2.iter().map(|m| m.id.as_str()).collect();
assert_ne!(
v1_ids, v2_ids,
"Provider::Databricks and Provider::DatabricksV2 must return different \
fallback catalogs — if they are equal, the W1 arm split has been reverted"
);
}
}
crates/buzz-agent/src/llm.rs
+655 -47
View File
@@ -5,7 +5,10 @@ use reqwest::Client;
use serde_json::{json, Value};
use crate::auth::{PkceOAuthConfig, PkceOAuthTokenSource, StaticTokenSource, TokenSource};
use crate::config::{is_openai_host, Config, OpenAiApi, Provider};
use crate::config::{
is_openai_host, normalize_effort_for_anthropic_route, normalize_effort_for_openai_route,
Config, OpenAiApi, Provider, ThinkingEffort,
};
use crate::types::{
AgentError, HistoryItem, LlmResponse, ProviderStop, ToolCall, ToolDef, ToolResultContent,
};
@@ -66,24 +69,40 @@ impl Llm {
system_prompt: &str,
history: &[HistoryItem],
tools: &[ToolDef],
effective_model: &str,
) -> Result<LlmResponse, AgentError> {
let effort = cfg.thinking_effort;
match cfg.provider {
Provider::Anthropic => {
let v = self
.post_anthropic(cfg, &anthropic_body(cfg, system_prompt, history, tools))
.post_anthropic(
cfg,
&anthropic_body(
cfg,
system_prompt,
history,
tools,
effective_model,
effort,
),
)
.await?;
parse_anthropic(v)
}
Provider::OpenAi | Provider::Databricks => {
self.openai_request(cfg, |use_responses| {
self.openai_request(cfg, effective_model, |use_responses| {
// Normalize effort for model-specific availability (per-model table; max is
// already rejected at startup for pure OpenAI, but other per-model corrections
// like none→minimal on gpt-5 base still apply).
let e = effort.map(|ef| normalize_effort_for_openai_route(ef, effective_model));
if use_responses {
(
responses_body(cfg, system_prompt, history, tools),
responses_body(cfg, system_prompt, history, tools, effective_model, e),
parse_responses as OpenAiParse,
)
} else {
(
openai_body(cfg, system_prompt, history, tools),
openai_body(cfg, system_prompt, history, tools, effective_model, e),
parse_openai as OpenAiParse,
)
}
@@ -91,19 +110,33 @@ impl Llm {
.await
}
Provider::DatabricksV2 => {
self.databricks_v2_request(cfg, |route| match route {
DatabricksV2Route::OpenAiResponses => (
responses_body(cfg, system_prompt, history, tools),
parse_responses as OpenAiParse,
),
DatabricksV2Route::AnthropicMessages => (
anthropic_body(cfg, system_prompt, history, tools),
parse_anthropic as OpenAiParse,
),
DatabricksV2Route::MlflowChatCompletions => (
openai_body(cfg, system_prompt, history, tools),
parse_openai as OpenAiParse,
),
self.databricks_v2_request(cfg, effective_model, |route| match route {
DatabricksV2Route::OpenAiResponses => {
// OpenAI Responses path: normalize effort (max → xhigh, per-model table).
let e =
effort.map(|ef| normalize_effort_for_openai_route(ef, effective_model));
(
responses_body(cfg, system_prompt, history, tools, effective_model, e),
parse_responses as OpenAiParse,
)
}
DatabricksV2Route::AnthropicMessages => {
// Anthropic Messages path: normalize effort (none|minimal → omit).
let e = effort.and_then(normalize_effort_for_anthropic_route);
(
anthropic_body(cfg, system_prompt, history, tools, effective_model, e),
parse_anthropic as OpenAiParse,
)
}
DatabricksV2Route::MlflowChatCompletions => {
// MLflow Chat path (OpenAI-shaped): normalize effort (max → xhigh, per-model table).
let e =
effort.map(|ef| normalize_effort_for_openai_route(ef, effective_model));
(
openai_body(cfg, system_prompt, history, tools, effective_model, e),
parse_openai as OpenAiParse,
)
}
})
.await
}
@@ -116,11 +149,12 @@ impl Llm {
system_prompt: &str,
user_prompt: &str,
max_output_tokens: u32,
effective_model: &str,
) -> Result<String, AgentError> {
match cfg.provider {
Provider::Anthropic => {
let body = json!({
"model": cfg.model,
"model": effective_model,
"max_tokens": max_output_tokens,
"system": system_prompt,
"messages": [{
@@ -132,11 +166,11 @@ impl Llm {
}
Provider::OpenAi | Provider::Databricks => {
let r = self
.openai_request(cfg, |use_responses| {
.openai_request(cfg, effective_model, |use_responses| {
if use_responses {
(
json!({
"model": cfg.model,
"model": effective_model,
"max_output_tokens": max_output_tokens,
"instructions": system_prompt,
"input": user_prompt,
@@ -146,7 +180,7 @@ impl Llm {
} else {
(
json!({
"model": cfg.model,
"model": effective_model,
"stream": false,
"max_completion_tokens": max_output_tokens,
"messages": [
@@ -163,10 +197,10 @@ impl Llm {
}
Provider::DatabricksV2 => {
let r = self
.databricks_v2_request(cfg, |route| match route {
.databricks_v2_request(cfg, effective_model, |route| match route {
DatabricksV2Route::OpenAiResponses => (
json!({
"model": cfg.model,
"model": effective_model,
"max_output_tokens": max_output_tokens,
"instructions": system_prompt,
"input": user_prompt,
@@ -175,7 +209,7 @@ impl Llm {
),
DatabricksV2Route::AnthropicMessages => (
json!({
"model": cfg.model,
"model": effective_model,
"max_tokens": max_output_tokens,
"system": system_prompt,
"messages": [{
@@ -187,7 +221,7 @@ impl Llm {
),
DatabricksV2Route::MlflowChatCompletions => (
json!({
"model": cfg.model,
"model": effective_model,
"stream": false,
"max_completion_tokens": max_output_tokens,
"messages": [
@@ -217,7 +251,12 @@ impl Llm {
/// host), POST, and on `auto` retry once on Responses if the provider
/// asks for it. `build` is called with `use_responses` so callers
/// only construct the body actually needed.
async fn openai_request<F>(&self, cfg: &Config, mut build: F) -> Result<LlmResponse, AgentError>
async fn openai_request<F>(
&self,
cfg: &Config,
effective_model: &str,
mut build: F,
) -> Result<LlmResponse, AgentError>
where
F: FnMut(bool) -> (Value, OpenAiParse) + Send,
{
@@ -227,14 +266,21 @@ impl Llm {
if use_responses {
let (b, p) = build(true);
return p(self.post_openai(cfg, "/responses", &b).await?);
return p(self
.post_openai(cfg, "/responses", &b, effective_model)
.await?);
}
let (b, p) = build(false);
match self.post_openai(cfg, "/chat/completions", &b).await {
match self
.post_openai(cfg, "/chat/completions", &b, effective_model)
.await
{
Ok(v) => p(v),
Err(e) if cfg.openai_api == OpenAiApi::Auto && self.try_upgrade(&e) => {
let (b, p) = build(true);
p(self.post_openai(cfg, "/responses", &b).await?)
p(self
.post_openai(cfg, "/responses", &b, effective_model)
.await?)
}
Err(e) => Err(e),
}
@@ -243,15 +289,16 @@ impl Llm {
async fn databricks_v2_request<F>(
&self,
cfg: &Config,
effective_model: &str,
build: F,
) -> Result<LlmResponse, AgentError>
where
F: FnOnce(DatabricksV2Route) -> (Value, OpenAiParse) + Send,
{
let route = databricks_v2_route_for_model(&cfg.model);
let route = databricks_v2_route_for_model(effective_model);
let (body, parse) = build(route);
parse(
self.post_openai(cfg, databricks_v2_path(route), &body)
self.post_openai(cfg, databricks_v2_path(route), &body, effective_model)
.await?,
)
}
@@ -266,6 +313,7 @@ impl Llm {
cfg: &Config,
path: &str,
body: &Value,
effective_model: &str,
) -> Result<Value, AgentError> {
let (url, body_owned);
let body_ref: &Value = match cfg.provider {
@@ -273,7 +321,7 @@ impl Llm {
url = format!(
"{}/serving-endpoints/{}/invocations",
cfg.base_url.trim_end_matches('/'),
cfg.model
effective_model
);
body_owned = strip_model(body);
&body_owned
@@ -330,6 +378,8 @@ fn anthropic_body(
system_prompt: &str,
history: &[HistoryItem],
tools: &[ToolDef],
effective_model: &str,
effort: Option<ThinkingEffort>,
) -> Value {
let mut messages: Vec<Value> = Vec::new();
let mut pending: Vec<Value> = Vec::new();
@@ -377,8 +427,18 @@ fn anthropic_body(
"name": t.name, "description": t.description, "input_schema": t.input_schema })
})
.collect();
let mut body = json!({ "model": cfg.model, "max_tokens": cfg.max_output_tokens,
let mut body = json!({ "model": effective_model, "max_tokens": cfg.max_output_tokens,
"system": system_prompt, "messages": messages });
if let Some(e) = effort {
let (thinking, output_config) =
crate::config::anthropic_thinking_config(effective_model, e, cfg.max_output_tokens);
if let Some(t) = thinking {
body["thinking"] = t;
}
if let Some(oc) = output_config {
body["output_config"] = oc;
}
}
if !tools_json.is_empty() {
body["tools"] = Value::Array(tools_json);
}
@@ -403,6 +463,8 @@ fn openai_body(
system_prompt: &str,
history: &[HistoryItem],
tools: &[ToolDef],
effective_model: &str,
effort: Option<ThinkingEffort>,
) -> Value {
let mut messages: Vec<Value> = vec![json!({ "role": "system", "content": system_prompt })];
// Images returned from tool calls ride on a trailing `role:"user"`
@@ -463,8 +525,11 @@ fn openai_body(
"parameters": t.input_schema } })
})
.collect();
let mut body = json!({ "model": cfg.model, "stream": false,
let mut body = json!({ "model": effective_model, "stream": false,
"max_completion_tokens": cfg.max_output_tokens, "messages": messages });
if let Some(e) = effort {
body["reasoning_effort"] = json!(e.openai_effort_str());
}
if !tools_json.is_empty() {
body["tools"] = Value::Array(tools_json);
body["tool_choice"] = json!("auto");
@@ -511,6 +576,8 @@ fn responses_body(
system_prompt: &str,
history: &[HistoryItem],
tools: &[ToolDef],
effective_model: &str,
effort: Option<ThinkingEffort>,
) -> Value {
let mut input: Vec<Value> = Vec::with_capacity(history.len());
for item in history {
@@ -574,11 +641,14 @@ fn responses_body(
.collect();
let mut body = json!({
"model": cfg.model,
"model": effective_model,
"instructions": system_prompt,
"max_output_tokens": cfg.max_output_tokens,
"input": input,
});
if let Some(e) = effort {
body["reasoning"] = json!({ "effort": e.openai_effort_str() });
}
if !tools_json.is_empty() {
body["tools"] = Value::Array(tools_json);
body["tool_choice"] = json!("auto");
@@ -1106,6 +1176,7 @@ mod tests {
anthropic_api_version: "2023-06-01".into(),
openai_api: OpenAiApi::Chat,
hints_enabled: true,
thinking_effort: None,
}
}
@@ -1136,7 +1207,14 @@ mod tests {
#[test]
fn anthropic_tool_result_preserves_image_block() {
let body = anthropic_body(&cfg(Provider::Anthropic), "system", &image_history(), &[]);
let body = anthropic_body(
&cfg(Provider::Anthropic),
"system",
&image_history(),
&[],
"model",
None,
);
let content = &body["messages"][2]["content"][0]["content"];
assert_eq!(content[0]["type"], "text");
assert_eq!(content[1]["type"], "image");
@@ -1185,6 +1263,8 @@ mod tests {
"system",
&[HistoryItem::User("hi".into())],
&tools,
"model",
None,
);
assert_eq!(body["model"], "model");
assert_eq!(body["instructions"], "system");
@@ -1213,7 +1293,14 @@ mod tests {
// function_call item *must* appear in `input[]` before its matching
// function_call_output, otherwise the API rejects with
// "No tool call found for call_id ...".
let body = responses_body(&cfg_responses(), "system", &tool_call_history(), &[]);
let body = responses_body(
&cfg_responses(),
"system",
&tool_call_history(),
&[],
"model",
None,
);
let input = body["input"].as_array().unwrap();
// [0] user, [1] assistant text, [2] function_call, [3] function_call_output
@@ -1252,7 +1339,7 @@ mod tests {
}],
},
];
let body = responses_body(&cfg_responses(), "system", &history, &[]);
let body = responses_body(&cfg_responses(), "system", &history, &[], "model", None);
let input = body["input"].as_array().unwrap();
assert_eq!(input.len(), 2);
assert_eq!(input[0]["role"], "user");
@@ -1261,7 +1348,14 @@ mod tests {
#[test]
fn responses_body_image_tool_result_attaches_input_image() {
let body = responses_body(&cfg_responses(), "system", &image_history(), &[]);
let body = responses_body(
&cfg_responses(),
"system",
&image_history(),
&[],
"model",
None,
);
let input = body["input"].as_array().unwrap();
// function_call_output carries the text part; image rides on a
// trailing user message as `input_image`.
@@ -1388,7 +1482,14 @@ mod tests {
#[test]
fn openai_tool_result_adds_followup_image_user_message() {
let body = openai_body(&cfg(Provider::OpenAi), "system", &image_history(), &[]);
let body = openai_body(
&cfg(Provider::OpenAi),
"system",
&image_history(),
&[],
"model",
None,
);
assert_eq!(body["messages"][3]["role"], "tool");
assert!(body["messages"][3]["content"]
.as_str()
@@ -1459,7 +1560,14 @@ mod tests {
is_error: false,
}),
];
let body = openai_body(&cfg(Provider::OpenAi), "system", &history, &[]);
let body = openai_body(
&cfg(Provider::OpenAi),
"system",
&history,
&[],
"model",
None,
);
let messages = body["messages"].as_array().unwrap();
// [0] system, [1] user, [2] assistant(tool_calls), [3] tool A, [4] tool B, [5] user(images)
assert_eq!(messages.len(), 6, "messages: {messages:#?}");
@@ -1477,6 +1585,497 @@ mod tests {
assert_eq!(imgs[1]["image_url"]["url"], "data:image/png;base64,bbb");
}
// ---- ThinkingEffort body-shape tests ----
#[test]
fn anthropic_body_omits_thinking_when_effort_none() {
let body = anthropic_body(
&cfg(Provider::Anthropic),
"system",
&[HistoryItem::User("hi".into())],
&[],
"model",
None,
);
assert!(
body.get("thinking").is_none(),
"thinking must be absent when effort is None"
);
}
#[test]
fn anthropic_body_emits_thinking_when_effort_high() {
// claude-3.x model → manual budget_tokens shape.
// Use max_output_tokens = 4096 so budget fits: headroom = 4096 - 1024 = 3072.
let mut c = cfg(Provider::Anthropic);
c.max_output_tokens = 4096;
let body = anthropic_body(
&c,
"system",
&[HistoryItem::User("hi".into())],
&[],
"claude-3-7-sonnet-20250219",
Some(ThinkingEffort::High),
);
assert_eq!(body["thinking"]["type"], "enabled");
// budget_tokens = min(32768, 4096-1024) = 3072
assert_eq!(body["thinking"]["budget_tokens"], 3072);
assert!(body.get("output_config").is_none());
}
#[test]
fn anthropic_body_omits_thinking_when_max_output_too_small() {
// max_output_tokens = 2047: headroom = 2047 - 1024 = 1023 < 1024 → omit thinking.
let mut c = cfg(Provider::Anthropic);
c.max_output_tokens = 2047;
let body = anthropic_body(
&c,
"system",
&[HistoryItem::User("hi".into())],
&[],
"claude-3-7-sonnet-20250219",
Some(ThinkingEffort::High),
);
assert!(
body.get("thinking").is_none(),
"thinking must be omitted when max_output_tokens leaves < 1024 for budget"
);
}
#[test]
fn anthropic_body_emits_thinking_at_boundary_2048() {
// max_output_tokens = 2048: headroom = 2048 - 1024 = 1024 ≥ 1024 → emit.
let mut c = cfg(Provider::Anthropic);
c.max_output_tokens = 2048;
let body = anthropic_body(
&c,
"system",
&[HistoryItem::User("hi".into())],
&[],
"claude-3-7-sonnet-20250219",
Some(ThinkingEffort::High),
);
let t = body
.get("thinking")
.expect("thinking must be present at boundary 2048");
assert_eq!(t["budget_tokens"], 1024); // min(32768, 2048-1024)
}
#[test]
fn anthropic_body_emits_thinking_high_uncapped_when_budget_fits() {
// When max_output_tokens is large enough, budget_tokens is not capped.
let mut c = cfg(Provider::Anthropic);
c.max_output_tokens = 65_536;
let body = anthropic_body(
&c,
"system",
&[HistoryItem::User("hi".into())],
&[],
"claude-3-7-sonnet-20250219",
Some(ThinkingEffort::High),
);
assert_eq!(body["thinking"]["budget_tokens"], 32_768);
}
#[test]
fn anthropic_body_emits_thinking_low_budget() {
// Low budget (1024 tokens) exactly fits when max_output_tokens = 2048.
// headroom = 2048 - 1024 = 1024; min(1024, 1024) = 1024 ≥ 1024 → emit.
let mut c = cfg(Provider::Anthropic);
c.max_output_tokens = 2048;
let body = anthropic_body(
&c,
"system",
&[HistoryItem::User("hi".into())],
&[],
"claude-3-7-sonnet-20250219",
Some(ThinkingEffort::Low),
);
// Low budget (1024) fits exactly at the boundary — emitted without capping.
assert_eq!(body["thinking"]["budget_tokens"], 1024);
}
#[test]
fn anthropic_body_emits_adaptive_thinking_for_opus_4() {
// Adaptive Claude (claude-opus-4-6/4.7/4.8) → thinking:{type:"adaptive"} + output_config.effort.
// Note: Opus 4.5 is NOT adaptive — it uses manual budget.
let mut c = cfg(Provider::Anthropic);
c.max_output_tokens = 32_768;
let body = anthropic_body(
&c,
"system",
&[HistoryItem::User("hi".into())],
&[],
"claude-opus-4-7",
Some(ThinkingEffort::High),
);
assert_eq!(
body["thinking"]["type"], "adaptive",
"thinking must be {{type:adaptive}} for claude-opus-4-7"
);
assert_eq!(body["output_config"]["effort"], "high");
}
#[test]
fn anthropic_body_emits_manual_budget_for_opus_4_5() {
// Opus 4.5 uses manual budget (effort page: "uses manual thinking").
// max_output_tokens = 32768; headroom = 32768 - 1024 = 31744; min(32768, 31744) = 31744.
let mut c = cfg(Provider::Anthropic);
c.max_output_tokens = 32_768;
let body = anthropic_body(
&c,
"system",
&[HistoryItem::User("hi".into())],
&[],
"claude-opus-4-5",
Some(ThinkingEffort::High),
);
assert_eq!(body["thinking"]["type"], "enabled");
assert_eq!(body["thinking"]["budget_tokens"], 31_744); // min(32768, 32768-1024)
assert!(
body.get("output_config").is_none(),
"output_config must be absent for claude-opus-4-5 (manual budget)"
);
}
#[test]
fn anthropic_body_omits_both_fields_for_unrecognized_model() {
// Non-Anthropic models (gpt-5, llama, etc.) → omit both fields rather than guess.
let body = anthropic_body(
&cfg(Provider::Anthropic),
"system",
&[HistoryItem::User("hi".into())],
&[],
"gpt-4o",
Some(ThinkingEffort::High),
);
assert!(body.get("thinking").is_none(), "thinking must be absent");
assert!(
body.get("output_config").is_none(),
"output_config must be absent"
);
}
#[test]
fn openai_body_omits_reasoning_effort_when_none() {
let body = openai_body(
&cfg(Provider::OpenAi),
"system",
&[HistoryItem::User("hi".into())],
&[],
"model",
None,
);
assert!(
body.get("reasoning_effort").is_none(),
"reasoning_effort must be absent when effort is None"
);
}
#[test]
fn openai_body_emits_reasoning_effort_medium() {
let body = openai_body(
&cfg(Provider::OpenAi),
"system",
&[HistoryItem::User("hi".into())],
&[],
"model",
Some(ThinkingEffort::Medium),
);
assert_eq!(body["reasoning_effort"], "medium");
}
#[test]
fn responses_body_omits_reasoning_when_effort_none() {
let body = responses_body(
&cfg_responses(),
"system",
&[HistoryItem::User("hi".into())],
&[],
"model",
None,
);
assert!(
body.get("reasoning").is_none(),
"reasoning must be absent when effort is None"
);
}
#[test]
fn responses_body_emits_reasoning_effort_low() {
let body = responses_body(
&cfg_responses(),
"system",
&[HistoryItem::User("hi".into())],
&[],
"model",
Some(ThinkingEffort::Low),
);
assert_eq!(body["reasoning"]["effort"], "low");
}
#[test]
fn effective_model_overrides_cfg_model_in_anthropic_body() {
let body = anthropic_body(
&cfg(Provider::Anthropic),
"system",
&[HistoryItem::User("hi".into())],
&[],
"override-model",
None,
);
assert_eq!(body["model"], "override-model");
}
#[test]
fn effective_model_overrides_cfg_model_in_openai_body() {
let body = openai_body(
&cfg(Provider::OpenAi),
"system",
&[HistoryItem::User("hi".into())],
&[],
"override-model",
None,
);
assert_eq!(body["model"], "override-model");
}
#[test]
fn anthropic_body_opus_4_8_xhigh_emits_xhigh_effort() {
// Body-shape regression: xhigh on Opus 4.8 must emit output_config.effort="xhigh".
let mut c = cfg(Provider::Anthropic);
c.max_output_tokens = 32_768;
let body = anthropic_body(
&c,
"system",
&[HistoryItem::User("hi".into())],
&[],
"claude-opus-4-8",
Some(ThinkingEffort::XHigh),
);
assert_eq!(body["thinking"]["type"], "adaptive");
assert_eq!(body["output_config"]["effort"], "xhigh");
}
#[test]
fn anthropic_body_opus_4_8_max_emits_max_effort() {
// Body-shape regression: max on Opus 4.8 must emit output_config.effort="max".
let mut c = cfg(Provider::Anthropic);
c.max_output_tokens = 32_768;
let body = anthropic_body(
&c,
"system",
&[HistoryItem::User("hi".into())],
&[],
"claude-opus-4-8",
Some(ThinkingEffort::Max),
);
assert_eq!(body["thinking"]["type"], "adaptive");
assert_eq!(body["output_config"]["effort"], "max");
}
#[test]
fn openai_body_emits_xhigh_effort() {
// xhigh is a valid OpenAI effort value — must pass through.
let body = openai_body(
&cfg(Provider::OpenAi),
"system",
&[HistoryItem::User("hi".into())],
&[],
"model",
Some(ThinkingEffort::XHigh),
);
assert_eq!(body["reasoning_effort"], "xhigh");
}
#[test]
fn openai_body_emits_none_effort() {
// none is a valid OpenAI effort value.
let body = openai_body(
&cfg(Provider::OpenAi),
"system",
&[HistoryItem::User("hi".into())],
&[],
"model",
Some(ThinkingEffort::None),
);
assert_eq!(body["reasoning_effort"], "none");
}
#[test]
fn responses_body_emits_xhigh_effort() {
// xhigh is a valid Responses API effort value.
let body = responses_body(
&cfg_responses(),
"system",
&[HistoryItem::User("hi".into())],
&[],
"model",
Some(ThinkingEffort::XHigh),
);
assert_eq!(body["reasoning"]["effort"], "xhigh");
}
#[test]
fn responses_body_emits_minimal_effort() {
// minimal is a valid Responses API effort value.
let body = responses_body(
&cfg_responses(),
"system",
&[HistoryItem::User("hi".into())],
&[],
"model",
Some(ThinkingEffort::Minimal),
);
assert_eq!(body["reasoning"]["effort"], "minimal");
}
// ---- DatabricksV2 route-aware effort normalization (body-level assertions) ----
//
// The DBv2 `complete()` dispatch applies `normalize_effort_for_openai_route` /
// `normalize_effort_for_anthropic_route` before calling body builders. These tests
// verify the body shape that results from the already-normalized effort values — i.e.,
// they confirm the body builders correctly serialize the values the dispatch passes them.
#[test]
fn dbv2_openai_route_max_effort_clamped_to_xhigh_in_responses_body() {
// DBv2 GPT-5.5 route: max → clamped to xhigh by normalize_effort_for_openai_route
// before reaching responses_body. gpt-5.5 supports xhigh so the final value is xhigh.
let clamped =
crate::config::normalize_effort_for_openai_route(ThinkingEffort::Max, "gpt-5.5");
let body = responses_body(
&cfg_responses(),
"system",
&[HistoryItem::User("hi".into())],
&[],
"gpt-5.5",
Some(clamped),
);
assert_eq!(
body["reasoning"]["effort"], "xhigh",
"DBv2 GPT-5.5 route: max must be clamped to xhigh before responses_body"
);
}
#[test]
fn dbv2_mlflow_route_max_effort_clamped_to_xhigh_in_openai_body() {
// DBv2 MLflow route (unknown model): max → clamped to xhigh by normalize_effort_for_openai_route.
// Unknown models pass through after the max→xhigh clamp.
let clamped =
crate::config::normalize_effort_for_openai_route(ThinkingEffort::Max, "llama-4");
let body = openai_body(
&cfg(Provider::OpenAi),
"system",
&[HistoryItem::User("hi".into())],
&[],
"llama-4",
Some(clamped),
);
assert_eq!(
body["reasoning_effort"], "xhigh",
"DBv2 MLflow route: max must be clamped to xhigh before openai_body"
);
}
#[test]
fn dbv2_openai_route_none_minimal_pass_through_in_responses_body() {
// Verify that supported values pass through for the respective model families.
// gpt-5.5 supports none (but not minimal); gpt-5 base supports minimal (but not none).
let none_normalized =
crate::config::normalize_effort_for_openai_route(ThinkingEffort::None, "gpt-5.5");
assert_eq!(
none_normalized,
ThinkingEffort::None,
"OpenAI normalizer must not touch none for gpt-5.5"
);
let minimal_normalized =
crate::config::normalize_effort_for_openai_route(ThinkingEffort::Minimal, "gpt-5");
assert_eq!(
minimal_normalized,
ThinkingEffort::Minimal,
"OpenAI normalizer must not touch minimal for gpt-5 base"
);
let body = responses_body(
&cfg_responses(),
"system",
&[HistoryItem::User("hi".into())],
&[],
"gpt-5.5",
Some(none_normalized),
);
assert_eq!(
body["reasoning"]["effort"], "none",
"DBv2 GPT-5.5 route: none must be emitted as-is"
);
}
#[test]
fn dbv2_claude_route_none_effort_omits_thinking_fields() {
// DBv2 Claude route: none → normalize_effort_for_anthropic_route returns None → omit.
let normalized = crate::config::normalize_effort_for_anthropic_route(ThinkingEffort::None);
assert_eq!(
normalized, None,
"Anthropic normalizer must return None for ThinkingEffort::None"
);
let mut c = cfg(Provider::Anthropic);
c.max_output_tokens = 32_768;
let body = anthropic_body(
&c,
"system",
&[HistoryItem::User("hi".into())],
&[],
"claude-opus-4-8",
normalized, // None → omit thinking fields
);
assert!(
body.get("thinking").is_none(),
"DBv2 Claude route: none effort must omit thinking fields"
);
assert!(
body.get("output_config").is_none(),
"DBv2 Claude route: none effort must omit output_config"
);
}
#[test]
fn dbv2_route_switch_max_body_level_simulation() {
// Body-level simulation of a session/set_model switch from a Claude model to a GPT-5
// model when thinking_effort=max. Calls body builders and normalizers directly (not
// through the ACP session/set_model path or DatabricksV2 dispatch) to verify the
// correct output shape for each side of the route switch.
// Before the switch: Claude route → max passes through as Anthropic "max".
// After the switch: GPT-5 route → max clamped to xhigh.
let mut c = cfg(Provider::Anthropic);
c.max_output_tokens = 32_768;
// Before switch: claude-opus-4-8 with effort=max → adaptive shape, effort="max"
let (thinking_before, oc_before) = crate::config::anthropic_thinking_config(
"claude-opus-4-8",
ThinkingEffort::Max,
32_768,
);
assert_eq!(thinking_before.unwrap()["type"], "adaptive");
assert_eq!(oc_before.unwrap()["effort"], "max");
// After switch to GPT-5.5 route: normalize max → xhigh for responses_body
// (gpt-5.5 supports xhigh, so the clamp result is xhigh, not further reduced)
let clamped =
crate::config::normalize_effort_for_openai_route(ThinkingEffort::Max, "gpt-5.5");
assert_eq!(clamped, ThinkingEffort::XHigh);
let body_after = responses_body(
&cfg_responses(),
"system",
&[HistoryItem::User("hi".into())],
&[],
"gpt-5.5",
Some(clamped),
);
assert_eq!(
body_after["reasoning"]["effort"], "xhigh",
"After set_model to GPT-5.5: max must be clamped to xhigh"
);
}
/// Regression: a connection that is accepted and then dropped before any
/// HTTP response bytes are written surfaces as a reqwest request-class
/// error (not `is_connect()`, not `is_timeout()`). The retry predicate
@@ -1761,7 +2360,7 @@ mod tests {
c.base_url = base;
let out = llm
.post_openai(&c, "/v1/x", &json!({}))
.post_openai(&c, "/v1/x", &json!({}), "model")
.await
.expect("retry with fresh token should succeed");
assert_eq!(out, json!({ "ok": true }));
@@ -1769,7 +2368,10 @@ mod tests {
// Second call's 401 must trigger its own refresh — the guard cannot
// be a stored flag that an earlier turn already tripped.
let out2 = llm.post_openai(&c, "/v1/x", &json!({})).await.unwrap();
let out2 = llm
.post_openai(&c, "/v1/x", &json!({}), "model")
.await
.unwrap();
assert_eq!(out2, json!({ "ok": true }));
assert_eq!(
auth.refreshes.load(Ordering::SeqCst),
@@ -1793,7 +2395,10 @@ mod tests {
let mut c = cfg(Provider::OpenAi);
c.base_url = base;
let err = llm.post_openai(&c, "/v1/x", &json!({})).await.unwrap_err();
let err = llm
.post_openai(&c, "/v1/x", &json!({}), "model")
.await
.unwrap_err();
assert!(matches!(err, AgentError::LlmAuth(_)), "got {err:?}");
assert_eq!(
auth.refreshes.load(Ordering::SeqCst),
@@ -1818,7 +2423,10 @@ mod tests {
let mut c = cfg(Provider::OpenAi);
c.base_url = base;
let err = llm.post_openai(&c, "/v1/x", &json!({})).await.unwrap_err();
let err = llm
.post_openai(&c, "/v1/x", &json!({}), "model")
.await
.unwrap_err();
assert!(matches!(err, AgentError::LlmAuth(_)), "got {err:?}");
assert_eq!(
auth.refreshes.load(Ordering::SeqCst),
@@ -1844,7 +2452,7 @@ mod tests {
c.base_url = base;
let out = llm
.post_openai(&c, "/v1/x", &json!({}))
.post_openai(&c, "/v1/x", &json!({}), "model")
.await
.expect("retry with fresh token should clear the 403");
assert_eq!(out, json!({ "ok": true }));
crates/buzz-agent/src/wire.rs
+10
View File
@@ -80,6 +80,16 @@ pub struct SessionSteerParams {
pub expected_run_id: String,
}
/// Params for `session/set_model`: override the active model for an existing
/// session without respawning. Applied immediately; subsequent prompts on this
/// session use `model_id` instead of the configured `BUZZ_AGENT_MODEL`.
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SessionSetModelParams {
pub session_id: String,
pub model_id: String,
}
pub fn classify(msg: &Value) -> Inbound {
if !msg.is_object() || msg.get("jsonrpc").and_then(Value::as_str) != Some("2.0") {
return Inbound::Invalid {
crates/buzz-agent/tests/databricks_oauth.rs
+253 -4
View File
@@ -544,14 +544,33 @@ async fn run_single_prompt(provider: &str, base: &str, model: &str) {
async fn run_captured_prompt(
provider: &str,
model: &str,
canned: Vec<serde_json::Value>,
llm_canned: Vec<serde_json::Value>,
) -> CapturedRequest {
let (base, captured) = spawn_capturing_server(canned).await;
// session/new triggers model catalog discovery against the same stub server.
// Prepend a minimal valid discovery response (empty endpoints list) so the
// discovery call is served cleanly and the LLM canned responses follow.
// Legacy Databricks discovery hits /api/2.0/serving-endpoints;
// Databricks v2 discovery hits /api/ai-gateway/v2/endpoints.
let discovery_resp = json!({ "endpoints": [], "next_page_token": null });
let mut all_canned = vec![discovery_resp];
all_canned.extend(llm_canned);
let (base, captured) = spawn_capturing_server(all_canned).await;
run_single_prompt(provider, &base, model).await;
// Filter out discovery requests — keep only the LLM invocation(s).
// Discovery paths: /api/2.0/serving-endpoints, /api/ai-gateway/v2/endpoints*
// LLM paths: /serving-endpoints/*, /ai-gateway/anthropic/*, /ai-gateway/openai/*, /ai-gateway/mlflow/*
let reqs = captured.lock().await;
assert_eq!(reqs.len(), 1, "expected exactly one LLM request");
reqs[0].clone()
let llm_reqs: Vec<_> = reqs
.iter()
.filter(|r| {
!r.path.starts_with("/api/2.0/serving-endpoints")
&& !r.path.starts_with("/api/ai-gateway/v2/endpoints")
})
.collect();
assert_eq!(llm_reqs.len(), 1, "expected exactly one LLM request");
llm_reqs[0].clone()
}
#[tokio::test]
@@ -689,3 +708,233 @@ async fn databricks_v2_other_models_route_through_ai_gateway_mlflow_chat() {
req.body
);
}
// ---------- session/set_model integration tests ----------
/// Helper: run initialize + session/new + optional set_model + session/prompt on a
/// freshly-spawned harness against the given stub server base URL. Returns the
/// session/prompt response and the session ID so callers can also call set_model.
async fn run_with_set_model(
provider: &str,
base: &str,
initial_model: &str,
switch_to_model: Option<&str>,
) -> (String, serde_json::Value) {
let mut h = AgentHarness::spawn_provider(provider, base, initial_model).await;
h.send(
"initialize",
json!({ "protocolVersion": 1, "clientCapabilities": {} }),
)
.await;
h.recv_for(1).await;
h.send("session/new", json!({ "cwd": "/tmp", "mcpServers": [] }))
.await;
let r = h.recv_for(2).await;
let sid = r["result"]["sessionId"].as_str().unwrap().to_string();
if let Some(new_model) = switch_to_model {
h.send(
"session/set_model",
json!({ "sessionId": sid, "modelId": new_model }),
)
.await;
let set_r = h.recv_for(3).await;
// Verify the response carries the expected modelId.
assert_eq!(
set_r["result"]["modelId"],
json!(new_model),
"set_model response must echo the new modelId"
);
h.send(
"session/prompt",
json!({ "sessionId": sid, "prompt": [{ "type": "text", "text": "say ok" }] }),
)
.await;
let prompt_r = h.recv_for(4).await;
(sid, prompt_r)
} else {
h.send(
"session/prompt",
json!({ "sessionId": sid, "prompt": [{ "type": "text", "text": "say ok" }] }),
)
.await;
let prompt_r = h.recv_for(3).await;
(sid, prompt_r)
}
}
/// After session/set_model switches from model A to model B, the next
/// session/prompt must route to B's Databricks serving-endpoint URL and
/// strip the `model` field from the body (legacy Databricks behaviour).
#[tokio::test]
async fn session_set_model_switches_databricks_legacy_route() {
let initial_model = "initial-model";
let switched_model = "switched-model";
// Two canned responses: one for the discovery call (session/new),
// one for the LLM call after the switch.
let canned = vec![
json!({ "endpoints": [], "next_page_token": null }), // discovery
json!({ // LLM response
"id": "x",
"object": "chat.completion",
"choices": [{
"index": 0,
"message": { "role": "assistant", "content": "ok" },
"finish_reason": "stop"
}]
}),
];
let (base, captured) = spawn_capturing_server(canned).await;
run_with_set_model("databricks", &base, initial_model, Some(switched_model)).await;
let reqs = captured.lock().await;
let llm_reqs: Vec<_> = reqs
.iter()
.filter(|r| {
!r.path.starts_with("/api/2.0/serving-endpoints")
&& !r.path.starts_with("/api/ai-gateway/v2/endpoints")
})
.collect();
assert_eq!(
llm_reqs.len(),
1,
"expected exactly one LLM request after switch"
);
let req = &llm_reqs[0];
// The request must go to the SWITCHED model endpoint, not the initial one.
assert_eq!(
req.path.as_str(),
format!("/serving-endpoints/{switched_model}/invocations"),
"Databricks legacy must route to the switched model endpoint"
);
// The body must not include `model` (Databricks rejects it).
assert!(
req.body.get("model").is_none(),
"request body must NOT include `model` after switch: {:?}",
req.body
);
}
/// After session/set_model switches a Databricks v2 session from a GPT-5 model
/// (OpenAI Responses route) to a Claude model (Anthropic Messages route),
/// the next prompt must hit the Anthropic AI Gateway path.
#[tokio::test]
async fn session_set_model_switches_databricks_v2_route() {
let initial_model = "databricks-gpt-5-5"; // → OpenAI Responses
let switched_model = "databricks-claude-opus-4-7"; // → Anthropic Messages
let canned = vec![
json!({ "endpoints": [], "next_page_token": null }), // discovery (v2: /api/ai-gateway/v2/endpoints)
json!({ // LLM response (Anthropic Messages shape)
"stop_reason": "end_turn",
"content": [{ "type": "text", "text": "ok" }]
}),
];
let (base, captured) = spawn_capturing_server(canned).await;
run_with_set_model("databricks_v2", &base, initial_model, Some(switched_model)).await;
let reqs = captured.lock().await;
let llm_reqs: Vec<_> = reqs
.iter()
.filter(|r| {
!r.path.starts_with("/api/2.0/serving-endpoints")
&& !r.path.starts_with("/api/ai-gateway/v2/endpoints")
})
.collect();
assert_eq!(
llm_reqs.len(),
1,
"expected exactly one LLM request after v2 route switch"
);
let req = &llm_reqs[0];
assert_eq!(
req.path.as_str(),
"/ai-gateway/anthropic/v1/messages",
"After switching to a Claude model, Databricks v2 must route to Anthropic Messages"
);
assert_eq!(
req.body["model"],
json!(switched_model),
"body must carry the switched model ID"
);
}
/// session/set_model with an unknown session ID must return an invalid_params
/// error without touching any LLM endpoint.
#[tokio::test]
async fn session_set_model_unknown_session_returns_error() {
// Spawn with a single discovery canned response; no LLM response needed.
let canned = vec![json!({ "endpoints": [], "next_page_token": null })];
let (base, _captured) = spawn_capturing_server(canned).await;
let mut h = AgentHarness::spawn_provider("databricks", &base, "some-model").await;
h.send(
"initialize",
json!({ "protocolVersion": 1, "clientCapabilities": {} }),
)
.await;
h.recv_for(1).await;
h.send("session/new", json!({ "cwd": "/tmp", "mcpServers": [] }))
.await;
h.recv_for(2).await;
// Call set_model with a bogus session ID.
h.send(
"session/set_model",
json!({ "sessionId": "nonexistent-session-id", "modelId": "new-model" }),
)
.await;
let r = h.recv_for(3).await;
assert!(
r.get("error").is_some(),
"set_model with unknown session must return an error: {:?}",
r
);
let msg = r["error"]["message"].as_str().unwrap_or("");
assert!(
msg.contains("unknown session"),
"error message must mention unknown session, got: {msg}"
);
}
/// session/set_model with an empty modelId must return an invalid_params error.
#[tokio::test]
async fn session_set_model_empty_model_id_returns_error() {
let canned = vec![json!({ "endpoints": [], "next_page_token": null })];
let (base, _captured) = spawn_capturing_server(canned).await;
let mut h = AgentHarness::spawn_provider("databricks", &base, "some-model").await;
h.send(
"initialize",
json!({ "protocolVersion": 1, "clientCapabilities": {} }),
)
.await;
h.recv_for(1).await;
h.send("session/new", json!({ "cwd": "/tmp", "mcpServers": [] }))
.await;
let r = h.recv_for(2).await;
let sid = r["result"]["sessionId"].as_str().unwrap().to_string();
// Empty string modelId.
h.send(
"session/set_model",
json!({ "sessionId": sid, "modelId": "" }),
)
.await;
let r = h.recv_for(3).await;
assert!(
r.get("error").is_some(),
"set_model with empty modelId must return an error: {:?}",
r
);
let msg = r["error"]["message"].as_str().unwrap_or("");
assert!(
msg.contains("modelId"),
"error message must mention modelId, got: {msg}"
);
}
desktop/scripts/check-file-sizes.mjs
+4 -1
View File
@@ -102,7 +102,10 @@ const overrides = new Map([
// agents keep an installed runtime alias when the primary command is absent.
// Load-bearing, not generic debt.
// config-bridge: schema-driven field extraction adds ~26 lines. Queued to split.
["src-tauri/src/managed_agents/discovery.rs", 1111],
// config-parity: max_tokens_env_var + context_limit_env_var fields added to
// KnownAcpRuntime (2 fields × 4 runtimes + discovery tests = ~13 lines).
// Load-bearing — required for buzz-agent normalized config parity.
["src-tauri/src/managed_agents/discovery.rs", 1124],
// migration_tests.rs carries the harness-sync migration coverage plus the
// patch_json_records owner-only writeback regression test (SECURITY.md:90
// crash-safe 0o600 fallback). Load-bearing security + feature coverage, not
desktop/src-tauri/src/commands/agent_config.rs
+2
View File
@@ -404,6 +404,8 @@ mod tests {
config_file_format: Some("yaml"),
supports_acp_native_config: true,
thinking_env_var: Some("GOOSE_THINKING_EFFORT"),
max_tokens_env_var: Some("GOOSE_MAX_TOKENS"),
context_limit_env_var: Some("GOOSE_CONTEXT_LIMIT"),
required_normalized_fields: &["model", "provider"],
}
}
desktop/src-tauri/src/managed_agents/config_bridge/reader.rs
+47 -19
View File
@@ -44,6 +44,8 @@ pub(crate) fn read_config_surface(
let required_fields: &[&str] = runtime_meta
.map(|m| m.required_normalized_fields)
.unwrap_or(&[]);
let max_tokens_env_var = runtime_meta.and_then(|m| m.max_tokens_env_var);
let context_limit_env_var = runtime_meta.and_then(|m| m.context_limit_env_var);
// Tier 1b: ACP configOptions from session cache.
// For unstable/switchable agents, current_model comes from the `models`
@@ -95,25 +97,16 @@ pub(crate) fn read_config_surface(
is_pre_spawn,
session_cache,
),
max_output_tokens: file_config
.max_output_tokens
.as_ref()
.map(|v| NormalizedField {
value: Some(v.clone()),
origin: ConfigOrigin::ConfigFile,
write_via: ConfigWriteMechanism::ReadOnly,
overridden_value: None,
overridden_origin: None,
is_required: false,
}),
context_limit: file_config.context_limit.as_ref().map(|v| NormalizedField {
value: Some(v.clone()),
origin: ConfigOrigin::ConfigFile,
write_via: ConfigWriteMechanism::ReadOnly,
overridden_value: None,
overridden_origin: None,
is_required: false,
}),
max_output_tokens: build_numeric_env_field(
max_tokens_env_var,
&record.env_vars,
&file_config.max_output_tokens,
),
context_limit: build_numeric_env_field(
context_limit_env_var,
&record.env_vars,
&file_config.context_limit,
),
system_prompt: build_system_prompt_field(
&record
.system_prompt
@@ -142,6 +135,8 @@ pub(crate) fn read_config_surface(
model_env_var,
provider_env_var,
thinking_env_var,
max_tokens_env_var,
context_limit_env_var,
Some("BUZZ_ACP_SYSTEM_PROMPT"),
]
.into_iter()
@@ -452,6 +447,39 @@ fn build_thinking_field(
})
}
/// Numeric fields (max_output_tokens, context_limit) — env-var tier wins over
/// config-file tier. When an env var key is given and present in the record's
/// env_vars map the field is BuzzExplicit + RespawnWithEnvVar; otherwise if the
/// config file supplied a value it is ConfigFile + ReadOnly; otherwise None.
fn build_numeric_env_field(
env_var: Option<&'static str>,
record_env: &std::collections::BTreeMap<String, String>,
file_value: &Option<String>,
) -> Option<NormalizedField> {
if let Some(key) = env_var {
if let Some(v) = record_env.get(key) {
return Some(NormalizedField {
value: Some(v.clone()),
origin: ConfigOrigin::BuzzExplicit,
write_via: ConfigWriteMechanism::RespawnWithEnvVar {
env_key: key.to_string(),
},
overridden_value: file_value.clone(),
overridden_origin: file_value.as_ref().map(|_| ConfigOrigin::ConfigFile),
is_required: false,
});
}
}
file_value.as_ref().map(|v| NormalizedField {
value: Some(v.clone()),
origin: ConfigOrigin::ConfigFile,
write_via: ConfigWriteMechanism::ReadOnly,
overridden_value: None,
overridden_origin: None,
is_required: false,
})
}
/// Record/env prompt wins (BuzzExplicit, respawnable); a config-file prompt it
/// shadows is reported as the overridden secondary. A config-file-only prompt
/// — no record/env value to shadow it — is surfaced directly (read-only)
desktop/src-tauri/src/managed_agents/config_bridge/reader_tests.rs
+164
View File
@@ -33,6 +33,8 @@ fn test_runtime() -> &'static KnownAcpRuntime {
config_file_format: Some("yaml"),
supports_acp_native_config: true,
thinking_env_var: Some("GOOSE_THINKING_EFFORT"),
max_tokens_env_var: Some("GOOSE_MAX_TOKENS"),
context_limit_env_var: Some("GOOSE_CONTEXT_LIMIT"),
required_normalized_fields: &["model", "provider"],
}
}
@@ -495,3 +497,165 @@ fn extra_env_var_skipped_when_already_in_file_config_extra() {
"normalized thinking key must not appear in advanced"
);
}
// ── buzz-agent normalized env-var field tests ───────────────────────────────
//
// buzz-agent uses env vars (not a config file) for max_output_tokens and
// context_limit. build_numeric_env_field must surface these as BuzzExplicit
// when the env var is present in record.env_vars, and must not double-surface
// them in the advanced tier.
fn buzz_agent_runtime() -> &'static KnownAcpRuntime {
&KnownAcpRuntime {
id: "buzz-agent",
label: "Buzz Agent",
commands: &["buzz-agent"],
aliases: &[],
avatar_url: "",
mcp_command: None,
mcp_hooks: false,
underlying_cli: None,
cli_install_commands: &[],
adapter_install_commands: &[],
install_instructions_url: "",
cli_install_hint: "",
adapter_install_hint: "",
skill_dir: None,
supports_acp_model_switching: true,
model_env_var: Some("BUZZ_AGENT_MODEL"),
provider_env_var: Some("BUZZ_AGENT_PROVIDER"),
provider_locked: false,
default_env: &[],
config_file_path: None,
config_file_format: None,
supports_acp_native_config: false,
thinking_env_var: Some("BUZZ_AGENT_THINKING_EFFORT"),
max_tokens_env_var: Some("BUZZ_AGENT_MAX_OUTPUT_TOKENS"),
context_limit_env_var: Some("BUZZ_AGENT_MAX_CONTEXT_TOKENS"),
required_normalized_fields: &["model", "provider"],
}
}
#[test]
fn buzz_agent_max_output_tokens_from_env_is_buzz_explicit() {
let mut record = test_record();
record.env_vars.insert(
"BUZZ_AGENT_MAX_OUTPUT_TOKENS".to_string(),
"8192".to_string(),
);
let runtime = buzz_agent_runtime();
let surface = read_config_surface(&record, Some(runtime), None, None);
let field = surface.normalized.max_output_tokens.unwrap();
assert_eq!(field.value.as_deref(), Some("8192"));
assert_eq!(field.origin, ConfigOrigin::BuzzExplicit);
assert!(matches!(
field.write_via,
ConfigWriteMechanism::RespawnWithEnvVar { ref env_key }
if env_key == "BUZZ_AGENT_MAX_OUTPUT_TOKENS"
));
}
#[test]
fn buzz_agent_context_limit_from_env_is_buzz_explicit() {
let mut record = test_record();
record.env_vars.insert(
"BUZZ_AGENT_MAX_CONTEXT_TOKENS".to_string(),
"100000".to_string(),
);
let runtime = buzz_agent_runtime();
let surface = read_config_surface(&record, Some(runtime), None, None);
let field = surface.normalized.context_limit.unwrap();
assert_eq!(field.value.as_deref(), Some("100000"));
assert_eq!(field.origin, ConfigOrigin::BuzzExplicit);
assert!(matches!(
field.write_via,
ConfigWriteMechanism::RespawnWithEnvVar { ref env_key }
if env_key == "BUZZ_AGENT_MAX_CONTEXT_TOKENS"
));
}
#[test]
fn buzz_agent_max_tokens_absent_when_no_env_var_or_file() {
// buzz-agent has no config file, and env var is not set.
let record = test_record();
let runtime = buzz_agent_runtime();
let surface = read_config_surface(&record, Some(runtime), None, None);
assert!(
surface.normalized.max_output_tokens.is_none(),
"max_output_tokens must be None when env var not set and no config file"
);
assert!(
surface.normalized.context_limit.is_none(),
"context_limit must be None when env var not set and no config file"
);
}
#[test]
fn buzz_agent_max_tokens_env_var_not_double_surfaced_in_advanced() {
let mut record = test_record();
record.env_vars.insert(
"BUZZ_AGENT_MAX_OUTPUT_TOKENS".to_string(),
"4096".to_string(),
);
record.env_vars.insert(
"BUZZ_AGENT_MAX_CONTEXT_TOKENS".to_string(),
"50000".to_string(),
);
let runtime = buzz_agent_runtime();
let surface = read_config_surface(&record, Some(runtime), None, None);
let advanced_keys: Vec<&str> = surface.advanced.iter().map(|f| f.key.as_str()).collect();
assert!(
!advanced_keys.contains(&"BUZZ_AGENT_MAX_OUTPUT_TOKENS"),
"max_output_tokens must not appear in advanced when normalized"
);
assert!(
!advanced_keys.contains(&"BUZZ_AGENT_MAX_CONTEXT_TOKENS"),
"context_limit must not appear in advanced when normalized"
);
}
#[test]
fn buzz_agent_thinking_effort_from_env_is_buzz_explicit() {
let mut record = test_record();
record
.env_vars
.insert("BUZZ_AGENT_THINKING_EFFORT".to_string(), "high".to_string());
let runtime = buzz_agent_runtime();
let surface = read_config_surface(&record, Some(runtime), None, None);
let field = surface.normalized.thinking_effort.unwrap();
assert_eq!(field.value.as_deref(), Some("high"));
assert_eq!(field.origin, ConfigOrigin::BuzzExplicit);
assert!(matches!(
field.write_via,
ConfigWriteMechanism::RespawnWithEnvVar { ref env_key }
if env_key == "BUZZ_AGENT_THINKING_EFFORT"
));
}
#[test]
fn buzz_agent_thinking_effort_env_var_not_double_surfaced_in_advanced() {
let mut record = test_record();
record.env_vars.insert(
"BUZZ_AGENT_THINKING_EFFORT".to_string(),
"medium".to_string(),
);
let runtime = buzz_agent_runtime();
let surface = read_config_surface(&record, Some(runtime), None, None);
let advanced_keys: Vec<&str> = surface.advanced.iter().map(|f| f.key.as_str()).collect();
assert!(
!advanced_keys.contains(&"BUZZ_AGENT_THINKING_EFFORT"),
"thinking_effort must not appear in advanced when normalized"
);
}
desktop/src-tauri/src/managed_agents/discovery.rs
+14 -1
View File
@@ -47,6 +47,11 @@ pub(crate) struct KnownAcpRuntime {
pub config_file_format: Option<&'static str>,
pub supports_acp_native_config: bool, // tier 1a: config/read+write
pub thinking_env_var: Option<&'static str>,
/// Env var for normalizing `max_output_tokens`. `None` when the harness
/// does not have a first-class env var for this field (config-file only).
pub max_tokens_env_var: Option<&'static str>,
/// Env var for normalizing `context_limit`. `None` when not applicable.
pub context_limit_env_var: Option<&'static str>,
/// Normalized field keys that must be set for this harness to function.
/// Used by the config bridge to mark fields as required in the UI.
/// Keys match the camelCase names used in `NormalizedConfig` (e.g. "model", "provider").
@@ -106,6 +111,8 @@ const KNOWN_ACP_RUNTIMES: &[KnownAcpRuntime] = &[
config_file_format: Some("yaml"),
supports_acp_native_config: true,
thinking_env_var: Some("GOOSE_THINKING_EFFORT"),
max_tokens_env_var: Some("GOOSE_MAX_TOKENS"),
context_limit_env_var: Some("GOOSE_CONTEXT_LIMIT"),
required_normalized_fields: &["model", "provider"],
},
KnownAcpRuntime {
@@ -132,6 +139,8 @@ const KNOWN_ACP_RUNTIMES: &[KnownAcpRuntime] = &[
config_file_format: Some("json"),
supports_acp_native_config: false,
thinking_env_var: None,
max_tokens_env_var: None,
context_limit_env_var: None,
required_normalized_fields: &[],
},
KnownAcpRuntime {
@@ -158,6 +167,8 @@ const KNOWN_ACP_RUNTIMES: &[KnownAcpRuntime] = &[
config_file_format: Some("toml"),
supports_acp_native_config: false,
thinking_env_var: None,
max_tokens_env_var: None,
context_limit_env_var: None,
required_normalized_fields: &[],
},
KnownAcpRuntime {
@@ -183,7 +194,9 @@ const KNOWN_ACP_RUNTIMES: &[KnownAcpRuntime] = &[
config_file_path: None,
config_file_format: None,
supports_acp_native_config: false,
thinking_env_var: None,
thinking_env_var: Some("BUZZ_AGENT_THINKING_EFFORT"),
max_tokens_env_var: Some("BUZZ_AGENT_MAX_OUTPUT_TOKENS"),
context_limit_env_var: Some("BUZZ_AGENT_MAX_CONTEXT_TOKENS"),
required_normalized_fields: &["model", "provider"],
},
];