refactor: conform backend to authoritative architecture

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-11 00:43:47 +08:00
co-authored by Codex
parent fe5494a614
commit 7cb186c4ae
63 changed files with 5719 additions and 5604 deletions
File diff suppressed because it is too large Load Diff
@@ -1,3 +1,5 @@
//! 账号获取过滤条件、池配置与选择结果类型。
use super::*;
#[derive(Debug, Clone, PartialEq, Eq)]
-715
View File
@@ -1,715 +0,0 @@
use super::*;
/// 账号运行参数中不支持热更新的部分。
#[derive(Debug, Clone)]
pub struct AccountPoolStaticSettings {
pub skip_quota_limited: bool,
pub tier_priority: Vec<String>,
}
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<tokio::sync::Mutex<AccountPool>>,
store: Arc<dyn AccountStore>,
usage_store: Arc<dyn AccountUsageStore>,
request_interval: Arc<std::sync::RwLock<StdDuration>>,
}
impl RuntimeAccountPoolService {
/// 构造运行时账号池服务。
pub fn new(
store: Arc<dyn AccountStore>,
usage_store: Arc<dyn AccountUsageStore>,
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<Self>,
mut receiver: tokio::sync::watch::Receiver<crate::settings::SettingsSnapshot>,
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<String, Vec<String>>,
fetched_model_plan_types: BTreeSet<String>,
) {
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<usize, RuntimeAccountPoolError> {
let mut accounts = self
.store
.list_pool_accounts()
.await
.map_err(|_| RuntimeAccountPoolError::Generic)?;
let account_ids = accounts
.iter()
.map(|account| account.id.clone())
.collect::<Vec<_>>();
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<Utc>) -> 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<Utc>) -> 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<AcquiredAccount> {
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<Utc>) -> Vec<DistinctPlanAccount> {
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<RefreshedAccountState>) {
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::<serde_json::Value>(&quota_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, &quota).await;
}
/// 读取运行时账号快照。
pub async fn account_snapshot(&self, account_id: &str) -> Option<Account> {
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, &quota_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<bool, RuntimeAccountPoolError> {
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<Utc>,
) -> 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<Utc>) {
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<DateTime<Utc>>,
limit_window_seconds: Option<u64>,
now: DateTime<Utc>,
) -> Option<DateTime<Utc>> {
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<Utc>) {
if account
.cloudflare_cooldown_until
.is_some_and(|cooldown_until| now >= cooldown_until)
{
account.cloudflare_cooldown_until = None;
}
}
+712
View File
@@ -0,0 +1,712 @@
//! 账号池存储态、槽位与状态变更。
use super::*;
#[derive(Debug)]
pub(super) struct AccountAcquireWithStatusRefresh {
pub(super) acquired: Option<AcquiredAccount>,
pub(super) refreshed_accounts: Vec<RefreshedAccountState>,
}
#[derive(Debug)]
pub(super) struct DistinctPlanAccountsWithStatusRefresh {
pub(super) accounts: Vec<DistinctPlanAccount>,
pub(super) refreshed_accounts: Vec<RefreshedAccountState>,
}
#[derive(Debug)]
pub(super) struct AccountCapacitySummaryWithStatusRefresh {
pub(super) summary: AccountCapacitySummary,
pub(super) refreshed_accounts: Vec<RefreshedAccountState>,
}
/// 账号池容量摘要。
#[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<DateTime<Utc>>,
cloudflare_cooldown_until: Option<DateTime<Utc>>,
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<DateTime<Utc>>,
window_reset_at: Option<DateTime<Utc>>,
limit_window_seconds: Option<u64>,
}
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<String, Account>,
slots: BTreeMap<String, VecDeque<AccountSlot>>,
pub(super) options: AccountPoolOptions,
scheduler: AccountScheduler,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct AccountSlot {
created_at: DateTime<Utc>,
model: Option<String>,
}
/// 账号释放结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ReleasedAccountUsage {
/// 被释放槽位对应的模型。槽位已被过期清理时为空。
pub model: Option<String>,
}
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<Account> {
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<u64>) {
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<Utc>,
) -> 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<DateTime<Utc>>,
) -> 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<Utc>,
limit_window_seconds: Option<u64>,
) -> 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<Utc>,
) -> 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<AcquiredAccount> {
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<Utc>) -> Vec<RefreshedAccountState> {
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<ReleasedAccountUsage> {
self.release_slot(account_id, true)
}
/// 释放指定账号的一个在途槽位,不累计请求用量。
pub fn release_without_request_usage(
&mut self,
account_id: &str,
) -> Option<ReleasedAccountUsage> {
self.release_slot(account_id, false)
}
fn release_slot(
&mut self,
account_id: &str,
record_request_usage: bool,
) -> Option<ReleasedAccountUsage> {
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<Utc>) -> AccountCapacitySummary {
self.capacity_summary_with_status_refresh(now).summary
}
pub(super) fn capacity_summary_with_status_refresh(
&mut self,
now: DateTime<Utc>,
) -> 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<Utc>) -> Vec<DistinctPlanAccount> {
self.distinct_plan_accounts_with_status_refresh(now)
.accounts
}
pub(super) fn distinct_plan_accounts_with_status_refresh(
&mut self,
now: DateTime<Utc>,
) -> DistinctPlanAccountsWithStatusRefresh {
self.cleanup_stale_slots(now);
let refreshed_accounts = self.refresh_account_statuses(now);
let mut by_plan = IndexMap::<String, Vec<Account>>::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<DateTime<Utc>> {
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<Utc>, 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<Utc>) -> Option<Account> {
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<Utc>,
limit_window_seconds: Option<u64>,
now: DateTime<Utc>,
) -> 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<Utc>) {
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<Account> {
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<Utc>) -> Option<Account> {
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<Utc>) -> 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(),
}
}
+100 -4
View File
@@ -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<Utc>,
}
/// 账号调度器:账号选择的对外门面。
///
@@ -78,7 +118,7 @@ impl AccountScheduler {
slot_count: &dyn Fn(&str) -> usize,
now: DateTime<Utc>,
) -> Option<Account> {
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<Account> {
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<Account> {
let mut sorted = candidates.iter().collect::<Vec<_>>();
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<DateTime<Utc>>, b: Option<DateTime<Utc>>) -> 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)
}
@@ -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};
@@ -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<Account> {
@@ -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<Account> {
@@ -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],权重和通常为
/// 个位数),保证未封顶账号总是优先于封顶账号。
@@ -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` 最大者)。
///
@@ -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<Utc>,
}
/// 按指定策略从候选集合中选择一个账号。
///
/// `cursor` 是账号池持有的轮转游标,用于在并列最优账号之间做公平轮转;不同策略对它
/// 的读写方式不同(详见各策略模块),本函数只负责透传可变引用。
pub fn select(
strategy: RotationStrategy,
input: &SelectionInput<'_>,
cursor: &mut usize,
) -> Option<Account> {
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<Account> {
let mut sorted = candidates.iter().collect::<Vec<_>>();
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<DateTime<Utc>>, b: Option<DateTime<Utc>>) -> 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)
}
+10
View File
@@ -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,
+161 -1
View File
@@ -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<Option<StoredAccount>> {
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<Option<StoredAccount>> {
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<String>,
limit: u32,
) -> PgAccountStoreResult<Page<StoredAccount>> {
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<String>,
limit: u32,
) -> PgAccountStoreResult<Page<StoredAccountMetadata>> {
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<NumberedPage<StoredAccountMetadata>> {
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::<Postgres>::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::<PgAccountStoreResult<Vec<_>>>()?;
Ok(NumberedPage {
items,
total,
page: page.max(1),
page_size,
})
}
/// 读取单个账号元数据(不含 token)。
pub async fn get_metadata(
&self,
account_id: &str,
) -> PgAccountStoreResult<Option<StoredAccountMetadata>> {
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<Vec<AccountQuotaSnapshot>> {
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<Option<String>> {
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::<Option<sqlx::types::Json<Value>>, _>("quota_json"))
.map(|value| value.0.to_string()))
}
}
-166
View File
@@ -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<Option<StoredAccount>> {
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<Option<StoredAccount>> {
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<String>,
limit: u32,
) -> PgAccountStoreResult<Page<StoredAccount>> {
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<String>,
limit: u32,
) -> PgAccountStoreResult<Page<StoredAccountMetadata>> {
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<NumberedPage<StoredAccountMetadata>> {
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::<Postgres>::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::<PgAccountStoreResult<Vec<_>>>()?;
Ok(NumberedPage {
items,
total,
page: page.max(1),
page_size,
})
}
/// 读取单个账号元数据(不含 token)。
pub async fn get_metadata(
&self,
account_id: &str,
) -> PgAccountStoreResult<Option<StoredAccountMetadata>> {
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<Vec<AccountQuotaSnapshot>> {
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<Option<String>> {
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::<Option<sqlx::types::Json<Value>>, _>("quota_json"))
.map(|value| value.0.to_string()))
}
/// 更新配额 JSON。
pub async fn update_quota_json(
&self,
+14 -4
View File
@@ -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::*;
@@ -1,321 +0,0 @@
use super::*;
/// `GET /api/admin/accounts/export`
pub(crate) async fn export_accounts(
State(state): State<AppState>,
_auth: AdminAuth,
Query(query): Query<AccountExportQuery>,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
Json(payload): Json<AccountActionRequest>,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
Json(payload): Json<AccountHealthCheckRequest>,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
Query(query): Query<AccountIdQuery>,
) -> Result<impl IntoResponse, AdminError> {
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(&quota_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<AdminAccountData, AdminError> {
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), &quota_by_account).await?;
Ok(stats.data_for(account))
}
/// `POST /api/admin/accounts/import`
pub(crate) async fn import_accounts(
State(state): State<AppState>,
_auth: AdminAuth,
Json(payload): Json<Value>,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
Json(payload): Json<AccountOAuthExchangeRequest>,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
Query(query): Query<AccountTestQuery>,
_auth: AdminAuth,
) -> Result<Response, AdminError> {
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<AppState>,
Query(query): Query<AccountIdQuery>,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
Json(payload): Json<BatchDeleteAccountsRequest>,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
Json(payload): Json<Value>,
) -> Result<impl IntoResponse, AdminError> {
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)),
}
}
@@ -0,0 +1,24 @@
use super::*;
/// `GET /api/admin/accounts/export`
pub(crate) async fn export_accounts(
State(state): State<AppState>,
_auth: AdminAuth,
Query(query): Query<AccountExportQuery>,
) -> Result<impl IntoResponse, AdminError> {
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)),
}
}
@@ -0,0 +1,20 @@
use super::*;
/// `POST /api/admin/accounts/import`
pub(crate) async fn import_accounts(
State(state): State<AppState>,
_auth: AdminAuth,
Json(payload): Json<Value>,
) -> Result<impl IntoResponse, AdminError> {
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)),
}
}
@@ -0,0 +1,110 @@
use super::*;
/// `POST /api/admin/accounts/refresh`
pub(crate) async fn refresh_account(
State(state): State<AppState>,
_auth: AdminAuth,
Json(payload): Json<AccountActionRequest>,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
Json(payload): Json<AccountHealthCheckRequest>,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
Json(payload): Json<BatchDeleteAccountsRequest>,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
Json(payload): Json<Value>,
) -> Result<impl IntoResponse, AdminError> {
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)),
}
}
@@ -0,0 +1,49 @@
use super::*;
/// `POST /api/admin/accounts/oauth/authorize`
pub(crate) async fn oauth_authorize_account(
State(state): State<AppState>,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
Json(payload): Json<AccountOAuthExchangeRequest>,
) -> Result<impl IntoResponse, AdminError> {
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)),
}
}
@@ -0,0 +1,62 @@
use super::*;
/// `GET /api/admin/accounts/test?id=...&modelId=...`
pub(crate) async fn test_account_connection(
State(state): State<AppState>,
Query(query): Query<AccountTestQuery>,
_auth: AdminAuth,
) -> Result<Response, AdminError> {
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<AppState>,
Query(query): Query<AccountIdQuery>,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
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)),
}
}
@@ -0,0 +1,61 @@
use super::*;
/// `GET /api/admin/accounts/quota`
pub(crate) async fn account_quota(
State(state): State<AppState>,
_auth: AdminAuth,
Query(query): Query<AccountIdQuery>,
) -> Result<impl IntoResponse, AdminError> {
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(&quota_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<AdminAccountData, AdminError> {
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), &quota_by_account).await?;
Ok(stats.data_for(account))
}
+143
View File
@@ -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<String>,
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<AppState>,
client_ip: Option<Extension<ClientIp>>,
Json(payload): Json<AdminLoginRequest>,
) -> Result<Response, AdminError> {
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<AppState>,
headers: HeaderMap,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
headers: HeaderMap,
) -> Result<Response, AdminError> {
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<Extension<ClientIp>>) -> Option<String> {
let Extension(client_ip) = client_ip?;
Some(client_ip.as_str().to_string())
}
fn admin_login_source(fallback_client_ip: Option<String>) -> String {
fallback_client_ip.unwrap_or_else(|| "unknown".to_string())
}
+300 -4
View File
@@ -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<Utc>,
) -> Vec<AccountUsageTimeBucket> {
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<String>,
) -> 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<String, f64>,
now: DateTime<Utc>,
) -> Vec<DashboardAccountUsageData> {
let account_by_id = accounts
.iter()
.map(|account| (account.id.as_str(), account))
.collect::<HashMap<_, _>>();
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<String, f64> {
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(&quota_json) {
quota_used_by_account.insert(usage.account_id.clone(), used_percent);
}
}
quota_used_by_account
}
fn quota_used_percent(quota_json: &str) -> Option<f64> {
let quota = serde_json::from_str::<Value>(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<f64>,
) {
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<f64> {
let percent = value
.as_f64()
.or_else(|| value.as_str().and_then(|value| value.parse::<f64>().ok()))?;
percent.is_finite().then_some(percent.clamp(0.0, 100.0))
}
fn dashboard_usage_record_items(
items: Vec<UsageRecord>,
account_emails: &HashMap<String, String>,
) -> Vec<DashboardUsageRecordData> {
items
.into_iter()
.map(|record| dashboard_usage_record_data(record, account_emails))
.collect()
}
fn dashboard_usage_record_data(
record: UsageRecord,
account_emails: &HashMap<String, String>,
) -> 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<DashboardServiceStatusData> {
let current_fingerprint = state.services.fingerprint.snapshot();
let fingerprint = fingerprint_diagnostics(&current_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>) -> 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)
}
@@ -1,304 +0,0 @@
use super::*;
pub(super) async fn dashboard_time_buckets(
state: &AppState,
now: DateTime<Utc>,
) -> Vec<AccountUsageTimeBucket> {
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<String>,
) -> 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<String, f64>,
now: DateTime<Utc>,
) -> Vec<DashboardAccountUsageData> {
let account_by_id = accounts
.iter()
.map(|account| (account.id.as_str(), account))
.collect::<HashMap<_, _>>();
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<String, f64> {
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(&quota_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<f64> {
let quota = serde_json::from_str::<Value>(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<f64>,
) {
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<f64> {
let percent = value
.as_f64()
.or_else(|| value.as_str().and_then(|value| value.parse::<f64>().ok()))?;
percent.is_finite().then_some(percent.clamp(0.0, 100.0))
}
pub(super) fn dashboard_usage_record_items(
items: Vec<UsageRecord>,
account_emails: &HashMap<String, String>,
) -> Vec<DashboardUsageRecordData> {
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<String, String>,
) -> 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<DashboardServiceStatusData> {
let current_fingerprint = state.services.fingerprint.snapshot();
let fingerprint = fingerprint_diagnostics(&current_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>) -> 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)
}
+1
View File
@@ -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;
+1 -1
View File
@@ -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,
+7 -138
View File
@@ -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<AppState> for AdminAuth {
}
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct AdminLoginRequest {
username: Option<String>,
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<AppState>,
client_ip: Option<Extension<ClientIp>>,
Json(payload): Json<AdminLoginRequest>,
) -> Result<Response, AdminError> {
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<AppState>,
headers: HeaderMap,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
headers: HeaderMap,
) -> Result<Response, AdminError> {
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<String> {
pub(super) fn admin_session_cookie(headers: &HeaderMap) -> Option<String> {
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<String> {
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<Extension<ClientIp>>) -> Option<String> {
let Extension(client_ip) = client_ip?;
Some(client_ip.as_str().to_string())
}
fn admin_login_source(fallback_client_ip: Option<String>) -> String {
fallback_client_ip.unwrap_or_else(|| "unknown".to_string())
}
+4 -4
View File
@@ -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},
},
-222
View File
@@ -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<String> {
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 = "<system-reminder>";
const LEADING_SYSTEM_REMINDER_CLOSE: &str = "</system-reminder>";
/// 从 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<String> {
let instructions = request
.instructions()
.chars()
.take(2000)
.collect::<String>();
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::<String>();
}
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)
}
+1 -2
View File
@@ -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::*;
+373 -4
View File
@@ -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<String>,
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<Utc>,
) -> 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<Utc>,
) -> Option<String> {
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<Utc>,
) -> Option<String> {
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<ImplicitResumeSnapshot>,
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<ImplicitResumeSnapshot>,
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<RuntimeSessionAffinityService>,
reasoning_replay: &Arc<Mutex<ReasoningReplayCache>>,
@@ -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<String> {
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 = "<system-reminder>";
const LEADING_SYSTEM_REMINDER_CLOSE: &str = "</system-reminder>";
/// 从 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<String> {
let instructions = request
.instructions()
.chars()
.take(2000)
.collect::<String>();
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::<String>();
}
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)
}
+1 -1
View File
@@ -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},
},
@@ -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<ImplicitResumeSnapshot> {
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<ImplicitResumeSnapshot> {
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<ImplicitResumeSnapshot>,
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<String> {
input
.iter()
+589 -317
View File
@@ -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<RuntimeAccountPoolService>,
models: Arc<ModelService>,
codex: Arc<CodexBackendClient>,
session_affinity: Arc<RuntimeSessionAffinityService>,
reasoning_replay: Arc<Mutex<ReasoningReplayCache>>,
recorder: Arc<Recorder>,
installation_id: Option<String>,
cloudflare: CloudflareRecovery,
pub(in crate::dispatch) account_pool: Arc<RuntimeAccountPoolService>,
pub(in crate::dispatch) models: Arc<ModelService>,
pub(in crate::dispatch) codex: Arc<CodexBackendClient>,
pub(in crate::dispatch) session_affinity: Arc<RuntimeSessionAffinityService>,
pub(in crate::dispatch) reasoning_replay: Arc<Mutex<ReasoningReplayCache>>,
pub(in crate::dispatch) recorder: Arc<Recorder>,
pub(in crate::dispatch) installation_id: Option<String>,
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<String>,
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<Box<dyn Stream<Item = Result<Bytes, ResponseDispatchStreamError>> + Send + 'static>>;
@@ -182,246 +132,11 @@ impl ResponseDispatchService {
}
}
async fn prepare_response_session(
&self,
request: &mut CodexResponsesRequest,
) -> Option<ImplicitResumeSnapshot> {
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<ImplicitResumeSnapshot> {
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<Utc>,
) -> 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<Utc>,
) -> Option<String> {
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<Utc>,
) -> Option<String> {
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<ImplicitResumeSnapshot>,
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<ImplicitResumeSnapshot>,
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<ImplicitResumeSnapshot>,
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<Value, ResponseDispatchError> {
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<String> = 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;
}
}
-290
View File
@@ -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<Value, ResponseDispatchError> {
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;
}
}
-575
View File
@@ -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<Value, ResponseDispatchError> {
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<String> = 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;
}
}
-619
View File
@@ -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<ResponseDispatchStream, ResponseDispatchError> {
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<String> = 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));
}
}
}
}
}
+645 -237
View File
@@ -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<RuntimeAccountPoolService>,
pub(in crate::dispatch) session_affinity: Arc<RuntimeSessionAffinityService>,
pub(in crate::dispatch) reasoning_replay: Arc<Mutex<ReasoningReplayCache>>,
pub(in crate::dispatch) recorder: Arc<Recorder>,
pub(in crate::dispatch) cloudflare: CloudflareRecovery,
pub(in crate::dispatch) account_id: String,
pub(in crate::dispatch) account_plan_type: Option<String>,
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<String>,
pub(in crate::dispatch) request: CodexResponsesRequest,
pub(in crate::dispatch) tuple_schema: Option<Value>,
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<CodexRateLimitHeaderUpdates>,
pub(in crate::dispatch) turn_state_update: Option<CodexTurnStateUpdate>,
pub(in crate::dispatch) websocket_pool_decision: Option<WebSocketPoolDecision>,
pub(in crate::dispatch) turn_state: Option<String>,
pub(in crate::dispatch) diagnostics: CodexUpstreamDiagnostics,
pub(in crate::dispatch) attempt: ResponseDispatchAttempt,
pub(in crate::dispatch) attempts: Vec<ResponseDispatchAttempt>,
pub(in crate::dispatch) started_at: Instant,
}
pub(in crate::dispatch) fn latest_response_id(body: &str) -> Option<String> {
parse_sse_events(body).ok().and_then(|events| {
events.iter().rev().find_map(|event| {
serde_json::from_str::<Value>(&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<i64>,
) {
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<ResponseDispatchStream, ResponseDispatchError> {
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<String> = 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;
}
+252 -11
View File
@@ -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<Result<Bytes, ResponseDispatchStreamError>>,
cancel: Option<oneshot::Sender<()>>,
@@ -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<RuntimeAccountPoolService>,
pub(in crate::dispatch) session_affinity: Arc<RuntimeSessionAffinityService>,
pub(in crate::dispatch) reasoning_replay: Arc<Mutex<ReasoningReplayCache>>,
pub(in crate::dispatch) recorder: Arc<Recorder>,
pub(in crate::dispatch) cloudflare: CloudflareRecovery,
pub(in crate::dispatch) account_id: String,
pub(in crate::dispatch) account_plan_type: Option<String>,
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<String>,
pub(in crate::dispatch) request: CodexResponsesRequest,
pub(in crate::dispatch) tuple_schema: Option<Value>,
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<CodexRateLimitHeaderUpdates>,
pub(in crate::dispatch) turn_state_update: Option<CodexTurnStateUpdate>,
pub(in crate::dispatch) websocket_pool_decision: Option<WebSocketPoolDecision>,
pub(in crate::dispatch) turn_state: Option<String>,
pub(in crate::dispatch) diagnostics: CodexUpstreamDiagnostics,
pub(in crate::dispatch) attempt: ResponseDispatchAttempt,
pub(in crate::dispatch) attempts: Vec<ResponseDispatchAttempt>,
pub(in crate::dispatch) started_at: Instant,
}
pub(in crate::dispatch) fn latest_response_id(body: &str) -> Option<String> {
parse_sse_events(body).ok().and_then(|events| {
events.iter().rev().find_map(|event| {
serde_json::from_str::<Value>(&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<i64>,
) {
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;
}
+327 -8
View File
@@ -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<Value, ResponseDispatchError> {
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;
}
}
+127
View File
@@ -28,6 +28,13 @@ pub struct PgRequestBucketStore {
pool: PgPool,
}
/// 可由事实表重建的请求时间桶范围结果。
pub(crate) struct RebuiltRequestBuckets {
pub cutoff: DateTime<Utc>,
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<RebuiltRequestBuckets, PgRequestBucketStoreError> {
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<Utc> = 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 {
+9 -113
View File
@@ -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<RebuildBucketsReport, RebuildBucketsError> {
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<Utc> = 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,
})
}
+438 -5
View File
@@ -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<u64> {
let mut builder = QueryBuilder::<Postgres>::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<UsageRecordSummary> {
let mut builder = QueryBuilder::<Postgres>::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<Vec<UsageRecordAccountUsage>> {
let mut builder = QueryBuilder::<Postgres>::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<HashMap<String, String>> {
let mut account_ids = items
.iter()
.map(|item| item.account_id.clone())
.collect::<Vec<_>>();
account_ids.sort_unstable();
account_ids.dedup();
if account_ids.is_empty() {
return Ok(HashMap::new());
}
let mut builder = QueryBuilder::<Postgres>::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::<Option<String>, _>("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<Vec<UsageRecordBreakdown>> {
let mut builder = QueryBuilder::<Postgres>::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::<String, BreakdownAccumulator>::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::<String, _>("billing_model"),
row.get::<Option<String>, _>("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::<Vec<_>>();
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<Vec<UsageRecordTrendPoint>> {
let mut builder = QueryBuilder::<Postgres>::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::<String, UsageTrendAccumulator>::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::<String, _>("billing_model"),
row.get::<Option<String>, _>("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::<Vec<_>>();
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::<i32, _>("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::<sqlx::types::Json<serde_json::Value>, _>("metadata_json")
.0,
created_at: row.get("created_at"),
}
}
fn nonnegative(value: Option<i64>) -> 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<f64>,
}
#[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<Utc>,
}
#[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<f64>,
}
#[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<f64>,
}
+8 -13
View File
@@ -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<Postgres>,
filter: &UsageRecordFilter,
cursor: Option<&str>,
-432
View File
@@ -1,432 +0,0 @@
use super::*;
pub(super) async fn count_usage_records(
pool: &PgPool,
filter: &UsageRecordFilter,
) -> PgUsageRecordStoreResult<u64> {
let mut builder = QueryBuilder::<Postgres>::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<UsageRecordSummary> {
let mut builder = QueryBuilder::<Postgres>::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<Vec<UsageRecordAccountUsage>> {
let mut builder = QueryBuilder::<Postgres>::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<HashMap<String, String>> {
let mut account_ids = items
.iter()
.map(|item| item.account_id.clone())
.collect::<Vec<_>>();
account_ids.sort_unstable();
account_ids.dedup();
if account_ids.is_empty() {
return Ok(HashMap::new());
}
let mut builder = QueryBuilder::<Postgres>::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::<Option<String>, _>("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<Vec<UsageRecordBreakdown>> {
let mut builder = QueryBuilder::<Postgres>::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::<String, BreakdownAccumulator>::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::<String, _>("billing_model"),
row.get::<Option<String>, _>("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::<Vec<_>>();
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<Vec<UsageRecordTrendPoint>> {
let mut builder = QueryBuilder::<Postgres>::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::<String, UsageTrendAccumulator>::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::<String, _>("billing_model"),
row.get::<Option<String>, _>("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::<Vec<_>>();
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::<i32, _>("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::<sqlx::types::Json<serde_json::Value>, _>("metadata_json")
.0,
created_at: row.get("created_at"),
}
}
fn nonnegative(value: Option<i64>) -> 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<f64>,
}
#[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<Utc>,
}
#[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<f64>,
}
#[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<f64>,
}
@@ -216,8 +216,9 @@ pub fn websocket_incomplete_response_reason(raw: &str) -> Option<String> {
.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 {
@@ -337,6 +337,7 @@ pub struct CodexBackendClient {
websocket_first_token_timeout: Option<Duration>,
}
#[path = "client_sse.rs"]
mod requests;
#[async_trait]
@@ -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<Option<PathBuf>> {
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<Option<PathBuf>> {
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 {
/// 返回连接描述。
@@ -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<Option<PathBuf>> {
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<Option<PathBuf>> {
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<TokenUsage>,
/// 上游 metadata 帧中的最新 turn state。
pub turn_state: Option<String>,
/// 上游握手响应里的 `set-cookie` 列表。
pub set_cookie_headers: Vec<String>,
/// 上游握手响应里的限流头。
pub rate_limit_headers: Vec<(String, String)>,
/// 首个有效上游 WebSocket 事件到达代理的耗时。
pub first_token_ms: Option<i64>,
/// WebSocket 连接池决策。
pub pool_decision: Option<WebSocketPoolDecision>,
/// 上游诊断元数据。
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<String>,
/// 上游握手响应里的 `set-cookie` 列表。
pub set_cookie_headers: Vec<String>,
/// 上游握手响应里的限流头。
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<WebSocketPoolDecision>,
/// 上游诊断元数据。
pub diagnostics: CodexUpstreamDiagnostics,
}
/// Responses WebSocket live SSE byte stream.
pub type CodexWebSocketSseStream =
Pin<Box<dyn Stream<Item = Result<Bytes, CodexWebSocketExchangeError>> + Send + 'static>>;
/// WebSocket live stream rate-limit header updates.
pub type CodexWebSocketRateLimitHeaderUpdates = Arc<Mutex<Vec<(String, String)>>>;
/// WebSocket live stream turn-state update.
pub type CodexWebSocketTurnStateUpdate = Arc<Mutex<Option<String>>>;
/// 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<CodexWebSocketUpstreamError>),
/// 上游返回 `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<u64>,
/// 原始错误帧。
pub body: String,
/// 上游透传的 `set-cookie` 列表。
pub set_cookie_headers: Vec<String>,
/// 上游诊断元数据。
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<u64>,
body: String,
set_cookie_headers: Vec<String>,
diagnostics: CodexUpstreamDiagnostics,
) -> Self {
Self::Upstream(Box::new(CodexWebSocketUpstreamError {
status_code,
retry_after_seconds,
body,
set_cookie_headers,
diagnostics,
}))
}
}
@@ -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<TokenUsage>,
/// 上游 metadata 帧中的最新 turn state。
pub turn_state: Option<String>,
/// 上游握手响应里的 `set-cookie` 列表。
pub set_cookie_headers: Vec<String>,
/// 上游握手响应里的限流头。
pub rate_limit_headers: Vec<(String, String)>,
/// 首个有效上游 WebSocket 事件到达代理的耗时。
pub first_token_ms: Option<i64>,
/// WebSocket 连接池决策。
pub pool_decision: Option<WebSocketPoolDecision>,
/// 上游诊断元数据。
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<String>,
/// 上游握手响应里的 `set-cookie` 列表。
pub set_cookie_headers: Vec<String>,
/// 上游握手响应里的限流头。
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<WebSocketPoolDecision>,
/// 上游诊断元数据。
pub diagnostics: CodexUpstreamDiagnostics,
}
/// Responses WebSocket live SSE byte stream.
pub type CodexWebSocketSseStream =
Pin<Box<dyn Stream<Item = Result<Bytes, CodexWebSocketExchangeError>> + Send + 'static>>;
/// WebSocket live stream rate-limit header updates.
pub type CodexWebSocketRateLimitHeaderUpdates = Arc<Mutex<Vec<(String, String)>>>;
/// WebSocket live stream turn-state update.
pub type CodexWebSocketTurnStateUpdate = Arc<Mutex<Option<String>>>;
/// 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<CodexWebSocketUpstreamError>),
/// 上游返回 `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<u64>,
/// 原始错误帧。
pub body: String,
/// 上游透传的 `set-cookie` 列表。
pub set_cookie_headers: Vec<String>,
/// 上游诊断元数据。
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<u64>,
body: String,
set_cookie_headers: Vec<String>,
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,
@@ -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,
@@ -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;
+83 -19
View File
@@ -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 <描述> <grep-pattern> <路径...>
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 来源并标注"成功累计"。
---