Tighten OAuth refresh consistency

apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/refresh/request.rs

@@ -87,7 +87,7 @@ pub(super) async fn parse_admin_provider_oauth_refresh_request(
         )));
     };
     let Some(transport) = state
-        .read_provider_transport_snapshot(&provider_id, &endpoint.id, &key_id)
+        .read_provider_transport_snapshot_uncached(&provider_id, &endpoint.id, &key_id)
         .await?
     else {
         return Ok(RefreshDispatch::Respond(response::control_error_response(

apps/aether-gateway/src/handlers/admin/request/provider/transport.rs

@@ -4,6 +4,17 @@ use serde_json::{Map, Value};
 use url::Url;
 
 impl<'a> AdminAppState<'a> {
+    pub(crate) async fn read_provider_transport_snapshot_uncached(
+        &self,
+        provider_id: &str,
+        endpoint_id: &str,
+        key_id: &str,
+    ) -> Result<Option<AdminGatewayProviderTransportSnapshot>, GatewayError> {
+        self.app
+            .read_provider_transport_snapshot_uncached(provider_id, endpoint_id, key_id)
+            .await
+    }
+
     pub(crate) async fn read_provider_transport_snapshot(
         &self,
         provider_id: &str,

apps/aether-gateway/src/state/oauth.rs

@@ -491,7 +491,7 @@ impl AppState {
         );
     }
 
-    async fn read_provider_transport_snapshot_uncached(
+    pub(crate) async fn read_provider_transport_snapshot_uncached(
         &self,
         provider_id: &str,
         endpoint_id: &str,
@@ -892,6 +892,16 @@ impl AppState {
                         error = ?err,
                         "gateway local oauth refresh persistence failed"
                     );
+                    let _ = self
+                        .invalidate_local_oauth_refresh_entry(&current_transport.key.id)
+                        .await;
+                } else {
+                    self.oauth_refresh
+                        .store_cached_entry(
+                            current_transport.key.id.trim(),
+                            refreshed_entry.clone(),
+                        )
+                        .await;
                 }
             }
 
@@ -977,7 +987,17 @@ impl AppState {
                         error = ?err,
                         "gateway manual oauth refresh persistence failed"
                     );
+                    let _ = self
+                        .invalidate_local_oauth_refresh_entry(&current_transport.key.id)
+                        .await;
+                    return Err(provider_transport::LocalOAuthRefreshError::InvalidResponse {
+                        provider_type: "gateway",
+                        message: format!("local oauth refresh persistence failed: {err:?}"),
+                    });
                 }
+                self.oauth_refresh
+                    .store_cached_entry(current_transport.key.id.trim(), refreshed_entry.clone())
+                    .await;
                 return Ok(Some(refreshed_entry.clone()));
             }
 

apps/aether-gateway/src/tests/control/admin/oauth.rs

@@ -5149,6 +5149,227 @@ async fn gateway_concurrent_manual_oauth_refresh_uses_rotated_refresh_token_afte
     execution_runtime_handle.abort();
 }
 
+#[tokio::test]
+async fn gateway_manual_oauth_refresh_prefers_fresher_transport_auth_config_over_stale_runtime_cache()
+{
+    let refresh_request_bodies = Arc::new(Mutex::new(Vec::<String>::new()));
+    let refresh_request_bodies_clone = Arc::clone(&refresh_request_bodies);
+    let execution_runtime = Router::new().route(
+        "/v1/execute/sync",
+        any(move |Json(plan): Json<ExecutionPlan>| {
+            let refresh_request_bodies_inner = Arc::clone(&refresh_request_bodies_clone);
+            async move {
+                if plan.request_id == "provider-oauth:local-refresh-token" {
+                    use base64::Engine as _;
+
+                    let body_text = plan
+                        .body
+                        .body_bytes_b64
+                        .as_deref()
+                        .and_then(|body| {
+                            base64::engine::general_purpose::STANDARD.decode(body).ok()
+                        })
+                        .and_then(|body| String::from_utf8(body).ok())
+                        .unwrap_or_default();
+                    refresh_request_bodies_inner
+                        .lock()
+                        .expect("mutex should lock")
+                        .push(body_text.clone());
+
+                    if body_text.contains("refresh_token=old-codex-refresh-token") {
+                        return Json(json!({
+                            "request_id": plan.request_id,
+                            "status_code": 200,
+                            "headers": {
+                                "content-type": "application/json"
+                            },
+                            "body": {
+                                "json_body": {
+                                    "access_token": "cached-codex-access-token",
+                                    "refresh_token": "cached-codex-refresh-token",
+                                    "token_type": "Bearer",
+                                    "expires_in": 1800,
+                                    "scope": "openid email profile offline_access",
+                                    "email": "alice@example.com",
+                                    "account_id": "acct-codex-123",
+                                    "plan_type": "plus"
+                                }
+                            }
+                        }));
+                    }
+
+                    if body_text.contains("refresh_token=fresh-codex-refresh-token") {
+                        return Json(json!({
+                            "request_id": plan.request_id,
+                            "status_code": 200,
+                            "headers": {
+                                "content-type": "application/json"
+                            },
+                            "body": {
+                                "json_body": {
+                                    "access_token": "fresh-codex-access-token-2",
+                                    "refresh_token": "fresh-codex-refresh-token-2",
+                                    "token_type": "Bearer",
+                                    "expires_in": 1800,
+                                    "scope": "openid email profile offline_access",
+                                    "email": "alice@example.com",
+                                    "account_id": "acct-codex-123",
+                                    "plan_type": "plus"
+                                }
+                            }
+                        }));
+                    }
+
+                    Json(json!({
+                        "request_id": plan.request_id,
+                        "status_code": 401,
+                        "headers": {
+                            "content-type": "application/json"
+                        },
+                        "body": {
+                            "json_body": {
+                                "error": {
+                                    "message": "Could not validate your refresh token. Please try signing in again."
+                                }
+                            }
+                        }
+                    }))
+                } else {
+                    Json(json!({
+                        "request_id": plan.request_id,
+                        "status_code": 200,
+                        "headers": {
+                            "content-type": "application/json"
+                        },
+                        "body": {
+                            "json_body": {}
+                        }
+                    }))
+                }
+            }
+        }),
+    );
+
+    let mut provider = sample_provider("provider-codex", "codex", 10);
+    provider.provider_type = "codex".to_string();
+    let endpoint = sample_endpoint(
+        "endpoint-codex-cli",
+        "provider-codex",
+        "openai:responses",
+        "https://chatgpt.com/backend-api/codex",
+    );
+
+    let mut key = sample_key(
+        "key-codex-oauth-stale-cache",
+        "provider-codex",
+        "openai:responses",
+        "stale-codex-access-token",
+    );
+    key.auth_type = "oauth".to_string();
+    key.encrypted_auth_config = Some(
+        encrypt_python_fernet_plaintext(
+            DEVELOPMENT_ENCRYPTION_KEY,
+            r#"{"provider_type":"codex","refresh_token":"old-codex-refresh-token","email":"alice@example.com","account_id":"acct-codex-123","plan_type":"plus","expires_at":1,"updated_at":1700000001}"#,
+        )
+        .expect("auth config ciphertext should build"),
+    );
+
+    let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
+        vec![provider],
+        vec![endpoint],
+        vec![key],
+    ));
+
+    let oauth_refresh =
+        crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
+            Arc::new(
+                crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
+                    .with_token_url_for_tests("codex", "https://oauth.example/oauth/token"),
+            ),
+        ]);
+
+    let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
+    let app_state = build_state_with_execution_runtime_override(execution_runtime_url)
+        .with_data_state_for_tests(
+            GatewayDataState::with_provider_catalog_repository_for_tests(
+                provider_catalog_repository.clone(),
+            )
+            .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
+        )
+        .with_oauth_refresh_coordinator_for_tests(oauth_refresh);
+
+    let stale_transport = app_state
+        .read_provider_transport_snapshot(
+            "provider-codex",
+            "endpoint-codex-cli",
+            "key-codex-oauth-stale-cache",
+        )
+        .await
+        .expect("transport should load")
+        .expect("transport should exist");
+    let cached_entry = app_state
+        .force_local_oauth_refresh_entry(&stale_transport)
+        .await
+        .expect("initial refresh should succeed")
+        .expect("initial refresh should return cached entry");
+    assert_eq!(cached_entry.auth_header_value, "Bearer cached-codex-access-token");
+
+    let mut updated_key = provider_catalog_repository
+        .list_keys_by_ids(&["key-codex-oauth-stale-cache".to_string()])
+        .await
+        .expect("keys should list")
+        .into_iter()
+        .next()
+        .expect("key should exist");
+    updated_key.encrypted_auth_config = Some(
+        encrypt_python_fernet_plaintext(
+            DEVELOPMENT_ENCRYPTION_KEY,
+            r#"{"provider_type":"codex","refresh_token":"fresh-codex-refresh-token","email":"alice@example.com","account_id":"acct-codex-123","plan_type":"plus","expires_at":1,"updated_at":4102444810}"#,
+        )
+        .expect("updated auth config ciphertext should build"),
+    );
+    provider_catalog_repository
+        .update_key(&updated_key)
+        .await
+        .expect("key should update");
+
+    let gateway = build_router_with_state(app_state);
+    let (gateway_url, gateway_handle) = start_server(gateway).await;
+
+    let response = reqwest::Client::new()
+        .post(format!(
+            "{gateway_url}/api/admin/provider-oauth/keys/key-codex-oauth-stale-cache/refresh"
+        ))
+        .header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
+        .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
+        .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
+        .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
+        .send()
+        .await
+        .expect("request should succeed");
+
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let bodies = refresh_request_bodies
+        .lock()
+        .expect("mutex should lock")
+        .clone();
+    assert_eq!(bodies.len(), 2);
+    assert!(
+        bodies[0].contains("refresh_token=old-codex-refresh-token"),
+        "unexpected first refresh body: {}",
+        bodies[0]
+    );
+    assert!(
+        bodies[1].contains("refresh_token=fresh-codex-refresh-token"),
+        "unexpected second refresh body: {}",
+        bodies[1]
+    );
+
+    gateway_handle.abort();
+    execution_runtime_handle.abort();
+}
+
 #[tokio::test]
 async fn gateway_refreshes_admin_provider_oauth_key_locally_via_execution_runtime_key_proxy_before_system_proxy(
 ) {

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
+        let cached = entry.and_then(|cached| Self::auth_config_from_entry(transport, cached));
-            .and_then(|cached| Self::auth_config_from_entry(transport, cached))
+        let transport_auth = Self::auth_config_from_transport(transport);
-            .or_else(|| 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) =
+        let cached_auth_config = entry.and_then(|cached| Self::auth_config_from_entry(transport, cached));
-            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);
-            ("cached_entry", Some(value))
+        let base_auth_config_source = match (
-        } else {
+            base_auth_config.as_ref(),
-            let value = Self::auth_config_from_transport(transport);
+            cached_auth_config.as_ref(),
-            let source = if value.is_some() {
+            transport_auth_config.as_ref(),
+        ) {
+            (Some(selected), Some(cached), Some(transport_auth))
+                if selected == transport_auth && selected != cached =>
+            {
                 "transport_auth_config"
-            } else {
+            }
-                "none"
+            (Some(selected), Some(cached), Some(transport_auth))
-            };
+                if selected == cached && selected != transport_auth =>
-            (source, value)
+            {
+                "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);
     }
 }