fix(provider): harden Agent Identity OAuth lifecycle

apps/aether-gateway/src/execution_runtime/oauth_retry.rs

@@ -1,6 +1,7 @@
 use aether_contracts::ExecutionPlan;
 use tracing::warn;
 
+use crate::state::AgentIdentityAuthConfigFence;
 use crate::{provider_transport::LocalOAuthRefreshError, AppState};
 
 pub(crate) async fn refresh_oauth_plan_auth_for_retry(
@@ -13,8 +14,11 @@ pub(crate) async fn refresh_oauth_plan_auth_for_retry(
     if !status_may_be_oauth_invalid(status_code, response_text) {
         return false;
     }
-    let access_token_invalid_proven =
-        status_proves_access_token_invalid(status_code, response_text);
+    let request_authorization = execution_plan_authorization(plan);
+    let request_uses_agent_identity = request_authorization
+        .is_some_and(aether_provider_transport::is_codex_agent_identity_authorization);
+    let access_token_invalid_proven = !request_uses_agent_identity
+        && status_proves_access_token_invalid(status_code, response_text);
 
     let transport = match state
         .read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
@@ -37,11 +41,44 @@ pub(crate) async fn refresh_oauth_plan_auth_for_retry(
         }
     };
 
-    if aether_provider_transport::is_codex_agent_identity_transport(&transport)
-        && !aether_provider_transport::is_codex_agent_identity_invalid_task_response(
-            status_code,
-            response_text,
-        )
+    let current_uses_agent_identity =
+        aether_provider_transport::is_codex_agent_identity_transport(&transport);
+    if request_uses_agent_identity {
+        if !current_uses_agent_identity
+            || !aether_provider_transport::is_codex_agent_identity_invalid_task_response(
+                status_code,
+                response_text,
+            )
+            || !request_authorization.is_some_and(|authorization| {
+                aether_provider_transport::codex_agent_identity_authorization_matches_transport(
+                    &transport,
+                    authorization,
+                )
+            })
+        {
+            return false;
+        }
+        if !matches!(
+            state
+                .capture_agent_identity_auth_config_fence(&transport)
+                .await,
+            Ok(AgentIdentityAuthConfigFence::Current(_))
+        ) {
+            return false;
+        }
+    } else if current_uses_agent_identity {
+        // A bearer-token response cannot authorize refreshing an Agent Identity
+        // installed under the same key id while the request was in flight.
+        return false;
+    } else if transport
+        .provider
+        .provider_type
+        .trim()
+        .eq_ignore_ascii_case("codex")
+        && transport.key.auth_type.trim().eq_ignore_ascii_case("oauth")
+        && !request_authorization.is_some_and(|authorization| {
+            bearer_authorization_matches_transport(authorization, &transport)
+        })
     {
         return false;
     }
@@ -120,6 +157,26 @@ pub(crate) async fn refresh_oauth_plan_auth_for_retry(
     }
 }
 
+fn execution_plan_authorization(plan: &ExecutionPlan) -> Option<&str> {
+    plan.headers
+        .iter()
+        .find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
+        .map(|(_, value)| value.as_str())
+}
+
+fn bearer_authorization_matches_transport(
+    authorization: &str,
+    transport: &aether_provider_transport::GatewayProviderTransportSnapshot,
+) -> bool {
+    let current_token = transport.key.decrypted_api_key.trim();
+    !current_token.is_empty()
+        && authorization
+            .trim()
+            .strip_prefix("Bearer ")
+            .map(str::trim)
+            .is_some_and(|token| token == current_token)
+}
+
 fn status_may_be_oauth_invalid(status_code: u16, response_text: Option<&str>) -> bool {
     if status_code == 401 {
         return true;