add databricks ai gateway provider (#9274)

Signed-off-by: Bradley Axen <baxen@squareup.com>

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,

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<String>,
-    },
-}
-
-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<Mutex<Option<String>>>,
-}
-
-#[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::<String>("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)]

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<String>,
+    },
+}
+
+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<Mutex<Option<String>>>,
+}
+
+#[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::<String>("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)))
+    }
+}

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<Mutex<Option<String>>>,
+}
+
+impl DatabricksV2Provider {
+    pub async fn cleanup() -> Result<()> {
+        super::oauth::cleanup_oauth_cache()
+    }
+
+    pub async fn from_env(model: ModelConfig) -> Result<Self> {
+        let config = crate::config::Config::global();
+
+        let mut host: Result<String, ConfigError> = 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> {
+        Self::new(
+            host,
+            DatabricksAuth::token(api_key),
+            model,
+            RetryConfig::default(),
+        )
+    }
+
+    fn new(
+        host: String,
+        auth: DatabricksAuth,
+        model: ModelConfig,
+        retry_config: RetryConfig,
+    ) -> Result<Self> {
+        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::<usize>().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::<u64>().ok())
+            .unwrap_or(DEFAULT_INITIAL_RETRY_INTERVAL_MS);
+
+        let backoff_multiplier = config
+            .get_param("DATABRICKS_BACKOFF_MULTIPLIER")
+            .ok()
+            .and_then(|v: String| v.parse::<f64>().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::<u64>().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<Vec<String>, 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<String> = 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<MessageStream, ProviderError> {
+        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<MessageStream, ProviderError> {
+        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<MessageStream, ProviderError> {
+        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<crate::config::ExtensionConfig>,
+    ) -> BoxFuture<'static, Result<Self::Provider>> {
+        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<MessageStream, ProviderError> {
+        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<Vec<String>, 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"));
+    }
+}

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<ProviderRegistry> {
         registry.register::<CodexProvider>(true);
         registry.register::<CursorAgentProvider>(false);
         registry.register::<DatabricksProvider>(true);
+        registry.register::<DatabricksV2Provider>(false);
         registry.register::<GcpVertexAIProvider>(false);
         registry.register::<GeminiCliProvider>(false);
         registry.register::<GeminiOAuthProvider>(true);
@@ -95,6 +97,10 @@ async fn init_registry() -> RwLock<ProviderRegistry> {
         "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())),

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;

ui/desktop/src/utils/configUtils.ts

@@ -79,6 +79,7 @@ export const providerPrefixes: Record<string, string[]> = {
   google: ['GOOGLE_'],
   groq: ['GROQ_'],
   databricks: ['DATABRICKS_'],
+  databricks_v2: ['DATABRICKS_'],
   openrouter: ['OPENROUTER_'],
   ollama: ['OLLAMA_'],
   azure_openai: ['AZURE_'],