mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 18:07:47 +08:00
Merge remote-tracking branch 'origin/main' into fix/gemini-cli-v1internal
# Conflicts: # apps/aether-gateway/src/ai_serving/transport.rs # apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/parse.rs # apps/aether-gateway/src/handlers/shared/catalog.rs # crates/aether-admin/src/provider/quota.rs # crates/aether-model-fetch/src/strategy.rs # crates/aether-provider-pool/src/lib.rs # crates/aether-provider-pool/src/service.rs # crates/aether-provider-transport/src/provider_types.rs # frontend/src/features/providers/components/ProviderDetailDrawer.vue # frontend/src/utils/__tests__/providerKeyQuota.spec.ts # frontend/src/utils/providerKeyQuota.ts # frontend/src/views/admin/PoolManagement.vue
This commit is contained in:
@@ -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)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -49,6 +49,7 @@ use crate::maintenance::spawn_pending_cleanup_worker;
|
||||
use crate::maintenance::spawn_pool_monitor_worker;
|
||||
use crate::maintenance::spawn_pool_score_rebuild_worker;
|
||||
use crate::maintenance::spawn_provider_checkin_worker;
|
||||
use crate::maintenance::spawn_provider_quota_alert_worker;
|
||||
use crate::maintenance::spawn_proxy_node_metrics_cleanup_worker;
|
||||
use crate::maintenance::spawn_proxy_node_stale_cleanup_worker;
|
||||
use crate::maintenance::spawn_proxy_upgrade_rollout_worker;
|
||||
@@ -89,24 +90,17 @@ impl AppState {
|
||||
fn usage_worker_queue_for(
|
||||
runtime_state: &Arc<RuntimeState>,
|
||||
) -> Option<Arc<dyn RuntimeQueueStore>> {
|
||||
if runtime_state.is_redis() {
|
||||
let queue: Arc<dyn RuntimeQueueStore> = runtime_state.clone();
|
||||
Some(queue)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
let queue: Arc<dyn RuntimeQueueStore> = runtime_state.clone();
|
||||
Some(queue)
|
||||
}
|
||||
|
||||
fn spawn_scheduler_affinity_redis_write(
|
||||
fn spawn_scheduler_affinity_runtime_write(
|
||||
&self,
|
||||
cache_key: &str,
|
||||
target: &SchedulerAffinityTarget,
|
||||
ttl: Duration,
|
||||
epoch: u64,
|
||||
) {
|
||||
if self.runtime_state.is_memory() {
|
||||
return;
|
||||
}
|
||||
let Ok(handle) = tokio::runtime::Handle::try_current() else {
|
||||
return;
|
||||
};
|
||||
@@ -578,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();
|
||||
}
|
||||
@@ -1088,7 +1087,7 @@ impl AppState {
|
||||
if self.scheduler_affinity_epoch() != epoch {
|
||||
return false;
|
||||
}
|
||||
self.spawn_scheduler_affinity_redis_write(cache_key, &target, ttl, epoch);
|
||||
self.spawn_scheduler_affinity_runtime_write(cache_key, &target, ttl, epoch);
|
||||
self.scheduler_affinity_cache.insert_for_epoch(
|
||||
cache_key.to_string(),
|
||||
target,
|
||||
@@ -1210,6 +1209,10 @@ impl AppState {
|
||||
crate::task_runtime::TASK_KEY_PROVIDER_CHECKIN,
|
||||
spawn_provider_checkin_worker(self.clone()),
|
||||
);
|
||||
supervise_worker(
|
||||
crate::task_runtime::TASK_KEY_PROVIDER_QUOTA_ALERT,
|
||||
spawn_provider_quota_alert_worker(self.clone()),
|
||||
);
|
||||
supervise_worker(
|
||||
crate::task_runtime::TASK_KEY_OAUTH_TOKEN_REFRESH,
|
||||
spawn_oauth_token_refresh_worker(self.clone()),
|
||||
|
||||
Reference in New Issue
Block a user