From 7cb186c4ae08b1ff87e8a0266e48343f62b31686 Mon Sep 17 00:00:00 2001 From: zyycn Date: Sat, 11 Jul 2026 00:43:47 +0800 Subject: [PATCH] refactor: conform backend to authoritative architecture Co-authored-by: Codex --- backend/src/accounts/pool.rs | 1371 ++++++++--------- .../accounts/pool/{types.rs => filters.rs} | 2 + backend/src/accounts/pool/runtime.rs | 715 --------- backend/src/accounts/pool/state.rs | 712 +++++++++ backend/src/accounts/scheduler.rs | 104 +- .../src/accounts/scheduler/strategy/mod.rs | 5 - .../scheduler/strategy/quota_reset.rs | 5 +- .../scheduler/strategy/round_robin.rs | 3 +- .../src/accounts/scheduler/strategy/smart.rs | 3 +- .../src/accounts/scheduler/strategy/sticky.rs | 3 +- .../src/accounts/scheduler/strategy/types.rs | 137 -- backend/src/accounts/store.rs | 10 + backend/src/accounts/store/queries.rs | 162 +- backend/src/accounts/store/write.rs | 166 -- backend/src/api/admin/accounts_routes.rs | 18 +- .../src/api/admin/accounts_routes/actions.rs | 321 ---- .../admin/accounts_routes/export_routes.rs | 24 + .../admin/accounts_routes/import_routes.rs | 20 + .../admin/accounts_routes/lifecycle_routes.rs | 110 ++ .../api/admin/accounts_routes/oauth_routes.rs | 49 + .../api/admin/accounts_routes/probe_routes.rs | 62 + .../accounts_routes/{stats.rs => query.rs} | 0 .../api/admin/accounts_routes/quota_routes.rs | 61 + backend/src/api/admin/auth_routes.rs | 143 ++ backend/src/api/admin/dashboard_routes.rs | 304 +++- .../src/api/admin/dashboard_routes/data.rs | 304 ---- backend/src/api/admin/mod.rs | 1 + backend/src/api/admin/router.rs | 2 +- backend/src/api/admin/session.rs | 145 +- backend/src/api/admin/usage_routes.rs | 8 +- backend/src/dispatch/affinity/identity.rs | 222 --- backend/src/dispatch/affinity/mod.rs | 3 +- backend/src/dispatch/affinity/resolve.rs | 377 ++++- backend/src/dispatch/recording.rs | 2 +- .../src/dispatch/recovery/implicit_resume.rs | 153 +- backend/src/dispatch/service.rs | 906 +++++++---- backend/src/dispatch/service/compact.rs | 290 ---- backend/src/dispatch/service/complete.rs | 575 ------- backend/src/dispatch/service/stream.rs | 619 -------- backend/src/dispatch/stream/lifecycle.rs | 882 ++++++++--- backend/src/dispatch/stream/live.rs | 263 +++- backend/src/dispatch/upstream_call.rs | 335 +++- backend/src/telemetry/buckets/store.rs | 127 ++ backend/src/telemetry/rebuild.rs | 122 +- backend/src/telemetry/usage/query.rs | 443 +++++- backend/src/telemetry/usage/store.rs | 21 +- backend/src/telemetry/usage/store/query.rs | 432 ------ .../src/upstream/openai/protocol/websocket.rs | 5 +- .../validation.rs => websocket_errors.rs} | 0 .../src/upstream/openai/transport/client.rs | 1 + .../{client/requests.rs => client_sse.rs} | 0 .../upstream/openai/transport/websocket.rs | 57 +- .../openai/transport/websocket/types.rs | 189 --- .../stream.rs => websocket_frames.rs} | 147 ++ .../{export.rs => export_routes.rs} | 0 .../{import.rs => import_routes.rs} | 0 .../{lifecycle.rs => lifecycle_routes.rs} | 0 .../tests/api/admin/accounts_routes/mod.rs | 12 +- .../{oauth.rs => oauth_routes.rs} | 0 .../{probe.rs => probe_routes.rs} | 0 .../{quota.rs => quota_routes.rs} | 0 .../dispatch/{service/mod.rs => service.rs} | 68 +- docs/architecture.md | 102 +- 63 files changed, 5719 insertions(+), 5604 deletions(-) rename backend/src/accounts/pool/{types.rs => filters.rs} (98%) delete mode 100644 backend/src/accounts/pool/runtime.rs create mode 100644 backend/src/accounts/pool/state.rs delete mode 100644 backend/src/accounts/scheduler/strategy/types.rs delete mode 100644 backend/src/api/admin/accounts_routes/actions.rs create mode 100644 backend/src/api/admin/accounts_routes/export_routes.rs create mode 100644 backend/src/api/admin/accounts_routes/import_routes.rs create mode 100644 backend/src/api/admin/accounts_routes/lifecycle_routes.rs create mode 100644 backend/src/api/admin/accounts_routes/oauth_routes.rs create mode 100644 backend/src/api/admin/accounts_routes/probe_routes.rs rename backend/src/api/admin/accounts_routes/{stats.rs => query.rs} (100%) create mode 100644 backend/src/api/admin/accounts_routes/quota_routes.rs create mode 100644 backend/src/api/admin/auth_routes.rs delete mode 100644 backend/src/api/admin/dashboard_routes/data.rs delete mode 100644 backend/src/dispatch/affinity/identity.rs delete mode 100644 backend/src/dispatch/service/compact.rs delete mode 100644 backend/src/dispatch/service/complete.rs delete mode 100644 backend/src/dispatch/service/stream.rs delete mode 100644 backend/src/telemetry/usage/store/query.rs rename backend/src/upstream/openai/protocol/{websocket/validation.rs => websocket_errors.rs} (100%) rename backend/src/upstream/openai/transport/{client/requests.rs => client_sse.rs} (100%) delete mode 100644 backend/src/upstream/openai/transport/websocket/types.rs rename backend/src/upstream/openai/transport/{websocket/stream.rs => websocket_frames.rs} (81%) rename backend/tests/api/admin/accounts_routes/{export.rs => export_routes.rs} (100%) rename backend/tests/api/admin/accounts_routes/{import.rs => import_routes.rs} (100%) rename backend/tests/api/admin/accounts_routes/{lifecycle.rs => lifecycle_routes.rs} (100%) rename backend/tests/api/admin/accounts_routes/{oauth.rs => oauth_routes.rs} (100%) rename backend/tests/api/admin/accounts_routes/{probe.rs => probe_routes.rs} (100%) rename backend/tests/api/admin/accounts_routes/{quota.rs => quota_routes.rs} (100%) rename backend/tests/dispatch/{service/mod.rs => service.rs} (95%) diff --git a/backend/src/accounts/pool.rs b/backend/src/accounts/pool.rs index f9da2c973..8319e2281 100644 --- a/backend/src/accounts/pool.rs +++ b/backend/src/accounts/pool.rs @@ -31,723 +31,722 @@ use crate::upstream::openai::protocol::events::{ pub use crate::accounts::scheduler::RotationStrategy; -/// 账号池配置。 -mod types; -pub use types::*; +mod filters; +mod state; -#[derive(Debug)] -struct AccountAcquireWithStatusRefresh { - acquired: Option, - refreshed_accounts: Vec, -} - -#[derive(Debug)] -struct DistinctPlanAccountsWithStatusRefresh { - accounts: Vec, - refreshed_accounts: Vec, -} - -#[derive(Debug)] -struct AccountCapacitySummaryWithStatusRefresh { - summary: AccountCapacitySummary, - refreshed_accounts: Vec, -} - -/// 账号池容量摘要。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct AccountCapacitySummary { - /// 单账号最大并发请求数。 - pub max_concurrent_per_account: usize, - /// 所有可用账号合计槽位数。 - pub total_slots: usize, - /// 当前已占用槽位数。 - pub used_slots: usize, - /// 当前可用槽位数。 - pub available_slots: usize, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -struct RuntimeAccountStateSnapshot { - status: AccountStatus, - quota_limit_reached: bool, - quota_verify_required: bool, - quota_cooldown_until: Option>, - cloudflare_cooldown_until: Option>, - window_request_count: u64, - window_input_tokens: u64, - window_output_tokens: u64, - window_cached_tokens: u64, - window_image_input_tokens: u64, - window_image_output_tokens: u64, - window_image_request_count: u64, - window_image_request_failed_count: u64, - window_started_at: Option>, - window_reset_at: Option>, - limit_window_seconds: Option, -} - -impl From<&Account> for RuntimeAccountStateSnapshot { - fn from(account: &Account) -> Self { - Self { - status: account.status, - quota_limit_reached: account.quota_limit_reached, - quota_verify_required: account.quota_verify_required, - quota_cooldown_until: account.quota_cooldown_until, - cloudflare_cooldown_until: account.cloudflare_cooldown_until, - window_request_count: account.window_request_count, - window_input_tokens: account.window_input_tokens, - window_output_tokens: account.window_output_tokens, - window_cached_tokens: account.window_cached_tokens, - window_image_input_tokens: account.window_image_input_tokens, - window_image_output_tokens: account.window_image_output_tokens, - window_image_request_count: account.window_image_request_count, - window_image_request_failed_count: account.window_image_request_failed_count, - window_started_at: account.window_started_at, - window_reset_at: account.window_reset_at, - limit_window_seconds: account.limit_window_seconds, - } - } -} - -impl RuntimeAccountStateSnapshot { - fn usage_window_changed(&self, other: &Self) -> bool { - self.window_request_count != other.window_request_count - || self.window_input_tokens != other.window_input_tokens - || self.window_output_tokens != other.window_output_tokens - || self.window_cached_tokens != other.window_cached_tokens - || self.window_image_input_tokens != other.window_image_input_tokens - || self.window_image_output_tokens != other.window_image_output_tokens - || self.window_image_request_count != other.window_image_request_count - || self.window_image_request_failed_count != other.window_image_request_failed_count - || self.window_started_at != other.window_started_at - || self.window_reset_at != other.window_reset_at - || self.limit_window_seconds != other.limit_window_seconds - } -} +pub use filters::*; +pub use state::*; +/// 账号运行参数中不支持热更新的部分。 #[derive(Debug, Clone)] -struct RefreshedAccountState { - account: Account, - sync_usage_window: bool, +pub struct AccountPoolStaticSettings { + pub skip_quota_limited: bool, + pub tier_priority: Vec, } -/// 运行时窗口用量增量。 -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct AccountWindowUsageDelta { - /// 输入 token 增量。 - pub input_tokens: u64, - /// 输出 token 增量。 - pub output_tokens: u64, - /// 缓存 token 增量。 - pub cached_tokens: u64, - /// 图片输入 token 增量。 - pub image_input_tokens: u64, - /// 图片输出 token 增量。 - pub image_output_tokens: u64, - /// 是否记录一次成功图片请求。 - pub image_request_succeeded: bool, - /// 是否记录一次失败图片请求。 - pub image_request_failed: bool, -} - -/// 纯内存账号池：负责账号存储、在途槽位与运行时状态维护。 -/// -/// 账号**选择**已剥离到 [`crate::accounts::scheduler`]：候选过滤走 -/// [`candidates::filter`]，策略选择走 [`AccountScheduler::select`]。账号池只提供选择 -/// 所需的只读视图（候选切片 + `slot_count`），并持有调度器实例。 -#[derive(Debug)] -pub struct AccountPool { - accounts: IndexMap, - slots: BTreeMap>, - options: AccountPoolOptions, - scheduler: AccountScheduler, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -struct AccountSlot { - created_at: DateTime, - model: Option, -} - -/// 账号释放结果。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ReleasedAccountUsage { - /// 被释放槽位对应的模型。槽位已被过期清理时为空。 - pub model: Option, -} - -impl Default for AccountPool { - fn default() -> Self { - Self::with_options(AccountPoolOptions::default()) - } -} - -impl AccountPool { - /// 使用指定配置创建账号池。 - pub fn with_options(options: AccountPoolOptions) -> Self { - Self { - accounts: IndexMap::new(), - slots: BTreeMap::new(), - options, - scheduler: AccountScheduler::new(ScoreWeights::default()), - } - } - - /// 插入或替换账号。 - pub fn insert(&mut self, account: Account) { - self.accounts.insert(account.id.clone(), account); - } - - /// 获取账号池中的账号快照。 - pub fn get(&self, account_id: &str) -> Option { - self.accounts.get(account_id).cloned() - } - - /// 移除账号及其在途槽位与运行时反馈。 - pub fn remove(&mut self, account_id: &str) -> bool { - let removed = self.accounts.shift_remove(account_id).is_some(); - self.slots.remove(account_id); - self.scheduler.forget_feedback(account_id); - removed - } - - /// 清空账号池。 - pub fn clear(&mut self) { - self.accounts.clear(); - self.slots.clear(); - self.scheduler.reset_cursor(); - self.scheduler.clear_feedback(); - } - - /// 替换账号池运行参数。 - pub fn set_options(&mut self, options: AccountPoolOptions) { - if self.options.rotation_strategy != options.rotation_strategy { - self.scheduler.reset_cursor(); - } - self.options = options; - } - - /// 回灌一次请求结果到调度器的运行时反馈（供 Smart 打分）。 - pub fn report_feedback(&self, account_id: &str, success: bool, first_token_ms: Option) { - self.scheduler - .report_feedback(account_id, success, first_token_ms); - } - - /// 更新账号状态。 - pub fn set_status(&mut self, account_id: &str, status: AccountStatus) -> bool { - let Some(account) = self.accounts.get_mut(account_id) else { - return false; - }; - let status = status_after_quota_limit(status, account.quota_limit_reached); - account.status = status; - if status != AccountStatus::Active { - self.slots.remove(account_id); - } - true - } - - /// 标记账号配额限流直到指定时间。 - pub fn mark_quota_limited_until( - &mut self, - account_id: &str, - cooldown_until: DateTime, - ) -> bool { - let Some(account) = self.accounts.get_mut(account_id) else { - return false; - }; - let final_cooldown_until = account - .quota_cooldown_until - .filter(|existing| *existing > cooldown_until) - .unwrap_or(cooldown_until); - account.quota_limit_reached = true; - account.quota_cooldown_until = Some(final_cooldown_until); - account.status = status_after_quota_limit(account.status, true); - account.window_reset_at = account - .window_reset_at - .filter(|existing| *existing > final_cooldown_until) - .or(Some(final_cooldown_until)); - if let Ok(seconds) = (final_cooldown_until - Utc::now()) - .to_std() - .map(|duration| duration.as_secs()) - { - account.limit_window_seconds.get_or_insert(seconds); - } - self.slots.remove(account_id); - true - } - - /// 应用已经验证过的配额状态。 - pub fn apply_quota_state( - &mut self, - account_id: &str, - limit_reached: bool, - cooldown_until: Option>, - ) -> bool { - let Some(account) = self.accounts.get_mut(account_id) else { - return false; - }; - account.quota_verify_required = false; - account.quota_limit_reached = limit_reached; - account.quota_cooldown_until = limit_reached.then_some(cooldown_until).flatten(); - match (account.status, limit_reached) { - (AccountStatus::Active, true) => account.status = AccountStatus::QuotaExhausted, - (AccountStatus::QuotaExhausted, false) => account.status = AccountStatus::Active, - _ => {} - } - if let Some(cooldown_until) = account.quota_cooldown_until { - account.window_reset_at = account - .window_reset_at - .filter(|existing| *existing > cooldown_until) - .or(Some(cooldown_until)); - } - if limit_reached { - self.slots.remove(account_id); - } - true - } - - /// 同步上游返回的限流窗口信息。 - pub fn sync_rate_limit_window( - &mut self, - account_id: &str, - new_reset_at: DateTime, - limit_window_seconds: Option, - ) -> bool { - self.sync_rate_limit_window_at(account_id, new_reset_at, limit_window_seconds, Utc::now()) - } - - /// 记录账号运行时窗口 token 用量。 - pub fn record_window_token_usage( - &mut self, - account_id: &str, - usage: AccountWindowUsageDelta, - ) -> bool { - let Some(account) = self.accounts.get_mut(account_id) else { - return false; - }; - account.image_input_tokens = account - .image_input_tokens - .saturating_add(usage.image_input_tokens); - account.image_output_tokens = account - .image_output_tokens - .saturating_add(usage.image_output_tokens); - if usage.image_request_succeeded { - account.image_request_count = account.image_request_count.saturating_add(1); - } - if usage.image_request_failed { - account.image_request_failed_count = - account.image_request_failed_count.saturating_add(1); - } - account.window_input_tokens = account - .window_input_tokens - .saturating_add(usage.input_tokens); - account.window_output_tokens = account - .window_output_tokens - .saturating_add(usage.output_tokens); - account.window_cached_tokens = account - .window_cached_tokens - .saturating_add(usage.cached_tokens); - account.window_image_input_tokens = account - .window_image_input_tokens - .saturating_add(usage.image_input_tokens); - account.window_image_output_tokens = account - .window_image_output_tokens - .saturating_add(usage.image_output_tokens); - if usage.image_request_succeeded { - account.window_image_request_count = - account.window_image_request_count.saturating_add(1); - } - if usage.image_request_failed { - account.window_image_request_failed_count = - account.window_image_request_failed_count.saturating_add(1); - } - true - } - - /// 标记账号处于 Cloudflare 冷却期。 - pub fn set_cloudflare_cooldown_until( - &mut self, - account_id: &str, - cooldown_until: DateTime, - ) -> bool { - let Some(account) = self.accounts.get_mut(account_id) else { - return false; - }; - account.cloudflare_cooldown_until = Some(cooldown_until); - self.slots.remove(account_id); - true - } - - /// 使用完整调度请求获取账号。 - pub fn acquire_with(&mut self, request: &AccountAcquireRequest) -> Option { - self.acquire_with_status_refresh(request).acquired - } - - fn acquire_with_status_refresh( - &mut self, - request: &AccountAcquireRequest, - ) -> AccountAcquireWithStatusRefresh { - self.cleanup_stale_slots(request.now); - let refreshed_accounts = self.refresh_account_statuses(request.now); - let candidates = self.filter_candidates(request); - // 亲和优先：required 硬锁定该账号；preferred 若在候选内则直接选中； - // 二者都不命中时，回退到调度器按当前策略选择。 - let selected = if let Some(required_account_id) = &request.required_account_id { - candidates.first().cloned().map(|account| { - ( - account, - self.previous_slot_at(required_account_id), - "required", - ) - }) - } else if let Some(preferred_account_id) = &request.preferred_account_id { - candidates - .iter() - .find(|account| account.id == *preferred_account_id) - .cloned() - .map(|account| { - ( - account, - self.previous_slot_at(preferred_account_id), - "preferred", - ) - }) - } else { - None - } - .or_else(|| { - let account = self.select_candidate(&candidates, request.now)?; - let previous = self.previous_slot_at(&account.id); - Some((account, previous, self.options.rotation_strategy.as_str())) - }); - let Some((selected, previous_slot_at, selection_source)) = selected else { - return AccountAcquireWithStatusRefresh { - acquired: None, - refreshed_accounts, - }; - }; - tracing::info!( - model = %request.model, - rotation_strategy = self.options.rotation_strategy.as_str(), - selection_source, - account_id = %selected.id, - candidate_count = candidates.len(), - slot_count = self.slot_count(&selected.id), - request_count = selected.request_count, - window_request_count = selected.window_request_count, - window_reset_at = ?selected.window_reset_at, - last_used_at = selected.last_used_at.as_deref().unwrap_or_default(), - quota_limit_reached = selected.quota_limit_reached, - quota_cooldown_until = ?selected.quota_cooldown_until, - previous_slot_at = ?previous_slot_at, - "account selected for upstream request" - ); - - let selected_id = selected.id.clone(); - self.push_slot(&selected_id, request.now, Some(&request.model)); - AccountAcquireWithStatusRefresh { - acquired: Some(AcquiredAccount { - account: selected, - previous_slot_at, - }), - refreshed_accounts, - } - } - - fn refresh_account_statuses(&mut self, now: DateTime) -> Vec { - let mut refreshed_accounts = Vec::new(); - let mut expired_account_ids = Vec::new(); - for (account_id, account) in &mut self.accounts { - let before = RuntimeAccountStateSnapshot::from(&*account); - if account.status == AccountStatus::Active && access_token_expired(account, now) { - account.status = AccountStatus::Expired; - expired_account_ids.push(account_id.clone()); - } else { - refresh_quota_window(account, now); - refresh_cloudflare_cooldown(account, now); - } - let after = RuntimeAccountStateSnapshot::from(&*account); - if before != after { - let sync_usage_window = before.usage_window_changed(&after); - refreshed_accounts.push(RefreshedAccountState { - account: account.clone(), - sync_usage_window, - }); - } - } - for account_id in &expired_account_ids { - self.slots.remove(account_id); - } - refreshed_accounts - } - - /// 释放指定账号的一个在途槽位。 - pub fn release(&mut self, account_id: &str) -> Option { - self.release_slot(account_id, true) - } - - /// 释放指定账号的一个在途槽位，不累计请求用量。 - pub fn release_without_request_usage( - &mut self, - account_id: &str, - ) -> Option { - self.release_slot(account_id, false) - } - - fn release_slot( - &mut self, - account_id: &str, - record_request_usage: bool, - ) -> Option { - self.accounts.get(account_id)?; - - let mut remove_slots = false; - let slot = self.slots.get_mut(account_id).and_then(|slots| { - let slot = slots.pop_front(); - remove_slots = slots.is_empty(); - slot - }); - if remove_slots { - self.slots.remove(account_id); - } - if record_request_usage { - self.mark_request_usage(account_id, Utc::now()); - } - Some(ReleasedAccountUsage { - model: slot.and_then(|slot| slot.model), - }) - } - - /// 计算账号池容量摘要。 - pub fn capacity_summary(&mut self, now: DateTime) -> AccountCapacitySummary { - self.capacity_summary_with_status_refresh(now).summary - } - - fn capacity_summary_with_status_refresh( - &mut self, - now: DateTime, - ) -> AccountCapacitySummaryWithStatusRefresh { - self.cleanup_stale_slots(now); - let refreshed_accounts = self.refresh_account_statuses(now); - let is_capacity_account = |account: &Account| { - account.status == AccountStatus::Active - && candidates::quota_available_at(account, now, self.options.skip_quota_limited) - }; - let active_accounts = self - .accounts - .values() - .filter(|account| is_capacity_account(account)) - .count(); - let total_slots = active_accounts * self.options.max_concurrent_per_account; - let used_slots = self - .slots - .iter() - .filter(|(account_id, _)| { - self.accounts - .get(*account_id) - .is_some_and(&is_capacity_account) - }) - .map(|(_, slots)| slots.len().min(self.options.max_concurrent_per_account)) - .sum(); - - AccountCapacitySummaryWithStatusRefresh { - summary: AccountCapacitySummary { - max_concurrent_per_account: self.options.max_concurrent_per_account, - total_slots, - used_slots, - available_slots: total_slots.saturating_sub(used_slots), +impl AccountPoolStaticSettings { + /// 将设置快照转换为账号池选项。 + pub fn pool_options(&self, settings: &crate::settings::SettingsSnapshot) -> AccountPoolOptions { + AccountPoolOptions { + max_concurrent_per_account: settings.max_concurrent_per_account, + rotation_strategy: match settings.rotation_strategy.as_str() { + "smart" => RotationStrategy::Smart, + "quota_reset_priority" => RotationStrategy::QuotaResetPriority, + "round_robin" => RotationStrategy::RoundRobin, + "sticky" => RotationStrategy::Sticky, + _ => RotationStrategy::Smart, }, - refreshed_accounts, + skip_quota_limited: self.skip_quota_limited, + tier_priority: self.tier_priority.clone(), + ..AccountPoolOptions::default() } } +} + +#[derive(Clone)] +pub struct RuntimeAccountPoolService { + pool: Arc>, + store: Arc, + usage_store: Arc, + request_interval: Arc>, +} + +impl RuntimeAccountPoolService { + /// 构造运行时账号池服务。 + pub fn new( + store: Arc, + usage_store: Arc, + options: AccountPoolOptions, + request_interval_ms: u64, + ) -> Self { + Self { + pool: Arc::new(tokio::sync::Mutex::new(AccountPool::with_options(options))), + store, + usage_store, + request_interval: Arc::new(std::sync::RwLock::new(StdDuration::from_millis( + request_interval_ms, + ))), + } + } + + /// 更新账号池运行参数。 + pub async fn apply_options(&self, mut options: AccountPoolOptions, request_interval_ms: u64) { + let mut pool = self.pool.lock().await; + options.model_plan_allowlist = pool.options.model_plan_allowlist.clone(); + options.fetched_model_plan_types = pool.options.fetched_model_plan_types.clone(); + pool.set_options(options); + *self + .request_interval + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + StdDuration::from_millis(request_interval_ms); + } + + /// 持续接收运行时设置并更新账号池参数。 + pub async fn subscribe_settings( + self: Arc, + mut receiver: tokio::sync::watch::Receiver, + static_settings: AccountPoolStaticSettings, + ) { + while receiver.changed().await.is_ok() { + let settings = receiver.borrow_and_update().clone(); + self.apply_options( + static_settings.pool_options(&settings), + settings.request_interval_ms, + ) + .await; + } + } + + /// 更新模型到可用账号计划的调度约束。 + pub(crate) async fn apply_model_plan_routing( + &self, + model_plan_allowlist: BTreeMap>, + fetched_model_plan_types: BTreeSet, + ) { + let mut pool = self.pool.lock().await; + pool.options.model_plan_allowlist = model_plan_allowlist; + pool.options.fetched_model_plan_types = fetched_model_plan_types; + } + + /// 从持久化账号恢复运行时账号池。 + pub async fn restore_from_store(&self) -> Result { + let mut accounts = self + .store + .list_pool_accounts() + .await + .map_err(|_| RuntimeAccountPoolError::Generic)?; + let account_ids = accounts + .iter() + .map(|account| account.id.clone()) + .collect::>(); + let mut usage = self + .usage_store + .snapshots(&account_ids) + .await + .map_err(|_| RuntimeAccountPoolError::Generic)?; + for account in &mut accounts { + if let Some(snapshot) = usage.remove(&account.id) { + apply_usage_snapshot(account, snapshot); + } + } + let count = accounts.len(); + let mut pool = self.pool.lock().await; + pool.clear(); + for account in accounts { + pool.insert(account); + } + Ok(count) + } + + /// 读取账号池容量摘要。 + pub async fn capacity_summary(&self, now: DateTime) -> AccountCapacitySummary { + let refresh = self + .pool + .lock() + .await + .capacity_summary_with_status_refresh(now); + self.persist_runtime_account_states(refresh.refreshed_accounts) + .await; + refresh.summary + } + + /// 使用当前时间读取账号池容量摘要。 + pub async fn capacity_summary_now(&self) -> AccountCapacitySummary { + self.capacity_summary(Utc::now()).await + } + + /// 释放账号在途槽位。 + pub async fn release(&self, account_id: &str) { + self.release_slot(account_id, true).await; + } + + /// 释放账号在途槽位，不累计请求用量。 + pub async fn release_without_request_usage(&self, account_id: &str) { + self.release_slot(account_id, false).await; + } + + async fn release_slot(&self, account_id: &str, record_request_usage: bool) { + let released = if record_request_usage { + self.pool.lock().await.release(account_id) + } else { + self.pool + .lock() + .await + .release_without_request_usage(account_id) + }; + let Some(released) = released else { + return; + }; + if !record_request_usage { + return; + } + if let Err(error) = self.usage_store.record_request(account_id).await { + tracing::warn!( + account_id, + error = %error, + "failed to persist account request usage" + ); + } + let Some(model) = released.model else { + tracing::debug!( + account_id, + "released account without active slot; skipped model request usage persistence" + ); + return; + }; + if let Err(error) = self + .usage_store + .record_model_usage_delta( + account_id, + &model, + AccountModelUsageDelta { + requests: 1, + ..AccountModelUsageDelta::default() + }, + ) + .await + { + tracing::warn!( + account_id, + model, + error = %error, + "failed to persist account model request usage" + ); + } + } + + /// 更新账号状态并同步内存池。 + pub async fn set_status(&self, account_id: &str, status: AccountStatus) -> bool { + let persisted = match self.store.set_status(account_id, status).await { + Ok(persisted) => persisted, + Err(error) => { + tracing::warn!( + account_id, + error = %error, + "failed to persist account status" + ); + false + } + }; + let in_memory = self.pool.lock().await.set_status(account_id, status); + persisted || in_memory + } + + /// 标记账号 quota 冷却状态。 + pub async fn mark_quota_limited_until(&self, account_id: &str, until: DateTime) -> bool { + let persisted = match self.store.mark_quota_limited_until(account_id, until).await { + Ok(persisted) => persisted, + Err(error) => { + tracing::warn!( + account_id, + error = %error, + "failed to persist quota cooldown state" + ); + false + } + }; + let in_memory = self + .pool + .lock() + .await + .mark_quota_limited_until(account_id, until); + persisted || in_memory + } + + /// 按请求上下文获取可用账号。 + pub async fn acquire_with(&self, request: &AccountAcquireRequest) -> Option { + let refresh = self.pool.lock().await.acquire_with_status_refresh(request); + self.persist_runtime_account_states(refresh.refreshed_accounts) + .await; + refresh.acquired + } /// 按订阅计划各选一个可用账号，用于刷新模型列表。 - pub fn distinct_plan_accounts(&mut self, now: DateTime) -> Vec { - self.distinct_plan_accounts_with_status_refresh(now) - .accounts + pub async fn distinct_plan_accounts(&self, now: DateTime) -> Vec { + let refresh = self + .pool + .lock() + .await + .distinct_plan_accounts_with_status_refresh(now); + self.persist_runtime_account_states(refresh.refreshed_accounts) + .await; + refresh.accounts } - fn distinct_plan_accounts_with_status_refresh( - &mut self, - now: DateTime, - ) -> DistinctPlanAccountsWithStatusRefresh { - self.cleanup_stale_slots(now); - let refreshed_accounts = self.refresh_account_statuses(now); - let mut by_plan = IndexMap::>::new(); - - let max_concurrent = self.options.max_concurrent_per_account; - let skip_quota_limited = self.options.skip_quota_limited; - { - let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len); - for account in self.accounts.values() { - if !candidates::is_model_refresh_available( - account, - max_concurrent, - skip_quota_limited, - &slot_count, - now, - ) { - continue; + async fn persist_runtime_account_states(&self, refreshed_states: Vec) { + for refreshed in refreshed_states { + if let Err(error) = self + .store + .sync_runtime_account_state(&refreshed.account) + .await + { + tracing::warn!( + account_id = %refreshed.account.id, + error = %error, + "failed to persist refreshed runtime account state" + ); + } + if refreshed.sync_usage_window { + if let Err(error) = self + .usage_store + .sync_runtime_window( + &refreshed.account.id, + account_usage_window(&refreshed.account), + ) + .await + { + tracing::warn!( + account_id = %refreshed.account.id, + error = %error, + "failed to persist refreshed runtime account usage window" + ); } - let Some(plan_type) = account.plan_type.as_ref() else { - continue; - }; - by_plan - .entry(plan_type.clone()) - .or_default() - .push(account.clone()); } } + } - let strategy = self.options.rotation_strategy; - let mut accounts = Vec::new(); - for (plan_type, group) in by_plan { - let Some(account) = self.select_candidate(&group, now) else { - continue; - }; - tracing::info!( - plan_type, - rotation_strategy = strategy.as_str(), - account_id = %account.id, - candidate_count = group.len(), - slot_count = self.slot_count(&account.id), - request_count = account.request_count, - window_request_count = account.window_request_count, - window_reset_at = ?account.window_reset_at, - last_used_at = account.last_used_at.as_deref().unwrap_or_default(), - quota_limit_reached = account.quota_limit_reached, - quota_cooldown_until = ?account.quota_cooldown_until, - "account selected for model refresh" - ); - self.push_slot(&account.id, now, None); - accounts.push(DistinctPlanAccount { plan_type, account }); + /// 按账号上一个在途槽位控制请求间隔。 + pub async fn wait_for_request_interval(&self, acquired: &AcquiredAccount) { + let request_interval = *self + .request_interval + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if request_interval.is_zero() { + return; } - - DistinctPlanAccountsWithStatusRefresh { - accounts, - refreshed_accounts, - } - } - - fn slot_count(&self, account_id: &str) -> usize { - self.slots.get(account_id).map_or(0, VecDeque::len) - } - - fn previous_slot_at(&self, account_id: &str) -> Option> { - self.slots - .get(account_id) - .and_then(|slots| slots.back()) - .map(|slot| slot.created_at) - } - - fn push_slot(&mut self, account_id: &str, now: DateTime, model: Option<&str>) { - self.slots - .entry(account_id.to_string()) - .or_default() - .push_back(AccountSlot { - created_at: now, - model: model.map(str::to_string), - }); - } - - fn mark_request_usage(&mut self, account_id: &str, now: DateTime) -> Option { - let account = self.accounts.get_mut(account_id)?; - account.last_used_at = Some(now.to_rfc3339()); - account.request_count = account.request_count.saturating_add(1); - account.window_request_count = account.window_request_count.saturating_add(1); - if account.window_started_at.is_none() { - account.window_started_at = Some(now); - if let Some(seconds) = account.limit_window_seconds { - let seconds = i64::try_from(seconds).unwrap_or(i64::MAX); - account.window_reset_at = Some(now + Duration::seconds(seconds)); - } - } - Some(account.clone()) - } - - fn sync_rate_limit_window_at( - &mut self, - account_id: &str, - new_reset_at: DateTime, - limit_window_seconds: Option, - now: DateTime, - ) -> bool { - let Some(account) = self.accounts.get_mut(account_id) else { - return false; + let Some(previous_slot_at) = acquired.previous_slot_at else { + return; }; - if should_reset_usage_window( - account.window_reset_at, - account.limit_window_seconds, - new_reset_at, - limit_window_seconds, - ) { - reset_window_counters(account); - account.window_started_at = Some(now); + let elapsed = Utc::now() + .signed_duration_since(previous_slot_at) + .to_std() + .unwrap_or_default(); + if elapsed < request_interval { + tokio::time::sleep(request_interval - elapsed).await; } - account.window_reset_at = Some(new_reset_at); - if let Some(limit_window_seconds) = limit_window_seconds { - account.limit_window_seconds = Some(limit_window_seconds); - } - true } - fn cleanup_stale_slots(&mut self, now: DateTime) { - let ttl = self.options.stale_slot_ttl; - self.slots.retain(|_, slots| { - slots.retain(|slot| now.signed_duration_since(slot.created_at) <= ttl); - !slots.is_empty() - }); + /// 记录上游 token 用量。 + pub async fn record_token_usage(&self, account_id: &str, model: &str, usage: &CodexTokenUsage) { + self.record_response_usage(account_id, model, *usage, false) + .await; } - /// 过滤出本次请求的可用候选集合（委托 [`candidates::filter`]）。 - fn filter_candidates(&self, request: &AccountAcquireRequest) -> Vec { - let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len); - candidates::filter( - self.accounts.values(), - &CandidateFilter { - max_concurrent_per_account: self.options.max_concurrent_per_account, - skip_quota_limited: self.options.skip_quota_limited, - tier_priority: &self.options.tier_priority, - model_plan_allowlist: &self.options.model_plan_allowlist, - fetched_model_plan_types: &self.options.fetched_model_plan_types, - slot_count: &slot_count, - }, - &CandidateRequest { - model: &request.model, - exclude_account_ids: &request.exclude_account_ids, - required_account_id: request.required_account_id.as_deref(), - now: request.now, + /// 记录 Responses 请求用量。 + pub async fn record_response_usage( + &self, + account_id: &str, + model: &str, + usage: CodexTokenUsage, + image_generation_requested: bool, + ) { + let image_request_succeeded = image_generation_requested && usage.image_output_tokens > 0; + let image_request_failed = image_generation_requested && !image_request_succeeded; + let persisted_usage = AccountUsageDelta { + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cached_tokens: usage.cached_tokens, + reasoning_tokens: usage.reasoning_tokens, + total_tokens: usage.total_tokens, + image_input_tokens: usage.image_input_tokens, + image_output_tokens: usage.image_output_tokens, + image_requests: u64::from(image_request_succeeded), + image_request_failures: u64::from(image_request_failed), + ..AccountUsageDelta::default() + }; + if let Err(error) = self + .usage_store + .record_usage_delta(account_id, persisted_usage) + .await + { + tracing::warn!( + account_id, + error = %error, + "failed to persist account token usage" + ); + } + if let Err(error) = self + .usage_store + .record_model_usage_delta( + account_id, + model, + AccountModelUsageDelta { + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cached_tokens: usage.cached_tokens, + ..AccountModelUsageDelta::default() + }, + ) + .await + { + tracing::warn!( + account_id, + model, + error = %error, + "failed to persist account model token usage" + ); + } + self.pool.lock().await.record_window_token_usage( + account_id, + AccountWindowUsageDelta { + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cached_tokens: usage.cached_tokens, + image_input_tokens: usage.image_input_tokens, + image_output_tokens: usage.image_output_tokens, + image_request_succeeded, + image_request_failed, }, + ); + } + + /// 记录空响应尝试。 + pub async fn record_empty_response_attempt( + &self, + account_id: &str, + model: &str, + image_generation_requested: bool, + ) { + let usage = AccountUsageDelta { + empty_responses: 1, + image_request_failures: u64::from(image_generation_requested), + ..AccountUsageDelta::default() + }; + if let Err(error) = self.usage_store.record_usage_delta(account_id, usage).await { + tracing::warn!( + account_id, + error = %error, + "failed to persist empty response usage" + ); + } + if let Err(error) = self + .usage_store + .record_model_usage_delta( + account_id, + model, + AccountModelUsageDelta { + errors: 1, + ..AccountModelUsageDelta::default() + }, + ) + .await + { + tracing::warn!( + account_id, + model, + error = %error, + "failed to persist account model empty response usage" + ); + } + if image_generation_requested { + self.pool.lock().await.record_window_token_usage( + account_id, + AccountWindowUsageDelta { + image_request_failed: true, + ..AccountWindowUsageDelta::default() + }, + ); + } + } + + /// 根据上游返回的 rate-limit header 被动同步 quota 状态。 + pub async fn sync_passive_rate_limit_headers( + &self, + account: &Account, + headers: &[(String, String)], + ) { + self.sync_passive_rate_limit_headers_for_account( + &account.id, + account.plan_type.as_deref(), + headers, ) + .await; } - /// 按当前策略从候选集合中选择一个账号（委托 [`AccountScheduler::select`]）。 - fn select_candidate(&self, candidates: &[Account], now: DateTime) -> Option { - let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len); - self.scheduler - .select(self.options.rotation_strategy, candidates, &slot_count, now) + /// 根据上游返回的 rate-limit header 被动同步指定账号 quota 状态。 + pub async fn sync_passive_rate_limit_headers_for_account( + &self, + account_id: &str, + plan_type: Option<&str>, + headers: &[(String, String)], + ) { + let Some(rate_limits) = parse_rate_limit_headers(headers) else { + return; + }; + let existing_quota = match self.store.get_quota_json(account_id).await { + Ok(Some(quota_json)) => serde_json::from_str::("a_json).ok(), + Ok(None) => None, + Err(error) => { + tracing::warn!( + account_id = %account_id, + error = %error, + "failed to read existing quota json before passive rate-limit sync" + ); + None + } + }; + let quota = rate_limit_quota(&rate_limits, plan_type, existing_quota.as_ref()); + self.apply_quota_snapshot(account_id, "a).await; + } + + /// 读取运行时账号快照。 + pub async fn account_snapshot(&self, account_id: &str) -> Option { + self.pool.lock().await.get(account_id) + } + + /// 应用 quota 快照到持久化存储和运行时账号池。 + pub async fn apply_quota_snapshot(&self, account_id: &str, quota: &serde_json::Value) -> bool { + let limit_reached = quota_snapshot_limit_reached(quota); + let reset_at = quota_snapshot_reset_at(quota); + let cooldown_until = limit_reached.then_some(reset_at).flatten(); + let limit_window_seconds = reset_at + .and_then(|_| crate::accounts::quota::quota_snapshot_limit_window_seconds(quota)); + let quota_json = quota.to_string(); + let persisted = match self + .store + .apply_quota_snapshot(account_id, "a_json, limit_reached, cooldown_until) + .await + { + Ok(persisted) => persisted, + Err(error) => { + tracing::warn!( + account_id, + error = %error, + "failed to persist verified quota snapshot" + ); + false + } + }; + + if let Some(reset_at) = reset_at { + if let Err(error) = self + .usage_store + .sync_rate_limit_window(account_id, reset_at, limit_window_seconds) + .await + { + tracing::warn!( + account_id, + error = %error, + "failed to persist verified quota window" + ); + } + } + + let in_memory = { + let mut pool = self.pool.lock().await; + let in_memory = pool.apply_quota_state(account_id, limit_reached, cooldown_until); + if let Some(reset_at) = reset_at { + pool.sync_rate_limit_window(account_id, reset_at, limit_window_seconds); + } + in_memory + }; + + persisted || in_memory + } + + /// 从运行时账号池移除账号。 + pub async fn remove_account(&self, account_id: &str) -> bool { + self.pool.lock().await.remove(account_id) + } + + /// 从仓储同步单个账号到运行时账号池。 + pub async fn sync_account_from_store( + &self, + account_id: &str, + ) -> Result { + let account = self + .store + .get_pool_account(account_id) + .await + .map_err(|_| RuntimeAccountPoolError::Generic)?; + if let Some(mut account) = account { + let account_ids = vec![account.id.clone()]; + let mut usage = self + .usage_store + .snapshots(&account_ids) + .await + .map_err(|_| RuntimeAccountPoolError::Generic)?; + if let Some(snapshot) = usage.remove(&account.id) { + apply_usage_snapshot(&mut account, snapshot); + } + self.pool.lock().await.insert(account); + return Ok(true); + } + Ok(self.pool.lock().await.remove(account_id)) + } + + /// 设置 Cloudflare 冷却状态。 + pub async fn set_cloudflare_cooldown_until( + &self, + account_id: &str, + cooldown_until: DateTime, + ) -> bool { + let persisted = match self + .store + .set_cloudflare_cooldown_until(account_id, cooldown_until) + .await + { + Ok(persisted) => persisted, + Err(error) => { + tracing::warn!( + account_id, + error = %error, + "failed to persist Cloudflare cooldown state" + ); + false + } + }; + let in_memory = self + .pool + .lock() + .await + .set_cloudflare_cooldown_until(account_id, cooldown_until); + persisted || in_memory } } -fn access_token_expired(account: &Account, now: DateTime) -> bool { +fn apply_usage_snapshot(account: &mut Account, snapshot: AccountUsageSnapshot) { + account.request_count = snapshot.request_count; + account.empty_response_count = snapshot.empty_response_count; + account.image_input_tokens = snapshot.image_input_tokens; + account.image_output_tokens = snapshot.image_output_tokens; + account.image_request_count = snapshot.image_request_count; + account.image_request_failed_count = snapshot.image_request_failed_count; + account.window_request_count = snapshot.window_request_count; + account.window_input_tokens = snapshot.window_input_tokens; + account.window_output_tokens = snapshot.window_output_tokens; + account.window_cached_tokens = snapshot.window_cached_tokens; + account.window_image_input_tokens = snapshot.window_image_input_tokens; + account.window_image_output_tokens = snapshot.window_image_output_tokens; + account.window_image_request_count = snapshot.window_image_request_count; + account.window_image_request_failed_count = snapshot.window_image_request_failed_count; + account.window_started_at = snapshot.window_started_at; + account.window_reset_at = snapshot.window_reset_at.or(account.window_reset_at); + account.limit_window_seconds = snapshot + .limit_window_seconds + .or(account.limit_window_seconds); + account.last_used_at = snapshot.last_used_at.map(|value| value.to_rfc3339()); +} + +fn account_usage_window(account: &Account) -> AccountUsageWindow { + AccountUsageWindow { + request_count: account.window_request_count, + input_tokens: account.window_input_tokens, + output_tokens: account.window_output_tokens, + cached_tokens: account.window_cached_tokens, + image_input_tokens: account.window_image_input_tokens, + image_output_tokens: account.window_image_output_tokens, + image_request_count: account.window_image_request_count, + image_request_failed_count: account.window_image_request_failed_count, + started_at: account.window_started_at, + reset_at: account.window_reset_at, + limit_window_seconds: account.limit_window_seconds, + } +} + +/// 运行时账号池错误。 +#[derive(Debug, Error)] +pub enum RuntimeAccountPoolError { + /// 通用账号池错误。 + #[error("account pool error")] + Generic, +} + +pub(super) fn refresh_quota_window(account: &mut Account, now: DateTime) { + let window_expired = account + .window_reset_at + .is_some_and(|reset_at| now >= reset_at); + let cooldown_expired = account + .quota_cooldown_until + .is_some_and(|cooldown_until| now >= cooldown_until); + + if window_expired { + reset_window_counters(account); + account.window_started_at = Some(now); + account.window_reset_at = + next_window_reset_at(account.window_reset_at, account.limit_window_seconds, now); + } + + if account.quota_limit_reached && (cooldown_expired || window_expired) { + account.quota_verify_required = true; + account.quota_limit_reached = false; + account.quota_cooldown_until = None; + if account.status == AccountStatus::QuotaExhausted { + account.status = AccountStatus::Active; + } + } +} + +pub(super) fn status_after_quota_limit( + status: AccountStatus, + quota_limit_reached: bool, +) -> AccountStatus { + match (status, quota_limit_reached) { + (AccountStatus::Active, true) => AccountStatus::QuotaExhausted, + (status, _) => status, + } +} + +pub(super) fn reset_window_counters(account: &mut Account) { + account.window_request_count = 0; + account.window_input_tokens = 0; + account.window_output_tokens = 0; + account.window_cached_tokens = 0; + account.window_image_input_tokens = 0; + account.window_image_output_tokens = 0; + account.window_image_request_count = 0; + account.window_image_request_failed_count = 0; +} + +fn next_window_reset_at( + reset_at: Option>, + limit_window_seconds: Option, + now: DateTime, +) -> Option> { + let reset_at = reset_at?; + let window_seconds = limit_window_seconds?; + if window_seconds == 0 { + return None; + } + + let elapsed_seconds = now + .signed_duration_since(reset_at) + .num_seconds() + .max(0) + .cast_unsigned(); + let windows_to_advance = elapsed_seconds / window_seconds + 1; + let advance_seconds = window_seconds + .saturating_mul(windows_to_advance) + .min(i64::MAX as u64); + Some(reset_at + Duration::seconds(advance_seconds as i64)) +} + +pub(super) fn refresh_cloudflare_cooldown(account: &mut Account, now: DateTime) { if account - .access_token_expires_at - .is_some_and(|expires_at| now >= expires_at) + .cloudflare_cooldown_until + .is_some_and(|cooldown_until| now >= cooldown_until) { - return true; - } - - match jwt_expiry(&account.access_token, now) { - JwtExpiry::Expired => true, - JwtExpiry::Valid => false, - JwtExpiry::MissingOrInvalid => account.access_token_expires_at.is_none(), + account.cloudflare_cooldown_until = None; } } - -/// 运行时账号池服务。 -mod runtime; -pub use runtime::*; -use runtime::{ - refresh_cloudflare_cooldown, refresh_quota_window, reset_window_counters, - status_after_quota_limit, -}; diff --git a/backend/src/accounts/pool/types.rs b/backend/src/accounts/pool/filters.rs similarity index 98% rename from backend/src/accounts/pool/types.rs rename to backend/src/accounts/pool/filters.rs index 01f4951d3..1cc17433d 100644 --- a/backend/src/accounts/pool/types.rs +++ b/backend/src/accounts/pool/filters.rs @@ -1,3 +1,5 @@ +//! 账号获取过滤条件、池配置与选择结果类型。 + use super::*; #[derive(Debug, Clone, PartialEq, Eq)] diff --git a/backend/src/accounts/pool/runtime.rs b/backend/src/accounts/pool/runtime.rs deleted file mode 100644 index 987f563b2..000000000 --- a/backend/src/accounts/pool/runtime.rs +++ /dev/null @@ -1,715 +0,0 @@ -use super::*; - -/// 账号运行参数中不支持热更新的部分。 -#[derive(Debug, Clone)] -pub struct AccountPoolStaticSettings { - pub skip_quota_limited: bool, - pub tier_priority: Vec, -} - -impl AccountPoolStaticSettings { - /// 将设置快照转换为账号池选项。 - pub fn pool_options(&self, settings: &crate::settings::SettingsSnapshot) -> AccountPoolOptions { - AccountPoolOptions { - max_concurrent_per_account: settings.max_concurrent_per_account, - rotation_strategy: match settings.rotation_strategy.as_str() { - "smart" => RotationStrategy::Smart, - "quota_reset_priority" => RotationStrategy::QuotaResetPriority, - "round_robin" => RotationStrategy::RoundRobin, - "sticky" => RotationStrategy::Sticky, - _ => RotationStrategy::Smart, - }, - skip_quota_limited: self.skip_quota_limited, - tier_priority: self.tier_priority.clone(), - ..AccountPoolOptions::default() - } - } -} - -#[derive(Clone)] -pub struct RuntimeAccountPoolService { - pool: Arc>, - store: Arc, - usage_store: Arc, - request_interval: Arc>, -} - -impl RuntimeAccountPoolService { - /// 构造运行时账号池服务。 - pub fn new( - store: Arc, - usage_store: Arc, - options: AccountPoolOptions, - request_interval_ms: u64, - ) -> Self { - Self { - pool: Arc::new(tokio::sync::Mutex::new(AccountPool::with_options(options))), - store, - usage_store, - request_interval: Arc::new(std::sync::RwLock::new(StdDuration::from_millis( - request_interval_ms, - ))), - } - } - - /// 更新账号池运行参数。 - pub async fn apply_options(&self, mut options: AccountPoolOptions, request_interval_ms: u64) { - let mut pool = self.pool.lock().await; - options.model_plan_allowlist = pool.options.model_plan_allowlist.clone(); - options.fetched_model_plan_types = pool.options.fetched_model_plan_types.clone(); - pool.set_options(options); - *self - .request_interval - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner) = - StdDuration::from_millis(request_interval_ms); - } - - /// 持续接收运行时设置并更新账号池参数。 - pub async fn subscribe_settings( - self: Arc, - mut receiver: tokio::sync::watch::Receiver, - static_settings: AccountPoolStaticSettings, - ) { - while receiver.changed().await.is_ok() { - let settings = receiver.borrow_and_update().clone(); - self.apply_options( - static_settings.pool_options(&settings), - settings.request_interval_ms, - ) - .await; - } - } - - /// 更新模型到可用账号计划的调度约束。 - pub(crate) async fn apply_model_plan_routing( - &self, - model_plan_allowlist: BTreeMap>, - fetched_model_plan_types: BTreeSet, - ) { - let mut pool = self.pool.lock().await; - pool.options.model_plan_allowlist = model_plan_allowlist; - pool.options.fetched_model_plan_types = fetched_model_plan_types; - } - - /// 从持久化账号恢复运行时账号池。 - pub async fn restore_from_store(&self) -> Result { - let mut accounts = self - .store - .list_pool_accounts() - .await - .map_err(|_| RuntimeAccountPoolError::Generic)?; - let account_ids = accounts - .iter() - .map(|account| account.id.clone()) - .collect::>(); - let mut usage = self - .usage_store - .snapshots(&account_ids) - .await - .map_err(|_| RuntimeAccountPoolError::Generic)?; - for account in &mut accounts { - if let Some(snapshot) = usage.remove(&account.id) { - apply_usage_snapshot(account, snapshot); - } - } - let count = accounts.len(); - let mut pool = self.pool.lock().await; - pool.clear(); - for account in accounts { - pool.insert(account); - } - Ok(count) - } - - /// 读取账号池容量摘要。 - pub async fn capacity_summary(&self, now: DateTime) -> AccountCapacitySummary { - let refresh = self - .pool - .lock() - .await - .capacity_summary_with_status_refresh(now); - self.persist_runtime_account_states(refresh.refreshed_accounts) - .await; - refresh.summary - } - - /// 使用当前时间读取账号池容量摘要。 - pub async fn capacity_summary_now(&self) -> AccountCapacitySummary { - self.capacity_summary(Utc::now()).await - } - - /// 释放账号在途槽位。 - pub async fn release(&self, account_id: &str) { - self.release_slot(account_id, true).await; - } - - /// 释放账号在途槽位，不累计请求用量。 - pub async fn release_without_request_usage(&self, account_id: &str) { - self.release_slot(account_id, false).await; - } - - async fn release_slot(&self, account_id: &str, record_request_usage: bool) { - let released = if record_request_usage { - self.pool.lock().await.release(account_id) - } else { - self.pool - .lock() - .await - .release_without_request_usage(account_id) - }; - let Some(released) = released else { - return; - }; - if !record_request_usage { - return; - } - if let Err(error) = self.usage_store.record_request(account_id).await { - tracing::warn!( - account_id, - error = %error, - "failed to persist account request usage" - ); - } - let Some(model) = released.model else { - tracing::debug!( - account_id, - "released account without active slot; skipped model request usage persistence" - ); - return; - }; - if let Err(error) = self - .usage_store - .record_model_usage_delta( - account_id, - &model, - AccountModelUsageDelta { - requests: 1, - ..AccountModelUsageDelta::default() - }, - ) - .await - { - tracing::warn!( - account_id, - model, - error = %error, - "failed to persist account model request usage" - ); - } - } - - /// 更新账号状态并同步内存池。 - pub async fn set_status(&self, account_id: &str, status: AccountStatus) -> bool { - let persisted = match self.store.set_status(account_id, status).await { - Ok(persisted) => persisted, - Err(error) => { - tracing::warn!( - account_id, - error = %error, - "failed to persist account status" - ); - false - } - }; - let in_memory = self.pool.lock().await.set_status(account_id, status); - persisted || in_memory - } - - /// 标记账号 quota 冷却状态。 - pub async fn mark_quota_limited_until(&self, account_id: &str, until: DateTime) -> bool { - let persisted = match self.store.mark_quota_limited_until(account_id, until).await { - Ok(persisted) => persisted, - Err(error) => { - tracing::warn!( - account_id, - error = %error, - "failed to persist quota cooldown state" - ); - false - } - }; - let in_memory = self - .pool - .lock() - .await - .mark_quota_limited_until(account_id, until); - persisted || in_memory - } - - /// 按请求上下文获取可用账号。 - pub async fn acquire_with(&self, request: &AccountAcquireRequest) -> Option { - let refresh = self.pool.lock().await.acquire_with_status_refresh(request); - self.persist_runtime_account_states(refresh.refreshed_accounts) - .await; - refresh.acquired - } - - /// 按订阅计划各选一个可用账号，用于刷新模型列表。 - pub async fn distinct_plan_accounts(&self, now: DateTime) -> Vec { - let refresh = self - .pool - .lock() - .await - .distinct_plan_accounts_with_status_refresh(now); - self.persist_runtime_account_states(refresh.refreshed_accounts) - .await; - refresh.accounts - } - - async fn persist_runtime_account_states(&self, refreshed_states: Vec) { - for refreshed in refreshed_states { - if let Err(error) = self - .store - .sync_runtime_account_state(&refreshed.account) - .await - { - tracing::warn!( - account_id = %refreshed.account.id, - error = %error, - "failed to persist refreshed runtime account state" - ); - } - if refreshed.sync_usage_window { - if let Err(error) = self - .usage_store - .sync_runtime_window( - &refreshed.account.id, - account_usage_window(&refreshed.account), - ) - .await - { - tracing::warn!( - account_id = %refreshed.account.id, - error = %error, - "failed to persist refreshed runtime account usage window" - ); - } - } - } - } - - /// 按账号上一个在途槽位控制请求间隔。 - pub async fn wait_for_request_interval(&self, acquired: &AcquiredAccount) { - let request_interval = *self - .request_interval - .read() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if request_interval.is_zero() { - return; - } - let Some(previous_slot_at) = acquired.previous_slot_at else { - return; - }; - let elapsed = Utc::now() - .signed_duration_since(previous_slot_at) - .to_std() - .unwrap_or_default(); - if elapsed < request_interval { - tokio::time::sleep(request_interval - elapsed).await; - } - } - - /// 记录上游 token 用量。 - pub async fn record_token_usage(&self, account_id: &str, model: &str, usage: &CodexTokenUsage) { - self.record_response_usage(account_id, model, *usage, false) - .await; - } - - /// 记录 Responses 请求用量。 - pub async fn record_response_usage( - &self, - account_id: &str, - model: &str, - usage: CodexTokenUsage, - image_generation_requested: bool, - ) { - let image_request_succeeded = image_generation_requested && usage.image_output_tokens > 0; - let image_request_failed = image_generation_requested && !image_request_succeeded; - let persisted_usage = AccountUsageDelta { - input_tokens: usage.input_tokens, - output_tokens: usage.output_tokens, - cached_tokens: usage.cached_tokens, - reasoning_tokens: usage.reasoning_tokens, - total_tokens: usage.total_tokens, - image_input_tokens: usage.image_input_tokens, - image_output_tokens: usage.image_output_tokens, - image_requests: u64::from(image_request_succeeded), - image_request_failures: u64::from(image_request_failed), - ..AccountUsageDelta::default() - }; - if let Err(error) = self - .usage_store - .record_usage_delta(account_id, persisted_usage) - .await - { - tracing::warn!( - account_id, - error = %error, - "failed to persist account token usage" - ); - } - if let Err(error) = self - .usage_store - .record_model_usage_delta( - account_id, - model, - AccountModelUsageDelta { - input_tokens: usage.input_tokens, - output_tokens: usage.output_tokens, - cached_tokens: usage.cached_tokens, - ..AccountModelUsageDelta::default() - }, - ) - .await - { - tracing::warn!( - account_id, - model, - error = %error, - "failed to persist account model token usage" - ); - } - self.pool.lock().await.record_window_token_usage( - account_id, - AccountWindowUsageDelta { - input_tokens: usage.input_tokens, - output_tokens: usage.output_tokens, - cached_tokens: usage.cached_tokens, - image_input_tokens: usage.image_input_tokens, - image_output_tokens: usage.image_output_tokens, - image_request_succeeded, - image_request_failed, - }, - ); - } - - /// 记录空响应尝试。 - pub async fn record_empty_response_attempt( - &self, - account_id: &str, - model: &str, - image_generation_requested: bool, - ) { - let usage = AccountUsageDelta { - empty_responses: 1, - image_request_failures: u64::from(image_generation_requested), - ..AccountUsageDelta::default() - }; - if let Err(error) = self.usage_store.record_usage_delta(account_id, usage).await { - tracing::warn!( - account_id, - error = %error, - "failed to persist empty response usage" - ); - } - if let Err(error) = self - .usage_store - .record_model_usage_delta( - account_id, - model, - AccountModelUsageDelta { - errors: 1, - ..AccountModelUsageDelta::default() - }, - ) - .await - { - tracing::warn!( - account_id, - model, - error = %error, - "failed to persist account model empty response usage" - ); - } - if image_generation_requested { - self.pool.lock().await.record_window_token_usage( - account_id, - AccountWindowUsageDelta { - image_request_failed: true, - ..AccountWindowUsageDelta::default() - }, - ); - } - } - - /// 根据上游返回的 rate-limit header 被动同步 quota 状态。 - pub async fn sync_passive_rate_limit_headers( - &self, - account: &Account, - headers: &[(String, String)], - ) { - self.sync_passive_rate_limit_headers_for_account( - &account.id, - account.plan_type.as_deref(), - headers, - ) - .await; - } - - /// 根据上游返回的 rate-limit header 被动同步指定账号 quota 状态。 - pub async fn sync_passive_rate_limit_headers_for_account( - &self, - account_id: &str, - plan_type: Option<&str>, - headers: &[(String, String)], - ) { - let Some(rate_limits) = parse_rate_limit_headers(headers) else { - return; - }; - let existing_quota = match self.store.get_quota_json(account_id).await { - Ok(Some(quota_json)) => serde_json::from_str::("a_json).ok(), - Ok(None) => None, - Err(error) => { - tracing::warn!( - account_id = %account_id, - error = %error, - "failed to read existing quota json before passive rate-limit sync" - ); - None - } - }; - let quota = rate_limit_quota(&rate_limits, plan_type, existing_quota.as_ref()); - self.apply_quota_snapshot(account_id, "a).await; - } - - /// 读取运行时账号快照。 - pub async fn account_snapshot(&self, account_id: &str) -> Option { - self.pool.lock().await.get(account_id) - } - - /// 应用 quota 快照到持久化存储和运行时账号池。 - pub async fn apply_quota_snapshot(&self, account_id: &str, quota: &serde_json::Value) -> bool { - let limit_reached = quota_snapshot_limit_reached(quota); - let reset_at = quota_snapshot_reset_at(quota); - let cooldown_until = limit_reached.then_some(reset_at).flatten(); - let limit_window_seconds = reset_at - .and_then(|_| crate::accounts::quota::quota_snapshot_limit_window_seconds(quota)); - let quota_json = quota.to_string(); - let persisted = match self - .store - .apply_quota_snapshot(account_id, "a_json, limit_reached, cooldown_until) - .await - { - Ok(persisted) => persisted, - Err(error) => { - tracing::warn!( - account_id, - error = %error, - "failed to persist verified quota snapshot" - ); - false - } - }; - - if let Some(reset_at) = reset_at { - if let Err(error) = self - .usage_store - .sync_rate_limit_window(account_id, reset_at, limit_window_seconds) - .await - { - tracing::warn!( - account_id, - error = %error, - "failed to persist verified quota window" - ); - } - } - - let in_memory = { - let mut pool = self.pool.lock().await; - let in_memory = pool.apply_quota_state(account_id, limit_reached, cooldown_until); - if let Some(reset_at) = reset_at { - pool.sync_rate_limit_window(account_id, reset_at, limit_window_seconds); - } - in_memory - }; - - persisted || in_memory - } - - /// 从运行时账号池移除账号。 - pub async fn remove_account(&self, account_id: &str) -> bool { - self.pool.lock().await.remove(account_id) - } - - /// 从仓储同步单个账号到运行时账号池。 - pub async fn sync_account_from_store( - &self, - account_id: &str, - ) -> Result { - let account = self - .store - .get_pool_account(account_id) - .await - .map_err(|_| RuntimeAccountPoolError::Generic)?; - if let Some(mut account) = account { - let account_ids = vec![account.id.clone()]; - let mut usage = self - .usage_store - .snapshots(&account_ids) - .await - .map_err(|_| RuntimeAccountPoolError::Generic)?; - if let Some(snapshot) = usage.remove(&account.id) { - apply_usage_snapshot(&mut account, snapshot); - } - self.pool.lock().await.insert(account); - return Ok(true); - } - Ok(self.pool.lock().await.remove(account_id)) - } - - /// 设置 Cloudflare 冷却状态。 - pub async fn set_cloudflare_cooldown_until( - &self, - account_id: &str, - cooldown_until: DateTime, - ) -> bool { - let persisted = match self - .store - .set_cloudflare_cooldown_until(account_id, cooldown_until) - .await - { - Ok(persisted) => persisted, - Err(error) => { - tracing::warn!( - account_id, - error = %error, - "failed to persist Cloudflare cooldown state" - ); - false - } - }; - let in_memory = self - .pool - .lock() - .await - .set_cloudflare_cooldown_until(account_id, cooldown_until); - persisted || in_memory - } -} - -fn apply_usage_snapshot(account: &mut Account, snapshot: AccountUsageSnapshot) { - account.request_count = snapshot.request_count; - account.empty_response_count = snapshot.empty_response_count; - account.image_input_tokens = snapshot.image_input_tokens; - account.image_output_tokens = snapshot.image_output_tokens; - account.image_request_count = snapshot.image_request_count; - account.image_request_failed_count = snapshot.image_request_failed_count; - account.window_request_count = snapshot.window_request_count; - account.window_input_tokens = snapshot.window_input_tokens; - account.window_output_tokens = snapshot.window_output_tokens; - account.window_cached_tokens = snapshot.window_cached_tokens; - account.window_image_input_tokens = snapshot.window_image_input_tokens; - account.window_image_output_tokens = snapshot.window_image_output_tokens; - account.window_image_request_count = snapshot.window_image_request_count; - account.window_image_request_failed_count = snapshot.window_image_request_failed_count; - account.window_started_at = snapshot.window_started_at; - account.window_reset_at = snapshot.window_reset_at.or(account.window_reset_at); - account.limit_window_seconds = snapshot - .limit_window_seconds - .or(account.limit_window_seconds); - account.last_used_at = snapshot.last_used_at.map(|value| value.to_rfc3339()); -} - -fn account_usage_window(account: &Account) -> AccountUsageWindow { - AccountUsageWindow { - request_count: account.window_request_count, - input_tokens: account.window_input_tokens, - output_tokens: account.window_output_tokens, - cached_tokens: account.window_cached_tokens, - image_input_tokens: account.window_image_input_tokens, - image_output_tokens: account.window_image_output_tokens, - image_request_count: account.window_image_request_count, - image_request_failed_count: account.window_image_request_failed_count, - started_at: account.window_started_at, - reset_at: account.window_reset_at, - limit_window_seconds: account.limit_window_seconds, - } -} - -/// 运行时账号池错误。 -#[derive(Debug, Error)] -pub enum RuntimeAccountPoolError { - /// 通用账号池错误。 - #[error("account pool error")] - Generic, -} - -pub(super) fn refresh_quota_window(account: &mut Account, now: DateTime) { - let window_expired = account - .window_reset_at - .is_some_and(|reset_at| now >= reset_at); - let cooldown_expired = account - .quota_cooldown_until - .is_some_and(|cooldown_until| now >= cooldown_until); - - if window_expired { - reset_window_counters(account); - account.window_started_at = Some(now); - account.window_reset_at = - next_window_reset_at(account.window_reset_at, account.limit_window_seconds, now); - } - - if account.quota_limit_reached && (cooldown_expired || window_expired) { - account.quota_verify_required = true; - account.quota_limit_reached = false; - account.quota_cooldown_until = None; - if account.status == AccountStatus::QuotaExhausted { - account.status = AccountStatus::Active; - } - } -} - -pub(super) fn status_after_quota_limit( - status: AccountStatus, - quota_limit_reached: bool, -) -> AccountStatus { - match (status, quota_limit_reached) { - (AccountStatus::Active, true) => AccountStatus::QuotaExhausted, - (status, _) => status, - } -} - -pub(super) fn reset_window_counters(account: &mut Account) { - account.window_request_count = 0; - account.window_input_tokens = 0; - account.window_output_tokens = 0; - account.window_cached_tokens = 0; - account.window_image_input_tokens = 0; - account.window_image_output_tokens = 0; - account.window_image_request_count = 0; - account.window_image_request_failed_count = 0; -} - -fn next_window_reset_at( - reset_at: Option>, - limit_window_seconds: Option, - now: DateTime, -) -> Option> { - let reset_at = reset_at?; - let window_seconds = limit_window_seconds?; - if window_seconds == 0 { - return None; - } - - let elapsed_seconds = now - .signed_duration_since(reset_at) - .num_seconds() - .max(0) - .cast_unsigned(); - let windows_to_advance = elapsed_seconds / window_seconds + 1; - let advance_seconds = window_seconds - .saturating_mul(windows_to_advance) - .min(i64::MAX as u64); - Some(reset_at + Duration::seconds(advance_seconds as i64)) -} - -pub(super) fn refresh_cloudflare_cooldown(account: &mut Account, now: DateTime) { - if account - .cloudflare_cooldown_until - .is_some_and(|cooldown_until| now >= cooldown_until) - { - account.cloudflare_cooldown_until = None; - } -} diff --git a/backend/src/accounts/pool/state.rs b/backend/src/accounts/pool/state.rs new file mode 100644 index 000000000..9be6fd5d1 --- /dev/null +++ b/backend/src/accounts/pool/state.rs @@ -0,0 +1,712 @@ +//! 账号池存储态、槽位与状态变更。 + +use super::*; + +#[derive(Debug)] +pub(super) struct AccountAcquireWithStatusRefresh { + pub(super) acquired: Option, + pub(super) refreshed_accounts: Vec, +} + +#[derive(Debug)] +pub(super) struct DistinctPlanAccountsWithStatusRefresh { + pub(super) accounts: Vec, + pub(super) refreshed_accounts: Vec, +} + +#[derive(Debug)] +pub(super) struct AccountCapacitySummaryWithStatusRefresh { + pub(super) summary: AccountCapacitySummary, + pub(super) refreshed_accounts: Vec, +} + +/// 账号池容量摘要。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AccountCapacitySummary { + /// 单账号最大并发请求数。 + pub max_concurrent_per_account: usize, + /// 所有可用账号合计槽位数。 + pub total_slots: usize, + /// 当前已占用槽位数。 + pub used_slots: usize, + /// 当前可用槽位数。 + pub available_slots: usize, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct RuntimeAccountStateSnapshot { + status: AccountStatus, + quota_limit_reached: bool, + quota_verify_required: bool, + quota_cooldown_until: Option>, + cloudflare_cooldown_until: Option>, + window_request_count: u64, + window_input_tokens: u64, + window_output_tokens: u64, + window_cached_tokens: u64, + window_image_input_tokens: u64, + window_image_output_tokens: u64, + window_image_request_count: u64, + window_image_request_failed_count: u64, + window_started_at: Option>, + window_reset_at: Option>, + limit_window_seconds: Option, +} + +impl From<&Account> for RuntimeAccountStateSnapshot { + fn from(account: &Account) -> Self { + Self { + status: account.status, + quota_limit_reached: account.quota_limit_reached, + quota_verify_required: account.quota_verify_required, + quota_cooldown_until: account.quota_cooldown_until, + cloudflare_cooldown_until: account.cloudflare_cooldown_until, + window_request_count: account.window_request_count, + window_input_tokens: account.window_input_tokens, + window_output_tokens: account.window_output_tokens, + window_cached_tokens: account.window_cached_tokens, + window_image_input_tokens: account.window_image_input_tokens, + window_image_output_tokens: account.window_image_output_tokens, + window_image_request_count: account.window_image_request_count, + window_image_request_failed_count: account.window_image_request_failed_count, + window_started_at: account.window_started_at, + window_reset_at: account.window_reset_at, + limit_window_seconds: account.limit_window_seconds, + } + } +} + +impl RuntimeAccountStateSnapshot { + fn usage_window_changed(&self, other: &Self) -> bool { + self.window_request_count != other.window_request_count + || self.window_input_tokens != other.window_input_tokens + || self.window_output_tokens != other.window_output_tokens + || self.window_cached_tokens != other.window_cached_tokens + || self.window_image_input_tokens != other.window_image_input_tokens + || self.window_image_output_tokens != other.window_image_output_tokens + || self.window_image_request_count != other.window_image_request_count + || self.window_image_request_failed_count != other.window_image_request_failed_count + || self.window_started_at != other.window_started_at + || self.window_reset_at != other.window_reset_at + || self.limit_window_seconds != other.limit_window_seconds + } +} + +#[derive(Debug, Clone)] +pub(super) struct RefreshedAccountState { + pub(super) account: Account, + pub(super) sync_usage_window: bool, +} + +/// 运行时窗口用量增量。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct AccountWindowUsageDelta { + /// 输入 token 增量。 + pub input_tokens: u64, + /// 输出 token 增量。 + pub output_tokens: u64, + /// 缓存 token 增量。 + pub cached_tokens: u64, + /// 图片输入 token 增量。 + pub image_input_tokens: u64, + /// 图片输出 token 增量。 + pub image_output_tokens: u64, + /// 是否记录一次成功图片请求。 + pub image_request_succeeded: bool, + /// 是否记录一次失败图片请求。 + pub image_request_failed: bool, +} + +/// 纯内存账号池：负责账号存储、在途槽位与运行时状态维护。 +/// +/// 账号**选择**已剥离到 [`crate::accounts::scheduler`]：候选过滤走 +/// [`candidates::filter`]，策略选择走 [`AccountScheduler::select`]。账号池只提供选择 +/// 所需的只读视图（候选切片 + `slot_count`），并持有调度器实例。 +#[derive(Debug)] +pub struct AccountPool { + accounts: IndexMap, + slots: BTreeMap>, + pub(super) options: AccountPoolOptions, + scheduler: AccountScheduler, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct AccountSlot { + created_at: DateTime, + model: Option, +} + +/// 账号释放结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ReleasedAccountUsage { + /// 被释放槽位对应的模型。槽位已被过期清理时为空。 + pub model: Option, +} + +impl Default for AccountPool { + fn default() -> Self { + Self::with_options(AccountPoolOptions::default()) + } +} + +impl AccountPool { + /// 使用指定配置创建账号池。 + pub fn with_options(options: AccountPoolOptions) -> Self { + Self { + accounts: IndexMap::new(), + slots: BTreeMap::new(), + options, + scheduler: AccountScheduler::new(ScoreWeights::default()), + } + } + + /// 插入或替换账号。 + pub fn insert(&mut self, account: Account) { + self.accounts.insert(account.id.clone(), account); + } + + /// 获取账号池中的账号快照。 + pub fn get(&self, account_id: &str) -> Option { + self.accounts.get(account_id).cloned() + } + + /// 移除账号及其在途槽位与运行时反馈。 + pub fn remove(&mut self, account_id: &str) -> bool { + let removed = self.accounts.shift_remove(account_id).is_some(); + self.slots.remove(account_id); + self.scheduler.forget_feedback(account_id); + removed + } + + /// 清空账号池。 + pub fn clear(&mut self) { + self.accounts.clear(); + self.slots.clear(); + self.scheduler.reset_cursor(); + self.scheduler.clear_feedback(); + } + + /// 替换账号池运行参数。 + pub fn set_options(&mut self, options: AccountPoolOptions) { + if self.options.rotation_strategy != options.rotation_strategy { + self.scheduler.reset_cursor(); + } + self.options = options; + } + + /// 回灌一次请求结果到调度器的运行时反馈（供 Smart 打分）。 + pub fn report_feedback(&self, account_id: &str, success: bool, first_token_ms: Option) { + self.scheduler + .report_feedback(account_id, success, first_token_ms); + } + + /// 更新账号状态。 + pub fn set_status(&mut self, account_id: &str, status: AccountStatus) -> bool { + let Some(account) = self.accounts.get_mut(account_id) else { + return false; + }; + let status = status_after_quota_limit(status, account.quota_limit_reached); + account.status = status; + if status != AccountStatus::Active { + self.slots.remove(account_id); + } + true + } + + /// 标记账号配额限流直到指定时间。 + pub fn mark_quota_limited_until( + &mut self, + account_id: &str, + cooldown_until: DateTime, + ) -> bool { + let Some(account) = self.accounts.get_mut(account_id) else { + return false; + }; + let final_cooldown_until = account + .quota_cooldown_until + .filter(|existing| *existing > cooldown_until) + .unwrap_or(cooldown_until); + account.quota_limit_reached = true; + account.quota_cooldown_until = Some(final_cooldown_until); + account.status = status_after_quota_limit(account.status, true); + account.window_reset_at = account + .window_reset_at + .filter(|existing| *existing > final_cooldown_until) + .or(Some(final_cooldown_until)); + if let Ok(seconds) = (final_cooldown_until - Utc::now()) + .to_std() + .map(|duration| duration.as_secs()) + { + account.limit_window_seconds.get_or_insert(seconds); + } + self.slots.remove(account_id); + true + } + + /// 应用已经验证过的配额状态。 + pub fn apply_quota_state( + &mut self, + account_id: &str, + limit_reached: bool, + cooldown_until: Option>, + ) -> bool { + let Some(account) = self.accounts.get_mut(account_id) else { + return false; + }; + account.quota_verify_required = false; + account.quota_limit_reached = limit_reached; + account.quota_cooldown_until = limit_reached.then_some(cooldown_until).flatten(); + match (account.status, limit_reached) { + (AccountStatus::Active, true) => account.status = AccountStatus::QuotaExhausted, + (AccountStatus::QuotaExhausted, false) => account.status = AccountStatus::Active, + _ => {} + } + if let Some(cooldown_until) = account.quota_cooldown_until { + account.window_reset_at = account + .window_reset_at + .filter(|existing| *existing > cooldown_until) + .or(Some(cooldown_until)); + } + if limit_reached { + self.slots.remove(account_id); + } + true + } + + /// 同步上游返回的限流窗口信息。 + pub fn sync_rate_limit_window( + &mut self, + account_id: &str, + new_reset_at: DateTime, + limit_window_seconds: Option, + ) -> bool { + self.sync_rate_limit_window_at(account_id, new_reset_at, limit_window_seconds, Utc::now()) + } + + /// 记录账号运行时窗口 token 用量。 + pub fn record_window_token_usage( + &mut self, + account_id: &str, + usage: AccountWindowUsageDelta, + ) -> bool { + let Some(account) = self.accounts.get_mut(account_id) else { + return false; + }; + account.image_input_tokens = account + .image_input_tokens + .saturating_add(usage.image_input_tokens); + account.image_output_tokens = account + .image_output_tokens + .saturating_add(usage.image_output_tokens); + if usage.image_request_succeeded { + account.image_request_count = account.image_request_count.saturating_add(1); + } + if usage.image_request_failed { + account.image_request_failed_count = + account.image_request_failed_count.saturating_add(1); + } + account.window_input_tokens = account + .window_input_tokens + .saturating_add(usage.input_tokens); + account.window_output_tokens = account + .window_output_tokens + .saturating_add(usage.output_tokens); + account.window_cached_tokens = account + .window_cached_tokens + .saturating_add(usage.cached_tokens); + account.window_image_input_tokens = account + .window_image_input_tokens + .saturating_add(usage.image_input_tokens); + account.window_image_output_tokens = account + .window_image_output_tokens + .saturating_add(usage.image_output_tokens); + if usage.image_request_succeeded { + account.window_image_request_count = + account.window_image_request_count.saturating_add(1); + } + if usage.image_request_failed { + account.window_image_request_failed_count = + account.window_image_request_failed_count.saturating_add(1); + } + true + } + + /// 标记账号处于 Cloudflare 冷却期。 + pub fn set_cloudflare_cooldown_until( + &mut self, + account_id: &str, + cooldown_until: DateTime, + ) -> bool { + let Some(account) = self.accounts.get_mut(account_id) else { + return false; + }; + account.cloudflare_cooldown_until = Some(cooldown_until); + self.slots.remove(account_id); + true + } + + /// 使用完整调度请求获取账号。 + pub fn acquire_with(&mut self, request: &AccountAcquireRequest) -> Option { + self.acquire_with_status_refresh(request).acquired + } + + pub(super) fn acquire_with_status_refresh( + &mut self, + request: &AccountAcquireRequest, + ) -> AccountAcquireWithStatusRefresh { + self.cleanup_stale_slots(request.now); + let refreshed_accounts = self.refresh_account_statuses(request.now); + let candidates = self.filter_candidates(request); + // 亲和优先：required 硬锁定该账号；preferred 若在候选内则直接选中； + // 二者都不命中时，回退到调度器按当前策略选择。 + let selected = if let Some(required_account_id) = &request.required_account_id { + candidates.first().cloned().map(|account| { + ( + account, + self.previous_slot_at(required_account_id), + "required", + ) + }) + } else if let Some(preferred_account_id) = &request.preferred_account_id { + candidates + .iter() + .find(|account| account.id == *preferred_account_id) + .cloned() + .map(|account| { + ( + account, + self.previous_slot_at(preferred_account_id), + "preferred", + ) + }) + } else { + None + } + .or_else(|| { + let account = self.select_candidate(&candidates, request.now)?; + let previous = self.previous_slot_at(&account.id); + Some((account, previous, self.options.rotation_strategy.as_str())) + }); + let Some((selected, previous_slot_at, selection_source)) = selected else { + return AccountAcquireWithStatusRefresh { + acquired: None, + refreshed_accounts, + }; + }; + tracing::info!( + model = %request.model, + rotation_strategy = self.options.rotation_strategy.as_str(), + selection_source, + account_id = %selected.id, + candidate_count = candidates.len(), + slot_count = self.slot_count(&selected.id), + request_count = selected.request_count, + window_request_count = selected.window_request_count, + window_reset_at = ?selected.window_reset_at, + last_used_at = selected.last_used_at.as_deref().unwrap_or_default(), + quota_limit_reached = selected.quota_limit_reached, + quota_cooldown_until = ?selected.quota_cooldown_until, + previous_slot_at = ?previous_slot_at, + "account selected for upstream request" + ); + + let selected_id = selected.id.clone(); + self.push_slot(&selected_id, request.now, Some(&request.model)); + AccountAcquireWithStatusRefresh { + acquired: Some(AcquiredAccount { + account: selected, + previous_slot_at, + }), + refreshed_accounts, + } + } + + fn refresh_account_statuses(&mut self, now: DateTime) -> Vec { + let mut refreshed_accounts = Vec::new(); + let mut expired_account_ids = Vec::new(); + for (account_id, account) in &mut self.accounts { + let before = RuntimeAccountStateSnapshot::from(&*account); + if account.status == AccountStatus::Active && access_token_expired(account, now) { + account.status = AccountStatus::Expired; + expired_account_ids.push(account_id.clone()); + } else { + refresh_quota_window(account, now); + refresh_cloudflare_cooldown(account, now); + } + let after = RuntimeAccountStateSnapshot::from(&*account); + if before != after { + let sync_usage_window = before.usage_window_changed(&after); + refreshed_accounts.push(RefreshedAccountState { + account: account.clone(), + sync_usage_window, + }); + } + } + for account_id in &expired_account_ids { + self.slots.remove(account_id); + } + refreshed_accounts + } + + /// 释放指定账号的一个在途槽位。 + pub fn release(&mut self, account_id: &str) -> Option { + self.release_slot(account_id, true) + } + + /// 释放指定账号的一个在途槽位，不累计请求用量。 + pub fn release_without_request_usage( + &mut self, + account_id: &str, + ) -> Option { + self.release_slot(account_id, false) + } + + fn release_slot( + &mut self, + account_id: &str, + record_request_usage: bool, + ) -> Option { + self.accounts.get(account_id)?; + + let mut remove_slots = false; + let slot = self.slots.get_mut(account_id).and_then(|slots| { + let slot = slots.pop_front(); + remove_slots = slots.is_empty(); + slot + }); + if remove_slots { + self.slots.remove(account_id); + } + if record_request_usage { + self.mark_request_usage(account_id, Utc::now()); + } + Some(ReleasedAccountUsage { + model: slot.and_then(|slot| slot.model), + }) + } + + /// 计算账号池容量摘要。 + pub fn capacity_summary(&mut self, now: DateTime) -> AccountCapacitySummary { + self.capacity_summary_with_status_refresh(now).summary + } + + pub(super) fn capacity_summary_with_status_refresh( + &mut self, + now: DateTime, + ) -> AccountCapacitySummaryWithStatusRefresh { + self.cleanup_stale_slots(now); + let refreshed_accounts = self.refresh_account_statuses(now); + let is_capacity_account = |account: &Account| { + account.status == AccountStatus::Active + && candidates::quota_available_at(account, now, self.options.skip_quota_limited) + }; + let active_accounts = self + .accounts + .values() + .filter(|account| is_capacity_account(account)) + .count(); + let total_slots = active_accounts * self.options.max_concurrent_per_account; + let used_slots = self + .slots + .iter() + .filter(|(account_id, _)| { + self.accounts + .get(*account_id) + .is_some_and(&is_capacity_account) + }) + .map(|(_, slots)| slots.len().min(self.options.max_concurrent_per_account)) + .sum(); + + AccountCapacitySummaryWithStatusRefresh { + summary: AccountCapacitySummary { + max_concurrent_per_account: self.options.max_concurrent_per_account, + total_slots, + used_slots, + available_slots: total_slots.saturating_sub(used_slots), + }, + refreshed_accounts, + } + } + + /// 按订阅计划各选一个可用账号，用于刷新模型列表。 + pub fn distinct_plan_accounts(&mut self, now: DateTime) -> Vec { + self.distinct_plan_accounts_with_status_refresh(now) + .accounts + } + + pub(super) fn distinct_plan_accounts_with_status_refresh( + &mut self, + now: DateTime, + ) -> DistinctPlanAccountsWithStatusRefresh { + self.cleanup_stale_slots(now); + let refreshed_accounts = self.refresh_account_statuses(now); + let mut by_plan = IndexMap::>::new(); + + let max_concurrent = self.options.max_concurrent_per_account; + let skip_quota_limited = self.options.skip_quota_limited; + { + let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len); + for account in self.accounts.values() { + if !candidates::is_model_refresh_available( + account, + max_concurrent, + skip_quota_limited, + &slot_count, + now, + ) { + continue; + } + let Some(plan_type) = account.plan_type.as_ref() else { + continue; + }; + by_plan + .entry(plan_type.clone()) + .or_default() + .push(account.clone()); + } + } + + let strategy = self.options.rotation_strategy; + let mut accounts = Vec::new(); + for (plan_type, group) in by_plan { + let Some(account) = self.select_candidate(&group, now) else { + continue; + }; + tracing::info!( + plan_type, + rotation_strategy = strategy.as_str(), + account_id = %account.id, + candidate_count = group.len(), + slot_count = self.slot_count(&account.id), + request_count = account.request_count, + window_request_count = account.window_request_count, + window_reset_at = ?account.window_reset_at, + last_used_at = account.last_used_at.as_deref().unwrap_or_default(), + quota_limit_reached = account.quota_limit_reached, + quota_cooldown_until = ?account.quota_cooldown_until, + "account selected for model refresh" + ); + self.push_slot(&account.id, now, None); + accounts.push(DistinctPlanAccount { plan_type, account }); + } + + DistinctPlanAccountsWithStatusRefresh { + accounts, + refreshed_accounts, + } + } + + fn slot_count(&self, account_id: &str) -> usize { + self.slots.get(account_id).map_or(0, VecDeque::len) + } + + fn previous_slot_at(&self, account_id: &str) -> Option> { + self.slots + .get(account_id) + .and_then(|slots| slots.back()) + .map(|slot| slot.created_at) + } + + fn push_slot(&mut self, account_id: &str, now: DateTime, model: Option<&str>) { + self.slots + .entry(account_id.to_string()) + .or_default() + .push_back(AccountSlot { + created_at: now, + model: model.map(str::to_string), + }); + } + + fn mark_request_usage(&mut self, account_id: &str, now: DateTime) -> Option { + let account = self.accounts.get_mut(account_id)?; + account.last_used_at = Some(now.to_rfc3339()); + account.request_count = account.request_count.saturating_add(1); + account.window_request_count = account.window_request_count.saturating_add(1); + if account.window_started_at.is_none() { + account.window_started_at = Some(now); + if let Some(seconds) = account.limit_window_seconds { + let seconds = i64::try_from(seconds).unwrap_or(i64::MAX); + account.window_reset_at = Some(now + Duration::seconds(seconds)); + } + } + Some(account.clone()) + } + + fn sync_rate_limit_window_at( + &mut self, + account_id: &str, + new_reset_at: DateTime, + limit_window_seconds: Option, + now: DateTime, + ) -> bool { + let Some(account) = self.accounts.get_mut(account_id) else { + return false; + }; + if should_reset_usage_window( + account.window_reset_at, + account.limit_window_seconds, + new_reset_at, + limit_window_seconds, + ) { + reset_window_counters(account); + account.window_started_at = Some(now); + } + account.window_reset_at = Some(new_reset_at); + if let Some(limit_window_seconds) = limit_window_seconds { + account.limit_window_seconds = Some(limit_window_seconds); + } + true + } + + fn cleanup_stale_slots(&mut self, now: DateTime) { + let ttl = self.options.stale_slot_ttl; + self.slots.retain(|_, slots| { + slots.retain(|slot| now.signed_duration_since(slot.created_at) <= ttl); + !slots.is_empty() + }); + } + + /// 过滤出本次请求的可用候选集合（委托 [`candidates::filter`]）。 + fn filter_candidates(&self, request: &AccountAcquireRequest) -> Vec { + let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len); + candidates::filter( + self.accounts.values(), + &CandidateFilter { + max_concurrent_per_account: self.options.max_concurrent_per_account, + skip_quota_limited: self.options.skip_quota_limited, + tier_priority: &self.options.tier_priority, + model_plan_allowlist: &self.options.model_plan_allowlist, + fetched_model_plan_types: &self.options.fetched_model_plan_types, + slot_count: &slot_count, + }, + &CandidateRequest { + model: &request.model, + exclude_account_ids: &request.exclude_account_ids, + required_account_id: request.required_account_id.as_deref(), + now: request.now, + }, + ) + } + + /// 按当前策略从候选集合中选择一个账号（委托 [`AccountScheduler::select`]）。 + fn select_candidate(&self, candidates: &[Account], now: DateTime) -> Option { + let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len); + self.scheduler + .select(self.options.rotation_strategy, candidates, &slot_count, now) + } +} + +fn access_token_expired(account: &Account, now: DateTime) -> bool { + if account + .access_token_expires_at + .is_some_and(|expires_at| now >= expires_at) + { + return true; + } + + match jwt_expiry(&account.access_token, now) { + JwtExpiry::Expired => true, + JwtExpiry::Valid => false, + JwtExpiry::MissingOrInvalid => account.access_token_expires_at.is_none(), + } +} diff --git a/backend/src/accounts/scheduler.rs b/backend/src/accounts/scheduler.rs index 5e914f96e..32ac974a5 100644 --- a/backend/src/accounts/scheduler.rs +++ b/backend/src/accounts/scheduler.rs @@ -15,16 +15,56 @@ pub mod candidates; pub mod feedback; pub mod strategy; -use std::sync::Mutex; +use std::{cmp::Ordering, sync::Mutex}; use chrono::{DateTime, Utc}; use crate::accounts::account::Account; +use crate::infra::time::parse_rfc3339_utc; pub use candidates::{CandidateFilter, CandidateRequest}; pub use feedback::{FeedbackSample, FeedbackStats}; pub use strategy::smart::rank_candidates; -pub use strategy::{RotationStrategy, ScoreBreakdown, ScoreWeights}; +pub use strategy::{ScoreBreakdown, ScoreWeights}; + +/// 账号轮转策略。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RotationStrategy { + /// 加权线性打分调度。 + Smart, + /// 优先消耗即将重置的额度窗口。 + QuotaResetPriority, + /// 按候选账号顺序循环选择。 + RoundRobin, + /// 优先选择最近使用过的账号。 + Sticky, +} + +impl RotationStrategy { + /// 返回策略的持久化与日志字符串。 + pub fn as_str(self) -> &'static str { + match self { + Self::Smart => "smart", + Self::QuotaResetPriority => "quota_reset_priority", + Self::RoundRobin => "round_robin", + Self::Sticky => "sticky", + } + } +} + +/// 策略选择的运行时输入。 +pub struct SelectionInput<'a> { + /// 已过滤的候选账号集合。 + pub candidates: &'a [Account], + /// 按账号 ID 读取当前在途槽位数。 + pub slot_count: &'a dyn Fn(&str) -> usize, + /// 账号 EWMA 反馈存储。 + pub feedback: &'a FeedbackStats, + /// Smart 策略打分权重。 + pub weights: &'a ScoreWeights, + /// 当前时刻。 + pub now: DateTime, +} /// 账号调度器：账号选择的对外门面。 /// @@ -78,7 +118,7 @@ impl AccountScheduler { slot_count: &dyn Fn(&str) -> usize, now: DateTime, ) -> Option { - let input = strategy::SelectionInput { + let input = SelectionInput { candidates, slot_count, feedback: &self.feedback, @@ -86,7 +126,7 @@ impl AccountScheduler { now, }; let mut cursor = self.lock_cursor(); - strategy::select(strategy, &input, &mut cursor) + select_with_strategy(strategy, &input, &mut cursor) } /// 回灌一次请求结果到运行时 EWMA 反馈（错误率 / TTFT），供 Smart 打分。 @@ -110,3 +150,59 @@ impl AccountScheduler { .unwrap_or_else(std::sync::PoisonError::into_inner) } } + +fn select_with_strategy( + strategy: RotationStrategy, + input: &SelectionInput<'_>, + cursor: &mut usize, +) -> Option { + match strategy { + RotationStrategy::Smart => strategy::smart::select(input, cursor), + RotationStrategy::QuotaResetPriority => strategy::quota_reset::select(input, cursor), + RotationStrategy::RoundRobin => strategy::round_robin::select(input, cursor), + RotationStrategy::Sticky => strategy::sticky::select(input, cursor), + } +} + +pub(crate) fn select_by( + candidates: &[Account], + cursor: &mut usize, + compare: impl Fn(&Account, &Account) -> Ordering, +) -> Option { + let mut sorted = candidates.iter().collect::>(); + sorted.sort_by(|a, b| compare(a, b)); + let best = *sorted.first()?; + let tied_count = sorted + .iter() + .take_while(|account| compare(best, account) == Ordering::Equal) + .count(); + let index = *cursor % tied_count; + *cursor = cursor.wrapping_add(1); + Some((*sorted[index]).clone()) +} + +pub(crate) fn account_window_token_count(account: &Account) -> u64 { + account + .window_input_tokens + .saturating_add(account.window_output_tokens) + .saturating_add(account.window_image_input_tokens) + .saturating_add(account.window_image_output_tokens) +} + +pub(crate) fn compare_window_reset(a: Option>, b: Option>) -> Ordering { + match (a, b) { + (Some(a), Some(b)) => a.cmp(&b), + (Some(_) | None, None) | (None, Some(_)) => Ordering::Equal, + } +} + +pub(crate) fn compare_last_used(a: Option<&str>, b: Option<&str>) -> Ordering { + last_used_millis(a).cmp(&last_used_millis(b)) +} + +fn last_used_millis(value: Option<&str>) -> i64 { + value + .and_then(|value| parse_rfc3339_utc(value).ok()) + .map(|datetime| datetime.timestamp_millis()) + .unwrap_or(0) +} diff --git a/backend/src/accounts/scheduler/strategy/mod.rs b/backend/src/accounts/scheduler/strategy/mod.rs index ba837de05..ccbc6ea78 100644 --- a/backend/src/accounts/scheduler/strategy/mod.rs +++ b/backend/src/accounts/scheduler/strategy/mod.rs @@ -4,10 +4,5 @@ pub mod quota_reset; pub mod round_robin; pub mod smart; pub mod sticky; -mod types; pub use smart::{ScoreBreakdown, ScoreWeights}; -pub(crate) use types::{ - account_window_token_count, compare_last_used, compare_window_reset, select_by, -}; -pub use types::{select, RotationStrategy, SelectionInput}; diff --git a/backend/src/accounts/scheduler/strategy/quota_reset.rs b/backend/src/accounts/scheduler/strategy/quota_reset.rs index b8d143e63..8ec857141 100644 --- a/backend/src/accounts/scheduler/strategy/quota_reset.rs +++ b/backend/src/accounts/scheduler/strategy/quota_reset.rs @@ -11,8 +11,9 @@ use std::cmp::Ordering; use crate::accounts::account::Account; - -use super::{compare_last_used, compare_window_reset, select_by, SelectionInput}; +use crate::accounts::scheduler::{ + compare_last_used, compare_window_reset, select_by, SelectionInput, +}; /// QuotaResetPriority 策略选择。 pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option { diff --git a/backend/src/accounts/scheduler/strategy/round_robin.rs b/backend/src/accounts/scheduler/strategy/round_robin.rs index a3f89e41b..d93b0d240 100644 --- a/backend/src/accounts/scheduler/strategy/round_robin.rs +++ b/backend/src/accounts/scheduler/strategy/round_robin.rs @@ -4,8 +4,7 @@ //! 保证请求在候选账号间均匀轮转。 use crate::accounts::account::Account; - -use super::SelectionInput; +use crate::accounts::scheduler::SelectionInput; /// 按候选顺序做纯轮转选择。 pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option { diff --git a/backend/src/accounts/scheduler/strategy/smart.rs b/backend/src/accounts/scheduler/strategy/smart.rs index 483afc6a6..695da1718 100644 --- a/backend/src/accounts/scheduler/strategy/smart.rs +++ b/backend/src/accounts/scheduler/strategy/smart.rs @@ -11,8 +11,7 @@ use crate::accounts::account::Account; use crate::accounts::scheduler::feedback::{FeedbackSample, FeedbackStats}; - -use super::{account_window_token_count, SelectionInput}; +use crate::accounts::scheduler::{account_window_token_count, SelectionInput}; /// 配额封顶账号的惩罚分。远大于各因子加权和的理论上限（各因子 ∈ [0,1]，权重和通常为 /// 个位数），保证未封顶账号总是优先于封顶账号。 diff --git a/backend/src/accounts/scheduler/strategy/sticky.rs b/backend/src/accounts/scheduler/strategy/sticky.rs index 38b118c0d..868a20ed1 100644 --- a/backend/src/accounts/scheduler/strategy/sticky.rs +++ b/backend/src/accounts/scheduler/strategy/sticky.rs @@ -6,8 +6,7 @@ use std::cmp::Ordering; use crate::accounts::account::Account; - -use super::{compare_last_used, SelectionInput}; +use crate::accounts::scheduler::{compare_last_used, SelectionInput}; /// 选中最近使用过的账号（`last_used_at` 最大者）。 /// diff --git a/backend/src/accounts/scheduler/strategy/types.rs b/backend/src/accounts/scheduler/strategy/types.rs deleted file mode 100644 index 3b17eb940..000000000 --- a/backend/src/accounts/scheduler/strategy/types.rs +++ /dev/null @@ -1,137 +0,0 @@ -//! 账号调度策略层。 -//! -//! 所有账号选择策略集中于此，每个策略一个文件，互不依赖： -//! -//! - [`smart`] —— 加权线性打分 + EWMA 反馈（“智能调度”）。 -//! - [`quota_reset`] —— 优先消耗即将重置的额度窗口。 -//! - [`round_robin`] —— 按候选顺序循环。 -//! - [`sticky`] —— 优先最近使用过的账号。 -//! -//! 本模块只承担三件事：定义 [`RotationStrategy`] 枚举、把选择请求分派到对应策略、 -//! 提供各策略共享的 tie-break 比较原语（[`select_by`] / [`compare_window_reset`] / -//! [`compare_last_used`]）。策略本身不触碰账号池存储：在途槽位数经 `slot_count` -//! 闭包按需读取，EWMA 反馈与打分权重经 [`SelectionInput`] 传入。 - -use std::cmp::Ordering; - -use chrono::{DateTime, Utc}; - -use crate::accounts::account::Account; -use crate::accounts::scheduler::feedback::FeedbackStats; -use crate::infra::time::parse_rfc3339_utc; - -use super::smart::ScoreWeights; - -/// 账号轮转策略。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum RotationStrategy { - /// 加权线性打分调度：多因子（负载/窗口/错误率/TTFT/重置）归一化后加权求和， - /// 配合 EWMA 反馈。 - Smart, - /// 优先消耗即将重置的额度窗口。 - QuotaResetPriority, - /// 按候选账号顺序循环选择。 - RoundRobin, - /// 优先选择最近使用过的账号。 - Sticky, -} - -impl RotationStrategy { - /// 返回策略的持久化/日志字符串。 - pub fn as_str(self) -> &'static str { - match self { - Self::Smart => "smart", - Self::QuotaResetPriority => "quota_reset_priority", - Self::RoundRobin => "round_robin", - Self::Sticky => "sticky", - } - } -} - -/// 策略选择的运行时输入。 -/// -/// `slot_count` 按账号 ID 读取当前在途槽位数；`feedback` / `weights` 仅 [`smart`] -/// 使用，其余策略忽略；`now` 用于窗口重置相关的打分。 -pub struct SelectionInput<'a> { - /// 已过滤的候选账号集合。 - pub candidates: &'a [Account], - /// 按账号 ID 读取当前在途槽位数。 - pub slot_count: &'a dyn Fn(&str) -> usize, - /// 账号 EWMA 反馈存储（Smart 打分用）。 - pub feedback: &'a FeedbackStats, - /// 打分权重（Smart 打分用）。 - pub weights: &'a ScoreWeights, - /// 当前时刻（窗口重置因子用）。 - pub now: DateTime, -} - -/// 按指定策略从候选集合中选择一个账号。 -/// -/// `cursor` 是账号池持有的轮转游标，用于在并列最优账号之间做公平轮转；不同策略对它 -/// 的读写方式不同（详见各策略模块），本函数只负责透传可变引用。 -pub fn select( - strategy: RotationStrategy, - input: &SelectionInput<'_>, - cursor: &mut usize, -) -> Option { - match strategy { - RotationStrategy::Smart => super::smart::select(input, cursor), - RotationStrategy::QuotaResetPriority => super::quota_reset::select(input, cursor), - RotationStrategy::RoundRobin => super::round_robin::select(input, cursor), - RotationStrategy::Sticky => super::sticky::select(input, cursor), - } -} - -// ==================================================================== -// 共享 tie-break 原语 -// ==================================================================== - -/// 按比较函数升序排序后，在并列最优（`Ordering::Equal`）的账号之间用轮转游标破并列。 -/// -/// 这是 [`smart`] 与 [`quota_reset`] 共享的选择骨架：先按各自的比较规则排序，再在得分 -/// 相同的账号之间轮转，保证并列账号被公平地依次选中。 -pub(crate) fn select_by( - candidates: &[Account], - cursor: &mut usize, - compare: impl Fn(&Account, &Account) -> Ordering, -) -> Option { - let mut sorted = candidates.iter().collect::>(); - sorted.sort_by(|a, b| compare(a, b)); - let best = *sorted.first()?; - let tied_count = sorted - .iter() - .take_while(|account| compare(best, account) == Ordering::Equal) - .count(); - let index = *cursor % tied_count; - *cursor = cursor.wrapping_add(1); - Some((*sorted[index]).clone()) -} - -/// 计算账号当前窗口的总 token 数（输入 + 输出 + 图片输入 + 图片输出）。 -pub(crate) fn account_window_token_count(account: &Account) -> u64 { - account - .window_input_tokens - .saturating_add(account.window_output_tokens) - .saturating_add(account.window_image_input_tokens) - .saturating_add(account.window_image_output_tokens) -} - -/// 比较两个窗口重置时间：越早重置越优先；缺失重置时间视为并列（不惩罚）。 -pub(crate) fn compare_window_reset(a: Option>, b: Option>) -> Ordering { - match (a, b) { - (Some(a), Some(b)) => a.cmp(&b), - (Some(_) | None, None) | (None, Some(_)) => Ordering::Equal, - } -} - -/// 比较两个账号的最近使用时间：越久未用越优先（升序，最小者排前）。 -pub(crate) fn compare_last_used(a: Option<&str>, b: Option<&str>) -> Ordering { - last_used_millis(a).cmp(&last_used_millis(b)) -} - -fn last_used_millis(value: Option<&str>) -> i64 { - value - .and_then(|value| parse_rfc3339_utc(value).ok()) - .map(|datetime| datetime.timestamp_millis()) - .unwrap_or(0) -} diff --git a/backend/src/accounts/store.rs b/backend/src/accounts/store.rs index 08ccd714d..022451a1e 100644 --- a/backend/src/accounts/store.rs +++ b/backend/src/accounts/store.rs @@ -282,6 +282,16 @@ pub struct PgAccountStore { mod write; impl PgAccountStore { + /// 构造存储。 + pub fn new(pool: PgPool) -> Self { + Self { pool } + } + + /// 返回底层连接池。 + pub fn pool(&self) -> &PgPool { + &self.pool + } + /// 设置账号标签。 pub async fn set_label( &self, diff --git a/backend/src/accounts/store/queries.rs b/backend/src/accounts/store/queries.rs index a00a5f324..4d50974fd 100644 --- a/backend/src/accounts/store/queries.rs +++ b/backend/src/accounts/store/queries.rs @@ -1,4 +1,6 @@ -//! PostgreSQL 账号仓储 SQL。 +//! PostgreSQL 账号只读查询与 SQL。 + +use super::*; // ============================================================================ // SQL 常量 @@ -387,3 +389,161 @@ set where id = $7"; pub(super) const DELETE_ACCOUNT_SQL: &str = "delete from accounts where id = $1"; + +impl PgAccountStore { + /// 读取单个账号。 + pub async fn get(&self, account_id: &str) -> PgAccountStoreResult> { + let row = sqlx::query(SELECT_STORED_ACCOUNT_BY_ID_SQL) + .bind(account_id) + .fetch_optional(&self.pool) + .await?; + row.as_ref().map(stored_account_from_row).transpose() + } + + /// 通过 ChatGPT 身份查找账号。 + pub async fn find_by_chatgpt_identity( + &self, + chatgpt_account_id: &str, + chatgpt_user_id: Option<&str>, + ) -> PgAccountStoreResult> { + let row = sqlx::query(SELECT_STORED_ACCOUNT_BY_CHATGPT_IDENTITY_SQL) + .bind(chatgpt_account_id) + .bind(chatgpt_user_id) + .bind(chatgpt_user_id) + .fetch_optional(&self.pool) + .await?; + row.as_ref().map(stored_account_from_row).transpose() + } + + /// 分页列出所有账号（含 token）。 + pub async fn list( + &self, + cursor: Option, + limit: u32, + ) -> PgAccountStoreResult> { + let limit = limit.clamp(1, 200); + if let Some(cursor) = cursor { + let (added_at, id) = + decode_cursor(&cursor).ok_or(PgAccountStoreError::InvalidCursor)?; + let added_at = + parse_rfc3339_utc(&added_at).map_err(|_| PgAccountStoreError::InvalidCursor)?; + let rows = sqlx::query(LIST_STORED_ACCOUNTS_AFTER_CURSOR_SQL) + .bind(added_at) + .bind(added_at) + .bind(&id) + .bind(i64::from(limit) + 1) + .fetch_all(&self.pool) + .await?; + Ok(to_page( + &rows, + limit, + stored_account_from_row, + ("added_at", "id"), + )) + } else { + let rows = sqlx::query(LIST_STORED_ACCOUNTS_SQL) + .bind(i64::from(limit) + 1) + .fetch_all(&self.pool) + .await?; + Ok(to_page( + &rows, + limit, + stored_account_from_row, + ("added_at", "id"), + )) + } + } + + /// 分页列出账号元数据（不含 token）。 + pub async fn list_metadata( + &self, + cursor: Option, + limit: u32, + ) -> PgAccountStoreResult> { + let limit = limit.clamp(1, 200); + if let Some(cursor) = cursor { + let (added_at, id) = + decode_cursor(&cursor).ok_or(PgAccountStoreError::InvalidCursor)?; + let added_at = + parse_rfc3339_utc(&added_at).map_err(|_| PgAccountStoreError::InvalidCursor)?; + let rows = sqlx::query(LIST_ACCOUNT_METADATA_AFTER_CURSOR_SQL) + .bind(added_at) + .bind(added_at) + .bind(&id) + .bind(i64::from(limit) + 1) + .fetch_all(&self.pool) + .await?; + Ok(to_page(&rows, limit, metadata_from_row, ("added_at", "id"))) + } else { + let rows = sqlx::query(LIST_ACCOUNT_METADATA_SQL) + .bind(i64::from(limit) + 1) + .fetch_all(&self.pool) + .await?; + Ok(to_page(&rows, limit, metadata_from_row, ("added_at", "id"))) + } + } + + /// 按页码列出账号元数据（不含 token）。 + pub async fn list_metadata_page( + &self, + page: u32, + page_size: u32, + search: Option<&str>, + ) -> PgAccountStoreResult> { + let page_size = page_size.clamp(1, 200); + let search = search.map(str::trim).filter(|value| !value.is_empty()); + let total = count_account_metadata(&self.pool, search).await?; + let offset = page_offset(page, page_size); + + let mut builder = QueryBuilder::::new(LIST_ACCOUNT_METADATA_SELECT_SQL); + push_account_metadata_search(&mut builder, search); + builder.push(" order by added_at desc, id desc limit "); + builder.push_bind(i64::from(page_size)); + builder.push(" offset "); + builder.push_bind(offset.min(i64::MAX as u64) as i64); + + let rows = builder.build().fetch_all(&self.pool).await?; + let items = rows + .iter() + .map(metadata_from_row) + .collect::>>()?; + + Ok(NumberedPage { + items, + total, + page: page.max(1), + page_size, + }) + } + + /// 读取单个账号元数据（不含 token）。 + pub async fn get_metadata( + &self, + account_id: &str, + ) -> PgAccountStoreResult> { + let row = sqlx::query(SELECT_ACCOUNT_METADATA_BY_ID_SQL) + .bind(account_id) + .fetch_optional(&self.pool) + .await?; + row.as_ref().map(metadata_from_row).transpose() + } + + /// 读取配额快照列表。 + pub async fn list_quota_snapshots(&self) -> PgAccountStoreResult> { + let rows = sqlx::query(LIST_QUOTA_SNAPSHOTS_SQL) + .fetch_all(&self.pool) + .await?; + rows.iter().map(quota_snapshot_from_row).collect() + } + + /// 读取单账号配额 JSON。 + pub async fn get_quota_json(&self, account_id: &str) -> PgAccountStoreResult> { + let row = sqlx::query("select quota_json from accounts where id = $1") + .bind(account_id) + .fetch_optional(&self.pool) + .await?; + Ok(row + .and_then(|row| row.get::>, _>("quota_json")) + .map(|value| value.0.to_string())) + } +} diff --git a/backend/src/accounts/store/write.rs b/backend/src/accounts/store/write.rs index 0a4eced55..09bf4d7e9 100644 --- a/backend/src/accounts/store/write.rs +++ b/backend/src/accounts/store/write.rs @@ -1,16 +1,6 @@ use super::*; impl PgAccountStore { - /// 构造存储。 - pub fn new(pool: PgPool) -> Self { - Self { pool } - } - - /// 返回底层连接池。 - pub fn pool(&self) -> &PgPool { - &self.pool - } - /// 插入新账号。 pub async fn insert(&self, account: NewAccount) -> PgAccountStoreResult<()> { let now = Utc::now(); @@ -36,143 +26,6 @@ impl PgAccountStore { Ok(()) } - /// 读取单个账号。 - pub async fn get(&self, account_id: &str) -> PgAccountStoreResult> { - let row = sqlx::query(SELECT_STORED_ACCOUNT_BY_ID_SQL) - .bind(account_id) - .fetch_optional(&self.pool) - .await?; - row.as_ref().map(stored_account_from_row).transpose() - } - - /// 通过 ChatGPT 身份查找账号。 - pub async fn find_by_chatgpt_identity( - &self, - chatgpt_account_id: &str, - chatgpt_user_id: Option<&str>, - ) -> PgAccountStoreResult> { - let row = sqlx::query(SELECT_STORED_ACCOUNT_BY_CHATGPT_IDENTITY_SQL) - .bind(chatgpt_account_id) - .bind(chatgpt_user_id) - .bind(chatgpt_user_id) - .fetch_optional(&self.pool) - .await?; - row.as_ref().map(stored_account_from_row).transpose() - } - - /// 分页列出所有账号（含 token）。 - pub async fn list( - &self, - cursor: Option, - limit: u32, - ) -> PgAccountStoreResult> { - let limit = limit.clamp(1, 200); - if let Some(cursor) = cursor { - let (added_at, id) = - decode_cursor(&cursor).ok_or(PgAccountStoreError::InvalidCursor)?; - let added_at = - parse_rfc3339_utc(&added_at).map_err(|_| PgAccountStoreError::InvalidCursor)?; - let rows = sqlx::query(LIST_STORED_ACCOUNTS_AFTER_CURSOR_SQL) - .bind(added_at) - .bind(added_at) - .bind(&id) - .bind(i64::from(limit) + 1) - .fetch_all(&self.pool) - .await?; - Ok(to_page( - &rows, - limit, - stored_account_from_row, - ("added_at", "id"), - )) - } else { - let rows = sqlx::query(LIST_STORED_ACCOUNTS_SQL) - .bind(i64::from(limit) + 1) - .fetch_all(&self.pool) - .await?; - Ok(to_page( - &rows, - limit, - stored_account_from_row, - ("added_at", "id"), - )) - } - } - - /// 分页列出账号元数据（不含 token）。 - pub async fn list_metadata( - &self, - cursor: Option, - limit: u32, - ) -> PgAccountStoreResult> { - let limit = limit.clamp(1, 200); - if let Some(cursor) = cursor { - let (added_at, id) = - decode_cursor(&cursor).ok_or(PgAccountStoreError::InvalidCursor)?; - let added_at = - parse_rfc3339_utc(&added_at).map_err(|_| PgAccountStoreError::InvalidCursor)?; - let rows = sqlx::query(LIST_ACCOUNT_METADATA_AFTER_CURSOR_SQL) - .bind(added_at) - .bind(added_at) - .bind(&id) - .bind(i64::from(limit) + 1) - .fetch_all(&self.pool) - .await?; - Ok(to_page(&rows, limit, metadata_from_row, ("added_at", "id"))) - } else { - let rows = sqlx::query(LIST_ACCOUNT_METADATA_SQL) - .bind(i64::from(limit) + 1) - .fetch_all(&self.pool) - .await?; - Ok(to_page(&rows, limit, metadata_from_row, ("added_at", "id"))) - } - } - - /// 按页码列出账号元数据（不含 token）。 - pub async fn list_metadata_page( - &self, - page: u32, - page_size: u32, - search: Option<&str>, - ) -> PgAccountStoreResult> { - let page_size = page_size.clamp(1, 200); - let search = search.map(str::trim).filter(|value| !value.is_empty()); - let total = count_account_metadata(&self.pool, search).await?; - let offset = page_offset(page, page_size); - - let mut builder = QueryBuilder::::new(LIST_ACCOUNT_METADATA_SELECT_SQL); - push_account_metadata_search(&mut builder, search); - builder.push(" order by added_at desc, id desc limit "); - builder.push_bind(i64::from(page_size)); - builder.push(" offset "); - builder.push_bind(offset.min(i64::MAX as u64) as i64); - - let rows = builder.build().fetch_all(&self.pool).await?; - let items = rows - .iter() - .map(metadata_from_row) - .collect::>>()?; - - Ok(NumberedPage { - items, - total, - page: page.max(1), - page_size, - }) - } - - /// 读取单个账号元数据（不含 token）。 - pub async fn get_metadata( - &self, - account_id: &str, - ) -> PgAccountStoreResult> { - let row = sqlx::query(SELECT_ACCOUNT_METADATA_BY_ID_SQL) - .bind(account_id) - .fetch_optional(&self.pool) - .await?; - row.as_ref().map(metadata_from_row).transpose() - } - /// 更新单账号元数据（不含 token）。 pub async fn update_metadata( &self, @@ -419,25 +272,6 @@ where id = $4", Ok(result.rows_affected() > 0) } - /// 读取配额快照列表。 - pub async fn list_quota_snapshots(&self) -> PgAccountStoreResult> { - let rows = sqlx::query(LIST_QUOTA_SNAPSHOTS_SQL) - .fetch_all(&self.pool) - .await?; - rows.iter().map(quota_snapshot_from_row).collect() - } - - /// 读取单账号配额 JSON。 - pub async fn get_quota_json(&self, account_id: &str) -> PgAccountStoreResult> { - let row = sqlx::query("select quota_json from accounts where id = $1") - .bind(account_id) - .fetch_optional(&self.pool) - .await?; - Ok(row - .and_then(|row| row.get::>, _>("quota_json")) - .map(|value| value.0.to_string())) - } - /// 更新配额 JSON。 pub async fn update_quota_json( &self, diff --git a/backend/src/api/admin/accounts_routes.rs b/backend/src/api/admin/accounts_routes.rs index 4485d4c71..321601455 100644 --- a/backend/src/api/admin/accounts_routes.rs +++ b/backend/src/api/admin/accounts_routes.rs @@ -599,8 +599,18 @@ pub(crate) async fn create_account( } } -mod actions; -mod stats; +mod export_routes; +mod import_routes; +mod lifecycle_routes; +mod oauth_routes; +mod probe_routes; +mod query; +mod quota_routes; -pub(crate) use actions::*; -use stats::*; +pub(crate) use export_routes::*; +pub(crate) use import_routes::*; +pub(crate) use lifecycle_routes::*; +pub(crate) use oauth_routes::*; +pub(crate) use probe_routes::*; +use query::*; +pub(crate) use quota_routes::*; diff --git a/backend/src/api/admin/accounts_routes/actions.rs b/backend/src/api/admin/accounts_routes/actions.rs deleted file mode 100644 index 813c7426d..000000000 --- a/backend/src/api/admin/accounts_routes/actions.rs +++ /dev/null @@ -1,321 +0,0 @@ -use super::*; - -/// `GET /api/admin/accounts/export` -pub(crate) async fn export_accounts( - State(state): State, - _auth: AdminAuth, - Query(query): Query, -) -> Result { - let ids = account_export_ids(query.ids.as_deref()); - - if query.confirm.as_deref() != Some(ACCOUNT_EXPORT_CONFIRMATION) { - return Err(AdminError::bad_request( - "account export requires confirm=export_sensitive_accounts", - )); - } - - match state.services.admin_accounts.export(ids).await { - Ok(result) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(result), - )), - Err(error) => Err(account_error(&error)), - } -} - -/// `POST /api/admin/accounts/refresh` -pub(crate) async fn refresh_account( - State(state): State, - _auth: AdminAuth, - Json(payload): Json, -) -> Result { - let account_id = payload.id; - match state - .services - .admin_accounts - .refresh_account(&account_id) - .await - { - Ok(result) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountRefreshData::from(result)), - )), - Err(error) => Err(account_error(&error)), - } -} - -/// `POST /api/admin/accounts/health-check` -pub(crate) async fn health_check_accounts( - State(state): State, - _auth: AdminAuth, - Json(payload): Json, -) -> Result { - let stagger_ms = payload.stagger_ms.unwrap_or(3_000); - if !(500..=30_000).contains(&stagger_ms) && stagger_ms != 0 { - return Err(AdminError::bad_request( - "stagger_ms must be between 500 and 30000", - )); - } - if let Some(concurrency) = payload.concurrency { - if !(1..=10).contains(&concurrency) { - return Err(AdminError::bad_request( - "concurrency must be between 1 and 10", - )); - } - } - match state - .services - .admin_accounts - .health_check_accounts(payload.ids, stagger_ms, payload.concurrency.unwrap_or(2)) - .await - { - Ok(result) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountHealthCheckData::from(result)), - )), - Err(error) => Err(account_error(&error)), - } -} - -/// `GET /api/admin/accounts/quota` -pub(crate) async fn account_quota( - State(state): State, - _auth: AdminAuth, - Query(query): Query, -) -> Result { - let account_id = query.id; - match state - .services - .admin_accounts - .account_quota(&account_id) - .await - { - Ok(data) => { - let quota = data.get("quota").cloned().unwrap_or(Value::Null); - let raw = data.get("raw").cloned().unwrap_or(Value::Null); - let quota_json = quota.to_string(); - let mut quota_data = quota_data("a_json, Some(Utc::now())); - apply_account_quota_window_local_usage(&state, &account_id, &mut quota_data).await; - let account = - account_data_for_quota_refresh(&state, &account_id, quota_data.clone()).await?; - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AdminAccountQuotaResponseData::from_account( - quota, raw, quota_data, account, - )), - )) - } - Err(AccountManageError::NotFound) => Err(account_not_found()), - Err(AccountManageError::Inactive(status)) => Err(AdminError::conflict(format!( - "Account is {}, cannot query quota", - status.as_str() - ))), - Err(AccountManageError::FetchQuota(msg)) => Err(AdminError::bad_gateway(format!( - "Failed to fetch quota from Codex API: {msg}" - ))), - Err(e) => Err(account_error(&e)), - } -} - -async fn account_data_for_quota_refresh( - state: &AppState, - account_id: &str, - quota_data: AccountQuotaData, -) -> Result { - let Some(account) = state - .services - .admin_accounts - .get(account_id) - .await - .map_err(|error| account_error(&error))? - else { - return Err(account_not_found()); - }; - let quota_by_account = HashMap::from([(account_id.to_string(), quota_data)]); - let stats = - account_list_stats(state, std::slice::from_ref(&account), "a_by_account).await?; - Ok(stats.data_for(account)) -} - -/// `POST /api/admin/accounts/import` -pub(crate) async fn import_accounts( - State(state): State, - _auth: AdminAuth, - Json(payload): Json, -) -> Result { - match state.services.admin_accounts.import(payload).await { - Ok(result) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountImportData { - imported: result.imported, - skipped: result.skipped, - source_format: result.source_format.to_string(), - }), - )), - Err(error) => Err(account_error(&error)), - } -} - -/// `POST /api/admin/accounts/oauth/authorize` -pub(crate) async fn oauth_authorize_account( - State(state): State, - _auth: AdminAuth, -) -> Result { - match state.services.admin_accounts.oauth_authorize().await { - Ok(result) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountOAuthAuthorizeData { - session_id: result.session_id, - auth_url: result.auth_url, - expires_at: china_rfc3339(&result.expires_at), - expires_at_display: china_datetime(&result.expires_at), - }), - )), - Err(error) => Err(account_error(&error)), - } -} - -/// `POST /api/admin/accounts/oauth/exchange` -pub(crate) async fn oauth_exchange_account( - State(state): State, - _auth: AdminAuth, - Json(payload): Json, -) -> Result { - match state - .services - .admin_accounts - .oauth_exchange(OAuthExchangeInput { - session_id: payload.session_id, - callback_url: payload.callback_url, - code: payload.code, - state: payload.state, - }) - .await - { - Ok(result) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountImportData { - imported: result.imported, - skipped: result.skipped, - source_format: result.source_format.to_string(), - }), - )), - Err(error) => Err(account_error(&error)), - } -} - -/// `GET /api/admin/accounts/test?id=...&modelId=...` -pub(crate) async fn test_account_connection( - State(state): State, - Query(query): Query, - _auth: AdminAuth, -) -> Result { - let account_id = query.id; - - let model = query - .model_id - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToString::to_string) - .ok_or_else(|| AdminError::bad_request("Model is required"))?; - let stream = state - .services - .admin_accounts - .test_connection_stream(&account_id, model) - .await - .map_err(|error| account_error(&error))?; - - Response::builder() - .status(StatusCode::OK) - .header(CONTENT_TYPE, "text/event-stream") - .header(CACHE_CONTROL, "no-cache") - .header(CONNECTION, "keep-alive") - .header("x-accel-buffering", "no") - .body(Body::from_stream(stream)) - .map_err(|_| AdminError::internal("Failed to build account test stream")) -} - -/// `GET /api/admin/accounts/models?id=...` -pub(crate) async fn account_models( - State(state): State, - Query(query): Query, - _auth: AdminAuth, -) -> Result { - let account_id = query.id; - match state - .services - .admin_accounts - .account_models(&account_id) - .await - { - Ok(models) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountModelsData { - models: models - .into_iter() - .map(|model| AccountModelData { - id: model.id, - label: model.label, - }) - .collect(), - }), - )), - Err(error) => Err(account_error(&error)), - } -} - -/// `POST /api/admin/accounts/delete` -pub(crate) async fn batch_delete_accounts( - State(state): State, - _auth: AdminAuth, - Json(payload): Json, -) -> Result { - match state - .services - .admin_accounts - .batch_delete(payload.ids) - .await - { - Ok(result) => { - for account_id in &result.deleted_ids { - state - .services - .session_affinity - .forget_account(account_id) - .await; - } - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(BatchDeleteData { - deleted: result.deleted, - not_found: result.not_found, - }), - )) - } - Err(error) => Err(account_error(&error)), - } -} - -/// `POST /api/admin/accounts/update` -pub(crate) async fn update_account( - State(state): State, - _auth: AdminAuth, - Json(payload): Json, -) -> Result { - let ParsedAccountUpdate { id, update } = parse_account_update(&payload)?; - - match state - .services - .admin_accounts - .update_account(&id, update) - .await - { - Ok(Some(account)) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountUpdateData::Account(Box::new(account.into()))), - )), - Ok(None) => Err(account_not_found()), - Err(error) => Err(account_error(&error)), - } -} diff --git a/backend/src/api/admin/accounts_routes/export_routes.rs b/backend/src/api/admin/accounts_routes/export_routes.rs new file mode 100644 index 000000000..84fb0ae68 --- /dev/null +++ b/backend/src/api/admin/accounts_routes/export_routes.rs @@ -0,0 +1,24 @@ +use super::*; + +/// `GET /api/admin/accounts/export` +pub(crate) async fn export_accounts( + State(state): State, + _auth: AdminAuth, + Query(query): Query, +) -> Result { + let ids = account_export_ids(query.ids.as_deref()); + + if query.confirm.as_deref() != Some(ACCOUNT_EXPORT_CONFIRMATION) { + return Err(AdminError::bad_request( + "account export requires confirm=export_sensitive_accounts", + )); + } + + match state.services.admin_accounts.export(ids).await { + Ok(result) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(result), + )), + Err(error) => Err(account_error(&error)), + } +} diff --git a/backend/src/api/admin/accounts_routes/import_routes.rs b/backend/src/api/admin/accounts_routes/import_routes.rs new file mode 100644 index 000000000..e1b588e58 --- /dev/null +++ b/backend/src/api/admin/accounts_routes/import_routes.rs @@ -0,0 +1,20 @@ +use super::*; + +/// `POST /api/admin/accounts/import` +pub(crate) async fn import_accounts( + State(state): State, + _auth: AdminAuth, + Json(payload): Json, +) -> Result { + match state.services.admin_accounts.import(payload).await { + Ok(result) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountImportData { + imported: result.imported, + skipped: result.skipped, + source_format: result.source_format.to_string(), + }), + )), + Err(error) => Err(account_error(&error)), + } +} diff --git a/backend/src/api/admin/accounts_routes/lifecycle_routes.rs b/backend/src/api/admin/accounts_routes/lifecycle_routes.rs new file mode 100644 index 000000000..635015ce1 --- /dev/null +++ b/backend/src/api/admin/accounts_routes/lifecycle_routes.rs @@ -0,0 +1,110 @@ +use super::*; + +/// `POST /api/admin/accounts/refresh` +pub(crate) async fn refresh_account( + State(state): State, + _auth: AdminAuth, + Json(payload): Json, +) -> Result { + let account_id = payload.id; + match state + .services + .admin_accounts + .refresh_account(&account_id) + .await + { + Ok(result) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountRefreshData::from(result)), + )), + Err(error) => Err(account_error(&error)), + } +} + +/// `POST /api/admin/accounts/health-check` +pub(crate) async fn health_check_accounts( + State(state): State, + _auth: AdminAuth, + Json(payload): Json, +) -> Result { + let stagger_ms = payload.stagger_ms.unwrap_or(3_000); + if !(500..=30_000).contains(&stagger_ms) && stagger_ms != 0 { + return Err(AdminError::bad_request( + "stagger_ms must be between 500 and 30000", + )); + } + if let Some(concurrency) = payload.concurrency { + if !(1..=10).contains(&concurrency) { + return Err(AdminError::bad_request( + "concurrency must be between 1 and 10", + )); + } + } + match state + .services + .admin_accounts + .health_check_accounts(payload.ids, stagger_ms, payload.concurrency.unwrap_or(2)) + .await + { + Ok(result) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountHealthCheckData::from(result)), + )), + Err(error) => Err(account_error(&error)), + } +} + +/// `POST /api/admin/accounts/delete` +pub(crate) async fn batch_delete_accounts( + State(state): State, + _auth: AdminAuth, + Json(payload): Json, +) -> Result { + match state + .services + .admin_accounts + .batch_delete(payload.ids) + .await + { + Ok(result) => { + for account_id in &result.deleted_ids { + state + .services + .session_affinity + .forget_account(account_id) + .await; + } + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(BatchDeleteData { + deleted: result.deleted, + not_found: result.not_found, + }), + )) + } + Err(error) => Err(account_error(&error)), + } +} + +/// `POST /api/admin/accounts/update` +pub(crate) async fn update_account( + State(state): State, + _auth: AdminAuth, + Json(payload): Json, +) -> Result { + let ParsedAccountUpdate { id, update } = parse_account_update(&payload)?; + + match state + .services + .admin_accounts + .update_account(&id, update) + .await + { + Ok(Some(account)) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountUpdateData::Account(Box::new(account.into()))), + )), + Ok(None) => Err(account_not_found()), + Err(error) => Err(account_error(&error)), + } +} diff --git a/backend/src/api/admin/accounts_routes/oauth_routes.rs b/backend/src/api/admin/accounts_routes/oauth_routes.rs new file mode 100644 index 000000000..796f2083a --- /dev/null +++ b/backend/src/api/admin/accounts_routes/oauth_routes.rs @@ -0,0 +1,49 @@ +use super::*; + +/// `POST /api/admin/accounts/oauth/authorize` +pub(crate) async fn oauth_authorize_account( + State(state): State, + _auth: AdminAuth, +) -> Result { + match state.services.admin_accounts.oauth_authorize().await { + Ok(result) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountOAuthAuthorizeData { + session_id: result.session_id, + auth_url: result.auth_url, + expires_at: china_rfc3339(&result.expires_at), + expires_at_display: china_datetime(&result.expires_at), + }), + )), + Err(error) => Err(account_error(&error)), + } +} + +/// `POST /api/admin/accounts/oauth/exchange` +pub(crate) async fn oauth_exchange_account( + State(state): State, + _auth: AdminAuth, + Json(payload): Json, +) -> Result { + match state + .services + .admin_accounts + .oauth_exchange(OAuthExchangeInput { + session_id: payload.session_id, + callback_url: payload.callback_url, + code: payload.code, + state: payload.state, + }) + .await + { + Ok(result) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountImportData { + imported: result.imported, + skipped: result.skipped, + source_format: result.source_format.to_string(), + }), + )), + Err(error) => Err(account_error(&error)), + } +} diff --git a/backend/src/api/admin/accounts_routes/probe_routes.rs b/backend/src/api/admin/accounts_routes/probe_routes.rs new file mode 100644 index 000000000..63f05a1d8 --- /dev/null +++ b/backend/src/api/admin/accounts_routes/probe_routes.rs @@ -0,0 +1,62 @@ +use super::*; + +/// `GET /api/admin/accounts/test?id=...&modelId=...` +pub(crate) async fn test_account_connection( + State(state): State, + Query(query): Query, + _auth: AdminAuth, +) -> Result { + let account_id = query.id; + + let model = query + .model_id + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToString::to_string) + .ok_or_else(|| AdminError::bad_request("Model is required"))?; + let stream = state + .services + .admin_accounts + .test_connection_stream(&account_id, model) + .await + .map_err(|error| account_error(&error))?; + + Response::builder() + .status(StatusCode::OK) + .header(CONTENT_TYPE, "text/event-stream") + .header(CACHE_CONTROL, "no-cache") + .header(CONNECTION, "keep-alive") + .header("x-accel-buffering", "no") + .body(Body::from_stream(stream)) + .map_err(|_| AdminError::internal("Failed to build account test stream")) +} + +/// `GET /api/admin/accounts/models?id=...` +pub(crate) async fn account_models( + State(state): State, + Query(query): Query, + _auth: AdminAuth, +) -> Result { + let account_id = query.id; + match state + .services + .admin_accounts + .account_models(&account_id) + .await + { + Ok(models) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountModelsData { + models: models + .into_iter() + .map(|model| AccountModelData { + id: model.id, + label: model.label, + }) + .collect(), + }), + )), + Err(error) => Err(account_error(&error)), + } +} diff --git a/backend/src/api/admin/accounts_routes/stats.rs b/backend/src/api/admin/accounts_routes/query.rs similarity index 100% rename from backend/src/api/admin/accounts_routes/stats.rs rename to backend/src/api/admin/accounts_routes/query.rs diff --git a/backend/src/api/admin/accounts_routes/quota_routes.rs b/backend/src/api/admin/accounts_routes/quota_routes.rs new file mode 100644 index 000000000..d0625bf49 --- /dev/null +++ b/backend/src/api/admin/accounts_routes/quota_routes.rs @@ -0,0 +1,61 @@ +use super::*; + +/// `GET /api/admin/accounts/quota` +pub(crate) async fn account_quota( + State(state): State, + _auth: AdminAuth, + Query(query): Query, +) -> Result { + let account_id = query.id; + match state + .services + .admin_accounts + .account_quota(&account_id) + .await + { + Ok(data) => { + let quota = data.get("quota").cloned().unwrap_or(Value::Null); + let raw = data.get("raw").cloned().unwrap_or(Value::Null); + let quota_json = quota.to_string(); + let mut quota_data = quota_data("a_json, Some(Utc::now())); + apply_account_quota_window_local_usage(&state, &account_id, &mut quota_data).await; + let account = + account_data_for_quota_refresh(&state, &account_id, quota_data.clone()).await?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AdminAccountQuotaResponseData::from_account( + quota, raw, quota_data, account, + )), + )) + } + Err(AccountManageError::NotFound) => Err(account_not_found()), + Err(AccountManageError::Inactive(status)) => Err(AdminError::conflict(format!( + "Account is {}, cannot query quota", + status.as_str() + ))), + Err(AccountManageError::FetchQuota(msg)) => Err(AdminError::bad_gateway(format!( + "Failed to fetch quota from Codex API: {msg}" + ))), + Err(e) => Err(account_error(&e)), + } +} + +async fn account_data_for_quota_refresh( + state: &AppState, + account_id: &str, + quota_data: AccountQuotaData, +) -> Result { + let Some(account) = state + .services + .admin_accounts + .get(account_id) + .await + .map_err(|error| account_error(&error))? + else { + return Err(account_not_found()); + }; + let quota_by_account = HashMap::from([(account_id.to_string(), quota_data)]); + let stats = + account_list_stats(state, std::slice::from_ref(&account), "a_by_account).await?; + Ok(stats.data_for(account)) +} diff --git a/backend/src/api/admin/auth_routes.rs b/backend/src/api/admin/auth_routes.rs new file mode 100644 index 000000000..b202a4ee4 --- /dev/null +++ b/backend/src/api/admin/auth_routes.rs @@ -0,0 +1,143 @@ +//! 管理端认证路由。 + +use axum::{ + extract::State, + http::{header::SET_COOKIE, HeaderMap, HeaderValue, StatusCode}, + response::{IntoResponse, Response}, + Extension, Json, +}; +use serde::{Deserialize, Serialize}; + +use crate::{ + api::{ + admin::{ + response::{AdminEnvelope, AdminError, AdminResponse}, + session::{admin_session_cookie, ADMIN_SESSION_COOKIE}, + }, + middleware::request_id::ClientIp, + AppState, + }, + auth::types::SessionError, + infra::time::china_rfc3339, +}; + +const ADMIN_SESSION_COOKIE_ATTRS: &str = "Path=/; Secure; HttpOnly; SameSite=Lax"; + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AdminLoginRequest { + username: Option, + password: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct AdminLoginData { + expires_at: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct AdminSessionStatusData { + authenticated: bool, +} + +/// `POST /api/admin/login` +pub(crate) async fn login( + State(state): State, + client_ip: Option>, + Json(payload): Json, +) -> Result { + let source = admin_login_source(optional_client_ip(client_ip)); + let login_result = state + .services + .admin_sessions + .login( + &source, + payload.username.as_deref(), + payload.password.as_str(), + ) + .await; + let session = match login_result { + Ok(Some(session)) => session, + Ok(None) => return Err(AdminError::invalid_admin_credentials()), + Err(SessionError::LoginThrottled) => return Err(AdminError::too_many_login_attempts()), + Err(_) => return Err(AdminError::internal("Failed to create admin session")), + }; + + let mut response = AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AdminLoginData { + expires_at: china_rfc3339(&session.expires_at), + }), + ) + .into_response(); + let cookie = format!( + "{ADMIN_SESSION_COOKIE}={}; {ADMIN_SESSION_COOKIE_ATTRS}", + session.session_id + ); + response.headers_mut().insert( + SET_COOKIE, + HeaderValue::from_str(&cookie) + .map_err(|_| AdminError::internal("Failed to create admin session cookie"))?, + ); + Ok(response) +} + +/// `GET /api/admin/auth/status` +pub(crate) async fn session_status( + State(state): State, + headers: HeaderMap, +) -> Result { + let authenticated = state + .services + .admin_sessions + .validate(admin_session_cookie(&headers).as_deref()) + .await + .map_err(|_| AdminError::internal("Failed to validate admin session"))?; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AdminSessionStatusData { authenticated }), + )) +} + +/// `POST /api/admin/logout` +pub(crate) async fn logout( + State(state): State, + headers: HeaderMap, +) -> Result { + if let Some(session_id) = admin_session_cookie(&headers) { + let _ = state + .services + .admin_sessions + .delete_session(&session_id) + .await; + } + + let mut response = AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(serde_json::json!({ + "message": "Logged out successfully" + })), + ) + .into_response(); + + let cookie = format!("{ADMIN_SESSION_COOKIE}=; {ADMIN_SESSION_COOKIE_ATTRS}; Max-Age=0"); + response.headers_mut().insert( + SET_COOKIE, + HeaderValue::from_str(&cookie) + .map_err(|_| AdminError::internal("Failed to clear admin session cookie"))?, + ); + + Ok(response) +} + +fn optional_client_ip(client_ip: Option>) -> Option { + let Extension(client_ip) = client_ip?; + Some(client_ip.as_str().to_string()) +} + +fn admin_login_source(fallback_client_ip: Option) -> String { + fallback_client_ip.unwrap_or_else(|| "unknown".to_string()) +} diff --git a/backend/src/api/admin/dashboard_routes.rs b/backend/src/api/admin/dashboard_routes.rs index 365d1831d..49acab73a 100644 --- a/backend/src/api/admin/dashboard_routes.rs +++ b/backend/src/api/admin/dashboard_routes.rs @@ -41,8 +41,7 @@ use crate::{ DashboardTrendData, DashboardTrendKind, }, usage::{ - query::UsageQueryFilter, - store::UsageRecordAccountUsage, + query::{UsageQueryFilter, UsageRecordAccountUsage}, types::{metadata_string, UsageRecord}, }, }, @@ -293,5 +292,302 @@ pub(crate) async fn dashboard_trend( )) } -mod data; -use data::*; +async fn dashboard_time_buckets( + state: &AppState, + now: DateTime, +) -> Vec { + let current_slot = china_quarter_hour_start(now); + let start = current_slot + - Duration::minutes(DASHBOARD_TIME_BUCKET_MINUTES * (DASHBOARD_TIME_BUCKET_SLOTS - 1)); + state + .services + .usage + .time_buckets(start, now) + .await + .unwrap_or_default() +} + +fn dashboard_account_counts(accounts: &[Account]) -> DashboardAccountCounts { + DashboardAccountCounts { + total: accounts.len() as u64, + enabled: accounts + .iter() + .filter(|account| account.status == AccountStatus::Active) + .count() as u64, + abnormal: accounts + .iter() + .filter(|account| account.status != AccountStatus::Active) + .count() as u64, + } +} + +fn account_pool_summary( + accounts: &[Account], + refreshing_account_ids: &HashSet, +) -> AccountPoolDiagnostics { + let mut summary = AccountPoolDiagnostics { + total: accounts.len(), + ..AccountPoolDiagnostics::default() + }; + for account in accounts { + match account.status { + status + if token_refresh_status_eligible(status) + && refreshing_account_ids.contains(&account.id) => + { + summary.refreshing += 1; + } + AccountStatus::Active => summary.active += 1, + AccountStatus::Expired => summary.expired += 1, + AccountStatus::QuotaExhausted => summary.quota_exhausted += 1, + AccountStatus::Disabled => summary.disabled += 1, + AccountStatus::Banned => summary.banned += 1, + } + } + summary +} + +fn account_usage_data( + accounts: &[Account], + usage_records: &[UsageRecordAccountUsage], + quota_used_by_account: &HashMap, + now: DateTime, +) -> Vec { + let account_by_id = accounts + .iter() + .map(|account| (account.id.as_str(), account)) + .collect::>(); + usage_records + .iter() + .map(|usage| { + let account = account_by_id.get(usage.account_id.as_str()).copied(); + let quota_used_percent = quota_used_by_account + .get(&usage.account_id) + .copied() + .or_else(|| { + matches!( + account.map(|account| account.status), + Some(AccountStatus::QuotaExhausted) + ) + .then_some(100.0) + }); + DashboardAccountUsageData { + id: usage.account_id.clone(), + email: account + .and_then(|account| account.email.as_ref()) + .cloned() + .unwrap_or_else(|| usage.account_id.clone()), + plan_type: account + .and_then(|account| account.plan_type.as_ref()) + .cloned(), + tokens: format_tokens(usage.total_tokens), + quota_used_percent, + last_used: china_relative_time(Some(usage.last_used_at), now), + } + }) + .collect() +} + +async fn account_quota_used_percent_by_id( + state: &AppState, + usage_records: &[UsageRecordAccountUsage], +) -> HashMap { + let mut quota_used_by_account = HashMap::with_capacity(usage_records.len()); + for usage in usage_records { + let Ok(Some(quota_json)) = state + .services + .accounts + .get_quota_json(&usage.account_id) + .await + else { + continue; + }; + if let Some(used_percent) = quota_used_percent("a_json) { + quota_used_by_account.insert(usage.account_id.clone(), used_percent); + } + } + quota_used_by_account +} + +fn quota_used_percent(quota_json: &str) -> Option { + let quota = serde_json::from_str::(quota_json).ok()?; + let mut selected = None; + if let Some(monthly_limit) = quota.get("monthly_limit") { + select_quota_used_percent( + &mut selected, + quota_window_priority(monthly_limit, QuotaWindowPriority::Monthly), + monthly_limit.get("used_percent").and_then(percent_value), + ); + } + if let Some(snapshots) = quota.get("snapshots").and_then(Value::as_array) { + for snapshot in snapshots { + for role in ["primary", "secondary"] { + if let Some(window) = snapshot.get(role) { + select_quota_used_percent( + &mut selected, + quota_window_priority(window, QuotaWindowPriority::Other), + window.get("used_percent").and_then(percent_value), + ); + } + } + } + } + selected.map(|(_, used_percent)| used_percent) +} + +fn select_quota_used_percent( + selected: &mut Option<(QuotaWindowPriority, f64)>, + priority: QuotaWindowPriority, + used_percent: Option, +) { + let Some(used_percent) = used_percent else { + return; + }; + match selected { + Some((selected_priority, selected_percent)) + if priority > *selected_priority + || (priority == *selected_priority && used_percent <= *selected_percent) => {} + _ => *selected = Some((priority, used_percent)), + } +} + +fn quota_window_priority(window: &Value, fallback: QuotaWindowPriority) -> QuotaWindowPriority { + let window_seconds = window + .get("window_minutes") + .and_then(Value::as_u64) + .and_then(|minutes| minutes.checked_mul(60)); + match window_seconds { + Some(seconds) if quota_window_matches(seconds, FIVE_HOUR_WINDOW_SECONDS) => { + QuotaWindowPriority::FiveHour + } + Some(seconds) if quota_window_matches(seconds, WEEK_WINDOW_SECONDS) => { + QuotaWindowPriority::Weekly + } + Some(seconds) if quota_window_matches(seconds, MONTH_WINDOW_SECONDS) => { + QuotaWindowPriority::Monthly + } + _ => fallback, + } +} + +fn quota_window_matches(actual: u64, expected: u64) -> bool { + actual > 0 && actual.abs_diff(expected) <= expected / 20 +} + +fn percent_value(value: &Value) -> Option { + let percent = value + .as_f64() + .or_else(|| value.as_str().and_then(|value| value.parse::().ok()))?; + percent.is_finite().then_some(percent.clamp(0.0, 100.0)) +} + +fn dashboard_usage_record_items( + items: Vec, + account_emails: &HashMap, +) -> Vec { + items + .into_iter() + .map(|record| dashboard_usage_record_data(record, account_emails)) + .collect() +} + +fn dashboard_usage_record_data( + record: UsageRecord, + account_emails: &HashMap, +) -> DashboardUsageRecordData { + let account_email = account_emails + .get(&record.account_id) + .cloned() + .or_else(|| metadata_string(&record.metadata, &["accountEmail", "account_email"])); + let (requested_model, upstream_model) = usage_record_models(&record); + let client_ip = metadata_string(&record.metadata, &["clientIp", "ipAddress", "ip_address"]); + let user_agent = metadata_string(&record.metadata, &["userAgent", "user_agent"]); + let reasoning_effort = + metadata_string(&record.metadata, &["reasoningEffort", "reasoning_effort"]); + let token_details = usage_token_details(&record); + let cost_details = usage_cost_details(&record, upstream_model.as_deref(), &token_details); + let first_token_latency_ms = record.first_token_ms; + + DashboardUsageRecordData { + id: record.id, + route: record.route, + model: record.model, + status_code: record.status_code, + transport: record.transport, + created_at_display: china_datetime(&record.created_at), + account_email, + requested_model, + upstream_model, + client_ip, + user_agent, + reasoning_effort, + token_details, + cost_details, + first_token_latency_ms_display: format_duration_ms(first_token_latency_ms), + latency_ms_display: format_duration_ms(record.latency_ms), + } +} + +fn service_status_data(state: &AppState) -> Vec { + let current_fingerprint = state.services.fingerprint.snapshot(); + let fingerprint = fingerprint_diagnostics(¤t_fingerprint); + vec![ + DashboardServiceStatusData { + label: "客户端版本".to_string(), + value: empty_dash(fingerprint.app_version), + detail: if fingerprint.build_number.trim().is_empty() { + "-".to_string() + } else { + format!("Build {}", fingerprint.build_number) + }, + tone: "info".to_string(), + }, + DashboardServiceStatusData { + label: "平台架构".to_string(), + value: empty_dash(fingerprint.platform), + detail: empty_dash(fingerprint.arch), + tone: "info".to_string(), + }, + DashboardServiceStatusData { + label: "Chromium".to_string(), + value: if fingerprint.chromium_version.trim().is_empty() { + "-".to_string() + } else { + format!("v{}", fingerprint.chromium_version) + }, + detail: empty_dash(fingerprint.originator), + tone: "normal".to_string(), + }, + DashboardServiceStatusData { + label: "更新时间".to_string(), + value: format_fingerprint_updated_at(fingerprint.updated_at), + detail: String::new(), + tone: "normal".to_string(), + }, + DashboardServiceStatusData { + label: "User Agent".to_string(), + value: fingerprint.user_agent, + detail: String::new(), + tone: "normal".to_string(), + }, + ] +} + +fn empty_dash(value: String) -> String { + if value.trim().is_empty() { + "-".to_string() + } else { + value + } +} + +fn format_fingerprint_updated_at(value: Option) -> String { + let Some(value) = value else { + return "-".to_string(); + }; + let value = value.trim(); + if value.is_empty() { + return "-".to_string(); + } + china_datetime_rfc3339_str(value) +} diff --git a/backend/src/api/admin/dashboard_routes/data.rs b/backend/src/api/admin/dashboard_routes/data.rs deleted file mode 100644 index eb113891b..000000000 --- a/backend/src/api/admin/dashboard_routes/data.rs +++ /dev/null @@ -1,304 +0,0 @@ -use super::*; - -pub(super) async fn dashboard_time_buckets( - state: &AppState, - now: DateTime, -) -> Vec { - let current_slot = china_quarter_hour_start(now); - let start = current_slot - - Duration::minutes(DASHBOARD_TIME_BUCKET_MINUTES * (DASHBOARD_TIME_BUCKET_SLOTS - 1)); - state - .services - .usage - .time_buckets(start, now) - .await - .unwrap_or_default() -} - -pub(super) fn dashboard_account_counts(accounts: &[Account]) -> DashboardAccountCounts { - DashboardAccountCounts { - total: accounts.len() as u64, - enabled: accounts - .iter() - .filter(|account| account.status == AccountStatus::Active) - .count() as u64, - abnormal: accounts - .iter() - .filter(|account| account.status != AccountStatus::Active) - .count() as u64, - } -} - -pub(super) fn account_pool_summary( - accounts: &[Account], - refreshing_account_ids: &HashSet, -) -> AccountPoolDiagnostics { - let mut summary = AccountPoolDiagnostics { - total: accounts.len(), - ..AccountPoolDiagnostics::default() - }; - for account in accounts { - match account.status { - status - if token_refresh_status_eligible(status) - && refreshing_account_ids.contains(&account.id) => - { - summary.refreshing += 1; - } - AccountStatus::Active => summary.active += 1, - AccountStatus::Expired => summary.expired += 1, - AccountStatus::QuotaExhausted => summary.quota_exhausted += 1, - AccountStatus::Disabled => summary.disabled += 1, - AccountStatus::Banned => summary.banned += 1, - } - } - summary -} - -pub(super) fn account_usage_data( - accounts: &[Account], - usage_records: &[UsageRecordAccountUsage], - quota_used_by_account: &HashMap, - now: DateTime, -) -> Vec { - let account_by_id = accounts - .iter() - .map(|account| (account.id.as_str(), account)) - .collect::>(); - usage_records - .iter() - .map(|usage| { - let account = account_by_id.get(usage.account_id.as_str()).copied(); - let quota_used_percent = quota_used_by_account - .get(&usage.account_id) - .copied() - .or_else(|| { - matches!( - account.map(|account| account.status), - Some(AccountStatus::QuotaExhausted) - ) - .then_some(100.0) - }); - DashboardAccountUsageData { - id: usage.account_id.clone(), - email: account - .and_then(|account| account.email.as_ref()) - .cloned() - .unwrap_or_else(|| usage.account_id.clone()), - plan_type: account - .and_then(|account| account.plan_type.as_ref()) - .cloned(), - tokens: format_tokens(usage.total_tokens), - quota_used_percent, - last_used: china_relative_time(Some(usage.last_used_at), now), - } - }) - .collect() -} - -pub(super) async fn account_quota_used_percent_by_id( - state: &AppState, - usage_records: &[UsageRecordAccountUsage], -) -> HashMap { - let mut quota_used_by_account = HashMap::with_capacity(usage_records.len()); - for usage in usage_records { - let Ok(Some(quota_json)) = state - .services - .accounts - .get_quota_json(&usage.account_id) - .await - else { - continue; - }; - if let Some(used_percent) = quota_used_percent("a_json) { - quota_used_by_account.insert(usage.account_id.clone(), used_percent); - } - } - quota_used_by_account -} - -pub(super) fn quota_used_percent(quota_json: &str) -> Option { - let quota = serde_json::from_str::(quota_json).ok()?; - let mut selected = None; - if let Some(monthly_limit) = quota.get("monthly_limit") { - select_quota_used_percent( - &mut selected, - quota_window_priority(monthly_limit, QuotaWindowPriority::Monthly), - monthly_limit.get("used_percent").and_then(percent_value), - ); - } - if let Some(snapshots) = quota.get("snapshots").and_then(Value::as_array) { - for snapshot in snapshots { - for role in ["primary", "secondary"] { - if let Some(window) = snapshot.get(role) { - select_quota_used_percent( - &mut selected, - quota_window_priority(window, QuotaWindowPriority::Other), - window.get("used_percent").and_then(percent_value), - ); - } - } - } - } - selected.map(|(_, used_percent)| used_percent) -} - -pub(super) fn select_quota_used_percent( - selected: &mut Option<(QuotaWindowPriority, f64)>, - priority: QuotaWindowPriority, - used_percent: Option, -) { - let Some(used_percent) = used_percent else { - return; - }; - match selected { - Some((selected_priority, selected_percent)) - if priority > *selected_priority - || (priority == *selected_priority && used_percent <= *selected_percent) => {} - _ => *selected = Some((priority, used_percent)), - } -} - -pub(super) fn quota_window_priority( - window: &Value, - fallback: QuotaWindowPriority, -) -> QuotaWindowPriority { - let window_seconds = window - .get("window_minutes") - .and_then(Value::as_u64) - .and_then(|minutes| minutes.checked_mul(60)); - match window_seconds { - Some(seconds) if quota_window_matches(seconds, FIVE_HOUR_WINDOW_SECONDS) => { - QuotaWindowPriority::FiveHour - } - Some(seconds) if quota_window_matches(seconds, WEEK_WINDOW_SECONDS) => { - QuotaWindowPriority::Weekly - } - Some(seconds) if quota_window_matches(seconds, MONTH_WINDOW_SECONDS) => { - QuotaWindowPriority::Monthly - } - _ => fallback, - } -} - -pub(super) fn quota_window_matches(actual: u64, expected: u64) -> bool { - actual > 0 && actual.abs_diff(expected) <= expected / 20 -} - -pub(super) fn percent_value(value: &Value) -> Option { - let percent = value - .as_f64() - .or_else(|| value.as_str().and_then(|value| value.parse::().ok()))?; - percent.is_finite().then_some(percent.clamp(0.0, 100.0)) -} - -pub(super) fn dashboard_usage_record_items( - items: Vec, - account_emails: &HashMap, -) -> Vec { - items - .into_iter() - .map(|record| dashboard_usage_record_data(record, account_emails)) - .collect() -} - -pub(super) fn dashboard_usage_record_data( - record: UsageRecord, - account_emails: &HashMap, -) -> DashboardUsageRecordData { - let account_email = account_emails - .get(&record.account_id) - .cloned() - .or_else(|| metadata_string(&record.metadata, &["accountEmail", "account_email"])); - let (requested_model, upstream_model) = usage_record_models(&record); - let client_ip = metadata_string(&record.metadata, &["clientIp", "ipAddress", "ip_address"]); - let user_agent = metadata_string(&record.metadata, &["userAgent", "user_agent"]); - let reasoning_effort = - metadata_string(&record.metadata, &["reasoningEffort", "reasoning_effort"]); - let token_details = usage_token_details(&record); - let cost_details = usage_cost_details(&record, upstream_model.as_deref(), &token_details); - let first_token_latency_ms = record.first_token_ms; - - DashboardUsageRecordData { - id: record.id, - route: record.route, - model: record.model, - status_code: record.status_code, - transport: record.transport, - created_at_display: china_datetime(&record.created_at), - account_email, - requested_model, - upstream_model, - client_ip, - user_agent, - reasoning_effort, - token_details, - cost_details, - first_token_latency_ms_display: format_duration_ms(first_token_latency_ms), - latency_ms_display: format_duration_ms(record.latency_ms), - } -} - -pub(super) fn service_status_data(state: &AppState) -> Vec { - let current_fingerprint = state.services.fingerprint.snapshot(); - let fingerprint = fingerprint_diagnostics(¤t_fingerprint); - vec![ - DashboardServiceStatusData { - label: "客户端版本".to_string(), - value: empty_dash(fingerprint.app_version), - detail: if fingerprint.build_number.trim().is_empty() { - "-".to_string() - } else { - format!("Build {}", fingerprint.build_number) - }, - tone: "info".to_string(), - }, - DashboardServiceStatusData { - label: "平台架构".to_string(), - value: empty_dash(fingerprint.platform), - detail: empty_dash(fingerprint.arch), - tone: "info".to_string(), - }, - DashboardServiceStatusData { - label: "Chromium".to_string(), - value: if fingerprint.chromium_version.trim().is_empty() { - "-".to_string() - } else { - format!("v{}", fingerprint.chromium_version) - }, - detail: empty_dash(fingerprint.originator), - tone: "normal".to_string(), - }, - DashboardServiceStatusData { - label: "更新时间".to_string(), - value: format_fingerprint_updated_at(fingerprint.updated_at), - detail: String::new(), - tone: "normal".to_string(), - }, - DashboardServiceStatusData { - label: "User Agent".to_string(), - value: fingerprint.user_agent, - detail: String::new(), - tone: "normal".to_string(), - }, - ] -} - -pub(super) fn empty_dash(value: String) -> String { - if value.trim().is_empty() { - "-".to_string() - } else { - value - } -} - -pub(super) fn format_fingerprint_updated_at(value: Option) -> String { - let Some(value) = value else { - return "-".to_string(); - }; - let value = value.trim(); - if value.is_empty() { - return "-".to_string(); - } - china_datetime_rfc3339_str(value) -} diff --git a/backend/src/api/admin/mod.rs b/backend/src/api/admin/mod.rs index e657c55f6..e50691b74 100644 --- a/backend/src/api/admin/mod.rs +++ b/backend/src/api/admin/mod.rs @@ -1,6 +1,7 @@ //! 管理端 HTTP API。 pub mod accounts_routes; +pub mod auth_routes; pub mod dashboard_routes; pub mod keys_routes; pub mod ops_routes; diff --git a/backend/src/api/admin/router.rs b/backend/src/api/admin/router.rs index 9f0ccf0ed..9aefced22 100644 --- a/backend/src/api/admin/router.rs +++ b/backend/src/api/admin/router.rs @@ -13,10 +13,10 @@ use super::{ export_accounts, health_check_accounts, import_accounts, oauth_authorize_account, oauth_exchange_account, refresh_account, test_account_connection, update_account, }, + auth_routes::{login, logout, session_status}, dashboard_routes::{dashboard_summary, dashboard_trend}, keys_routes::{api_keys, batch_delete_api_keys, create_api_key, update_api_key}, ops_routes::ops_errors, - session::{login, logout, session_status}, settings_routes::{ admin_api_key_status, delete_admin_api_key, regenerate_admin_api_key, settings, update_settings, diff --git a/backend/src/api/admin/session.rs b/backend/src/api/admin/session.rs index 37376c395..17d87c29a 100644 --- a/backend/src/api/admin/session.rs +++ b/backend/src/api/admin/session.rs @@ -1,25 +1,13 @@ -//! 管理端会话处理器。 +//! 管理端认证提取器与会话鉴权。 use axum::{ - extract::{FromRequestParts, State}, - http::{header::SET_COOKIE, request::Parts, HeaderMap, HeaderValue, StatusCode}, - response::{IntoResponse, Response}, - Extension, Json, -}; -use serde::{Deserialize, Serialize}; - -use crate::{ - api::AppState, - api::{ - admin::response::{AdminEnvelope, AdminError, AdminResponse}, - middleware::request_id::ClientIp, - }, - auth::types::SessionError, - infra::time::china_rfc3339, + extract::FromRequestParts, + http::{request::Parts, HeaderMap}, }; -const ADMIN_SESSION_COOKIE: &str = "cpr_admin_session"; -const ADMIN_SESSION_COOKIE_ATTRS: &str = "Path=/; Secure; HttpOnly; SameSite=Lax"; +use crate::api::{admin::response::AdminError, AppState}; + +pub(super) const ADMIN_SESSION_COOKIE: &str = "cpr_admin_session"; /// 已通过管理员会话或管理员 API Key 鉴权的请求。 pub(crate) struct AdminAuth; @@ -36,116 +24,6 @@ impl FromRequestParts for AdminAuth { } } -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct AdminLoginRequest { - username: Option, - password: String, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct AdminLoginData { - expires_at: String, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct AdminSessionStatusData { - authenticated: bool, -} - -/// `POST /api/admin/login` -pub(crate) async fn login( - State(state): State, - client_ip: Option>, - Json(payload): Json, -) -> Result { - let source = admin_login_source(optional_client_ip(client_ip)); - let login_result = state - .services - .admin_sessions - .login( - &source, - payload.username.as_deref(), - payload.password.as_str(), - ) - .await; - let session = match login_result { - Ok(Some(session)) => session, - Ok(None) => return Err(AdminError::invalid_admin_credentials()), - Err(SessionError::LoginThrottled) => return Err(AdminError::too_many_login_attempts()), - Err(_) => return Err(AdminError::internal("Failed to create admin session")), - }; - - let mut response = AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AdminLoginData { - expires_at: china_rfc3339(&session.expires_at), - }), - ) - .into_response(); - let cookie = format!( - "{ADMIN_SESSION_COOKIE}={}; {ADMIN_SESSION_COOKIE_ATTRS}", - session.session_id - ); - response.headers_mut().insert( - SET_COOKIE, - HeaderValue::from_str(&cookie) - .map_err(|_| AdminError::internal("Failed to create admin session cookie"))?, - ); - Ok(response) -} - -/// `GET /api/admin/auth/status` -pub(crate) async fn session_status( - State(state): State, - headers: HeaderMap, -) -> Result { - let authenticated = state - .services - .admin_sessions - .validate(admin_session_cookie(&headers).as_deref()) - .await - .map_err(|_| AdminError::internal("Failed to validate admin session"))?; - - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AdminSessionStatusData { authenticated }), - )) -} - -/// `POST /api/admin/logout` -pub(crate) async fn logout( - State(state): State, - headers: HeaderMap, -) -> Result { - if let Some(session_id) = admin_session_cookie(&headers) { - let _ = state - .services - .admin_sessions - .delete_session(&session_id) - .await; - } - - let mut response = AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(serde_json::json!({ - "message": "Logged out successfully" - })), - ) - .into_response(); - - let cookie = format!("{ADMIN_SESSION_COOKIE}=; {ADMIN_SESSION_COOKIE_ATTRS}; Max-Age=0"); - response.headers_mut().insert( - SET_COOKIE, - HeaderValue::from_str(&cookie) - .map_err(|_| AdminError::internal("Failed to clear admin session cookie"))?, - ); - - Ok(response) -} - /// 要求请求携带有效管理员会话。 pub(crate) async fn require_admin_session( state: &AppState, @@ -179,7 +57,7 @@ pub(crate) async fn require_admin_auth( require_admin_session(state, headers).await } -fn admin_session_cookie(headers: &HeaderMap) -> Option { +pub(super) fn admin_session_cookie(headers: &HeaderMap) -> Option { let cookie = headers.get("cookie")?.to_str().ok()?; cookie.split(';').find_map(|part| { let (name, value) = part.trim().split_once('=')?; @@ -191,12 +69,3 @@ fn admin_api_key_header(headers: &HeaderMap) -> Option { let value = headers.get("x-api-key")?.to_str().ok()?.trim(); (!value.is_empty()).then(|| value.to_string()) } - -fn optional_client_ip(client_ip: Option>) -> Option { - let Extension(client_ip) = client_ip?; - Some(client_ip.as_str().to_string()) -} - -fn admin_login_source(fallback_client_ip: Option) -> String { - fallback_client_ip.unwrap_or_else(|| "unknown".to_string()) -} diff --git a/backend/src/api/admin/usage_routes.rs b/backend/src/api/admin/usage_routes.rs index aff3fe0c9..029450a7c 100644 --- a/backend/src/api/admin/usage_routes.rs +++ b/backend/src/api/admin/usage_routes.rs @@ -24,10 +24,10 @@ use crate::{ time::china_datetime, }, telemetry::{ - usage::query::{usage_record_cost_details, UsageQueryError, UsageQueryFilter}, - usage::store::{ - UsageRecordBreakdown, UsageRecordEndpointSource, UsageRecordModelSource, - UsageRecordSummary, UsageRecordTrendPoint, + usage::query::{ + usage_record_cost_details, UsageQueryError, UsageQueryFilter, UsageRecordBreakdown, + UsageRecordEndpointSource, UsageRecordModelSource, UsageRecordSummary, + UsageRecordTrendPoint, }, usage::types::{metadata_string, UsageRecord}, }, diff --git a/backend/src/dispatch/affinity/identity.rs b/backend/src/dispatch/affinity/identity.rs deleted file mode 100644 index ed41093d9..000000000 --- a/backend/src/dispatch/affinity/identity.rs +++ /dev/null @@ -1,222 +0,0 @@ -//! 会话亲和变体与 Conversation identity 构建。 - -use serde_json::Value; -use sha2::{Digest, Sha256}; -use uuid::Uuid; - -use super::types::ConversationIdentity; -use crate::upstream::openai::protocol::responses::CodexResponsesRequest; - -// ==================================================================== -// Variant identity helpers -// ==================================================================== - -/// 准备用于区分并发分支的变体身份。 -pub fn prepare_variant_identity(request: &mut CodexResponsesRequest) { - request.variant_identity = build_variant_identity(request); -} - -/// 计算请求变体哈希。 -pub fn compute_variant_hash(request: &CodexResponsesRequest) -> String { - compute_variant_hash_with_identity(request, request.variant_identity.as_deref()) -} - -fn compute_variant_hash_with_identity( - request: &CodexResponsesRequest, - identity: Option<&str>, -) -> String { - let tools_json = request.tools().map_or_else(|| "[]".to_string(), tools_json); - let mut hasher = Sha256::new(); - hasher.update(request.instructions().as_bytes()); - hasher.update(b"\0"); - hasher.update(tools_json.as_bytes()); - if let Some(identity) = identity - .map(str::trim) - .filter(|identity| !identity.is_empty()) - { - hasher.update(b"\0"); - hasher.update(identity.as_bytes()); - } - hex::encode(hasher.finalize()).chars().take(12).collect() -} - -fn tools_json(tools: &[Value]) -> String { - serde_json::to_string(tools).unwrap_or_else(|_| "[]".to_string()) -} - -fn build_variant_identity(request: &CodexResponsesRequest) -> Option { - let mut parts = Vec::with_capacity(2); - if let Some(window_id) = non_empty_str(request.codex_window_id.as_deref()) { - parts.push(format!("window:{window_id}")); - } - if request.explicit_prompt_cache_key - || non_empty_str(request.client_conversation_id.as_deref()).is_some() - { - if let Some(anchor) = derive_stable_conversation_key(request) { - parts.push(format!("anchor:{anchor}")); - } - } - - (!parts.is_empty()).then(|| parts.join("\0")) -} - -fn non_empty_str(value: Option<&str>) -> Option<&str> { - value.map(str::trim).filter(|value| !value.is_empty()) -} - -/// 计算 instructions 哈希。 -pub fn hash_instructions(instructions: Option<&str>) -> String { - let mut hasher = Sha256::new(); - hasher.update(instructions.unwrap_or_default().as_bytes()); - hex::encode(hasher.finalize()) -} - -// ==================================================================== -// Conversation identity 构建器 -// ==================================================================== - -const LEADING_SYSTEM_REMINDER_OPEN: &str = ""; -const LEADING_SYSTEM_REMINDER_CLOSE: &str = ""; - -/// 从 prompt_cache_key 和可选的 window_id 构建 conversation identity -pub fn build_conversation_identity( - prompt_cache_key: Option<&str>, - client_window_id: Option<&str>, - account_scope: &str, -) -> ConversationIdentity { - let conversation_id = prompt_cache_key - .filter(|s| !s.trim().is_empty()) - .map(|key| build_account_scoped_identity("conversation", account_scope, key)); - - let window_id = if let Some(client_win) = client_window_id.filter(|s| !s.trim().is_empty()) { - Some(build_account_scoped_identity( - "window", - account_scope, - client_win, - )) - } else { - conversation_id - .as_ref() - .map(|conv_id| format!("{}:0", conv_id)) - }; - - ConversationIdentity { - conversation_id, - window_id, - } -} - -/// 确保请求拥有上游可复用的 prompt cache key。 -pub fn ensure_prompt_cache_key(request: &mut CodexResponsesRequest) { - if let Some(existing) = request - .prompt_cache_key() - .map(str::trim) - .filter(|value| !value.is_empty()) - { - let existing = existing.to_string(); - request.set_prompt_cache_key(Some(existing)); - return; - } - - request.set_prompt_cache_key(Some( - derive_stable_conversation_key(request).unwrap_or_else(|| Uuid::new_v4().to_string()), - )); -} - -/// 按原版 `stable-conversation-key.ts` 的规则派生稳定 conversation key。 -pub fn derive_stable_conversation_key(request: &CodexResponsesRequest) -> Option { - let instructions = request - .instructions() - .chars() - .take(2000) - .collect::(); - let first_user_text = first_user_text(request.input()); - let normalized_first_user_text = normalize_conversation_anchor_text(&first_user_text); - let first_user_text = if normalized_first_user_text.is_empty() { - first_user_text - } else { - normalized_first_user_text - }; - if instructions.is_empty() && first_user_text.is_empty() { - return None; - } - - let mut hasher = Sha256::new(); - hasher.update(request.model().as_bytes()); - hasher.update(b"\0"); - hasher.update(instructions.as_bytes()); - hasher.update(b"\0"); - hasher.update(first_user_text.as_bytes()); - let hex = hex::encode(hasher.finalize()); - - Some(format!( - "{}-{}-{}-{}-{}", - &hex[0..8], - &hex[8..12], - &hex[12..16], - &hex[16..20], - &hex[20..32] - )) -} - -fn first_user_text(input: &[Value]) -> String { - for item in input { - if item.get("role").and_then(Value::as_str) != Some("user") { - continue; - } - let Some(content) = item.get("content") else { - return String::new(); - }; - if let Some(text) = content.as_str() { - return text.to_string(); - } - if let Some(parts) = content.as_array() { - return parts - .iter() - .filter(|part| part.get("type").and_then(Value::as_str) == Some("input_text")) - .filter_map(|part| part.get("text").and_then(Value::as_str)) - .collect::(); - } - return String::new(); - } - - String::new() -} - -fn normalize_conversation_anchor_text(text: &str) -> String { - let mut rest = text.trim_start(); - loop { - let lower = rest.to_ascii_lowercase(); - if !lower.starts_with(LEADING_SYSTEM_REMINDER_OPEN) { - break; - } - let Some(close_start) = lower.find(LEADING_SYSTEM_REMINDER_CLOSE) else { - break; - }; - let close_end = close_start + LEADING_SYSTEM_REMINDER_CLOSE.len(); - rest = rest[close_end..].trim_start(); - } - rest.to_string() -} - -/// 构建账号作用域的身份哈希。 -fn build_account_scoped_identity(kind: &str, account_scope: &str, client_value: &str) -> String { - let mut hasher = Sha256::new(); - hasher.update(kind.as_bytes()); - hasher.update(b"\0"); - hasher.update(account_scope.as_bytes()); - hasher.update(b"\0"); - hasher.update(client_value.as_bytes()); - - let digest = hasher.finalize(); - let hex = hex::encode(digest); - let truncated = &hex[..32]; - - let prefix = match kind { - "conversation" => "cp", - "window" => "cw", - _ => "cx", - }; - - format!("{}_{}", prefix, truncated) -} diff --git a/backend/src/dispatch/affinity/mod.rs b/backend/src/dispatch/affinity/mod.rs index 576dd89d6..de11bb80a 100644 --- a/backend/src/dispatch/affinity/mod.rs +++ b/backend/src/dispatch/affinity/mod.rs @@ -1,12 +1,11 @@ //! 会话亲和身份、存储、解析与运行时服务。 -pub mod identity; pub mod resolve; pub mod service; pub mod store; pub mod types; -pub use identity::*; +pub use resolve::*; pub use service::*; pub use store::*; pub use types::*; diff --git a/backend/src/dispatch/affinity/resolve.rs b/backend/src/dispatch/affinity/resolve.rs index 2d3615cad..52fb30453 100644 --- a/backend/src/dispatch/affinity/resolve.rs +++ b/backend/src/dispatch/affinity/resolve.rs @@ -1,12 +1,22 @@ use std::sync::Arc; -use chrono::Utc; +use chrono::{DateTime, Utc}; +use serde_json::Value; +use sha2::{Digest, Sha256}; use tokio::sync::Mutex; +use uuid::Uuid; use crate::{ + accounts::account::AccountStatus, dispatch::{ - affinity::{compute_variant_hash, hash_instructions, RuntimeSessionAffinityService}, - recovery::reasoning_replay::ReasoningReplayCache, + affinity::{ + types::ConversationIdentity, RuntimeSessionAffinityService, SessionAffinityEntry, + }, + recovery::{ + implicit_resume::{restore_full_request_without_history, ImplicitResumeSnapshot}, + reasoning_replay::ReasoningReplayCache, + }, + service::ResponseDispatchService, }, upstream::openai::protocol::{ events::TokenUsage, @@ -14,6 +24,151 @@ use crate::{ }, }; +#[derive(Debug, Clone, Default)] +pub(in crate::dispatch) struct AccountAffinityDecision { + preferred_account_id: Option, + request_history_owner: RequestHistoryOwner, +} + +#[derive(Debug, Clone, Default)] +enum RequestHistoryOwner { + #[default] + Absent, + Known(String), + Unknown, +} + +impl AccountAffinityDecision { + pub(in crate::dispatch) fn preferred_account_id(&self) -> Option<&str> { + self.preferred_account_id.as_deref() + } + + fn should_strip_request_history_for(&self, account_id: &str) -> bool { + match &self.request_history_owner { + RequestHistoryOwner::Absent => false, + RequestHistoryOwner::Known(history_account_id) => history_account_id != account_id, + RequestHistoryOwner::Unknown => true, + } + } +} + +impl ResponseDispatchService { + pub(in crate::dispatch) async fn account_affinity_for_request( + &self, + request: &CodexResponsesRequest, + now: DateTime, + ) -> AccountAffinityDecision { + let has_previous_response_id = request.previous_response_id().is_some(); + if let Some(account_id) = self.previous_response_account_id(request, now).await { + return AccountAffinityDecision { + preferred_account_id: Some(account_id.clone()), + request_history_owner: RequestHistoryOwner::Known(account_id), + }; + } + + let conversation_account_id = self.conversation_account_id(request, now).await; + AccountAffinityDecision { + preferred_account_id: conversation_account_id, + request_history_owner: if has_previous_response_id { + RequestHistoryOwner::Unknown + } else { + RequestHistoryOwner::Absent + }, + } + } + + async fn previous_response_account_id( + &self, + request: &CodexResponsesRequest, + now: DateTime, + ) -> Option { + let previous_response_id = request.previous_response_id()?; + self.session_affinity + .lookup_account(previous_response_id, now) + .await + } + + async fn conversation_account_id( + &self, + request: &CodexResponsesRequest, + now: DateTime, + ) -> Option { + let conversation_id = request + .prompt_cache_key() + .map(str::trim) + .filter(|value| !value.is_empty())?; + let variant_hash = compute_variant_hash(request); + self.session_affinity + .lookup_latest_account_by_conversation(conversation_id, None, Some(&variant_hash), now) + .await + } + + pub(in crate::dispatch) fn strip_history_if_account_changed( + request: &mut CodexResponsesRequest, + implicit_resume: &mut Option, + account_affinity: &AccountAffinityDecision, + acquired_account_id: &str, + ) { + if request.previous_response_id().is_none() + || !account_affinity.should_strip_request_history_for(acquired_account_id) + { + return; + } + + let Some(snapshot) = implicit_resume.take() else { + return; + }; + restore_full_request_without_history(request, snapshot); + } + + pub(in crate::dispatch) async fn apply_cascading_ban_defense( + &self, + request: &mut CodexResponsesRequest, + implicit_resume: &mut Option, + preferred_account_id: Option<&str>, + acquired_account_id: &str, + ) -> bool { + let Some(preferred_account_id) = + preferred_account_id.filter(|account_id| *account_id != acquired_account_id) + else { + return false; + }; + let has_history = request.previous_response_id().is_some() + || request + .turn_state + .as_deref() + .is_some_and(|value| !value.is_empty()); + if !has_history { + return false; + } + let Some(preferred_account) = self + .account_pool + .account_snapshot(preferred_account_id) + .await + else { + return false; + }; + if !matches!( + preferred_account.status, + AccountStatus::Banned | AccountStatus::Disabled + ) { + return false; + } + + let response_id_to_forget = request.previous_response_id().map(str::to_string); + let restored_full_request = if let Some(snapshot) = implicit_resume.take() { + restore_full_request_without_history(request, snapshot); + true + } else { + false + }; + if let Some(response_id) = response_id_to_forget { + self.session_affinity.forget(&response_id).await; + } + restored_full_request + } +} + pub(crate) async fn record_response_affinity( session_affinity: &Arc, reasoning_replay: &Arc>, @@ -43,7 +198,7 @@ pub(crate) async fn record_response_affinity( }; let variant_hash = compute_variant_hash(request); - let entry = crate::dispatch::affinity::SessionAffinityEntry { + let entry = SessionAffinityEntry { account_id: account_id.to_string(), conversation_id: conversation_id.to_string(), turn_state: turn_state @@ -91,3 +246,217 @@ pub(crate) async fn evict_reasoning_replay( Utc::now(), ); } + +// ==================================================================== +// Variant identity helpers +// ==================================================================== + +/// 准备用于区分并发分支的变体身份。 +pub fn prepare_variant_identity(request: &mut CodexResponsesRequest) { + request.variant_identity = build_variant_identity(request); +} + +/// 计算请求变体哈希。 +pub fn compute_variant_hash(request: &CodexResponsesRequest) -> String { + compute_variant_hash_with_identity(request, request.variant_identity.as_deref()) +} + +fn compute_variant_hash_with_identity( + request: &CodexResponsesRequest, + identity: Option<&str>, +) -> String { + let tools_json = request.tools().map_or_else(|| "[]".to_string(), tools_json); + let mut hasher = Sha256::new(); + hasher.update(request.instructions().as_bytes()); + hasher.update(b"\0"); + hasher.update(tools_json.as_bytes()); + if let Some(identity) = identity + .map(str::trim) + .filter(|identity| !identity.is_empty()) + { + hasher.update(b"\0"); + hasher.update(identity.as_bytes()); + } + hex::encode(hasher.finalize()).chars().take(12).collect() +} + +fn tools_json(tools: &[Value]) -> String { + serde_json::to_string(tools).unwrap_or_else(|_| "[]".to_string()) +} + +fn build_variant_identity(request: &CodexResponsesRequest) -> Option { + let mut parts = Vec::with_capacity(2); + if let Some(window_id) = non_empty_str(request.codex_window_id.as_deref()) { + parts.push(format!("window:{window_id}")); + } + if request.explicit_prompt_cache_key + || non_empty_str(request.client_conversation_id.as_deref()).is_some() + { + if let Some(anchor) = derive_stable_conversation_key(request) { + parts.push(format!("anchor:{anchor}")); + } + } + + (!parts.is_empty()).then(|| parts.join("\0")) +} + +fn non_empty_str(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} + +/// 计算 instructions 哈希。 +pub fn hash_instructions(instructions: Option<&str>) -> String { + let mut hasher = Sha256::new(); + hasher.update(instructions.unwrap_or_default().as_bytes()); + hex::encode(hasher.finalize()) +} + +// ==================================================================== +// Conversation identity 构建器 +// ==================================================================== + +const LEADING_SYSTEM_REMINDER_OPEN: &str = ""; +const LEADING_SYSTEM_REMINDER_CLOSE: &str = ""; + +/// 从 prompt_cache_key 和可选的 window_id 构建 conversation identity +pub fn build_conversation_identity( + prompt_cache_key: Option<&str>, + client_window_id: Option<&str>, + account_scope: &str, +) -> ConversationIdentity { + let conversation_id = prompt_cache_key + .filter(|s| !s.trim().is_empty()) + .map(|key| build_account_scoped_identity("conversation", account_scope, key)); + + let window_id = if let Some(client_win) = client_window_id.filter(|s| !s.trim().is_empty()) { + Some(build_account_scoped_identity( + "window", + account_scope, + client_win, + )) + } else { + conversation_id + .as_ref() + .map(|conv_id| format!("{}:0", conv_id)) + }; + + ConversationIdentity { + conversation_id, + window_id, + } +} + +/// 确保请求拥有上游可复用的 prompt cache key。 +pub fn ensure_prompt_cache_key(request: &mut CodexResponsesRequest) { + if let Some(existing) = request + .prompt_cache_key() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + let existing = existing.to_string(); + request.set_prompt_cache_key(Some(existing)); + return; + } + + request.set_prompt_cache_key(Some( + derive_stable_conversation_key(request).unwrap_or_else(|| Uuid::new_v4().to_string()), + )); +} + +/// 按原版 `stable-conversation-key.ts` 的规则派生稳定 conversation key。 +pub fn derive_stable_conversation_key(request: &CodexResponsesRequest) -> Option { + let instructions = request + .instructions() + .chars() + .take(2000) + .collect::(); + let first_user_text = first_user_text(request.input()); + let normalized_first_user_text = normalize_conversation_anchor_text(&first_user_text); + let first_user_text = if normalized_first_user_text.is_empty() { + first_user_text + } else { + normalized_first_user_text + }; + if instructions.is_empty() && first_user_text.is_empty() { + return None; + } + + let mut hasher = Sha256::new(); + hasher.update(request.model().as_bytes()); + hasher.update(b"\0"); + hasher.update(instructions.as_bytes()); + hasher.update(b"\0"); + hasher.update(first_user_text.as_bytes()); + let hex = hex::encode(hasher.finalize()); + + Some(format!( + "{}-{}-{}-{}-{}", + &hex[0..8], + &hex[8..12], + &hex[12..16], + &hex[16..20], + &hex[20..32] + )) +} + +fn first_user_text(input: &[Value]) -> String { + for item in input { + if item.get("role").and_then(Value::as_str) != Some("user") { + continue; + } + let Some(content) = item.get("content") else { + return String::new(); + }; + if let Some(text) = content.as_str() { + return text.to_string(); + } + if let Some(parts) = content.as_array() { + return parts + .iter() + .filter(|part| part.get("type").and_then(Value::as_str) == Some("input_text")) + .filter_map(|part| part.get("text").and_then(Value::as_str)) + .collect::(); + } + return String::new(); + } + + String::new() +} + +fn normalize_conversation_anchor_text(text: &str) -> String { + let mut rest = text.trim_start(); + loop { + let lower = rest.to_ascii_lowercase(); + if !lower.starts_with(LEADING_SYSTEM_REMINDER_OPEN) { + break; + } + let Some(close_start) = lower.find(LEADING_SYSTEM_REMINDER_CLOSE) else { + break; + }; + let close_end = close_start + LEADING_SYSTEM_REMINDER_CLOSE.len(); + rest = rest[close_end..].trim_start(); + } + rest.to_string() +} + +/// 构建账号作用域的身份哈希。 +fn build_account_scoped_identity(kind: &str, account_scope: &str, client_value: &str) -> String { + let mut hasher = Sha256::new(); + hasher.update(kind.as_bytes()); + hasher.update(b"\0"); + hasher.update(account_scope.as_bytes()); + hasher.update(b"\0"); + hasher.update(client_value.as_bytes()); + + let digest = hasher.finalize(); + let hex = hex::encode(digest); + let truncated = &hex[..32]; + + let prefix = match kind { + "conversation" => "cp", + "window" => "cw", + _ => "cx", + }; + + format!("{}_{}", prefix, truncated) +} diff --git a/backend/src/dispatch/recording.rs b/backend/src/dispatch/recording.rs index c83f3c150..2854c7682 100644 --- a/backend/src/dispatch/recording.rs +++ b/backend/src/dispatch/recording.rs @@ -31,7 +31,7 @@ use super::{ dispatch_error_metadata, enrich_response_dispatch_error_metadata, ResponseDispatchError, }, stream::{ - lifecycle::{latest_response_id, LiveResponseStreamContext}, + live::{latest_response_id, LiveResponseStreamContext}, sse_failure::{status_code_for_stream_failure, stream_failure_metadata}, trace::{ResponseDispatchAttempt, ResponseDispatchTrace}, }, diff --git a/backend/src/dispatch/recovery/implicit_resume.rs b/backend/src/dispatch/recovery/implicit_resume.rs index 7d8efbde9..4dd1a0c1c 100644 --- a/backend/src/dispatch/recovery/implicit_resume.rs +++ b/backend/src/dispatch/recovery/implicit_resume.rs @@ -2,9 +2,21 @@ use std::collections::HashSet; +use chrono::{Duration, Utc}; use serde_json::Value; -use crate::upstream::openai::protocol::responses::CodexResponsesRequest; +use crate::{ + dispatch::{ + affinity::{ + compute_variant_hash, ensure_prompt_cache_key, hash_instructions, + prepare_variant_identity, + }, + service::ResponseDispatchService, + }, + upstream::openai::protocol::responses::CodexResponsesRequest, +}; + +const IMPLICIT_RESUME_MAX_AGE_SECS: i64 = 55 * 60; /// 隐式续接前的可恢复请求状态。 #[derive(Debug, Clone, PartialEq)] @@ -38,6 +50,132 @@ impl ImplicitResumeSnapshot { } } +impl ResponseDispatchService { + pub(in crate::dispatch) async fn prepare_response_session( + &self, + request: &mut CodexResponsesRequest, + ) -> Option { + prepare_variant_identity(request); + if let Some(previous_response_id) = request.previous_response_id().map(ToString::to_string) + { + if request.prompt_cache_key().is_none() { + let conversation_id = self + .session_affinity + .lookup_conversation_id(&previous_response_id, Utc::now()) + .await; + request.set_prompt_cache_key(conversation_id); + } + if request.turn_state.as_deref().is_none_or(str::is_empty) { + request.turn_state = self + .session_affinity + .lookup_turn_state(&previous_response_id, Utc::now()) + .await; + } + ensure_prompt_cache_key(request); + return None; + } + + ensure_prompt_cache_key(request); + self.apply_implicit_resume(request).await + } + + async fn apply_implicit_resume( + &self, + request: &mut CodexResponsesRequest, + ) -> Option { + let continuation_start = continuation_input_start(request.input()); + if continuation_start == 0 || continuation_start >= request.input().len() { + return None; + } + let conversation_id = request + .prompt_cache_key() + .map(str::trim) + .filter(|value| !value.is_empty())? + .to_string(); + let snapshot = ImplicitResumeSnapshot::capture(request); + let variant_hash = compute_variant_hash(request); + let now = Utc::now(); + let previous_response_id = self + .session_affinity + .lookup_latest_response_by_conversation( + &conversation_id, + Some(Duration::seconds(IMPLICIT_RESUME_MAX_AGE_SECS)), + Some(&variant_hash), + now, + ) + .await?; + let current_instructions_hash = hash_instructions(Some(request.instructions())); + if self + .session_affinity + .lookup_instructions_hash(&previous_response_id, now) + .await + .as_deref() + != Some(current_instructions_hash.as_str()) + { + return None; + } + let stored_function_call_ids = self + .session_affinity + .lookup_function_call_ids(&previous_response_id, now) + .await; + if !implicit_resume_allowed( + &request.input()[continuation_start..], + request.input(), + &stored_function_call_ids, + ) { + return None; + } + let account_id = self + .session_affinity + .lookup_account(&previous_response_id, now) + .await?; + let replay_items = self.reasoning_replay.lock().await.lookup( + &previous_response_id, + &account_id, + &conversation_id, + &variant_hash, + now, + ); + let continuation = request.input()[continuation_start..].to_vec(); + let mut input = replay_items; + input.extend(continuation); + + request.set_previous_response_id(Some(previous_response_id.clone())); + request.use_websocket = true; + request.force_http_sse = false; + request.set_input(input); + if let Some(turn_state) = self + .session_affinity + .lookup_turn_state(&previous_response_id, now) + .await + { + request.turn_state = Some(turn_state); + } + + Some(snapshot) + } + + pub(in crate::dispatch) async fn try_recover_implicit_resume( + &self, + request: &mut CodexResponsesRequest, + implicit_resume: &mut Option, + account_id: &str, + evict_reasoning_replay: bool, + ) -> bool { + let Some(snapshot) = implicit_resume.take() else { + return false; + }; + if evict_reasoning_replay { + self.evict_reasoning_replay(request, account_id).await; + } + if let Some(previous_response_id) = request.previous_response_id() { + self.session_affinity.forget(previous_response_id).await; + } + restore_full_request_without_history(request, snapshot); + true + } +} + /// 返回续接输入在完整输入中的起始位置。 pub fn continuation_input_start(input: &[Value]) -> usize { let mut last_model_output_index = None; @@ -84,6 +222,19 @@ pub fn implicit_resume_allowed( .all(|call_id| required_call_ids.contains(*call_id)) } +fn strip_request_history(request: &mut CodexResponsesRequest) { + request.set_previous_response_id(None); + request.turn_state = None; +} + +pub(in crate::dispatch) fn restore_full_request_without_history( + request: &mut CodexResponsesRequest, + snapshot: ImplicitResumeSnapshot, +) { + snapshot.restore(request); + strip_request_history(request); +} + fn function_call_output_ids(input: &[Value]) -> Vec { input .iter() diff --git a/backend/src/dispatch/service.rs b/backend/src/dispatch/service.rs index 41ea9fa17..787c0970f 100644 --- a/backend/src/dispatch/service.rs +++ b/backend/src/dispatch/service.rs @@ -10,7 +10,7 @@ use std::{pin::Pin, sync::Arc, time::Instant}; use bytes::Bytes; -use chrono::{DateTime, Duration, Utc}; +use chrono::{Duration, Utc}; use futures::stream::Stream; use serde_json::{json, Value}; use tokio::sync::Mutex; @@ -22,10 +22,7 @@ use crate::{ }, dispatch::{ affinity::resolve::{evict_reasoning_replay, record_response_affinity}, - affinity::{ - compute_variant_hash, ensure_prompt_cache_key, hash_instructions, - prepare_variant_identity, RuntimeSessionAffinityService, - }, + affinity::RuntimeSessionAffinityService, errors::{ backend_transport_name, is_history_recovery_upstream_error, is_model_unsupported_upstream_error, is_quota_exhausted_upstream_error, @@ -43,8 +40,6 @@ use crate::{ reasoning_replay::ReasoningReplayCache, }, upstream_call::{ - create_compact_response_with_account_retrying_5xx, - create_response_stream_with_account_retrying_5xx, create_response_with_account_retrying_5xx, verify_acquired_quota_if_required, QuotaVerificationContext, QuotaVerificationDecision, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, @@ -52,19 +47,13 @@ use crate::{ }, models::service::ModelService, telemetry::{ - recorder::{ - reasoning_effort_from_compact_request, reasoning_effort_from_request, - record_response_event, Recorder, - }, + recorder::{reasoning_effort_from_request, record_response_event, Recorder}, usage::types::ResponseUsageRecord, }, upstream::openai::{ protocol::{ - events::{extract_usage, TokenUsage}, - responses::{ - response_from_codex_sse, CodexCompactRequest, CodexResponsesRequest, - CollectedResponse, - }, + events::TokenUsage, + responses::{response_from_codex_sse, CodexResponsesRequest, CollectedResponse}, }, transport::{ backend_transport_for_response_request, is_banned_upstream_error, CodexBackendClient, @@ -73,30 +62,21 @@ use crate::{ }, }; -use crate::dispatch::recovery::implicit_resume::{ - continuation_input_start, implicit_resume_allowed, ImplicitResumeSnapshot, -}; - use super::{ errors::{ResponseDispatchError, ResponseDispatchStreamError}, recording::{ insert_response_status_metadata, insert_response_trace_metadata, insert_response_upstream_diagnostics, insert_websocket_pool_decision, - record_prefetched_response_stream_failure_event, record_response_dispatch_error_event, - record_response_upstream_error_event, ResponseDispatchErrorDetails, - ResponseDispatchErrorEventRecord, ResponseStreamFailureEventRecord, + record_response_dispatch_error_event, record_response_upstream_error_event, + ResponseDispatchErrorDetails, ResponseDispatchErrorEventRecord, ResponseUpstreamErrorEventRecord, }, stream::{ - lifecycle::LiveResponseStreamContext, - live::spawn_live_response_stream, - prefetch::prefetch_first_sse_chunk, sse_failure::{ auth_sse_failure_account_status, client_error_invalid_reasoning_replay, - first_sse_failure, is_auth_sse_failure, is_history_recovery_sse_failure, - is_model_unsupported_sse_failure, is_quota_exhausted_sse_failure, - sse_failure_error_body, sse_failure_invalid_reasoning_replay, - stream_failure_http_status, + is_auth_sse_failure, is_history_recovery_sse_failure, is_model_unsupported_sse_failure, + is_quota_exhausted_sse_failure, sse_failure_error_body, + sse_failure_invalid_reasoning_replay, stream_failure_http_status, }, trace::{ResponseDispatchAttempt, ResponseDispatchTrace}, }, @@ -105,14 +85,14 @@ use super::{ /// OpenAI Responses 调度服务。 #[derive(Clone)] pub struct ResponseDispatchService { - account_pool: Arc, - models: Arc, - codex: Arc, - session_affinity: Arc, - reasoning_replay: Arc>, - recorder: Arc, - installation_id: Option, - cloudflare: CloudflareRecovery, + pub(in crate::dispatch) account_pool: Arc, + pub(in crate::dispatch) models: Arc, + pub(in crate::dispatch) codex: Arc, + pub(in crate::dispatch) session_affinity: Arc, + pub(in crate::dispatch) reasoning_replay: Arc>, + pub(in crate::dispatch) recorder: Arc, + pub(in crate::dispatch) installation_id: Option, + pub(in crate::dispatch) cloudflare: CloudflareRecovery, } pub(crate) struct ResponseDispatchServiceParts { @@ -127,36 +107,6 @@ pub(crate) struct ResponseDispatchServiceParts { /// 默认 reasoning replay TTL 秒数。 const DEFAULT_REASONING_REPLAY_TTL_SECS: i64 = 55 * 60; -const IMPLICIT_RESUME_MAX_AGE_SECS: i64 = DEFAULT_REASONING_REPLAY_TTL_SECS; - -#[derive(Debug, Clone, Default)] -struct AccountAffinityDecision { - preferred_account_id: Option, - request_history_owner: RequestHistoryOwner, -} - -#[derive(Debug, Clone, Default)] -enum RequestHistoryOwner { - #[default] - Absent, - Known(String), - Unknown, -} - -impl AccountAffinityDecision { - fn preferred_account_id(&self) -> Option<&str> { - self.preferred_account_id.as_deref() - } - - fn should_strip_request_history_for(&self, account_id: &str) -> bool { - match &self.request_history_owner { - RequestHistoryOwner::Absent => false, - RequestHistoryOwner::Known(history_account_id) => history_account_id != account_id, - RequestHistoryOwner::Unknown => true, - } - } -} - /// Responses live SSE 响应体流。 pub type ResponseBodyStream = Pin> + Send + 'static>>; @@ -182,246 +132,11 @@ impl ResponseDispatchService { } } - async fn prepare_response_session( - &self, - request: &mut CodexResponsesRequest, - ) -> Option { - prepare_variant_identity(request); - if let Some(previous_response_id) = request.previous_response_id().map(ToString::to_string) - { - if request.prompt_cache_key().is_none() { - let conversation_id = self - .session_affinity - .lookup_conversation_id(&previous_response_id, Utc::now()) - .await; - request.set_prompt_cache_key(conversation_id); - } - if request.turn_state.as_deref().is_none_or(str::is_empty) { - request.turn_state = self - .session_affinity - .lookup_turn_state(&previous_response_id, Utc::now()) - .await; - } - ensure_prompt_cache_key(request); - return None; - } - - ensure_prompt_cache_key(request); - self.apply_implicit_resume(request).await - } - - async fn apply_implicit_resume( - &self, - request: &mut CodexResponsesRequest, - ) -> Option { - let continuation_start = continuation_input_start(request.input()); - if continuation_start == 0 || continuation_start >= request.input().len() { - return None; - } - let conversation_id = request - .prompt_cache_key() - .map(str::trim) - .filter(|value| !value.is_empty())? - .to_string(); - let snapshot = ImplicitResumeSnapshot::capture(request); - let variant_hash = compute_variant_hash(request); - let now = Utc::now(); - let previous_response_id = self - .session_affinity - .lookup_latest_response_by_conversation( - &conversation_id, - Some(Duration::seconds(IMPLICIT_RESUME_MAX_AGE_SECS)), - Some(&variant_hash), - now, - ) - .await?; - let current_instructions_hash = hash_instructions(Some(request.instructions())); - if self - .session_affinity - .lookup_instructions_hash(&previous_response_id, now) - .await - .as_deref() - != Some(current_instructions_hash.as_str()) - { - return None; - } - let stored_function_call_ids = self - .session_affinity - .lookup_function_call_ids(&previous_response_id, now) - .await; - if !implicit_resume_allowed( - &request.input()[continuation_start..], - request.input(), - &stored_function_call_ids, - ) { - return None; - } - let account_id = self - .session_affinity - .lookup_account(&previous_response_id, now) - .await?; - let replay_items = self.reasoning_replay.lock().await.lookup( - &previous_response_id, - &account_id, - &conversation_id, - &variant_hash, - now, - ); - let continuation = request.input()[continuation_start..].to_vec(); - let mut input = replay_items; - input.extend(continuation); - - request.set_previous_response_id(Some(previous_response_id.clone())); - request.use_websocket = true; - request.force_http_sse = false; - request.set_input(input); - if let Some(turn_state) = self - .session_affinity - .lookup_turn_state(&previous_response_id, now) - .await - { - request.turn_state = Some(turn_state); - } - - Some(snapshot) - } - - async fn account_affinity_for_request( + pub(in crate::dispatch) async fn evict_reasoning_replay( &self, request: &CodexResponsesRequest, - now: DateTime, - ) -> AccountAffinityDecision { - let has_previous_response_id = request.previous_response_id().is_some(); - if let Some(account_id) = self.previous_response_account_id(request, now).await { - return AccountAffinityDecision { - preferred_account_id: Some(account_id.clone()), - request_history_owner: RequestHistoryOwner::Known(account_id), - }; - } - - let conversation_account_id = self.conversation_account_id(request, now).await; - AccountAffinityDecision { - preferred_account_id: conversation_account_id, - request_history_owner: if has_previous_response_id { - RequestHistoryOwner::Unknown - } else { - RequestHistoryOwner::Absent - }, - } - } - - async fn previous_response_account_id( - &self, - request: &CodexResponsesRequest, - now: DateTime, - ) -> Option { - let previous_response_id = request.previous_response_id()?; - self.session_affinity - .lookup_account(previous_response_id, now) - .await - } - - async fn conversation_account_id( - &self, - request: &CodexResponsesRequest, - now: DateTime, - ) -> Option { - let conversation_id = request - .prompt_cache_key() - .map(str::trim) - .filter(|value| !value.is_empty())?; - let variant_hash = compute_variant_hash(request); - self.session_affinity - .lookup_latest_account_by_conversation(conversation_id, None, Some(&variant_hash), now) - .await - } - - async fn try_recover_implicit_resume( - &self, - request: &mut CodexResponsesRequest, - implicit_resume: &mut Option, account_id: &str, - evict_reasoning_replay: bool, - ) -> bool { - let Some(snapshot) = implicit_resume.take() else { - return false; - }; - if evict_reasoning_replay { - self.evict_reasoning_replay(request, account_id).await; - } - if let Some(previous_response_id) = request.previous_response_id() { - self.session_affinity.forget(previous_response_id).await; - } - restore_full_request_without_history(request, snapshot); - true - } - - fn strip_history_if_account_changed( - request: &mut CodexResponsesRequest, - implicit_resume: &mut Option, - account_affinity: &AccountAffinityDecision, - acquired_account_id: &str, ) { - if request.previous_response_id().is_none() - || !account_affinity.should_strip_request_history_for(acquired_account_id) - { - return; - } - - let Some(snapshot) = implicit_resume.take() else { - return; - }; - restore_full_request_without_history(request, snapshot); - } - - async fn apply_cascading_ban_defense( - &self, - request: &mut CodexResponsesRequest, - implicit_resume: &mut Option, - preferred_account_id: Option<&str>, - acquired_account_id: &str, - ) -> bool { - let Some(preferred_account_id) = - preferred_account_id.filter(|account_id| *account_id != acquired_account_id) - else { - return false; - }; - let has_history = request.previous_response_id().is_some() - || request - .turn_state - .as_deref() - .is_some_and(|value| !value.is_empty()); - if !has_history { - return false; - } - let Some(preferred_account) = self - .account_pool - .account_snapshot(preferred_account_id) - .await - else { - return false; - }; - if !matches!( - preferred_account.status, - AccountStatus::Banned | AccountStatus::Disabled - ) { - return false; - } - - let response_id_to_forget = request.previous_response_id().map(str::to_string); - let restored_full_request = if let Some(snapshot) = implicit_resume.take() { - restore_full_request_without_history(request, snapshot); - true - } else { - false - }; - if let Some(response_id) = response_id_to_forget { - self.session_affinity.forget(&response_id).await; - } - restored_full_request - } - - async fn evict_reasoning_replay(&self, request: &CodexResponsesRequest, account_id: &str) { evict_reasoning_replay(&self.reasoning_replay, request, account_id).await; } @@ -446,19 +161,576 @@ impl ResponseDispatchService { } } -mod compact; -mod complete; -mod stream; +impl ResponseDispatchService { + /// 调度非流式 Responses 请求到 Codex Responses 上游。 + pub async fn complete( + &self, + request_id: &str, + route: &str, + mut request: CodexResponsesRequest, + requested_model: &str, + ) -> Result { + let started_at = Instant::now(); + let catalog = self.models.catalog().await; + let display_model = catalog.resolve_model_id(requested_model); + request.set_model(display_model.clone()); + let tuple_schema = request.tuple_schema.clone(); + let image_generation_requested = request.expects_image_generation(); + let now = Utc::now(); + let mut implicit_resume = self.prepare_response_session(&mut request).await; + let account_affinity = self.account_affinity_for_request(&request, now).await; + let mut acquire_request = AccountAcquireRequest::new(request.model(), now); + if let Some(preferred_account_id) = account_affinity.preferred_account_id() { + acquire_request = acquire_request.with_preferred_account_id(preferred_account_id); + } + let mut excluded_account_ids = Vec::new(); + let mut exhausted_accounts = AccountExhaustionTracker::default(); + let mut next_required_account_id: Option = None; + let mut empty_response_retries = 0u8; + let mut quota_verify_attempts = 0usize; + let mut trace = ResponseDispatchTrace::default(); + const MAX_EMPTY_RESPONSE_RETRIES: u8 = 2; + let (account, response, collected_response, attempt): ( + Account, + CodexBackendResponse, + CollectedResponse, + ResponseDispatchAttempt, + ) = loop { + let mut attempt_acquire_request = acquire_request + .clone() + .with_exclude_account_ids(excluded_account_ids.iter().cloned()); + if let Some(account_id) = next_required_account_id.take() { + attempt_acquire_request = + attempt_acquire_request.with_required_account_id(account_id); + } + let allow_quota_verify_account_retry = + attempt_acquire_request.required_account_id.is_none(); + attempt_acquire_request.now = Utc::now(); + let Some(acquired) = self + .account_pool + .acquire_with(&attempt_acquire_request) + .await + else { + let error = exhausted_accounts + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount); + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: exhausted_accounts.last_account_id(), + stream: false, + compact: false, + transport: Some(backend_transport_name( + backend_transport_for_response_request(&request), + )), + }, + &error, + ) + .await; + return Err(error); + }; + let acquired_account_id = acquired.account.id.clone(); -fn strip_request_history(request: &mut CodexResponsesRequest) { - request.set_previous_response_id(None); - request.turn_state = None; -} + // 配额验证 + let acquired = match verify_acquired_quota_if_required( + QuotaVerificationContext { + account_pool: self.account_pool.as_ref(), + codex: self.codex.as_ref(), + cloudflare: &self.cloudflare, + installation_id: self.installation_id.as_deref(), + request_id, + excluded_account_ids: &mut excluded_account_ids, + verify_attempts: &mut quota_verify_attempts, + allow_retry_with_another_account: allow_quota_verify_account_retry, + }, + acquired, + ) + .await + { + QuotaVerificationDecision::Ready(acquired) => *acquired, + QuotaVerificationDecision::RetryWithAnotherAccount => { + exhausted_accounts + .record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE); + continue; + } + QuotaVerificationDecision::MaxAttemptsReached + | QuotaVerificationDecision::RequiredAccountUnavailable => { + exhausted_accounts.record_rate_limited( + Some(&acquired_account_id), + QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, + ); + let error = exhausted_accounts + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount); + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: Some(&acquired_account_id), + stream: false, + compact: false, + transport: Some(backend_transport_name( + backend_transport_for_response_request(&request), + )), + }, + &error, + ) + .await; + return Err(error); + } + }; -fn restore_full_request_without_history( - request: &mut CodexResponsesRequest, - snapshot: ImplicitResumeSnapshot, -) { - snapshot.restore(request); - strip_request_history(request); + self.apply_cascading_ban_defense( + &mut request, + &mut implicit_resume, + account_affinity.preferred_account_id(), + &acquired.account.id, + ) + .await; + Self::strip_history_if_account_changed( + &mut request, + &mut implicit_resume, + &account_affinity, + &acquired.account.id, + ); + self.account_pool.wait_for_request_interval(&acquired).await; + let account = acquired.account; + let release_account_id = account.id.clone(); + let attempt = trace.start_attempt(&release_account_id); + let response_result = create_response_with_account_retrying_5xx( + &self.codex, + self.installation_id.as_deref(), + &self.cloudflare, + &request, + request_id, + &account, + started_at, + ) + .await; + self.account_pool.release(&release_account_id).await; + if let Err(error) = &response_result { + self.cloudflare + .capture_set_cookie_headers( + &release_account_id, + upstream_error_set_cookie_headers(error), + ) + .await; + } + + match response_result { + Ok(response) => { + self.cloudflare + .capture_set_cookie_headers( + &release_account_id, + &response.set_cookie_headers, + ) + .await; + self.account_pool + .sync_passive_rate_limit_headers(&account, &response.rate_limit_headers) + .await; + let collected_response = + match response_from_codex_sse(&response.body, tuple_schema.as_ref()) { + Ok(collected_response) => collected_response, + Err(error) => { + let error = ResponseDispatchError::InvalidSse(error); + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: Some(&release_account_id), + stream: false, + compact: false, + transport: Some(backend_transport_name(response.transport)), + }, + &error, + ) + .await; + return Err(error); + } + }; + if matches!(collected_response, CollectedResponse::Empty) { + self.account_pool + .record_empty_response_attempt( + &release_account_id, + request.model(), + image_generation_requested, + ) + .await; + empty_response_retries += 1; + if empty_response_retries <= MAX_EMPTY_RESPONSE_RETRIES { + continue; + } + } + if let CollectedResponse::Failed(failure) = &collected_response { + if is_history_recovery_sse_failure(failure) + && self + .try_recover_implicit_resume( + &mut request, + &mut implicit_resume, + &release_account_id, + sse_failure_invalid_reasoning_replay(failure), + ) + .await + { + next_required_account_id = Some(release_account_id); + continue; + } + if is_model_unsupported_sse_failure(failure) { + let upstream_error = sse_failure_error_body(failure); + if let Some(exhausted) = exhausted_accounts + .model_unsupported_retry_exhausted(upstream_error.clone()) + { + let error = + ResponseDispatchError::from_exhausted_account(exhausted); + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: Some(&release_account_id), + stream: false, + compact: false, + transport: Some(backend_transport_name(response.transport)), + }, + &error, + ) + .await; + return Err(error); + } + exhausted_accounts.record_model_unsupported( + Some(&release_account_id), + upstream_error, + ); + excluded_account_ids.push(release_account_id); + continue; + } + if is_quota_exhausted_sse_failure(failure) { + exhausted_accounts.record_quota_exhausted( + Some(&release_account_id), + failure.message.clone(), + ); + self.account_pool + .set_status(&release_account_id, AccountStatus::QuotaExhausted) + .await; + excluded_account_ids.push(release_account_id); + continue; + } + if is_auth_sse_failure(failure) { + let upstream_error = sse_failure_error_body(failure); + let account_status = auth_sse_failure_account_status(failure); + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + account_status, + upstream_error, + Some(stream_failure_http_status(failure)), + ); + self.account_pool + .set_status(&release_account_id, account_status) + .await; + excluded_account_ids.push(release_account_id); + continue; + } + } + break (account, response, collected_response, attempt); + } + Err(error) if is_rate_limit_upstream_error(&error) => { + exhausted_accounts.record_rate_limited( + Some(&release_account_id), + upstream_error_body(&error), + ); + let cooldown_until = rate_limit_cooldown_until(&error, Utc::now()); + self.account_pool + .mark_quota_limited_until(&release_account_id, cooldown_until) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_quota_exhausted_upstream_error(&error) => { + exhausted_accounts.record_quota_exhausted( + Some(&release_account_id), + upstream_error_body(&error), + ); + self.account_pool + .set_status(&release_account_id, AccountStatus::QuotaExhausted) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_auth_upstream_error(&error) => { + let upstream_error = upstream_error_body(&error); + let account_status = auth_failure_account_status(&error); + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + account_status, + upstream_error, + Some(upstream_error_http_status(&error)), + ); + self.account_pool + .set_status(&release_account_id, account_status) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_cloudflare_challenge_upstream_error(&error) => { + exhausted_accounts.record_cloudflare_challenge( + Some(&release_account_id), + cloudflare_challenge_error_message(), + ); + self.cloudflare + .apply_challenge(self.account_pool.as_ref(), &release_account_id) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_cloudflare_path_block_upstream_error(&error) => { + exhausted_accounts.record_cloudflare_path_blocked( + Some(&release_account_id), + cloudflare_path_block_error_message(), + ); + self.cloudflare + .apply_path_block(self.account_pool.as_ref(), &release_account_id) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_model_unsupported_upstream_error(&error) => { + let upstream_error = upstream_error_body(&error); + if let Some(exhausted) = + exhausted_accounts.model_unsupported_retry_exhausted(upstream_error.clone()) + { + let error = ResponseDispatchError::from_exhausted_account(exhausted); + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: Some(&release_account_id), + stream: false, + compact: false, + transport: Some(backend_transport_name( + backend_transport_for_response_request(&request), + )), + }, + &error, + ) + .await; + return Err(error); + } + exhausted_accounts + .record_model_unsupported(Some(&release_account_id), upstream_error); + excluded_account_ids.push(release_account_id); + } + Err(error) if is_banned_upstream_error(&error) => { + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + AccountStatus::Banned, + upstream_error_body(&error), + Some(upstream_error_http_status(&error)), + ); + self.account_pool + .set_status(&release_account_id, AccountStatus::Banned) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) => { + let history_unavailable = is_history_recovery_upstream_error(&error); + if history_unavailable + && self + .try_recover_implicit_resume( + &mut request, + &mut implicit_resume, + &release_account_id, + client_error_invalid_reasoning_replay(&error), + ) + .await + { + next_required_account_id = Some(release_account_id); + continue; + } + record_response_upstream_error_event(ResponseUpstreamErrorEventRecord { + recorder: &self.recorder, + request_id, + account_id: &release_account_id, + account_email: account.email.as_deref(), + route, + model: requested_model, + started_at, + stream: false, + transport: backend_transport_for_response_request(&request), + request: &request, + error: &error, + trace: &trace, + attempt: Some(&attempt), + }) + .await; + if history_unavailable { + return Err(ResponseDispatchError::HistoryUnavailable { + upstream_error: upstream_error_body(&error), + }); + } + return Err(ResponseDispatchError::Upstream(error)); + } + } + }; + + match collected_response { + CollectedResponse::Completed(body) => { + let response_id = body.get("id").and_then(Value::as_str); + self.cloudflare.reset_account_recovery(&account.id).await; + if let Some(usage) = response.usage { + self.account_pool + .record_response_usage( + &account.id, + request.model(), + usage, + image_generation_requested, + ) + .await; + } + self.record_response_affinity( + &request, + &account.id, + &response.body, + response.turn_state.clone(), + response.usage, + ) + .await; + let mut metadata = json!({ + "responseId": response_id, + "stream": false, + "transport": backend_transport_name(response.transport), + "firstTokenMs": response.first_token_ms, + "usage": response.usage, + }); + insert_response_status_metadata( + &mut metadata, + 200, + 200, + response.diagnostics.status_code.map(i64::from), + ); + insert_response_upstream_diagnostics(&mut metadata, &response.diagnostics); + insert_response_trace_metadata(&mut metadata, &trace, Some(&attempt)); + insert_websocket_pool_decision(&mut metadata, response.websocket_pool_decision); + record_response_event(ResponseUsageRecord { + recorder: &self.recorder, + request_id, + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: &account.id, + route, + model: &display_model, + requested_model: Some(requested_model), + client_ip: request.client_ip.as_deref(), + client_user_agent: request.client_user_agent.as_deref(), + reasoning_effort: reasoning_effort_from_request(&request), + service_tier: request.service_tier(), + started_at, + status_code: 200, + message: "v1 responses completed", + metadata, + rate_limit_headers: &response.rate_limit_headers, + }) + .await; + Ok(body) + } + CollectedResponse::Failed(failure) => { + let error = if is_history_recovery_sse_failure(&failure) { + ResponseDispatchError::HistoryUnavailable { + upstream_error: sse_failure_error_body(&failure), + } + } else { + ResponseDispatchError::Failed(failure) + }; + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: Some(&account.id), + stream: false, + compact: false, + transport: Some(backend_transport_name(response.transport)), + }, + &error, + ) + .await; + Err(error) + } + CollectedResponse::MissingCompleted => { + let error = ResponseDispatchError::MissingCompleted; + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: Some(&account.id), + stream: false, + compact: false, + transport: Some(backend_transport_name(response.transport)), + }, + &error, + ) + .await; + Err(error) + } + CollectedResponse::Empty => { + let error = ResponseDispatchError::EmptyUpstreamResponse; + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: Some(&account.id), + stream: false, + compact: false, + transport: Some(backend_transport_name(response.transport)), + }, + &error, + ) + .await; + Err(error) + } + } + } + + pub(super) async fn record_response_dispatch_error( + &self, + request_id: &str, + route: &str, + requested_model: &str, + started_at: Instant, + details: ResponseDispatchErrorDetails<'_>, + error: &ResponseDispatchError, + ) { + record_response_dispatch_error_event(ResponseDispatchErrorEventRecord { + recorder: &self.recorder, + request_id, + client_api_key_id: details.client_api_key_id, + account_id: details.account_id, + route, + model: requested_model, + started_at, + stream: details.stream, + compact: details.compact, + transport: details.transport, + error, + }) + .await; + } } diff --git a/backend/src/dispatch/service/compact.rs b/backend/src/dispatch/service/compact.rs deleted file mode 100644 index 7b4e81fd5..000000000 --- a/backend/src/dispatch/service/compact.rs +++ /dev/null @@ -1,290 +0,0 @@ -use super::*; - -impl ResponseDispatchService { - /// 调度 Responses compact 请求到 Codex compact 上游。 - pub async fn compact( - &self, - request_id: &str, - mut request: CodexCompactRequest, - requested_model: &str, - ) -> Result { - let started_at = Instant::now(); - let client_api_key_id = request.client_api_key_id.clone(); - let catalog = self.models.catalog().await; - let display_model = catalog.resolve_model_id(requested_model); - request.set_model(display_model.clone()); - let mut excluded_account_ids = Vec::new(); - let mut exhausted_accounts = AccountExhaustionTracker::default(); - let mut quota_verify_attempts = 0usize; - let mut trace = ResponseDispatchTrace::default(); - - loop { - let acquire_request = AccountAcquireRequest::new(request.model(), Utc::now()) - .with_exclude_account_ids(excluded_account_ids.iter().cloned()); - let acquired = match self.account_pool.acquire_with(&acquire_request).await { - Some(acquired) => acquired, - None => { - let error = exhausted_accounts - .last_exhausted() - .map(ResponseDispatchError::from_exhausted_account) - .unwrap_or(ResponseDispatchError::NoActiveAccount); - self.record_compact_dispatch_error( - request_id, - client_api_key_id.as_deref(), - requested_model, - started_at, - exhausted_accounts.last_account_id(), - &error, - ) - .await; - return Err(error); - } - }; - let acquired_account_id = acquired.account.id.clone(); - let acquired = match verify_acquired_quota_if_required( - QuotaVerificationContext { - account_pool: self.account_pool.as_ref(), - codex: self.codex.as_ref(), - cloudflare: &self.cloudflare, - installation_id: self.installation_id.as_deref(), - request_id, - excluded_account_ids: &mut excluded_account_ids, - verify_attempts: &mut quota_verify_attempts, - allow_retry_with_another_account: true, - }, - acquired, - ) - .await - { - QuotaVerificationDecision::Ready(acquired) => *acquired, - QuotaVerificationDecision::RetryWithAnotherAccount => { - exhausted_accounts - .record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE); - continue; - } - QuotaVerificationDecision::MaxAttemptsReached - | QuotaVerificationDecision::RequiredAccountUnavailable => { - exhausted_accounts.record_rate_limited( - Some(&acquired_account_id), - QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, - ); - let error = exhausted_accounts - .last_exhausted() - .map(ResponseDispatchError::from_exhausted_account) - .unwrap_or(ResponseDispatchError::NoActiveAccount); - self.record_compact_dispatch_error( - request_id, - client_api_key_id.as_deref(), - requested_model, - started_at, - Some(&acquired_account_id), - &error, - ) - .await; - return Err(error); - } - }; - let account = acquired.account; - let release_account_id = account.id.clone(); - let attempt = trace.start_attempt(&release_account_id); - let response_result = create_compact_response_with_account_retrying_5xx( - &self.codex, - self.installation_id.as_deref(), - &self.cloudflare, - &request, - request_id, - &account, - ) - .await; - self.account_pool.release(&release_account_id).await; - if let Err(error) = &response_result { - self.cloudflare - .capture_set_cookie_headers( - &release_account_id, - upstream_error_set_cookie_headers(error), - ) - .await; - } - - match response_result { - Ok(response) => { - self.cloudflare - .capture_set_cookie_headers( - &release_account_id, - &response.set_cookie_headers, - ) - .await; - self.cloudflare.reset_account_recovery(&account.id).await; - self.account_pool - .sync_passive_rate_limit_headers(&account, &response.rate_limit_headers) - .await; - let usage = extract_usage(&response.body); - if let Some(usage) = usage { - self.account_pool - .record_token_usage(&account.id, request.model(), &usage) - .await; - } - let mut metadata = json!({ - "stream": false, - "compact": true, - "usage": usage, - }); - insert_response_status_metadata( - &mut metadata, - 200, - 200, - response.diagnostics.status_code.map(i64::from), - ); - insert_response_upstream_diagnostics(&mut metadata, &response.diagnostics); - insert_response_trace_metadata(&mut metadata, &trace, Some(&attempt)); - record_response_event(ResponseUsageRecord { - recorder: &self.recorder, - request_id, - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: &account.id, - route: "/v1/responses/compact", - model: &display_model, - requested_model: Some(requested_model), - client_ip: request.client_ip.as_deref(), - client_user_agent: request.client_user_agent.as_deref(), - reasoning_effort: reasoning_effort_from_compact_request(&request), - service_tier: None, - started_at, - status_code: 200, - message: "v1 responses compact completed", - metadata, - rate_limit_headers: &response.rate_limit_headers, - }) - .await; - return Ok(response.body); - } - Err(error) if is_rate_limit_upstream_error(&error) => { - exhausted_accounts.record_rate_limited( - Some(&release_account_id), - upstream_error_body(&error), - ); - let cooldown_until = rate_limit_cooldown_until(&error, Utc::now()); - self.account_pool - .mark_quota_limited_until(&release_account_id, cooldown_until) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_quota_exhausted_upstream_error(&error) => { - exhausted_accounts.record_quota_exhausted( - Some(&release_account_id), - upstream_error_body(&error), - ); - self.account_pool - .set_status(&release_account_id, AccountStatus::QuotaExhausted) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_auth_upstream_error(&error) => { - let upstream_error = upstream_error_body(&error); - let account_status = auth_failure_account_status(&error); - exhausted_accounts.record_auth_failure( - Some(&release_account_id), - account_status, - upstream_error, - Some(upstream_error_http_status(&error)), - ); - self.account_pool - .set_status(&release_account_id, account_status) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_cloudflare_challenge_upstream_error(&error) => { - exhausted_accounts.record_cloudflare_challenge( - Some(&release_account_id), - cloudflare_challenge_error_message(), - ); - self.cloudflare - .apply_challenge(self.account_pool.as_ref(), &release_account_id) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_cloudflare_path_block_upstream_error(&error) => { - exhausted_accounts.record_cloudflare_path_blocked( - Some(&release_account_id), - cloudflare_path_block_error_message(), - ); - self.cloudflare - .apply_path_block(self.account_pool.as_ref(), &release_account_id) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_model_unsupported_upstream_error(&error) => { - let upstream_error = upstream_error_body(&error); - if let Some(exhausted) = - exhausted_accounts.model_unsupported_retry_exhausted(upstream_error.clone()) - { - let error = ResponseDispatchError::from_exhausted_account(exhausted); - self.record_compact_dispatch_error( - request_id, - client_api_key_id.as_deref(), - requested_model, - started_at, - Some(&release_account_id), - &error, - ) - .await; - return Err(error); - } - exhausted_accounts - .record_model_unsupported(Some(&release_account_id), upstream_error); - excluded_account_ids.push(release_account_id); - } - Err(error) if is_banned_upstream_error(&error) => { - exhausted_accounts.record_auth_failure( - Some(&release_account_id), - AccountStatus::Banned, - upstream_error_body(&error), - Some(upstream_error_http_status(&error)), - ); - self.account_pool - .set_status(&release_account_id, AccountStatus::Banned) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) => { - let error = ResponseDispatchError::Upstream(error); - self.record_compact_dispatch_error( - request_id, - client_api_key_id.as_deref(), - requested_model, - started_at, - Some(&release_account_id), - &error, - ) - .await; - return Err(error); - } - } - } - } - - async fn record_compact_dispatch_error( - &self, - request_id: &str, - client_api_key_id: Option<&str>, - requested_model: &str, - started_at: Instant, - account_id: Option<&str>, - error: &ResponseDispatchError, - ) { - record_response_dispatch_error_event(ResponseDispatchErrorEventRecord { - recorder: &self.recorder, - request_id, - client_api_key_id, - account_id, - route: "/v1/responses/compact", - model: requested_model, - started_at, - stream: false, - compact: true, - transport: Some("http"), - error, - }) - .await; - } -} diff --git a/backend/src/dispatch/service/complete.rs b/backend/src/dispatch/service/complete.rs deleted file mode 100644 index 8a2cb5f41..000000000 --- a/backend/src/dispatch/service/complete.rs +++ /dev/null @@ -1,575 +0,0 @@ -use super::*; - -impl ResponseDispatchService { - /// 调度非流式 Responses 请求到 Codex Responses 上游。 - pub async fn complete( - &self, - request_id: &str, - route: &str, - mut request: CodexResponsesRequest, - requested_model: &str, - ) -> Result { - let started_at = Instant::now(); - let catalog = self.models.catalog().await; - let display_model = catalog.resolve_model_id(requested_model); - request.set_model(display_model.clone()); - let tuple_schema = request.tuple_schema.clone(); - let image_generation_requested = request.expects_image_generation(); - let now = Utc::now(); - let mut implicit_resume = self.prepare_response_session(&mut request).await; - let account_affinity = self.account_affinity_for_request(&request, now).await; - let mut acquire_request = AccountAcquireRequest::new(request.model(), now); - if let Some(preferred_account_id) = account_affinity.preferred_account_id() { - acquire_request = acquire_request.with_preferred_account_id(preferred_account_id); - } - let mut excluded_account_ids = Vec::new(); - let mut exhausted_accounts = AccountExhaustionTracker::default(); - let mut next_required_account_id: Option = None; - let mut empty_response_retries = 0u8; - let mut quota_verify_attempts = 0usize; - let mut trace = ResponseDispatchTrace::default(); - const MAX_EMPTY_RESPONSE_RETRIES: u8 = 2; - let (account, response, collected_response, attempt): ( - Account, - CodexBackendResponse, - CollectedResponse, - ResponseDispatchAttempt, - ) = loop { - let mut attempt_acquire_request = acquire_request - .clone() - .with_exclude_account_ids(excluded_account_ids.iter().cloned()); - if let Some(account_id) = next_required_account_id.take() { - attempt_acquire_request = - attempt_acquire_request.with_required_account_id(account_id); - } - let allow_quota_verify_account_retry = - attempt_acquire_request.required_account_id.is_none(); - attempt_acquire_request.now = Utc::now(); - let Some(acquired) = self - .account_pool - .acquire_with(&attempt_acquire_request) - .await - else { - let error = exhausted_accounts - .last_exhausted() - .map(ResponseDispatchError::from_exhausted_account) - .unwrap_or(ResponseDispatchError::NoActiveAccount); - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: exhausted_accounts.last_account_id(), - stream: false, - compact: false, - transport: Some(backend_transport_name( - backend_transport_for_response_request(&request), - )), - }, - &error, - ) - .await; - return Err(error); - }; - let acquired_account_id = acquired.account.id.clone(); - - // 配额验证 - let acquired = match verify_acquired_quota_if_required( - QuotaVerificationContext { - account_pool: self.account_pool.as_ref(), - codex: self.codex.as_ref(), - cloudflare: &self.cloudflare, - installation_id: self.installation_id.as_deref(), - request_id, - excluded_account_ids: &mut excluded_account_ids, - verify_attempts: &mut quota_verify_attempts, - allow_retry_with_another_account: allow_quota_verify_account_retry, - }, - acquired, - ) - .await - { - QuotaVerificationDecision::Ready(acquired) => *acquired, - QuotaVerificationDecision::RetryWithAnotherAccount => { - exhausted_accounts - .record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE); - continue; - } - QuotaVerificationDecision::MaxAttemptsReached - | QuotaVerificationDecision::RequiredAccountUnavailable => { - exhausted_accounts.record_rate_limited( - Some(&acquired_account_id), - QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, - ); - let error = exhausted_accounts - .last_exhausted() - .map(ResponseDispatchError::from_exhausted_account) - .unwrap_or(ResponseDispatchError::NoActiveAccount); - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: Some(&acquired_account_id), - stream: false, - compact: false, - transport: Some(backend_transport_name( - backend_transport_for_response_request(&request), - )), - }, - &error, - ) - .await; - return Err(error); - } - }; - - self.apply_cascading_ban_defense( - &mut request, - &mut implicit_resume, - account_affinity.preferred_account_id(), - &acquired.account.id, - ) - .await; - Self::strip_history_if_account_changed( - &mut request, - &mut implicit_resume, - &account_affinity, - &acquired.account.id, - ); - self.account_pool.wait_for_request_interval(&acquired).await; - let account = acquired.account; - let release_account_id = account.id.clone(); - let attempt = trace.start_attempt(&release_account_id); - let response_result = create_response_with_account_retrying_5xx( - &self.codex, - self.installation_id.as_deref(), - &self.cloudflare, - &request, - request_id, - &account, - started_at, - ) - .await; - self.account_pool.release(&release_account_id).await; - if let Err(error) = &response_result { - self.cloudflare - .capture_set_cookie_headers( - &release_account_id, - upstream_error_set_cookie_headers(error), - ) - .await; - } - - match response_result { - Ok(response) => { - self.cloudflare - .capture_set_cookie_headers( - &release_account_id, - &response.set_cookie_headers, - ) - .await; - self.account_pool - .sync_passive_rate_limit_headers(&account, &response.rate_limit_headers) - .await; - let collected_response = - match response_from_codex_sse(&response.body, tuple_schema.as_ref()) { - Ok(collected_response) => collected_response, - Err(error) => { - let error = ResponseDispatchError::InvalidSse(error); - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: Some(&release_account_id), - stream: false, - compact: false, - transport: Some(backend_transport_name(response.transport)), - }, - &error, - ) - .await; - return Err(error); - } - }; - if matches!(collected_response, CollectedResponse::Empty) { - self.account_pool - .record_empty_response_attempt( - &release_account_id, - request.model(), - image_generation_requested, - ) - .await; - empty_response_retries += 1; - if empty_response_retries <= MAX_EMPTY_RESPONSE_RETRIES { - continue; - } - } - if let CollectedResponse::Failed(failure) = &collected_response { - if is_history_recovery_sse_failure(failure) - && self - .try_recover_implicit_resume( - &mut request, - &mut implicit_resume, - &release_account_id, - sse_failure_invalid_reasoning_replay(failure), - ) - .await - { - next_required_account_id = Some(release_account_id); - continue; - } - if is_model_unsupported_sse_failure(failure) { - let upstream_error = sse_failure_error_body(failure); - if let Some(exhausted) = exhausted_accounts - .model_unsupported_retry_exhausted(upstream_error.clone()) - { - let error = - ResponseDispatchError::from_exhausted_account(exhausted); - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: Some(&release_account_id), - stream: false, - compact: false, - transport: Some(backend_transport_name(response.transport)), - }, - &error, - ) - .await; - return Err(error); - } - exhausted_accounts.record_model_unsupported( - Some(&release_account_id), - upstream_error, - ); - excluded_account_ids.push(release_account_id); - continue; - } - if is_quota_exhausted_sse_failure(failure) { - exhausted_accounts.record_quota_exhausted( - Some(&release_account_id), - failure.message.clone(), - ); - self.account_pool - .set_status(&release_account_id, AccountStatus::QuotaExhausted) - .await; - excluded_account_ids.push(release_account_id); - continue; - } - if is_auth_sse_failure(failure) { - let upstream_error = sse_failure_error_body(failure); - let account_status = auth_sse_failure_account_status(failure); - exhausted_accounts.record_auth_failure( - Some(&release_account_id), - account_status, - upstream_error, - Some(stream_failure_http_status(failure)), - ); - self.account_pool - .set_status(&release_account_id, account_status) - .await; - excluded_account_ids.push(release_account_id); - continue; - } - } - break (account, response, collected_response, attempt); - } - Err(error) if is_rate_limit_upstream_error(&error) => { - exhausted_accounts.record_rate_limited( - Some(&release_account_id), - upstream_error_body(&error), - ); - let cooldown_until = rate_limit_cooldown_until(&error, Utc::now()); - self.account_pool - .mark_quota_limited_until(&release_account_id, cooldown_until) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_quota_exhausted_upstream_error(&error) => { - exhausted_accounts.record_quota_exhausted( - Some(&release_account_id), - upstream_error_body(&error), - ); - self.account_pool - .set_status(&release_account_id, AccountStatus::QuotaExhausted) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_auth_upstream_error(&error) => { - let upstream_error = upstream_error_body(&error); - let account_status = auth_failure_account_status(&error); - exhausted_accounts.record_auth_failure( - Some(&release_account_id), - account_status, - upstream_error, - Some(upstream_error_http_status(&error)), - ); - self.account_pool - .set_status(&release_account_id, account_status) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_cloudflare_challenge_upstream_error(&error) => { - exhausted_accounts.record_cloudflare_challenge( - Some(&release_account_id), - cloudflare_challenge_error_message(), - ); - self.cloudflare - .apply_challenge(self.account_pool.as_ref(), &release_account_id) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_cloudflare_path_block_upstream_error(&error) => { - exhausted_accounts.record_cloudflare_path_blocked( - Some(&release_account_id), - cloudflare_path_block_error_message(), - ); - self.cloudflare - .apply_path_block(self.account_pool.as_ref(), &release_account_id) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_model_unsupported_upstream_error(&error) => { - let upstream_error = upstream_error_body(&error); - if let Some(exhausted) = - exhausted_accounts.model_unsupported_retry_exhausted(upstream_error.clone()) - { - let error = ResponseDispatchError::from_exhausted_account(exhausted); - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: Some(&release_account_id), - stream: false, - compact: false, - transport: Some(backend_transport_name( - backend_transport_for_response_request(&request), - )), - }, - &error, - ) - .await; - return Err(error); - } - exhausted_accounts - .record_model_unsupported(Some(&release_account_id), upstream_error); - excluded_account_ids.push(release_account_id); - } - Err(error) if is_banned_upstream_error(&error) => { - exhausted_accounts.record_auth_failure( - Some(&release_account_id), - AccountStatus::Banned, - upstream_error_body(&error), - Some(upstream_error_http_status(&error)), - ); - self.account_pool - .set_status(&release_account_id, AccountStatus::Banned) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) => { - let history_unavailable = is_history_recovery_upstream_error(&error); - if history_unavailable - && self - .try_recover_implicit_resume( - &mut request, - &mut implicit_resume, - &release_account_id, - client_error_invalid_reasoning_replay(&error), - ) - .await - { - next_required_account_id = Some(release_account_id); - continue; - } - record_response_upstream_error_event(ResponseUpstreamErrorEventRecord { - recorder: &self.recorder, - request_id, - account_id: &release_account_id, - account_email: account.email.as_deref(), - route, - model: requested_model, - started_at, - stream: false, - transport: backend_transport_for_response_request(&request), - request: &request, - error: &error, - trace: &trace, - attempt: Some(&attempt), - }) - .await; - if history_unavailable { - return Err(ResponseDispatchError::HistoryUnavailable { - upstream_error: upstream_error_body(&error), - }); - } - return Err(ResponseDispatchError::Upstream(error)); - } - } - }; - - match collected_response { - CollectedResponse::Completed(body) => { - let response_id = body.get("id").and_then(Value::as_str); - self.cloudflare.reset_account_recovery(&account.id).await; - if let Some(usage) = response.usage { - self.account_pool - .record_response_usage( - &account.id, - request.model(), - usage, - image_generation_requested, - ) - .await; - } - self.record_response_affinity( - &request, - &account.id, - &response.body, - response.turn_state.clone(), - response.usage, - ) - .await; - let mut metadata = json!({ - "responseId": response_id, - "stream": false, - "transport": backend_transport_name(response.transport), - "firstTokenMs": response.first_token_ms, - "usage": response.usage, - }); - insert_response_status_metadata( - &mut metadata, - 200, - 200, - response.diagnostics.status_code.map(i64::from), - ); - insert_response_upstream_diagnostics(&mut metadata, &response.diagnostics); - insert_response_trace_metadata(&mut metadata, &trace, Some(&attempt)); - insert_websocket_pool_decision(&mut metadata, response.websocket_pool_decision); - record_response_event(ResponseUsageRecord { - recorder: &self.recorder, - request_id, - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: &account.id, - route, - model: &display_model, - requested_model: Some(requested_model), - client_ip: request.client_ip.as_deref(), - client_user_agent: request.client_user_agent.as_deref(), - reasoning_effort: reasoning_effort_from_request(&request), - service_tier: request.service_tier(), - started_at, - status_code: 200, - message: "v1 responses completed", - metadata, - rate_limit_headers: &response.rate_limit_headers, - }) - .await; - Ok(body) - } - CollectedResponse::Failed(failure) => { - let error = if is_history_recovery_sse_failure(&failure) { - ResponseDispatchError::HistoryUnavailable { - upstream_error: sse_failure_error_body(&failure), - } - } else { - ResponseDispatchError::Failed(failure) - }; - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: Some(&account.id), - stream: false, - compact: false, - transport: Some(backend_transport_name(response.transport)), - }, - &error, - ) - .await; - Err(error) - } - CollectedResponse::MissingCompleted => { - let error = ResponseDispatchError::MissingCompleted; - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: Some(&account.id), - stream: false, - compact: false, - transport: Some(backend_transport_name(response.transport)), - }, - &error, - ) - .await; - Err(error) - } - CollectedResponse::Empty => { - let error = ResponseDispatchError::EmptyUpstreamResponse; - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: Some(&account.id), - stream: false, - compact: false, - transport: Some(backend_transport_name(response.transport)), - }, - &error, - ) - .await; - Err(error) - } - } - } - - pub(super) async fn record_response_dispatch_error( - &self, - request_id: &str, - route: &str, - requested_model: &str, - started_at: Instant, - details: ResponseDispatchErrorDetails<'_>, - error: &ResponseDispatchError, - ) { - record_response_dispatch_error_event(ResponseDispatchErrorEventRecord { - recorder: &self.recorder, - request_id, - client_api_key_id: details.client_api_key_id, - account_id: details.account_id, - route, - model: requested_model, - started_at, - stream: details.stream, - compact: details.compact, - transport: details.transport, - error, - }) - .await; - } -} diff --git a/backend/src/dispatch/service/stream.rs b/backend/src/dispatch/service/stream.rs deleted file mode 100644 index 2aa9e3843..000000000 --- a/backend/src/dispatch/service/stream.rs +++ /dev/null @@ -1,619 +0,0 @@ -use super::*; - -impl ResponseDispatchService { - /// 调度流式 Responses 请求到 Codex Responses 上游。 - pub async fn stream( - &self, - request_id: &str, - route: &str, - mut request: CodexResponsesRequest, - requested_model: &str, - ) -> Result { - let started_at = Instant::now(); - let catalog = self.models.catalog().await; - let display_model = catalog.resolve_model_id(requested_model); - request.set_model(display_model.clone()); - request.set_stream(true); - let tuple_schema = request.tuple_schema.clone(); - let now = Utc::now(); - let mut implicit_resume = self.prepare_response_session(&mut request).await; - let account_affinity = self.account_affinity_for_request(&request, now).await; - let mut acquire_request = AccountAcquireRequest::new(request.model(), now); - if let Some(preferred_account_id) = account_affinity.preferred_account_id() { - acquire_request = acquire_request.with_preferred_account_id(preferred_account_id); - } - let mut excluded_account_ids = Vec::new(); - let mut exhausted_accounts = AccountExhaustionTracker::default(); - let mut next_required_account_id: Option = None; - let mut quota_verify_attempts = 0usize; - let mut trace = ResponseDispatchTrace::default(); - macro_rules! return_stream_dispatch_error { - ($error:expr) => {{ - let error = $error; - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: exhausted_accounts.last_account_id(), - stream: true, - compact: false, - transport: Some(backend_transport_name( - backend_transport_for_response_request(&request), - )), - }, - &error, - ) - .await; - return Err(error); - }}; - ($error:expr, account_id: $account_id:expr, transport: $transport:expr) => {{ - let error = $error; - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - client_api_key_id: request.client_api_key_id.as_deref(), - account_id: $account_id, - stream: true, - compact: false, - transport: $transport, - }, - &error, - ) - .await; - return Err(error); - }}; - } - loop { - let mut attempt_acquire_request = acquire_request - .clone() - .with_exclude_account_ids(excluded_account_ids.iter().cloned()); - if let Some(account_id) = next_required_account_id.take() { - attempt_acquire_request = - attempt_acquire_request.with_required_account_id(account_id); - } - let allow_quota_verify_account_retry = - attempt_acquire_request.required_account_id.is_none(); - attempt_acquire_request.now = Utc::now(); - let Some(acquired) = self - .account_pool - .acquire_with(&attempt_acquire_request) - .await - else { - let error = exhausted_accounts - .last_exhausted() - .map(ResponseDispatchError::from_exhausted_account) - .unwrap_or(ResponseDispatchError::NoActiveAccount); - return_stream_dispatch_error!(error); - }; - let acquired_account_id = acquired.account.id.clone(); - let acquired = match verify_acquired_quota_if_required( - QuotaVerificationContext { - account_pool: self.account_pool.as_ref(), - codex: self.codex.as_ref(), - cloudflare: &self.cloudflare, - installation_id: self.installation_id.as_deref(), - request_id, - excluded_account_ids: &mut excluded_account_ids, - verify_attempts: &mut quota_verify_attempts, - allow_retry_with_another_account: allow_quota_verify_account_retry, - }, - acquired, - ) - .await - { - QuotaVerificationDecision::Ready(acquired) => *acquired, - QuotaVerificationDecision::RetryWithAnotherAccount => { - exhausted_accounts - .record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE); - continue; - } - QuotaVerificationDecision::MaxAttemptsReached - | QuotaVerificationDecision::RequiredAccountUnavailable => { - exhausted_accounts.record_rate_limited( - Some(&acquired_account_id), - QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, - ); - let error = exhausted_accounts - .last_exhausted() - .map(ResponseDispatchError::from_exhausted_account) - .unwrap_or(ResponseDispatchError::NoActiveAccount); - return_stream_dispatch_error!( - error, - account_id: Some(&acquired_account_id), - transport: Some(backend_transport_name(backend_transport_for_response_request( - &request - ))) - ); - } - }; - - self.apply_cascading_ban_defense( - &mut request, - &mut implicit_resume, - account_affinity.preferred_account_id(), - &acquired.account.id, - ) - .await; - Self::strip_history_if_account_changed( - &mut request, - &mut implicit_resume, - &account_affinity, - &acquired.account.id, - ); - - self.account_pool.wait_for_request_interval(&acquired).await; - let account = acquired.account; - let release_account_id = account.id.clone(); - let attempt = trace.start_attempt(&release_account_id); - let response_result = create_response_stream_with_account_retrying_5xx( - &self.codex, - self.installation_id.as_deref(), - &self.cloudflare, - &request, - request_id, - &account, - ) - .await; - if let Err(error) = &response_result { - self.cloudflare - .capture_set_cookie_headers( - &release_account_id, - upstream_error_set_cookie_headers(error), - ) - .await; - } - - match response_result { - Ok(response) => { - let transport = response.transport; - let set_cookie_headers = response.set_cookie_headers; - let rate_limit_headers = response.rate_limit_headers; - let rate_limit_header_updates = response.rate_limit_header_updates; - let turn_state_update = response.turn_state_update; - let websocket_pool_decision = response.websocket_pool_decision; - let turn_state = response.turn_state; - let diagnostics = response.diagnostics; - self.cloudflare - .capture_set_cookie_headers(&release_account_id, &set_cookie_headers) - .await; - self.account_pool - .sync_passive_rate_limit_headers(&account, &rate_limit_headers) - .await; - let (prefetched, body) = match prefetch_first_sse_chunk(response.body).await { - Ok(prefetched) => prefetched, - Err(ResponseDispatchError::Upstream(error)) - if is_rate_limit_upstream_error(&error) => - { - self.account_pool.release(&release_account_id).await; - exhausted_accounts.record_rate_limited( - Some(&release_account_id), - upstream_error_body(&error), - ); - let cooldown_until = rate_limit_cooldown_until(&error, Utc::now()); - self.account_pool - .mark_quota_limited_until(&release_account_id, cooldown_until) - .await; - excluded_account_ids.push(release_account_id); - continue; - } - Err(ResponseDispatchError::Upstream(error)) - if is_quota_exhausted_upstream_error(&error) => - { - self.account_pool.release(&release_account_id).await; - exhausted_accounts.record_quota_exhausted( - Some(&release_account_id), - upstream_error_body(&error), - ); - self.account_pool - .set_status(&release_account_id, AccountStatus::QuotaExhausted) - .await; - excluded_account_ids.push(release_account_id); - continue; - } - Err(ResponseDispatchError::Upstream(error)) - if is_auth_upstream_error(&error) => - { - self.account_pool.release(&release_account_id).await; - let upstream_error = upstream_error_body(&error); - let account_status = auth_failure_account_status(&error); - exhausted_accounts.record_auth_failure( - Some(&release_account_id), - account_status, - upstream_error, - Some(upstream_error_http_status(&error)), - ); - self.account_pool - .set_status(&release_account_id, account_status) - .await; - excluded_account_ids.push(release_account_id); - continue; - } - Err(ResponseDispatchError::Upstream(error)) - if is_cloudflare_challenge_upstream_error(&error) => - { - self.account_pool.release(&release_account_id).await; - exhausted_accounts.record_cloudflare_challenge( - Some(&release_account_id), - cloudflare_challenge_error_message(), - ); - self.cloudflare - .apply_challenge(self.account_pool.as_ref(), &release_account_id) - .await; - excluded_account_ids.push(release_account_id); - continue; - } - Err(ResponseDispatchError::Upstream(error)) - if is_cloudflare_path_block_upstream_error(&error) => - { - self.account_pool.release(&release_account_id).await; - exhausted_accounts.record_cloudflare_path_blocked( - Some(&release_account_id), - cloudflare_path_block_error_message(), - ); - self.cloudflare - .apply_path_block(self.account_pool.as_ref(), &release_account_id) - .await; - excluded_account_ids.push(release_account_id); - continue; - } - Err(ResponseDispatchError::Upstream(error)) - if is_model_unsupported_upstream_error(&error) => - { - self.account_pool.release(&release_account_id).await; - let upstream_error = upstream_error_body(&error); - if let Some(exhausted) = exhausted_accounts - .model_unsupported_retry_exhausted(upstream_error.clone()) - { - return_stream_dispatch_error!( - ResponseDispatchError::from_exhausted_account(exhausted), - account_id: Some(&release_account_id), - transport: Some(backend_transport_name(transport)) - ); - } - exhausted_accounts.record_model_unsupported( - Some(&release_account_id), - upstream_error, - ); - excluded_account_ids.push(release_account_id); - continue; - } - Err(ResponseDispatchError::Upstream(error)) - if is_banned_upstream_error(&error) => - { - self.account_pool.release(&release_account_id).await; - exhausted_accounts.record_auth_failure( - Some(&release_account_id), - AccountStatus::Banned, - upstream_error_body(&error), - Some(upstream_error_http_status(&error)), - ); - self.account_pool - .set_status(&release_account_id, AccountStatus::Banned) - .await; - excluded_account_ids.push(release_account_id); - continue; - } - Err(error) => { - self.account_pool.release(&release_account_id).await; - if let ResponseDispatchError::Upstream(upstream_error) = &error { - let history_unavailable = - is_history_recovery_upstream_error(upstream_error); - if history_unavailable - && self - .try_recover_implicit_resume( - &mut request, - &mut implicit_resume, - &release_account_id, - client_error_invalid_reasoning_replay(upstream_error), - ) - .await - { - next_required_account_id = Some(release_account_id); - continue; - } - record_response_upstream_error_event( - ResponseUpstreamErrorEventRecord { - recorder: &self.recorder, - request_id, - account_id: &release_account_id, - account_email: account.email.as_deref(), - route, - model: requested_model, - started_at, - stream: true, - transport, - request: &request, - error: upstream_error, - trace: &trace, - attempt: Some(&attempt), - }, - ) - .await; - if history_unavailable { - return Err(ResponseDispatchError::HistoryUnavailable { - upstream_error: upstream_error_body(upstream_error), - }); - } - return Err(error); - } - return_stream_dispatch_error!( - error, - account_id: Some(&release_account_id), - transport: Some(backend_transport_name(transport)) - ); - } - }; - let first_failure = match first_sse_failure(&prefetched) { - Ok(failure) => failure, - Err(error) => { - self.account_pool.release(&release_account_id).await; - return_stream_dispatch_error!( - ResponseDispatchError::InvalidSse(error), - account_id: Some(&release_account_id), - transport: Some(backend_transport_name(transport)) - ); - } - }; - if let Some(failure) = first_failure { - if is_history_recovery_sse_failure(&failure) - && self - .try_recover_implicit_resume( - &mut request, - &mut implicit_resume, - &release_account_id, - sse_failure_invalid_reasoning_replay(&failure), - ) - .await - { - self.account_pool.release(&release_account_id).await; - next_required_account_id = Some(release_account_id); - continue; - } - if is_model_unsupported_sse_failure(&failure) { - let upstream_error = sse_failure_error_body(&failure); - if let Some(exhausted) = exhausted_accounts - .model_unsupported_retry_exhausted(upstream_error.clone()) - { - self.account_pool.release(&release_account_id).await; - return_stream_dispatch_error!( - ResponseDispatchError::from_exhausted_account(exhausted), - account_id: Some(&release_account_id), - transport: Some(backend_transport_name(transport)) - ); - } - exhausted_accounts.record_model_unsupported( - Some(&release_account_id), - upstream_error, - ); - excluded_account_ids.push(release_account_id); - self.account_pool.release(&account.id).await; - continue; - } - if is_quota_exhausted_sse_failure(&failure) { - exhausted_accounts.record_quota_exhausted( - Some(&release_account_id), - failure.message.clone(), - ); - self.account_pool - .set_status(&release_account_id, AccountStatus::QuotaExhausted) - .await; - excluded_account_ids.push(release_account_id); - self.account_pool.release(&account.id).await; - continue; - } - if is_auth_sse_failure(&failure) { - let upstream_error = sse_failure_error_body(&failure); - let account_status = auth_sse_failure_account_status(&failure); - exhausted_accounts.record_auth_failure( - Some(&release_account_id), - account_status, - upstream_error, - Some(stream_failure_http_status(&failure)), - ); - self.account_pool - .set_status(&release_account_id, account_status) - .await; - excluded_account_ids.push(release_account_id); - self.account_pool.release(&account.id).await; - continue; - } - self.account_pool.release(&release_account_id).await; - record_prefetched_response_stream_failure_event( - ResponseStreamFailureEventRecord { - recorder: &self.recorder, - request_id, - account_id: &release_account_id, - route, - model: &display_model, - requested_model, - started_at, - transport, - request: &request, - failure: &failure, - diagnostics: &diagnostics, - rate_limit_headers: &rate_limit_headers, - prefetched: &prefetched, - trace: &trace, - attempt: &attempt, - }, - ) - .await; - let error = if is_history_recovery_sse_failure(&failure) { - ResponseDispatchError::HistoryUnavailable { - upstream_error: sse_failure_error_body(&failure), - } - } else { - ResponseDispatchError::Failed(failure.clone()) - }; - return Err(error); - } - - let context = LiveResponseStreamContext { - account_pool: Arc::clone(&self.account_pool), - session_affinity: Arc::clone(&self.session_affinity), - reasoning_replay: Arc::clone(&self.reasoning_replay), - recorder: Arc::clone(&self.recorder), - cloudflare: self.cloudflare.clone(), - account_id: account.id, - account_plan_type: account.plan_type, - request_id: request_id.to_string(), - route: route.to_string(), - model: request.model().to_string(), - display_model: display_model.clone(), - requested_model: requested_model.to_string(), - client_ip: request.client_ip.clone(), - request, - tuple_schema, - transport, - rate_limit_headers, - rate_limit_header_updates, - turn_state_update, - websocket_pool_decision, - turn_state, - diagnostics, - attempt: attempt.clone(), - attempts: trace.attempts().to_vec(), - started_at, - }; - return Ok(spawn_live_response_stream(context, prefetched, body)); - } - Err(error) if is_rate_limit_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; - exhausted_accounts.record_rate_limited( - Some(&release_account_id), - upstream_error_body(&error), - ); - let cooldown_until = rate_limit_cooldown_until(&error, Utc::now()); - self.account_pool - .mark_quota_limited_until(&release_account_id, cooldown_until) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_quota_exhausted_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; - exhausted_accounts.record_quota_exhausted( - Some(&release_account_id), - upstream_error_body(&error), - ); - self.account_pool - .set_status(&release_account_id, AccountStatus::QuotaExhausted) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_auth_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; - let upstream_error = upstream_error_body(&error); - let account_status = auth_failure_account_status(&error); - exhausted_accounts.record_auth_failure( - Some(&release_account_id), - account_status, - upstream_error, - Some(upstream_error_http_status(&error)), - ); - self.account_pool - .set_status(&release_account_id, account_status) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_cloudflare_challenge_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; - exhausted_accounts.record_cloudflare_challenge( - Some(&release_account_id), - cloudflare_challenge_error_message(), - ); - self.cloudflare - .apply_challenge(self.account_pool.as_ref(), &release_account_id) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_cloudflare_path_block_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; - exhausted_accounts.record_cloudflare_path_blocked( - Some(&release_account_id), - cloudflare_path_block_error_message(), - ); - self.cloudflare - .apply_path_block(self.account_pool.as_ref(), &release_account_id) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_model_unsupported_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; - let upstream_error = upstream_error_body(&error); - if let Some(exhausted) = - exhausted_accounts.model_unsupported_retry_exhausted(upstream_error.clone()) - { - return_stream_dispatch_error!( - ResponseDispatchError::from_exhausted_account(exhausted), - account_id: Some(&release_account_id), - transport: Some(backend_transport_name(backend_transport_for_response_request( - &request - ))) - ); - } - exhausted_accounts - .record_model_unsupported(Some(&release_account_id), upstream_error); - excluded_account_ids.push(release_account_id); - } - Err(error) if is_banned_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; - exhausted_accounts.record_auth_failure( - Some(&release_account_id), - AccountStatus::Banned, - upstream_error_body(&error), - Some(upstream_error_http_status(&error)), - ); - self.account_pool - .set_status(&release_account_id, AccountStatus::Banned) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) => { - self.account_pool.release(&release_account_id).await; - let history_unavailable = is_history_recovery_upstream_error(&error); - if history_unavailable - && self - .try_recover_implicit_resume( - &mut request, - &mut implicit_resume, - &release_account_id, - client_error_invalid_reasoning_replay(&error), - ) - .await - { - next_required_account_id = Some(release_account_id); - continue; - } - record_response_upstream_error_event(ResponseUpstreamErrorEventRecord { - recorder: &self.recorder, - request_id, - account_id: &release_account_id, - account_email: account.email.as_deref(), - route, - model: requested_model, - started_at, - stream: true, - transport: backend_transport_for_response_request(&request), - request: &request, - error: &error, - trace: &trace, - attempt: Some(&attempt), - }) - .await; - if history_unavailable { - return Err(ResponseDispatchError::HistoryUnavailable { - upstream_error: upstream_error_body(&error), - }); - } - return Err(ResponseDispatchError::Upstream(error)); - } - } - } - } -} diff --git a/backend/src/dispatch/stream/lifecycle.rs b/backend/src/dispatch/stream/lifecycle.rs index 5abdb0cd6..bf4b42754 100644 --- a/backend/src/dispatch/stream/lifecycle.rs +++ b/backend/src/dispatch/stream/lifecycle.rs @@ -1,260 +1,668 @@ use std::{sync::Arc, time::Instant}; -use serde_json::Value; -use tokio::sync::Mutex; +use chrono::Utc; use crate::{ - accounts::pool::RuntimeAccountPoolService, + accounts::{account::AccountStatus, pool::AccountAcquireRequest}, dispatch::{ - affinity::{ - resolve::{evict_reasoning_replay, record_response_affinity}, - RuntimeSessionAffinityService, + errors::{ + backend_transport_name, is_history_recovery_upstream_error, + is_model_unsupported_upstream_error, is_quota_exhausted_upstream_error, + is_rate_limit_upstream_error, rate_limit_cooldown_until, upstream_error_body, + upstream_error_http_status, upstream_error_set_cookie_headers, ResponseDispatchError, }, - errors::backend_transport_name, recording::{ - insert_first_token_ms, live_response_rate_limit_headers, live_response_turn_state, - record_live_response_stream_event, + record_prefetched_response_stream_failure_event, record_response_upstream_error_event, + ResponseDispatchErrorDetails, ResponseStreamFailureEventRecord, + ResponseUpstreamErrorEventRecord, + }, + recovery::{ + auth::{auth_failure_account_status, is_auth_upstream_error}, + cloudflare::{ + cloudflare_challenge_error_message, cloudflare_path_block_error_message, + is_cloudflare_challenge_upstream_error, is_cloudflare_path_block_upstream_error, + }, + exhaustion::AccountExhaustionTracker, + }, + service::{ResponseDispatchService, ResponseDispatchStream}, + stream::{ + live::{spawn_live_response_stream, LiveResponseStreamContext}, + prefetch::prefetch_first_sse_chunk, + sse_failure::{ + auth_sse_failure_account_status, client_error_invalid_reasoning_replay, + first_sse_failure, is_auth_sse_failure, is_history_recovery_sse_failure, + is_model_unsupported_sse_failure, is_quota_exhausted_sse_failure, + sse_failure_error_body, sse_failure_invalid_reasoning_replay, + stream_failure_http_status, + }, + trace::ResponseDispatchTrace, + }, + upstream_call::{ + create_response_stream_with_account_retrying_5xx, verify_acquired_quota_if_required, + QuotaVerificationContext, QuotaVerificationDecision, + QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, }, - recovery::reasoning_replay::ReasoningReplayCache, }, - infra::time::elapsed_millis_i64, - telemetry::{recorder::Recorder, usage::types::UsageRecordLevel}, upstream::openai::{ - protocol::{ - events::extract_sse_usage, - responses::{response_from_codex_sse, CodexResponsesRequest, CollectedResponse}, - sse::{parse_sse_events, response_failed_sse_event_with_id}, - }, - transport::{ - CodexBackendTransport, CodexRateLimitHeaderUpdates, CodexTurnStateUpdate, - CodexUpstreamDiagnostics, WebSocketPoolDecision, - }, + protocol::responses::CodexResponsesRequest, + transport::{backend_transport_for_response_request, is_banned_upstream_error}, }, }; -use super::{ - sse_failure::{ - sse_failure_invalid_reasoning_replay, status_code_for_stream_failure, - stream_failure_metadata, stream_failure_source, synthetic_stream_disconnected_detail, - STREAM_DISCONNECTED_CODE, STREAM_DISCONNECTED_MESSAGE, - }, - trace::ResponseDispatchAttempt, -}; -use crate::dispatch::recovery::cloudflare::CloudflareRecovery; - -pub(in crate::dispatch) struct LiveResponseStreamContext { - pub(in crate::dispatch) account_pool: Arc, - pub(in crate::dispatch) session_affinity: Arc, - pub(in crate::dispatch) reasoning_replay: Arc>, - pub(in crate::dispatch) recorder: Arc, - pub(in crate::dispatch) cloudflare: CloudflareRecovery, - pub(in crate::dispatch) account_id: String, - pub(in crate::dispatch) account_plan_type: Option, - pub(in crate::dispatch) request_id: String, - pub(in crate::dispatch) route: String, - pub(in crate::dispatch) model: String, - pub(in crate::dispatch) display_model: String, - pub(in crate::dispatch) requested_model: String, - pub(in crate::dispatch) client_ip: Option, - pub(in crate::dispatch) request: CodexResponsesRequest, - pub(in crate::dispatch) tuple_schema: Option, - pub(in crate::dispatch) transport: CodexBackendTransport, - pub(in crate::dispatch) rate_limit_headers: Vec<(String, String)>, - pub(in crate::dispatch) rate_limit_header_updates: Option, - pub(in crate::dispatch) turn_state_update: Option, - pub(in crate::dispatch) websocket_pool_decision: Option, - pub(in crate::dispatch) turn_state: Option, - pub(in crate::dispatch) diagnostics: CodexUpstreamDiagnostics, - pub(in crate::dispatch) attempt: ResponseDispatchAttempt, - pub(in crate::dispatch) attempts: Vec, - pub(in crate::dispatch) started_at: Instant, -} - -pub(in crate::dispatch) fn latest_response_id(body: &str) -> Option { - parse_sse_events(body).ok().and_then(|events| { - events.iter().rev().find_map(|event| { - serde_json::from_str::(&event.data) - .ok() - .and_then(|data| { - data.pointer("/response/id") - .and_then(Value::as_str) - .filter(|id| !id.trim().is_empty()) - .map(ToString::to_string) - }) - }) - }) -} - -pub(in crate::dispatch) fn premature_close_failed_event( - response_id: Option<&str>, - detail: Option<&str>, -) -> String { - let message = match detail.filter(|value| !value.trim().is_empty()) { - Some(detail) => format!("{STREAM_DISCONNECTED_MESSAGE}: {detail}"), - None => STREAM_DISCONNECTED_MESSAGE.to_string(), - }; - response_failed_sse_event_with_id( - response_id, - "server_error", - STREAM_DISCONNECTED_CODE, - &message, - ) -} - -pub(in crate::dispatch) async fn finalize_live_response_stream( - context: LiveResponseStreamContext, - body: String, - first_token_ms: Option, -) { - let rate_limit_headers = live_response_rate_limit_headers(&context).await; - context - .account_pool - .sync_passive_rate_limit_headers_for_account( - &context.account_id, - context.account_plan_type.as_deref(), - &rate_limit_headers, - ) - .await; - let turn_state = live_response_turn_state(&context).await; - let usage = match extract_sse_usage(&body) { - Ok(Some(usage)) => { - context - .account_pool - .record_token_usage(&context.account_id, &context.model, &usage) - .await; - Some(usage) +impl ResponseDispatchService { + /// 调度流式 Responses 请求到 Codex Responses 上游。 + pub async fn stream( + &self, + request_id: &str, + route: &str, + mut request: CodexResponsesRequest, + requested_model: &str, + ) -> Result { + let started_at = Instant::now(); + let catalog = self.models.catalog().await; + let display_model = catalog.resolve_model_id(requested_model); + request.set_model(display_model.clone()); + request.set_stream(true); + let tuple_schema = request.tuple_schema.clone(); + let now = Utc::now(); + let mut implicit_resume = self.prepare_response_session(&mut request).await; + let account_affinity = self.account_affinity_for_request(&request, now).await; + let mut acquire_request = AccountAcquireRequest::new(request.model(), now); + if let Some(preferred_account_id) = account_affinity.preferred_account_id() { + acquire_request = acquire_request.with_preferred_account_id(preferred_account_id); } - Ok(None) => None, - Err(error) => { - tracing::warn!(account_id = %context.account_id, error = %error, "failed to extract streaming token usage"); - None - } - }; - - match response_from_codex_sse(&body, context.tuple_schema.as_ref()) { - Ok(CollectedResponse::Completed(completed)) => { - context - .cloudflare - .reset_account_recovery(&context.account_id) - .await; - let response_id = completed.get("id").and_then(Value::as_str); - record_response_affinity( - &context.session_affinity, - &context.reasoning_replay, - &context.request, - &context.account_id, - &body, - turn_state, - usage, - ) - .await; - record_live_response_stream_event( - &context, - 200, - UsageRecordLevel::Info, - "v1 responses stream completed", - serde_json::json!({ - "stream": true, - "completed": true, - "responseId": response_id, - "firstTokenMs": first_token_ms, - "usage": usage, - }), - &rate_limit_headers, - &body, - ) - .await; - } - Ok(CollectedResponse::Failed(failure)) => { - if sse_failure_invalid_reasoning_replay(&failure) { - evict_reasoning_replay( - &context.reasoning_replay, - &context.request, - &context.account_id, + let mut excluded_account_ids = Vec::new(); + let mut exhausted_accounts = AccountExhaustionTracker::default(); + let mut next_required_account_id: Option = None; + let mut quota_verify_attempts = 0usize; + let mut trace = ResponseDispatchTrace::default(); + macro_rules! return_stream_dispatch_error { + ($error:expr) => {{ + let error = $error; + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: exhausted_accounts.last_account_id(), + stream: true, + compact: false, + transport: Some(backend_transport_name( + backend_transport_for_response_request(&request), + )), + }, + &error, ) .await; + return Err(error); + }}; + ($error:expr, account_id: $account_id:expr, transport: $transport:expr) => {{ + let error = $error; + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: $account_id, + stream: true, + compact: false, + transport: $transport, + }, + &error, + ) + .await; + return Err(error); + }}; + } + loop { + let mut attempt_acquire_request = acquire_request + .clone() + .with_exclude_account_ids(excluded_account_ids.iter().cloned()); + if let Some(account_id) = next_required_account_id.take() { + attempt_acquire_request = + attempt_acquire_request.with_required_account_id(account_id); } - let response_id = latest_response_id(&body); - let latency_ms = elapsed_millis_i64(context.started_at); - let failure_source = stream_failure_source(&failure); - let failure_detail = synthetic_stream_disconnected_detail(&failure); - let websocket_pool_kind = context - .websocket_pool_decision - .map(|decision| decision.kind()); - let websocket_pool_reason = context - .websocket_pool_decision - .and_then(|decision| decision.reason()); - tracing::warn!( - account_id = %context.account_id, - request_id = %context.request_id, - response_id = response_id.as_deref().unwrap_or(""), - transport = %backend_transport_name(context.transport), - websocket_pool_kind = ?websocket_pool_kind, - websocket_pool_reason = ?websocket_pool_reason, - first_token_ms = ?first_token_ms, - latency_ms, - event = %failure.event, - code = ?failure.upstream_code.as_deref(), - failure_source = %failure_source, - failure_detail = ?failure_detail.as_deref(), - "live upstream stream ended with response.failed" + let allow_quota_verify_account_retry = + attempt_acquire_request.required_account_id.is_none(); + attempt_acquire_request.now = Utc::now(); + let Some(acquired) = self + .account_pool + .acquire_with(&attempt_acquire_request) + .await + else { + let error = exhausted_accounts + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount); + return_stream_dispatch_error!(error); + }; + let acquired_account_id = acquired.account.id.clone(); + let acquired = match verify_acquired_quota_if_required( + QuotaVerificationContext { + account_pool: self.account_pool.as_ref(), + codex: self.codex.as_ref(), + cloudflare: &self.cloudflare, + installation_id: self.installation_id.as_deref(), + request_id, + excluded_account_ids: &mut excluded_account_ids, + verify_attempts: &mut quota_verify_attempts, + allow_retry_with_another_account: allow_quota_verify_account_retry, + }, + acquired, + ) + .await + { + QuotaVerificationDecision::Ready(acquired) => *acquired, + QuotaVerificationDecision::RetryWithAnotherAccount => { + exhausted_accounts + .record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE); + continue; + } + QuotaVerificationDecision::MaxAttemptsReached + | QuotaVerificationDecision::RequiredAccountUnavailable => { + exhausted_accounts.record_rate_limited( + Some(&acquired_account_id), + QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, + ); + let error = exhausted_accounts + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount); + return_stream_dispatch_error!( + error, + account_id: Some(&acquired_account_id), + transport: Some(backend_transport_name(backend_transport_for_response_request( + &request + ))) + ); + } + }; + + self.apply_cascading_ban_defense( + &mut request, + &mut implicit_resume, + account_affinity.preferred_account_id(), + &acquired.account.id, + ) + .await; + Self::strip_history_if_account_changed( + &mut request, + &mut implicit_resume, + &account_affinity, + &acquired.account.id, ); - let mut metadata = stream_failure_metadata(&failure, usage); - insert_first_token_ms(&mut metadata, first_token_ms); - record_live_response_stream_event( - &context, - status_code_for_stream_failure(&failure), - UsageRecordLevel::Error, - "v1 responses stream failed", - metadata, - &rate_limit_headers, - &body, - ) - .await; - } - Ok(CollectedResponse::MissingCompleted | CollectedResponse::Empty) => { - tracing::warn!( - account_id = %context.account_id, - "live upstream stream ended without response.completed" - ); - let mut metadata = serde_json::json!({ - "stream": true, - "failed": true, - "upstreamCode": "missing_completed", - "usage": usage, - }); - insert_first_token_ms(&mut metadata, first_token_ms); - record_live_response_stream_event( - &context, - 502, - UsageRecordLevel::Error, - "v1 responses stream ended without response.completed", - metadata, - &rate_limit_headers, - &body, - ) - .await; - } - Err(error) => { - tracing::warn!(account_id = %context.account_id, error = %error, "failed to parse completed live stream"); - let mut metadata = serde_json::json!({ - "stream": true, - "sseParseError": error.to_string(), - "usage": usage, - }); - insert_first_token_ms(&mut metadata, first_token_ms); - record_live_response_stream_event( - &context, - 502, - UsageRecordLevel::Warn, - "v1 responses stream SSE response invalid", - metadata, - &rate_limit_headers, - &body, + + self.account_pool.wait_for_request_interval(&acquired).await; + let account = acquired.account; + let release_account_id = account.id.clone(); + let attempt = trace.start_attempt(&release_account_id); + let response_result = create_response_stream_with_account_retrying_5xx( + &self.codex, + self.installation_id.as_deref(), + &self.cloudflare, + &request, + request_id, + &account, ) .await; + if let Err(error) = &response_result { + self.cloudflare + .capture_set_cookie_headers( + &release_account_id, + upstream_error_set_cookie_headers(error), + ) + .await; + } + + match response_result { + Ok(response) => { + let transport = response.transport; + let set_cookie_headers = response.set_cookie_headers; + let rate_limit_headers = response.rate_limit_headers; + let rate_limit_header_updates = response.rate_limit_header_updates; + let turn_state_update = response.turn_state_update; + let websocket_pool_decision = response.websocket_pool_decision; + let turn_state = response.turn_state; + let diagnostics = response.diagnostics; + self.cloudflare + .capture_set_cookie_headers(&release_account_id, &set_cookie_headers) + .await; + self.account_pool + .sync_passive_rate_limit_headers(&account, &rate_limit_headers) + .await; + let (prefetched, body) = match prefetch_first_sse_chunk(response.body).await { + Ok(prefetched) => prefetched, + Err(ResponseDispatchError::Upstream(error)) + if is_rate_limit_upstream_error(&error) => + { + self.account_pool.release(&release_account_id).await; + exhausted_accounts.record_rate_limited( + Some(&release_account_id), + upstream_error_body(&error), + ); + let cooldown_until = rate_limit_cooldown_until(&error, Utc::now()); + self.account_pool + .mark_quota_limited_until(&release_account_id, cooldown_until) + .await; + excluded_account_ids.push(release_account_id); + continue; + } + Err(ResponseDispatchError::Upstream(error)) + if is_quota_exhausted_upstream_error(&error) => + { + self.account_pool.release(&release_account_id).await; + exhausted_accounts.record_quota_exhausted( + Some(&release_account_id), + upstream_error_body(&error), + ); + self.account_pool + .set_status(&release_account_id, AccountStatus::QuotaExhausted) + .await; + excluded_account_ids.push(release_account_id); + continue; + } + Err(ResponseDispatchError::Upstream(error)) + if is_auth_upstream_error(&error) => + { + self.account_pool.release(&release_account_id).await; + let upstream_error = upstream_error_body(&error); + let account_status = auth_failure_account_status(&error); + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + account_status, + upstream_error, + Some(upstream_error_http_status(&error)), + ); + self.account_pool + .set_status(&release_account_id, account_status) + .await; + excluded_account_ids.push(release_account_id); + continue; + } + Err(ResponseDispatchError::Upstream(error)) + if is_cloudflare_challenge_upstream_error(&error) => + { + self.account_pool.release(&release_account_id).await; + exhausted_accounts.record_cloudflare_challenge( + Some(&release_account_id), + cloudflare_challenge_error_message(), + ); + self.cloudflare + .apply_challenge(self.account_pool.as_ref(), &release_account_id) + .await; + excluded_account_ids.push(release_account_id); + continue; + } + Err(ResponseDispatchError::Upstream(error)) + if is_cloudflare_path_block_upstream_error(&error) => + { + self.account_pool.release(&release_account_id).await; + exhausted_accounts.record_cloudflare_path_blocked( + Some(&release_account_id), + cloudflare_path_block_error_message(), + ); + self.cloudflare + .apply_path_block(self.account_pool.as_ref(), &release_account_id) + .await; + excluded_account_ids.push(release_account_id); + continue; + } + Err(ResponseDispatchError::Upstream(error)) + if is_model_unsupported_upstream_error(&error) => + { + self.account_pool.release(&release_account_id).await; + let upstream_error = upstream_error_body(&error); + if let Some(exhausted) = exhausted_accounts + .model_unsupported_retry_exhausted(upstream_error.clone()) + { + return_stream_dispatch_error!( + ResponseDispatchError::from_exhausted_account(exhausted), + account_id: Some(&release_account_id), + transport: Some(backend_transport_name(transport)) + ); + } + exhausted_accounts.record_model_unsupported( + Some(&release_account_id), + upstream_error, + ); + excluded_account_ids.push(release_account_id); + continue; + } + Err(ResponseDispatchError::Upstream(error)) + if is_banned_upstream_error(&error) => + { + self.account_pool.release(&release_account_id).await; + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + AccountStatus::Banned, + upstream_error_body(&error), + Some(upstream_error_http_status(&error)), + ); + self.account_pool + .set_status(&release_account_id, AccountStatus::Banned) + .await; + excluded_account_ids.push(release_account_id); + continue; + } + Err(error) => { + self.account_pool.release(&release_account_id).await; + if let ResponseDispatchError::Upstream(upstream_error) = &error { + let history_unavailable = + is_history_recovery_upstream_error(upstream_error); + if history_unavailable + && self + .try_recover_implicit_resume( + &mut request, + &mut implicit_resume, + &release_account_id, + client_error_invalid_reasoning_replay(upstream_error), + ) + .await + { + next_required_account_id = Some(release_account_id); + continue; + } + record_response_upstream_error_event( + ResponseUpstreamErrorEventRecord { + recorder: &self.recorder, + request_id, + account_id: &release_account_id, + account_email: account.email.as_deref(), + route, + model: requested_model, + started_at, + stream: true, + transport, + request: &request, + error: upstream_error, + trace: &trace, + attempt: Some(&attempt), + }, + ) + .await; + if history_unavailable { + return Err(ResponseDispatchError::HistoryUnavailable { + upstream_error: upstream_error_body(upstream_error), + }); + } + return Err(error); + } + return_stream_dispatch_error!( + error, + account_id: Some(&release_account_id), + transport: Some(backend_transport_name(transport)) + ); + } + }; + let first_failure = match first_sse_failure(&prefetched) { + Ok(failure) => failure, + Err(error) => { + self.account_pool.release(&release_account_id).await; + return_stream_dispatch_error!( + ResponseDispatchError::InvalidSse(error), + account_id: Some(&release_account_id), + transport: Some(backend_transport_name(transport)) + ); + } + }; + if let Some(failure) = first_failure { + if is_history_recovery_sse_failure(&failure) + && self + .try_recover_implicit_resume( + &mut request, + &mut implicit_resume, + &release_account_id, + sse_failure_invalid_reasoning_replay(&failure), + ) + .await + { + self.account_pool.release(&release_account_id).await; + next_required_account_id = Some(release_account_id); + continue; + } + if is_model_unsupported_sse_failure(&failure) { + let upstream_error = sse_failure_error_body(&failure); + if let Some(exhausted) = exhausted_accounts + .model_unsupported_retry_exhausted(upstream_error.clone()) + { + self.account_pool.release(&release_account_id).await; + return_stream_dispatch_error!( + ResponseDispatchError::from_exhausted_account(exhausted), + account_id: Some(&release_account_id), + transport: Some(backend_transport_name(transport)) + ); + } + exhausted_accounts.record_model_unsupported( + Some(&release_account_id), + upstream_error, + ); + excluded_account_ids.push(release_account_id); + self.account_pool.release(&account.id).await; + continue; + } + if is_quota_exhausted_sse_failure(&failure) { + exhausted_accounts.record_quota_exhausted( + Some(&release_account_id), + failure.message.clone(), + ); + self.account_pool + .set_status(&release_account_id, AccountStatus::QuotaExhausted) + .await; + excluded_account_ids.push(release_account_id); + self.account_pool.release(&account.id).await; + continue; + } + if is_auth_sse_failure(&failure) { + let upstream_error = sse_failure_error_body(&failure); + let account_status = auth_sse_failure_account_status(&failure); + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + account_status, + upstream_error, + Some(stream_failure_http_status(&failure)), + ); + self.account_pool + .set_status(&release_account_id, account_status) + .await; + excluded_account_ids.push(release_account_id); + self.account_pool.release(&account.id).await; + continue; + } + self.account_pool.release(&release_account_id).await; + record_prefetched_response_stream_failure_event( + ResponseStreamFailureEventRecord { + recorder: &self.recorder, + request_id, + account_id: &release_account_id, + route, + model: &display_model, + requested_model, + started_at, + transport, + request: &request, + failure: &failure, + diagnostics: &diagnostics, + rate_limit_headers: &rate_limit_headers, + prefetched: &prefetched, + trace: &trace, + attempt: &attempt, + }, + ) + .await; + let error = if is_history_recovery_sse_failure(&failure) { + ResponseDispatchError::HistoryUnavailable { + upstream_error: sse_failure_error_body(&failure), + } + } else { + ResponseDispatchError::Failed(failure.clone()) + }; + return Err(error); + } + + let context = LiveResponseStreamContext { + account_pool: Arc::clone(&self.account_pool), + session_affinity: Arc::clone(&self.session_affinity), + reasoning_replay: Arc::clone(&self.reasoning_replay), + recorder: Arc::clone(&self.recorder), + cloudflare: self.cloudflare.clone(), + account_id: account.id, + account_plan_type: account.plan_type, + request_id: request_id.to_string(), + route: route.to_string(), + model: request.model().to_string(), + display_model: display_model.clone(), + requested_model: requested_model.to_string(), + client_ip: request.client_ip.clone(), + request, + tuple_schema, + transport, + rate_limit_headers, + rate_limit_header_updates, + turn_state_update, + websocket_pool_decision, + turn_state, + diagnostics, + attempt: attempt.clone(), + attempts: trace.attempts().to_vec(), + started_at, + }; + return Ok(spawn_live_response_stream(context, prefetched, body)); + } + Err(error) if is_rate_limit_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + exhausted_accounts.record_rate_limited( + Some(&release_account_id), + upstream_error_body(&error), + ); + let cooldown_until = rate_limit_cooldown_until(&error, Utc::now()); + self.account_pool + .mark_quota_limited_until(&release_account_id, cooldown_until) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_quota_exhausted_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + exhausted_accounts.record_quota_exhausted( + Some(&release_account_id), + upstream_error_body(&error), + ); + self.account_pool + .set_status(&release_account_id, AccountStatus::QuotaExhausted) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_auth_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + let upstream_error = upstream_error_body(&error); + let account_status = auth_failure_account_status(&error); + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + account_status, + upstream_error, + Some(upstream_error_http_status(&error)), + ); + self.account_pool + .set_status(&release_account_id, account_status) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_cloudflare_challenge_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + exhausted_accounts.record_cloudflare_challenge( + Some(&release_account_id), + cloudflare_challenge_error_message(), + ); + self.cloudflare + .apply_challenge(self.account_pool.as_ref(), &release_account_id) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_cloudflare_path_block_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + exhausted_accounts.record_cloudflare_path_blocked( + Some(&release_account_id), + cloudflare_path_block_error_message(), + ); + self.cloudflare + .apply_path_block(self.account_pool.as_ref(), &release_account_id) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_model_unsupported_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + let upstream_error = upstream_error_body(&error); + if let Some(exhausted) = + exhausted_accounts.model_unsupported_retry_exhausted(upstream_error.clone()) + { + return_stream_dispatch_error!( + ResponseDispatchError::from_exhausted_account(exhausted), + account_id: Some(&release_account_id), + transport: Some(backend_transport_name(backend_transport_for_response_request( + &request + ))) + ); + } + exhausted_accounts + .record_model_unsupported(Some(&release_account_id), upstream_error); + excluded_account_ids.push(release_account_id); + } + Err(error) if is_banned_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + AccountStatus::Banned, + upstream_error_body(&error), + Some(upstream_error_http_status(&error)), + ); + self.account_pool + .set_status(&release_account_id, AccountStatus::Banned) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) => { + self.account_pool.release(&release_account_id).await; + let history_unavailable = is_history_recovery_upstream_error(&error); + if history_unavailable + && self + .try_recover_implicit_resume( + &mut request, + &mut implicit_resume, + &release_account_id, + client_error_invalid_reasoning_replay(&error), + ) + .await + { + next_required_account_id = Some(release_account_id); + continue; + } + record_response_upstream_error_event(ResponseUpstreamErrorEventRecord { + recorder: &self.recorder, + request_id, + account_id: &release_account_id, + account_email: account.email.as_deref(), + route, + model: requested_model, + started_at, + stream: true, + transport: backend_transport_for_response_request(&request), + request: &request, + error: &error, + trace: &trace, + attempt: Some(&attempt), + }) + .await; + if history_unavailable { + return Err(ResponseDispatchError::HistoryUnavailable { + upstream_error: upstream_error_body(&error), + }); + } + return Err(ResponseDispatchError::Upstream(error)); + } + } } } - - context.account_pool.release(&context.account_id).await; } diff --git a/backend/src/dispatch/stream/live.rs b/backend/src/dispatch/stream/live.rs index 105dece8c..adfb4e9b3 100644 --- a/backend/src/dispatch/stream/live.rs +++ b/backend/src/dispatch/stream/live.rs @@ -1,37 +1,61 @@ use std::{ pin::Pin, + sync::Arc, task::{Context, Poll}, + time::Instant, }; use bytes::Bytes; use futures::{stream::Stream, StreamExt}; use serde_json::Value; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::{mpsc, oneshot, Mutex}; use crate::{ + accounts::pool::RuntimeAccountPoolService, dispatch::{ - errors::ResponseDispatchStreamError, - service::ResponseDispatchStream, - stream::lifecycle::{ - finalize_live_response_stream, latest_response_id, premature_close_failed_event, - LiveResponseStreamContext, + affinity::{ + resolve::{evict_reasoning_replay, record_response_affinity}, + RuntimeSessionAffinityService, }, + errors::{backend_transport_name, ResponseDispatchStreamError}, + recording::{ + insert_first_token_ms, live_response_rate_limit_headers, live_response_turn_state, + record_live_response_stream_event, + }, + recovery::{cloudflare::CloudflareRecovery, reasoning_replay::ReasoningReplayCache}, + service::ResponseDispatchStream, }, + infra::time::elapsed_millis_i64, + telemetry::{recorder::Recorder, usage::types::UsageRecordLevel}, upstream::openai::{ protocol::{ + events::extract_sse_usage, responses::{ - reconvert_responses_sse_event_tuple_values, response_sse_event_is_terminal, - update_first_response_output_ms, + reconvert_responses_sse_event_tuple_values, response_from_codex_sse, + response_sse_event_is_terminal, update_first_response_output_ms, + CodexResponsesRequest, CollectedResponse, }, sse::{ - encode_sse_event, parse_sse_events, sse_body_has_done, sse_frame_end, - DONE_SSE_FRAME, + encode_sse_event, parse_sse_events, response_failed_sse_event_with_id, + sse_body_has_done, sse_frame_end, DONE_SSE_FRAME, }, }, - transport::CodexBackendSseStream, + transport::{ + CodexBackendSseStream, CodexBackendTransport, CodexRateLimitHeaderUpdates, + CodexTurnStateUpdate, CodexUpstreamDiagnostics, WebSocketPoolDecision, + }, }, }; +use super::{ + sse_failure::{ + sse_failure_invalid_reasoning_replay, status_code_for_stream_failure, + stream_failure_metadata, stream_failure_source, synthetic_stream_disconnected_detail, + STREAM_DISCONNECTED_CODE, STREAM_DISCONNECTED_MESSAGE, + }, + trace::ResponseDispatchAttempt, +}; + struct MpscResponseBodyStream { receiver: mpsc::Receiver>, cancel: Option>, @@ -313,3 +337,220 @@ fn missing_sse_event_separator(body: &str) -> Option<&'static str> { Some("\n\n") } } + +pub(in crate::dispatch) struct LiveResponseStreamContext { + pub(in crate::dispatch) account_pool: Arc, + pub(in crate::dispatch) session_affinity: Arc, + pub(in crate::dispatch) reasoning_replay: Arc>, + pub(in crate::dispatch) recorder: Arc, + pub(in crate::dispatch) cloudflare: CloudflareRecovery, + pub(in crate::dispatch) account_id: String, + pub(in crate::dispatch) account_plan_type: Option, + pub(in crate::dispatch) request_id: String, + pub(in crate::dispatch) route: String, + pub(in crate::dispatch) model: String, + pub(in crate::dispatch) display_model: String, + pub(in crate::dispatch) requested_model: String, + pub(in crate::dispatch) client_ip: Option, + pub(in crate::dispatch) request: CodexResponsesRequest, + pub(in crate::dispatch) tuple_schema: Option, + pub(in crate::dispatch) transport: CodexBackendTransport, + pub(in crate::dispatch) rate_limit_headers: Vec<(String, String)>, + pub(in crate::dispatch) rate_limit_header_updates: Option, + pub(in crate::dispatch) turn_state_update: Option, + pub(in crate::dispatch) websocket_pool_decision: Option, + pub(in crate::dispatch) turn_state: Option, + pub(in crate::dispatch) diagnostics: CodexUpstreamDiagnostics, + pub(in crate::dispatch) attempt: ResponseDispatchAttempt, + pub(in crate::dispatch) attempts: Vec, + pub(in crate::dispatch) started_at: Instant, +} + +pub(in crate::dispatch) fn latest_response_id(body: &str) -> Option { + parse_sse_events(body).ok().and_then(|events| { + events.iter().rev().find_map(|event| { + serde_json::from_str::(&event.data) + .ok() + .and_then(|data| { + data.pointer("/response/id") + .and_then(Value::as_str) + .filter(|id| !id.trim().is_empty()) + .map(ToString::to_string) + }) + }) + }) +} + +pub(in crate::dispatch) fn premature_close_failed_event( + response_id: Option<&str>, + detail: Option<&str>, +) -> String { + let message = match detail.filter(|value| !value.trim().is_empty()) { + Some(detail) => format!("{STREAM_DISCONNECTED_MESSAGE}: {detail}"), + None => STREAM_DISCONNECTED_MESSAGE.to_string(), + }; + response_failed_sse_event_with_id( + response_id, + "server_error", + STREAM_DISCONNECTED_CODE, + &message, + ) +} + +pub(in crate::dispatch) async fn finalize_live_response_stream( + context: LiveResponseStreamContext, + body: String, + first_token_ms: Option, +) { + let rate_limit_headers = live_response_rate_limit_headers(&context).await; + context + .account_pool + .sync_passive_rate_limit_headers_for_account( + &context.account_id, + context.account_plan_type.as_deref(), + &rate_limit_headers, + ) + .await; + let turn_state = live_response_turn_state(&context).await; + let usage = match extract_sse_usage(&body) { + Ok(Some(usage)) => { + context + .account_pool + .record_token_usage(&context.account_id, &context.model, &usage) + .await; + Some(usage) + } + Ok(None) => None, + Err(error) => { + tracing::warn!(account_id = %context.account_id, error = %error, "failed to extract streaming token usage"); + None + } + }; + + match response_from_codex_sse(&body, context.tuple_schema.as_ref()) { + Ok(CollectedResponse::Completed(completed)) => { + context + .cloudflare + .reset_account_recovery(&context.account_id) + .await; + let response_id = completed.get("id").and_then(Value::as_str); + record_response_affinity( + &context.session_affinity, + &context.reasoning_replay, + &context.request, + &context.account_id, + &body, + turn_state, + usage, + ) + .await; + record_live_response_stream_event( + &context, + 200, + UsageRecordLevel::Info, + "v1 responses stream completed", + serde_json::json!({ + "stream": true, + "completed": true, + "responseId": response_id, + "firstTokenMs": first_token_ms, + "usage": usage, + }), + &rate_limit_headers, + &body, + ) + .await; + } + Ok(CollectedResponse::Failed(failure)) => { + if sse_failure_invalid_reasoning_replay(&failure) { + evict_reasoning_replay( + &context.reasoning_replay, + &context.request, + &context.account_id, + ) + .await; + } + let response_id = latest_response_id(&body); + let latency_ms = elapsed_millis_i64(context.started_at); + let failure_source = stream_failure_source(&failure); + let failure_detail = synthetic_stream_disconnected_detail(&failure); + let websocket_pool_kind = context + .websocket_pool_decision + .map(|decision| decision.kind()); + let websocket_pool_reason = context + .websocket_pool_decision + .and_then(|decision| decision.reason()); + tracing::warn!( + account_id = %context.account_id, + request_id = %context.request_id, + response_id = response_id.as_deref().unwrap_or(""), + transport = %backend_transport_name(context.transport), + websocket_pool_kind = ?websocket_pool_kind, + websocket_pool_reason = ?websocket_pool_reason, + first_token_ms = ?first_token_ms, + latency_ms, + event = %failure.event, + code = ?failure.upstream_code.as_deref(), + failure_source = %failure_source, + failure_detail = ?failure_detail.as_deref(), + "live upstream stream ended with response.failed" + ); + let mut metadata = stream_failure_metadata(&failure, usage); + insert_first_token_ms(&mut metadata, first_token_ms); + record_live_response_stream_event( + &context, + status_code_for_stream_failure(&failure), + UsageRecordLevel::Error, + "v1 responses stream failed", + metadata, + &rate_limit_headers, + &body, + ) + .await; + } + Ok(CollectedResponse::MissingCompleted | CollectedResponse::Empty) => { + tracing::warn!( + account_id = %context.account_id, + "live upstream stream ended without response.completed" + ); + let mut metadata = serde_json::json!({ + "stream": true, + "failed": true, + "upstreamCode": "missing_completed", + "usage": usage, + }); + insert_first_token_ms(&mut metadata, first_token_ms); + record_live_response_stream_event( + &context, + 502, + UsageRecordLevel::Error, + "v1 responses stream ended without response.completed", + metadata, + &rate_limit_headers, + &body, + ) + .await; + } + Err(error) => { + tracing::warn!(account_id = %context.account_id, error = %error, "failed to parse completed live stream"); + let mut metadata = serde_json::json!({ + "stream": true, + "sseParseError": error.to_string(), + "usage": usage, + }); + insert_first_token_ms(&mut metadata, first_token_ms); + record_live_response_stream_event( + &context, + 502, + UsageRecordLevel::Warn, + "v1 responses stream SSE response invalid", + metadata, + &rate_limit_headers, + &body, + ) + .await; + } + } + + context.account_pool.release(&context.account_id).await; +} diff --git a/backend/src/dispatch/upstream_call.rs b/backend/src/dispatch/upstream_call.rs index e8b0d321c..f45e5fb0a 100644 --- a/backend/src/dispatch/upstream_call.rs +++ b/backend/src/dispatch/upstream_call.rs @@ -2,26 +2,56 @@ use std::{future::Future, time::Instant}; -use serde_json::Value; +use chrono::Utc; +use serde_json::{json, Value}; use crate::{ accounts::{ - account::Account, - pool::{AcquiredAccount, RuntimeAccountPoolService}, + account::{Account, AccountStatus}, + pool::{AccountAcquireRequest, AcquiredAccount, RuntimeAccountPoolService}, quota::{ quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds, quota_snapshot_reset_at, }, }, dispatch::{ - affinity::build_conversation_identity, errors::is_retryable_upstream_5xx_error, - recovery::cloudflare::CloudflareRecovery, + affinity::build_conversation_identity, + errors::{ + is_model_unsupported_upstream_error, is_quota_exhausted_upstream_error, + is_rate_limit_upstream_error, is_retryable_upstream_5xx_error, + rate_limit_cooldown_until, upstream_error_body, upstream_error_http_status, + upstream_error_set_cookie_headers, ResponseDispatchError, + }, + recording::{ + insert_response_status_metadata, insert_response_trace_metadata, + insert_response_upstream_diagnostics, record_response_dispatch_error_event, + ResponseDispatchErrorEventRecord, + }, + recovery::{ + auth::{auth_failure_account_status, is_auth_upstream_error}, + cloudflare::{ + cloudflare_challenge_error_message, cloudflare_path_block_error_message, + is_cloudflare_challenge_upstream_error, is_cloudflare_path_block_upstream_error, + CloudflareRecovery, + }, + exhaustion::AccountExhaustionTracker, + }, + service::ResponseDispatchService, + stream::trace::ResponseDispatchTrace, + }, + telemetry::{ + recorder::{reasoning_effort_from_compact_request, record_response_event}, + usage::types::ResponseUsageRecord, }, upstream::openai::{ - protocol::responses::{CodexCompactRequest, CodexResponsesRequest}, + protocol::{ + events::extract_usage, + responses::{CodexCompactRequest, CodexResponsesRequest}, + }, transport::{ - CodexBackendClient, CodexBackendResponse, CodexBackendStreamingResponse, - CodexClientError, CodexCompactResponse, CodexRequestContext, + is_banned_upstream_error, CodexBackendClient, CodexBackendResponse, + CodexBackendStreamingResponse, CodexClientError, CodexCompactResponse, + CodexRequestContext, }, }, }; @@ -449,3 +479,292 @@ pub(crate) async fn create_compact_response_with_account_retrying_5xx( ) .await } + +impl ResponseDispatchService { + /// 调度 Responses compact 请求到 Codex compact 上游。 + pub async fn compact( + &self, + request_id: &str, + mut request: CodexCompactRequest, + requested_model: &str, + ) -> Result { + let started_at = Instant::now(); + let client_api_key_id = request.client_api_key_id.clone(); + let catalog = self.models.catalog().await; + let display_model = catalog.resolve_model_id(requested_model); + request.set_model(display_model.clone()); + let mut excluded_account_ids = Vec::new(); + let mut exhausted_accounts = AccountExhaustionTracker::default(); + let mut quota_verify_attempts = 0usize; + let mut trace = ResponseDispatchTrace::default(); + + loop { + let acquire_request = AccountAcquireRequest::new(request.model(), Utc::now()) + .with_exclude_account_ids(excluded_account_ids.iter().cloned()); + let acquired = match self.account_pool.acquire_with(&acquire_request).await { + Some(acquired) => acquired, + None => { + let error = exhausted_accounts + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount); + self.record_compact_dispatch_error( + request_id, + client_api_key_id.as_deref(), + requested_model, + started_at, + exhausted_accounts.last_account_id(), + &error, + ) + .await; + return Err(error); + } + }; + let acquired_account_id = acquired.account.id.clone(); + let acquired = match verify_acquired_quota_if_required( + QuotaVerificationContext { + account_pool: self.account_pool.as_ref(), + codex: self.codex.as_ref(), + cloudflare: &self.cloudflare, + installation_id: self.installation_id.as_deref(), + request_id, + excluded_account_ids: &mut excluded_account_ids, + verify_attempts: &mut quota_verify_attempts, + allow_retry_with_another_account: true, + }, + acquired, + ) + .await + { + QuotaVerificationDecision::Ready(acquired) => *acquired, + QuotaVerificationDecision::RetryWithAnotherAccount => { + exhausted_accounts + .record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE); + continue; + } + QuotaVerificationDecision::MaxAttemptsReached + | QuotaVerificationDecision::RequiredAccountUnavailable => { + exhausted_accounts.record_rate_limited( + Some(&acquired_account_id), + QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, + ); + let error = exhausted_accounts + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount); + self.record_compact_dispatch_error( + request_id, + client_api_key_id.as_deref(), + requested_model, + started_at, + Some(&acquired_account_id), + &error, + ) + .await; + return Err(error); + } + }; + let account = acquired.account; + let release_account_id = account.id.clone(); + let attempt = trace.start_attempt(&release_account_id); + let response_result = create_compact_response_with_account_retrying_5xx( + &self.codex, + self.installation_id.as_deref(), + &self.cloudflare, + &request, + request_id, + &account, + ) + .await; + self.account_pool.release(&release_account_id).await; + if let Err(error) = &response_result { + self.cloudflare + .capture_set_cookie_headers( + &release_account_id, + upstream_error_set_cookie_headers(error), + ) + .await; + } + + match response_result { + Ok(response) => { + self.cloudflare + .capture_set_cookie_headers( + &release_account_id, + &response.set_cookie_headers, + ) + .await; + self.cloudflare.reset_account_recovery(&account.id).await; + self.account_pool + .sync_passive_rate_limit_headers(&account, &response.rate_limit_headers) + .await; + let usage = extract_usage(&response.body); + if let Some(usage) = usage { + self.account_pool + .record_token_usage(&account.id, request.model(), &usage) + .await; + } + let mut metadata = json!({ + "stream": false, + "compact": true, + "usage": usage, + }); + insert_response_status_metadata( + &mut metadata, + 200, + 200, + response.diagnostics.status_code.map(i64::from), + ); + insert_response_upstream_diagnostics(&mut metadata, &response.diagnostics); + insert_response_trace_metadata(&mut metadata, &trace, Some(&attempt)); + record_response_event(ResponseUsageRecord { + recorder: &self.recorder, + request_id, + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: &account.id, + route: "/v1/responses/compact", + model: &display_model, + requested_model: Some(requested_model), + client_ip: request.client_ip.as_deref(), + client_user_agent: request.client_user_agent.as_deref(), + reasoning_effort: reasoning_effort_from_compact_request(&request), + service_tier: None, + started_at, + status_code: 200, + message: "v1 responses compact completed", + metadata, + rate_limit_headers: &response.rate_limit_headers, + }) + .await; + return Ok(response.body); + } + Err(error) if is_rate_limit_upstream_error(&error) => { + exhausted_accounts.record_rate_limited( + Some(&release_account_id), + upstream_error_body(&error), + ); + let cooldown_until = rate_limit_cooldown_until(&error, Utc::now()); + self.account_pool + .mark_quota_limited_until(&release_account_id, cooldown_until) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_quota_exhausted_upstream_error(&error) => { + exhausted_accounts.record_quota_exhausted( + Some(&release_account_id), + upstream_error_body(&error), + ); + self.account_pool + .set_status(&release_account_id, AccountStatus::QuotaExhausted) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_auth_upstream_error(&error) => { + let upstream_error = upstream_error_body(&error); + let account_status = auth_failure_account_status(&error); + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + account_status, + upstream_error, + Some(upstream_error_http_status(&error)), + ); + self.account_pool + .set_status(&release_account_id, account_status) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_cloudflare_challenge_upstream_error(&error) => { + exhausted_accounts.record_cloudflare_challenge( + Some(&release_account_id), + cloudflare_challenge_error_message(), + ); + self.cloudflare + .apply_challenge(self.account_pool.as_ref(), &release_account_id) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_cloudflare_path_block_upstream_error(&error) => { + exhausted_accounts.record_cloudflare_path_blocked( + Some(&release_account_id), + cloudflare_path_block_error_message(), + ); + self.cloudflare + .apply_path_block(self.account_pool.as_ref(), &release_account_id) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_model_unsupported_upstream_error(&error) => { + let upstream_error = upstream_error_body(&error); + if let Some(exhausted) = + exhausted_accounts.model_unsupported_retry_exhausted(upstream_error.clone()) + { + let error = ResponseDispatchError::from_exhausted_account(exhausted); + self.record_compact_dispatch_error( + request_id, + client_api_key_id.as_deref(), + requested_model, + started_at, + Some(&release_account_id), + &error, + ) + .await; + return Err(error); + } + exhausted_accounts + .record_model_unsupported(Some(&release_account_id), upstream_error); + excluded_account_ids.push(release_account_id); + } + Err(error) if is_banned_upstream_error(&error) => { + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + AccountStatus::Banned, + upstream_error_body(&error), + Some(upstream_error_http_status(&error)), + ); + self.account_pool + .set_status(&release_account_id, AccountStatus::Banned) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) => { + let error = ResponseDispatchError::Upstream(error); + self.record_compact_dispatch_error( + request_id, + client_api_key_id.as_deref(), + requested_model, + started_at, + Some(&release_account_id), + &error, + ) + .await; + return Err(error); + } + } + } + } + + async fn record_compact_dispatch_error( + &self, + request_id: &str, + client_api_key_id: Option<&str>, + requested_model: &str, + started_at: Instant, + account_id: Option<&str>, + error: &ResponseDispatchError, + ) { + record_response_dispatch_error_event(ResponseDispatchErrorEventRecord { + recorder: &self.recorder, + request_id, + client_api_key_id, + account_id, + route: "/v1/responses/compact", + model: requested_model, + started_at, + stream: false, + compact: true, + transport: Some("http"), + error, + }) + .await; + } +} diff --git a/backend/src/telemetry/buckets/store.rs b/backend/src/telemetry/buckets/store.rs index 26dafa7fd..471bebf43 100644 --- a/backend/src/telemetry/buckets/store.rs +++ b/backend/src/telemetry/buckets/store.rs @@ -28,6 +28,13 @@ pub struct PgRequestBucketStore { pool: PgPool, } +/// 可由事实表重建的请求时间桶范围结果。 +pub(crate) struct RebuiltRequestBuckets { + pub cutoff: DateTime, + pub deleted_rows: u64, + pub rebuilt_rows: u64, +} + impl PgRequestBucketStore { /// 构造请求时间桶存储。 pub fn new(pool: PgPool) -> Self { @@ -133,6 +140,126 @@ on conflict (bucket_start, provider, account_id, model, service_tier) do update .await?; Ok(result.rows_affected()) } + + /// 在单个事务内重建两类事实都仍处于保留期内的请求时间桶。 + pub(crate) async fn rebuild_reconstructible_range( + &self, + ) -> Result { + let mut transaction = self.pool.begin().await?; + sqlx::query( + "lock table usage_records, ops_error_logs, request_time_buckets + in share row exclusive mode", + ) + .execute(&mut *transaction) + .await?; + + let retention_days: i64 = sqlx::query_scalar( + "select least(usage_retention_days, ops_error_retention_days) + from runtime_settings where id = 1", + ) + .fetch_one(&mut *transaction) + .await?; + let cutoff: DateTime = sqlx::query_scalar( + "select to_timestamp( + floor(extract(epoch from now() - ($1::bigint * interval '1 day')) / 900) * 900 + )", + ) + .bind(retention_days) + .fetch_one(&mut *transaction) + .await?; + + let deleted_rows = sqlx::query("delete from request_time_buckets where bucket_start >= $1") + .bind(cutoff) + .execute(&mut *transaction) + .await? + .rows_affected(); + + let rebuilt_rows = sqlx::query( + r" +with facts as ( + select + to_timestamp(floor(extract(epoch from created_at) / 900) * 900) as bucket_start, + coalesce(nullif(trim(provider), ''), '__unknown__') as provider, + coalesce(nullif(trim(account_id), ''), '__unknown__') as account_id, + coalesce(nullif(trim(model), ''), '__unknown__') as model, + coalesce(nullif(trim(service_tier), ''), '__unknown__') as service_tier, + 1::bigint as success_count, + 0::bigint as error_count, + coalesce(input_tokens, 0)::bigint as input_tokens, + coalesce(output_tokens, 0)::bigint as output_tokens, + coalesce(cached_tokens, 0)::bigint as cached_tokens, + coalesce(first_token_ms, 0)::bigint as first_token_latency_sum, + (first_token_ms is not null)::integer::bigint as first_token_latency_count, + coalesce(latency_ms, 0)::bigint as latency_sum, + (latency_ms is not null)::integer::bigint as latency_count, + coalesce(latency_ms, 0)::bigint as max_latency_ms, + latency_ms::bigint as min_latency_ms + from usage_records + where created_at >= $1 + + union all + + select + to_timestamp(floor(extract(epoch from created_at) / 900) * 900), + coalesce(nullif(trim(provider), ''), '__unknown__'), + coalesce(nullif(trim(account_id), ''), '__unknown__'), + coalesce(nullif(trim(model), ''), '__unknown__'), + '__unknown__', + 0::bigint, + 1::bigint, + 0::bigint, + 0::bigint, + 0::bigint, + 0::bigint, + 0::bigint, + 0::bigint, + 0::bigint, + 0::bigint, + null::bigint + from ops_error_logs + where created_at >= $1 +), aggregated as ( + select + bucket_start, provider, account_id, model, service_tier, + sum(success_count) as success_count, + sum(error_count) as error_count, + sum(input_tokens) as input_tokens, + sum(output_tokens) as output_tokens, + sum(cached_tokens) as cached_tokens, + sum(first_token_latency_sum) as first_token_latency_sum, + sum(first_token_latency_count) as first_token_latency_count, + sum(latency_sum) as latency_sum, + sum(latency_count) as latency_count, + max(max_latency_ms) as max_latency_ms, + min(min_latency_ms) as min_latency_ms + from facts + group by bucket_start, provider, account_id, model, service_tier +) +insert into request_time_buckets ( + bucket_start, provider, account_id, model, service_tier, + success_count, error_count, input_tokens, output_tokens, cached_tokens, + first_token_latency_sum, first_token_latency_count, latency_sum, latency_count, + max_latency_ms, min_latency_ms, updated_at +) +select + bucket_start, provider, account_id, model, service_tier, + success_count, error_count, input_tokens, output_tokens, cached_tokens, + first_token_latency_sum, first_token_latency_count, latency_sum, latency_count, + max_latency_ms, min_latency_ms, now() +from aggregated", + ) + .bind(cutoff) + .execute(&mut *transaction) + .await? + .rows_affected(); + + transaction.commit().await?; + Ok(RebuiltRequestBuckets { + cutoff, + deleted_rows, + rebuilt_rows, + }) + } } fn dimension(value: Option<&str>) -> &str { diff --git a/backend/src/telemetry/rebuild.rs b/backend/src/telemetry/rebuild.rs index f5d7bf9fb..954aa398d 100644 --- a/backend/src/telemetry/rebuild.rs +++ b/backend/src/telemetry/rebuild.rs @@ -4,6 +4,8 @@ use chrono::{DateTime, Utc}; use sqlx::PgPool; use thiserror::Error; +use crate::telemetry::buckets::store::{PgRequestBucketStore, PgRequestBucketStoreError}; + /// 桶重算结果。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct RebuildBucketsReport { @@ -16,123 +18,17 @@ pub struct RebuildBucketsReport { #[derive(Debug, Error)] pub enum RebuildBucketsError { #[error("failed to rebuild request time buckets: {0}")] - Database(#[from] sqlx::Error), + Store(#[from] PgRequestBucketStoreError), } /// 删除两类事实都仍在保留期内的桶，并从成功/失败事实重新聚合。 pub async fn rebuild_buckets(pool: &PgPool) -> Result { - let mut transaction = pool.begin().await?; - sqlx::query( - "lock table usage_records, ops_error_logs, request_time_buckets - in share row exclusive mode", - ) - .execute(&mut *transaction) - .await?; - - let retention_days: i64 = sqlx::query_scalar( - "select least(usage_retention_days, ops_error_retention_days) - from runtime_settings where id = 1", - ) - .fetch_one(&mut *transaction) - .await?; - let cutoff: DateTime = sqlx::query_scalar( - "select to_timestamp( - floor(extract(epoch from now() - ($1::bigint * interval '1 day')) / 900) * 900 - )", - ) - .bind(retention_days) - .fetch_one(&mut *transaction) - .await?; - - let deleted_rows = sqlx::query("delete from request_time_buckets where bucket_start >= $1") - .bind(cutoff) - .execute(&mut *transaction) - .await? - .rows_affected(); - - let rebuilt_rows = sqlx::query( - r" -with facts as ( - select - to_timestamp(floor(extract(epoch from created_at) / 900) * 900) as bucket_start, - coalesce(nullif(trim(provider), ''), '__unknown__') as provider, - coalesce(nullif(trim(account_id), ''), '__unknown__') as account_id, - coalesce(nullif(trim(model), ''), '__unknown__') as model, - coalesce(nullif(trim(service_tier), ''), '__unknown__') as service_tier, - 1::bigint as success_count, - 0::bigint as error_count, - coalesce(input_tokens, 0)::bigint as input_tokens, - coalesce(output_tokens, 0)::bigint as output_tokens, - coalesce(cached_tokens, 0)::bigint as cached_tokens, - coalesce(first_token_ms, 0)::bigint as first_token_latency_sum, - (first_token_ms is not null)::integer::bigint as first_token_latency_count, - coalesce(latency_ms, 0)::bigint as latency_sum, - (latency_ms is not null)::integer::bigint as latency_count, - coalesce(latency_ms, 0)::bigint as max_latency_ms, - latency_ms::bigint as min_latency_ms - from usage_records - where created_at >= $1 - - union all - - select - to_timestamp(floor(extract(epoch from created_at) / 900) * 900), - coalesce(nullif(trim(provider), ''), '__unknown__'), - coalesce(nullif(trim(account_id), ''), '__unknown__'), - coalesce(nullif(trim(model), ''), '__unknown__'), - '__unknown__', - 0::bigint, - 1::bigint, - 0::bigint, - 0::bigint, - 0::bigint, - 0::bigint, - 0::bigint, - 0::bigint, - 0::bigint, - 0::bigint, - null::bigint - from ops_error_logs - where created_at >= $1 -), aggregated as ( - select - bucket_start, provider, account_id, model, service_tier, - sum(success_count) as success_count, - sum(error_count) as error_count, - sum(input_tokens) as input_tokens, - sum(output_tokens) as output_tokens, - sum(cached_tokens) as cached_tokens, - sum(first_token_latency_sum) as first_token_latency_sum, - sum(first_token_latency_count) as first_token_latency_count, - sum(latency_sum) as latency_sum, - sum(latency_count) as latency_count, - max(max_latency_ms) as max_latency_ms, - min(min_latency_ms) as min_latency_ms - from facts - group by bucket_start, provider, account_id, model, service_tier -) -insert into request_time_buckets ( - bucket_start, provider, account_id, model, service_tier, - success_count, error_count, input_tokens, output_tokens, cached_tokens, - first_token_latency_sum, first_token_latency_count, latency_sum, latency_count, - max_latency_ms, min_latency_ms, updated_at -) -select - bucket_start, provider, account_id, model, service_tier, - success_count, error_count, input_tokens, output_tokens, cached_tokens, - first_token_latency_sum, first_token_latency_count, latency_sum, latency_count, - max_latency_ms, min_latency_ms, now() -from aggregated", - ) - .bind(cutoff) - .execute(&mut *transaction) - .await? - .rows_affected(); - - transaction.commit().await?; + let rebuilt = PgRequestBucketStore::new(pool.clone()) + .rebuild_reconstructible_range() + .await?; Ok(RebuildBucketsReport { - cutoff, - deleted_rows, - rebuilt_rows, + cutoff: rebuilt.cutoff, + deleted_rows: rebuilt.deleted_rows, + rebuilt_rows: rebuilt.rebuilt_rows, }) } diff --git a/backend/src/telemetry/usage/query.rs b/backend/src/telemetry/usage/query.rs index 3b3c76586..c11205fdb 100644 --- a/backend/src/telemetry/usage/query.rs +++ b/backend/src/telemetry/usage/query.rs @@ -1,18 +1,20 @@ //! 管理端使用记录服务。 -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use chrono::{DateTime, Utc}; +use sqlx::{PgPool, Postgres, QueryBuilder, Row}; use thiserror::Error; use crate::{ - infra::json::{NumberedPage, Page}, + infra::{ + format::optional_nonnegative_i64_to_u64, + json::{NumberedPage, Page}, + }, telemetry::{ billing, usage::store::{ - PgUsageRecordStore, UsageRecordAccountUsage, UsageRecordBreakdown, - UsageRecordEndpointSource, UsageRecordFilter, UsageRecordModelSource, - UsageRecordSummary, UsageRecordTrendPoint, + push_filter, PgUsageRecordStore, PgUsageRecordStoreResult, UsageRecordFilter, }, usage::types::UsageRecord, }, @@ -290,3 +292,434 @@ fn format_service_tier(value: &str) -> String { other => other.to_string(), } } + +pub(super) async fn count_usage_records( + pool: &PgPool, + filter: &UsageRecordFilter, +) -> PgUsageRecordStoreResult { + let mut builder = QueryBuilder::::new("select count(*) from usage_records"); + push_filter(&mut builder, filter, None)?; + let total: i64 = builder.build_query_scalar().fetch_one(pool).await?; + Ok(total.max(0) as u64) +} + +pub(super) async fn usage_summary( + pool: &PgPool, + filter: &UsageRecordFilter, +) -> PgUsageRecordStoreResult { + let mut builder = QueryBuilder::::new( + r" +select + count(*) as total_requests, + coalesce(sum(input_tokens), 0)::bigint as input_tokens, + coalesce(sum(output_tokens), 0)::bigint as output_tokens, + coalesce(sum(cached_tokens), 0)::bigint as cached_tokens, + avg(latency_ms::double precision) as average_latency_ms +from usage_records", + ); + push_filter(&mut builder, filter, None)?; + let row = builder.build().fetch_one(pool).await?; + let input_tokens = nonnegative(row.get("input_tokens")); + let output_tokens = nonnegative(row.get("output_tokens")); + let cached_tokens = nonnegative(row.get("cached_tokens")); + Ok(UsageRecordSummary { + total_requests: nonnegative(row.get("total_requests")), + input_tokens, + output_tokens, + cached_tokens, + total_tokens: input_tokens + output_tokens, + average_latency_ms: row.get("average_latency_ms"), + }) +} + +pub(super) async fn usage_account_usage( + pool: &PgPool, + filter: &UsageRecordFilter, + limit: u32, +) -> PgUsageRecordStoreResult> { + let mut builder = QueryBuilder::::new( + r" +select + account_id, + coalesce(sum(input_tokens), 0)::bigint as input_tokens, + coalesce(sum(output_tokens), 0)::bigint as output_tokens, + max(created_at) as last_used_at +from usage_records", + ); + push_filter(&mut builder, filter, None)?; + builder.push(" group by account_id order by last_used_at desc, account_id limit "); + builder.push_bind(i64::from(limit.clamp(1, 50))); + Ok(builder + .build() + .fetch_all(pool) + .await? + .into_iter() + .map(|row| { + let input_tokens = nonnegative(row.get("input_tokens")); + let output_tokens = nonnegative(row.get("output_tokens")); + UsageRecordAccountUsage { + account_id: row.get("account_id"), + input_tokens, + output_tokens, + total_tokens: input_tokens + output_tokens, + last_used_at: row.get("last_used_at"), + } + }) + .collect()) +} + +pub(super) async fn usage_record_account_email_map( + pool: &PgPool, + items: &[UsageRecord], +) -> PgUsageRecordStoreResult> { + let mut account_ids = items + .iter() + .map(|item| item.account_id.clone()) + .collect::>(); + account_ids.sort_unstable(); + account_ids.dedup(); + if account_ids.is_empty() { + return Ok(HashMap::new()); + } + let mut builder = QueryBuilder::::new("select id, email from accounts where id in ("); + let mut values = builder.separated(", "); + for account_id in account_ids { + values.push_bind(account_id); + } + values.push_unseparated(")"); + Ok(builder + .build() + .fetch_all(pool) + .await? + .into_iter() + .filter_map(|row| { + row.get::, _>("email") + .filter(|email| !email.trim().is_empty()) + .map(|email| (row.get("id"), email)) + }) + .collect()) +} + +pub(super) async fn usage_breakdown( + pool: &PgPool, + filter: &UsageRecordFilter, + group_expression: &str, + limit: u32, +) -> PgUsageRecordStoreResult> { + let mut builder = QueryBuilder::::new("select "); + builder.push(group_expression); + builder.push( + r" as name, + coalesce(nullif(upstream_model, ''), model) as billing_model, + service_tier, + count(*) as request_count, + coalesce(sum(input_tokens), 0)::bigint as input_tokens, + coalesce(sum(output_tokens), 0)::bigint as output_tokens, + coalesce(sum(cached_tokens), 0)::bigint as cached_tokens, + coalesce(sum(latency_ms), 0)::bigint as latency_sum, + count(latency_ms) as latency_count +from usage_records", + ); + push_filter(&mut builder, filter, None)?; + builder.push(" group by name, billing_model, service_tier order by request_count desc limit "); + builder.push_bind(i64::from(limit.clamp(1, 50) * 8)); + + let mut grouped = BTreeMap::::new(); + for row in builder.build().fetch_all(pool).await? { + let name: String = row.get("name"); + let input_tokens = nonnegative(row.get("input_tokens")); + let output_tokens = nonnegative(row.get("output_tokens")); + let cached_tokens = nonnegative(row.get("cached_tokens")); + let request_count = nonnegative(row.get("request_count")); + let latency_sum = nonnegative(row.get("latency_sum")); + let latency_count = nonnegative(row.get("latency_count")); + let cost = usage_breakdown_cost( + input_tokens, + output_tokens, + cached_tokens, + &row.get::("billing_model"), + row.get::, _>("service_tier").as_deref(), + ); + grouped + .entry(name.clone()) + .or_insert_with(|| BreakdownAccumulator::new(name)) + .push( + request_count, + input_tokens, + output_tokens, + cached_tokens, + cost, + latency_sum, + latency_count, + ); + } + let mut items = grouped + .into_values() + .map(BreakdownAccumulator::finish) + .collect::>(); + items.sort_by(|left, right| { + right + .request_count + .cmp(&left.request_count) + .then_with(|| left.name.cmp(&right.name)) + }); + items.truncate(limit.clamp(1, 50) as usize); + Ok(items) +} + +pub(super) async fn usage_trend( + pool: &PgPool, + filter: &UsageRecordFilter, +) -> PgUsageRecordStoreResult> { + let mut builder = QueryBuilder::::new( + r" +select + to_char(created_at at time zone 'Asia/Shanghai', 'YYYY-MM-DD') as date, + coalesce(nullif(upstream_model, ''), model) as billing_model, + service_tier, + coalesce(sum(input_tokens), 0)::bigint as input_tokens, + coalesce(sum(output_tokens), 0)::bigint as output_tokens, + coalesce(sum(cached_tokens), 0)::bigint as cached_tokens, + coalesce(sum(latency_ms), 0)::bigint as latency_sum, + count(latency_ms) as latency_count +from usage_records", + ); + push_filter(&mut builder, filter, None)?; + builder.push(" group by date, billing_model, service_tier order by date asc"); + + let mut days = BTreeMap::::new(); + for row in builder.build().fetch_all(pool).await? { + let date: String = row.get("date"); + let input_tokens = nonnegative(row.get("input_tokens")); + let output_tokens = nonnegative(row.get("output_tokens")); + let cached_tokens = nonnegative(row.get("cached_tokens")); + let cost = usage_breakdown_cost( + input_tokens, + output_tokens, + cached_tokens, + &row.get::("billing_model"), + row.get::, _>("service_tier").as_deref(), + ); + days.entry(date.clone()) + .or_insert_with(|| UsageTrendAccumulator::new(date)) + .push( + input_tokens, + output_tokens, + cached_tokens, + cost, + nonnegative(row.get("latency_sum")), + nonnegative(row.get("latency_count")), + ); + } + let mut points = days + .into_values() + .map(UsageTrendAccumulator::finish) + .collect::>(); + if points.len() > 60 { + points = points.split_off(points.len() - 60); + } + Ok(points) +} + +pub(super) fn usage_record_from_row(row: &sqlx::postgres::PgRow) -> UsageRecord { + UsageRecord { + id: row.get("id"), + request_id: row.get("request_id"), + client_api_key_id: row.get("client_api_key_id"), + kind: row.get("kind"), + provider: row.get("provider"), + account_id: row.get("account_id"), + route: row.get("route"), + model: row.get("model"), + requested_model: row.get("requested_model"), + upstream_model: row.get("upstream_model"), + service_tier: row.get("service_tier"), + status_code: row.get::("status_code").into(), + transport: row.get("transport"), + attempt_index: row.get("attempt_index"), + response_id: row.get("response_id"), + upstream_request_id: row.get("upstream_request_id"), + latency_ms: row.get("latency_ms"), + first_token_ms: row.get("first_token_ms"), + input_tokens: row.get("input_tokens"), + output_tokens: row.get("output_tokens"), + cached_tokens: row.get("cached_tokens"), + reasoning_tokens: row.get("reasoning_tokens"), + message: row.get("message"), + metadata: row + .get::, _>("metadata_json") + .0, + created_at: row.get("created_at"), + } +} + +fn nonnegative(value: Option) -> u64 { + optional_nonnegative_i64_to_u64(value) +} + +fn usage_breakdown_cost( + input_tokens: u64, + output_tokens: u64, + cached_tokens: u64, + model: &str, + service_tier: Option<&str>, +) -> f64 { + billing::calculate_cost( + input_tokens, + output_tokens, + cached_tokens, + model, + service_tier, + ) +} + +struct BreakdownAccumulator { + item: UsageRecordBreakdown, + latency_sum: u64, + latency_count: u64, +} + +impl BreakdownAccumulator { + fn new(name: String) -> Self { + Self { + item: UsageRecordBreakdown { + name, + ..Default::default() + }, + latency_sum: 0, + latency_count: 0, + } + } + + #[allow(clippy::too_many_arguments)] + fn push( + &mut self, + requests: u64, + input: u64, + output: u64, + cached: u64, + cost: f64, + latency_sum: u64, + latency_count: u64, + ) { + self.item.request_count += requests; + self.item.input_tokens += input; + self.item.output_tokens += output; + self.item.cached_tokens += cached; + self.item.total_tokens += input + output; + self.item.cost += cost; + self.item.actual_cost += cost; + self.item.account_cost += cost; + self.latency_sum += latency_sum; + self.latency_count += latency_count; + } + + fn finish(mut self) -> UsageRecordBreakdown { + self.item.average_latency_ms = + (self.latency_count > 0).then(|| self.latency_sum as f64 / self.latency_count as f64); + self.item + } +} + +struct UsageTrendAccumulator { + point: UsageRecordTrendPoint, + latency_sum: u64, + latency_count: u64, +} + +impl UsageTrendAccumulator { + fn new(date: String) -> Self { + Self { + point: UsageRecordTrendPoint { + date, + ..Default::default() + }, + latency_sum: 0, + latency_count: 0, + } + } + + fn push( + &mut self, + input: u64, + output: u64, + cached: u64, + cost: f64, + latency_sum: u64, + latency_count: u64, + ) { + self.point.input_tokens += input; + self.point.output_tokens += output; + self.point.cached_tokens += cached; + self.point.total_tokens += input + output; + self.point.cost += cost; + self.point.actual_cost += cost; + self.latency_sum += latency_sum; + self.latency_count += latency_count; + } + + fn finish(mut self) -> UsageRecordTrendPoint { + self.point.average_latency_ms = + (self.latency_count > 0).then(|| self.latency_sum as f64 / self.latency_count as f64); + self.point + } +} + +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub struct UsageRecordSummary { + pub total_requests: u64, + pub input_tokens: u64, + pub output_tokens: u64, + pub cached_tokens: u64, + pub total_tokens: u64, + pub average_latency_ms: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageRecordAccountUsage { + pub account_id: String, + pub input_tokens: u64, + pub output_tokens: u64, + pub total_tokens: u64, + pub last_used_at: DateTime, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum UsageRecordModelSource { + Requested, + Upstream, + Mapping, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum UsageRecordEndpointSource { + Inbound, + Upstream, + Path, +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub struct UsageRecordBreakdown { + pub name: String, + pub request_count: u64, + pub input_tokens: u64, + pub output_tokens: u64, + pub cached_tokens: u64, + pub total_tokens: u64, + pub cost: f64, + pub actual_cost: f64, + pub account_cost: f64, + pub average_latency_ms: Option, +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub struct UsageRecordTrendPoint { + pub date: String, + pub input_tokens: u64, + pub output_tokens: u64, + pub cached_tokens: u64, + pub total_tokens: u64, + pub cost: f64, + pub actual_cost: f64, + pub average_latency_ms: Option, +} diff --git a/backend/src/telemetry/usage/store.rs b/backend/src/telemetry/usage/store.rs index 2cccdf32e..39c0ee0f9 100644 --- a/backend/src/telemetry/usage/store.rs +++ b/backend/src/telemetry/usage/store.rs @@ -1,29 +1,24 @@ //! PostgreSQL 成功使用事实存储。 -mod query; - -use query::{ +use super::query::{ count_usage_records, usage_account_usage, usage_breakdown, usage_record_account_email_map, - usage_record_from_row, usage_summary, usage_trend, -}; -pub use query::{ - UsageRecordAccountUsage, UsageRecordBreakdown, UsageRecordEndpointSource, - UsageRecordModelSource, UsageRecordSummary, UsageRecordTrendPoint, + usage_record_from_row, usage_summary, usage_trend, UsageRecordAccountUsage, + UsageRecordBreakdown, UsageRecordEndpointSource, UsageRecordModelSource, UsageRecordSummary, + UsageRecordTrendPoint, }; -use std::collections::{BTreeMap, HashMap}; +use std::collections::HashMap; use chrono::{DateTime, Duration, Utc}; -use sqlx::{PgPool, Postgres, QueryBuilder, Row}; +use sqlx::{PgPool, Postgres, QueryBuilder}; use thiserror::Error; use crate::{ infra::{ - format::optional_nonnegative_i64_to_u64, json::{decode_cursor, encode_cursor, page_offset, NumberedPage, Page}, time::parse_rfc3339_utc, }, - telemetry::{billing, buckets::store::PgRequestBucketStore, usage::types::UsageRecord}, + telemetry::{buckets::store::PgRequestBucketStore, usage::types::UsageRecord}, }; pub const USAGE_RECORD_RETENTION_DAYS: i64 = 30; @@ -286,7 +281,7 @@ insert into usage_records ( } } -fn push_filter( +pub(super) fn push_filter( builder: &mut QueryBuilder, filter: &UsageRecordFilter, cursor: Option<&str>, diff --git a/backend/src/telemetry/usage/store/query.rs b/backend/src/telemetry/usage/store/query.rs deleted file mode 100644 index 6f862e33d..000000000 --- a/backend/src/telemetry/usage/store/query.rs +++ /dev/null @@ -1,432 +0,0 @@ -use super::*; - -pub(super) async fn count_usage_records( - pool: &PgPool, - filter: &UsageRecordFilter, -) -> PgUsageRecordStoreResult { - let mut builder = QueryBuilder::::new("select count(*) from usage_records"); - push_filter(&mut builder, filter, None)?; - let total: i64 = builder.build_query_scalar().fetch_one(pool).await?; - Ok(total.max(0) as u64) -} - -pub(super) async fn usage_summary( - pool: &PgPool, - filter: &UsageRecordFilter, -) -> PgUsageRecordStoreResult { - let mut builder = QueryBuilder::::new( - r" -select - count(*) as total_requests, - coalesce(sum(input_tokens), 0)::bigint as input_tokens, - coalesce(sum(output_tokens), 0)::bigint as output_tokens, - coalesce(sum(cached_tokens), 0)::bigint as cached_tokens, - avg(latency_ms::double precision) as average_latency_ms -from usage_records", - ); - push_filter(&mut builder, filter, None)?; - let row = builder.build().fetch_one(pool).await?; - let input_tokens = nonnegative(row.get("input_tokens")); - let output_tokens = nonnegative(row.get("output_tokens")); - let cached_tokens = nonnegative(row.get("cached_tokens")); - Ok(UsageRecordSummary { - total_requests: nonnegative(row.get("total_requests")), - input_tokens, - output_tokens, - cached_tokens, - total_tokens: input_tokens + output_tokens, - average_latency_ms: row.get("average_latency_ms"), - }) -} - -pub(super) async fn usage_account_usage( - pool: &PgPool, - filter: &UsageRecordFilter, - limit: u32, -) -> PgUsageRecordStoreResult> { - let mut builder = QueryBuilder::::new( - r" -select - account_id, - coalesce(sum(input_tokens), 0)::bigint as input_tokens, - coalesce(sum(output_tokens), 0)::bigint as output_tokens, - max(created_at) as last_used_at -from usage_records", - ); - push_filter(&mut builder, filter, None)?; - builder.push(" group by account_id order by last_used_at desc, account_id limit "); - builder.push_bind(i64::from(limit.clamp(1, 50))); - Ok(builder - .build() - .fetch_all(pool) - .await? - .into_iter() - .map(|row| { - let input_tokens = nonnegative(row.get("input_tokens")); - let output_tokens = nonnegative(row.get("output_tokens")); - UsageRecordAccountUsage { - account_id: row.get("account_id"), - input_tokens, - output_tokens, - total_tokens: input_tokens + output_tokens, - last_used_at: row.get("last_used_at"), - } - }) - .collect()) -} - -pub(super) async fn usage_record_account_email_map( - pool: &PgPool, - items: &[UsageRecord], -) -> PgUsageRecordStoreResult> { - let mut account_ids = items - .iter() - .map(|item| item.account_id.clone()) - .collect::>(); - account_ids.sort_unstable(); - account_ids.dedup(); - if account_ids.is_empty() { - return Ok(HashMap::new()); - } - let mut builder = QueryBuilder::::new("select id, email from accounts where id in ("); - let mut values = builder.separated(", "); - for account_id in account_ids { - values.push_bind(account_id); - } - values.push_unseparated(")"); - Ok(builder - .build() - .fetch_all(pool) - .await? - .into_iter() - .filter_map(|row| { - row.get::, _>("email") - .filter(|email| !email.trim().is_empty()) - .map(|email| (row.get("id"), email)) - }) - .collect()) -} - -pub(super) async fn usage_breakdown( - pool: &PgPool, - filter: &UsageRecordFilter, - group_expression: &str, - limit: u32, -) -> PgUsageRecordStoreResult> { - let mut builder = QueryBuilder::::new("select "); - builder.push(group_expression); - builder.push( - r" as name, - coalesce(nullif(upstream_model, ''), model) as billing_model, - service_tier, - count(*) as request_count, - coalesce(sum(input_tokens), 0)::bigint as input_tokens, - coalesce(sum(output_tokens), 0)::bigint as output_tokens, - coalesce(sum(cached_tokens), 0)::bigint as cached_tokens, - coalesce(sum(latency_ms), 0)::bigint as latency_sum, - count(latency_ms) as latency_count -from usage_records", - ); - push_filter(&mut builder, filter, None)?; - builder.push(" group by name, billing_model, service_tier order by request_count desc limit "); - builder.push_bind(i64::from(limit.clamp(1, 50) * 8)); - - let mut grouped = BTreeMap::::new(); - for row in builder.build().fetch_all(pool).await? { - let name: String = row.get("name"); - let input_tokens = nonnegative(row.get("input_tokens")); - let output_tokens = nonnegative(row.get("output_tokens")); - let cached_tokens = nonnegative(row.get("cached_tokens")); - let request_count = nonnegative(row.get("request_count")); - let latency_sum = nonnegative(row.get("latency_sum")); - let latency_count = nonnegative(row.get("latency_count")); - let cost = usage_breakdown_cost( - input_tokens, - output_tokens, - cached_tokens, - &row.get::("billing_model"), - row.get::, _>("service_tier").as_deref(), - ); - grouped - .entry(name.clone()) - .or_insert_with(|| BreakdownAccumulator::new(name)) - .push( - request_count, - input_tokens, - output_tokens, - cached_tokens, - cost, - latency_sum, - latency_count, - ); - } - let mut items = grouped - .into_values() - .map(BreakdownAccumulator::finish) - .collect::>(); - items.sort_by(|left, right| { - right - .request_count - .cmp(&left.request_count) - .then_with(|| left.name.cmp(&right.name)) - }); - items.truncate(limit.clamp(1, 50) as usize); - Ok(items) -} - -pub(super) async fn usage_trend( - pool: &PgPool, - filter: &UsageRecordFilter, -) -> PgUsageRecordStoreResult> { - let mut builder = QueryBuilder::::new( - r" -select - to_char(created_at at time zone 'Asia/Shanghai', 'YYYY-MM-DD') as date, - coalesce(nullif(upstream_model, ''), model) as billing_model, - service_tier, - coalesce(sum(input_tokens), 0)::bigint as input_tokens, - coalesce(sum(output_tokens), 0)::bigint as output_tokens, - coalesce(sum(cached_tokens), 0)::bigint as cached_tokens, - coalesce(sum(latency_ms), 0)::bigint as latency_sum, - count(latency_ms) as latency_count -from usage_records", - ); - push_filter(&mut builder, filter, None)?; - builder.push(" group by date, billing_model, service_tier order by date asc"); - - let mut days = BTreeMap::::new(); - for row in builder.build().fetch_all(pool).await? { - let date: String = row.get("date"); - let input_tokens = nonnegative(row.get("input_tokens")); - let output_tokens = nonnegative(row.get("output_tokens")); - let cached_tokens = nonnegative(row.get("cached_tokens")); - let cost = usage_breakdown_cost( - input_tokens, - output_tokens, - cached_tokens, - &row.get::("billing_model"), - row.get::, _>("service_tier").as_deref(), - ); - days.entry(date.clone()) - .or_insert_with(|| UsageTrendAccumulator::new(date)) - .push( - input_tokens, - output_tokens, - cached_tokens, - cost, - nonnegative(row.get("latency_sum")), - nonnegative(row.get("latency_count")), - ); - } - let mut points = days - .into_values() - .map(UsageTrendAccumulator::finish) - .collect::>(); - if points.len() > 60 { - points = points.split_off(points.len() - 60); - } - Ok(points) -} - -pub(super) fn usage_record_from_row(row: &sqlx::postgres::PgRow) -> UsageRecord { - UsageRecord { - id: row.get("id"), - request_id: row.get("request_id"), - client_api_key_id: row.get("client_api_key_id"), - kind: row.get("kind"), - provider: row.get("provider"), - account_id: row.get("account_id"), - route: row.get("route"), - model: row.get("model"), - requested_model: row.get("requested_model"), - upstream_model: row.get("upstream_model"), - service_tier: row.get("service_tier"), - status_code: row.get::("status_code").into(), - transport: row.get("transport"), - attempt_index: row.get("attempt_index"), - response_id: row.get("response_id"), - upstream_request_id: row.get("upstream_request_id"), - latency_ms: row.get("latency_ms"), - first_token_ms: row.get("first_token_ms"), - input_tokens: row.get("input_tokens"), - output_tokens: row.get("output_tokens"), - cached_tokens: row.get("cached_tokens"), - reasoning_tokens: row.get("reasoning_tokens"), - message: row.get("message"), - metadata: row - .get::, _>("metadata_json") - .0, - created_at: row.get("created_at"), - } -} - -fn nonnegative(value: Option) -> u64 { - optional_nonnegative_i64_to_u64(value) -} - -fn usage_breakdown_cost( - input_tokens: u64, - output_tokens: u64, - cached_tokens: u64, - model: &str, - service_tier: Option<&str>, -) -> f64 { - billing::calculate_cost( - input_tokens, - output_tokens, - cached_tokens, - model, - service_tier, - ) -} - -struct BreakdownAccumulator { - item: UsageRecordBreakdown, - latency_sum: u64, - latency_count: u64, -} - -impl BreakdownAccumulator { - fn new(name: String) -> Self { - Self { - item: UsageRecordBreakdown { - name, - ..Default::default() - }, - latency_sum: 0, - latency_count: 0, - } - } - - #[allow(clippy::too_many_arguments)] - fn push( - &mut self, - requests: u64, - input: u64, - output: u64, - cached: u64, - cost: f64, - latency_sum: u64, - latency_count: u64, - ) { - self.item.request_count += requests; - self.item.input_tokens += input; - self.item.output_tokens += output; - self.item.cached_tokens += cached; - self.item.total_tokens += input + output; - self.item.cost += cost; - self.item.actual_cost += cost; - self.item.account_cost += cost; - self.latency_sum += latency_sum; - self.latency_count += latency_count; - } - - fn finish(mut self) -> UsageRecordBreakdown { - self.item.average_latency_ms = - (self.latency_count > 0).then(|| self.latency_sum as f64 / self.latency_count as f64); - self.item - } -} - -struct UsageTrendAccumulator { - point: UsageRecordTrendPoint, - latency_sum: u64, - latency_count: u64, -} - -impl UsageTrendAccumulator { - fn new(date: String) -> Self { - Self { - point: UsageRecordTrendPoint { - date, - ..Default::default() - }, - latency_sum: 0, - latency_count: 0, - } - } - - fn push( - &mut self, - input: u64, - output: u64, - cached: u64, - cost: f64, - latency_sum: u64, - latency_count: u64, - ) { - self.point.input_tokens += input; - self.point.output_tokens += output; - self.point.cached_tokens += cached; - self.point.total_tokens += input + output; - self.point.cost += cost; - self.point.actual_cost += cost; - self.latency_sum += latency_sum; - self.latency_count += latency_count; - } - - fn finish(mut self) -> UsageRecordTrendPoint { - self.point.average_latency_ms = - (self.latency_count > 0).then(|| self.latency_sum as f64 / self.latency_count as f64); - self.point - } -} - -#[derive(Debug, Clone, Copy, Default, PartialEq)] -pub struct UsageRecordSummary { - pub total_requests: u64, - pub input_tokens: u64, - pub output_tokens: u64, - pub cached_tokens: u64, - pub total_tokens: u64, - pub average_latency_ms: Option, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct UsageRecordAccountUsage { - pub account_id: String, - pub input_tokens: u64, - pub output_tokens: u64, - pub total_tokens: u64, - pub last_used_at: DateTime, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum UsageRecordModelSource { - Requested, - Upstream, - Mapping, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum UsageRecordEndpointSource { - Inbound, - Upstream, - Path, -} - -#[derive(Debug, Clone, Default, PartialEq)] -pub struct UsageRecordBreakdown { - pub name: String, - pub request_count: u64, - pub input_tokens: u64, - pub output_tokens: u64, - pub cached_tokens: u64, - pub total_tokens: u64, - pub cost: f64, - pub actual_cost: f64, - pub account_cost: f64, - pub average_latency_ms: Option, -} - -#[derive(Debug, Clone, Default, PartialEq)] -pub struct UsageRecordTrendPoint { - pub date: String, - pub input_tokens: u64, - pub output_tokens: u64, - pub cached_tokens: u64, - pub total_tokens: u64, - pub cost: f64, - pub actual_cost: f64, - pub average_latency_ms: Option, -} diff --git a/backend/src/upstream/openai/protocol/websocket.rs b/backend/src/upstream/openai/protocol/websocket.rs index 2f2378e2c..922b4eadb 100644 --- a/backend/src/upstream/openai/protocol/websocket.rs +++ b/backend/src/upstream/openai/protocol/websocket.rs @@ -216,8 +216,9 @@ pub fn websocket_incomplete_response_reason(raw: &str) -> Option { .map(ToString::to_string) } -mod validation; -pub use validation::*; +#[path = "websocket_errors.rs"] +mod errors; +pub use errors::*; /// 判断 WebSocket 事件是否会结束当前响应流。 pub fn is_terminal_websocket_event(event: &str) -> bool { diff --git a/backend/src/upstream/openai/protocol/websocket/validation.rs b/backend/src/upstream/openai/protocol/websocket_errors.rs similarity index 100% rename from backend/src/upstream/openai/protocol/websocket/validation.rs rename to backend/src/upstream/openai/protocol/websocket_errors.rs diff --git a/backend/src/upstream/openai/transport/client.rs b/backend/src/upstream/openai/transport/client.rs index cf6ea595e..8bf36e579 100644 --- a/backend/src/upstream/openai/transport/client.rs +++ b/backend/src/upstream/openai/transport/client.rs @@ -337,6 +337,7 @@ pub struct CodexBackendClient { websocket_first_token_timeout: Option, } +#[path = "client_sse.rs"] mod requests; #[async_trait] diff --git a/backend/src/upstream/openai/transport/client/requests.rs b/backend/src/upstream/openai/transport/client_sse.rs similarity index 100% rename from backend/src/upstream/openai/transport/client/requests.rs rename to backend/src/upstream/openai/transport/client_sse.rs diff --git a/backend/src/upstream/openai/transport/websocket.rs b/backend/src/upstream/openai/transport/websocket.rs index c33be4532..3122967f7 100644 --- a/backend/src/upstream/openai/transport/websocket.rs +++ b/backend/src/upstream/openai/transport/websocket.rs @@ -61,18 +61,63 @@ pub const WS_AUDIT_DIR_ENV: &str = "CODEX_PROXY_WS_AUDIT_DIR"; // Types // --------------------------------------------------------------------------- -/// WebSocket 连接适配器。 -mod stream; -mod types; +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CodexWebSocketConnection { + pub(super) endpoint: String, + pub(super) headers: Vec<(String, String)>, +} -pub use stream::execute_response_create_request; -use stream::{ +/// 显式写入 WebSocket audit artifact。 +pub async fn write_websocket_audit_artifact_for_dir( + dir: Option<&Path>, + artifact: &WebSocketAuditArtifact, +) -> io::Result> { + let Some(dir) = dir.filter(|dir| !dir.as_os_str().is_empty()) else { + return Ok(None); + }; + + tokio::fs::create_dir_all(dir).await?; + let path = dir.join(websocket_audit_file_name()); + let body = serde_json::to_vec_pretty(artifact).map_err(io::Error::other)?; + tokio::fs::write(&path, body).await?; + Ok(Some(path)) +} + +/// 按环境变量配置写入 WebSocket audit artifact。 +pub async fn write_websocket_audit_artifact_from_env( + artifact: &WebSocketAuditArtifact, +) -> io::Result> { + let Some(dir) = std::env::var_os(WS_AUDIT_DIR_ENV) + .filter(|value| !value.is_empty()) + .map(PathBuf::from) + else { + return Ok(None); + }; + + write_websocket_audit_artifact_for_dir(Some(&dir), artifact).await +} + +/// Prepared Responses WebSocket request descriptor. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CodexWebSocketRequest { + pub(super) connection: CodexWebSocketConnection, + pub(super) payload_text: String, +} + +#[path = "websocket_frames.rs"] +mod frames; + +use frames::{ audit_header_value, collect_websocket_response, is_first_token_timeout, prefetch_stream_frames_until_output_or_terminal, reusable_websocket_metadata, reused_stream_prefetch_error, stream_websocket_response, websocket_audit_file_name, websocket_connection_metadata, WebSocketStreamPoolReturn, }; -pub use types::*; +pub use frames::{ + execute_response_create_request, CodexWebSocketExchange, CodexWebSocketExchangeError, + CodexWebSocketRateLimitHeaderUpdates, CodexWebSocketSseStream, CodexWebSocketStreamingExchange, + CodexWebSocketTurnStateUpdate, CodexWebSocketUpstreamError, +}; impl CodexWebSocketRequest { /// 返回连接描述。 diff --git a/backend/src/upstream/openai/transport/websocket/types.rs b/backend/src/upstream/openai/transport/websocket/types.rs deleted file mode 100644 index ecaef922c..000000000 --- a/backend/src/upstream/openai/transport/websocket/types.rs +++ /dev/null @@ -1,189 +0,0 @@ -use super::*; - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CodexWebSocketConnection { - pub(super) endpoint: String, - pub(super) headers: Vec<(String, String)>, -} - -/// 显式写入 WebSocket audit artifact。 -pub async fn write_websocket_audit_artifact_for_dir( - dir: Option<&Path>, - artifact: &WebSocketAuditArtifact, -) -> io::Result> { - let Some(dir) = dir.filter(|dir| !dir.as_os_str().is_empty()) else { - return Ok(None); - }; - - tokio::fs::create_dir_all(dir).await?; - let path = dir.join(websocket_audit_file_name()); - let body = serde_json::to_vec_pretty(artifact).map_err(io::Error::other)?; - tokio::fs::write(&path, body).await?; - Ok(Some(path)) -} - -/// 按环境变量配置写入 WebSocket audit artifact。 -pub async fn write_websocket_audit_artifact_from_env( - artifact: &WebSocketAuditArtifact, -) -> io::Result> { - let Some(dir) = std::env::var_os(WS_AUDIT_DIR_ENV) - .filter(|value| !value.is_empty()) - .map(PathBuf::from) - else { - return Ok(None); - }; - - write_websocket_audit_artifact_for_dir(Some(&dir), artifact).await -} - -/// Prepared Responses WebSocket request descriptor. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CodexWebSocketRequest { - pub(super) connection: CodexWebSocketConnection, - pub(super) payload_text: String, -} - -/// Responses WebSocket exchange result. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CodexWebSocketExchange { - /// 完整 SSE 文本。 - pub body: String, - /// 从 SSE 中提取的 usage。 - pub usage: Option, - /// 上游 metadata 帧中的最新 turn state。 - pub turn_state: Option, - /// 上游握手响应里的 `set-cookie` 列表。 - pub set_cookie_headers: Vec, - /// 上游握手响应里的限流头。 - pub rate_limit_headers: Vec<(String, String)>, - /// 首个有效上游 WebSocket 事件到达代理的耗时。 - pub first_token_ms: Option, - /// WebSocket 连接池决策。 - pub pool_decision: Option, - /// 上游诊断元数据。 - pub diagnostics: CodexUpstreamDiagnostics, -} - -/// Responses WebSocket live SSE exchange result. -pub struct CodexWebSocketStreamingExchange { - /// Live SSE bytes converted from WebSocket events. - pub body: CodexWebSocketSseStream, - /// 上游 metadata 帧中的最新 turn state。 - pub turn_state: Option, - /// 上游握手响应里的 `set-cookie` 列表。 - pub set_cookie_headers: Vec, - /// 上游握手响应里的限流头。 - pub rate_limit_headers: Vec<(String, String)>, - /// 上游内部 rate-limit 事件里的动态更新。 - pub rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates, - /// 上游内部 metadata 事件里的动态 turn state。 - pub turn_state_update: CodexWebSocketTurnStateUpdate, - /// WebSocket 连接池决策。 - pub pool_decision: Option, - /// 上游诊断元数据。 - pub diagnostics: CodexUpstreamDiagnostics, -} - -/// Responses WebSocket live SSE byte stream. -pub type CodexWebSocketSseStream = - Pin> + Send + 'static>>; -/// WebSocket live stream rate-limit header updates. -pub type CodexWebSocketRateLimitHeaderUpdates = Arc>>; -/// WebSocket live stream turn-state update. -pub type CodexWebSocketTurnStateUpdate = Arc>>; - -/// Responses WebSocket exchange error. -#[derive(Debug, Error)] -pub enum CodexWebSocketExchangeError { - /// opening request 无法构造。 - #[error("invalid websocket request: {0}")] - InvalidRequest(#[from] tungstenite::http::Error), - /// WebSocket 传输失败。 - #[error("websocket transport error: {0}")] - Transport(#[from] tungstenite::Error), - /// SSE 聚合结果无法解析。 - #[error("invalid websocket SSE response: {0}")] - InvalidSse(#[from] SseError), - /// 上游 WebSocket 错误帧。 - #[error("{0}")] - Upstream(Box), - /// 上游返回 `response.incomplete`。 - #[error("Incomplete response returned, reason: {reason}")] - IncompleteResponse { - /// incomplete_details.reason。 - reason: String, - }, - /// 上游返回无法按官方形状解析的 `response.completed`。 - #[error("{message}")] - InvalidCompletedResponse { - /// 解析失败说明。 - message: String, - }, - /// 上游在 terminal 事件前关闭。 - #[error("websocket closed before terminal event")] - ClosedBeforeTerminal, - /// 上游在指定时间内没有发送任何事件。 - #[error("websocket receive idle timeout after {timeout:?}")] - ReceiveIdleTimeout { - /// 超时时长。 - timeout: Duration, - }, - /// 上游返回非文本事件帧。 - #[error("unexpected binary websocket event")] - UnexpectedBinaryEvent, - /// 复用的池连接在首个真实输出前失效。 - #[error("reused websocket connection died before first output: {message}")] - ReusedConnectionDiedBeforeFirstOutput { - /// 底层失效原因。 - message: String, - }, - /// 首个真实输出在配置的绝对超时内未到达（连接落到病态上游后端）。 - #[error("websocket first output not received within {timeout:?}")] - FirstTokenTimeout { - /// 首 token 超时时长。 - timeout: Duration, - }, -} - -/// WebSocket 上游错误帧载荷。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CodexWebSocketUpstreamError { - /// HTTP-style upstream status code. - pub status_code: u16, - /// 推导出的重试秒数。 - pub retry_after_seconds: Option, - /// 原始错误帧。 - pub body: String, - /// 上游透传的 `set-cookie` 列表。 - pub set_cookie_headers: Vec, - /// 上游诊断元数据。 - pub diagnostics: CodexUpstreamDiagnostics, -} - -impl fmt::Display for CodexWebSocketUpstreamError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!( - f, - "websocket upstream error {}: {}", - self.status_code, self.body - ) - } -} - -impl CodexWebSocketExchangeError { - pub(super) fn upstream( - status_code: u16, - retry_after_seconds: Option, - body: String, - set_cookie_headers: Vec, - diagnostics: CodexUpstreamDiagnostics, - ) -> Self { - Self::Upstream(Box::new(CodexWebSocketUpstreamError { - status_code, - retry_after_seconds, - body, - set_cookie_headers, - diagnostics, - })) - } -} diff --git a/backend/src/upstream/openai/transport/websocket/stream.rs b/backend/src/upstream/openai/transport/websocket_frames.rs similarity index 81% rename from backend/src/upstream/openai/transport/websocket/stream.rs rename to backend/src/upstream/openai/transport/websocket_frames.rs index 1fc684c18..0b7db70a8 100644 --- a/backend/src/upstream/openai/transport/websocket/stream.rs +++ b/backend/src/upstream/openai/transport/websocket_frames.rs @@ -1,5 +1,152 @@ +//! WebSocket 帧交换、流转发与错误类型。 + use super::*; +/// Responses WebSocket exchange result. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CodexWebSocketExchange { + /// 完整 SSE 文本。 + pub body: String, + /// 从 SSE 中提取的 usage。 + pub usage: Option, + /// 上游 metadata 帧中的最新 turn state。 + pub turn_state: Option, + /// 上游握手响应里的 `set-cookie` 列表。 + pub set_cookie_headers: Vec, + /// 上游握手响应里的限流头。 + pub rate_limit_headers: Vec<(String, String)>, + /// 首个有效上游 WebSocket 事件到达代理的耗时。 + pub first_token_ms: Option, + /// WebSocket 连接池决策。 + pub pool_decision: Option, + /// 上游诊断元数据。 + pub diagnostics: CodexUpstreamDiagnostics, +} + +/// Responses WebSocket live SSE exchange result. +pub struct CodexWebSocketStreamingExchange { + /// Live SSE bytes converted from WebSocket events. + pub body: CodexWebSocketSseStream, + /// 上游 metadata 帧中的最新 turn state。 + pub turn_state: Option, + /// 上游握手响应里的 `set-cookie` 列表。 + pub set_cookie_headers: Vec, + /// 上游握手响应里的限流头。 + pub rate_limit_headers: Vec<(String, String)>, + /// 上游内部 rate-limit 事件里的动态更新。 + pub rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates, + /// 上游内部 metadata 事件里的动态 turn state。 + pub turn_state_update: CodexWebSocketTurnStateUpdate, + /// WebSocket 连接池决策。 + pub pool_decision: Option, + /// 上游诊断元数据。 + pub diagnostics: CodexUpstreamDiagnostics, +} + +/// Responses WebSocket live SSE byte stream. +pub type CodexWebSocketSseStream = + Pin> + Send + 'static>>; +/// WebSocket live stream rate-limit header updates. +pub type CodexWebSocketRateLimitHeaderUpdates = Arc>>; +/// WebSocket live stream turn-state update. +pub type CodexWebSocketTurnStateUpdate = Arc>>; + +/// Responses WebSocket exchange error. +#[derive(Debug, Error)] +pub enum CodexWebSocketExchangeError { + /// opening request 无法构造。 + #[error("invalid websocket request: {0}")] + InvalidRequest(#[from] tungstenite::http::Error), + /// WebSocket 传输失败。 + #[error("websocket transport error: {0}")] + Transport(#[from] tungstenite::Error), + /// SSE 聚合结果无法解析。 + #[error("invalid websocket SSE response: {0}")] + InvalidSse(#[from] SseError), + /// 上游 WebSocket 错误帧。 + #[error("{0}")] + Upstream(Box), + /// 上游返回 `response.incomplete`。 + #[error("Incomplete response returned, reason: {reason}")] + IncompleteResponse { + /// incomplete_details.reason。 + reason: String, + }, + /// 上游返回无法按官方形状解析的 `response.completed`。 + #[error("{message}")] + InvalidCompletedResponse { + /// 解析失败说明。 + message: String, + }, + /// 上游在 terminal 事件前关闭。 + #[error("websocket closed before terminal event")] + ClosedBeforeTerminal, + /// 上游在指定时间内没有发送任何事件。 + #[error("websocket receive idle timeout after {timeout:?}")] + ReceiveIdleTimeout { + /// 超时时长。 + timeout: Duration, + }, + /// 上游返回非文本事件帧。 + #[error("unexpected binary websocket event")] + UnexpectedBinaryEvent, + /// 复用的池连接在首个真实输出前失效。 + #[error("reused websocket connection died before first output: {message}")] + ReusedConnectionDiedBeforeFirstOutput { + /// 底层失效原因。 + message: String, + }, + /// 首个真实输出在配置的绝对超时内未到达（连接落到病态上游后端）。 + #[error("websocket first output not received within {timeout:?}")] + FirstTokenTimeout { + /// 首 token 超时时长。 + timeout: Duration, + }, +} + +/// WebSocket 上游错误帧载荷。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CodexWebSocketUpstreamError { + /// HTTP-style upstream status code. + pub status_code: u16, + /// 推导出的重试秒数。 + pub retry_after_seconds: Option, + /// 原始错误帧。 + pub body: String, + /// 上游透传的 `set-cookie` 列表。 + pub set_cookie_headers: Vec, + /// 上游诊断元数据。 + pub diagnostics: CodexUpstreamDiagnostics, +} + +impl fmt::Display for CodexWebSocketUpstreamError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "websocket upstream error {}: {}", + self.status_code, self.body + ) + } +} + +impl CodexWebSocketExchangeError { + pub(super) fn upstream( + status_code: u16, + retry_after_seconds: Option, + body: String, + set_cookie_headers: Vec, + diagnostics: CodexUpstreamDiagnostics, + ) -> Self { + Self::Upstream(Box::new(CodexWebSocketUpstreamError { + status_code, + retry_after_seconds, + body, + set_cookie_headers, + diagnostics, + })) + } +} + /// 执行一次 prepared Responses WebSocket 请求并聚合为 SSE。 pub async fn execute_response_create_request( request: &CodexWebSocketRequest, diff --git a/backend/tests/api/admin/accounts_routes/export.rs b/backend/tests/api/admin/accounts_routes/export_routes.rs similarity index 100% rename from backend/tests/api/admin/accounts_routes/export.rs rename to backend/tests/api/admin/accounts_routes/export_routes.rs diff --git a/backend/tests/api/admin/accounts_routes/import.rs b/backend/tests/api/admin/accounts_routes/import_routes.rs similarity index 100% rename from backend/tests/api/admin/accounts_routes/import.rs rename to backend/tests/api/admin/accounts_routes/import_routes.rs diff --git a/backend/tests/api/admin/accounts_routes/lifecycle.rs b/backend/tests/api/admin/accounts_routes/lifecycle_routes.rs similarity index 100% rename from backend/tests/api/admin/accounts_routes/lifecycle.rs rename to backend/tests/api/admin/accounts_routes/lifecycle_routes.rs diff --git a/backend/tests/api/admin/accounts_routes/mod.rs b/backend/tests/api/admin/accounts_routes/mod.rs index eab86bf48..b0e08d51a 100644 --- a/backend/tests/api/admin/accounts_routes/mod.rs +++ b/backend/tests/api/admin/accounts_routes/mod.rs @@ -28,13 +28,13 @@ use crate::support::{ storage::{background_task_stores, create_test_redis, init_test_db, test_database_url}, }; -mod export; -mod import; -mod lifecycle; +mod export_routes; +mod import_routes; +mod lifecycle_routes; mod list; -mod oauth; -mod probe; -mod quota; +mod oauth_routes; +mod probe_routes; +mod quota_routes; struct AdminAccountsTestState { app_state: AppState, diff --git a/backend/tests/api/admin/accounts_routes/oauth.rs b/backend/tests/api/admin/accounts_routes/oauth_routes.rs similarity index 100% rename from backend/tests/api/admin/accounts_routes/oauth.rs rename to backend/tests/api/admin/accounts_routes/oauth_routes.rs diff --git a/backend/tests/api/admin/accounts_routes/probe.rs b/backend/tests/api/admin/accounts_routes/probe_routes.rs similarity index 100% rename from backend/tests/api/admin/accounts_routes/probe.rs rename to backend/tests/api/admin/accounts_routes/probe_routes.rs diff --git a/backend/tests/api/admin/accounts_routes/quota.rs b/backend/tests/api/admin/accounts_routes/quota_routes.rs similarity index 100% rename from backend/tests/api/admin/accounts_routes/quota.rs rename to backend/tests/api/admin/accounts_routes/quota_routes.rs diff --git a/backend/tests/dispatch/service/mod.rs b/backend/tests/dispatch/service.rs similarity index 95% rename from backend/tests/dispatch/service/mod.rs rename to backend/tests/dispatch/service.rs index 46f63f0e7..3d7033b87 100644 --- a/backend/tests/dispatch/service/mod.rs +++ b/backend/tests/dispatch/service.rs @@ -106,124 +106,124 @@ where } const RESPONSES_SUCCESS_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/success.sse"), + include_str!("../fixtures/responses/http_sse/success.sse"), "\n" ); const RESPONSES_COMPLETED_USAGE_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/completed_usage.sse"), + include_str!("../fixtures/responses/http_sse/completed_usage.sse"), "\n" ); const RESPONSES_COMPLETED_IMAGE_USAGE_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/completed_image_usage.sse"), + include_str!("../fixtures/responses/http_sse/completed_image_usage.sse"), "\n" ); const RESPONSES_EMPTY_COMPLETED_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/empty_completed.sse"), + include_str!("../fixtures/responses/http_sse/empty_completed.sse"), "\n" ); const RESPONSES_TEXT_DELTAS_COMPLETED_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/text_deltas_completed.sse"), + include_str!("../fixtures/responses/http_sse/text_deltas_completed.sse"), "\n" ); const RESPONSES_DONE_ITEM_COMPLETED_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/done_item_completed.sse"), + include_str!("../fixtures/responses/http_sse/done_item_completed.sse"), "\n" ); const RESPONSES_TUPLE_OBJECT_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/tuple_object.sse"), + include_str!("../fixtures/responses/http_sse/tuple_object.sse"), "\n" ); const RESPONSES_AFTER_5XX_RETRY_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/after_5xx_retry.sse"), + include_str!("../fixtures/responses/http_sse/after_5xx_retry.sse"), "\n" ); const RESPONSES_AFTER_402_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/after_402.sse"), + include_str!("../fixtures/responses/http_sse/after_402.sse"), "\n" ); const RESPONSES_FAILED_QUOTA_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/failed_quota.sse"), + include_str!("../fixtures/responses/http_sse/failed_quota.sse"), "\n" ); const RESPONSES_FAILED_AUTH_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/failed_auth.sse"), + include_str!("../fixtures/responses/http_sse/failed_auth.sse"), "\n" ); const RESPONSES_AFTER_401_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/after_401.sse"), + include_str!("../fixtures/responses/http_sse/after_401.sse"), "\n" ); const RESPONSES_AFTER_403_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/after_403.sse"), + include_str!("../fixtures/responses/http_sse/after_403.sse"), "\n" ); const RESPONSES_STREAM_USAGE_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/stream_usage.sse"), + include_str!("../fixtures/responses/http_sse/stream_usage.sse"), "\n" ); const RESPONSES_INCOMPLETE_USAGE_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/chat_delta_incomplete_usage.sse"), + include_str!("../fixtures/responses/http_sse/chat_delta_incomplete_usage.sse"), "\n" ); const RESPONSES_STREAM_AFTER_429_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/stream_after_429.sse"), + include_str!("../fixtures/responses/http_sse/stream_after_429.sse"), "\n" ); const RESPONSES_STREAM_AFTER_5XX_RETRY_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/stream_after_5xx_retry.sse"), + include_str!("../fixtures/responses/http_sse/stream_after_5xx_retry.sse"), "\n" ); const RESPONSES_STREAM_AFTER_402_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/stream_after_402.sse"), + include_str!("../fixtures/responses/http_sse/stream_after_402.sse"), "\n" ); const RESPONSES_STREAM_AFTER_401_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/stream_after_401.sse"), + include_str!("../fixtures/responses/http_sse/stream_after_401.sse"), "\n" ); const RESPONSES_STREAM_AFTER_403_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/stream_after_403.sse"), + include_str!("../fixtures/responses/http_sse/stream_after_403.sse"), "\n" ); const RESPONSES_AFTER_CLOUDFLARE_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/after_cloudflare.sse"), + include_str!("../fixtures/responses/http_sse/after_cloudflare.sse"), "\n" ); const RESPONSES_STREAM_AFTER_CLOUDFLARE_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/stream_after_cloudflare.sse"), + include_str!("../fixtures/responses/http_sse/stream_after_cloudflare.sse"), "\n" ); const RESPONSES_AFTER_MODEL_UNSUPPORTED_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/after_model_unsupported.sse"), + include_str!("../fixtures/responses/http_sse/after_model_unsupported.sse"), "\n" ); const RESPONSES_STREAM_AFTER_MODEL_UNSUPPORTED_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/stream_after_model_unsupported.sse"), + include_str!("../fixtures/responses/http_sse/stream_after_model_unsupported.sse"), "\n" ); const RESPONSES_FAILED_MODEL_UNSUPPORTED_SSE: &str = concat!( - include_str!("../../fixtures/responses/http_sse/failed_model_unsupported.sse"), + include_str!("../fixtures/responses/http_sse/failed_model_unsupported.sse"), "\n" ); const WEBSOCKET_COMPLETED_WITH_REASONING_REPLAY: &str = - include_str!("../../fixtures/responses/websocket/completed_with_reasoning_replay.json"); + include_str!("../fixtures/responses/websocket/completed_with_reasoning_replay.json"); const WEBSOCKET_HISTORY_RATE_LIMITED: &str = - include_str!("../../fixtures/responses/websocket/history_rate_limited.json"); + include_str!("../fixtures/responses/websocket/history_rate_limited.json"); const WEBSOCKET_RATE_LIMITED: &str = - include_str!("../../fixtures/responses/websocket/rate_limited.json"); + include_str!("../fixtures/responses/websocket/rate_limited.json"); const WEBSOCKET_TOKEN_REVOKED: &str = - include_str!("../../fixtures/responses/websocket/token_revoked.json"); + include_str!("../fixtures/responses/websocket/token_revoked.json"); const WEBSOCKET_FIRST_ACCOUNT_LIMITED: &str = - include_str!("../../fixtures/responses/websocket/first_account_limited.json"); + include_str!("../fixtures/responses/websocket/first_account_limited.json"); const WEBSOCKET_SECOND_ACCOUNT_LIMITED: &str = - include_str!("../../fixtures/responses/websocket/second_account_limited.json"); + include_str!("../fixtures/responses/websocket/second_account_limited.json"); const WEBSOCKET_INVALID_ENCRYPTED_CONTENT: &str = - include_str!("../../fixtures/responses/websocket/invalid_encrypted_content.json"); + include_str!("../fixtures/responses/websocket/invalid_encrypted_content.json"); const WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND: &str = - include_str!("../../fixtures/responses/websocket/previous_response_not_found.json"); + include_str!("../fixtures/responses/websocket/previous_response_not_found.json"); const REASONING_REPLAY_REQUEST_GOLDEN: &str = - include_str!("../../fixtures/responses/golden/reasoning_replay_request.json"); + include_str!("../fixtures/responses/golden/reasoning_replay_request.json"); mod compact_routes; mod responses_http; diff --git a/docs/architecture.md b/docs/architecture.md index 4650ab537..a8057ad9c 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -40,7 +40,7 @@ | L0 | `infra` | PostgreSQL 连接/迁移、Redis 连接/键前缀、时间、JSON、日志、身份生成、路径、格式化 | 出现任何业务概念 | | L1 | `upstream` | 纯上游客户端：协议编解码、HTTP/SSE/WebSocket 传输、指纹、OAuth token client。**每家 provider 一个子目录**（`openai/`，规划中 `cloudflare/`） | 依赖任何领域模块；出现账号池/调度/遥测概念 | | L1 | `telemetry` | 遥测域：usage_records / ops_error_logs / request_time_buckets / account_usage 的**写入与查询**、billing、dashboard 聚合、rebuild | 被冠以 admin 语义；HTTP handler | -| L1 | `keys` | 客户端 API key 域：store、内存缓存、鉴权校验、管理操作 | — | +| L1 | `keys` | 客户端 API key 域：PG 明文 store、唯一索引鉴权点查、管理操作 | 不得建立进程内鉴权缓存；完整 key 只向已认证管理端返回 | | L1 | `auth` | 管理端认证域：admin 用户（PG）、会话（Redis，token_hash 键，database.md §4B.1） | 客户端 key 概念（那是 `keys`） | | L1 | `settings` | 运行时设置域：持久化、快照、变更广播（watch）、admin API key 校验 | 直接调用上层服务推送变更（§2.3） | | L1 | `update` | 自更新域：release 查询、下载、校验、解压、替换、状态 | — | @@ -60,7 +60,7 @@ | `accounts` → `settings` | 池参数/刷新策略订阅设置快照（watch 类型）；单向 | | `auth` → `settings` | admin API key 哈希存放在 runtime_settings 表，经 settings 的 store 方法读取；单向 | -**全库硬禁**： +**全库硬禁**（迁移验收时由 §9 临时审计断言）： - L1–L3 任何模块 import `crate::api` / `crate::bootstrap`。 - `dispatch`、`accounts`、`telemetry` 等领域模块 import `axum`。 @@ -93,11 +93,13 @@ backend/src/ │ ├── mod.rs │ ├── token_client.rs │ ├── protocol/ -│ │ ├── mod.rs schema.rs responses.rs events.rs sse.rs websocket.rs +│ │ ├── mod.rs schema.rs responses.rs events.rs sse.rs +│ │ ├── websocket.rs websocket_errors.rs │ ├── transport/ -│ │ ├── mod.rs client.rs endpoints.rs headers.rs tls.rs +│ │ ├── mod.rs client.rs client_sse.rs endpoints.rs headers.rs tls.rs │ │ ├── response_meta.rs usage.rs diagnostics.rs -│ │ ├── websocket.rs websocket_pool.rs websocket_pump.rs +│ │ ├── websocket.rs websocket_frames.rs +│ │ ├── websocket_pool.rs websocket_pump.rs │ └── fingerprint/ │ ├── mod.rs types.rs store.rs runtime.rs updater.rs ├── telemetry/ # L1 遥测域（原 admin/monitoring + proxy/dispatch/usage_events） @@ -128,6 +130,8 @@ backend/src/ │ ├── account.rs # 实体（原 model.rs） │ ├── window.rs │ ├── pool.rs # 池：存储态/槽位/过滤（§5.4 行数预算内需拆出 pool/ 子文件） +│ ├── pool/ +│ │ ├── state.rs filters.rs │ ├── scheduler.rs # 调度门面（原 scheduler/mod.rs 的类型） │ ├── scheduler/ │ │ ├── candidates.rs feedback.rs @@ -135,7 +139,7 @@ backend/src/ │ │ ├── mod.rs smart.rs quota_reset.rs round_robin.rs sticky.rs │ ├── store.rs # AccountStore trait + PG 门面 │ ├── store/ -│ │ ├── queries.rs rows.rs +│ │ ├── queries.rs rows.rs write.rs │ ├── quota/ │ │ ├── mod.rs service.rs runtime.rs │ ├── refresh/ # 原 token_refresh @@ -155,8 +159,7 @@ backend/src/ │ ├── errors.rs # 原 dispatch/errors.rs + responses/errors.rs 合并 │ ├── recording.rs # 原 responses/event_recording.rs，调 telemetry::recorder │ ├── affinity/ # 原 session_affinity.rs（683 行三合一）+ responses/affinity.rs -│ │ ├── mod.rs types.rs identity.rs store.rs service.rs resolve.rs -│ │ │ # store = Redis 键 + 二级索引（§4B.3），无内存映射副本 +│ │ ├── mod.rs types.rs store.rs service.rs resolve.rs # store = Redis 键 + 二级索引（§4B.3），无内存映射副本 │ ├── recovery/ │ │ ├── mod.rs auth.rs cloudflare.rs exhaustion.rs │ │ ├── implicit_resume.rs reasoning_replay.rs @@ -175,12 +178,17 @@ backend/src/ │ ├── mod.rs router.rs response.rs session.rs # session = 鉴权提取器/中间件 │ ├── auth_routes.rs settings_routes.rs system_routes.rs │ ├── accounts_routes.rs keys_routes.rs +│ ├── accounts_routes/ +│ │ ├── query.rs lifecycle_routes.rs oauth_routes.rs quota_routes.rs +│ │ ├── import_routes.rs export_routes.rs probe_routes.rs │ ├── usage_routes.rs ops_routes.rs dashboard_routes.rs # ops_routes = /api/admin/ops/errors └── bootstrap/ # L5 进程装配（原 runtime + config/schema,loader） ├── mod.rs ├── config.rs # 启动配置 schema + loader（原 config/schema.rs + loader.rs；含 redis.url） ├── services.rs state.rs shutdown.rs ├── import_sqlite.rs # 一次性 SQLite v3 → PG 导入命令（database.md §6；§5.5 的唯一跨域写豁免） + ├── import_sqlite/ + │ ├── core.rs telemetry.rs # 导入表组拆分；SQL 仍受 §5.5 唯一豁免约束 └── tasks/ ├── mod.rs coordinator.rs periodic.rs cleanup.rs ├── cookie_cleanup.rs retention_trim.rs # retention_trim 新增（database.md §7） @@ -234,7 +242,7 @@ backend/src/ | `Admin*` | **仅** `api/admin` 下的 DTO/提取器 | `AdminAccountPayload` | | 无前缀 | 领域服务与实体 | `Recorder`、`AccountScheduler`、`SettingsService` | -由此更名（现状 → 终态）：`Sqlite*Store` 全系 → `Pg*Store` 或 `Redis*Store`（按 database.md 的归属，引擎切换 PR0 完成）；`AdminUsageRecordService` → `telemetry::usage::UsageQueryService`（读）+ `telemetry::Recorder`（写）；`AdminOpsErrorLogService` → `telemetry::ops` 同拆；`AdminUsageService`（实际管 account_usage）→ `telemetry::account_usage::AccountUsageQueryService`；`AdminSessionService` → `auth::SessionService`；`AdminClientKeyService` → `keys::KeyManageService`；`ClientKeyService` → `keys::KeyVerifier`（完整 key → PG unique 点查，旧 `RuntimeClientKeyStore` 内存鉴权表消亡）。**类型名与文件名对不上的历史包袱全部在搬家时清算，不留别名 re-export。** +由此更名（现状 → 终态）：`Sqlite*Store` 全系 → `Pg*Store` 或 `Redis*Store`（按 database.md 的归属，引擎切换 PR0 完成）；`AdminUsageRecordService` → `telemetry::usage::UsageQueryService`（读）+ `telemetry::Recorder`（写）；`AdminOpsErrorLogService` → `telemetry::ops` 同拆；`AdminUsageService`（实际管 account_usage）→ `telemetry::account_usage::AccountUsageQueryService`；`AdminSessionService` → `auth::SessionService`；`AdminClientKeyService` → `keys::KeyManageService`；`ClientKeyService` → `keys::KeyVerifier`（完整 key → PG unique 点查，`RuntimeClientKeyStore` 进程内鉴权缓存消亡）。客户端 API key 的 PG 明文列、管理端完整取回、复制与 CCSwitch 导入是产品需求，不得改成不可逆哈希。**类型名与文件名对不上的历史包袱全部在搬家时清算，不留别名 re-export。** 禁止无意义别名导入：`use std::sync::Arc as StdArc` 这类一律 `use std::sync::Arc`。 @@ -256,7 +264,7 @@ axum/tower 类型只允许出现在 `api` 与 `bootstrap`。领域函数签名 ### 5.4 文件行数预算 -- 软上限 **400 行**（超过即在 review 中说明理由）；硬上限 **800 行**（超过必须拆）。 +- 软上限 **400 行**（超过即在 review 中说明理由）；硬上限 **800 行**（超过必须拆，§9 临时审计直接失败）。 - 现状超硬线清单（11 个，全部有既定拆法）： | 文件（现路径） | 行数 | 拆法 | @@ -266,13 +274,15 @@ axum/tower 类型只允许出现在 `api` 与 `bootstrap`。领域函数签名 | `admin/accounts/routes.rs` | 1455 | → `api/admin/accounts_routes.rs` 按子资源拆（accounts / oauth / quota / import-export） | | `upstream/accounts/pool.rs` | 1420 | → `accounts/pool.rs`（门面）+ `pool/state.rs`（槽位/在途）+ `pool/filters.rs` | | `upstream/transport/client.rs` | 1251 | → `client.rs` + `client_sse.rs`（SSE 请求路径） | -| `upstream/accounts/store/mod.rs` | 1126 | → `store.rs` 门面 + `store/queries.rs` 按读写再分 | +| `upstream/accounts/store/mod.rs` | 1126 | → `store.rs` 门面 + `store/{queries,rows,write}.rs` 按读写与行映射拆分 | | `admin/monitoring/usage_record_store.rs` | 1114 | PR0 原地重写（拆表后 metadata 聚合全数删除，行数大减），PR1 搬家时拆 `telemetry/usage/{store,query}.rs` | | `admin/monitoring/dashboard.rs` | 1080 | → `telemetry/dashboard.rs`（聚合）+ `api/admin/dashboard_routes.rs`（HTTP） | | `upstream/protocol/websocket.rs` | 1033 | → `websocket.rs`（事件转换）+ `websocket_errors.rs`（错误帧解析） | | `upstream/accounts/token_refresh/runtime.rs` | 953 | → `accounts/refresh/{service,lease,policy}.rs` | | `admin/system/updater/mod.rs` | 811 | → `update/{service,state}.rs` | +新增的 `bootstrap/import_sqlite.rs` 按导入表组拆为 `import_sqlite/{core,telemetry}.rs`，避免一次性迁移命令突破 800 行硬线；它仍是 §5.5 唯一允许跨域写入的入口。 + ### 5.5 一域一账本 每张 PG 表 / 每个 Redis 键空间恰好属于一个域，只有该域的 store 可以写它： @@ -365,7 +375,7 @@ axum/tower 类型只允许出现在 `api` 与 `bootstrap`。领域函数签名 | `admin/monitoring/dashboard.rs` | `telemetry/dashboard.rs` + `api/admin/dashboard_routes.rs` † | | `admin/monitoring/diagnostics.rs` | `telemetry/diagnostics.rs` | | `admin/monitoring/usage_record_routes.rs` | `api/admin/usage_routes.rs` | -| `admin/keys/service.rs` | `keys/{types,store,service,manage}.rs` †（进程内鉴权缓存消亡，database.md §4.3） | +| `admin/keys/service.rs` | `keys/{types,store,service,manage}.rs` †（明文内存缓存消亡，database.md §4.3） | | `admin/keys/routes.rs` | `api/admin/keys_routes.rs` | | `admin/auth/service.rs` | `auth/{types,store,service}.rs` †（会话改 Redis，database.md §4B.1） | | `admin/auth/session.rs` | `api/admin/session.rs` | @@ -393,30 +403,84 @@ axum/tower 类型只允许出现在 `api` 与 `bootstrap`。领域函数签名 --- -## 8. 执行记录（PR0 + 六个顺序搬家阶段） +## 8. 执行计划（PR0 + 六个顺序搬家 PR，每个独立编译绿、测试绿） -以下阶段已于 **2026-07-10** 全部完成。固定验收：`cargo fmt --check`、`cargo clippy -- -D warnings`、`cargo test`、`backend/tests` 镜像目录同步搬移。**PR0 是唯一的迁移行为变更阶段**（存储引擎与 schema 语义一次到位）；PR1 起为搬家、更名与拆分。 +每个 PR 的固定验收：`cargo fmt --check`、`cargo clippy -- -D warnings`、`cargo test`、§9 临时架构审计、`backend/tests` 镜像目录同步搬移。**PR0 是唯一的行为变更 PR**（存储引擎与 schema 语义一次到位）；PR1 起全部为零行为搬家（更名与移动，不改逻辑）。 | PR | 内容 | 关键动作 | | --- | --- | --- | -| **PR0** | **存储终态切换**（database.md §7 全部 12 条，在现有目录结构内完成） | Cargo：sqlx 加 postgres、新增 redis；`infra/database.rs` 改 PG + 0001 终态基线、新增 `infra/redis.rs`；全部 store 换 PG 方言 / 运行态四键空间改 Redis；成功/失败事实拆表 + provider 维度 + 两条写入事务；client key 可取回、session/admin-key 哈希化 + `client_api_key_id` 归因贯通；Dashboard 口径换源；trim 周期任务化；`import-sqlite` / `rebuild-buckets` 子命令；tests support 换 PG+Redis;deploy compose 加双服务。验收对齐 database.md §6.4 | +| **PR0** | **存储终态切换**（database.md §7 全部 12 条，在现有目录结构内完成） | Cargo：sqlx 加 postgres、新增 redis；`infra/database.rs` 改 PG + 0001 终态基线、新增 `infra/redis.rs`；全部 store 换 PG 方言 / 运行态四键空间改 Redis；成功/失败事实拆表 + provider 维度 + 两条写入事务；客户端 key 明文可取回、session/admin-key 哈希化 + `client_api_key_id` 归因贯通；Dashboard 口径换源；trim 周期任务化；`import-sqlite` / `rebuild-buckets` 子命令；tests support 换 PG+Redis;deploy compose 加双服务。验收对齐 database.md §6.4 | | **PR1** | `telemetry` 域成立 | monitoring 全部迁出 admin；`usage_events.rs` → `recorder.rs`；Admin* 遥测类型按 §4.3 更名；dispatch/api 改依赖 telemetry；启用"禁止 `crate::admin::monitoring`"检查 | | **PR2** | `keys` / `auth` / `settings` / `update` 域成立 | 四个域从 admin/config 拆出；settings 改 watch 广播（§2.3）；config 只剩 schema+loader | | **PR3** | `accounts` / `models` 域成立 | `upstream/accounts`、`upstream/scheduler`、`admin/accounts/service` 三处合并为 `accounts`；`upstream` 缩为纯客户端并落成 `openai/` 子目录；动名词文件更名；Repository→Store | | **PR4** | `dispatch` 域成立 | proxy/dispatch 整体迁出并按目标树重组；affinity store 独立成文件；`proxy` 只剩 api/client 的料 | | **PR5** | `api` / `bootstrap` 壳成立 | http + web + proxy 入口 + admin routes 合并为 `api`；runtime + config 装配部分改名 `bootstrap`；删除 `crate::{admin,proxy,runtime,web,http,config}` 六个旧根模块；启用全部依赖检查 | -| **PR6** | 行数预算清账 | §5.4 清单内 11 个文件拆到硬线以下 | +| **PR6** | 行数预算清账 | §5.4 清单内 11 个文件拆到硬线以下；纳入 §9 临时审计 | -顺序依据：存储切换触及所有 store、无法按域切片，必须整体先行（PR0）——先切引擎再搬家，每个 store 文件只被"重写"一次、"移动"一次，不叠加；PR1–4 自底向上消除旧模块的存在必要；PR5 收壳；PR6 纯内部拆分无跨域影响。原"PR1 与数据库改造互锁"的约定由 PR0 替代，PR2 中的哈希化条目亦已并入 PR0。 +顺序依据：存储切换触及所有 store、无法按域切片，必须整体先行（PR0）——先切引擎再搬家，每个 store 文件只被"重写"一次、"移动"一次，不叠加；PR1–4 自底向上消除旧模块的存在必要；PR5 收壳；PR6 纯内部拆分无跨域影响。原"PR1 与数据库改造互锁"的约定由 PR0 替代，PR2 中的 session/admin-key 哈希化条目亦已并入 PR0；客户端 API key 保持 database.md §4.3 的明文可取回设计。 --- -## 9. 仓库级约定 +## 9. 临时架构审计 + +迁移验收期间可从历史提交临时取出以下门禁逻辑在本地执行，用于检查依赖方向、命名和行数预算。该脚本不恢复到仓库，也不加入 CI；正式门禁仍由现有编译、测试与 Clippy 工作流承担。 + +```bash +#!/usr/bin/env bash +# 架构依赖与命名断言。违规即失败。 +set -euo pipefail +cd "$(dirname "$0")/.." +fail=0 +deny() { # deny <描述> <路径...> + if grep -rn --include='*.rs' -E "$2" "${@:3}" ; then + echo "ARCH VIOLATION: $1" >&2; fail=1 + fi +} + +# 分层：领域层不得依赖 api / bootstrap +deny "domain -> api/bootstrap" 'crate::(api|bootstrap)\b' \ + src/infra src/upstream src/telemetry src/keys src/auth src/settings src/update \ + src/accounts src/models src/dispatch +# upstream 是纯客户端 +deny "upstream -> domain" 'crate::(accounts|dispatch|telemetry|keys|auth|settings|update|models)\b' src/upstream +# HTTP 类型只在 api / bootstrap +deny "axum outside api" '\baxum(::|\b)' \ + src/infra src/upstream src/telemetry src/keys src/auth src/settings src/update \ + src/accounts src/models src/dispatch +# SQL 与 Redis 命令只在 store / query +deny "sqlx outside store/query" 'sqlx::query' \ + $(find src/api -name '*.rs') +deny "redis outside store" '\bredis::' \ + $(find src/api -name '*.rs') +# 旧模块根不得复活 +deny "legacy module path" 'crate::(admin|proxy|runtime|web|http|config)::' src +# 禁用文件名 +banned=$(find src -name 'importing.rs' -o -name 'exporting.rs' -o -name 'testing.rs' \ + -o -name 'util.rs' -o -name 'helper.rs' -o -name 'common.rs' -o -name 'misc.rs') +[ -n "$banned" ] && { echo "ARCH VIOLATION: banned file names: $banned" >&2; fail=1; } +# mod.rs 纪律 +while read -r f; do + lines=$(wc -l < "$f") + [ "$lines" -gt 50 ] && { echo "ARCH VIOLATION: $f has $lines lines (mod.rs <= 50)" >&2; fail=1; } +done < <(find src -name mod.rs) +# 行数硬上限 +while read -r f; do + lines=$(wc -l < "$f") + [ "$lines" -gt 800 ] && { echo "ARCH VIOLATION: $f has $lines lines (> 800)" >&2; fail=1; } +done < <(find src -name '*.rs' ! -name mod.rs) +exit $fail +``` + +临时审计对迁移完成后的全部目录执行；检查结束后不保留生成的脚本文件。 + +--- + +## 10. 仓库级约定 - `backend/tests/` 镜像 `src/`，不在 `src/` 写测试模块。测试依赖本机（或 CI 服务容器）的 PostgreSQL 与 Redis：`CPR_TEST_DATABASE_URL` 下每测试建独立数据库、`CPR_TEST_REDIS_URL` 下每测试随机键前缀（database.md §7）。 - `docs/` 只保留权威文档：`architecture.md`（本文）、`database.md`。行为细节（调度策略、WS 池、更新流程等）以模块文档注释为准（`scheduler/mod.rs` 是范本），不再在 docs 里维护第二份会腐化的副本。 - `deploy/` 随 PR0 扩容（compose 增 postgres/redis 服务与健康检查依赖）；`release/`、`skills/` 现状合格，不动。 -- 前端联动（全部随 PR0）：keys 列表继续返回完整 key，并长期提供复制与 CCSwitch 导入；admin API key 状态不再有脱敏值（哈希不可逆，只余"已启用"）；成功记录响应不再有 `level` 字段；错误明细/筛选改走 `/api/admin/ops/errors`；Dashboard "总请求/总 token" 卡片改为 account_usage 来源并标注"成功累计"。 +- 前端联动（全部随 PR0）：已认证管理端的 keys 列表继续返回完整 key，并保留长期复制与 CCSwitch 导入；admin API key 状态不再有脱敏值（哈希不可逆，只余"已启用"）；成功记录响应不再有 `level` 字段；错误明细/筛选改走 `/api/admin/ops/errors`；Dashboard "总请求/总 token" 卡片改为 account_usage 来源并标注"成功累计"。 ---