mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 16:06:08 +00:00
refactor: conform backend to authoritative architecture
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
+685
-686
File diff suppressed because it is too large
Load Diff
@@ -1,3 +1,5 @@
|
||||
//! 账号获取过滤条件、池配置与选择结果类型。
|
||||
|
||||
use super::*;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
@@ -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>("a_json).ok(),
|
||||
Ok(None) => None,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
account_id = %account_id,
|
||||
error = %error,
|
||||
"failed to read existing quota json before passive rate-limit sync"
|
||||
);
|
||||
None
|
||||
}
|
||||
};
|
||||
let quota = rate_limit_quota(&rate_limits, plan_type, existing_quota.as_ref());
|
||||
self.apply_quota_snapshot(account_id, "a).await;
|
||||
}
|
||||
|
||||
/// 读取运行时账号快照。
|
||||
pub async fn account_snapshot(&self, account_id: &str) -> Option<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, "a_json, limit_reached, cooldown_until)
|
||||
.await
|
||||
{
|
||||
Ok(persisted) => persisted,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
account_id,
|
||||
error = %error,
|
||||
"failed to persist verified quota snapshot"
|
||||
);
|
||||
false
|
||||
}
|
||||
};
|
||||
|
||||
if let Some(reset_at) = reset_at {
|
||||
if let Err(error) = self
|
||||
.usage_store
|
||||
.sync_rate_limit_window(account_id, reset_at, limit_window_seconds)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
account_id,
|
||||
error = %error,
|
||||
"failed to persist verified quota window"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let in_memory = {
|
||||
let mut pool = self.pool.lock().await;
|
||||
let in_memory = pool.apply_quota_state(account_id, limit_reached, cooldown_until);
|
||||
if let Some(reset_at) = reset_at {
|
||||
pool.sync_rate_limit_window(account_id, reset_at, limit_window_seconds);
|
||||
}
|
||||
in_memory
|
||||
};
|
||||
|
||||
persisted || in_memory
|
||||
}
|
||||
|
||||
/// 从运行时账号池移除账号。
|
||||
pub async fn remove_account(&self, account_id: &str) -> bool {
|
||||
self.pool.lock().await.remove(account_id)
|
||||
}
|
||||
|
||||
/// 从仓储同步单个账号到运行时账号池。
|
||||
pub async fn sync_account_from_store(
|
||||
&self,
|
||||
account_id: &str,
|
||||
) -> Result<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;
|
||||
}
|
||||
}
|
||||
@@ -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(),
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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()))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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("a_json, Some(Utc::now()));
|
||||
apply_account_quota_window_local_usage(&state, &account_id, &mut quota_data).await;
|
||||
let account =
|
||||
account_data_for_quota_refresh(&state, &account_id, quota_data.clone()).await?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AdminAccountQuotaResponseData::from_account(
|
||||
quota, raw, quota_data, account,
|
||||
)),
|
||||
))
|
||||
}
|
||||
Err(AccountManageError::NotFound) => Err(account_not_found()),
|
||||
Err(AccountManageError::Inactive(status)) => Err(AdminError::conflict(format!(
|
||||
"Account is {}, cannot query quota",
|
||||
status.as_str()
|
||||
))),
|
||||
Err(AccountManageError::FetchQuota(msg)) => Err(AdminError::bad_gateway(format!(
|
||||
"Failed to fetch quota from Codex API: {msg}"
|
||||
))),
|
||||
Err(e) => Err(account_error(&e)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn account_data_for_quota_refresh(
|
||||
state: &AppState,
|
||||
account_id: &str,
|
||||
quota_data: AccountQuotaData,
|
||||
) -> Result<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), "a_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("a_json, Some(Utc::now()));
|
||||
apply_account_quota_window_local_usage(&state, &account_id, &mut quota_data).await;
|
||||
let account =
|
||||
account_data_for_quota_refresh(&state, &account_id, quota_data.clone()).await?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AdminAccountQuotaResponseData::from_account(
|
||||
quota, raw, quota_data, account,
|
||||
)),
|
||||
))
|
||||
}
|
||||
Err(AccountManageError::NotFound) => Err(account_not_found()),
|
||||
Err(AccountManageError::Inactive(status)) => Err(AdminError::conflict(format!(
|
||||
"Account is {}, cannot query quota",
|
||||
status.as_str()
|
||||
))),
|
||||
Err(AccountManageError::FetchQuota(msg)) => Err(AdminError::bad_gateway(format!(
|
||||
"Failed to fetch quota from Codex API: {msg}"
|
||||
))),
|
||||
Err(e) => Err(account_error(&e)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn account_data_for_quota_refresh(
|
||||
state: &AppState,
|
||||
account_id: &str,
|
||||
quota_data: AccountQuotaData,
|
||||
) -> Result<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), "a_by_account).await?;
|
||||
Ok(stats.data_for(account))
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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("a_json) {
|
||||
quota_used_by_account.insert(usage.account_id.clone(), used_percent);
|
||||
}
|
||||
}
|
||||
quota_used_by_account
|
||||
}
|
||||
|
||||
fn quota_used_percent(quota_json: &str) -> Option<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(¤t_fingerprint);
|
||||
vec![
|
||||
DashboardServiceStatusData {
|
||||
label: "客户端版本".to_string(),
|
||||
value: empty_dash(fingerprint.app_version),
|
||||
detail: if fingerprint.build_number.trim().is_empty() {
|
||||
"-".to_string()
|
||||
} else {
|
||||
format!("Build {}", fingerprint.build_number)
|
||||
},
|
||||
tone: "info".to_string(),
|
||||
},
|
||||
DashboardServiceStatusData {
|
||||
label: "平台架构".to_string(),
|
||||
value: empty_dash(fingerprint.platform),
|
||||
detail: empty_dash(fingerprint.arch),
|
||||
tone: "info".to_string(),
|
||||
},
|
||||
DashboardServiceStatusData {
|
||||
label: "Chromium".to_string(),
|
||||
value: if fingerprint.chromium_version.trim().is_empty() {
|
||||
"-".to_string()
|
||||
} else {
|
||||
format!("v{}", fingerprint.chromium_version)
|
||||
},
|
||||
detail: empty_dash(fingerprint.originator),
|
||||
tone: "normal".to_string(),
|
||||
},
|
||||
DashboardServiceStatusData {
|
||||
label: "更新时间".to_string(),
|
||||
value: format_fingerprint_updated_at(fingerprint.updated_at),
|
||||
detail: String::new(),
|
||||
tone: "normal".to_string(),
|
||||
},
|
||||
DashboardServiceStatusData {
|
||||
label: "User Agent".to_string(),
|
||||
value: fingerprint.user_agent,
|
||||
detail: String::new(),
|
||||
tone: "normal".to_string(),
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
fn empty_dash(value: String) -> String {
|
||||
if value.trim().is_empty() {
|
||||
"-".to_string()
|
||||
} else {
|
||||
value
|
||||
}
|
||||
}
|
||||
|
||||
fn format_fingerprint_updated_at(value: Option<String>) -> 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("a_json) {
|
||||
quota_used_by_account.insert(usage.account_id.clone(), used_percent);
|
||||
}
|
||||
}
|
||||
quota_used_by_account
|
||||
}
|
||||
|
||||
pub(super) fn quota_used_percent(quota_json: &str) -> Option<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(¤t_fingerprint);
|
||||
vec![
|
||||
DashboardServiceStatusData {
|
||||
label: "客户端版本".to_string(),
|
||||
value: empty_dash(fingerprint.app_version),
|
||||
detail: if fingerprint.build_number.trim().is_empty() {
|
||||
"-".to_string()
|
||||
} else {
|
||||
format!("Build {}", fingerprint.build_number)
|
||||
},
|
||||
tone: "info".to_string(),
|
||||
},
|
||||
DashboardServiceStatusData {
|
||||
label: "平台架构".to_string(),
|
||||
value: empty_dash(fingerprint.platform),
|
||||
detail: empty_dash(fingerprint.arch),
|
||||
tone: "info".to_string(),
|
||||
},
|
||||
DashboardServiceStatusData {
|
||||
label: "Chromium".to_string(),
|
||||
value: if fingerprint.chromium_version.trim().is_empty() {
|
||||
"-".to_string()
|
||||
} else {
|
||||
format!("v{}", fingerprint.chromium_version)
|
||||
},
|
||||
detail: empty_dash(fingerprint.originator),
|
||||
tone: "normal".to_string(),
|
||||
},
|
||||
DashboardServiceStatusData {
|
||||
label: "更新时间".to_string(),
|
||||
value: format_fingerprint_updated_at(fingerprint.updated_at),
|
||||
detail: String::new(),
|
||||
tone: "normal".to_string(),
|
||||
},
|
||||
DashboardServiceStatusData {
|
||||
label: "User Agent".to_string(),
|
||||
value: fingerprint.user_agent,
|
||||
detail: String::new(),
|
||||
tone: "normal".to_string(),
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
pub(super) fn empty_dash(value: String) -> String {
|
||||
if value.trim().is_empty() {
|
||||
"-".to_string()
|
||||
} else {
|
||||
value
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn format_fingerprint_updated_at(value: Option<String>) -> 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,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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
|
||||
@@ -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},
|
||||
},
|
||||
|
||||
@@ -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,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::*;
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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>,
|
||||
}
|
||||
|
||||
@@ -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>,
|
||||
|
||||
@@ -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,
|
||||
}))
|
||||
}
|
||||
}
|
||||
+147
@@ -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
@@ -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 来源并标注"成功累计"。
|
||||
|
||||
---
|
||||
|
||||
|
||||
Reference in New Issue
Block a user