feat: Add databricks oauth (#313)

2026-07-03 14:10:03 +02:00 · 2024-11-24 08:02:53 +11:00
parent 7601cdc7e3
commit 339f5f4b9e
11 changed files with 564 additions and 52 deletions
@@ -1,5 +1,7 @@
 use crate::inputs::inputs::get_env_value_or_input;
-use goose::providers::configs::{DatabricksProviderConfig, OllamaProviderConfig, OpenAiProviderConfig, ProviderConfig};
+use goose::providers::configs::{
+    DatabricksAuth, DatabricksProviderConfig, OpenAiProviderConfig, OllamaProviderConfig, ProviderConfig
+};
 use goose::providers::factory::ProviderType;
 use goose::providers::ollama::OLLAMA_HOST;
 use strum::IntoEnumIterator;
@@ -35,21 +37,21 @@ pub fn set_provider_config(provider_name: &str, model: String) -> ProviderConfig
            temperature: None,
            max_tokens: None,
        }),
-        PROVIDER_DATABRICKS => ProviderConfig::Databricks(DatabricksProviderConfig {
-            host: get_env_value_or_input(
+        PROVIDER_DATABRICKS => {
+            let host = get_env_value_or_input(
                "DATABRICKS_HOST",
                "Please enter your Databricks host:",
                false,
-            ),
-            token: get_env_value_or_input(
-                "DATABRICKS_TOKEN",
-                "Please enter your Databricks token:",
-                true,
-            ),
-            model,
-            temperature: None,
-            max_tokens: None,
-        }),
+            );
+            ProviderConfig::Databricks(DatabricksProviderConfig {
+                host: host.clone(),
+                // TODO revisit configuration
+                auth: DatabricksAuth::oauth(host),
+                model,
+                temperature: None,
+                max_tokens: None,
+            })
+        }
        PROVIDER_OLLAMA => ProviderConfig::Ollama(OllamaProviderConfig {
            host: std::env::var("OLLAMA_HOST")
                .unwrap_or_else(|_| String::from(OLLAMA_HOST)),
@@ -1,7 +1,7 @@
 use crate::error::{to_env_var, ConfigError};
 use config::{Config, Environment};
 use goose::providers::{
-    configs::{DatabricksProviderConfig, OllamaProviderConfig, OpenAiProviderConfig, ProviderConfig},
+    configs::{DatabricksAuth, DatabricksProviderConfig, OllamaProviderConfig, OpenAiProviderConfig, ProviderConfig},
    factory::ProviderType,
    ollama,
 };
@@ -41,7 +41,6 @@ pub enum ProviderSettings {
    Databricks {
        #[serde(default = "default_databricks_host")]
        host: String,
-        token: String,
        #[serde(default = "default_model")]
        model: String,
        #[serde(default)]
@@ -90,13 +89,12 @@ impl ProviderSettings {
            }),
            ProviderSettings::Databricks {
                host,
-                token,
                model,
                temperature,
                max_tokens,
            } => ProviderConfig::Databricks(DatabricksProviderConfig {
-                host,
-                token,
+                host: host.clone(),
+                auth: DatabricksAuth::oauth(host),
                model,
                temperature,
                max_tokens,
@@ -257,7 +255,6 @@ mod tests {
    fn test_databricks_settings() {
        clean_env();
        env::set_var("GOOSE_PROVIDER__TYPE", "databricks");
-        env::set_var("GOOSE_PROVIDER__TOKEN", "test-token");
        env::set_var("GOOSE_PROVIDER__HOST", "https://custom.databricks.com");
        env::set_var("GOOSE_PROVIDER__MODEL", "llama-2-70b");
        env::set_var("GOOSE_PROVIDER__TEMPERATURE", "0.7");
@@ -266,14 +263,12 @@ mod tests {
        let settings = Settings::new().unwrap();
        if let ProviderSettings::Databricks {
            host,
-            token,
            model,
            temperature,
            max_tokens,
        } = settings.provider
        {
            assert_eq!(host, "https://custom.databricks.com");
-            assert_eq!(token, "test-token");
            assert_eq!(model, "llama-2-70b");
            assert_eq!(temperature, Some(0.7));
            assert_eq!(max_tokens, Some(2000));
@@ -283,7 +278,6 @@ mod tests {

        // Clean up
        env::remove_var("GOOSE_PROVIDER__TYPE");
-        env::remove_var("GOOSE_PROVIDER__TOKEN");
        env::remove_var("GOOSE_PROVIDER__HOST");
        env::remove_var("GOOSE_PROVIDER__MODEL");
        env::remove_var("GOOSE_PROVIDER__TEMPERATURE");
@@ -372,4 +366,4 @@ mod tests {
        let addr = server_settings.socket_addr();
        assert_eq!(addr.to_string(), "127.0.0.1:3000");
    }
-}
+}
@@ -22,7 +22,7 @@ impl Clone for AppState {
                ProviderConfig::Databricks(config) => ProviderConfig::Databricks(
                    goose::providers::configs::DatabricksProviderConfig {
                        host: config.host.clone(),
-                        token: config.token.clone(),
+                        auth: config.auth.clone(),
                        model: config.model.clone(),
                        temperature: config.temperature,
                        max_tokens: config.max_tokens,
@@ -16,6 +16,7 @@ reqwest = { version = "0.11", features = ["json"] }
 tokio = { version = "1.0", features = ["full"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
+serde_urlencoded = "0.7"
 uuid = { version = "1.0", features = ["v4"] }
 regex = "1.11.1"
 async-trait = "0.1"
@@ -25,11 +26,22 @@ strum_macros = "0.26"
 tera = "1.20.0"
 tokenizers = "0.20.3"
 include_dir = "0.7.4"
-chrono = "0.4.38"
+chrono = { version = "0.4.38", features = ["serde"] }
 indoc = "2.0.5"
+nanoid = "0.4"
+sha2 = "0.10"
+base64 = "0.21"
+url = "2.5"
+axum = "0.7"
+tower-http = { version = "0.5", features = ["cors"] }
+webbrowser = "0.8"
+dotenv = "0.15"

 [dev-dependencies]
 wiremock = "0.6.0"
 mockito = "1.2"
 tempfile = "3.8"
-dotenv = "0.15"
+
+[[example]]
+name = "databricks_oauth"
+path = "examples/databricks_oauth.rs"
@@ -0,0 +1,3 @@
+# Databricks OAuth Configuration
+DATABRICKS_HOST=https://your-workspace.cloud.databricks.com
+DATABRICKS_MODEL=your-model-name
@@ -0,0 +1,49 @@
+use anyhow::Result;
+use dotenv::dotenv;
+use goose::{
+    models::message::Message,
+    providers::{
+        configs::{DatabricksProviderConfig, ProviderConfig},
+        factory::get_provider,
+    },
+};
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    // Load environment variables from .env file
+    dotenv().ok();
+
+    // Get required environment variables
+    let host =
+        std::env::var("DATABRICKS_HOST").expect("DATABRICKS_HOST environment variable is required");
+    let model = std::env::var("DATABRICKS_MODEL")
+        .expect("DATABRICKS_MODEL environment variable is required");
+
+    // Create the Databricks provider configuration with OAuth
+    let config = ProviderConfig::Databricks(DatabricksProviderConfig::with_oauth(host, model));
+
+    // Create the provider
+    let provider = get_provider(config)?;
+
+    // Create a simple message
+    let message = Message::user().with_text("Tell me a short joke about programming.");
+
+    // Get a response
+    let (response, usage) = provider
+        .complete("You are a helpful assistant.", &[message], &[])
+        .await?;
+
+    // Print the response and usage statistics
+    println!("\nResponse from AI:");
+    println!("---------------");
+    for content in response.content {
+        dbg!(content);
+    }
+    println!("\nToken Usage:");
+    println!("------------");
+    println!("Input tokens: {:?}", usage.input_tokens);
+    println!("Output tokens: {:?}", usage.output_tokens);
+    println!("Total tokens: {:?}", usage.total_tokens);
+
+    Ok(())
+}
@@ -2,9 +2,10 @@ pub mod base;
 pub mod configs;
 pub mod databricks;
 pub mod factory;
+pub mod oauth;
 pub mod ollama;
 pub mod openai;
 pub mod utils;

 #[cfg(test)]
-pub mod mock;
+pub mod mock;
@@ -1,11 +1,73 @@
-// Unified enum to wrap different provider configurations
+use serde::{Deserialize, Serialize};
+
+const DEFAULT_CLIENT_ID: &str = "databricks-cli";
+const DEFAULT_REDIRECT_URL: &str = "http://localhost:8020";
+const DEFAULT_SCOPES: &[&str] = &["all-apis"];
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum ProviderConfig {
    OpenAi(OpenAiProviderConfig),
    Databricks(DatabricksProviderConfig),
    Ollama(OllamaProviderConfig),
 }

-// Define specific config structs for each provider
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum DatabricksAuth {
+    Token(String),
+    OAuth {
+        host: String,
+        client_id: String,
+        redirect_url: String,
+        scopes: Vec<String>,
+    },
+}
+
+impl DatabricksAuth {
+    /// Create a new OAuth configuration with default values
+    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(),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DatabricksProviderConfig {
+    pub host: String,
+    pub model: String,
+    pub auth: DatabricksAuth,
+    pub temperature: Option<f32>,
+    pub max_tokens: Option<i32>,
+}
+
+impl DatabricksProviderConfig {
+    /// Create a new configuration with token authentication
+    pub fn with_token(host: String, model: String, token: String) -> Self {
+        Self {
+            host,
+            model,
+            auth: DatabricksAuth::Token(token),
+            temperature: None,
+            max_tokens: None,
+        }
+    }
+
+    /// Create a new configuration with OAuth authentication using default settings
+    pub fn with_oauth(host: String, model: String) -> Self {
+        Self {
+            host: host.clone(),
+            model,
+            auth: DatabricksAuth::oauth(host),
+            temperature: None,
+            max_tokens: None,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct OpenAiProviderConfig {
    pub host: String,
    pub api_key: String,
@@ -14,14 +76,7 @@ pub struct OpenAiProviderConfig {
    pub max_tokens: Option<i32>,
 }

-pub struct DatabricksProviderConfig {
-    pub host: String,
-    pub token: String,
-    pub model: String,
-    pub temperature: Option<f32>,
-    pub max_tokens: Option<i32>,
-}
-
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct OllamaProviderConfig {
    pub host: String,
    pub model: String,
@@ -5,7 +5,8 @@ use serde_json::{json, Value};
 use std::time::Duration;

 use super::base::{Provider, Usage};
-use super::configs::DatabricksProviderConfig;
+use super::configs::{DatabricksAuth, DatabricksProviderConfig};
+use super::oauth;
 use super::utils::{
    check_openai_context_length_error, messages_to_openai_spec, openai_response_to_message,
    tools_to_openai_spec,
@@ -22,16 +23,27 @@ impl DatabricksProvider {
    pub fn new(config: DatabricksProviderConfig) -> Result<Self> {
        let client = Client::builder()
            .timeout(Duration::from_secs(600)) // 10 minutes timeout
-            .default_headers({
-                let mut headers = reqwest::header::HeaderMap::new();
-                headers.insert("Authorization", format!("Bearer {}", config.token).parse()?);
-                headers
-            })
            .build()?;

        Ok(Self { client, config })
    }

+    async fn ensure_auth_header(&self) -> Result<String> {
+        match &self.config.auth {
+            DatabricksAuth::Token(token) => Ok(format!("Bearer {}", token)),
+            DatabricksAuth::OAuth {
+                host,
+                client_id,
+                redirect_url,
+                scopes,
+            } => {
+                let token =
+                    oauth::get_oauth_token_async(host, client_id, redirect_url, scopes).await?;
+                Ok(format!("Bearer {}", token))
+            }
+        }
+    }
+
    fn get_usage(data: &Value) -> Result<Usage> {
        let usage = data
            .get("usage")
@@ -66,7 +78,14 @@ impl DatabricksProvider {
            self.config.model
        );

-        let response = self.client.post(&url).json(&payload).send().await?;
+        let auth_header = self.ensure_auth_header().await?;
+        let response = self
+            .client
+            .post(&url)
+            .header("Authorization", auth_header)
+            .json(&payload)
+            .send()
+            .await?;

        match response.status() {
            StatusCode::OK => Ok(response.json().await?),
@@ -152,13 +171,11 @@ impl Provider for DatabricksProvider {
 mod tests {
    use super::*;
    use crate::models::message::MessageContent;
-    use anyhow::Result;
-    use serde_json::json;
    use wiremock::matchers::{body_json, header, method, path};
    use wiremock::{Mock, MockServer, ResponseTemplate};

    #[tokio::test]
-    async fn test_databricks_completion() -> Result<()> {
+    async fn test_databricks_completion_with_token() -> Result<()> {
        // Start a mock server
        let mock_server = MockServer::start().await;

@@ -199,8 +216,8 @@ mod tests {
        // Create the DatabricksProvider with the mock server's URL as the host
        let config = DatabricksProviderConfig {
            host: mock_server.uri(),
-            token: "test_token".to_string(),
            model: "my-databricks-model".to_string(),
+            auth: DatabricksAuth::Token("test_token".to_string()),
            temperature: None,
            max_tokens: None,
        };
@@ -0,0 +1,357 @@
+use anyhow::Result;
+use axum::{extract::Query, response::Html, routing::get, Router};
+use base64::Engine;
+use chrono::{DateTime, Utc};
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use sha2::Digest;
+use std::{collections::HashMap, fs, net::SocketAddr, path::PathBuf, sync::Arc};
+use tokio::sync::oneshot;
+use url::Url;
+
+#[derive(Debug, Clone)]
+struct OidcEndpoints {
+    authorization_endpoint: String,
+    token_endpoint: String,
+}
+
+#[derive(Serialize, Deserialize)]
+struct TokenData {
+    access_token: String,
+    expires_at: Option<DateTime<Utc>>,
+}
+
+struct TokenCache {
+    cache_path: PathBuf,
+}
+
+const BASE_PATH: &str = concat!(env!("HOME"), "/.config/goose/databricks/oauth");
+
+impl TokenCache {
+    fn new(host: &str, client_id: &str, scopes: &[String]) -> Self {
+        let mut hasher = sha2::Sha256::new();
+        hasher.update(host.as_bytes());
+        hasher.update(client_id.as_bytes());
+        hasher.update(scopes.join(",").as_bytes());
+        let hash = format!("{:x}", hasher.finalize());
+
+        fs::create_dir_all(BASE_PATH).unwrap();
+        let cache_path = PathBuf::from(BASE_PATH).join(format!("{}.json", hash));
+
+        Self { cache_path }
+    }
+
+    fn load_token(&self) -> Option<TokenData> {
+        if let Ok(contents) = fs::read_to_string(&self.cache_path) {
+            if let Ok(token_data) = serde_json::from_str::<TokenData>(&contents) {
+                if let Some(expires_at) = token_data.expires_at {
+                    if expires_at > Utc::now() {
+                        return Some(token_data);
+                    }
+                } else {
+                    return Some(token_data);
+                }
+            }
+        }
+        None
+    }
+
+    fn save_token(&self, token_data: &TokenData) -> Result<()> {
+        if let Some(parent) = self.cache_path.parent() {
+            fs::create_dir_all(parent)?;
+        }
+        let contents = serde_json::to_string(token_data)?;
+        fs::write(&self.cache_path, contents)?;
+        Ok(())
+    }
+}
+
+async fn get_workspace_endpoints(host: &str) -> Result<OidcEndpoints> {
+    let host = host.trim_end_matches('/');
+    let oidc_url = format!("{}/oidc/.well-known/oauth-authorization-server", host);
+
+    let client = reqwest::Client::new();
+    let resp = client.get(&oidc_url).send().await?;
+
+    if !resp.status().is_success() {
+        return Err(anyhow::anyhow!(
+            "Failed to get OIDC configuration from {}",
+            oidc_url
+        ));
+    }
+
+    let oidc_config: Value = resp.json().await?;
+
+    let authorization_endpoint = oidc_config
+        .get("authorization_endpoint")
+        .and_then(|v| v.as_str())
+        .ok_or_else(|| anyhow::anyhow!("authorization_endpoint not found in OIDC configuration"))?
+        .to_string();
+
+    let token_endpoint = oidc_config
+        .get("token_endpoint")
+        .and_then(|v| v.as_str())
+        .ok_or_else(|| anyhow::anyhow!("token_endpoint not found in OIDC configuration"))?
+        .to_string();
+
+    Ok(OidcEndpoints {
+        authorization_endpoint,
+        token_endpoint,
+    })
+}
+
+struct OAuthFlow {
+    endpoints: OidcEndpoints,
+    client_id: String,
+    redirect_url: String,
+    scopes: Vec<String>,
+    state: String,
+    verifier: String,
+}
+
+impl OAuthFlow {
+    fn new(
+        endpoints: OidcEndpoints,
+        client_id: String,
+        redirect_url: String,
+        scopes: Vec<String>,
+    ) -> Self {
+        Self {
+            endpoints,
+            client_id,
+            redirect_url,
+            scopes,
+            state: nanoid::nanoid!(16),
+            verifier: nanoid::nanoid!(64),
+        }
+    }
+
+    fn get_authorization_url(&self) -> String {
+        let challenge = {
+            let digest = sha2::Sha256::digest(self.verifier.as_bytes());
+            base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(digest)
+        };
+
+        let params = [
+            ("response_type", "code"),
+            ("client_id", &self.client_id),
+            ("redirect_uri", &self.redirect_url),
+            ("scope", &self.scopes.join(" ")),
+            ("state", &self.state),
+            ("code_challenge", &challenge),
+            ("code_challenge_method", "S256"),
+        ];
+
+        format!(
+            "{}?{}",
+            self.endpoints.authorization_endpoint,
+            serde_urlencoded::to_string(params).unwrap()
+        )
+    }
+
+    async fn exchange_code_for_token(&self, code: &str) -> Result<TokenData> {
+        let params = [
+            ("grant_type", "authorization_code"),
+            ("code", code),
+            ("redirect_uri", &self.redirect_url),
+            ("code_verifier", &self.verifier),
+            ("client_id", &self.client_id),
+        ];
+
+        let client = reqwest::Client::new();
+        let resp = client
+            .post(&self.endpoints.token_endpoint)
+            .header("Content-Type", "application/x-www-form-urlencoded")
+            .form(&params)
+            .send()
+            .await?;
+
+        if !resp.status().is_success() {
+            let err_text = resp.text().await?;
+            return Err(anyhow::anyhow!(
+                "Failed to exchange code for token: {}",
+                err_text
+            ));
+        }
+
+        let token_response: Value = resp.json().await?;
+        let access_token = token_response
+            .get("access_token")
+            .and_then(|v| v.as_str())
+            .ok_or_else(|| anyhow::anyhow!("access_token not found in token response"))?
+            .to_string();
+
+        let expires_in = token_response
+            .get("expires_in")
+            .and_then(|v| v.as_u64())
+            .unwrap_or(3600);
+
+        let expires_at = Utc::now() + chrono::Duration::seconds(expires_in as i64);
+
+        Ok(TokenData {
+            access_token,
+            expires_at: Some(expires_at),
+        })
+    }
+
+    async fn execute(&self) -> Result<TokenData> {
+        // Create a channel that will send the auth code from the app process
+        let (tx, rx) = oneshot::channel();
+        let state = self.state.clone();
+        // Axum can theoretically spawn multiple threads, so we need this to be in an Arc even
+        // though it will ultimately only get used once
+        let tx = Arc::new(tokio::sync::Mutex::new(Some(tx)));
+
+        // Setup a server that will recieve the redirect, capture the code, and display success/failure
+        let app = Router::new().route(
+            "/",
+            get(move |Query(params): Query<HashMap<String, String>>| {
+                let tx = Arc::clone(&tx);
+                let state = state.clone();
+                async move {
+                    let code = params.get("code").cloned();
+                    let received_state = params.get("state").cloned();
+
+                    if let (Some(code), Some(received_state)) = (code, received_state) {
+                        if received_state == state {
+                            if let Some(sender) = tx.lock().await.take() {
+                                if sender.send(code).is_ok() {
+                                    // Use the improved HTML response
+                                    return Html(
+                                        "<h2>Login Success</h2><p>You can close this window</p>",
+                                    );
+                                }
+                            }
+                            Html("<h2>Error</h2><p>Authentication already completed.</p>")
+                        } else {
+                            Html("<h2>Error</h2><p>State mismatch.</p>")
+                        }
+                    } else {
+                        Html("<h2>Error</h2><p>Authentication failed.</p>")
+                    }
+                }
+            }),
+        );
+
+        // Start the server to accept the oauth code
+        let redirect_url = Url::parse(&self.redirect_url)?;
+        let port = redirect_url.port().unwrap_or(80);
+        let addr = SocketAddr::from(([127, 0, 0, 1], port));
+
+        let listener = tokio::net::TcpListener::bind(addr).await?;
+
+        let server_handle = tokio::spawn(async move {
+            let server = axum::serve(listener, app);
+            server.await.unwrap();
+        });
+
+        // Open the browser which will redirect with the code to the server
+        let authorization_url = self.get_authorization_url();
+        if webbrowser::open(&authorization_url).is_err() {
+            println!(
+                "Please open this URL in your browser:\n{}",
+                authorization_url
+            );
+        }
+
+        // Wait for the authorization code with a timeout
+        let code = tokio::time::timeout(
+            std::time::Duration::from_secs(60), // 1 minute timeout
+            rx,
+        )
+        .await
+        .map_err(|_| anyhow::anyhow!("Authentication timed out"))??;
+
+        // Stop the server
+        server_handle.abort();
+
+        // Exchange the code for a token
+        self.exchange_code_for_token(&code).await
+    }
+}
+
+pub(crate) async fn get_oauth_token_async(
+    host: &str,
+    client_id: &str,
+    redirect_url: &str,
+    scopes: &[String],
+) -> Result<String> {
+    let token_cache = TokenCache::new(host, client_id, scopes);
+
+    // Try cache first
+    if let Some(token) = token_cache.load_token() {
+        return Ok(token.access_token);
+    }
+
+    // Get endpoints and execute flow
+    let endpoints = get_workspace_endpoints(host).await?;
+    let flow = OAuthFlow::new(
+        endpoints,
+        client_id.to_string(),
+        redirect_url.to_string(),
+        scopes.to_vec(),
+    );
+
+    // Execute the OAuth flow and get token
+    let token = flow.execute().await?;
+
+    // Cache and return
+    token_cache.save_token(&token)?;
+    Ok(token.access_token)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use wiremock::{
+        matchers::{method, path},
+        Mock, MockServer, ResponseTemplate,
+    };
+
+    #[tokio::test]
+    async fn test_get_workspace_endpoints() -> Result<()> {
+        let mock_server = MockServer::start().await;
+
+        let mock_response = serde_json::json!({
+            "authorization_endpoint": "https://example.com/oauth2/authorize",
+            "token_endpoint": "https://example.com/oauth2/token"
+        });
+
+        Mock::given(method("GET"))
+            .and(path("/oidc/.well-known/oauth-authorization-server"))
+            .respond_with(ResponseTemplate::new(200).set_body_json(&mock_response))
+            .mount(&mock_server)
+            .await;
+
+        let endpoints = get_workspace_endpoints(&mock_server.uri()).await?;
+
+        assert_eq!(
+            endpoints.authorization_endpoint,
+            "https://example.com/oauth2/authorize"
+        );
+        assert_eq!(endpoints.token_endpoint, "https://example.com/oauth2/token");
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_token_cache() -> Result<()> {
+        let cache = TokenCache::new(
+            "https://example.com",
+            "test-client",
+            &["scope1".to_string()],
+        );
+
+        let token_data = TokenData {
+            access_token: "test-token".to_string(),
+            expires_at: Some(Utc::now() + chrono::Duration::hours(1)),
+        };
+
+        cache.save_token(&token_data)?;
+
+        let loaded_token = cache.load_token().unwrap();
+        assert_eq!(loaded_token.access_token, token_data.access_token);
+
+        Ok(())
+    }
+}
@@ -7,7 +7,7 @@ use goose::{
    },
    providers::{
        base::Provider,
-        configs::{DatabricksProviderConfig, OpenAiProviderConfig, ProviderConfig},
+        configs::{DatabricksAuth, DatabricksProviderConfig, OpenAiProviderConfig, ProviderConfig},
        factory::get_provider,
    },
 };
@@ -91,9 +91,7 @@ impl ProviderTester {

    /// Run all provider tests
    async fn run_test_suite(&self) -> Result<()> {
-        println!("Running basic response test...");
        self.test_basic_response().await?;
-        println!("Running tool usage test...");
        self.test_tool_usage().await?;
        Ok(())
    }
@@ -144,8 +142,32 @@ async fn test_databricks_provider() -> Result<()> {

    let config = ProviderConfig::Databricks(DatabricksProviderConfig {
        host: std::env::var("DATABRICKS_HOST")?,
-        token: std::env::var("DATABRICKS_TOKEN")?,
        model: std::env::var("DATABRICKS_MODEL")?,
+        auth: DatabricksAuth::Token(std::env::var("DATABRICKS_TOKEN")?),
+        temperature: None,
+        max_tokens: None,
+    });
+
+    let tester = ProviderTester::new(config)?;
+    tester.run_test_suite().await?;
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_databricks_provider_oauth() -> Result<()> {
+    load_env();
+
+    // Skip if credentials aren't available
+    if std::env::var("DATABRICKS_HOST").is_err() || std::env::var("DATABRICKS_MODEL").is_err() {
+        println!("Skipping Databricks OAuth tests - credentials not configured");
+        return Ok(());
+    }
+
+    let config = ProviderConfig::Databricks(DatabricksProviderConfig {
+        host: std::env::var("DATABRICKS_HOST")?,
+        model: std::env::var("DATABRICKS_MODEL")?,
+        auth: DatabricksAuth::oauth(std::env::var("DATABRICKS_HOST")?),
        temperature: None,
        max_tokens: None,
    });