refactor: 抽离 AI pipeline 与调度共享能力逻辑

2026-09-02 17:30:23 +08:00 · 2026-04-10 01:46:14 +08:00
parent b901a6ffc7
commit 5014e2f5fd
255 changed files with 15057 additions and 3115 deletions
--- a/apps/aether-gateway/src/state/core.rs
+++ b/apps/aether-gateway/src/state/core.rs
@@ -23,7 +23,7 @@ use super::super::async_task::{
 };
 use super::super::cache::{
    AuthApiKeyLastUsedCache, AuthContextCache, DirectPlanBypassCache, SchedulerAffinityCache,
-    SchedulerAffinityTarget,
+    SchedulerAffinitySnapshotEntry, SchedulerAffinityTarget,
 };
 use super::super::data::{GatewayDataConfig, GatewayDataState};
 use super::super::fallback_metrics;
@@ -47,6 +47,45 @@ use crate::maintenance::spawn_usage_cleanup_worker;
 use crate::maintenance::spawn_wallet_daily_usage_aggregation_worker;

 impl AppState {
+    fn spawn_scheduler_affinity_redis_write(
+        &self,
+        cache_key: &str,
+        target: &SchedulerAffinityTarget,
+        ttl: Duration,
+    ) {
+        let Some(runner) = self.redis_kv_runner() else {
+            return;
+        };
+        let Ok(handle) = tokio::runtime::Handle::try_current() else {
+            return;
+        };
+
+        let cache_key = cache_key.to_string();
+        let provider_id = target.provider_id.clone();
+        let endpoint_id = target.endpoint_id.clone();
+        let key_id = target.key_id.clone();
+        let ttl_seconds = ttl.as_secs();
+        let now_unix_secs = chrono::Utc::now().timestamp().max(0) as u64;
+        let expire_at = now_unix_secs.saturating_add(ttl_seconds);
+
+        handle.spawn(async move {
+            let payload = serde_json::json!({
+                "provider_id": provider_id,
+                "endpoint_id": endpoint_id,
+                "key_id": key_id,
+                "created_at": now_unix_secs,
+                "expire_at": expire_at,
+                "request_count": 0,
+            });
+            let Ok(serialized) = serde_json::to_string(&payload) else {
+                return;
+            };
+            let _ = runner
+                .setex(&cache_key, &serialized, Some(ttl_seconds))
+                .await;
+        });
+    }
+
    pub(crate) fn replace_data_state(&mut self, data: Arc<GatewayDataState>) {
        self.clear_provider_transport_snapshot_cache();
        self.tunnel = crate::tunnel::EmbeddedTunnelState::with_data(Arc::clone(&data));
@@ -572,10 +611,18 @@ impl AppState {
        ttl: Duration,
        max_entries: usize,
    ) {
+        self.spawn_scheduler_affinity_redis_write(cache_key, &target, ttl);
        self.scheduler_affinity_cache
            .insert(cache_key.to_string(), target, ttl, max_entries);
    }

+    pub(crate) fn list_scheduler_affinity_entries(
+        &self,
+        ttl: Duration,
+    ) -> Vec<SchedulerAffinitySnapshotEntry> {
+        self.scheduler_affinity_cache.fresh_entries(ttl)
+    }
+
    pub fn with_video_task_store_path(
        mut self,
        path: impl Into<std::path::PathBuf>,
--- a/apps/aether-gateway/src/state/integrations.rs
+++ b/apps/aether-gateway/src/state/integrations.rs
@@ -28,7 +28,8 @@ use crate::provider_transport::{
    LocalResolvedOAuthRequestAuth,
 };
 use crate::request_candidate_runtime::{
-    RequestCandidateRuntimeReader, RequestCandidateRuntimeWriter,
+    RequestCandidateRuntimeCapabilityReader, RequestCandidateRuntimeReader,
+    RequestCandidateRuntimeWriter,
 };
 use crate::scheduler::state::SchedulerRuntimeState;
 use crate::{execution_runtime, provider_transport};
@@ -246,6 +247,24 @@ impl RequestCandidateRuntimeReader for AppState {
    }
 }

+#[async_trait]
+impl RequestCandidateRuntimeCapabilityReader for AppState {
+    async fn read_request_candidate_user_model_capability_settings(
+        &self,
+        user_id: &str,
+    ) -> Result<Option<Value>, GatewayError> {
+        AppState::read_user_model_capability_settings(self, user_id).await
+    }
+
+    async fn read_request_candidate_api_key_force_capabilities(
+        &self,
+        user_id: &str,
+        api_key_id: &str,
+    ) -> Result<Option<Value>, GatewayError> {
+        AppState::read_auth_api_key_force_capabilities(self, user_id, api_key_id).await
+    }
+}
+
 #[async_trait]
 impl RequestCandidateRuntimeWriter for AppState {
    fn has_request_candidate_data_writer(&self) -> bool {
@@ -311,4 +330,10 @@ impl SchedulerRuntimeState for AppState {
    ) {
        AppState::remember_scheduler_affinity_target(self, cache_key, target, ttl, max_entries);
    }
+
+    async fn read_scheduler_ordering_config(
+        &self,
+    ) -> Result<crate::scheduler::config::SchedulerOrderingConfig, GatewayError> {
+        crate::scheduler::config::read_scheduler_ordering_config(self).await
+    }
 }
--- a/apps/aether-gateway/src/state/runtime/api_key_exports.rs
+++ b/apps/aether-gateway/src/state/runtime/api_key_exports.rs
@@ -1,6 +1,19 @@
 use crate::{AppState, GatewayError};

 impl AppState {
+    pub(crate) async fn read_auth_api_key_force_capabilities(
+        &self,
+        user_id: &str,
+        api_key_id: &str,
+    ) -> Result<Option<serde_json::Value>, GatewayError> {
+        Ok(self
+            .list_auth_api_key_export_records_by_ids(&[api_key_id.to_string()])
+            .await?
+            .into_iter()
+            .find(|record| record.api_key_id == api_key_id && record.user_id == user_id)
+            .and_then(|record| record.force_capabilities))
+    }
+
    pub(crate) async fn list_auth_api_key_export_records_by_user_ids(
        &self,
        user_ids: &[String],
--- a/apps/aether-gateway/src/state/runtime/billing/admin.rs
+++ b/apps/aether-gateway/src/state/runtime/billing/admin.rs
@@ -60,7 +60,7 @@ impl AppState {
                variables: input.variables.clone(),
                dimension_mappings: input.dimension_mappings.clone(),
                is_enabled: input.is_enabled,
-                created_at_unix_secs: now_unix_secs,
+                created_at_unix_ms: now_unix_secs,
                updated_at_unix_secs: now_unix_secs,
            };
            store
@@ -187,7 +187,7 @@ impl AppState {
                default_value: input.default_value.clone(),
                priority: input.priority,
                is_enabled: input.is_enabled,
-                created_at_unix_secs: now_unix_secs,
+                created_at_unix_ms: now_unix_secs,
                updated_at_unix_secs: now_unix_secs,
            };
            store
@@ -369,7 +369,7 @@ impl AppState {
                            default_value: collector.default_value.clone(),
                            priority: collector.priority,
                            is_enabled: collector.is_enabled,
-                            created_at_unix_secs: now_unix_secs,
+                            created_at_unix_ms: now_unix_secs,
                            updated_at_unix_secs: now_unix_secs,
                        };
                        guard.insert(record.id.clone(), record);
--- a/apps/aether-gateway/src/state/runtime/billing/finance_queries.rs
+++ b/apps/aether-gateway/src/state/runtime/billing/finance_queries.rs
@@ -86,8 +86,8 @@ impl AppState {
                .collect::<Vec<_>>();
            items.sort_by(|left, right| {
                right
-                    .created_at_unix_secs
-                    .cmp(&left.created_at_unix_secs)
+                    .created_at_unix_ms
+                    .cmp(&left.created_at_unix_ms)
                    .then_with(|| right.id.cmp(&left.id))
            });
            let total = items.len() as u64;
@@ -137,8 +137,8 @@ impl AppState {
                .collect::<Vec<_>>();
            items.sort_by(|left, right| {
                right
-                    .created_at_unix_secs
-                    .cmp(&left.created_at_unix_secs)
+                    .created_at_unix_ms
+                    .cmp(&left.created_at_unix_ms)
                    .then_with(|| right.id.cmp(&left.id))
            });
            let total = items.len() as u64;
@@ -182,8 +182,8 @@ impl AppState {
                .collect::<Vec<_>>();
            items.sort_by(|left, right| {
                right
-                    .created_at_unix_secs
-                    .cmp(&left.created_at_unix_secs)
+                    .created_at_unix_ms
+                    .cmp(&left.created_at_unix_ms)
                    .then_with(|| right.id.cmp(&left.id))
            });
            let total = items.len() as u64;
@@ -209,7 +209,7 @@ impl AppState {
                    operator_name: None,
                    operator_email: None,
                    description: record.description,
-                    created_at_unix_secs: Some(record.created_at_unix_secs),
+                    created_at_unix_ms: Some(record.created_at_unix_ms),
                })
                .collect::<Vec<_>>();
            return Ok((items, total));
@@ -240,8 +240,8 @@ impl AppState {
                .collect::<Vec<_>>();
            items.sort_by(|left, right| {
                right
-                    .created_at_unix_secs
-                    .cmp(&left.created_at_unix_secs)
+                    .created_at_unix_ms
+                    .cmp(&left.created_at_unix_ms)
                    .then_with(|| right.id.cmp(&left.id))
            });
            let total = items.len() as u64;
@@ -296,8 +296,8 @@ impl AppState {
                .collect::<Vec<_>>();
            items.sort_by(|left, right| {
                right
-                    .created_at_unix_secs
-                    .cmp(&left.created_at_unix_secs)
+                    .created_at_unix_ms
+                    .cmp(&left.created_at_unix_ms)
                    .then_with(|| right.id.cmp(&left.id))
            });
            let total = items.len() as u64;
@@ -332,7 +332,7 @@ impl AppState {
                            wallet_api_key_id: wallet.api_key_id.clone(),
                            api_key_name: None,
                            wallet_status: wallet.status.clone(),
-                            created_at_unix_secs: Some(refund.created_at_unix_secs),
+                            created_at_unix_ms: Some(refund.created_at_unix_ms),
                            updated_at_unix_secs: Some(refund.updated_at_unix_secs),
                            processed_at_unix_secs: refund.processed_at_unix_secs,
                            completed_at_unix_secs: refund.completed_at_unix_secs,
@@ -406,7 +406,7 @@ fn stored_admin_payment_order_to_gateway(
        gateway_order_id: record.gateway_order_id,
        status: record.status,
        gateway_response: record.gateway_response,
-        created_at_unix_secs: record.created_at_unix_secs,
+        created_at_unix_ms: record.created_at_unix_ms,
        paid_at_unix_secs: record.paid_at_unix_secs,
        credited_at_unix_secs: record.credited_at_unix_secs,
        expires_at_unix_secs: record.expires_at_unix_secs,
@@ -428,7 +428,7 @@ fn stored_admin_payment_callback_to_gateway(
        status: record.status,
        payload: record.payload,
        error_message: record.error_message,
-        created_at_unix_secs: record.created_at_unix_secs,
+        created_at_unix_ms: record.created_at_unix_ms,
        processed_at_unix_secs: record.processed_at_unix_secs,
    }
 }
@@ -456,7 +456,7 @@ fn stored_admin_wallet_refund_to_gateway(
        requested_by: record.requested_by,
        approved_by: record.approved_by,
        processed_by: record.processed_by,
-        created_at_unix_secs: record.created_at_unix_secs,
+        created_at_unix_ms: record.created_at_unix_ms,
        updated_at_unix_secs: record.updated_at_unix_secs,
        processed_at_unix_secs: record.processed_at_unix_secs,
        completed_at_unix_secs: record.completed_at_unix_secs,
--- a/apps/aether-gateway/src/state/runtime/payments.rs
+++ b/apps/aether-gateway/src/state/runtime/payments.rs
@@ -249,7 +249,7 @@ fn stored_admin_payment_order_to_gateway(
        gateway_order_id: order.gateway_order_id,
        status: order.status,
        gateway_response: order.gateway_response,
-        created_at_unix_secs: order.created_at_unix_secs,
+        created_at_unix_ms: order.created_at_unix_ms,
        paid_at_unix_secs: order.paid_at_unix_secs,
        credited_at_unix_secs: order.credited_at_unix_secs,
        expires_at_unix_secs: order.expires_at_unix_secs,
--- a/apps/aether-gateway/src/state/runtime/wallet/balance_mutations.rs
+++ b/apps/aether-gateway/src/state/runtime/wallet/balance_mutations.rs
@@ -85,7 +85,7 @@ impl AppState {
                        .unwrap_or("管理员调账")
                        .to_string(),
                ),
-                created_at_unix_secs: chrono::Utc::now().timestamp().max(0) as u64,
+                created_at_unix_ms: chrono::Utc::now().timestamp().max(0) as u64,
            };
            return Ok(Some((wallet.clone(), transaction)));
        }
@@ -148,7 +148,7 @@ impl AppState {
                    "operator_id": operator_id,
                    "description": description,
                })),
-                created_at_unix_secs: created_at,
+                created_at_unix_ms: created_at,
                paid_at_unix_secs: Some(created_at),
                credited_at_unix_secs: Some(created_at),
                expires_at_unix_secs: None,
@@ -193,7 +193,7 @@ fn stored_wallet_transaction_to_gateway(
        link_id: transaction.link_id,
        operator_id: transaction.operator_id,
        description: transaction.description,
-        created_at_unix_secs: transaction.created_at_unix_secs.unwrap_or_default(),
+        created_at_unix_ms: transaction.created_at_unix_ms.unwrap_or_default(),
    }
 }

@@ -215,7 +215,7 @@ fn stored_admin_payment_order_to_gateway(
        gateway_order_id: order.gateway_order_id,
        status: order.status,
        gateway_response: order.gateway_response,
-        created_at_unix_secs: order.created_at_unix_secs,
+        created_at_unix_ms: order.created_at_unix_ms,
        paid_at_unix_secs: order.paid_at_unix_secs,
        credited_at_unix_secs: order.credited_at_unix_secs,
        expires_at_unix_secs: order.expires_at_unix_secs,
--- a/apps/aether-gateway/src/state/runtime/wallet/billing.rs
+++ b/apps/aether-gateway/src/state/runtime/wallet/billing.rs
@@ -102,8 +102,8 @@ pub(super) fn admin_wallet_payment_order_from_row(
        gateway_response: row
            .try_get("gateway_response")
            .map_err(|err| GatewayError::Internal(err.to_string()))?,
-        created_at_unix_secs: row
-            .try_get::<i64, _>("created_at_unix_secs")
+        created_at_unix_ms: row
+            .try_get::<i64, _>("created_at_unix_ms")
            .map_err(|err| GatewayError::Internal(err.to_string()))?
            .max(0) as u64,
        paid_at_unix_secs: row
@@ -182,8 +182,8 @@ pub(super) fn admin_wallet_refund_from_row(
        processed_by: row
            .try_get("processed_by")
            .map_err(|err| GatewayError::Internal(err.to_string()))?,
-        created_at_unix_secs: row
-            .try_get::<i64, _>("created_at_unix_secs")
+        created_at_unix_ms: row
+            .try_get::<i64, _>("created_at_unix_ms")
            .map_err(|err| GatewayError::Internal(err.to_string()))?
            .max(0) as u64,
        updated_at_unix_secs: row
@@ -234,8 +234,8 @@ pub(super) fn admin_billing_rule_from_row(
        is_enabled: row
            .try_get("is_enabled")
            .map_err(|err| GatewayError::Internal(err.to_string()))?,
-        created_at_unix_secs: row
-            .try_get::<i64, _>("created_at_unix_secs")
+        created_at_unix_ms: row
+            .try_get::<i64, _>("created_at_unix_ms")
            .map_err(|err| GatewayError::Internal(err.to_string()))?
            .max(0) as u64,
        updated_at_unix_secs: row
@@ -282,8 +282,8 @@ pub(super) fn admin_billing_collector_from_row(
        is_enabled: row
            .try_get("is_enabled")
            .map_err(|err| GatewayError::Internal(err.to_string()))?,
-        created_at_unix_secs: row
-            .try_get::<i64, _>("created_at_unix_secs")
+        created_at_unix_ms: row
+            .try_get::<i64, _>("created_at_unix_ms")
            .map_err(|err| GatewayError::Internal(err.to_string()))?
            .max(0) as u64,
        updated_at_unix_secs: row
--- a/apps/aether-gateway/src/state/runtime/wallet/refund_lifecycle.rs
+++ b/apps/aether-gateway/src/state/runtime/wallet/refund_lifecycle.rs
@@ -104,7 +104,7 @@ impl AppState {
                link_id: Some(refund.id.clone()),
                operator_id: operator_id.map(ToOwned::to_owned),
                description: Some("退款占款".to_string()),
-                created_at_unix_secs: now_unix_secs,
+                created_at_unix_ms: now_unix_secs,
            };

            let mut updated_refund = refund.clone();
@@ -318,7 +318,7 @@ impl AppState {
                link_id: Some(refund.id.clone()),
                operator_id: operator_id.map(ToOwned::to_owned),
                description: Some("退款失败回补".to_string()),
-                created_at_unix_secs: now_unix_secs,
+                created_at_unix_ms: now_unix_secs,
            };

            if let Some(payment_order_id) = refund.payment_order_id.clone() {
@@ -415,7 +415,7 @@ fn stored_admin_wallet_refund_to_gateway(
        requested_by: refund.requested_by,
        approved_by: refund.approved_by,
        processed_by: refund.processed_by,
-        created_at_unix_secs: refund.created_at_unix_secs,
+        created_at_unix_ms: refund.created_at_unix_ms,
        updated_at_unix_secs: refund.updated_at_unix_secs,
        processed_at_unix_secs: refund.processed_at_unix_secs,
        completed_at_unix_secs: refund.completed_at_unix_secs,
@@ -441,6 +441,6 @@ fn stored_admin_wallet_transaction_to_gateway(
        link_id: transaction.link_id,
        operator_id: transaction.operator_id,
        description: transaction.description,
-        created_at_unix_secs: transaction.created_at_unix_secs.unwrap_or_default(),
+        created_at_unix_ms: transaction.created_at_unix_ms.unwrap_or_default(),
    }
 }
--- a/apps/aether-gateway/src/state/types.rs
+++ b/apps/aether-gateway/src/state/types.rs
@@ -315,7 +315,7 @@ pub(crate) struct GatewayAdminPaymentCallbackView {
    pub(crate) status: String,
    pub(crate) payload: Option<serde_json::Value>,
    pub(crate) error_message: Option<String>,
-    pub(crate) created_at_unix_secs: u64,
+    pub(crate) created_at_unix_ms: u64,
    pub(crate) processed_at_unix_secs: Option<u64>,
 }

@@ -333,7 +333,7 @@ impl From<super::AdminPaymentCallbackRecord> for GatewayAdminPaymentCallbackView
            status: value.status,
            payload: value.payload,
            error_message: value.error_message,
-            created_at_unix_secs: value.created_at_unix_secs,
+            created_at_unix_ms: value.created_at_unix_ms,
            processed_at_unix_secs: value.processed_at_unix_secs,
        }
    }