jerry/buzz
1
0
Fork 0
mirror of https://github.com/block/buzz.git synced 2026-10-30 22:11:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(desktop): add databricks v1+v2 model discovery via buzz-agent lib

Browse Source 
Implements in-process Databricks model catalog discovery (Option C) so
PKCE token-cache hits are covered without spawning a subprocess.

buzz-agent changes:
- catalog.rs: new pub module with discover_databricks_models(cfg) —
  v1 (serving-endpoints) and v2 (ai-gateway/v2/endpoints, paginated)
  HTTP fetch + parse helpers; 7 unit tests
- auth.rs: bearer_no_browser() trait default + PkceOAuthTokenSource
  override that stops before the browser step (returns Err(LlmAuth)
  when cache is empty)
- config.rs: Config::for_discovery() minimal constructor — avoids
  requiring DATABRICKS_MODEL/other fields irrelevant to discovery
- llm.rs: build_token_source() pub(crate) for catalog.rs access
- lib.rs: re-exports discover_databricks_models, ModelEntry,
  DATABRICKS_V2_KNOWN_MODELS, Provider, AgentError

desktop changes:
- Cargo.toml: buzz_agent_pkg path dep added
- agent_models.rs: discover_databricks_models wrapper calls
  buzz_agent_pkg::discover_databricks_models directly; LlmAuth ->
  Ok(None) fallthrough; other errors redacted; ~270 lines removed
- agent_models_tests.rs: stale parse/URL tests removed (moved to
  catalog::tests); is_databricks_provider test retained
- check-file-sizes.mjs: TEMP override lowered from 1244 -> 1066

Co-authored-by: Will Pfleger <pfleger.will@gmail.com>
Signed-off-by: Will Pfleger <pfleger.will@gmail.com>
This commit is contained in:
npub1mn7jgtj4w2pd0g0zeuhxsa6jy6p0rewxz4kujt98my82ahfmp72sxjexk7
2026-06-29 21:44:27 -04:00
parent fd9270e8fd
commit d659e90ffe
10 changed files with 675 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

71
crates/buzz-agent/src/auth.rs
View File

@@ -44,6 +44,15 @@ const BROWSER_AUTH_TIMEOUT: Duration = Duration::from_secs(60);
pub trait TokenSource: Send + Sync {
async fn bearer(&self) -> Result<String, AgentError>;
/// Return a bearer token from cache or refresh, **never** opening a browser.
///
/// The default delegates to [`bearer`](Self::bearer) — correct for token
/// sources (e.g. static API keys) that can never trigger a browser flow.
/// [`PkceOAuthTokenSource`] overrides this to stop before the browser step.
async fn bearer_no_browser(&self) -> Result<String, AgentError> {
self.bearer().await
}
/// Force a fresh bearer after the server rejected the current one (401).
///
/// `rejected` is the exact access token that just got the 401. Unlike
@@ -287,6 +296,10 @@ impl TokenSource for PkceOAuthTokenSource {
Ok(bearer)
}
async fn bearer_no_browser(&self) -> Result<String, AgentError> {
self.try_bearer_no_browser().await
}
/// Force-refresh after a 401, never touching the browser flow.
///
/// `rejected` is the access token the server just 401'd. Coalescing keys
@@ -344,6 +357,64 @@ impl TokenSource for PkceOAuthTokenSource {
}
}
impl PkceOAuthTokenSource {
/// Return a bearer token from cache or refresh, **never** opening a browser.
///
/// Follows the same steps as [`bearer`](TokenSource::bearer) but stops at
/// step 4 — if no usable token is available after cache + refresh attempts,
/// returns `Err(LlmAuth(...))` instead of launching the browser PKCE flow.
/// Used by model-discovery paths that must not block on user interaction.
pub(crate) async fn try_bearer_no_browser(&self) -> Result<String, AgentError> {
let mut state = self.state.lock().await;
// 1. In-memory cache hit, still fresh.
if let Some(tok) = state.as_ref() {
if !is_expired(tok) {
return Ok(tok.access_token.clone());
}
}
// 2. Re-read disk — another process may have refreshed already.
if let Some(disk_tok) = read_cache(&self.cache_path) {
if !is_expired(&disk_tok) {
let bearer = disk_tok.access_token.clone();
*state = Some(disk_tok);
return Ok(bearer);
}
}
// 3. Try refresh if we have a refresh token.
let endpoints = self.endpoints().await?;
let refresh = state.as_ref().and_then(|t| t.refresh_token.clone());
if let Some(rt) = refresh {
match self.refresh(&endpoints, &rt).await {
Ok(fresh) => {
let bearer = fresh.access_token.clone();
self.save(&mut state, fresh)?;
return Ok(bearer);
}
Err(e) => {
tracing::warn!(error = %e, "oauth refresh failed during model discovery");
}
}
// 4. Re-read disk after refresh failure.
if let Some(disk_tok) = read_cache(&self.cache_path) {
if !is_expired(&disk_tok) {
let bearer = disk_tok.access_token.clone();
*state = Some(disk_tok);
return Ok(bearer);
}
}
}
// No usable token — return error instead of opening a browser.
Err(AgentError::LlmAuth(
"no cached Databricks token; run `buzz-agent auth databricks` first".into(),
))
}
}
// ---- helpers -------------------------------------------------------------
/// Aborts a spawned task when dropped. Used to guarantee the localhost

369
crates/buzz-agent/src/catalog.rs Normal file
View File

@@ -0,0 +1,369 @@
//! Databricks model catalog discovery.
//!
//! Exposes [`discover_databricks_models`] — an async helper that lists
//! available models for the `databricks` and `databricks_v2` providers
//! without triggering a browser OAuth flow. Auth is acquired in-process via
//! [`build_token_source`](crate::llm::build_token_source):
//!
//! - Static bearer (`DATABRICKS_TOKEN`): returned immediately.
//! - PKCE cache hit: returned from disk without a network round-trip.
//! - PKCE cache empty / no token: returns `Err(AgentError::LlmAuth)` — the
//! caller degrades gracefully; no browser, no hang.
use reqwest::Client;
use crate::{
config::{Config, Provider},
llm::build_token_source,
types::AgentError,
};
/// A discovered model entry: `id` is the picker value, `name` is the display
/// label (same as `id` for Databricks — the API has no separate display name).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelEntry {
pub id: String,
pub name: String,
}
/// Known Databricks AI Gateway v2 models — used as a fallback when the
/// `api/ai-gateway/v2/endpoints` call returns an empty list.
/// Mirrors goose's `DATABRICKS_V2_KNOWN_MODELS`.
pub const DATABRICKS_V2_KNOWN_MODELS: &[&str] =
&["databricks-gpt-5-5", "databricks-claude-opus-4-7"];
/// Discover available models for a Databricks provider.
///
/// Returns a non-empty `Vec<ModelEntry>` on success. Returns
/// `Err(AgentError::LlmAuth)` when no token is available (no static token,
/// no PKCE cache) — callers should degrade gracefully rather than hanging.
///
/// # Panics
/// Never panics.
pub async fn discover_databricks_models(cfg: &Config) -> Result<Vec<ModelEntry>, AgentError> {
let token_source = build_token_source(cfg)?;
let bearer = token_source.bearer_no_browser().await?;
let http = Client::new();
let host = cfg.base_url.trim_end_matches('/');
match cfg.provider {
Provider::Databricks => fetch_v1_models(&http, host, &bearer).await,
Provider::DatabricksV2 => fetch_v2_models(&http, host, &bearer).await,
_ => Err(AgentError::InvalidParams(
"discover_databricks_models called for non-Databricks provider".into(),
)),
}
}
// ---------------------------------------------------------------------------
// v1 — api/2.0/serving-endpoints
// ---------------------------------------------------------------------------
async fn fetch_v1_models(
http: &Client,
host: &str,
bearer: &str,
) -> Result<Vec<ModelEntry>, AgentError> {
let url = format!("{host}/api/2.0/serving-endpoints");
let response = http
.get(&url)
.bearer_auth(bearer)
.send()
.await
.map_err(|e| AgentError::Llm(format!("Databricks model discovery request failed: {e}")))?;
let status = response.status();
if !status.is_success() {
let body = response.text().await.unwrap_or_default();
return Err(AgentError::Llm(format!(
"Databricks model discovery HTTP {status}: {body}"
)));
}
let json: serde_json::Value = response.json().await.map_err(|e| {
AgentError::Llm(format!(
"Databricks model discovery response parse failed: {e}"
))
})?;
parse_v1_endpoints(&json)
}
/// Parse a `GET api/2.0/serving-endpoints` response.
///
/// Filters to endpoints that are READY and serve an LLM chat/completions task.
/// When `state.ready` or `task` is absent the endpoint is included — prefer
/// including over silently dropping, per spec.
pub(crate) fn parse_v1_endpoints(json: &serde_json::Value) -> Result<Vec<ModelEntry>, AgentError> {
let endpoints = json
.get("endpoints")
.and_then(|v| v.as_array())
.ok_or_else(|| {
AgentError::Llm(
"Databricks model discovery: unexpected response (missing 'endpoints' array)"
.into(),
)
})?;
let models = endpoints
.iter()
.filter_map(|endpoint| {
let name = endpoint.get("name")?.as_str()?.to_string();
// Require READY state when present; include when absent.
let state_ready = endpoint
.get("state")
.and_then(|s| s.get("ready"))
.and_then(|r| r.as_str())
.map(|r| r == "READY")
.unwrap_or(true);
if !state_ready {
return None;
}
// Require LLM chat or completions task when present.
let task_ok = endpoint
.get("task")
.and_then(|t| t.as_str())
.map(|t| t == "llm/v1/chat" || t == "llm/v1/completions")
.unwrap_or(true);
if !task_ok {
return None;
}
Some(ModelEntry {
id: name.clone(),
name,
})
})
.collect();
Ok(models)
}
// ---------------------------------------------------------------------------
// v2 — api/ai-gateway/v2/endpoints (paginated)
// ---------------------------------------------------------------------------
/// Percent-encode a string for use as a URL query parameter value.
/// Only encodes characters that are not unreserved (RFC 3986).
fn percent_encode(s: &str) -> String {
s.bytes()
.flat_map(|b| match b {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
vec![b as char]
}
_ => format!("%{b:02X}").chars().collect(),
})
.collect()
}
async fn fetch_v2_models(
http: &Client,
host: &str,
bearer: &str,
) -> Result<Vec<ModelEntry>, AgentError> {
let mut all_models: Vec<ModelEntry> = Vec::new();
let mut page_token: Option<String> = None;
let base_url = format!("{host}/api/ai-gateway/v2/endpoints");
// Cap at 20 pages (2 000 endpoints) to bound execution time.
for _ in 0..20 {
// Build URL with query params manually — avoids requiring the `query`
// reqwest feature in buzz-agent's Cargo.toml.
let url = match &page_token {
Some(tok) => format!(
"{base_url}?page_size=100&page_token={}",
percent_encode(tok)
),
None => format!("{base_url}?page_size=100"),
};
let response = http
.get(&url)
.bearer_auth(bearer)
.send()
.await
.map_err(|e| {
AgentError::Llm(format!("Databricks v2 model discovery request failed: {e}"))
})?;
let status = response.status();
if !status.is_success() {
let body = response.text().await.unwrap_or_default();
return Err(AgentError::Llm(format!(
"Databricks v2 model discovery HTTP {status}: {body}"
)));
}
let json: serde_json::Value = response.json().await.map_err(|e| {
AgentError::Llm(format!(
"Databricks v2 model discovery response parse failed: {e}"
))
})?;
let (page_models, next) = parse_v2_endpoints_page(&json)?;
all_models.extend(page_models);
match next {
Some(tok) if Some(&tok) != page_token.as_ref() => page_token = Some(tok),
_ => break,
}
}
// Fall back to known-model list if the API returned nothing.
if all_models.is_empty() {
all_models = DATABRICKS_V2_KNOWN_MODELS
.iter()
.map(|id| ModelEntry {
id: id.to_string(),
name: id.to_string(),
})
.collect();
}
Ok(all_models)
}
/// Parse one page of a `GET api/ai-gateway/v2/endpoints` response.
///
/// Returns `(models, next_page_token)`. An empty or absent `next_page_token`
/// signals the last page.
pub(crate) fn parse_v2_endpoints_page(
json: &serde_json::Value,
) -> Result<(Vec<ModelEntry>, Option<String>), AgentError> {
let endpoints = json
.get("endpoints")
.and_then(|v| v.as_array())
.ok_or_else(|| {
AgentError::Llm(
"Databricks v2 model discovery: unexpected response (missing 'endpoints' array)"
.into(),
)
})?;
let models = endpoints
.iter()
.filter_map(|endpoint| {
let name = endpoint.get("name")?.as_str()?.to_string();
Some(ModelEntry {
id: name.clone(),
name,
})
})
.collect();
let next_page_token = json
.get("next_page_token")
.and_then(|v| v.as_str())
.filter(|token| !token.is_empty())
.map(str::to_string);
Ok((models, next_page_token))
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn v1_parse_filters_ready_chat_endpoints() {
let json = serde_json::json!({
"endpoints": [
// included: READY + llm/v1/chat
{"name": "my-llm", "state": {"ready": "READY"}, "task": "llm/v1/chat"},
// included: READY + llm/v1/completions
{"name": "my-completions", "state": {"ready": "READY"}, "task": "llm/v1/completions"},
// excluded: NOT_READY
{"name": "dead-endpoint", "state": {"ready": "NOT_READY"}, "task": "llm/v1/chat"},
// excluded: wrong task
{"name": "embedding-ep", "state": {"ready": "READY"}, "task": "llm/v1/embedding"},
// included: no state field → include by default
{"name": "no-state", "task": "llm/v1/chat"},
// included: no task field → include by default
{"name": "no-task", "state": {"ready": "READY"}},
]
});
let models = parse_v1_endpoints(&json).unwrap();
let ids: Vec<&str> = models.iter().map(|m| m.id.as_str()).collect();
assert_eq!(ids, vec!["my-llm", "my-completions", "no-state", "no-task"]);
}
#[test]
fn v1_parse_errors_on_missing_endpoints_array() {
let json = serde_json::json!({"data": []});
let err = parse_v1_endpoints(&json).unwrap_err();
assert!(
err.to_string().contains("missing 'endpoints' array"),
"got: {err}"
);
}
#[test]
fn v1_parse_empty_endpoints_returns_empty_vec() {
let json = serde_json::json!({"endpoints": []});
let models = parse_v1_endpoints(&json).unwrap();
assert!(models.is_empty());
}
#[test]
fn v2_parse_extracts_names_and_page_token() {
let json = serde_json::json!({
"endpoints": [
{"name": "databricks-claude-opus-4-7"},
{"name": "databricks-gpt-5-5"},
{"name": "custom-model"}
],
"next_page_token": "tok123"
});
let (models, next) = parse_v2_endpoints_page(&json).unwrap();
let ids: Vec<&str> = models.iter().map(|m| m.id.as_str()).collect();
assert_eq!(
ids,
vec![
"databricks-claude-opus-4-7",
"databricks-gpt-5-5",
"custom-model"
]
);
assert_eq!(next.as_deref(), Some("tok123"));
}
#[test]
fn v2_parse_empty_token_signals_last_page() {
let json = serde_json::json!({
"endpoints": [{"name": "only-model"}],
"next_page_token": ""
});
let (models, next) = parse_v2_endpoints_page(&json).unwrap();
assert_eq!(models.len(), 1);
assert!(
next.is_none(),
"empty token should be treated as no more pages"
);
}
#[test]
fn v2_parse_absent_token_signals_last_page() {
let json = serde_json::json!({"endpoints": [{"name": "only-model"}]});
let (_, next) = parse_v2_endpoints_page(&json).unwrap();
assert!(next.is_none());
}
#[test]
fn v2_parse_errors_on_missing_endpoints_array() {
let json = serde_json::json!({"data": []});
let err = parse_v2_endpoints_page(&json).unwrap_err();
assert!(
err.to_string().contains("missing 'endpoints' array"),
"got: {err}"
);
}
}

37
crates/buzz-agent/src/config.rs
View File

@@ -192,6 +192,43 @@ impl Config {
Ok(cfg)
}
/// Construct a minimal `Config` for model-catalog discovery.
///
/// Only the fields used by [`build_token_source`](crate::llm::build_token_source)
/// and the catalog HTTP helpers are meaningful; all others are set to
/// inert defaults. Never call `from_env` for discovery — it requires
/// `DATABRICKS_MODEL` and other fields that are irrelevant here.
pub fn for_discovery(provider: Provider, api_key: String, base_url: String) -> Self {
Self {
provider,
api_key,
base_url,
model: String::new(),
system_prompt: String::new(),
anthropic_api_version: "2023-06-01".into(),
openai_api: OpenAiApi::Chat,
max_rounds: 0,
max_output_tokens: 1,
llm_timeout: Duration::from_secs(30),
tool_timeout: Duration::from_secs(30),
mcp_init_timeout: Duration::from_secs(30),
mcp_max_restart_attempts: 0,
mcp_restart_base_ms: 0,
mcp_restart_max_ms: 0,
max_sessions: 1,
max_line_bytes: 4 * 1024 * 1024,
max_history_bytes: 16 * 1024 * 1024,
max_tool_result_text_bytes: 50 * 1024,
max_context_tokens: 200_001,
max_handoffs: 0,
max_parallel_tools: 1,
hook_timeout: Duration::from_secs(1),
stop_max_rejections: 0,
hook_servers: HookServers::None,
hints_enabled: false,
}
}
fn validate(&self) -> Result<(), String> {
const MIN_HISTORY_BYTES: usize = 4096;
const MIN_LINE_BYTES: usize = 1024;

9
crates/buzz-agent/src/lib.rs
View File

@@ -2,14 +2,19 @@
mod agent;
pub mod auth;
mod builtin;
mod config;
pub mod catalog;
pub mod config;
mod handoff;
mod hints;
mod llm;
mod mcp;
mod types;
pub mod types;
mod wire;
pub use catalog::{discover_databricks_models, ModelEntry, DATABRICKS_V2_KNOWN_MODELS};
pub use config::Provider;
pub use types::AgentError;
use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;

2
crates/buzz-agent/src/llm.rs
View File

@@ -982,7 +982,7 @@ where
/// Otherwise a `PkceOAuthTokenSource` pointed at the workspace's OIDC
/// discovery URL. First request without a cached token triggers a browser
/// flow; subsequent requests use the cache + refresh transparently.
fn build_token_source(cfg: &Config) -> Result<Arc<dyn TokenSource>, AgentError> {
pub(crate) fn build_token_source(cfg: &Config) -> Result<Arc<dyn TokenSource>, AgentError> {
match cfg.provider {
Provider::Anthropic | Provider::OpenAi => {
Ok(Arc::new(StaticTokenSource::new(cfg.api_key.clone())))

5
desktop/scripts/check-file-sizes.mjs
View File

@@ -132,6 +132,11 @@ const overrides = new Map([
["src/shared/ui/markdown.tsx", 2082],
["src/shared/ui/VideoPlayer.tsx", 2199],
["src/shared/ui/sidebar.tsx", 1042],
// Option C databricks-model-discovery: parse/HTTP logic moved to buzz-agent
// catalog module; agent_models.rs retains the thin wrapper (~50 lines).
// File still exceeds 1000 due to OpenAI/Anthropic discovery + subprocess
// fallback. Queued to split into dedicated discovery modules.
["src-tauri/src/commands/agent_models.rs", 1066],
]);
await runFileSizeCheck({

56
desktop/src-tauri/Cargo.lock generated
View File

@@ -535,6 +535,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90"
dependencies = [
"axum-core",
"axum-macros",
"base64 0.22.1",
"bytes",
"form_urlencoded",
"futures-util",
@@ -553,8 +555,10 @@ dependencies = [
"serde_json",
"serde_path_to_error",
"serde_urlencoded",
"sha1 0.10.6",
"sync_wrapper",
"tokio",
"tokio-tungstenite 0.29.0",
"tower",
"tower-layer",
"tower-service",
@@ -580,6 +584,17 @@ dependencies = [
"tracing",
]
[[package]]
name = "axum-macros"
version = "0.5.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7aa268c23bfbbd2c4363b9cd302a4f504fb2a9dfe7e3451d66f35dd392e20aca"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.118",
]
[[package]]
name = "backon"
version = "1.6.0"
@@ -824,6 +839,30 @@ version = "3.20.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649"
[[package]]
name = "buzz-agent"
version = "0.1.0"
dependencies = [
"arc-swap",
"async-trait",
"axum",
"base64 0.22.1",
"getrandom 0.4.2",
"hex",
"nix 0.31.3",
"reqwest 0.13.4",
"rmcp",
"serde",
"serde_json",
"serde_yaml",
"sha2 0.11.0",
"tokio",
"tracing",
"tracing-subscriber",
"urlencoding",
"webbrowser",
]
[[package]]
name = "buzz-core"
version = "0.1.0"
@@ -853,6 +892,7 @@ dependencies = [
"audioadapter-buffers",
"axum",
"base64 0.22.1",
"buzz-agent",
"buzz-core",
"buzz-persona",
"buzz-sdk",
@@ -10471,6 +10511,22 @@ dependencies = [
"string_cache_codegen",
]
[[package]]
name = "webbrowser"
version = "1.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0fc95580916af1e68ff6a7be07446fc5db73ebf71cf092de939bbf5f7e189f72"
dependencies = [
"core-foundation 0.10.1",
"jni 0.22.4",
"log",
"ndk-context",
"objc2",
"objc2-foundation",
"url",
"web-sys",
]
[[package]]
name = "webkit2gtk"
version = "2.0.2"

1
desktop/src-tauri/Cargo.toml
View File

@@ -83,6 +83,7 @@ url = "2"
buzz_core_pkg = { package = "buzz-core", path = "../../crates/buzz-core" }
buzz_persona_pkg = { package = "buzz-persona", path = "../../crates/buzz-persona" }
buzz_sdk_pkg = { package = "buzz-sdk", path = "../../crates/buzz-sdk" }
buzz_agent_pkg = { package = "buzz-agent", path = "../../crates/buzz-agent" }
mesh-llm-sdk = { git = "https://github.com/Mesh-LLM/mesh-llm.git", rev = "ebc3ab4c4b5f4a45d5bc942c66c18583800c3f82", package = "mesh-llm-sdk", default-features = false, features = ["client", "serving"], optional = true }
mesh-llm-host-runtime = { git = "https://github.com/Mesh-LLM/mesh-llm.git", rev = "ebc3ab4c4b5f4a45d5bc942c66c18583800c3f82", package = "mesh-llm-host-runtime", default-features = false, features = ["dynamic-native-runtime"], optional = true }
base64 = "0.22"

112
desktop/src-tauri/src/commands/agent_models.rs
View File

@@ -112,6 +112,17 @@ pub async fn get_agent_models(
return Ok(models);
}
if let Some(models) = discover_databricks_models(
&state.http_client,
effective_provider.as_deref(),
&merged_env,
persisted_model.clone(),
)
.await?
{
return Ok(models);
}
run_agent_models_command(
resolved_acp,
agent_command,
@@ -234,6 +245,17 @@ pub async fn discover_agent_models(
return Ok(models);
}
if let Some(models) = discover_databricks_models(
&state.http_client,
input.provider.as_deref(),
&merged_env,
None,
)
.await?
{
return Ok(models);
}
run_agent_models_command(resolved_acp, resolved_agent, agent_args, None, merged_env).await
}
@@ -603,6 +625,96 @@ async fn discover_anthropic_models(
}))
}
// ---------------------------------------------------------------------------
// Databricks model discovery (v1 + v2)
// ---------------------------------------------------------------------------
//
// Delegates to buzz_agent_pkg::catalog::discover_databricks_models, which
// acquires auth in-process via build_token_source:
// - Static bearer (DATABRICKS_TOKEN): returned immediately.
// - PKCE cache hit: returned from disk without a browser flow.
// - No token, no cache: returns Err(LlmAuth) → we return Ok(None) and fall
// through to run_agent_models_command. Never hangs, never opens a browser.
fn is_databricks_provider(provider: Option<&str>) -> bool {
matches!(
provider
.map(str::trim)
.map(str::to_ascii_lowercase)
.as_deref(),
Some("databricks" | "databricks_v2")
)
}
fn databricks_agent_provider(provider: &str) -> buzz_agent_pkg::config::Provider {
if provider.trim().eq_ignore_ascii_case("databricks_v2") {
buzz_agent_pkg::config::Provider::DatabricksV2
} else {
buzz_agent_pkg::config::Provider::Databricks
}
}
async fn discover_databricks_models(
_client: &reqwest::Client,
provider: Option<&str>,
env: &BTreeMap<String, String>,
selected_model: Option<String>,
) -> Result<Option<AgentModelsResponse>, String> {
let provider_str = match provider {
Some(p) if is_databricks_provider(Some(p)) => p,
_ => return Ok(None),
};
let host = match env_or_process_value(env, "DATABRICKS_HOST") {
Some(h) => h,
None => return Ok(None), // no host → fall through to subprocess
};
// api_key = DATABRICKS_TOKEN (empty string = use PKCE cache).
let api_key = env_or_process_value(env, "DATABRICKS_TOKEN").unwrap_or_default();
let agent_provider = databricks_agent_provider(provider_str);
let cfg = buzz_agent_pkg::config::Config::for_discovery(agent_provider, api_key, host);
// Build a redaction env so the token never appears in surfaced errors.
let token_for_redact = env_or_process_value(env, "DATABRICKS_TOKEN").unwrap_or_default();
let redaction_env = redaction_env_with_value(env, "DATABRICKS_TOKEN", &token_for_redact);
let entries = match buzz_agent_pkg::discover_databricks_models(&cfg).await {
Ok(e) => e,
Err(buzz_agent_pkg::AgentError::LlmAuth(_)) => {
// No token + no PKCE cache → fall through to subprocess.
return Ok(None);
}
Err(e) => {
let msg = crate::managed_agents::redact_env_values_in(&e.to_string(), &redaction_env);
return Err(format!("Databricks model discovery failed: {msg}"));
}
};
if entries.is_empty() {
return Err("Databricks model discovery returned no models".to_string());
}
let models = entries
.into_iter()
.map(|e| AgentModelInfo {
id: e.id,
name: Some(e.name),
description: None,
})
.collect();
Ok(Some(AgentModelsResponse {
agent_name: provider_str.trim().to_string(),
agent_version: "models-api".to_string(),
models,
agent_default_model: None,
selected_model,
supports_switching: true,
}))
}
async fn run_agent_models_command(
resolved_acp: PathBuf,
agent_command: String,

16
desktop/src-tauri/src/commands/agent_models_tests.rs
View File

@@ -211,3 +211,19 @@ fn saved_agent_model_discovery_uses_record_snapshot() {
);
assert!(!config.env.contains_key("BUZZ_PRIVATE_KEY"));
}
// ---------------------------------------------------------------------------
// Databricks provider detection
// ---------------------------------------------------------------------------
//
// Parse/filter/pagination tests live in crates/buzz-agent/src/catalog.rs
// (they moved there with the Option C refactor).
#[test]
fn is_databricks_provider_matches_both_variants() {
assert!(is_databricks_provider(Some("databricks")));
assert!(is_databricks_provider(Some("databricks_v2")));
assert!(is_databricks_provider(Some(" DATABRICKS ")));
assert!(!is_databricks_provider(Some("anthropic")));
assert!(!is_databricks_provider(None));
}