refactor: align implementation with authoritative architecture

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-10 22:53:27 +08:00
co-authored by Codex
parent 61cdaf753e
commit fe5494a614
181 changed files with 4398 additions and 4194 deletions
-3
View File
@@ -61,9 +61,6 @@ jobs:
working-directory: backend
run: cargo clippy --all-targets --all-features --locked -- -D warnings
- name: Architecture boundaries
run: backend/scripts/check_architecture.sh
- name: Rust tests
working-directory: backend
run: cargo test --test main --locked
-91
View File
@@ -1,91 +0,0 @@
#!/usr/bin/env bash
set -euo pipefail
cd "$(dirname "$0")/.."
fail=0
deny() {
local description=$1
local pattern=$2
shift 2
if grep -rn --include='*.rs' -E "$pattern" "$@"; then
echo "ARCH VIOLATION: $description" >&2
fail=1
fi
}
domain_paths=(
src/infra src/upstream src/telemetry src/keys src/auth src/settings src/update
src/accounts src/models src/dispatch
)
deny "domain -> api/bootstrap" '(^|[^[:alnum:]_])(api|bootstrap)::' "${domain_paths[@]}"
deny "upstream -> domain" '(^|[^[:alnum:]_])(accounts|dispatch|telemetry|keys|auth|settings|update|models)::' src/upstream
deny "axum outside api/bootstrap" '\baxum(::|\b)' "${domain_paths[@]}"
deny "sqlx in api" '\bsqlx(::|\b)' src/api
deny "redis in api" '\bredis::' src/api
deny "legacy module path" 'crate::(admin|proxy|runtime|web|http|config)::' src
deny "Admin-prefixed domain type" '\bAdmin[A-Z][A-Za-z0-9]*\b' "${domain_paths[@]}"
deny "Repository naming" '\b[A-Za-z0-9]*Repository\b' src
deny "persistence repository naming" '\b[A-Za-z0-9_]*repository[A-Za-z0-9_]*\b' \
src/accounts src/auth src/dispatch src/keys src/models src/settings src/telemetry \
src/upstream/openai/fingerprint
deny "Arc alias" '\bArc[[:space:]]+as[[:space:]]+' src
storage_command_check() {
local kind=$1
local pattern=$2
local file
while IFS= read -r file; do
case "$file" in
*/store.rs | */store/*)
continue
;;
esac
if [ "$kind" = sql ]; then
case "$file" in
src/infra/database.rs | src/bootstrap/import_sqlite.rs | src/bootstrap/import_sqlite/* | src/telemetry/rebuild.rs | */query.rs)
continue
;;
esac
else
case "$file" in
src/infra/redis.rs | src/accounts/refresh/lease.rs)
continue
;;
esac
fi
echo "$file"
echo "ARCH VIOLATION: $kind command outside an allowed storage adapter" >&2
fail=1
done < <(grep -rl --include='*.rs' -E "$pattern" src || true)
}
storage_command_check sql 'sqlx::(query|query_as|query_scalar|raw_sql)'
storage_command_check redis 'redis::(cmd|pipe)|[.]query_async[[:space:]]*[(]|use redis::.*AsyncCommands'
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' \))
if [ -n "$banned" ]; then
echo "ARCH VIOLATION: banned file names: $banned" >&2
fail=1
fi
while read -r file; do
lines=$(wc -l < "$file")
if [ "$lines" -gt 50 ]; then
echo "ARCH VIOLATION: $file has $lines lines (mod.rs <= 50)" >&2
fail=1
fi
done < <(find src -name mod.rs)
while read -r file; do
lines=$(wc -l < "$file")
if [ "$lines" -gt 800 ]; then
echo "ARCH VIOLATION: $file has $lines lines (> 800)" >&2
fail=1
fi
done < <(find src -name '*.rs' ! -name mod.rs)
exit "$fail"
-42
View File
@@ -109,45 +109,3 @@ pub struct Account {
/// 最近一次使用时间。
pub last_used_at: Option<String>,
}
/// 账号用量增量。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountUsageDelta {
/// 请求数增量。
pub requests: u64,
/// 输入 token 增量。
pub input_tokens: u64,
/// 输出 token 增量。
pub output_tokens: u64,
/// 缓存 token 增量。
pub cached_tokens: u64,
/// reasoning token 增量。
pub reasoning_tokens: u64,
/// 上游返回的总 token 增量。
pub total_tokens: u64,
/// 空响应数增量。
pub empty_responses: u64,
/// 图片工具输入 token 增量。
pub image_input_tokens: u64,
/// 图片工具输出 token 增量。
pub image_output_tokens: u64,
/// 图片工具成功请求数增量。
pub image_requests: u64,
/// 图片工具失败请求数增量。
pub image_request_failures: u64,
}
/// 账号模型维度用量增量。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountModelUsageDelta {
/// 请求数增量。
pub requests: u64,
/// 错误数增量。
pub errors: u64,
/// 输入 token 增量。
pub input_tokens: u64,
/// 输出 token 增量。
pub output_tokens: u64,
/// 缓存 token 增量。
pub cached_tokens: u64,
}
-11
View File
@@ -1,11 +0,0 @@
use super::AccountManageService;
impl AccountManageService {
pub(crate) async fn usage_cookie_header(&self, account_id: &str) -> Option<String> {
self.cookies
.cookie_header_for_request(account_id, "chatgpt.com", "/codex/usage")
.await
.ok()
.flatten()
}
}
+7 -8
View File
@@ -238,25 +238,24 @@ impl AccountManageService {
return Err(AccountManageError::EmptyIds);
}
let mut deleted = 0u32;
let mut deleted_ids = Vec::new();
let mut not_found = Vec::new();
for id in ids {
match self.store.delete(&id).await {
Ok(true) => {
deleted += 1;
self.account_pool.remove_account(&id).await;
if let Err(error) = self.session_affinity.forget_account(&id).await {
tracing::warn!(
account_id = id,
error = %error,
"failed to remove Redis affinities for deleted account"
);
}
deleted_ids.push(id);
}
Ok(false) => not_found.push(id),
Err(_) => return Err(AccountManageError::Delete),
}
}
Ok(BatchDeleteAccounts { deleted, not_found })
Ok(BatchDeleteAccounts {
deleted,
not_found,
deleted_ids,
})
}
pub async fn refresh_account(
&self,
-1
View File
@@ -1,6 +1,5 @@
//! 账号管理用例。
mod cookies;
mod export;
mod import;
mod lifecycle;
+51 -59
View File
@@ -7,23 +7,16 @@ use serde_json::{json, Value};
use tokio::sync::mpsc;
use crate::{
accounts::{
account::{Account, AccountStatus},
store::StoredAccount,
},
dispatch::{
errors::{
auth_failure_account_status, is_auth_upstream_error, is_quota_exhausted_upstream_error,
},
responses::sse_failure::sse_failure_account_status,
},
models::{info::CodexModelInfo, service::ModelRefreshPlanAccount},
accounts::{account::AccountStatus, store::StoredAccount},
models::{service::ModelRefreshPlanAccount, types::CodexModelInfo},
upstream::openai::{
protocol::{
responses::{CodexResponsesRequest, ResponsesSseFailure},
sse::{encode_sse_event, parse_sse_events, SseEvent},
},
transport::{is_banned_upstream_error, CodexClientError, CodexRequestContext},
transport::{
is_banned_auth_signal, is_banned_upstream_error, CodexClientError, CodexRequestContext,
},
},
};
@@ -87,7 +80,8 @@ impl AccountManageService {
let request_id = uuid::Uuid::new_v4().to_string();
let plan_account = ModelRefreshPlanAccount {
plan_type: plan_type.to_string(),
account: stored_account_to_model_refresh_account(account.clone()),
access_token: account.access_token.expose_secret().to_string(),
account_id: account.account_id.clone(),
};
if let Err(error) = self
.models
@@ -257,47 +251,6 @@ fn account_plan_type(account: &StoredAccount) -> String {
.to_string()
}
fn stored_account_to_model_refresh_account(stored: StoredAccount) -> Account {
Account {
id: stored.id,
email: stored.email,
account_id: stored.account_id,
user_id: stored.user_id,
label: stored.label,
plan_type: stored.plan_type,
access_token: stored.access_token.expose_secret().to_string(),
refresh_token: stored
.refresh_token
.map(|token| token.expose_secret().to_string()),
access_token_expires_at: stored.access_token_expires_at,
next_refresh_at: stored.next_refresh_at,
status: stored.status,
quota_limit_reached: false,
quota_verify_required: false,
quota_cooldown_until: None,
cloudflare_cooldown_until: None,
request_count: 0,
empty_response_count: 0,
image_input_tokens: 0,
image_output_tokens: 0,
image_request_count: 0,
image_request_failed_count: 0,
window_request_count: 0,
window_input_tokens: 0,
window_output_tokens: 0,
window_cached_tokens: 0,
window_image_input_tokens: 0,
window_image_output_tokens: 0,
window_image_request_count: 0,
window_image_request_failed_count: 0,
window_started_at: None,
window_reset_at: None,
limit_window_seconds: None,
added_at: stored.added_at,
last_used_at: None,
}
}
fn test_responses_request(model: String) -> CodexResponsesRequest {
let mut request = CodexResponsesRequest::new_http_sse(
model,
@@ -396,15 +349,54 @@ async fn process_sse_event(
fn account_status_from_client_error(error: &CodexClientError) -> Option<AccountStatus> {
if is_banned_upstream_error(error) {
Some(AccountStatus::Banned)
} else if is_quota_exhausted_upstream_error(error) {
Some(AccountStatus::QuotaExhausted)
} else if is_auth_upstream_error(error) {
Some(auth_failure_account_status(error))
} else {
None
match error {
CodexClientError::Upstream { status, .. } if status.as_u16() == 402 => {
Some(AccountStatus::QuotaExhausted)
}
CodexClientError::Upstream { status, body, .. } if status.as_u16() == 401 => {
Some(if is_banned_auth_signal(body) {
AccountStatus::Banned
} else {
AccountStatus::Expired
})
}
_ => None,
}
}
}
fn sse_failure_account_status(failure: &ResponsesSseFailure) -> Option<AccountStatus> {
let code = failure
.upstream_code
.as_deref()
.unwrap_or_default()
.to_ascii_lowercase();
let message = failure.message.to_ascii_lowercase();
if matches!(code.as_str(), "quota_exceeded" | "insufficient_quota") || message.contains("quota")
{
return Some(AccountStatus::QuotaExhausted);
}
let authentication_failed = matches!(
code.as_str(),
"token_invalid"
| "token_expired"
| "token_revoked"
| "account_deactivated"
| "unauthorized"
| "invalid_api_key"
) || message.contains("token revoked")
|| message.contains("token invalid")
|| message.contains("token expired");
authentication_failed.then(|| {
if is_banned_auth_signal(&failure.message) {
AccountStatus::Banned
} else {
AccountStatus::Expired
}
})
}
fn take_sse_frame(buffer: &mut String) -> Option<String> {
let index = buffer.find("\n\n").or_else(|| buffer.find("\r\n\r\n"))?;
let delimiter_len = if buffer[index..].starts_with("\r\n\r\n") {
+17
View File
@@ -3,6 +3,23 @@ use secrecy::ExposeSecret;
use super::{types::AccountManageError, AccountManageService};
impl AccountManageService {
async fn usage_cookie_header(&self, account_id: &str) -> Option<String> {
self.cookies
.cookie_header_for_request(account_id, "chatgpt.com", "/codex/usage")
.await
.ok()
.flatten()
}
pub async fn quota_snapshots(
&self,
) -> Result<Vec<crate::accounts::store::AccountQuotaSnapshot>, AccountManageError> {
self.store
.list_quota_snapshots()
.await
.map_err(|_| AccountManageError::Inspect)
}
pub async fn account_quota(
&self,
account_id: &str,
-4
View File
@@ -14,7 +14,6 @@ use crate::{
refresh::{RedisRefreshLeaseStore, RuntimeRefreshPolicy},
store::PgAccountStore,
},
dispatch::affinity::RedisSessionAffinityStore,
models::service::ModelService,
upstream::openai::{token_client::TokenRefresher, transport::CodexBackendClient},
};
@@ -30,7 +29,6 @@ pub struct AccountManageService {
pub(crate) account_pool: Arc<RuntimeAccountPoolService>,
pub(crate) token_refresher: Arc<dyn TokenRefresher>,
pub(crate) refresh_leases: RedisRefreshLeaseStore,
pub(crate) session_affinity: RedisSessionAffinityStore,
pub(crate) refresh_lease_owner_prefix: String,
pub(crate) oauth: oauth::AccountOAuthService,
pub(crate) refresh_policy: RuntimeRefreshPolicy,
@@ -45,7 +43,6 @@ pub struct AccountManageServiceParts {
pub account_pool: Arc<RuntimeAccountPoolService>,
pub token_refresher: Arc<dyn TokenRefresher>,
pub refresh_leases: RedisRefreshLeaseStore,
pub session_affinity: RedisSessionAffinityStore,
pub oauth: oauth::AccountOAuthService,
pub refresh_policy: RuntimeRefreshPolicy,
pub installation_id: Option<String>,
@@ -61,7 +58,6 @@ impl AccountManageService {
account_pool: parts.account_pool,
token_refresher: parts.token_refresher,
refresh_leases: parts.refresh_leases,
session_affinity: parts.session_affinity,
refresh_lease_owner_prefix: format!(
"admin-account-refresh:{}",
Uuid::new_v4().simple()
+1
View File
@@ -176,6 +176,7 @@ impl AccountUpdate {
pub struct BatchDeleteAccounts {
pub deleted: u32,
pub not_found: Vec<String>,
pub(crate) deleted_ids: Vec<String>,
}
#[derive(Debug, Clone)]
-1
View File
@@ -8,6 +8,5 @@ pub mod pool;
pub mod quota;
pub mod refresh;
pub mod scheduler;
pub mod service;
pub mod store;
mod window;
+7 -4
View File
@@ -17,10 +17,13 @@ use crate::accounts::scheduler::{
AccountScheduler, ScoreWeights,
};
use crate::accounts::store::AccountStore;
use crate::accounts::window::should_reset_usage_window;
use crate::accounts::{
account::{Account, AccountModelUsageDelta, AccountStatus, AccountUsageDelta},
service::AccountService,
account::{Account, AccountStatus},
window::should_reset_usage_window,
};
use crate::telemetry::account_usage::store::{
AccountModelUsageDelta, AccountUsageDelta, AccountUsageSnapshot, AccountUsageStore,
AccountUsageWindow,
};
use crate::upstream::openai::protocol::events::{
parse_rate_limit_headers, rate_limit_quota, TokenUsage as CodexTokenUsage,
@@ -528,7 +531,7 @@ impl AccountPool {
let refreshed_accounts = self.refresh_account_statuses(now);
let is_capacity_account = |account: &Account| {
account.status == AccountStatus::Active
&& AccountService::quota_available_at(account, now, self.options.skip_quota_limited)
&& candidates::quota_available_at(account, now, self.options.skip_quota_limited)
};
let active_accounts = self
.accounts
+135 -13
View File
@@ -1,9 +1,36 @@
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>>,
}
@@ -11,12 +38,14 @@ 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,
))),
@@ -36,6 +65,22 @@ impl RuntimeAccountPoolService {
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,
@@ -49,11 +94,25 @@ impl RuntimeAccountPoolService {
/// 从持久化账号恢复运行时账号池。
pub async fn restore_from_store(&self) -> Result<usize, RuntimeAccountPoolError> {
let accounts = self
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();
@@ -105,7 +164,7 @@ impl RuntimeAccountPoolService {
if !record_request_usage {
return;
}
if let Err(error) = self.store.record_request(account_id).await {
if let Err(error) = self.usage_store.record_request(account_id).await {
tracing::warn!(
account_id,
error = %error,
@@ -120,7 +179,7 @@ impl RuntimeAccountPoolService {
return;
};
if let Err(error) = self
.store
.usage_store
.record_model_usage_delta(
account_id,
&model,
@@ -202,7 +261,7 @@ impl RuntimeAccountPoolService {
for refreshed in refreshed_states {
if let Err(error) = self
.store
.sync_runtime_account_state(&refreshed.account, refreshed.sync_usage_window)
.sync_runtime_account_state(&refreshed.account)
.await
{
tracing::warn!(
@@ -211,6 +270,22 @@ impl RuntimeAccountPoolService {
"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"
);
}
}
}
}
@@ -264,7 +339,7 @@ impl RuntimeAccountPoolService {
..AccountUsageDelta::default()
};
if let Err(error) = self
.store
.usage_store
.record_usage_delta(account_id, persisted_usage)
.await
{
@@ -275,7 +350,7 @@ impl RuntimeAccountPoolService {
);
}
if let Err(error) = self
.store
.usage_store
.record_model_usage_delta(
account_id,
model,
@@ -321,7 +396,7 @@ impl RuntimeAccountPoolService {
image_request_failures: u64::from(image_generation_requested),
..AccountUsageDelta::default()
};
if let Err(error) = self.store.record_usage_delta(account_id, usage).await {
if let Err(error) = self.usage_store.record_usage_delta(account_id, usage).await {
tracing::warn!(
account_id,
error = %error,
@@ -329,7 +404,7 @@ impl RuntimeAccountPoolService {
);
}
if let Err(error) = self
.store
.usage_store
.record_model_usage_delta(
account_id,
model,
@@ -429,7 +504,7 @@ impl RuntimeAccountPoolService {
if let Some(reset_at) = reset_at {
if let Err(error) = self
.store
.usage_store
.sync_rate_limit_window(account_id, reset_at, limit_window_seconds)
.await
{
@@ -468,12 +543,20 @@ impl RuntimeAccountPoolService {
.get_pool_account(account_id)
.await
.map_err(|_| RuntimeAccountPoolError::Generic)?;
let mut pool = self.pool.lock().await;
if let Some(account) = account {
pool.insert(account);
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(pool.remove(account_id))
Ok(self.pool.lock().await.remove(account_id))
}
/// 设置 Cloudflare 冷却状态。
@@ -506,6 +589,45 @@ impl RuntimeAccountPoolService {
}
}
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 {
+16 -2
View File
@@ -18,6 +18,9 @@ use crate::accounts::{
};
use crate::{
accounts::cookies::PgCookieStore,
telemetry::account_usage::store::{
AccountUsageStore, AccountUsageStoreError, PgAccountUsageStore,
},
upstream::openai::transport::{CodexBackendClient, CodexClientError, CodexRequestContext},
};
@@ -28,6 +31,7 @@ const COOLDOWN_REFRESH_GRACE_SECS: i64 = 30;
/// 运行时 quota refresh 服务。
pub struct RuntimeQuotaRefreshService {
store: PgAccountStore,
usage_store: PgAccountUsageStore,
codex: Arc<CodexBackendClient>,
min_refresh_interval_secs: u64,
request_spacing: Duration,
@@ -38,9 +42,14 @@ pub struct RuntimeQuotaRefreshService {
impl RuntimeQuotaRefreshService {
/// 构造默认 quota refresh 服务。
pub fn new(store: PgAccountStore, codex: Arc<CodexBackendClient>) -> Self {
pub fn new(
store: PgAccountStore,
usage_store: PgAccountUsageStore,
codex: Arc<CodexBackendClient>,
) -> Self {
Self {
store,
usage_store,
codex,
min_refresh_interval_secs: MIN_REFRESH_INTERVAL_SECS,
request_spacing: Self::default_request_spacing(),
@@ -57,11 +66,13 @@ impl RuntimeQuotaRefreshService {
/// 使用自定义最小刷新间隔构造 quota refresh 服务。
pub fn with_min_refresh_interval_secs(
store: PgAccountStore,
usage_store: PgAccountUsageStore,
codex: Arc<CodexBackendClient>,
min_refresh_interval_secs: u64,
) -> Self {
Self {
store,
usage_store,
codex,
min_refresh_interval_secs,
request_spacing: Self::default_request_spacing(),
@@ -182,7 +193,7 @@ impl RuntimeQuotaRefreshService {
summary.missing += 1;
}
if let Some(reset_at) = reset_at {
self.store
self.usage_store
.sync_rate_limit_window(
&account.id,
reset_at,
@@ -267,6 +278,9 @@ pub enum QuotaRefreshServiceError {
/// 账号写入失败。
#[error("failed to persist quota refresh result: {0}")]
Store(#[from] PgAccountStoreError),
/// 用量窗口写入失败。
#[error("failed to persist quota usage window: {0}")]
UsageStore(#[from] AccountUsageStoreError),
/// 上游 usage 请求失败。
#[error("failed to fetch quota usage: {0}")]
Codex(#[from] CodexClientError),
+2 -2
View File
@@ -1,9 +1,9 @@
//! Token 刷新策略、租约与运行时服务。
mod lease;
mod runtime;
mod policy;
mod service;
pub use lease::*;
pub use runtime::*;
pub use policy::*;
pub use service::*;
+485
View File
@@ -0,0 +1,485 @@
//! 刷新租约存储、JWT 解码与账号 claims 验证。
use std::{
collections::hash_map::DefaultHasher,
hash::{Hash, Hasher},
sync::{Arc, RwLock},
time::Duration as StdDuration,
};
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
use chrono::{DateTime, Duration, TimeZone, Utc};
use serde_json::{Map, Value};
use tokio::sync::{watch, Semaphore};
use crate::accounts::account::{Account, AccountStatus};
use crate::upstream::openai::token_client::{RefreshFailure, TokenPair, TokenRefresher};
pub(crate) const PERMANENT_FAILURE_CONFIRMATION_THRESHOLD: usize = 2;
pub(crate) const RECOVERY_DELAY_SECONDS: i64 = 10 * 60;
pub(crate) const RETRY_DELAY_JITTER: f64 = 0.30;
pub(crate) const RECOVERY_DELAY_JITTER: f64 = 0.20;
const MAX_REFRESH_ATTEMPTS: usize = 5;
const REFRESH_RETRY_BASE_DELAY_MILLIS: u64 = 5_000;
const REFRESH_RETRY_MAX_DELAY_MILLIS: u64 = 300_000;
// ---------------------------------------------------------------------------
// JWT 解码
// ---------------------------------------------------------------------------
/// JWT 过期状态。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JwtExpiry {
/// token 已过期。
Expired,
/// token 仍然有效。
Valid,
/// token 缺失、格式错误或不包含可解析的 exp。
MissingOrInvalid,
}
/// 按给定时间点判断 JWT 的 `exp` 是否已过期。
pub fn jwt_expiry(token: &str, now: DateTime<Utc>) -> JwtExpiry {
let Some(exp) = jwt_exp(token) else {
return JwtExpiry::MissingOrInvalid;
};
if now.timestamp() >= exp {
JwtExpiry::Expired
} else {
JwtExpiry::Valid
}
}
/// 读取 JWT `exp` 并转换成 UTC 时间。
pub fn jwt_expiration(token: &str) -> Option<DateTime<Utc>> {
Utc.timestamp_opt(jwt_exp(token)?, 0).single()
}
fn jwt_exp(token: &str) -> Option<i64> {
let mut parts = token.split('.');
let _header = parts.next()?;
let payload = parts.next()?;
let _signature = parts.next()?;
if parts.next().is_some() {
return None;
}
let decoded = URL_SAFE_NO_PAD.decode(payload).ok()?;
let value = serde_json::from_slice::<Value>(&decoded).ok()?;
value.get("exp")?.as_i64()
}
// ---------------------------------------------------------------------------
// 账号 claims 解码(手动创建/导入时验证 JWT)
// ---------------------------------------------------------------------------
/// 手动创建账号时从 JWT 提取的 claims。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ManualAccountClaims {
/// ChatGPT 账号 ID。
pub account_id: Option<String>,
/// ChatGPT 用户 ID。
pub user_id: Option<String>,
/// 邮箱。
pub email: Option<String>,
/// 订阅计划类型。
pub plan_type: Option<String>,
/// access token 过期时间。
pub expires_at: DateTime<Utc>,
}
/// 从 JWT 中解码 payload 部分。
pub fn decode_jwt_payload(token: &str) -> Option<Map<String, Value>> {
let mut parts = token.split('.');
let _header = parts.next()?;
let payload = parts.next()?;
if payload.is_empty() {
return None;
}
let bytes = URL_SAFE_NO_PAD.decode(payload).ok()?;
serde_json::from_slice::<Value>(&bytes)
.ok()?
.as_object()
.cloned()
}
/// 从 JWT 中提取手动创建账号所需的 claims。
pub fn manual_account_claims(
token: &str,
now: DateTime<Utc>,
) -> Result<ManualAccountClaims, &'static str> {
let payload = decode_jwt_payload(token).ok_or("Invalid JWT format")?;
let exp = payload
.get("exp")
.and_then(Value::as_i64)
.ok_or("Token is expired")?;
if now.timestamp() >= exp {
return Err("Token is expired");
}
let expires_at = DateTime::<Utc>::from_timestamp(exp, 0).ok_or("Invalid JWT exp claim")?;
let auth = payload
.get("https://api.openai.com/auth")
.and_then(Value::as_object)
.ok_or("Token missing OpenAI auth claim")?;
let account_id = string_claim(auth, "chatgpt_account_id");
let profile = payload
.get("https://api.openai.com/profile")
.and_then(Value::as_object);
let user_id = string_claim(auth, "chatgpt_user_id")
.or_else(|| profile.and_then(|profile| string_claim(profile, "chatgpt_user_id")))
.or_else(|| string_claim(auth, "user_id"))
.or_else(|| profile.and_then(|profile| string_claim(profile, "user_id")));
let plan_type = string_claim(auth, "chatgpt_plan_type")
.or_else(|| profile.and_then(|profile| string_claim(profile, "chatgpt_plan_type")));
let email = profile.and_then(|profile| string_claim(profile, "email"));
Ok(ManualAccountClaims {
account_id,
user_id,
email,
plan_type,
expires_at,
})
}
fn string_claim(map: &Map<String, Value>, key: &str) -> Option<String> {
map.get(key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
// ---------------------------------------------------------------------------
// 刷新策略和上游端口
// ---------------------------------------------------------------------------
/// 刷新任务的调度策略。
#[derive(Debug, Clone, Copy)]
pub struct RefreshPolicy {
/// 提前多久开始刷新访问令牌。
pub refresh_margin_seconds: u64,
/// 允许并发执行的刷新任务数。
pub refresh_concurrency: u32,
}
/// 单次刷新扫描摘要。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct TokenRefreshSummary {
/// 扫描账号数。
pub scanned: usize,
/// 成功刷新 token 的账号数。
pub refreshed: usize,
/// 仅更新状态的账号数。
pub status_updated: usize,
/// 无需刷新的账号数。
pub skipped: usize,
/// 传输失败账号数。
pub failed: usize,
}
/// 账号级刷新定时器调度摘要。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct TokenTimerSummary {
/// 扫描账号数。
pub scanned: usize,
/// 已调度的未来刷新定时器数。
pub scheduled: usize,
/// 已调度的立即刷新定时器数。
pub immediate: usize,
/// 无需调度的账号数。
pub skipped: usize,
/// 被替换的既有定时器数。
pub replaced: usize,
}
impl TokenTimerSummary {
/// 返回本轮调度发生变化的定时器数量。
pub fn changed(&self) -> usize {
self.scheduled + self.immediate
}
}
/// 运行时共享刷新策略。
#[derive(Clone)]
pub struct RuntimeRefreshPolicy {
policy: Arc<RwLock<RefreshPolicy>>,
semaphore: Arc<RwLock<Arc<Semaphore>>>,
}
impl RuntimeRefreshPolicy {
/// 构造共享刷新策略。
pub fn new(policy: RefreshPolicy) -> Self {
Self {
policy: Arc::new(RwLock::new(policy)),
semaphore: Arc::new(RwLock::new(Arc::new(Semaphore::new(
policy.refresh_concurrency.max(1) as usize,
)))),
}
}
/// 返回当前策略快照。
pub fn snapshot(&self) -> RefreshPolicy {
*self
.policy
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
/// 更新当前策略。
pub fn update(&self, policy: RefreshPolicy) {
*self
.policy
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner) = policy;
*self
.semaphore
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner) =
Arc::new(Semaphore::new(policy.refresh_concurrency.max(1) as usize));
}
/// 返回当前提前刷新秒数。
pub fn refresh_margin_seconds(&self) -> u64 {
self.snapshot().refresh_margin_seconds
}
/// 持续接收运行时设置并更新刷新策略。
pub async fn subscribe_settings(
self,
mut receiver: watch::Receiver<crate::settings::SettingsSnapshot>,
) {
while receiver.changed().await.is_ok() {
let settings = receiver.borrow_and_update();
self.update(RefreshPolicy {
refresh_margin_seconds: settings.refresh_margin_seconds,
refresh_concurrency: settings.refresh_concurrency,
});
}
}
fn semaphore(&self) -> Arc<Semaphore> {
self.semaphore
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone()
}
}
impl From<RefreshPolicy> for RuntimeRefreshPolicy {
fn from(policy: RefreshPolicy) -> Self {
Self::new(policy)
}
}
/// 按账号稳定扰动刷新提前量,范围与 TS 版 `jitter(margin, 0.15)` 一致。
pub(crate) fn jittered_refresh_at(
account_id: &str,
expires_at: DateTime<Utc>,
margin_seconds: u64,
) -> DateTime<Utc> {
expires_at - Duration::seconds(jittered_refresh_margin_seconds(account_id, margin_seconds))
}
fn jittered_refresh_margin_seconds(account_id: &str, margin_seconds: u64) -> i64 {
if margin_seconds == 0 {
return 0;
}
let mut hasher = DefaultHasher::new();
account_id.hash(&mut hasher);
let unit = hasher.finish() as f64 / u64::MAX as f64;
let factor = 0.85 + unit * 0.30;
(margin_seconds as f64 * factor)
.round()
.clamp(0.0, i64::MAX as f64) as i64
}
/// 调度器自身的执行错误。
#[derive(Debug, thiserror::Error)]
pub enum RefreshError {
/// 并发限制信号量已关闭。
#[error("refresh task semaphore closed")]
ConcurrencyClosed,
/// 刷新请求在到达服务端前失败,可安全复用当前 refresh token 重试。
#[error("refresh transport failed before server processing")]
RetryableTransport,
/// 刷新请求在传输层失败,refresh token 可能已经被服务端消费。
#[error("refresh transport failed after possible server processing")]
Transport,
}
/// 负责执行单账号刷新策略的调度器。
#[derive(Clone)]
pub struct RefreshScheduler<C> {
policy: RuntimeRefreshPolicy,
client: Arc<C>,
}
impl<C> RefreshScheduler<C>
where
C: TokenRefresher,
{
/// 使用策略和上游刷新端口构造调度器。
pub fn new(policy: impl Into<RuntimeRefreshPolicy>, client: C) -> Self {
Self {
policy: policy.into(),
client: Arc::new(client),
}
}
/// 更新刷新策略。
pub fn update_policy(&self, policy: RefreshPolicy) {
self.policy.update(policy);
}
/// 在给定时间点按到期策略刷新账号。
///
/// 当账号不需要刷新时返回原账号快照;当刷新失败但属于可映射的领域错误时,
/// 返回更新过状态的账号。
///
/// # Errors
///
/// 当并发控制已关闭或刷新传输失败时返回 [`RefreshError`]。
pub async fn refresh_account_at(
&self,
account: &Account,
now: DateTime<Utc>,
) -> Result<Account, RefreshError> {
if !self.should_refresh_account_at(account, now) {
return Ok(account.clone());
}
let Some(refresh_token) = account.refresh_token.as_deref() else {
let mut expired = account.clone();
expired.status = AccountStatus::Expired;
return Ok(expired);
};
let semaphore = self.policy.semaphore();
let _permit = semaphore
.acquire()
.await
.map_err(|_| RefreshError::ConcurrencyClosed)?;
match self.client.refresh(refresh_token).await {
Ok(token_pair) => Ok(apply_token_pair(account, token_pair)),
Err(RefreshFailure::RetryableTransport) => Err(RefreshError::RetryableTransport),
Err(RefreshFailure::Transport) => Err(RefreshError::Transport),
Err(error) => Ok(apply_refresh_failure(account, error)),
}
}
/// 判断账号在给定时间点是否需要刷新。
pub fn should_refresh_account_at(&self, account: &Account, now: DateTime<Utc>) -> bool {
token_refresh_status_eligible(account.status)
&& self.should_refresh_before_expiry(account, now)
}
fn should_refresh_before_expiry(&self, account: &Account, now: DateTime<Utc>) -> bool {
if let Some(next_refresh_at) = account.next_refresh_at {
return next_refresh_at <= now;
}
account
.access_token_expires_at
.is_some_and(|expires_at| self.refresh_at(account, expires_at) <= now)
}
fn refresh_at(&self, account: &Account, expires_at: DateTime<Utc>) -> DateTime<Utc> {
let policy = self.policy.snapshot();
jittered_refresh_at(
account.id.as_str(),
expires_at,
policy.refresh_margin_seconds,
)
}
}
pub(crate) fn token_refresh_status_eligible(status: AccountStatus) -> bool {
matches!(
status,
AccountStatus::Active | AccountStatus::QuotaExhausted
)
}
/// 将新的 token 对应用到账号快照上。
pub fn apply_token_pair(account: &Account, token_pair: TokenPair) -> Account {
let mut refreshed = account.clone();
refreshed.access_token = token_pair.access_token;
// 刷新响应不返回 refresh_token 时,继续保留已有值,避免永久失去刷新能力。
if let Some(refresh_token) = token_pair.refresh_token {
refreshed.refresh_token = Some(refresh_token);
}
refreshed
}
/// 将刷新失败映射为账号状态变更。
pub fn apply_refresh_failure(account: &Account, failure: RefreshFailure) -> Account {
let mut updated = account.clone();
updated.status = match failure {
RefreshFailure::InvalidGrant => AccountStatus::Expired,
RefreshFailure::Banned => AccountStatus::Banned,
RefreshFailure::RetryableTransport | RefreshFailure::Transport => account.status,
};
updated
}
pub(crate) fn is_permanent_refresh_failure_status(status: AccountStatus) -> bool {
matches!(status, AccountStatus::Expired | AccountStatus::Banned)
}
pub(crate) fn default_refresh_retry_delays() -> Vec<StdDuration> {
(0..MAX_REFRESH_ATTEMPTS.saturating_sub(1))
.map(|attempt_index| {
let multiplier = 3_u64.saturating_pow(attempt_index as u32);
let millis = REFRESH_RETRY_BASE_DELAY_MILLIS
.saturating_mul(multiplier)
.min(REFRESH_RETRY_MAX_DELAY_MILLIS);
StdDuration::from_millis(millis)
})
.collect()
}
pub(crate) fn stable_jittered_duration(
account_id: &str,
base: StdDuration,
variance: f64,
salt: &str,
) -> StdDuration {
if base.is_zero() {
return StdDuration::ZERO;
}
let mut hasher = DefaultHasher::new();
account_id.hash(&mut hasher);
salt.hash(&mut hasher);
let unit = hasher.finish() as f64 / u64::MAX as f64;
let factor = (1.0 - variance) + unit * variance * 2.0;
let millis = (base.as_millis() as f64 * factor)
.round()
.clamp(0.0, u64::MAX as f64) as u64;
StdDuration::from_millis(millis)
}
pub(crate) fn scheduled_at_from_delay(now: DateTime<Utc>, delay: StdDuration) -> DateTime<Utc> {
match chrono::Duration::from_std(delay) {
Ok(delay) => now + delay,
Err(_) => now,
}
}
pub(crate) fn duration_until(now: DateTime<Utc>, scheduled_at: DateTime<Utc>) -> StdDuration {
(scheduled_at - now).to_std().unwrap_or(StdDuration::ZERO)
}
pub(crate) fn same_refresh_time(left: Option<DateTime<Utc>>, right: Option<DateTime<Utc>>) -> bool {
match (left, right) {
(Some(left), Some(right)) => left.timestamp() == right.timestamp(),
(None, None) => true,
_ => false,
}
}
-777
View File
@@ -1,777 +0,0 @@
//! 运行时 token refresh 业务服务。
use std::{
collections::{hash_map::DefaultHasher, HashMap, HashSet},
future::Future,
hash::{Hash, Hasher},
pin::Pin,
sync::Arc,
time::Duration,
};
use chrono::{DateTime, Utc};
use secrecy::{ExposeSecret, SecretString};
use thiserror::Error;
use tokio::{sync::Mutex, task::JoinHandle, time::sleep};
use tracing::{debug, info, warn};
use uuid::Uuid;
use crate::accounts::{
account::{Account, AccountStatus},
refresh::{
jittered_refresh_at, jwt_expiration, token_refresh_status_eligible, RedisRefreshLeaseStore,
RedisRefreshLeaseStoreError, RefreshError, RefreshPolicy, RefreshScheduler,
RuntimeRefreshPolicy,
},
store::{
AccountClaimsUpdate, AccountStore, PgAccountStore, PgAccountStoreError, StoredAccount,
},
};
use crate::upstream::openai::token_client::TokenRefresher;
const REFRESH_LEASE_TTL_SECONDS: i64 = 5 * 60;
const MAX_REFRESH_ATTEMPTS: usize = 5;
const REFRESH_RETRY_BASE_DELAY_MILLIS: u64 = 5_000;
const REFRESH_RETRY_MAX_DELAY_MILLIS: u64 = 300_000;
const PERMANENT_FAILURE_CONFIRMATION_THRESHOLD: usize = 2;
const RECOVERY_DELAY_SECONDS: i64 = 10 * 60;
const RETRY_DELAY_JITTER: f64 = 0.30;
const RECOVERY_DELAY_JITTER: f64 = 0.20;
/// 运行时 token refresh 服务。
pub struct RuntimeTokenRefreshService<C>
where
C: TokenRefresher,
{
store: PgAccountStore,
scheduler: Arc<RefreshScheduler<C>>,
policy: RuntimeRefreshPolicy,
refresh_leases: Option<RedisRefreshLeaseStore>,
lease_owner: String,
retry_delays: Vec<Duration>,
in_flight: Arc<Mutex<HashSet<String>>>,
timers: Arc<Mutex<HashMap<String, ScheduledRefreshTimer>>>,
}
struct ScheduledRefreshTimer {
scheduled_at: DateTime<Utc>,
handle: JoinHandle<()>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ScheduleTimerResult {
Scheduled,
Replaced,
Unchanged,
}
impl<C> Clone for RuntimeTokenRefreshService<C>
where
C: TokenRefresher,
{
fn clone(&self) -> Self {
Self {
store: self.store.clone(),
scheduler: self.scheduler.clone(),
policy: self.policy.clone(),
refresh_leases: self.refresh_leases.clone(),
lease_owner: self.lease_owner.clone(),
retry_delays: self.retry_delays.clone(),
in_flight: self.in_flight.clone(),
timers: self.timers.clone(),
}
}
}
impl<C> RuntimeTokenRefreshService<C>
where
C: TokenRefresher,
{
/// 构造运行时 token refresh 服务。
pub fn new(store: PgAccountStore, policy: impl Into<RuntimeRefreshPolicy>, client: C) -> Self {
let policy = policy.into();
Self {
store,
scheduler: Arc::new(RefreshScheduler::new(policy.clone(), client)),
policy,
refresh_leases: None,
lease_owner: refresh_lease_owner(),
retry_delays: default_refresh_retry_delays(),
in_flight: Arc::new(Mutex::new(HashSet::new())),
timers: Arc::new(Mutex::new(HashMap::new())),
}
}
/// 使用刷新租约存储保护账号刷新。
pub fn with_refresh_lease_store(mut self, refresh_leases: RedisRefreshLeaseStore) -> Self {
self.refresh_leases = Some(refresh_leases);
self
}
/// 更新刷新策略。
pub fn update_policy(&self, policy: RefreshPolicy) {
self.policy.update(policy);
}
/// 使用自定义刷新重试延迟。
pub fn with_retry_delays<I>(mut self, retry_delays: I) -> Self
where
I: IntoIterator<Item = Duration>,
{
self.retry_delays = retry_delays.into_iter().collect();
self
}
/// 返回给定账号集合中当前正在执行 token 刷新的账号 ID。
pub async fn refreshing_account_ids(
&self,
account_ids: &[String],
now: DateTime<Utc>,
) -> TokenRefreshServiceResult<HashSet<String>> {
if account_ids.is_empty() {
return Ok(HashSet::new());
}
let requested_ids = account_ids
.iter()
.map(String::as_str)
.collect::<HashSet<_>>();
let mut refreshing = self
.in_flight
.lock()
.await
.iter()
.filter(|account_id| requested_ids.contains(account_id.as_str()))
.cloned()
.collect::<HashSet<_>>();
if let Some(refresh_leases) = self.refresh_leases.as_ref() {
refreshing.extend(
refresh_leases
.active_account_ids(account_ids, now)
.await
.map_err(TokenRefreshServiceError::Lease)?,
);
}
Ok(refreshing)
}
/// 执行一次账号级刷新定时器调度。
pub async fn schedule_account_timers_once(
&self,
) -> TokenRefreshServiceResult<TokenTimerSummary> {
self.schedule_account_timers_once_at(Utc::now()).await
}
/// 在指定时间点执行一次账号级刷新定时器调度。
pub async fn schedule_account_timers_once_at(
&self,
now: DateTime<Utc>,
) -> TokenRefreshServiceResult<TokenTimerSummary> {
let accounts = self
.store
.list_pool_accounts()
.await
.map_err(TokenRefreshServiceError::AccountStore)?;
let mut summary = TokenTimerSummary {
scanned: accounts.len(),
..TokenTimerSummary::default()
};
for account in accounts {
if self.clear_refresh_schedule_if_ineligible(&account).await? {
summary.skipped += 1;
continue;
}
let Some(delay) = self.timer_delay_for_account(&account, now) else {
summary.skipped += 1;
continue;
};
if delay.is_zero() {
if self.clear_scheduled_timer(&account.id).await {
summary.replaced += 1;
}
let outcome = self.refresh_scheduled_account(&account.id, now).await?;
if matches!(outcome, TokenRefreshOutcome::Skipped) {
summary.skipped += 1;
} else {
summary.immediate += 1;
}
continue;
}
let scheduled_at = scheduled_at_from_delay(now, delay);
self.persist_next_refresh_at_if_changed(&account, Some(scheduled_at))
.await?;
let schedule_result = self
.schedule_account_timer(account.id.clone(), delay, scheduled_at)
.await;
match schedule_result {
ScheduleTimerResult::Unchanged => {
summary.skipped += 1;
continue;
}
ScheduleTimerResult::Replaced => summary.replaced += 1,
ScheduleTimerResult::Scheduled => {}
}
if delay.is_zero() {
summary.immediate += 1;
} else {
summary.scheduled += 1;
}
}
Ok(summary)
}
/// 返回当前账号级刷新定时器数量。
pub async fn scheduled_timer_count(&self) -> usize {
self.timers.lock().await.len()
}
/// 在指定时间点执行一次到期刷新扫描。
pub async fn refresh_due_accounts_once_at(
&self,
now: DateTime<Utc>,
) -> TokenRefreshServiceResult<TokenRefreshSummary> {
let accounts = self
.store
.list_pool_accounts()
.await
.map_err(TokenRefreshServiceError::AccountStore)?;
let mut summary = TokenRefreshSummary {
scanned: accounts.len(),
..TokenRefreshSummary::default()
};
for account in accounts {
if self.clear_refresh_schedule_if_ineligible(&account).await? {
summary.skipped += 1;
continue;
}
if account.next_refresh_at.is_some_and(|value| value > now) {
summary.skipped += 1;
continue;
}
if !self.scheduler.should_refresh_account_at(&account, now) {
summary.skipped += 1;
continue;
}
if !self.try_mark_in_flight(&account.id).await {
summary.skipped += 1;
continue;
}
let lease_acquired = match self.try_acquire_refresh_lease(&account.id, now).await {
Ok(acquired) => acquired,
Err(error) => {
self.release_in_flight(&account.id).await;
return Err(error);
}
};
if !lease_acquired {
self.release_in_flight(&account.id).await;
summary.skipped += 1;
continue;
}
let outcome = self
.refresh_account_with_status_transitions(&account, now)
.await;
self.release_in_flight(&account.id).await;
self.release_refresh_lease(&account.id).await?;
let outcome = outcome?;
self.schedule_next_refresh_after_outcome(&outcome, Utc::now())
.await;
match outcome {
TokenRefreshOutcome::Refreshed(_) => summary.refreshed += 1,
TokenRefreshOutcome::StatusUpdated => summary.status_updated += 1,
TokenRefreshOutcome::Skipped => summary.skipped += 1,
TokenRefreshOutcome::Failed => summary.failed += 1,
};
}
Ok(summary)
}
fn timer_delay_for_account(&self, account: &Account, now: DateTime<Utc>) -> Option<Duration> {
if let Some(next_refresh_at) = account.next_refresh_at.filter(|value| *value > now) {
return Some(duration_until(now, next_refresh_at));
}
let refresh_at = self.computed_refresh_at(account)?;
if refresh_at <= now {
return Some(Duration::ZERO);
}
(refresh_at - now).to_std().ok()
}
async fn clear_refresh_schedule_if_ineligible(
&self,
account: &Account,
) -> TokenRefreshServiceResult<bool> {
if account.refresh_token.is_some() && token_refresh_status_eligible(account.status) {
return Ok(false);
}
self.clear_scheduled_timer(&account.id).await;
if account.next_refresh_at.is_some() {
self.persist_next_refresh_at_if_changed(account, None)
.await?;
}
Ok(true)
}
fn computed_refresh_at(&self, account: &Account) -> Option<DateTime<Utc>> {
let expires_at = account
.access_token_expires_at
.or_else(|| jwt_expiration(&account.access_token))?;
Some(jittered_refresh_at(
&account.id,
expires_at,
self.policy.refresh_margin_seconds(),
))
}
async fn persist_next_refresh_at_if_changed(
&self,
account: &Account,
next_refresh_at: Option<DateTime<Utc>>,
) -> TokenRefreshServiceResult<()> {
if same_refresh_time(account.next_refresh_at, next_refresh_at) {
return Ok(());
}
persist_next_refresh_at(&self.store, &account.id, next_refresh_at).await
}
fn schedule_account_timer(
&self,
account_id: String,
delay: Duration,
scheduled_at: DateTime<Utc>,
) -> Pin<Box<dyn Future<Output = ScheduleTimerResult> + Send + '_>> {
Box::pin(async move {
{
let timers = self.timers.lock().await;
if timers
.get(&account_id)
.is_some_and(|timer| timer.scheduled_at == scheduled_at)
{
return ScheduleTimerResult::Unchanged;
}
}
let replaced = self.clear_scheduled_timer(&account_id).await;
let task = self.clone();
let timer_account_id = account_id.clone();
let handle = tokio::spawn(async move {
if !delay.is_zero() {
sleep(delay).await;
}
task.timers.lock().await.remove(&timer_account_id);
if let Err(error) = task
.refresh_scheduled_account(&timer_account_id, scheduled_at)
.await
{
warn!(
account_id = %timer_account_id,
error = %error,
"scheduled token refresh failed"
);
}
});
self.timers.lock().await.insert(
account_id,
ScheduledRefreshTimer {
scheduled_at,
handle,
},
);
if replaced {
ScheduleTimerResult::Replaced
} else {
ScheduleTimerResult::Scheduled
}
})
}
async fn clear_scheduled_timer(&self, account_id: &str) -> bool {
if let Some(timer) = self.timers.lock().await.remove(account_id) {
timer.handle.abort();
true
} else {
false
}
}
/// 清理全部已调度的账号级刷新定时器。
pub async fn clear_scheduled_timers(&self) {
let mut timers = self.timers.lock().await;
for (_, timer) in timers.drain() {
timer.handle.abort();
}
}
async fn refresh_scheduled_account(
&self,
account_id: &str,
now: DateTime<Utc>,
) -> TokenRefreshServiceResult<TokenRefreshOutcome> {
let Some(account) = self.store.get(account_id).await? else {
return Ok(TokenRefreshOutcome::Skipped);
};
let account = stored_account_to_refresh_account(account);
if self.clear_refresh_schedule_if_ineligible(&account).await? {
return Ok(TokenRefreshOutcome::Skipped);
}
if account.next_refresh_at.is_some_and(|value| value > now) {
return Ok(TokenRefreshOutcome::Skipped);
}
if !self.try_mark_in_flight(&account.id).await {
return Ok(TokenRefreshOutcome::Skipped);
}
let lease_acquired = match self.try_acquire_refresh_lease(&account.id, now).await {
Ok(acquired) => acquired,
Err(error) => {
self.release_in_flight(&account.id).await;
return Err(error);
}
};
if !lease_acquired {
self.release_in_flight(&account.id).await;
return Ok(TokenRefreshOutcome::Skipped);
}
let outcome = self
.refresh_account_with_status_transitions(&account, now)
.await;
self.release_in_flight(&account.id).await;
self.release_refresh_lease(&account.id).await?;
let outcome = outcome?;
self.schedule_next_refresh_after_outcome(&outcome, Utc::now())
.await;
Ok(outcome)
}
async fn schedule_next_refresh_after_outcome(
&self,
outcome: &TokenRefreshOutcome,
now: DateTime<Utc>,
) {
let TokenRefreshOutcome::Refreshed(account) = outcome else {
return;
};
let Some(delay) = self.timer_delay_for_account(account, now) else {
return;
};
if delay.is_zero() {
debug!(
account_id = %account.id,
"refreshed token is already within refresh margin; skipping next timer"
);
return;
}
let scheduled_at = scheduled_at_from_delay(now, delay);
self.schedule_account_timer(account.id.clone(), delay, scheduled_at)
.await;
}
async fn refresh_account_with_status_transitions(
&self,
account: &Account,
now: DateTime<Utc>,
) -> TokenRefreshServiceResult<TokenRefreshOutcome> {
let max_attempts = self.retry_delays.len() + 1;
let mut permanent_failures = 0;
for attempt_index in 0..max_attempts {
let Some(attempt) = self.prepare_refresh_attempt(account, now).await? else {
return Ok(TokenRefreshOutcome::Skipped);
};
if attempt.refresh_token.is_none() {
persist_status(&self.store, &attempt.id, AccountStatus::Expired).await?;
persist_next_refresh_at(&self.store, &attempt.id, None).await?;
return Ok(TokenRefreshOutcome::StatusUpdated);
}
match self.scheduler.refresh_account_at(&attempt, now).await {
Ok(updated) if is_permanent_refresh_failure_status(updated.status) => {
permanent_failures += 1;
if permanent_failures < PERMANENT_FAILURE_CONFIRMATION_THRESHOLD
&& self.sleep_before_retry(&attempt.id, attempt_index).await
{
continue;
}
persist_status(&self.store, &attempt.id, updated.status).await?;
persist_next_refresh_at(&self.store, &attempt.id, None).await?;
return Ok(TokenRefreshOutcome::StatusUpdated);
}
Ok(mut updated) => {
updated.access_token_expires_at =
jwt_expiration(&updated.access_token).or(updated.access_token_expires_at);
updated.next_refresh_at = self.computed_refresh_at(&updated);
let refresh_token_rotated = updated.refresh_token != attempt.refresh_token;
persist_token_update(&self.store, &updated).await?;
info!(
account_id = %updated.id,
attempt = attempt_index + 1,
status = %updated.status,
access_token_expires_at = ?updated.access_token_expires_at,
next_refresh_at = ?updated.next_refresh_at,
refresh_token_present = updated.refresh_token.is_some(),
refresh_token_rotated,
"token refresh succeeded"
);
return Ok(TokenRefreshOutcome::Refreshed(Box::new(updated)));
}
Err(RefreshError::RetryableTransport | RefreshError::Transport)
if self.sleep_before_retry(&attempt.id, attempt_index).await => {}
Err(RefreshError::RetryableTransport | RefreshError::Transport) => {
persist_status(&self.store, &attempt.id, attempt.status).await?;
let recovery_delay = stable_jittered_duration(
&attempt.id,
Duration::from_secs(RECOVERY_DELAY_SECONDS as u64),
RECOVERY_DELAY_JITTER,
"recovery",
);
let next_refresh_at = scheduled_at_from_delay(now, recovery_delay);
persist_next_refresh_at(&self.store, &attempt.id, Some(next_refresh_at))
.await?;
return Ok(TokenRefreshOutcome::Failed);
}
Err(error) => return Err(TokenRefreshServiceError::Refresh(error)),
}
}
Ok(TokenRefreshOutcome::Skipped)
}
async fn prepare_refresh_attempt(
&self,
account: &Account,
now: DateTime<Utc>,
) -> TokenRefreshServiceResult<Option<Account>> {
let Some(stored) = self.store.get(&account.id).await? else {
return Err(TokenRefreshServiceError::AccountNotFound(
account.id.clone(),
));
};
let attempt = stored_account_to_refresh_account(stored);
if !token_refresh_status_eligible(attempt.status)
|| attempt.next_refresh_at.is_some_and(|value| value > now)
|| attempt.refresh_token != account.refresh_token
{
return Ok(None);
}
if !self.scheduler.should_refresh_account_at(&attempt, now) {
return Ok(None);
}
Ok(Some(attempt))
}
async fn sleep_before_retry(&self, account_id: &str, attempt_index: usize) -> bool {
let Some(delay) = self.retry_delays.get(attempt_index).copied() else {
return false;
};
let delay = stable_jittered_duration(
account_id,
delay,
RETRY_DELAY_JITTER,
&format!("retry:{attempt_index}"),
);
if !delay.is_zero() {
sleep(delay).await;
}
true
}
async fn try_mark_in_flight(&self, account_id: &str) -> bool {
self.in_flight.lock().await.insert(account_id.to_string())
}
async fn release_in_flight(&self, account_id: &str) {
self.in_flight.lock().await.remove(account_id);
}
async fn try_acquire_refresh_lease(
&self,
account_id: &str,
now: DateTime<Utc>,
) -> TokenRefreshServiceResult<bool> {
let Some(refresh_leases) = self.refresh_leases.as_ref() else {
return Ok(true);
};
refresh_leases
.try_acquire(
account_id,
&self.lease_owner,
now + chrono::Duration::seconds(REFRESH_LEASE_TTL_SECONDS),
now,
)
.await
.map_err(TokenRefreshServiceError::Lease)
}
async fn release_refresh_lease(&self, account_id: &str) -> TokenRefreshServiceResult<()> {
let Some(refresh_leases) = self.refresh_leases.as_ref() else {
return Ok(());
};
refresh_leases
.release(account_id, &self.lease_owner)
.await
.map(|_| ())
.map_err(TokenRefreshServiceError::Lease)
}
}
mod types;
use types::*;
async fn persist_token_update(
store: &PgAccountStore,
account: &Account,
) -> TokenRefreshServiceResult<()> {
let updated = store
.update_from_claims(
&account.id,
AccountClaimsUpdate {
email: account.email.clone(),
account_id: account.account_id.clone(),
user_id: account.user_id.clone(),
plan_type: account.plan_type.clone(),
access_token: SecretString::new(account.access_token.clone().into()),
refresh_token: account
.refresh_token
.clone()
.map(|token| SecretString::new(token.into())),
access_token_expires_at: jwt_expiration(&account.access_token)
.or(account.access_token_expires_at),
next_refresh_at: account.next_refresh_at,
status: account.status,
},
)
.await?;
if updated {
Ok(())
} else {
Err(TokenRefreshServiceError::AccountNotFound(
account.id.clone(),
))
}
}
async fn persist_status(
store: &PgAccountStore,
account_id: &str,
status: AccountStatus,
) -> TokenRefreshServiceResult<()> {
let updated = store.set_status(account_id, status).await?;
if updated {
Ok(())
} else {
Err(TokenRefreshServiceError::AccountNotFound(
account_id.to_string(),
))
}
}
async fn persist_next_refresh_at(
store: &PgAccountStore,
account_id: &str,
next_refresh_at: Option<DateTime<Utc>>,
) -> TokenRefreshServiceResult<()> {
let updated = store
.set_next_refresh_at(account_id, next_refresh_at)
.await?;
if updated {
Ok(())
} else {
Err(TokenRefreshServiceError::AccountNotFound(
account_id.to_string(),
))
}
}
fn is_permanent_refresh_failure_status(status: AccountStatus) -> bool {
matches!(status, AccountStatus::Expired | AccountStatus::Banned)
}
fn refresh_lease_owner() -> String {
format!("runtime-token-refresh:{}", Uuid::new_v4().simple())
}
fn default_refresh_retry_delays() -> Vec<Duration> {
(0..MAX_REFRESH_ATTEMPTS.saturating_sub(1))
.map(|attempt_index| {
let multiplier = 3_u64.saturating_pow(attempt_index as u32);
let millis = REFRESH_RETRY_BASE_DELAY_MILLIS
.saturating_mul(multiplier)
.min(REFRESH_RETRY_MAX_DELAY_MILLIS);
Duration::from_millis(millis)
})
.collect()
}
fn stable_jittered_duration(
account_id: &str,
base: Duration,
variance: f64,
salt: &str,
) -> Duration {
if base.is_zero() {
return Duration::ZERO;
}
let mut hasher = DefaultHasher::new();
account_id.hash(&mut hasher);
salt.hash(&mut hasher);
let unit = hasher.finish() as f64 / u64::MAX as f64;
let factor = (1.0 - variance) + unit * variance * 2.0;
let millis = (base.as_millis() as f64 * factor)
.round()
.clamp(0.0, u64::MAX as f64) as u64;
Duration::from_millis(millis)
}
fn scheduled_at_from_delay(now: DateTime<Utc>, delay: Duration) -> DateTime<Utc> {
match chrono::Duration::from_std(delay) {
Ok(delay) => now + delay,
Err(_) => now,
}
}
fn duration_until(now: DateTime<Utc>, scheduled_at: DateTime<Utc>) -> Duration {
(scheduled_at - now).to_std().unwrap_or(Duration::ZERO)
}
fn same_refresh_time(left: Option<DateTime<Utc>>, right: Option<DateTime<Utc>>) -> bool {
match (left, right) {
(Some(left), Some(right)) => left.timestamp() == right.timestamp(),
(None, None) => true,
_ => false,
}
}
@@ -1,109 +0,0 @@
use super::*;
pub(super) enum TokenRefreshOutcome {
Refreshed(Box<Account>),
StatusUpdated,
Skipped,
Failed,
}
/// 单次刷新扫描摘要。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct TokenRefreshSummary {
/// 扫描账号数。
pub scanned: usize,
/// 成功刷新 token 的账号数。
pub refreshed: usize,
/// 仅更新状态的账号数。
pub status_updated: usize,
/// 无需刷新的账号数。
pub skipped: usize,
/// 传输失败账号数。
pub failed: usize,
}
/// 账号级刷新定时器调度摘要。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct TokenTimerSummary {
/// 扫描账号数。
pub scanned: usize,
/// 已调度的未来刷新定时器数。
pub scheduled: usize,
/// 已调度的立即刷新定时器数。
pub immediate: usize,
/// 无需调度的账号数。
pub skipped: usize,
/// 被替换的既有定时器数。
pub replaced: usize,
}
impl TokenTimerSummary {
/// 返回本轮调度发生变化的定时器数量。
pub fn changed(&self) -> usize {
self.scheduled + self.immediate
}
}
/// 令牌刷新任务错误。
#[derive(Debug, Error)]
pub enum TokenRefreshServiceError {
/// 账号存储读取失败。
#[error("failed to list accounts for token refresh: {0}")]
AccountStore(#[from] crate::accounts::store::AccountStoreError),
/// 账号存储写入失败。
#[error("failed to persist token refresh result: {0}")]
Store(#[from] PgAccountStoreError),
/// 刷新租约存储失败。
#[error("failed to coordinate token refresh lease: {0}")]
Lease(#[from] RedisRefreshLeaseStoreError),
/// core 刷新调度失败。
#[error("token refresh scheduler failed: {0}")]
Refresh(#[from] RefreshError),
/// 刷新期间账号不存在。
#[error("account disappeared during token refresh: {0}")]
AccountNotFound(String),
}
/// 令牌刷新任务结果。
pub type TokenRefreshServiceResult<T> = Result<T, TokenRefreshServiceError>;
pub(super) fn stored_account_to_refresh_account(stored: StoredAccount) -> Account {
Account {
id: stored.id,
email: stored.email,
account_id: stored.account_id,
user_id: stored.user_id,
label: stored.label,
plan_type: stored.plan_type,
access_token: stored.access_token.expose_secret().to_string(),
refresh_token: stored
.refresh_token
.map(|token| token.expose_secret().to_string()),
access_token_expires_at: stored.access_token_expires_at,
next_refresh_at: stored.next_refresh_at,
status: stored.status,
quota_limit_reached: false,
quota_verify_required: false,
quota_cooldown_until: None,
cloudflare_cooldown_until: None,
request_count: 0,
empty_response_count: 0,
image_input_tokens: 0,
image_output_tokens: 0,
image_request_count: 0,
image_request_failed_count: 0,
window_request_count: 0,
window_input_tokens: 0,
window_output_tokens: 0,
window_cached_tokens: 0,
window_image_input_tokens: 0,
window_image_output_tokens: 0,
window_image_request_count: 0,
window_image_request_failed_count: 0,
window_started_at: None,
window_reset_at: None,
limit_window_seconds: None,
added_at: stored.added_at,
last_used_at: None,
}
}
File diff suppressed because it is too large Load Diff
+24 -6
View File
@@ -7,7 +7,7 @@
//!
//! 1. 账号状态为 [`AccountStatus::Active`];
//! 2. 在途槽位数未达单账号并发上限;
//! 3. 配额可用([`AccountService::quota_available_at`],受 `skip_quota_limited` 控制);
//! 3. 配额可用(受 `skip_quota_limited` 控制);
//! 4. 模型在账号订阅计划的允许列表内;
//! 5. 不在本次请求的排除列表内;
//! 6. 不处于 Cloudflare 冷却期。
@@ -18,7 +18,6 @@
use chrono::{DateTime, Utc};
use crate::accounts::account::{Account, AccountStatus};
use crate::accounts::service::AccountService;
/// 候选过滤所需的账号池策略参数(只读视图)。
pub struct CandidateFilter<'a> {
@@ -78,7 +77,7 @@ fn is_base_available(
) -> bool {
account.status == AccountStatus::Active
&& (filter.slot_count)(&account.id) < filter.max_concurrent_per_account
&& AccountService::quota_available_at(account, request.now, filter.skip_quota_limited)
&& quota_available_at(account, request.now, filter.skip_quota_limited)
&& is_model_allowed(
account,
request.model,
@@ -89,7 +88,7 @@ fn is_base_available(
.exclude_account_ids
.iter()
.any(|account_id| account_id == &account.id)
&& AccountService::cloudflare_available_at(account, request.now)
&& cloudflare_available_at(account, request.now)
}
/// 判断模型是否允许该账号的订阅计划使用。
@@ -132,6 +131,25 @@ pub fn is_model_refresh_available(
) -> bool {
account.status == AccountStatus::Active
&& slot_count(&account.id) < max_concurrent_per_account
&& AccountService::quota_available_at(account, now, skip_quota_limited)
&& AccountService::cloudflare_available_at(account, now)
&& quota_available_at(account, now, skip_quota_limited)
&& cloudflare_available_at(account, now)
}
pub(crate) fn quota_available_at(
account: &Account,
now: DateTime<Utc>,
skip_quota_limited: bool,
) -> bool {
if !skip_quota_limited || !account.quota_limit_reached {
return true;
}
account
.quota_cooldown_until
.is_some_and(|cooldown_until| now >= cooldown_until)
}
fn cloudflare_available_at(account: &Account, now: DateTime<Utc>) -> bool {
account
.cloudflare_cooldown_until
.is_none_or(|cooldown_until| now >= cooldown_until)
}
-32
View File
@@ -1,32 +0,0 @@
//! 账号管理用例辅助。
use chrono::{DateTime, Utc};
use crate::accounts::account::Account;
/// 账号服务。
#[derive(Debug, Clone, Default)]
pub struct AccountService;
impl AccountService {
/// 判断账号配额锁定是否已经释放。
pub fn quota_available_at(
account: &Account,
now: DateTime<Utc>,
skip_quota_limited: bool,
) -> bool {
if !skip_quota_limited || !account.quota_limit_reached {
return true;
}
account
.quota_cooldown_until
.is_some_and(|cooldown_until| now >= cooldown_until)
}
/// 判断 Cloudflare 冷却是否已经释放。
pub fn cloudflare_available_at(account: &Account, now: DateTime<Utc>) -> bool {
account
.cloudflare_cooldown_until
.is_none_or(|cooldown_until| now >= cooldown_until)
}
}
+6 -107
View File
@@ -7,8 +7,7 @@ use serde_json::Value;
use sqlx::{PgPool, Postgres, QueryBuilder, Row};
use thiserror::Error;
use crate::accounts::account::{Account, AccountModelUsageDelta, AccountStatus, AccountUsageDelta};
use crate::accounts::window::should_reset_usage_window;
use crate::accounts::account::{Account, AccountStatus};
use crate::infra::{
json::{decode_cursor, page_offset, NumberedPage, Page},
time::parse_rfc3339_utc,
@@ -20,9 +19,8 @@ mod rows;
use queries::*;
use rows::{
count_account_metadata, get_pool_account, list_pool_accounts, map_account_store_error,
metadata_from_row, optional_positive_i64_to_u64, optional_update_value, pg_usage_delta,
push_account_metadata_search, quota_plan_type, quota_snapshot_from_row, status_to_db,
stored_account_from_row, to_page, u64_to_i64_saturating,
metadata_from_row, optional_update_value, push_account_metadata_search, quota_plan_type,
quota_snapshot_from_row, status_to_db, stored_account_from_row, to_page,
};
// ============================================================================
@@ -210,30 +208,6 @@ pub struct AccountQuotaSnapshot {
pub quota_fetched_at: Option<DateTime<Utc>>,
}
/// 用量增量(PostgreSQL 内部表示)。
#[derive(Debug, Clone, Copy, Default)]
pub struct UsageDelta {
request_count: i64,
empty_response_count: i64,
input_tokens: i64,
output_tokens: i64,
cached_tokens: i64,
reasoning_tokens: i64,
total_tokens: i64,
image_input_tokens: i64,
image_output_tokens: i64,
image_request_count: i64,
image_request_failed_count: i64,
window_request_count: i64,
window_input_tokens: i64,
window_output_tokens: i64,
window_cached_tokens: i64,
window_image_input_tokens: i64,
window_image_output_tokens: i64,
window_image_request_count: i64,
window_image_request_failed_count: i64,
}
// ============================================================================
// AccountStore 端口 trait
// ============================================================================
@@ -284,21 +258,6 @@ pub trait AccountStore: Send + Sync + 'static {
async fn set_status(&self, account_id: &str, status: AccountStatus)
-> AccountStoreResult<bool>;
/// 记录账号用量增量。
async fn record_usage_delta(
&self,
account_id: &str,
usage: AccountUsageDelta,
) -> AccountStoreResult<()>;
/// 记录账号模型维度用量增量。
async fn record_model_usage_delta(
&self,
account_id: &str,
model: &str,
usage: AccountModelUsageDelta,
) -> AccountStoreResult<()>;
/// 读取账号当前配额 JSON。
async fn get_quota_json(&self, account_id: &str) -> AccountStoreResult<Option<String>>;
@@ -312,31 +271,7 @@ pub trait AccountStore: Send + Sync + 'static {
) -> AccountStoreResult<bool>;
/// 同步运行时自然刷新出来的账号状态。
async fn sync_runtime_account_state(
&self,
account: &Account,
sync_usage_window: bool,
) -> AccountStoreResult<bool>;
/// 同步账号当前 rate-limit 统计窗口。
async fn sync_rate_limit_window(
&self,
account_id: &str,
reset_at: DateTime<Utc>,
limit_window_seconds: Option<u64>,
) -> AccountStoreResult<()>;
/// 记录账号被用于一次外部请求。
async fn record_request(&self, account_id: &str) -> AccountStoreResult<()> {
self.record_usage_delta(
account_id,
AccountUsageDelta {
requests: 1,
..AccountUsageDelta::default()
},
)
.await
}
async fn sync_runtime_account_state(&self, account: &Account) -> AccountStoreResult<bool>;
}
#[derive(Clone)]
@@ -409,27 +344,6 @@ impl AccountStore for PgAccountStore {
.map_err(|error| map_account_store_error(&error))
}
async fn record_usage_delta(
&self,
account_id: &str,
usage: AccountUsageDelta,
) -> AccountStoreResult<()> {
self.record_usage(account_id, pg_usage_delta(usage))
.await
.map_err(|error| map_account_store_error(&error))
}
async fn record_model_usage_delta(
&self,
account_id: &str,
model: &str,
usage: AccountModelUsageDelta,
) -> AccountStoreResult<()> {
self.record_model_usage(account_id, model, usage)
.await
.map_err(|error| map_account_store_error(&error))
}
async fn get_quota_json(&self, account_id: &str) -> AccountStoreResult<Option<String>> {
PgAccountStore::get_quota_json(self, account_id)
.await
@@ -454,23 +368,8 @@ impl AccountStore for PgAccountStore {
.map_err(|error| map_account_store_error(&error))
}
async fn sync_runtime_account_state(
&self,
account: &Account,
sync_usage_window: bool,
) -> AccountStoreResult<bool> {
PgAccountStore::sync_runtime_account_state(self, account, sync_usage_window)
.await
.map_err(|error| map_account_store_error(&error))
}
async fn sync_rate_limit_window(
&self,
account_id: &str,
reset_at: DateTime<Utc>,
limit_window_seconds: Option<u64>,
) -> AccountStoreResult<()> {
PgAccountStore::sync_rate_limit_window(self, account_id, reset_at, limit_window_seconds)
async fn sync_runtime_account_state(&self, account: &Account) -> AccountStoreResult<bool> {
PgAccountStore::sync_runtime_account_state(self, account)
.await
.map_err(|error| map_account_store_error(&error))
}
+12 -219
View File
@@ -6,85 +6,21 @@
pub(super) const LIST_POOL_ACCOUNTS_SQL: &str = r"
select
a.id,
a.email,
a.chatgpt_account_id as account_id,
a.chatgpt_user_id as user_id,
a.label,
a.plan_type,
a.access_token,
a.refresh_token,
a.access_token_expires_at,
a.next_refresh_at,
a.status,
a.quota_limit_reached,
a.quota_verify_required,
a.quota_cooldown_until,
a.quota_json,
a.cloudflare_cooldown_until,
a.added_at,
coalesce(au.request_count, 0) as usage_request_count,
coalesce(au.empty_response_count, 0) as usage_empty_response_count,
coalesce(au.image_input_tokens, 0) as usage_image_input_tokens,
coalesce(au.image_output_tokens, 0) as usage_image_output_tokens,
coalesce(au.image_request_count, 0) as usage_image_request_count,
coalesce(au.image_request_failed_count, 0) as usage_image_request_failed_count,
coalesce(au.window_request_count, 0) as usage_window_request_count,
coalesce(au.window_input_tokens, 0) as usage_window_input_tokens,
coalesce(au.window_output_tokens, 0) as usage_window_output_tokens,
coalesce(au.window_cached_tokens, 0) as usage_window_cached_tokens,
coalesce(au.window_image_input_tokens, 0) as usage_window_image_input_tokens,
coalesce(au.window_image_output_tokens, 0) as usage_window_image_output_tokens,
coalesce(au.window_image_request_count, 0) as usage_window_image_request_count,
coalesce(au.window_image_request_failed_count, 0) as usage_window_image_request_failed_count,
au.window_started_at as usage_window_started_at,
au.window_reset_at as usage_window_reset_at,
au.limit_window_seconds as usage_limit_window_seconds,
au.last_used_at as usage_last_used_at
from accounts a
left join account_usage au on au.account_id = a.id
order by a.added_at asc, a.id asc";
id, email, chatgpt_account_id as account_id, chatgpt_user_id as user_id,
label, plan_type, access_token, refresh_token, access_token_expires_at,
next_refresh_at, status, quota_limit_reached, quota_verify_required,
quota_cooldown_until, quota_json, cloudflare_cooldown_until, added_at
from accounts
order by added_at asc, id asc";
pub(super) const GET_POOL_ACCOUNT_SQL: &str = r"
select
a.id,
a.email,
a.chatgpt_account_id as account_id,
a.chatgpt_user_id as user_id,
a.label,
a.plan_type,
a.access_token,
a.refresh_token,
a.access_token_expires_at,
a.next_refresh_at,
a.status,
a.quota_limit_reached,
a.quota_verify_required,
a.quota_cooldown_until,
a.quota_json,
a.cloudflare_cooldown_until,
a.added_at,
coalesce(au.request_count, 0) as usage_request_count,
coalesce(au.empty_response_count, 0) as usage_empty_response_count,
coalesce(au.image_input_tokens, 0) as usage_image_input_tokens,
coalesce(au.image_output_tokens, 0) as usage_image_output_tokens,
coalesce(au.image_request_count, 0) as usage_image_request_count,
coalesce(au.image_request_failed_count, 0) as usage_image_request_failed_count,
coalesce(au.window_request_count, 0) as usage_window_request_count,
coalesce(au.window_input_tokens, 0) as usage_window_input_tokens,
coalesce(au.window_output_tokens, 0) as usage_window_output_tokens,
coalesce(au.window_cached_tokens, 0) as usage_window_cached_tokens,
coalesce(au.window_image_input_tokens, 0) as usage_window_image_input_tokens,
coalesce(au.window_image_output_tokens, 0) as usage_window_image_output_tokens,
coalesce(au.window_image_request_count, 0) as usage_window_image_request_count,
coalesce(au.window_image_request_failed_count, 0) as usage_window_image_request_failed_count,
au.window_started_at as usage_window_started_at,
au.window_reset_at as usage_window_reset_at,
au.limit_window_seconds as usage_limit_window_seconds,
au.last_used_at as usage_last_used_at
from accounts a
left join account_usage au on au.account_id = a.id
where a.id = $1";
id, email, chatgpt_account_id as account_id, chatgpt_user_id as user_id,
label, plan_type, access_token, refresh_token, access_token_expires_at,
next_refresh_at, status, quota_limit_reached, quota_verify_required,
quota_cooldown_until, quota_json, cloudflare_cooldown_until, added_at
from accounts
where id = $1";
pub(super) const INSERT_ACCOUNT_SQL: &str = r"
insert into accounts (
@@ -266,73 +202,6 @@ select
updated_at
from accounts";
pub(super) const RECORD_USAGE_SQL: &str = r"
insert into account_usage (
account_id,
request_count,
empty_response_count,
input_tokens,
output_tokens,
cached_tokens,
reasoning_tokens,
total_tokens,
image_input_tokens,
image_output_tokens,
image_request_count,
image_request_failed_count,
window_request_count,
window_input_tokens,
window_output_tokens,
window_cached_tokens,
window_image_input_tokens,
window_image_output_tokens,
window_image_request_count,
window_image_request_failed_count,
window_started_at,
last_used_at
) values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22)
on conflict(account_id) do update set
request_count = account_usage.request_count + excluded.request_count,
empty_response_count = account_usage.empty_response_count + excluded.empty_response_count,
input_tokens = account_usage.input_tokens + excluded.input_tokens,
output_tokens = account_usage.output_tokens + excluded.output_tokens,
cached_tokens = account_usage.cached_tokens + excluded.cached_tokens,
reasoning_tokens = account_usage.reasoning_tokens + excluded.reasoning_tokens,
total_tokens = account_usage.total_tokens + excluded.total_tokens,
image_input_tokens = account_usage.image_input_tokens + excluded.image_input_tokens,
image_output_tokens = account_usage.image_output_tokens + excluded.image_output_tokens,
image_request_count = account_usage.image_request_count + excluded.image_request_count,
image_request_failed_count = account_usage.image_request_failed_count + excluded.image_request_failed_count,
window_request_count = account_usage.window_request_count + excluded.window_request_count,
window_input_tokens = account_usage.window_input_tokens + excluded.window_input_tokens,
window_output_tokens = account_usage.window_output_tokens + excluded.window_output_tokens,
window_cached_tokens = account_usage.window_cached_tokens + excluded.window_cached_tokens,
window_image_input_tokens = account_usage.window_image_input_tokens + excluded.window_image_input_tokens,
window_image_output_tokens = account_usage.window_image_output_tokens + excluded.window_image_output_tokens,
window_image_request_count = account_usage.window_image_request_count + excluded.window_image_request_count,
window_image_request_failed_count = account_usage.window_image_request_failed_count + excluded.window_image_request_failed_count,
window_started_at = coalesce(account_usage.window_started_at, excluded.window_started_at),
last_used_at = excluded.last_used_at";
pub(super) const RECORD_MODEL_USAGE_SQL: &str = r"
insert into account_model_usage (
account_id,
model,
request_count,
error_count,
input_tokens,
output_tokens,
cached_tokens,
last_used_at
) values ($1, $2, $3, $4, $5, $6, $7, $8)
on conflict(account_id, model) do update set
request_count = account_model_usage.request_count + excluded.request_count,
error_count = account_model_usage.error_count + excluded.error_count,
input_tokens = account_model_usage.input_tokens + excluded.input_tokens,
output_tokens = account_model_usage.output_tokens + excluded.output_tokens,
cached_tokens = account_model_usage.cached_tokens + excluded.cached_tokens,
last_used_at = excluded.last_used_at";
pub(super) const LIST_QUOTA_SNAPSHOTS_SQL: &str = r"
select
id,
@@ -370,51 +239,6 @@ set
updated_at = $7
where id = $8";
pub(super) const SELECT_RATE_LIMIT_WINDOW_SQL: &str = r"
select
window_reset_at,
limit_window_seconds
from account_usage
where account_id = $1";
pub(super) const SYNC_RATE_LIMIT_WINDOW_RESET_SQL: &str = r"
insert into account_usage (
account_id,
window_request_count,
window_input_tokens,
window_output_tokens,
window_cached_tokens,
window_image_input_tokens,
window_image_output_tokens,
window_image_request_count,
window_image_request_failed_count,
window_started_at,
window_reset_at,
limit_window_seconds
) values ($1, 0, 0, 0, 0, 0, 0, 0, 0, $2, $3, $4)
on conflict(account_id) do update set
window_request_count = 0,
window_input_tokens = 0,
window_output_tokens = 0,
window_cached_tokens = 0,
window_image_input_tokens = 0,
window_image_output_tokens = 0,
window_image_request_count = 0,
window_image_request_failed_count = 0,
window_started_at = excluded.window_started_at,
window_reset_at = excluded.window_reset_at,
limit_window_seconds = coalesce(excluded.limit_window_seconds, account_usage.limit_window_seconds)";
pub(super) const SYNC_RATE_LIMIT_WINDOW_SQL: &str = r"
insert into account_usage (
account_id,
window_reset_at,
limit_window_seconds
) values ($1, $2, $3)
on conflict(account_id) do update set
window_reset_at = excluded.window_reset_at,
limit_window_seconds = coalesce(excluded.limit_window_seconds, account_usage.limit_window_seconds)";
pub(super) const MARK_QUOTA_LIMITED_UNTIL_SQL: &str = r"
update accounts
set
@@ -460,37 +284,6 @@ set
updated_at = $17
where id = $18";
pub(super) const SYNC_RUNTIME_ACCOUNT_USAGE_WINDOW_SQL: &str = r"
insert into account_usage (
account_id,
window_request_count,
window_input_tokens,
window_output_tokens,
window_cached_tokens,
window_image_input_tokens,
window_image_output_tokens,
window_image_request_count,
window_image_request_failed_count,
window_started_at,
window_reset_at,
limit_window_seconds
) values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
on conflict(account_id) do update set
window_request_count = excluded.window_request_count,
window_input_tokens = excluded.window_input_tokens,
window_output_tokens = excluded.window_output_tokens,
window_cached_tokens = excluded.window_cached_tokens,
window_image_input_tokens = excluded.window_image_input_tokens,
window_image_output_tokens = excluded.window_image_output_tokens,
window_image_request_count = excluded.window_image_request_count,
window_image_request_failed_count = excluded.window_image_request_failed_count,
window_started_at = excluded.window_started_at,
window_reset_at = excluded.window_reset_at,
limit_window_seconds = excluded.limit_window_seconds
where account_usage.window_reset_at is null
or account_usage.window_reset_at <= excluded.window_reset_at
or account_usage.window_reset_at <= $13";
pub(super) const SET_CLOUDFLARE_COOLDOWN_UNTIL_SQL: &str = r"
update accounts
set
+20 -87
View File
@@ -6,7 +6,7 @@ use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
use crate::{
accounts::{
account::{Account, AccountStatus, AccountUsageDelta},
account::{Account, AccountStatus},
quota::{quota_snapshot_limit_window_seconds, quota_snapshot_reset_at},
},
infra::{format::nonnegative_i64_to_u64, json::Page},
@@ -15,7 +15,7 @@ use crate::{
use super::{
queries::{GET_POOL_ACCOUNT_SQL, LIST_POOL_ACCOUNTS_SQL},
AccountQuotaSnapshot, AccountStoreError, PgAccountStore, PgAccountStoreError,
PgAccountStoreResult, StoredAccount, StoredAccountMetadata, UsageDelta,
PgAccountStoreResult, StoredAccount, StoredAccountMetadata,
};
pub(super) async fn list_pool_accounts(
@@ -69,52 +69,27 @@ pub(super) fn pool_account_from_row(row: &PgRow) -> PgAccountStoreResult<Account
quota_verify_required: row.get("quota_verify_required"),
quota_cooldown_until: row.get("quota_cooldown_until"),
cloudflare_cooldown_until: row.get("cloudflare_cooldown_until"),
request_count: nonnegative_i64_to_u64(row.get::<i64, _>("usage_request_count")),
empty_response_count: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_empty_response_count"),
),
image_input_tokens: nonnegative_i64_to_u64(row.get::<i64, _>("usage_image_input_tokens")),
image_output_tokens: nonnegative_i64_to_u64(row.get::<i64, _>("usage_image_output_tokens")),
image_request_count: nonnegative_i64_to_u64(row.get::<i64, _>("usage_image_request_count")),
image_request_failed_count: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_image_request_failed_count"),
),
window_request_count: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_window_request_count"),
),
window_input_tokens: nonnegative_i64_to_u64(row.get::<i64, _>("usage_window_input_tokens")),
window_output_tokens: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_window_output_tokens"),
),
window_cached_tokens: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_window_cached_tokens"),
),
window_image_input_tokens: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_window_image_input_tokens"),
),
window_image_output_tokens: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_window_image_output_tokens"),
),
window_image_request_count: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_window_image_request_count"),
),
window_image_request_failed_count: nonnegative_i64_to_u64(
row.get::<i64, _>("usage_window_image_request_failed_count"),
),
window_started_at: row.get("usage_window_started_at"),
window_reset_at: row
.get::<Option<chrono::DateTime<chrono::Utc>>, _>("usage_window_reset_at")
.or(quota_window_reset_at),
limit_window_seconds: optional_positive_i64_to_u64(
row.get::<Option<i64>, _>("usage_limit_window_seconds"),
)
.or(quota_limit_window_seconds),
request_count: 0,
empty_response_count: 0,
image_input_tokens: 0,
image_output_tokens: 0,
image_request_count: 0,
image_request_failed_count: 0,
window_request_count: 0,
window_input_tokens: 0,
window_output_tokens: 0,
window_cached_tokens: 0,
window_image_input_tokens: 0,
window_image_output_tokens: 0,
window_image_request_count: 0,
window_image_request_failed_count: 0,
window_started_at: None,
window_reset_at: quota_window_reset_at,
limit_window_seconds: quota_limit_window_seconds,
added_at: row
.get::<chrono::DateTime<chrono::Utc>, _>("added_at")
.to_rfc3339(),
last_used_at: row
.get::<Option<chrono::DateTime<chrono::Utc>>, _>("usage_last_used_at")
.map(|value| value.to_rfc3339()),
last_used_at: None,
})
}
@@ -168,38 +143,6 @@ pub(super) fn map_account_store_error(error: &impl ToString) -> AccountStoreErro
}
}
pub(super) fn pg_usage_delta(usage: AccountUsageDelta) -> UsageDelta {
let request_count = u64_to_i64_saturating(usage.requests);
let input_tokens = u64_to_i64_saturating(usage.input_tokens);
let output_tokens = u64_to_i64_saturating(usage.output_tokens);
let cached_tokens = u64_to_i64_saturating(usage.cached_tokens);
let image_input_tokens = u64_to_i64_saturating(usage.image_input_tokens);
let image_output_tokens = u64_to_i64_saturating(usage.image_output_tokens);
let image_request_count = u64_to_i64_saturating(usage.image_requests);
let image_request_failed_count = u64_to_i64_saturating(usage.image_request_failures);
UsageDelta {
request_count,
empty_response_count: u64_to_i64_saturating(usage.empty_responses),
input_tokens,
output_tokens,
cached_tokens,
reasoning_tokens: u64_to_i64_saturating(usage.reasoning_tokens),
total_tokens: u64_to_i64_saturating(usage.total_tokens),
image_input_tokens,
image_output_tokens,
image_request_count,
image_request_failed_count,
window_request_count: request_count,
window_input_tokens: input_tokens,
window_output_tokens: output_tokens,
window_cached_tokens: cached_tokens,
window_image_input_tokens: image_input_tokens,
window_image_output_tokens: image_output_tokens,
window_image_request_count: image_request_count,
window_image_request_failed_count: image_request_failed_count,
}
}
pub(super) fn quota_plan_type(quota_json: &str) -> Option<String> {
serde_json::from_str::<Value>(quota_json)
.ok()?
@@ -245,16 +188,6 @@ pub(super) fn status_from_db(value: &str) -> PgAccountStoreResult<AccountStatus>
}
}
pub(super) fn u64_to_i64_saturating(value: u64) -> i64 {
value.min(i64::MAX as u64) as i64
}
pub(super) fn optional_positive_i64_to_u64(value: Option<i64>) -> Option<u64> {
value
.and_then(|value| u64::try_from(value).ok())
.filter(|value| *value > 0)
}
pub(super) async fn count_account_metadata(
pool: &PgPool,
search: Option<&str>,
+2 -139
View File
@@ -345,14 +345,12 @@ impl PgAccountStore {
pub async fn sync_runtime_account_state(
&self,
account: &Account,
sync_usage_window: bool,
) -> PgAccountStoreResult<bool> {
let now = Utc::now();
let quota_limit_reached = account.quota_limit_reached;
let quota_verify_required = account.quota_verify_required;
let quota_cooldown_until = account.quota_cooldown_until;
let cloudflare_cooldown_until = account.cloudflare_cooldown_until;
let mut tx = self.pool.begin().await?;
let result = sqlx::query(SYNC_RUNTIME_ACCOUNT_STATE_SQL)
.bind(quota_limit_reached)
.bind(now)
@@ -372,36 +370,9 @@ impl PgAccountStore {
.bind(cloudflare_cooldown_until)
.bind(now)
.bind(&account.id)
.execute(&mut *tx)
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
tx.commit().await?;
return Ok(false);
}
if sync_usage_window {
sqlx::query(SYNC_RUNTIME_ACCOUNT_USAGE_WINDOW_SQL)
.bind(&account.id)
.bind(u64_to_i64_saturating(account.window_request_count))
.bind(u64_to_i64_saturating(account.window_input_tokens))
.bind(u64_to_i64_saturating(account.window_output_tokens))
.bind(u64_to_i64_saturating(account.window_cached_tokens))
.bind(u64_to_i64_saturating(account.window_image_input_tokens))
.bind(u64_to_i64_saturating(account.window_image_output_tokens))
.bind(u64_to_i64_saturating(account.window_image_request_count))
.bind(u64_to_i64_saturating(
account.window_image_request_failed_count,
))
.bind(account.window_started_at)
.bind(account.window_reset_at)
.bind(account.limit_window_seconds.map(u64_to_i64_saturating))
.bind(now)
.execute(&mut *tx)
.await?;
}
tx.commit().await?;
Ok(true)
Ok(result.rows_affected() > 0)
}
/// 标记账号 Cloudflare 冷却期。
@@ -448,67 +419,6 @@ where id = $4",
Ok(result.rows_affected() > 0)
}
/// 记录用量。
pub async fn record_usage(
&self,
account_id: &str,
delta: UsageDelta,
) -> PgAccountStoreResult<()> {
let last_used_at = Utc::now();
sqlx::query(RECORD_USAGE_SQL)
.bind(account_id)
.bind(delta.request_count)
.bind(delta.empty_response_count)
.bind(delta.input_tokens)
.bind(delta.output_tokens)
.bind(delta.cached_tokens)
.bind(delta.reasoning_tokens)
.bind(delta.total_tokens)
.bind(delta.image_input_tokens)
.bind(delta.image_output_tokens)
.bind(delta.image_request_count)
.bind(delta.image_request_failed_count)
.bind(delta.window_request_count)
.bind(delta.window_input_tokens)
.bind(delta.window_output_tokens)
.bind(delta.window_cached_tokens)
.bind(delta.window_image_input_tokens)
.bind(delta.window_image_output_tokens)
.bind(delta.window_image_request_count)
.bind(delta.window_image_request_failed_count)
.bind(last_used_at)
.bind(last_used_at)
.execute(&self.pool)
.await?;
Ok(())
}
/// 记录模型维度用量。
pub async fn record_model_usage(
&self,
account_id: &str,
model: &str,
delta: AccountModelUsageDelta,
) -> PgAccountStoreResult<()> {
let model = model.trim();
if model.is_empty() {
return Ok(());
}
let last_used_at = Utc::now();
sqlx::query(RECORD_MODEL_USAGE_SQL)
.bind(account_id)
.bind(model)
.bind(u64_to_i64_saturating(delta.requests))
.bind(u64_to_i64_saturating(delta.errors))
.bind(u64_to_i64_saturating(delta.input_tokens))
.bind(u64_to_i64_saturating(delta.output_tokens))
.bind(u64_to_i64_saturating(delta.cached_tokens))
.bind(last_used_at)
.execute(&self.pool)
.await?;
Ok(())
}
/// 读取配额快照列表。
pub async fn list_quota_snapshots(&self) -> PgAccountStoreResult<Vec<AccountQuotaSnapshot>> {
let rows = sqlx::query(LIST_QUOTA_SNAPSHOTS_SQL)
@@ -602,53 +512,6 @@ where id = $4",
Ok(result.rows_affected() > 0)
}
/// 同步 rate-limit 窗口(含重置)。
pub async fn sync_rate_limit_window(
&self,
account_id: &str,
reset_at: DateTime<Utc>,
limit_window_seconds: Option<u64>,
) -> PgAccountStoreResult<()> {
let existing = sqlx::query(SELECT_RATE_LIMIT_WINDOW_SQL)
.bind(account_id)
.fetch_optional(&self.pool)
.await?;
let should_reset = existing
.as_ref()
.map(|row| {
let existing_reset_at = row.get::<Option<DateTime<Utc>>, _>("window_reset_at");
let existing_limit_window_seconds =
optional_positive_i64_to_u64(row.get::<Option<i64>, _>("limit_window_seconds"));
should_reset_usage_window(
existing_reset_at,
existing_limit_window_seconds,
reset_at,
limit_window_seconds,
)
})
.unwrap_or_default();
if should_reset {
let window_started_at = Utc::now();
sqlx::query(SYNC_RATE_LIMIT_WINDOW_RESET_SQL)
.bind(account_id)
.bind(window_started_at)
.bind(reset_at)
.bind(limit_window_seconds.map(u64_to_i64_saturating))
.execute(&self.pool)
.await?;
} else {
sqlx::query(SYNC_RATE_LIMIT_WINDOW_SQL)
.bind(account_id)
.bind(reset_at)
.bind(limit_window_seconds.map(u64_to_i64_saturating))
.execute(&self.pool)
.await?;
}
Ok(())
}
/// 删除单个账号。
pub async fn delete(&self, account_id: &str) -> PgAccountStoreResult<bool> {
let result = sqlx::query(DELETE_ACCOUNT_SQL)
+2 -30
View File
@@ -1,31 +1,3 @@
//! 账号用量窗口规则。
//! 账号配额窗口规则。
use chrono::{DateTime, Utc};
pub(crate) fn should_reset_usage_window(
existing_reset_at: Option<DateTime<Utc>>,
existing_limit_window_seconds: Option<u64>,
new_reset_at: DateTime<Utc>,
new_limit_window_seconds: Option<u64>,
) -> bool {
let Some(existing_reset_at) = existing_reset_at else {
return false;
};
if existing_reset_at == new_reset_at {
return false;
}
let drift = existing_reset_at
.signed_duration_since(new_reset_at)
.num_seconds()
.unsigned_abs();
let window_seconds = new_limit_window_seconds
.or(existing_limit_window_seconds)
.unwrap_or(0);
let threshold = if window_seconds > 0 {
window_seconds / 2
} else {
3_600
};
drift >= threshold
}
pub(crate) use crate::telemetry::account_usage::store::should_reset_usage_window;
+4 -4
View File
@@ -30,13 +30,13 @@ use crate::{
},
api::admin::{
response::{
AdminEnvelope, AdminError, AdminResponse, BatchDeleteData, CursorPageMeta,
NumberedPageMeta, PageMeta, ADMIN_OK_CODE, ADMIN_OK_MESSAGE,
parse_editable_update, AdminEnvelope, AdminError, AdminResponse, BatchDeleteData,
CursorPageMeta, EditableUpdateMessages, NumberedPageMeta, PageMeta, ADMIN_OK_CODE,
ADMIN_OK_MESSAGE,
},
session::AdminAuth,
update_payload::{parse_editable_update, EditableUpdateMessages},
},
bootstrap::state::AppState,
api::AppState,
infra::{
format::{
format_cost, format_percent, format_plain_number, format_tokens, nonnegative_i64_to_u64,
@@ -277,13 +277,22 @@ pub(crate) async fn batch_delete_accounts(
.batch_delete(payload.ids)
.await
{
Ok(result) => Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(BatchDeleteData {
deleted: result.deleted,
not_found: result.not_found,
}),
)),
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)),
}
}
@@ -238,9 +238,8 @@ pub(super) async fn quota_snapshots_by_account(
) -> HashMap<String, AccountQuotaData> {
state
.services
.background_tasks
.accounts
.list_quota_snapshots()
.admin_accounts
.quota_snapshots()
.await
.unwrap_or_default()
.into_iter()
+77 -163
View File
@@ -14,6 +14,7 @@ use serde_json::Value;
use crate::{
accounts::{
account::{Account, AccountStatus},
pool::AccountCapacitySummary,
refresh::token_refresh_status_eligible,
},
api::admin::{
@@ -24,21 +25,20 @@ use crate::{
UsageRecordCostDetailsData, UsageRecordTokenDetailsData,
},
},
bootstrap::state::AppState,
api::AppState,
infra::{
format::{
format_compact_number, format_cost, format_duration_ms, format_percent, format_rate,
format_tokens, nonnegative_i64_to_u64,
},
format::{format_duration_ms, format_tokens},
time::{
china_datetime, china_datetime_rfc3339_str, china_day_start, china_hour,
china_hour_start, china_quarter_hour_start, china_relative_time,
china_datetime, china_datetime_rfc3339_str, china_day_start, china_quarter_hour_start,
china_relative_time,
},
},
telemetry::{
account_usage::query::{AccountUsageSummary, AccountUsageTimeBucket},
diagnostics::{
fingerprint_diagnostics, AccountCapacityDiagnostics, AccountPoolDiagnostics,
account_usage::query::AccountUsageTimeBucket,
dashboard::{
dashboard_cards, dashboard_health_timeline_data, dashboard_trend_data,
DashboardAccountCounts, DashboardCardsData, DashboardHealthTimelineData,
DashboardTrendData, DashboardTrendKind,
},
usage::{
query::UsageQueryFilter,
@@ -46,14 +46,13 @@ use crate::{
types::{metadata_string, UsageRecord},
},
},
upstream::openai::fingerprint::Fingerprint,
};
const DASHBOARD_ACCOUNT_USAGE_LIMIT: u32 = 4;
const DASHBOARD_USAGE_RECORD_LIMIT: u32 = 10;
const HEALTH_TIMELINE_SLOT_MINUTES: i64 = 15;
const DASHBOARD_TIME_BUCKET_MINUTES: i64 = 15;
const DASHBOARD_TIME_BUCKET_SLOTS: i64 = 7 * 24 * 4;
const HEALTH_TIMELINE_SLOTS: i64 = 24 * 4;
const HEALTH_TIMELINE_STABLE_SUCCESS_THRESHOLD: u64 = 3;
const FIVE_HOUR_WINDOW_SECONDS: u64 = 18_000;
const WEEK_WINDOW_SECONDS: u64 = 604_800;
const MONTH_WINDOW_SECONDS: u64 = 2_592_000;
@@ -64,15 +63,6 @@ pub(crate) struct DashboardTrendQuery {
kind: Option<DashboardTrendKind>,
}
#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
enum DashboardTrendKind {
#[default]
Usage,
Latency,
Errors,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardSummaryData {
@@ -87,107 +77,6 @@ struct DashboardSummaryData {
rotation_strategy: Option<String>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardHealthTimelineData {
title: String,
description: String,
reliability_display: String,
points: String,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardCardsData {
accounts: DashboardAccountsCardData,
traffic: DashboardTrafficCardData,
tokens: DashboardTokenCardData,
cache: DashboardCacheCardData,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardAccountsCardData {
total: String,
total_value: u64,
enabled: String,
enabled_value: u64,
abnormal: String,
abnormal_value: u64,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardTrafficCardData {
today_requests: String,
today_requests_value: u64,
yesterday_requests_value: u64,
total_requests: String,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardTokenCardData {
today_tokens: String,
today_tokens_value: u64,
yesterday_tokens_value: u64,
total_tokens: String,
today_cost_usd: String,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardCacheCardData {
today_hit_rate: String,
today_hit_rate_value: Option<f64>,
yesterday_hit_rate_value: Option<f64>,
total_hit_rate: String,
total_cached_tokens: String,
average_first_token_latency_ms: String,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardTrendData {
kind: DashboardTrendKind,
points: Vec<DashboardTrendPointData>,
summary: Vec<DashboardTrendSummaryData>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardTrendPointData {
time: String,
requests: String,
requests_value: u64,
input_tokens: String,
input_tokens_value: u64,
output_tokens: String,
output_tokens_value: u64,
cached_tokens: String,
cached_tokens_value: u64,
cache_hit_rate_value: f64,
tokens_value: u64,
errors: String,
errors_value: u64,
latency: String,
latency_value: u64,
max_latency: String,
max_latency_value: u64,
min_latency: String,
min_latency_value: u64,
success_rate: String,
success_rate_value: f64,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardTrendSummaryData {
label: String,
value: String,
ratio: Option<String>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardAccountUsageData {
@@ -229,17 +118,65 @@ struct DashboardServiceStatusData {
tone: String,
}
#[derive(Debug, Clone, Copy, Default)]
struct UsageWindow {
requests: u64,
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
errors: u64,
first_token_latency_sum: u64,
first_token_latency_count: u64,
max_first_token_bucket_latency: u64,
min_first_token_bucket_latency: u64,
#[derive(Debug, Default, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountPoolDiagnostics {
total: usize,
active: usize,
expired: usize,
quota_exhausted: usize,
refreshing: usize,
disabled: usize,
banned: usize,
}
#[derive(Debug, Default, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountCapacityDiagnostics {
max_concurrent_per_account: usize,
total_slots: usize,
used_slots: usize,
available_slots: usize,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct FingerprintDiagnostics {
source: &'static str,
originator: String,
app_version: String,
build_number: String,
platform: String,
arch: String,
chromium_version: String,
user_agent: String,
#[serde(skip_serializing_if = "Option::is_none")]
updated_at: Option<String>,
}
fn fingerprint_diagnostics(fingerprint: &Fingerprint) -> FingerprintDiagnostics {
FingerprintDiagnostics {
source: "runtime",
originator: fingerprint.originator.clone(),
app_version: fingerprint.app_version.clone(),
build_number: fingerprint.build_number.clone(),
platform: fingerprint.platform.clone(),
arch: fingerprint.arch.clone(),
chromium_version: fingerprint.chromium_version.clone(),
user_agent: fingerprint.user_agent(),
updated_at: fingerprint.updated_at.clone(),
}
}
impl From<AccountCapacitySummary> for AccountCapacityDiagnostics {
fn from(summary: AccountCapacitySummary) -> Self {
Self {
max_concurrent_per_account: summary.max_concurrent_per_account,
total_slots: summary.total_slots,
used_slots: summary.used_slots,
available_slots: summary.available_slots,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
@@ -250,33 +187,6 @@ enum QuotaWindowPriority {
Other,
}
impl UsageWindow {
fn tokens(self) -> u64 {
self.input_tokens + self.output_tokens
}
fn cache_hit_rate(self) -> Option<f64> {
if self.input_tokens > 0 {
Some(self.cached_tokens as f64 / self.input_tokens as f64)
} else {
None
}
}
fn avg_first_token_latency(self) -> Option<u64> {
self.first_token_latency_sum
.checked_div(self.first_token_latency_count)
}
fn max_first_token_bucket_latency(self) -> Option<u64> {
(self.max_first_token_bucket_latency > 0).then_some(self.max_first_token_bucket_latency)
}
fn min_first_token_bucket_latency(self) -> Option<u64> {
(self.min_first_token_bucket_latency > 0).then_some(self.min_first_token_bucket_latency)
}
}
/// `GET /api/admin/dashboard/summary`
pub(crate) async fn dashboard_summary(
State(state): State<AppState>,
@@ -322,7 +232,7 @@ pub(crate) async fn dashboard_summary(
let trend = dashboard_trend_data(&time_buckets, query.kind.unwrap_or_default());
let quota_used_by_account =
account_quota_used_percent_by_id(&state, &account_usage_records).await;
let settings = state.services.settings.current().await;
let settings = state.services.settings.current();
let recent_usage_records = recent_events;
let account_emails = state
.services
@@ -336,7 +246,11 @@ pub(crate) async fn dashboard_summary(
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(DashboardSummaryData {
cards: dashboard_cards(&accounts, &time_buckets, &lifetime_usage),
cards: dashboard_cards(
dashboard_account_counts(&accounts),
&time_buckets,
&lifetime_usage,
),
trend,
health_timeline: dashboard_health_timeline_data(&time_buckets),
account_usage: account_usage_data(
+12 -416
View File
@@ -1,102 +1,12 @@
use super::*;
pub(super) fn dashboard_cards(
accounts: &[Account],
buckets: &[AccountUsageTimeBucket],
lifetime_usage: &AccountUsageSummary,
) -> DashboardCardsData {
dashboard_cards_at(accounts, buckets, lifetime_usage, Utc::now())
}
pub(super) fn dashboard_cards_at(
accounts: &[Account],
buckets: &[AccountUsageTimeBucket],
lifetime_usage: &AccountUsageSummary,
now: DateTime<Utc>,
) -> DashboardCardsData {
let today_start = china_day_start(now);
let yesterday_start = today_start - Duration::days(1);
let today = usage_window(buckets, today_start, now);
let yesterday = usage_window(buckets, yesterday_start, today_start);
let today_cost = cost_window(buckets, today_start, now).unwrap_or(0.0);
let total_requests = nonnegative_i64_to_u64(lifetime_usage.request_count);
let total_tokens = nonnegative_i64_to_u64(lifetime_usage.total_tokens);
let total_cached_tokens = nonnegative_i64_to_u64(lifetime_usage.cached_tokens);
let total_hit_rate = summary_cache_hit_rate(lifetime_usage);
let total_accounts = accounts.len() as u64;
let enabled_accounts = accounts
.iter()
.filter(|account| account.status == AccountStatus::Active)
.count() as u64;
let abnormal_accounts = accounts
.iter()
.filter(|account| {
matches!(
account.status,
AccountStatus::Expired
| AccountStatus::QuotaExhausted
| AccountStatus::Disabled
| AccountStatus::Banned
)
})
.count() as u64;
DashboardCardsData {
accounts: DashboardAccountsCardData {
total: format_compact_number(total_accounts),
total_value: total_accounts,
enabled: format_compact_number(enabled_accounts),
enabled_value: enabled_accounts,
abnormal: format_compact_number(abnormal_accounts),
abnormal_value: abnormal_accounts,
},
traffic: DashboardTrafficCardData {
today_requests: format_compact_number(today.requests),
today_requests_value: today.requests,
yesterday_requests_value: yesterday.requests,
total_requests: format_compact_number(total_requests),
},
tokens: DashboardTokenCardData {
today_tokens: format_tokens(today.tokens()),
today_tokens_value: today.tokens(),
yesterday_tokens_value: yesterday.tokens(),
total_tokens: format_tokens(total_tokens),
today_cost_usd: format_cost(today_cost),
},
cache: DashboardCacheCardData {
today_hit_rate: format_rate(today.cache_hit_rate()),
today_hit_rate_value: today.cache_hit_rate(),
yesterday_hit_rate_value: yesterday.cache_hit_rate(),
total_hit_rate: format_rate(total_hit_rate),
total_cached_tokens: format_tokens(total_cached_tokens),
average_first_token_latency_ms: format_optional_duration_ms(
today.avg_first_token_latency(),
),
},
}
}
pub(super) fn summary_cache_hit_rate(summary: &AccountUsageSummary) -> Option<f64> {
let input_tokens = nonnegative_i64_to_u64(summary.input_tokens);
if input_tokens > 0 {
Some(nonnegative_i64_to_u64(summary.cached_tokens) as f64 / input_tokens as f64)
} else {
None
}
}
pub(super) fn format_optional_duration_ms(value: Option<u64>) -> String {
format_duration_ms(value.and_then(|value| i64::try_from(value).ok()))
}
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(HEALTH_TIMELINE_SLOT_MINUTES * (DASHBOARD_TIME_BUCKET_SLOTS - 1));
- Duration::minutes(DASHBOARD_TIME_BUCKET_MINUTES * (DASHBOARD_TIME_BUCKET_SLOTS - 1));
state
.services
.usage
@@ -105,332 +15,18 @@ pub(super) async fn dashboard_time_buckets(
.unwrap_or_default()
}
pub(super) fn dashboard_trend_data(
buckets: &[AccountUsageTimeBucket],
kind: DashboardTrendKind,
) -> DashboardTrendData {
dashboard_trend_data_at(buckets, kind, Utc::now())
}
pub(super) fn dashboard_trend_data_at(
records: &[AccountUsageTimeBucket],
kind: DashboardTrendKind,
now: DateTime<Utc>,
) -> DashboardTrendData {
let current_hour = china_hour_start(now);
let start = china_day_start(now);
let elapsed_hours = current_hour.signed_duration_since(start).num_hours();
let mut buckets = (0..=elapsed_hours)
.map(|index| {
let bucket_start = start + Duration::hours(index);
(bucket_start, UsageWindow::default())
})
.collect::<Vec<_>>();
for record in records {
if record.bucket_start < start || record.bucket_start > now {
continue;
}
let record_hour = china_hour_start(record.bucket_start);
if let Some((_, bucket)) = buckets
.iter_mut()
.find(|(bucket_start, _)| *bucket_start == record_hour)
{
apply_bucket(bucket, record);
}
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,
}
let points = buckets
.iter()
.map(|(bucket_start, bucket)| {
let latency = bucket.avg_first_token_latency().unwrap_or(0);
let max_latency = bucket.max_first_token_bucket_latency().unwrap_or(0);
let min_latency = bucket.min_first_token_bucket_latency().unwrap_or(0);
let tokens = bucket.tokens();
let cache_hit_rate = bucket.cache_hit_rate().unwrap_or(0.0);
let success_rate = if bucket.requests > 0 {
((bucket.requests - bucket.errors) as f64 / bucket.requests as f64 * 1000.0).round()
/ 10.0
} else {
0.0
};
DashboardTrendPointData {
time: format!("{:02}", china_hour(bucket_start)),
requests: format_compact_number(bucket.requests),
requests_value: bucket.requests,
input_tokens: format_tokens(bucket.input_tokens),
input_tokens_value: bucket.input_tokens,
output_tokens: format_tokens(bucket.output_tokens),
output_tokens_value: bucket.output_tokens,
cached_tokens: format_tokens(bucket.cached_tokens),
cached_tokens_value: bucket.cached_tokens,
cache_hit_rate_value: cache_hit_rate,
tokens_value: tokens,
errors: format_compact_number(bucket.errors),
errors_value: bucket.errors,
latency: format_optional_duration_ms(Some(latency)),
latency_value: latency,
max_latency: format_optional_duration_ms(Some(max_latency)),
max_latency_value: max_latency,
min_latency: format_optional_duration_ms(Some(min_latency)),
min_latency_value: min_latency,
success_rate: format_percent(success_rate),
success_rate_value: success_rate,
}
})
.collect::<Vec<_>>();
DashboardTrendData {
kind,
summary: trend_summary(kind, &points),
points,
}
}
pub(super) fn dashboard_health_timeline_data(
buckets: &[AccountUsageTimeBucket],
) -> DashboardHealthTimelineData {
dashboard_health_timeline_data_at(buckets, Utc::now())
}
pub(super) fn dashboard_health_timeline_data_at(
records: &[AccountUsageTimeBucket],
now: DateTime<Utc>,
) -> DashboardHealthTimelineData {
let current_slot = china_quarter_hour_start(now);
let start = china_day_start(now);
let mut buckets = (0..HEALTH_TIMELINE_SLOTS)
.map(|index| {
let bucket_start = start + Duration::minutes(HEALTH_TIMELINE_SLOT_MINUTES * index);
(bucket_start, UsageWindow::default())
})
.collect::<Vec<_>>();
for record in records {
if record.bucket_start < start || record.bucket_start > now {
continue;
}
let record_slot = china_quarter_hour_start(record.bucket_start);
if let Some((_, bucket)) = buckets
.iter_mut()
.find(|(bucket_start, _)| *bucket_start == record_slot)
{
apply_bucket(bucket, record);
}
}
let requests = buckets
.iter()
.filter(|(bucket_start, _)| *bucket_start <= current_slot)
.map(|(_, bucket)| bucket.requests)
.sum::<u64>();
let errors = buckets
.iter()
.filter(|(bucket_start, _)| *bucket_start <= current_slot)
.map(|(_, bucket)| bucket.errors)
.sum::<u64>();
let successes = requests.saturating_sub(errors);
let reliability = if requests > 0 {
(successes as f64 / requests as f64 * 1000.0).round() / 10.0
} else {
0.0
};
DashboardHealthTimelineData {
title: "请求健康时间线".to_string(),
description: "请求可靠性".to_string(),
reliability_display: if requests > 0 {
format!("{reliability:.1}%")
} else {
"-".to_string()
},
points: buckets
.into_iter()
.map(|(bucket_start, bucket)| {
let success_rate = if bucket.requests > 0 {
(bucket.requests.saturating_sub(bucket.errors) as f64 / bucket.requests as f64
* 1000.0)
.round()
/ 10.0
} else {
0.0
};
health_tone(bucket, success_rate, bucket_start > current_slot)
})
.collect(),
}
}
pub(super) fn health_tone(bucket: UsageWindow, success_rate: f64, is_future: bool) -> char {
let successes = bucket.requests.saturating_sub(bucket.errors);
if is_future {
'0'
} else if bucket.requests == 0 {
'1'
} else if successes == 0 {
'2'
} else if success_rate < 90.0 {
'3'
} else if successes < HEALTH_TIMELINE_STABLE_SUCCESS_THRESHOLD {
'4'
} else {
'5'
}
}
pub(super) fn trend_summary(
kind: DashboardTrendKind,
points: &[DashboardTrendPointData],
) -> Vec<DashboardTrendSummaryData> {
match kind {
DashboardTrendKind::Usage => vec![
trend_summary_item(
kind,
"输入",
points.iter().map(|point| point.input_tokens_value).sum(),
None,
),
trend_summary_item(
kind,
"输出",
points.iter().map(|point| point.output_tokens_value).sum(),
None,
),
trend_summary_item(
kind,
"缓存",
points.iter().map(|point| point.cached_tokens_value).sum(),
None,
),
],
DashboardTrendKind::Latency => {
let samples = points
.iter()
.filter(|point| point.latency_value > 0)
.collect::<Vec<_>>();
let avg = if samples.is_empty() {
0
} else {
samples.iter().map(|point| point.latency_value).sum::<u64>() / samples.len() as u64
};
vec![
trend_summary_item(kind, "平均", avg, None),
trend_summary_item(
kind,
"最高",
samples
.iter()
.map(|point| point.max_latency_value)
.max()
.unwrap_or(0),
None,
),
trend_summary_item(
kind,
"最低",
samples
.iter()
.filter_map(|point| {
(point.min_latency_value > 0).then_some(point.min_latency_value)
})
.min()
.unwrap_or(0),
None,
),
]
}
DashboardTrendKind::Errors => {
let errors = points.iter().map(|point| point.errors_value).sum::<u64>();
let requests = points.iter().map(|point| point.requests_value).sum::<u64>();
let success_rate = if requests > 0 {
Some(((requests - errors) as f64 / requests as f64 * 1000.0).round() / 10.0)
} else {
None
};
vec![
trend_summary_item(kind, "错误数", errors, None),
trend_summary_item(kind, "成功率", 0, success_rate),
trend_summary_item(kind, "总请求", requests, None),
]
}
}
}
pub(super) fn trend_summary_item(
kind: DashboardTrendKind,
label: &str,
value: u64,
ratio: Option<f64>,
) -> DashboardTrendSummaryData {
DashboardTrendSummaryData {
label: label.to_string(),
value: trend_summary_value_display(kind, value),
ratio: ratio.map(format_percent),
}
}
pub(super) fn trend_summary_value_display(kind: DashboardTrendKind, value: u64) -> String {
match kind {
DashboardTrendKind::Usage => format_tokens(value),
DashboardTrendKind::Latency => format_optional_duration_ms(Some(value)),
DashboardTrendKind::Errors => format_compact_number(value),
}
}
pub(super) fn usage_window(
records: &[AccountUsageTimeBucket],
start: DateTime<Utc>,
end: DateTime<Utc>,
) -> UsageWindow {
let mut window = UsageWindow::default();
for record in records {
if record.bucket_start >= start && record.bucket_start < end {
apply_bucket(&mut window, record);
}
}
window
}
pub(super) fn cost_window(
records: &[AccountUsageTimeBucket],
start: DateTime<Utc>,
end: DateTime<Utc>,
) -> Option<f64> {
let mut total = 0.0;
let mut has_usage = false;
for record in records {
if record.bucket_start >= start && record.bucket_start < end {
if let Some(cost) = record.cost_usd {
total += cost;
has_usage = true;
}
}
}
has_usage.then_some(total)
}
pub(super) fn apply_bucket(window: &mut UsageWindow, record: &AccountUsageTimeBucket) {
window.requests += nonnegative_i64_to_u64(record.request_count);
window.input_tokens += nonnegative_i64_to_u64(record.input_tokens);
window.output_tokens += nonnegative_i64_to_u64(record.output_tokens);
window.cached_tokens += nonnegative_i64_to_u64(record.cached_tokens);
window.errors += nonnegative_i64_to_u64(record.error_count);
window.first_token_latency_sum += nonnegative_i64_to_u64(record.first_token_latency_sum);
window.first_token_latency_count += nonnegative_i64_to_u64(record.first_token_latency_count);
let bucket_first_token_latency = first_token_bucket_latency(record);
if let Some(latency) = bucket_first_token_latency {
window.max_first_token_bucket_latency = window.max_first_token_bucket_latency.max(latency);
window.min_first_token_bucket_latency = if window.min_first_token_bucket_latency == 0 {
latency
} else {
window.min_first_token_bucket_latency.min(latency)
};
}
}
pub(super) fn first_token_bucket_latency(record: &AccountUsageTimeBucket) -> Option<u64> {
let sum = nonnegative_i64_to_u64(record.first_token_latency_sum);
let count = nonnegative_i64_to_u64(record.first_token_latency_count);
sum.checked_div(count).filter(|latency| *latency > 0)
}
pub(super) fn account_pool_summary(
+3 -3
View File
@@ -11,11 +11,11 @@ use serde_json::Value;
use crate::{
api::admin::response::{
AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse, BatchDeleteData,
parse_editable_update, AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse,
BatchDeleteData, EditableUpdateMessages,
},
api::admin::session::AdminAuth,
api::admin::update_payload::{parse_editable_update, EditableUpdateMessages},
bootstrap::state::AppState,
api::AppState,
infra::{
json::{clamp_limit, Page},
time::{china_relative_time_str, china_rfc3339_str},
-1
View File
@@ -9,5 +9,4 @@ pub mod router;
pub mod session;
pub mod settings_routes;
pub mod system_routes;
pub mod update_payload;
pub mod usage_routes;
+1 -1
View File
@@ -12,7 +12,7 @@ use crate::{
response::{AdminError, AdminPageEnvelope, AdminResponse},
session::AdminAuth,
},
bootstrap::state::AppState,
api::AppState,
infra::{
json::{clamp_limit, clamp_page, NumberedPage, Page},
time::china_datetime,
+58 -1
View File
@@ -6,7 +6,8 @@ use axum::{
response::{IntoResponse, Response},
Json,
};
use serde::Serialize;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::fmt;
pub(crate) const ADMIN_OK_CODE: u32 = 200;
@@ -291,3 +292,59 @@ impl<T: Serialize> IntoResponse for AdminResponse<T> {
(self.status, Json(self.body)).into_response()
}
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct EditableUpdateRequest {
id: String,
label: Option<Option<String>>,
status: Option<String>,
}
pub(crate) struct ParsedEditableUpdate {
pub id: String,
pub label: Option<Option<String>>,
pub status: Option<String>,
}
pub(crate) struct EditableUpdateMessages<'a> {
pub object_required: &'a str,
pub invalid: &'a str,
pub empty_update: &'a str,
pub unknown_field_editable: bool,
}
pub(crate) fn parse_editable_update(
payload: &Value,
messages: EditableUpdateMessages<'_>,
) -> Result<ParsedEditableUpdate, AdminError> {
let object = payload
.as_object()
.ok_or_else(|| AdminError::bad_request(messages.object_required))?;
if messages.unknown_field_editable {
for field in object.keys() {
if !matches!(field.as_str(), "id" | "label" | "status") {
return Err(AdminError::bad_request(format!("{field} is not editable")));
}
}
}
let update = serde_json::from_value::<EditableUpdateRequest>(payload.clone())
.map_err(|_| AdminError::bad_request(messages.invalid))?;
if is_blank(&update.id) || update.status.as_deref().is_some_and(is_blank) {
return Err(AdminError::bad_request(messages.invalid));
}
if update.label.is_none() && update.status.is_none() {
return Err(AdminError::bad_request(messages.empty_update));
}
Ok(ParsedEditableUpdate {
id: update.id,
label: update.label,
status: update.status,
})
}
fn is_blank(value: &str) -> bool {
value.trim().is_empty()
}
+1 -1
View File
@@ -5,7 +5,7 @@ use axum::{
Router,
};
use crate::bootstrap::state::AppState;
use crate::api::AppState;
use super::{
accounts_routes::{
+1 -1
View File
@@ -9,12 +9,12 @@ use axum::{
use serde::{Deserialize, Serialize};
use crate::{
api::AppState,
api::{
admin::response::{AdminEnvelope, AdminError, AdminResponse},
middleware::request_id::ClientIp,
},
auth::types::SessionError,
bootstrap::state::AppState,
infra::time::china_rfc3339,
};
+4 -3
View File
@@ -10,7 +10,7 @@ use crate::{
response::{AdminEnvelope, AdminError, AdminResponse},
session::AdminAuth,
},
bootstrap::state::AppState,
api::AppState,
settings::{SettingsError, SettingsPatch, SettingsSnapshot, SettingsValidationError},
};
@@ -52,7 +52,7 @@ pub(crate) async fn settings(
State(state): State<AppState>,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
let config = state.services.settings.current().await;
let config = state.services.settings.current();
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(AdminSettingsData::from_settings(&config)),
@@ -68,7 +68,8 @@ pub(crate) async fn update_settings(
let patch = parse_settings_patch(&body)?;
let config = state
.services
.update_settings(patch)
.settings
.update(patch)
.await
.map_err(settings_service_error)?;
+13 -6
View File
@@ -3,7 +3,7 @@
use std::{convert::Infallible, time::Duration};
use axum::{
extract::{Json, Query},
extract::{Json, Query, State},
http::StatusCode,
response::{
sse::{Event, KeepAlive, Sse},
@@ -18,7 +18,7 @@ use crate::{
response::{AdminEnvelope, AdminError, AdminResponse},
session::AdminAuth,
},
bootstrap::shutdown::{request_process_restart, request_shutdown, shutdown_subscription},
api::AppState,
update::{
service::{self, RestartAction},
state::operation_id,
@@ -56,10 +56,11 @@ pub(crate) async fn update_detail(
}
pub(crate) async fn update_event_stream(
State(state): State<AppState>,
_auth: AdminAuth,
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, AdminError> {
let receiver = service::subscribe_update_events();
let shutdown = shutdown_subscription();
let shutdown = state.services.process_control.subscribe_shutdown();
let stream = futures::stream::unfold(
(receiver, shutdown, false),
|(mut receiver, mut shutdown, close_after_send)| async move {
@@ -122,14 +123,20 @@ pub(crate) async fn rollback(_auth: AdminAuth) -> Result<impl IntoResponse, Admi
))
}
pub(crate) async fn restart(_auth: AdminAuth) -> Result<impl IntoResponse, AdminError> {
pub(crate) async fn restart(
State(state): State<AppState>,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
let plan = service::restart_plan().map_err(update_error)?;
let message = plan.message;
let process_control = state.services.process_control.clone();
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(500)).await;
match plan.action {
RestartAction::Exec(executable_path) => request_process_restart(executable_path),
RestartAction::Shutdown => request_shutdown(),
RestartAction::Exec(executable_path) => {
process_control.request_restart(executable_path);
}
RestartAction::Shutdown => process_control.request_shutdown(),
}
});
Ok(AdminResponse::new(
-62
View File
@@ -1,62 +0,0 @@
//! 管理端通用更新请求解析。
use serde::Deserialize;
use serde_json::Value;
use crate::api::admin::response::AdminError;
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct EditableUpdateRequest {
id: String,
label: Option<Option<String>>,
status: Option<String>,
}
pub(crate) struct ParsedEditableUpdate {
pub id: String,
pub label: Option<Option<String>>,
pub status: Option<String>,
}
pub(crate) struct EditableUpdateMessages<'a> {
pub object_required: &'a str,
pub invalid: &'a str,
pub empty_update: &'a str,
pub unknown_field_editable: bool,
}
pub(crate) fn parse_editable_update(
payload: &Value,
messages: EditableUpdateMessages<'_>,
) -> Result<ParsedEditableUpdate, AdminError> {
let object = payload
.as_object()
.ok_or_else(|| AdminError::bad_request(messages.object_required))?;
if messages.unknown_field_editable {
for field in object.keys() {
if !matches!(field.as_str(), "id" | "label" | "status") {
return Err(AdminError::bad_request(format!("{field} is not editable")));
}
}
}
let update = serde_json::from_value::<EditableUpdateRequest>(payload.clone())
.map_err(|_| AdminError::bad_request(messages.invalid))?;
if is_blank(&update.id) || update.status.as_deref().is_some_and(is_blank) {
return Err(AdminError::bad_request(messages.invalid));
}
if update.label.is_none() && update.status.is_none() {
return Err(AdminError::bad_request(messages.empty_update));
}
Ok(ParsedEditableUpdate {
id: update.id,
label: update.label,
status: update.status,
})
}
fn is_blank(value: &str) -> bool {
value.trim().is_empty()
}
+1 -1
View File
@@ -14,7 +14,7 @@ use crate::{
response::{AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse},
session::AdminAuth,
},
bootstrap::state::AppState,
api::AppState,
infra::{
format::{
format_compact_number, format_cost, format_duration_ms, format_duration_ms_f64,
+1 -1
View File
@@ -4,7 +4,7 @@
use axum::http::{header::AUTHORIZATION, HeaderMap};
use crate::bootstrap::state::AppState;
use crate::api::AppState;
/// client API key 鉴权失败原因。
#[derive(Debug, Clone, PartialEq, Eq)]
+1 -1
View File
@@ -8,7 +8,7 @@ use axum::{
use serde_json::{json, Value};
use crate::{
dispatch::{errors::DispatchFailureClass, responses::errors::ResponseDispatchError},
dispatch::errors::{DispatchFailureClass, ResponseDispatchError},
upstream::openai::protocol::sse::response_failed_sse_event as encode_response_failed_sse_event,
};
+1 -2
View File
@@ -9,8 +9,7 @@ use axum::{
use serde_json::{json, Value};
use crate::{
api::client::auth::authorize_client_api_key, bootstrap::state::AppState,
models::catalog::ModelCatalog,
api::client::auth::authorize_client_api_key, api::AppState, models::catalog::ModelCatalog,
};
use super::errors::{missing_client_api_key_response, model_not_found_response};
+2 -2
View File
@@ -12,8 +12,8 @@ use serde_json::{Map, Value};
use crate::{
api::middleware::request_id::{ClientIp, RequestId},
bootstrap::state::AppState,
dispatch::responses::{errors::ResponseDispatchError, service::ResponseDispatchStream},
api::AppState,
dispatch::{errors::ResponseDispatchError, service::ResponseDispatchStream},
upstream::openai::protocol::{
responses::{CodexCompactRequest, CodexResponsesRequest},
sse::sse_body_has_done,
+1 -1
View File
@@ -6,7 +6,7 @@ use axum::{
Router,
};
use crate::bootstrap::state::AppState;
use crate::api::AppState;
use super::{
models::{model_catalog, model_detail, model_info, models},
+2
View File
@@ -5,3 +5,5 @@ pub mod assets;
pub mod client;
pub mod middleware;
pub mod router;
pub use router::AppState;
+62 -8
View File
@@ -1,17 +1,75 @@
//! 顶层 HTTP 路由 —— 组合 OpenAI API、管理端 API 和静态资源服务。
use std::path::Path;
use std::{
path::{Path, PathBuf},
sync::Arc,
};
use async_trait::async_trait;
use axum::{extract::State, http::StatusCode, middleware, routing::get, Router};
use crate::{
accounts::{
manage::AccountManageService, pool::RuntimeAccountPoolService,
refresh::RuntimeTokenRefreshService, store::AccountStore,
},
api::{
admin, assets, client,
middleware::{request_id::attach_request_id, trace::http_trace_layer},
},
bootstrap::state::AppState,
auth::service::SessionService,
dispatch::{affinity::RuntimeSessionAffinityService, service::ResponseDispatchService},
keys::{manage::KeyManageService, service::KeyVerifier},
models::service::ModelService,
settings::service::RuntimeSettingsService,
telemetry::{
account_usage::query::AccountUsageQueryService, ops::query::OpsQueryService,
usage::query::UsageQueryService,
},
upstream::openai::{fingerprint::RuntimeFingerprint, token_client::OpenAiTokenClient},
};
/// HTTP API 所需的领域服务集合。
#[derive(Clone)]
pub struct ApiServices {
pub health_probe: Arc<dyn HealthProbe>,
pub models: Arc<ModelService>,
pub accounts: Arc<dyn AccountStore>,
pub client_keys: Arc<KeyVerifier>,
pub admin_client_keys: Arc<KeyManageService>,
pub admin_sessions: Arc<SessionService>,
pub settings: Arc<RuntimeSettingsService>,
pub admin_accounts: Arc<AccountManageService>,
pub usage_records: Arc<UsageQueryService>,
pub ops_errors: Arc<OpsQueryService>,
pub usage: Arc<AccountUsageQueryService>,
pub account_pool: Arc<RuntimeAccountPoolService>,
pub token_refresh: Arc<RuntimeTokenRefreshService<OpenAiTokenClient>>,
pub responses: Arc<ResponseDispatchService>,
pub session_affinity: Arc<RuntimeSessionAffinityService>,
pub fingerprint: RuntimeFingerprint,
pub process_control: Arc<dyn ProcessControl>,
}
/// 由 bootstrap 实现的进程生命周期控制端口。
pub trait ProcessControl: Send + Sync + 'static {
fn request_shutdown(&self);
fn request_restart(&self, executable_path: PathBuf);
fn subscribe_shutdown(&self) -> tokio::sync::broadcast::Receiver<()>;
}
/// 由 bootstrap 实现的基础设施健康检查端口。
#[async_trait]
pub trait HealthProbe: Send + Sync + 'static {
async fn check(&self) -> Result<(), String>;
}
/// HTTP handler 共享状态。
#[derive(Clone)]
pub struct AppState {
pub services: ApiServices,
}
/// 默认前端构建产物目录。
pub const DEFAULT_ASSET_DIST_DIR: &str = "web/dist";
@@ -32,12 +90,8 @@ pub fn router_with_assets(dist_dir: impl AsRef<Path>) -> Router<AppState> {
}
async fn healthz(State(state): State<AppState>) -> StatusCode {
if let Err(error) = crate::infra::database::ping(&state.services.database).await {
tracing::warn!(error = %error, "PostgreSQL health check failed");
return StatusCode::SERVICE_UNAVAILABLE;
}
if let Err(error) = state.services.redis.ping().await {
tracing::warn!(error = %error, "Redis health check failed");
if let Err(error) = state.services.health_probe.check().await {
tracing::warn!(error, "health check failed");
return StatusCode::SERVICE_UNAVAILABLE;
}
StatusCode::NO_CONTENT
-168
View File
@@ -1,168 +0,0 @@
use std::{env, error::Error, time::Duration};
use axum::Router;
use crate::bootstrap::config::AppConfig;
use crate::bootstrap::services::{
apply_settings_to_config, fingerprint_from_config, settings_snapshot_from_config,
BackgroundTaskStores, Services,
};
use crate::bootstrap::shutdown::{restart_executable_path, shutdown_signal};
use crate::bootstrap::state::{AppState, RuntimeConfig};
use crate::bootstrap::tasks::coordinator::TaskCoordinator;
use crate::infra::database::connect;
use crate::infra::logging::{init_tracing, LogGuard, RotationConfig};
use crate::infra::paths::{ensure_data_dir, load_or_create_installation_id};
use crate::infra::redis::RedisConnection;
use crate::settings::service::RuntimeSettingsService;
use crate::upstream::openai::fingerprint::{PgFingerprintStore, RuntimeFingerprint};
const RESTART_DELAY_ENV: &str = "CPR_RESTART_DELAY_MS";
pub async fn run() -> Result<(), Box<dyn Error + Send + Sync>> {
wait_for_scheduled_restart().await;
let config = AppConfig::load()?;
let host = config.server.host.clone();
let port = config.server.port;
let log_guard = init_logging(&config)?;
let listener = tokio::net::TcpListener::bind((host.as_str(), port)).await?;
let (app, task_coordinator) = build_application(config).await?;
let serve_result = axum::serve(
listener,
app.into_make_service_with_connect_info::<std::net::SocketAddr>(),
)
.with_graceful_shutdown(shutdown_signal())
.await;
task_coordinator.shutdown().await;
serve_result?;
if let Some(executable_path) = restart_executable_path() {
tracing::info!(
executable = %executable_path.display(),
"正在以更新后的二进制替换当前进程"
);
drop(log_guard);
exec_replacement_process(executable_path)?;
}
Ok(())
}
async fn build_application(
mut config: AppConfig,
) -> Result<(Router, TaskCoordinator), Box<dyn Error + Send + Sync>> {
let pool = connect(&config.database.url).await?;
let redis = RedisConnection::connect(&config.redis.url, "cpr").await?;
let settings =
RuntimeSettingsService::load_or_initialize(settings_snapshot_from_config(&config), &pool)
.await?;
apply_settings_to_config(&mut config, &settings);
let fingerprint_store = PgFingerprintStore::new(pool.clone());
let stores = BackgroundTaskStores {
redis: redis.clone(),
accounts: crate::accounts::store::PgAccountStore::new(pool.clone()),
admin_users: crate::auth::store::PgAdminUserStore::new(pool.clone()),
admin_sessions: crate::auth::store::RedisAdminSessionStore::new(redis.clone()),
cookies: crate::accounts::cookies::PgCookieStore::new(pool.clone()),
fingerprints: fingerprint_store.clone(),
session_affinity: crate::dispatch::affinity::RedisSessionAffinityStore::new(redis.clone()),
refresh_leases: crate::accounts::refresh::RedisRefreshLeaseStore::new(redis.clone()),
model_snapshots: crate::models::store::RedisModelSnapshotStore::new(redis.clone()),
client_keys: crate::keys::store::PgClientKeyStore::new(pool.clone()),
usage_records: crate::telemetry::usage::store::PgUsageRecordStore::new(pool.clone()),
ops_errors: crate::telemetry::ops::store::PgOpsErrorLogStore::new(pool.clone()),
account_usage: crate::telemetry::account_usage::store::PgAccountUsageStore::new(
pool.clone(),
),
request_buckets: crate::telemetry::buckets::store::PgRequestBucketStore::new(pool),
};
config.admin.validate_default_password()?;
let default_admin_password = config.admin.default_password.clone();
let data_dir = ensure_data_dir()?;
let installation_id = load_or_create_installation_id(Some(&data_dir))?;
let default_fingerprint = fingerprint_from_config(&config.fingerprint);
let runtime_fingerprint = fingerprint_store
.ensure_current_seed(&default_fingerprint)
.await?;
let runtime_fingerprint = RuntimeFingerprint::new(runtime_fingerprint);
let runtime_config = RuntimeConfig::from(config.clone());
let services = Services::try_with_installation_id(
&config,
stores,
runtime_fingerprint,
Some(installation_id),
)?;
services.initialize_hot_path_state().await?;
let created_default_admin = services
.admin_sessions
.ensure_default_admin(&default_admin_password)
.await?;
tracing::info!(
created = created_default_admin,
"默认管理员账号已完成初始化检查"
);
let restored_accounts = services.account_pool.restore_from_store().await?;
tracing::info!(
count = restored_accounts,
"运行时账号池已从 PostgreSQL 恢复"
);
let task_coordinator = TaskCoordinator::start(&runtime_config, &services);
let app_state = AppState { services };
Ok((
crate::api::router::router().with_state(app_state),
task_coordinator,
))
}
fn init_logging(config: &AppConfig) -> Result<Option<LogGuard>, crate::infra::logging::LogError> {
if !config.logging.enabled {
return Ok(None);
}
let rotation = RotationConfig::new(&config.logging.directory, config.logging.retention_days);
Ok(Some(init_tracing(&rotation)?))
}
async fn wait_for_scheduled_restart() {
let Some(delay_ms) = env::var(RESTART_DELAY_ENV)
.ok()
.and_then(|value| value.parse::<u64>().ok())
.filter(|value| *value > 0)
else {
return;
};
tokio::time::sleep(Duration::from_millis(delay_ms)).await;
}
#[cfg(unix)]
fn exec_replacement_process(
executable_path: std::path::PathBuf,
) -> Result<(), Box<dyn Error + Send + Sync>> {
use std::os::unix::process::CommandExt;
use std::process::Command;
let error = Command::new(executable_path)
.args(env::args_os().skip(1))
.exec();
Err(Box::new(error))
}
#[cfg(not(unix))]
fn exec_replacement_process(
_executable_path: std::path::PathBuf,
) -> Result<(), Box<dyn Error + Send + Sync>> {
Err(Box::new(std::io::Error::new(
std::io::ErrorKind::Unsupported,
"process exec restart is only supported on Unix",
)))
}
-1
View File
@@ -1,6 +1,5 @@
//! 进程装配、共享状态与后台任务。
pub mod application;
pub mod config;
pub mod import_sqlite;
pub mod services;
+222 -76
View File
@@ -1,17 +1,18 @@
//! 运行时服务集合。
use std::sync::Arc;
use std::{env, error::Error, sync::Arc, time::Duration};
use axum::Router;
use crate::{
accounts::{
cookies::PgCookieStore,
manage::{AccountManageService, AccountManageServiceParts},
pool::RuntimeAccountPoolService,
pool::{AccountPoolOptions, RotationStrategy},
pool::{AccountPoolOptions, AccountPoolStaticSettings, RuntimeAccountPoolService},
refresh::{
RedisRefreshLeaseStore, RefreshPolicy, RuntimeRefreshPolicy, RuntimeTokenRefreshService,
},
store::{AccountStore as AccountStoreTrait, PgAccountStore},
store::{AccountStore, PgAccountStore},
},
auth::{
service::SessionService,
@@ -20,19 +21,24 @@ use crate::{
bootstrap::config::AppConfig,
dispatch::{
affinity::{RedisSessionAffinityStore, RuntimeSessionAffinityService},
responses::service::{ResponseDispatchService, ResponseDispatchServiceParts},
service::{ResponseDispatchService, ResponseDispatchServiceParts},
},
infra::{
database::connect,
logging::{init_tracing, LogGuard, RotationConfig},
paths::{ensure_data_dir, load_or_create_installation_id},
redis::RedisConnection,
},
infra::redis::RedisConnection,
keys::{manage::KeyManageService, service::KeyVerifier, store::PgClientKeyStore},
models::{
config::ModelConfig,
service::{ModelService, ModelServiceError},
store::{ModelSnapshotStore, RedisModelSnapshotStore},
types::ModelConfig,
},
settings::{service::RuntimeSettingsService, SettingsError, SettingsPatch, SettingsSnapshot},
settings::{service::RuntimeSettingsService, SettingsSnapshot},
telemetry::{
account_usage::query::AccountUsageQueryService,
account_usage::store::PgAccountUsageStore,
account_usage::store::{AccountUsageStore, PgAccountUsageStore},
buckets::store::PgRequestBucketStore,
ops::query::OpsQueryService,
ops::store::PgOpsErrorLogStore,
@@ -50,6 +56,8 @@ use crate::{
},
};
const RESTART_DELAY_ENV: &str = "CPR_RESTART_DELAY_MS";
/// 运行时热路径内存状态初始化错误。
#[derive(Debug, thiserror::Error)]
pub enum RuntimeServiceInitializationError {
@@ -90,7 +98,7 @@ pub struct Services {
pub database: sqlx::PgPool,
pub redis: RedisConnection,
pub models: Arc<ModelService>,
pub accounts: Arc<dyn AccountStoreTrait>,
pub accounts: Arc<dyn AccountStore>,
pub client_keys: Arc<KeyVerifier>,
pub admin_client_keys: Arc<KeyManageService>,
pub admin_sessions: Arc<SessionService>,
@@ -109,14 +117,7 @@ pub struct Services {
pub fingerprint: RuntimeFingerprint,
pub installation_id: Option<String>,
pub background_tasks: BackgroundTaskStores,
account_pool_static: AccountPoolStaticSettings,
settings_update_lock: Arc<tokio::sync::Mutex<()>>,
}
#[derive(Debug, Clone)]
struct AccountPoolStaticSettings {
skip_quota_limited: bool,
tier_priority: Vec<String>,
pub(crate) account_pool_static: AccountPoolStaticSettings,
}
/// 使用记录运行选项。
@@ -178,7 +179,9 @@ impl Services {
let installation_id = installation_id.filter(|id| !id.trim().is_empty());
let database = stores.accounts.pool().clone();
let redis = stores.redis.clone();
let account_store_trait = Arc::new(stores.accounts.clone()) as Arc<dyn AccountStoreTrait>;
let account_store = Arc::new(stores.accounts.clone()) as Arc<dyn AccountStore>;
let account_usage_store_trait =
Arc::new(stores.account_usage.clone()) as Arc<dyn AccountUsageStore>;
let websocket_pool = config.ws_pool.enabled.then(|| {
Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig {
enabled: true,
@@ -213,24 +216,23 @@ impl Services {
usage_record_options.enabled,
usage_record_options.capture_body,
));
let settings_snapshot = settings_snapshot_from_config(config);
let account_pool_static = AccountPoolStaticSettings {
skip_quota_limited: config.quota.skip_exhausted,
tier_priority: config.auth.tier_priority.clone(),
};
let account_pool = Arc::new(RuntimeAccountPoolService::new(
account_store_trait.clone(),
account_pool_options_from_settings(
&settings_snapshot_from_config(config),
&account_pool_static,
),
config.auth.request_interval_ms,
account_store.clone(),
account_usage_store_trait,
account_pool_static.pool_options(&settings_snapshot),
settings_snapshot.request_interval_ms,
));
let refresh_policy = RuntimeRefreshPolicy::new(RefreshPolicy {
refresh_margin_seconds: config.auth.refresh_margin_seconds,
refresh_concurrency: config.auth.refresh_concurrency,
});
let settings = Arc::new(RuntimeSettingsService::new(
settings_snapshot_from_config(config),
settings_snapshot,
stores.accounts.pool().clone(),
));
let admin_sessions = Arc::new(SessionService::new(
@@ -267,7 +269,6 @@ impl Services {
account_pool: account_pool.clone(),
token_refresher: Arc::new(token_client),
refresh_leases: stores.refresh_leases.clone(),
session_affinity: stores.session_affinity.clone(),
oauth: crate::accounts::manage::oauth::AccountOAuthService::new(
reqwest::Client::new(),
config.auth.oauth_client_id.clone(),
@@ -282,7 +283,7 @@ impl Services {
));
let cloudflare_recovery =
crate::dispatch::cloudflare::CloudflareRecovery::new(stores.cookies.clone());
crate::dispatch::recovery::cloudflare::CloudflareRecovery::new(stores.cookies.clone());
let responses = Arc::new(ResponseDispatchService::new(ResponseDispatchServiceParts {
account_pool: account_pool.clone(),
models: models.clone(),
@@ -297,7 +298,7 @@ impl Services {
database,
redis,
models,
accounts: account_store_trait,
accounts: account_store,
client_keys,
admin_client_keys,
admin_sessions,
@@ -317,7 +318,6 @@ impl Services {
installation_id,
background_tasks: stores,
account_pool_static,
settings_update_lock: Arc::new(tokio::sync::Mutex::new(())),
})
}
@@ -326,32 +326,46 @@ impl Services {
self.models.reload_from_store().await?;
Ok(())
}
}
/// 持久化设置并同步到所有运行时消费者。
pub async fn update_settings(
&self,
patch: SettingsPatch,
) -> Result<Arc<SettingsSnapshot>, SettingsError> {
let _guard = self.settings_update_lock.lock().await;
let settings = self.settings.update(patch).await?;
self.apply_settings(&settings).await;
Ok(settings)
impl From<&Services> for crate::api::router::ApiServices {
fn from(services: &Services) -> Self {
Self {
health_probe: Arc::new(crate::bootstrap::state::RuntimeHealthProbe::new(
services.database.clone(),
services.redis.clone(),
)),
models: services.models.clone(),
accounts: services.accounts.clone(),
client_keys: services.client_keys.clone(),
admin_client_keys: services.admin_client_keys.clone(),
admin_sessions: services.admin_sessions.clone(),
settings: services.settings.clone(),
admin_accounts: services.admin_accounts.clone(),
usage_records: services.usage_records.clone(),
ops_errors: services.ops_errors.clone(),
usage: services.usage.clone(),
account_pool: services.account_pool.clone(),
token_refresh: services.token_refresh.clone(),
responses: services.responses.clone(),
session_affinity: services.session_affinity.clone(),
fingerprint: services.fingerprint.clone(),
process_control: Arc::new(crate::bootstrap::shutdown::RuntimeProcessControl),
}
}
}
async fn apply_settings(&self, settings: &SettingsSnapshot) {
self.models.update_config(ModelConfig {
model_aliases: settings.model_aliases.clone(),
});
self.account_pool
.apply_options(
account_pool_options_from_settings(settings, &self.account_pool_static),
settings.request_interval_ms,
)
.await;
self.refresh_policy.update(RefreshPolicy {
refresh_margin_seconds: settings.refresh_margin_seconds,
refresh_concurrency: settings.refresh_concurrency,
});
impl From<&Services> for crate::api::router::AppState {
fn from(services: &Services) -> Self {
Self {
services: crate::api::router::ApiServices::from(services),
}
}
}
impl From<Services> for crate::api::router::AppState {
fn from(services: Services) -> Self {
Self::from(&services)
}
}
@@ -403,31 +417,163 @@ pub fn fingerprint_from_config(
}
}
fn account_pool_options_from_settings(
settings: &SettingsSnapshot,
static_settings: &AccountPoolStaticSettings,
) -> 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: static_settings.skip_quota_limited,
tier_priority: static_settings.tier_priority.clone(),
..AccountPoolOptions::default()
pub fn account_pool_options_from_config(config: &AppConfig) -> AccountPoolOptions {
AccountPoolStaticSettings {
skip_quota_limited: config.quota.skip_exhausted,
tier_priority: config.auth.tier_priority.clone(),
}
.pool_options(&settings_snapshot_from_config(config))
}
pub fn account_pool_options_from_config(config: &AppConfig) -> AccountPoolOptions {
account_pool_options_from_settings(
&settings_snapshot_from_config(config),
&AccountPoolStaticSettings {
skip_quota_limited: config.quota.skip_exhausted,
tier_priority: config.auth.tier_priority.clone(),
},
pub async fn run() -> Result<(), Box<dyn Error + Send + Sync>> {
wait_for_scheduled_restart().await;
let config = AppConfig::load()?;
let host = config.server.host.clone();
let port = config.server.port;
let log_guard = init_logging(&config)?;
let listener = tokio::net::TcpListener::bind((host.as_str(), port)).await?;
let (router, task_coordinator) = build_router(config).await?;
let serve_result = axum::serve(
listener,
router.into_make_service_with_connect_info::<std::net::SocketAddr>(),
)
.with_graceful_shutdown(crate::bootstrap::shutdown::shutdown_signal())
.await;
task_coordinator.shutdown().await;
serve_result?;
if let Some(executable_path) = crate::bootstrap::shutdown::restart_executable_path() {
tracing::info!(
executable = %executable_path.display(),
"正在以更新后的二进制替换当前进程"
);
drop(log_guard);
exec_replacement_process(executable_path)?;
}
Ok(())
}
async fn build_router(
mut config: AppConfig,
) -> Result<
(
Router,
crate::bootstrap::tasks::coordinator::TaskCoordinator,
),
Box<dyn Error + Send + Sync>,
> {
let pool = connect(&config.database.url).await?;
let redis = RedisConnection::connect(&config.redis.url, "cpr").await?;
let settings =
RuntimeSettingsService::load_or_initialize(settings_snapshot_from_config(&config), &pool)
.await?;
apply_settings_to_config(&mut config, &settings);
let fingerprint_store = PgFingerprintStore::new(pool.clone());
let stores = BackgroundTaskStores {
redis: redis.clone(),
accounts: PgAccountStore::new(pool.clone()),
admin_users: PgAdminUserStore::new(pool.clone()),
admin_sessions: RedisAdminSessionStore::new(redis.clone()),
cookies: PgCookieStore::new(pool.clone()),
fingerprints: fingerprint_store.clone(),
session_affinity: RedisSessionAffinityStore::new(redis.clone()),
refresh_leases: RedisRefreshLeaseStore::new(redis.clone()),
model_snapshots: RedisModelSnapshotStore::new(redis.clone()),
client_keys: PgClientKeyStore::new(pool.clone()),
usage_records: PgUsageRecordStore::new(pool.clone()),
ops_errors: PgOpsErrorLogStore::new(pool.clone()),
account_usage: PgAccountUsageStore::new(pool.clone()),
request_buckets: PgRequestBucketStore::new(pool),
};
config.admin.validate_default_password()?;
let default_admin_password = config.admin.default_password.clone();
let data_dir = ensure_data_dir()?;
let installation_id = load_or_create_installation_id(Some(&data_dir))?;
let default_fingerprint = fingerprint_from_config(&config.fingerprint);
let runtime_fingerprint = fingerprint_store
.ensure_current_seed(&default_fingerprint)
.await?;
let runtime_fingerprint = RuntimeFingerprint::new(runtime_fingerprint);
let runtime_config = crate::bootstrap::state::RuntimeConfig::from(config.clone());
let services = Services::try_with_installation_id(
&config,
stores,
runtime_fingerprint,
Some(installation_id),
)?;
services.initialize_hot_path_state().await?;
let created_default_admin = services
.admin_sessions
.ensure_default_admin(&default_admin_password)
.await?;
tracing::info!(
created = created_default_admin,
"默认管理员账号已完成初始化检查"
);
let restored_accounts = services.account_pool.restore_from_store().await?;
tracing::info!(
count = restored_accounts,
"运行时账号池已从 PostgreSQL 恢复"
);
let task_coordinator =
crate::bootstrap::tasks::coordinator::TaskCoordinator::start(&runtime_config, &services);
let app_state = crate::api::AppState::from(&services);
Ok((
crate::api::router::router().with_state(app_state),
task_coordinator,
))
}
fn init_logging(config: &AppConfig) -> Result<Option<LogGuard>, crate::infra::logging::LogError> {
if !config.logging.enabled {
return Ok(None);
}
let rotation = RotationConfig::new(&config.logging.directory, config.logging.retention_days);
Ok(Some(init_tracing(&rotation)?))
}
async fn wait_for_scheduled_restart() {
let Some(delay_ms) = env::var(RESTART_DELAY_ENV)
.ok()
.and_then(|value| value.parse::<u64>().ok())
.filter(|value| *value > 0)
else {
return;
};
tokio::time::sleep(Duration::from_millis(delay_ms)).await;
}
#[cfg(unix)]
fn exec_replacement_process(
executable_path: std::path::PathBuf,
) -> Result<(), Box<dyn Error + Send + Sync>> {
use std::os::unix::process::CommandExt;
use std::process::Command;
let error = Command::new(executable_path)
.args(env::args_os().skip(1))
.exec();
Err(Box::new(error))
}
#[cfg(not(unix))]
fn exec_replacement_process(
_executable_path: std::path::PathBuf,
) -> Result<(), Box<dyn Error + Send + Sync>> {
Err(Box::new(std::io::Error::new(
std::io::ErrorKind::Unsupported,
"process exec restart is only supported on Unix",
)))
}
+17
View File
@@ -15,6 +15,23 @@ static SHUTDOWN_REQUESTS: OnceLock<broadcast::Sender<()>> = OnceLock::new();
static RESTART_REQUESTED: AtomicBool = AtomicBool::new(false);
static RESTART_EXECUTABLE_PATH: OnceLock<PathBuf> = OnceLock::new();
#[derive(Debug, Clone, Copy, Default)]
pub struct RuntimeProcessControl;
impl crate::api::router::ProcessControl for RuntimeProcessControl {
fn request_shutdown(&self) {
request_shutdown();
}
fn request_restart(&self, executable_path: PathBuf) {
request_process_restart(executable_path);
}
fn subscribe_shutdown(&self) -> Receiver<()> {
shutdown_subscription()
}
}
/// 等待进程关闭信号(Ctrl+C 或 SIGTERM)。
pub async fn shutdown_signal() {
let mut requested = shutdown_sender().subscribe();
+30 -7
View File
@@ -1,13 +1,10 @@
//! 运行时共享状态。
use crate::bootstrap::config::AppConfig;
use crate::bootstrap::services::Services;
use async_trait::async_trait;
/// HTTP handler 共享的运行时状态。
#[derive(Clone)]
pub struct AppState {
pub services: Services,
}
use crate::{
api::router::HealthProbe, bootstrap::config::AppConfig, infra::redis::RedisConnection,
};
/// 运行时配置镜像(从 AppConfig 衍生)。
#[derive(Debug, Clone)]
@@ -17,6 +14,32 @@ pub struct RuntimeConfig {
pub quota: crate::bootstrap::config::QuotaConfig,
}
/// PostgreSQL 与 Redis 的进程健康检查实现。
#[derive(Clone)]
pub struct RuntimeHealthProbe {
database: sqlx::PgPool,
redis: RedisConnection,
}
impl RuntimeHealthProbe {
pub fn new(database: sqlx::PgPool, redis: RedisConnection) -> Self {
Self { database, redis }
}
}
#[async_trait]
impl HealthProbe for RuntimeHealthProbe {
async fn check(&self) -> Result<(), String> {
crate::infra::database::ping(&self.database)
.await
.map_err(|error| format!("PostgreSQL: {error}"))?;
self.redis
.ping()
.await
.map_err(|error| format!("Redis: {error}"))
}
}
impl From<AppConfig> for RuntimeConfig {
fn from(config: AppConfig) -> Self {
Self {
+45 -2
View File
@@ -1,6 +1,6 @@
//! 后台任务协调器。
use std::{path::PathBuf, sync::Arc, time::Duration};
use std::{future::Future, path::PathBuf, sync::Arc, time::Duration};
use tokio::task::JoinHandle;
@@ -33,7 +33,7 @@ impl TaskCoordinator {
config: &crate::bootstrap::state::RuntimeConfig,
services: &crate::bootstrap::services::Services,
) -> Self {
let mut coordinator = TaskCoordinator::default();
let mut coordinator = Self::start_settings_subscriptions(services);
let stores = &services.background_tasks;
coordinator.push(
@@ -65,6 +65,7 @@ impl TaskCoordinator {
"quota_refresh",
QuotaRefreshTask::with_intervals(
stores.accounts.clone(),
stores.account_usage.clone(),
services.codex.clone(),
config
.quota
@@ -99,6 +100,37 @@ impl TaskCoordinator {
coordinator
}
/// 启动各领域的设置订阅任务。
pub fn start_settings_subscriptions(services: &crate::bootstrap::services::Services) -> Self {
let mut coordinator = TaskCoordinator::default();
coordinator.push(
"model_settings",
SchedulerHandle::spawn(
services
.models
.clone()
.subscribe_settings(services.settings.subscribe()),
),
);
coordinator.push(
"account_pool_settings",
SchedulerHandle::spawn(services.account_pool.clone().subscribe_settings(
services.settings.subscribe(),
services.account_pool_static.clone(),
)),
);
coordinator.push(
"refresh_policy_settings",
SchedulerHandle::spawn(
services
.refresh_policy
.clone()
.subscribe_settings(services.settings.subscribe()),
),
);
coordinator
}
/// 关闭所有后台任务。
pub async fn shutdown(self) {
tracing::info!("正在关闭后台任务");
@@ -122,6 +154,17 @@ pub enum SchedulerHandle {
}
impl SchedulerHandle {
fn spawn(future: impl Future<Output = ()> + Send + 'static) -> Self {
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
let handle = tokio::spawn(async move {
tokio::select! {
() = future => {}
_ = shutdown_rx.recv() => {}
}
});
Self::new(shutdown_tx, handle)
}
/// 使用关闭 channel 构造句柄。
pub(crate) fn new(shutdown_tx: tokio::sync::mpsc::Sender<()>, handle: JoinHandle<()>) -> Self {
Self::Channel {
+2 -1
View File
@@ -75,7 +75,8 @@ impl ModelRefreshTask {
.iter()
.map(|selected| ModelRefreshPlanAccount {
plan_type: selected.plan_type.clone(),
account: selected.account.clone(),
access_token: selected.account.access_token.clone(),
account_id: selected.account.account_id.clone(),
})
.collect::<Vec<_>>();
@@ -12,6 +12,7 @@ use crate::{
quota::{QuotaRefreshSummary, RuntimeQuotaRefreshService},
store::PgAccountStore,
},
telemetry::account_usage::store::PgAccountUsageStore,
upstream::openai::transport::CodexBackendClient,
};
@@ -31,6 +32,7 @@ impl QuotaRefreshTask {
/// 使用自定义间隔构造配额刷新后台任务。
pub fn with_intervals(
store: PgAccountStore,
usage_store: PgAccountUsageStore,
codex: Arc<CodexBackendClient>,
interval_secs: u64,
min_refresh_interval_secs: u64,
@@ -38,6 +40,7 @@ impl QuotaRefreshTask {
Self {
service: RuntimeQuotaRefreshService::with_min_refresh_interval_secs(
store,
usage_store,
codex,
min_refresh_interval_secs,
),
+1 -1
View File
@@ -6,7 +6,7 @@ use tokio::sync::Mutex;
use crate::{
dispatch::{
affinity::{compute_variant_hash, hash_instructions, RuntimeSessionAffinityService},
reasoning_replay::ReasoningReplayCache,
recovery::reasoning_replay::ReasoningReplayCache,
},
upstream::openai::protocol::{
events::TokenUsage,
+8
View File
@@ -60,6 +60,14 @@ impl RuntimeSessionAffinityService {
.and_then(|entry| entry.turn_state)
}
pub async fn lookup(
&self,
response_id: &str,
now: DateTime<Utc>,
) -> Option<SessionAffinityEntry> {
self.entry(response_id, now).await
}
pub async fn lookup_instructions_hash(
&self,
response_id: &str,
+331 -20
View File
@@ -2,12 +2,16 @@
use chrono::{DateTime, Duration, Utc};
use serde_json::{json, Map, Value};
use thiserror::Error;
use crate::{
accounts::account::AccountStatus,
dispatch::exhaustion::{ExhaustedAccountKind, ExhaustedAccountRef},
upstream::openai::transport::{
is_banned_auth_signal, CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics,
dispatch::{
recovery::exhaustion::{ExhaustedAccount, ExhaustedAccountKind, ExhaustedAccountRef},
stream::sse_failure::{sse_failure_error_body, stream_failure_http_status},
},
upstream::openai::{
protocol::{responses::ResponsesSseFailure, sse::SseError},
transport::{CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics},
},
};
@@ -202,13 +206,6 @@ pub(crate) fn is_quota_exhausted_upstream_error(error: &CodexClientError) -> boo
)
}
pub(crate) fn is_auth_upstream_error(error: &CodexClientError) -> bool {
matches!(
error,
CodexClientError::Upstream { status, .. } if status.as_u16() == 401
)
}
pub(crate) fn is_model_unsupported_upstream_error(error: &CodexClientError) -> bool {
matches!(
error,
@@ -226,15 +223,6 @@ pub(crate) fn is_history_recovery_upstream_error(error: &CodexClientError) -> bo
)
}
pub(crate) fn auth_failure_account_status(error: &CodexClientError) -> AccountStatus {
match error {
CodexClientError::Upstream { body, .. } if is_banned_auth_signal(body) => {
AccountStatus::Banned
}
_ => AccountStatus::Expired,
}
}
pub(crate) fn upstream_error_body(error: &CodexClientError) -> String {
match error {
CodexClientError::Upstream { body, .. } => body.clone(),
@@ -315,3 +303,326 @@ fn status_code_is_quota_exhausted(status_code: u16) -> bool {
fn status_code_is_transient_upstream(status_code: u16) -> bool {
matches!(status_code, 500..=599)
}
/// Responses 调度错误。
#[derive(Debug, Error)]
pub enum ResponseDispatchError {
#[error("no active account is available")]
NoActiveAccount,
#[error("all accounts exhausted by quota")]
QuotaExhausted {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by rate limit")]
RateLimited {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by expired auth")]
Expired {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by disabled auth")]
Disabled {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by banned auth")]
Banned {
count: usize,
upstream_error: String,
status_code: u16,
},
#[error("all accounts exhausted by Cloudflare challenge")]
CloudflareChallenge {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by Cloudflare path-block")]
CloudflarePathBlocked {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by unsupported model")]
ModelUnsupported {
count: usize,
upstream_error: String,
},
#[error("previous response context is unavailable: {upstream_error}")]
HistoryUnavailable { upstream_error: String },
#[error("upstream request failed: {0}")]
Upstream(#[from] CodexClientError),
#[error("invalid upstream SSE response: {0}")]
InvalidSse(#[from] SseError),
#[error("upstream response did not include response.completed")]
MissingCompleted,
#[error("upstream response did not include visible output")]
EmptyUpstreamResponse,
#[error("upstream response failed: {0:?}")]
Failed(ResponsesSseFailure),
}
impl ResponseDispatchError {
pub(crate) fn from_exhausted_account(exhausted: ExhaustedAccount) -> Self {
match exhausted.kind {
ExhaustedAccountKind::QuotaExhausted => Self::QuotaExhausted {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::RateLimited => Self::RateLimited {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::Expired => Self::Expired {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::Disabled => Self::Disabled {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::Banned => Self::Banned {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
status_code: exhausted.status_code.unwrap_or(403),
},
ExhaustedAccountKind::CloudflareChallenge => Self::CloudflareChallenge {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::CloudflarePathBlocked => Self::CloudflarePathBlocked {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::ModelUnsupported => Self::ModelUnsupported {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
}
}
pub fn http_status_code(&self) -> u16 {
match self {
Self::NoActiveAccount => 503,
Self::QuotaExhausted { .. } => 429,
Self::RateLimited { .. } => 429,
Self::Expired { .. } | Self::Disabled { .. } => 401,
Self::Banned { status_code, .. } => *status_code,
Self::CloudflareChallenge { .. }
| Self::CloudflarePathBlocked { .. }
| Self::InvalidSse(_)
| Self::MissingCompleted
| Self::EmptyUpstreamResponse => 502,
Self::Failed(failure) => stream_failure_http_status(failure),
Self::ModelUnsupported { .. } | Self::HistoryUnavailable { .. } => 400,
Self::Upstream(error) => upstream_error_http_status(error),
}
}
pub(crate) fn metadata(&self) -> DispatchErrorMetadata {
match self {
Self::NoActiveAccount => DispatchErrorMetadata::no_available_accounts(),
Self::QuotaExhausted {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::QuotaExhausted,
*count,
upstream_error,
),
Self::RateLimited {
count,
upstream_error,
} => {
Self::exhausted_metadata(ExhaustedAccountKind::RateLimited, *count, upstream_error)
}
Self::Expired {
count,
upstream_error,
} => Self::exhausted_metadata(ExhaustedAccountKind::Expired, *count, upstream_error),
Self::Disabled {
count,
upstream_error,
} => Self::exhausted_metadata(ExhaustedAccountKind::Disabled, *count, upstream_error),
Self::Banned {
count,
upstream_error,
..
} => Self::exhausted_metadata(ExhaustedAccountKind::Banned, *count, upstream_error),
Self::CloudflareChallenge {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::CloudflareChallenge,
*count,
upstream_error,
),
Self::CloudflarePathBlocked {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::CloudflarePathBlocked,
*count,
upstream_error,
),
Self::ModelUnsupported {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::ModelUnsupported,
*count,
upstream_error,
),
Self::HistoryUnavailable { upstream_error } => DispatchErrorMetadata {
failure_class: DispatchFailureClass::HistoryUnavailable,
exhausted_count: None,
upstream_error: Some(upstream_error.clone()),
upstream_status: Some(400),
diagnostics: None,
},
Self::Upstream(error) => DispatchErrorMetadata::upstream(error),
Self::InvalidSse(_) => DispatchErrorMetadata::simple(DispatchFailureClass::InvalidSse),
Self::MissingCompleted => {
DispatchErrorMetadata::simple(DispatchFailureClass::MissingCompleted)
}
Self::EmptyUpstreamResponse => {
DispatchErrorMetadata::simple(DispatchFailureClass::EmptyUpstreamResponse)
}
Self::Failed(failure) => DispatchErrorMetadata {
failure_class: DispatchFailureClass::ResponseFailed,
exhausted_count: None,
upstream_error: Some(sse_failure_error_body(failure)),
upstream_status: None,
diagnostics: None,
},
}
}
pub(crate) fn exhausted_account(&self) -> Option<ExhaustedAccountRef<'_>> {
match self {
Self::QuotaExhausted {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::QuotaExhausted,
count: *count,
upstream_error,
}),
Self::RateLimited {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::RateLimited,
count: *count,
upstream_error,
}),
Self::Expired {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Expired,
count: *count,
upstream_error,
}),
Self::Disabled {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Disabled,
count: *count,
upstream_error,
}),
Self::Banned {
count,
upstream_error,
..
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Banned,
count: *count,
upstream_error,
}),
Self::CloudflareChallenge {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::CloudflareChallenge,
count: *count,
upstream_error,
}),
Self::CloudflarePathBlocked {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::CloudflarePathBlocked,
count: *count,
upstream_error,
}),
Self::ModelUnsupported {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::ModelUnsupported,
count: *count,
upstream_error,
}),
_ => None,
}
}
fn exhausted_metadata(
kind: ExhaustedAccountKind,
count: usize,
upstream_error: &str,
) -> DispatchErrorMetadata {
DispatchErrorMetadata::exhausted_ref(ExhaustedAccountRef {
kind,
count,
upstream_error,
})
}
}
/// Responses live SSE body stream error.
#[derive(Debug, Error)]
pub enum ResponseDispatchStreamError {
#[error("upstream stream failed: {0}")]
Upstream(#[from] CodexClientError),
}
pub(super) fn dispatch_error_metadata(
error: impl std::fmt::Display,
stream: bool,
compact: bool,
transport: Option<&str>,
) -> Value {
let mut metadata = serde_json::json!({
"stream": stream,
"failed": true,
"errorKind": "dispatch",
"error": error.to_string(),
});
let Some(object) = metadata.as_object_mut() else {
return metadata;
};
if compact {
object.insert("compact".to_string(), Value::Bool(true));
}
if let Some(transport) = transport {
object.insert(
"transport".to_string(),
Value::String(transport.to_string()),
);
}
metadata
}
pub(super) fn enrich_response_dispatch_error_metadata(
metadata: &mut Value,
error: &ResponseDispatchError,
) {
let Some(object) = metadata.as_object_mut() else {
return;
};
insert_dispatch_error_metadata(object, error.metadata());
}
+5 -6
View File
@@ -1,10 +1,9 @@
//! 请求调度模块(账号选择、回退、恢复、会话亲和性)。
pub mod affinity;
pub mod cloudflare;
pub mod errors;
mod exhaustion;
pub mod implicit_resume;
pub mod reasoning_replay;
pub mod responses;
pub mod upstream;
mod recording;
pub mod recovery;
pub(crate) mod service;
pub(crate) mod stream;
pub mod upstream_call;
@@ -30,9 +30,11 @@ use super::{
errors::{
dispatch_error_metadata, enrich_response_dispatch_error_metadata, ResponseDispatchError,
},
sse_failure::{status_code_for_stream_failure, stream_failure_metadata},
stream_lifecycle::{latest_response_id, LiveResponseStreamContext},
trace::{ResponseDispatchAttempt, ResponseDispatchTrace},
stream::{
lifecycle::{latest_response_id, LiveResponseStreamContext},
sse_failure::{status_code_for_stream_failure, stream_failure_metadata},
trace::{ResponseDispatchAttempt, ResponseDispatchTrace},
},
};
pub(super) struct ResponseUpstreamErrorEventRecord<'a> {
+22
View File
@@ -0,0 +1,22 @@
//! 上游认证失败分类。
use crate::{
accounts::account::AccountStatus,
upstream::openai::transport::{is_banned_auth_signal, CodexClientError},
};
pub(crate) fn is_auth_upstream_error(error: &CodexClientError) -> bool {
matches!(
error,
CodexClientError::Upstream { status, .. } if status.as_u16() == 401
)
}
pub(crate) fn auth_failure_account_status(error: &CodexClientError) -> AccountStatus {
match error {
CodexClientError::Upstream { body, .. } if is_banned_auth_signal(body) => {
AccountStatus::Banned
}
_ => AccountStatus::Expired,
}
}
+7
View File
@@ -0,0 +1,7 @@
//! 调度失败恢复策略。
pub(crate) mod auth;
pub mod cloudflare;
pub(crate) mod exhaustion;
pub mod implicit_resume;
pub mod reasoning_replay;
-342
View File
@@ -1,342 +0,0 @@
use serde_json::Value;
use thiserror::Error;
use crate::{
dispatch::{
errors::{
insert_dispatch_error_metadata, upstream_error_http_status, DispatchErrorMetadata,
DispatchFailureClass,
},
exhaustion::{ExhaustedAccount, ExhaustedAccountKind, ExhaustedAccountRef},
},
upstream::openai::{
protocol::{responses::ResponsesSseFailure, sse::SseError},
transport::CodexClientError,
},
};
use super::sse_failure::{sse_failure_error_body, stream_failure_http_status};
/// Responses 调度错误。
#[derive(Debug, Error)]
pub enum ResponseDispatchError {
#[error("no active account is available")]
NoActiveAccount,
#[error("all accounts exhausted by quota")]
QuotaExhausted {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by rate limit")]
RateLimited {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by expired auth")]
Expired {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by disabled auth")]
Disabled {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by banned auth")]
Banned {
count: usize,
upstream_error: String,
status_code: u16,
},
#[error("all accounts exhausted by Cloudflare challenge")]
CloudflareChallenge {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by Cloudflare path-block")]
CloudflarePathBlocked {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by unsupported model")]
ModelUnsupported {
count: usize,
upstream_error: String,
},
#[error("previous response context is unavailable: {upstream_error}")]
HistoryUnavailable { upstream_error: String },
#[error("upstream request failed: {0}")]
Upstream(#[from] CodexClientError),
#[error("invalid upstream SSE response: {0}")]
InvalidSse(#[from] SseError),
#[error("upstream response did not include response.completed")]
MissingCompleted,
#[error("upstream response did not include visible output")]
EmptyUpstreamResponse,
#[error("upstream response failed: {0:?}")]
Failed(ResponsesSseFailure),
}
impl ResponseDispatchError {
pub(crate) fn from_exhausted_account(exhausted: ExhaustedAccount) -> Self {
match exhausted.kind {
ExhaustedAccountKind::QuotaExhausted => Self::QuotaExhausted {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::RateLimited => Self::RateLimited {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::Expired => Self::Expired {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::Disabled => Self::Disabled {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::Banned => Self::Banned {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
status_code: exhausted.status_code.unwrap_or(403),
},
ExhaustedAccountKind::CloudflareChallenge => Self::CloudflareChallenge {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::CloudflarePathBlocked => Self::CloudflarePathBlocked {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::ModelUnsupported => Self::ModelUnsupported {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
}
}
pub fn http_status_code(&self) -> u16 {
match self {
Self::NoActiveAccount => 503,
Self::QuotaExhausted { .. } => 429,
Self::RateLimited { .. } => 429,
Self::Expired { .. } | Self::Disabled { .. } => 401,
Self::Banned { status_code, .. } => *status_code,
Self::CloudflareChallenge { .. }
| Self::CloudflarePathBlocked { .. }
| Self::InvalidSse(_)
| Self::MissingCompleted
| Self::EmptyUpstreamResponse => 502,
Self::Failed(failure) => stream_failure_http_status(failure),
Self::ModelUnsupported { .. } | Self::HistoryUnavailable { .. } => 400,
Self::Upstream(error) => upstream_error_http_status(error),
}
}
pub(crate) fn metadata(&self) -> DispatchErrorMetadata {
match self {
Self::NoActiveAccount => DispatchErrorMetadata::no_available_accounts(),
Self::QuotaExhausted {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::QuotaExhausted,
*count,
upstream_error,
),
Self::RateLimited {
count,
upstream_error,
} => {
Self::exhausted_metadata(ExhaustedAccountKind::RateLimited, *count, upstream_error)
}
Self::Expired {
count,
upstream_error,
} => Self::exhausted_metadata(ExhaustedAccountKind::Expired, *count, upstream_error),
Self::Disabled {
count,
upstream_error,
} => Self::exhausted_metadata(ExhaustedAccountKind::Disabled, *count, upstream_error),
Self::Banned {
count,
upstream_error,
..
} => Self::exhausted_metadata(ExhaustedAccountKind::Banned, *count, upstream_error),
Self::CloudflareChallenge {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::CloudflareChallenge,
*count,
upstream_error,
),
Self::CloudflarePathBlocked {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::CloudflarePathBlocked,
*count,
upstream_error,
),
Self::ModelUnsupported {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::ModelUnsupported,
*count,
upstream_error,
),
Self::HistoryUnavailable { upstream_error } => DispatchErrorMetadata {
failure_class: DispatchFailureClass::HistoryUnavailable,
exhausted_count: None,
upstream_error: Some(upstream_error.clone()),
upstream_status: Some(400),
diagnostics: None,
},
Self::Upstream(error) => DispatchErrorMetadata::upstream(error),
Self::InvalidSse(_) => DispatchErrorMetadata::simple(DispatchFailureClass::InvalidSse),
Self::MissingCompleted => {
DispatchErrorMetadata::simple(DispatchFailureClass::MissingCompleted)
}
Self::EmptyUpstreamResponse => {
DispatchErrorMetadata::simple(DispatchFailureClass::EmptyUpstreamResponse)
}
Self::Failed(failure) => DispatchErrorMetadata {
failure_class: DispatchFailureClass::ResponseFailed,
exhausted_count: None,
upstream_error: Some(sse_failure_error_body(failure)),
upstream_status: None,
diagnostics: None,
},
}
}
pub(crate) fn exhausted_account(&self) -> Option<ExhaustedAccountRef<'_>> {
match self {
Self::QuotaExhausted {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::QuotaExhausted,
count: *count,
upstream_error,
}),
Self::RateLimited {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::RateLimited,
count: *count,
upstream_error,
}),
Self::Expired {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Expired,
count: *count,
upstream_error,
}),
Self::Disabled {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Disabled,
count: *count,
upstream_error,
}),
Self::Banned {
count,
upstream_error,
..
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Banned,
count: *count,
upstream_error,
}),
Self::CloudflareChallenge {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::CloudflareChallenge,
count: *count,
upstream_error,
}),
Self::CloudflarePathBlocked {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::CloudflarePathBlocked,
count: *count,
upstream_error,
}),
Self::ModelUnsupported {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::ModelUnsupported,
count: *count,
upstream_error,
}),
_ => None,
}
}
fn exhausted_metadata(
kind: ExhaustedAccountKind,
count: usize,
upstream_error: &str,
) -> DispatchErrorMetadata {
DispatchErrorMetadata::exhausted_ref(ExhaustedAccountRef {
kind,
count,
upstream_error,
})
}
}
/// Responses live SSE body stream error.
#[derive(Debug, Error)]
pub enum ResponseDispatchStreamError {
#[error("upstream stream failed: {0}")]
Upstream(#[from] CodexClientError),
}
pub(super) fn dispatch_error_metadata(
error: impl std::fmt::Display,
stream: bool,
compact: bool,
transport: Option<&str>,
) -> Value {
let mut metadata = serde_json::json!({
"stream": stream,
"failed": true,
"errorKind": "dispatch",
"error": error.to_string(),
});
let Some(object) = metadata.as_object_mut() else {
return metadata;
};
if compact {
object.insert("compact".to_string(), Value::Bool(true));
}
if let Some(transport) = transport {
object.insert(
"transport".to_string(),
Value::String(transport.to_string()),
);
}
metadata
}
pub(super) fn enrich_response_dispatch_error_metadata(
metadata: &mut Value,
error: &ResponseDispatchError,
) {
let Some(object) = metadata.as_object_mut() else {
return;
};
insert_dispatch_error_metadata(object, error.metadata());
}
-10
View File
@@ -1,10 +0,0 @@
//! Responses 调度领域。
pub(crate) mod errors;
mod event_recording;
mod live_stream;
mod prefetch;
pub(crate) mod service;
pub(crate) mod sse_failure;
mod stream_lifecycle;
mod trace;
@@ -26,21 +26,23 @@ use crate::{
compute_variant_hash, ensure_prompt_cache_key, hash_instructions,
prepare_variant_identity, RuntimeSessionAffinityService,
},
cloudflare::{
cloudflare_challenge_error_message, cloudflare_path_block_error_message,
is_cloudflare_challenge_upstream_error, is_cloudflare_path_block_upstream_error,
CloudflareRecovery,
},
errors::{
auth_failure_account_status, backend_transport_name, is_auth_upstream_error,
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,
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,
},
exhaustion::AccountExhaustionTracker,
reasoning_replay::ReasoningReplayCache,
upstream::{
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,
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,
@@ -71,13 +73,13 @@ use crate::{
},
};
use crate::dispatch::implicit_resume::{
use crate::dispatch::recovery::implicit_resume::{
continuation_input_start, implicit_resume_allowed, ImplicitResumeSnapshot,
};
use super::{
errors::{ResponseDispatchError, ResponseDispatchStreamError},
event_recording::{
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,
@@ -85,16 +87,19 @@ use super::{
ResponseDispatchErrorEventRecord, ResponseStreamFailureEventRecord,
ResponseUpstreamErrorEventRecord,
},
live_stream::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,
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,
},
trace::{ResponseDispatchAttempt, ResponseDispatchTrace},
},
stream_lifecycle::LiveResponseStreamContext,
trace::{ResponseDispatchAttempt, ResponseDispatchTrace},
};
/// OpenAI Responses 调度服务。
@@ -11,7 +11,11 @@ use crate::{
RuntimeSessionAffinityService,
},
errors::backend_transport_name,
reasoning_replay::ReasoningReplayCache,
recording::{
insert_first_token_ms, live_response_rate_limit_headers, live_response_turn_state,
record_live_response_stream_event,
},
recovery::reasoning_replay::ReasoningReplayCache,
},
infra::time::elapsed_millis_i64,
telemetry::{recorder::Recorder, usage::types::UsageRecordLevel},
@@ -29,10 +33,6 @@ use crate::{
};
use super::{
event_recording::{
insert_first_token_ms, live_response_rate_limit_headers, live_response_turn_state,
record_live_response_stream_event,
},
sse_failure::{
sse_failure_invalid_reasoning_replay, status_code_for_stream_failure,
stream_failure_metadata, stream_failure_source, synthetic_stream_disconnected_detail,
@@ -40,37 +40,37 @@ use super::{
},
trace::ResponseDispatchAttempt,
};
use crate::dispatch::cloudflare::CloudflareRecovery;
use crate::dispatch::recovery::cloudflare::CloudflareRecovery;
pub(super) struct LiveResponseStreamContext {
pub(super) account_pool: Arc<RuntimeAccountPoolService>,
pub(super) session_affinity: Arc<RuntimeSessionAffinityService>,
pub(super) reasoning_replay: Arc<Mutex<ReasoningReplayCache>>,
pub(super) recorder: Arc<Recorder>,
pub(super) cloudflare: CloudflareRecovery,
pub(super) account_id: String,
pub(super) account_plan_type: Option<String>,
pub(super) request_id: String,
pub(super) route: String,
pub(super) model: String,
pub(super) display_model: String,
pub(super) requested_model: String,
pub(super) client_ip: Option<String>,
pub(super) request: CodexResponsesRequest,
pub(super) tuple_schema: Option<Value>,
pub(super) transport: CodexBackendTransport,
pub(super) rate_limit_headers: Vec<(String, String)>,
pub(super) rate_limit_header_updates: Option<CodexRateLimitHeaderUpdates>,
pub(super) turn_state_update: Option<CodexTurnStateUpdate>,
pub(super) websocket_pool_decision: Option<WebSocketPoolDecision>,
pub(super) turn_state: Option<String>,
pub(super) diagnostics: CodexUpstreamDiagnostics,
pub(super) attempt: ResponseDispatchAttempt,
pub(super) attempts: Vec<ResponseDispatchAttempt>,
pub(super) started_at: Instant,
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(super) fn latest_response_id(body: &str) -> Option<String> {
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)
@@ -85,7 +85,7 @@ pub(super) fn latest_response_id(body: &str) -> Option<String> {
})
}
pub(super) fn premature_close_failed_event(
pub(in crate::dispatch) fn premature_close_failed_event(
response_id: Option<&str>,
detail: Option<&str>,
) -> String {
@@ -101,7 +101,7 @@ pub(super) fn premature_close_failed_event(
)
}
pub(super) async fn finalize_live_response_stream(
pub(in crate::dispatch) async fn finalize_live_response_stream(
context: LiveResponseStreamContext,
body: String,
first_token_ms: Option<i64>,
@@ -9,10 +9,10 @@ use serde_json::Value;
use tokio::sync::{mpsc, oneshot};
use crate::{
dispatch::responses::{
dispatch::{
errors::ResponseDispatchStreamError,
service::ResponseDispatchStream,
stream_lifecycle::{
stream::lifecycle::{
finalize_live_response_stream, latest_response_id, premature_close_failed_event,
LiveResponseStreamContext,
},
@@ -53,7 +53,7 @@ impl Stream for MpscResponseBodyStream {
}
}
pub(super) fn spawn_live_response_stream(
pub(in crate::dispatch) fn spawn_live_response_stream(
context: LiveResponseStreamContext,
prefetched: Bytes,
mut body: CodexBackendSseStream,
+7
View File
@@ -0,0 +1,7 @@
//! Responses 流式转发与生命周期处理。
pub(crate) mod lifecycle;
pub(crate) mod live;
pub(crate) mod prefetch;
pub(crate) mod sse_failure;
pub(crate) mod trace;
@@ -6,11 +6,11 @@ use crate::upstream::openai::{
transport::CodexBackendSseStream,
};
use super::errors::ResponseDispatchError;
use crate::dispatch::errors::ResponseDispatchError;
const MAX_STREAM_PREFETCH_BYTES: usize = 64 * 1024;
pub(super) async fn prefetch_first_sse_chunk(
pub(in crate::dispatch) async fn prefetch_first_sse_chunk(
mut body: CodexBackendSseStream,
) -> Result<(Bytes, CodexBackendSseStream), ResponseDispatchError> {
let mut prefetched = Vec::new();
@@ -16,11 +16,11 @@ use crate::{
},
};
pub(super) const STREAM_DISCONNECTED_CODE: &str = "stream_disconnected";
pub(super) const STREAM_DISCONNECTED_MESSAGE: &str =
pub(in crate::dispatch) const STREAM_DISCONNECTED_CODE: &str = "stream_disconnected";
pub(in crate::dispatch) const STREAM_DISCONNECTED_MESSAGE: &str =
"Upstream stream closed before response.completed";
pub(super) fn sse_failure_error_body(failure: &ResponsesSseFailure) -> String {
pub(in crate::dispatch) fn sse_failure_error_body(failure: &ResponsesSseFailure) -> String {
match failure.upstream_code.as_deref() {
Some(code) => serde_json::json!({
"error": {
@@ -88,19 +88,21 @@ pub(crate) fn is_auth_sse_failure(failure: &ResponsesSseFailure) -> bool {
)
}
pub(super) fn is_model_unsupported_sse_failure(failure: &ResponsesSseFailure) -> bool {
pub(in crate::dispatch) fn is_model_unsupported_sse_failure(failure: &ResponsesSseFailure) -> bool {
sse_failure_matches(failure, is_model_unsupported_signal)
}
pub(super) fn is_history_recovery_sse_failure(failure: &ResponsesSseFailure) -> bool {
pub(in crate::dispatch) fn is_history_recovery_sse_failure(failure: &ResponsesSseFailure) -> bool {
sse_failure_matches(failure, is_history_recovery_signal)
}
pub(super) fn sse_failure_invalid_reasoning_replay(failure: &ResponsesSseFailure) -> bool {
pub(in crate::dispatch) fn sse_failure_invalid_reasoning_replay(
failure: &ResponsesSseFailure,
) -> bool {
sse_failure_matches(failure, is_invalid_encrypted_content_signal)
}
pub(super) fn client_error_invalid_reasoning_replay(error: &CodexClientError) -> bool {
pub(in crate::dispatch) fn client_error_invalid_reasoning_replay(error: &CodexClientError) -> bool {
matches!(
error,
CodexClientError::Upstream { body, .. } if is_invalid_encrypted_content_signal(body)
@@ -115,17 +117,7 @@ pub(crate) fn auth_sse_failure_account_status(failure: &ResponsesSseFailure) ->
}
}
pub(crate) fn sse_failure_account_status(failure: &ResponsesSseFailure) -> Option<AccountStatus> {
if is_quota_exhausted_sse_failure(failure) {
Some(AccountStatus::QuotaExhausted)
} else if is_auth_sse_failure(failure) {
Some(auth_sse_failure_account_status(failure))
} else {
None
}
}
pub(super) fn first_sse_failure(
pub(in crate::dispatch) fn first_sse_failure(
prefetched: &[u8],
) -> Result<Option<ResponsesSseFailure>, SseError> {
let body = String::from_utf8_lossy(prefetched);
@@ -137,7 +129,7 @@ pub(super) fn first_sse_failure(
}
}
pub(super) fn stream_failure_metadata(
pub(in crate::dispatch) fn stream_failure_metadata(
failure: &ResponsesSseFailure,
usage: Option<TokenUsage>,
) -> Value {
@@ -169,7 +161,7 @@ fn enrich_stream_failure_source_metadata(metadata: &mut Value, failure: &Respons
}
}
pub(super) fn stream_failure_source(failure: &ResponsesSseFailure) -> &'static str {
pub(in crate::dispatch) fn stream_failure_source(failure: &ResponsesSseFailure) -> &'static str {
if synthetic_stream_disconnected_detail(failure).is_some() {
"proxy"
} else {
@@ -177,7 +169,7 @@ pub(super) fn stream_failure_source(failure: &ResponsesSseFailure) -> &'static s
}
}
pub(super) fn synthetic_stream_disconnected_detail(
pub(in crate::dispatch) fn synthetic_stream_disconnected_detail(
failure: &ResponsesSseFailure,
) -> Option<String> {
if failure.upstream_code.as_deref() != Some(STREAM_DISCONNECTED_CODE) {
@@ -193,7 +185,7 @@ pub(super) fn synthetic_stream_disconnected_detail(
Some(detail)
}
pub(super) fn status_code_for_stream_failure(failure: &ResponsesSseFailure) -> i64 {
pub(in crate::dispatch) fn status_code_for_stream_failure(failure: &ResponsesSseFailure) -> i64 {
let code = failure
.upstream_code
.as_deref()
@@ -236,6 +228,6 @@ pub(super) fn status_code_for_stream_failure(failure: &ResponsesSseFailure) -> i
502
}
pub(super) fn stream_failure_http_status(failure: &ResponsesSseFailure) -> u16 {
pub(in crate::dispatch) fn stream_failure_http_status(failure: &ResponsesSseFailure) -> u16 {
u16::try_from(status_code_for_stream_failure(failure)).unwrap_or(502)
}
@@ -1,19 +1,22 @@
//! Request-level trace state for Responses dispatch.
#[derive(Debug, Default)]
pub(super) struct ResponseDispatchTrace {
pub(in crate::dispatch) struct ResponseDispatchTrace {
next_attempt_index: i64,
attempts: Vec<ResponseDispatchAttempt>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) struct ResponseDispatchAttempt {
pub(in crate::dispatch) struct ResponseDispatchAttempt {
index: i64,
account_id: String,
}
impl ResponseDispatchTrace {
pub(super) fn start_attempt(&mut self, account_id: &str) -> ResponseDispatchAttempt {
pub(in crate::dispatch) fn start_attempt(
&mut self,
account_id: &str,
) -> ResponseDispatchAttempt {
let attempt = ResponseDispatchAttempt {
index: self.next_attempt_index,
account_id: account_id.to_string(),
@@ -23,17 +26,17 @@ impl ResponseDispatchTrace {
attempt
}
pub(super) fn attempts(&self) -> &[ResponseDispatchAttempt] {
pub(in crate::dispatch) fn attempts(&self) -> &[ResponseDispatchAttempt] {
&self.attempts
}
}
impl ResponseDispatchAttempt {
pub(super) fn index(&self) -> i64 {
pub(in crate::dispatch) fn index(&self) -> i64 {
self.index
}
pub(super) fn account_id(&self) -> &str {
pub(in crate::dispatch) fn account_id(&self) -> &str {
&self.account_id
}
}
@@ -14,8 +14,8 @@ use crate::{
},
},
dispatch::{
affinity::build_conversation_identity, cloudflare::CloudflareRecovery,
errors::is_retryable_upstream_5xx_error,
affinity::build_conversation_identity, errors::is_retryable_upstream_5xx_error,
recovery::cloudflare::CloudflareRecovery,
},
upstream::openai::{
protocol::responses::{CodexCompactRequest, CodexResponsesRequest},
+1 -1
View File
@@ -6,7 +6,7 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
match arguments.next().as_deref().and_then(|value| value.to_str()) {
None | Some("serve") => {
reject_extra_arguments(arguments)?;
codex_proxy_rs::bootstrap::application::run().await
codex_proxy_rs::bootstrap::services::run().await
}
Some("import-sqlite") => {
let source = arguments
+17 -6
View File
@@ -2,12 +2,9 @@
use std::collections::{BTreeMap, BTreeSet};
use crate::upstream::openai::protocol::model_catalog::{BackendModelEntry, BackendReasoningEffort};
use super::{
config::ModelConfig,
info::{CodexModelInfo, ReasoningEffortInfo},
snapshot::ModelPlanSnapshot,
use super::types::{
BackendModelEntry, BackendReasoningEffort, CodexModelInfo, ModelConfig, ModelPlanSnapshot,
ReasoningEffortInfo,
};
const BUILTIN_CODEX_MODELS: [(&str, &str, bool); 8] = [
@@ -41,6 +38,20 @@ impl ModelPlanSnapshot {
models: entries.into_iter().map(normalize_backend_model).collect(),
}
}
/// 将上游 JSON 条目解析并标准化为可持久化快照。
pub fn from_backend_values(
plan_type: impl Into<String>,
entries: Vec<serde_json::Value>,
) -> Self {
Self::from_backend_entries(
plan_type,
entries
.into_iter()
.filter_map(|entry| serde_json::from_value(entry).ok())
.collect(),
)
}
}
impl ModelCatalog {
-12
View File
@@ -1,12 +0,0 @@
//! 模型目录配置。
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
/// 模型目录别名配置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
pub struct ModelConfig {
/// 模型别名映射。
pub model_aliases: BTreeMap<String, String>,
}
-53
View File
@@ -1,53 +0,0 @@
//! 标准化模型展示信息。
use serde::{Deserialize, Serialize};
/// 单个 reasoning effort 的标准化展示信息。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct ReasoningEffortInfo {
/// 推理强度名。
pub reasoning_effort: String,
/// 展示描述。
pub description: String,
}
/// 对外暴露的单个模型目录条目。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct CodexModelInfo {
/// 模型唯一 ID。
pub id: String,
/// 展示名。
pub display_name: String,
/// 描述文本。
pub description: String,
/// 是否默认模型。
pub is_default: bool,
/// 支持的推理强度列表。
pub supported_reasoning_efforts: Vec<ReasoningEffortInfo>,
/// 默认推理强度。
pub default_reasoning_effort: String,
/// 输入模态。
pub input_modalities: Vec<String>,
/// 输出模态。
pub output_modalities: Vec<String>,
/// 是否支持 personality。
pub supports_personality: bool,
/// 升级提示。
pub upgrade: Option<String>,
/// 来源标记。
pub source: String,
/// 当前上下文窗口。
#[serde(skip_serializing_if = "Option::is_none")]
pub context_window: Option<u64>,
/// 最大上下文窗口。
#[serde(skip_serializing_if = "Option::is_none")]
pub max_context_window: Option<u64>,
/// 最大输出 token 数。
#[serde(skip_serializing_if = "Option::is_none")]
pub max_output_tokens: Option<u64>,
/// 截断策略限制。
#[serde(skip_serializing_if = "Option::is_none")]
pub truncation_policy_limit: Option<u64>,
}
+1 -3
View File
@@ -1,8 +1,6 @@
//! 模型目录领域。
pub mod catalog;
pub mod config;
pub mod info;
pub mod service;
pub mod snapshot;
pub mod store;
pub mod types;
+27 -8
View File
@@ -6,17 +6,17 @@ use std::{
};
use thiserror::Error;
use tokio::sync::watch;
use crate::{
accounts::account::Account,
settings::SettingsSnapshot,
upstream::openai::transport::{CodexModelCatalogClient, CodexModelCatalogRequest},
};
use super::{
catalog::ModelCatalog,
config::ModelConfig,
snapshot::ModelPlanSnapshot,
store::{ModelSnapshotStore, ModelSnapshotStoreError},
types::{ModelConfig, ModelPlanSnapshot},
};
/// 模型刷新摘要。
@@ -76,8 +76,10 @@ pub struct ModelPlanRouting {
pub struct ModelRefreshPlanAccount {
/// 订阅计划类型。
pub plan_type: String,
/// 用于刷新该计划模型列表的账号。
pub account: Account,
/// 用于访问模型目录的令牌。
pub access_token: String,
/// 上游账号标识。
pub account_id: Option<String>,
}
/// 模型目录服务。
@@ -116,6 +118,19 @@ impl ModelService {
self.rebuild_catalog_from_memory();
}
/// 持续接收运行时设置并更新模型别名。
pub async fn subscribe_settings(
self: Arc<Self>,
mut receiver: watch::Receiver<SettingsSnapshot>,
) {
while receiver.changed().await.is_ok() {
let settings = receiver.borrow_and_update().clone();
self.update_config(ModelConfig {
model_aliases: settings.model_aliases,
});
}
}
/// 从快照存储加载运行时内存模型目录。
pub async fn reload_from_store(&self) -> Result<(), ModelServiceError> {
let Some(store) = self.store.as_ref() else {
@@ -163,8 +178,8 @@ impl ModelService {
for plan_account in plan_accounts {
let request = CodexModelCatalogRequest {
access_token: &plan_account.account.access_token,
account_id: plan_account.account.account_id.as_deref(),
access_token: &plan_account.access_token,
account_id: plan_account.account_id.as_deref(),
request_id,
installation_id,
plan_type: &plan_account.plan_type,
@@ -184,7 +199,11 @@ impl ModelService {
};
let snapshot =
ModelPlanSnapshot::from_backend_entries(plan_account.plan_type.clone(), entries);
ModelPlanSnapshot::from_backend_values(plan_account.plan_type.clone(), entries);
if snapshot.models.is_empty() {
result.failed_plans += 1;
continue;
}
result.model_count += snapshot.models.len();
store
.replace_plan_snapshot(&snapshot)
-15
View File
@@ -1,15 +0,0 @@
//! 模型计划快照。
use serde::{Deserialize, Serialize};
use super::info::CodexModelInfo;
/// 按计划类型持久化的模型快照。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct ModelPlanSnapshot {
/// 订阅计划类型。
pub plan_type: String,
/// 该计划可见的模型列表。
pub models: Vec<CodexModelInfo>,
}
+1 -1
View File
@@ -8,7 +8,7 @@ use thiserror::Error;
use crate::infra::redis::RedisConnection;
use super::{info::CodexModelInfo, snapshot::ModelPlanSnapshot};
use super::types::{CodexModelInfo, ModelPlanSnapshot};
/// 模型快照存储错误。
#[derive(Debug, Error, PartialEq, Eq)]
+119
View File
@@ -0,0 +1,119 @@
//! 模型目录领域类型。
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
/// 模型目录别名配置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
pub struct ModelConfig {
/// 模型别名映射。
pub model_aliases: BTreeMap<String, String>,
}
/// 单个 reasoning effort 的标准化展示信息。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct ReasoningEffortInfo {
/// 推理强度名。
pub reasoning_effort: String,
/// 展示描述。
pub description: String,
}
/// 对外暴露的单个模型目录条目。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct CodexModelInfo {
/// 模型唯一 ID。
pub id: String,
/// 展示名。
pub display_name: String,
/// 描述文本。
pub description: String,
/// 是否默认模型。
pub is_default: bool,
/// 支持的推理强度列表。
pub supported_reasoning_efforts: Vec<ReasoningEffortInfo>,
/// 默认推理强度。
pub default_reasoning_effort: String,
/// 输入模态。
pub input_modalities: Vec<String>,
/// 输出模态。
pub output_modalities: Vec<String>,
/// 是否支持 personality。
pub supports_personality: bool,
/// 升级提示。
pub upgrade: Option<String>,
/// 来源标记。
pub source: String,
/// 当前上下文窗口。
#[serde(skip_serializing_if = "Option::is_none")]
pub context_window: Option<u64>,
/// 最大上下文窗口。
#[serde(skip_serializing_if = "Option::is_none")]
pub max_context_window: Option<u64>,
/// 最大输出 token 数。
#[serde(skip_serializing_if = "Option::is_none")]
pub max_output_tokens: Option<u64>,
/// 截断策略限制。
#[serde(skip_serializing_if = "Option::is_none")]
pub truncation_policy_limit: Option<u64>,
}
/// 按计划类型持久化的模型快照。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct ModelPlanSnapshot {
/// 订阅计划类型。
pub plan_type: String,
/// 该计划可见的模型列表。
pub models: Vec<CodexModelInfo>,
}
/// 上游声明的 reasoning effort 条目。
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct BackendReasoningEffort {
#[serde(default, rename = "reasoning_effort", alias = "reasoningEffort")]
pub reasoning_effort: Option<String>,
#[serde(default)]
pub effort: Option<String>,
#[serde(default)]
pub description: Option<String>,
}
/// 上游截断策略定义。
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct BackendTruncationPolicy {
pub limit: Option<u64>,
}
/// 上游模型目录的原始模型条目。
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct BackendModelEntry {
pub slug: Option<String>,
pub id: Option<String>,
pub name: Option<String>,
pub display_name: Option<String>,
pub title: Option<String>,
pub description: Option<String>,
pub is_default: Option<bool>,
pub default_reasoning_effort: Option<String>,
pub default_reasoning_level: Option<String>,
#[serde(default)]
pub supported_reasoning_efforts: Vec<BackendReasoningEffort>,
#[serde(default)]
pub supported_reasoning_levels: Vec<BackendReasoningEffort>,
pub input_modalities: Option<Vec<String>>,
pub output_modalities: Option<Vec<String>>,
pub supports_personality: Option<bool>,
pub upgrade: Option<String>,
#[serde(alias = "contextWindow")]
pub context_window: Option<u64>,
#[serde(alias = "maxContextWindow")]
pub max_context_window: Option<u64>,
#[serde(alias = "maxOutputTokens")]
pub max_output_tokens: Option<u64>,
#[serde(alias = "truncationPolicy")]
pub truncation_policy: Option<BackendTruncationPolicy>,
}
+18 -14
View File
@@ -3,7 +3,7 @@
use std::{collections::BTreeMap, sync::Arc};
use subtle::ConstantTimeEq;
use tokio::sync::RwLock;
use tokio::sync::{watch, Mutex};
use crate::infra::identity::{generate_admin_api_key, hash_credential};
@@ -52,14 +52,17 @@ impl SettingsService {
#[derive(Clone)]
pub struct RuntimeSettingsService {
store: PgSettingsStore,
snapshot: Arc<RwLock<Arc<SettingsSnapshot>>>,
sender: watch::Sender<SettingsSnapshot>,
update_lock: Arc<Mutex<()>>,
}
impl RuntimeSettingsService {
pub fn new(settings: SettingsSnapshot, pool: sqlx::PgPool) -> Self {
let (sender, _) = watch::channel(settings);
Self {
store: PgSettingsStore::new(pool),
snapshot: Arc::new(RwLock::new(Arc::new(settings))),
sender,
update_lock: Arc::new(Mutex::new(())),
}
}
@@ -72,20 +75,21 @@ impl RuntimeSettingsService {
.await
}
pub async fn current(&self) -> Arc<SettingsSnapshot> {
self.snapshot.read().await.clone()
pub fn current(&self) -> SettingsSnapshot {
self.sender.borrow().clone()
}
pub async fn update(
&self,
patch: SettingsPatch,
) -> Result<Arc<SettingsSnapshot>, SettingsError> {
let mut current = self.snapshot.write().await;
let mut next = (**current).clone();
/// 订阅设置快照变更。
pub fn subscribe(&self) -> watch::Receiver<SettingsSnapshot> {
self.sender.subscribe()
}
pub async fn update(&self, patch: SettingsPatch) -> Result<SettingsSnapshot, SettingsError> {
let _guard = self.update_lock.lock().await;
let mut next = self.current();
SettingsService::apply_patch(&mut next, patch)?;
self.store.save(&next).await?;
let next = Arc::new(next);
*current = next.clone();
self.sender.send_replace(next.clone());
Ok(next)
}
@@ -94,7 +98,7 @@ impl RuntimeSettingsService {
}
pub async fn regenerate_admin_api_key(&self) -> Result<String, SettingsError> {
let current = self.current().await;
let current = self.current();
self.store.ensure(&current).await?;
let key = generate_admin_api_key();
self.store
+502 -1
View File
@@ -1,14 +1,114 @@
//! PostgreSQL 账号用量查询存储。
use std::collections::HashMap;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
use thiserror::Error;
use crate::infra::{
format::nonnegative_i64_to_u64,
json::{decode_cursor, encode_cursor, Page},
time::parse_rfc3339_utc as parse_rfc3339,
};
const RECORD_USAGE_SQL: &str = r"
insert into account_usage (
account_id, request_count, empty_response_count,
input_tokens, output_tokens, cached_tokens, reasoning_tokens, total_tokens,
image_input_tokens, image_output_tokens, image_request_count, image_request_failed_count,
window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens,
window_image_input_tokens, window_image_output_tokens,
window_image_request_count, window_image_request_failed_count,
window_started_at, last_used_at
) values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22)
on conflict(account_id) do update set
request_count = account_usage.request_count + excluded.request_count,
empty_response_count = account_usage.empty_response_count + excluded.empty_response_count,
input_tokens = account_usage.input_tokens + excluded.input_tokens,
output_tokens = account_usage.output_tokens + excluded.output_tokens,
cached_tokens = account_usage.cached_tokens + excluded.cached_tokens,
reasoning_tokens = account_usage.reasoning_tokens + excluded.reasoning_tokens,
total_tokens = account_usage.total_tokens + excluded.total_tokens,
image_input_tokens = account_usage.image_input_tokens + excluded.image_input_tokens,
image_output_tokens = account_usage.image_output_tokens + excluded.image_output_tokens,
image_request_count = account_usage.image_request_count + excluded.image_request_count,
image_request_failed_count = account_usage.image_request_failed_count + excluded.image_request_failed_count,
window_request_count = account_usage.window_request_count + excluded.window_request_count,
window_input_tokens = account_usage.window_input_tokens + excluded.window_input_tokens,
window_output_tokens = account_usage.window_output_tokens + excluded.window_output_tokens,
window_cached_tokens = account_usage.window_cached_tokens + excluded.window_cached_tokens,
window_image_input_tokens = account_usage.window_image_input_tokens + excluded.window_image_input_tokens,
window_image_output_tokens = account_usage.window_image_output_tokens + excluded.window_image_output_tokens,
window_image_request_count = account_usage.window_image_request_count + excluded.window_image_request_count,
window_image_request_failed_count = account_usage.window_image_request_failed_count + excluded.window_image_request_failed_count,
window_started_at = coalesce(account_usage.window_started_at, excluded.window_started_at),
last_used_at = excluded.last_used_at";
const RECORD_MODEL_USAGE_SQL: &str = r"
insert into account_model_usage (
account_id, model, request_count, error_count,
input_tokens, output_tokens, cached_tokens, last_used_at
) values ($1, $2, $3, $4, $5, $6, $7, $8)
on conflict(account_id, model) do update set
request_count = account_model_usage.request_count + excluded.request_count,
error_count = account_model_usage.error_count + excluded.error_count,
input_tokens = account_model_usage.input_tokens + excluded.input_tokens,
output_tokens = account_model_usage.output_tokens + excluded.output_tokens,
cached_tokens = account_model_usage.cached_tokens + excluded.cached_tokens,
last_used_at = excluded.last_used_at";
const SYNC_RUNTIME_WINDOW_SQL: &str = r"
insert into account_usage (
account_id, window_request_count, window_input_tokens, window_output_tokens,
window_cached_tokens, window_image_input_tokens, window_image_output_tokens,
window_image_request_count, window_image_request_failed_count,
window_started_at, window_reset_at, limit_window_seconds
) values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
on conflict(account_id) do update set
window_request_count = excluded.window_request_count,
window_input_tokens = excluded.window_input_tokens,
window_output_tokens = excluded.window_output_tokens,
window_cached_tokens = excluded.window_cached_tokens,
window_image_input_tokens = excluded.window_image_input_tokens,
window_image_output_tokens = excluded.window_image_output_tokens,
window_image_request_count = excluded.window_image_request_count,
window_image_request_failed_count = excluded.window_image_request_failed_count,
window_started_at = excluded.window_started_at,
window_reset_at = excluded.window_reset_at,
limit_window_seconds = excluded.limit_window_seconds
where account_usage.window_reset_at is null
or account_usage.window_reset_at <= excluded.window_reset_at
or account_usage.window_reset_at <= $13";
const SYNC_RATE_LIMIT_WINDOW_RESET_SQL: &str = r"
insert into account_usage (
account_id, window_request_count, window_input_tokens, window_output_tokens,
window_cached_tokens, window_image_input_tokens, window_image_output_tokens,
window_image_request_count, window_image_request_failed_count,
window_started_at, window_reset_at, limit_window_seconds
) values ($1, 0, 0, 0, 0, 0, 0, 0, 0, $2, $3, $4)
on conflict(account_id) do update set
window_request_count = 0,
window_input_tokens = 0,
window_output_tokens = 0,
window_cached_tokens = 0,
window_image_input_tokens = 0,
window_image_output_tokens = 0,
window_image_request_count = 0,
window_image_request_failed_count = 0,
window_started_at = excluded.window_started_at,
window_reset_at = excluded.window_reset_at,
limit_window_seconds = coalesce(excluded.limit_window_seconds, account_usage.limit_window_seconds)";
const SYNC_RATE_LIMIT_WINDOW_SQL: &str = r"
insert into account_usage (account_id, window_reset_at, limit_window_seconds)
values ($1, $2, $3)
on conflict(account_id) do update set
window_reset_at = excluded.window_reset_at,
limit_window_seconds = coalesce(excluded.limit_window_seconds, account_usage.limit_window_seconds)";
const LIST_USAGE_SELECT_SQL: &str = r"
select
au.account_id,
@@ -71,6 +171,127 @@ pub enum PgAccountUsageStoreError {
/// PostgreSQL 用量存储结果。
pub type PgAccountUsageStoreResult<T> = Result<T, PgAccountUsageStoreError>;
/// 账号用量端口错误。
#[derive(Debug, Error)]
#[error("account usage store operation failed: {message}")]
pub struct AccountUsageStoreError {
message: String,
}
impl From<PgAccountUsageStoreError> for AccountUsageStoreError {
fn from(error: PgAccountUsageStoreError) -> Self {
Self {
message: error.to_string(),
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountUsageDelta {
pub requests: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub reasoning_tokens: u64,
pub total_tokens: u64,
pub empty_responses: u64,
pub image_input_tokens: u64,
pub image_output_tokens: u64,
pub image_requests: u64,
pub image_request_failures: u64,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountModelUsageDelta {
pub requests: u64,
pub errors: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct AccountUsageSnapshot {
pub request_count: u64,
pub empty_response_count: u64,
pub image_input_tokens: u64,
pub image_output_tokens: u64,
pub image_request_count: u64,
pub image_request_failed_count: u64,
pub window_request_count: u64,
pub window_input_tokens: u64,
pub window_output_tokens: u64,
pub window_cached_tokens: u64,
pub window_image_input_tokens: u64,
pub window_image_output_tokens: u64,
pub window_image_request_count: u64,
pub window_image_request_failed_count: u64,
pub window_started_at: Option<DateTime<Utc>>,
pub window_reset_at: Option<DateTime<Utc>>,
pub limit_window_seconds: Option<u64>,
pub last_used_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountUsageWindow {
pub request_count: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub image_input_tokens: u64,
pub image_output_tokens: u64,
pub image_request_count: u64,
pub image_request_failed_count: u64,
pub started_at: Option<DateTime<Utc>>,
pub reset_at: Option<DateTime<Utc>>,
pub limit_window_seconds: Option<u64>,
}
#[async_trait]
pub trait AccountUsageStore: Send + Sync + 'static {
async fn snapshots(
&self,
account_ids: &[String],
) -> Result<HashMap<String, AccountUsageSnapshot>, AccountUsageStoreError>;
async fn record_usage_delta(
&self,
account_id: &str,
usage: AccountUsageDelta,
) -> Result<(), AccountUsageStoreError>;
async fn record_model_usage_delta(
&self,
account_id: &str,
model: &str,
usage: AccountModelUsageDelta,
) -> Result<(), AccountUsageStoreError>;
async fn sync_runtime_window(
&self,
account_id: &str,
window: AccountUsageWindow,
) -> Result<(), AccountUsageStoreError>;
async fn sync_rate_limit_window(
&self,
account_id: &str,
reset_at: DateTime<Utc>,
limit_window_seconds: Option<u64>,
) -> Result<(), AccountUsageStoreError>;
async fn record_request(&self, account_id: &str) -> Result<(), AccountUsageStoreError> {
self.record_usage_delta(
account_id,
AccountUsageDelta {
requests: 1,
..AccountUsageDelta::default()
},
)
.await
}
}
/// 账号用量列表记录。
#[derive(Debug, Clone)]
pub struct UsageListRecord {
@@ -158,7 +379,6 @@ pub struct PgAccountUsageStore {
}
impl PgAccountUsageStore {
/// 构造存储。
pub fn new(pool: PgPool) -> Self {
Self { pool }
}
@@ -167,6 +387,172 @@ impl PgAccountUsageStore {
&self.pool
}
async fn load_snapshots(
&self,
account_ids: &[String],
) -> PgAccountUsageStoreResult<HashMap<String, AccountUsageSnapshot>> {
if account_ids.is_empty() {
return Ok(HashMap::new());
}
let mut builder = QueryBuilder::<Postgres>::new(
r"select
account_id, request_count, empty_response_count,
image_input_tokens, image_output_tokens,
image_request_count, image_request_failed_count,
window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens,
window_image_input_tokens, window_image_output_tokens,
window_image_request_count, window_image_request_failed_count,
window_started_at, window_reset_at, limit_window_seconds, last_used_at
from account_usage
where account_id in (",
);
let mut separated = builder.separated(", ");
for account_id in account_ids {
separated.push_bind(account_id);
}
separated.push_unseparated(")");
let rows = builder.build().fetch_all(&self.pool).await?;
Ok(rows
.iter()
.map(|row| (row.get("account_id"), usage_snapshot_from_row(row)))
.collect())
}
async fn write_usage_delta(
&self,
account_id: &str,
delta: AccountUsageDelta,
) -> PgAccountUsageStoreResult<()> {
let request_count = u64_to_i64_saturating(delta.requests);
let input_tokens = u64_to_i64_saturating(delta.input_tokens);
let output_tokens = u64_to_i64_saturating(delta.output_tokens);
let cached_tokens = u64_to_i64_saturating(delta.cached_tokens);
let image_input_tokens = u64_to_i64_saturating(delta.image_input_tokens);
let image_output_tokens = u64_to_i64_saturating(delta.image_output_tokens);
let image_request_count = u64_to_i64_saturating(delta.image_requests);
let image_request_failed_count = u64_to_i64_saturating(delta.image_request_failures);
let last_used_at = Utc::now();
sqlx::query(RECORD_USAGE_SQL)
.bind(account_id)
.bind(request_count)
.bind(u64_to_i64_saturating(delta.empty_responses))
.bind(input_tokens)
.bind(output_tokens)
.bind(cached_tokens)
.bind(u64_to_i64_saturating(delta.reasoning_tokens))
.bind(u64_to_i64_saturating(delta.total_tokens))
.bind(image_input_tokens)
.bind(image_output_tokens)
.bind(image_request_count)
.bind(image_request_failed_count)
.bind(request_count)
.bind(input_tokens)
.bind(output_tokens)
.bind(cached_tokens)
.bind(image_input_tokens)
.bind(image_output_tokens)
.bind(image_request_count)
.bind(image_request_failed_count)
.bind(last_used_at)
.bind(last_used_at)
.execute(&self.pool)
.await?;
Ok(())
}
async fn write_model_usage_delta(
&self,
account_id: &str,
model: &str,
delta: AccountModelUsageDelta,
) -> PgAccountUsageStoreResult<()> {
let model = model.trim();
if model.is_empty() {
return Ok(());
}
sqlx::query(RECORD_MODEL_USAGE_SQL)
.bind(account_id)
.bind(model)
.bind(u64_to_i64_saturating(delta.requests))
.bind(u64_to_i64_saturating(delta.errors))
.bind(u64_to_i64_saturating(delta.input_tokens))
.bind(u64_to_i64_saturating(delta.output_tokens))
.bind(u64_to_i64_saturating(delta.cached_tokens))
.bind(Utc::now())
.execute(&self.pool)
.await?;
Ok(())
}
async fn replace_runtime_window(
&self,
account_id: &str,
window: AccountUsageWindow,
) -> PgAccountUsageStoreResult<()> {
sqlx::query(SYNC_RUNTIME_WINDOW_SQL)
.bind(account_id)
.bind(u64_to_i64_saturating(window.request_count))
.bind(u64_to_i64_saturating(window.input_tokens))
.bind(u64_to_i64_saturating(window.output_tokens))
.bind(u64_to_i64_saturating(window.cached_tokens))
.bind(u64_to_i64_saturating(window.image_input_tokens))
.bind(u64_to_i64_saturating(window.image_output_tokens))
.bind(u64_to_i64_saturating(window.image_request_count))
.bind(u64_to_i64_saturating(window.image_request_failed_count))
.bind(window.started_at)
.bind(window.reset_at)
.bind(window.limit_window_seconds.map(u64_to_i64_saturating))
.bind(Utc::now())
.execute(&self.pool)
.await?;
Ok(())
}
async fn update_rate_limit_window(
&self,
account_id: &str,
reset_at: DateTime<Utc>,
limit_window_seconds: Option<u64>,
) -> PgAccountUsageStoreResult<()> {
let existing = sqlx::query(
"select window_reset_at, limit_window_seconds from account_usage where account_id = $1",
)
.bind(account_id)
.fetch_optional(&self.pool)
.await?;
let should_reset = existing
.as_ref()
.map(|row| {
should_reset_usage_window(
row.get("window_reset_at"),
optional_positive_i64_to_u64(row.get("limit_window_seconds")),
reset_at,
limit_window_seconds,
)
})
.unwrap_or_default();
let limit_window_seconds = limit_window_seconds.map(u64_to_i64_saturating);
if should_reset {
sqlx::query(SYNC_RATE_LIMIT_WINDOW_RESET_SQL)
.bind(account_id)
.bind(Utc::now())
.bind(reset_at)
.bind(limit_window_seconds)
.execute(&self.pool)
.await?;
} else {
sqlx::query(SYNC_RATE_LIMIT_WINDOW_SQL)
.bind(account_id)
.bind(reset_at)
.bind(limit_window_seconds)
.execute(&self.pool)
.await?;
}
Ok(())
}
/// 分页列出账号用量。
pub async fn list_usage(
&self,
@@ -239,6 +625,58 @@ impl PgAccountUsageStore {
}
}
#[async_trait]
impl AccountUsageStore for PgAccountUsageStore {
async fn snapshots(
&self,
account_ids: &[String],
) -> Result<HashMap<String, AccountUsageSnapshot>, AccountUsageStoreError> {
self.load_snapshots(account_ids).await.map_err(Into::into)
}
async fn record_usage_delta(
&self,
account_id: &str,
usage: AccountUsageDelta,
) -> Result<(), AccountUsageStoreError> {
self.write_usage_delta(account_id, usage)
.await
.map_err(Into::into)
}
async fn record_model_usage_delta(
&self,
account_id: &str,
model: &str,
usage: AccountModelUsageDelta,
) -> Result<(), AccountUsageStoreError> {
self.write_model_usage_delta(account_id, model, usage)
.await
.map_err(Into::into)
}
async fn sync_runtime_window(
&self,
account_id: &str,
window: AccountUsageWindow,
) -> Result<(), AccountUsageStoreError> {
self.replace_runtime_window(account_id, window)
.await
.map_err(Into::into)
}
async fn sync_rate_limit_window(
&self,
account_id: &str,
reset_at: DateTime<Utc>,
limit_window_seconds: Option<u64>,
) -> Result<(), AccountUsageStoreError> {
self.update_rate_limit_window(account_id, reset_at, limit_window_seconds)
.await
.map_err(Into::into)
}
}
fn usage_list_from_row(row: &PgRow) -> PgAccountUsageStoreResult<UsageListRecord> {
Ok(UsageListRecord {
account_id: row.get("account_id"),
@@ -269,6 +707,69 @@ fn usage_list_from_row(row: &PgRow) -> PgAccountUsageStoreResult<UsageListRecord
})
}
fn usage_snapshot_from_row(row: &PgRow) -> AccountUsageSnapshot {
AccountUsageSnapshot {
request_count: nonnegative_i64_to_u64(row.get("request_count")),
empty_response_count: nonnegative_i64_to_u64(row.get("empty_response_count")),
image_input_tokens: nonnegative_i64_to_u64(row.get("image_input_tokens")),
image_output_tokens: nonnegative_i64_to_u64(row.get("image_output_tokens")),
image_request_count: nonnegative_i64_to_u64(row.get("image_request_count")),
image_request_failed_count: nonnegative_i64_to_u64(row.get("image_request_failed_count")),
window_request_count: nonnegative_i64_to_u64(row.get("window_request_count")),
window_input_tokens: nonnegative_i64_to_u64(row.get("window_input_tokens")),
window_output_tokens: nonnegative_i64_to_u64(row.get("window_output_tokens")),
window_cached_tokens: nonnegative_i64_to_u64(row.get("window_cached_tokens")),
window_image_input_tokens: nonnegative_i64_to_u64(row.get("window_image_input_tokens")),
window_image_output_tokens: nonnegative_i64_to_u64(row.get("window_image_output_tokens")),
window_image_request_count: nonnegative_i64_to_u64(row.get("window_image_request_count")),
window_image_request_failed_count: nonnegative_i64_to_u64(
row.get("window_image_request_failed_count"),
),
window_started_at: row.get("window_started_at"),
window_reset_at: row.get("window_reset_at"),
limit_window_seconds: optional_positive_i64_to_u64(row.get("limit_window_seconds")),
last_used_at: row.get("last_used_at"),
}
}
pub fn should_reset_usage_window(
existing_reset_at: Option<DateTime<Utc>>,
existing_limit_window_seconds: Option<u64>,
new_reset_at: DateTime<Utc>,
new_limit_window_seconds: Option<u64>,
) -> bool {
let Some(existing_reset_at) = existing_reset_at else {
return false;
};
if existing_reset_at == new_reset_at {
return false;
}
let drift = existing_reset_at
.signed_duration_since(new_reset_at)
.num_seconds()
.unsigned_abs();
let window_seconds = new_limit_window_seconds
.or(existing_limit_window_seconds)
.unwrap_or(0);
let threshold = if window_seconds > 0 {
window_seconds / 2
} else {
3_600
};
drift >= threshold
}
fn u64_to_i64_saturating(value: u64) -> i64 {
value.min(i64::MAX as u64) as i64
}
fn optional_positive_i64_to_u64(value: Option<i64>) -> Option<u64> {
value
.and_then(|value| u64::try_from(value).ok())
.filter(|value| *value > 0)
}
fn to_page(rows: &[PgRow], limit: u32) -> Page<UsageListRecord> {
let has_more = rows.len() > limit as usize;
let mut items = Vec::with_capacity(limit as usize);
+539
View File
@@ -0,0 +1,539 @@
//! Dashboard 遥测聚合与展示模型。
use chrono::{DateTime, Duration, Utc};
use serde::{Deserialize, Serialize};
use crate::infra::{
format::{
format_compact_number, format_cost, format_duration_ms, format_percent, format_rate,
format_tokens, nonnegative_i64_to_u64,
},
time::{china_day_start, china_hour, china_hour_start, china_quarter_hour_start},
};
use super::account_usage::query::{AccountUsageSummary, AccountUsageTimeBucket};
const HEALTH_TIMELINE_SLOT_MINUTES: i64 = 15;
const HEALTH_TIMELINE_SLOTS: i64 = 24 * 4;
const HEALTH_TIMELINE_STABLE_SUCCESS_THRESHOLD: u64 = 3;
#[derive(Debug, Clone, Copy, Default)]
pub struct DashboardAccountCounts {
pub total: u64,
pub enabled: u64,
pub abnormal: u64,
}
#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum DashboardTrendKind {
#[default]
Usage,
Latency,
Errors,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct DashboardHealthTimelineData {
title: String,
description: String,
reliability_display: String,
points: String,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct DashboardCardsData {
accounts: DashboardAccountsCardData,
traffic: DashboardTrafficCardData,
tokens: DashboardTokenCardData,
cache: DashboardCacheCardData,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardAccountsCardData {
total: String,
total_value: u64,
enabled: String,
enabled_value: u64,
abnormal: String,
abnormal_value: u64,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardTrafficCardData {
today_requests: String,
today_requests_value: u64,
yesterday_requests_value: u64,
total_requests: String,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardTokenCardData {
today_tokens: String,
today_tokens_value: u64,
yesterday_tokens_value: u64,
total_tokens: String,
today_cost_usd: String,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardCacheCardData {
today_hit_rate: String,
today_hit_rate_value: Option<f64>,
yesterday_hit_rate_value: Option<f64>,
total_hit_rate: String,
total_cached_tokens: String,
average_first_token_latency_ms: String,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct DashboardTrendData {
kind: DashboardTrendKind,
points: Vec<DashboardTrendPointData>,
summary: Vec<DashboardTrendSummaryData>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardTrendPointData {
time: String,
requests: String,
requests_value: u64,
input_tokens: String,
input_tokens_value: u64,
output_tokens: String,
output_tokens_value: u64,
cached_tokens: String,
cached_tokens_value: u64,
cache_hit_rate_value: f64,
tokens_value: u64,
errors: String,
errors_value: u64,
latency: String,
latency_value: u64,
max_latency: String,
max_latency_value: u64,
min_latency: String,
min_latency_value: u64,
success_rate: String,
success_rate_value: f64,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardTrendSummaryData {
label: String,
value: String,
ratio: Option<String>,
}
#[derive(Debug, Clone, Copy, Default)]
struct UsageWindow {
requests: u64,
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
errors: u64,
first_token_latency_sum: u64,
first_token_latency_count: u64,
max_first_token_bucket_latency: u64,
min_first_token_bucket_latency: u64,
}
impl UsageWindow {
fn tokens(self) -> u64 {
self.input_tokens + self.output_tokens
}
fn cache_hit_rate(self) -> Option<f64> {
(self.input_tokens > 0).then(|| self.cached_tokens as f64 / self.input_tokens as f64)
}
fn avg_first_token_latency(self) -> Option<u64> {
self.first_token_latency_sum
.checked_div(self.first_token_latency_count)
}
fn max_first_token_bucket_latency(self) -> Option<u64> {
(self.max_first_token_bucket_latency > 0).then_some(self.max_first_token_bucket_latency)
}
fn min_first_token_bucket_latency(self) -> Option<u64> {
(self.min_first_token_bucket_latency > 0).then_some(self.min_first_token_bucket_latency)
}
}
pub fn dashboard_cards(
accounts: DashboardAccountCounts,
buckets: &[AccountUsageTimeBucket],
lifetime_usage: &AccountUsageSummary,
) -> DashboardCardsData {
dashboard_cards_at(accounts, buckets, lifetime_usage, Utc::now())
}
pub fn dashboard_cards_at(
accounts: DashboardAccountCounts,
buckets: &[AccountUsageTimeBucket],
lifetime_usage: &AccountUsageSummary,
now: DateTime<Utc>,
) -> DashboardCardsData {
let today_start = china_day_start(now);
let yesterday_start = today_start - Duration::days(1);
let today = usage_window(buckets, today_start, now);
let yesterday = usage_window(buckets, yesterday_start, today_start);
let today_cost = cost_window(buckets, today_start, now).unwrap_or(0.0);
let total_requests = nonnegative_i64_to_u64(lifetime_usage.request_count);
let total_tokens = nonnegative_i64_to_u64(lifetime_usage.total_tokens);
let total_cached_tokens = nonnegative_i64_to_u64(lifetime_usage.cached_tokens);
DashboardCardsData {
accounts: DashboardAccountsCardData {
total: format_compact_number(accounts.total),
total_value: accounts.total,
enabled: format_compact_number(accounts.enabled),
enabled_value: accounts.enabled,
abnormal: format_compact_number(accounts.abnormal),
abnormal_value: accounts.abnormal,
},
traffic: DashboardTrafficCardData {
today_requests: format_compact_number(today.requests),
today_requests_value: today.requests,
yesterday_requests_value: yesterday.requests,
total_requests: format_compact_number(total_requests),
},
tokens: DashboardTokenCardData {
today_tokens: format_tokens(today.tokens()),
today_tokens_value: today.tokens(),
yesterday_tokens_value: yesterday.tokens(),
total_tokens: format_tokens(total_tokens),
today_cost_usd: format_cost(today_cost),
},
cache: DashboardCacheCardData {
today_hit_rate: format_rate(today.cache_hit_rate()),
today_hit_rate_value: today.cache_hit_rate(),
yesterday_hit_rate_value: yesterday.cache_hit_rate(),
total_hit_rate: format_rate(summary_cache_hit_rate(lifetime_usage)),
total_cached_tokens: format_tokens(total_cached_tokens),
average_first_token_latency_ms: format_optional_duration_ms(
today.avg_first_token_latency(),
),
},
}
}
pub fn dashboard_trend_data(
buckets: &[AccountUsageTimeBucket],
kind: DashboardTrendKind,
) -> DashboardTrendData {
dashboard_trend_data_at(buckets, kind, Utc::now())
}
pub fn dashboard_trend_data_at(
records: &[AccountUsageTimeBucket],
kind: DashboardTrendKind,
now: DateTime<Utc>,
) -> DashboardTrendData {
let current_hour = china_hour_start(now);
let start = china_day_start(now);
let elapsed_hours = current_hour.signed_duration_since(start).num_hours();
let mut buckets = (0..=elapsed_hours)
.map(|index| (start + Duration::hours(index), UsageWindow::default()))
.collect::<Vec<_>>();
for record in records {
if record.bucket_start < start || record.bucket_start > now {
continue;
}
let record_hour = china_hour_start(record.bucket_start);
if let Some((_, bucket)) = buckets
.iter_mut()
.find(|(bucket_start, _)| *bucket_start == record_hour)
{
apply_bucket(bucket, record);
}
}
let points = buckets
.iter()
.map(|(bucket_start, bucket)| {
let latency = bucket.avg_first_token_latency().unwrap_or(0);
let max_latency = bucket.max_first_token_bucket_latency().unwrap_or(0);
let min_latency = bucket.min_first_token_bucket_latency().unwrap_or(0);
let cache_hit_rate = bucket.cache_hit_rate().unwrap_or(0.0);
let success_rate = if bucket.requests > 0 {
((bucket.requests - bucket.errors) as f64 / bucket.requests as f64 * 1000.0).round()
/ 10.0
} else {
0.0
};
DashboardTrendPointData {
time: format!("{:02}", china_hour(bucket_start)),
requests: format_compact_number(bucket.requests),
requests_value: bucket.requests,
input_tokens: format_tokens(bucket.input_tokens),
input_tokens_value: bucket.input_tokens,
output_tokens: format_tokens(bucket.output_tokens),
output_tokens_value: bucket.output_tokens,
cached_tokens: format_tokens(bucket.cached_tokens),
cached_tokens_value: bucket.cached_tokens,
cache_hit_rate_value: cache_hit_rate,
tokens_value: bucket.tokens(),
errors: format_compact_number(bucket.errors),
errors_value: bucket.errors,
latency: format_optional_duration_ms(Some(latency)),
latency_value: latency,
max_latency: format_optional_duration_ms(Some(max_latency)),
max_latency_value: max_latency,
min_latency: format_optional_duration_ms(Some(min_latency)),
min_latency_value: min_latency,
success_rate: format_percent(success_rate),
success_rate_value: success_rate,
}
})
.collect::<Vec<_>>();
DashboardTrendData {
kind,
summary: trend_summary(kind, &points),
points,
}
}
pub fn dashboard_health_timeline_data(
buckets: &[AccountUsageTimeBucket],
) -> DashboardHealthTimelineData {
dashboard_health_timeline_data_at(buckets, Utc::now())
}
pub fn dashboard_health_timeline_data_at(
records: &[AccountUsageTimeBucket],
now: DateTime<Utc>,
) -> DashboardHealthTimelineData {
let current_slot = china_quarter_hour_start(now);
let start = china_day_start(now);
let mut buckets = (0..HEALTH_TIMELINE_SLOTS)
.map(|index| {
(
start + Duration::minutes(HEALTH_TIMELINE_SLOT_MINUTES * index),
UsageWindow::default(),
)
})
.collect::<Vec<_>>();
for record in records {
if record.bucket_start < start || record.bucket_start > now {
continue;
}
let record_slot = china_quarter_hour_start(record.bucket_start);
if let Some((_, bucket)) = buckets
.iter_mut()
.find(|(bucket_start, _)| *bucket_start == record_slot)
{
apply_bucket(bucket, record);
}
}
let (requests, errors) = buckets
.iter()
.filter(|(bucket_start, _)| *bucket_start <= current_slot)
.fold((0, 0), |(requests, errors), (_, bucket)| {
(requests + bucket.requests, errors + bucket.errors)
});
let reliability = (requests > 0)
.then(|| ((requests - errors) as f64 / requests as f64 * 1000.0).round() / 10.0);
DashboardHealthTimelineData {
title: "请求健康时间线".to_string(),
description: "请求可靠性".to_string(),
reliability_display: reliability
.map(|value| format!("{value:.1}%"))
.unwrap_or_else(|| "-".to_string()),
points: buckets
.into_iter()
.map(|(bucket_start, bucket)| {
let success_rate = if bucket.requests > 0 {
(bucket.requests.saturating_sub(bucket.errors) as f64 / bucket.requests as f64
* 1000.0)
.round()
/ 10.0
} else {
0.0
};
health_tone(bucket, success_rate, bucket_start > current_slot)
})
.collect(),
}
}
fn trend_summary(
kind: DashboardTrendKind,
points: &[DashboardTrendPointData],
) -> Vec<DashboardTrendSummaryData> {
match kind {
DashboardTrendKind::Usage => vec![
trend_summary_item(
kind,
"输入",
points.iter().map(|p| p.input_tokens_value).sum(),
None,
),
trend_summary_item(
kind,
"输出",
points.iter().map(|p| p.output_tokens_value).sum(),
None,
),
trend_summary_item(
kind,
"缓存",
points.iter().map(|p| p.cached_tokens_value).sum(),
None,
),
],
DashboardTrendKind::Latency => {
let samples = points
.iter()
.filter(|point| point.latency_value > 0)
.collect::<Vec<_>>();
let average = if samples.is_empty() {
0
} else {
samples.iter().map(|point| point.latency_value).sum::<u64>() / samples.len() as u64
};
vec![
trend_summary_item(kind, "平均", average, None),
trend_summary_item(
kind,
"最高",
samples
.iter()
.map(|p| p.max_latency_value)
.max()
.unwrap_or(0),
None,
),
trend_summary_item(
kind,
"最低",
samples
.iter()
.filter_map(|p| (p.min_latency_value > 0).then_some(p.min_latency_value))
.min()
.unwrap_or(0),
None,
),
]
}
DashboardTrendKind::Errors => {
let errors = points.iter().map(|point| point.errors_value).sum::<u64>();
let requests = points.iter().map(|point| point.requests_value).sum::<u64>();
let success_rate = (requests > 0)
.then(|| ((requests - errors) as f64 / requests as f64 * 1000.0).round() / 10.0);
vec![
trend_summary_item(kind, "错误数", errors, None),
trend_summary_item(kind, "成功率", 0, success_rate),
trend_summary_item(kind, "总请求", requests, None),
]
}
}
}
fn trend_summary_item(
kind: DashboardTrendKind,
label: &str,
value: u64,
ratio: Option<f64>,
) -> DashboardTrendSummaryData {
DashboardTrendSummaryData {
label: label.to_string(),
value: match kind {
DashboardTrendKind::Usage => format_tokens(value),
DashboardTrendKind::Latency => format_optional_duration_ms(Some(value)),
DashboardTrendKind::Errors => format_compact_number(value),
},
ratio: ratio.map(format_percent),
}
}
fn usage_window(
records: &[AccountUsageTimeBucket],
start: DateTime<Utc>,
end: DateTime<Utc>,
) -> UsageWindow {
records
.iter()
.filter(|record| record.bucket_start >= start && record.bucket_start < end)
.fold(UsageWindow::default(), |mut window, record| {
apply_bucket(&mut window, record);
window
})
}
fn cost_window(
records: &[AccountUsageTimeBucket],
start: DateTime<Utc>,
end: DateTime<Utc>,
) -> Option<f64> {
let costs = records
.iter()
.filter(|record| record.bucket_start >= start && record.bucket_start < end)
.filter_map(|record| record.cost_usd)
.collect::<Vec<_>>();
(!costs.is_empty()).then(|| costs.into_iter().sum())
}
fn apply_bucket(window: &mut UsageWindow, record: &AccountUsageTimeBucket) {
window.requests += nonnegative_i64_to_u64(record.request_count);
window.input_tokens += nonnegative_i64_to_u64(record.input_tokens);
window.output_tokens += nonnegative_i64_to_u64(record.output_tokens);
window.cached_tokens += nonnegative_i64_to_u64(record.cached_tokens);
window.errors += nonnegative_i64_to_u64(record.error_count);
window.first_token_latency_sum += nonnegative_i64_to_u64(record.first_token_latency_sum);
window.first_token_latency_count += nonnegative_i64_to_u64(record.first_token_latency_count);
if let Some(latency) = first_token_bucket_latency(record) {
window.max_first_token_bucket_latency = window.max_first_token_bucket_latency.max(latency);
window.min_first_token_bucket_latency = if window.min_first_token_bucket_latency == 0 {
latency
} else {
window.min_first_token_bucket_latency.min(latency)
};
}
}
fn first_token_bucket_latency(record: &AccountUsageTimeBucket) -> Option<u64> {
nonnegative_i64_to_u64(record.first_token_latency_sum)
.checked_div(nonnegative_i64_to_u64(record.first_token_latency_count))
.filter(|latency| *latency > 0)
}
fn health_tone(bucket: UsageWindow, success_rate: f64, is_future: bool) -> char {
let successes = bucket.requests.saturating_sub(bucket.errors);
if is_future {
'0'
} else if bucket.requests == 0 {
'1'
} else if successes == 0 {
'2'
} else if success_rate < 90.0 {
'3'
} else if successes < HEALTH_TIMELINE_STABLE_SUCCESS_THRESHOLD {
'4'
} else {
'5'
}
}
fn summary_cache_hit_rate(summary: &AccountUsageSummary) -> Option<f64> {
let input_tokens = nonnegative_i64_to_u64(summary.input_tokens);
(input_tokens > 0)
.then(|| nonnegative_i64_to_u64(summary.cached_tokens) as f64 / input_tokens as f64)
}
fn format_optional_duration_ms(value: Option<u64>) -> String {
format_duration_ms(value.and_then(|value| i64::try_from(value).ok()))
}
-90
View File
@@ -1,90 +0,0 @@
//! 诊断数据类型与聚合函数。
use serde::Serialize;
use crate::{accounts::pool::AccountCapacitySummary, upstream::openai::fingerprint::Fingerprint};
/// 账号池摘要。
#[derive(Debug, Default, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountPoolDiagnostics {
/// 总数。
pub total: usize,
/// 活跃数。
pub active: usize,
/// 过期数。
pub expired: usize,
/// 配额耗尽数。
pub quota_exhausted: usize,
/// 刷新中数。
pub refreshing: usize,
/// 禁用数。
pub disabled: usize,
/// 封禁数。
pub banned: usize,
}
/// Account-pool capacity diagnostics.
#[derive(Debug, Default, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountCapacityDiagnostics {
/// Maximum concurrent slots per account.
pub max_concurrent_per_account: usize,
/// Total slots available across active accounts.
pub total_slots: usize,
/// Currently occupied slots.
pub used_slots: usize,
/// Currently available slots.
pub available_slots: usize,
}
/// Runtime fingerprint summary.
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct FingerprintDiagnostics {
/// Fingerprint source label.
pub source: &'static str,
/// Client originator.
pub originator: String,
/// App version.
pub app_version: String,
/// Build number.
pub build_number: String,
/// Platform name.
pub platform: String,
/// Architecture name.
pub arch: String,
/// Chromium major version.
pub chromium_version: String,
/// Expanded user-agent.
pub user_agent: String,
/// 指纹最后更新时间。
#[serde(skip_serializing_if = "Option::is_none")]
pub updated_at: Option<String>,
}
/// 构造 fingerprint 诊断数据。
pub fn fingerprint_diagnostics(fingerprint: &Fingerprint) -> FingerprintDiagnostics {
FingerprintDiagnostics {
source: "runtime",
originator: fingerprint.originator.clone(),
app_version: fingerprint.app_version.clone(),
build_number: fingerprint.build_number.clone(),
platform: fingerprint.platform.clone(),
arch: fingerprint.arch.clone(),
chromium_version: fingerprint.chromium_version.clone(),
user_agent: fingerprint.user_agent(),
updated_at: fingerprint.updated_at.clone(),
}
}
impl From<AccountCapacitySummary> for AccountCapacityDiagnostics {
fn from(summary: AccountCapacitySummary) -> Self {
Self {
max_concurrent_per_account: summary.max_concurrent_per_account,
total_slots: summary.total_slots,
used_slots: summary.used_slots,
available_slots: summary.available_slots,
}
}
}
+1 -1
View File
@@ -3,7 +3,7 @@
pub mod account_usage;
pub(crate) mod billing;
pub mod buckets;
pub mod diagnostics;
pub mod dashboard;
pub mod ops;
pub mod rebuild;
pub mod recorder;
@@ -2,8 +2,6 @@
/// Codex 事件语义与用量提取。
pub mod events;
/// 上游模型目录响应类型。
pub mod model_catalog;
/// Codex Responses / Compact 请求体。
pub mod responses;
/// Codex JSON schema 处理。
@@ -1,73 +0,0 @@
//! 上游模型目录原始条目。
use serde::{Deserialize, Serialize};
/// 上游声明的 reasoning effort 条目。
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct BackendReasoningEffort {
/// reasoning_effort 字段。
#[serde(default, rename = "reasoning_effort", alias = "reasoningEffort")]
pub reasoning_effort: Option<String>,
/// 上游可能返回的 effort 字段别名。
#[serde(default)]
pub effort: Option<String>,
/// 展示描述。
#[serde(default)]
pub description: Option<String>,
}
/// 上游截断策略定义。
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct BackendTruncationPolicy {
/// 上下文截断限制。
pub limit: Option<u64>,
}
/// 上游 `/codex/models` 返回的原始模型条目。
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct BackendModelEntry {
/// slug。
pub slug: Option<String>,
/// id。
pub id: Option<String>,
/// name。
pub name: Option<String>,
/// display_name。
pub display_name: Option<String>,
/// title。
pub title: Option<String>,
/// description。
pub description: Option<String>,
/// 是否默认模型。
pub is_default: Option<bool>,
/// 默认 reasoning effort。
pub default_reasoning_effort: Option<String>,
/// 默认 reasoning level。
pub default_reasoning_level: Option<String>,
/// 支持的 reasoning effort 列表。
#[serde(default)]
pub supported_reasoning_efforts: Vec<BackendReasoningEffort>,
/// 支持的 reasoning level 列表。
#[serde(default)]
pub supported_reasoning_levels: Vec<BackendReasoningEffort>,
/// 输入模态。
pub input_modalities: Option<Vec<String>>,
/// 输出模态。
pub output_modalities: Option<Vec<String>>,
/// 是否支持 personality。
pub supports_personality: Option<bool>,
/// 升级提示。
pub upgrade: Option<String>,
/// contextWindow 别名。
#[serde(alias = "contextWindow")]
pub context_window: Option<u64>,
/// maxContextWindow 别名。
#[serde(alias = "maxContextWindow")]
pub max_context_window: Option<u64>,
/// maxOutputTokens 别名。
#[serde(alias = "maxOutputTokens")]
pub max_output_tokens: Option<u64>,
/// truncationPolicy 别名。
#[serde(alias = "truncationPolicy")]
pub truncation_policy: Option<BackendTruncationPolicy>,
}
+15 -15
View File
@@ -24,7 +24,6 @@ use uuid::Uuid;
use crate::upstream::openai::fingerprint::{Fingerprint, RuntimeFingerprint};
use crate::upstream::openai::protocol::events::{extract_sse_usage, retry_after_seconds_from_body};
use crate::upstream::openai::protocol::model_catalog::BackendModelEntry;
use crate::upstream::openai::protocol::responses::{
http_sse_fallback_allowed, transport_for_request, CodexCompactRequest, CodexResponsesRequest,
CodexTransport,
@@ -201,7 +200,7 @@ pub trait CodexModelCatalogClient: Send + Sync + 'static {
async fn fetch_models(
&self,
request: &CodexModelCatalogRequest<'_>,
) -> Result<Vec<BackendModelEntry>, CodexModelCatalogClientError>;
) -> Result<Vec<Value>, CodexModelCatalogClientError>;
}
// ---------------------------------------------------------------------------
@@ -345,7 +344,7 @@ impl CodexModelCatalogClient for CodexBackendClient {
async fn fetch_models(
&self,
request: &CodexModelCatalogRequest<'_>,
) -> Result<Vec<BackendModelEntry>, CodexModelCatalogClientError> {
) -> Result<Vec<Value>, CodexModelCatalogClientError> {
self.fetch_models_with_context(CodexRequestContext {
access_token: request.access_token,
account_id: request.account_id,
@@ -446,7 +445,7 @@ pub(super) async fn read_capped_error_body(
// Model entry extraction
// ---------------------------------------------------------------------------
fn extract_backend_model_entries(value: &Value) -> Vec<BackendModelEntry> {
fn extract_model_entries(value: &Value) -> Vec<Value> {
let Some(models) = value
.pointer("/chat_models/models")
.or_else(|| value.get("models"))
@@ -460,22 +459,23 @@ fn extract_backend_model_entries(value: &Value) -> Vec<BackendModelEntry> {
let mut entries = Vec::new();
for model in models {
if let Some(nested) = model.get("models").and_then(Value::as_array) {
entries.extend(nested.iter().filter_map(parse_backend_model_entry));
} else if let Some(entry) = parse_backend_model_entry(model) {
entries.push(entry);
entries.extend(nested.iter().filter(|entry| is_model_entry(entry)).cloned());
} else if is_model_entry(model) {
entries.push(model.clone());
}
}
entries
}
fn parse_backend_model_entry(value: &Value) -> Option<BackendModelEntry> {
let entry = serde_json::from_value::<BackendModelEntry>(value.clone()).ok()?;
(entry.slug.is_some()
|| entry.id.is_some()
|| entry.name.is_some()
|| entry.display_name.is_some()
|| entry.title.is_some())
.then_some(entry)
fn is_model_entry(value: &Value) -> bool {
["slug", "id", "name", "display_name", "title"]
.into_iter()
.any(|key| {
value
.get(key)
.and_then(Value::as_str)
.is_some_and(|value| !value.trim().is_empty())
})
}
// ---------------------------------------------------------------------------
@@ -463,7 +463,7 @@ impl CodexBackendClient {
pub(super) async fn fetch_models_with_context(
&self,
context: CodexRequestContext<'_>,
) -> CodexClientResult<Vec<BackendModelEntry>> {
) -> CodexClientResult<Vec<Value>> {
let fingerprint = self.fingerprint.current();
let endpoints = [
format!(
@@ -481,7 +481,7 @@ impl CodexBackendClient {
continue;
}
let parsed = response.json::<Value>().await?;
let models = extract_backend_model_entries(&parsed);
let models = extract_model_entries(&parsed);
if !models.is_empty() {
return Ok(models);
}
+2
View File
@@ -0,0 +1,2 @@
mod cooldown;
mod store;

Some files were not shown because too many files have changed in this diff Show More