From 8eb6e3eb601174249642373a6a367262fa476753 Mon Sep 17 00:00:00 2001 From: Atish Patel Date: Sat, 25 Jul 2026 19:56:18 -0500 Subject: [PATCH] fix(agents): run live Databricks discovery instead of the fallback list (#2890) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem The Databricks model dropdown offers a handful of stale models — and there's no way to tell that list apart from the real one. The AI Gateway exposes **66** chat/embedding endpoints on `block-lakehouse-production`, but the picker was showing a short list that includes models the gateway no longer serves and embedding endpoints that can't chat at all. Three independent defects, all on the discovery path: **1. Live discovery never ran for agents with no saved provider.** `get_agent_models` gates every in-process discovery attempt on the provider (`is_openai_compatible_provider` / `is_anthropic_provider` / `is_databricks_provider`), reading it straight from `record.provider`. That field is `null` for every agent record created before provider persistence — and for any agent that inherits its provider from the build. So all three gates saw `None`, no HTTP discovery ran, and the request fell through to the `buzz-acp models` subprocess. On the Databricks path that subprocess returns `discovery_failure_fallback` — the small hardcoded `DATABRICKS_V2_KNOWN_MODELS` catalog — which the frontend renders exactly like a live catalog. An internal DMG that bakes `BUZZ_AGENT_PROVIDER=databricks_v2` and a `DATABRICKS_HOST` still got the fallback. **2. The fallback list couldn't represent the running model.** When discovery genuinely fails, the picker should at minimum be able to show what the agent is actually configured with. For `DatabricksV2` it couldn't: the fallback returned only the hardcoded slate, so a model like `databricks-gpt-5-5` wasn't selectable in its own picker. **3. Embedding endpoints were offered as chat models.** `databricks-bge-large-en` was selectable (visible in the dialog today). The v2 endpoints payload carries no `task` or `state` field, so there is nothing to filter on but the name. ## Changes - **`effective_discovery_provider`** (new, `desktop/src-tauri/src/commands/agent_models_env.rs`) — an explicit provider (saved record value, or the create/edit dialog's current form value) still always wins; when there is none, discovery falls back to the runtime's own provider env var (`GOOSE_PROVIDER`, `BUZZ_AGENT_PROVIDER`, …) read off the merged env, which by that point already carries the baked build floor and the process env. Wired into both `get_agent_models` and `discover_agent_models`. `SavedAgentModelDiscoveryConfig` now carries `provider_env_var` from `known_acp_runtime`, so each runtime reads *its own* key rather than a shared guess. - The relay-mesh branches in `discover_agent_models` deliberately keep using `input.provider`: those key off a deliberate provider selection, never a baked default. - **Asserted vs inferred matters for missing credentials.** The OpenAI and Anthropic gates error on a missing API key, while the Databricks gate falls through; an inferred provider hitting the first two would have replaced a working subprocess catalog with `config: ANTHROPIC_API_KEY required` (`export GOOSE_PROVIDER=anthropic` is goose's documented way to pick a provider, and it keeps the key in its own keyring). So `effective_discovery_provider` returns a `DiscoveryProvider` that remembers how the value was resolved, and `required_env` only reports a missing credential for an asserted provider. A wrong guess declines and lets the subprocess answer. - **`is_chat_capable_endpoint`** (new, `crates/buzz-agent/src/catalog.rs`) — applied in `parse_v2_endpoints_page`. Drops `*embedding*` and segment-matched `bge` / `gte` endpoints; keeps everything unrecognised (fail-open, so a new model family is never hidden). Segment matching is why it's `split('-')` and not `contains`: a substring check would swallow legitimate names. - **`discovery_failure_fallback`** for `Provider::DatabricksV2` now leads with the configured model (deduped against the known slate, blank-tolerant), so a failed discovery still yields a picker that can show the running model. The configured model is trimmed once up front — `resolve_model` doesn't trim, so a padded `DATABRICKS_MODEL` used to slip past the dedupe and appear twice. - **`sort_v2_endpoints_newest_first`** (new, second commit) — the catalog is now ordered newest-first on each endpoint's `created_timestamp`, ties broken by name. Previously Buzz sorted nothing, so the gateway's own order reached the picker: it pages in two phases (Databricks-managed, then workspace-created — the page token decodes to `{"phase":"user"}`), each alphabetical, which buried `databricks-claude-opus-5` 8th behind five older Claude endpoints and `goose-claude-opus-5` — the newest endpoint in the catalog — 55th of 63. Sorting in `fetch_v2_models` means both discovery paths inherit it with no wire or type changes, and the combobox filter preserves incoming order. Endpoints with an absent or unparseable timestamp sort last rather than first, so a wire-shape change degrades to "unordered at the bottom" instead of "shuffled to the top". - The name tiebreak is load-bearing: eleven managed endpoints share one placeholder timestamp (`1699610000000`), so without it their relative order would vary between runs. That placeholder is also not always accurate — a few genuinely recent endpoints (`databricks-kimi-k2-7-code`, `databricks-llama-4-maverick`) land at the bottom with the 2023 batch. The gateway offers nothing better to sort on. - Env/provider lookup helpers moved out of `agent_models.rs` into `agent_models_env.rs`. This keeps the command module under the file-size limit **without ratcheting the override up** — the existing 1079 entry is untouched (file is now 1066 lines). ## Verification Live against `block-lakehouse-production`, release build: ``` BUZZ_ACP_AGENT_COMMAND=$PWD/target/release/buzz-agent \ BUZZ_AGENT_PROVIDER=databricks_v2 \ DATABRICKS_HOST=https://block-lakehouse-production.cloud.databricks.com \ DATABRICKS_MODEL=databricks-gpt-5-5 \ ./target/release/buzz-acp models --json ``` - before: 66 endpoints, including `databricks-bge-large-en`, `databricks-gte-large-en`, `databricks-qwen3-embedding-0-6b` - after: **63** endpoints, `[.models[] | select(.id | test("embedding|-bge-|-gte-"))]` → `[]` Top of the list after the sort commit: ``` goose-claude-opus-5 2026-07-24 databricks-claude-opus-5 2026-07-23 databricks-gemini-3-6-flash 2026-07-20 databricks-gemini-3-5-flash-lite 2026-07-20 databricks-inkling 2026-07-14 ``` Tests: 15 new (8 in `catalog.rs` — including the two-wire-shape timestamp parse, the sort's tiebreak/no-timestamp cases, and the padded-model dedupe — and 7 plus one assertion in `agent_models_tests.rs`, 3 of them covering the asserted/inferred credential split), two existing tests updated. `just check`, `just test-unit`, and `just desktop-tauri-test` all pass (1636 desktop-tauri tests, 274 buzz-agent lib tests). Not run locally: the Docker-backed integration suite (`just test`) — this diff touches neither `buzz-relay`, `buzz-db`, nor `buzz-auth`. ## Follow-ups (deliberately out of scope) Two inference-path defects found while investigating, both reproduced live against the gateway and both independent of discovery: 1. **Gemini thought signatures are dropped.** The gateway returns a bare `thoughtSignature` on tool calls; the external-model serving endpoints return it nested as `extra_content.google.thought_signature`. Neither shape is round-tripped, so multi-turn tool use on `databricks-gemini-*` fails with a 400 on the second turn. 2. **Array-shaped `content` is silently discarded.** Some models return OpenAI `content` as a block array rather than a string; `parse_openai`'s `str_field` returns `None` and the text is dropped. The legacy `serving-endpoints` path does not work around either one, and costs reasoning support on the GPT-5 family. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 5 (1M context) --- crates/buzz-agent/src/catalog.rs | 291 ++++++++++++++++-- crates/buzz-agent/src/lib.rs | 12 +- .../src-tauri/src/commands/agent_models.rs | 132 ++++---- .../src/commands/agent_models_env.rs | 114 +++++++ .../src/commands/agent_models_tests.rs | 129 ++++++++ desktop/src-tauri/src/commands/mod.rs | 1 + 6 files changed, 583 insertions(+), 96 deletions(-) create mode 100644 desktop/src-tauri/src/commands/agent_models_env.rs diff --git a/crates/buzz-agent/src/catalog.rs b/crates/buzz-agent/src/catalog.rs index fa5793571..aa2a121c9 100644 --- a/crates/buzz-agent/src/catalog.rs +++ b/crates/buzz-agent/src/catalog.rs @@ -37,8 +37,11 @@ pub const DATABRICKS_V2_KNOWN_MODELS: &[&str] = /// 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. +/// - `DatabricksV2` falls back to the configured model plus +/// [`DATABRICKS_V2_KNOWN_MODELS`] so the model-picker is always populated for +/// AI Gateway v2 users. The configured model leads: without it a fallback +/// catalog can omit the very model the agent is running, leaving the picker +/// unable to represent the current selection. /// - 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. @@ -46,25 +49,62 @@ pub const DATABRICKS_V2_KNOWN_MODELS: &[&str] = /// 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 { + // `resolve_model` does not trim, so a padded `DATABRICKS_MODEL` reaches here: + // normalize once, or the dedupe below misses and the picker lists the model + // twice (once padded, once from the known slate). + let configured_model = configured_model.trim(); + let configured = ModelEntry { + id: configured_model.to_string(), + name: configured_model.to_string(), + }; 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(), - }], + Provider::DatabricksV2 => { + let mut entries = Vec::with_capacity(DATABRICKS_V2_KNOWN_MODELS.len() + 1); + if !configured_model.is_empty() { + entries.push(configured); + } + entries.extend( + DATABRICKS_V2_KNOWN_MODELS + .iter() + .filter(|id| **id != configured_model) + .map(|id| ModelEntry { + id: id.to_string(), + name: id.to_string(), + }), + ); + entries + } + Provider::Databricks => vec![configured], + _ => vec![configured], } } +/// Heuristic: `true` when a v2 AI Gateway endpoint name looks like it serves +/// chat/completions traffic. +/// +/// The v1 `serving-endpoints` payload carries `task`, so [`parse_v1_endpoints`] +/// can filter on it directly. The v2 `ai-gateway/v2/endpoints` payload carries +/// no task or readiness field at all, so the only signal available here is the +/// endpoint name. Embedding endpoints are the one family that reliably cannot +/// serve a chat request — they reject it with +/// `API type 'mlflow/v1/chat/completions' is not supported by ''` — so +/// they are dropped rather than offered as selectable models. +/// +/// Deliberately narrow: image-capable endpoints (e.g. +/// `databricks-gemini-3-pro-image`) do answer chat requests, so they stay. Any +/// name this heuristic does not recognise is kept — preferring to include over +/// silently dropping, matching [`parse_v1_endpoints`]. +pub(crate) fn is_chat_capable_endpoint(name: &str) -> bool { + let lower = name.to_ascii_lowercase(); + if lower.contains("embedding") { + return false; + } + // Segment match so `bge`/`gte` cannot fire on a substring of a longer word. + !lower + .split('-') + .any(|segment| matches!(segment, "bge" | "gte")) +} + /// Discover available models for a Databricks provider. /// /// Returns a non-empty `Vec` on success. Returns @@ -197,7 +237,7 @@ async fn fetch_v2_models( host: &str, bearer: &str, ) -> Result, AgentError> { - let mut all_models: Vec = Vec::new(); + let mut all_endpoints: Vec = Vec::new(); let mut page_token: Option = None; let base_url = format!("{host}/api/ai-gateway/v2/endpoints"); @@ -235,8 +275,8 @@ async fn fetch_v2_models( )) })?; - let (page_models, next) = parse_v2_endpoints_page(&json)?; - all_models.extend(page_models); + let (page_endpoints, next) = parse_v2_endpoints_page(&json)?; + all_endpoints.extend(page_endpoints); match next { Some(tok) if Some(&tok) != page_token.as_ref() => page_token = Some(tok), @@ -245,26 +285,74 @@ async fn fetch_v2_models( } // Fall back to known-model list if the API returned nothing. - if all_models.is_empty() { - all_models = DATABRICKS_V2_KNOWN_MODELS + if all_endpoints.is_empty() { + return Ok(DATABRICKS_V2_KNOWN_MODELS .iter() .map(|id| ModelEntry { id: id.to_string(), name: id.to_string(), }) - .collect(); + .collect()); } - Ok(all_models) + sort_v2_endpoints_newest_first(&mut all_endpoints); + + Ok(all_endpoints + .into_iter() + .map(|endpoint| endpoint.entry) + .collect()) +} + +/// A v2 gateway endpoint plus the key discovery orders the catalog by. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct V2Endpoint { + pub(crate) entry: ModelEntry, + /// `created_timestamp` as epoch milliseconds. `None` when the field is + /// absent or unparseable — those sort last rather than jumping the queue. + pub(crate) created_ms: Option, +} + +/// Read `created_timestamp` from one endpoint object. +/// +/// The gateway sends epoch milliseconds as a JSON *string* +/// (`"created_timestamp": "1699610000000"`); accept a bare number too, so a +/// wire-shape change doesn't silently drop every endpoint to the bottom. +fn endpoint_created_ms(endpoint: &serde_json::Value) -> Option { + let value = endpoint.get("created_timestamp")?; + value + .as_i64() + .or_else(|| value.as_str()?.trim().parse::().ok()) +} + +/// Order the catalog newest-first, breaking ties by name. +/// +/// The gateway returns endpoints in two phases — Databricks-managed first, then +/// workspace-created — each alphabetical by name, which buries a brand-new +/// frontier model deep in the list. Newest-first puts the models people are +/// reaching for at the top of the picker. +/// +/// Endpoints with no usable timestamp sort last, and the name tiebreak keeps the +/// result stable: several managed endpoints share one placeholder timestamp, so +/// without it their relative order would be arbitrary. +pub(crate) fn sort_v2_endpoints_newest_first(endpoints: &mut [V2Endpoint]) { + endpoints.sort_by(|a, b| { + // `None` < `Some(_)`, so reversing puts timestamped endpoints first. + b.created_ms + .cmp(&a.created_ms) + .then_with(|| a.entry.name.cmp(&b.entry.name)) + }); } /// 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. +/// Returns `(endpoints, next_page_token)`. An empty or absent `next_page_token` +/// signals the last page. Endpoints that cannot serve chat traffic are dropped +/// (see [`is_chat_capable_endpoint`]) so the model picker only offers models the +/// agent can actually run. Page order is preserved here; the caller sorts once +/// every page is in (see [`sort_v2_endpoints_newest_first`]). pub(crate) fn parse_v2_endpoints_page( json: &serde_json::Value, -) -> Result<(Vec, Option), AgentError> { +) -> Result<(Vec, Option), AgentError> { let endpoints = json .get("endpoints") .and_then(|v| v.as_array()) @@ -279,9 +367,15 @@ pub(crate) fn parse_v2_endpoints_page( .iter() .filter_map(|endpoint| { let name = endpoint.get("name")?.as_str()?.to_string(); - Some(ModelEntry { - id: name.clone(), - name, + if !is_chat_capable_endpoint(&name) { + return None; + } + Some(V2Endpoint { + entry: ModelEntry { + id: name.clone(), + name, + }, + created_ms: endpoint_created_ms(endpoint), }) }) .collect(); @@ -356,7 +450,7 @@ mod tests { }); let (models, next) = parse_v2_endpoints_page(&json).unwrap(); - let ids: Vec<&str> = models.iter().map(|m| m.id.as_str()).collect(); + let ids: Vec<&str> = models.iter().map(|m| m.entry.id.as_str()).collect(); assert_eq!( ids, vec![ @@ -399,4 +493,139 @@ mod tests { "got: {err}" ); } + + #[test] + fn v2_parse_drops_embedding_endpoints() { + // The v2 payload carries no `task`, so embedding endpoints are only + // recognisable by name. They reject chat requests, so offering them in + // the picker can only produce a 400 at send time. + let json = serde_json::json!({ + "endpoints": [ + {"name": "databricks-bge-large-en"}, + {"name": "databricks-gte-large-en"}, + {"name": "databricks-qwen3-embedding-0-6b"}, + {"name": "databricks-claude-opus-5"}, + {"name": "databricks-gemini-3-pro-image"}, + ] + }); + + let (models, _) = parse_v2_endpoints_page(&json).unwrap(); + let ids: Vec<&str> = models.iter().map(|m| m.entry.id.as_str()).collect(); + // Image endpoints DO answer chat requests, so they are retained. + assert_eq!( + ids, + vec!["databricks-claude-opus-5", "databricks-gemini-3-pro-image"] + ); + } + + #[test] + fn v2_parse_reads_created_timestamp_in_either_wire_shape() { + // The gateway sends epoch ms as a string; a bare number must work too. + let json = serde_json::json!({ + "endpoints": [ + {"name": "string-ts", "created_timestamp": "1784932442251"}, + {"name": "number-ts", "created_timestamp": 1784932442251i64}, + {"name": "junk-ts", "created_timestamp": "not-a-number"}, + {"name": "no-ts"}, + ] + }); + + let (models, _) = parse_v2_endpoints_page(&json).unwrap(); + let stamps: Vec> = models.iter().map(|m| m.created_ms).collect(); + assert_eq!( + stamps, + vec![Some(1784932442251), Some(1784932442251), None, None,] + ); + } + + #[test] + fn v2_endpoints_sort_newest_first_then_by_name() { + // Mirrors the real catalog: the gateway pages Databricks-managed + // endpoints first, then workspace-created ones, each alphabetical — so + // the newest model is buried mid-list until this sort runs. + let json = serde_json::json!({ + "endpoints": [ + {"name": "databricks-claude-opus-5", "created_timestamp": "1784851200000"}, + {"name": "databricks-gpt-5-6-sol", "created_timestamp": "1784073600000"}, + {"name": "databricks-gpt-5-6-luna", "created_timestamp": "1784073600000"}, + {"name": "databricks-llama-4-maverick", "created_timestamp": "1699610000000"}, + {"name": "goose-claude-opus-5", "created_timestamp": "1784932442251"}, + {"name": "endpoint-without-timestamp"}, + ] + }); + + let (mut models, _) = parse_v2_endpoints_page(&json).unwrap(); + sort_v2_endpoints_newest_first(&mut models); + + let ids: Vec<&str> = models.iter().map(|m| m.entry.id.as_str()).collect(); + assert_eq!( + ids, + vec![ + // Newest first, across both pagination phases. + "goose-claude-opus-5", + "databricks-claude-opus-5", + // Same timestamp — the name tiebreak keeps this deterministic. + "databricks-gpt-5-6-luna", + "databricks-gpt-5-6-sol", + "databricks-llama-4-maverick", + // No usable timestamp sorts last, never first. + "endpoint-without-timestamp", + ] + ); + } + + #[test] + fn is_chat_capable_endpoint_keeps_unrecognised_names() { + // Prefer including over silently dropping — an unknown family is kept. + assert!(is_chat_capable_endpoint("databricks-glm-5-2")); + assert!(is_chat_capable_endpoint("some-teams-custom-endpoint")); + // `bge`/`gte` match as whole segments only, never as substrings. + assert!(is_chat_capable_endpoint("databricks-budget-gtex-model")); + assert!(!is_chat_capable_endpoint("databricks-bge-large-en")); + assert!(!is_chat_capable_endpoint("databricks-gte-large-en")); + assert!(!is_chat_capable_endpoint("databricks-qwen3-embedding-0-6b")); + } + + #[test] + fn v2_discovery_failure_fallback_leads_with_configured_model() { + let result = discovery_failure_fallback(Provider::DatabricksV2, "databricks-claude-opus-5"); + let ids: Vec<&str> = result.iter().map(|m| m.id.as_str()).collect(); + + // The running model must be representable in the picker even when + // discovery failed, so it leads the fallback catalog. + assert_eq!(ids.first(), Some(&"databricks-claude-opus-5")); + for known in DATABRICKS_V2_KNOWN_MODELS { + assert!(ids.contains(known), "fallback must retain '{known}'"); + } + } + + #[test] + fn v2_discovery_failure_fallback_does_not_duplicate_configured_model() { + let configured = DATABRICKS_V2_KNOWN_MODELS[0]; + let result = discovery_failure_fallback(Provider::DatabricksV2, configured); + let occurrences = result.iter().filter(|m| m.id == configured).count(); + assert_eq!(occurrences, 1, "got: {result:?}"); + assert_eq!(result.len(), DATABRICKS_V2_KNOWN_MODELS.len()); + } + + #[test] + fn v2_discovery_failure_fallback_tolerates_blank_configured_model() { + for configured in ["", " "] { + let result = discovery_failure_fallback(Provider::DatabricksV2, configured); + let ids: Vec<&str> = result.iter().map(|m| m.id.as_str()).collect(); + assert_eq!(ids, DATABRICKS_V2_KNOWN_MODELS.to_vec()); + } + } + + #[test] + fn v2_discovery_failure_fallback_dedupes_a_padded_configured_model() { + // `DATABRICKS_MODEL=" databricks-gpt-5-5 "` reaches here untrimmed, and an + // untrimmed comparison would list the model twice — once padded, once from + // the known slate. + let configured = DATABRICKS_V2_KNOWN_MODELS[0]; + let result = + discovery_failure_fallback(Provider::DatabricksV2, &format!(" {configured} ")); + let ids: Vec<&str> = result.iter().map(|m| m.id.as_str()).collect(); + assert_eq!(ids, DATABRICKS_V2_KNOWN_MODELS.to_vec()); + } } diff --git a/crates/buzz-agent/src/lib.rs b/crates/buzz-agent/src/lib.rs index f5faafe85..e141b9860 100644 --- a/crates/buzz-agent/src/lib.rs +++ b/crates/buzz-agent/src/lib.rs @@ -920,11 +920,13 @@ mod tests { let configured = "my-configured-model"; let result = discovery_failure_fallback(Provider::DatabricksV2, configured); - // DatabricksV2 must return the full DATABRICKS_V2_KNOWN_MODELS list. + // DatabricksV2 must return the full DATABRICKS_V2_KNOWN_MODELS list, + // plus the configured model so the picker can still represent the model + // the agent is actually running. assert_eq!( result.len(), - DATABRICKS_V2_KNOWN_MODELS.len(), - "DatabricksV2 fallback must return all known models" + DATABRICKS_V2_KNOWN_MODELS.len() + 1, + "DatabricksV2 fallback must return all known models plus the configured model" ); let result_ids: Vec<&str> = result.iter().map(|m| m.id.as_str()).collect(); for known_id in DATABRICKS_V2_KNOWN_MODELS { @@ -933,6 +935,10 @@ mod tests { "DatabricksV2 fallback must include known model '{known_id}'" ); } + assert!( + result_ids.contains(&configured), + "DatabricksV2 fallback must include the configured model" + ); } #[test] diff --git a/desktop/src-tauri/src/commands/agent_models.rs b/desktop/src-tauri/src/commands/agent_models.rs index 425eadb5f..112f297fa 100644 --- a/desktop/src-tauri/src/commands/agent_models.rs +++ b/desktop/src-tauri/src/commands/agent_models.rs @@ -5,6 +5,13 @@ use serde::Deserialize; use tauri::{AppHandle, State}; use super::agent_model_process::run_agent_models_command; +// The map-only lookup is reached solely from the base-URL helpers that exist for +// their unit tests; discovery itself always goes through the process-env variant. +#[cfg(test)] +use super::agent_models_env::env_value; +use super::agent_models_env::{ + effective_discovery_provider, env_or_process_value, redaction_env_with_value, DiscoveryProvider, +}; use super::agent_update_rollback::{rollback_failed_agent_update, AgentUpdateRollback}; use crate::{ @@ -31,7 +38,15 @@ pub async fn get_agent_models( app: AppHandle, state: State<'_, AppState>, ) -> Result { - let (resolved_acp, agent_command, agent_args, persisted_model, effective_provider, merged_env) = { + let ( + resolved_acp, + agent_command, + agent_args, + persisted_model, + saved_provider, + provider_env_var, + merged_env, + ) = { let _store_guard = state .managed_agents_store_lock .lock() @@ -82,14 +97,19 @@ pub async fn get_agent_models( args, discovery.model, discovery.provider, + discovery.provider_env_var, discovery.env, ) }; // store lock released — subprocess runs without holding the lock let merged_env = discovery_env_with_baked_floor(merged_env); + // Resolve against the baked/process env when the record saved no provider, + // so a build-provided provider still gets live discovery. + let effective_provider = + effective_discovery_provider(saved_provider.as_deref(), provider_env_var, &merged_env); if let Some(models) = discover_openai_compatible_models( &state.http_client, - effective_provider.as_deref(), + &effective_provider, &merged_env, persisted_model.clone(), ) @@ -100,7 +120,7 @@ pub async fn get_agent_models( if let Some(models) = discover_anthropic_models( &state.http_client, - effective_provider.as_deref(), + &effective_provider, &merged_env, persisted_model.clone(), ) @@ -111,7 +131,7 @@ pub async fn get_agent_models( if let Some(models) = discover_databricks_models( &state.http_client, - effective_provider.as_deref(), + &effective_provider, &merged_env, persisted_model.clone(), ) @@ -134,6 +154,10 @@ pub async fn get_agent_models( struct SavedAgentModelDiscoveryConfig { model: Option, provider: Option, + /// The runtime's provider env var (e.g. `BUZZ_AGENT_PROVIDER`), so discovery + /// can recover the provider from the env when the record has none. `None` + /// for runtimes that do not take a provider, or an unknown command. + provider_env_var: Option<&'static str>, env: BTreeMap, } @@ -141,8 +165,9 @@ fn saved_agent_model_discovery_config( record: &crate::managed_agents::ManagedAgentRecord, agent_command: &str, ) -> SavedAgentModelDiscoveryConfig { + let runtime_meta = known_acp_runtime(agent_command); let mut derived_env = BTreeMap::new(); - if let Some(meta) = known_acp_runtime(agent_command) { + if let Some(meta) = runtime_meta { for (key, value) in crate::managed_agents::runtime_metadata_env_vars( meta.model_env_var, meta.provider_env_var, @@ -157,6 +182,7 @@ fn saved_agent_model_discovery_config( SavedAgentModelDiscoveryConfig { model: record.model.clone(), provider: record.provider.clone(), + provider_env_var: runtime_meta.and_then(|meta| meta.provider_env_var), env: crate::managed_agents::merged_user_env(&derived_env, &record.env_vars), } } @@ -205,8 +231,9 @@ pub async fn discover_agent_models( .map(|p| p.display().to_string()) .unwrap_or_else(|| agent_command.to_string()); + let runtime_meta = known_acp_runtime(agent_command); let mut derived_env = BTreeMap::new(); - if let Some(meta) = known_acp_runtime(agent_command) { + if let Some(meta) = runtime_meta { let provider = input .provider .as_deref() @@ -220,6 +247,13 @@ pub async fn discover_agent_models( } let merged_env = crate::managed_agents::merged_user_env(&derived_env, &input.env_vars); let merged_env = discovery_env_with_baked_floor(merged_env); + // Recover a build-provided provider when the form has none, so the create + // dialog discovers live models instead of falling through to the subprocess. + let effective_provider = effective_discovery_provider( + input.provider.as_deref(), + runtime_meta.and_then(|meta| meta.provider_env_var), + &merged_env, + ); // Buzz shared compute discovery must not depend on the local OpenAI ingress: that // client endpoint is started only after a live target is selected. @@ -268,7 +302,7 @@ pub async fn discover_agent_models( if let Some(models) = discover_openai_compatible_models( &state.http_client, - input.provider.as_deref(), + &effective_provider, &merged_env, None, ) @@ -277,24 +311,16 @@ pub async fn discover_agent_models( return Ok(models); } - if let Some(models) = discover_anthropic_models( - &state.http_client, - input.provider.as_deref(), - &merged_env, - None, - ) - .await? + if let Some(models) = + discover_anthropic_models(&state.http_client, &effective_provider, &merged_env, None) + .await? { return Ok(models); } - if let Some(models) = discover_databricks_models( - &state.http_client, - input.provider.as_deref(), - &merged_env, - None, - ) - .await? + if let Some(models) = + discover_databricks_models(&state.http_client, &effective_provider, &merged_env, None) + .await? { return Ok(models); } @@ -337,33 +363,6 @@ fn openai_compatible_models_url_for_discovery(env: &BTreeMap) -> format!("{}/models", base_url.trim_end_matches('/')) } -fn env_value(env: &BTreeMap, key: &str) -> Option { - env.get(key) - .map(String::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string) -} - -fn env_or_process_value(env: &BTreeMap, key: &str) -> Option { - env_value(env, key).or_else(|| { - std::env::var(key) - .ok() - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()) - }) -} - -fn redaction_env_with_value( - env: &BTreeMap, - key: &str, - value: &str, -) -> BTreeMap { - let mut redaction_env = env.clone(); - redaction_env.insert(key.to_string(), value.to_string()); - redaction_env -} - fn is_agent_text_model_id(id: &str) -> bool { let lower = id.to_ascii_lowercase(); if [ @@ -482,20 +481,23 @@ fn normalize_openai_compatible_models( async fn discover_openai_compatible_models( client: &reqwest::Client, - provider: Option<&str>, + provider: &DiscoveryProvider, env: &BTreeMap, selected_model: Option, ) -> Result, String> { - let relay_mesh = provider.map(str::trim) == Some(crate::managed_agents::RELAY_MESH_PROVIDER_ID); - if !relay_mesh && !is_openai_compatible_provider(provider) { + let relay_mesh = + provider.as_deref().map(str::trim) == Some(crate::managed_agents::RELAY_MESH_PROVIDER_ID); + if !relay_mesh && !is_openai_compatible_provider(provider.as_deref()) { return Ok(None); } let api_key = if relay_mesh { crate::managed_agents::RELAY_MESH_API_KEY_PLACEHOLDER.to_string() } else { - env_or_process_value(env, "OPENAI_COMPAT_API_KEY") - .ok_or_else(|| "config: OPENAI_COMPAT_API_KEY required".to_string())? + match provider.required_env(env, "OPENAI_COMPAT_API_KEY")? { + Some(api_key) => api_key, + None => return Ok(None), + } }; let redaction_env = redaction_env_with_value(env, "OPENAI_COMPAT_API_KEY", &api_key); let url = if relay_mesh { @@ -520,13 +522,13 @@ async fn discover_openai_compatible_models( .json::() .await .map_err(|error| format!("OpenAI model discovery response parse failed: {error}"))?; - let models = normalize_openai_compatible_models(response, provider); + let models = normalize_openai_compatible_models(response, provider.as_deref()); if models.is_empty() { return Err("OpenAI model discovery returned no compatible text models".to_string()); } Ok(Some(AgentModelsResponse { - agent_name: provider.unwrap_or("openai").trim().to_string(), + agent_name: provider.as_deref().unwrap_or("openai").trim().to_string(), agent_version: "models-api".to_string(), models, agent_default_model: None, @@ -631,16 +633,18 @@ async fn fetch_anthropic_model_page( async fn discover_anthropic_models( client: &reqwest::Client, - provider: Option<&str>, + provider: &DiscoveryProvider, env: &BTreeMap, selected_model: Option, ) -> Result, String> { - if !is_anthropic_provider(provider) { + if !is_anthropic_provider(provider.as_deref()) { return Ok(None); } - let api_key = env_or_process_value(env, "ANTHROPIC_API_KEY") - .ok_or_else(|| "config: ANTHROPIC_API_KEY required".to_string())?; + let api_key = match provider.required_env(env, "ANTHROPIC_API_KEY")? { + Some(api_key) => api_key, + None => return Ok(None), + }; let redaction_env = redaction_env_with_value(env, "ANTHROPIC_API_KEY", &api_key); let url = anthropic_models_url_for_discovery(env); let mut models = Vec::new(); @@ -666,7 +670,11 @@ async fn discover_anthropic_models( } Ok(Some(AgentModelsResponse { - agent_name: provider.unwrap_or("anthropic").trim().to_string(), + agent_name: provider + .as_deref() + .unwrap_or("anthropic") + .trim() + .to_string(), agent_version: "models-api".to_string(), models, agent_default_model: None, @@ -708,11 +716,11 @@ fn databricks_agent_provider(provider: &str) -> buzz_agent_pkg::config::Provider async fn discover_databricks_models( _client: &reqwest::Client, - provider: Option<&str>, + provider: &DiscoveryProvider, env: &BTreeMap, selected_model: Option, ) -> Result, String> { - let provider_str = match provider { + let provider_str = match provider.as_deref() { Some(p) if is_databricks_provider(Some(p)) => p, _ => return Ok(None), }; diff --git a/desktop/src-tauri/src/commands/agent_models_env.rs b/desktop/src-tauri/src/commands/agent_models_env.rs new file mode 100644 index 000000000..0a40b6bd8 --- /dev/null +++ b/desktop/src-tauri/src/commands/agent_models_env.rs @@ -0,0 +1,114 @@ +//! Env lookups shared by the in-process model discovery paths. +//! +//! Discovery reads configuration from a merged env map (record/form env over the +//! runtime-derived env, under the baked build floor) and falls back to the +//! process env for values a GUI launch never wrote into the map. + +use std::collections::BTreeMap; + +/// Read a non-blank value from the merged discovery env. +pub(super) fn env_value(env: &BTreeMap, key: &str) -> Option { + env.get(key) + .map(String::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) +} + +/// Read a non-blank value from the merged discovery env, then the process env. +pub(super) fn env_or_process_value(env: &BTreeMap, key: &str) -> Option { + env_value(env, key).or_else(|| { + std::env::var(key) + .ok() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) + }) +} + +/// Clone `env` with `key` set to the value a request actually used, so error +/// redaction masks the inherited process value and not just the mapped one. +pub(super) fn redaction_env_with_value( + env: &BTreeMap, + key: &str, + value: &str, +) -> BTreeMap { + let mut redaction_env = env.clone(); + redaction_env.insert(key.to_string(), value.to_string()); + redaction_env +} + +/// A provider resolved for discovery, plus how it was resolved. +/// +/// The distinction decides what a missing credential means. An explicit provider +/// is an assertion — the record or the dialog says this agent runs on Anthropic, +/// so a missing `ANTHROPIC_API_KEY` is a real misconfiguration and the user +/// should see it. An inferred provider is only a guess read off the environment, +/// and a wrong guess must not replace a working catalog with an error: a +/// `GOOSE_PROVIDER=anthropic` export (goose's documented way to pick a provider, +/// with the key in goose's own keyring rather than Buzz's env) would otherwise +/// turn the subprocess catalog into `config: ANTHROPIC_API_KEY required`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) struct DiscoveryProvider { + value: Option, + inferred: bool, +} + +impl DiscoveryProvider { + pub(super) fn as_deref(&self) -> Option<&str> { + self.value.as_deref() + } + + /// Read a credential the provider's discovery cannot run without. + /// + /// `Ok(None)` means "not configured, and the provider was only inferred" — + /// the caller must fall through to the `buzz-acp models` subprocess rather + /// than surface an error. + pub(super) fn required_env( + &self, + env: &BTreeMap, + key: &str, + ) -> Result, String> { + match env_or_process_value(env, key) { + Some(value) => Ok(Some(value)), + None if self.inferred => Ok(None), + None => Err(format!("config: {key} required")), + } + } +} + +/// Resolve the provider that live model discovery should run against. +/// +/// An explicit provider — the agent record's saved provider, or the create/edit +/// dialog's current form value — always wins. When there is none, fall back to +/// the provider the agent will actually launch with: the runtime's provider env +/// var read out of `env`, which by this point already carries the baked build +/// floor (see `discovery_env_with_baked_floor`) and the process env. +/// +/// Without this fallback, every provider gate sees `None` for an agent whose +/// record predates provider persistence, so no in-process discovery runs at all +/// — even on an internal build that bakes `BUZZ_AGENT_PROVIDER=databricks_v2` +/// and a `DATABRICKS_HOST`. Discovery then degrades to the `buzz-acp models` +/// subprocess, which on a Databricks failure path surfaces the small +/// known-models fallback catalog instead of the live gateway list — +/// indistinguishable, from the picker's side, from the real thing. +pub(super) fn effective_discovery_provider( + provider: Option<&str>, + provider_env_var: Option<&str>, + env: &BTreeMap, +) -> DiscoveryProvider { + if let Some(explicit) = provider + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + { + return DiscoveryProvider { + value: Some(explicit), + inferred: false, + }; + } + + DiscoveryProvider { + value: provider_env_var.and_then(|key| env_or_process_value(env, key)), + inferred: true, + } +} diff --git a/desktop/src-tauri/src/commands/agent_models_tests.rs b/desktop/src-tauri/src/commands/agent_models_tests.rs index 0f99927c9..0d6449e78 100644 --- a/desktop/src-tauri/src/commands/agent_models_tests.rs +++ b/desktop/src-tauri/src/commands/agent_models_tests.rs @@ -210,6 +210,135 @@ fn saved_agent_model_discovery_uses_record_snapshot() { Some("record-key") ); assert!(!config.env.contains_key("BUZZ_PRIVATE_KEY")); + assert_eq!(config.provider_env_var, Some("GOOSE_PROVIDER")); +} + +// --------------------------------------------------------------------------- +// Provider resolution for discovery +// --------------------------------------------------------------------------- + +#[test] +fn effective_discovery_provider_prefers_the_explicit_provider() { + let env = BTreeMap::from([( + "BUZZ_AGENT_PROVIDER".to_string(), + "databricks_v2".to_string(), + )]); + + // A saved/selected provider is a deliberate choice and must win over the + // build-provided default, so discovery matches what spawn will use. + assert_eq!( + effective_discovery_provider(Some("anthropic"), Some("BUZZ_AGENT_PROVIDER"), &env) + .as_deref(), + Some("anthropic") + ); +} + +#[test] +fn effective_discovery_provider_recovers_baked_provider_when_record_has_none() { + let env = BTreeMap::from([( + "BUZZ_AGENT_PROVIDER".to_string(), + "databricks_v2".to_string(), + )]); + + // The regression this guards: records predating provider persistence carry + // `provider: null`, so every discovery gate saw None and no live Databricks + // catalog was ever fetched on builds that bake the provider in. + for provider in [None, Some(""), Some(" ")] { + assert_eq!( + effective_discovery_provider(provider, Some("BUZZ_AGENT_PROVIDER"), &env).as_deref(), + Some("databricks_v2"), + "provider input {provider:?} must fall back to the env value" + ); + } +} + +#[test] +fn effective_discovery_provider_is_none_without_an_explicit_or_env_provider() { + let env = BTreeMap::new(); + assert_eq!( + effective_discovery_provider(None, Some("BUZZ_AGENT_PROVIDER"), &env).as_deref(), + None + ); + // A runtime that takes no provider env var has nothing to recover from. + assert_eq!( + effective_discovery_provider( + None, + None, + &BTreeMap::from([( + "BUZZ_AGENT_PROVIDER".to_string(), + "databricks_v2".to_string() + )]) + ) + .as_deref(), + None + ); +} + +/// A credential name no environment sets, so `required_env` is exercised without +/// depending on what the developer happens to have exported. +const UNSET_CREDENTIAL: &str = "BUZZ_TEST_UNSET_DISCOVERY_CREDENTIAL"; + +#[test] +fn env_derived_provider_falls_through_when_its_credential_is_missing() { + let env = BTreeMap::from([("GOOSE_PROVIDER".to_string(), "anthropic".to_string())]); + let inferred = effective_discovery_provider(None, Some("GOOSE_PROVIDER"), &env); + assert_eq!(inferred.as_deref(), Some("anthropic")); + + // `export GOOSE_PROVIDER=anthropic` is goose's documented way to pick a + // provider, and it keeps the API key in its own config/keyring rather than in + // Buzz's env — so the provider is visible here and the credential is not. + // Erroring would swap the working subprocess catalog for a hard + // "config: ... required" on exactly the null-provider records this fallback + // exists to serve; the gate has to decline instead. + assert_eq!(inferred.required_env(&env, UNSET_CREDENTIAL), Ok(None)); +} + +#[test] +fn explicit_provider_still_reports_a_missing_credential() { + // An explicit provider is an assertion about this agent, so a missing + // credential is a real misconfiguration and stays user-visible. + let env = BTreeMap::new(); + let explicit = effective_discovery_provider(Some("anthropic"), Some("GOOSE_PROVIDER"), &env); + assert_eq!( + explicit.required_env(&env, UNSET_CREDENTIAL), + Err(format!("config: {UNSET_CREDENTIAL} required")) + ); +} + +#[test] +fn required_env_returns_a_configured_credential_however_the_provider_was_resolved() { + let env = BTreeMap::from([ + ("GOOSE_PROVIDER".to_string(), "anthropic".to_string()), + ( + UNSET_CREDENTIAL.to_string(), + " sk-configured ".to_string(), + ), + ]); + for provider in [Some("anthropic"), None] { + let resolved = effective_discovery_provider(provider, Some("GOOSE_PROVIDER"), &env); + assert_eq!( + resolved.required_env(&env, UNSET_CREDENTIAL), + Ok(Some("sk-configured".to_string())), + "provider input {provider:?} must read the configured credential" + ); + } +} + +#[test] +fn effective_discovery_provider_reads_the_runtimes_own_env_var() { + // goose keys its provider off GOOSE_PROVIDER, so a BUZZ_AGENT_PROVIDER in + // the env must not be mistaken for this runtime's provider. + let env = BTreeMap::from([ + ("GOOSE_PROVIDER".to_string(), "databricks".to_string()), + ( + "BUZZ_AGENT_PROVIDER".to_string(), + "databricks_v2".to_string(), + ), + ]); + assert_eq!( + effective_discovery_provider(None, Some("GOOSE_PROVIDER"), &env).as_deref(), + Some("databricks") + ); } // --------------------------------------------------------------------------- diff --git a/desktop/src-tauri/src/commands/mod.rs b/desktop/src-tauri/src/commands/mod.rs index 7e9d916be..a048ad24a 100644 --- a/desktop/src-tauri/src/commands/mod.rs +++ b/desktop/src-tauri/src/commands/mod.rs @@ -5,6 +5,7 @@ mod agent_logs; mod agent_metric_archive; mod agent_model_process; mod agent_models; +mod agent_models_env; mod agent_providers; mod agent_settings; mod agent_update_rollback;