From 339f5f4b9ec643cef2f3668990bc02d9b1f6aa9d Mon Sep 17 00:00:00 2001 From: Bradley Axen Date: Sun, 24 Nov 2024 08:02:53 +1100 Subject: [PATCH] feat: Add databricks oauth (#313) --- .../goose-cli/src/profile/provider_helper.rs | 28 +- crates/goose-server/src/configuration.rs | 14 +- crates/goose-server/src/state.rs | 2 +- crates/goose/Cargo.toml | 16 +- crates/goose/examples/.env.example | 3 + crates/goose/examples/databricks_oauth.rs | 49 +++ crates/goose/src/providers.rs | 3 +- crates/goose/src/providers/configs.rs | 75 +++- crates/goose/src/providers/databricks.rs | 39 +- crates/goose/src/providers/oauth.rs | 357 ++++++++++++++++++ crates/goose/tests/providers.rs | 30 +- 11 files changed, 564 insertions(+), 52 deletions(-) create mode 100644 crates/goose/examples/.env.example create mode 100644 crates/goose/examples/databricks_oauth.rs create mode 100644 crates/goose/src/providers/oauth.rs diff --git a/crates/goose-cli/src/profile/provider_helper.rs b/crates/goose-cli/src/profile/provider_helper.rs index 47463c2a08..dcf5b3bd50 100644 --- a/crates/goose-cli/src/profile/provider_helper.rs +++ b/crates/goose-cli/src/profile/provider_helper.rs @@ -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)), diff --git a/crates/goose-server/src/configuration.rs b/crates/goose-server/src/configuration.rs index e321030287..178a204cf2 100644 --- a/crates/goose-server/src/configuration.rs +++ b/crates/goose-server/src/configuration.rs @@ -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"); } -} \ No newline at end of file +} diff --git a/crates/goose-server/src/state.rs b/crates/goose-server/src/state.rs index 204b8d6d12..18f8974a5b 100644 --- a/crates/goose-server/src/state.rs +++ b/crates/goose-server/src/state.rs @@ -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, diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index 4f7d1fc30a..6143da75a3 100644 --- a/crates/goose/Cargo.toml +++ b/crates/goose/Cargo.toml @@ -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" \ No newline at end of file diff --git a/crates/goose/examples/.env.example b/crates/goose/examples/.env.example new file mode 100644 index 0000000000..8154b22bbf --- /dev/null +++ b/crates/goose/examples/.env.example @@ -0,0 +1,3 @@ +# Databricks OAuth Configuration +DATABRICKS_HOST=https://your-workspace.cloud.databricks.com +DATABRICKS_MODEL=your-model-name \ No newline at end of file diff --git a/crates/goose/examples/databricks_oauth.rs b/crates/goose/examples/databricks_oauth.rs new file mode 100644 index 0000000000..aa0df1ea9a --- /dev/null +++ b/crates/goose/examples/databricks_oauth.rs @@ -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(()) +} diff --git a/crates/goose/src/providers.rs b/crates/goose/src/providers.rs index 172b812098..45969601c9 100644 --- a/crates/goose/src/providers.rs +++ b/crates/goose/src/providers.rs @@ -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; \ No newline at end of file +pub mod mock; diff --git a/crates/goose/src/providers/configs.rs b/crates/goose/src/providers/configs.rs index 6ca6a08864..5cc92ee2a6 100644 --- a/crates/goose/src/providers/configs.rs +++ b/crates/goose/src/providers/configs.rs @@ -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, + }, +} + +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, + pub max_tokens: Option, +} + +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, } -pub struct DatabricksProviderConfig { - pub host: String, - pub token: String, - pub model: String, - pub temperature: Option, - pub max_tokens: Option, -} - +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct OllamaProviderConfig { pub host: String, pub model: String, diff --git a/crates/goose/src/providers/databricks.rs b/crates/goose/src/providers/databricks.rs index 64798bf7c9..905dac90e8 100644 --- a/crates/goose/src/providers/databricks.rs +++ b/crates/goose/src/providers/databricks.rs @@ -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 { 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 { + 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 { 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, }; diff --git a/crates/goose/src/providers/oauth.rs b/crates/goose/src/providers/oauth.rs new file mode 100644 index 0000000000..a6badf94f2 --- /dev/null +++ b/crates/goose/src/providers/oauth.rs @@ -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>, +} + +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 { + if let Ok(contents) = fs::read_to_string(&self.cache_path) { + if let Ok(token_data) = serde_json::from_str::(&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 { + 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, + state: String, + verifier: String, +} + +impl OAuthFlow { + fn new( + endpoints: OidcEndpoints, + client_id: String, + redirect_url: String, + scopes: Vec, + ) -> 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 { + 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(¶ms) + .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 { + // 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>| { + 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( + "

Login Success

You can close this window

", + ); + } + } + Html("

Error

Authentication already completed.

") + } else { + Html("

Error

State mismatch.

") + } + } else { + Html("

Error

Authentication failed.

") + } + } + }), + ); + + // 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 { + 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(()) + } +} diff --git a/crates/goose/tests/providers.rs b/crates/goose/tests/providers.rs index 0c75237159..a5d891c608 100644 --- a/crates/goose/tests/providers.rs +++ b/crates/goose/tests/providers.rs @@ -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, });