diff --git a/backend/crates/gateway-admin/src/model/settings.rs b/backend/crates/gateway-admin/src/model/settings.rs index 74b68615d..6f002cded 100644 --- a/backend/crates/gateway-admin/src/model/settings.rs +++ b/backend/crates/gateway-admin/src/model/settings.rs @@ -38,6 +38,7 @@ pub struct RuntimeSettings { pub max_waiting_per_account: u32, pub concurrency_wait_timeout_seconds: u32, pub responses_max_decompressed_body_bytes: u64, + pub smart_scheduling: gateway_core::account::SmartSchedulingConfig, pub rotation_strategy: RotationStrategy, pub min_codex_desktop_version: Option, pub min_codex_cli_version: Option, @@ -73,6 +74,7 @@ pub struct ReplaceRuntimeSettings { pub max_waiting_per_account: u32, pub concurrency_wait_timeout_seconds: u32, pub responses_max_decompressed_body_bytes: u64, + pub smart_scheduling: gateway_core::account::SmartSchedulingConfig, pub rotation_strategy: RotationStrategy, pub min_codex_desktop_version: Option, pub min_codex_cli_version: Option, diff --git a/backend/crates/gateway-admin/tests/use_case/accounts.rs b/backend/crates/gateway-admin/tests/use_case/accounts.rs index 8c4e2aff7..70031c0ef 100644 --- a/backend/crates/gateway-admin/tests/use_case/accounts.rs +++ b/backend/crates/gateway-admin/tests/use_case/accounts.rs @@ -1105,6 +1105,7 @@ impl SettingsStore for StaticSettingsStore { max_waiting_per_account: 0, concurrency_wait_timeout_seconds: 30, responses_max_decompressed_body_bytes: 64 * 1024 * 1024, + smart_scheduling: gateway_core::account::SmartSchedulingConfig::default(), rotation_strategy: RotationStrategy::Smart, min_codex_desktop_version: None, min_codex_cli_version: None, diff --git a/backend/crates/gateway-admin/tests/use_case/freeze_recovery.rs b/backend/crates/gateway-admin/tests/use_case/freeze_recovery.rs index a8f367f84..8058daae2 100644 --- a/backend/crates/gateway-admin/tests/use_case/freeze_recovery.rs +++ b/backend/crates/gateway-admin/tests/use_case/freeze_recovery.rs @@ -38,6 +38,7 @@ fn runtime_settings(enabled: bool, probe_enabled: bool, adaptive: bool) -> Runti max_waiting_per_account: 0, concurrency_wait_timeout_seconds: 30, responses_max_decompressed_body_bytes: 64 * 1024 * 1024, + smart_scheduling: gateway_core::account::SmartSchedulingConfig::default(), rotation_strategy: gateway_admin::model::settings::RotationStrategy::Smart, min_codex_desktop_version: None, min_codex_cli_version: None, diff --git a/backend/crates/gateway-admin/tests/use_case/observability.rs b/backend/crates/gateway-admin/tests/use_case/observability.rs index 5a0292514..f3a317ccf 100644 --- a/backend/crates/gateway-admin/tests/use_case/observability.rs +++ b/backend/crates/gateway-admin/tests/use_case/observability.rs @@ -845,6 +845,7 @@ impl SettingsStore for FixtureSettingsStore { max_waiting_per_account: 0, concurrency_wait_timeout_seconds: 30, responses_max_decompressed_body_bytes: 64 * 1024 * 1024, + smart_scheduling: gateway_core::account::SmartSchedulingConfig::default(), rotation_strategy: RotationStrategy::Smart, min_codex_desktop_version: None, min_codex_cli_version: None, diff --git a/backend/crates/gateway-admin/tests/use_case/settings.rs b/backend/crates/gateway-admin/tests/use_case/settings.rs index 30ea2a820..5ab386caf 100644 --- a/backend/crates/gateway-admin/tests/use_case/settings.rs +++ b/backend/crates/gateway-admin/tests/use_case/settings.rs @@ -225,6 +225,7 @@ async fn settings_should_reject_zero_refresh_margin_before_store_call() { max_waiting_per_account: 0, concurrency_wait_timeout_seconds: 30, responses_max_decompressed_body_bytes: 64 * 1024 * 1024, + smart_scheduling: gateway_core::account::SmartSchedulingConfig::default(), rotation_strategy: RotationStrategy::Smart, min_codex_desktop_version: None, min_codex_cli_version: None, diff --git a/backend/crates/gateway-api/src/admin/settings.rs b/backend/crates/gateway-api/src/admin/settings.rs index c646392b9..7c8ae1aac 100644 --- a/backend/crates/gateway-api/src/admin/settings.rs +++ b/backend/crates/gateway-api/src/admin/settings.rs @@ -37,6 +37,7 @@ pub type ProviderRequestProfileUpdates = #[derive(Debug, Clone, PartialEq, Eq, Serialize)] #[serde(rename_all = "camelCase")] pub struct RuntimeSettingsView { + pub smart_scheduling_defaults: gateway_core::account::SmartSchedulingConfig, pub provider_request_profiles: ProviderRequestProfiles, /// 固定兼容字段；值始终从 provider_request_profiles 派生。 pub openai_client_profile: Option>, @@ -53,6 +54,7 @@ pub struct RuntimeSettingsView { pub max_waiting_per_account: u32, pub concurrency_wait_timeout_seconds: u32, pub responses_max_decompressed_body_bytes: u64, + pub smart_scheduling: gateway_core::account::SmartSchedulingConfig, pub rotation_strategy: String, pub min_codex_desktop_version: Option, pub min_codex_cli_version: Option, @@ -95,6 +97,7 @@ pub struct UpdateRuntimeSettingsRequest { pub max_waiting_per_account: u32, pub concurrency_wait_timeout_seconds: u32, pub responses_max_decompressed_body_bytes: u64, + pub smart_scheduling: gateway_core::account::SmartSchedulingConfig, pub rotation_strategy: String, pub min_codex_desktop_version: Option, pub min_codex_cli_version: Option, @@ -227,6 +230,7 @@ impl UpdateRuntimeSettingsRequest { max_waiting_per_account: self.max_waiting_per_account, concurrency_wait_timeout_seconds: self.concurrency_wait_timeout_seconds, responses_max_decompressed_body_bytes: self.responses_max_decompressed_body_bytes, + smart_scheduling: self.smart_scheduling, rotation_strategy: RotationStrategy::parse(&self.rotation_strategy) .ok_or_else(|| WireValidationError::new("rotationStrategy"))?, min_codex_desktop_version: self.min_codex_desktop_version, @@ -274,6 +278,8 @@ impl From for RuntimeSettingsView { max_waiting_per_account: settings.max_waiting_per_account, concurrency_wait_timeout_seconds: settings.concurrency_wait_timeout_seconds, responses_max_decompressed_body_bytes: settings.responses_max_decompressed_body_bytes, + smart_scheduling: settings.smart_scheduling, + smart_scheduling_defaults: gateway_core::account::SmartSchedulingConfig::default(), rotation_strategy: settings.rotation_strategy.as_str().to_owned(), min_codex_desktop_version: settings.min_codex_desktop_version, min_codex_cli_version: settings.min_codex_cli_version, diff --git a/backend/crates/gateway-api/tests/admin/mod.rs b/backend/crates/gateway-api/tests/admin/mod.rs index 0616a518a..11f5710aa 100644 --- a/backend/crates/gateway-api/tests/admin/mod.rs +++ b/backend/crates/gateway-api/tests/admin/mod.rs @@ -557,6 +557,7 @@ impl SettingsStore for MemorySettingsStore { max_waiting_per_account: command.max_waiting_per_account, concurrency_wait_timeout_seconds: command.concurrency_wait_timeout_seconds, responses_max_decompressed_body_bytes: command.responses_max_decompressed_body_bytes, + smart_scheduling: command.smart_scheduling, rotation_strategy: command.rotation_strategy, min_codex_desktop_version: command.min_codex_desktop_version, min_codex_cli_version: command.min_codex_cli_version, @@ -1519,6 +1520,7 @@ fn test_runtime_settings() -> RuntimeSettings { max_waiting_per_account: 0, concurrency_wait_timeout_seconds: 30, responses_max_decompressed_body_bytes: 64 * 1024 * 1024, + smart_scheduling: gateway_core::account::SmartSchedulingConfig::default(), rotation_strategy: RotationStrategy::Smart, min_codex_desktop_version: None, min_codex_cli_version: None, diff --git a/backend/crates/gateway-api/tests/admin/settings.rs b/backend/crates/gateway-api/tests/admin/settings.rs index c2646249c..b72bfd0d4 100644 --- a/backend/crates/gateway-api/tests/admin/settings.rs +++ b/backend/crates/gateway-api/tests/admin/settings.rs @@ -52,6 +52,7 @@ fn update_body() -> Value { "concurrencyWaitTimeoutSeconds": 30, "responsesMaxDecompressedBodyBytes": 67108864, "rotationStrategy": "round_robin", + "smartScheduling": gateway_core::account::SmartSchedulingConfig::default(), "minCodexDesktopVersion": "26.825.6671", "minCodexCliVersion": "0.40.0", "usageRetentionDays": 32, @@ -70,6 +71,76 @@ fn update_body() -> Value { }) } +#[tokio::test] +async fn smart_settings_round_trip_and_invalid_updates_leave_the_saved_value_intact() { + let fixture = AdminTestFixture::new().await; + fixture.auth.insert_session("valid-session"); + let mut body = update_body(); + let custom = json!({"loadWeight": 2.0, "quotaWeight": 0.0, "healthWeight": 1.0, "latencyWeight": 0.5, "resetWeight": 1.2, "queueWeight": 2.3, "preferHigherWeight": true}); + body["smartScheduling"] = custom.clone(); + let response = app(fixture.state()) + .oneshot(request( + Method::POST, + "/api/admin/settings/update", + Some(body.clone()), + )) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let response = response_json(response).await; + assert_eq!(response["data"]["smartScheduling"], custom); + assert_eq!( + response["data"]["smartSchedulingDefaults"], + json!(gateway_core::account::SmartSchedulingConfig::default()) + ); + let mut invalid_values = vec![ + json!(null), + json!({}), + json!({"loadWeight":0,"quotaWeight":0,"healthWeight":0,"latencyWeight":0,"resetWeight":0,"queueWeight":0,"preferHigherWeight":true}), + json!({"loadWeight":0.01,"quotaWeight":1,"healthWeight":1,"latencyWeight":1,"resetWeight":0,"queueWeight":0,"preferHigherWeight":false}), + ]; + for field in ["resetWeight", "queueWeight"] { + for value in [json!(-0.1), json!(10.1), json!(0.01), json!(null)] { + let mut invalid = custom.clone(); + invalid[field] = value; + invalid_values.push(invalid); + } + let mut missing = custom.clone(); + missing.as_object_mut().unwrap().remove(field); + invalid_values.push(missing); + } + for invalid in invalid_values { + body["smartScheduling"] = invalid; + let response = app(fixture.state()) + .oneshot(request( + Method::POST, + "/api/admin/settings/update", + Some(body.clone()), + )) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY); + } + body.as_object_mut().unwrap().remove("smartScheduling"); + let response = app(fixture.state()) + .oneshot(request( + Method::POST, + "/api/admin/settings/update", + Some(body), + )) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY); + let response = app(fixture.state()) + .oneshot(request(Method::GET, "/api/admin/settings", None)) + .await + .unwrap(); + assert_eq!( + response_json(response).await["data"]["smartScheduling"], + custom + ); +} + #[test] fn settings_request_should_reject_unknown_rotation_strategy() { let mut body = update_body(); @@ -149,6 +220,7 @@ fn settings_response_should_cover_the_full_runtime_settings_contract() { max_waiting_per_account: 0, concurrency_wait_timeout_seconds: 30, responses_max_decompressed_body_bytes: 64 * 1024 * 1024, + smart_scheduling: gateway_core::account::SmartSchedulingConfig::default(), rotation_strategy: RotationStrategy::RoundRobin, min_codex_desktop_version: Some("26.825.6671".to_owned()), min_codex_cli_version: Some("0.40.0".to_owned()), @@ -193,6 +265,8 @@ fn settings_response_should_cover_the_full_runtime_settings_contract() { "concurrencyWaitTimeoutSeconds": 30, "responsesMaxDecompressedBodyBytes": 67108864, "rotationStrategy": "round_robin", + "smartScheduling": gateway_core::account::SmartSchedulingConfig::default(), + "smartSchedulingDefaults": gateway_core::account::SmartSchedulingConfig::default(), "minCodexDesktopVersion": "26.825.6671", "minCodexCliVersion": "0.40.0", "usageRetentionDays": 32, @@ -257,6 +331,7 @@ fn settings_request_and_response_fields_should_stay_in_lockstep() { max_waiting_per_account: 0, concurrency_wait_timeout_seconds: 30, responses_max_decompressed_body_bytes: 64 * 1024 * 1024, + smart_scheduling: request.smart_scheduling, rotation_strategy: RotationStrategy::parse(&request.rotation_strategy) .expect("fixture rotation strategy"), min_codex_desktop_version: request.min_codex_desktop_version, @@ -290,6 +365,7 @@ fn settings_request_and_response_fields_should_stay_in_lockstep() { expected_fields.insert("openaiClientProfile".to_owned()); expected_fields.insert("xaiClientProfile".to_owned()); expected_fields.insert("updatedAt".to_owned()); + expected_fields.insert("smartSchedulingDefaults".to_owned()); assert_eq!(response_fields, expected_fields); } diff --git a/backend/crates/gateway-core/src/account/mod.rs b/backend/crates/gateway-core/src/account/mod.rs index 66c48781d..3b6d0d42f 100644 --- a/backend/crates/gateway-core/src/account/mod.rs +++ b/backend/crates/gateway-core/src/account/mod.rs @@ -13,10 +13,12 @@ mod proxy; pub use proxy::{InvalidOutboundProxy, OutboundProxy}; pub mod scope; mod selection; +mod smart_scheduling; mod store; pub use error::CredentialError; pub use model::*; +pub(crate) use selection::smart_score; pub use selection::*; -pub(crate) use selection::{SMART_SCORE_TOLERANCE, smart_score}; +pub use smart_scheduling::{SmartSchedulingConfig, SmartSchedulingConfigError}; pub use store::ProviderAccountStore; diff --git a/backend/crates/gateway-core/src/account/selection.rs b/backend/crates/gateway-core/src/account/selection.rs index 80fa79ec5..b0c015801 100644 --- a/backend/crates/gateway-core/src/account/selection.rs +++ b/backend/crates/gateway-core/src/account/selection.rs @@ -6,10 +6,12 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Mutex, RwLock}; use std::time::{Duration, Instant, SystemTime}; -use crate::concurrency::ConcurrencyQueuePolicy; +use crate::concurrency::{CapacityWait, ConcurrencyQueuePolicy, QueueRejection}; use crate::identity::ProviderKind; -use super::{AccountConcurrency, AccountStatus, ProviderAccount, ProviderAccountId}; +use super::{ + AccountConcurrency, AccountStatus, ProviderAccount, ProviderAccountId, SmartSchedulingConfig, +}; /// `runtime_settings.rotation_strategy` 的稳定值。 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -47,6 +49,7 @@ impl RotationStrategy { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct AccountSelectionPolicy { strategy: RotationStrategy, + smart_scheduling: SmartSchedulingConfig, max_concurrent_per_account: AccountConcurrency, request_interval: Duration, queue_policy: ConcurrencyQueuePolicy, @@ -61,6 +64,7 @@ impl AccountSelectionPolicy { ) -> Self { Self { strategy, + smart_scheduling: SmartSchedulingConfig::default(), max_concurrent_per_account: max_concurrent_per_account.into(), request_interval, queue_policy: ConcurrencyQueuePolicy { @@ -86,6 +90,17 @@ impl AccountSelectionPolicy { self.strategy } + #[must_use] + pub const fn with_smart_scheduling(mut self, config: SmartSchedulingConfig) -> Self { + self.smart_scheduling = config; + self + } + + #[must_use] + pub const fn smart_scheduling(self) -> SmartSchedulingConfig { + self.smart_scheduling + } + #[must_use] pub const fn max_concurrent_per_account(self) -> AccountConcurrency { self.max_concurrent_per_account @@ -510,12 +525,7 @@ impl AccountSelector { if self.scheduling_blocker(candidate, context).is_some() { return None; } - let highest_weight = candidates - .iter() - .filter(|candidate| self.scheduling_blocker(candidate, context).is_none()) - .map(|candidate| candidate.account.weight()) - .max()?; - let (preferred, _) = self.preferred_decision(candidates, context, highest_weight); + let (preferred, _) = self.preferred_decision(candidates, context); let preferred = if preferred == PreferredAccountSelection::Hit && context.preferred_account.as_ref() != Some(account_id) { @@ -585,13 +595,22 @@ impl AccountSelector { .iter() .map(|candidate| candidate.account.weight()) .max()?; - let (preferred, preferred_candidate) = - self.preferred_decision(candidates, context, highest_weight); + let (mut preferred, preferred_candidate) = self.preferred_decision(candidates, context); if let Some(candidate) = preferred_candidate { - return Some(AccountSelection { - candidate, - preferred: PreferredAccountSelection::Hit, - }); + // 权重回切只裁决内置策略的软亲和，插件显式选号按实际选择记录结果。 + let prefer_higher_weight = context.policy.strategy() == RotationStrategy::Smart + && context.policy.smart_scheduling().prefer_higher_weight(); + if (!context.preferred_account_overrides_weight || prefer_higher_weight) + && candidate.account.weight() < highest_weight + { + preferred = + PreferredAccountSelection::Blocked(AccountSchedulingBlocker::LowerWeight); + } else { + return Some(AccountSelection { + candidate, + preferred: PreferredAccountSelection::Hit, + }); + } } eligible.retain(|candidate| candidate.account.weight() == highest_weight); @@ -629,6 +648,8 @@ impl AccountSelector { &eligible, context.policy.max_concurrent_per_account(), context.round_robin_cursor, + context.policy.smart_scheduling(), + context.now, )?, RotationStrategy::Sticky => { eligible.sort_by_key(|candidate| { @@ -650,7 +671,6 @@ impl AccountSelector { &self, candidates: &'a [AccountCandidate], context: &AccountSelectionContext, - highest_weight: super::AccountWeight, ) -> (PreferredAccountSelection, Option<&'a AccountCandidate>) { let Some(preferred) = context.preferred_account.as_ref() else { return (PreferredAccountSelection::NotRequested, None); @@ -663,14 +683,6 @@ impl AccountSelector { }; match self.scheduling_blocker(candidate, context) { Some(blocker) => (PreferredAccountSelection::Blocked(blocker), None), - None if !context.preferred_account_overrides_weight - && candidate.account.weight() < highest_weight => - { - ( - PreferredAccountSelection::Blocked(AccountSchedulingBlocker::LowerWeight), - None, - ) - } None => (PreferredAccountSelection::Hit, Some(candidate)), } } @@ -707,6 +719,41 @@ impl AccountSelector { .collect() } + /// Provider 先投影续写范围与资格，再将等待候选交给共用评分规则选择队列。 + pub async fn wait_for_capacity( + &self, + waiting: &mut CapacityWait<'_, ProviderAccountId>, + keys: &[ProviderAccountId], + candidates: &[AccountCandidate], + context: &AccountSelectionContext, + ) -> Result<(), QueueRejection> { + let config = context.policy.smart_scheduling(); + let queue_weight = config.weights()[5]; + if context.policy.strategy() != RotationStrategy::Smart || queue_weight == 0.0 { + return waiting.wait(keys).await; + } + // 账号信号在队列锁外评分，锁内只读取实时队长并查表，避免扫描候选阻塞其他等待者。 + let scores = candidates + .iter() + .map(|candidate| { + ( + candidate.account.id(), + smart_score( + candidate, + context.policy.max_concurrent_per_account(), + config, + context.now, + ), + ) + }) + .collect::>(); + waiting + .wait_with_priority(keys, |key, count| { + scores[key] + queue_weight / (1.0 + count as f64) + }) + .await + } + pub(crate) fn scheduling_blocker( &self, candidate: &AccountCandidate, @@ -755,14 +802,10 @@ impl AccountSelector { } } -const SMART_LOAD_WEIGHT: f64 = 1.0; -const SMART_QUOTA_WEIGHT: f64 = 0.8; -const SMART_FAILURE_WEIGHT: f64 = 1.0; -const SMART_LATENCY_WEIGHT: f64 = 0.5; -// 容忍 5 个百分点的单项负载/失败率差异，避免微小信号波动独占新会话。 -pub(crate) const SMART_SCORE_TOLERANCE: f64 = 0.05; // 首输出 10 秒时延迟得分减半；固定尺度不随其他候选账号变化。 const SMART_LATENCY_HALF_SCORE_MS: f64 = 10_000.0; +// 距重置一小时时得分减半；未知和已过期时间不提供重置奖励。 +const SMART_RESET_HALF_SCORE_SECONDS: f64 = 3_600.0; fn capacity_utilization( candidate: &AccountCandidate, @@ -781,16 +824,23 @@ fn select_smart_candidate<'a>( candidates: &[&'a AccountCandidate], default_concurrency: AccountConcurrency, cursor: u64, + config: SmartSchedulingConfig, + now: SystemTime, ) -> Option<&'a AccountCandidate> { let mut ranked = candidates .iter() - .map(|candidate| (*candidate, smart_score(candidate, default_concurrency))) + .map(|candidate| { + ( + *candidate, + smart_score(candidate, default_concurrency, config, now), + ) + }) .collect::>(); let best_score = ranked .iter() .map(|(_, score)| *score) .max_by(f64::total_cmp)?; - ranked.retain(|(_, score)| best_score - score <= SMART_SCORE_TOLERANCE); + ranked.retain(|(_, score)| best_score - score <= config.score_tolerance()); // 轮换顺序保持稳定，避免分数轻微交错与 cursor 同步后仍反复命中同一账号。 ranked.sort_unstable_by(|(left, _), (right, _)| left.account.id().cmp(right.account.id())); let index = (cursor % ranked.len() as u64) as usize; @@ -800,6 +850,8 @@ fn select_smart_candidate<'a>( pub(crate) fn smart_score( candidate: &AccountCandidate, default_concurrency: AccountConcurrency, + config: SmartSchedulingConfig, + now: SystemTime, ) -> f64 { let load = 1.0 - capacity_utilization(candidate, default_concurrency).clamp(0.0, 1.0); let quota = candidate @@ -822,8 +874,26 @@ pub(crate) fn smart_score( SMART_LATENCY_HALF_SCORE_MS / (SMART_LATENCY_HALF_SCORE_MS + latency as f64) }); - SMART_LOAD_WEIGHT * load - + SMART_QUOTA_WEIGHT * quota - + SMART_FAILURE_WEIGHT * failure - + SMART_LATENCY_WEIGHT * latency + let reset = candidate + .signals + .quota_reset_at + .and_then(|reset| reset.duration_since(now).ok()) + .filter(|remaining| !remaining.is_zero()) + .map_or(0.0, |remaining| { + SMART_RESET_HALF_SCORE_SECONDS + / (SMART_RESET_HALF_SCORE_SECONDS + remaining.as_secs_f64()) + }); + let [ + load_weight, + quota_weight, + health_weight, + latency_weight, + reset_weight, + _, + ] = config.weights(); + load_weight * load + + quota_weight * quota + + health_weight * failure + + latency_weight * latency + + reset_weight * reset } diff --git a/backend/crates/gateway-core/src/account/smart_scheduling.rs b/backend/crates/gateway-core/src/account/smart_scheduling.rs new file mode 100644 index 000000000..02fc6ac8b --- /dev/null +++ b/backend/crates/gateway-core/src/account/smart_scheduling.rs @@ -0,0 +1,138 @@ +//! 智能调度参数；固定十分位保证快照可精确比较，浮点数只用于 wire 与评分。 + +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(try_from = "SmartSchedulingValues", into = "SmartSchedulingValues")] +pub struct SmartSchedulingConfig { + weights: [u8; 6], + prefer_higher_weight: bool, +} + +impl Default for SmartSchedulingConfig { + fn default() -> Self { + Self { + weights: [10, 8, 10, 5, 0, 0], + prefer_higher_weight: false, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum SmartSchedulingConfigError { + #[error("{0} must be between 0 and 10 with at most one decimal place")] + InvalidWeight(&'static str), + #[error("at least one smart scheduling weight must be positive")] + EmptyWeights, +} + +impl SmartSchedulingConfig { + /// 顺序为负载、剩余额度、健康、首输出延迟、额度重置和排队压力。 + pub fn new( + weights: [f64; 6], + prefer_higher_weight: bool, + ) -> Result { + let mut tenths = [0; 6]; + for (index, field) in [ + "loadWeight", + "quotaWeight", + "healthWeight", + "latencyWeight", + "resetWeight", + "queueWeight", + ] + .into_iter() + .enumerate() + { + let value = weights[index]; + if !value.is_finite() + || !(0.0..=10.0).contains(&value) + || (value * 10.0).round() / 10.0 != value + { + return Err(SmartSchedulingConfigError::InvalidWeight(field)); + } + tenths[index] = (value * 10.0).round() as u8; + } + if tenths == [0; 6] { + return Err(SmartSchedulingConfigError::EmptyWeights); + } + Ok(Self { + weights: tenths, + prefer_higher_weight, + }) + } + + #[must_use] + pub fn weights(self) -> [f64; 6] { + self.weights.map(|weight| f64::from(weight) / 10.0) + } + + #[must_use] + pub const fn prefer_higher_weight(self) -> bool { + self.prefer_higher_weight + } + + pub(crate) fn score_tolerance(self) -> f64 { + // 与系数同步缩放，避免仅放大相同比例的系数就改变近似最优候选集合。 + // 排队系数只参与入队选择，不能放大立即选号的候选容差。 + 0.05 * f64::from( + self.weights[..5] + .iter() + .copied() + .map(u16::from) + .sum::(), + ) / 33.0 + } +} + +#[derive(Serialize, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct SmartSchedulingValues { + load_weight: f64, + quota_weight: f64, + health_weight: f64, + latency_weight: f64, + reset_weight: f64, + queue_weight: f64, + prefer_higher_weight: bool, +} + +impl TryFrom for SmartSchedulingConfig { + type Error = SmartSchedulingConfigError; + + fn try_from(value: SmartSchedulingValues) -> Result { + Self::new( + [ + value.load_weight, + value.quota_weight, + value.health_weight, + value.latency_weight, + value.reset_weight, + value.queue_weight, + ], + value.prefer_higher_weight, + ) + } +} + +impl From for SmartSchedulingValues { + fn from(value: SmartSchedulingConfig) -> Self { + let [ + load_weight, + quota_weight, + health_weight, + latency_weight, + reset_weight, + queue_weight, + ] = value.weights(); + Self { + load_weight, + quota_weight, + health_weight, + latency_weight, + reset_weight, + queue_weight, + prefer_higher_weight: value.prefer_higher_weight, + } + } +} diff --git a/backend/crates/gateway-core/src/concurrency.rs b/backend/crates/gateway-core/src/concurrency.rs index 033607748..005a5a4c4 100644 --- a/backend/crates/gateway-core/src/concurrency.rs +++ b/backend/crates/gateway-core/src/concurrency.rs @@ -80,6 +80,18 @@ impl ConcurrencyWaitQueue { keys: &[K], max_waiting: u32, deadline: Instant, + ) -> Result, QueueRejection> { + self.enqueue_with_priority(keys, max_waiting, deadline, |_, waiting| -(waiting as f64)) + } + + /// 在同一锁内读取队长并入队，分数越大越优先，同分沿用候选顺序。 + /// 回调只做内存评分，不得阻塞或再次访问本队列。 + pub fn enqueue_with_priority( + &self, + keys: &[K], + max_waiting: u32, + deadline: Instant, + mut priority: impl FnMut(&K, usize) -> f64, ) -> Result, QueueRejection> { if Instant::now() >= deadline { return Err(QueueRejection::Timeout); @@ -95,9 +107,9 @@ impl ConcurrencyWaitQueue { .iter() .filter_map(|key| { let len = state.queues.get(key).map_or(0, VecDeque::len); - (len < max_waiting as usize).then_some((key, len)) + (len < max_waiting as usize).then(|| (key, priority(key, len))) }) - .min_by_key(|(_, len)| *len) + .min_by(|(_, left), (_, right)| right.total_cmp(left)) .map(|(key, _)| key.clone()) .ok_or(QueueRejection::Full)?; let waker = Arc::new(AtomicWaker::new()); @@ -237,6 +249,16 @@ impl<'a, K: Clone + Eq + Hash> CapacityWait<'a, K> { } pub async fn wait(&mut self, keys: &[K]) -> Result<(), QueueRejection> { + self.wait_with_priority(keys, |_, waiting| -(waiting as f64)) + .await + } + + /// 只在首次入队或原账号不再合格时重新选队列，已取得的位置不因评分变化而移动。 + pub async fn wait_with_priority( + &mut self, + keys: &[K], + priority: impl FnMut(&K, usize) -> f64, + ) -> Result<(), QueueRejection> { self.started_at.get_or_insert_with(Instant::now); let deadline = self .budget @@ -249,10 +271,12 @@ impl<'a, K: Clone + Eq + Hash> CapacityWait<'a, K> { self.ticket = None; } if self.ticket.is_none() { - self.ticket = Some( - self.queue - .enqueue(keys, self.policy.max_waiting, deadline)?, - ); + self.ticket = Some(self.queue.enqueue_with_priority( + keys, + self.policy.max_waiting, + deadline, + priority, + )?); } if let Some(ticket) = &self.ticket { ticket.retry().await?; diff --git a/backend/crates/gateway-core/src/diagnostics/selection.rs b/backend/crates/gateway-core/src/diagnostics/selection.rs index 7d146771e..7d805818a 100644 --- a/backend/crates/gateway-core/src/diagnostics/selection.rs +++ b/backend/crates/gateway-core/src/diagnostics/selection.rs @@ -5,7 +5,7 @@ use serde_json::json; use super::TraceContext; use crate::account::{ AccountCandidate, AccountSelection, AccountSelectionContext, AccountSelector, RotationStrategy, - SMART_SCORE_TOLERANCE, smart_score, + smart_score, }; // 单个 trace 事件最多 4 KiB；优先保留实际选中的账号。 @@ -50,7 +50,7 @@ impl TraceContext { "failureRateBasisPoints": signals.failure_rate_basis_points, "firstOutputLatencyMs": signals.first_output_latency_ms, "smartScore": smart.then(|| smart_score( - candidate, context.policy.max_concurrent_per_account() + candidate, context.policy.max_concurrent_per_account(), context.policy.smart_scheduling(), context.now )), }) }) @@ -63,7 +63,7 @@ impl TraceContext { "preferredAccountId": context.preferred_account.as_ref().map(|id| id.as_str()), "preferredResult": selection.as_ref().map(|s| format!("{:?}", s.preferred())), "roundRobinCursor": context.round_robin_cursor, - "smartScoreTolerance": smart.then_some(SMART_SCORE_TOLERANCE), + "smartScoreTolerance": smart.then(|| context.policy.smart_scheduling().score_tolerance()), "candidateCount": candidates.len(), "omittedCandidates": candidates.len().saturating_sub(observations.len()), "candidates": observations, diff --git a/backend/crates/gateway-core/src/routing/snapshot.rs b/backend/crates/gateway-core/src/routing/snapshot.rs index d244e1463..d3b89c118 100644 --- a/backend/crates/gateway-core/src/routing/snapshot.rs +++ b/backend/crates/gateway-core/src/routing/snapshot.rs @@ -42,6 +42,7 @@ pub struct SnapshotSettingsFacts { concurrency_wait_timeout_seconds: u32, responses_max_decompressed_body_bytes: u64, request_interval_ms: u64, + smart_scheduling: crate::account::SmartSchedulingConfig, rotation_strategy: String, model_mappings: BTreeMap, min_codex_desktop_version: Option, @@ -49,6 +50,15 @@ pub struct SnapshotSettingsFacts { } impl SnapshotSettingsFacts { + #[must_use] + pub const fn with_smart_scheduling( + mut self, + config: crate::account::SmartSchedulingConfig, + ) -> Self { + self.smart_scheduling = config; + self + } + #[must_use] pub fn with_pricing(mut self, pricing: crate::metering::PricingOverrides) -> Self { self.pricing = Arc::new(pricing); @@ -114,6 +124,7 @@ impl SnapshotSettingsFacts { concurrency_wait_timeout_seconds: 30, responses_max_decompressed_body_bytes: 64 * 1024 * 1024, request_interval_ms, + smart_scheduling: crate::account::SmartSchedulingConfig::default(), rotation_strategy: rotation_strategy.into(), model_mappings, min_codex_desktop_version, @@ -592,6 +603,7 @@ async fn compile_runtime_snapshot( AccountConcurrency::new(facts.settings.max_concurrent_per_account), Duration::from_millis(facts.settings.request_interval_ms), ) + .with_smart_scheduling(facts.settings.smart_scheduling) .with_queue(ConcurrencyQueuePolicy { max_waiting: facts.settings.max_waiting_per_account, timeout: queue_timeout, diff --git a/backend/crates/gateway-core/tests/account/mod.rs b/backend/crates/gateway-core/tests/account/mod.rs index edf98bcc2..ec72805a2 100644 --- a/backend/crates/gateway-core/tests/account/mod.rs +++ b/backend/crates/gateway-core/tests/account/mod.rs @@ -3,6 +3,7 @@ mod model_access; mod proxy; mod selection; mod selection_diagnostics; +mod smart_scheduling; use std::collections::{BTreeMap, BTreeSet}; use std::num::NonZeroU32; diff --git a/backend/crates/gateway-core/tests/account/smart_scheduling.rs b/backend/crates/gateway-core/tests/account/smart_scheduling.rs new file mode 100644 index 000000000..c6aa6cab4 --- /dev/null +++ b/backend/crates/gateway-core/tests/account/smart_scheduling.rs @@ -0,0 +1,349 @@ +use gateway_core::account::{ + AccountSelector, PreferredAccountSelection, RotationStrategy, SmartSchedulingConfig, +}; +use serde_json::json; +use std::time::{Duration, Instant, SystemTime}; + +use futures::FutureExt; +use gateway_core::concurrency::{ + CapacityWait, ConcurrencyQueuePolicy, ConcurrencyWaitBudget, ConcurrencyWaitQueue, +}; + +use super::{candidate_with_concurrency, context, weighted_candidate}; + +#[test] +fn smart_config_rejects_invalid_weights_and_incomplete_wire_values() { + for value in [-0.1, 10.1, 0.01, f64::NAN, f64::INFINITY] { + for dimension in 0..6 { + let mut weights = [1.0; 6]; + weights[dimension] = value; + assert!(SmartSchedulingConfig::new(weights, false).is_err()); + } + } + assert!(SmartSchedulingConfig::new([0.0; 6], true).is_err()); + for value in [json!({}), json!(null), json!({"loadWeight": 1.0})] { + assert!(serde_json::from_value::(value).is_err()); + } + let config = SmartSchedulingConfig::new([0.0, 0.1, 9.9, 10.0, 1.2, 2.3], true).unwrap(); + let encoded = serde_json::to_value(config).unwrap(); + assert_eq!( + serde_json::from_value::(encoded).unwrap(), + config + ); +} + +#[test] +fn reset_weight_rewards_only_future_resets_and_preserves_capacity_checks() { + let mut context = context(RotationStrategy::Smart); + context.policy = context.policy.with_smart_scheduling( + SmartSchedulingConfig::new([0.0, 0.0, 0.0, 0.0, 1.0, 0.0], false).unwrap(), + ); + let mut candidates = [ + candidate_with_concurrency("acct_unknown", 0, 2), + candidate_with_concurrency("acct_past", 0, 2), + candidate_with_concurrency("acct_now", 0, 2), + candidate_with_concurrency("acct_near", 0, 2), + candidate_with_concurrency("acct_far", 0, 2), + ]; + candidates[1].signals.quota_reset_at = Some(context.now - Duration::from_secs(1)); + candidates[2].signals.quota_reset_at = Some(context.now); + candidates[3].signals.quota_reset_at = Some(context.now + Duration::from_secs(3600)); + candidates[4].signals.quota_reset_at = Some(context.now + Duration::from_secs(86_400)); + assert_eq!( + AccountSelector + .select(&candidates, &context) + .unwrap() + .candidate() + .account + .id() + .as_str(), + "acct_near" + ); + let trace = gateway_core::diagnostics::TraceContext::new("req_reset_scores"); + trace.account_selection(&candidates, &context, None); + let snapshot = trace.snapshot().unwrap(); + let scores = snapshot["events"][0]["data"]["candidates"] + .as_array() + .unwrap(); + for (candidate, expected) in scores.iter().zip([0.0, 0.0, 0.0, 0.5, 0.04]) { + assert_eq!(candidate["smartScore"].as_f64().unwrap(), expected); + } + candidates[3].signals.in_flight = 2; + assert_eq!( + AccountSelector + .select(&candidates, &context) + .unwrap() + .candidate() + .account + .id() + .as_str(), + "acct_far" + ); + candidates[3].signals.in_flight = 0; + context.now += Duration::from_secs(3600); + assert_eq!( + AccountSelector + .select(&candidates, &context) + .unwrap() + .candidate() + .account + .id() + .as_str(), + "acct_far" + ); +} + +#[test] +fn queue_weight_balances_live_pressure_and_health_only_when_enabled_for_smart() { + for (strategy, queue_weight, expects_shortest) in [ + (RotationStrategy::Smart, 0.0, true), + (RotationStrategy::Smart, 0.2, false), + (RotationStrategy::Smart, 2.0, true), + (RotationStrategy::RoundRobin, 0.2, true), + (RotationStrategy::QuotaResetPriority, 0.2, true), + (RotationStrategy::Sticky, 0.2, true), + ] { + let mut candidates = [ + candidate_with_concurrency("acct_healthy", 2, 2), + candidate_with_concurrency("acct_shortest", 2, 2), + ]; + candidates[1].signals.failure_rate_basis_points = Some(6000); + let mut context = context(strategy); + context.policy = context.policy.with_smart_scheduling( + SmartSchedulingConfig::new([0.0, 0.0, 1.0, 0.0, 0.0, queue_weight], false).unwrap(), + ); + let queue = ConcurrencyWaitQueue::default(); + let head = queue + .enqueue( + &[candidates[0].account.id().clone()], + 2, + Instant::now() + Duration::from_secs(2), + ) + .unwrap(); + let budget = ConcurrencyWaitBudget::default(); + let policy = ConcurrencyQueuePolicy { + max_waiting: 2, + timeout: Duration::from_secs(2), + }; + let mut waiting = CapacityWait::new( + &queue, + policy, + SystemTime::now() + Duration::from_secs(2), + &budget, + ); + let keys = AccountSelector.wait_candidates(&candidates, &context); + assert!( + AccountSelector + .wait_for_capacity(&mut waiting, &keys, &candidates, &context) + .now_or_never() + .is_none() + ); + assert_eq!( + queue.has_waiters(candidates[1].account.id()), + expects_shortest, + "{strategy:?}, {queue_weight}" + ); + // 新请求不能绕过已经选定的队列，取消请求同步归还位置。 + let newcomer = CapacityWait::new( + &queue, + policy, + SystemTime::now() + Duration::from_secs(2), + &budget, + ); + assert!(!newcomer.can_try(candidates[usize::from(expects_shortest)].account.id())); + drop(waiting); + assert!(!queue.has_waiters(candidates[1].account.id())); + drop(head); + assert!(!queue.has_waiters(candidates[0].account.id())); + } +} + +#[test] +fn queue_weight_does_not_widen_immediate_selection_tolerance() { + let mut candidates = [ + candidate_with_concurrency("acct_a", 0, 2), + candidate_with_concurrency("acct_b", 0, 2), + ]; + candidates[0].signals.quota_remaining_rank = Some(10000); + candidates[1].signals.quota_remaining_rank = Some(9000); + let mut context = context(RotationStrategy::Smart); + context.round_robin_cursor = 1; + for queue_weight in [0.0, 10.0] { + context.policy = context.policy.with_smart_scheduling( + SmartSchedulingConfig::new([0.0, 1.0, 0.0, 0.0, 0.0, queue_weight], false).unwrap(), + ); + assert_eq!( + AccountSelector + .select(&candidates, &context) + .unwrap() + .candidate() + .account + .id() + .as_str(), + "acct_a" + ); + } + context.policy = context.policy.with_smart_scheduling( + SmartSchedulingConfig::new([0.0, 0.0, 0.0, 0.0, 0.0, 1.0], false).unwrap(), + ); + // 仅开启排队系数时，立即可用的同权重账号按既有游标轮换。 + assert_eq!( + AccountSelector + .select(&candidates, &context) + .unwrap() + .candidate() + .account + .id() + .as_str(), + "acct_b" + ); +} + +#[test] +fn each_smart_weight_controls_its_signal_without_relaxing_capacity() { + let mut candidates = [ + candidate_with_concurrency("acct_idle", 0, 10), + candidate_with_concurrency("acct_healthy", 5, 10), + ]; + candidates[0].signals.quota_remaining_rank = Some(1_000); + candidates[1].signals.quota_remaining_rank = Some(9_000); + candidates[0].signals.failure_rate_basis_points = Some(5_000); + candidates[1].signals.failure_rate_basis_points = Some(0); + candidates[0].signals.first_output_latency_ms = Some(30_000); + candidates[1].signals.first_output_latency_ms = Some(1_000); + for dimension in 0..4 { + let mut weights = [0.0; 6]; + weights[dimension] = 1.0; + let mut context = context(RotationStrategy::Smart); + context.policy = context + .policy + .with_smart_scheduling(SmartSchedulingConfig::new(weights, false).unwrap()); + let selected = AccountSelector.select(&candidates, &context).unwrap(); + assert_eq!( + selected.candidate().account.id().as_str(), + if dimension == 0 { + "acct_idle" + } else { + "acct_healthy" + } + ); + let mut saturated = candidates.clone(); + saturated[1].signals.in_flight = 10; + assert_eq!( + AccountSelector + .select(&saturated, &context) + .unwrap() + .candidate() + .account + .id() + .as_str(), + "acct_idle" + ); + } +} + +#[test] +fn proportional_coefficients_preserve_near_best_rotation_and_diagnostics() { + let mut candidates = [ + candidate_with_concurrency("acct_a", 0, 100), + candidate_with_concurrency("acct_b", 0, 100), + candidate_with_concurrency("acct_c", 0, 100), + ]; + // 默认额度权重 0.8 下，9375 正好位于容差边界，9374 刚好超出。 + for (candidate, quota) in candidates.iter_mut().zip([10_000, 9_375, 9_374]) { + candidate.signals.quota_remaining_rank = Some(quota); + } + for scale in [1.0, 2.0, 5.0] { + let mut context = context(RotationStrategy::Smart); + context.policy = context.policy.with_smart_scheduling( + SmartSchedulingConfig::new([scale, 0.8 * scale, scale, 0.5 * scale, 0.0, 0.0], false) + .unwrap(), + ); + let mut ids = Vec::new(); + for cursor in 0..4 { + candidates.reverse(); + context.round_robin_cursor = cursor; + ids.push( + AccountSelector + .select(&candidates, &context) + .unwrap() + .candidate() + .account + .id() + .as_str() + .to_owned(), + ); + } + assert_eq!(ids, ["acct_a", "acct_b", "acct_a", "acct_b"]); + let trace = gateway_core::diagnostics::TraceContext::new("req_scaled_weights"); + trace.account_selection(&candidates, &context, None); + let snapshot = trace.snapshot().unwrap(); + let tolerance = snapshot["events"][0]["data"]["smartScoreTolerance"] + .as_f64() + .unwrap(); + assert!((tolerance - 0.05 * scale).abs() < 1e-12); + } +} + +#[test] +fn switchback_is_opt_in_and_only_affects_smart_soft_affinity() { + let mut candidates = [ + weighted_candidate("acct_affinity", 10, 0), + weighted_candidate("acct_primary", 100, 0), + ]; + for strategy in [ + RotationStrategy::Smart, + RotationStrategy::Sticky, + RotationStrategy::RoundRobin, + RotationStrategy::QuotaResetPriority, + ] { + let mut context = context(strategy); + context.preferred_account = Some(candidates[0].account.id().clone()); + context.preferred_account_overrides_weight = true; + for enabled in [false, true] { + context.policy = context.policy.with_smart_scheduling( + SmartSchedulingConfig::new([1.0, 0.8, 1.0, 0.5, 0.0, 0.0], enabled).unwrap(), + ); + let selected = AccountSelector.select(&candidates, &context).unwrap(); + let should_switch = enabled && strategy == RotationStrategy::Smart; + assert_eq!( + selected.candidate().account.id(), + candidates[usize::from(should_switch)].account.id() + ); + } + } + let mut context = context(RotationStrategy::Smart); + context.policy = context.policy.with_smart_scheduling( + SmartSchedulingConfig::new([1.0, 0.8, 1.0, 0.5, 0.0, 0.0], true).unwrap(), + ); + context.preferred_account_overrides_weight = true; + context.preferred_account = Some(candidates[0].account.id().clone()); + candidates[1].signals.in_flight = 3; + assert_eq!( + AccountSelector + .select(&candidates, &context) + .unwrap() + .preferred(), + PreferredAccountSelection::Hit + ); + candidates[1].signals.in_flight = 0; + context + .excluded_accounts + .insert(candidates[1].account.id().clone()); + assert_eq!( + AccountSelector + .select(&candidates, &context) + .unwrap() + .preferred(), + PreferredAccountSelection::Hit + ); + context.excluded_accounts.clear(); + candidates[0] = weighted_candidate("acct_affinity", 100, 2); + assert_eq!( + AccountSelector + .select(&candidates, &context) + .unwrap() + .preferred(), + PreferredAccountSelection::Hit + ); +} diff --git a/backend/crates/gateway-core/tests/concurrency.rs b/backend/crates/gateway-core/tests/concurrency.rs index 8c22c5522..6731f6abf 100644 --- a/backend/crates/gateway-core/tests/concurrency.rs +++ b/backend/crates/gateway-core/tests/concurrency.rs @@ -159,3 +159,75 @@ fn creating_a_waiter_does_not_start_the_budget_before_queueing() { waiting.wait(&["account"]).await.unwrap(); }); } + +#[test] +fn custom_priority_reads_live_counts_and_preserves_limits_and_fifo() { + block_on(async { + let queue = ConcurrencyWaitQueue::default(); + let deadline = Instant::now() + Duration::from_secs(2); + let head = queue.enqueue(&["a"], 2, deadline).unwrap(); + let mut observed = Vec::new(); + let tail = queue + .enqueue_with_priority(&["a", "b"], 2, deadline, |key, count| { + observed.push((*key, count)); + if *key == "a" { 10.0 } else { 0.0 } + }) + .unwrap(); + assert_eq!(observed, [("a", 1), ("b", 0)]); + assert_eq!(*tail.key(), "a"); + assert!(tail.turn().now_or_never().is_none()); + // 满队列不参与评分，不能因偏好而突破上限。 + let fallback = queue + .enqueue_with_priority(&["a", "b"], 2, deadline, |key, count| { + assert_eq!((*key, count), ("b", 0)); + 0.0 + }) + .unwrap(); + assert_eq!(*fallback.key(), "b"); + drop(head); + tail.turn().await.unwrap(); + drop((tail, fallback)); + let tie = queue + .enqueue_with_priority(&["b", "a"], 2, deadline, |_, _| 1.0) + .unwrap(); + assert_eq!(*tie.key(), "b"); + }); +} + +#[test] +fn changed_scores_keep_existing_position_until_account_becomes_ineligible() { + block_on(async { + let queue = ConcurrencyWaitQueue::default(); + let budget = ConcurrencyWaitBudget::default(); + let policy = ConcurrencyQueuePolicy { + max_waiting: 1, + timeout: Duration::from_secs(2), + }; + let mut waiting = CapacityWait::new( + &queue, + policy, + SystemTime::now() + Duration::from_secs(2), + &budget, + ); + waiting + .wait_with_priority(&["a", "b"], |key, _| if *key == "a" { 1.0 } else { 0.0 }) + .await + .unwrap(); + waiting + .wait_with_priority(&["a", "b"], |_, _| { + panic!("existing ticket must not be rescored") + }) + .await + .unwrap(); + assert!(queue.has_waiters(&"a")); + assert!(!queue.has_waiters(&"b")); + waiting + .wait_with_priority(&["b"], |_, _| 1.0) + .await + .unwrap(); + assert!(!queue.has_waiters(&"a")); + assert!(queue.has_waiters(&"b")); + drop(waiting); + assert!(!queue.has_waiters(&"b")); + }); +} diff --git a/backend/crates/gateway-core/tests/engine/policy.rs b/backend/crates/gateway-core/tests/engine/policy.rs index f460ce7d6..a113c5ef4 100644 --- a/backend/crates/gateway-core/tests/engine/policy.rs +++ b/backend/crates/gateway-core/tests/engine/policy.rs @@ -169,6 +169,52 @@ fn delegated_scheduler_keeps_builtin_affinity() { }); } +#[test] +fn smart_switchback_does_not_override_explicit_plugin_choices_or_their_telemetry() { + block_on(async { + let mut candidates = [candidate("acct_a"), candidate("acct_b")]; + candidates[0].account = candidates[0] + .account + .clone() + .with_scheduling(None, AccountWeight::new(10).unwrap()); + let mut selection_context = context(); + selection_context.policy = AccountSelectionPolicy::new( + RotationStrategy::Smart, + NonZeroU32::new(2).unwrap(), + Duration::ZERO, + ) + .with_smart_scheduling( + gateway_core::account::SmartSchedulingConfig::new([1.0, 0.8, 1.0, 0.5, 0.0, 0.0], true) + .unwrap(), + ); + for chosen in [Some("acct_a"), Some("acct_b"), None] { + let selected = policy(chosen) + .select_account( + NonZeroU32::new(1).unwrap(), + &ProviderKind::new("openai").unwrap(), + None, + &candidates, + &selection_context, + ) + .await + .unwrap() + .unwrap(); + assert_eq!( + selected.candidate().account.id().as_str(), + chosen.unwrap_or("acct_b") + ); + let expected = match chosen { + Some("acct_a") => PreferredAccountSelection::Hit, + Some(_) => PreferredAccountSelection::OverriddenByPolicy, + None => PreferredAccountSelection::Blocked( + gateway_core::account::AccountSchedulingBlocker::LowerWeight, + ), + }; + assert_eq!(selected.preferred(), expected); + } + }); +} + #[test] fn scheduler_cannot_select_outside_the_hard_bound_candidate_pool() { block_on(async { diff --git a/backend/crates/gateway-core/tests/routing/snapshot.rs b/backend/crates/gateway-core/tests/routing/snapshot.rs index 0e1cf9568..60c3bd23c 100644 --- a/backend/crates/gateway-core/tests/routing/snapshot.rs +++ b/backend/crates/gateway-core/tests/routing/snapshot.rs @@ -362,6 +362,49 @@ fn compiler_accepts_unlimited_default_account_concurrency() { ); } +#[test] +fn compiled_plans_keep_their_smart_config_after_a_new_snapshot_is_built() { + use gateway_core::account::SmartSchedulingConfig; + let configs = [ + SmartSchedulingConfig::default(), + SmartSchedulingConfig::new([0.0, 2.0, 1.0, 0.5, 1.2, 2.3], true).unwrap(), + ]; + let mut plans = Vec::new(); + for (index, config) in configs.into_iter().enumerate() { + let facts = SnapshotFacts::new( + revision(index as u64 + 1), + revision(index as u64 + 1), + SnapshotSettingsFacts::new(3, 0, "smart", BTreeMap::new(), None, None) + .with_smart_scheduling(config), + vec![], + vec![], + vec![SnapshotProviderAccountFacts::new( + gateway_core::account::ProviderAccountId::new("acct_config").unwrap(), + "alpha", + )], + vec![], + ); + let compiler = RuntimeSnapshotCompiler::new( + Arc::new(TestSnapshotStore::new(Ok(facts))), + Arc::new(TestCatalog::Unavailable), + ); + let snapshot = block_on(compiler.compile()).unwrap(); + plans.push( + snapshot + .plan( + &PublicModelId::new("any-model").unwrap(), + &super::operation(), + snapshot.all_account_scope(), + &Default::default(), + ) + .unwrap(), + ); + } + for (plan, config) in plans.iter().zip(configs) { + assert_eq!(plan.account_selection_policy().smart_scheduling(), config); + } +} + #[test] fn withdrawn_provider_keeps_accounts_without_becoming_a_route_or_expanding_group_scope() { use gateway_core::{ diff --git a/backend/crates/gateway-store/src/admin_adapter.rs b/backend/crates/gateway-store/src/admin_adapter.rs index 5f340fbee..aa46521e1 100644 --- a/backend/crates/gateway-store/src/admin_adapter.rs +++ b/backend/crates/gateway-store/src/admin_adapter.rs @@ -99,6 +99,7 @@ impl SettingsStore for AdminSettingsStoreAdapter { concurrency_wait_timeout_seconds: command.concurrency_wait_timeout_seconds, responses_max_decompressed_body_bytes: command .responses_max_decompressed_body_bytes, + smart_scheduling: command.smart_scheduling, rotation_strategy: command.rotation_strategy.as_str().to_owned(), model_mappings: store_model_mappings(command.model_mappings), min_codex_desktop_version: command.min_codex_desktop_version, @@ -137,6 +138,7 @@ impl SettingsStore for AdminSettingsStoreAdapter { "concurrency_wait_timeout_seconds".to_owned(), "responses_max_decompressed_body_bytes".to_owned(), "rotation_strategy".to_owned(), + "smart_scheduling_json".to_owned(), "min_codex_desktop_version".to_owned(), "min_codex_cli_version".to_owned(), "retention".to_owned(), @@ -247,6 +249,7 @@ pub(crate) fn admin_runtime_settings( max_waiting_per_account: settings.max_waiting_per_account, concurrency_wait_timeout_seconds: settings.concurrency_wait_timeout_seconds, responses_max_decompressed_body_bytes: settings.responses_max_decompressed_body_bytes, + smart_scheduling: settings.smart_scheduling, rotation_strategy, min_codex_desktop_version: settings.min_codex_desktop_version, min_codex_cli_version: settings.min_codex_cli_version, diff --git a/backend/crates/gateway-store/src/postgres/runtime_settings.rs b/backend/crates/gateway-store/src/postgres/runtime_settings.rs index 8d4c28d94..349f81df7 100644 --- a/backend/crates/gateway-store/src/postgres/runtime_settings.rs +++ b/backend/crates/gateway-store/src/postgres/runtime_settings.rs @@ -31,6 +31,7 @@ pub struct RuntimeSettings { pub max_waiting_per_account: u32, pub concurrency_wait_timeout_seconds: u32, pub responses_max_decompressed_body_bytes: u64, + pub smart_scheduling: gateway_core::account::SmartSchedulingConfig, pub rotation_strategy: String, pub request_location_enabled: bool, pub request_location: gateway_core::account::RequestLocation, @@ -133,6 +134,7 @@ pub struct RuntimeSettingsUpdate { pub max_waiting_per_account: u32, pub concurrency_wait_timeout_seconds: u32, pub responses_max_decompressed_body_bytes: u64, + pub smart_scheduling: gateway_core::account::SmartSchedulingConfig, pub rotation_strategy: String, pub request_location_enabled: bool, pub request_location: gateway_core::account::RequestLocation, @@ -263,7 +265,7 @@ pub(crate) async fn load_runtime_settings_from_pool(pool: &PgPool) -> StoreResul let row = sqlx::query_as::<_, RuntimeSettingsRow>( "select provider_request_profiles_json, config_revision, admin_api_key, refresh_margin_seconds, request_location_json, request_location_enabled, refresh_concurrency, max_concurrent_per_account, request_interval_ms, - rotation_strategy, model_mappings_json, usage_retention_days, ops_event_retention_days, + rotation_strategy, smart_scheduling_json, model_mappings_json, usage_retention_days, ops_event_retention_days, audit_retention_days, min_codex_desktop_version, min_codex_cli_version, updated_at, responses_max_decompressed_body_bytes, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds, account_auto_freeze_enabled, account_auto_freeze_threshold, @@ -415,7 +417,7 @@ pub(crate) async fn load_runtime_settings_in_transaction( let row = sqlx::query_as::<_, RuntimeSettingsRow>( "select provider_request_profiles_json, config_revision, admin_api_key, refresh_margin_seconds, request_location_json, request_location_enabled, refresh_concurrency, max_concurrent_per_account, request_interval_ms, - rotation_strategy, model_mappings_json, usage_retention_days, ops_event_retention_days, + rotation_strategy, smart_scheduling_json, model_mappings_json, usage_retention_days, ops_event_retention_days, audit_retention_days, min_codex_desktop_version, min_codex_cli_version, updated_at, responses_max_decompressed_body_bytes, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds, account_auto_freeze_enabled, account_auto_freeze_threshold, @@ -492,6 +494,7 @@ pub(crate) async fn update_runtime_settings_in_transaction( account_warmup_enabled = $28, account_warmup_schedule_time = $29, account_warmup_model = $30, + smart_scheduling_json = $31, updated_at = now() where id = 1 returning config_revision", @@ -538,6 +541,7 @@ pub(crate) async fn update_runtime_settings_in_transaction( .bind(update.account_warmup_enabled) .bind(&update.account_warmup_schedule_time) .bind(update.account_warmup_model.as_deref()) + .bind(sqlx::types::Json(update.smart_scheduling)) .fetch_optional(&mut **transaction) .await .map_err(|_| postgres_unavailable("update runtime settings in transaction"))? @@ -596,6 +600,7 @@ struct RuntimeSettingsRow { refresh_concurrency: i64, max_concurrent_per_account: i64, request_interval_ms: i64, + smart_scheduling_json: sqlx::types::Json, rotation_strategy: String, request_location_enabled: bool, request_location_json: sqlx::types::Json, @@ -644,6 +649,7 @@ fn runtime_settings_from_row(row: RuntimeSettingsRow) -> StoreResult, pub min_codex_desktop_version: Option, @@ -162,6 +163,7 @@ impl SnapshotStorePort for PgRuntimeSnapshotRepository { .with_responses_max_decompressed_body_bytes( data.settings.responses_max_decompressed_body_bytes, ) + .with_smart_scheduling(data.settings.smart_scheduling) .with_request_profiles(data.settings.request_profiles) .with_pricing(data.settings.pricing) .with_request_location( @@ -254,6 +256,7 @@ struct SnapshotSettingsRow { refresh_concurrency: i64, max_concurrent_per_account: i64, request_interval_ms: i64, + smart_scheduling_json: sqlx::types::Json, rotation_strategy: String, model_mappings_json: sqlx::types::Json>, min_codex_desktop_version: Option, @@ -272,7 +275,7 @@ async fn load_settings( transaction: &mut Transaction<'_, Postgres>, ) -> StoreResult<(Revision, SnapshotRuntimeSettings)> { let row = sqlx::query_as::<_, SnapshotSettingsRow>( - "select config_revision, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, model_mappings_json, min_codex_desktop_version, min_codex_cli_version, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds, request_location_json, request_location_enabled, responses_max_decompressed_body_bytes, provider_request_profiles_json, pricing_overrides_json, pricing_synced_json from runtime_settings where id = 1", + "select config_revision, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, smart_scheduling_json, model_mappings_json, min_codex_desktop_version, min_codex_cli_version, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds, request_location_json, request_location_enabled, responses_max_decompressed_body_bytes, provider_request_profiles_json, pricing_overrides_json, pricing_synced_json from runtime_settings where id = 1", ) .fetch_optional(&mut **transaction) .await @@ -302,6 +305,7 @@ async fn load_settings( refresh_concurrency: to_u32(row.refresh_concurrency)?, max_concurrent_per_account: to_u32(row.max_concurrent_per_account)?, request_interval_ms: to_u64(row.request_interval_ms)?, + smart_scheduling: row.smart_scheduling_json.0, rotation_strategy: row.rotation_strategy, model_mappings: row.model_mappings_json.0, min_codex_desktop_version: row.min_codex_desktop_version, diff --git a/backend/crates/gateway-store/tests/postgres/runtime_settings.rs b/backend/crates/gateway-store/tests/postgres/runtime_settings.rs index f80ebaa75..6819749f0 100644 --- a/backend/crates/gateway-store/tests/postgres/runtime_settings.rs +++ b/backend/crates/gateway-store/tests/postgres/runtime_settings.rs @@ -21,6 +21,7 @@ fn settings_with_margin(refresh_margin_seconds: u64) -> RuntimeSettingsUpdate { max_waiting_per_account: 0, concurrency_wait_timeout_seconds: 30, responses_max_decompressed_body_bytes: 64 * 1024 * 1024, + smart_scheduling: gateway_core::account::SmartSchedulingConfig::default(), rotation_strategy: "smart".to_owned(), model_mappings: BTreeMap::from([ ("gpt-5.4".to_owned(), "gpt-5.5".to_owned()), @@ -50,6 +51,39 @@ fn runtime_settings_keep_account_rotation_global() { assert!(settings.validate().is_ok()); } +#[tokio::test] +async fn smart_settings_upgrade_preserves_selection_and_publishes_custom_config() { + use gateway_core::account::SmartSchedulingConfig; + use gateway_store::postgres::{PgRuntimeSnapshotRepository, RuntimeSnapshotRepository}; + let Some(database) = TestDatabase::create_through("smart_config", 19).await else { + return; + }; + // 升级前不能用包含新列的 Repository，直接写入旧版本已有字段。 + sqlx::query("update runtime_settings set rotation_strategy = 'sticky', refresh_margin_seconds = 3600 where id = 1") + .execute(&database.pool).await.unwrap(); + super::TEST_MIGRATOR.run(&database.pool).await.unwrap(); + let repository = PgRuntimeSettingsRepository::new(database.pool.clone()); + let before = repository.load_runtime_settings().await.unwrap(); + assert_eq!(before.rotation_strategy, "sticky"); + assert_eq!(before.refresh_margin_seconds, 3600); + assert_eq!(before.smart_scheduling, SmartSchedulingConfig::default()); + let mut update = settings_with_margin(3600); + update.smart_scheduling = + SmartSchedulingConfig::new([0.0, 2.0, 1.0, 0.5, 1.2, 2.3], true).unwrap(); + let expected = update.smart_scheduling; + repository.update_runtime_settings(update).await.unwrap(); + let reloaded = repository.load_runtime_settings().await.unwrap(); + assert_eq!(reloaded.smart_scheduling, expected); + let snapshot = PgRuntimeSnapshotRepository::new(database.pool.clone()) + .load_runtime_snapshot() + .await + .unwrap(); + assert_eq!(snapshot.settings.smart_scheduling, expected); + assert!(snapshot.config_revision > before.config_revision); + assert_eq!(before.smart_scheduling, SmartSchedulingConfig::default()); + database.close().await; +} + #[test] fn runtime_settings_require_model_when_warmup_is_enabled() { let settings = RuntimeSettingsUpdate { diff --git a/backend/crates/providers/openai/src/credential/selector.rs b/backend/crates/providers/openai/src/credential/selector.rs index ce7069695..a48d2f209 100644 --- a/backend/crates/providers/openai/src/credential/selector.rs +++ b/backend/crates/providers/openai/src/credential/selector.rs @@ -639,16 +639,24 @@ impl CodexCredentialSelector { ); let Some(selection) = selection else { if !diagnostic && queue_policy.max_waiting > 0 && !wait_candidates.is_empty() { - waiting.wait(&wait_candidates).await.map_err(|error| { - tracing::info!( - request_id = request.attempt.request_id().as_str(), - queue_layer = "account", - queue_wait_ms = waiting.elapsed().as_millis() as u64, - reason = %error, - "OpenAI 账号排队请求被拒绝" - ); - CredentialSelectionError::QueueRejected(error) - })?; + AccountSelector + .wait_for_capacity( + &mut waiting, + &wait_candidates, + &candidates, + &wait_context, + ) + .await + .map_err(|error| { + tracing::info!( + request_id = request.attempt.request_id().as_str(), + queue_layer = "account", + queue_wait_ms = waiting.elapsed().as_millis() as u64, + reason = %error, + "OpenAI 账号排队请求被拒绝" + ); + CredentialSelectionError::QueueRejected(error) + })?; continue 'capacity; } // 只判断本次账号范围；空池、认证失效和租约失败不能伪装成额度耗尽。 diff --git a/backend/crates/providers/openai/tests/credential/contract.rs b/backend/crates/providers/openai/tests/credential/contract.rs index b6b9f09fb..05ee8c584 100644 --- a/backend/crates/providers/openai/tests/credential/contract.rs +++ b/backend/crates/providers/openai/tests/credential/contract.rs @@ -124,6 +124,80 @@ fn round_robin_attempt() -> AttemptContext { ) } +#[tokio::test] +async fn smart_reselection_respects_native_and_required_account_boundaries() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_original", "test-original"); + create_account(&store, "acct_primary", "test-primary"); + store.set_scheduling("acct_original", None, AccountWeight::new(10).unwrap()); + store.set_scheduling("acct_primary", None, AccountWeight::new(100).unwrap()); + let original_id = ProviderAccountId::new("acct_original").unwrap(); + let provider = ProviderKind::new("openai").unwrap(); + let affinity = Arc::new(MemorySessionAffinity::default()); + let key = ProviderSessionAffinityKey::try_new("smart-config-test").unwrap(); + affinity + .bind(&provider, &key, &original_id, Duration::from_secs(60)) + .await + .unwrap(); + let selector = + selector_with_affinity(&store, Arc::new(TestLeaseCoordinator::default()), affinity); + let request_url = Url::parse(OFFICIAL_CODEX_BASE_URL).unwrap(); + for enabled in [false, true] { + for binding in ["soft", "required", "native"] { + let continuation = (binding == "native").then(|| { + ContinuationBinding::Pinned(NativeContinuationPin::new( + PreviousResponseId::new("client_response"), + PreviousResponseId::new("upstream_response"), + ClientApiKeyId::new("key_codex_contract").unwrap(), + provider.clone(), + original_id.clone(), + )) + }); + let attempt = AttemptContext::new( + RequestAttemptContext::new( + ModelRequestId::new("req_smart_scope").unwrap(), + ClientApiKeyId::new("key_codex_contract").unwrap(), + ), + NonZeroU32::new(1).unwrap(), + SystemTime::now() + Duration::from_secs(30), + account_policy().with_smart_scheduling( + gateway_core::account::SmartSchedulingConfig::new( + [1.0, 0.8, 1.0, 0.5, 1.0, 1.0], + enabled, + ) + .unwrap(), + ), + AccountAttemptContext::new( + BTreeSet::new(), + (binding == "required").then(|| original_id.clone()), + None, + ) + .with_account_scope(contract_account_scope()), + continuation, + CancellationToken::new(), + ); + let selected = selector + .select(&SelectCodexCredential { + upstream_model: "gpt-5.4", + request_url: &request_url, + attempt: &attempt, + session_affinity_key: Some(&key), + }) + .await + .unwrap(); + assert_eq!( + selected.account_id().as_str(), + if enabled && binding == "soft" { + "acct_primary" + } else { + "acct_original" + }, + "{binding}, switchback={enabled}" + ); + } + } +} + fn selector( store: &Arc, leases: Arc, @@ -1776,7 +1850,14 @@ fn queued_attempt(timeout: Duration) -> AttemptContext { .with_queue(gateway_core::concurrency::ConcurrencyQueuePolicy { max_waiting: 1, timeout, - }), + }) + .with_smart_scheduling( + gateway_core::account::SmartSchedulingConfig::new( + [1.0, 0.8, 1.0, 0.5, 1.0, 1.0], + false, + ) + .unwrap(), + ), AccountAttemptContext::new(BTreeSet::new(), None, None) .with_account_scope(contract_account_scope()), None, @@ -1784,6 +1865,155 @@ fn queued_attempt(timeout: Duration) -> AttemptContext { ) } +#[test] +fn smart_queue_weight_changes_the_selected_wait_queue() { + use futures::FutureExt; + for queue_weight in [0.0, 0.2] { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_primary", "test-healthy"); + create_account(&store, "acct_other", "test-unhealthy"); + let leases = Arc::new(TestLeaseCoordinator::default()); + *leases.busy.lock().unwrap() = true; + let feedback = Arc::new(AccountFeedbackStats::default()); + for _ in 0..4 { + feedback.report( + &ProviderKind::new("openai").unwrap(), + &ProviderAccountId::new("acct_other").unwrap(), + AccountAttemptFeedback::Failed { + first_output_ms: None, + }, + ); + } + let selector = selector_with_runtime( + &store, + leases, + Arc::new(MemorySessionAffinity::default()), + feedback, + Arc::new(MemoryCooldownPort::new()), + ); + let make_attempt = |required| { + AttemptContext::new( + RequestAttemptContext::new( + ModelRequestId::new("req_queue_weight").unwrap(), + ClientApiKeyId::new("key_codex_contract").unwrap(), + ), + NonZeroU32::new(1).unwrap(), + SystemTime::now() + Duration::from_secs(5), + account_policy() + .with_queue(gateway_core::concurrency::ConcurrencyQueuePolicy { + max_waiting: 2, + timeout: Duration::from_secs(2), + }) + .with_smart_scheduling( + gateway_core::account::SmartSchedulingConfig::new( + [0.0, 0.0, 1.0, 0.0, 0.0, queue_weight], + false, + ) + .unwrap(), + ), + AccountAttemptContext::new(BTreeSet::new(), required, None) + .with_account_scope(contract_account_scope()), + None, + CancellationToken::new(), + ) + }; + let head_attempt = make_attempt(Some(ProviderAccountId::new("acct_primary").unwrap())); + let next_attempt = make_attempt(None); + let url = Url::parse(OFFICIAL_CODEX_BASE_URL).unwrap(); + let head_request = SelectCodexCredential { + upstream_model: "gpt-5.4", + request_url: &url, + attempt: &head_attempt, + session_affinity_key: None, + }; + let next_request = SelectCodexCredential { + attempt: &next_attempt, + ..head_request + }; + let mut head = Box::pin(selector.select(&head_request)); + assert!(head.as_mut().now_or_never().is_none()); + let mut next = Box::pin(selector.select(&next_request)); + assert!(next.as_mut().now_or_never().is_none()); + let mut probe = Box::pin(selector.select(&head_request)); + match probe.as_mut().now_or_never() { + Some(Err(CredentialSelectionError::QueueRejected( + gateway_core::concurrency::QueueRejection::Full, + ))) => assert!(queue_weight > 0.0), + None => assert_eq!(queue_weight, 0.0), + other => panic!("unexpected queue probe: {other:?}"), + } + } +} + +#[test] +fn smart_queue_keeps_native_and_required_account_pins_when_another_account_is_free() { + use futures::FutureExt; + for native in [false, true] { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_original", "test-pinned"); + create_account(&store, "acct_primary", "test-free"); + let original = ProviderAccountId::new("acct_original").unwrap(); + let leases = Arc::new(TestLeaseCoordinator::default()); + leases + .busy_accounts + .lock() + .unwrap() + .insert(original.clone()); + let selector = selector(&store, leases.clone()); + let attempt = AttemptContext::new( + RequestAttemptContext::new( + ModelRequestId::new("req_queue_pin").unwrap(), + ClientApiKeyId::new("key_codex_contract").unwrap(), + ), + NonZeroU32::new(1).unwrap(), + SystemTime::now() + Duration::from_secs(5), + account_policy() + .with_queue(gateway_core::concurrency::ConcurrencyQueuePolicy { + max_waiting: 1, + timeout: Duration::from_secs(2), + }) + .with_smart_scheduling( + gateway_core::account::SmartSchedulingConfig::new( + [1.0, 0.8, 1.0, 0.5, 10.0, 10.0], + true, + ) + .unwrap(), + ), + AccountAttemptContext::new(BTreeSet::new(), (!native).then(|| original.clone()), None) + .with_account_scope(contract_account_scope()), + native.then(|| { + ContinuationBinding::Pinned(NativeContinuationPin::new( + PreviousResponseId::new("client_response"), + PreviousResponseId::new("upstream_response"), + ClientApiKeyId::new("key_codex_contract").unwrap(), + ProviderKind::new("openai").unwrap(), + original.clone(), + )) + }), + CancellationToken::new(), + ); + let url = Url::parse(OFFICIAL_CODEX_BASE_URL).unwrap(); + let request = SelectCodexCredential { + upstream_model: "gpt-5.4", + request_url: &url, + attempt: &attempt, + session_affinity_key: None, + }; + let mut pending = Box::pin(selector.select(&request)); + assert!(pending.as_mut().now_or_never().is_none()); + leases.busy_accounts.lock().unwrap().clear(); + assert_eq!(block_on(pending).unwrap().account_id(), &original); + assert!( + leases + .requests + .lock() + .unwrap() + .iter() + .all(|request| request.account_id() == &original) + ); + } +} + #[test] fn account_queue_is_bounded_and_resumes_after_capacity_is_released() { use futures::FutureExt; diff --git a/backend/crates/providers/xai/src/credential/selector.rs b/backend/crates/providers/xai/src/credential/selector.rs index 7f8f7689f..735c50af9 100644 --- a/backend/crates/providers/xai/src/credential/selector.rs +++ b/backend/crates/providers/xai/src/credential/selector.rs @@ -354,15 +354,18 @@ impl GrokAccountSessionSelector { } if !diagnostic && queue_policy.max_waiting > 0 && !wait_candidates.is_empty() { - waiting.wait(&wait_candidates).await.map_err(|error| { - tracing::info!( - queue_layer = "account", - queue_wait_ms = waiting.elapsed().as_millis() as u64, - reason = %error, - "xAI 账号排队请求被拒绝" - ); - GrokSessionSelectorError::QueueRejected(error) - })?; + AccountSelector + .wait_for_capacity(&mut waiting, &wait_candidates, &candidates, &context) + .await + .map_err(|error| { + tracing::info!( + queue_layer = "account", + queue_wait_ms = waiting.elapsed().as_millis() as u64, + reason = %error, + "xAI 账号排队请求被拒绝" + ); + GrokSessionSelectorError::QueueRejected(error) + })?; continue; } if capacity_denied { diff --git a/backend/crates/providers/xai/tests/credential/selector.rs b/backend/crates/providers/xai/tests/credential/selector.rs index e693b74f6..680be2c86 100644 --- a/backend/crates/providers/xai/tests/credential/selector.rs +++ b/backend/crates/providers/xai/tests/credential/selector.rs @@ -1086,18 +1086,103 @@ fn queue_request(fixture: &SelectorFixture, timeout: Duration) -> GrokSessionSel request.upstream_model().clone(), BTreeSet::new(), None, - request.account_selection_policy().with_queue( - gateway_core::concurrency::ConcurrencyQueuePolicy { + request + .account_selection_policy() + .with_queue(gateway_core::concurrency::ConcurrencyQueuePolicy { max_waiting: 1, timeout, - }, - ), + }) + .with_smart_scheduling( + gateway_core::account::SmartSchedulingConfig::new( + [1.0, 0.8, 1.0, 0.5, 1.0, 1.0], + false, + ) + .unwrap(), + ), request.deadline(), request.account_scope().clone(), request.client_api_key_id().clone(), ) } +#[tokio::test] +async fn smart_queue_weight_changes_the_selected_wait_queue_and_respects_required_account() { + use futures::FutureExt; + for queue_weight in [0.0, 0.2] { + let fixture = SelectorFixture::new(&["healthy", "unhealthy"]).await; + let healthy = account_id("healthy"); + for signal in fixture.coordinator.signals.lock().unwrap().values_mut() { + signal.in_flight = 2; + } + for _ in 0..4 { + fixture.feedback.report( + &gateway_core::routing::ProviderKind::new("xai").unwrap(), + &account_id("unhealthy"), + AccountAttemptFeedback::Failed { + first_output_ms: None, + }, + ); + } + let make_request = |required| { + let request = fixture.request(BTreeSet::new()); + GrokSessionSelection::new( + request.upstream_model().clone(), + BTreeSet::new(), + required, + request + .account_selection_policy() + .with_queue(gateway_core::concurrency::ConcurrencyQueuePolicy { + max_waiting: 2, + timeout: Duration::from_secs(2), + }) + .with_smart_scheduling( + gateway_core::account::SmartSchedulingConfig::new( + [0.0, 0.0, 1.0, 0.0, 0.0, queue_weight], + false, + ) + .unwrap(), + ), + request.deadline(), + request.account_scope().clone(), + request.client_api_key_id().clone(), + ) + }; + let mut head = fixture.selector.select(make_request(Some(healthy.clone()))); + assert!(head.as_mut().now_or_never().is_none()); + let mut next = fixture.selector.select(make_request(None)); + assert!(next.as_mut().now_or_never().is_none()); + let mut probe = fixture.selector.select(make_request(Some(healthy.clone()))); + match probe.as_mut().now_or_never() { + Some(Err(GrokSessionSelectorError::QueueRejected( + gateway_core::concurrency::QueueRejection::Full, + ))) => assert!(queue_weight > 0.0), + None => assert_eq!(queue_weight, 0.0), + other => panic!("unexpected queue probe: {other:?}"), + } + drop((next, probe)); + // 其他账号先空闲不能解除 required account 限制。 + fixture + .coordinator + .signals + .lock() + .unwrap() + .get_mut(&account_id("unhealthy")) + .unwrap() + .in_flight = 0; + tokio::time::sleep(Duration::from_millis(110)).await; + assert!(head.as_mut().now_or_never().is_none()); + fixture + .coordinator + .signals + .lock() + .unwrap() + .get_mut(&healthy) + .unwrap() + .in_flight = 0; + assert_eq!(head.await.unwrap().account_id(), &healthy); + } +} + #[tokio::test] async fn saturated_account_snapshot_queues_and_rechecks_live_capacity() { use futures::FutureExt; diff --git a/backend/migrations/.frozen-sha256 b/backend/migrations/.frozen-sha256 index e2ffe16dc..f87244e5a 100644 --- a/backend/migrations/.frozen-sha256 +++ b/backend/migrations/.frozen-sha256 @@ -17,3 +17,4 @@ a1e23b8535ed41de27861833cd2d5ee7498a9bb0e523156ba3f011f9c6b3966b 0016_client_pr fcfae912ddefd54551ebf8e7e733442519d2862d2977134d548ebc1e75c1dc34 0017_unlimited_default_account_concurrency.sql 229c70296764b94fe69c3f19889e566fa290c6c44c64970efe96d7fd2a0008d5 0018_plugins.sql 4daf61846ce59c449865523eb2f7439627dbb65076e56f20de8a74d50905bbf0 0019_plugin_resources.sql +7fcbf25ef5fac8a5d7ae3e3a7adf54cc5c8b9753be48fae5e1d6de8e160a359c 0020_smart_scheduling_settings.sql diff --git a/backend/migrations/0020_smart_scheduling_settings.sql b/backend/migrations/0020_smart_scheduling_settings.sql new file mode 100644 index 000000000..5df7b4170 --- /dev/null +++ b/backend/migrations/0020_smart_scheduling_settings.sql @@ -0,0 +1,4 @@ +-- 保留默认调度行为；评分系数与回切偏好随运行设置原子发布。 +alter table runtime_settings add column smart_scheduling_json jsonb not null default + '{"loadWeight":1.0,"quotaWeight":0.8,"healthWeight":1.0,"latencyWeight":0.5,"resetWeight":0.0,"queueWeight":0.0,"preferHigherWeight":false}'::jsonb + check (jsonb_typeof(smart_scheduling_json) = 'object'); diff --git a/docs/api.md b/docs/api.md index 63d3232c9..8d6e86d83 100644 --- a/docs/api.md +++ b/docs/api.md @@ -1163,6 +1163,7 @@ concurrencyWaitTimeoutSeconds responsesMaxDecompressedBodyBytes requestIntervalMs rotationStrategy +smartScheduling minCodexDesktopVersion minCodexCliVersion usageRetentionDays @@ -1212,6 +1213,35 @@ accountWarmupModel `rotationStrategy` 可取 `smart`、`quota_reset_priority`、`round_robin`、`sticky`。 两个 `minCodex*Version` 字段为 `string | null`，只设置最低版本，不存在最大版本字段。 +`smartScheduling` 是必填的完整对象，仅在内置 `smart` 策略下生效，切换其他策略时仍保存其值： + +```json +{ + "loadWeight": 1.0, + "quotaWeight": 0.8, + "healthWeight": 1.0, + "latencyWeight": 0.5, + "resetWeight": 0.0, + "queueWeight": 0.0, + "preferHigherWeight": false +} +``` + +六项系数分别调整负载、剩余额度、健康、首个有效输出延迟、额度重置和排队压力的评分偏好，范围 `0～10`、最多一位小数， +至少一项大于 `0`。`0` 仅关闭对应评分维度，不放宽账号资格、额度或并发限制；系数不表示流量百分比。 +缺字段、`null`、未知字段及无效数值均返回 `422`，不写入配置。读取设置额外返回只读 +`smartSchedulingDefaults`，内容为上述默认对象，用于恢复默认与自定义状态比较，不可提交到更新接口。 + +`resetWeight` 越大越偏向即将重置额度的账号，复用已有有效重置时间，未知或已过期时不加分。 +`queueWeight` 只在没有可立即使用的账号、需要选择等待队列时参与评分，越大越偏向等待人数少的账号。 +设为 `0` 时沿用最短队列规则；启用后结合其余五项评分选择队列。队列人数限于当前进程，只有开启账号排队才会生效， +不影响已入队请求的位置、队内 FIFO、容量上限或等待超时。两项默认均为 `0`。 + +`preferHigherWeight` 默认关闭。开启后，更高权重账号恢复可用时，后续允许重新选号的请求优先回切； +同权重且可用的会话亲和继续保留。没有可用亲和时，在最高可用权重层内按配置评分。 +原生续写账号绑定仍是硬约束，不因回切而主动换号或触发历史重放。配置随运行设置原子保存和发布， +新请求使用新值，已开始请求及其重试沿用原快照，无需重启。 + ### 模型定价 定价接口独立于运行设置的整体替换，仅允许管理员访问。改价影响本地估算费用与 Client Key 金额限额， diff --git a/docs/architecture.md b/docs/architecture.md index f01a5d138..e7e7b800e 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -269,6 +269,16 @@ sequenceDiagram 候选顺序和调度策略。 运行中的请求始终使用该快照，不拼接不同版本的配置，也不在热路径查询分组关系。 +智能调度配置由 Core 定义和校验，Store 随运行设置及配置 revision 原子保存，快照编译器将其纳入 +`AccountSelectionPolicy`。内置选号与诊断复用同一评分函数，近似最优容差随参与立即选号的五项系数总和缩放。 +额度重置复用 Provider 的有效重置时间投影，未知或已过期时间不加分。排队系数不影响立即选号的分数和容差， +只在选择账号等待队列时结合其他评分与当前进程内的实时等待人数。评分规则由 Core 账号模块统一拥有， +`ConcurrencyWaitQueue` 在同一锁内读取队长并入队；Provider 只传递已限定续写范围的候选。 +排队系数为零时保留最短队列规则；已入队位置、队内 FIFO、取消回收、容量上限和请求共享等待预算保持原合同。 +高权重回切仅影响智能调度的软亲和：先验证候选资格，再决定是否由更高权重层覆盖亲和；同层亲和保留。 +Provider 限定的原生续写账号范围仍是硬约束。回切在请求选号时发生，不新增后台任务或主动历史重放。 +插件显式选号沿用既有候选校验和实际亲和结果，委托内置选号时才应用智能调度配置。 + Client Key 鉴权完成后，API adapter 从有界请求头识别 Codex Desktop/CLI，Core 使用同一请求冻结的 `RuntimeSnapshot` 比较对应最低版本。Desktop 优先于其 User-Agent 内嵌的 CLI/Core 标记；未知客户端不 应用门禁。API 识别手机远程 UA 后缀；这类 Desktop 请求缺失应用版本时不应用门禁，不以 Core 或远程 diff --git a/frontend/src/api/modules/settings.ts b/frontend/src/api/modules/settings.ts index 5955487b6..b34bd6351 100644 --- a/frontend/src/api/modules/settings.ts +++ b/frontend/src/api/modules/settings.ts @@ -5,7 +5,19 @@ import request from '../request' export type RotationStrategy = 'smart' | 'quota_reset_priority' | 'round_robin' | 'sticky' +export interface SmartSchedulingConfig { + loadWeight: number + quotaWeight: number + healthWeight: number + latencyWeight: number + resetWeight: number + queueWeight: number + preferHigherWeight: boolean +} + export interface RuntimeSettings { + smartScheduling: SmartSchedulingConfig + smartSchedulingDefaults: SmartSchedulingConfig providerRequestProfiles: ProviderRequestProfiles openaiClientProfile: ClientProfileSelection | null xaiClientProfile: XaiClientProfileSelection | null @@ -80,7 +92,7 @@ export function getSettings(options: RequestOptions = {}) { }) } -type UpdateSettingsParam = Omit & { +type UpdateSettingsParam = Omit & { providerRequestProfiles: ProviderRequestProfileUpdates } diff --git a/frontend/src/views/settings/components/RotationStrategyCard.vue b/frontend/src/views/settings/components/RotationStrategyCard.vue index 733f55cc3..d8d70a46a 100644 --- a/frontend/src/views/settings/components/RotationStrategyCard.vue +++ b/frontend/src/views/settings/components/RotationStrategyCard.vue @@ -1,48 +1,84 @@ diff --git a/frontend/src/views/settings/components/SmartSchedulingModal.vue b/frontend/src/views/settings/components/SmartSchedulingModal.vue new file mode 100644 index 000000000..1421b3d38 --- /dev/null +++ b/frontend/src/views/settings/components/SmartSchedulingModal.vue @@ -0,0 +1,124 @@ + + + diff --git a/frontend/src/views/settings/composables/useSettingsForm.ts b/frontend/src/views/settings/composables/useSettingsForm.ts index cf8ec9aa3..7c2568d06 100644 --- a/frontend/src/views/settings/composables/useSettingsForm.ts +++ b/frontend/src/views/settings/composables/useSettingsForm.ts @@ -1,5 +1,5 @@ import type { rotationOptions } from '../constants' -import type { RequestLocation } from '@/api' +import type { RequestLocation, SmartSchedulingConfig } from '@/api' import type { ProviderRequestProfiles, ProviderRequestProfileUpdates } from '@/api/modules/client-profiles' import { toast } from '@codex-proxy/ui' import { isEqual } from 'es-toolkit' @@ -22,7 +22,9 @@ export function useSettingsForm() { const error = shallowRef('') const mappings = ref>([]) const savedRequestLocation = shallowRef() + const smartSchedulingDefaults = shallowRef() const form = reactive({ + smartScheduling: undefined as SmartSchedulingConfig | undefined, providerRequestProfiles: {} as ProviderRequestProfiles, requestLocationEnabled: false, requestLocation: { country: '', region: '', city: '', timezone: '' }, @@ -58,6 +60,7 @@ export function useSettingsForm() { return { form: { ...form, + smartScheduling: form.smartScheduling ? { ...form.smartScheduling } : undefined, providerRequestProfiles: cloneProfiles(form.providerRequestProfiles), requestLocation: { ...form.requestLocation }, }, @@ -73,6 +76,7 @@ export function useSettingsForm() { if (!saved.value || saving.value) return Object.assign(form, saved.value.form, { + smartScheduling: saved.value.form.smartScheduling ? { ...saved.value.form.smartScheduling } : undefined, providerRequestProfiles: cloneProfiles(saved.value.form.providerRequestProfiles), requestLocation: { ...saved.value.form.requestLocation }, }) @@ -126,6 +130,8 @@ export function useSettingsForm() { form.concurrencyWaitTimeoutSeconds = data.concurrencyWaitTimeoutSeconds form.responsesMaxDecompressedBodyMiB = data.responsesMaxDecompressedBodyBytes / MIB + form.smartScheduling = { ...data.smartScheduling } + smartSchedulingDefaults.value = { ...data.smartSchedulingDefaults } form.rotationStrategy = data.rotationStrategy form.minCodexDesktopVersion = data.minCodexDesktopVersion ?? '' form.providerRequestProfiles = cloneProfiles(data.providerRequestProfiles) @@ -197,6 +203,9 @@ export function useSettingsForm() { } async function saveSettings() { + const smartScheduling = form.smartScheduling + if (!smartScheduling) + return const savedSettings = saved.value if (saving.value || loading.value || !savedRequestLocation.value || !savedSettings) return @@ -280,6 +289,7 @@ export function useSettingsForm() { concurrencyWaitTimeoutSeconds, responsesMaxDecompressedBodyBytes: responsesMaxDecompressedBodyMiB * MIB, rotationStrategy, + smartScheduling: { ...smartScheduling }, minCodexDesktopVersion: form.minCodexDesktopVersion.trim() || null, minCodexCliVersion: form.minCodexCliVersion.trim() || null, usageRetentionDays: form.usageRetentionDays, @@ -313,6 +323,7 @@ export function useSettingsForm() { resetSettings, error, form, + smartSchedulingDefaults, mappings, addMapping, updateMapping, diff --git a/frontend/src/views/settings/constants.ts b/frontend/src/views/settings/constants.ts index 353392ff7..ea6617605 100644 --- a/frontend/src/views/settings/constants.ts +++ b/frontend/src/views/settings/constants.ts @@ -2,8 +2,9 @@ export const rotationOptions = [ { label: '智能调度（推荐）', value: 'smart', - description: '按负载、窗口用量、请求数和健康反馈评分，优先选择更空闲的账号', + description: '综合负载、剩余额度、健康和延迟评分，支持自定义偏好与权重回切', }, + { label: '额度重置优先', value: 'quota_reset_priority', diff --git a/frontend/src/views/settings/index.vue b/frontend/src/views/settings/index.vue index c92c72f25..64ea43f57 100644 --- a/frontend/src/views/settings/index.vue +++ b/frontend/src/views/settings/index.vue @@ -53,6 +53,7 @@ const { resetSettings, error, form, + smartSchedulingDefaults, mappings, addMapping, updateMapping, @@ -147,7 +148,13 @@ watch(section, (value) => { v-model:max-concurrent-per-account="maxConcurrentPerAccountValue" v-model:request-interval-ms="requestIntervalMsValue" /> - +