feat(scheduling): add configurable smart scheduling preferences (#295)

Make smart scheduling coefficients configurable and expose optional
higher-weight switchback while preserving the original defaults.

Reuse current quota reset signals, process-local queues, and frozen request
settings while preserving continuation ownership and queue fairness.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Charmsky
2026-09-26 21:51:28 +08:00
committed by GitHub
co-authored by Codex
parent d1bf90d5fd
commit 33cc427b86
37 changed files with 1557 additions and 101 deletions
@@ -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<String>,
pub min_codex_cli_version: Option<String>,
@@ -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<String>,
pub min_codex_cli_version: Option<String>,
@@ -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,
@@ -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,
@@ -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,
@@ -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,
@@ -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<serde_json::Map<String, serde_json::Value>>,
@@ -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<String>,
pub min_codex_cli_version: Option<String>,
@@ -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<String>,
pub min_codex_cli_version: Option<String>,
@@ -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<RuntimeSettings> 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,
@@ -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,
@@ -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);
}
@@ -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;
@@ -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::<HashMap<_, _>>();
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::<Vec<_>>();
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
}
@@ -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<Self, SmartSchedulingConfigError> {
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::<u16>(),
) / 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<SmartSchedulingValues> for SmartSchedulingConfig {
type Error = SmartSchedulingConfigError;
fn try_from(value: SmartSchedulingValues) -> Result<Self, Self::Error> {
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<SmartSchedulingConfig> 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,
}
}
}
+30 -6
View File
@@ -80,6 +80,18 @@ impl<K: Clone + Eq + Hash> ConcurrencyWaitQueue<K> {
keys: &[K],
max_waiting: u32,
deadline: Instant,
) -> Result<WaitTicket<K>, 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<WaitTicket<K>, QueueRejection> {
if Instant::now() >= deadline {
return Err(QueueRejection::Timeout);
@@ -95,9 +107,9 @@ impl<K: Clone + Eq + Hash> ConcurrencyWaitQueue<K> {
.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?;
@@ -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,
@@ -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<String, String>,
min_codex_desktop_version: Option<String>,
@@ -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,
@@ -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;
@@ -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::<SmartSchedulingConfig>(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::<SmartSchedulingConfig>(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
);
}
@@ -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"));
});
}
@@ -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 {
@@ -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::{
@@ -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,
@@ -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<gateway_core::account::SmartSchedulingConfig>,
rotation_strategy: String,
request_location_enabled: bool,
request_location_json: sqlx::types::Json<gateway_core::account::RequestLocation>,
@@ -644,6 +649,7 @@ fn runtime_settings_from_row(row: RuntimeSettingsRow) -> StoreResult<RuntimeSett
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,
request_location_enabled: row.request_location_enabled,
request_location: row
@@ -33,6 +33,7 @@ pub struct SnapshotRuntimeSettings {
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 model_mappings: BTreeMap<String, String>,
pub min_codex_desktop_version: Option<String>,
@@ -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<gateway_core::account::SmartSchedulingConfig>,
rotation_strategy: String,
model_mappings_json: sqlx::types::Json<BTreeMap<String, String>>,
min_codex_desktop_version: Option<String>,
@@ -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,
@@ -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 {
@@ -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;
}
// 只判断本次账号范围;空池、认证失效和租约失败不能伪装成额度耗尽。
@@ -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<MemoryAccountStore>,
leases: Arc<TestLeaseCoordinator>,
@@ -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;
@@ -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 {
@@ -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;
+1
View File
@@ -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
@@ -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');
+30
View File
@@ -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 金额限额,
+10
View File
@@ -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 或远程
+13 -1
View File
@@ -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<RuntimeSettings, 'updatedAt' | 'openaiClientProfile' | 'xaiClientProfile' | 'providerRequestProfiles'> & {
type UpdateSettingsParam = Omit<RuntimeSettings, 'updatedAt' | 'smartSchedulingDefaults' | 'openaiClientProfile' | 'xaiClientProfile' | 'providerRequestProfiles'> & {
providerRequestProfiles: ProviderRequestProfileUpdates
}
@@ -1,48 +1,84 @@
<script setup lang="ts">
import type { rotationOptions } from '../constants'
import { BaseCard } from '@codex-proxy/ui'
import type { SmartSchedulingConfig } from '@/api'
import { BaseCard, BaseIconButton } from '@codex-proxy/ui'
import { Settings2 } from '@lucide/vue'
import { isEqual } from 'es-toolkit'
import { computed, shallowRef } from 'vue'
import SmartSchedulingModal from './SmartSchedulingModal.vue'
type RotationOption = (typeof rotationOptions)[number]
type RotationStrategy = RotationOption['value']
defineProps<{
const props = defineProps<{
options: readonly RotationOption[]
smartDefaults?: SmartSchedulingConfig
disabled: boolean
}>()
const model = defineModel<RotationStrategy | ''>({ required: true })
const smartScheduling = defineModel<SmartSchedulingConfig | undefined>('smartScheduling', { required: true })
const settingsOpen = shallowRef(false)
const customized = computed(() => smartScheduling.value && props.smartDefaults && !isEqual(smartScheduling.value, props.smartDefaults))
</script>
<template>
<BaseCard title="调度策略">
<div class="grid max-w-6xl gap-3 sm:grid-cols-2 lg:grid-cols-4">
<button
<div
v-for="option in options"
:key="option.value"
type="button"
class="min-h-25 cursor-pointer rounded-cp border-0 px-4 py-3.5 text-left shadow-cp-input outline-none transition-[background-color,box-shadow,color] duration-160 focus-visible:ring-2 focus-visible:ring-cp-control-outline"
:class="
model === option.value
? 'bg-cp-control-item-bg-active text-cp-primary-text shadow-cp-tertiary'
: 'bg-(--cp-input-bg,var(--cp-input-bg)) text-cp-text hover:bg-(--cp-input-hover-bg,var(--cp-input-hover-bg)) hover:shadow-cp-input-hover'
"
:aria-pressed="model === option.value"
@click="model = option.value"
class="relative flex"
>
<span class="flex items-center gap-2">
<span
class="inline-flex size-4 shrink-0 items-center justify-center rounded-full bg-cp-bg-container shadow-[inset_0_0_0_1px_var(--cp-color-border)]"
>
<button
type="button"
class="min-h-25 w-full cursor-pointer rounded-cp border-0 px-4 py-3.5 text-left shadow-cp-input outline-none transition-[background-color,box-shadow,color] duration-160 focus-visible:ring-2 focus-visible:ring-cp-control-outline"
:class="
model === option.value
? 'bg-cp-control-item-bg-active text-cp-primary-text shadow-cp-tertiary'
: 'bg-(--cp-input-bg,var(--cp-input-bg)) text-cp-text hover:bg-(--cp-input-hover-bg,var(--cp-input-hover-bg)) hover:shadow-cp-input-hover'
"
:aria-pressed="model === option.value"
:disabled="disabled"
@click="model = option.value"
>
<span class="flex items-center gap-2" :class="option.value === 'smart' ? 'pr-7' : undefined">
<span
class="size-2 rounded-full transition-opacity duration-150"
:class="model === option.value ? 'bg-cp-primary opacity-100' : 'opacity-0'"
/>
class="inline-flex size-4 shrink-0 items-center justify-center rounded-full bg-cp-bg-container shadow-[inset_0_0_0_1px_var(--cp-color-border)]"
>
<span
class="size-2 rounded-full transition-opacity duration-150"
:class="model === option.value ? 'bg-cp-primary opacity-100' : 'opacity-0'"
/>
</span>
<span class="inline-flex items-baseline gap-1 text-cp leading-snug font-emphasis">
{{ option.label }}
<span v-if="option.value === 'smart' && customized" class="shrink-0 text-cp-xs font-normal text-cp-text-secondary">自定义</span>
</span>
</span>
<span class="text-cp leading-snug font-emphasis">{{ option.label }}</span>
</span>
<span class="mt-2 block text-cp leading-normal font-emphasis text-cp-text-secondary">
{{ option.description }}
</span>
</button>
<span class="mt-2 block text-cp leading-normal font-emphasis text-cp-text-secondary">
{{ option.description }}
</span>
</button>
<BaseIconButton
v-if="option.value === 'smart'"
class="absolute top-3 right-3 size-6! rounded-cp-sm p-0"
label="智能调度设置"
size="sm"
:disabled="disabled || !smartScheduling || !smartDefaults"
@click="settingsOpen = true"
>
<Settings2 class="size-3.5" />
</BaseIconButton>
</div>
</div>
</BaseCard>
<SmartSchedulingModal
v-if="smartScheduling && smartDefaults"
v-model="settingsOpen"
:config="smartScheduling"
:defaults="smartDefaults"
:active="model === 'smart'"
@confirm="smartScheduling = $event"
/>
</template>
@@ -0,0 +1,124 @@
<script setup lang="ts">
import type { SmartSchedulingConfig } from '@/api'
import { BaseButton, BaseFormItem, BaseInput, BaseModal, BasePopover, BaseSwitch } from '@codex-proxy/ui'
import { CircleAlert, CircleHelp } from '@lucide/vue'
import { computed, reactive, shallowRef, useId, watch } from 'vue'
const props = defineProps<{
config: SmartSchedulingConfig
defaults: SmartSchedulingConfig
active: boolean
}>()
const emit = defineEmits<{ confirm: [config: SmartSchedulingConfig] }>()
const open = defineModel<boolean>({ required: true })
const formId = useId()
const fields = [
{ key: 'loadWeight', label: '负载系数', help: '越大越偏向并发占用比例较低的账号' },
{ key: 'quotaWeight', label: '剩余额度系数', help: '越大越偏向剩余额度比例较高的账号' },
{ key: 'healthWeight', label: '健康系数', help: '越大越偏向近期失败率较低的账号' },
{ key: 'latencyWeight', label: '延迟系数', help: '越大越偏向首个有效输出较快的账号' },
{ key: 'resetWeight', label: '额度重置系数', help: '越大越偏向即将重置额度的账号,重置时间未知或已过期时不加分' },
{ key: 'queueWeight', label: '排队系数', help: '需要排队时,越大越偏向等待请求较少的账号,设为 0 时沿用最短队列规则' },
] as const
const draft = reactive({ loadWeight: '', quotaWeight: '', healthWeight: '', latencyWeight: '', resetWeight: '', queueWeight: '', preferHigherWeight: false })
const submitted = shallowRef(false)
const errors = computed(() => Object.fromEntries(fields.map(({ key }) => {
const value = Number(draft[key])
const valid = draft[key].trim() && Number.isFinite(value) && value >= 0 && value <= 10 && Math.round(value * 10) / 10 === value
return [key, valid ? undefined : '请输入 0~10 的数值,最多一位小数']
})))
const emptyWeights = computed(() => fields.every(({ key }) => Number(draft[key]) === 0))
function replaceDraft(config: SmartSchedulingConfig) {
for (const { key } of fields)
draft[key] = String(config[key])
draft.preferHigherWeight = config.preferHigherWeight
submitted.value = false
}
watch(open, (value) => {
if (value)
replaceDraft(props.config)
}, { immediate: true })
function confirm() {
submitted.value = true
if (Object.values(errors.value).some(Boolean) || emptyWeights.value)
return
emit('confirm', {
loadWeight: Number(draft.loadWeight),
quotaWeight: Number(draft.quotaWeight),
healthWeight: Number(draft.healthWeight),
latencyWeight: Number(draft.latencyWeight),
resetWeight: Number(draft.resetWeight),
queueWeight: Number(draft.queueWeight),
preferHigherWeight: draft.preferHigherWeight,
})
open.value = false
}
</script>
<template>
<BaseModal v-model="open" title="智能调度设置" description="调整账号选择偏好与回切方式" size="md">
<form :id="formId" class="grid gap-5" novalidate @submit.prevent="confirm">
<p v-if="!active" class="m-0 text-cp text-cp-text-secondary">
仅在启用智能调度后生效
</p>
<fieldset class="m-0 min-w-0 border-0 p-0">
<legend class="mb-3 text-cp font-emphasis">
评分系数
</legend>
<div class="grid gap-4 sm:grid-cols-2">
<BaseFormItem v-for="field in fields" :key="field.key" :error="submitted ? errors[field.key] : undefined">
<template #label>
<span class="inline-flex items-center gap-1">
{{ field.label }}
<BasePopover trigger="hover-click" placement="top-start">
<template #trigger="{ open: helpOpen }">
<button type="button" class="inline-flex size-6 cursor-pointer items-center justify-center rounded-cp-sm border-0 bg-transparent p-0 text-cp-text-tertiary outline-none hover:text-cp-text focus-visible:ring-2 focus-visible:ring-cp-control-outline" :aria-label="`${field.label}说明`" :aria-expanded="helpOpen">
<CircleHelp class="size-3.5" />
</button>
</template>
<p class="m-0 max-w-64 px-3 py-2 text-cp-sm leading-relaxed">{{ field.help }}</p>
</BasePopover>
</span>
</template>
<BaseInput v-model="draft[field.key]" type="number" min="0" max="10" step="0.1" placeholder="0 表示不参与评分" :aria-label="field.label" />
</BaseFormItem>
</div>
<p v-if="submitted && emptyWeights" role="alert" class="mt-2 mb-0 text-cp-sm text-cp-danger">
至少一项评分系数需要大于 0
</p>
</fieldset>
<fieldset class="m-0 min-w-0 border-0 p-0">
<legend class="mb-3 text-cp font-emphasis">
<span class="inline-flex items-center gap-1">
回切行为
<BasePopover trigger="hover-click" placement="top-start">
<template #trigger="{ open: helpOpen }">
<button type="button" class="inline-flex size-6 cursor-pointer items-center justify-center rounded-cp-sm border-0 bg-transparent p-0 text-cp-text-tertiary outline-none hover:text-cp-text focus-visible:ring-2 focus-visible:ring-cp-control-outline" aria-label="回切行为说明" :aria-expanded="helpOpen">
<CircleAlert class="size-3.5" />
</button>
</template>
<p class="m-0 max-w-72 px-3 py-2 text-cp-sm leading-relaxed">
高权重账号恢复可用后,后续允许重新选号的请求优先回切<br>原生续写仍绑定原账号,同权重账号保留已有会话亲和
</p>
</BasePopover>
</span>
</legend>
<BaseSwitch v-model="draft.preferHigherWeight" label="自动回切" show-label />
</fieldset>
</form>
<template #footer>
<BaseButton class="mr-auto" variant="secondary" @click="replaceDraft(defaults)">
恢复默认
</BaseButton>
<BaseButton variant="secondary" @click="open = false">
取消
</BaseButton>
<BaseButton type="submit" :form="formId" variant="primary">
确定
</BaseButton>
</template>
</BaseModal>
</template>
@@ -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<Array<{ requestedModel: string, upstreamModel: string }>>([])
const savedRequestLocation = shallowRef<RequestLocation>()
const smartSchedulingDefaults = shallowRef<SmartSchedulingConfig>()
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,
+2 -1
View File
@@ -2,8 +2,9 @@ export const rotationOptions = [
{
label: '智能调度(推荐)',
value: 'smart',
description: '按负载、窗口用量、请求数和健康反馈评分,优先选择更空闲的账号',
description: '综合负载、剩余额度、健康和延迟评分,支持自定义偏好与权重回切',
},
{
label: '额度重置优先',
value: 'quota_reset_priority',
+8 -1
View File
@@ -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"
/>
<RotationStrategyCard v-model="form.rotationStrategy" :options="rotationOptions" />
<RotationStrategyCard
v-model="form.rotationStrategy"
v-model:smart-scheduling="form.smartScheduling"
:smart-defaults="smartSchedulingDefaults"
:disabled="disabled"
:options="rotationOptions"
/>
<RequestQueueCard
v-model:max-waiting-per-key="maxWaitingPerKeyValue"
v-model:max-waiting-per-account="maxWaitingPerAccountValue"