mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
refactor: align implementation with authoritative architecture
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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"
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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,6 +1,5 @@
|
||||
//! 账号管理用例。
|
||||
|
||||
mod cookies;
|
||||
mod export;
|
||||
mod import;
|
||||
mod lifecycle;
|
||||
|
||||
@@ -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") {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -8,6 +8,5 @@ pub mod pool;
|
||||
pub mod quota;
|
||||
pub mod refresh;
|
||||
pub mod scheduler;
|
||||
pub mod service;
|
||||
pub mod store;
|
||||
mod window;
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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::*;
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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>,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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},
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -5,7 +5,7 @@ use axum::{
|
||||
Router,
|
||||
};
|
||||
|
||||
use crate::bootstrap::state::AppState;
|
||||
use crate::api::AppState;
|
||||
|
||||
use super::{
|
||||
accounts_routes::{
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
|
||||
@@ -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)?;
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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},
|
||||
|
||||
@@ -5,3 +5,5 @@ pub mod assets;
|
||||
pub mod client;
|
||||
pub mod middleware;
|
||||
pub mod router;
|
||||
|
||||
pub use router::AppState;
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,6 +1,5 @@
|
||||
//! 进程装配、共享状态与后台任务。
|
||||
|
||||
pub mod application;
|
||||
pub mod config;
|
||||
pub mod import_sqlite;
|
||||
pub mod services;
|
||||
|
||||
@@ -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",
|
||||
)))
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
@@ -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());
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
+5
-3
@@ -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> {
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
//! 调度失败恢复策略。
|
||||
|
||||
pub(crate) mod auth;
|
||||
pub mod cloudflare;
|
||||
pub(crate) mod exhaustion;
|
||||
pub mod implicit_resume;
|
||||
pub mod reasoning_replay;
|
||||
@@ -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());
|
||||
}
|
||||
@@ -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 调度服务。
|
||||
+35
-35
@@ -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,
|
||||
@@ -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;
|
||||
+2
-2
@@ -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();
|
||||
+15
-23
@@ -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
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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>,
|
||||
}
|
||||
@@ -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,8 +1,6 @@
|
||||
//! 模型目录领域。
|
||||
|
||||
pub mod catalog;
|
||||
pub mod config;
|
||||
pub mod info;
|
||||
pub mod service;
|
||||
pub mod snapshot;
|
||||
pub mod store;
|
||||
pub mod types;
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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>,
|
||||
}
|
||||
@@ -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)]
|
||||
|
||||
@@ -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>,
|
||||
}
|
||||
@@ -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(¤t).await?;
|
||||
let key = generate_admin_api_key();
|
||||
self.store
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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()))
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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>,
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
mod cooldown;
|
||||
mod store;
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user