Merge pull request #534 from RWDai/fix/issue-528-cache-affinity-health

fix(gateway): preserve cache affinity during health updates

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

@@ -615,6 +615,21 @@ impl AppState {
         Ok(updated)
     }
 
+    pub(crate) async fn update_provider_catalog_key_runtime_state(
+        &self,
+        key: &provider_catalog::StoredProviderCatalogKey,
+    ) -> Result<Option<provider_catalog::StoredProviderCatalogKey>, GatewayError> {
+        let updated = self
+            .data
+            .update_provider_catalog_key(key)
+            .await
+            .map_err(|err| GatewayError::Internal(err.to_string()))?;
+        if updated.is_some() {
+            self.invalidate_provider_health_routing_caches();
+        }
+        Ok(updated)
+    }
+
     pub(crate) async fn update_provider_catalog_key_upstream_metadata(
         &self,
         key_id: &str,
@@ -806,7 +821,7 @@ impl AppState {
             .await
             .map_err(|err| GatewayError::Internal(err.to_string()))?;
         if updated {
-            self.invalidate_provider_routing_caches();
+            self.invalidate_provider_health_routing_caches();
         }
         Ok(updated)
     }
@@ -930,4 +945,47 @@ mod tests {
             .expect("provider transport should exist after update");
         assert!(snapshot.provider.keep_priority_on_conversion);
     }
+
+    #[tokio::test]
+    async fn provider_catalog_health_update_keeps_scheduler_affinity_cache() {
+        let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
+            vec![sample_provider()],
+            vec![sample_endpoint()],
+            vec![sample_key()],
+        ));
+        let state = AppState::new()
+            .expect("app state should build")
+            .with_data_state_for_tests(
+                GatewayDataState::with_provider_catalog_repository_for_tests(repository)
+                    .with_encryption_key_for_tests("test-encryption-key"),
+            );
+
+        let cache_key = "scheduler_affinity:api-key-1:openai:chat:gpt-5";
+        let ttl = Duration::from_secs(300);
+        let target = SchedulerAffinityTarget {
+            provider_id: "provider-1".to_string(),
+            endpoint_id: "endpoint-1".to_string(),
+            key_id: "key-1".to_string(),
+        };
+        state.remember_scheduler_affinity_target(cache_key, target.clone(), ttl, 128);
+        let initial_epoch = state.scheduler_affinity_epoch();
+
+        let health_by_format = serde_json::json!({
+            "openai:chat": {
+                "last_success_at_unix_secs": 1,
+                "consecutive_failures": 0
+            }
+        });
+        let updated = state
+            .update_provider_catalog_key_health_state("key-1", true, Some(&health_by_format), None)
+            .await
+            .expect("key health update should succeed");
+
+        assert!(updated);
+        assert_eq!(state.scheduler_affinity_epoch(), initial_epoch);
+        assert_eq!(
+            state.read_scheduler_affinity_target(cache_key, ttl),
+            Some(target)
+        );
+    }
 }

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

@@ -572,6 +572,11 @@ impl AppState {
         self.invalidate_scheduler_affinity_cache();
     }
 
+    pub(crate) fn invalidate_provider_health_routing_caches(&self) {
+        self.data.clear_minimal_candidate_selection_cache();
+        self.clear_provider_transport_snapshot_cache();
+    }
+
     pub(crate) fn invalidate_auth_context_cache(&self) {
         self.auth_context_cache.clear();
     }