From c82c431c70e1afa31c9393cf24cdf354db8f0e2e Mon Sep 17 00:00:00 2001 From: Jack Amadeo Date: Thu, 2 Jul 2026 09:18:57 -0700 Subject: [PATCH] move databricks providers into goose-providers crate (#10162) --- Cargo.lock | 1 + crates/goose-provider-types/src/formats.rs | 1 + .../src}/formats/databricks.rs | 16 +- crates/goose-providers/Cargo.toml | 1 + .../src}/databricks.rs | 162 ++++++++---------- .../src}/databricks_auth.rs | 39 ++++- .../src}/databricks_v2.rs | 135 ++++++--------- crates/goose-providers/src/lib.rs | 3 + crates/goose/examples/agent.rs | 2 +- crates/goose/examples/databricks_oauth.rs | 2 +- crates/goose/examples/image_tool.rs | 2 +- crates/goose/src/providers/databricks_def.rs | 111 ++++++++++++ .../goose/src/providers/databricks_v2_def.rs | 97 +++++++++++ crates/goose/src/providers/formats/mod.rs | 4 +- crates/goose/src/providers/init.rs | 12 +- crates/goose/src/providers/mod.rs | 5 +- crates/goose/tests/providers.rs | 2 +- 17 files changed, 385 insertions(+), 210 deletions(-) rename crates/{goose/src/providers => goose-provider-types/src}/formats/databricks.rs (99%) rename crates/{goose/src/providers => goose-providers/src}/databricks.rs (89%) rename crates/{goose/src/providers => goose-providers/src}/databricks_auth.rs (53%) rename crates/{goose/src/providers => goose-providers/src}/databricks_v2.rs (80%) create mode 100644 crates/goose/src/providers/databricks_def.rs create mode 100644 crates/goose/src/providers/databricks_v2_def.rs diff --git a/Cargo.lock b/Cargo.lock index bb47795c9d..596a64d8b5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5200,6 +5200,7 @@ dependencies = [ "tokio-util", "tracing", "url", + "urlencoding", "utoipa 4.2.3", "wiremock", ] diff --git a/crates/goose-provider-types/src/formats.rs b/crates/goose-provider-types/src/formats.rs index 5abaf4c3e9..208835687e 100644 --- a/crates/goose-provider-types/src/formats.rs +++ b/crates/goose-provider-types/src/formats.rs @@ -1,4 +1,5 @@ pub mod anthropic; +pub mod databricks; pub mod ollama; pub mod openai; pub mod openai_responses; diff --git a/crates/goose/src/providers/formats/databricks.rs b/crates/goose-provider-types/src/formats/databricks.rs similarity index 99% rename from crates/goose/src/providers/formats/databricks.rs rename to crates/goose-provider-types/src/formats/databricks.rs index 7ba1141a0b..886ef28f03 100644 --- a/crates/goose/src/providers/formats/databricks.rs +++ b/crates/goose-provider-types/src/formats/databricks.rs @@ -1,16 +1,16 @@ use crate::conversation::message::{Message, MessageContent}; -use crate::providers::formats::anthropic::{ +use crate::formats::anthropic::{ adaptive_output_effort, model_supports_temperature, thinking_budget_tokens, thinking_type_for_provider, ThinkingType, }; -use goose_providers::model::ModelConfig; +use crate::model::ModelConfig; -use anyhow::{anyhow, Error}; -use goose_providers::formats::openai::{ +use crate::formats::openai::{ extract_reasoning_effort, is_openai_responses_model, is_valid_function_name, openai_reasoning_effort_for_thinking, sanitize_function_name, }; -use goose_providers::images::{convert_image, detect_image_path, load_image_file, ImageFormat}; +use crate::images::{convert_image, detect_image_path, load_image_file, ImageFormat}; +use anyhow::{anyhow, Error}; use rmcp::model::{ object, AnnotateAble, CallToolRequestParams, Content, ErrorCode, ErrorData, RawContent, ResourceContents, Role, Tool, @@ -19,7 +19,7 @@ use serde::Serialize; use serde_json::{json, Value}; use std::borrow::Cow; -pub(crate) const DATABRICKS_PROVIDER_NAME: &str = "databricks"; +pub const DATABRICKS_PROVIDER_NAME: &str = "databricks"; #[derive(Serialize)] struct DatabricksMessage { @@ -424,7 +424,7 @@ pub fn response_to_message(response: &Value) -> anyhow::Result { }; content.push(MessageContent::tool_request(id, Err(error))); } else { - match goose_providers::json::parse_tool_arguments(&arguments_str) { + match crate::json::parse_tool_arguments(&arguments_str) { Some(params) if params.is_object() => { content.push(MessageContent::tool_request( id, @@ -448,7 +448,7 @@ pub fn response_to_message(response: &Value) -> anyhow::Result { } None => { let message_text = - goose_providers::json::truncation_error_message(&arguments_str) + crate::json::truncation_error_message(&arguments_str) .unwrap_or_else(|| { format!( "Could not interpret tool use parameters for id {id}" diff --git a/crates/goose-providers/Cargo.toml b/crates/goose-providers/Cargo.toml index d7e0351110..4f203080a2 100644 --- a/crates/goose-providers/Cargo.toml +++ b/crates/goose-providers/Cargo.toml @@ -48,6 +48,7 @@ tokio = { workspace = true } tokio-stream = { workspace = true, features = ["io-util"] } tokio-util = { workspace = true, features = ["compat"] } url = { workspace = true } +urlencoding = { workspace = true } pem = { version = "3.0.2", default-features = false, features = ["std"], optional = true } pkcs1 = { version = "0.7.5", default-features = false, features = ["pkcs8", "std"], optional = true } # v0.10 matches the der/const-oid series used by sec1 v0.7 and pkcs1 v0.7; upgrading to v0.11 causes type mismatches across those crates. diff --git a/crates/goose/src/providers/databricks.rs b/crates/goose-providers/src/databricks.rs similarity index 89% rename from crates/goose/src/providers/databricks.rs rename to crates/goose-providers/src/databricks.rs index b3e9ce7dd0..b857b91371 100644 --- a/crates/goose/src/providers/databricks.rs +++ b/crates/goose-providers/src/databricks.rs @@ -1,39 +1,38 @@ -use anyhow::Result; -use async_trait::async_trait; -use futures::future::BoxFuture; -use goose_providers::formats::openai::{ +use crate::formats::openai::{ extract_reasoning_effort, is_openai_responses_model, openai_reasoning_effort_for_thinking, }; -use goose_providers::images::ImageFormat; +use crate::images::ImageFormat; +use anyhow::Result; +use async_trait::async_trait; 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}; -use super::base::{ - ConfigKey, MessageStream, ModelInfo, Provider, ProviderDef, ProviderMetadata, - DEFAULT_PROVIDER_TIMEOUT_SECS, +use crate::api_client::{ApiClient, AuthMethod, TlsConfig}; +use crate::base::{ConfigKey, MessageStream, ModelInfo, Provider, ProviderMetadata}; +const DEFAULT_PROVIDER_TIMEOUT_SECS: u64 = 600; +use crate::conversation::message::Message; +use crate::databricks_auth::{ + DatabricksAuth, DatabricksAuthProvider, DatabricksOauthTokenProvider, DatabricksRefreshHook, + DatabricksSessionIdProvider, DatabricksTokenResolver, }; -use super::databricks_auth::{DatabricksAuth, DatabricksAuthProvider}; -use super::formats::databricks::{create_request_for_provider, DATABRICKS_PROVIDER_NAME}; -use super::openai_compatible::{ +use crate::errors::ProviderError; +use crate::formats::databricks::create_request_for_provider; +pub use crate::formats::databricks::DATABRICKS_PROVIDER_NAME; +use crate::formats::openai_responses::create_responses_request; +use crate::model::ModelConfig; +use crate::openai_compatible::{ handle_status, map_http_error_to_provider_error, sanitize_url, stream_openai_compat, stream_responses_compat, }; -use super::retry::ProviderRetry; -use crate::config::ConfigError; -use crate::conversation::message::Message; -use crate::instance_id::get_instance_id; -use crate::providers::retry::{ +use crate::request_log::{start_log, LoggerHandleExt}; +use crate::retry::ProviderRetry; +use crate::retry::{ RetryConfig, DEFAULT_BACKOFF_MULTIPLIER, DEFAULT_INITIAL_RETRY_INTERVAL_MS, DEFAULT_MAX_RETRIES, DEFAULT_MAX_RETRY_INTERVAL_MS, }; -use goose_providers::errors::ProviderError; -use goose_providers::formats::openai_responses::create_responses_request; -use goose_providers::model::ModelConfig; -use goose_providers::request_log::{start_log, LoggerHandleExt}; use rmcp::model::Tool; use serde_json::json; @@ -73,7 +72,7 @@ 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, serde::Serialize)] +#[derive(serde::Serialize)] pub struct DatabricksProvider { #[serde(skip)] api_client: ApiClient, @@ -89,39 +88,30 @@ pub struct DatabricksProvider { token_cache: Arc>>, #[serde(skip)] instance_id: Option, + #[serde(skip)] + refresh_hook: Option, + #[serde(skip)] + session_id_provider: Option, } impl DatabricksProvider { pub async fn cleanup() -> Result<()> { - super::oauth::cleanup_oauth_cache() + Ok(()) } - pub async fn from_env( - tls_config: Option, + #[allow(clippy::too_many_arguments)] + pub fn new( + host: String, + auth: DatabricksAuth, + retry_config: RetryConfig, + tls_config: Option, + oauth_token_provider: Option, + token_resolver: Option, + request_builder: Option, + instance_id: Option, + refresh_hook: Option, + session_id_provider: Option, ) -> 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()) - }; - let token_cache = Arc::new(Mutex::new(match &auth { DatabricksAuth::Token(t) => Some(t.clone()), _ => None, @@ -130,15 +120,19 @@ impl DatabricksProvider { let auth_method = AuthMethod::Custom(Box::new(DatabricksAuthProvider { auth: auth.clone(), token_cache: token_cache.clone(), + oauth_token_provider, + token_resolver, })); - let api_client = ApiClient::with_timeout_and_tls( + let mut api_client = ApiClient::with_timeout_and_tls( host.clone(), auth_method, Duration::from_secs(DEFAULT_PROVIDER_TIMEOUT_SECS), - tls_config.clone(), - )? - .with_request_builder(crate::session_context::session_id_request_builder()); + tls_config, + )?; + if let Some(request_builder) = request_builder { + api_client = api_client.with_request_builder(request_builder); + } Ok(Self { api_client, @@ -148,33 +142,27 @@ impl DatabricksProvider { retry_config, name: DATABRICKS_PROVIDER_NAME.to_string(), token_cache, - instance_id: Self::resolve_instance_id(), + instance_id, + refresh_hook, + session_id_provider, }) } - 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()) + pub fn load_retry_config(get_param: impl Fn(&str) -> Option) -> RetryConfig { + let max_retries = get_param("DATABRICKS_MAX_RETRIES") + .and_then(|v| 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()) + let initial_interval_ms = get_param("DATABRICKS_INITIAL_RETRY_INTERVAL_MS") + .and_then(|v| 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()) + let backoff_multiplier = get_param("DATABRICKS_BACKOFF_MULTIPLIER") + .and_then(|v| 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()) + let max_interval_ms = get_param("DATABRICKS_MAX_RETRY_INTERVAL_MS") + .and_then(|v| v.parse::().ok()) .unwrap_or(DEFAULT_MAX_RETRY_INTERVAL_MS); RetryConfig::new( @@ -185,17 +173,6 @@ impl DatabricksProvider { ) } - fn resolve_instance_id() -> Option { - let enabled = crate::config::Config::global() - .get_param::("GOOSE_DATABRICKS_CLIENT_REQUEST_ID") - .unwrap_or(false); - if enabled { - Some(get_instance_id().to_string()) - } else { - None - } - } - fn is_claude_model(model_name: &str) -> bool { model_name.to_lowercase().contains("claude") } @@ -514,7 +491,7 @@ impl DatabricksProvider { } } -impl goose_providers::base::ProviderDescriptor for DatabricksProvider { +impl crate::base::ProviderDescriptor for DatabricksProvider { fn metadata() -> ProviderMetadata { ProviderMetadata::new( DATABRICKS_PROVIDER_NAME, @@ -532,17 +509,6 @@ impl goose_providers::base::ProviderDescriptor for DatabricksProvider { } } -impl ProviderDef for DatabricksProvider { - type Provider = Self; - - fn from_env( - _extensions: Vec, - tls_config: Option, - ) -> BoxFuture<'static, Result> { - Box::pin(Self::from_env(tls_config)) - } -} - #[async_trait] impl Provider for DatabricksProvider { fn get_name(&self) -> &str { @@ -554,7 +520,9 @@ impl Provider for DatabricksProvider { } async fn refresh_credentials(&self) -> Result<(), ProviderError> { - crate::config::Config::global().invalidate_secrets_cache(); + if let Some(refresh_hook) = &self.refresh_hook { + refresh_hook(); + } *self.token_cache.lock().unwrap() = None; tracing::info!("Invalidated secrets cache and token cache for credential refresh"); Ok(()) @@ -567,7 +535,11 @@ impl Provider for DatabricksProvider { messages: &[Message], tools: &[Tool], ) -> Result { - let session_id = crate::session_context::current_session_id().unwrap_or_default(); + let session_id = self + .session_id_provider + .as_ref() + .and_then(|provider| provider()) + .unwrap_or_default(); let (endpoint_name, _) = extract_reasoning_effort(&model_config.model_name); let endpoint_info = self.resolve_endpoint_info_cached(&endpoint_name).await.ok(); let effective_model_name = endpoint_info diff --git a/crates/goose/src/providers/databricks_auth.rs b/crates/goose-providers/src/databricks_auth.rs similarity index 53% rename from crates/goose/src/providers/databricks_auth.rs rename to crates/goose-providers/src/databricks_auth.rs index e2795e0ef5..46973c4481 100644 --- a/crates/goose/src/providers/databricks_auth.rs +++ b/crates/goose-providers/src/databricks_auth.rs @@ -1,15 +1,23 @@ use anyhow::Result; use async_trait::async_trait; use serde::{Deserialize, Serialize}; +use std::future::Future; +use std::pin::Pin; use std::sync::{Arc, Mutex}; -use super::api_client::AuthProvider; -use super::oauth; +use crate::api_client::AuthProvider; const DEFAULT_CLIENT_ID: &str = "databricks-cli"; const DEFAULT_REDIRECT_URL: &str = "http://localhost"; const DEFAULT_SCOPES: &[&str] = &["all-apis", "offline_access"]; +pub type DatabricksOauthTokenFuture = Pin> + Send>>; +pub type DatabricksOauthTokenProvider = + Arc) -> DatabricksOauthTokenFuture + Send + Sync>; +pub type DatabricksTokenResolver = Arc Option + Send + Sync>; +pub type DatabricksRefreshHook = Arc; +pub type DatabricksSessionIdProvider = Arc Option + Send + Sync>; + #[derive(Debug, Clone, Serialize, Deserialize)] pub enum DatabricksAuth { Token(String), @@ -36,9 +44,11 @@ impl DatabricksAuth { } } -pub(crate) struct DatabricksAuthProvider { +pub struct DatabricksAuthProvider { pub auth: DatabricksAuth, pub token_cache: Arc>>, + pub oauth_token_provider: Option, + pub token_resolver: Option, } #[async_trait] @@ -50,9 +60,11 @@ impl AuthProvider for DatabricksAuthProvider { match cached { Some(t) => t, None => { - let fresh = crate::config::Config::global() - .get_secret::("DATABRICKS_TOKEN") - .unwrap_or_else(|_| original.clone()); + let fresh = self + .token_resolver + .as_ref() + .and_then(|resolve| resolve()) + .unwrap_or_else(|| original.clone()); *self.token_cache.lock().unwrap() = Some(fresh.clone()); fresh } @@ -63,8 +75,19 @@ impl AuthProvider for DatabricksAuthProvider { client_id, redirect_url, scopes, - } => oauth::get_oauth_token_async(host, client_id, redirect_url, scopes).await?, + } => { + let Some(provider) = &self.oauth_token_provider else { + anyhow::bail!("Databricks OAuth token provider is not configured") + }; + provider( + host.clone(), + client_id.clone(), + redirect_url.clone(), + scopes.clone(), + ) + .await? + } }; - Ok(("Authorization".to_string(), format!("Bearer {}", token))) + Ok(("Authorization".to_string(), format!("Bearer {token}"))) } } diff --git a/crates/goose/src/providers/databricks_v2.rs b/crates/goose-providers/src/databricks_v2.rs similarity index 80% rename from crates/goose/src/providers/databricks_v2.rs rename to crates/goose-providers/src/databricks_v2.rs index fb289b13a0..af30140590 100644 --- a/crates/goose/src/providers/databricks_v2.rs +++ b/crates/goose-providers/src/databricks_v2.rs @@ -1,11 +1,10 @@ +use crate::formats::anthropic::{AnthropicFormatOptions, ANTHROPIC_PROVIDER_NAME}; +use crate::formats::openai::{self, extract_reasoning_effort, is_openai_responses_model}; +use crate::images::ImageFormat; use anyhow::Result; use async_stream::try_stream; use async_trait::async_trait; -use futures::future::BoxFuture; use futures::TryStreamExt; -use goose_providers::formats::anthropic::{AnthropicFormatOptions, ANTHROPIC_PROVIDER_NAME}; -use goose_providers::formats::openai::{self, extract_reasoning_effort, is_openai_responses_model}; -use goose_providers::images::ImageFormat; use serde::Serialize; use serde_json::Value; use std::io; @@ -14,25 +13,25 @@ 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::formats::anthropic; -use super::openai_compatible::{handle_status, stream_openai_compat, stream_responses_compat}; -use super::retry::ProviderRetry; -use crate::config::ConfigError; +use crate::api_client::{ApiClient, AuthMethod, TlsConfig}; +use crate::base::{ConfigKey, MessageStream, Provider, ProviderMetadata}; +const DEFAULT_PROVIDER_TIMEOUT_SECS: u64 = 600; use crate::conversation::message::Message; -use crate::providers::retry::{ +use crate::databricks_auth::{ + DatabricksAuth, DatabricksAuthProvider, DatabricksOauthTokenProvider, DatabricksRefreshHook, + DatabricksTokenResolver, +}; +use crate::errors::ProviderError; +use crate::formats::anthropic; +use crate::formats::openai_responses; +use crate::model::ModelConfig; +use crate::openai_compatible::{handle_status, stream_openai_compat, stream_responses_compat}; +use crate::request_log::{start_log, LoggerHandleExt}; +use crate::retry::ProviderRetry; +use crate::retry::{ RetryConfig, DEFAULT_BACKOFF_MULTIPLIER, DEFAULT_INITIAL_RETRY_INTERVAL_MS, DEFAULT_MAX_RETRIES, DEFAULT_MAX_RETRY_INTERVAL_MS, }; -use goose_providers::errors::ProviderError; -use goose_providers::formats::openai_responses; -use goose_providers::model::ModelConfig; -use goose_providers::request_log::{start_log, LoggerHandleExt}; use rmcp::model::Tool; const DATABRICKS_V2_PROVIDER_NAME: &str = "databricks_v2"; @@ -51,7 +50,7 @@ enum DatabricksV2Route { MlflowChatCompletions, } -#[derive(Debug, Serialize)] +#[derive(Serialize)] pub struct DatabricksV2Provider { #[serde(skip)] api_client: ApiClient, @@ -61,47 +60,25 @@ pub struct DatabricksV2Provider { name: String, #[serde(skip)] token_cache: Arc>>, + #[serde(skip)] + refresh_hook: Option, } impl DatabricksV2Provider { pub async fn cleanup() -> Result<()> { - super::oauth::cleanup_oauth_cache() + Ok(()) } - pub async fn from_env( - tls_config: Option, - ) -> 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, retry_config, tls_config) - } - - fn new( + #[allow(clippy::too_many_arguments)] + pub fn new( host: String, auth: DatabricksAuth, retry_config: RetryConfig, - tls_config: Option, + tls_config: Option, + oauth_token_provider: Option, + token_resolver: Option, + request_builder: Option, + refresh_hook: Option, ) -> Result { let token_cache = Arc::new(Mutex::new(match &auth { DatabricksAuth::Token(t) => Some(t.clone()), @@ -111,47 +88,44 @@ impl DatabricksV2Provider { let auth_method = AuthMethod::Custom(Box::new(DatabricksAuthProvider { auth: auth.clone(), token_cache: token_cache.clone(), + oauth_token_provider, + token_resolver, })); - let api_client = ApiClient::with_timeout_and_tls( + let mut api_client = ApiClient::with_timeout_and_tls( host, auth_method, Duration::from_secs(DEFAULT_PROVIDER_TIMEOUT_SECS), tls_config, - )? - .with_request_builder(crate::session_context::session_id_request_builder()); + )?; + if let Some(request_builder) = request_builder { + api_client = api_client.with_request_builder(request_builder); + } Ok(Self { api_client, retry_config, name: DATABRICKS_V2_PROVIDER_NAME.to_string(), token_cache, + refresh_hook, }) } - 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()) + pub fn load_retry_config(get_param: impl Fn(&str) -> Option) -> RetryConfig { + let max_retries = get_param("DATABRICKS_MAX_RETRIES") + .and_then(|v| 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()) + let initial_interval_ms = get_param("DATABRICKS_INITIAL_RETRY_INTERVAL_MS") + .and_then(|v| 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()) + let backoff_multiplier = get_param("DATABRICKS_BACKOFF_MULTIPLIER") + .and_then(|v| 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()) + let max_interval_ms = get_param("DATABRICKS_MAX_RETRY_INTERVAL_MS") + .and_then(|v| v.parse::().ok()) .unwrap_or(DEFAULT_MAX_RETRY_INTERVAL_MS); RetryConfig::new( @@ -328,7 +302,7 @@ impl DatabricksV2Provider { } } -impl goose_providers::base::ProviderDescriptor for DatabricksV2Provider { +impl crate::base::ProviderDescriptor for DatabricksV2Provider { fn metadata() -> ProviderMetadata { ProviderMetadata::new( DATABRICKS_V2_PROVIDER_NAME, @@ -345,17 +319,6 @@ impl goose_providers::base::ProviderDescriptor for DatabricksV2Provider { } } -impl ProviderDef for DatabricksV2Provider { - type Provider = Self; - - fn from_env( - _extensions: Vec, - tls_config: Option, - ) -> BoxFuture<'static, Result> { - Box::pin(Self::from_env(tls_config)) - } -} - #[async_trait] impl Provider for DatabricksV2Provider { fn get_name(&self) -> &str { @@ -367,7 +330,9 @@ impl Provider for DatabricksV2Provider { } async fn refresh_credentials(&self) -> Result<(), ProviderError> { - crate::config::Config::global().invalidate_secrets_cache(); + if let Some(refresh_hook) = &self.refresh_hook { + refresh_hook(); + } *self.token_cache.lock().unwrap() = None; Ok(()) } diff --git a/crates/goose-providers/src/lib.rs b/crates/goose-providers/src/lib.rs index 8584a1b955..4eebda49a4 100644 --- a/crates/goose-providers/src/lib.rs +++ b/crates/goose-providers/src/lib.rs @@ -1,5 +1,8 @@ pub mod anthropic; pub mod api_client; +pub mod databricks; +pub mod databricks_auth; +pub mod databricks_v2; pub use goose_provider_types::{ base, canonical, conversation, errors, formats, goose_mode, images, json, model, permission, request_log, retry, thinking, utils, diff --git a/crates/goose/examples/agent.rs b/crates/goose/examples/agent.rs index 8a045304bc..9e12db22d3 100644 --- a/crates/goose/examples/agent.rs +++ b/crates/goose/examples/agent.rs @@ -4,8 +4,8 @@ use goose::agents::{Agent, AgentEvent, ExtensionConfig, SessionConfig}; use goose::config::{GooseMode, DEFAULT_EXTENSION_DESCRIPTION, DEFAULT_EXTENSION_TIMEOUT}; use goose::conversation::message::Message; use goose::providers::create_with_named_model; -use goose::providers::databricks::DATABRICKS_DEFAULT_MODEL; use goose::session::session_manager::SessionType; +use goose_providers::databricks::DATABRICKS_DEFAULT_MODEL; use std::path::PathBuf; #[tokio::main] diff --git a/crates/goose/examples/databricks_oauth.rs b/crates/goose/examples/databricks_oauth.rs index e1883f248e..904829939e 100644 --- a/crates/goose/examples/databricks_oauth.rs +++ b/crates/goose/examples/databricks_oauth.rs @@ -2,7 +2,7 @@ use anyhow::Result; use dotenvy::dotenv; use goose::conversation::message::Message; use goose::providers::create_with_named_model; -use goose::providers::databricks::DATABRICKS_DEFAULT_MODEL; +use goose_providers::databricks::DATABRICKS_DEFAULT_MODEL; #[tokio::main] async fn main() -> Result<()> { diff --git a/crates/goose/examples/image_tool.rs b/crates/goose/examples/image_tool.rs index d2ba3a1456..90979cac6e 100644 --- a/crates/goose/examples/image_tool.rs +++ b/crates/goose/examples/image_tool.rs @@ -4,8 +4,8 @@ use dotenvy::dotenv; use goose::conversation::message::Message; use goose::providers::anthropic::ANTHROPIC_DEFAULT_MODEL; use goose::providers::create_with_named_model; -use goose::providers::databricks::DATABRICKS_DEFAULT_MODEL; use goose::providers::openai::OPEN_AI_DEFAULT_MODEL; +use goose_providers::databricks::DATABRICKS_DEFAULT_MODEL; use rmcp::model::{CallToolRequestParams, Content, Tool}; use rmcp::object; use std::fs; diff --git a/crates/goose/src/providers/databricks_def.rs b/crates/goose/src/providers/databricks_def.rs new file mode 100644 index 0000000000..4096aa5b41 --- /dev/null +++ b/crates/goose/src/providers/databricks_def.rs @@ -0,0 +1,111 @@ +use anyhow::Result; +use futures::future::BoxFuture; +use goose_providers::api_client::TlsConfig; +use goose_providers::base::ProviderDescriptor; +use goose_providers::databricks::DatabricksProvider; +use goose_providers::databricks_auth::{ + DatabricksAuth, DatabricksOauthTokenProvider, DatabricksRefreshHook, + DatabricksSessionIdProvider, DatabricksTokenResolver, +}; +use std::sync::Arc; + +use crate::config::{Config, ConfigError, ExtensionConfig}; +use crate::providers::base::ProviderDef; + +pub struct DatabricksProviderDef; + +impl ProviderDescriptor for DatabricksProviderDef { + fn metadata() -> goose_providers::base::ProviderMetadata { + DatabricksProvider::metadata() + } +} + +impl ProviderDef for DatabricksProviderDef { + type Provider = DatabricksProvider; + + fn from_env( + _extensions: Vec, + tls_config: Option, + ) -> BoxFuture<'static, Result> { + Box::pin(from_env(tls_config)) + } +} + +pub async fn from_env(tls_config: Option) -> Result { + let config = Config::global(); + let host = load_host(config)?; + let retry_config = DatabricksProvider::load_retry_config(|key| config.get_param(key).ok()); + let auth = if let Ok(api_key) = config.get_secret("DATABRICKS_TOKEN") { + DatabricksAuth::token(api_key) + } else { + DatabricksAuth::oauth(host.clone()) + }; + + DatabricksProvider::new( + host, + auth, + retry_config, + tls_config, + Some(oauth_token_provider()), + Some(token_resolver()), + Some(crate::session_context::session_id_request_builder()), + resolve_instance_id(), + Some(refresh_hook()), + Some(session_id_provider()), + ) +} + +pub async fn cleanup() -> Result<()> { + crate::providers::oauth::cleanup_oauth_cache() +} + +fn load_host(config: &Config) -> Result { + let mut host: Result = config.get_param("DATABRICKS_HOST"); + if host.is_err() { + host = config.get_secret("DATABRICKS_HOST") + } + + host.map_err(|_| { + ConfigError::NotFound( + "Did not find DATABRICKS_HOST in either config file or keyring".to_string(), + ) + .into() + }) +} + +fn resolve_instance_id() -> Option { + let enabled = Config::global() + .get_param::("GOOSE_DATABRICKS_CLIENT_REQUEST_ID") + .unwrap_or(false); + enabled.then(|| crate::instance_id::get_instance_id().to_string()) +} + +fn token_resolver() -> DatabricksTokenResolver { + Arc::new(|| { + Config::global() + .get_secret::("DATABRICKS_TOKEN") + .ok() + }) +} + +fn refresh_hook() -> DatabricksRefreshHook { + Arc::new(|| Config::global().invalidate_secrets_cache()) +} + +fn session_id_provider() -> DatabricksSessionIdProvider { + Arc::new(crate::session_context::current_session_id) +} + +fn oauth_token_provider() -> DatabricksOauthTokenProvider { + Arc::new(|host, client_id, redirect_url, scopes| { + Box::pin(async move { + crate::providers::oauth::get_oauth_token_async( + &host, + &client_id, + &redirect_url, + &scopes, + ) + .await + }) + }) +} diff --git a/crates/goose/src/providers/databricks_v2_def.rs b/crates/goose/src/providers/databricks_v2_def.rs new file mode 100644 index 0000000000..a3a35cbc54 --- /dev/null +++ b/crates/goose/src/providers/databricks_v2_def.rs @@ -0,0 +1,97 @@ +use anyhow::Result; +use futures::future::BoxFuture; +use goose_providers::api_client::TlsConfig; +use goose_providers::base::ProviderDescriptor; +use goose_providers::databricks_auth::{ + DatabricksAuth, DatabricksOauthTokenProvider, DatabricksRefreshHook, DatabricksTokenResolver, +}; +use goose_providers::databricks_v2::DatabricksV2Provider; +use std::sync::Arc; + +use crate::config::{Config, ConfigError, ExtensionConfig}; +use crate::providers::base::ProviderDef; + +pub struct DatabricksV2ProviderDef; + +impl ProviderDescriptor for DatabricksV2ProviderDef { + fn metadata() -> goose_providers::base::ProviderMetadata { + DatabricksV2Provider::metadata() + } +} + +impl ProviderDef for DatabricksV2ProviderDef { + type Provider = DatabricksV2Provider; + + fn from_env( + _extensions: Vec, + tls_config: Option, + ) -> BoxFuture<'static, Result> { + Box::pin(from_env(tls_config)) + } +} + +pub async fn from_env(tls_config: Option) -> Result { + let config = Config::global(); + let host = load_host(config)?; + let retry_config = DatabricksV2Provider::load_retry_config(|key| config.get_param(key).ok()); + let auth = if let Ok(api_key) = config.get_secret("DATABRICKS_TOKEN") { + DatabricksAuth::token(api_key) + } else { + DatabricksAuth::oauth(host.clone()) + }; + + DatabricksV2Provider::new( + host, + auth, + retry_config, + tls_config, + Some(oauth_token_provider()), + Some(token_resolver()), + Some(crate::session_context::session_id_request_builder()), + Some(refresh_hook()), + ) +} + +pub async fn cleanup() -> Result<()> { + crate::providers::oauth::cleanup_oauth_cache() +} + +fn load_host(config: &Config) -> Result { + let mut host: Result = config.get_param("DATABRICKS_HOST"); + if host.is_err() { + host = config.get_secret("DATABRICKS_HOST") + } + + host.map_err(|_| { + ConfigError::NotFound( + "Did not find DATABRICKS_HOST in either config file or keyring".to_string(), + ) + .into() + }) +} + +fn token_resolver() -> DatabricksTokenResolver { + Arc::new(|| { + Config::global() + .get_secret::("DATABRICKS_TOKEN") + .ok() + }) +} + +fn refresh_hook() -> DatabricksRefreshHook { + Arc::new(|| Config::global().invalidate_secrets_cache()) +} + +fn oauth_token_provider() -> DatabricksOauthTokenProvider { + Arc::new(|host, client_id, redirect_url, scopes| { + Box::pin(async move { + crate::providers::oauth::get_oauth_token_async( + &host, + &client_id, + &redirect_url, + &scopes, + ) + .await + }) + }) +} diff --git a/crates/goose/src/providers/formats/mod.rs b/crates/goose/src/providers/formats/mod.rs index f8b3d506ea..bec979c250 100644 --- a/crates/goose/src/providers/formats/mod.rs +++ b/crates/goose/src/providers/formats/mod.rs @@ -3,7 +3,9 @@ pub mod anthropic { } #[cfg(feature = "aws-providers")] pub mod bedrock; -pub mod databricks; +pub mod databricks { + pub use goose_providers::formats::databricks::*; +} pub mod gcpvertexai; pub mod google; pub mod openrouter; diff --git a/crates/goose/src/providers/init.rs b/crates/goose/src/providers/init.rs index 3752715fbf..41079ba836 100644 --- a/crates/goose/src/providers/init.rs +++ b/crates/goose/src/providers/init.rs @@ -19,8 +19,6 @@ use super::{ codex_acp::CodexAcpProvider, copilot_acp::CopilotAcpProvider, cursor_agent::CursorAgentProvider, - databricks::DatabricksProvider, - databricks_v2::DatabricksV2Provider, gcpvertexai::GcpVertexAIProvider, gemini_cli::GeminiCliProvider, gemini_oauth::GeminiOAuthProvider, @@ -41,6 +39,8 @@ use super::{ use crate::config::ExtensionConfig; use crate::providers::anthropic_def::AnthropicProviderDef; use crate::providers::base::ProviderType; +use crate::providers::databricks_def::{self, DatabricksProviderDef}; +use crate::providers::databricks_v2_def::{self, DatabricksV2ProviderDef}; use crate::providers::ollama_def::OllamaProviderDef; use crate::providers::openai_def::OpenAiProviderDef; use crate::{ @@ -92,11 +92,11 @@ async fn init_registry() -> RwLock { ); registry.register::(true); registry.register::(false); - registry.register_with_inventory::( + registry.register_with_inventory::( true, Some(registrations::refresh_only()), ); - registry.register_with_inventory::( + registry.register_with_inventory::( false, Some(registrations::refresh_only()), ); @@ -145,11 +145,11 @@ async fn init_registry() -> RwLock { ); registry.set_cleanup( "databricks", - Arc::new(|| Box::pin(DatabricksProvider::cleanup())), + Arc::new(|| Box::pin(databricks_def::cleanup())), ); registry.set_cleanup( "databricks_v2", - Arc::new(|| Box::pin(DatabricksV2Provider::cleanup())), + Arc::new(|| Box::pin(databricks_v2_def::cleanup())), ); registry.set_cleanup( "kimi_code", diff --git a/crates/goose/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index dd91b1d945..20b883c8bb 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -29,9 +29,8 @@ pub mod codex_acp; pub mod copilot_acp; pub mod cursor_agent; pub mod custom_provider_config; -pub mod databricks; -pub mod databricks_auth; -pub mod databricks_v2; +pub mod databricks_def; +pub mod databricks_v2_def; pub mod formats; mod gcpauth; pub mod gcpvertexai; diff --git a/crates/goose/tests/providers.rs b/crates/goose/tests/providers.rs index 9b1411563d..ab5753b41a 100644 --- a/crates/goose/tests/providers.rs +++ b/crates/goose/tests/providers.rs @@ -15,7 +15,6 @@ use goose::providers::bedrock::BEDROCK_DEFAULT_MODEL; use goose::providers::claude_code::CLAUDE_CODE_DEFAULT_MODEL; use goose::providers::codex::CODEX_DEFAULT_MODEL; use goose::providers::create_with_named_model; -use goose::providers::databricks::DATABRICKS_DEFAULT_MODEL; use goose::providers::google::GOOGLE_DEFAULT_MODEL; use goose::providers::litellm::LITELLM_DEFAULT_MODEL; use goose::providers::openai::OPEN_AI_DEFAULT_MODEL; @@ -24,6 +23,7 @@ use goose::providers::sagemaker_tgi::SAGEMAKER_TGI_DEFAULT_MODEL; use goose::providers::snowflake::SNOWFLAKE_DEFAULT_MODEL; use goose::providers::xai::XAI_DEFAULT_MODEL; use goose::session::{SessionManager, SessionType}; +use goose_providers::databricks::DATABRICKS_DEFAULT_MODEL; use goose_providers::errors::ProviderError; use goose_test_support::{ EnforceSessionId, ExpectedSessionId, IgnoreSessionId, McpFixture, FAKE_CODE,