fix(gateway): evict pooled affinity after sibling key failures

apps/aether-gateway/src/orchestration/effects.rs

@@ -243,6 +243,35 @@ fn local_scheduler_affinity_target(plan: &ExecutionPlan) -> Option<SchedulerAffi
     })
 }
 
+async fn local_execution_plan_uses_pool(state: &AppState, plan: &ExecutionPlan) -> bool {
+    let Ok(Some(transport)) = state
+        .read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
+        .await
+    else {
+        return false;
+    };
+
+    admin_provider_pool_config_from_config_value(transport.provider.config.as_ref()).is_some()
+}
+
+async fn local_scheduler_affinity_matches_failed_target(
+    state: &AppState,
+    plan: &ExecutionPlan,
+    cached_target: &SchedulerAffinityTarget,
+    failed_target: &SchedulerAffinityTarget,
+) -> bool {
+    if cached_target == failed_target {
+        return true;
+    }
+    if cached_target.provider_id != failed_target.provider_id
+        || cached_target.endpoint_id != failed_target.endpoint_id
+    {
+        return false;
+    }
+
+    local_execution_plan_uses_pool(state, plan).await
+}
+
 async fn scheduler_cache_affinity_enabled(state: &AppState) -> bool {
     match read_scheduler_ordering_config(state).await {
         Ok(config) => config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity,
@@ -363,10 +392,18 @@ async fn record_attempt_failure_effect(
         let Some(failed_target) = local_scheduler_affinity_target(context.plan) else {
             return;
         };
-        if state
-            .read_scheduler_affinity_target(&cache_key, SCHEDULER_AFFINITY_TTL)
-            .as_ref()
-            == Some(&failed_target)
+        let Some(cached_target) =
+            state.read_scheduler_affinity_target(&cache_key, SCHEDULER_AFFINITY_TTL)
+        else {
+            return;
+        };
+        if local_scheduler_affinity_matches_failed_target(
+            state,
+            context.plan,
+            &cached_target,
+            &failed_target,
+        )
+        .await
         {
             let _ = state.remove_scheduler_affinity_cache_entry(&cache_key);
         }
@@ -1219,6 +1256,20 @@ mod tests {
         .expect("provider should build")
     }
 
+    fn sample_pool_health_provider() -> StoredProviderCatalogProvider {
+        sample_health_provider().with_transport_fields(
+            true,
+            false,
+            false,
+            None,
+            None,
+            None,
+            None,
+            None,
+            Some(json!({"pool_advanced": {}})),
+        )
+    }
+
     fn sample_health_endpoint() -> StoredProviderCatalogEndpoint {
         StoredProviderCatalogEndpoint::new(
             "ep-1".to_string(),
@@ -1281,6 +1332,20 @@ mod tests {
             )
     }
 
+    fn pool_health_state() -> AppState {
+        let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
+            vec![sample_pool_health_provider()],
+            vec![sample_health_endpoint()],
+            vec![sample_health_key()],
+        ));
+        AppState::new()
+            .expect("gateway state should build")
+            .with_data_state_for_tests(
+                GatewayDataState::with_provider_catalog_repository_for_tests(repository)
+                    .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
+            )
+    }
+
     fn health_state_with_key(key: StoredProviderCatalogKey) -> AppState {
         let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
             vec![sample_health_provider()],
@@ -1488,6 +1553,93 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn attempt_failure_keeps_scheduler_affinity_for_non_pool_sibling_key() {
+        let state = health_state();
+        let plan = sample_plan();
+        let report_context = json!({
+            "api_key_id": "api-key-1",
+            "client_api_format": "openai:chat",
+            "model": "gpt-5",
+        });
+        let cache_key =
+            build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
+                .expect("scheduler affinity cache key should build");
+        let affinity_target = SchedulerAffinityTarget {
+            provider_id: "prov-1".to_string(),
+            endpoint_id: "ep-1".to_string(),
+            key_id: "key-2".to_string(),
+        };
+
+        state.remember_scheduler_affinity_target(
+            &cache_key,
+            affinity_target.clone(),
+            SCHEDULER_AFFINITY_TTL,
+            16,
+        );
+
+        apply_local_execution_effect(
+            &state,
+            LocalExecutionEffectContext {
+                plan: &plan,
+                report_context: Some(&report_context),
+            },
+            LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
+                status_code: 524,
+                classification: LocalFailoverClassification::RetryUpstreamFailure,
+            }),
+        )
+        .await;
+
+        assert_eq!(
+            state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
+            Some(affinity_target)
+        );
+    }
+
+    #[tokio::test]
+    async fn attempt_failure_invalidates_scheduler_affinity_for_same_pool_candidate() {
+        let state = pool_health_state();
+        let plan = sample_plan();
+        let report_context = json!({
+            "api_key_id": "api-key-1",
+            "client_api_format": "openai:chat",
+            "model": "gpt-5",
+        });
+        let cache_key =
+            build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
+                .expect("scheduler affinity cache key should build");
+
+        state.remember_scheduler_affinity_target(
+            &cache_key,
+            SchedulerAffinityTarget {
+                provider_id: "prov-1".to_string(),
+                endpoint_id: "ep-1".to_string(),
+                key_id: "key-2".to_string(),
+            },
+            SCHEDULER_AFFINITY_TTL,
+            16,
+        );
+
+        apply_local_execution_effect(
+            &state,
+            LocalExecutionEffectContext {
+                plan: &plan,
+                report_context: Some(&report_context),
+            },
+            LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
+                status_code: 524,
+                classification: LocalFailoverClassification::RetryUpstreamFailure,
+            }),
+        )
+        .await;
+
+        assert_eq!(
+            state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
+            None
+        );
+    }
+
     #[tokio::test]
     async fn attempt_failure_keeps_scheduler_affinity_for_non_failure_status() {
         let state = AppState::new().expect("gateway state should build");