Add usage queue worker autoscaling

2026-10-04 16:37:46 +08:00 · 2026-06-26 14:02:57 +08:00
parent 6c5e70ccb1
commit 7e9424008f
21 changed files with 1801 additions and 85 deletions
@@ -15,7 +15,8 @@ use self::plan::{build_codex_quota_request_spec, execute_codex_quota_plan};
 use super::shared::{
    build_quota_snapshot_payload, extract_execution_error_message,
    oauth_refresh_auto_removed_result, persist_provider_quota_refresh_state,
-    provider_auto_remove_banned_keys, quota_key_auto_removed, quota_refresh_success_invalid_state,
+    provider_auto_remove_banned_keys, provider_auto_remove_quota_exhausted_keys,
+    quota_key_auto_removed, quota_refresh_success_invalid_state,
    should_auto_remove_structured_reason, ProviderQuotaExecutionOutcome,
 };
 use crate::handlers::admin::request::AdminAppState;
@@ -61,6 +62,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
    proxy_override: Option<ProxySnapshot>,
 ) -> Result<Option<serde_json::Value>, GatewayError> {
    let auto_remove_abnormal_keys = provider_auto_remove_banned_keys(provider.config.as_ref());
+    let auto_remove_quota_exhausted_keys =
+        provider_auto_remove_quota_exhausted_keys(provider.config.as_ref());
    let mut results = Vec::new();
    let mut success_count = 0usize;
    let mut failed_count = 0usize;
@@ -312,7 +315,7 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
            }));
            continue;
        }
-        let auto_removed = if auto_remove_candidate {
+        let auto_removed_hard_banned = if auto_remove_candidate {
            state
                .cleanup_provider_catalog_key_if_current(provider, &key.id, |latest_key| {
                    should_auto_remove_structured_reason(latest_key.oauth_invalid_reason.as_deref())
@@ -321,10 +324,28 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
        } else {
            false
        };
-        if auto_removed {
+        if auto_removed_hard_banned {
            auto_removed_count += 1;
            auto_removed_hard_banned_count += 1;
        }
+        let auto_removed_quota_exhausted =
+            if !auto_removed_hard_banned && auto_remove_quota_exhausted_keys {
+                state
+                    .cleanup_provider_catalog_key_if_current(provider, &key.id, |latest_key| {
+                        aether_admin::provider::pool::admin_pool_key_account_quota_exhausted(
+                            latest_key,
+                            provider.provider_type.as_str(),
+                        )
+                    })
+                    .await?
+            } else {
+                false
+            };
+        if auto_removed_quota_exhausted {
+            auto_removed_count += 1;
+            status = "quota_exhausted".to_string();
+        }
+        let auto_removed = auto_removed_hard_banned || auto_removed_quota_exhausted;
        let refresh_fixed =
            status == "success" && had_oauth_refresh_issue && oauth_invalid_reason.is_none();
        if refresh_fixed {
@@ -370,8 +391,13 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
        }
        if auto_removed {
            payload.insert("auto_removed".to_string(), json!(true));
+        }
+        if auto_removed_hard_banned {
            payload.insert("auto_removed_hard_banned".to_string(), json!(true));
        }
+        if auto_removed_quota_exhausted {
+            payload.insert("auto_removed_quota_exhausted".to_string(), json!(true));
+        }
        if refresh_fixed {
            payload.insert("refresh_fixed".to_string(), json!(true));
        }
@@ -6,8 +6,8 @@ use self::plan::execute_kiro_quota_plan;
 use super::shared::{
    build_quota_snapshot_payload, extract_execution_error_message,
    oauth_refresh_auto_removed_result, persist_provider_quota_refresh_state,
-    persist_quota_oauth_refresh_failure_state, quota_refresh_success_invalid_state,
-    ProviderQuotaExecutionOutcome,
+    persist_quota_oauth_refresh_failure_state, provider_auto_remove_quota_exhausted_keys,
+    quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
 };
 use crate::handlers::admin::request::{AdminAppState, AdminLocalOAuthRefreshError};
 use crate::GatewayError;
@@ -97,6 +97,8 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
    let mut success_count = 0usize;
    let mut failed_count = 0usize;
    let mut auto_removed_count = 0usize;
+    let auto_remove_quota_exhausted_keys =
+        provider_auto_remove_quota_exhausted_keys(provider.config.as_ref());

    for key in keys {
        let transport = match state
@@ -304,6 +306,23 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
            continue;
        }

+        let auto_removed_quota_exhausted = if auto_remove_quota_exhausted_keys {
+            state
+                .cleanup_provider_catalog_key_if_current(provider, &key.id, |latest_key| {
+                    aether_admin::provider::pool::admin_pool_key_account_quota_exhausted(
+                        latest_key,
+                        provider.provider_type.as_str(),
+                    )
+                })
+                .await?
+        } else {
+            false
+        };
+        if auto_removed_quota_exhausted {
+            auto_removed_count += 1;
+            status = "quota_exhausted".to_string();
+        }
+
        if status == "success" {
            success_count += 1;
        } else {
@@ -331,6 +350,10 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
        ) {
            payload.insert("quota_snapshot".to_string(), quota_snapshot);
        }
+        if auto_removed_quota_exhausted {
+            payload.insert("auto_removed".to_string(), json!(true));
+            payload.insert("auto_removed_quota_exhausted".to_string(), json!(true));
+        }
        results.push(serde_json::Value::Object(payload));
    }

@@ -67,6 +67,12 @@ pub(crate) fn provider_auto_remove_banned_keys(config: Option<&serde_json::Value
    admin_provider_quota_pure::provider_auto_remove_banned_keys(config)
 }

+pub(crate) fn provider_auto_remove_quota_exhausted_keys(
+    config: Option<&serde_json::Value>,
+) -> bool {
+    admin_provider_quota_pure::provider_auto_remove_quota_exhausted_keys(config)
+}
+
 pub(super) fn should_auto_remove_structured_reason(reason: Option<&str>) -> bool {
    admin_provider_quota_pure::should_auto_remove_structured_reason(reason)
 }
@@ -16,6 +16,8 @@ use std::collections::{BTreeMap, BTreeSet};
 const FIXED_PROVIDER_TEMPLATE_METADATA_KEY: &str = "_aether_fixed_provider_template";
 const OVERRIDE_BODY_RULES: &str = "body_rules";
 const OVERRIDE_FORMAT_ACCEPTANCE_CONFIG: &str = "format_acceptance_config";
+const OVERRIDE_BASE_URL: &str = "base_url";
+const OVERRIDE_CUSTOM_PATH: &str = "custom_path";
 const OVERRIDE_HEADER_RULES: &str = "header_rules";
 const OVERRIDE_IS_ACTIVE: &str = "is_active";
 const OVERRIDE_MAX_RETRIES: &str = "max_retries";
@@ -158,6 +160,20 @@ pub(crate) fn apply_admin_fixed_provider_endpoint_template_overrides(
        .unwrap_or_else(|| managed_fixed_provider_endpoint_metadata(template, endpoint_template));
    let mut overrides = metadata.overrides.clone();

+    sync_override_if_changed(
+        &mut overrides,
+        OVERRIDE_BASE_URL,
+        &existing_endpoint.base_url,
+        &updated_endpoint.base_url,
+        &defaults.base_url,
+    );
+    sync_override_if_changed(
+        &mut overrides,
+        OVERRIDE_CUSTOM_PATH,
+        &existing_endpoint.custom_path,
+        &updated_endpoint.custom_path,
+        &defaults.custom_path,
+    );
    sync_override_if_changed(
        &mut overrides,
        OVERRIDE_HEADER_RULES,
@@ -249,9 +265,13 @@ fn reconcile_fixed_provider_endpoint(
    updated.api_format = defaults.api_format.clone();
    updated.api_family = Some(defaults.api_family.clone());
    updated.endpoint_kind = Some(defaults.endpoint_kind.clone());
-    updated.base_url = defaults.base_url;
-    updated.custom_path = defaults.custom_path;

+    if !metadata.overrides.contains(OVERRIDE_BASE_URL) {
+        updated.base_url = defaults.base_url;
+    }
+    if !metadata.overrides.contains(OVERRIDE_CUSTOM_PATH) {
+        updated.custom_path = defaults.custom_path;
+    }
    if !metadata.overrides.contains(OVERRIDE_HEADER_RULES) {
        updated.header_rules = defaults.header_rules;
    }
@@ -534,3 +554,74 @@ fn sync_override_if_changed<T>(
    }
    sync_override(overrides, key, actual, desired);
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{
+        apply_admin_fixed_provider_endpoint_template_overrides, fixed_provider_endpoint_metadata,
+        reconcile_fixed_provider_endpoint,
+    };
+    use aether_data_contracts::repository::provider_catalog::{
+        StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
+    };
+    use aether_provider_transport::provider_types::fixed_provider_template;
+
+    fn sample_codex_provider() -> StoredProviderCatalogProvider {
+        StoredProviderCatalogProvider::new(
+            "provider-codex".to_string(),
+            "Codex".to_string(),
+            None,
+            "codex".to_string(),
+        )
+        .expect("provider should build")
+    }
+
+    fn sample_codex_endpoint(base_url: &str) -> StoredProviderCatalogEndpoint {
+        StoredProviderCatalogEndpoint::new(
+            "endpoint-codex-responses".to_string(),
+            "provider-codex".to_string(),
+            "openai:responses".to_string(),
+            Some("openai".to_string()),
+            Some("responses".to_string()),
+            true,
+        )
+        .expect("endpoint should build")
+        .with_transport_fields(
+            base_url.to_string(),
+            None,
+            None,
+            Some(2),
+            None,
+            None,
+            None,
+            None,
+        )
+        .expect("endpoint transport should build")
+    }
+
+    #[test]
+    fn fixed_provider_endpoint_reconcile_preserves_base_url_override() {
+        let provider = sample_codex_provider();
+        let template = fixed_provider_template("codex").expect("codex template should exist");
+        let endpoint_template = template
+            .endpoints
+            .iter()
+            .find(|endpoint| endpoint.api_format == "openai:responses")
+            .expect("responses endpoint template should exist");
+
+        let existing = sample_codex_endpoint("https://chatgpt.com/backend-api/codex");
+        let mut updated = existing.clone();
+        updated.base_url = "http://127.0.0.1:18181/v1".to_string();
+        apply_admin_fixed_provider_endpoint_template_overrides(&provider, &existing, &mut updated)
+            .expect("override metadata should apply");
+
+        let metadata = fixed_provider_endpoint_metadata(&updated)
+            .expect("fixed provider metadata should exist");
+        assert!(metadata.overrides.contains("base_url"));
+
+        let reconciled =
+            reconcile_fixed_provider_endpoint(&provider, &updated, template, endpoint_template)
+                .expect("endpoint should reconcile");
+        assert_eq!(reconciled.base_url, "http://127.0.0.1:18181/v1");
+    }
+}
@@ -315,14 +315,6 @@ impl<'a> AdminAppState<'a> {
            return Err("Gemini CLI Endpoint 由系统固定管理，不允许修改".to_string());
        }

-        if self.provider_type_is_fixed(&provider.provider_type)
-            && (fields.contains("base_url") || fields.contains("custom_path"))
-        {
-            return Err(
-                "固定类型 Provider 的 Endpoint 不允许修改 base_url/custom_path".to_string(),
-            );
-        }
-
        let mut update_fields = admin_provider_endpoints_pure::AdminProviderEndpointUpdateFields {
            base_url: payload.base_url,
            custom_path: payload.custom_path,
@@ -1,4 +1,8 @@
 use super::*;
+use crate::ai_serving::provider_key_pool_score_scope;
+use aether_data_contracts::repository::pool_scores::{
+    ListPoolMemberScoresQuery, PoolMemberHardState, POOL_KIND_PROVIDER_KEY_POOL,
+};
 use aether_data_contracts::repository::provider_catalog::{
    StoredProviderCatalogKey, StoredProviderCatalogProvider,
 };
@@ -266,6 +270,100 @@ impl<'a> AdminAppState<'a> {
        Ok(affected)
    }

+    pub(crate) async fn cleanup_quota_exhausted_provider_catalog_keys(
+        &self,
+        provider: &StoredProviderCatalogProvider,
+        provider_type: &str,
+    ) -> Result<usize, GatewayError> {
+        use aether_admin::provider::pool as admin_provider_pool_pure;
+
+        let keys = self
+            .list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
+            .await?;
+        if keys.is_empty() {
+            return Ok(0);
+        }
+
+        let known_key_ids = keys
+            .iter()
+            .map(|key| key.id.as_str())
+            .collect::<std::collections::BTreeSet<_>>();
+        let mut exhausted_key_ids = keys
+            .iter()
+            .filter(|key| {
+                admin_provider_pool_pure::admin_pool_key_account_quota_exhausted(key, provider_type)
+            })
+            .map(|key| key.id.clone())
+            .collect::<std::collections::BTreeSet<_>>();
+
+        if self.app().data.has_pool_score_reader() {
+            let scope = provider_key_pool_score_scope();
+            let page_size = 10_000usize;
+            let mut offset = 0usize;
+            loop {
+                let scores = self
+                    .app()
+                    .data
+                    .list_pool_member_scores(&ListPoolMemberScoresQuery {
+                        pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
+                        pool_id: provider.id.clone(),
+                        capability: Some(scope.capability.clone()),
+                        scope_kind: Some(scope.scope_kind.clone()),
+                        scope_id: scope.scope_id.clone(),
+                        hard_states: vec![PoolMemberHardState::QuotaExhausted],
+                        probe_statuses: None,
+                        offset,
+                        limit: page_size,
+                    })
+                    .await
+                    .map_err(|err| GatewayError::Internal(err.to_string()))?;
+                if scores.is_empty() {
+                    break;
+                }
+                let page_len = scores.len();
+                for score in scores {
+                    if known_key_ids.contains(score.member_id.as_str()) {
+                        exhausted_key_ids.insert(score.member_id);
+                    }
+                }
+                if page_len < page_size {
+                    break;
+                }
+                offset = offset.saturating_add(page_size);
+            }
+        }
+
+        let exhausted_keys = keys
+            .iter()
+            .filter(|key| exhausted_key_ids.contains(&key.id))
+            .collect::<Vec<_>>();
+        if exhausted_keys.is_empty() {
+            return Ok(0);
+        }
+
+        let deleted_key_ids = exhausted_keys
+            .iter()
+            .map(|key| key.id.clone())
+            .collect::<Vec<_>>();
+        for key in exhausted_keys {
+            self.clear_admin_provider_pool_cooldown(&provider.id, &key.id)
+                .await;
+            self.reset_admin_provider_pool_cost(&provider.id, &key.id)
+                .await;
+        }
+
+        let mut affected = 0usize;
+        for key_id in &deleted_key_ids {
+            if self.delete_provider_catalog_key(key_id).await? {
+                affected += 1;
+            }
+        }
+        self.cleanup_deleted_provider_catalog_refs(&provider.id, false, &[], &deleted_key_ids)
+            .await?;
+
+        Ok(affected)
+    }
+
    pub(crate) async fn cleanup_provider_catalog_key_if_current<F>(
        &self,
        provider: &StoredProviderCatalogProvider,