diff --git a/crates/goose/src/config/base.rs b/crates/goose/src/config/base.rs index 51a5f460b0..2553f315ff 100644 --- a/crates/goose/src/config/base.rs +++ b/crates/goose/src/config/base.rs @@ -941,7 +941,7 @@ impl Config { Ok(()) } - fn invalidate_secrets_cache(&self) { + pub fn invalidate_secrets_cache(&self) { let mut cache = self.secrets_cache.lock().unwrap(); *cache = None; } diff --git a/crates/goose/src/providers/base.rs b/crates/goose/src/providers/base.rs index f9968b7601..5b66939493 100644 --- a/crates/goose/src/providers/base.rs +++ b/crates/goose/src/providers/base.rs @@ -707,6 +707,12 @@ pub trait Provider: Send + Sync { )) } + async fn refresh_credentials(&self) -> Result<(), ProviderError> { + Err(ProviderError::NotImplemented( + "credential refresh not supported by this provider".to_string(), + )) + } + async fn update_mode(&self, _session_id: &str, _mode: GooseMode) -> Result<(), ProviderError> { Ok(()) } diff --git a/crates/goose/src/providers/databricks.rs b/crates/goose/src/providers/databricks.rs index f5c764940a..c89aefe77e 100644 --- a/crates/goose/src/providers/databricks.rs +++ b/crates/goose/src/providers/databricks.rs @@ -6,6 +6,7 @@ use futures::{StreamExt, TryStreamExt}; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::io; +use std::sync::{Arc, Mutex}; use std::time::Duration; use tokio::pin; use tokio_util::codec::{FramedRead, LinesCodec}; @@ -81,13 +82,30 @@ impl DatabricksAuth { struct DatabricksAuthProvider { auth: DatabricksAuth, + token_cache: Arc>>, } #[async_trait] impl AuthProvider for DatabricksAuthProvider { async fn get_auth_header(&self) -> Result<(String, String)> { let token = match &self.auth { - DatabricksAuth::Token(token) => token.clone(), + DatabricksAuth::Token(original) => { + let cached = self.token_cache.lock().unwrap().clone(); + match cached { + Some(t) => t, + None => { + // Cache was cleared by refresh_credentials(); re-read + // from config which may have a sidecar-rotated token. + // Fall back to the constructor-provided token if config + // lookup fails (e.g. from_params usage). + let fresh = crate::config::Config::global() + .get_secret::("DATABRICKS_TOKEN") + .unwrap_or_else(|_| original.clone()); + *self.token_cache.lock().unwrap() = Some(fresh.clone()); + fresh + } + } + } DatabricksAuth::OAuth { host, client_id, @@ -112,6 +130,8 @@ pub struct DatabricksProvider { fast_retry_config: RetryConfig, #[serde(skip)] name: String, + #[serde(skip)] + token_cache: Arc>>, } impl DatabricksProvider { @@ -140,8 +160,15 @@ impl DatabricksProvider { DatabricksAuth::oauth(host.clone()) }; - let auth_method = - AuthMethod::Custom(Box::new(DatabricksAuthProvider { auth: auth.clone() })); + 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_TIMEOUT_SECS))?; @@ -154,6 +181,7 @@ impl DatabricksProvider { retry_config, fast_retry_config, name: DATABRICKS_PROVIDER_NAME.to_string(), + token_cache, }; provider.model = model.with_fast(DATABRICKS_DEFAULT_FAST_MODEL, DATABRICKS_PROVIDER_NAME)?; @@ -199,9 +227,12 @@ impl DatabricksProvider { } pub fn from_params(host: String, api_key: String, model: ModelConfig) -> Result { + let token_cache = Arc::new(Mutex::new(Some(api_key.clone()))); let auth = DatabricksAuth::token(api_key); - let auth_method = - AuthMethod::Custom(Box::new(DatabricksAuthProvider { auth: auth.clone() })); + 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(600))?; @@ -213,6 +244,7 @@ impl DatabricksProvider { retry_config: RetryConfig::default(), fast_retry_config: RetryConfig::new(0, 0, 1.0, 0), name: DATABRICKS_PROVIDER_NAME.to_string(), + token_cache, }) } @@ -285,6 +317,13 @@ impl Provider for DatabricksProvider { self.retry_config.clone() } + async fn refresh_credentials(&self) -> Result<(), ProviderError> { + crate::config::Config::global().invalidate_secrets_cache(); + *self.token_cache.lock().unwrap() = None; + tracing::info!("Invalidated secrets cache and token cache for credential refresh"); + Ok(()) + } + fn get_model_config(&self) -> ModelConfig { self.model.clone() } diff --git a/crates/goose/src/providers/retry.rs b/crates/goose/src/providers/retry.rs index 4a8a17a92f..0ddd093f39 100644 --- a/crates/goose/src/providers/retry.rs +++ b/crates/goose/src/providers/retry.rs @@ -122,7 +122,9 @@ where } } -/// Trait for retry functionality to keep Provider dyn-compatible +/// Trait for retry functionality to keep Provider dyn-compatible. +/// +/// All `Provider` implementors get this via the blanket impl below. #[async_trait] pub trait ProviderRetry { fn retry_config(&self) -> RetryConfig { @@ -138,6 +140,23 @@ pub trait ProviderRetry { self.with_retry_config(operation, self.retry_config()).await } + async fn with_retry_config( + &self, + operation: F, + config: RetryConfig, + ) -> Result + where + F: Fn() -> Fut + Send, + Fut: Future> + Send, + T: Send; +} + +#[async_trait] +impl ProviderRetry for P { + fn retry_config(&self) -> RetryConfig { + Provider::retry_config(self) + } + async fn with_retry_config( &self, operation: F, @@ -149,11 +168,33 @@ pub trait ProviderRetry { T: Send, { let mut attempts = 0; + let mut auth_retried = false; loop { return match operation().await { Ok(result) => Ok(result), Err(error) => { + // Auth retry is separate from transient-error retries: we get + // at most 1 credential refresh, independent of max_retries. + if matches!(error, ProviderError::Authentication(_)) && !auth_retried { + auth_retried = true; + match self.refresh_credentials().await { + Ok(()) => { + tracing::warn!( + "Credentials refreshed after auth error, retrying: {:?}", + error + ); + continue; + } + Err(refresh_err) => { + tracing::warn!( + "Credential refresh failed, returning original auth error: {:?}", + refresh_err + ); + } + } + } + if should_retry(&error) && attempts < config.max_retries { attempts += 1; tracing::warn!( @@ -191,9 +232,3 @@ pub trait ProviderRetry { } } } - -impl ProviderRetry for P { - fn retry_config(&self) -> RetryConfig { - Provider::retry_config(self) - } -} diff --git a/ui/desktop/src/App.tsx b/ui/desktop/src/App.tsx index 05945e05ba..c0e55e4f7e 100644 --- a/ui/desktop/src/App.tsx +++ b/ui/desktop/src/App.tsx @@ -174,7 +174,13 @@ const SettingsRoute = ({ activeSessionId }: { activeSessionId?: string }) => { viewOptions.section = sectionFromUrl; } - return navigate('/')} setView={setView} viewOptions={{...viewOptions, sessionId: activeSessionId}} />; + return ( + navigate('/')} + setView={setView} + viewOptions={{ ...viewOptions, sessionId: activeSessionId }} + /> + ); }; const SessionsRoute = () => { @@ -667,7 +673,14 @@ export function AppInner() { /> } /> - } /> + + } + />