fix: retry on authentication failure with credential refresh (#7812)

This commit is contained in:
Will Pfleger
2026-03-16 12:23:12 -04:00
committed by GitHub
parent 3f661c3334
commit b3510bad25
5 changed files with 108 additions and 15 deletions
+1 -1
View File
@@ -941,7 +941,7 @@ impl Config {
Ok(()) Ok(())
} }
fn invalidate_secrets_cache(&self) { pub fn invalidate_secrets_cache(&self) {
let mut cache = self.secrets_cache.lock().unwrap(); let mut cache = self.secrets_cache.lock().unwrap();
*cache = None; *cache = None;
} }
+6
View File
@@ -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> { async fn update_mode(&self, _session_id: &str, _mode: GooseMode) -> Result<(), ProviderError> {
Ok(()) Ok(())
} }
+44 -5
View File
@@ -6,6 +6,7 @@ use futures::{StreamExt, TryStreamExt};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use serde_json::Value; use serde_json::Value;
use std::io; use std::io;
use std::sync::{Arc, Mutex};
use std::time::Duration; use std::time::Duration;
use tokio::pin; use tokio::pin;
use tokio_util::codec::{FramedRead, LinesCodec}; use tokio_util::codec::{FramedRead, LinesCodec};
@@ -81,13 +82,30 @@ impl DatabricksAuth {
struct DatabricksAuthProvider { struct DatabricksAuthProvider {
auth: DatabricksAuth, auth: DatabricksAuth,
token_cache: Arc<Mutex<Option<String>>>,
} }
#[async_trait] #[async_trait]
impl AuthProvider for DatabricksAuthProvider { impl AuthProvider for DatabricksAuthProvider {
async fn get_auth_header(&self) -> Result<(String, String)> { async fn get_auth_header(&self) -> Result<(String, String)> {
let token = match &self.auth { 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::<String>("DATABRICKS_TOKEN")
.unwrap_or_else(|_| original.clone());
*self.token_cache.lock().unwrap() = Some(fresh.clone());
fresh
}
}
}
DatabricksAuth::OAuth { DatabricksAuth::OAuth {
host, host,
client_id, client_id,
@@ -112,6 +130,8 @@ pub struct DatabricksProvider {
fast_retry_config: RetryConfig, fast_retry_config: RetryConfig,
#[serde(skip)] #[serde(skip)]
name: String, name: String,
#[serde(skip)]
token_cache: Arc<Mutex<Option<String>>>,
} }
impl DatabricksProvider { impl DatabricksProvider {
@@ -140,8 +160,15 @@ impl DatabricksProvider {
DatabricksAuth::oauth(host.clone()) DatabricksAuth::oauth(host.clone())
}; };
let auth_method = let token_cache = Arc::new(Mutex::new(match &auth {
AuthMethod::Custom(Box::new(DatabricksAuthProvider { auth: auth.clone() })); 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 = let api_client =
ApiClient::with_timeout(host, auth_method, Duration::from_secs(DEFAULT_TIMEOUT_SECS))?; ApiClient::with_timeout(host, auth_method, Duration::from_secs(DEFAULT_TIMEOUT_SECS))?;
@@ -154,6 +181,7 @@ impl DatabricksProvider {
retry_config, retry_config,
fast_retry_config, fast_retry_config,
name: DATABRICKS_PROVIDER_NAME.to_string(), name: DATABRICKS_PROVIDER_NAME.to_string(),
token_cache,
}; };
provider.model = provider.model =
model.with_fast(DATABRICKS_DEFAULT_FAST_MODEL, DATABRICKS_PROVIDER_NAME)?; 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<Self> { pub fn from_params(host: String, api_key: String, model: ModelConfig) -> Result<Self> {
let token_cache = Arc::new(Mutex::new(Some(api_key.clone())));
let auth = DatabricksAuth::token(api_key); let auth = DatabricksAuth::token(api_key);
let auth_method = let auth_method = AuthMethod::Custom(Box::new(DatabricksAuthProvider {
AuthMethod::Custom(Box::new(DatabricksAuthProvider { auth: auth.clone() })); auth: auth.clone(),
token_cache: token_cache.clone(),
}));
let api_client = ApiClient::with_timeout(host, auth_method, Duration::from_secs(600))?; let api_client = ApiClient::with_timeout(host, auth_method, Duration::from_secs(600))?;
@@ -213,6 +244,7 @@ impl DatabricksProvider {
retry_config: RetryConfig::default(), retry_config: RetryConfig::default(),
fast_retry_config: RetryConfig::new(0, 0, 1.0, 0), fast_retry_config: RetryConfig::new(0, 0, 1.0, 0),
name: DATABRICKS_PROVIDER_NAME.to_string(), name: DATABRICKS_PROVIDER_NAME.to_string(),
token_cache,
}) })
} }
@@ -285,6 +317,13 @@ impl Provider for DatabricksProvider {
self.retry_config.clone() 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 { fn get_model_config(&self) -> ModelConfig {
self.model.clone() self.model.clone()
} }
+42 -7
View File
@@ -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] #[async_trait]
pub trait ProviderRetry { pub trait ProviderRetry {
fn retry_config(&self) -> RetryConfig { fn retry_config(&self) -> RetryConfig {
@@ -138,6 +140,23 @@ pub trait ProviderRetry {
self.with_retry_config(operation, self.retry_config()).await self.with_retry_config(operation, self.retry_config()).await
} }
async fn with_retry_config<F, Fut, T>(
&self,
operation: F,
config: RetryConfig,
) -> Result<T, ProviderError>
where
F: Fn() -> Fut + Send,
Fut: Future<Output = Result<T, ProviderError>> + Send,
T: Send;
}
#[async_trait]
impl<P: Provider> ProviderRetry for P {
fn retry_config(&self) -> RetryConfig {
Provider::retry_config(self)
}
async fn with_retry_config<F, Fut, T>( async fn with_retry_config<F, Fut, T>(
&self, &self,
operation: F, operation: F,
@@ -149,11 +168,33 @@ pub trait ProviderRetry {
T: Send, T: Send,
{ {
let mut attempts = 0; let mut attempts = 0;
let mut auth_retried = false;
loop { loop {
return match operation().await { return match operation().await {
Ok(result) => Ok(result), Ok(result) => Ok(result),
Err(error) => { 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 { if should_retry(&error) && attempts < config.max_retries {
attempts += 1; attempts += 1;
tracing::warn!( tracing::warn!(
@@ -191,9 +232,3 @@ pub trait ProviderRetry {
} }
} }
} }
impl<P: Provider> ProviderRetry for P {
fn retry_config(&self) -> RetryConfig {
Provider::retry_config(self)
}
}
+15 -2
View File
@@ -174,7 +174,13 @@ const SettingsRoute = ({ activeSessionId }: { activeSessionId?: string }) => {
viewOptions.section = sectionFromUrl; viewOptions.section = sectionFromUrl;
} }
return <SettingsView onClose={() => navigate('/')} setView={setView} viewOptions={{...viewOptions, sessionId: activeSessionId}} />; return (
<SettingsView
onClose={() => navigate('/')}
setView={setView}
viewOptions={{ ...viewOptions, sessionId: activeSessionId }}
/>
);
}; };
const SessionsRoute = () => { const SessionsRoute = () => {
@@ -667,7 +673,14 @@ export function AppInner() {
/> />
} }
/> />
<Route path="settings" element={<SettingsRoute activeSessionId={activeSessions[activeSessions.length - 1]?.sessionId} />} /> <Route
path="settings"
element={
<SettingsRoute
activeSessionId={activeSessions[activeSessions.length - 1]?.sessionId}
/>
}
/>
<Route <Route
path="extensions" path="extensions"
element={ element={