fix: retry on authentication failure with credential refresh (#7812)
This commit is contained in:
@@ -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;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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(())
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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()
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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
@@ -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={
|
||||||
|
|||||||
Reference in New Issue
Block a user