diff --git a/crates/goose/src/providers/catalog.rs b/crates/goose/src/providers/catalog.rs index e1dfe6d0..2cfe698d 100644 --- a/crates/goose/src/providers/catalog.rs +++ b/crates/goose/src/providers/catalog.rs @@ -495,6 +495,36 @@ const SETUP_METADATA: &[CuratedSetupMetadata] = &[ }, ], }, + CuratedSetupMetadata { + provider_id: "databricks_v2", + category: ProviderSetupCategory::Model, + setup_method: ProviderSetupMethod::HostWithOauthFallback, + group: ProviderSetupGroup::Additional, + display_name: Some("Databricks AI Gateway"), + description: Some("Models on Databricks AI Gateway v2"), + docs_url: Some("https://docs.databricks.com/en/generative-ai/ai-gateway/"), + aliases: &["databricks_ai_gateway"], + native_connect_query: None, + binary_name: None, + setup_capabilities: setup_capabilities(false, true, false), + show_only_when_installed: false, + synthetic: false, + secret_field_default: None, + field_overrides: &[ + CuratedFieldMetadata { + key: "DATABRICKS_HOST", + label: "Host URL", + placeholder: Some("https://dbc-...cloud.databricks.com"), + default_value: None, + }, + CuratedFieldMetadata { + key: "DATABRICKS_TOKEN", + label: "Access Token", + placeholder: Some("Paste your access token"), + default_value: None, + }, + ], + }, CuratedSetupMetadata { provider_id: "github_copilot", category: ProviderSetupCategory::Model, diff --git a/crates/goose/src/providers/databricks.rs b/crates/goose/src/providers/databricks.rs index 0eff8b19..76b9bc45 100644 --- a/crates/goose/src/providers/databricks.rs +++ b/crates/goose/src/providers/databricks.rs @@ -1,23 +1,22 @@ use anyhow::Result; use async_trait::async_trait; use futures::future::BoxFuture; -use serde::{Deserialize, Serialize}; use serde_json::Value; use std::collections::HashSet; use std::sync::LazyLock; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; -use super::api_client::{ApiClient, AuthMethod, AuthProvider}; +use super::api_client::{ApiClient, AuthMethod}; use super::base::{ ConfigKey, MessageStream, ModelInfo, Provider, ProviderDef, ProviderMetadata, DEFAULT_PROVIDER_TIMEOUT_SECS, }; +use super::databricks_auth::{DatabricksAuth, DatabricksAuthProvider}; use super::embedding::EmbeddingCapable; use super::errors::ProviderError; use super::formats::databricks::create_request; use super::formats::openai_responses::create_responses_request; -use super::oauth; use super::openai_compatible::{ handle_response_openai_compat, handle_status, map_http_error_to_provider_error, sanitize_url, stream_openai_compat, stream_responses_compat, @@ -55,10 +54,6 @@ struct CachedDatabricksEndpointInfo { fetched_at: Instant, } -const DEFAULT_CLIENT_ID: &str = "databricks-cli"; -const DEFAULT_REDIRECT_URL: &str = "http://localhost"; -const DEFAULT_SCOPES: &[&str] = &["all-apis", "offline_access"]; - const DATABRICKS_PROVIDER_NAME: &str = "databricks"; const DATABRICKS_ENDPOINT_METADATA_TTL_SECS: u64 = 60; static DATABRICKS_ENDPOINT_INFO_CACHE: LazyLock< @@ -75,69 +70,6 @@ pub const DATABRICKS_KNOWN_MODELS: &[&str] = &[ pub const DATABRICKS_DOC_URL: &str = "https://docs.databricks.com/en/generative-ai/external-models/index.html"; -#[derive(Debug, Clone, Serialize, Deserialize)] -pub enum DatabricksAuth { - Token(String), - OAuth { - host: String, - client_id: String, - redirect_url: String, - scopes: Vec, - }, -} - -impl DatabricksAuth { - pub fn oauth(host: String) -> Self { - Self::OAuth { - host, - client_id: DEFAULT_CLIENT_ID.to_string(), - redirect_url: DEFAULT_REDIRECT_URL.to_string(), - scopes: DEFAULT_SCOPES.iter().map(|s| s.to_string()).collect(), - } - } - - pub fn token(token: String) -> Self { - Self::Token(token) - } -} - -struct DatabricksAuthProvider { - auth: DatabricksAuth, - token_cache: Arc>>, -} - -#[async_trait] -impl AuthProvider for DatabricksAuthProvider { - async fn get_auth_header(&self) -> Result<(String, String)> { - let token = match &self.auth { - DatabricksAuth::Token(original) => { - let cached = self.token_cache.lock().unwrap().clone(); - match cached { - Some(t) => t, - None => { - // Cache was cleared by refresh_credentials(); re-read - // from config which may have a sidecar-rotated token. - // Fall back to the constructor-provided token if config - // lookup fails (e.g. from_params usage). - let fresh = crate::config::Config::global() - .get_secret::("DATABRICKS_TOKEN") - .unwrap_or_else(|_| original.clone()); - *self.token_cache.lock().unwrap() = Some(fresh.clone()); - fresh - } - } - } - DatabricksAuth::OAuth { - host, - client_id, - redirect_url, - scopes, - } => oauth::get_oauth_token_async(host, client_id, redirect_url, scopes).await?, - }; - Ok(("Authorization".to_string(), format!("Bearer {}", token))) - } -} - #[derive(Debug, serde::Serialize)] pub struct DatabricksProvider { #[serde(skip)] diff --git a/crates/goose/src/providers/databricks_auth.rs b/crates/goose/src/providers/databricks_auth.rs new file mode 100644 index 00000000..e2795e0e --- /dev/null +++ b/crates/goose/src/providers/databricks_auth.rs @@ -0,0 +1,70 @@ +use anyhow::Result; +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use std::sync::{Arc, Mutex}; + +use super::api_client::AuthProvider; +use super::oauth; + +const DEFAULT_CLIENT_ID: &str = "databricks-cli"; +const DEFAULT_REDIRECT_URL: &str = "http://localhost"; +const DEFAULT_SCOPES: &[&str] = &["all-apis", "offline_access"]; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum DatabricksAuth { + Token(String), + OAuth { + host: String, + client_id: String, + redirect_url: String, + scopes: Vec, + }, +} + +impl DatabricksAuth { + pub fn oauth(host: String) -> Self { + Self::OAuth { + host, + client_id: DEFAULT_CLIENT_ID.to_string(), + redirect_url: DEFAULT_REDIRECT_URL.to_string(), + scopes: DEFAULT_SCOPES.iter().map(|s| s.to_string()).collect(), + } + } + + pub fn token(token: String) -> Self { + Self::Token(token) + } +} + +pub(crate) struct DatabricksAuthProvider { + pub auth: DatabricksAuth, + pub token_cache: Arc>>, +} + +#[async_trait] +impl AuthProvider for DatabricksAuthProvider { + async fn get_auth_header(&self) -> Result<(String, String)> { + let token = match &self.auth { + DatabricksAuth::Token(original) => { + let cached = self.token_cache.lock().unwrap().clone(); + match cached { + Some(t) => t, + None => { + let fresh = crate::config::Config::global() + .get_secret::("DATABRICKS_TOKEN") + .unwrap_or_else(|_| original.clone()); + *self.token_cache.lock().unwrap() = Some(fresh.clone()); + fresh + } + } + } + DatabricksAuth::OAuth { + host, + client_id, + redirect_url, + scopes, + } => oauth::get_oauth_token_async(host, client_id, redirect_url, scopes).await?, + }; + Ok(("Authorization".to_string(), format!("Bearer {}", token))) + } +} diff --git a/crates/goose/src/providers/databricks_v2.rs b/crates/goose/src/providers/databricks_v2.rs new file mode 100644 index 00000000..2385ae9b --- /dev/null +++ b/crates/goose/src/providers/databricks_v2.rs @@ -0,0 +1,499 @@ +use anyhow::Result; +use async_stream::try_stream; +use async_trait::async_trait; +use futures::future::BoxFuture; +use futures::TryStreamExt; +use serde::Serialize; +use serde_json::Value; +use std::io; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use tokio::pin; +use tokio_util::io::StreamReader; + +use super::api_client::{ApiClient, AuthMethod}; +use super::base::{ + ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata, + DEFAULT_PROVIDER_TIMEOUT_SECS, +}; +use super::databricks_auth::{DatabricksAuth, DatabricksAuthProvider}; +use super::errors::ProviderError; +use super::formats::{anthropic, openai, openai_responses}; +use super::openai_compatible::{handle_status, stream_openai_compat, stream_responses_compat}; +use super::retry::ProviderRetry; +use super::utils::{extract_reasoning_effort, is_openai_responses_model, ImageFormat, RequestLog}; +use crate::config::ConfigError; +use crate::conversation::message::Message; +use crate::model::ModelConfig; +use crate::providers::retry::{ + RetryConfig, DEFAULT_BACKOFF_MULTIPLIER, DEFAULT_INITIAL_RETRY_INTERVAL_MS, + DEFAULT_MAX_RETRIES, DEFAULT_MAX_RETRY_INTERVAL_MS, +}; +use rmcp::model::Tool; + +const DATABRICKS_V2_PROVIDER_NAME: &str = "databricks_v2"; +const DATABRICKS_V2_LIST_ENDPOINTS_PATH: &str = "api/ai-gateway/v2/endpoints"; +pub const DATABRICKS_V2_DEFAULT_MODEL: &str = "databricks-gpt-5-5"; +pub const DATABRICKS_V2_KNOWN_MODELS: &[&str] = + &["databricks-gpt-5-5", "databricks-claude-opus-4-7"]; + +pub const DATABRICKS_V2_DOC_URL: &str = "https://docs.databricks.com/en/generative-ai/ai-gateway/"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum DatabricksV2Route { + OpenAiResponses, + AnthropicMessages, + MlflowChatCompletions, +} + +#[derive(Debug, Serialize)] +pub struct DatabricksV2Provider { + #[serde(skip)] + api_client: ApiClient, + model: ModelConfig, + #[serde(skip)] + retry_config: RetryConfig, + #[serde(skip)] + name: String, + #[serde(skip)] + token_cache: Arc>>, +} + +impl DatabricksV2Provider { + pub async fn cleanup() -> Result<()> { + super::oauth::cleanup_oauth_cache() + } + + pub async fn from_env(model: ModelConfig) -> Result { + let config = crate::config::Config::global(); + + let mut host: Result = config.get_param("DATABRICKS_HOST"); + if host.is_err() { + host = config.get_secret("DATABRICKS_HOST") + } + + if host.is_err() { + return Err(ConfigError::NotFound( + "Did not find DATABRICKS_HOST in either config file or keyring".to_string(), + ) + .into()); + } + + let host = host?; + let retry_config = Self::load_retry_config(config); + + let auth = if let Ok(api_key) = config.get_secret("DATABRICKS_TOKEN") { + DatabricksAuth::token(api_key) + } else { + DatabricksAuth::oauth(host.clone()) + }; + + Self::new(host, auth, model, retry_config) + } + + pub fn from_params(host: String, api_key: String, model: ModelConfig) -> Result { + Self::new( + host, + DatabricksAuth::token(api_key), + model, + RetryConfig::default(), + ) + } + + fn new( + host: String, + auth: DatabricksAuth, + model: ModelConfig, + retry_config: RetryConfig, + ) -> Result { + let token_cache = Arc::new(Mutex::new(match &auth { + DatabricksAuth::Token(t) => Some(t.clone()), + _ => None, + })); + + let auth_method = AuthMethod::Custom(Box::new(DatabricksAuthProvider { + auth: auth.clone(), + token_cache: token_cache.clone(), + })); + + let api_client = ApiClient::with_timeout( + host, + auth_method, + Duration::from_secs(DEFAULT_PROVIDER_TIMEOUT_SECS), + )?; + + Ok(Self { + api_client, + model, + retry_config, + name: DATABRICKS_V2_PROVIDER_NAME.to_string(), + token_cache, + }) + } + + fn load_retry_config(config: &crate::config::Config) -> RetryConfig { + let max_retries = config + .get_param("DATABRICKS_MAX_RETRIES") + .ok() + .and_then(|v: String| v.parse::().ok()) + .unwrap_or(DEFAULT_MAX_RETRIES); + + let initial_interval_ms = config + .get_param("DATABRICKS_INITIAL_RETRY_INTERVAL_MS") + .ok() + .and_then(|v: String| v.parse::().ok()) + .unwrap_or(DEFAULT_INITIAL_RETRY_INTERVAL_MS); + + let backoff_multiplier = config + .get_param("DATABRICKS_BACKOFF_MULTIPLIER") + .ok() + .and_then(|v: String| v.parse::().ok()) + .unwrap_or(DEFAULT_BACKOFF_MULTIPLIER); + + let max_interval_ms = config + .get_param("DATABRICKS_MAX_RETRY_INTERVAL_MS") + .ok() + .and_then(|v: String| v.parse::().ok()) + .unwrap_or(DEFAULT_MAX_RETRY_INTERVAL_MS); + + RetryConfig::new( + max_retries, + initial_interval_ms, + backoff_multiplier, + max_interval_ms, + ) + } + + fn route_for_model(model_name: &str) -> DatabricksV2Route { + let (clean_name, _) = extract_reasoning_effort(model_name); + let lower = clean_name.to_lowercase(); + + if is_openai_responses_model(&clean_name) || Self::looks_like_gpt5(&lower) { + DatabricksV2Route::OpenAiResponses + } else if Self::is_claude_model(&lower) { + DatabricksV2Route::AnthropicMessages + } else { + DatabricksV2Route::MlflowChatCompletions + } + } + + fn looks_like_gpt5(model_name: &str) -> bool { + model_name.contains("gpt-5") || model_name.contains("gpt5") + } + + fn is_claude_model(model_name: &str) -> bool { + model_name.contains("claude") + } + + fn parse_list_endpoints_response(json: &Value) -> Result, ProviderError> { + let endpoints = json + .get("endpoints") + .and_then(|v| v.as_array()) + .ok_or_else(|| { + ProviderError::RequestFailed( + "Unexpected response format from Databricks AI Gateway endpoints API" + .to_string(), + ) + })?; + + let mut models: Vec = endpoints + .iter() + .filter_map(|endpoint| { + endpoint + .get("name") + .and_then(|v| v.as_str()) + .map(str::to_string) + }) + .collect(); + models.sort(); + Ok(models) + } + + async fn stream_openai_responses( + &self, + model_config: &ModelConfig, + session_id: &str, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result { + let mut payload = + openai_responses::create_responses_request(model_config, system, messages, tools)?; + payload["stream"] = Value::Bool(true); + let mut log = RequestLog::start(model_config, &payload)?; + + let response = self + .with_retry(|| async { + let resp = self + .api_client + .response_post(Some(session_id), "ai-gateway/openai/v1/responses", &payload) + .await?; + handle_status(resp).await + }) + .await + .inspect_err(|e| { + let _ = log.error(e); + })?; + + stream_responses_compat(response, log) + } + + async fn stream_mlflow_chat_completions( + &self, + model_config: &ModelConfig, + session_id: &str, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result { + let mut payload = openai::create_request( + model_config, + system, + messages, + tools, + &ImageFormat::OpenAi, + true, + )?; + if payload.get("max_tokens").is_none() { + payload["max_tokens"] = Value::from(model_config.max_output_tokens()); + } + let mut log = RequestLog::start(model_config, &payload)?; + + let response = self + .with_retry(|| async { + let resp = self + .api_client + .response_post( + Some(session_id), + "ai-gateway/mlflow/v1/chat/completions", + &payload, + ) + .await?; + handle_status(resp).await + }) + .await + .inspect_err(|e| { + let _ = log.error(e); + })?; + + stream_openai_compat(response, log) + } + + async fn stream_anthropic_messages( + &self, + model_config: &ModelConfig, + session_id: &str, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result { + let mut payload = anthropic::create_request(model_config, system, messages, tools)?; + payload["stream"] = Value::Bool(true); + let mut log = RequestLog::start(model_config, &payload)?; + + let response = self + .with_retry(|| async { + let resp = self + .api_client + .response_post( + Some(session_id), + "ai-gateway/anthropic/v1/messages", + &payload, + ) + .await?; + handle_status(resp).await + }) + .await + .inspect_err(|e| { + let _ = log.error(e); + })?; + + let stream = response.bytes_stream().map_err(io::Error::other); + + Ok(Box::pin(try_stream! { + let stream_reader = StreamReader::new(stream); + let framed = tokio_util::codec::FramedRead::new(stream_reader, tokio_util::codec::LinesCodec::new()) + .map_err(anyhow::Error::from); + + let message_stream = anthropic::response_to_streaming_message(framed); + pin!(message_stream); + while let Some(message) = futures::StreamExt::next(&mut message_stream).await { + let (message, usage) = message.map_err(|e| ProviderError::RequestFailed(format!("Stream decode error: {e}")))?; + log.write(&message, usage.as_ref().map(|f| f.usage).as_ref())?; + yield (message, usage); + } + })) + } +} + +impl ProviderDef for DatabricksV2Provider { + type Provider = Self; + + fn metadata() -> ProviderMetadata { + ProviderMetadata::new( + DATABRICKS_V2_PROVIDER_NAME, + "Databricks AI Gateway", + "Models on Databricks AI Gateway v2", + DATABRICKS_V2_DEFAULT_MODEL, + DATABRICKS_V2_KNOWN_MODELS.to_vec(), + DATABRICKS_V2_DOC_URL, + vec![ + ConfigKey::new("DATABRICKS_HOST", true, false, None, true), + ConfigKey::new("DATABRICKS_TOKEN", false, true, None, true), + ], + ) + } + + fn from_env( + model: ModelConfig, + _extensions: Vec, + ) -> BoxFuture<'static, Result> { + Box::pin(Self::from_env(model)) + } + + fn supports_inventory_refresh() -> bool { + true + } +} + +#[async_trait] +impl Provider for DatabricksV2Provider { + fn get_name(&self) -> &str { + &self.name + } + + fn retry_config(&self) -> RetryConfig { + self.retry_config.clone() + } + + async fn refresh_credentials(&self) -> Result<(), ProviderError> { + crate::config::Config::global().invalidate_secrets_cache(); + *self.token_cache.lock().unwrap() = None; + Ok(()) + } + + fn get_model_config(&self) -> ModelConfig { + self.model.clone() + } + + async fn stream( + &self, + model_config: &ModelConfig, + session_id: &str, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result { + match Self::route_for_model(&model_config.model_name) { + DatabricksV2Route::OpenAiResponses => { + self.stream_openai_responses(model_config, session_id, system, messages, tools) + .await + } + DatabricksV2Route::AnthropicMessages => { + self.stream_anthropic_messages(model_config, session_id, system, messages, tools) + .await + } + DatabricksV2Route::MlflowChatCompletions => { + self.stream_mlflow_chat_completions( + model_config, + session_id, + system, + messages, + tools, + ) + .await + } + } + } + + async fn fetch_supported_models(&self) -> Result, ProviderError> { + let response = self + .api_client + .response_get(None, DATABRICKS_V2_LIST_ENDPOINTS_PATH) + .await + .map_err(|e| { + ProviderError::RequestFailed(format!( + "Failed to fetch Databricks AI Gateway endpoints: {e}" + )) + })?; + + if !response.status().is_success() { + let status = response.status(); + let detail = response.text().await.unwrap_or_default(); + return Err(ProviderError::RequestFailed(format!( + "Failed to fetch Databricks AI Gateway endpoints: {status} {detail}" + ))); + } + + let json: Value = response.json().await.map_err(|e| { + ProviderError::RequestFailed(format!( + "Failed to parse Databricks AI Gateway endpoints response: {e}" + )) + })?; + + Self::parse_list_endpoints_response(&json) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn routes_known_model_families() { + for model in ["databricks-gpt-5-5", "databricks-gpt5"] { + assert_eq!( + DatabricksV2Provider::route_for_model(model), + DatabricksV2Route::OpenAiResponses, + "unexpected route for {model}" + ); + } + + for model in ["databricks-claude-opus-4-7", "databricks-claude-sonnet-4-6"] { + assert_eq!( + DatabricksV2Provider::route_for_model(model), + DatabricksV2Route::AnthropicMessages, + "unexpected route for {model}" + ); + } + + assert_eq!( + DatabricksV2Provider::route_for_model("custom-model"), + DatabricksV2Route::MlflowChatCompletions + ); + } + + #[test] + fn parses_list_endpoints_response() { + let json = serde_json::json!({ + "endpoints": [ + {"name": "databricks-claude-opus-4-7"}, + {"name": "databricks-gpt-5-5"}, + {"name": "custom-model"} + ] + }); + + let models = DatabricksV2Provider::parse_list_endpoints_response(&json).unwrap(); + + assert_eq!( + models, + vec![ + "custom-model".to_string(), + "databricks-claude-opus-4-7".to_string(), + "databricks-gpt-5-5".to_string(), + ] + ); + } + + #[test] + fn errors_when_list_endpoints_response_has_no_endpoints_array() { + let json = serde_json::json!({"data": []}); + + let error = DatabricksV2Provider::parse_list_endpoints_response(&json).unwrap_err(); + + assert!(matches!(error, ProviderError::RequestFailed(_))); + assert!(error + .to_string() + .contains("Unexpected response format from Databricks AI Gateway endpoints API")); + } +} diff --git a/crates/goose/src/providers/init.rs b/crates/goose/src/providers/init.rs index 24364399..cc4c20b6 100644 --- a/crates/goose/src/providers/init.rs +++ b/crates/goose/src/providers/init.rs @@ -21,6 +21,7 @@ use super::{ copilot_acp::CopilotAcpProvider, cursor_agent::CursorAgentProvider, databricks::DatabricksProvider, + databricks_v2::DatabricksV2Provider, gcpvertexai::GcpVertexAIProvider, gemini_cli::GeminiCliProvider, gemini_oauth::GeminiOAuthProvider, @@ -68,6 +69,7 @@ async fn init_registry() -> RwLock { registry.register::(true); registry.register::(false); registry.register::(true); + registry.register::(false); registry.register::(false); registry.register::(false); registry.register::(true); @@ -95,6 +97,10 @@ async fn init_registry() -> RwLock { "databricks", Arc::new(|| Box::pin(DatabricksProvider::cleanup())), ); + registry.set_cleanup( + "databricks_v2", + Arc::new(|| Box::pin(DatabricksV2Provider::cleanup())), + ); registry.set_cleanup( "kimi_code", Arc::new(|| Box::pin(KimiCodeProvider::cleanup())), diff --git a/crates/goose/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index 44f7dd0a..f5e4f952 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -19,6 +19,8 @@ pub mod codex_acp; pub mod copilot_acp; pub mod cursor_agent; pub mod databricks; +pub mod databricks_auth; +pub mod databricks_v2; pub mod embedding; pub mod errors; pub mod formats; diff --git a/ui/desktop/src/utils/configUtils.ts b/ui/desktop/src/utils/configUtils.ts index 9e0a5e55..58f866e5 100644 --- a/ui/desktop/src/utils/configUtils.ts +++ b/ui/desktop/src/utils/configUtils.ts @@ -79,6 +79,7 @@ export const providerPrefixes: Record = { google: ['GOOGLE_'], groq: ['GROQ_'], databricks: ['DATABRICKS_'], + databricks_v2: ['DATABRICKS_'], openrouter: ['OPENROUTER_'], ollama: ['OLLAMA_'], azure_openai: ['AZURE_'],