Tighten OAuth refresh consistency

crates/aether-provider-transport/src/generic_oauth.rs

@@ -126,14 +126,35 @@ impl GenericOAuthRefreshAdapter {
             .cloned()
     }
 
+    fn auth_config_updated_at(auth_config: &Value) -> Option<u64> {
+        auth_config
+            .as_object()
+            .and_then(|object| object.get("updated_at"))
+            .and_then(|value| parse_u64_value(Some(value)))
+    }
+
     fn base_auth_config(
         &self,
         transport: &GatewayProviderTransportSnapshot,
         entry: Option<&CachedOAuthEntry>,
     ) -> Option<Value> {
-        entry
-            .and_then(|cached| Self::auth_config_from_entry(transport, cached))
-            .or_else(|| Self::auth_config_from_transport(transport))
+        let cached = entry.and_then(|cached| Self::auth_config_from_entry(transport, cached));
+        let transport_auth = Self::auth_config_from_transport(transport);
+
+        match (cached, transport_auth) {
+            (Some(cached), Some(transport_auth)) => {
+                let cached_updated_at = Self::auth_config_updated_at(&cached);
+                let transport_updated_at = Self::auth_config_updated_at(&transport_auth);
+                if transport_updated_at > cached_updated_at {
+                    Some(transport_auth)
+                } else {
+                    Some(cached)
+                }
+            }
+            (Some(cached), None) => Some(cached),
+            (None, Some(transport_auth)) => Some(transport_auth),
+            (None, None) => None,
+        }
     }
 
     fn resolve_direct_header(
@@ -257,18 +278,28 @@ impl LocalOAuthRefreshAdapter for GenericOAuthRefreshAdapter {
         else {
             return Ok(None);
         };
-        let (base_auth_config_source, base_auth_config) = if let Some(value) =
-            entry.and_then(|cached| Self::auth_config_from_entry(transport, cached))
-        {
-            ("cached_entry", Some(value))
-        } else {
-            let value = Self::auth_config_from_transport(transport);
-            let source = if value.is_some() {
+        let cached_auth_config = entry.and_then(|cached| Self::auth_config_from_entry(transport, cached));
+        let transport_auth_config = Self::auth_config_from_transport(transport);
+        let base_auth_config = self.base_auth_config(transport, entry);
+        let base_auth_config_source = match (
+            base_auth_config.as_ref(),
+            cached_auth_config.as_ref(),
+            transport_auth_config.as_ref(),
+        ) {
+            (Some(selected), Some(cached), Some(transport_auth))
+                if selected == transport_auth && selected != cached =>
+            {
                 "transport_auth_config"
-            } else {
-                "none"
-            };
-            (source, value)
+            }
+            (Some(selected), Some(cached), Some(transport_auth))
+                if selected == cached && selected != transport_auth =>
+            {
+                "cached_entry"
+            }
+            (Some(_), Some(_), Some(_)) => "cached_entry",
+            (Some(_), Some(_), None) => "cached_entry",
+            (Some(_), None, Some(_)) => "transport_auth_config",
+            _ => "none",
         };
         let mut metadata = base_auth_config
             .and_then(|value| value.as_object().cloned())

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

@@ -231,6 +231,10 @@ impl LocalOAuthRefreshCoordinator {
         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 {
         self.cache.lock().await.remove(key_id).is_some()
     }
@@ -374,8 +378,6 @@ impl LocalOAuthRefreshCoordinator {
         let Some(refreshed_entry) = refresh_result? else {
             return Ok(None);
         };
-        self.insert_cached_entry(key_id, refreshed_entry.clone())
-            .await;
         Ok(adapter
             .resolve_cached(transport, &refreshed_entry)
             .map(|auth| LocalOAuthResolution::resolved(auth, Some(refreshed_entry))))
@@ -567,6 +569,15 @@ mod tests {
             .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
@@ -619,6 +630,15 @@ mod tests {
             .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
@@ -627,6 +647,6 @@ mod tests {
         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), 0);
+        assert_eq!(refresh_with_entry_hits.load(Ordering::SeqCst), 1);
     }
 }