feat(gateway): harden provider request execution

Preserve exact request payloads and model client surface and API operation explicitly. Add Anthropic compatibility profiles, bounded stream commitment, and scoped OAuth retry behavior across provider transports.
2026-10-10 19:29:50 +08:00 · 2026-07-27 09:36:31 +08:00
parent 79b70f7b5c
commit 531cf11025
152 changed files with 13984 additions and 2075 deletions
@@ -12,7 +12,7 @@ use aether_runtime_state::{RuntimeLockLease, RuntimeState};
 use async_trait::async_trait;
 use serde_json::Value;
 use thiserror::Error;
-use tokio::sync::Mutex;
+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;
@@ -47,6 +47,35 @@ pub struct LocalOAuthResolution {
    /// 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)]
@@ -369,6 +398,12 @@ pub trait LocalOAuthRefreshAdapter: Send + Sync {
        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,
@@ -388,6 +423,7 @@ pub struct LocalOAuthRefreshCoordinator {
 struct RefreshBackoffState {
    failures: u32,
    retry_after: Instant,
+    refresh_fingerprint: Option<String>,
 }

 impl fmt::Debug for LocalOAuthRefreshCoordinator {
@@ -424,6 +460,19 @@ impl LocalOAuthRefreshCoordinator {
        }
    }

+    /// 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
@@ -530,6 +579,11 @@ impl LocalOAuthRefreshCoordinator {
        } 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()
@@ -548,7 +602,7 @@ impl LocalOAuthRefreshCoordinator {
            return Ok(None);
        }

-        if force_refresh {
+        if force_refresh && shares_refresh_through_transport_persistence {
            if let Some(resolution) = Self::resolve_if_refresh_fence_advanced(
                adapter.as_ref(),
                transport,
@@ -558,25 +612,46 @@ impl LocalOAuthRefreshCoordinator {
                return Ok(Some(resolution));
            }
        }
-        if let Some(error) = self.backoff_error(key_id, adapter.provider_type()).await {
+        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().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(),
-                expected_refresh_fingerprint,
+                winner_fingerprint,
            ) {
                return Ok(Some(resolution));
            }
        }
-        if let Some(error) = self.backoff_error(key_id, adapter.provider_type()).await {
+        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 {
@@ -594,40 +669,48 @@ impl LocalOAuthRefreshCoordinator {
            }
        }

-        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;
+        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);
                            }
-                            return Err(error);
+                            None
                        }
-                        None
                    }
                }
+                _ => None,
            }
-            _ => None,
        };

        // Forced refresh still needs the latest rotated refresh_token as input.
@@ -653,7 +736,8 @@ impl LocalOAuthRefreshCoordinator {
            }
            Err(error) => {
                if adapter.should_backoff_after_error(&error) {
-                    self.record_refresh_failure(key_id).await;
+                    self.record_refresh_failure(key_id, refresh_fingerprint.as_deref())
+                        .await;
                }
                Self::release_distributed_lease(
                    distributed_lock,
@@ -665,14 +749,9 @@ impl LocalOAuthRefreshCoordinator {
                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;
-        }
+        // 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,
@@ -683,10 +762,16 @@ impl LocalOAuthRefreshCoordinator {
            .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,
        )))
    }

@@ -728,8 +813,16 @@ impl LocalOAuthRefreshCoordinator {
        &self,
        key_id: &str,
        provider_type: &'static str,
+        refresh_fingerprint: Option<&str>,
    ) -> Option<LocalOAuthRefreshError> {
-        let backoff = self.refresh_backoff.lock().await;
+        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 {
@@ -742,14 +835,19 @@ impl LocalOAuthRefreshCoordinator {
        })
    }

-    async fn record_refresh_failure(&self, key_id: &str) {
+    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));
@@ -800,6 +898,7 @@ impl LocalOAuthResolution {
            refresh_in_flight: false,
            reused_refresh: false,
            distributed_lease: None,
+            local_refresh_guard: None,
        }
    }

@@ -807,6 +906,7 @@ impl LocalOAuthResolution {
        auth: LocalResolvedOAuthRequestAuth,
        refreshed_entry: CachedOAuthEntry,
        distributed_lease: Option<RuntimeLockLease>,
+        local_refresh_guard: Option<LocalOAuthRefreshCommitGuard>,
    ) -> Self {
        Self {
            auth: Some(auth),
@@ -814,6 +914,7 @@ impl LocalOAuthResolution {
            refresh_in_flight: false,
            reused_refresh: false,
            distributed_lease,
+            local_refresh_guard,
        }
    }

@@ -829,6 +930,7 @@ impl LocalOAuthResolution {
            refresh_in_flight: false,
            reused_refresh: true,
            distributed_lease: None,
+            local_refresh_guard: None,
        }
    }

@@ -839,6 +941,7 @@ impl LocalOAuthResolution {
            refresh_in_flight: true,
            reused_refresh: false,
            distributed_lease: None,
+            local_refresh_guard: None,
        }
    }
 }
@@ -855,6 +958,7 @@ pub fn supports_local_oauth_request_auth_resolution(
 #[cfg(test)]
 mod tests {
    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+    use std::time::Duration;

    use super::super::snapshot::{
        GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
@@ -878,6 +982,13 @@ mod tests {
    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]
@@ -978,7 +1089,12 @@ mod tests {
        ) -> Option<String> {
            entry
                .and_then(|entry| entry.source_fingerprint.clone())
-                .or_else(|| Some("generation-1".to_string()))
+                .or_else(|| {
+                    Some(format!(
+                        "generation-{}",
+                        self.generation.load(Ordering::SeqCst)
+                    ))
+                })
        }

        fn should_backoff_after_error(&self, _error: &LocalOAuthRefreshError) -> bool {
@@ -1004,7 +1120,75 @@ mod tests {
                auth_header_value: "stale-winner-cache-value".to_string(),
                expires_at_unix_secs: None,
                metadata: None,
-                source_fingerprint: Some("generation-2".to_string()),
+                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}")),
            }))
        }
    }
@@ -1082,8 +1266,12 @@ mod tests {
            .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());
        coordinator
-            .insert_cached_entry(
+            .store_cached_entry(
                transport.key.id.as_str(),
                first
                    .as_ref()
@@ -1115,6 +1303,7 @@ mod tests {
                refresh_in_flight: false,
                reused_refresh: false,
                distributed_lease: None,
+                local_refresh_guard: None,
            })
        );
        assert_eq!(
@@ -1128,6 +1317,7 @@ mod tests {
                refresh_in_flight: false,
                reused_refresh: false,
                distributed_lease: None,
+                local_refresh_guard: None,
            })
        );
    }
@@ -1149,7 +1339,7 @@ mod tests {
            .await
            .expect("initial resolve should succeed");
        coordinator
-            .insert_cached_entry(
+            .store_cached_entry(
                transport.key.id.as_str(),
                first
                    .as_ref()
@@ -1168,6 +1358,98 @@ mod tests {
        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));
@@ -1175,6 +1457,7 @@ mod tests {
            FencedTestAdapter {
                refresh_hits: Arc::clone(&refresh_hits),
                fail_refresh: Arc::new(AtomicBool::new(false)),
+                generation: Arc::new(AtomicUsize::new(1)),
            },
        )]);
        let transport = sample_transport();
@@ -1191,6 +1474,15 @@ mod tests {
            .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,
@@ -1221,10 +1513,12 @@ mod tests {
    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();
@@ -1241,11 +1535,11 @@ mod tests {
        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;
+        generation.store(2, Ordering::SeqCst);
        assert!(coordinator
            .force_refresh_with_result(&executor, &transport, None, None)
            .await
-            .expect("replacement should refresh immediately")
+            .expect("new credential generation should bypass old backoff")
            .is_some());
        assert_eq!(refresh_hits.load(Ordering::SeqCst), 2);
    }