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

use std::collections::BTreeMap;
use std::fmt;
use std::sync::Arc;
use std::time::{Duration, Instant};

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;

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;
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>,
}

#[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,
    },
    #[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,
        })
    }
}

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
    }

    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,
}

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()),
                Arc::new(VertexServiceAccountRefreshAdapter),
                Arc::new(GenericOAuthRefreshAdapter::default()),
            ],
            cache: Mutex::new(BTreeMap::new()),
            key_locks: Mutex::new(BTreeMap::new()),
            refresh_backoff: Mutex::new(BTreeMap::new()),
        }
    }

    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);
    }

    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,
        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,
        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,
        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
        };
        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 {
            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));
            }
        }
        if let Some(error) = self.backoff_error(key_id, adapter.provider_type()).await {
            return Err(error);
        }

        let key_lock = self.lock_for_key(key_id).await;
        let _key_guard = key_lock.lock().await;

        let cached_entry = self.cached_entry(key_id).await;
        if force_refresh {
            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));
            }
        }
        if let Some(error) = self.backoff_error(key_id, adapter.provider_type()).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 = 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).await;
                            }
                            return Err(error);
                        }
                        None
                    }
                }
            }
            _ => None,
        };

        // 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();
        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).await;
                }
                Self::release_distributed_lease(
                    distributed_lock,
                    distributed_lease.as_ref(),
                    key_id,
                    adapter.provider_type(),
                )
                .await;
                return Err(error);
            }
        };
        // In production the distributed lease is held through the gateway's
        // DB CAS. Do not publish a provisional task before that CAS succeeds;
        // otherwise a waiter could consume an assertion that loses the CAS.
        // Lock-free/test callers retain the historical in-memory behavior.
        if distributed_lease.is_none() {
            self.insert_cached_entry(key_id, refreshed_entry.clone())
                .await;
        }
        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);
        };
        Ok(Some(LocalOAuthResolution::refreshed(
            auth,
            refreshed_entry,
            distributed_lease,
        )))
    }

    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,
    ) -> Option<LocalOAuthRefreshError> {
        let backoff = self.refresh_backoff.lock().await;
        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) {
        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(),
            });
        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,
        }
    }

    fn refreshed(
        auth: LocalResolvedOAuthRequestAuth,
        refreshed_entry: CachedOAuthEntry,
        distributed_lease: Option<RuntimeLockLease>,
    ) -> Self {
        Self {
            auth: Some(auth),
            refreshed_entry: Some(refreshed_entry),
            refresh_in_flight: false,
            reused_refresh: false,
            distributed_lease,
        }
    }

    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,
        }
    }

    fn refresh_in_flight() -> Self {
        Self {
            auth: None,
            refreshed_entry: None,
            refresh_in_flight: true,
            reused_refresh: false,
            distributed_lease: 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)
        || 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 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>,
        refresh_with_entry_hits: Arc<AtomicUsize>,
    }

    #[derive(Debug)]
    struct FencedTestAdapter {
        refresh_hits: Arc<AtomicUsize>,
        fail_refresh: Arc<AtomicBool>,
    }

    #[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,
            entry: Option<&CachedOAuthEntry>,
        ) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError> {
            self.refresh_hits.fetch_add(1, Ordering::SeqCst);
            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("generation-1".to_string()))
        }

        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("generation-2".to_string()),
            }))
        }
    }

    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(),
                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,
                auth_type_by_format: None,
                allow_auth_channel_mismatch_formats: None,

                allowed_models: None,
                capabilities: None,
                rate_multipliers: None,
                global_priority_by_format: None,
                expires_at_unix_secs: None,
                proxy: None,
                fingerprint: None,
                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));
        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),
                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");
        coordinator
            .insert_cached_entry(
                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,
            })
        );
        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,
            })
        );
    }

    #[tokio::test]
    async fn coordinator_force_refresh_bypasses_runtime_cache() {
        let refresh_hits = Arc::new(AtomicUsize::new(0));
        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),
                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");
        coordinator
            .insert_cached_entry(
                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);
        assert_eq!(refresh_with_entry_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)),
            },
        )]);
        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");
        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 coordinator = LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![Arc::new(
            FencedTestAdapter {
                refresh_hits: Arc::clone(&refresh_hits),
                fail_refresh: Arc::clone(&fail_refresh),
            },
        )]);
        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);
        coordinator.invalidate_cached_entry("key-1").await;
        assert!(coordinator
            .force_refresh_with_result(&executor, &transport, None, None)
            .await
            .expect("replacement should refresh immediately")
            .is_some());
        assert_eq!(refresh_hits.load(Ordering::SeqCst), 2);
    }
}