From 4ff6ed25ee1caa57c360585c802fe36572724a42 Mon Sep 17 00:00:00 2001 From: Kalvin Chau Date: Wed, 26 Aug 2026 15:58:46 -0700 Subject: [PATCH 1/4] feat(agent): add unity catalog model discovery Extend Databricks v2 discovery with Unity Catalog model services and a shared visibility filter. Route model-service FQNs through MLflow Chat Completions while preserving Desktop and ACP fallback parity. Signed-off-by: Kalvin Chau Signed-off-by: am <6e30cd56c30e030cd31bb0939b94a7c257c9a09d5ba2d92cf2735da45629f248@buzz.block.builderlab.xyz> --- .env.example | 6 + crates/buzz-agent/README.md | 5 +- crates/buzz-agent/src/catalog.rs | 899 +++++++++++++++--- crates/buzz-agent/src/config.rs | 162 +++- crates/buzz-agent/src/lib.rs | 37 +- crates/buzz-agent/src/llm.rs | 160 +++- crates/buzz-agent/tests/databricks_oauth.rs | 67 +- .../src-tauri/src/commands/agent_models.rs | 8 +- .../src/commands/agent_models_databricks.rs | 34 +- .../src/commands/agent_models_env.rs | 16 + .../src/commands/agent_models_tests.rs | 41 +- .../src-tauri/src/managed_agents/env_vars.rs | 3 +- .../spawn_snapshot/diff/tests.rs | 21 + desktop/src/features/agents/AGENTS.md | 2 + desktop/src/shared/api/tauri.ts | 2 +- 15 files changed, 1241 insertions(+), 222 deletions(-) diff --git a/.env.example b/.env.example index a6740f7a7d8..14134418955 100644 --- a/.env.example +++ b/.env.example @@ -172,6 +172,12 @@ RUST_LOG=buzz_relay=debug,buzz_datastore=info,buzz_db=debug,buzz_auth=debug,buzz # Use `buzz-acp models` to discover available model IDs. # BUZZ_ACP_MODEL= +# Optional Databricks model-picker visibility filter. Discovery-only; this does +# not grant inference access. Comma-separated full-string * / ? patterns are +# OR-matched against raw workspace endpoint and Unity Catalog model-service IDs. +# Unset or blank shows every catalog entry. A nonblank value with no usable patterns is invalid. +# DATABRICKS_MODEL_FILTER=databricks-*,data_tools.goose.* + # ── Timeouts & sessions ────────────────────────────────────────────────────── # Max seconds per agent turn before timeout (default 320 = ~5 min). # BUZZ_ACP_TURN_TIMEOUT=320 diff --git a/crates/buzz-agent/README.md b/crates/buzz-agent/README.md index 0bc03db7813..56e62cf9e79 100644 --- a/crates/buzz-agent/README.md +++ b/crates/buzz-agent/README.md @@ -149,6 +149,7 @@ Everything is environment variables. No flags, no config files. (We are a subpro | `OPENROUTER_BASE_URL` | `https://openrouter.ai/api/v1` | | | `DATABRICKS_HOST` | — | Required when provider=databricks or provider=databricks_v2. | | `DATABRICKS_MODEL` | — | Required when provider=databricks or provider=databricks_v2. | +| `DATABRICKS_MODEL_FILTER` | — | Optional discovery-only, comma-separated full-string `*`/`?` patterns OR-matched against raw Databricks endpoint and Unity Catalog model-service IDs. Blank/unset shows all; this is visibility filtering, not an authorization boundary. | | `DATABRICKS_TOKEN` | — | Optional static bearer escape hatch. If unset, Databricks uses browser OAuth + refresh cache. | | `BUZZ_AGENT_SYSTEM_PROMPT` | built-in | Inline system prompt. | | `BUZZ_AGENT_SYSTEM_PROMPT_FILE` | — | File path. Mutually exclusive with the above. | @@ -241,7 +242,9 @@ lifecycle hook — see [MCP_DRIVEN_HOOKS.md](../../docs/MCP_DRIVEN_HOOKS.md). | Block Gateway | `openai` | `POST {base}/chat/completions` | gpt-5, claude | | OpenRouter | `openrouter` | `POST {base}/chat/completions` | anything they route (extended-thinking replay, provider-agnostic tool calling) | | Databricks | `databricks` | `POST {host}/serving-endpoints/{model}/invocations` | goose-claude-4-6-sonnet | -| Databricks AI Gateway v2 | `databricks_v2` | `POST {host}/ai-gateway/{provider}/v1/...` | databricks-gpt-5-5, databricks-claude-opus-4-7 | +| Databricks AI Gateway v2 | `databricks_v2` | `POST {host}/ai-gateway/{provider}/v1/...` | workspace endpoints and Unity Catalog model-service FQNs; UC FQNs use MLflow Chat Completions | + +The optional `DATABRICKS_MODEL_FILTER` applies only to model discovery. Each comma-separated entry is trimmed and matched against the complete raw ID with case-sensitive `*` (zero or more characters) and `?` (one Unicode character) semantics; patterns are OR-ed. Unset or blank preserves the full authenticated catalog. A nonblank value containing no usable patterns is rejected. This controls picker visibility only; Databricks and Unity Catalog permissions remain the authorization boundary. A filtered-empty result is authoritative and does not restore the built-in fallback models. If `BUZZ_AGENT_PROVIDER=anthropic` is selected without `ANTHROPIC_API_KEY`, `BUZZ_AGENT_PROVIDER=openai` is selected without `OPENAI_COMPAT_API_KEY`, or `BUZZ_AGENT_PROVIDER=openrouter` is selected without `OPENROUTER_API_KEY`, the agent returns an error — there is no implicit fallback to another provider. diff --git a/crates/buzz-agent/src/catalog.rs b/crates/buzz-agent/src/catalog.rs index 69714b145c5..9a8a0e36c0f 100644 --- a/crates/buzz-agent/src/catalog.rs +++ b/crates/buzz-agent/src/catalog.rs @@ -12,23 +12,23 @@ //! This helper never opens a browser. Callers choose whether to reject, degrade, //! or start a separate interactive authentication flow. -use std::sync::Arc; +use std::{collections::HashSet, sync::Arc}; use reqwest::Client; +use serde_json::Value; use crate::{ auth::TokenSource, - config::{Config, Provider}, + config::{Config, DatabricksModelFilter, Provider}, llm::build_token_source, types::AgentError, }; -/// A discovered model entry: `id` is the picker value (the raw endpoint id, and -/// the wire/config value), `name` is the display label. The Databricks API has -/// no display-name field, so discovery curates `name` from the capability -/// manifest ([`model_capabilities::databricks_registry_label`]) — a known id -/// yields its curated label (e.g. `GPT-5.5`), an unknown id falls back to the -/// raw id. +/// A discovered model entry: `id` is the picker value (the raw endpoint id or +/// Unity Catalog model-service FQN, and the wire/config value), `name` is the +/// display label. Databricks catalog APIs do not provide a consistently useful +/// picker label, so discovery curates names from the capability manifest when +/// an exact known id exists and otherwise uses the raw id. #[derive(Debug, Clone, PartialEq, Eq)] pub struct ModelEntry { pub id: String, @@ -36,20 +36,23 @@ pub struct ModelEntry { } const AUTHENTICATED_EMPTY_CATALOG_SUFFIX: &str = " (default catalog)"; - -/// Curated display label for a discovered Databricks endpoint id: the manifest's -/// exact-record label when one exists, otherwise the raw id. The API returns no -/// display name, so this is the single seam that turns a raw endpoint id into a -/// human label for the picker. +const MAX_CATALOG_PAGES: usize = 20; +const MAX_CATALOG_ERROR_BODY_BYTES: usize = 4 * 1024; +const WORKSPACE_CATALOG_QUERY: &str = "?page_size=100"; +const UNITY_CATALOG_QUERY: &str = "?page_size=100&view=FULL"; +type CatalogPage = Result<(Vec, Option), AgentError>; + +/// Curated display label for a discovered Databricks endpoint or model-service +/// id. Unknown ids deliberately pass through unchanged. fn curated_model_name(id: &str) -> String { crate::model_capabilities::databricks_registry_label(id) .unwrap_or(id) .to_string() } -/// Fallback catalog used only when an authenticated `api/ai-gateway/v2/endpoints` -/// call succeeds with an empty list. The known-model ids come from the manifest -/// ([`model_capabilities::databricks_v2_known_models`]), the single runtime source. +/// Fallback catalog used only when both authenticated Databricks v2 catalogs +/// successfully respond with no entries and no visibility filter is active. +/// The known-model ids come from the manifest, the single runtime source. fn authenticated_empty_v2_catalog() -> Vec { crate::model_capabilities::databricks_v2_known_models() .iter() @@ -63,37 +66,16 @@ fn authenticated_empty_v2_catalog() -> Vec { .collect() } -/// 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 -/// `Err(AgentError::LlmAuth)` when no token is available (no static token, -/// no PKCE cache). The helper itself never starts interactive authentication. +/// Returns an empty vector when an authenticated catalog is valid but no +/// visible entries remain after filtering. Returns `Err(AgentError::LlmAuth)` +/// when no token is available (no static token, no PKCE cache). The helper +/// itself never starts interactive authentication. +/// +/// For v2, the known-model fallback is used only when both catalog requests +/// succeed empty and no filter is active. A filter is applied to v1 results +/// after its existing endpoint capability filtering. /// /// # Panics /// Never panics. @@ -112,8 +94,19 @@ async fn discover_databricks_models_with_token_source( loop { let result = match cfg.provider { - Provider::Databricks => fetch_v1_models(&http, host, &bearer).await, - Provider::DatabricksV2 => fetch_v2_models(&http, host, &bearer).await, + Provider::Databricks => fetch_v1_models(&http, host, &bearer) + .await + .map(|models| apply_model_filter(models, cfg.databricks_model_filter.as_ref())), + Provider::DatabricksV2 => { + fetch_v2_models( + &http, + host, + &bearer, + cfg.databricks_model_filter.as_ref(), + refreshed, + ) + .await + } _ => { return Err(AgentError::InvalidParams( "discover_databricks_models called for non-Databricks provider".into(), @@ -137,6 +130,19 @@ async fn discover_databricks_models_with_token_source( } } +fn apply_model_filter( + models: Vec, + filter: Option<&DatabricksModelFilter>, +) -> Vec { + match filter { + Some(filter) => models + .into_iter() + .filter(|model| filter.matches(&model.id)) + .collect(), + None => models, + } +} + // --------------------------------------------------------------------------- // v1 — api/2.0/serving-endpoints // --------------------------------------------------------------------------- @@ -152,24 +158,26 @@ async fn fetch_v1_models( .bearer_auth(bearer) .send() .await - .map_err(|e| AgentError::Llm(format!("Databricks model discovery request failed: {e}")))?; + .map_err(|e| { + AgentError::Llm(format!( + "Databricks serving-endpoints catalog request failed: {e}" + )) + })?; let status = response.status(); if !status.is_success() { - let body = response.text().await.unwrap_or_default(); - if status.as_u16() == 401 { - return Err(AgentError::LlmAuth(format!( - "Databricks model discovery HTTP {status}" - ))); - } - return Err(AgentError::Llm(format!( - "Databricks model discovery HTTP {status}: {body}" - ))); + return Err(catalog_http_error( + "Databricks serving-endpoints catalog", + status, + response, + bearer, + ) + .await); } - let json: serde_json::Value = response.json().await.map_err(|e| { + let json: Value = response.json().await.map_err(|e| { AgentError::Llm(format!( - "Databricks model discovery response parse failed: {e}" + "Databricks serving-endpoints catalog response parse failed: {e}" )) })?; @@ -180,11 +188,11 @@ async fn fetch_v1_models( /// /// 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, AgentError> { +/// including over silently dropping, per the existing v1 contract. +pub(crate) fn parse_v1_endpoints(json: &Value) -> Result, AgentError> { let endpoints = json .get("endpoints") - .and_then(|v| v.as_array()) + .and_then(Value::as_array) .ok_or_else(|| { AgentError::Llm( "Databricks model discovery: unexpected response (missing 'endpoints' array)" @@ -201,7 +209,7 @@ pub(crate) fn parse_v1_endpoints(json: &serde_json::Value) -> Result Result Result String { .collect() } +/// Fetch both Databricks v2 catalogs concurrently and merge them into the +/// selectable model list. One catalog may be unavailable; an empty result is +/// still authoritative and never falls through to the known-model fallback +/// when a visibility filter is active. async fn fetch_v2_models( http: &Client, host: &str, bearer: &str, + filter: Option<&DatabricksModelFilter>, + allow_partial_auth_failure: bool, ) -> Result, AgentError> { - let mut all_endpoints: Vec = Vec::new(); + let workspace = fetch_catalog_pages( + http, + host, + bearer, + "Databricks workspace endpoint catalog", + "/api/ai-gateway/v2/endpoints", + WORKSPACE_CATALOG_QUERY, + parse_v2_endpoints_page, + ); + let unity_catalog = fetch_catalog_pages( + http, + host, + bearer, + "Databricks Unity Catalog model-service catalog", + "/api/2.1/unity-catalog/model-services", + UNITY_CATALOG_QUERY, + parse_uc_model_services_page, + ); + + let (workspace, unity_catalog) = tokio::join!(workspace, unity_catalog); + let (workspace, unity_catalog, both_succeeded) = match (workspace, unity_catalog) { + (Ok(workspace), Ok(unity_catalog)) => (workspace, unity_catalog, true), + (Ok(workspace), Err(error)) => { + if matches!(&error, AgentError::LlmAuth(_)) && !allow_partial_auth_failure { + return Err(error); + } + tracing::warn!( + catalog = "unity-catalog model-services", + error = %error, + "Databricks model discovery degraded: catalog unavailable" + ); + (workspace, Vec::new(), false) + } + (Err(error), Ok(unity_catalog)) => { + if matches!(&error, AgentError::LlmAuth(_)) && !allow_partial_auth_failure { + return Err(error); + } + tracing::warn!( + catalog = "workspace ai-gateway v2 endpoints", + error = %error, + "Databricks model discovery degraded: catalog unavailable" + ); + (Vec::new(), unity_catalog, false) + } + (Err(workspace_error), Err(unity_catalog_error)) => { + return Err(combined_catalog_error(workspace_error, unity_catalog_error)); + } + }; + + Ok(merge_v2_models( + workspace, + unity_catalog, + filter, + both_succeeded && filter.is_none(), + )) +} + +fn combined_catalog_error(workspace: AgentError, unity_catalog: AgentError) -> AgentError { + let auth_failure = matches!(&workspace, AgentError::LlmAuth(_)) + || matches!(&unity_catalog, AgentError::LlmAuth(_)); + let message = format!( + "Databricks v2 model discovery failed: workspace endpoint catalog: {workspace}; Unity Catalog model-service catalog: {unity_catalog}" + ); + if auth_failure { + AgentError::LlmAuth(message) + } else { + AgentError::Llm(message) + } +} + +fn merge_v2_models( + workspace: Vec, + mut unity_catalog: Vec, + filter: Option<&DatabricksModelFilter>, + allow_known_model_fallback: bool, +) -> Vec { + let mut seen_ids = HashSet::new(); + let mut merged = Vec::with_capacity(workspace.len() + unity_catalog.len()); + + // Workspace endpoints are ordered newest-first across all pages. + let mut workspace = workspace; + sort_v2_endpoints_newest_first(&mut workspace); + for endpoint in workspace { + if seen_ids.insert(endpoint.entry.id.clone()) { + merged.push(endpoint.entry); + } + } + + // UC has no user-facing recency contract. Sort by the raw FQN for stable + // picker order, then deduplicate only by raw selectable id. + unity_catalog.sort_unstable_by(|a, b| a.id.cmp(&b.id)); + for entry in unity_catalog { + if seen_ids.insert(entry.id.clone()) { + merged.push(entry); + } + } + + if merged.is_empty() && allow_known_model_fallback && filter.is_none() { + merged = authenticated_empty_v2_catalog(); + } + + apply_model_filter(merged, filter) +} + +async fn fetch_catalog_pages( + http: &Client, + host: &str, + bearer: &str, + catalog: &'static str, + path: &'static str, + initial_query: &str, + parse_page: fn(&Value) -> CatalogPage, +) -> Result, AgentError> { + let base_url = format!("{host}{path}"); + let mut all_items = Vec::new(); let mut page_token: Option = None; - let base_url = format!("{host}/api/ai-gateway/v2/endpoints"); + let mut seen_tokens = HashSet::new(); - // 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. + for _page in 0..MAX_CATALOG_PAGES { let url = match &page_token { - Some(tok) => format!( - "{base_url}?page_size=100&page_token={}", - percent_encode(tok) + Some(token) => format!( + "{base_url}{initial_query}&page_token={}", + percent_encode(token) ), - None => format!("{base_url}?page_size=100"), + None => format!("{base_url}{initial_query}"), }; let response = http .get(&url) .bearer_auth(bearer) .send() .await - .map_err(|e| { - AgentError::Llm(format!("Databricks v2 model discovery request failed: {e}")) - })?; + .map_err(|e| AgentError::Llm(format!("{catalog} request failed: {e}")))?; let status = response.status(); if !status.is_success() { - let body = response.text().await.unwrap_or_default(); - if status.as_u16() == 401 { - return Err(AgentError::LlmAuth(format!( - "Databricks v2 model discovery HTTP {status}" + return Err(catalog_http_error(catalog, status, response, bearer).await); + } + + let json: Value = response + .json() + .await + .map_err(|e| AgentError::Llm(format!("{catalog} response parse failed: {e}")))?; + let (items, next_token) = parse_page(&json) + .map_err(|error| catalog_context_error(catalog, error, "response parse failed"))?; + all_items.extend(items); + + match next_token { + None => return Ok(all_items), + Some(next_token) if seen_tokens.insert(next_token.clone()) => { + page_token = Some(next_token); + } + Some(next_token) => { + return Err(AgentError::Llm(format!( + "{catalog} pagination repeated page token {next_token:?}" ))); } - 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}" - )) - })?; + Err(AgentError::Llm(format!( + "{catalog} pagination exhausted after {MAX_CATALOG_PAGES} pages" + ))) +} - let (page_endpoints, next) = parse_v2_endpoints_page(&json)?; - all_endpoints.extend(page_endpoints); +async fn catalog_http_error( + catalog: &str, + status: reqwest::StatusCode, + response: reqwest::Response, + bearer: &str, +) -> AgentError { + if status.as_u16() == 401 { + // Do not read or include a provider body for auth failures. Some + // gateways echo authorization material in diagnostic payloads. + return AgentError::LlmAuth(format!("{catalog} HTTP {status}")); + } - match next { - Some(tok) if Some(&tok) != page_token.as_ref() => page_token = Some(tok), + let body = if bearer.len() > MAX_CATALOG_ERROR_BODY_BYTES { + // A token longer than the diagnostic bound cannot be safely searched in + // a bounded prefix. Do not return a partial provider body that might + // expose any part of it. + String::new() + } else { + let read_limit = MAX_CATALOG_ERROR_BODY_BYTES.saturating_add(bearer.len()); + let body = read_catalog_error_body(response, read_limit).await; + if bearer.is_empty() { + body + } else { + body.replace(bearer, "[redacted]") + } + }; + let body = truncate_utf8_bytes(&body, MAX_CATALOG_ERROR_BODY_BYTES); + let classification = if status.as_u16() == 499 || status.is_server_error() { + "transient" + } else { + "failed" + }; + AgentError::Llm(format!("{catalog} {classification} HTTP {status}: {body}")) +} + +async fn read_catalog_error_body(mut response: reqwest::Response, limit: usize) -> String { + let mut body = Vec::with_capacity(limit); + while body.len() < limit { + match response.chunk().await { + Ok(Some(chunk)) => { + let take = chunk.len().min(limit - body.len()); + body.extend_from_slice(&chunk[..take]); + if take < chunk.len() { + break; + } + } _ => break, } } + String::from_utf8_lossy(&body).into_owned() +} - // Fall back to known-model list if the API returned nothing. - if all_endpoints.is_empty() { - return Ok(authenticated_empty_v2_catalog()); +fn truncate_utf8_bytes(value: &str, max_bytes: usize) -> String { + if value.len() <= max_bytes { + return value.to_string(); } + let mut end = max_bytes; + while !value.is_char_boundary(end) { + end -= 1; + } + value[..end].to_string() +} - sort_v2_endpoints_newest_first(&mut all_endpoints); - - Ok(all_endpoints - .into_iter() - .map(|endpoint| endpoint.entry) - .collect()) +fn catalog_context_error(catalog: &str, error: AgentError, context: &str) -> AgentError { + match error { + AgentError::LlmAuth(message) => { + AgentError::LlmAuth(format!("{catalog} {context}: {message}")) + } + AgentError::Llm(message) => AgentError::Llm(format!("{catalog} {context}: {message}")), + other => AgentError::Llm(format!("{catalog} {context}: {other}")), + } } -/// A v2 gateway endpoint plus the key discovery orders the catalog by. +/// A v2 gateway endpoint plus the key discovery order field. #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct V2Endpoint { pub(crate) entry: ModelEntry, @@ -328,24 +514,15 @@ pub(crate) struct V2Endpoint { /// /// 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 { +/// wire-shape change does not silently drop every endpoint to the bottom. +fn endpoint_created_ms(endpoint: &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. +/// Order workspace endpoints newest-first, breaking ties by name. pub(crate) fn sort_v2_endpoints_newest_first(endpoints: &mut [V2Endpoint]) { endpoints.sort_by(|a, b| { // `None` < `Some(_)`, so reversing puts timestamped endpoints first. @@ -357,29 +534,20 @@ pub(crate) fn sort_v2_endpoints_newest_first(endpoints: &mut [V2Endpoint]) { /// Parse one page of a `GET api/ai-gateway/v2/endpoints` response. /// -/// 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`]). +/// Page order is preserved here; the caller sorts once every page is in. pub(crate) fn parse_v2_endpoints_page( - json: &serde_json::Value, + json: &Value, ) -> Result<(Vec, Option), 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(), - ) - })?; + .and_then(Value::as_array) + .ok_or_else(|| AgentError::Llm("unexpected response (missing 'endpoints' array)".into()))?; let models = endpoints .iter() .filter_map(|endpoint| { let name = endpoint.get("name")?.as_str()?.to_string(); - if !is_chat_capable_endpoint(&name) { + if name.is_empty() { return None; } Some(V2Endpoint { @@ -392,15 +560,51 @@ pub(crate) fn parse_v2_endpoints_page( }) .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); - + let next_page_token = next_page_token(json); Ok((models, next_page_token)) } +/// Parse one page of a `GET api/2.1/unity-catalog/model-services` response. +/// +/// Unity Catalog resource names are returned as `model-services/..`. +/// Only the exact resource prefix and a structurally valid three-component FQN +/// are selectable. All other resources are ignored without capability/name +/// heuristics; the positive visibility filter is the only further restriction. +pub(crate) fn parse_uc_model_services_page( + json: &Value, +) -> Result<(Vec, Option), AgentError> { + let services = json + .get("model_services") + .and_then(Value::as_array) + .ok_or_else(|| { + AgentError::Llm("unexpected response (missing 'model_services' array)".into()) + })?; + + let models = services + .iter() + .filter_map(|service| { + let resource_name = service.get("name")?.as_str()?; + let fqn = resource_name.strip_prefix("model-services/")?; + if !crate::llm::is_model_service_fqn(fqn) { + return None; + } + Some(ModelEntry { + id: fqn.to_string(), + name: curated_model_name(fqn), + }) + }) + .collect(); + + Ok((models, next_page_token(json))) +} + +fn next_page_token(json: &Value) -> Option { + json.get("next_page_token") + .and_then(Value::as_str) + .filter(|token| !token.is_empty()) + .map(str::to_string) +} + // --------------------------------------------------------------------------- // Tests // --------------------------------------------------------------------------- @@ -409,6 +613,8 @@ pub(crate) fn parse_v2_endpoints_page( mod tests { use super::*; use async_trait::async_trait; + use axum::{extract::Query, http::StatusCode, routing::get, Json, Router}; + use std::collections::HashMap; use std::sync::atomic::{AtomicUsize, Ordering}; struct RefreshingTestTokenSource { @@ -470,7 +676,7 @@ mod tests { let source = Arc::new(RefreshingTestTokenSource { refreshes: AtomicUsize::new(0), }); - let cfg = Config::for_discovery(Provider::DatabricksV2, String::new(), host); + let cfg = Config::for_discovery(Provider::DatabricksV2, String::new(), host, None); let models = discover_databricks_models_with_token_source(&cfg, source.clone()) .await .unwrap(); @@ -480,6 +686,275 @@ mod tests { assert_eq!(requests.load(Ordering::SeqCst), 2); } + #[tokio::test] + async fn v2_discovery_merges_workspace_and_unity_catalog_after_filtering() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let host = format!("http://{}", listener.local_addr().unwrap()); + let app = Router::new() + .route( + "/api/ai-gateway/v2/endpoints", + get(|Query(query): Query>| async move { + assert_eq!(query.get("page_size").map(String::as_str), Some("100")); + Json(serde_json::json!({ + "endpoints": [ + {"name": "blocked-workspace", "created_timestamp": 3}, + {"name": "allowed-workspace", "created_timestamp": 2}, + ], + "next_page_token": null, + })) + }), + ) + .route( + "/api/2.1/unity-catalog/model-services", + get(|Query(query): Query>| async move { + assert_eq!(query.get("page_size").map(String::as_str), Some("100")); + assert_eq!(query.get("view").map(String::as_str), Some("FULL")); + Json(serde_json::json!({ + "model_services": [ + {"name": "model-services/catalog.schema.blocked-service"}, + {"name": "model-services/catalog.schema.allowed-service"}, + {"name": "model-services/catalog.schema.allowed-service"}, + ], + "next_page_token": null, + })) + }), + ); + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + + let filter = + DatabricksModelFilter::parse(Some("allowed-*,catalog.schema.allowed-*")).unwrap(); + let cfg = Config::for_discovery(Provider::DatabricksV2, "token".into(), host, filter); + let models = discover_databricks_models(&cfg).await.unwrap(); + assert_eq!( + models + .iter() + .map(|model| model.id.as_str()) + .collect::>(), + vec!["allowed-workspace", "catalog.schema.allowed-service"] + ); + } + + #[tokio::test] + async fn v2_discovery_keeps_unity_catalog_when_workspace_catalog_fails() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let host = format!("http://{}", listener.local_addr().unwrap()); + let app = Router::new() + .route( + "/api/ai-gateway/v2/endpoints", + get(|| async { (StatusCode::SERVICE_UNAVAILABLE, "workspace unavailable") }), + ) + .route( + "/api/2.1/unity-catalog/model-services", + get(|| async { + Json(serde_json::json!({ + "model_services": [ + {"name": "model-services/catalog.schema.uc-service"} + ], + "next_page_token": null, + })) + }), + ); + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + + let cfg = Config::for_discovery(Provider::DatabricksV2, "token".into(), host, None); + let models = discover_databricks_models(&cfg).await.unwrap(); + assert_eq!( + models + .iter() + .map(|model| model.id.as_str()) + .collect::>(), + vec!["catalog.schema.uc-service"] + ); + } + + #[tokio::test] + async fn v2_empty_catalog_fallback_is_disabled_by_filter() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let host = format!("http://{}", listener.local_addr().unwrap()); + let app = Router::new() + .route( + "/api/ai-gateway/v2/endpoints", + get(|| async { + Json(serde_json::json!({ + "endpoints": [], + "next_page_token": null, + })) + }), + ) + .route( + "/api/2.1/unity-catalog/model-services", + get(|| async { + Json(serde_json::json!({ + "model_services": [], + "next_page_token": null, + })) + }), + ); + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + + let unfiltered = + Config::for_discovery(Provider::DatabricksV2, "token".into(), host.clone(), None); + let fallback = discover_databricks_models(&unfiltered).await.unwrap(); + assert_eq!( + fallback + .iter() + .map(|model| model.id.as_str()) + .collect::>(), + crate::model_capabilities::databricks_v2_known_models() + .iter() + .map(String::as_str) + .collect::>() + ); + + let filter = DatabricksModelFilter::parse(Some("no-match")).unwrap(); + let filtered = Config::for_discovery(Provider::DatabricksV2, "token".into(), host, filter); + assert!(discover_databricks_models(&filtered) + .await + .unwrap() + .is_empty()); + } + + #[tokio::test] + async fn catalog_pagination_encodes_tokens_and_rejects_repeated_tokens() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let host = format!("http://{}", listener.local_addr().unwrap()); + let app = Router::new().route( + "/catalog", + get(|Query(query): Query>| async move { + match query.get("page_token").map(String::as_str) { + None => Json(serde_json::json!({ + "endpoints": [{"name": "first"}], + "next_page_token": "token with/slash", + })), + Some("token with/slash") => Json(serde_json::json!({ + "endpoints": [{"name": "second"}], + })), + Some(other) => panic!("unexpected decoded page token: {other}"), + } + }), + ); + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + + let entries = fetch_catalog_pages( + &Client::new(), + &host, + "token", + "test catalog", + "/catalog", + "?page_size=100", + parse_v2_endpoints_page, + ) + .await + .unwrap(); + assert_eq!(entries.len(), 2); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let host = format!("http://{}", listener.local_addr().unwrap()); + let app = Router::new().route( + "/catalog", + get(|| async { + Json(serde_json::json!({ + "endpoints": [{"name": "loop"}], + "next_page_token": "same-token", + })) + }), + ); + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + let error = fetch_catalog_pages( + &Client::new(), + &host, + "token", + "test catalog", + "/catalog", + "?page_size=100", + parse_v2_endpoints_page, + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("repeated page token")); + } + + #[tokio::test] + async fn catalog_pagination_errors_after_the_finite_page_cap() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let host = format!("http://{}", listener.local_addr().unwrap()); + let requests = Arc::new(AtomicUsize::new(0)); + let requests_for_handler = requests.clone(); + let app = Router::new().route( + "/catalog", + get(move |Query(_query): Query>| { + let page = requests_for_handler.fetch_add(1, Ordering::SeqCst) + 1; + async move { + Json(serde_json::json!({ + "endpoints": [{"name": format!("model-{page}")}], + "next_page_token": format!("token-{page}"), + })) + } + }), + ); + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + + let error = fetch_catalog_pages( + &Client::new(), + &host, + "token", + "test catalog", + "/catalog", + "?page_size=100", + parse_v2_endpoints_page, + ) + .await + .unwrap_err(); + assert!(error + .to_string() + .contains("pagination exhausted after 20 pages")); + assert_eq!(requests.load(Ordering::SeqCst), 20); + } + + #[test] + fn v1_filter_applies_to_raw_ids_after_endpoint_filtering() { + let filter = DatabricksModelFilter::parse(Some("allowed-*")).unwrap(); + let models = apply_model_filter( + vec![ + ModelEntry { + id: "allowed-model".into(), + name: "Allowed".into(), + }, + ModelEntry { + id: "blocked-model".into(), + name: "Blocked".into(), + }, + ], + filter.as_ref(), + ); + assert_eq!(models.len(), 1); + assert_eq!(models[0].id, "allowed-model"); + } + + #[test] + fn catalog_error_body_is_bounded_and_redacts_bearer() { + // This pure assertion documents the byte-bound helper used after the + // provider response is read. The network path exercises the same + // redaction before truncation; keeping the helper pure makes the UTF-8 + // boundary behavior explicit. + let value = format!("{}é", "x".repeat(MAX_CATALOG_ERROR_BODY_BYTES)); + let truncated = truncate_utf8_bytes(&value, MAX_CATALOG_ERROR_BODY_BYTES); + assert_eq!(truncated.len(), MAX_CATALOG_ERROR_BODY_BYTES); + assert!(truncated.is_char_boundary(truncated.len())); + } + #[test] fn v1_parse_filters_ready_chat_endpoints() { let json = serde_json::json!({ @@ -578,10 +1053,7 @@ mod tests { } #[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. + fn v2_parse_keeps_all_nonempty_endpoint_names_without_keyword_filtering() { let json = serde_json::json!({ "endpoints": [ {"name": "databricks-bge-large-en"}, @@ -594,10 +1066,139 @@ mod tests { 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"] + vec![ + "databricks-bge-large-en", + "databricks-gte-large-en", + "databricks-qwen3-embedding-0-6b", + "databricks-claude-opus-5", + "databricks-gemini-3-pro-image", + ] + ); + } + + #[test] + fn uc_parse_requires_exact_prefix_and_structural_fqn() { + let json = serde_json::json!({ + "model_services": [ + {"name": "model-services/data_tools.goose.kimi-k3"}, + {"name": "model-services/catalog.schema.claude-gpt-5"}, + {"name": "model-services/two.parts"}, + {"name": "model-services/too.many.parts.here"}, + {"name": "Model-services/wrong.case.service"}, + {"name": "models/data_tools.goose.other"}, + {"name": "model-services/.schema.service"}, + {"name": "model-services/catalog..service"}, + {"name": "model-services/catalog.schema."}, + {"name": "model-services/catalog.schema/service"}, + ], + "next_page_token": "next token/1" + }); + + let (models, next) = parse_uc_model_services_page(&json).unwrap(); + let ids: Vec<&str> = models.iter().map(|m| m.id.as_str()).collect(); + assert_eq!( + ids, + vec!["data_tools.goose.kimi-k3", "catalog.schema.claude-gpt-5"] + ); + assert_eq!(next.as_deref(), Some("next token/1")); + } + + #[test] + fn uc_parse_requires_model_services_array() { + let err = parse_uc_model_services_page(&serde_json::json!({"data": []})).unwrap_err(); + assert!(err.to_string().contains("missing 'model_services' array")); + } + + #[test] + fn merge_deduplicates_raw_ids_and_preserves_workspace_then_lexical_uc_order() { + let workspace = vec![ + V2Endpoint { + entry: ModelEntry { + id: "workspace-new".into(), + name: "workspace-new".into(), + }, + created_ms: Some(2), + }, + V2Endpoint { + entry: ModelEntry { + id: "duplicate".into(), + name: "duplicate".into(), + }, + created_ms: Some(1), + }, + ]; + let uc = vec![ + ModelEntry { + id: "z.schema.service".into(), + name: "z.schema.service".into(), + }, + ModelEntry { + id: "a.schema.service".into(), + name: "a.schema.service".into(), + }, + ModelEntry { + id: "duplicate".into(), + name: "same leaf".into(), + }, + ModelEntry { + id: "a.other.service".into(), + name: "same leaf".into(), + }, + ]; + + let models = merge_v2_models(workspace, uc, None, false); + let ids: Vec<&str> = models.iter().map(|model| model.id.as_str()).collect(); + assert_eq!( + ids, + vec![ + "workspace-new", + "duplicate", + "a.other.service", + "a.schema.service", + "z.schema.service", + ] + ); + } + + #[test] + fn merge_applies_filter_after_union_and_does_not_restore_fallback() { + let filter = DatabricksModelFilter::parse(Some("allowed.*")).unwrap(); + let filter = filter.as_ref(); + let workspace = vec![V2Endpoint { + entry: ModelEntry { + id: "blocked-workspace".into(), + name: "blocked-workspace".into(), + }, + created_ms: Some(1), + }]; + let uc = vec![ModelEntry { + id: "allowed.schema.service".into(), + name: "allowed.schema.service".into(), + }]; + let models = merge_v2_models(workspace, uc, filter, false); + assert_eq!( + models.iter().map(|m| m.id.as_str()).collect::>(), + vec!["allowed.schema.service"] + ); + + let no_match = DatabricksModelFilter::parse(Some("no-match")).unwrap(); + assert!(merge_v2_models(Vec::new(), Vec::new(), no_match.as_ref(), true).is_empty()); + } + + #[test] + fn merge_uses_known_fallback_only_for_unfiltered_successful_empty_union() { + let models = merge_v2_models(Vec::new(), Vec::new(), None, true); + assert_eq!( + models + .iter() + .map(|model| model.id.as_str()) + .collect::>(), + crate::model_capabilities::databricks_v2_known_models() + .iter() + .map(String::as_str) + .collect::>() ); } @@ -716,16 +1317,4 @@ mod tests { "custom-unlisted-endpoint" ); } - - #[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")); - } } diff --git a/crates/buzz-agent/src/config.rs b/crates/buzz-agent/src/config.rs index abf17004e8a..202d73e5548 100644 --- a/crates/buzz-agent/src/config.rs +++ b/crates/buzz-agent/src/config.rs @@ -429,6 +429,96 @@ pub enum Provider { OpenRouter, } +/// Optional visibility filter for the Databricks model catalog. +/// +/// Each comma-separated pattern is trimmed and matched against the complete, +/// case-sensitive model id. Only `*` (zero or more characters) and `?` (one +/// character) have wildcard semantics; all other characters are literals. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DatabricksModelFilter { + patterns: Vec, +} + +impl DatabricksModelFilter { + /// Parse `DATABRICKS_MODEL_FILTER`-style input. + /// + /// Unset or whitespace-only input disables filtering. A nonblank value must + /// contain at least one nonblank comma-separated pattern. + pub fn parse(raw: Option<&str>) -> Result, String> { + let Some(raw) = raw else { + return Ok(None); + }; + + if raw.trim().is_empty() { + return Ok(None); + } + + let patterns: Vec = raw + .split(',') + .map(str::trim) + .filter(|pattern| !pattern.is_empty()) + .map(str::to_owned) + .collect(); + if patterns.is_empty() { + return Err( + "config: DATABRICKS_MODEL_FILTER must contain at least one nonblank pattern".into(), + ); + } + + Ok(Some(Self { patterns })) + } + + /// Return whether the complete model id matches at least one pattern. + pub fn matches(&self, model_id: &str) -> bool { + self.patterns + .iter() + .any(|pattern| glob_matches(pattern, model_id)) + } +} + +/// Match one full-string `*`/`?` pattern without treating any other character +/// as syntax. The inputs are converted to Unicode scalar values so `?` means +/// one character rather than one UTF-8 byte. +fn glob_matches(pattern: &str, value: &str) -> bool { + let pattern: Vec = pattern.chars().collect(); + let value: Vec = value.chars().collect(); + let mut pattern_index = 0; + let mut value_index = 0; + let mut star_index = None; + let mut star_value_index = 0; + + while value_index < value.len() { + match pattern.get(pattern_index) { + Some('?') => { + pattern_index += 1; + value_index += 1; + } + Some('*') => { + star_index = Some(pattern_index); + star_value_index = value_index; + pattern_index += 1; + } + Some(character) if *character == value[value_index] => { + pattern_index += 1; + value_index += 1; + } + _ if star_index.is_some() => { + if let Some(star_index) = star_index { + pattern_index = star_index + 1; + } + star_value_index += 1; + value_index = star_value_index; + } + _ => return false, + } + } + + while matches!(pattern.get(pattern_index), Some('*')) { + pattern_index += 1; + } + pattern_index == pattern.len() +} + /// Which OpenAI-family HTTP API to call. Set via `OPENAI_COMPAT_API` /// (`auto|chat|responses`); ignored when `provider = Anthropic`. `Auto` /// picks Responses for `*.openai.com`, Chat Completions otherwise, and @@ -509,6 +599,9 @@ pub struct Config { /// Default (env unset/empty) is `None` — hooks are off unless the /// operator explicitly opts in. pub hook_servers: HookServers, + /// The effective `DATABRICKS_MODEL_FILTER` value. This is parsed by the + /// caller and passed explicitly so discovery never consults process env. + pub databricks_model_filter: Option, pub api_key: String, pub model: String, pub base_url: String, @@ -643,6 +736,9 @@ impl Config { stop_max_rejections: parse_env("BUZZ_AGENT_STOP_MAX_REJECTIONS", 3u32)?, require_reply: parse_env("BUZZ_AGENT_REQUIRE_REPLY", 0u8)? != 0, hook_servers: parse_hook_servers_env("MCP_HOOK_SERVERS"), + databricks_model_filter: DatabricksModelFilter::parse( + env("DATABRICKS_MODEL_FILTER").as_deref(), + )?, hints_enabled: parse_env("BUZZ_AGENT_NO_HINTS", 0u8)? == 0, thinking_effort: parse_thinking_effort(env("BUZZ_AGENT_THINKING_EFFORT").as_deref())?, thinking_summary: parse_thinking_summary( @@ -660,7 +756,12 @@ impl Config { /// 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 { + pub fn for_discovery( + provider: Provider, + api_key: String, + base_url: String, + databricks_model_filter: Option, + ) -> Self { Self { provider, api_key, @@ -691,6 +792,7 @@ impl Config { stop_max_rejections: 0, require_reply: false, hook_servers: HookServers::None, + databricks_model_filter, hints_enabled: false, thinking_effort: None, thinking_summary: ThinkingSummary::Auto, @@ -1024,6 +1126,61 @@ fn parse_hook_servers(raw: Option<&str>) -> HookServers { mod tests { use super::*; + #[test] + fn databricks_model_filter_unset_and_blank_disable_filtering() { + for raw in [None, Some(""), Some(" ")] { + assert_eq!(DatabricksModelFilter::parse(raw).unwrap(), None); + } + } + + #[test] + fn databricks_model_filter_rejects_nonblank_input_without_patterns() { + let error = DatabricksModelFilter::parse(Some(" , , ")).unwrap_err(); + assert!(error.contains("DATABRICKS_MODEL_FILTER"), "{error}"); + } + + #[test] + fn databricks_model_filter_matches_exact_full_string_case_sensitively() { + let filter = DatabricksModelFilter::parse(Some("data_tools.goose.kimi-k3")).unwrap(); + assert!(filter.as_ref().unwrap().matches("data_tools.goose.kimi-k3")); + assert!(!filter + .as_ref() + .unwrap() + .matches("prefix.data_tools.goose.kimi-k3")); + assert!(!filter.as_ref().unwrap().matches("data_tools.goose.Kimi-k3")); + } + + #[test] + fn databricks_model_filter_matches_star_and_question_mark() { + let filter = + DatabricksModelFilter::parse(Some("databricks-*,data_tools.goose.????-k3")).unwrap(); + let filter = filter.as_ref().unwrap(); + assert!(filter.matches("databricks-gpt-5")); + assert!(filter.matches("data_tools.goose.kimi-k3")); + assert!(!filter.matches("data_tools.goose.kimi-k33")); + assert!(!filter.matches("other-model")); + } + + #[test] + fn databricks_model_filter_trims_multiple_patterns_and_preserves_no_match() { + let filter = DatabricksModelFilter::parse(Some(" first , second-model , third-* ")) + .unwrap() + .unwrap(); + assert!(filter.matches("first")); + assert!(filter.matches("second-model")); + assert!(filter.matches("third-model")); + assert!(!filter.matches("fourth-model")); + } + + #[test] + fn databricks_model_filter_question_mark_matches_one_unicode_character() { + let filter = DatabricksModelFilter::parse(Some("goose-? ")) + .unwrap() + .unwrap(); + assert!(filter.matches("goose-é")); + assert!(!filter.matches("goose-eé")); + } + #[test] fn hook_servers_unset_is_none() { assert!(matches!(parse_hook_servers(None), HookServers::None)); @@ -1834,7 +1991,8 @@ mod tests { provider: Provider, thinking_effort: Option, ) -> Config { - let mut cfg = Config::for_discovery(provider, "key".into(), "https://example.com".into()); + let mut cfg = + Config::for_discovery(provider, "key".into(), "https://example.com".into(), None); cfg.model = "some-model".into(); cfg.thinking_effort = thinking_effort; // for_discovery sets max_output_tokens=1 and max_context_tokens=200_001 which satisfies diff --git a/crates/buzz-agent/src/lib.rs b/crates/buzz-agent/src/lib.rs index 208acc34692..b094a0f9fd7 100644 --- a/crates/buzz-agent/src/lib.rs +++ b/crates/buzz-agent/src/lib.rs @@ -368,13 +368,17 @@ async fn initialize(app: &Arc, id: Value, params: Value, wire_tx: &WireSend .await; } -/// Resolve the Databricks model catalog for one `session/new` call. +/// Resolve a 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 error is returned and the cell is intentionally left empty so the next session retries. +/// The active filter is part of the result's authority: discovery failure may +/// not fall back to a configured model when it is present, because that would +/// bypass the same restriction applied to a successful catalog. /// -/// Extracted from `session_new` so that tests can drive this path with an injected +/// Tries to use a previously cached successful discovery result. If the cache +/// is empty, runs `discover` and — on success — populates the cache. On failure +/// the error is returned and the cell remains empty so the next session retries. +/// +/// Extracted from `session_new` so 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>, @@ -383,7 +387,7 @@ async fn resolve_models_catalog( cache.get_or_try_init(|| discover).await.cloned() } -/// Return the configured model as a one-entry catalog for this response. +/// Return the configured model as an unfiltered discovery fallback. /// /// This value is never written to `models_cache`; failed discovery must be retried by /// the next session rather than pinning degraded state for the process lifetime. @@ -398,6 +402,17 @@ fn configured_model_fallback(model: &str) -> Vec { vec![ModelEntry { id: model, name }] } +/// A discovery failure may use the configured model only when no visibility +/// filter is active. Returning that model under an active filter would silently +/// bypass the operator's authoritative catalog restriction. +fn discovery_error_fallback(cfg: &Config) -> Vec { + if cfg.databricks_model_filter.is_some() { + Vec::new() + } else { + configured_model_fallback(&cfg.model) + } +} + async fn session_new(app: &Arc, id: Value, params: Value, wire_tx: &WireSender) { let p: SessionNewParams = match decode(params, "session/new") { Ok(p) => p, @@ -482,16 +497,18 @@ async fn session_new(app: &Arc, id: Value, params: Value, wire_tx: &WireSen Err(error @ AgentError::LlmAuth(_)) => { tracing::warn!( error = %error, - "Databricks OAuth model catalog unavailable; using configured model" + filter_active = app.cfg.databricks_model_filter.is_some(), + "Databricks OAuth model catalog unavailable; using filter-aware fallback" ); - configured_model_fallback(&app.cfg.model) + discovery_error_fallback(&app.cfg) } Err(error) => { tracing::warn!( error = %error, - "Databricks model catalog unavailable; using configured model" + filter_active = app.cfg.databricks_model_filter.is_some(), + "Databricks model catalog unavailable; using filter-aware fallback" ); - configured_model_fallback(&app.cfg.model) + discovery_error_fallback(&app.cfg) } }; models diff --git a/crates/buzz-agent/src/llm.rs b/crates/buzz-agent/src/llm.rs index 599954292aa..35a71717497 100644 --- a/crates/buzz-agent/src/llm.rs +++ b/crates/buzz-agent/src/llm.rs @@ -170,13 +170,17 @@ impl Llm { ) } DatabricksV2Route::MlflowChatCompletions => { - // MLflow Chat path (OpenAI-shaped): normalize effort via manifest. + // UC model-service FQNs are not native provider model names. Build + // against a neutral capability identity so family-name text in a + // catalog/schema/service component cannot change request semantics, + // then restore the actual FQN in the payload. + let capability_model = databricks_v2_capability_model(effective_model); let e = effort - .map(|ef| normalize_effort_for_databricks_v2(ef, effective_model)); - ( - openai_body(cfg, system_prompt, history, tools, effective_model, e), - parse_openai as OpenAiParse, - ) + .map(|ef| normalize_effort_for_databricks_v2(ef, capability_model)); + let mut body = + openai_body(cfg, system_prompt, history, tools, capability_model, e); + body["model"] = json!(effective_model); + (body, parse_openai as OpenAiParse) } }) .await @@ -325,18 +329,21 @@ impl Llm { }), parse_anthropic as OpenAiParse, ), - DatabricksV2Route::MlflowChatCompletions => ( - json!({ - "model": effective_model, + DatabricksV2Route::MlflowChatCompletions => { + let capability_model = + databricks_v2_capability_model(effective_model); + let mut body = json!({ + "model": capability_model, "stream": false, "max_completion_tokens": max_output_tokens, "messages": [ { "role": "system", "content": system_prompt }, { "role": "user", "content": user_prompt }, ], - }), - parse_openai as OpenAiParse, - ), + }); + body["model"] = json!(effective_model); + (body, parse_openai as OpenAiParse) + } }) .await?; Ok(r.text) @@ -967,26 +974,47 @@ fn is_responses_required_error(body: &str) -> bool { || b.contains("use the responses api") } -/// Resolve the Databricks v2 AI Gateway wire route for `model` from the manifest. -/// -/// The route is a capability of the `(databricks_v2, model)` pair, owned by -/// `scripts/model-capabilities.json` and resolved by the shared interpreter — the -/// same authority that drives effort/label resolution. This function only maps the -/// manifest's route enum onto the three concrete wire routes this dispatch path can -/// serve; it holds no routing knowledge of its own. +/// Resolve the Databricks v2 AI Gateway wire route for `model`. /// -/// The manifest enum carries two non-wire variants that cannot occur here for a -/// concrete Databricks v2 model at dispatch time: -/// - `NotApplicable` is produced only for non-`databricks_v2` providers, and this -/// seam is reached only under `Provider::DatabricksV2`. -/// - `RouteUnknown` is produced only for a blank model id, which `Config` rejects at -/// startup (`DATABRICKS_MODEL` required) and `session/set_model` rejects at runtime -/// (empty `modelId` → `invalid_params`), so `effective_model` is never blank here. +/// Strict three-component Unity Catalog model-service FQNs are checked before +/// the manifest, because their namespace is catalog data rather than a model +/// family hint. Other IDs retain the manifest-owned route. +fn databricks_v2_capability_model(model: &str) -> &str { + if is_model_service_fqn(model) { + "model-service" + } else { + model + } +} + +/// Return whether `model` is exactly three non-empty dot-separated components. +/// This is deliberately structural: model-service names are catalog data, so +/// family words such as `claude` and `gpt-5` have no routing meaning here. +pub(crate) fn is_model_service_fqn(model: &str) -> bool { + let mut components = model.split('.'); + let (Some(catalog), Some(schema), Some(service)) = + (components.next(), components.next(), components.next()) + else { + return false; + }; + [catalog, schema, service].into_iter().all(|component| { + !component.is_empty() + && !component.chars().any(char::is_whitespace) + && !component.contains('/') + }) && components.next().is_none() +} + +/// Resolve the Databricks v2 AI Gateway wire route for `model`. /// -/// Both are folded into `MlflowChatCompletions` — the manifest's own concrete-unknown -/// fallback and the route a blank id would historically have taken — so an unforeseen -/// reshape degrades to the safe OpenAI-wire route rather than panicking. +/// A Unity Catalog model-service FQN always uses MLflow Chat Completions. The +/// namespace is data, not a provider-family hint, so names containing `claude` +/// or `gpt` must not reach native Anthropic/Responses routes. Non-FQN ids keep +/// the existing manifest projection unchanged. fn databricks_v2_route(model: &str) -> DatabricksV2Route { + if is_model_service_fqn(model) { + return DatabricksV2Route::MlflowChatCompletions; + } + use crate::model_capabilities::DatabricksV2Route as Manifest; match crate::model_capabilities::resolve("databricks_v2", model).databricks_v2_wire_route { Manifest::OpenaiResponses => DatabricksV2Route::OpenAiResponses, @@ -2615,6 +2643,7 @@ mod tests { stop_max_rejections: 0, require_reply: false, hook_servers: HookServers::None, + databricks_model_filter: None, api_key: "key".into(), model: "model".into(), base_url: "http://example.invalid".into(), @@ -2832,6 +2861,33 @@ mod tests { } } + #[tokio::test] + async fn databricks_v2_model_service_fqn_summary_uses_mlflow_chat() { + let model = "catalog.schema.claude-gpt-5"; + let (base_url, captured) = + spawn_sequence_stub(vec![StubHttpResponse::ok(chat_response("summary"))]).await; + let mut config = cfg(Provider::DatabricksV2); + config.base_url = base_url; + let llm = Llm::new(&config).unwrap(); + + let summary = llm + .summarize(&config, "system", "history", 128, model) + .await + .unwrap(); + assert_eq!(summary, "summary"); + + let requests = captured.lock().await; + let request = requests + .iter() + .find(|request| request.method == "POST") + .expect("summary must issue one POST"); + assert_eq!(request.path, "/v1/ai-gateway/mlflow/v1/chat/completions"); + let body = request.body.as_ref().expect("summary body"); + assert_eq!(body["model"], model); + assert!(body["messages"].is_array()); + assert_eq!(body["max_completion_tokens"], 128); + } + fn image_history() -> Vec { vec![ HistoryItem::User("describe the image".into()), @@ -3241,6 +3297,52 @@ mod tests { } } + #[test] + fn databricks_v2_model_service_fqn_shape_is_strict_and_precedes_manifest() { + use crate::model_capabilities::{resolve, DatabricksV2Route as Manifest}; + + for model in [ + "catalog.schema.service", + "catalog.schema.claude-gpt-5", + "data_tools.goose.kimi-k3", + ] { + assert!(is_model_service_fqn(model), "expected FQN shape: {model}"); + assert_eq!( + databricks_v2_route(model), + DatabricksV2Route::MlflowChatCompletions, + "FQN route must precede manifest family inference: {model}" + ); + } + + let manifest_route = + |model: &str| match resolve("databricks_v2", model).databricks_v2_wire_route { + Manifest::OpenaiResponses => DatabricksV2Route::OpenAiResponses, + Manifest::AnthropicMessages => DatabricksV2Route::AnthropicMessages, + Manifest::MlflowChat | Manifest::NotApplicable | Manifest::RouteUnknown => { + DatabricksV2Route::MlflowChatCompletions + } + }; + for model in [ + "catalog.schema", + "catalog..service", + ".schema.service", + "catalog.schema.", + "catalog.schema.service.extra", + "catalog/schema/service", + "catalog.schema service", + ] { + assert!( + !is_model_service_fqn(model), + "unexpected FQN shape: {model}" + ); + assert_eq!( + databricks_v2_route(model), + manifest_route(model), + "malformed/partial IDs must retain manifest routing: {model}" + ); + } + } + #[test] fn databricks_v2_dispatch_is_pure_manifest_projection() { // Mutation-bypass guard: the dispatch seam must be a pure projection of diff --git a/crates/buzz-agent/tests/databricks_oauth.rs b/crates/buzz-agent/tests/databricks_oauth.rs index fbe0dc1f862..f7d41acf46b 100644 --- a/crates/buzz-agent/tests/databricks_oauth.rs +++ b/crates/buzz-agent/tests/databricks_oauth.rs @@ -429,17 +429,25 @@ async fn spawn_capturing_server( let body: serde_json::Value = serde_json::from_slice(&buf[header_end..header_end + body_len]) .unwrap_or(json!(null)); + let is_unity_catalog = path.starts_with("/api/2.1/unity-catalog/model-services"); captured.lock().await.push(CapturedRequest { path, authorization, body, }); - let body = queue - .lock() - .await - .pop_front() - .unwrap_or_else(|| json!({ "error": "no canned response" })); - let body_s = serde_json::to_string(&body).unwrap(); + let response_body = if is_unity_catalog { + // v2 discovery probes both catalogs concurrently. Existing + // request-shape tests need only the workspace fixture, so the + // UC side is explicitly successful and empty. + json!({ "model_services": [], "next_page_token": null }) + } else { + queue + .lock() + .await + .pop_front() + .unwrap_or_else(|| json!({ "error": "no canned response" })) + }; + let body_s = serde_json::to_string(&response_body).unwrap(); let resp = format!( "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", body_s.len(), @@ -622,6 +630,7 @@ async fn run_captured_prompt( .filter(|r| { !r.path.starts_with("/api/2.0/serving-endpoints") && !r.path.starts_with("/api/ai-gateway/v2/endpoints") + && !r.path.starts_with("/api/2.1/unity-catalog/model-services") }) .collect(); assert_eq!(llm_reqs.len(), 1, "expected exactly one LLM request"); @@ -764,6 +773,37 @@ async fn databricks_v2_other_models_route_through_ai_gateway_mlflow_chat() { ); } +#[tokio::test] +async fn databricks_v2_model_service_fqn_uses_mlflow_chat_and_preserves_full_id() { + let canned = vec![json!({ + "id": "x", + "object": "chat.completion", + "choices": [{ + "index": 0, + "message": { "role": "assistant", "content": "ok" }, + "finish_reason": "stop" + }] + })]; + // Family-looking text in a Unity Catalog namespace is data, not route + // authority. The full raw FQN must reach the MLflow model field. + let model = "catalog.schema.claude-gpt-5"; + let req = run_captured_prompt("databricks_v2", model, canned).await; + + assert_eq!( + req.path.as_str(), + "/ai-gateway/mlflow/v1/chat/completions", + "Unity Catalog model-service FQNs must always use MLflow Chat" + ); + assert_eq!(req.body["model"], model); + assert!( + req.body + .get("messages") + .and_then(|value| value.as_array()) + .is_some(), + "model-service FQN requests must use the Chat Completions envelope" + ); +} + // ---------- session/set_model integration tests ---------- /// Helper: run initialize + session/new + optional set_model + session/prompt on a @@ -849,6 +889,7 @@ async fn session_set_model_switches_databricks_legacy_route() { .filter(|r| { !r.path.starts_with("/api/2.0/serving-endpoints") && !r.path.starts_with("/api/ai-gateway/v2/endpoints") + && !r.path.starts_with("/api/2.1/unity-catalog/model-services") }) .collect(); assert_eq!( @@ -896,6 +937,7 @@ async fn session_set_model_switches_databricks_v2_route() { .filter(|r| { !r.path.starts_with("/api/2.0/serving-endpoints") && !r.path.starts_with("/api/ai-gateway/v2/endpoints") + && !r.path.starts_with("/api/2.1/unity-catalog/model-services") }) .collect(); assert_eq!( @@ -1020,7 +1062,7 @@ async fn model_discovery_surfaces_rejected_static_token_as_auth_failure() { let _ = axum::serve(listener, app).await; }); - let cfg = Config::for_discovery(Provider::DatabricksV2, "rejected".into(), host); + let cfg = Config::for_discovery(Provider::DatabricksV2, "rejected".into(), host, None); let error = discover_databricks_models(&cfg).await.unwrap_err(); assert!( @@ -1031,10 +1073,13 @@ async fn model_discovery_surfaces_rejected_static_token_as_auth_failure() { !error.to_string().contains("rejected bearer"), "auth errors must not propagate provider bodies that may echo credentials: {error}" ); - assert_eq!( - requests.load(Ordering::SeqCst), - 1, - "a static token cannot refresh, so discovery must not issue a duplicate request" + // The independent catalog requests run concurrently; the first auth + // failure can short-circuit the joined result before the peer finishes. + // Assert the contract at the behavior boundary rather than assuming both + // in-flight requests always reach the stub. + assert!( + requests.load(Ordering::SeqCst) >= 1, + "static-token auth failure must issue at least one catalog request" ); } diff --git a/desktop/src-tauri/src/commands/agent_models.rs b/desktop/src-tauri/src/commands/agent_models.rs index cb809b6c04a..05b1abad90d 100644 --- a/desktop/src-tauri/src/commands/agent_models.rs +++ b/desktop/src-tauri/src/commands/agent_models.rs @@ -8,11 +8,11 @@ use super::agent_model_process::run_agent_models_command; use super::managed_agent_definition::apply_model_provider_prompt_update; // 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, }; +#[cfg(test)] +use super::agent_models_env::{env_value, env_value_or_process_if_absent}; use super::agent_update_rollback::{rollback_failed_agent_update, AgentUpdateRollback}; use crate::{ @@ -692,8 +692,8 @@ async fn discover_anthropic_models( mod databricks; #[cfg(test)] use databricks::{ - databricks_sign_in_required_error, databricks_static_token_error, is_databricks_provider, - should_start_interactive_auth, + databricks_models_response, databricks_sign_in_required_error, databricks_static_token_error, + is_databricks_provider, should_start_interactive_auth, }; use databricks::{discover_databricks_models, DatabricksAuthIntent}; diff --git a/desktop/src-tauri/src/commands/agent_models_databricks.rs b/desktop/src-tauri/src/commands/agent_models_databricks.rs index 4b6e512c059..1f66f24c6a3 100644 --- a/desktop/src-tauri/src/commands/agent_models_databricks.rs +++ b/desktop/src-tauri/src/commands/agent_models_databricks.rs @@ -5,7 +5,8 @@ use std::sync::{LazyLock, Mutex, MutexGuard}; use std::time::{Duration, Instant}; use crate::commands::agent_models_env::{ - env_or_process_value, redaction_env_with_value, DiscoveryProvider, + env_or_process_value, env_value_or_process_if_absent, redaction_env_with_value, + DiscoveryProvider, }; use crate::managed_agents::AgentModelInfo; use crate::managed_agents::AgentModelsResponse; @@ -167,10 +168,14 @@ pub(super) async fn discover_databricks_models( None => return Ok(None), }; let api_key = env_or_process_value(env, "DATABRICKS_TOKEN").unwrap_or_default(); + let filter = env_value_or_process_if_absent(env, "DATABRICKS_MODEL_FILTER"); + let parsed_filter = buzz_agent_pkg::config::DatabricksModelFilter::parse(filter.as_deref()) + .map_err(|error| format!("invalid DATABRICKS_MODEL_FILTER: {error}"))?; let config = buzz_agent_pkg::config::Config::for_discovery( databricks_agent_provider(provider_name), api_key.clone(), host.clone(), + parsed_filter.clone(), ); let redaction_env = redaction_env_with_value(env, "DATABRICKS_TOKEN", &api_key); @@ -230,11 +235,30 @@ pub(super) async fn discover_databricks_models( } }; - if entries.is_empty() { + databricks_models_response( + provider_name, + entries, + selected_model, + parsed_filter.as_ref(), + ) + .map(Some) +} + +/// When a catalog query fails, Desktop reports the catalog error to the UI and +/// does not fall through to subprocess discovery, so the filter cannot be +/// bypassed by a second model source. +pub(super) fn databricks_models_response( + provider_name: &str, + entries: Vec, + selected_model: Option, + filter: Option<&buzz_agent_pkg::config::DatabricksModelFilter>, +) -> Result { + let entries_are_empty = entries.is_empty(); + if entries_are_empty && filter.is_none() { return Err("Databricks model discovery returned no models".to_string()); } - Ok(Some(AgentModelsResponse { + Ok(AgentModelsResponse { agent_name: provider_name.trim().to_string(), agent_version: "models-api".to_string(), models: entries @@ -247,8 +271,8 @@ pub(super) async fn discover_databricks_models( .collect(), agent_default_model: None, selected_model, - supports_switching: true, - })) + supports_switching: !entries_are_empty, + }) } fn format_redacted_error( diff --git a/desktop/src-tauri/src/commands/agent_models_env.rs b/desktop/src-tauri/src/commands/agent_models_env.rs index 0a40b6bd8ff..06840c90f71 100644 --- a/desktop/src-tauri/src/commands/agent_models_env.rs +++ b/desktop/src-tauri/src/commands/agent_models_env.rs @@ -25,6 +25,22 @@ pub(super) fn env_or_process_value(env: &BTreeMap, key: &str) -> }) } +/// Read a value from the merged discovery env, preserving an explicit blank +/// override. Only when the merged map has no such key does the inherited +/// process environment provide a fallback. This mirrors the child process, +/// where a merged key overrides the inherited environment even when blank. +pub(super) fn env_value_or_process_if_absent( + env: &BTreeMap, + key: &str, +) -> Option { + match env.get(key) { + Some(value) => Some(value.trim().to_string()), + None => std::env::var(key) + .ok() + .map(|value| value.trim().to_string()), + } +} + /// 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( diff --git a/desktop/src-tauri/src/commands/agent_models_tests.rs b/desktop/src-tauri/src/commands/agent_models_tests.rs index df3849de4a4..8d0ec899c7d 100644 --- a/desktop/src-tauri/src/commands/agent_models_tests.rs +++ b/desktop/src-tauri/src/commands/agent_models_tests.rs @@ -476,11 +476,46 @@ fn model_discovery_ignores_stale_record_for_linked_agent() { // --------------------------------------------------------------------------- // Databricks provider detection -// --------------------------------------------------------------------------- -// + +#[test] +fn merged_filter_value_overrides_inherited_process_value_even_when_blank() { + let env = BTreeMap::from([("DATABRICKS_MODEL_FILTER".to_string(), " ".to_string())]); + assert_eq!( + env_value_or_process_if_absent(&env, "DATABRICKS_MODEL_FILTER"), + Some(String::new()) + ); +} + +#[test] +fn absent_filter_value_uses_process_value_when_available() { + const TEST_FILTER_ENV: &str = "BUZZ_TEST_DATABRICKS_MODEL_FILTER"; + let original = std::env::var(TEST_FILTER_ENV).ok(); + std::env::set_var(TEST_FILTER_ENV, "process-*"); + let value = env_value_or_process_if_absent(&BTreeMap::new(), TEST_FILTER_ENV); + match original { + Some(value) => std::env::set_var(TEST_FILTER_ENV, value), + None => std::env::remove_var(TEST_FILTER_ENV), + } + assert_eq!(value.as_deref(), Some("process-*")); +} + +#[test] +fn databricks_filtered_empty_response_is_authoritative() { + let filter = buzz_agent_pkg::config::DatabricksModelFilter::parse(Some("allowed-*")).unwrap(); + let response = databricks_models_response( + "databricks_v2", + Vec::new(), + Some("configured".into()), + filter.as_ref(), + ) + .expect("active filter permits an empty authoritative catalog"); + assert!(response.models.is_empty()); + assert!(!response.supports_switching); + assert_eq!(response.selected_model.as_deref(), Some("configured")); +} + // Parse/filter/pagination tests live in crates/buzz-agent/src/catalog.rs // (they moved there with the Option C refactor). - // --------------------------------------------------------------------------- // Dead-knob guards: mcp_command and turn_timeout_seconds // --------------------------------------------------------------------------- diff --git a/desktop/src-tauri/src/managed_agents/env_vars.rs b/desktop/src-tauri/src/managed_agents/env_vars.rs index de6ec28c41a..6b12fbcd2be 100644 --- a/desktop/src-tauri/src/managed_agents/env_vars.rs +++ b/desktop/src-tauri/src/managed_agents/env_vars.rs @@ -180,7 +180,7 @@ pub fn validate_user_env_keys(env_vars: &BTreeMap) -> Result<(), /// - `BUZZ_AGENT_PROVIDER`, `BUZZ_AGENT_MODEL` — agent runtime selection /// - `BUZZ_AGENT_THINKING_EFFORT` — non-secret enum (none/minimal/low/medium/high/xhigh/max) /// - `BUZZ_AGENT_THINKING_SUMMARY` — non-secret enum (auto/concise/detailed) -/// - `DATABRICKS_HOST`, `DATABRICKS_MODEL` — Block non-secret defaults +/// - `DATABRICKS_HOST`, `DATABRICKS_MODEL`, `DATABRICKS_MODEL_FILTER` — Block non-secret defaults pub(crate) fn is_safe_to_reveal(key: &str) -> bool { const SAFE_KEYS: &[&str] = &[ "BUZZ_AGENT_PROVIDER", @@ -189,6 +189,7 @@ pub(crate) fn is_safe_to_reveal(key: &str) -> bool { "BUZZ_AGENT_THINKING_SUMMARY", "DATABRICKS_HOST", "DATABRICKS_MODEL", + "DATABRICKS_MODEL_FILTER", ]; let upper = key.to_ascii_uppercase(); SAFE_KEYS.iter().any(|safe| upper == *safe) diff --git a/desktop/src-tauri/src/managed_agents/spawn_snapshot/diff/tests.rs b/desktop/src-tauri/src/managed_agents/spawn_snapshot/diff/tests.rs index e21dc4735c7..79bf4f77e79 100644 --- a/desktop/src-tauri/src/managed_agents/spawn_snapshot/diff/tests.rs +++ b/desktop/src-tauri/src/managed_agents/spawn_snapshot/diff/tests.rs @@ -236,6 +236,27 @@ fn allowlisted_env_key_is_case_insensitive() { ); } +#[test] +fn allowlisted_databricks_filter_shows_plain_value() { + let mut before = base(); + before + .env + .insert("DATABRICKS_MODEL_FILTER".into(), "old-*".into()); + let mut after = before.clone(); + after + .env + .insert("DATABRICKS_MODEL_FILTER".into(), "new-*".into()); + + assert_eq!( + change_at(&diff(&before, &after), "env.DATABRICKS_MODEL_FILTER"), + &RestartChange::Value { + before: Value::String("old-*".into()), + after: Value::String("new-*".into()), + }, + "the discovery filter is non-secret and should be reviewable" + ); +} + #[test] fn non_allowlisted_env_key_stays_masked() { // A key not in the allowlist must remain masked regardless of its name. diff --git a/desktop/src/features/agents/AGENTS.md b/desktop/src/features/agents/AGENTS.md index 7822211541b..d9df8c164db 100644 --- a/desktop/src/features/agents/AGENTS.md +++ b/desktop/src/features/agents/AGENTS.md @@ -250,6 +250,8 @@ with a TypeScript lookup table or an id comparison in a component. refresh only local persona/team/managed-agent caches; they must never invalidate the remote relay directory. +15. **Databricks model discovery has one shared catalog authority.** Desktop and ACP call the shared `buzz-agent` discovery library; Desktop passes the effective merged `DATABRICKS_MODEL_FILTER` explicitly, and the library applies it to raw workspace endpoint IDs and Unity Catalog model-service FQNs after the additive union. A successful filtered-empty catalog is authoritative: it stays empty, disables switching, and never falls through to configured or known-model fallback. UC FQNs are catalog data and always use the MLflow Chat Completions route, regardless of family-looking text in their components. + ## The tests that enforce this - `lib/agentConfigCore.test.mjs` — field model per harness × scope, clearing diff --git a/desktop/src/shared/api/tauri.ts b/desktop/src/shared/api/tauri.ts index 6c5b59c9837..65d5f7bb7aa 100644 --- a/desktop/src/shared/api/tauri.ts +++ b/desktop/src/shared/api/tauri.ts @@ -988,7 +988,7 @@ export async function getBakedBuildEnvKeys(): Promise { * * The value is already masked in Rust for secret keys (keys not in the * explicit safe-to-reveal allowlist: `BUZZ_AGENT_PROVIDER`, `BUZZ_AGENT_MODEL`, - * `DATABRICKS_HOST`, `DATABRICKS_MODEL`). Non-allowlisted keys have their + * `DATABRICKS_HOST`, `DATABRICKS_MODEL`, `DATABRICKS_MODEL_FILTER`). Non-allowlisted keys have their * values replaced with `••••••`. Non-secret values are shown as-is. * Empty-value keys are filtered out. */ From 2cbe359747fa16a7e02268f971bc4871ee80755b Mon Sep 17 00:00:00 2001 From: Kalvin Chau Date: Wed, 26 Aug 2026 17:56:26 -0700 Subject: [PATCH 2/4] fix(agent): bound databricks catalog discovery retries Enclose each catalog attempt in one deadline covering response headers, bounded body consumption, and JSON parsing. Retry transient statuses, body failures, malformed JSON, transport failures, and attempt timeouts without logging provider bodies or bearer-derived data. Add deterministic stalled-body, transient-recovery, exhaustion, malformed-response, and redaction coverage. Signed-off-by: Kalvin Chau Signed-off-by: am <6e30cd56c30e030cd31bb0939b94a7c257c9a09d5ba2d92cf2735da45629f248@buzz.block.builderlab.xyz> --- crates/buzz-agent/src/catalog.rs | 769 +++++++++++++++++--- crates/buzz-agent/tests/databricks_oauth.rs | 2 +- 2 files changed, 657 insertions(+), 114 deletions(-) diff --git a/crates/buzz-agent/src/catalog.rs b/crates/buzz-agent/src/catalog.rs index 9a8a0e36c0f..70ff3162290 100644 --- a/crates/buzz-agent/src/catalog.rs +++ b/crates/buzz-agent/src/catalog.rs @@ -12,7 +12,7 @@ //! This helper never opens a browser. Callers choose whether to reject, degrade, //! or start a separate interactive authentication flow. -use std::{collections::HashSet, sync::Arc}; +use std::{collections::HashSet, sync::Arc, time::Duration}; use reqwest::Client; use serde_json::Value; @@ -38,10 +38,48 @@ pub struct ModelEntry { const AUTHENTICATED_EMPTY_CATALOG_SUFFIX: &str = " (default catalog)"; const MAX_CATALOG_PAGES: usize = 20; const MAX_CATALOG_ERROR_BODY_BYTES: usize = 4 * 1024; +const MAX_CATALOG_RESPONSE_BODY_BYTES: usize = 2 * 1024 * 1024; +const CATALOG_REQUEST_TIMEOUT: Duration = Duration::from_secs(10); +const CATALOG_MAX_RETRIES: usize = 3; +const CATALOG_RETRY_BACKOFF: Duration = Duration::from_millis(100); + +#[derive(Clone, Copy)] +struct CatalogRequestPolicy { + timeout: Duration, + max_retries: usize, + retry_backoff: Duration, +} + +const DEFAULT_CATALOG_REQUEST_POLICY: CatalogRequestPolicy = CatalogRequestPolicy { + timeout: CATALOG_REQUEST_TIMEOUT, + max_retries: CATALOG_MAX_RETRIES, + retry_backoff: CATALOG_RETRY_BACKOFF, +}; const WORKSPACE_CATALOG_QUERY: &str = "?page_size=100"; const UNITY_CATALOG_QUERY: &str = "?page_size=100&view=FULL"; type CatalogPage = Result<(Vec, Option), AgentError>; +#[derive(Clone, Copy)] +struct CatalogDescriptor { + name: &'static str, + path: &'static str, + initial_query: &'static str, + parse_page: fn(&Value) -> CatalogPage, +} + +const WORKSPACE_CATALOG_DESCRIPTOR: CatalogDescriptor = CatalogDescriptor { + name: "Databricks workspace endpoint catalog", + path: "/api/ai-gateway/v2/endpoints", + initial_query: WORKSPACE_CATALOG_QUERY, + parse_page: parse_v2_endpoints_page, +}; +const UNITY_CATALOG_DESCRIPTOR: CatalogDescriptor = CatalogDescriptor { + name: "Databricks Unity Catalog model-service catalog", + path: "/api/2.1/unity-catalog/model-services", + initial_query: UNITY_CATALOG_QUERY, + parse_page: parse_uc_model_services_page, +}; + /// Curated display label for a discovered Databricks endpoint or model-service /// id. Unknown ids deliberately pass through unchanged. fn curated_model_name(id: &str) -> String { @@ -153,33 +191,14 @@ async fn fetch_v1_models( bearer: &str, ) -> Result, 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 serving-endpoints catalog request failed: {e}" - )) - })?; - - let status = response.status(); - if !status.is_success() { - return Err(catalog_http_error( - "Databricks serving-endpoints catalog", - status, - response, - bearer, - ) - .await); - } - - let json: Value = response.json().await.map_err(|e| { - AgentError::Llm(format!( - "Databricks serving-endpoints catalog response parse failed: {e}" - )) - })?; + let json = fetch_catalog_page( + http, + &url, + "Databricks serving-endpoints catalog", + bearer, + DEFAULT_CATALOG_REQUEST_POLICY, + ) + .await?; parse_v1_endpoints(&json) } @@ -264,24 +283,29 @@ async fn fetch_v2_models( filter: Option<&DatabricksModelFilter>, allow_partial_auth_failure: bool, ) -> Result, AgentError> { - let workspace = fetch_catalog_pages( - http, - host, - bearer, - "Databricks workspace endpoint catalog", - "/api/ai-gateway/v2/endpoints", - WORKSPACE_CATALOG_QUERY, - parse_v2_endpoints_page, - ); - let unity_catalog = fetch_catalog_pages( + fetch_v2_models_with_policy( http, host, bearer, - "Databricks Unity Catalog model-service catalog", - "/api/2.1/unity-catalog/model-services", - UNITY_CATALOG_QUERY, - parse_uc_model_services_page, - ); + filter, + allow_partial_auth_failure, + DEFAULT_CATALOG_REQUEST_POLICY, + ) + .await +} + +async fn fetch_v2_models_with_policy( + http: &Client, + host: &str, + bearer: &str, + filter: Option<&DatabricksModelFilter>, + allow_partial_auth_failure: bool, + policy: CatalogRequestPolicy, +) -> Result, AgentError> { + let workspace = + fetch_catalog_pages_with_policy(http, host, bearer, WORKSPACE_CATALOG_DESCRIPTOR, policy); + let unity_catalog = + fetch_catalog_pages_with_policy(http, host, bearer, UNITY_CATALOG_DESCRIPTOR, policy); let (workspace, unity_catalog) = tokio::join!(workspace, unity_catalog); let (workspace, unity_catalog, both_succeeded) = match (workspace, unity_catalog) { @@ -292,7 +316,7 @@ async fn fetch_v2_models( } tracing::warn!( catalog = "unity-catalog model-services", - error = %error, + error_kind = catalog_error_kind(&error), "Databricks model discovery degraded: catalog unavailable" ); (workspace, Vec::new(), false) @@ -303,7 +327,7 @@ async fn fetch_v2_models( } tracing::warn!( catalog = "workspace ai-gateway v2 endpoints", - error = %error, + error_kind = catalog_error_kind(&error), "Databricks model discovery degraded: catalog unavailable" ); (Vec::new(), unity_catalog, false) @@ -321,6 +345,19 @@ async fn fetch_v2_models( )) } +fn catalog_error_kind(error: &AgentError) -> &'static str { + match error { + AgentError::InvalidParams(_) => "invalid-params", + AgentError::Llm(_) => "llm", + AgentError::LlmAuth(_) => "auth", + AgentError::LlmModelNotFound(_) => "model-not-found", + AgentError::LlmContextExceeded(_) => "context-exceeded", + AgentError::UnsupportedImageInput(_) => "unsupported-image", + AgentError::Mcp(_) => "mcp", + AgentError::Cancelled => "cancelled", + } +} + fn combined_catalog_error(workspace: AgentError, unity_catalog: AgentError) -> AgentError { let auth_failure = matches!(&workspace, AgentError::LlmAuth(_)) || matches!(&unity_catalog, AgentError::LlmAuth(_)); @@ -368,15 +405,19 @@ fn merge_v2_models( apply_model_filter(merged, filter) } -async fn fetch_catalog_pages( +async fn fetch_catalog_pages_with_policy( http: &Client, host: &str, bearer: &str, - catalog: &'static str, - path: &'static str, - initial_query: &str, - parse_page: fn(&Value) -> CatalogPage, + descriptor: CatalogDescriptor, + policy: CatalogRequestPolicy, ) -> Result, AgentError> { + let CatalogDescriptor { + name: catalog, + path, + initial_query, + parse_page, + } = descriptor; let base_url = format!("{host}{path}"); let mut all_items = Vec::new(); let mut page_token: Option = None; @@ -390,22 +431,7 @@ async fn fetch_catalog_pages( ), None => format!("{base_url}{initial_query}"), }; - let response = http - .get(&url) - .bearer_auth(bearer) - .send() - .await - .map_err(|e| AgentError::Llm(format!("{catalog} request failed: {e}")))?; - - let status = response.status(); - if !status.is_success() { - return Err(catalog_http_error(catalog, status, response, bearer).await); - } - - let json: Value = response - .json() - .await - .map_err(|e| AgentError::Llm(format!("{catalog} response parse failed: {e}")))?; + let json = fetch_catalog_page(http, &url, catalog, bearer, policy).await?; let (items, next_token) = parse_page(&json) .map_err(|error| catalog_context_error(catalog, error, "response parse failed"))?; all_items.extend(items); @@ -428,31 +454,223 @@ async fn fetch_catalog_pages( ))) } -async fn catalog_http_error( +struct ReadResponseBody { + bytes: Vec, + truncated: bool, +} + +enum CatalogRequestError { + Auth, + Status { + status: reqwest::StatusCode, + body: String, + }, + Transport(reqwest::Error), + Body(reqwest::Error), + InvalidJson(serde_json::Error), + BodyTooLarge, +} + +async fn fetch_catalog_page( + http: &Client, + url: &str, + catalog: &str, + bearer: &str, + policy: CatalogRequestPolicy, +) -> Result { + let max_retries = policy.max_retries.max(1); + let error_body_limit = if bearer.len() > MAX_CATALOG_ERROR_BODY_BYTES { + 0 + } else { + MAX_CATALOG_ERROR_BODY_BYTES.saturating_add(bearer.len()) + }; + + for attempt in 0..max_retries { + let result = tokio::time::timeout(policy.timeout, async { + let response = http + .get(url) + .bearer_auth(bearer) + .send() + .await + .map_err(CatalogRequestError::Transport)?; + let status = response.status(); + if status == reqwest::StatusCode::UNAUTHORIZED { + // Preserve the auth contract: do not consume an auth-failure + // body because gateways may echo credential material. The + // bounded attempt ends at headers for this intentionally + // redacted branch; all other status/body paths below consume + // their response body inside the same deadline. + return Err(CatalogRequestError::Auth); + } + if !status.is_success() { + let mut response = response; + let body = read_catalog_error_body(&mut response, error_body_limit) + .await + .map_err(CatalogRequestError::Body)?; + return Err(CatalogRequestError::Status { status, body }); + } + + let mut response = response; + if response + .content_length() + .is_some_and(|length| length > MAX_CATALOG_RESPONSE_BODY_BYTES as u64) + { + return Err(CatalogRequestError::BodyTooLarge); + } + let body = read_response_body(&mut response, MAX_CATALOG_RESPONSE_BODY_BYTES) + .await + .map_err(CatalogRequestError::Body)?; + if body.truncated { + return Err(CatalogRequestError::BodyTooLarge); + } + serde_json::from_slice(&body.bytes).map_err(CatalogRequestError::InvalidJson) + }) + .await; + + match result { + Ok(Ok(json)) => return Ok(json), + Ok(Err(CatalogRequestError::Auth)) => { + return Err(AgentError::LlmAuth(format!("{catalog} HTTP 401"))); + } + Ok(Err(CatalogRequestError::Status { status, body })) => { + if (status.as_u16() == 499 || status.is_server_error()) + && retry_catalog_attempt( + catalog, + attempt, + max_retries, + policy.retry_backoff, + Some(status.as_u16()), + "transient status", + ) + .await + { + continue; + } + return Err(catalog_http_error_body(catalog, status, &body, bearer)); + } + Ok(Err(CatalogRequestError::Transport(error))) => { + if (error.is_timeout() || error.is_connect() || error.is_request()) + && retry_catalog_attempt( + catalog, + attempt, + max_retries, + policy.retry_backoff, + None, + "transport error", + ) + .await + { + continue; + } + return Err(AgentError::Llm(format!( + "{catalog} request failed: {error}" + ))); + } + Ok(Err(CatalogRequestError::Body(error))) => { + if retry_catalog_attempt( + catalog, + attempt, + max_retries, + policy.retry_backoff, + None, + "response body error", + ) + .await + { + continue; + } + return Err(AgentError::Llm(format!( + "{catalog} response body read failed: {error}" + ))); + } + Ok(Err(CatalogRequestError::InvalidJson(error))) => { + if retry_catalog_attempt( + catalog, + attempt, + max_retries, + policy.retry_backoff, + None, + "invalid JSON response", + ) + .await + { + continue; + } + return Err(AgentError::Llm(format!( + "{catalog} response parse failed: {error}" + ))); + } + Ok(Err(CatalogRequestError::BodyTooLarge)) => { + return Err(AgentError::Llm(format!( + "{catalog} response exceeded {MAX_CATALOG_RESPONSE_BODY_BYTES} bytes" + ))); + } + Err(_) => { + if retry_catalog_attempt( + catalog, + attempt, + max_retries, + policy.retry_backoff, + None, + "attempt timeout", + ) + .await + { + continue; + } + return Err(AgentError::Llm(format!( + "{catalog} request timed out after {:?}", + policy.timeout + ))); + } + } + } + + Err(AgentError::Llm(format!( + "{catalog} request failed after {max_retries} attempts" + ))) +} + +async fn retry_catalog_attempt( + catalog: &str, + attempt: usize, + max_attempts: usize, + backoff: Duration, + status: Option, + reason: &'static str, +) -> bool { + if attempt + 1 >= max_attempts { + return false; + } + + tracing::warn!( + catalog, + attempt = attempt + 1, + max_attempts, + status = ?status, + reason, + "Databricks model discovery catalog request retrying" + ); + tokio::time::sleep(backoff).await; + true +} + +fn catalog_http_error_body( catalog: &str, status: reqwest::StatusCode, - response: reqwest::Response, + body: &str, bearer: &str, ) -> AgentError { - if status.as_u16() == 401 { - // Do not read or include a provider body for auth failures. Some - // gateways echo authorization material in diagnostic payloads. + if status == reqwest::StatusCode::UNAUTHORIZED { return AgentError::LlmAuth(format!("{catalog} HTTP {status}")); } let body = if bearer.len() > MAX_CATALOG_ERROR_BODY_BYTES { - // A token longer than the diagnostic bound cannot be safely searched in - // a bounded prefix. Do not return a partial provider body that might - // expose any part of it. String::new() + } else if bearer.is_empty() { + body.to_string() } else { - let read_limit = MAX_CATALOG_ERROR_BODY_BYTES.saturating_add(bearer.len()); - let body = read_catalog_error_body(response, read_limit).await; - if bearer.is_empty() { - body - } else { - body.replace(bearer, "[redacted]") - } + body.replace(bearer, "[redacted]") }; let body = truncate_utf8_bytes(&body, MAX_CATALOG_ERROR_BODY_BYTES); let classification = if status.as_u16() == 499 || status.is_server_error() { @@ -463,21 +681,51 @@ async fn catalog_http_error( AgentError::Llm(format!("{catalog} {classification} HTTP {status}: {body}")) } -async fn read_catalog_error_body(mut response: reqwest::Response, limit: usize) -> String { - let mut body = Vec::with_capacity(limit); - while body.len() < limit { - match response.chunk().await { - Ok(Some(chunk)) => { - let take = chunk.len().min(limit - body.len()); - body.extend_from_slice(&chunk[..take]); - if take < chunk.len() { - break; - } - } - _ => break, +async fn read_response_body( + response: &mut reqwest::Response, + limit: usize, +) -> Result { + let mut bytes = Vec::with_capacity(limit.min(16 * 1024)); + if limit == 0 { + return Ok(ReadResponseBody { + bytes, + truncated: true, + }); + } + + loop { + if bytes.len() == limit { + // Probe one frame past the bound. Without this read, a chunked body + // whose first chunk lands exactly on `limit` would be accepted + // without noticing the next frame. + let truncated = response.chunk().await?.is_some(); + return Ok(ReadResponseBody { bytes, truncated }); + } + + let Some(chunk) = response.chunk().await? else { + return Ok(ReadResponseBody { + bytes, + truncated: false, + }); + }; + let remaining = limit - bytes.len(); + if chunk.len() > remaining { + bytes.extend_from_slice(&chunk[..remaining]); + return Ok(ReadResponseBody { + bytes, + truncated: true, + }); } + bytes.extend_from_slice(&chunk); } - String::from_utf8_lossy(&body).into_owned() +} + +async fn read_catalog_error_body( + response: &mut reqwest::Response, + limit: usize, +) -> Result { + let body = read_response_body(response, limit).await?; + Ok(String::from_utf8_lossy(&body.bytes).into_owned()) } fn truncate_utf8_bytes(value: &str, max_bytes: usize) -> String { @@ -617,6 +865,21 @@ mod tests { use std::collections::HashMap; use std::sync::atomic::{AtomicUsize, Ordering}; + const TEST_CATALOG_DESCRIPTOR: CatalogDescriptor = CatalogDescriptor { + name: "test catalog", + path: "/catalog", + initial_query: "?page_size=100", + parse_page: parse_v2_endpoints_page, + }; + + fn test_policy(timeout: Duration, max_retries: usize) -> CatalogRequestPolicy { + CatalogRequestPolicy { + timeout, + max_retries, + retry_backoff: Duration::ZERO, + } + } + struct RefreshingTestTokenSource { refreshes: AtomicUsize, } @@ -843,14 +1106,12 @@ mod tests { let _ = axum::serve(listener, app).await; }); - let entries = fetch_catalog_pages( + let entries = fetch_catalog_pages_with_policy( &Client::new(), &host, "token", - "test catalog", - "/catalog", - "?page_size=100", - parse_v2_endpoints_page, + TEST_CATALOG_DESCRIPTOR, + DEFAULT_CATALOG_REQUEST_POLICY, ) .await .unwrap(); @@ -870,14 +1131,12 @@ mod tests { tokio::spawn(async move { let _ = axum::serve(listener, app).await; }); - let error = fetch_catalog_pages( + let error = fetch_catalog_pages_with_policy( &Client::new(), &host, "token", - "test catalog", - "/catalog", - "?page_size=100", - parse_v2_endpoints_page, + TEST_CATALOG_DESCRIPTOR, + DEFAULT_CATALOG_REQUEST_POLICY, ) .await .unwrap_err(); @@ -906,14 +1165,12 @@ mod tests { let _ = axum::serve(listener, app).await; }); - let error = fetch_catalog_pages( + let error = fetch_catalog_pages_with_policy( &Client::new(), &host, "token", - "test catalog", - "/catalog", - "?page_size=100", - parse_v2_endpoints_page, + TEST_CATALOG_DESCRIPTOR, + DEFAULT_CATALOG_REQUEST_POLICY, ) .await .unwrap_err(); @@ -923,6 +1180,269 @@ mod tests { assert_eq!(requests.load(Ordering::SeqCst), 20); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn v2_discovery_degrades_a_stalled_secondary_catalog() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let host = format!("http://{}", listener.local_addr().unwrap()); + let app = Router::new() + .route( + "/api/ai-gateway/v2/endpoints", + get(|| async { + Json(serde_json::json!({ + "endpoints": [{"name": "workspace-only"}], + "next_page_token": null, + })) + }), + ) + .route( + "/api/2.1/unity-catalog/model-services", + get(|| async { + // The handler never sends headers. The catalog attempt + // deadline must still let the workspace result win. + tokio::time::sleep(Duration::from_secs(60)).await; + Json(serde_json::json!({ + "model_services": [], + "next_page_token": null, + })) + }), + ); + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + + let started = std::time::Instant::now(); + let models = fetch_v2_models_with_policy( + &Client::new(), + &host, + "token", + None, + false, + test_policy(Duration::from_millis(40), 1), + ) + .await + .unwrap(); + + assert!( + started.elapsed() < Duration::from_secs(1), + "stalled catalog exceeded its request deadline: {:?}", + started.elapsed() + ); + assert_eq!( + models + .iter() + .map(|model| model.id.as_str()) + .collect::>(), + vec!["workspace-only"] + ); + } + + #[tokio::test] + async fn catalog_retries_499_and_5xx_then_recovers() { + for status in [ + StatusCode::from_u16(499).unwrap(), + StatusCode::SERVICE_UNAVAILABLE, + ] { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let host = format!("http://{}", listener.local_addr().unwrap()); + let requests = Arc::new(AtomicUsize::new(0)); + let requests_for_route = requests.clone(); + let app = Router::new().route( + "/catalog", + get(move || { + let attempt = requests_for_route.fetch_add(1, Ordering::SeqCst); + async move { + if attempt == 0 { + Err((status, "provider body secret-token")) + } else { + Ok(Json(serde_json::json!({ + "endpoints": [{"name": "recovered"}], + "next_page_token": null, + }))) + } + } + }), + ); + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + + let entries = fetch_catalog_pages_with_policy( + &Client::new(), + &host, + "secret-token", + TEST_CATALOG_DESCRIPTOR, + test_policy(Duration::from_secs(1), 3), + ) + .await + .unwrap(); + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].entry.id, "recovered"); + assert_eq!(requests.load(Ordering::SeqCst), 2); + } + } + + #[tokio::test] + async fn catalog_retries_malformed_json_then_recovers() { + use axum::response::IntoResponse; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let host = format!("http://{}", listener.local_addr().unwrap()); + let requests = Arc::new(AtomicUsize::new(0)); + let requests_for_route = requests.clone(); + let app = Router::new().route( + "/catalog", + get(move || { + let attempt = requests_for_route.fetch_add(1, Ordering::SeqCst); + async move { + if attempt == 0 { + (StatusCode::OK, "not-json").into_response() + } else { + Json(serde_json::json!({ + "endpoints": [{"name": "json-recovered"}], + "next_page_token": null, + })) + .into_response() + } + } + }), + ); + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + + let entries = fetch_catalog_pages_with_policy( + &Client::new(), + &host, + "token", + TEST_CATALOG_DESCRIPTOR, + test_policy(Duration::from_secs(1), 3), + ) + .await + .unwrap(); + + assert_eq!(requests.load(Ordering::SeqCst), 2); + assert_eq!(entries[0].entry.id, "json-recovered"); + } + + #[tokio::test] + async fn catalog_transient_failure_exhausts_exactly_three_attempts_without_bearer_leak() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let host = format!("http://{}", listener.local_addr().unwrap()); + let requests = Arc::new(AtomicUsize::new(0)); + let requests_for_route = requests.clone(); + let app = Router::new().route( + "/catalog", + get(move || { + requests_for_route.fetch_add(1, Ordering::SeqCst); + async { + ( + StatusCode::SERVICE_UNAVAILABLE, + "provider body secret-token", + ) + } + }), + ); + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + + let error = fetch_catalog_pages_with_policy( + &Client::new(), + &host, + "secret-token", + TEST_CATALOG_DESCRIPTOR, + test_policy(Duration::from_secs(1), 3), + ) + .await + .unwrap_err(); + + assert_eq!(requests.load(Ordering::SeqCst), 3); + let message = error.to_string(); + assert!( + message.contains("transient HTTP 503"), + "unexpected error: {message}" + ); + assert!( + message.contains("provider body"), + "body context was lost: {message}" + ); + assert!( + !message.contains("secret-token"), + "bearer leaked through catalog error: {message}" + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn catalog_retries_when_headers_arrive_but_response_body_stalls() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let host = format!("http://{}", listener.local_addr().unwrap()); + let requests = Arc::new(AtomicUsize::new(0)); + let headers_sent = Arc::new(AtomicUsize::new(0)); + let requests_for_server = requests.clone(); + let headers_for_server = headers_sent.clone(); + tokio::spawn(async move { + loop { + let Ok((mut socket, _)) = listener.accept().await else { + return; + }; + let attempt = requests_for_server.fetch_add(1, Ordering::SeqCst); + let headers_sent = headers_for_server.clone(); + tokio::spawn(async move { + let mut request = Vec::new(); + let mut chunk = [0u8; 1024]; + while !request.windows(4).any(|window| window == b"\r\n\r\n") { + match socket.read(&mut chunk).await { + Ok(0) | Err(_) => return, + Ok(read) => request.extend_from_slice(&chunk[..read]), + } + } + + if attempt == 0 { + socket + .write_all( + b"HTTP/1.1 200 OK\r\n\ + Content-Type: application/json\r\n\ + Content-Length: 64\r\n\ + Connection: close\r\n\r\n\ + {\"endpoints\": [", + ) + .await + .ok(); + headers_sent.store(1, Ordering::SeqCst); + // Keep the declared body incomplete. The outer attempt + // timeout, not reqwest::send(), must terminate this read. + tokio::time::sleep(Duration::from_secs(60)).await; + } else { + let body = + r#"{"endpoints":[{"name":"body-recovered"}],"next_page_token":null}"#; + let response = format!( + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\ + Content-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ); + socket.write_all(response.as_bytes()).await.ok(); + } + }); + } + }); + + let entries = fetch_catalog_pages_with_policy( + &Client::new(), + &host, + "token", + TEST_CATALOG_DESCRIPTOR, + test_policy(Duration::from_millis(40), 2), + ) + .await + .unwrap(); + + assert_eq!(headers_sent.load(Ordering::SeqCst), 1); + assert_eq!(requests.load(Ordering::SeqCst), 2); + assert_eq!(entries[0].entry.id, "body-recovered"); + } #[test] fn v1_filter_applies_to_raw_ids_after_endpoint_filtering() { let filter = DatabricksModelFilter::parse(Some("allowed-*")).unwrap(); @@ -945,10 +1465,33 @@ mod tests { #[test] fn catalog_error_body_is_bounded_and_redacts_bearer() { - // This pure assertion documents the byte-bound helper used after the - // provider response is read. The network path exercises the same - // redaction before truncation; keeping the helper pure makes the UTF-8 - // boundary behavior explicit. + let bearer = "secret-token"; + let provider_body = format!("prefix {bearer} {}", "x".repeat(8_192)); + let status = reqwest::StatusCode::SERVICE_UNAVAILABLE; + let error = catalog_http_error_body("test catalog", status, &provider_body, bearer); + let message = error.to_string(); + assert!( + message.contains("transient HTTP 503"), + "unexpected error: {message}" + ); + assert!( + message.contains("[redacted]"), + "bearer was not redacted: {message}" + ); + assert!(!message.contains(bearer), "bearer leaked: {message}"); + let prefix = format!("llm: test catalog transient HTTP {status}: "); + assert!( + message.starts_with(&prefix), + "unexpected catalog error prefix: message={message:?}, prefix={prefix:?}" + ); + let diagnostic = &message[prefix.len()..]; + assert!( + diagnostic.len() <= MAX_CATALOG_ERROR_BODY_BYTES, + "error body exceeded diagnostic bound: {}", + diagnostic.len() + ); + + // Keep the UTF-8 boundary behavior explicit as well. let value = format!("{}é", "x".repeat(MAX_CATALOG_ERROR_BODY_BYTES)); let truncated = truncate_utf8_bytes(&value, MAX_CATALOG_ERROR_BODY_BYTES); assert_eq!(truncated.len(), MAX_CATALOG_ERROR_BODY_BYTES); diff --git a/crates/buzz-agent/tests/databricks_oauth.rs b/crates/buzz-agent/tests/databricks_oauth.rs index f7d41acf46b..ac2b9578626 100644 --- a/crates/buzz-agent/tests/databricks_oauth.rs +++ b/crates/buzz-agent/tests/databricks_oauth.rs @@ -1225,7 +1225,7 @@ async fn non_auth_discovery_failure_uses_configured_model_without_caching_fallba .await; assert!(h.recv_for(initialize).await.get("result").is_some()); - for expected_attempts in 1..=2 { + for expected_attempts in [3, 6] { let request = h .send("session/new", json!({ "cwd": "/tmp", "mcpServers": [] })) .await; From 93dd353f7e47488878484c416743ec02024b4861 Mon Sep 17 00:00:00 2001 From: am <6e30cd56c30e030cd31bb0939b94a7c257c9a09d5ba2d92cf2735da45629f248@buzz.block.builderlab.xyz> Date: Thu, 27 Aug 2026 08:44:58 -0700 Subject: [PATCH 3/4] fix(agent): align unity catalog model capabilities Co-authored-by: Kalvin Chau Signed-off-by: Kalvin Chau Signed-off-by: am <6e30cd56c30e030cd31bb0939b94a7c257c9a09d5ba2d92cf2735da45629f248@buzz.block.builderlab.xyz> --- crates/buzz-agent/src/catalog.rs | 29 +++++--- crates/buzz-agent/src/llm.rs | 66 ++++--------------- crates/buzz-agent/src/model_capabilities.rs | 46 ++++++++++++- .../features/agents/ui/modelCapabilities.ts | 21 +++++- .../ui/modelCapabilitiesCorpus.test.mjs | 12 ++++ 5 files changed, 106 insertions(+), 68 deletions(-) diff --git a/crates/buzz-agent/src/catalog.rs b/crates/buzz-agent/src/catalog.rs index 70ff3162290..634255476df 100644 --- a/crates/buzz-agent/src/catalog.rs +++ b/crates/buzz-agent/src/catalog.rs @@ -104,6 +104,21 @@ fn authenticated_empty_v2_catalog() -> Vec { .collect() } +/// Heuristic chat-capability filter for v2 workspace endpoints. +/// +/// The v2 catalog omits task metadata. Known embedding endpoint families cannot +/// answer chat-completions requests, so do not offer them as selectable models. +/// Unknown names remain visible; this filter is intentionally narrow. +pub(crate) fn is_chat_capable_endpoint(name: &str) -> bool { + let lower = name.to_ascii_lowercase(); + if lower.contains("embedding") { + return false; + } + !lower + .split('-') + .any(|segment| matches!(segment, "bge" | "gte")) +} + /// Discover available models for a Databricks provider. /// /// Returns an empty vector when an authenticated catalog is valid but no @@ -795,7 +810,7 @@ pub(crate) fn parse_v2_endpoints_page( .iter() .filter_map(|endpoint| { let name = endpoint.get("name")?.as_str()?.to_string(); - if name.is_empty() { + if name.is_empty() || !is_chat_capable_endpoint(&name) { return None; } Some(V2Endpoint { @@ -833,7 +848,7 @@ pub(crate) fn parse_uc_model_services_page( .filter_map(|service| { let resource_name = service.get("name")?.as_str()?; let fqn = resource_name.strip_prefix("model-services/")?; - if !crate::llm::is_model_service_fqn(fqn) { + if !crate::model_capabilities::is_databricks_model_service_fqn(fqn) { return None; } Some(ModelEntry { @@ -1596,7 +1611,7 @@ mod tests { } #[test] - fn v2_parse_keeps_all_nonempty_endpoint_names_without_keyword_filtering() { + fn v2_parse_drops_embedding_endpoints() { let json = serde_json::json!({ "endpoints": [ {"name": "databricks-bge-large-en"}, @@ -1611,13 +1626,7 @@ mod tests { let ids: Vec<&str> = models.iter().map(|m| m.entry.id.as_str()).collect(); assert_eq!( ids, - vec![ - "databricks-bge-large-en", - "databricks-gte-large-en", - "databricks-qwen3-embedding-0-6b", - "databricks-claude-opus-5", - "databricks-gemini-3-pro-image", - ] + vec!["databricks-claude-opus-5", "databricks-gemini-3-pro-image",] ); } diff --git a/crates/buzz-agent/src/llm.rs b/crates/buzz-agent/src/llm.rs index 35a71717497..55c85bb0c5a 100644 --- a/crates/buzz-agent/src/llm.rs +++ b/crates/buzz-agent/src/llm.rs @@ -170,16 +170,10 @@ impl Llm { ) } DatabricksV2Route::MlflowChatCompletions => { - // UC model-service FQNs are not native provider model names. Build - // against a neutral capability identity so family-name text in a - // catalog/schema/service component cannot change request semantics, - // then restore the actual FQN in the payload. - let capability_model = databricks_v2_capability_model(effective_model); let e = effort - .map(|ef| normalize_effort_for_databricks_v2(ef, capability_model)); - let mut body = - openai_body(cfg, system_prompt, history, tools, capability_model, e); - body["model"] = json!(effective_model); + .map(|ef| normalize_effort_for_databricks_v2(ef, effective_model)); + let body = + openai_body(cfg, system_prompt, history, tools, effective_model, e); (body, parse_openai as OpenAiParse) } }) @@ -330,10 +324,8 @@ impl Llm { parse_anthropic as OpenAiParse, ), DatabricksV2Route::MlflowChatCompletions => { - let capability_model = - databricks_v2_capability_model(effective_model); - let mut body = json!({ - "model": capability_model, + let body = json!({ + "model": effective_model, "stream": false, "max_completion_tokens": max_output_tokens, "messages": [ @@ -341,7 +333,6 @@ impl Llm { { "role": "user", "content": user_prompt }, ], }); - body["model"] = json!(effective_model); (body, parse_openai as OpenAiParse) } }) @@ -976,45 +967,9 @@ fn is_responses_required_error(body: &str) -> bool { /// Resolve the Databricks v2 AI Gateway wire route for `model`. /// -/// Strict three-component Unity Catalog model-service FQNs are checked before -/// the manifest, because their namespace is catalog data rather than a model -/// family hint. Other IDs retain the manifest-owned route. -fn databricks_v2_capability_model(model: &str) -> &str { - if is_model_service_fqn(model) { - "model-service" - } else { - model - } -} - -/// Return whether `model` is exactly three non-empty dot-separated components. -/// This is deliberately structural: model-service names are catalog data, so -/// family words such as `claude` and `gpt-5` have no routing meaning here. -pub(crate) fn is_model_service_fqn(model: &str) -> bool { - let mut components = model.split('.'); - let (Some(catalog), Some(schema), Some(service)) = - (components.next(), components.next(), components.next()) - else { - return false; - }; - [catalog, schema, service].into_iter().all(|component| { - !component.is_empty() - && !component.chars().any(char::is_whitespace) - && !component.contains('/') - }) && components.next().is_none() -} - -/// Resolve the Databricks v2 AI Gateway wire route for `model`. -/// -/// A Unity Catalog model-service FQN always uses MLflow Chat Completions. The -/// namespace is data, not a provider-family hint, so names containing `claude` -/// or `gpt` must not reach native Anthropic/Responses routes. Non-FQN ids keep -/// the existing manifest projection unchanged. +/// The capability resolver owns Unity Catalog FQN classification so the Rust +/// request path and desktop effort picker cannot disagree. fn databricks_v2_route(model: &str) -> DatabricksV2Route { - if is_model_service_fqn(model) { - return DatabricksV2Route::MlflowChatCompletions; - } - use crate::model_capabilities::DatabricksV2Route as Manifest; match crate::model_capabilities::resolve("databricks_v2", model).databricks_v2_wire_route { Manifest::OpenaiResponses => DatabricksV2Route::OpenAiResponses, @@ -3306,7 +3261,10 @@ mod tests { "catalog.schema.claude-gpt-5", "data_tools.goose.kimi-k3", ] { - assert!(is_model_service_fqn(model), "expected FQN shape: {model}"); + assert!( + crate::model_capabilities::is_databricks_model_service_fqn(model), + "expected FQN shape: {model}" + ); assert_eq!( databricks_v2_route(model), DatabricksV2Route::MlflowChatCompletions, @@ -3332,7 +3290,7 @@ mod tests { "catalog.schema service", ] { assert!( - !is_model_service_fqn(model), + !crate::model_capabilities::is_databricks_model_service_fqn(model), "unexpected FQN shape: {model}" ); assert_eq!( diff --git a/crates/buzz-agent/src/model_capabilities.rs b/crates/buzz-agent/src/model_capabilities.rs index b299fa61179..f96838ecd86 100644 --- a/crates/buzz-agent/src/model_capabilities.rs +++ b/crates/buzz-agent/src/model_capabilities.rs @@ -284,14 +284,41 @@ fn prefix_matches(token: &str, s: &str) -> bool { } } +/// Return whether `model` is exactly three non-empty dot-separated components. +/// +/// Databricks Unity Catalog model-service names are catalog data, not model +/// family hints. Both capability interpreters use this shape check before +/// family matching so suffixes such as `kimi-k3` cannot inherit endpoint +/// capabilities accidentally. +pub(crate) fn is_databricks_model_service_fqn(model: &str) -> bool { + let mut components = model.split('.'); + let (Some(catalog), Some(schema), Some(service)) = + (components.next(), components.next(), components.next()) + else { + return false; + }; + [catalog, schema, service].into_iter().all(|component| { + !component.is_empty() + && !component.chars().any(char::is_whitespace) + && !component.contains('/') + }) && components.next().is_none() +} + /// Resolve the capability profile for a `(provider, raw_model_id)` pair. pub fn resolve(provider: &str, raw_model_id: &str) -> CapabilityResult { let m = manifest(); let canon = canonical_provider(provider); let blank = raw_model_id.trim().is_empty(); + // Unity Catalog FQNs are neutral model-service identities. Resolve them + // through the concrete-unknown fallback before any suffix can match a + // provider family rule. Routing and effort normalization then share this + // one answer in Rust and TypeScript. + let model_service_fqn = + canon == "databricks_v2" && is_databricks_model_service_fqn(raw_model_id); + // 1. Provider-qualified exact-record lookup (case-insensitive on the id). - if !blank { + if !blank && !model_service_fqn { for rec in &m.exact_records { if rec.provider == canon && rec.raw_model_id.eq_ignore_ascii_case(raw_model_id) { return CapabilityResult { @@ -307,7 +334,7 @@ pub fn resolve(provider: &str, raw_model_id: &str) -> CapabilityResult { } // 2. Boundary-aware family match: longest token wins, lexicographic tie-break. - if !blank { + if !blank && !model_service_fqn { let model_lower = raw_model_id.to_ascii_lowercase(); let stripped = strip_catalog_prefix(&model_lower, &m.family_tokens); let mut best: Option<(usize, &FamilyRule)> = None; @@ -809,6 +836,21 @@ mod tests { // --- Migrated relational/invariant tests (see 42-test inventory) --- // These assert cross-input properties a single corpus vector cannot express. + #[test] + fn databricks_v2_fqn_uses_neutral_concrete_unknown_capabilities() { + let fqn = resolve("databricks_v2", "data_workflow_tools.goose.goose-kimi-k3"); + let fallback = resolve("databricks_v2", "some-unknown-xyz"); + assert_eq!(fqn.thinking_mode, fallback.thinking_mode); + assert_eq!(fqn.supported_efforts, fallback.supported_efforts); + assert_eq!(fqn.default_effort, fallback.default_effort); + assert_eq!( + fqn.databricks_v2_wire_route, + fallback.databricks_v2_wire_route + ); + assert_eq!(fqn.normalization_policy, fallback.normalization_policy); + assert_eq!(fqn.registry_label, None); + } + #[test] fn test_gpt5_numeric_date_suffix_matches_base_not_version() { // A 4-digit date-like suffix on a non-boundary must fall to the gpt-5 base, diff --git a/desktop/src/features/agents/ui/modelCapabilities.ts b/desktop/src/features/agents/ui/modelCapabilities.ts index bce4af829ac..ddfc9889aeb 100644 --- a/desktop/src/features/agents/ui/modelCapabilities.ts +++ b/desktop/src/features/agents/ui/modelCapabilities.ts @@ -287,6 +287,19 @@ function toResult( }; } +function isDatabricksModelServiceFqn(model: string): boolean { + const components = model.split("."); + return ( + components.length === 3 && + components.every( + (component) => + component.length > 0 && + !/\s/.test(component) && + !component.includes("/"), + ) + ); +} + /** * Resolve the capability profile for a `(provider, rawModelId)` pair. * @@ -300,9 +313,13 @@ export function resolveModelCapabilities( ): CapabilityResult { const canon = canonicalizeProvider(provider); const blank = rawModelId.trim().length === 0; + // Unity Catalog FQNs are neutral model-service identities. Resolve them + // through the concrete-unknown fallback before suffix family matching. + const modelServiceFqn = + canon === "databricks_v2" && isDatabricksModelServiceFqn(rawModelId); // 1. Provider-qualified exact-record lookup (case-insensitive on the id). - if (!blank) { + if (!blank && !modelServiceFqn) { const idLower = rawModelId.toLowerCase(); for (const rec of MANIFEST.exact_records) { if ( @@ -315,7 +332,7 @@ export function resolveModelCapabilities( } // 2. Boundary-aware family match: longest token wins, lexicographic tie-break. - if (!blank) { + if (!blank && !modelServiceFqn) { const modelLower = rawModelId.toLowerCase(); const stripped = stripCatalogPrefix(modelLower, MANIFEST.family_tokens); let best: { len: number; rule: FamilyRule } | null = null; diff --git a/desktop/src/features/agents/ui/modelCapabilitiesCorpus.test.mjs b/desktop/src/features/agents/ui/modelCapabilitiesCorpus.test.mjs index 78c05a4df4b..43c45169500 100644 --- a/desktop/src/features/agents/ui/modelCapabilitiesCorpus.test.mjs +++ b/desktop/src/features/agents/ui/modelCapabilitiesCorpus.test.mjs @@ -63,6 +63,18 @@ test("registry label aliases refuse ambiguous stripped record keys", () => { ); }); +test("Unity Catalog FQNs use neutral concrete-unknown capabilities", () => { + const fqn = resolveModelCapabilities( + "databricks_v2", + "data_workflow_tools.goose.goose-kimi-k3", + ); + const fallback = resolveModelCapabilities( + "databricks_v2", + "some-unknown-xyz", + ); + assert.deepEqual(fqn, fallback); +}); + test("every executable corpus vector resolves to its expected six-axis profile", () => { for (const entry of executable) { const id = entry.id ?? ""; From 33ab5db6961e521944411f0cb1abf1c5cb4a1e93 Mon Sep 17 00:00:00 2001 From: am <6e30cd56c30e030cd31bb0939b94a7c257c9a09d5ba2d92cf2735da45629f248@buzz.block.builderlab.xyz> Date: Thu, 27 Aug 2026 10:28:44 -0700 Subject: [PATCH 4/4] fix(agent): filter non-chat unity catalog services Co-authored-by: Kalvin Chau Signed-off-by: Kalvin Chau Signed-off-by: am <6e30cd56c30e030cd31bb0939b94a7c257c9a09d5ba2d92cf2735da45629f248@buzz.block.builderlab.xyz> --- crates/buzz-agent/src/catalog.rs | 62 +++++++++++++++++++++++++++++--- 1 file changed, 58 insertions(+), 4 deletions(-) diff --git a/crates/buzz-agent/src/catalog.rs b/crates/buzz-agent/src/catalog.rs index 634255476df..82f3b086cd6 100644 --- a/crates/buzz-agent/src/catalog.rs +++ b/crates/buzz-agent/src/catalog.rs @@ -830,9 +830,11 @@ pub(crate) fn parse_v2_endpoints_page( /// Parse one page of a `GET api/2.1/unity-catalog/model-services` response. /// /// Unity Catalog resource names are returned as `model-services/..`. -/// Only the exact resource prefix and a structurally valid three-component FQN -/// are selectable. All other resources are ignored without capability/name -/// heuristics; the positive visibility filter is the only further restriction. +/// Only the exact resource prefix, a structurally valid three-component FQN, +/// and chat-capable service metadata are selectable. Missing or empty capability +/// metadata is retained for compatibility with older Databricks workspaces; a +/// non-empty capability list must advertise the MLflow chat API used for model- +/// service inference. The positive visibility filter is applied later. pub(crate) fn parse_uc_model_services_page( json: &Value, ) -> Result<(Vec, Option), AgentError> { @@ -848,7 +850,9 @@ pub(crate) fn parse_uc_model_services_page( .filter_map(|service| { let resource_name = service.get("name")?.as_str()?; let fqn = resource_name.strip_prefix("model-services/")?; - if !crate::model_capabilities::is_databricks_model_service_fqn(fqn) { + if !crate::model_capabilities::is_databricks_model_service_fqn(fqn) + || !uc_model_service_supports_chat(service) + { return None; } Some(ModelEntry { @@ -861,6 +865,17 @@ pub(crate) fn parse_uc_model_services_page( Ok((models, next_page_token(json))) } +fn uc_model_service_supports_chat(service: &Value) -> bool { + let Some(api_types) = service.get("supported_api_types").and_then(Value::as_array) else { + return true; + }; + + api_types.is_empty() + || api_types + .iter() + .any(|api_type| api_type.as_str() == Some("mlflow/v1/chat/completions")) +} + fn next_page_token(json: &Value) -> Option { json.get("next_page_token") .and_then(Value::as_str) @@ -1657,6 +1672,45 @@ mod tests { assert_eq!(next.as_deref(), Some("next token/1")); } + #[test] + fn uc_parse_filters_known_non_chat_services_and_preserves_unknown_capabilities() { + let json = serde_json::json!({ + "model_services": [ + { + "name": "model-services/system.ai.chat-model", + "supported_api_types": [ + "mlflow/v1/chat/completions", + "mlflow/v1/responses" + ] + }, + { + "name": "model-services/system.ai.embedding-model", + "supported_api_types": ["mlflow/v1/embeddings"] + }, + { + "name": "model-services/system.ai.responses-only-model", + "supported_api_types": ["mlflow/v1/responses"] + }, + { + "name": "model-services/catalog.schema.empty-capabilities", + "supported_api_types": [] + }, + {"name": "model-services/catalog.schema.absent-capabilities"}, + ] + }); + + let (models, _) = parse_uc_model_services_page(&json).unwrap(); + let ids: Vec<&str> = models.iter().map(|model| model.id.as_str()).collect(); + assert_eq!( + ids, + vec![ + "system.ai.chat-model", + "catalog.schema.empty-capabilities", + "catalog.schema.absent-capabilities", + ] + ); + } + #[test] fn uc_parse_requires_model_services_array() { let err = parse_uc_model_services_page(&serde_json::json!({"data": []})).unwrap_err();