Files
Aether/crates/aether-provider/transport/src/oauth_refresh/mod.rs
T

1547 lines
55 KiB
Rust
Raw Normal View History

use std::collections::BTreeMap;
use std::fmt;
use std::sync::Arc;
use std::time::{Duration, Instant};
2026-04-28 15:46:21 +08:00
use aether_oauth::core::OAuthError;
use aether_oauth::network::{
OAuthHttpExecutor, OAuthHttpRequest, OAuthHttpResponse, OAuthNetworkContext,
};
use aether_oauth::provider::ProviderOAuthTransportContext;
use aether_runtime_state::{RuntimeLockLease, RuntimeState};
use async_trait::async_trait;
use serde_json::Value;
use thiserror::Error;
use tokio::sync::{Mutex, OwnedMutexGuard};
use super::agent_identity::{is_codex_agent_identity_transport, CodexAgentIdentityRefreshAdapter};
use super::generic_oauth::supports_local_generic_oauth_request_auth_resolution;
pub use super::generic_oauth::GenericOAuthRefreshAdapter;
use super::kiro::{
supports_local_kiro_request_auth_resolution, KiroOAuthRefreshAdapter, KiroRequestAuth,
};
use super::snapshot::GatewayProviderTransportSnapshot;
2026-05-10 15:51:38 +08:00
use super::vertex::{
supports_local_vertex_service_account_auth_resolution, VertexServiceAccountRefreshAdapter,
};
#[derive(Debug, Clone, PartialEq, Eq)]
#[allow(clippy::large_enum_variant)]
pub enum LocalResolvedOAuthRequestAuth {
#[allow(dead_code)]
Header {
name: String,
value: String,
},
Kiro(KiroRequestAuth),
}
#[derive(Debug, Clone, PartialEq)]
pub struct LocalOAuthResolution {
pub auth: Option<LocalResolvedOAuthRequestAuth>,
pub refreshed_entry: Option<CachedOAuthEntry>,
pub refresh_in_flight: bool,
/// Indicates that a forced caller reused a newer completed refresh rather
/// than producing a new entry that needs persistence.
pub reused_refresh: bool,
/// Held until the caller persists `refreshed_entry`. The lease TTL remains
/// the cancellation fallback if the caller is dropped.
pub distributed_lease: Option<RuntimeLockLease>,
/// Keeps memory-only refreshes singleflight until the caller validates the
/// credential fence and publishes or discards `refreshed_entry`.
#[doc(hidden)]
pub local_refresh_guard: Option<LocalOAuthRefreshCommitGuard>,
}
#[derive(Clone)]
pub struct LocalOAuthRefreshCommitGuard {
guard: Arc<OwnedMutexGuard<()>>,
}
impl LocalOAuthRefreshCommitGuard {
fn new(guard: OwnedMutexGuard<()>) -> Self {
Self {
guard: Arc::new(guard),
}
}
}
impl fmt::Debug for LocalOAuthRefreshCommitGuard {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("LocalOAuthRefreshCommitGuard")
}
}
impl PartialEq for LocalOAuthRefreshCommitGuard {
fn eq(&self, other: &Self) -> bool {
Arc::ptr_eq(&self.guard, &other.guard)
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct CachedOAuthEntry {
pub provider_type: String,
pub auth_header_name: String,
pub auth_header_value: String,
pub expires_at_unix_secs: Option<u64>,
pub metadata: Option<Value>,
/// Non-secret fingerprint of the credential/configuration that produced it.
pub source_fingerprint: Option<String>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct LocalOAuthHttpRequest {
pub request_id: &'static str,
pub method: reqwest::Method,
pub url: String,
pub headers: BTreeMap<String, String>,
pub json_body: Option<Value>,
pub body_bytes: Option<Vec<u8>>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LocalOAuthHttpResponse {
pub status_code: u16,
pub body_text: String,
}
#[derive(Debug, Error)]
pub enum LocalOAuthRefreshError {
#[error("{provider_type} oauth refresh request failed: {source}")]
Transport {
provider_type: &'static str,
#[source]
source: reqwest::Error,
},
#[error("{provider_type} oauth refresh returned HTTP {status_code}: {body_excerpt}")]
HttpStatus {
provider_type: &'static str,
status_code: u16,
body_excerpt: String,
},
2026-04-28 15:46:21 +08:00
#[error("{provider_type} oauth refresh transport failed: {message}")]
TransportMessage {
provider_type: &'static str,
message: String,
},
#[error("{provider_type} oauth refresh returned invalid response: {message}")]
InvalidResponse {
provider_type: &'static str,
message: String,
},
}
#[async_trait]
pub trait LocalOAuthHttpExecutor: Send + Sync {
async fn execute(
&self,
provider_type: &'static str,
transport: &GatewayProviderTransportSnapshot,
request: &LocalOAuthHttpRequest,
) -> Result<LocalOAuthHttpResponse, LocalOAuthRefreshError>;
}
#[derive(Debug, Clone)]
pub struct ReqwestLocalOAuthHttpExecutor {
client: reqwest::Client,
}
impl ReqwestLocalOAuthHttpExecutor {
pub fn new(client: reqwest::Client) -> Self {
Self { client }
}
}
#[async_trait]
impl LocalOAuthHttpExecutor for ReqwestLocalOAuthHttpExecutor {
async fn execute(
&self,
provider_type: &'static str,
_transport: &GatewayProviderTransportSnapshot,
request: &LocalOAuthHttpRequest,
) -> Result<LocalOAuthHttpResponse, LocalOAuthRefreshError> {
let mut builder = self
.client
.request(request.method.clone(), request.url.as_str());
for (name, value) in &request.headers {
builder = builder.header(name, value);
}
if let Some(json_body) = request.json_body.as_ref() {
builder = builder.json(json_body);
} else if let Some(body_bytes) = request.body_bytes.as_ref() {
builder = builder.body(body_bytes.clone());
}
let response =
builder
.send()
.await
.map_err(|source| LocalOAuthRefreshError::Transport {
provider_type,
source,
})?;
let status_code = response.status().as_u16();
let body_text =
response
.text()
.await
.map_err(|source| LocalOAuthRefreshError::Transport {
provider_type,
source,
})?;
Ok(LocalOAuthHttpResponse {
status_code,
body_text,
})
}
}
2026-04-28 15:46:21 +08:00
pub(crate) struct ProviderOAuthLocalHttpExecutor<'a> {
provider_type: &'static str,
transport: &'a GatewayProviderTransportSnapshot,
inner: &'a dyn LocalOAuthHttpExecutor,
}
impl<'a> ProviderOAuthLocalHttpExecutor<'a> {
pub(crate) fn new(
provider_type: &'static str,
transport: &'a GatewayProviderTransportSnapshot,
inner: &'a dyn LocalOAuthHttpExecutor,
) -> Self {
Self {
provider_type,
transport,
inner,
}
}
}
#[async_trait]
impl OAuthHttpExecutor for ProviderOAuthLocalHttpExecutor<'_> {
async fn execute(&self, request: OAuthHttpRequest) -> Result<OAuthHttpResponse, OAuthError> {
let response = self
.inner
.execute(
self.provider_type,
self.transport,
&LocalOAuthHttpRequest {
request_id: "provider-oauth:local-refresh-token",
method: request.method,
url: request.url,
headers: request.headers,
json_body: request.json_body,
body_bytes: request.body_bytes,
},
)
.await
.map_err(local_refresh_error_to_oauth_error)?;
let json_body = serde_json::from_str::<Value>(&response.body_text).ok();
Ok(OAuthHttpResponse {
status_code: response.status_code,
body_text: response.body_text,
json_body,
})
}
}
pub(crate) fn provider_oauth_transport_context_from_snapshot(
transport: &GatewayProviderTransportSnapshot,
) -> ProviderOAuthTransportContext {
ProviderOAuthTransportContext {
provider_id: transport.provider.id.clone(),
provider_type: transport.provider.provider_type.clone(),
endpoint_id: Some(transport.endpoint.id.clone()),
key_id: Some(transport.key.id.clone()),
auth_type: Some(transport.key.auth_type.clone()),
decrypted_api_key: Some(transport.key.decrypted_api_key.clone()),
decrypted_auth_config: transport.key.decrypted_auth_config.clone(),
provider_config: transport.provider.config.clone(),
endpoint_config: transport.endpoint.config.clone(),
key_config: None,
network: OAuthNetworkContext::provider_operation(None),
}
}
pub(crate) fn oauth_error_to_local_refresh_error(
provider_type: &'static str,
error: OAuthError,
) -> LocalOAuthRefreshError {
match error {
OAuthError::HttpStatus {
status_code,
body_excerpt,
} => LocalOAuthRefreshError::HttpStatus {
provider_type,
status_code,
body_excerpt,
},
OAuthError::Transport(message) => LocalOAuthRefreshError::TransportMessage {
provider_type,
message,
},
OAuthError::InvalidRequest(message)
| OAuthError::InvalidResponse(message)
| OAuthError::Storage(message)
| OAuthError::UnsupportedProvider(message) => LocalOAuthRefreshError::InvalidResponse {
provider_type,
message,
},
OAuthError::InvalidState => LocalOAuthRefreshError::InvalidResponse {
provider_type,
message: "oauth state is invalid or expired".to_string(),
},
OAuthError::EncryptionUnavailable => LocalOAuthRefreshError::InvalidResponse {
provider_type,
message: "oauth encryption unavailable".to_string(),
},
}
}
fn local_refresh_error_to_oauth_error(error: LocalOAuthRefreshError) -> OAuthError {
match error {
LocalOAuthRefreshError::Transport { source, .. } => {
OAuthError::Transport(source.to_string())
}
LocalOAuthRefreshError::TransportMessage { message, .. } => OAuthError::Transport(message),
LocalOAuthRefreshError::HttpStatus {
status_code,
body_excerpt,
..
} => OAuthError::HttpStatus {
status_code,
body_excerpt,
},
LocalOAuthRefreshError::InvalidResponse { message, .. } => {
OAuthError::InvalidResponse(message)
}
}
}
#[async_trait]
pub trait LocalOAuthRefreshAdapter: Send + Sync {
fn provider_type(&self) -> &'static str;
fn supports(&self, transport: &GatewayProviderTransportSnapshot) -> bool {
transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case(self.provider_type())
}
fn resolve_cached(
&self,
transport: &GatewayProviderTransportSnapshot,
entry: &CachedOAuthEntry,
) -> Option<LocalResolvedOAuthRequestAuth>;
/// Resolves a cache entry that is known to have advanced the caller's
/// refresh fence. Agent task rotation can safely use the winner even while
/// the caller still holds the pre-refresh transport snapshot.
fn resolve_fenced_cached(
&self,
transport: &GatewayProviderTransportSnapshot,
entry: &CachedOAuthEntry,
) -> Option<LocalResolvedOAuthRequestAuth> {
self.resolve_cached(transport, entry)
}
/// Resolves the entry returned by this adapter's immediately preceding
/// refresh. Unlike a reusable cache entry, this entry is expected to have
/// advanced the transport generation.
fn resolve_refreshed(
&self,
transport: &GatewayProviderTransportSnapshot,
entry: &CachedOAuthEntry,
) -> Option<LocalResolvedOAuthRequestAuth> {
self.resolve_cached(transport, entry)
}
fn resolve_without_refresh(
&self,
transport: &GatewayProviderTransportSnapshot,
) -> Option<LocalResolvedOAuthRequestAuth>;
fn should_refresh(
&self,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> bool;
/// Identifies the credential/configuration generation used by a refresh.
/// Adapters that support fencing override this method.
fn refresh_fingerprint(
&self,
_transport: &GatewayProviderTransportSnapshot,
_entry: Option<&CachedOAuthEntry>,
) -> Option<String> {
None
}
/// Reconstructs a cache entry from an already-persisted transport after a
/// distributed refresh waiter reloads the winner.
fn cached_entry_from_transport(
&self,
_transport: &GatewayProviderTransportSnapshot,
) -> Option<CachedOAuthEntry> {
None
}
/// Enables bounded negative backoff for transient refresh failures.
fn should_backoff_after_error(&self, _error: &LocalOAuthRefreshError) -> bool {
false
}
/// Agent task registration is a non-idempotent external mutation and must
/// not continue unlocked when a configured distributed lock is unavailable.
fn requires_distributed_refresh_lock(&self) -> bool {
false
}
/// Whether another gateway instance can observe this refresh after the
/// caller persists its result into the provider transport record.
fn shares_refresh_through_transport_persistence(&self) -> bool {
true
}
async fn refresh(
&self,
executor: &dyn LocalOAuthHttpExecutor,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError>;
}
pub struct LocalOAuthRefreshCoordinator {
adapters: Vec<Arc<dyn LocalOAuthRefreshAdapter>>,
cache: Mutex<BTreeMap<String, CachedOAuthEntry>>,
key_locks: Mutex<BTreeMap<String, Arc<Mutex<()>>>>,
refresh_backoff: Mutex<BTreeMap<String, RefreshBackoffState>>,
}
#[derive(Debug, Clone)]
struct RefreshBackoffState {
failures: u32,
retry_after: Instant,
refresh_fingerprint: Option<String>,
}
impl fmt::Debug for LocalOAuthRefreshCoordinator {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("LocalOAuthRefreshCoordinator")
.field("adapter_count", &self.adapters.len())
.finish()
}
}
impl Default for LocalOAuthRefreshCoordinator {
fn default() -> Self {
Self::new()
}
}
impl LocalOAuthRefreshCoordinator {
// Keep the lease alive through the 30s upstream HTTP timeout and the
// subsequent encrypted DB CAS/persistence step. Cancellation still relies
// on expiry as the last-resort release path.
const DISTRIBUTED_REFRESH_LOCK_TTL_MS: u64 = 120_000;
pub fn new() -> Self {
Self {
adapters: vec![
Arc::new(CodexAgentIdentityRefreshAdapter::default()),
Arc::new(KiroOAuthRefreshAdapter::default()),
2026-05-10 15:51:38 +08:00
Arc::new(VertexServiceAccountRefreshAdapter),
Arc::new(GenericOAuthRefreshAdapter::default()),
],
cache: Mutex::new(BTreeMap::new()),
key_locks: Mutex::new(BTreeMap::new()),
refresh_backoff: Mutex::new(BTreeMap::new()),
}
}
/// Captures the refresh generation represented by this transport snapshot.
/// The coordinator cache is intentionally excluded: callers use this value
/// as the fence for the credential generation that produced their request.
pub fn refresh_fingerprint_for_transport(
&self,
transport: &GatewayProviderTransportSnapshot,
) -> Option<String> {
self.adapters
.iter()
.find(|adapter| adapter.supports(transport))
.and_then(|adapter| adapter.refresh_fingerprint(transport, None))
}
async fn lock_for_key(&self, key_id: &str) -> Arc<Mutex<()>> {
let mut key_locks = self.key_locks.lock().await;
key_locks
.entry(key_id.to_string())
.or_insert_with(|| Arc::new(Mutex::new(())))
.clone()
}
async fn cached_entry(&self, key_id: &str) -> Option<CachedOAuthEntry> {
self.cache.lock().await.get(key_id).cloned()
}
async fn insert_cached_entry(&self, key_id: &str, entry: CachedOAuthEntry) {
self.cache.lock().await.insert(key_id.to_string(), entry);
}
2026-04-28 09:35:23 +08:00
pub async fn store_cached_entry(&self, key_id: &str, entry: CachedOAuthEntry) {
self.insert_cached_entry(key_id, entry).await;
}
pub async fn invalidate_cached_entry(&self, key_id: &str) -> bool {
let removed = self.cache.lock().await.remove(key_id).is_some();
self.clear_refresh_backoff(key_id).await;
removed
}
pub async fn resolve_with_result(
&self,
executor: &dyn LocalOAuthHttpExecutor,
transport: &GatewayProviderTransportSnapshot,
2026-05-08 00:18:12 +08:00
distributed_lock: Option<&RuntimeState>,
distributed_owner: Option<&str>,
) -> Result<Option<LocalOAuthResolution>, LocalOAuthRefreshError> {
self.resolve_with_result_mode(
executor,
transport,
distributed_lock,
distributed_owner,
false,
None,
)
.await
}
pub async fn force_refresh_with_result(
&self,
executor: &dyn LocalOAuthHttpExecutor,
transport: &GatewayProviderTransportSnapshot,
2026-05-08 00:18:12 +08:00
distributed_lock: Option<&RuntimeState>,
distributed_owner: Option<&str>,
) -> Result<Option<LocalOAuthResolution>, LocalOAuthRefreshError> {
self.force_refresh_with_result_fenced(
executor,
transport,
distributed_lock,
distributed_owner,
None,
)
.await
}
/// Force a refresh unless another request has already advanced the supplied
/// refresh fence. This prevents a distributed waiter from re-registering a
/// task after the winner has persisted it.
pub async fn force_refresh_with_result_fenced(
&self,
executor: &dyn LocalOAuthHttpExecutor,
transport: &GatewayProviderTransportSnapshot,
distributed_lock: Option<&RuntimeState>,
distributed_owner: Option<&str>,
expected_refresh_fingerprint: Option<&str>,
) -> Result<Option<LocalOAuthResolution>, LocalOAuthRefreshError> {
self.resolve_with_result_mode(
executor,
transport,
distributed_lock,
distributed_owner,
true,
expected_refresh_fingerprint,
)
.await
}
async fn resolve_with_result_mode(
&self,
executor: &dyn LocalOAuthHttpExecutor,
transport: &GatewayProviderTransportSnapshot,
2026-05-08 00:18:12 +08:00
distributed_lock: Option<&RuntimeState>,
distributed_owner: Option<&str>,
force_refresh: bool,
expected_refresh_fingerprint: Option<&str>,
) -> Result<Option<LocalOAuthResolution>, LocalOAuthRefreshError> {
let Some(adapter) = self
.adapters
.iter()
.find(|adapter| adapter.supports(transport))
else {
return Ok(None);
};
let key_id = transport.key.id.trim();
let cached_entry = if key_id.is_empty() {
None
} else {
self.cached_entry(key_id).await
};
let shares_refresh_through_transport_persistence =
adapter.shares_refresh_through_transport_persistence();
let pre_lock_local_refresh_fingerprint = force_refresh
.then(|| adapter.refresh_fingerprint(transport, cached_entry.as_ref()))
.flatten();
if !force_refresh {
if let Some(auth) = cached_entry
.as_ref()
.and_then(|entry| adapter.resolve_cached(transport, entry))
{
return Ok(Some(LocalOAuthResolution::resolved(auth, None)));
}
if let Some(auth) = adapter.resolve_without_refresh(transport) {
return Ok(Some(LocalOAuthResolution::resolved(auth, None)));
}
if !adapter.should_refresh(transport, cached_entry.as_ref()) {
return Ok(None);
}
}
if key_id.is_empty() {
return Ok(None);
}
if force_refresh && shares_refresh_through_transport_persistence {
if let Some(resolution) = Self::resolve_if_refresh_fence_advanced(
adapter.as_ref(),
transport,
cached_entry.as_ref(),
expected_refresh_fingerprint,
) {
return Ok(Some(resolution));
}
}
let refresh_fingerprint = adapter.refresh_fingerprint(transport, cached_entry.as_ref());
if let Some(error) = self
.backoff_error(
key_id,
adapter.provider_type(),
refresh_fingerprint.as_deref(),
)
.await
{
return Err(error);
}
let key_lock = self.lock_for_key(key_id).await;
let key_guard = key_lock.lock_owned().await;
let cached_entry = self.cached_entry(key_id).await;
if force_refresh {
let winner_fingerprint = if shares_refresh_through_transport_persistence {
expected_refresh_fingerprint
} else {
pre_lock_local_refresh_fingerprint.as_deref()
};
if let Some(resolution) = Self::resolve_if_refresh_fence_advanced(
adapter.as_ref(),
transport,
cached_entry.as_ref(),
winner_fingerprint,
) {
return Ok(Some(resolution));
}
}
let refresh_fingerprint = adapter.refresh_fingerprint(transport, cached_entry.as_ref());
if let Some(error) = self
.backoff_error(
key_id,
adapter.provider_type(),
refresh_fingerprint.as_deref(),
)
.await
{
return Err(error);
}
if !force_refresh {
if let Some(auth) = cached_entry
.as_ref()
.and_then(|entry| adapter.resolve_cached(transport, entry))
{
return Ok(Some(LocalOAuthResolution::resolved(auth, None)));
}
if let Some(auth) = adapter.resolve_without_refresh(transport) {
return Ok(Some(LocalOAuthResolution::resolved(auth, None)));
}
if !adapter.should_refresh(transport, cached_entry.as_ref()) {
return Ok(None);
}
}
let distributed_lease = if !shares_refresh_through_transport_persistence {
None
} else {
match (distributed_lock, distributed_owner) {
(Some(lock), Some(owner)) if !owner.trim().is_empty() => {
match lock
.lock_try_acquire(
&format!("provider_oauth_refresh_lock:{key_id}"),
owner,
std::time::Duration::from_millis(Self::DISTRIBUTED_REFRESH_LOCK_TTL_MS),
)
.await
{
Ok(Some(lease)) => Some(lease),
Ok(None) => return Ok(Some(LocalOAuthResolution::refresh_in_flight())),
Err(err) => {
tracing::warn!(
key_id = %key_id,
provider_type = adapter.provider_type(),
error = ?err,
"gateway local oauth refresh distributed lock unavailable"
);
if adapter.requires_distributed_refresh_lock() {
let error = LocalOAuthRefreshError::TransportMessage {
provider_type: adapter.provider_type(),
message: "distributed refresh lock is unavailable".to_string(),
};
if adapter.should_backoff_after_error(&error) {
self.record_refresh_failure(
key_id,
refresh_fingerprint.as_deref(),
)
.await;
}
return Err(error);
}
None
}
}
}
_ => None,
}
};
2026-04-28 02:00:35 +08:00
// Forced refresh still needs the latest rotated refresh_token as input.
// Otherwise a second overlapping refresh can acquire the lock after the
// first one completes, then immediately retry with the stale token that
// came from the original transport snapshot.
let refresh_entry = cached_entry.as_ref();
2026-04-27 23:04:31 +08:00
let refresh_result = adapter.refresh(executor, transport, refresh_entry).await;
let refreshed_entry = match refresh_result {
Ok(Some(entry)) => {
self.clear_refresh_backoff(key_id).await;
entry
}
Ok(None) => {
Self::release_distributed_lease(
distributed_lock,
distributed_lease.as_ref(),
key_id,
adapter.provider_type(),
)
.await;
return Ok(None);
}
Err(error) => {
if adapter.should_backoff_after_error(&error) {
self.record_refresh_failure(key_id, refresh_fingerprint.as_deref())
.await;
}
Self::release_distributed_lease(
distributed_lock,
distributed_lease.as_ref(),
key_id,
adapter.provider_type(),
)
.await;
return Err(error);
}
};
// Cache publication belongs to the caller after durable persistence.
// The result still carries the entry so lock-free callers can inspect
// it or explicitly commit it with `store_cached_entry`.
let Some(auth) = adapter.resolve_refreshed(transport, &refreshed_entry) else {
Self::release_distributed_lease(
distributed_lock,
distributed_lease.as_ref(),
key_id,
adapter.provider_type(),
)
.await;
return Ok(None);
};
let local_refresh_guard = if shares_refresh_through_transport_persistence {
None
} else {
Some(LocalOAuthRefreshCommitGuard::new(key_guard))
};
Ok(Some(LocalOAuthResolution::refreshed(
auth,
refreshed_entry,
distributed_lease,
local_refresh_guard,
)))
}
fn resolve_if_refresh_fence_advanced(
adapter: &dyn LocalOAuthRefreshAdapter,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
expected_refresh_fingerprint: Option<&str>,
) -> Option<LocalOAuthResolution> {
let expected = expected_refresh_fingerprint?;
if adapter.refresh_fingerprint(transport, entry).as_deref() == Some(expected) {
return None;
}
entry
.and_then(|entry| adapter.resolve_fenced_cached(transport, entry))
.map(|auth| {
LocalOAuthResolution::reused(
auth,
entry.expect("cached auth is required when a refresh fence advanced"),
)
})
.or_else(|| {
adapter
.cached_entry_from_transport(transport)
.and_then(|entry| {
adapter
.resolve_cached(transport, &entry)
.map(|auth| LocalOAuthResolution::reused(auth, &entry))
})
})
.or_else(|| {
adapter
.resolve_without_refresh(transport)
.map(|auth| LocalOAuthResolution::resolved(auth, None))
})
}
async fn backoff_error(
&self,
key_id: &str,
provider_type: &'static str,
refresh_fingerprint: Option<&str>,
) -> Option<LocalOAuthRefreshError> {
let mut backoff = self.refresh_backoff.lock().await;
if backoff
.get(key_id)
.is_some_and(|state| state.refresh_fingerprint.as_deref() != refresh_fingerprint)
{
backoff.remove(key_id);
return None;
}
let state = backoff.get(key_id)?;
let remaining = state.retry_after.checked_duration_since(Instant::now())?;
Some(LocalOAuthRefreshError::InvalidResponse {
provider_type,
message: format!(
"refresh temporarily backed off after {} failed attempts (retry in {}ms)",
state.failures,
remaining.as_millis()
),
})
}
async fn record_refresh_failure(&self, key_id: &str, refresh_fingerprint: Option<&str>) {
let mut backoff = self.refresh_backoff.lock().await;
let state = backoff
.entry(key_id.to_string())
.or_insert(RefreshBackoffState {
failures: 0,
retry_after: Instant::now(),
refresh_fingerprint: refresh_fingerprint.map(ToOwned::to_owned),
});
if state.refresh_fingerprint.as_deref() != refresh_fingerprint {
state.failures = 0;
state.refresh_fingerprint = refresh_fingerprint.map(ToOwned::to_owned);
}
state.failures = state.failures.saturating_add(1);
let exponent = state.failures.saturating_sub(1).min(4);
let delay = Duration::from_millis(500u64.saturating_mul(1u64 << exponent));
state.retry_after = Instant::now() + delay.min(Duration::from_secs(8));
}
async fn clear_refresh_backoff(&self, key_id: &str) {
self.refresh_backoff.lock().await.remove(key_id);
}
async fn release_distributed_lease(
distributed_lock: Option<&RuntimeState>,
lease: Option<&RuntimeLockLease>,
key_id: &str,
provider_type: &'static str,
) {
let (Some(lock), Some(lease)) = (distributed_lock, lease) else {
return;
};
if let Err(err) = lock.lock_release(lease).await {
tracing::warn!(
key_id = %key_id,
provider_type,
error = ?err,
"gateway local oauth refresh distributed lock release failed"
);
}
}
pub fn with_adapters_for_tests(adapters: Vec<Arc<dyn LocalOAuthRefreshAdapter>>) -> Self {
Self {
adapters,
cache: Mutex::new(BTreeMap::new()),
key_locks: Mutex::new(BTreeMap::new()),
refresh_backoff: Mutex::new(BTreeMap::new()),
}
}
}
impl LocalOAuthResolution {
fn resolved(
auth: LocalResolvedOAuthRequestAuth,
refreshed_entry: Option<CachedOAuthEntry>,
) -> Self {
Self {
auth: Some(auth),
refreshed_entry,
refresh_in_flight: false,
reused_refresh: false,
distributed_lease: None,
local_refresh_guard: None,
}
}
fn refreshed(
auth: LocalResolvedOAuthRequestAuth,
refreshed_entry: CachedOAuthEntry,
distributed_lease: Option<RuntimeLockLease>,
local_refresh_guard: Option<LocalOAuthRefreshCommitGuard>,
) -> Self {
Self {
auth: Some(auth),
refreshed_entry: Some(refreshed_entry),
refresh_in_flight: false,
reused_refresh: false,
distributed_lease,
local_refresh_guard,
}
}
fn reused(auth: LocalResolvedOAuthRequestAuth, entry: &CachedOAuthEntry) -> Self {
let mut refreshed_entry = entry.clone();
if let LocalResolvedOAuthRequestAuth::Header { name, value } = &auth {
refreshed_entry.auth_header_name = name.clone();
refreshed_entry.auth_header_value = value.clone();
}
Self {
auth: Some(auth),
refreshed_entry: Some(refreshed_entry),
refresh_in_flight: false,
reused_refresh: true,
distributed_lease: None,
local_refresh_guard: None,
}
}
fn refresh_in_flight() -> Self {
Self {
auth: None,
refreshed_entry: None,
refresh_in_flight: true,
reused_refresh: false,
distributed_lease: None,
local_refresh_guard: None,
}
}
}
pub fn supports_local_oauth_request_auth_resolution(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
is_codex_agent_identity_transport(transport)
|| supports_local_kiro_request_auth_resolution(transport)
2026-05-10 15:51:38 +08:00
|| supports_local_vertex_service_account_auth_resolution(transport)
|| supports_local_generic_oauth_request_auth_resolution(transport)
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::time::Duration;
use super::super::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use super::{
CachedOAuthEntry, LocalOAuthHttpExecutor, LocalOAuthRefreshAdapter,
LocalOAuthRefreshCoordinator, LocalOAuthRefreshError, LocalOAuthResolution,
LocalResolvedOAuthRequestAuth, ReqwestLocalOAuthHttpExecutor,
};
use async_trait::async_trait;
use std::sync::Arc;
#[derive(Debug)]
struct TestAdapter {
refresh_hits: Arc<AtomicUsize>,
2026-04-27 23:04:31 +08:00
refresh_with_entry_hits: Arc<AtomicUsize>,
}
#[derive(Debug)]
struct FencedTestAdapter {
refresh_hits: Arc<AtomicUsize>,
fail_refresh: Arc<AtomicBool>,
generation: Arc<AtomicUsize>,
}
#[derive(Debug)]
struct MemoryOnlyFencedTestAdapter {
refresh_hits: Arc<AtomicUsize>,
fingerprint_hits: Arc<AtomicUsize>,
}
#[async_trait]
impl LocalOAuthRefreshAdapter for TestAdapter {
fn provider_type(&self) -> &'static str {
"test-oauth"
}
fn resolve_cached(
&self,
_transport: &GatewayProviderTransportSnapshot,
entry: &CachedOAuthEntry,
) -> Option<LocalResolvedOAuthRequestAuth> {
(entry.provider_type == "test-oauth").then(|| LocalResolvedOAuthRequestAuth::Header {
name: entry.auth_header_name.clone(),
value: entry.auth_header_value.clone(),
})
}
fn resolve_without_refresh(
&self,
transport: &GatewayProviderTransportSnapshot,
) -> Option<LocalResolvedOAuthRequestAuth> {
let secret = transport.key.decrypted_api_key.trim();
(!secret.is_empty() && secret != "__placeholder__").then(|| {
LocalResolvedOAuthRequestAuth::Header {
name: "authorization".to_string(),
value: format!("Bearer {secret}"),
}
})
}
fn should_refresh(
&self,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> bool {
entry.is_none() && transport.key.decrypted_api_key.trim() == "__placeholder__"
}
async fn refresh(
&self,
_executor: &dyn LocalOAuthHttpExecutor,
_transport: &GatewayProviderTransportSnapshot,
2026-04-27 23:04:31 +08:00
entry: Option<&CachedOAuthEntry>,
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError> {
self.refresh_hits.fetch_add(1, Ordering::SeqCst);
2026-04-27 23:04:31 +08:00
if entry.is_some() {
self.refresh_with_entry_hits.fetch_add(1, Ordering::SeqCst);
}
Ok(Some(CachedOAuthEntry {
provider_type: "test-oauth".to_string(),
auth_header_name: "authorization".to_string(),
auth_header_value: "Bearer refreshed-token".to_string(),
expires_at_unix_secs: Some(4_102_444_800),
metadata: None,
source_fingerprint: None,
}))
}
}
#[async_trait]
impl LocalOAuthRefreshAdapter for FencedTestAdapter {
fn provider_type(&self) -> &'static str {
"test-oauth"
}
fn resolve_cached(
&self,
_transport: &GatewayProviderTransportSnapshot,
entry: &CachedOAuthEntry,
) -> Option<LocalResolvedOAuthRequestAuth> {
Some(LocalResolvedOAuthRequestAuth::Header {
name: entry.auth_header_name.clone(),
value: "fresh-winner-assertion".to_string(),
})
}
fn resolve_without_refresh(
&self,
_transport: &GatewayProviderTransportSnapshot,
) -> Option<LocalResolvedOAuthRequestAuth> {
None
}
fn should_refresh(
&self,
_transport: &GatewayProviderTransportSnapshot,
_entry: Option<&CachedOAuthEntry>,
) -> bool {
true
}
fn refresh_fingerprint(
&self,
_transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> Option<String> {
entry
.and_then(|entry| entry.source_fingerprint.clone())
.or_else(|| {
Some(format!(
"generation-{}",
self.generation.load(Ordering::SeqCst)
))
})
}
fn should_backoff_after_error(&self, _error: &LocalOAuthRefreshError) -> bool {
true
}
async fn refresh(
&self,
_executor: &dyn LocalOAuthHttpExecutor,
_transport: &GatewayProviderTransportSnapshot,
_entry: Option<&CachedOAuthEntry>,
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError> {
self.refresh_hits.fetch_add(1, Ordering::SeqCst);
if self.fail_refresh.load(Ordering::SeqCst) {
return Err(LocalOAuthRefreshError::TransportMessage {
provider_type: "test-oauth",
message: "temporary failure".to_string(),
});
}
Ok(Some(CachedOAuthEntry {
provider_type: "test-oauth".to_string(),
auth_header_name: "authorization".to_string(),
auth_header_value: "stale-winner-cache-value".to_string(),
expires_at_unix_secs: None,
metadata: None,
source_fingerprint: Some(format!(
"generation-{}",
self.generation.load(Ordering::SeqCst).saturating_add(1)
)),
}))
}
}
#[async_trait]
impl LocalOAuthRefreshAdapter for MemoryOnlyFencedTestAdapter {
fn provider_type(&self) -> &'static str {
"test-oauth"
}
fn resolve_cached(
&self,
_transport: &GatewayProviderTransportSnapshot,
entry: &CachedOAuthEntry,
) -> Option<LocalResolvedOAuthRequestAuth> {
Some(LocalResolvedOAuthRequestAuth::Header {
name: entry.auth_header_name.clone(),
value: entry.auth_header_value.clone(),
})
}
fn resolve_without_refresh(
&self,
_transport: &GatewayProviderTransportSnapshot,
) -> Option<LocalResolvedOAuthRequestAuth> {
None
}
fn should_refresh(
&self,
_transport: &GatewayProviderTransportSnapshot,
_entry: Option<&CachedOAuthEntry>,
) -> bool {
true
}
fn refresh_fingerprint(
&self,
_transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> Option<String> {
self.fingerprint_hits.fetch_add(1, Ordering::SeqCst);
entry
.and_then(|entry| entry.source_fingerprint.clone())
.or_else(|| Some("transport-generation".to_string()))
}
fn shares_refresh_through_transport_persistence(&self) -> bool {
false
}
async fn refresh(
&self,
_executor: &dyn LocalOAuthHttpExecutor,
_transport: &GatewayProviderTransportSnapshot,
_entry: Option<&CachedOAuthEntry>,
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError> {
let hit = self.refresh_hits.fetch_add(1, Ordering::SeqCst) + 1;
Ok(Some(CachedOAuthEntry {
provider_type: "test-oauth".to_string(),
auth_header_name: "authorization".to_string(),
auth_header_value: format!("Bearer refreshed-token-{hit}"),
expires_at_unix_secs: Some(4_102_444_800),
metadata: None,
source_fingerprint: Some(format!("local-generation-{hit}")),
}))
}
}
fn sample_transport() -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "test".to_string(),
provider_type: "test-oauth".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
2026-04-29 09:25:19 +08:00
api_format: "claude:messages".to_string(),
api_family: Some("claude".to_string()),
endpoint_kind: Some("cli".to_string()),
is_active: true,
base_url: "https://example.test".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "bearer".to_string(),
is_active: true,
api_formats: None,
2026-04-29 15:46:50 +08:00
auth_type_by_format: None,
allow_auth_channel_mismatch_formats: None,
2026-04-29 15:46:50 +08:00
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
2026-05-21 15:38:10 +08:00
upstream_metadata: None,
decrypted_api_key: "__placeholder__".to_string(),
decrypted_auth_config: Some("{\"refresh_token\":\"rt-1\"}".to_string()),
},
}
}
#[tokio::test]
async fn coordinator_reuses_runtime_cached_refresh_result() {
let refresh_hits = Arc::new(AtomicUsize::new(0));
2026-04-27 23:04:31 +08:00
let refresh_with_entry_hits = Arc::new(AtomicUsize::new(0));
let coordinator =
LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![Arc::new(TestAdapter {
refresh_hits: Arc::clone(&refresh_hits),
2026-04-27 23:04:31 +08:00
refresh_with_entry_hits: Arc::clone(&refresh_with_entry_hits),
})]);
let transport = sample_transport();
let executor = ReqwestLocalOAuthHttpExecutor::new(reqwest::Client::new());
let first = coordinator
.resolve_with_result(&executor, &transport, None, None)
.await
.expect("first resolve should succeed");
assert!(coordinator
.cached_entry(transport.key.id.as_str())
.await
.is_none());
2026-04-28 09:35:23 +08:00
coordinator
.store_cached_entry(
2026-04-28 09:35:23 +08:00
transport.key.id.as_str(),
first
.as_ref()
.and_then(|result| result.refreshed_entry.clone())
.expect("first resolve should provide cached entry"),
)
.await;
let second = coordinator
.resolve_with_result(&executor, &transport, None, None)
.await
.expect("second resolve should succeed");
assert_eq!(refresh_hits.load(Ordering::SeqCst), 1);
assert_eq!(
first,
Some(LocalOAuthResolution {
auth: Some(LocalResolvedOAuthRequestAuth::Header {
name: "authorization".to_string(),
value: "Bearer refreshed-token".to_string(),
}),
refreshed_entry: Some(CachedOAuthEntry {
provider_type: "test-oauth".to_string(),
auth_header_name: "authorization".to_string(),
auth_header_value: "Bearer refreshed-token".to_string(),
expires_at_unix_secs: Some(4_102_444_800),
metadata: None,
source_fingerprint: None,
}),
refresh_in_flight: false,
reused_refresh: false,
distributed_lease: None,
local_refresh_guard: None,
})
);
assert_eq!(
second,
Some(LocalOAuthResolution {
auth: Some(LocalResolvedOAuthRequestAuth::Header {
name: "authorization".to_string(),
value: "Bearer refreshed-token".to_string(),
}),
refreshed_entry: None,
refresh_in_flight: false,
reused_refresh: false,
distributed_lease: None,
local_refresh_guard: None,
})
);
}
#[tokio::test]
async fn coordinator_force_refresh_bypasses_runtime_cache() {
let refresh_hits = Arc::new(AtomicUsize::new(0));
2026-04-27 23:04:31 +08:00
let refresh_with_entry_hits = Arc::new(AtomicUsize::new(0));
let coordinator =
LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![Arc::new(TestAdapter {
refresh_hits: Arc::clone(&refresh_hits),
2026-04-27 23:04:31 +08:00
refresh_with_entry_hits: Arc::clone(&refresh_with_entry_hits),
})]);
let transport = sample_transport();
let executor = ReqwestLocalOAuthHttpExecutor::new(reqwest::Client::new());
let first = coordinator
.resolve_with_result(&executor, &transport, None, None)
.await
.expect("initial resolve should succeed");
2026-04-28 09:35:23 +08:00
coordinator
.store_cached_entry(
2026-04-28 09:35:23 +08:00
transport.key.id.as_str(),
first
.as_ref()
.and_then(|result| result.refreshed_entry.clone())
.expect("first resolve should provide cached entry"),
)
.await;
let forced = coordinator
.force_refresh_with_result(&executor, &transport, None, None)
.await
.expect("forced refresh should succeed");
assert!(first.and_then(|result| result.refreshed_entry).is_some());
assert!(forced.and_then(|result| result.refreshed_entry).is_some());
assert_eq!(refresh_hits.load(Ordering::SeqCst), 2);
2026-04-28 09:35:23 +08:00
assert_eq!(refresh_with_entry_hits.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn memory_only_force_refresh_does_not_reuse_preexisting_cache_as_winner() {
let refresh_hits = Arc::new(AtomicUsize::new(0));
let fingerprint_hits = Arc::new(AtomicUsize::new(0));
let coordinator = Arc::new(LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(MemoryOnlyFencedTestAdapter {
refresh_hits: Arc::clone(&refresh_hits),
fingerprint_hits: Arc::clone(&fingerprint_hits),
}),
]));
let transport = sample_transport();
let executor = ReqwestLocalOAuthHttpExecutor::new(reqwest::Client::new());
coordinator
.store_cached_entry(
transport.key.id.as_str(),
CachedOAuthEntry {
provider_type: "test-oauth".to_string(),
auth_header_name: "authorization".to_string(),
auth_header_value: "Bearer rejected-token".to_string(),
expires_at_unix_secs: Some(4_102_444_800),
metadata: None,
source_fingerprint: Some("preexisting-local-generation".to_string()),
},
)
.await;
let mut forced = coordinator
.force_refresh_with_result_fenced(
&executor,
&transport,
None,
None,
Some("transport-generation"),
)
.await
.expect("memory-only force refresh should succeed")
.expect("memory-only force refresh should resolve");
assert_eq!(refresh_hits.load(Ordering::SeqCst), 1);
assert!(!forced.reused_refresh);
assert!(forced.local_refresh_guard.is_some());
assert_eq!(
forced.auth,
Some(LocalResolvedOAuthRequestAuth::Header {
name: "authorization".to_string(),
value: "Bearer refreshed-token-1".to_string(),
})
);
fingerprint_hits.store(0, Ordering::SeqCst);
let follower_coordinator = Arc::clone(&coordinator);
let follower_transport = transport.clone();
let follower = tokio::spawn(async move {
follower_coordinator
.force_refresh_with_result_fenced(
&ReqwestLocalOAuthHttpExecutor::new(reqwest::Client::new()),
&follower_transport,
None,
None,
Some("transport-generation"),
)
.await
.expect("memory-only follower should succeed")
.expect("memory-only follower should resolve")
});
tokio::time::timeout(Duration::from_secs(1), async {
while fingerprint_hits.load(Ordering::SeqCst) == 0 {
tokio::task::yield_now().await;
}
})
.await
.expect("follower should capture its pre-lock fingerprint");
coordinator
.store_cached_entry(
transport.key.id.as_str(),
forced
.refreshed_entry
.clone()
.expect("leader should provide the memory-only entry"),
)
.await;
forced.local_refresh_guard.take();
let follower = tokio::time::timeout(Duration::from_secs(1), follower)
.await
.expect("follower should unblock after cache publication")
.expect("follower task should join");
assert!(follower.reused_refresh);
assert_eq!(refresh_hits.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn fenced_force_refresh_reuses_the_winner() {
let refresh_hits = Arc::new(AtomicUsize::new(0));
let coordinator = LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![Arc::new(
FencedTestAdapter {
refresh_hits: Arc::clone(&refresh_hits),
fail_refresh: Arc::new(AtomicBool::new(false)),
generation: Arc::new(AtomicUsize::new(1)),
},
)]);
let transport = sample_transport();
let executor = ReqwestLocalOAuthHttpExecutor::new(reqwest::Client::new());
let first = coordinator
.force_refresh_with_result_fenced(
&executor,
&transport,
None,
None,
Some("generation-1"),
)
.await
.expect("first refresh should succeed")
.expect("first refresh should resolve");
coordinator
.store_cached_entry(
transport.key.id.as_str(),
first
.refreshed_entry
.clone()
.expect("first refresh should return an entry to persist"),
)
.await;
let waiter = coordinator
.force_refresh_with_result_fenced(
&executor,
&transport,
None,
None,
Some("generation-1"),
)
.await
.expect("waiter should reuse winner")
.expect("waiter should resolve");
assert!(first.refreshed_entry.is_some());
assert!(waiter.refreshed_entry.is_some());
assert!(waiter.reused_refresh);
assert_eq!(
waiter
.refreshed_entry
.as_ref()
.expect("reused entry")
.auth_header_value,
"fresh-winner-assertion"
);
assert_eq!(refresh_hits.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn refresh_failure_enters_bounded_negative_backoff() {
let refresh_hits = Arc::new(AtomicUsize::new(0));
let fail_refresh = Arc::new(AtomicBool::new(true));
let generation = Arc::new(AtomicUsize::new(1));
let coordinator = LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![Arc::new(
FencedTestAdapter {
refresh_hits: Arc::clone(&refresh_hits),
fail_refresh: Arc::clone(&fail_refresh),
generation: Arc::clone(&generation),
},
)]);
let transport = sample_transport();
let executor = ReqwestLocalOAuthHttpExecutor::new(reqwest::Client::new());
assert!(coordinator
.force_refresh_with_result(&executor, &transport, None, None)
.await
.is_err());
let second = coordinator
.force_refresh_with_result(&executor, &transport, None, None)
.await
.expect_err("second refresh should be backed off");
assert!(second.to_string().contains("temporarily backed off"));
assert_eq!(refresh_hits.load(Ordering::SeqCst), 1);
fail_refresh.store(false, Ordering::SeqCst);
generation.store(2, Ordering::SeqCst);
assert!(coordinator
.force_refresh_with_result(&executor, &transport, None, None)
.await
.expect("new credential generation should bypass old backoff")
.is_some());
assert_eq!(refresh_hits.load(Ordering::SeqCst), 2);
}
}