mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 16:06:08 +00:00
feat(openai): add scheduled account warmup (#192)
Co-authored-by: Zyy <zyycn@users.noreply.github.com> Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
committed by
Codex
co-authored by
Zyy
Codex
parent
7a0502cabd
commit
ea31bf0cf7
@@ -51,6 +51,9 @@ pub struct RuntimeSettings {
|
||||
pub account_auto_freeze_probe_enabled: bool,
|
||||
pub account_auto_freeze_probe_model: Option<String>,
|
||||
pub account_auto_freeze_adaptive_concurrency: bool,
|
||||
pub account_warmup_enabled: bool,
|
||||
pub account_warmup_schedule_time: String,
|
||||
pub account_warmup_model: Option<String>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
@@ -83,6 +86,9 @@ pub struct ReplaceRuntimeSettings {
|
||||
pub account_auto_freeze_probe_enabled: bool,
|
||||
pub account_auto_freeze_probe_model: Option<String>,
|
||||
pub account_auto_freeze_adaptive_concurrency: bool,
|
||||
pub account_warmup_enabled: bool,
|
||||
pub account_warmup_schedule_time: String,
|
||||
pub account_warmup_model: Option<String>,
|
||||
}
|
||||
|
||||
/// 明文管理员 API Key;按产品约束明文落库,但禁止 Debug 泄漏。
|
||||
|
||||
@@ -422,6 +422,11 @@ fn validate_settings(command: &ReplaceRuntimeSettings) -> Result<(), AdminError>
|
||||
&& valid_client_version(command.min_codex_desktop_version.as_deref())
|
||||
&& valid_client_version(command.min_codex_cli_version.as_deref())
|
||||
&& valid_probe_model(command.account_auto_freeze_probe_model.as_deref())
|
||||
&& gateway_core::provider_ports::valid_warmup_schedule_time(
|
||||
&command.account_warmup_schedule_time,
|
||||
)
|
||||
&& valid_probe_model(command.account_warmup_model.as_deref())
|
||||
&& (!command.account_warmup_enabled || command.account_warmup_model.is_some())
|
||||
&& i64::try_from(command.request_interval_ms).is_ok()
|
||||
&& (2..=1_000).contains(&command.account_auto_freeze_threshold)
|
||||
&& (60..=3_600).contains(&command.account_auto_freeze_window_seconds)
|
||||
|
||||
@@ -1118,6 +1118,9 @@ impl SettingsStore for StaticSettingsStore {
|
||||
account_auto_freeze_probe_enabled: true,
|
||||
account_auto_freeze_probe_model: None,
|
||||
account_auto_freeze_adaptive_concurrency: true,
|
||||
account_warmup_enabled: false,
|
||||
account_warmup_schedule_time: "08:00".to_owned(),
|
||||
account_warmup_model: None,
|
||||
updated_at: Utc::now(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -51,6 +51,9 @@ fn runtime_settings(enabled: bool, probe_enabled: bool, adaptive: bool) -> Runti
|
||||
account_auto_freeze_probe_enabled: probe_enabled,
|
||||
account_auto_freeze_probe_model: Some("gpt-5.5".to_owned()),
|
||||
account_auto_freeze_adaptive_concurrency: adaptive,
|
||||
account_warmup_enabled: false,
|
||||
account_warmup_schedule_time: "08:00".to_owned(),
|
||||
account_warmup_model: None,
|
||||
updated_at: Utc::now(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -835,6 +835,9 @@ impl SettingsStore for FixtureSettingsStore {
|
||||
account_auto_freeze_probe_enabled: true,
|
||||
account_auto_freeze_probe_model: None,
|
||||
account_auto_freeze_adaptive_concurrency: true,
|
||||
account_warmup_enabled: false,
|
||||
account_warmup_schedule_time: "08:00".to_owned(),
|
||||
account_warmup_model: None,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -238,6 +238,9 @@ async fn settings_should_reject_zero_refresh_margin_before_store_call() {
|
||||
account_auto_freeze_probe_enabled: true,
|
||||
account_auto_freeze_probe_model: None,
|
||||
account_auto_freeze_adaptive_concurrency: true,
|
||||
account_warmup_enabled: false,
|
||||
account_warmup_schedule_time: "08:00".to_owned(),
|
||||
account_warmup_model: None,
|
||||
},
|
||||
)
|
||||
.await
|
||||
|
||||
@@ -66,6 +66,9 @@ pub struct RuntimeSettingsView {
|
||||
pub account_auto_freeze_probe_enabled: bool,
|
||||
pub account_auto_freeze_probe_model: Option<String>,
|
||||
pub account_auto_freeze_adaptive_concurrency: bool,
|
||||
pub account_warmup_enabled: bool,
|
||||
pub account_warmup_schedule_time: String,
|
||||
pub account_warmup_model: Option<String>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
@@ -105,6 +108,9 @@ pub struct UpdateRuntimeSettingsRequest {
|
||||
pub account_auto_freeze_probe_enabled: bool,
|
||||
pub account_auto_freeze_probe_model: Option<String>,
|
||||
pub account_auto_freeze_adaptive_concurrency: bool,
|
||||
pub account_warmup_enabled: bool,
|
||||
pub account_warmup_schedule_time: String,
|
||||
pub account_warmup_model: Option<String>,
|
||||
}
|
||||
|
||||
impl UpdateRuntimeSettingsRequest {
|
||||
@@ -184,6 +190,15 @@ impl UpdateRuntimeSettingsRequest {
|
||||
self.account_auto_freeze_probe_model.as_deref(),
|
||||
"accountAutoFreezeProbeModel",
|
||||
)?;
|
||||
if !gateway_core::provider_ports::valid_warmup_schedule_time(
|
||||
&self.account_warmup_schedule_time,
|
||||
) {
|
||||
return Err(WireValidationError::new("accountWarmupScheduleTime"));
|
||||
}
|
||||
validate_optional_probe_model(self.account_warmup_model.as_deref(), "accountWarmupModel")?;
|
||||
if self.account_warmup_enabled && self.account_warmup_model.is_none() {
|
||||
return Err(WireValidationError::new("accountWarmupModel"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -230,6 +245,9 @@ impl UpdateRuntimeSettingsRequest {
|
||||
account_auto_freeze_probe_enabled: self.account_auto_freeze_probe_enabled,
|
||||
account_auto_freeze_probe_model: self.account_auto_freeze_probe_model,
|
||||
account_auto_freeze_adaptive_concurrency: self.account_auto_freeze_adaptive_concurrency,
|
||||
account_warmup_enabled: self.account_warmup_enabled,
|
||||
account_warmup_schedule_time: self.account_warmup_schedule_time,
|
||||
account_warmup_model: self.account_warmup_model,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -270,6 +288,9 @@ impl From<RuntimeSettings> for RuntimeSettingsView {
|
||||
account_auto_freeze_probe_model: settings.account_auto_freeze_probe_model,
|
||||
account_auto_freeze_adaptive_concurrency: settings
|
||||
.account_auto_freeze_adaptive_concurrency,
|
||||
account_warmup_enabled: settings.account_warmup_enabled,
|
||||
account_warmup_schedule_time: settings.account_warmup_schedule_time,
|
||||
account_warmup_model: settings.account_warmup_model,
|
||||
updated_at: settings.updated_at,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -569,6 +569,9 @@ impl SettingsStore for MemorySettingsStore {
|
||||
account_auto_freeze_probe_enabled: true,
|
||||
account_auto_freeze_probe_model: None,
|
||||
account_auto_freeze_adaptive_concurrency: true,
|
||||
account_warmup_enabled: false,
|
||||
account_warmup_schedule_time: "08:00".to_owned(),
|
||||
account_warmup_model: None,
|
||||
updated_at: Utc::now(),
|
||||
};
|
||||
*settings = updated.clone();
|
||||
@@ -1523,6 +1526,9 @@ fn test_runtime_settings() -> RuntimeSettings {
|
||||
account_auto_freeze_probe_enabled: true,
|
||||
account_auto_freeze_probe_model: None,
|
||||
account_auto_freeze_adaptive_concurrency: true,
|
||||
account_warmup_enabled: false,
|
||||
account_warmup_schedule_time: "08:00".to_owned(),
|
||||
account_warmup_model: None,
|
||||
updated_at: Utc::now(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -63,7 +63,10 @@ fn update_body() -> Value {
|
||||
"accountAutoFreezeDurationSeconds": 7200,
|
||||
"accountAutoFreezeProbeEnabled": true,
|
||||
"accountAutoFreezeProbeModel": null,
|
||||
"accountAutoFreezeAdaptiveConcurrency": true
|
||||
"accountAutoFreezeAdaptiveConcurrency": true,
|
||||
"accountWarmupEnabled": false,
|
||||
"accountWarmupScheduleTime": "08:00",
|
||||
"accountWarmupModel": null
|
||||
})
|
||||
}
|
||||
|
||||
@@ -86,6 +89,19 @@ fn settings_request_accepts_unlimited_default_account_concurrency() {
|
||||
request.validate().expect("zero means unlimited");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settings_request_requires_model_when_warmup_is_enabled() {
|
||||
let mut body = update_body();
|
||||
body["accountWarmupEnabled"] = json!(true);
|
||||
let request: UpdateRuntimeSettingsRequest =
|
||||
serde_json::from_value(body).expect("decode settings");
|
||||
|
||||
assert_eq!(
|
||||
request.validate().unwrap_err().field(),
|
||||
"accountWarmupModel"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settings_request_should_reject_non_semver_client_min() {
|
||||
let mut body = update_body();
|
||||
@@ -146,6 +162,9 @@ fn settings_response_should_cover_the_full_runtime_settings_contract() {
|
||||
account_auto_freeze_probe_enabled: true,
|
||||
account_auto_freeze_probe_model: None,
|
||||
account_auto_freeze_adaptive_concurrency: true,
|
||||
account_warmup_enabled: false,
|
||||
account_warmup_schedule_time: "08:00".to_owned(),
|
||||
account_warmup_model: None,
|
||||
updated_at: Utc
|
||||
.with_ymd_and_hms(2026, 8, 2, 10, 30, 0)
|
||||
.single()
|
||||
@@ -183,10 +202,13 @@ fn settings_response_should_cover_the_full_runtime_settings_contract() {
|
||||
"accountAutoFreezeThreshold": 12,
|
||||
"accountAutoFreezeWindowSeconds": 600,
|
||||
"accountAutoFreezeDurationSeconds": 7200,
|
||||
"accountAutoFreezeProbeEnabled": true,
|
||||
"accountAutoFreezeProbeModel": null,
|
||||
"accountAutoFreezeAdaptiveConcurrency": true,
|
||||
"updatedAt": "2026-08-02T10:30:00Z"
|
||||
"accountAutoFreezeProbeEnabled": true,
|
||||
"accountAutoFreezeProbeModel": null,
|
||||
"accountAutoFreezeAdaptiveConcurrency": true,
|
||||
"accountWarmupEnabled": false,
|
||||
"accountWarmupScheduleTime": "08:00",
|
||||
"accountWarmupModel": null,
|
||||
"updatedAt": "2026-08-02T10:30:00Z"
|
||||
})
|
||||
);
|
||||
}
|
||||
@@ -249,6 +271,9 @@ fn settings_request_and_response_fields_should_stay_in_lockstep() {
|
||||
account_auto_freeze_probe_enabled: true,
|
||||
account_auto_freeze_probe_model: None,
|
||||
account_auto_freeze_adaptive_concurrency: true,
|
||||
account_warmup_enabled: false,
|
||||
account_warmup_schedule_time: "08:00".to_owned(),
|
||||
account_warmup_model: None,
|
||||
updated_at: chrono::Utc::now(),
|
||||
};
|
||||
|
||||
|
||||
@@ -967,6 +967,13 @@ pub trait ProviderRuntimePolicyPort: Send + Sync {
|
||||
) -> BoxFuture<'_, Result<ProviderFreezePolicy, ProviderStoreError>> {
|
||||
Box::pin(async move { Ok(ProviderFreezePolicy::disabled()) })
|
||||
}
|
||||
|
||||
/// 读取账号模型预激活策略;默认关闭,只有实现运行时设置的存储需要覆盖。
|
||||
fn load_warmup_policy(
|
||||
&self,
|
||||
) -> BoxFuture<'_, Result<ProviderWarmupPolicy, ProviderStoreError>> {
|
||||
Box::pin(async move { Ok(ProviderWarmupPolicy::disabled()) })
|
||||
}
|
||||
}
|
||||
|
||||
/// 账号容量熔断(自动冻结)策略;来源于 `runtime_settings`,
|
||||
@@ -1074,6 +1081,109 @@ impl ProviderFreezePolicy {
|
||||
}
|
||||
}
|
||||
|
||||
/// 校验每日预激活时间格式,如 "08:00" 或 "08:00,13:00"。
|
||||
#[must_use]
|
||||
pub fn valid_warmup_schedule_time(value: &str) -> bool {
|
||||
if value.is_empty()
|
||||
|| value.len() > 255
|
||||
|| value != value.trim()
|
||||
|| value.chars().any(char::is_control)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
for part in value.split(',') {
|
||||
let bytes = part.as_bytes();
|
||||
if bytes.len() != 5 || bytes[2] != b':' {
|
||||
return false;
|
||||
}
|
||||
let Ok(hour) = part[0..2].parse::<u32>() else {
|
||||
return false;
|
||||
};
|
||||
let Ok(minute) = part[3..5].parse::<u32>() else {
|
||||
return false;
|
||||
};
|
||||
if hour > 23 || minute > 59 {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
/// 账号模型预激活(预热)策略;来源于 `runtime_settings`。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderWarmupPolicy {
|
||||
enabled: bool,
|
||||
schedule_time: String,
|
||||
model: Option<String>,
|
||||
}
|
||||
|
||||
impl ProviderWarmupPolicy {
|
||||
pub fn try_new(
|
||||
enabled: bool,
|
||||
schedule_time: String,
|
||||
model: Option<String>,
|
||||
) -> Result<Self, ProviderStoreError> {
|
||||
if !valid_warmup_schedule_time(&schedule_time)
|
||||
|| (enabled && model.is_none())
|
||||
|| model.as_deref().is_some_and(|m| {
|
||||
m.is_empty()
|
||||
|| m.len() > 128
|
||||
|| m != m.trim()
|
||||
|| m.bytes().any(|byte| byte.is_ascii_control())
|
||||
})
|
||||
{
|
||||
return Err(ProviderStoreError::new(
|
||||
ProviderStoreErrorKind::InvalidData,
|
||||
"validate warmup policy",
|
||||
));
|
||||
}
|
||||
Ok(Self {
|
||||
enabled,
|
||||
schedule_time,
|
||||
model,
|
||||
})
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn disabled() -> Self {
|
||||
Self {
|
||||
enabled: false,
|
||||
schedule_time: String::new(),
|
||||
model: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn enabled(&self) -> bool {
|
||||
self.enabled
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn schedule_time(&self) -> &str {
|
||||
&self.schedule_time
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn model(&self) -> Option<&str> {
|
||||
self.model.as_deref()
|
||||
}
|
||||
|
||||
/// 解析每日时间列表,返回如 `vec![(8, 0)]`。
|
||||
#[must_use]
|
||||
pub fn scheduled_times(&self) -> Vec<(u32, u32)> {
|
||||
self.schedule_time
|
||||
.split(',')
|
||||
.filter_map(|part| {
|
||||
let part = part.trim();
|
||||
let mut iter = part.split(':');
|
||||
let hour = iter.next()?.parse::<u32>().ok()?;
|
||||
let minute = iter.next()?.parse::<u32>().ok()?;
|
||||
Some((hour, minute))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// OAuth pending flow 的原始绑定只在 Provider 与 Store 边界内短暂存在。
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub struct OAuthPendingBinding(String);
|
||||
|
||||
@@ -110,3 +110,32 @@ fn scheduling_state_preserves_provider_neutral_signals() {
|
||||
1
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn warmup_schedule_time_and_policy_validation() {
|
||||
use gateway_core::provider_ports::{ProviderWarmupPolicy, valid_warmup_schedule_time};
|
||||
|
||||
assert!(valid_warmup_schedule_time("08:00"));
|
||||
assert!(valid_warmup_schedule_time("08:00,13:00"));
|
||||
assert!(valid_warmup_schedule_time("00:00,23:59"));
|
||||
assert!(!valid_warmup_schedule_time(""));
|
||||
assert!(!valid_warmup_schedule_time("8:00"));
|
||||
assert!(!valid_warmup_schedule_time("24:00"));
|
||||
assert!(!valid_warmup_schedule_time("08:60"));
|
||||
assert!(!valid_warmup_schedule_time("08:00,"));
|
||||
assert!(!valid_warmup_schedule_time("08:00, 13:00"));
|
||||
|
||||
let policy = ProviderWarmupPolicy::try_new(
|
||||
true,
|
||||
"08:00,13:00".to_owned(),
|
||||
Some("test-model".to_owned()),
|
||||
)
|
||||
.expect("valid warmup policy");
|
||||
assert!(policy.enabled());
|
||||
assert_eq!(policy.schedule_time(), "08:00,13:00");
|
||||
assert_eq!(policy.model(), Some("test-model"));
|
||||
assert_eq!(policy.scheduled_times(), vec![(8, 0), (13, 0)]);
|
||||
|
||||
assert!(ProviderWarmupPolicy::try_new(true, "invalid".to_owned(), None).is_err());
|
||||
assert!(ProviderWarmupPolicy::try_new(true, "08:00".to_owned(), None).is_err());
|
||||
}
|
||||
|
||||
@@ -114,6 +114,9 @@ impl SettingsStore for AdminSettingsStoreAdapter {
|
||||
account_auto_freeze_probe_model: command.account_auto_freeze_probe_model,
|
||||
account_auto_freeze_adaptive_concurrency: command
|
||||
.account_auto_freeze_adaptive_concurrency,
|
||||
account_warmup_enabled: command.account_warmup_enabled,
|
||||
account_warmup_schedule_time: command.account_warmup_schedule_time,
|
||||
account_warmup_model: command.account_warmup_model,
|
||||
},
|
||||
audit: mutation_audit(
|
||||
context,
|
||||
@@ -257,6 +260,9 @@ pub(crate) fn admin_runtime_settings(
|
||||
account_auto_freeze_probe_enabled: settings.account_auto_freeze_probe_enabled,
|
||||
account_auto_freeze_probe_model: settings.account_auto_freeze_probe_model,
|
||||
account_auto_freeze_adaptive_concurrency: settings.account_auto_freeze_adaptive_concurrency,
|
||||
account_warmup_enabled: settings.account_warmup_enabled,
|
||||
account_warmup_schedule_time: settings.account_warmup_schedule_time,
|
||||
account_warmup_model: settings.account_warmup_model,
|
||||
updated_at: settings.updated_at,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -12,7 +12,7 @@ use gateway_core::account::RotationStrategy;
|
||||
use gateway_core::policy::CodexClientVersion;
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderFreezePolicy, ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderStoreError,
|
||||
ProviderStoreErrorKind,
|
||||
ProviderStoreErrorKind, ProviderWarmupPolicy,
|
||||
};
|
||||
|
||||
use crate::{Revision, StoreError, StoreResult, postgres_unavailable};
|
||||
@@ -47,6 +47,9 @@ pub struct RuntimeSettings {
|
||||
pub account_auto_freeze_probe_enabled: bool,
|
||||
pub account_auto_freeze_probe_model: Option<String>,
|
||||
pub account_auto_freeze_adaptive_concurrency: bool,
|
||||
pub account_warmup_enabled: bool,
|
||||
pub account_warmup_schedule_time: String,
|
||||
pub account_warmup_model: Option<String>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
@@ -104,6 +107,12 @@ impl fmt::Debug for RuntimeSettings {
|
||||
"account_auto_freeze_adaptive_concurrency",
|
||||
&self.account_auto_freeze_adaptive_concurrency,
|
||||
)
|
||||
.field("account_warmup_enabled", &self.account_warmup_enabled)
|
||||
.field(
|
||||
"account_warmup_schedule_time",
|
||||
&self.account_warmup_schedule_time,
|
||||
)
|
||||
.field("account_warmup_model", &self.account_warmup_model)
|
||||
.field("updated_at", &self.updated_at)
|
||||
.finish()
|
||||
}
|
||||
@@ -140,6 +149,9 @@ pub struct RuntimeSettingsUpdate {
|
||||
pub account_auto_freeze_probe_enabled: bool,
|
||||
pub account_auto_freeze_probe_model: Option<String>,
|
||||
pub account_auto_freeze_adaptive_concurrency: bool,
|
||||
pub account_warmup_enabled: bool,
|
||||
pub account_warmup_schedule_time: String,
|
||||
pub account_warmup_model: Option<String>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for RuntimeSettingsUpdate {
|
||||
@@ -178,6 +190,11 @@ impl RuntimeSettingsUpdate {
|
||||
|| !valid_client_version(self.min_codex_desktop_version.as_deref())
|
||||
|| !valid_client_version(self.min_codex_cli_version.as_deref())
|
||||
|| !valid_probe_model(self.account_auto_freeze_probe_model.as_deref())
|
||||
|| !gateway_core::provider_ports::valid_warmup_schedule_time(
|
||||
&self.account_warmup_schedule_time,
|
||||
)
|
||||
|| !valid_probe_model(self.account_warmup_model.as_deref())
|
||||
|| (self.account_warmup_enabled && self.account_warmup_model.is_none())
|
||||
|| RotationStrategy::parse(&self.rotation_strategy).is_none()
|
||||
|| self.request_profile_updates.len() > 256
|
||||
|| self
|
||||
@@ -252,7 +269,8 @@ pub(crate) async fn load_runtime_settings_from_pool(pool: &PgPool) -> StoreResul
|
||||
account_auto_freeze_enabled, account_auto_freeze_threshold,
|
||||
account_auto_freeze_window_seconds, account_auto_freeze_duration_seconds,
|
||||
account_auto_freeze_probe_enabled, account_auto_freeze_probe_model,
|
||||
account_auto_freeze_adaptive_concurrency
|
||||
account_auto_freeze_adaptive_concurrency,
|
||||
account_warmup_enabled, account_warmup_schedule_time, account_warmup_model
|
||||
from runtime_settings where id = 1",
|
||||
)
|
||||
.fetch_optional(pool)
|
||||
@@ -374,6 +392,21 @@ impl ProviderRuntimePolicyPort for PgRuntimeSettingsRepository {
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn load_warmup_policy(
|
||||
&self,
|
||||
) -> futures::future::BoxFuture<'_, Result<ProviderWarmupPolicy, ProviderStoreError>> {
|
||||
Box::pin(async move {
|
||||
let settings = RuntimeSettingsRepository::load_runtime_settings(self)
|
||||
.await
|
||||
.map_err(|_| provider_unavailable("load warmup policy"))?;
|
||||
ProviderWarmupPolicy::try_new(
|
||||
settings.account_warmup_enabled,
|
||||
settings.account_warmup_schedule_time,
|
||||
settings.account_warmup_model,
|
||||
)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn load_runtime_settings_in_transaction(
|
||||
@@ -388,7 +421,8 @@ pub(crate) async fn load_runtime_settings_in_transaction(
|
||||
account_auto_freeze_enabled, account_auto_freeze_threshold,
|
||||
account_auto_freeze_window_seconds, account_auto_freeze_duration_seconds,
|
||||
account_auto_freeze_probe_enabled, account_auto_freeze_probe_model,
|
||||
account_auto_freeze_adaptive_concurrency
|
||||
account_auto_freeze_adaptive_concurrency,
|
||||
account_warmup_enabled, account_warmup_schedule_time, account_warmup_model
|
||||
from runtime_settings where id = 1",
|
||||
)
|
||||
.fetch_optional(&mut **transaction)
|
||||
@@ -455,6 +489,9 @@ pub(crate) async fn update_runtime_settings_in_transaction(
|
||||
request_location_enabled = $24,
|
||||
responses_max_decompressed_body_bytes = $25,
|
||||
provider_request_profiles_json = (provider_request_profiles_json - $26::text[]) || $27::jsonb,
|
||||
account_warmup_enabled = $28,
|
||||
account_warmup_schedule_time = $29,
|
||||
account_warmup_model = $30,
|
||||
updated_at = now()
|
||||
where id = 1
|
||||
returning config_revision",
|
||||
@@ -498,6 +535,9 @@ pub(crate) async fn update_runtime_settings_in_transaction(
|
||||
)
|
||||
.bind(request_profile_deletions)
|
||||
.bind(sqlx::types::Json(request_profile_updates))
|
||||
.bind(update.account_warmup_enabled)
|
||||
.bind(&update.account_warmup_schedule_time)
|
||||
.bind(update.account_warmup_model.as_deref())
|
||||
.fetch_optional(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("update runtime settings in transaction"))?
|
||||
@@ -577,6 +617,9 @@ struct RuntimeSettingsRow {
|
||||
account_auto_freeze_probe_enabled: bool,
|
||||
account_auto_freeze_probe_model: Option<String>,
|
||||
account_auto_freeze_adaptive_concurrency: bool,
|
||||
account_warmup_enabled: bool,
|
||||
account_warmup_schedule_time: String,
|
||||
account_warmup_model: Option<String>,
|
||||
}
|
||||
|
||||
fn runtime_settings_from_row(row: RuntimeSettingsRow) -> StoreResult<RuntimeSettings> {
|
||||
@@ -626,6 +669,9 @@ fn runtime_settings_from_row(row: RuntimeSettingsRow) -> StoreResult<RuntimeSett
|
||||
account_auto_freeze_probe_enabled: row.account_auto_freeze_probe_enabled,
|
||||
account_auto_freeze_probe_model: row.account_auto_freeze_probe_model,
|
||||
account_auto_freeze_adaptive_concurrency: row.account_auto_freeze_adaptive_concurrency,
|
||||
account_warmup_enabled: row.account_warmup_enabled,
|
||||
account_warmup_schedule_time: row.account_warmup_schedule_time,
|
||||
account_warmup_model: row.account_warmup_model,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -38,6 +38,9 @@ fn settings_with_margin(refresh_margin_seconds: u64) -> RuntimeSettingsUpdate {
|
||||
account_auto_freeze_probe_enabled: true,
|
||||
account_auto_freeze_probe_model: None,
|
||||
account_auto_freeze_adaptive_concurrency: true,
|
||||
account_warmup_enabled: false,
|
||||
account_warmup_schedule_time: "08:00".to_owned(),
|
||||
account_warmup_model: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -47,6 +50,52 @@ fn runtime_settings_keep_account_rotation_global() {
|
||||
assert!(settings.validate().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runtime_settings_require_model_when_warmup_is_enabled() {
|
||||
let settings = RuntimeSettingsUpdate {
|
||||
account_warmup_enabled: true,
|
||||
..settings_with_margin(3_600)
|
||||
};
|
||||
assert!(settings.validate().is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn warmup_settings_round_trip_with_database_constraint() {
|
||||
let Some(database) = TestDatabase::create("warmup_settings").await else {
|
||||
return;
|
||||
};
|
||||
let repository = PgRuntimeSettingsRepository::new(database.pool.clone());
|
||||
let mut update = settings_with_margin(3_600);
|
||||
update.account_warmup_enabled = true;
|
||||
update.account_warmup_schedule_time = "08:00,13:00".to_owned();
|
||||
update.account_warmup_model = Some("test-model".to_owned());
|
||||
repository
|
||||
.update_runtime_settings(update)
|
||||
.await
|
||||
.expect("save warmup settings");
|
||||
|
||||
let settings = repository
|
||||
.load_runtime_settings()
|
||||
.await
|
||||
.expect("load warmup settings");
|
||||
assert!(settings.account_warmup_enabled);
|
||||
assert_eq!(settings.account_warmup_schedule_time, "08:00,13:00");
|
||||
assert_eq!(settings.account_warmup_model.as_deref(), Some("test-model"));
|
||||
|
||||
let error = sqlx::query("update runtime_settings set account_warmup_model = null where id = 1")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect_err("enabled warmup requires a model");
|
||||
assert_eq!(
|
||||
error
|
||||
.as_database_error()
|
||||
.and_then(|error| error.code())
|
||||
.as_deref(),
|
||||
Some("23514")
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unlimited_default_account_concurrency_round_trips_without_relaxing_other_limits() {
|
||||
let Some(database) = TestDatabase::create("unlimited_default_concurrency").await else {
|
||||
|
||||
@@ -18,6 +18,7 @@ use std::collections::{BTreeMap, BTreeSet, HashMap};
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::{Duration, Instant, SystemTime};
|
||||
|
||||
use futures::StreamExt as _;
|
||||
use gateway_core::account::{
|
||||
AccountErrorReason, AccountQuotaSignals, CredentialRevision, CredentialState,
|
||||
OpaqueProviderData, ProviderAccount, ProviderAccountId, ProviderAccountStore,
|
||||
@@ -31,6 +32,7 @@ use gateway_core::provider_ports::{
|
||||
use gateway_protocol::openai::events::{
|
||||
ParsedRateLimits, RateLimitDetails, RateLimitWindow, parse_rate_limit_headers,
|
||||
};
|
||||
use gateway_protocol::openai::sse::{SseEvent, SseEventDecoder};
|
||||
use reqwest::Client;
|
||||
use secrecy::ExposeSecret;
|
||||
use serde_json::{Map, Value};
|
||||
@@ -39,9 +41,10 @@ use tokio::sync::Mutex;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::transport::profile::CodexWireProfileState;
|
||||
use crate::transport::protocol::responses::CodexResponsesRequest;
|
||||
use crate::transport::{
|
||||
CodexBackendClient, CodexClientError, CodexRateLimitResetCredits,
|
||||
CodexRateLimitResetCreditsConsumeResult, CodexRequestContext,
|
||||
CodexBackendClient, CodexBackendStreamingResponse, CodexClientError,
|
||||
CodexRateLimitResetCredits, CodexRateLimitResetCreditsConsumeResult, CodexRequestContext,
|
||||
};
|
||||
|
||||
use super::repository::{CodexCredentialRepository, CredentialRepositoryError};
|
||||
@@ -67,6 +70,8 @@ const INITIAL_QUOTA_SYNC_BATCH: usize = 100;
|
||||
// 5xx 上游拒绝的短退避重试预算;指数退避 1s/2s,吞掉瞬时抖动。
|
||||
const QUOTA_FETCH_5XX_MAX_RETRIES: u32 = 2;
|
||||
const QUOTA_FETCH_5XX_BASE_DELAY: Duration = Duration::from_secs(1);
|
||||
const WARMUP_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
const WARMUP_STREAM_MAX_BYTES: usize = 256 * 1024;
|
||||
|
||||
/// OpenAI Provider 主动额度刷新的调度策略。
|
||||
///
|
||||
@@ -106,6 +111,82 @@ impl CodexQuotaSyncSummary {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct CodexWarmupSummary {
|
||||
pub warmed_up: u64,
|
||||
pub skipped_active: u64,
|
||||
pub skipped_exhausted: u64,
|
||||
pub failed: u64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
enum WarmupTerminal {
|
||||
Missing,
|
||||
Completed,
|
||||
Failed,
|
||||
}
|
||||
|
||||
fn observe_warmup_events(events: Vec<SseEvent>, terminal: &mut WarmupTerminal) {
|
||||
for event in events {
|
||||
let parsed = serde_json::from_str::<Value>(&event.data).ok();
|
||||
let kind = parsed
|
||||
.as_ref()
|
||||
.and_then(|data| data.get("type"))
|
||||
.and_then(Value::as_str)
|
||||
.or(event.event.as_deref());
|
||||
match kind {
|
||||
Some("response.completed") => {
|
||||
// 只有明确完成的响应才算预热成功;失败事件即使随后出现完成帧也优先。
|
||||
if *terminal != WarmupTerminal::Failed {
|
||||
*terminal = if parsed.is_some()
|
||||
&& parsed
|
||||
.as_ref()
|
||||
.and_then(|data| data.pointer("/response/status"))
|
||||
.and_then(Value::as_str)
|
||||
.is_none_or(|status| status == "completed")
|
||||
{
|
||||
WarmupTerminal::Completed
|
||||
} else {
|
||||
WarmupTerminal::Failed
|
||||
};
|
||||
}
|
||||
}
|
||||
Some("response.failed" | "response.incomplete" | "error") => {
|
||||
*terminal = WarmupTerminal::Failed;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn consume_warmup_sse(
|
||||
response: &mut CodexBackendStreamingResponse,
|
||||
) -> Result<(), &'static str> {
|
||||
let mut decoder = SseEventDecoder::default();
|
||||
let mut terminal = WarmupTerminal::Missing;
|
||||
let mut received_bytes = 0usize;
|
||||
while let Some(chunk) = response.body.next().await {
|
||||
let chunk = chunk.map_err(|_| "stream_read_failed")?;
|
||||
received_bytes = received_bytes.saturating_add(chunk.len());
|
||||
if received_bytes > WARMUP_STREAM_MAX_BYTES {
|
||||
return Err("stream_size_limit_exceeded");
|
||||
}
|
||||
let events = decoder.push(&chunk).map_err(|_| "invalid_sse")?;
|
||||
observe_warmup_events(events, &mut terminal);
|
||||
if terminal != WarmupTerminal::Missing {
|
||||
break;
|
||||
}
|
||||
}
|
||||
if terminal == WarmupTerminal::Missing {
|
||||
observe_warmup_events(decoder.finish().map_err(|_| "invalid_sse")?, &mut terminal);
|
||||
}
|
||||
match terminal {
|
||||
WarmupTerminal::Completed => Ok(()),
|
||||
WarmupTerminal::Missing => Err("completion_event_missing"),
|
||||
WarmupTerminal::Failed => Err("response_failed"),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum CodexCredentialQuotaError {
|
||||
#[error("Codex quota response is invalid")]
|
||||
@@ -898,6 +979,198 @@ impl CodexCredentialQuotaService {
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn runtime_policy(&self) -> &Arc<dyn ProviderRuntimePolicyPort> {
|
||||
&self.runtime_policy
|
||||
}
|
||||
|
||||
/// 批量预激活 OAuth 账号的 5h 配额滑动窗口。
|
||||
pub async fn execute_warmup(
|
||||
&self,
|
||||
model: &str,
|
||||
) -> Result<CodexWarmupSummary, CodexCredentialQuotaError> {
|
||||
let mut accounts = self.repository.list_for_provider().await?;
|
||||
accounts.retain(|account| {
|
||||
account.authentication_kind() == crate::credential::CODEX_AUTHENTICATION_KIND_OAUTH
|
||||
&& account.enabled()
|
||||
&& matches!(
|
||||
account.credential_state(),
|
||||
CredentialState::Unknown | CredentialState::Ready
|
||||
)
|
||||
});
|
||||
let mut summary = CodexWarmupSummary::default();
|
||||
if accounts.is_empty() {
|
||||
return Ok(summary);
|
||||
}
|
||||
let account_ids = accounts
|
||||
.iter()
|
||||
.map(|account| account.id().clone())
|
||||
.collect::<Vec<_>>();
|
||||
let observed = self.store.get_quotas(&account_ids).await?;
|
||||
let observed_snapshots = observed
|
||||
.iter()
|
||||
.filter_map(|obs| {
|
||||
quota_snapshot_from_observation(obs)
|
||||
.map(|snapshot| (obs.account_id.clone(), snapshot))
|
||||
})
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
|
||||
let client = CodexBackendClient::new(
|
||||
self.http.clone(),
|
||||
self.base_url.clone(),
|
||||
self.profile.clone(),
|
||||
);
|
||||
|
||||
let now_utc = chrono::Utc::now();
|
||||
for account in accounts {
|
||||
if let Some(snapshot) = observed_snapshots.get(account.id()) {
|
||||
// 1. 周线已触顶或耗尽跳过
|
||||
let weekly_exhausted = snapshot
|
||||
.windows()
|
||||
.iter()
|
||||
.any(|w| w.kind() == CodexQuotaWindowKind::Weekly && w.limit_reached());
|
||||
if weekly_exhausted {
|
||||
summary.skipped_exhausted += 1;
|
||||
continue;
|
||||
}
|
||||
// 2. 5h 窗口当前活跃且距重置时间 > 30 分钟跳过
|
||||
let has_active_5h = snapshot.windows().iter().any(|w| {
|
||||
w.kind() == CodexQuotaWindowKind::ShortTerm
|
||||
&& w.reset_at().is_some_and(|reset_at| {
|
||||
reset_at > now_utc + chrono::Duration::minutes(30)
|
||||
})
|
||||
});
|
||||
if has_active_5h {
|
||||
summary.skipped_active += 1;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
let credential = match self.repository.load_runtime_credential(&account).await {
|
||||
Ok(cred) => cred,
|
||||
Err(_) => {
|
||||
summary.failed += 1;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let authorization = match credential.authentication.authorization_header() {
|
||||
Ok(auth) => auth,
|
||||
Err(_) => {
|
||||
summary.failed += 1;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let mut body = Map::new();
|
||||
body.insert("model".to_owned(), Value::String(model.to_owned()));
|
||||
body.insert(
|
||||
"input".to_owned(),
|
||||
serde_json::json!([{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [{"type": "input_text", "text": "hello"}]
|
||||
}]),
|
||||
);
|
||||
body.insert("stream".to_owned(), Value::Bool(true));
|
||||
body.insert("store".to_owned(), Value::Bool(false));
|
||||
body.insert(
|
||||
"service_tier".to_owned(),
|
||||
Value::String("default".to_owned()),
|
||||
);
|
||||
body.insert(
|
||||
"reasoning".to_owned(),
|
||||
serde_json::json!({"effort": "none"}),
|
||||
);
|
||||
body.insert("text".to_owned(), serde_json::json!({"verbosity": "low"}));
|
||||
|
||||
let upstream_request = CodexResponsesRequest::from_body(body);
|
||||
let request_id = format!("warmup_{}", Uuid::now_v7().simple());
|
||||
let client_for_account = match client.for_account(&account) {
|
||||
Ok(c) => c,
|
||||
Err(error) => {
|
||||
tracing::warn!(account_id = %account.id(), error = %error, "warmup client for account failed");
|
||||
summary.failed += 1;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let context = crate::transport::CodexRequestContext::auxiliary(
|
||||
authorization.expose_secret(),
|
||||
account.upstream_account_id(),
|
||||
&request_id,
|
||||
None,
|
||||
);
|
||||
|
||||
// HTTP 200 只确认响应头;必须消费 SSE 直到终态才能确认预热结果和额度事件。
|
||||
let attempt = tokio::time::timeout(WARMUP_REQUEST_TIMEOUT, async {
|
||||
let mut response = client_for_account
|
||||
.create_response_stream_http_sse(&upstream_request, context)
|
||||
.await
|
||||
.map_err(|_| "request_rejected")?;
|
||||
let terminal = consume_warmup_sse(&mut response).await;
|
||||
let rate_limit_updates = match response.rate_limit_updates.as_ref() {
|
||||
Some(updates) => std::mem::take(&mut *updates.lock().await),
|
||||
None => Vec::new(),
|
||||
};
|
||||
Ok::<_, &'static str>((terminal, response.rate_limit_headers, rate_limit_updates))
|
||||
})
|
||||
.await;
|
||||
match attempt {
|
||||
Ok(Ok((terminal, headers, rate_limit_updates))) => {
|
||||
match terminal {
|
||||
Ok(()) => {
|
||||
// 被动额度同步会把成功请求的窗口事实标为可用,失败流不能复用该结论。
|
||||
if !headers.is_empty()
|
||||
&& let Err(error) =
|
||||
self.synchronize_passive_headers(&account, &headers).await
|
||||
{
|
||||
tracing::warn!(account_id = %account.id(), error = %error, "OpenAI warmup rate-limit header sync failed");
|
||||
}
|
||||
if !rate_limit_updates.is_empty()
|
||||
&& let Err(error) = self
|
||||
.synchronize_passive_rate_limits(&account, &rate_limit_updates)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(account_id = %account.id(), error = %error, "OpenAI warmup rate-limit event sync failed");
|
||||
}
|
||||
summary.warmed_up += 1;
|
||||
tracing::info!(
|
||||
account_id = %account.id(),
|
||||
model,
|
||||
"OpenAI account warmed up successfully"
|
||||
);
|
||||
}
|
||||
Err(reason) => {
|
||||
summary.failed += 1;
|
||||
tracing::warn!(
|
||||
account_id = %account.id(),
|
||||
reason,
|
||||
"OpenAI account warmup stream failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(Err(reason)) => {
|
||||
summary.failed += 1;
|
||||
tracing::warn!(
|
||||
account_id = %account.id(),
|
||||
reason,
|
||||
"OpenAI account warmup request rejected by upstream"
|
||||
);
|
||||
}
|
||||
Err(_) => {
|
||||
summary.failed += 1;
|
||||
tracing::warn!(
|
||||
account_id = %account.id(),
|
||||
"OpenAI account warmup request timed out"
|
||||
);
|
||||
}
|
||||
}
|
||||
tokio::time::sleep(Duration::from_secs(2)).await;
|
||||
}
|
||||
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
/// `quota_observed_at` 为空代表首次异步观察尚未成功;不另建同步状态表。
|
||||
fn initial_quota_sync_accounts(
|
||||
accounts: &[ProviderAccount],
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
//! OpenAI Provider 向 Host 贡献的后台 worker。
|
||||
|
||||
use std::sync::Mutex;
|
||||
|
||||
use super::*;
|
||||
use crate::transport::profile::cli_release::CliReleaseService;
|
||||
use crate::transport::profile::platform_release::PlatformDesktopReleaseService;
|
||||
@@ -19,6 +21,8 @@ pub(super) const QUOTA_CHECK_INTERVAL: Duration = Duration::from_secs(30);
|
||||
pub(super) const DESKTOP_RELEASE_WORKER_OWNER: &str = "openai-desktop-release";
|
||||
pub(super) const MODEL_ETAG_WORKER_OWNER: &str = "openai-model-etag";
|
||||
pub(super) const MODEL_CATALOG_WORKER_OWNER: &str = "openai-model-catalog";
|
||||
pub(super) const WARMUP_WORKER_OWNER: &str = "openai-account-warmup";
|
||||
pub(super) const WARMUP_CHECK_INTERVAL: Duration = Duration::from_secs(30);
|
||||
|
||||
pub(crate) fn worker_contributions(
|
||||
refresh: Arc<CodexCredentialRefreshService>,
|
||||
@@ -35,6 +39,7 @@ pub(crate) fn worker_contributions(
|
||||
let desktop_release_id =
|
||||
WorkerId::try_new(WorkerKind::QuotaCatalogHealth, DESKTOP_RELEASE_WORKER_OWNER)?;
|
||||
let cli_release_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, "openai-cli-release")?;
|
||||
let warmup_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, WARMUP_WORKER_OWNER)?;
|
||||
let mut contributions = Vec::new();
|
||||
if oauth_refresh_enabled {
|
||||
contributions.push(WorkerContribution::Registration(scheduled_registration(
|
||||
@@ -64,7 +69,14 @@ pub(crate) fn worker_contributions(
|
||||
WorkerContribution::Registration(scheduled_registration(
|
||||
quota_id,
|
||||
QUOTA_CHECK_INTERVAL,
|
||||
Box::new(OpenAiQuotaTask { quota }),
|
||||
Box::new(OpenAiQuotaTask {
|
||||
quota: Arc::clone("a),
|
||||
}),
|
||||
)?),
|
||||
WorkerContribution::Registration(scheduled_registration(
|
||||
warmup_id,
|
||||
WARMUP_CHECK_INTERVAL,
|
||||
Box::new(OpenAiWarmupTask::new(Arc::clone("a))),
|
||||
)?),
|
||||
WorkerContribution::Registration(scheduled_registration(
|
||||
catalog_id,
|
||||
@@ -329,3 +341,103 @@ impl ScheduledTask for OpenAiPlatformDesktopReleaseTask {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct OpenAiWarmupTask {
|
||||
quota: Arc<CodexCredentialQuotaService>,
|
||||
last_slot_executed: Mutex<Option<String>>,
|
||||
}
|
||||
|
||||
impl OpenAiWarmupTask {
|
||||
pub(super) fn new(quota: Arc<CodexCredentialQuotaService>) -> Self {
|
||||
Self {
|
||||
quota,
|
||||
last_slot_executed: Mutex::new(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ScheduledTask for OpenAiWarmupTask {
|
||||
fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
if context.cancellation().is_cancelled() {
|
||||
return Ok(());
|
||||
}
|
||||
let policy = self
|
||||
.quota
|
||||
.runtime_policy()
|
||||
.load_warmup_policy()
|
||||
.await
|
||||
.map_err(|error| {
|
||||
tracing::warn!(error = %error, "OpenAI warmup policy load failed");
|
||||
WorkerTaskError::safe("OpenAI warmup policy load failed")
|
||||
})?;
|
||||
if !policy.enabled() {
|
||||
return Ok(());
|
||||
}
|
||||
use chrono::{Datelike as _, Timelike as _};
|
||||
let china_now = chrono::Utc::now() + chrono::Duration::hours(8);
|
||||
let hour = china_now.time().hour();
|
||||
let minute = china_now.time().minute();
|
||||
let scheduled_times = policy.scheduled_times();
|
||||
let matched = scheduled_times
|
||||
.iter()
|
||||
.any(|&(h, m)| h == hour && m == minute);
|
||||
if !matched {
|
||||
return Ok(());
|
||||
}
|
||||
let slot_key = format!(
|
||||
"{:04}-{:02}-{:02} {:02}:{:02}",
|
||||
china_now.date_naive().year(),
|
||||
china_now.date_naive().month(),
|
||||
china_now.date_naive().day(),
|
||||
hour,
|
||||
minute
|
||||
);
|
||||
{
|
||||
let mut last = self
|
||||
.last_slot_executed
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if last.as_deref() == Some(&slot_key) {
|
||||
return Ok(());
|
||||
}
|
||||
*last = Some(slot_key);
|
||||
}
|
||||
let Some(model) = policy.model() else {
|
||||
return Err(WorkerTaskError::safe("OpenAI warmup model is missing"));
|
||||
};
|
||||
tracing::info!(hour, minute, model, "OpenAI account warmup cycle started");
|
||||
let outcome = tokio::select! {
|
||||
() = context.cancellation().cancelled() => return Ok(()),
|
||||
outcome = self.quota.execute_warmup(model) => outcome,
|
||||
};
|
||||
match outcome {
|
||||
Ok(summary) => {
|
||||
tracing::info!(
|
||||
warmed_up = summary.warmed_up,
|
||||
skipped_active = summary.skipped_active,
|
||||
skipped_exhausted = summary.skipped_exhausted,
|
||||
failed = summary.failed,
|
||||
"OpenAI account warmup cycle completed"
|
||||
);
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(error = %error, "OpenAI account warmup cycle failed");
|
||||
}
|
||||
}
|
||||
let china_after = chrono::Utc::now() + chrono::Duration::hours(8);
|
||||
if china_after.date_naive() == china_now.date_naive()
|
||||
&& china_after.hour() == hour
|
||||
&& china_after.minute() == minute
|
||||
{
|
||||
// 持有本轮 leader lease 到时间槽结束,避免其他实例在同一分钟再次执行。
|
||||
let hold = Duration::from_secs(u64::from(61 - china_after.second()));
|
||||
tokio::select! {
|
||||
() = context.cancellation().cancelled() => return Ok(()),
|
||||
() = tokio::time::sleep(hold) => {}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -97,7 +97,7 @@ async fn openai_bundle_exposes_one_core_provider_and_drains_worker_contributions
|
||||
assert_eq!(bundle.core_provider().name(), "openai");
|
||||
assert_eq!(bundle.admin_provider().provider_kind().as_str(), "openai");
|
||||
let contributions = bundle.take_worker_contributions();
|
||||
assert_eq!(contributions.len(), 7);
|
||||
assert_eq!(contributions.len(), 8);
|
||||
assert!(
|
||||
contributions
|
||||
.iter()
|
||||
@@ -128,6 +128,7 @@ async fn openai_bundle_exposes_one_core_provider_and_drains_worker_contributions
|
||||
("openai-cli-release", APPCAST_POLL_INTERVAL),
|
||||
("openai-platform-desktop-release", APPCAST_POLL_INTERVAL),
|
||||
("openai", Duration::from_secs(30)),
|
||||
("openai-account-warmup", Duration::from_secs(30)),
|
||||
(
|
||||
"openai-model-catalog",
|
||||
config.config.quota_refresh_policy().interval(),
|
||||
|
||||
@@ -7,6 +7,7 @@ mod scheduling;
|
||||
mod slots;
|
||||
mod snapshot;
|
||||
mod subscription;
|
||||
mod warmup;
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::SystemTime;
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn warmup_consumes_completed_sse_and_persists_rate_limit_event() {
|
||||
let server = MockServer::start().await;
|
||||
let sse = concat!(
|
||||
"event: codex.rate_limits\n",
|
||||
"data: {\"type\":\"codex.rate_limits\",\"plan_type\":\"team\",\"rate_limits\":{\"allowed\":true,\"limit_reached\":false,\"primary\":{\"used_percent\":5,\"window_minutes\":300,\"reset_at\":1900000000}}}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_warmup\",\"status\":\"completed\",\"output\":[]}}\n\n",
|
||||
"data: [DONE]\n\n",
|
||||
);
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream"))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_warmup_completed").await;
|
||||
let service = quota_service_with_base_url(
|
||||
&store,
|
||||
reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.build()
|
||||
.expect("client"),
|
||||
server.uri(),
|
||||
);
|
||||
|
||||
let summary = service.execute_warmup("test-model").await.expect("warmup");
|
||||
|
||||
assert_eq!(summary.warmed_up, 1);
|
||||
assert_eq!(summary.failed, 0);
|
||||
assert_eq!(
|
||||
store.account("acct_warmup_completed").unwrap().plan_type(),
|
||||
Some("team")
|
||||
);
|
||||
assert!(store.quota_json("acct_warmup_completed").is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn warmup_rejects_failed_sse_despite_successful_http_headers() {
|
||||
let server = MockServer::start().await;
|
||||
let sse = concat!(
|
||||
"event: codex.rate_limits\n",
|
||||
"data: {\"type\":\"codex.rate_limits\",\"plan_type\":\"team\",\"rate_limits\":{\"allowed\":false,\"limit_reached\":true,\"primary\":{\"used_percent\":100,\"window_minutes\":300,\"reset_at\":1900000000}}}\n\n",
|
||||
"event: response.failed\n",
|
||||
"data: {\"type\":\"response.failed\",\"response\":{\"status\":\"failed\"}}\n\n",
|
||||
"data: [DONE]\n\n",
|
||||
);
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("x-codex-allowed", "false")
|
||||
.set_body_raw(sse, "text/event-stream"),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_warmup_failed").await;
|
||||
let service = quota_service_with_base_url(
|
||||
&store,
|
||||
reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.build()
|
||||
.expect("client"),
|
||||
server.uri(),
|
||||
);
|
||||
|
||||
let summary = service.execute_warmup("test-model").await.expect("warmup");
|
||||
|
||||
assert_eq!(summary.warmed_up, 0);
|
||||
assert_eq!(summary.failed, 1);
|
||||
assert!(store.quota_json("acct_warmup_failed").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn warmup_skips_expired_credential() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.respond_with(ResponseTemplate::new(200))
|
||||
.expect(0)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_warmup_expired").await;
|
||||
let account = store.account("acct_warmup_expired").expect("account");
|
||||
persist_credential_state(&store, &account, CredentialState::Expired).await;
|
||||
let service = quota_service_with_base_url(
|
||||
&store,
|
||||
reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.build()
|
||||
.expect("client"),
|
||||
server.uri(),
|
||||
);
|
||||
|
||||
let summary = service.execute_warmup("test-model").await.expect("warmup");
|
||||
|
||||
assert_eq!(summary.warmed_up, 0);
|
||||
assert_eq!(summary.failed, 0);
|
||||
}
|
||||
@@ -18,3 +18,4 @@ fcfae912ddefd54551ebf8e7e733442519d2862d2977134d548ebc1e75c1dc34 0017_unlimited
|
||||
deb4ef30b3a8a1c7a76055dbd09e5d232d9bb85b712205f5fb9c38c28417900d 0018_plugins.sql
|
||||
667fa2969ab9fc714a74e1a53470f738298d78ed6a248e07172c828bd702452c 0019_plugin_version_configurations.sql
|
||||
3c34fca5a04118684399d50e6923513e14478e7cf3e87f5ff78090ef9468918c 0020_proxy_auto_location.sql
|
||||
75acee6b60ce7f11bd7a370b0a67c9e5c3907318cdfbfdb3e155ae529ac97e94 0021_account_warmup.sql
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
-- 账号模型预激活/预热:支持按每日指定时间对 OAuth 账号发起极微小请求,
|
||||
-- 提前并行启动 5h 配额滑动窗口,最大化全天可用额度。默认关闭。
|
||||
alter table runtime_settings
|
||||
add column account_warmup_enabled boolean not null default false,
|
||||
add column account_warmup_schedule_time text not null default '08:00'
|
||||
check (
|
||||
octet_length(account_warmup_schedule_time) between 5 and 255
|
||||
and account_warmup_schedule_time = btrim(account_warmup_schedule_time)
|
||||
and account_warmup_schedule_time !~ '[[:cntrl:]]'
|
||||
and account_warmup_schedule_time ~ '^([01][0-9]|2[0-3]):[0-5][0-9](,([01][0-9]|2[0-3]):[0-5][0-9])*$'
|
||||
),
|
||||
add column account_warmup_model text
|
||||
check (
|
||||
account_warmup_model is null
|
||||
or (
|
||||
octet_length(account_warmup_model) between 1 and 128
|
||||
and account_warmup_model = btrim(account_warmup_model)
|
||||
and account_warmup_model !~ '[[:cntrl:]]'
|
||||
)
|
||||
),
|
||||
add constraint account_warmup_requires_model
|
||||
check (not account_warmup_enabled or account_warmup_model is not null);
|
||||
@@ -34,6 +34,9 @@ export interface RuntimeSettings {
|
||||
accountAutoFreezeProbeEnabled: boolean
|
||||
accountAutoFreezeProbeModel: string | null
|
||||
accountAutoFreezeAdaptiveConcurrency: boolean
|
||||
accountWarmupEnabled: boolean
|
||||
accountWarmupScheduleTime: string
|
||||
accountWarmupModel: string | null
|
||||
updatedAt: string
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
<script setup lang="ts">
|
||||
import { BaseCard, BaseForm, BaseFormItem, BaseInput, BaseSwitch } from '@codex-proxy/ui'
|
||||
import { Clock, Sparkles } from '@lucide/vue'
|
||||
|
||||
const enabled = defineModel<boolean>('enabled', { required: true })
|
||||
const scheduleTime = defineModel<string>('scheduleTime', { required: true })
|
||||
const model = defineModel<string>('model', { required: true })
|
||||
</script>
|
||||
|
||||
<template>
|
||||
<BaseCard
|
||||
title="账号预激活"
|
||||
description="在指定时间向空闲的 Codex OAuth 账号发送一次请求,尝试启动额度窗口"
|
||||
>
|
||||
<BaseForm class="max-w-6xl sm:grid-cols-2">
|
||||
<BaseSwitch
|
||||
v-model="enabled"
|
||||
class="col-span-full justify-self-start"
|
||||
label="启用每日预激活"
|
||||
show-label
|
||||
/>
|
||||
|
||||
<BaseFormItem
|
||||
label="执行时间(北京时间)"
|
||||
description="HH:MM;多个时间用英文逗号分隔"
|
||||
>
|
||||
<BaseInput
|
||||
v-model="scheduleTime"
|
||||
:disabled="!enabled"
|
||||
aria-label="每日激活时间"
|
||||
placeholder="08:00"
|
||||
>
|
||||
<template #prefix>
|
||||
<Clock class="size-4" />
|
||||
</template>
|
||||
</BaseInput>
|
||||
</BaseFormItem>
|
||||
|
||||
<BaseFormItem
|
||||
label="预激活模型"
|
||||
description="启用时必填,请填写账号可用的模型"
|
||||
:required="enabled"
|
||||
>
|
||||
<BaseInput
|
||||
v-model="model"
|
||||
:disabled="!enabled"
|
||||
aria-label="预激活模型"
|
||||
placeholder="请输入模型名称"
|
||||
>
|
||||
<template #prefix>
|
||||
<Sparkles class="size-4" />
|
||||
</template>
|
||||
</BaseInput>
|
||||
</BaseFormItem>
|
||||
</BaseForm>
|
||||
</BaseCard>
|
||||
</template>
|
||||
@@ -49,6 +49,9 @@ export function useSettingsForm() {
|
||||
accountAutoFreezeProbeEnabled: true,
|
||||
accountAutoFreezeProbeModel: '',
|
||||
accountAutoFreezeAdaptiveConcurrency: true,
|
||||
accountWarmupEnabled: false,
|
||||
accountWarmupScheduleTime: '08:00',
|
||||
accountWarmupModel: '',
|
||||
})
|
||||
|
||||
function snapshot() {
|
||||
@@ -137,6 +140,9 @@ export function useSettingsForm() {
|
||||
form.accountAutoFreezeProbeEnabled = data.accountAutoFreezeProbeEnabled
|
||||
form.accountAutoFreezeProbeModel = data.accountAutoFreezeProbeModel ?? ''
|
||||
form.accountAutoFreezeAdaptiveConcurrency = data.accountAutoFreezeAdaptiveConcurrency
|
||||
form.accountWarmupEnabled = data.accountWarmupEnabled
|
||||
form.accountWarmupScheduleTime = data.accountWarmupScheduleTime ?? '08:00'
|
||||
form.accountWarmupModel = data.accountWarmupModel ?? ''
|
||||
mappings.value = Object.entries(data.modelMappings || {}).map(([requestedModel, upstreamModel]) => ({
|
||||
requestedModel,
|
||||
upstreamModel: String(upstreamModel),
|
||||
@@ -241,6 +247,21 @@ export function useSettingsForm() {
|
||||
toast.warning('探测模型名称不能超过 128 个字符')
|
||||
return
|
||||
}
|
||||
const scheduleTime = form.accountWarmupScheduleTime.trim()
|
||||
const timeRegex = /^(?:[01]\d|2[0-3]):[0-5]\d(?:,(?:[01]\d|2[0-3]):[0-5]\d)*$/
|
||||
if (!scheduleTime || !timeRegex.test(scheduleTime)) {
|
||||
toast.warning('预激活时间格式无效,请输入 HH:MM 格式(如 08:00 或 08:00,13:00)')
|
||||
return
|
||||
}
|
||||
const warmupModel = form.accountWarmupModel.trim()
|
||||
if (form.accountWarmupEnabled && !warmupModel) {
|
||||
toast.warning('启用预激活时请选择模型')
|
||||
return
|
||||
}
|
||||
if (warmupModel.length > 128) {
|
||||
toast.warning('预激活模型名称不能超过 128 个字符')
|
||||
return
|
||||
}
|
||||
await saveAction.run(async () => {
|
||||
const result = await updateSettings({
|
||||
providerRequestProfiles: requestProfileUpdates(
|
||||
@@ -271,6 +292,9 @@ export function useSettingsForm() {
|
||||
accountAutoFreezeProbeEnabled: form.accountAutoFreezeProbeEnabled,
|
||||
accountAutoFreezeProbeModel: probeModel || null,
|
||||
accountAutoFreezeAdaptiveConcurrency: form.accountAutoFreezeAdaptiveConcurrency,
|
||||
accountWarmupEnabled: form.accountWarmupEnabled,
|
||||
accountWarmupScheduleTime: scheduleTime,
|
||||
accountWarmupModel: warmupModel || null,
|
||||
})
|
||||
applySettings(result)
|
||||
toast.success('设置已保存')
|
||||
|
||||
@@ -6,6 +6,7 @@ import { computed, reactive, shallowRef, watch } from 'vue'
|
||||
import { useRoute, useRouter } from 'vue-router'
|
||||
|
||||
import AccountAutoFreezeCard from './components/AccountAutoFreezeCard.vue'
|
||||
import AccountWarmupCard from './components/AccountWarmupCard.vue'
|
||||
import SettingsBackupSection from './components/backup/SettingsBackupSection.vue'
|
||||
import ClientProfileCard from './components/ClientProfileCard.vue'
|
||||
import ModelAliasesCard from './components/ModelAliasesCard.vue'
|
||||
@@ -161,6 +162,11 @@ watch(section, (value) => {
|
||||
v-model:probe-model="form.accountAutoFreezeProbeModel"
|
||||
v-model:adaptive-concurrency="form.accountAutoFreezeAdaptiveConcurrency"
|
||||
/>
|
||||
<AccountWarmupCard
|
||||
v-model:enabled="form.accountWarmupEnabled"
|
||||
v-model:schedule-time="form.accountWarmupScheduleTime"
|
||||
v-model:model="form.accountWarmupModel"
|
||||
/>
|
||||
</template>
|
||||
|
||||
<div v-if="visited.has('upstream')" v-show="section === 'upstream'" class="grid min-w-0 gap-5">
|
||||
|
||||
Reference in New Issue
Block a user