mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-13 08:06:41 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
cbf9ce91a1 | ||
|
|
71e51a868f | ||
|
|
3c9ec50bf2 | ||
|
|
c1b81809a6 | ||
|
|
939f392438 | ||
|
|
ff0de4742d | ||
|
|
44f45101b8 | ||
|
|
a240668a47 | ||
|
|
30f354965c | ||
|
|
3b462f7ecf | ||
|
|
5682d4fbe9 | ||
|
|
146776c742 | ||
|
|
e8c14d4324 | ||
|
|
5483e37b00 | ||
|
|
daad7c746c |
@@ -7,7 +7,7 @@ use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::infra::{
|
||||
format::{format_compact_percent, format_tokens},
|
||||
format::{format_compact_percent, format_plain_number, format_tokens, nonnegative_i64_to_u64},
|
||||
time::{china_datetime, china_relative_time},
|
||||
};
|
||||
|
||||
@@ -31,13 +31,29 @@ struct AdminAccountQuotaWindowData {
|
||||
label_display: String,
|
||||
used_percent: Option<f64>,
|
||||
used_percent_display: String,
|
||||
token_usage_display: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
local_usage: Option<AdminAccountQuotaWindowLocalUsageData>,
|
||||
reset_at_display: String,
|
||||
window_used_display: String,
|
||||
#[serde(skip)]
|
||||
reset_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AdminAccountQuotaWindowLocalUsageData {
|
||||
request_count: i64,
|
||||
request_count_display: String,
|
||||
input_tokens: i64,
|
||||
input_tokens_display: String,
|
||||
output_tokens: i64,
|
||||
output_tokens_display: String,
|
||||
cached_tokens: i64,
|
||||
cached_tokens_display: String,
|
||||
total_tokens: i64,
|
||||
total_tokens_display: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct AdminAccountQuotaUsageWindow {
|
||||
pub(crate) key: String,
|
||||
@@ -56,17 +72,37 @@ impl AdminAccountQuotaUsageWindow {
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub(crate) struct AdminAccountQuotaWindowTokenUsage {
|
||||
total_tokens: u64,
|
||||
pub(crate) struct AdminAccountQuotaWindowLocalUsage {
|
||||
pub(crate) request_count: i64,
|
||||
pub(crate) input_tokens: i64,
|
||||
pub(crate) output_tokens: i64,
|
||||
pub(crate) cached_tokens: i64,
|
||||
}
|
||||
|
||||
impl AdminAccountQuotaWindowTokenUsage {
|
||||
pub(crate) fn add_tokens(&mut self, tokens: u64) {
|
||||
self.total_tokens = self.total_tokens.saturating_add(tokens);
|
||||
impl AdminAccountQuotaWindowLocalUsage {
|
||||
pub(crate) fn add(&mut self, other: Self) {
|
||||
self.request_count = self.request_count.saturating_add(other.request_count);
|
||||
self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
|
||||
self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
|
||||
self.cached_tokens = self.cached_tokens.saturating_add(other.cached_tokens);
|
||||
}
|
||||
}
|
||||
|
||||
fn display(self) -> String {
|
||||
format_tokens(self.total_tokens)
|
||||
impl From<AdminAccountQuotaWindowLocalUsage> for AdminAccountQuotaWindowLocalUsageData {
|
||||
fn from(usage: AdminAccountQuotaWindowLocalUsage) -> Self {
|
||||
let total_tokens = usage.input_tokens.saturating_add(usage.output_tokens);
|
||||
Self {
|
||||
request_count: usage.request_count,
|
||||
request_count_display: format_plain_number(nonnegative_i64_to_u64(usage.request_count)),
|
||||
input_tokens: usage.input_tokens,
|
||||
input_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.input_tokens)),
|
||||
output_tokens: usage.output_tokens,
|
||||
output_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.output_tokens)),
|
||||
cached_tokens: usage.cached_tokens,
|
||||
cached_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.cached_tokens)),
|
||||
total_tokens,
|
||||
total_tokens_display: format_tokens(nonnegative_i64_to_u64(total_tokens)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -93,17 +129,18 @@ impl AdminAccountQuotaData {
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn apply_token_usage(
|
||||
pub(crate) fn apply_local_usage(
|
||||
&mut self,
|
||||
usage_by_window: &HashMap<String, AdminAccountQuotaWindowTokenUsage>,
|
||||
usage_by_window: &HashMap<String, AdminAccountQuotaWindowLocalUsage>,
|
||||
) {
|
||||
for window in &mut self.windows {
|
||||
window.token_usage_display =
|
||||
match (window.usage_window(), usage_by_window.get(&window.key)) {
|
||||
(Some(_), Some(usage)) => usage.display(),
|
||||
(Some(_), None) => AdminAccountQuotaWindowTokenUsage::default().display(),
|
||||
(None, _) => "-".to_string(),
|
||||
};
|
||||
window.local_usage = window.usage_window().map(|_| {
|
||||
usage_by_window
|
||||
.get(&window.key)
|
||||
.copied()
|
||||
.unwrap_or_default()
|
||||
.into()
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -185,7 +222,7 @@ fn push_monthly_quota_window(
|
||||
label_display: "月限额".to_string(),
|
||||
used_percent,
|
||||
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_compact_percent),
|
||||
token_usage_display: "-".to_string(),
|
||||
local_usage: None,
|
||||
reset_at_display: reset_at
|
||||
.as_ref()
|
||||
.map_or_else(|| "-".to_string(), china_datetime),
|
||||
@@ -274,7 +311,7 @@ fn push_quota_window(
|
||||
label_display,
|
||||
used_percent,
|
||||
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_compact_percent),
|
||||
token_usage_display: "-".to_string(),
|
||||
local_usage: None,
|
||||
reset_at_display: reset_at
|
||||
.as_ref()
|
||||
.map_or_else(|| "-".to_string(), china_datetime),
|
||||
|
||||
@@ -27,7 +27,7 @@ use crate::{
|
||||
admin::{
|
||||
accounts::quota_view::{
|
||||
quota_data, AdminAccountQuotaData, AdminAccountQuotaUsageWindow,
|
||||
AdminAccountQuotaWindowTokenUsage,
|
||||
AdminAccountQuotaWindowLocalUsage,
|
||||
},
|
||||
accounts::service::{
|
||||
AdminAccountError, AdminAccountHealthCheck, AdminAccountMetadata,
|
||||
@@ -135,6 +135,8 @@ struct AdminAccountData {
|
||||
plan_type: Option<String>,
|
||||
has_refresh_token: bool,
|
||||
status: String,
|
||||
display_status: String,
|
||||
token_refreshing: bool,
|
||||
access_token_expires_at: Option<String>,
|
||||
access_token_expires_at_display: Option<String>,
|
||||
added_at: String,
|
||||
@@ -147,7 +149,7 @@ struct AdminAccountData {
|
||||
|
||||
impl From<AdminAccountMetadata> for AdminAccountData {
|
||||
fn from(a: AdminAccountMetadata) -> Self {
|
||||
Self::from_parts(a, None, None, Vec::new())
|
||||
Self::from_parts(a, None, None, Vec::new(), false)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -157,10 +159,12 @@ impl AdminAccountData {
|
||||
usage: Option<&AdminUsageRecord>,
|
||||
quota: Option<AdminAccountQuotaData>,
|
||||
models: Vec<AdminAccountModelUsageData>,
|
||||
token_refreshing: bool,
|
||||
) -> Self {
|
||||
let access_token_expires_at = a.access_token_expires_at.as_ref().map(china_rfc3339);
|
||||
let access_token_expires_at_display =
|
||||
a.access_token_expires_at.as_ref().map(china_datetime);
|
||||
let display_status = account_display_status(a.status, token_refreshing).to_string();
|
||||
Self {
|
||||
id: a.id,
|
||||
email: a.email,
|
||||
@@ -170,6 +174,8 @@ impl AdminAccountData {
|
||||
plan_type: a.plan_type,
|
||||
has_refresh_token: a.has_refresh_token,
|
||||
status: a.status.as_str().to_string(),
|
||||
display_status,
|
||||
token_refreshing,
|
||||
access_token_expires_at,
|
||||
access_token_expires_at_display,
|
||||
added_at: china_rfc3339(&a.added_at),
|
||||
@@ -289,6 +295,7 @@ struct AccountListStats {
|
||||
usage_by_account: HashMap<String, AdminUsageRecord>,
|
||||
quota_by_account: HashMap<String, AdminAccountQuotaData>,
|
||||
models_by_account: HashMap<String, Vec<AdminAccountModelUsageData>>,
|
||||
refreshing_account_ids: HashSet<String>,
|
||||
}
|
||||
|
||||
struct AccountModelUsageRecord {
|
||||
@@ -303,14 +310,6 @@ struct AccountModelUsageRecord {
|
||||
last_used_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
struct AccountUsageWindowStats {
|
||||
request_count: i64,
|
||||
input_tokens: i64,
|
||||
output_tokens: i64,
|
||||
cached_tokens: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct AccountQuotaUsageWindow {
|
||||
account_id: String,
|
||||
@@ -328,6 +327,7 @@ impl AccountListStats {
|
||||
.get(&account_id)
|
||||
.cloned()
|
||||
.unwrap_or_default(),
|
||||
self.refreshing_account_ids.contains(&account_id),
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -531,7 +531,7 @@ pub(crate) async fn accounts(
|
||||
.await
|
||||
{
|
||||
Ok(page) => {
|
||||
let stats = account_list_stats(&state, &page.items, "a_by_account).await;
|
||||
let stats = account_list_stats(&state, &page.items, "a_by_account).await?;
|
||||
let page = crate::infra::json::NumberedPage {
|
||||
items: page
|
||||
.items
|
||||
@@ -558,7 +558,7 @@ pub(crate) async fn accounts(
|
||||
.await
|
||||
{
|
||||
Ok(page) => {
|
||||
let stats = account_list_stats(&state, &page.items, "a_by_account).await;
|
||||
let stats = account_list_stats(&state, &page.items, "a_by_account).await?;
|
||||
let page = Page {
|
||||
items: page
|
||||
.items
|
||||
@@ -691,7 +691,7 @@ pub(crate) async fn account_quota(
|
||||
let raw = data.get("raw").cloned().unwrap_or(Value::Null);
|
||||
let quota_json = quota.to_string();
|
||||
let mut quota_data = quota_data("a_json, Some(Utc::now()));
|
||||
apply_account_quota_window_token_usage(&state, &account_id, &mut quota_data).await;
|
||||
apply_account_quota_window_local_usage(&state, &account_id, &mut quota_data).await;
|
||||
let account =
|
||||
account_data_for_quota_refresh(&state, &account_id, quota_data.clone()).await?;
|
||||
Ok(AdminResponse::new(
|
||||
@@ -728,7 +728,8 @@ async fn account_data_for_quota_refresh(
|
||||
return Err(account_not_found());
|
||||
};
|
||||
let quota_by_account = HashMap::from([(account_id.to_string(), quota_data)]);
|
||||
let stats = account_list_stats(state, std::slice::from_ref(&account), "a_by_account).await;
|
||||
let stats =
|
||||
account_list_stats(state, std::slice::from_ref(&account), "a_by_account).await?;
|
||||
Ok(stats.data_for(account))
|
||||
}
|
||||
|
||||
@@ -910,7 +911,7 @@ async fn account_list_stats(
|
||||
state: &AppState,
|
||||
accounts: &[AdminAccountMetadata],
|
||||
quota_by_account: &HashMap<String, AdminAccountQuotaData>,
|
||||
) -> AccountListStats {
|
||||
) -> Result<AccountListStats, AdminError> {
|
||||
let account_ids = accounts
|
||||
.iter()
|
||||
.map(|account| account.id.clone())
|
||||
@@ -922,10 +923,11 @@ async fn account_list_stats(
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
let mut quota_by_account = quota_by_account.clone();
|
||||
apply_quota_window_token_usage(state, accounts, &mut quota_by_account).await;
|
||||
let selected_quota_windows = selected_quota_windows_by_account(accounts, "a_by_account);
|
||||
let all_quota_windows = quota_usage_windows_by_account(accounts, "a_by_account);
|
||||
let quota_window_usage_stats =
|
||||
account_usage_stats_by_quota_window(state, &selected_quota_windows).await;
|
||||
quota_window_local_usage_by_account(state, &all_quota_windows).await;
|
||||
apply_quota_window_local_usage(accounts, &mut quota_by_account, "a_window_usage_stats);
|
||||
let selected_quota_windows = selected_quota_windows_by_account(accounts, "a_by_account);
|
||||
let usage_records = usage_records
|
||||
.into_iter()
|
||||
.map(|usage| {
|
||||
@@ -937,14 +939,29 @@ async fn account_list_stats(
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let models_by_account = list_current_window_model_usage(state, &usage_records).await;
|
||||
let refreshing_account_ids = state
|
||||
.services
|
||||
.token_refresh
|
||||
.refreshing_account_ids(&account_ids, Utc::now())
|
||||
.await
|
||||
.map_err(|error| AdminError::internal(error.to_string()))?;
|
||||
|
||||
AccountListStats {
|
||||
Ok(AccountListStats {
|
||||
usage_by_account: usage_records
|
||||
.into_iter()
|
||||
.map(|usage| (usage.account_id.clone(), usage))
|
||||
.collect(),
|
||||
quota_by_account,
|
||||
models_by_account,
|
||||
refreshing_account_ids,
|
||||
})
|
||||
}
|
||||
|
||||
fn account_display_status(status: AccountStatus, token_refreshing: bool) -> &'static str {
|
||||
if token_refreshing && matches!(status, AccountStatus::Active | AccountStatus::Expired) {
|
||||
"refreshing"
|
||||
} else {
|
||||
status.as_str()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -972,105 +989,11 @@ fn selected_quota_window(quota: &AdminAccountQuotaData) -> Option<AdminAccountQu
|
||||
})
|
||||
}
|
||||
|
||||
async fn account_usage_stats_by_quota_window(
|
||||
state: &AppState,
|
||||
windows_by_account: &HashMap<String, AdminAccountQuotaUsageWindow>,
|
||||
) -> HashMap<String, AccountUsageWindowStats> {
|
||||
if windows_by_account.is_empty() {
|
||||
return HashMap::new();
|
||||
}
|
||||
|
||||
let mut builder = QueryBuilder::<Sqlite>::new(
|
||||
"select
|
||||
account_id,
|
||||
coalesce(sum(request_count), 0) as request_count,
|
||||
coalesce(sum(input_tokens), 0) as input_tokens,
|
||||
coalesce(sum(output_tokens), 0) as output_tokens,
|
||||
coalesce(sum(cached_tokens), 0) as cached_tokens
|
||||
from usage_time_buckets
|
||||
where ",
|
||||
);
|
||||
for (index, (account_id, window)) in windows_by_account.iter().enumerate() {
|
||||
if index > 0 {
|
||||
builder.push(" or ");
|
||||
}
|
||||
let bucket_start = china_quarter_hour_start(window.start);
|
||||
builder.push("(account_id = ");
|
||||
builder.push_bind(account_id);
|
||||
builder.push(" and bucket_start >= ");
|
||||
builder.push_bind(bucket_start.to_rfc3339());
|
||||
builder.push(" and bucket_start <= ");
|
||||
builder.push_bind(window.end.to_rfc3339());
|
||||
builder.push(")");
|
||||
}
|
||||
builder.push(" group by account_id");
|
||||
|
||||
let rows = builder
|
||||
.build()
|
||||
.fetch_all(state.services.background_tasks.accounts.pool())
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
rows.into_iter()
|
||||
.map(|row| {
|
||||
(
|
||||
row.get::<String, _>("account_id"),
|
||||
AccountUsageWindowStats {
|
||||
request_count: row.get("request_count"),
|
||||
input_tokens: row.get("input_tokens"),
|
||||
output_tokens: row.get("output_tokens"),
|
||||
cached_tokens: row.get("cached_tokens"),
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn apply_selected_quota_window_usage(
|
||||
mut usage: AdminUsageRecord,
|
||||
selected_quota_windows: &HashMap<String, AdminAccountQuotaUsageWindow>,
|
||||
quota_window_usage_stats: &HashMap<String, AccountUsageWindowStats>,
|
||||
) -> AdminUsageRecord {
|
||||
let Some(window) = selected_quota_windows.get(&usage.account_id) else {
|
||||
return usage;
|
||||
};
|
||||
let stats = quota_window_usage_stats
|
||||
.get(&usage.account_id)
|
||||
.copied()
|
||||
.unwrap_or_default();
|
||||
usage.window_request_count = stats.request_count;
|
||||
usage.window_input_tokens = stats.input_tokens;
|
||||
usage.window_output_tokens = stats.output_tokens;
|
||||
usage.window_cached_tokens = stats.cached_tokens;
|
||||
usage.window_started_at = Some(window.start);
|
||||
usage.window_reset_at = Some(window.end);
|
||||
usage.limit_window_seconds = Some(window.window_seconds);
|
||||
usage
|
||||
}
|
||||
|
||||
async fn apply_account_quota_window_token_usage(
|
||||
state: &AppState,
|
||||
account_id: &str,
|
||||
quota: &mut AdminAccountQuotaData,
|
||||
) {
|
||||
let windows = quota
|
||||
.usage_windows()
|
||||
.into_iter()
|
||||
.map(|window| AccountQuotaUsageWindow {
|
||||
account_id: account_id.to_string(),
|
||||
window,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let usage_by_account = quota_window_token_usage_by_account(state, &windows).await;
|
||||
let empty_usage = HashMap::new();
|
||||
quota.apply_token_usage(usage_by_account.get(account_id).unwrap_or(&empty_usage));
|
||||
}
|
||||
|
||||
async fn apply_quota_window_token_usage(
|
||||
state: &AppState,
|
||||
fn quota_usage_windows_by_account(
|
||||
accounts: &[AdminAccountMetadata],
|
||||
quota_by_account: &mut HashMap<String, AdminAccountQuotaData>,
|
||||
) {
|
||||
let windows = accounts
|
||||
quota_by_account: &HashMap<String, AdminAccountQuotaData>,
|
||||
) -> Vec<AccountQuotaUsageWindow> {
|
||||
accounts
|
||||
.iter()
|
||||
.filter_map(|account| {
|
||||
quota_by_account
|
||||
@@ -1086,21 +1009,44 @@ async fn apply_quota_window_token_usage(
|
||||
window,
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let usage_by_account = quota_window_token_usage_by_account(state, &windows).await;
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn apply_account_quota_window_local_usage(
|
||||
state: &AppState,
|
||||
account_id: &str,
|
||||
quota: &mut AdminAccountQuotaData,
|
||||
) {
|
||||
let windows = quota
|
||||
.usage_windows()
|
||||
.into_iter()
|
||||
.map(|window| AccountQuotaUsageWindow {
|
||||
account_id: account_id.to_string(),
|
||||
window,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let usage_by_account = quota_window_local_usage_by_account(state, &windows).await;
|
||||
let empty_usage = HashMap::new();
|
||||
quota.apply_local_usage(usage_by_account.get(account_id).unwrap_or(&empty_usage));
|
||||
}
|
||||
|
||||
fn apply_quota_window_local_usage(
|
||||
accounts: &[AdminAccountMetadata],
|
||||
quota_by_account: &mut HashMap<String, AdminAccountQuotaData>,
|
||||
usage_by_account: &HashMap<String, HashMap<String, AdminAccountQuotaWindowLocalUsage>>,
|
||||
) {
|
||||
for account in accounts {
|
||||
if let Some(quota) = quota_by_account.get_mut(&account.id) {
|
||||
let empty_usage = HashMap::new();
|
||||
quota.apply_token_usage(usage_by_account.get(&account.id).unwrap_or(&empty_usage));
|
||||
quota.apply_local_usage(usage_by_account.get(&account.id).unwrap_or(&empty_usage));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn quota_window_token_usage_by_account(
|
||||
async fn quota_window_local_usage_by_account(
|
||||
state: &AppState,
|
||||
windows: &[AccountQuotaUsageWindow],
|
||||
) -> HashMap<String, HashMap<String, AdminAccountQuotaWindowTokenUsage>> {
|
||||
) -> HashMap<String, HashMap<String, AdminAccountQuotaWindowLocalUsage>> {
|
||||
let Some(min_start) = windows.iter().map(|window| window.window.start).min() else {
|
||||
return HashMap::new();
|
||||
};
|
||||
@@ -1123,8 +1069,10 @@ async fn quota_window_token_usage_by_account(
|
||||
"select
|
||||
account_id,
|
||||
bucket_start,
|
||||
coalesce(sum(request_count), 0) as request_count,
|
||||
coalesce(sum(input_tokens), 0) as input_tokens,
|
||||
coalesce(sum(output_tokens), 0) as output_tokens
|
||||
coalesce(sum(output_tokens), 0) as output_tokens,
|
||||
coalesce(sum(cached_tokens), 0) as cached_tokens
|
||||
from usage_time_buckets
|
||||
where account_id in (",
|
||||
);
|
||||
@@ -1134,7 +1082,7 @@ async fn quota_window_token_usage_by_account(
|
||||
}
|
||||
separated.push_unseparated(")");
|
||||
builder.push(" and bucket_start >= ");
|
||||
builder.push_bind(min_start.to_rfc3339());
|
||||
builder.push_bind(china_quarter_hour_start(min_start).to_rfc3339());
|
||||
builder.push(" and bucket_start <= ");
|
||||
builder.push_bind(max_end.to_rfc3339());
|
||||
builder.push(" group by account_id, bucket_start");
|
||||
@@ -1145,35 +1093,61 @@ async fn quota_window_token_usage_by_account(
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
let mut usage_by_account =
|
||||
HashMap::<String, HashMap<String, AdminAccountQuotaWindowTokenUsage>>::new();
|
||||
HashMap::<String, HashMap<String, AdminAccountQuotaWindowLocalUsage>>::new();
|
||||
|
||||
for row in rows {
|
||||
let account_id: String = row.get("account_id");
|
||||
let Ok(bucket_start) = parse_rfc3339_utc(row.get::<&str, _>("bucket_start")) else {
|
||||
continue;
|
||||
};
|
||||
let total_tokens = nonnegative_i64_to_u64(row.get::<i64, _>("input_tokens"))
|
||||
.saturating_add(nonnegative_i64_to_u64(row.get::<i64, _>("output_tokens")));
|
||||
let bucket_usage = AdminAccountQuotaWindowLocalUsage {
|
||||
request_count: row.get("request_count"),
|
||||
input_tokens: row.get("input_tokens"),
|
||||
output_tokens: row.get("output_tokens"),
|
||||
cached_tokens: row.get("cached_tokens"),
|
||||
};
|
||||
let Some(account_windows) = windows_by_account.get(account_id.as_str()) else {
|
||||
continue;
|
||||
};
|
||||
|
||||
for window in account_windows
|
||||
.iter()
|
||||
.filter(|window| bucket_start >= window.start && bucket_start <= window.end)
|
||||
{
|
||||
for window in account_windows.iter().filter(|window| {
|
||||
bucket_start >= china_quarter_hour_start(window.start) && bucket_start <= window.end
|
||||
}) {
|
||||
usage_by_account
|
||||
.entry(account_id.clone())
|
||||
.or_default()
|
||||
.entry(window.key.clone())
|
||||
.or_default()
|
||||
.add_tokens(total_tokens);
|
||||
.add(bucket_usage);
|
||||
}
|
||||
}
|
||||
|
||||
usage_by_account
|
||||
}
|
||||
|
||||
fn apply_selected_quota_window_usage(
|
||||
mut usage: AdminUsageRecord,
|
||||
selected_quota_windows: &HashMap<String, AdminAccountQuotaUsageWindow>,
|
||||
quota_window_usage_stats: &HashMap<String, HashMap<String, AdminAccountQuotaWindowLocalUsage>>,
|
||||
) -> AdminUsageRecord {
|
||||
let Some(window) = selected_quota_windows.get(&usage.account_id) else {
|
||||
return usage;
|
||||
};
|
||||
let stats = quota_window_usage_stats
|
||||
.get(&usage.account_id)
|
||||
.and_then(|usage_by_window| usage_by_window.get(&window.key))
|
||||
.copied()
|
||||
.unwrap_or_default();
|
||||
usage.window_request_count = stats.request_count;
|
||||
usage.window_input_tokens = stats.input_tokens;
|
||||
usage.window_output_tokens = stats.output_tokens;
|
||||
usage.window_cached_tokens = stats.cached_tokens;
|
||||
usage.window_started_at = Some(window.start);
|
||||
usage.window_reset_at = Some(window.end);
|
||||
usage.limit_window_seconds = Some(window.window_seconds);
|
||||
usage
|
||||
}
|
||||
|
||||
async fn account_summary_data(
|
||||
state: &AppState,
|
||||
quota_by_account: &HashMap<String, AdminAccountQuotaData>,
|
||||
|
||||
@@ -75,6 +75,8 @@ pub enum AdminAccountError {
|
||||
Delete,
|
||||
#[error("failed to update claims")]
|
||||
UpdateClaims,
|
||||
#[error("failed to coordinate token refresh lease")]
|
||||
RefreshLease,
|
||||
#[error("failed to sync account pool")]
|
||||
SyncAccountPool,
|
||||
#[error("failed to store quota")]
|
||||
|
||||
@@ -3,10 +3,11 @@ use std::{
|
||||
hash::{Hash, Hasher},
|
||||
};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use futures::{stream, StreamExt};
|
||||
use secrecy::{ExposeSecret, SecretString};
|
||||
use tokio::time::sleep;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::{
|
||||
infra::{
|
||||
@@ -25,6 +26,8 @@ use super::{
|
||||
AdminAccountService,
|
||||
};
|
||||
|
||||
const ADMIN_REFRESH_LEASE_TTL_SECONDS: i64 = 5 * 60;
|
||||
|
||||
impl AdminAccountService {
|
||||
pub async fn list(
|
||||
&self,
|
||||
@@ -281,20 +284,44 @@ impl AdminAccountService {
|
||||
duration_ms: elapsed_millis_i64(started_at),
|
||||
};
|
||||
|
||||
let Some(refresh_token) = account.refresh_token.as_ref() else {
|
||||
let Some(refresh_token) = account
|
||||
.refresh_token
|
||||
.as_ref()
|
||||
.map(|token| token.expose_secret().to_string())
|
||||
else {
|
||||
return Ok(skipped("no refresh token"));
|
||||
};
|
||||
if account.status == crate::upstream::accounts::model::AccountStatus::Disabled {
|
||||
return Ok(skipped("manually disabled"));
|
||||
}
|
||||
if account.status == crate::upstream::accounts::model::AccountStatus::Refreshing {
|
||||
|
||||
let lease_now = Utc::now();
|
||||
let lease_owner = format!(
|
||||
"{}:{}",
|
||||
self.refresh_lease_owner_prefix,
|
||||
Uuid::new_v4().simple()
|
||||
);
|
||||
let lease_acquired = self
|
||||
.refresh_leases
|
||||
.try_acquire(
|
||||
account_id,
|
||||
&lease_owner,
|
||||
lease_now + Duration::seconds(ADMIN_REFRESH_LEASE_TTL_SECONDS),
|
||||
lease_now,
|
||||
)
|
||||
.await
|
||||
.map_err(|_| AdminAccountError::RefreshLease)?;
|
||||
if !lease_acquired {
|
||||
return Ok(skipped("refresh already in progress"));
|
||||
}
|
||||
|
||||
match self
|
||||
.refresh_tokens_from_refresh_token(refresh_token.expose_secret())
|
||||
.await
|
||||
{
|
||||
let refresh_result = self.refresh_tokens_from_refresh_token(&refresh_token).await;
|
||||
let release_result = self.refresh_leases.release(account_id, &lease_owner).await;
|
||||
if release_result.is_err() {
|
||||
return Err(AdminAccountError::RefreshLease);
|
||||
}
|
||||
|
||||
match refresh_result {
|
||||
Ok(refreshed) => {
|
||||
let updated = self
|
||||
.store
|
||||
|
||||
@@ -12,13 +12,14 @@ mod testing;
|
||||
use std::sync::Arc as StdArc;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::{
|
||||
upstream::accounts::{
|
||||
cookies::SqliteCookieStore,
|
||||
pool::RuntimeAccountPoolService,
|
||||
store::SqliteAccountStore,
|
||||
token_refresh::{RuntimeRefreshPolicy, TokenRefresher},
|
||||
token_refresh::{RefreshLeaseStore, RuntimeRefreshPolicy, TokenRefresher},
|
||||
},
|
||||
upstream::models::service::ModelService,
|
||||
upstream::transport::CodexBackendClient,
|
||||
@@ -38,6 +39,8 @@ pub struct AdminAccountService {
|
||||
pub(crate) models: StdArc<ModelService>,
|
||||
pub(crate) account_pool: StdArc<RuntimeAccountPoolService>,
|
||||
pub(crate) token_refresher: StdArc<dyn TokenRefresher>,
|
||||
pub(crate) refresh_leases: RefreshLeaseStore,
|
||||
pub(crate) refresh_lease_owner_prefix: String,
|
||||
pub(crate) oauth: oauth::AccountOAuthService,
|
||||
pub(crate) refresh_policy: RuntimeRefreshPolicy,
|
||||
pub(crate) installation_id: Option<String>,
|
||||
@@ -50,6 +53,7 @@ pub struct AdminAccountServiceParts {
|
||||
pub models: StdArc<ModelService>,
|
||||
pub account_pool: StdArc<RuntimeAccountPoolService>,
|
||||
pub token_refresher: StdArc<dyn TokenRefresher>,
|
||||
pub refresh_leases: RefreshLeaseStore,
|
||||
pub oauth: oauth::AccountOAuthService,
|
||||
pub refresh_policy: RuntimeRefreshPolicy,
|
||||
pub installation_id: Option<String>,
|
||||
@@ -64,6 +68,11 @@ impl AdminAccountService {
|
||||
models: parts.models,
|
||||
account_pool: parts.account_pool,
|
||||
token_refresher: parts.token_refresher,
|
||||
refresh_leases: parts.refresh_leases,
|
||||
refresh_lease_owner_prefix: format!(
|
||||
"admin-account-refresh:{}",
|
||||
Uuid::new_v4().simple()
|
||||
),
|
||||
oauth: parts.oauth,
|
||||
refresh_policy: parts.refresh_policy,
|
||||
installation_id: parts.installation_id,
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
//! 管理端 Dashboard 聚合视图。
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
use axum::{
|
||||
extract::{Query, State},
|
||||
@@ -15,7 +15,7 @@ use crate::{
|
||||
admin::{
|
||||
auth::session::AdminAuth,
|
||||
monitoring::{
|
||||
account_usage_service::{AdminUsageSummary, AdminUsageTimeBucketRecord},
|
||||
account_usage_service::AdminUsageTimeBucketRecord,
|
||||
diagnostics::{
|
||||
fingerprint_diagnostics, AccountCapacityDiagnostics, AccountPoolDiagnostics,
|
||||
},
|
||||
@@ -27,7 +27,7 @@ use crate::{
|
||||
UsageRecordCostDetailsData, UsageRecordTokenDetailsData,
|
||||
},
|
||||
usage_record_service::AdminUsageRecordFilter,
|
||||
usage_record_store::UsageRecordAccountUsage,
|
||||
usage_record_store::{UsageRecordAccountUsage, UsageRecordSummary},
|
||||
},
|
||||
response::{AdminEnvelope, AdminError, AdminResponse},
|
||||
},
|
||||
@@ -292,10 +292,10 @@ pub(crate) async fn dashboard_summary(
|
||||
let today_filter = today_usage_record_filter(now);
|
||||
let usage_summary = state
|
||||
.services
|
||||
.usage
|
||||
.summary()
|
||||
.usage_records
|
||||
.summary(AdminUsageRecordFilter::default())
|
||||
.await
|
||||
.unwrap_or_else(|_| empty_usage_summary());
|
||||
.unwrap_or_else(|_| UsageRecordSummary::default());
|
||||
let account_usage_records = state
|
||||
.services
|
||||
.usage_records
|
||||
@@ -311,7 +311,17 @@ pub(crate) async fn dashboard_summary(
|
||||
.unwrap_or_default();
|
||||
|
||||
let time_buckets = dashboard_time_buckets(&state, now).await;
|
||||
let pool_summary = account_pool_summary(&accounts);
|
||||
let account_ids = accounts
|
||||
.iter()
|
||||
.map(|account| account.id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let refreshing_account_ids = state
|
||||
.services
|
||||
.token_refresh
|
||||
.refreshing_account_ids(&account_ids, now)
|
||||
.await
|
||||
.map_err(|error| AdminError::internal(error.to_string()))?;
|
||||
let pool_summary = account_pool_summary(&accounts, &refreshing_account_ids);
|
||||
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;
|
||||
@@ -375,7 +385,7 @@ pub(crate) async fn dashboard_trend(
|
||||
fn dashboard_cards(
|
||||
accounts: &[Account],
|
||||
buckets: &[AdminUsageTimeBucketRecord],
|
||||
usage_summary: &AdminUsageSummary,
|
||||
usage_summary: &UsageRecordSummary,
|
||||
) -> DashboardCardsData {
|
||||
dashboard_cards_at(accounts, buckets, usage_summary, Utc::now())
|
||||
}
|
||||
@@ -383,7 +393,7 @@ fn dashboard_cards(
|
||||
fn dashboard_cards_at(
|
||||
accounts: &[Account],
|
||||
buckets: &[AdminUsageTimeBucketRecord],
|
||||
usage_summary: &AdminUsageSummary,
|
||||
usage_summary: &UsageRecordSummary,
|
||||
now: DateTime<Utc>,
|
||||
) -> DashboardCardsData {
|
||||
let today_start = china_day_start(now);
|
||||
@@ -391,9 +401,9 @@ fn dashboard_cards_at(
|
||||
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(usage_summary.request_count);
|
||||
let total_tokens = nonnegative_i64_to_u64(usage_summary.total_tokens);
|
||||
let total_cached_tokens = nonnegative_i64_to_u64(usage_summary.cached_tokens);
|
||||
let total_requests = usage_summary.total_requests;
|
||||
let total_tokens = usage_summary.total_tokens;
|
||||
let total_cached_tokens = usage_summary.cached_tokens;
|
||||
let total_hit_rate = summary_cache_hit_rate(usage_summary);
|
||||
|
||||
let total_accounts = accounts.len() as u64;
|
||||
@@ -449,27 +459,10 @@ fn dashboard_cards_at(
|
||||
}
|
||||
}
|
||||
|
||||
fn empty_usage_summary() -> AdminUsageSummary {
|
||||
AdminUsageSummary {
|
||||
account_count: 0,
|
||||
request_count: 0,
|
||||
empty_response_count: 0,
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
reasoning_tokens: 0,
|
||||
total_tokens: 0,
|
||||
image_input_tokens: 0,
|
||||
image_output_tokens: 0,
|
||||
image_request_count: 0,
|
||||
image_request_failed_count: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn summary_cache_hit_rate(summary: &AdminUsageSummary) -> Option<f64> {
|
||||
let input_tokens = nonnegative_i64_to_u64(summary.input_tokens);
|
||||
fn summary_cache_hit_rate(summary: &UsageRecordSummary) -> Option<f64> {
|
||||
let input_tokens = summary.input_tokens;
|
||||
if input_tokens > 0 {
|
||||
Some(nonnegative_i64_to_u64(summary.cached_tokens) as f64 / input_tokens as f64)
|
||||
Some(summary.cached_tokens as f64 / input_tokens as f64)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
@@ -822,17 +815,24 @@ fn first_token_bucket_latency(record: &AdminUsageTimeBucketRecord) -> Option<u64
|
||||
sum.checked_div(count).filter(|latency| *latency > 0)
|
||||
}
|
||||
|
||||
fn account_pool_summary(accounts: &[Account]) -> AccountPoolDiagnostics {
|
||||
fn account_pool_summary(
|
||||
accounts: &[Account],
|
||||
refreshing_account_ids: &HashSet<String>,
|
||||
) -> AccountPoolDiagnostics {
|
||||
let mut summary = AccountPoolDiagnostics {
|
||||
total: accounts.len(),
|
||||
..AccountPoolDiagnostics::default()
|
||||
};
|
||||
for account in accounts {
|
||||
match account.status {
|
||||
AccountStatus::Active | AccountStatus::Expired
|
||||
if refreshing_account_ids.contains(&account.id) =>
|
||||
{
|
||||
summary.refreshing += 1;
|
||||
}
|
||||
AccountStatus::Active => summary.active += 1,
|
||||
AccountStatus::Expired => summary.expired += 1,
|
||||
AccountStatus::QuotaExhausted => summary.quota_exhausted += 1,
|
||||
AccountStatus::Refreshing => summary.refreshing += 1,
|
||||
AccountStatus::Disabled => summary.disabled += 1,
|
||||
AccountStatus::Banned => summary.banned += 1,
|
||||
}
|
||||
|
||||
@@ -5,6 +5,9 @@ pub mod account_usage_store;
|
||||
pub(crate) mod billing;
|
||||
pub mod dashboard;
|
||||
pub mod diagnostics;
|
||||
pub mod ops_error_model;
|
||||
pub mod ops_error_service;
|
||||
pub mod ops_error_store;
|
||||
pub mod usage_record_model;
|
||||
pub mod usage_record_routes;
|
||||
pub mod usage_record_service;
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
//! 运维错误事件模型。
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
use uuid::Uuid;
|
||||
|
||||
/// 结构化运维错误事件。
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct OpsErrorLog {
|
||||
/// 事件 ID。
|
||||
pub id: String,
|
||||
/// 请求 ID。
|
||||
pub request_id: Option<String>,
|
||||
/// 事件类别。
|
||||
pub kind: String,
|
||||
/// 账号 ID。预账号调度失败时为空。
|
||||
pub account_id: Option<String>,
|
||||
/// HTTP 路由。
|
||||
pub route: Option<String>,
|
||||
/// 模型名。
|
||||
pub model: Option<String>,
|
||||
/// 事件状态码。
|
||||
pub status_code: Option<i64>,
|
||||
/// 客户端看到的状态码。
|
||||
pub client_status_code: Option<i64>,
|
||||
/// 上游 HTTP 状态码。
|
||||
pub upstream_status_code: Option<i64>,
|
||||
/// 上游传输方式,例如 http_sse 或 websocket。
|
||||
pub transport: Option<String>,
|
||||
/// 同一请求内的上游尝试序号。
|
||||
pub attempt_index: Option<i64>,
|
||||
/// 失败分类。
|
||||
pub failure_class: Option<String>,
|
||||
/// 上游响应 ID。
|
||||
pub response_id: Option<String>,
|
||||
/// 上游请求 ID。
|
||||
pub upstream_request_id: Option<String>,
|
||||
/// 请求耗时毫秒。
|
||||
pub latency_ms: Option<i64>,
|
||||
/// 展示消息。
|
||||
pub message: String,
|
||||
/// 结构化元数据。
|
||||
pub metadata: Value,
|
||||
/// 创建时间。
|
||||
#[serde(serialize_with = "crate::infra::time::serialize_china_rfc3339")]
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
impl OpsErrorLog {
|
||||
/// 构造运维错误事件。
|
||||
pub fn new(kind: impl Into<String>, message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
id: format!("ops_error_{}", Uuid::new_v4().simple()),
|
||||
request_id: None,
|
||||
kind: kind.into(),
|
||||
account_id: None,
|
||||
route: None,
|
||||
model: None,
|
||||
status_code: None,
|
||||
client_status_code: None,
|
||||
upstream_status_code: None,
|
||||
transport: None,
|
||||
attempt_index: None,
|
||||
failure_class: None,
|
||||
response_id: None,
|
||||
upstream_request_id: None,
|
||||
latency_ms: None,
|
||||
message: message.into(),
|
||||
metadata: Value::Object(Map::default()),
|
||||
created_at: Utc::now(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
//! 管理端运维错误事件服务。
|
||||
|
||||
use chrono::Utc;
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::admin::monitoring::{
|
||||
ops_error_model::OpsErrorLog, ops_error_store::SqliteOpsErrorLogStore,
|
||||
};
|
||||
|
||||
/// 管理端运维错误事件错误。
|
||||
#[derive(Debug, Error)]
|
||||
pub enum AdminOpsErrorLogError {
|
||||
/// 写入失败。
|
||||
#[error("failed to append ops error log")]
|
||||
Append,
|
||||
/// 保留期清理失败。
|
||||
#[error("failed to trim expired ops error logs")]
|
||||
Retention,
|
||||
}
|
||||
|
||||
/// 管理端运维错误事件服务。
|
||||
#[derive(Clone)]
|
||||
pub struct AdminOpsErrorLogService {
|
||||
store: SqliteOpsErrorLogStore,
|
||||
capture_body: bool,
|
||||
}
|
||||
|
||||
impl AdminOpsErrorLogService {
|
||||
/// 构造管理端运维错误事件服务。
|
||||
pub fn new(store: SqliteOpsErrorLogStore, capture_body: bool) -> Self {
|
||||
Self {
|
||||
store,
|
||||
capture_body,
|
||||
}
|
||||
}
|
||||
|
||||
/// 记录运维错误事件。
|
||||
pub async fn record(&self, mut event: OpsErrorLog) -> Result<(), AdminOpsErrorLogError> {
|
||||
apply_capture_body_policy(&mut event, self.capture_body);
|
||||
self.store
|
||||
.append(&event)
|
||||
.await
|
||||
.map_err(|_| AdminOpsErrorLogError::Append)?;
|
||||
self.store
|
||||
.trim_to_retention(Utc::now())
|
||||
.await
|
||||
.map_err(|_| AdminOpsErrorLogError::Retention)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_capture_body_policy(event: &mut OpsErrorLog, capture_body: bool) {
|
||||
if capture_body {
|
||||
return;
|
||||
}
|
||||
let Some(metadata) = event.metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
for key in [
|
||||
"body",
|
||||
"rawBody",
|
||||
"requestBody",
|
||||
"responseBody",
|
||||
"upstreamBody",
|
||||
] {
|
||||
metadata.remove(key);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,261 @@
|
||||
//! 运维错误事件存储实现(SQLite)。
|
||||
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use serde_json::Value;
|
||||
use sqlx::SqlitePool;
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::{
|
||||
admin::monitoring::{
|
||||
ops_error_model::OpsErrorLog,
|
||||
usage_record_model::{metadata_i64, metadata_service_tier, metadata_string},
|
||||
usage_record_store::USAGE_RECORD_RETENTION_DAYS,
|
||||
},
|
||||
infra::time::china_quarter_hour_start,
|
||||
};
|
||||
|
||||
/// SQLite 运维错误事件错误。
|
||||
#[derive(Debug, Error)]
|
||||
pub enum SqliteOpsErrorLogStoreError {
|
||||
/// 数据库错误。
|
||||
#[error("sqlite ops error log database error: {0}")]
|
||||
Database(#[from] sqlx::Error),
|
||||
}
|
||||
|
||||
/// SQLite 运维错误事件结果。
|
||||
pub type SqliteOpsErrorLogStoreResult<T> = Result<T, SqliteOpsErrorLogStoreError>;
|
||||
|
||||
/// SQLite 运维错误事件存储。
|
||||
#[derive(Clone)]
|
||||
pub struct SqliteOpsErrorLogStore {
|
||||
pool: SqlitePool,
|
||||
}
|
||||
|
||||
impl SqliteOpsErrorLogStore {
|
||||
/// 构造存储。
|
||||
pub fn new(pool: SqlitePool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
/// 写入运维错误事件。
|
||||
pub async fn append(&self, event: &OpsErrorLog) -> SqliteOpsErrorLogStoreResult<()> {
|
||||
append_event(&self.pool, event).await
|
||||
}
|
||||
|
||||
/// 按保留期清理运维错误事件。
|
||||
pub async fn trim_to_retention(&self, now: DateTime<Utc>) -> SqliteOpsErrorLogStoreResult<u64> {
|
||||
trim_to_retention(&self.pool, now).await
|
||||
}
|
||||
|
||||
/// 清空运维错误事件。
|
||||
pub async fn clear(&self) -> SqliteOpsErrorLogStoreResult<u64> {
|
||||
let result = sqlx::query("delete from ops_error_logs")
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected())
|
||||
}
|
||||
|
||||
/// 读取错误事件数量。
|
||||
pub async fn count(&self) -> SqliteOpsErrorLogStoreResult<u64> {
|
||||
let (total,): (i64,) = sqlx::query_as("select count(*) from ops_error_logs")
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
Ok(total.max(0) as u64)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct OpsErrorSummaryFields {
|
||||
client_status_code: Option<i64>,
|
||||
upstream_status_code: Option<i64>,
|
||||
transport: Option<String>,
|
||||
attempt_index: Option<i64>,
|
||||
failure_class: Option<String>,
|
||||
response_id: Option<String>,
|
||||
upstream_request_id: Option<String>,
|
||||
}
|
||||
|
||||
fn ops_error_summary_fields(event: &OpsErrorLog) -> OpsErrorSummaryFields {
|
||||
OpsErrorSummaryFields {
|
||||
client_status_code: event.client_status_code.or_else(|| {
|
||||
metadata_i64(
|
||||
&event.metadata,
|
||||
&["clientStatusCode", "client_status_code", "clientStatus"],
|
||||
)
|
||||
}),
|
||||
upstream_status_code: event.upstream_status_code.or_else(|| {
|
||||
metadata_i64(
|
||||
&event.metadata,
|
||||
&[
|
||||
"upstreamStatus",
|
||||
"upstreamStatusCode",
|
||||
"upstream_status_code",
|
||||
],
|
||||
)
|
||||
}),
|
||||
transport: event
|
||||
.transport
|
||||
.clone()
|
||||
.or_else(|| metadata_string(&event.metadata, &["transport"])),
|
||||
attempt_index: event
|
||||
.attempt_index
|
||||
.or_else(|| metadata_i64(&event.metadata, &["attemptIndex", "attempt_index"])),
|
||||
failure_class: event
|
||||
.failure_class
|
||||
.clone()
|
||||
.or_else(|| metadata_string(&event.metadata, &["failureClass", "failure_class"])),
|
||||
response_id: event
|
||||
.response_id
|
||||
.clone()
|
||||
.or_else(|| metadata_string(&event.metadata, &["responseId", "response_id"])),
|
||||
upstream_request_id: event.upstream_request_id.clone().or_else(|| {
|
||||
metadata_string(
|
||||
&event.metadata,
|
||||
&[
|
||||
"upstreamRequestId",
|
||||
"upstream_request_id",
|
||||
"openaiRequestId",
|
||||
],
|
||||
)
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
async fn append_event(pool: &SqlitePool, event: &OpsErrorLog) -> SqliteOpsErrorLogStoreResult<()> {
|
||||
let summary = ops_error_summary_fields(event);
|
||||
sqlx::query(
|
||||
"insert into ops_error_logs (id, request_id, kind, account_id, route, model, status_code, client_status_code, upstream_status_code, transport, attempt_index, failure_class, response_id, upstream_request_id, latency_ms, message, metadata_json, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(&event.id)
|
||||
.bind(&event.request_id)
|
||||
.bind(&event.kind)
|
||||
.bind(&event.account_id)
|
||||
.bind(&event.route)
|
||||
.bind(&event.model)
|
||||
.bind(event.status_code)
|
||||
.bind(summary.client_status_code)
|
||||
.bind(summary.upstream_status_code)
|
||||
.bind(summary.transport)
|
||||
.bind(summary.attempt_index)
|
||||
.bind(summary.failure_class)
|
||||
.bind(summary.response_id)
|
||||
.bind(summary.upstream_request_id)
|
||||
.bind(event.latency_ms)
|
||||
.bind(&event.message)
|
||||
.bind(event.metadata.to_string())
|
||||
.bind(event.created_at.to_rfc3339())
|
||||
.execute(pool)
|
||||
.await?;
|
||||
append_error_time_bucket(pool, event).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn append_error_time_bucket(
|
||||
pool: &SqlitePool,
|
||||
event: &OpsErrorLog,
|
||||
) -> SqliteOpsErrorLogStoreResult<()> {
|
||||
let bucket_start = china_quarter_hour_start(event.created_at);
|
||||
let input_tokens = metadata_usage_i64(&event.metadata, "inputTokens");
|
||||
let output_tokens = metadata_usage_i64(&event.metadata, "outputTokens");
|
||||
let cached_tokens = metadata_usage_i64(&event.metadata, "cachedTokens");
|
||||
let first_token_latency = metadata_first_token_latency(&event.metadata).unwrap_or(0);
|
||||
let first_token_latency_count = i64::from(first_token_latency > 0);
|
||||
let latency = event.latency_ms.filter(|value| *value > 0).unwrap_or(0);
|
||||
let latency_count = i64::from(latency > 0);
|
||||
let now = Utc::now().to_rfc3339();
|
||||
|
||||
sqlx::query(
|
||||
r"
|
||||
insert into usage_time_buckets (
|
||||
bucket_start,
|
||||
account_id,
|
||||
model,
|
||||
service_tier,
|
||||
request_count,
|
||||
error_count,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cached_tokens,
|
||||
first_token_latency_sum,
|
||||
first_token_latency_count,
|
||||
latency_sum,
|
||||
latency_count,
|
||||
max_latency_ms,
|
||||
min_latency_ms,
|
||||
updated_at
|
||||
) values (?, ?, ?, ?, 1, 1, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
on conflict(bucket_start, account_id, model, service_tier) do update set
|
||||
request_count = request_count + excluded.request_count,
|
||||
error_count = error_count + excluded.error_count,
|
||||
input_tokens = input_tokens + excluded.input_tokens,
|
||||
output_tokens = output_tokens + excluded.output_tokens,
|
||||
cached_tokens = cached_tokens + excluded.cached_tokens,
|
||||
first_token_latency_sum = first_token_latency_sum + excluded.first_token_latency_sum,
|
||||
first_token_latency_count = first_token_latency_count + excluded.first_token_latency_count,
|
||||
latency_sum = latency_sum + excluded.latency_sum,
|
||||
latency_count = latency_count + excluded.latency_count,
|
||||
max_latency_ms = max(max_latency_ms, excluded.max_latency_ms),
|
||||
min_latency_ms = case
|
||||
when min_latency_ms = 0 then excluded.min_latency_ms
|
||||
when excluded.min_latency_ms = 0 then min_latency_ms
|
||||
else min(min_latency_ms, excluded.min_latency_ms)
|
||||
end,
|
||||
updated_at = excluded.updated_at",
|
||||
)
|
||||
.bind(bucket_start.to_rfc3339())
|
||||
.bind(event.account_id.as_deref().unwrap_or_default())
|
||||
.bind(event.model.as_deref().unwrap_or_default())
|
||||
.bind(metadata_service_tier(&event.metadata).unwrap_or_default())
|
||||
.bind(input_tokens)
|
||||
.bind(output_tokens)
|
||||
.bind(cached_tokens)
|
||||
.bind(first_token_latency)
|
||||
.bind(first_token_latency_count)
|
||||
.bind(latency)
|
||||
.bind(latency_count)
|
||||
.bind(latency)
|
||||
.bind(latency)
|
||||
.bind(now)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn metadata_usage_i64(metadata: &Value, field: &str) -> i64 {
|
||||
metadata
|
||||
.get("usage")
|
||||
.and_then(|usage| usage.get(field))
|
||||
.or_else(|| metadata.get(field))
|
||||
.and_then(Value::as_i64)
|
||||
.filter(|value| *value > 0)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn metadata_first_token_latency(metadata: &Value) -> Option<i64> {
|
||||
[
|
||||
"firstTokenMs",
|
||||
"first_token_ms",
|
||||
"firstTokenLatencyMs",
|
||||
"first_token_latency_ms",
|
||||
]
|
||||
.into_iter()
|
||||
.find_map(|field| {
|
||||
metadata
|
||||
.get(field)
|
||||
.or_else(|| metadata.get("usage").and_then(|usage| usage.get(field)))
|
||||
.and_then(Value::as_i64)
|
||||
.filter(|value| *value > 0)
|
||||
})
|
||||
}
|
||||
|
||||
async fn trim_to_retention(
|
||||
pool: &SqlitePool,
|
||||
now: DateTime<Utc>,
|
||||
) -> SqliteOpsErrorLogStoreResult<u64> {
|
||||
let cutoff = now - Duration::days(USAGE_RECORD_RETENTION_DAYS);
|
||||
let result = sqlx::query("delete from ops_error_logs where created_at < ?")
|
||||
.bind(cutoff.to_rfc3339())
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected())
|
||||
}
|
||||
@@ -265,7 +265,19 @@ impl AdminUsageRecordService {
|
||||
/// 记录使用记录。
|
||||
pub async fn record(&self, event: UsageRecord) -> Result<(), AdminUsageRecordError> {
|
||||
let settings = *self.settings.read().await;
|
||||
if !settings.enabled && event.level != UsageRecordLevel::Error {
|
||||
if !settings.enabled {
|
||||
return Ok(());
|
||||
}
|
||||
if !is_usage_log_event(&event) {
|
||||
tracing::warn!(
|
||||
usage_record_id = %event.id,
|
||||
request_id = event.request_id.as_deref().unwrap_or(""),
|
||||
account_id = event.account_id.as_deref().unwrap_or(""),
|
||||
model = event.model.as_deref().unwrap_or(""),
|
||||
level = ?event.level,
|
||||
status_code = ?event.status_code,
|
||||
"skipped non-usage event in usage record service"
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
self.append_with_settings(event, settings).await
|
||||
@@ -289,6 +301,21 @@ impl AdminUsageRecordService {
|
||||
}
|
||||
}
|
||||
|
||||
fn is_usage_log_event(event: &UsageRecord) -> bool {
|
||||
event.level != UsageRecordLevel::Error
|
||||
&& event.status_code.is_none_or(|status| status < 400)
|
||||
&& event
|
||||
.account_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.is_some_and(|account_id| !account_id.is_empty())
|
||||
&& event
|
||||
.model
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.is_some_and(|model| !model.is_empty())
|
||||
}
|
||||
|
||||
fn apply_capture_body_policy(event: &mut UsageRecord, capture_body: bool) {
|
||||
if capture_body {
|
||||
return;
|
||||
|
||||
@@ -100,6 +100,11 @@ impl SqliteUsageRecordStore {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
/// 暴露底层连接池,供运行时组合层复用。
|
||||
pub fn pool(&self) -> &SqlitePool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
/// 写入使用记录。
|
||||
pub async fn append(&self, event: &UsageRecord) -> SqliteUsageRecordStoreResult<()> {
|
||||
append_event(&self.pool, event).await
|
||||
|
||||
@@ -110,6 +110,8 @@ struct SystemUpdateEvent {
|
||||
level: UpdateLogLevel,
|
||||
step: Option<String>,
|
||||
message: String,
|
||||
#[serde(skip_serializing_if = "is_false")]
|
||||
terminal: bool,
|
||||
at: String,
|
||||
}
|
||||
|
||||
@@ -157,18 +159,23 @@ pub(crate) async fn update_event_stream(
|
||||
let receiver = update_event_sender().subscribe();
|
||||
let shutdown = shutdown_subscription();
|
||||
let stream = futures::stream::unfold(
|
||||
(receiver, shutdown),
|
||||
|(mut receiver, mut shutdown)| async move {
|
||||
(receiver, shutdown, false),
|
||||
|(mut receiver, mut shutdown, close_after_send)| async move {
|
||||
if close_after_send {
|
||||
return None;
|
||||
}
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = shutdown.recv() => return None,
|
||||
received = receiver.recv() => {
|
||||
match received {
|
||||
Ok(message) => {
|
||||
let close_after_send = message.terminal;
|
||||
let id = message.id.clone();
|
||||
let data = serde_json::to_string(&message).unwrap_or_else(|_| "{}".to_string());
|
||||
let event = Event::default().event("update").id(id).data(data);
|
||||
return Some((Ok(event), (receiver, shutdown)));
|
||||
return Some((Ok(event), (receiver, shutdown, close_after_send)));
|
||||
}
|
||||
Err(broadcast::error::RecvError::Lagged(_)) => continue,
|
||||
Err(broadcast::error::RecvError::Closed) => return None,
|
||||
@@ -192,7 +199,7 @@ pub(crate) async fn perform_update(
|
||||
.map_err(|_| AdminError::conflict("System update already running"))?;
|
||||
let config = SystemUpdateConfig::from_env();
|
||||
if let Some(reason) = config.update_support_error() {
|
||||
emit_update_event(
|
||||
emit_terminal_update_event(
|
||||
UpdateLogLevel::Error,
|
||||
None,
|
||||
Some("preflight"),
|
||||
@@ -217,7 +224,7 @@ pub(crate) async fn perform_update(
|
||||
)
|
||||
.await
|
||||
.map_err(|error| {
|
||||
emit_update_event(UpdateLogLevel::Error, None, Some("release"), error.clone());
|
||||
emit_terminal_update_event(UpdateLogLevel::Error, None, Some("release"), error.clone());
|
||||
AdminError::bad_gateway(error)
|
||||
})?;
|
||||
let info = update_info_from_release(&config, release.clone());
|
||||
@@ -226,7 +233,7 @@ pub(crate) async fn perform_update(
|
||||
"远端最新版本已变更为 v{},请重新检查并确认",
|
||||
info.latest_version
|
||||
);
|
||||
emit_update_event(
|
||||
emit_terminal_update_event(
|
||||
UpdateLogLevel::Warning,
|
||||
None,
|
||||
Some("release"),
|
||||
@@ -235,7 +242,7 @@ pub(crate) async fn perform_update(
|
||||
return Err(AdminError::conflict(message));
|
||||
}
|
||||
if !info.has_update {
|
||||
emit_update_event(
|
||||
emit_terminal_update_event(
|
||||
UpdateLogLevel::Warning,
|
||||
None,
|
||||
Some("release"),
|
||||
@@ -261,14 +268,14 @@ pub(crate) async fn perform_update(
|
||||
);
|
||||
let result = perform_release_update(&config, &release, &target_version, &operation_id).await;
|
||||
if let Err(error) = &result {
|
||||
emit_update_event(
|
||||
emit_terminal_update_event(
|
||||
UpdateLogLevel::Error,
|
||||
Some(&operation_id),
|
||||
Some("failed"),
|
||||
error.to_string(),
|
||||
);
|
||||
} else {
|
||||
emit_update_event(
|
||||
emit_terminal_update_event(
|
||||
UpdateLogLevel::Success,
|
||||
Some(&operation_id),
|
||||
Some("done"),
|
||||
@@ -732,6 +739,25 @@ fn emit_update_event(
|
||||
operation_id: Option<&str>,
|
||||
step: Option<&str>,
|
||||
message: impl Into<String>,
|
||||
) {
|
||||
emit_update_event_with_terminal(level, operation_id, step, message, false);
|
||||
}
|
||||
|
||||
fn emit_terminal_update_event(
|
||||
level: UpdateLogLevel,
|
||||
operation_id: Option<&str>,
|
||||
step: Option<&str>,
|
||||
message: impl Into<String>,
|
||||
) {
|
||||
emit_update_event_with_terminal(level, operation_id, step, message, true);
|
||||
}
|
||||
|
||||
fn emit_update_event_with_terminal(
|
||||
level: UpdateLogLevel,
|
||||
operation_id: Option<&str>,
|
||||
step: Option<&str>,
|
||||
message: impl Into<String>,
|
||||
terminal: bool,
|
||||
) {
|
||||
let now = Utc::now();
|
||||
let event = SystemUpdateEvent {
|
||||
@@ -744,11 +770,16 @@ fn emit_update_event(
|
||||
level,
|
||||
step: step.map(ToString::to_string),
|
||||
message: message.into(),
|
||||
terminal,
|
||||
at: now.to_rfc3339(),
|
||||
};
|
||||
let _ = update_event_sender().send(event);
|
||||
}
|
||||
|
||||
fn is_false(value: &bool) -> bool {
|
||||
!*value
|
||||
}
|
||||
|
||||
fn format_bytes(bytes: u64) -> String {
|
||||
const KIB: u64 = 1024;
|
||||
const MIB: u64 = KIB * 1024;
|
||||
|
||||
@@ -57,9 +57,6 @@ pub struct ServerConfig {
|
||||
pub host: String,
|
||||
/// 监听端口。
|
||||
pub port: u16,
|
||||
/// 可信反向代理 IP 或 CIDR。非空时只有这些 peer 的转发 IP 头会被采信。
|
||||
#[serde(default)]
|
||||
pub trusted_proxies: Vec<String>,
|
||||
}
|
||||
|
||||
/// 上游 API 地址配置。
|
||||
|
||||
@@ -4,7 +4,7 @@ use std::net::{IpAddr, SocketAddr};
|
||||
use std::sync::LazyLock;
|
||||
|
||||
use axum::{
|
||||
extract::{connect_info::ConnectInfo, Request, State},
|
||||
extract::{connect_info::ConnectInfo, Request},
|
||||
http::{HeaderMap, HeaderValue},
|
||||
middleware::Next,
|
||||
response::Response,
|
||||
@@ -19,7 +19,7 @@ const CF_CONNECTING_IP_HEADER: &str = "cf-connecting-ip";
|
||||
const REAL_IP_HEADER: &str = "x-real-ip";
|
||||
const X_FORWARDED_FOR_HEADER: &str = "x-forwarded-for";
|
||||
|
||||
/// 私有 / 回环 / 唯一本地地址网段,用于在自动模式下从 `X-Forwarded-For` 中跳过内网跳板 IP。
|
||||
/// 私有 / 回环 / 唯一本地地址网段,用于从 `X-Forwarded-For` 中跳过内网跳板 IP。
|
||||
static PRIVATE_NETWORKS: LazyLock<Vec<IpNet>> = LazyLock::new(|| {
|
||||
[
|
||||
"10.0.0.0/8",
|
||||
@@ -38,40 +38,6 @@ fn is_private_ip(ip: IpAddr) -> bool {
|
||||
PRIVATE_NETWORKS.iter().any(|network| network.contains(&ip))
|
||||
}
|
||||
|
||||
/// 可信反向代理配置。
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct TrustedProxyConfig {
|
||||
networks: Vec<IpNet>,
|
||||
}
|
||||
|
||||
impl TrustedProxyConfig {
|
||||
/// 从 IP 或 CIDR 字符串构造可信代理配置。
|
||||
pub fn from_entries(entries: &[String]) -> Result<Self, TrustedProxyConfigError> {
|
||||
let networks = entries
|
||||
.iter()
|
||||
.map(|entry| parse_trusted_proxy_entry(entry))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
Ok(Self { networks })
|
||||
}
|
||||
|
||||
fn is_trusted(&self, ip: IpAddr) -> bool {
|
||||
self.networks.iter().any(|network| network.contains(&ip))
|
||||
}
|
||||
|
||||
/// 是否未配置任何可信代理(决定是否启用自动真实 IP 探测)。
|
||||
fn is_empty(&self) -> bool {
|
||||
self.networks.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
/// 可信代理配置错误。
|
||||
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
|
||||
pub enum TrustedProxyConfigError {
|
||||
/// 配置项不是 IP 地址或 CIDR。
|
||||
#[error("invalid trusted proxy entry `{entry}`; expected an IP address or CIDR network")]
|
||||
InvalidEntry { entry: String },
|
||||
}
|
||||
|
||||
/// 请求 ID。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RequestId(String);
|
||||
@@ -97,7 +63,7 @@ impl RequestId {
|
||||
}
|
||||
}
|
||||
|
||||
/// 由连接来源和可信代理链解析出的客户端 IP。
|
||||
/// 由连接来源和转发头解析出的客户端 IP。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ClientIp(String);
|
||||
|
||||
@@ -118,17 +84,7 @@ fn valid_request_id(value: &str) -> bool {
|
||||
|
||||
/// 为请求附加请求 ID,并在响应头中回写。
|
||||
pub async fn attach_request_id(mut request: Request, next: Next) -> Response {
|
||||
attach_request_context(&TrustedProxyConfig::default(), &mut request);
|
||||
run_with_request_id(request, next).await
|
||||
}
|
||||
|
||||
/// 为请求附加请求上下文,并按可信代理配置解析真实客户端 IP。
|
||||
pub async fn attach_request_id_with_proxy_config(
|
||||
State(trusted_proxies): State<TrustedProxyConfig>,
|
||||
mut request: Request,
|
||||
next: Next,
|
||||
) -> Response {
|
||||
attach_request_context(&trusted_proxies, &mut request);
|
||||
attach_request_context(&mut request);
|
||||
run_with_request_id(request, next).await
|
||||
}
|
||||
|
||||
@@ -147,17 +103,16 @@ async fn run_with_request_id(mut request: Request, next: Next) -> Response {
|
||||
response
|
||||
}
|
||||
|
||||
fn attach_request_context(trusted_proxies: &TrustedProxyConfig, request: &mut Request) {
|
||||
attach_client_ip_from_connection(trusted_proxies, request);
|
||||
fn attach_request_context(request: &mut Request) {
|
||||
attach_client_ip_from_connection(request);
|
||||
}
|
||||
|
||||
fn attach_client_ip_from_connection(trusted_proxies: &TrustedProxyConfig, request: &mut Request) {
|
||||
fn attach_client_ip_from_connection(request: &mut Request) {
|
||||
let peer_ip = request
|
||||
.extensions()
|
||||
.get::<ConnectInfo<SocketAddr>>()
|
||||
.map(|ConnectInfo(addr)| addr.ip());
|
||||
let client_ip =
|
||||
peer_ip.map(|peer_ip| client_ip_from_peer(peer_ip, request.headers(), trusted_proxies));
|
||||
let client_ip = peer_ip.map(|peer_ip| client_ip_from_peer(peer_ip, request.headers()));
|
||||
|
||||
let Some(client_ip) = client_ip else {
|
||||
return;
|
||||
@@ -165,28 +120,12 @@ fn attach_client_ip_from_connection(trusted_proxies: &TrustedProxyConfig, reques
|
||||
request.extensions_mut().insert(ClientIp(client_ip));
|
||||
}
|
||||
|
||||
fn client_ip_from_peer(
|
||||
peer_ip: IpAddr,
|
||||
headers: &HeaderMap,
|
||||
trusted_proxies: &TrustedProxyConfig,
|
||||
) -> String {
|
||||
if trusted_proxies.is_empty() {
|
||||
// 未配置可信代理:自动模式,按 CF-Connecting-IP → X-Real-IP →
|
||||
// X-Forwarded-For(首个公网 IP)→ peer IP 的优先级解析真实客户端 IP。
|
||||
return auto_forwarded_client_ip(headers).unwrap_or_else(|| peer_ip.to_string());
|
||||
}
|
||||
if trusted_proxies.is_trusted(peer_ip) {
|
||||
// 已配置可信代理且直连来源可信:只采信转发头。
|
||||
trusted_forwarded_client_ip(headers).unwrap_or_else(|| peer_ip.to_string())
|
||||
} else {
|
||||
// 已配置可信代理但直连来源不可信:忽略转发头,使用 peer IP。
|
||||
peer_ip.to_string()
|
||||
}
|
||||
fn client_ip_from_peer(peer_ip: IpAddr, headers: &HeaderMap) -> String {
|
||||
auto_forwarded_client_ip(headers).unwrap_or_else(|| peer_ip.to_string())
|
||||
}
|
||||
|
||||
/// 自动模式下解析转发头中的真实客户端 IP。
|
||||
/// 解析转发头中的真实客户端 IP。
|
||||
///
|
||||
/// 未配置可信代理时使用,对齐常见反代(Cloudflare / Nginx)默认行为:
|
||||
/// `CF-Connecting-IP` → `X-Real-IP` → `X-Forwarded-For` 首个公网 IP。
|
||||
fn auto_forwarded_client_ip(headers: &HeaderMap) -> Option<String> {
|
||||
if let Some(cf_ip) = header_string(headers, CF_CONNECTING_IP_HEADER) {
|
||||
@@ -212,13 +151,6 @@ fn forwarded_for_public_ip(headers: &HeaderMap) -> Option<String> {
|
||||
.map(ToString::to_string)
|
||||
}
|
||||
|
||||
fn trusted_forwarded_client_ip(headers: &HeaderMap) -> Option<String> {
|
||||
header_string(headers, CF_CONNECTING_IP_HEADER)
|
||||
.or_else(|| header_string(headers, REAL_IP_HEADER))
|
||||
.or_else(|| forwarded_for_public_ip(headers))
|
||||
.or_else(|| first_forwarded_for_value(headers))
|
||||
}
|
||||
|
||||
fn header_string(headers: &HeaderMap, name: &str) -> Option<String> {
|
||||
headers
|
||||
.get(name)
|
||||
@@ -237,19 +169,3 @@ fn first_forwarded_for_value(headers: &HeaderMap) -> Option<String> {
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToString::to_string)
|
||||
}
|
||||
|
||||
fn parse_trusted_proxy_entry(entry: &str) -> Result<IpNet, TrustedProxyConfigError> {
|
||||
if let Ok(network) = entry.parse::<IpNet>() {
|
||||
return Ok(network);
|
||||
}
|
||||
entry
|
||||
.parse::<IpAddr>()
|
||||
.map(IpNet::from)
|
||||
.map_err(|_| invalid_trusted_proxy_entry(entry))
|
||||
}
|
||||
|
||||
fn invalid_trusted_proxy_entry(entry: &str) -> TrustedProxyConfigError {
|
||||
TrustedProxyConfigError::InvalidEntry {
|
||||
entry: entry.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,10 +5,7 @@ use std::path::Path;
|
||||
use axum::{extract::State, http::StatusCode, middleware, routing::get, Router};
|
||||
|
||||
use crate::admin;
|
||||
use crate::http::middleware::{
|
||||
request_id::{attach_request_id_with_proxy_config, TrustedProxyConfig},
|
||||
trace::http_trace_layer,
|
||||
};
|
||||
use crate::http::middleware::{request_id::attach_request_id, trace::http_trace_layer};
|
||||
use crate::proxy;
|
||||
use crate::runtime::state::AppState;
|
||||
use crate::web;
|
||||
@@ -18,34 +15,18 @@ pub const DEFAULT_ASSET_DIST_DIR: &str = "web/dist";
|
||||
|
||||
/// 构造整个 HTTP 服务路由。
|
||||
pub fn router() -> Router<AppState> {
|
||||
router_with_trusted_proxies(TrustedProxyConfig::default())
|
||||
}
|
||||
|
||||
/// 使用指定可信代理配置构造整个 HTTP 服务路由。
|
||||
pub fn router_with_trusted_proxies(trusted_proxies: TrustedProxyConfig) -> Router<AppState> {
|
||||
router_with_assets_and_trusted_proxies(DEFAULT_ASSET_DIST_DIR, trusted_proxies)
|
||||
router_with_assets(DEFAULT_ASSET_DIST_DIR)
|
||||
}
|
||||
|
||||
/// 使用指定前端构建产物目录构造整个 HTTP 服务路由。
|
||||
pub fn router_with_assets(dist_dir: impl AsRef<Path>) -> Router<AppState> {
|
||||
router_with_assets_and_trusted_proxies(dist_dir, TrustedProxyConfig::default())
|
||||
}
|
||||
|
||||
/// 使用指定前端构建产物目录和可信代理配置构造整个 HTTP 服务路由。
|
||||
pub fn router_with_assets_and_trusted_proxies(
|
||||
dist_dir: impl AsRef<Path>,
|
||||
trusted_proxies: TrustedProxyConfig,
|
||||
) -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/healthz", get(healthz))
|
||||
.merge(proxy::router::router())
|
||||
.merge(admin::router::router())
|
||||
.fallback_service(web::assets::spa_router(dist_dir))
|
||||
.layer(http_trace_layer())
|
||||
.layer(middleware::from_fn_with_state(
|
||||
trusted_proxies,
|
||||
attach_request_id_with_proxy_config,
|
||||
))
|
||||
.layer(middleware::from_fn(attach_request_id))
|
||||
}
|
||||
|
||||
async fn healthz(State(state): State<AppState>) -> StatusCode {
|
||||
|
||||
@@ -4,7 +4,7 @@ use std::{fs, path::Path, str::FromStr, time::Duration};
|
||||
|
||||
use sqlx::{sqlite::SqliteConnectOptions, SqlitePool};
|
||||
|
||||
const CURRENT_SCHEMA_VERSION: i64 = 2;
|
||||
const CURRENT_SCHEMA_VERSION: i64 = 3;
|
||||
const MIGRATIONS: &[Migration] = &[
|
||||
Migration {
|
||||
version: 1,
|
||||
@@ -16,6 +16,11 @@ const MIGRATIONS: &[Migration] = &[
|
||||
name: "rename_rotation_strategies",
|
||||
sql: include_str!("migrations/0002_rename_rotation_strategies.sql"),
|
||||
},
|
||||
Migration {
|
||||
version: 3,
|
||||
name: "ops_error_logs",
|
||||
sql: include_str!("migrations/0003_ops_error_logs.sql"),
|
||||
},
|
||||
];
|
||||
|
||||
struct Migration {
|
||||
|
||||
@@ -68,7 +68,7 @@ create table accounts (
|
||||
refresh_token text,
|
||||
access_token_expires_at text,
|
||||
next_refresh_at text,
|
||||
status text not null check (status in ('active', 'expired', 'quota_exhausted', 'refreshing', 'disabled', 'banned')),
|
||||
status text not null check (status in ('active', 'expired', 'quota_exhausted', 'disabled', 'banned')),
|
||||
quota_json text,
|
||||
quota_fetched_at text,
|
||||
quota_limit_reached integer not null default 0 check (quota_limit_reached in (0, 1)),
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
create table ops_error_logs (
|
||||
id text primary key,
|
||||
request_id text,
|
||||
kind text not null,
|
||||
account_id text,
|
||||
route text,
|
||||
model text,
|
||||
status_code integer check (status_code is null or (status_code >= 100 and status_code <= 599)),
|
||||
client_status_code integer check (client_status_code is null or (client_status_code >= 100 and client_status_code <= 599)),
|
||||
upstream_status_code integer check (upstream_status_code is null or (upstream_status_code >= 100 and upstream_status_code <= 599)),
|
||||
transport text,
|
||||
attempt_index integer check (attempt_index is null or attempt_index >= 0),
|
||||
failure_class text,
|
||||
response_id text,
|
||||
upstream_request_id text,
|
||||
latency_ms integer check (latency_ms is null or latency_ms >= 0),
|
||||
message text not null,
|
||||
metadata_json text not null,
|
||||
created_at text not null
|
||||
);
|
||||
|
||||
create index idx_ops_error_logs_created_id on ops_error_logs(created_at desc, id desc);
|
||||
create index idx_ops_error_logs_request_id on ops_error_logs(request_id) where request_id is not null;
|
||||
create index idx_ops_error_logs_account on ops_error_logs(account_id, created_at desc) where account_id is not null;
|
||||
create index idx_ops_error_logs_route_created on ops_error_logs(route, created_at desc) where route is not null;
|
||||
create index idx_ops_error_logs_model_created on ops_error_logs(model, created_at desc) where model is not null;
|
||||
create index idx_ops_error_logs_status_created on ops_error_logs(status_code, created_at desc) where status_code is not null;
|
||||
create index idx_ops_error_logs_transport_created on ops_error_logs(transport, created_at desc) where transport is not null;
|
||||
create index idx_ops_error_logs_failure_class on ops_error_logs(failure_class, created_at desc) where failure_class is not null;
|
||||
create index idx_ops_error_logs_response_id on ops_error_logs(response_id) where response_id is not null;
|
||||
create index idx_ops_error_logs_upstream_request_id on ops_error_logs(upstream_request_id) where upstream_request_id is not null;
|
||||
@@ -4,8 +4,9 @@ use serde_json::{json, Map, Value};
|
||||
|
||||
use crate::{
|
||||
admin::monitoring::{
|
||||
usage_record_model::{ResponseUsageRecord, UsageRecord, UsageRecordLevel},
|
||||
usage_record_service::AdminUsageRecordService,
|
||||
ops_error_model::OpsErrorLog,
|
||||
ops_error_service::AdminOpsErrorLogService,
|
||||
usage_record_model::{UsageRecord, UsageRecordLevel},
|
||||
},
|
||||
infra::time::elapsed_millis_i64,
|
||||
proxy::dispatch::{
|
||||
@@ -17,7 +18,7 @@ use crate::{
|
||||
usage_events::{
|
||||
enrich_event_route_metadata, enrich_usage_record_identity,
|
||||
event_kind as response_event_kind, reasoning_effort_from_request,
|
||||
record_dispatch_error_event, record_response_event, DispatchErrorUsageRecord,
|
||||
record_dispatch_error_event, DispatchErrorLogRecord,
|
||||
},
|
||||
},
|
||||
upstream::{
|
||||
@@ -39,7 +40,7 @@ use super::{
|
||||
};
|
||||
|
||||
pub(super) struct ResponseUpstreamErrorEventRecord<'a> {
|
||||
pub(super) usage_records: &'a AdminUsageRecordService,
|
||||
pub(super) ops_errors: &'a AdminOpsErrorLogService,
|
||||
pub(super) request_id: &'a str,
|
||||
pub(super) account_id: &'a str,
|
||||
pub(super) account_email: Option<&'a str>,
|
||||
@@ -55,7 +56,7 @@ pub(super) struct ResponseUpstreamErrorEventRecord<'a> {
|
||||
}
|
||||
|
||||
pub(super) struct ResponseStreamFailureEventRecord<'a> {
|
||||
pub(super) usage_records: &'a AdminUsageRecordService,
|
||||
pub(super) ops_errors: &'a AdminOpsErrorLogService,
|
||||
pub(super) request_id: &'a str,
|
||||
pub(super) account_id: &'a str,
|
||||
pub(super) route: &'a str,
|
||||
@@ -73,7 +74,7 @@ pub(super) struct ResponseStreamFailureEventRecord<'a> {
|
||||
}
|
||||
|
||||
pub(super) struct ResponseDispatchErrorEventRecord<'a> {
|
||||
pub(super) usage_records: &'a AdminUsageRecordService,
|
||||
pub(super) ops_errors: &'a AdminOpsErrorLogService,
|
||||
pub(super) request_id: &'a str,
|
||||
pub(super) account_id: Option<&'a str>,
|
||||
pub(super) route: &'a str,
|
||||
@@ -110,8 +111,8 @@ pub(super) async fn record_response_dispatch_error_event(
|
||||
dispatch_error_upstream_status_code(record.error),
|
||||
);
|
||||
}
|
||||
record_dispatch_error_event(DispatchErrorUsageRecord {
|
||||
usage_records: record.usage_records,
|
||||
record_dispatch_error_event(DispatchErrorLogRecord {
|
||||
ops_errors: record.ops_errors,
|
||||
request_id: record.request_id,
|
||||
account_id: record.account_id,
|
||||
route: record.route,
|
||||
@@ -135,11 +136,7 @@ pub(super) async fn record_response_upstream_error_event(
|
||||
Some(backend_transport_name(record.transport)),
|
||||
);
|
||||
enrich_event_route_metadata(&mut metadata, record.route);
|
||||
let mut event = UsageRecord::new(
|
||||
"v1.response",
|
||||
UsageRecordLevel::Error,
|
||||
"v1 responses upstream request failed",
|
||||
);
|
||||
let mut event = OpsErrorLog::new("v1.response", "v1 responses upstream request failed");
|
||||
event.request_id = Some(record.request_id.to_string());
|
||||
event.account_id = Some(record.account_id.to_string());
|
||||
event.route = Some(record.route.to_string());
|
||||
@@ -167,8 +164,17 @@ pub(super) async fn record_response_upstream_error_event(
|
||||
);
|
||||
}
|
||||
}
|
||||
enrich_usage_record_identity(
|
||||
&mut metadata,
|
||||
Some(record.request.model()),
|
||||
record.model,
|
||||
record.request.client_ip.as_deref(),
|
||||
record.request.client_user_agent.as_deref(),
|
||||
reasoning_effort_from_request(record.request),
|
||||
record.request.service_tier(),
|
||||
);
|
||||
event.metadata = metadata;
|
||||
if let Err(error) = record.usage_records.record(event).await {
|
||||
if let Err(error) = record.ops_errors.record(event).await {
|
||||
tracing::warn!(account_id = %record.account_id, error = %error, "failed to record upstream error event");
|
||||
}
|
||||
}
|
||||
@@ -204,25 +210,38 @@ pub(super) async fn record_prefetched_response_stream_failure_event(
|
||||
Some(record.attempt),
|
||||
);
|
||||
}
|
||||
record_response_event(ResponseUsageRecord {
|
||||
usage_records: record.usage_records,
|
||||
request_id: record.request_id,
|
||||
account_id: record.account_id,
|
||||
route: record.route,
|
||||
model: record.model,
|
||||
requested_model: Some(record.requested_model),
|
||||
client_ip: record.request.client_ip.as_deref(),
|
||||
client_user_agent: record.request.client_user_agent.as_deref(),
|
||||
reasoning_effort: reasoning_effort_from_request(record.request),
|
||||
service_tier: record.request.service_tier(),
|
||||
started_at: record.started_at,
|
||||
status_code: event_status_code,
|
||||
level: UsageRecordLevel::Error,
|
||||
message: "v1 responses stream failed",
|
||||
metadata,
|
||||
rate_limit_headers: record.rate_limit_headers,
|
||||
})
|
||||
.await;
|
||||
enrich_event_route_metadata(&mut metadata, record.route);
|
||||
enrich_usage_record_identity(
|
||||
&mut metadata,
|
||||
Some(record.requested_model),
|
||||
record.model,
|
||||
record.request.client_ip.as_deref(),
|
||||
record.request.client_user_agent.as_deref(),
|
||||
reasoning_effort_from_request(record.request),
|
||||
record.request.service_tier(),
|
||||
);
|
||||
let mut event = OpsErrorLog::new(
|
||||
response_event_kind(record.route),
|
||||
"v1 responses stream failed",
|
||||
);
|
||||
event.request_id = Some(record.request_id.to_string());
|
||||
event.account_id = Some(record.account_id.to_string());
|
||||
event.route = Some(record.route.to_string());
|
||||
event.model = Some(record.model.to_string());
|
||||
event.status_code = Some(event_status_code);
|
||||
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
|
||||
if !record.rate_limit_headers.is_empty() {
|
||||
if let Some(object) = metadata.as_object_mut() {
|
||||
object.insert(
|
||||
"rateLimitHeaders".to_string(),
|
||||
serde_json::json!(record.rate_limit_headers),
|
||||
);
|
||||
}
|
||||
}
|
||||
event.metadata = metadata;
|
||||
if let Err(error) = record.ops_errors.record(event).await {
|
||||
tracing::warn!(account_id = %record.account_id, error = %error, "failed to record prefetched stream error event");
|
||||
}
|
||||
}
|
||||
|
||||
fn ensure_stream_metadata_flag(metadata: &mut Value) {
|
||||
@@ -469,13 +488,6 @@ pub(super) async fn record_live_response_stream_event(
|
||||
body,
|
||||
);
|
||||
log_live_response_stream_finalized(context, status_code, level, message, &metadata, body);
|
||||
let mut event = UsageRecord::new(response_event_kind(&context.route), level, message);
|
||||
event.request_id = Some(context.request_id.clone());
|
||||
event.account_id = Some(context.account_id.clone());
|
||||
event.route = Some(context.route.clone());
|
||||
event.model = Some(context.display_model.clone());
|
||||
event.status_code = Some(status_code);
|
||||
event.latency_ms = Some(elapsed_millis_i64(context.started_at));
|
||||
enrich_usage_record_identity(
|
||||
&mut metadata,
|
||||
Some(&context.requested_model),
|
||||
@@ -485,12 +497,38 @@ pub(super) async fn record_live_response_stream_event(
|
||||
reasoning_effort_from_request(&context.request),
|
||||
context.request.service_tier(),
|
||||
);
|
||||
event.metadata = metadata;
|
||||
if let Err(error) = context.usage_records.record(event).await {
|
||||
tracing::warn!(account_id = %context.account_id, error = %error, "failed to record live response stream event");
|
||||
|
||||
if is_success_usage_event(level, status_code) {
|
||||
let mut event = UsageRecord::new(response_event_kind(&context.route), level, message);
|
||||
event.request_id = Some(context.request_id.clone());
|
||||
event.account_id = Some(context.account_id.clone());
|
||||
event.route = Some(context.route.clone());
|
||||
event.model = Some(context.display_model.clone());
|
||||
event.status_code = Some(status_code);
|
||||
event.latency_ms = Some(elapsed_millis_i64(context.started_at));
|
||||
event.metadata = metadata;
|
||||
if let Err(error) = context.usage_records.record(event).await {
|
||||
tracing::warn!(account_id = %context.account_id, error = %error, "failed to record live response stream event");
|
||||
}
|
||||
} else {
|
||||
let mut event = OpsErrorLog::new(response_event_kind(&context.route), message);
|
||||
event.request_id = Some(context.request_id.clone());
|
||||
event.account_id = Some(context.account_id.clone());
|
||||
event.route = Some(context.route.clone());
|
||||
event.model = Some(context.display_model.clone());
|
||||
event.status_code = Some(status_code);
|
||||
event.latency_ms = Some(elapsed_millis_i64(context.started_at));
|
||||
event.metadata = metadata;
|
||||
if let Err(error) = context.ops_errors.record(event).await {
|
||||
tracing::warn!(account_id = %context.account_id, error = %error, "failed to record live response stream error event");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn is_success_usage_event(level: UsageRecordLevel, status_code: i64) -> bool {
|
||||
level != UsageRecordLevel::Error && status_code < 400
|
||||
}
|
||||
|
||||
fn log_live_response_stream_finalized(
|
||||
context: &LiveResponseStreamContext,
|
||||
status_code: i64,
|
||||
|
||||
@@ -17,6 +17,7 @@ use tokio::sync::Mutex;
|
||||
|
||||
use crate::{
|
||||
admin::monitoring::{
|
||||
ops_error_service::AdminOpsErrorLogService,
|
||||
usage_record_model::{ResponseUsageRecord, UsageRecordLevel},
|
||||
usage_record_service::AdminUsageRecordService,
|
||||
},
|
||||
@@ -110,6 +111,7 @@ pub struct ResponseDispatchService {
|
||||
session_affinity: Arc<RuntimeSessionAffinityService>,
|
||||
reasoning_replay: Arc<Mutex<ReasoningReplayCache>>,
|
||||
usage_records: Arc<AdminUsageRecordService>,
|
||||
ops_errors: Arc<AdminOpsErrorLogService>,
|
||||
token_refresh: Arc<RuntimeTokenRefreshService<OpenAiTokenClient>>,
|
||||
installation_id: Option<String>,
|
||||
cloudflare: CloudflareRecovery,
|
||||
@@ -121,6 +123,7 @@ pub(crate) struct ResponseDispatchServiceParts {
|
||||
pub codex: Arc<CodexBackendClient>,
|
||||
pub session_affinity: Arc<RuntimeSessionAffinityService>,
|
||||
pub usage_records: Arc<AdminUsageRecordService>,
|
||||
pub ops_errors: Arc<AdminOpsErrorLogService>,
|
||||
pub token_refresh: Arc<RuntimeTokenRefreshService<OpenAiTokenClient>>,
|
||||
pub installation_id: Option<String>,
|
||||
pub cloudflare: CloudflareRecovery,
|
||||
@@ -130,6 +133,34 @@ pub(crate) struct ResponseDispatchServiceParts {
|
||||
const DEFAULT_REASONING_REPLAY_TTL_SECS: i64 = 55 * 60;
|
||||
const IMPLICIT_RESUME_MAX_AGE_SECS: i64 = DEFAULT_REASONING_REPLAY_TTL_SECS;
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
struct AccountAffinityDecision {
|
||||
preferred_account_id: Option<String>,
|
||||
request_history_owner: RequestHistoryOwner,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
enum RequestHistoryOwner {
|
||||
#[default]
|
||||
Absent,
|
||||
Known(String),
|
||||
Unknown,
|
||||
}
|
||||
|
||||
impl AccountAffinityDecision {
|
||||
fn preferred_account_id(&self) -> Option<&str> {
|
||||
self.preferred_account_id.as_deref()
|
||||
}
|
||||
|
||||
fn should_strip_request_history_for(&self, account_id: &str) -> bool {
|
||||
match &self.request_history_owner {
|
||||
RequestHistoryOwner::Absent => false,
|
||||
RequestHistoryOwner::Known(history_account_id) => history_account_id != account_id,
|
||||
RequestHistoryOwner::Unknown => true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Responses live SSE 响应体流。
|
||||
pub type ResponseBodyStream =
|
||||
Pin<Box<dyn Stream<Item = Result<Bytes, ResponseDispatchStreamError>> + Send + 'static>>;
|
||||
@@ -150,6 +181,7 @@ impl ResponseDispatchService {
|
||||
DEFAULT_REASONING_REPLAY_TTL_SECS,
|
||||
)))),
|
||||
usage_records: parts.usage_records,
|
||||
ops_errors: parts.ops_errors,
|
||||
token_refresh: parts.token_refresh,
|
||||
installation_id: parts.installation_id,
|
||||
cloudflare: parts.cloudflare,
|
||||
@@ -260,7 +292,31 @@ impl ResponseDispatchService {
|
||||
Some(snapshot)
|
||||
}
|
||||
|
||||
async fn preferred_account_id_for_request(
|
||||
async fn account_affinity_for_request(
|
||||
&self,
|
||||
request: &CodexResponsesRequest,
|
||||
now: DateTime<Utc>,
|
||||
) -> AccountAffinityDecision {
|
||||
let has_previous_response_id = request.previous_response_id().is_some();
|
||||
if let Some(account_id) = self.previous_response_account_id(request, now).await {
|
||||
return AccountAffinityDecision {
|
||||
preferred_account_id: Some(account_id.clone()),
|
||||
request_history_owner: RequestHistoryOwner::Known(account_id),
|
||||
};
|
||||
}
|
||||
|
||||
let conversation_account_id = self.conversation_account_id(request, now).await;
|
||||
AccountAffinityDecision {
|
||||
preferred_account_id: conversation_account_id,
|
||||
request_history_owner: if has_previous_response_id {
|
||||
RequestHistoryOwner::Unknown
|
||||
} else {
|
||||
RequestHistoryOwner::Absent
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
async fn previous_response_account_id(
|
||||
&self,
|
||||
request: &CodexResponsesRequest,
|
||||
now: DateTime<Utc>,
|
||||
@@ -271,6 +327,21 @@ impl ResponseDispatchService {
|
||||
.await
|
||||
}
|
||||
|
||||
async fn conversation_account_id(
|
||||
&self,
|
||||
request: &CodexResponsesRequest,
|
||||
now: DateTime<Utc>,
|
||||
) -> Option<String> {
|
||||
let conversation_id = request
|
||||
.prompt_cache_key()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())?;
|
||||
let variant_hash = compute_variant_hash(request);
|
||||
self.session_affinity
|
||||
.lookup_latest_account_by_conversation(conversation_id, None, Some(&variant_hash), now)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn recover_request_history(
|
||||
&self,
|
||||
request: &mut CodexResponsesRequest,
|
||||
@@ -279,13 +350,22 @@ impl ResponseDispatchService {
|
||||
if let Some(previous_response_id) = request.previous_response_id() {
|
||||
self.session_affinity.forget(previous_response_id).await;
|
||||
}
|
||||
if let Some(snapshot) = implicit_resume.take() {
|
||||
snapshot.restore(request);
|
||||
request.set_previous_response_id(None);
|
||||
request.turn_state = None;
|
||||
} else {
|
||||
strip_request_history(request);
|
||||
restore_request_without_history(request, implicit_resume);
|
||||
}
|
||||
|
||||
fn strip_history_if_account_changed(
|
||||
request: &mut CodexResponsesRequest,
|
||||
implicit_resume: &mut Option<ImplicitResumeSnapshot>,
|
||||
account_affinity: &AccountAffinityDecision,
|
||||
acquired_account_id: &str,
|
||||
) {
|
||||
if request.previous_response_id().is_none()
|
||||
|| !account_affinity.should_strip_request_history_for(acquired_account_id)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
restore_request_without_history(request, implicit_resume);
|
||||
}
|
||||
|
||||
async fn apply_cascading_ban_defense(
|
||||
@@ -326,13 +406,7 @@ impl ResponseDispatchService {
|
||||
let response_id_to_forget = explicit_previous_response_id
|
||||
.or(request.previous_response_id())
|
||||
.map(str::to_string);
|
||||
if let Some(snapshot) = implicit_resume.take() {
|
||||
snapshot.restore(request);
|
||||
request.set_previous_response_id(None);
|
||||
request.turn_state = None;
|
||||
} else {
|
||||
strip_request_history(request);
|
||||
}
|
||||
restore_request_without_history(request, implicit_resume);
|
||||
if let Some(response_id) = response_id_to_forget {
|
||||
self.session_affinity.forget(&response_id).await;
|
||||
}
|
||||
@@ -380,9 +454,9 @@ impl ResponseDispatchService {
|
||||
let now = Utc::now();
|
||||
let explicit_previous_response_id = request.previous_response_id().map(ToString::to_string);
|
||||
let mut implicit_resume = self.prepare_response_session(&mut request).await;
|
||||
let preferred_account_id = self.preferred_account_id_for_request(&request, now).await;
|
||||
let account_affinity = self.account_affinity_for_request(&request, now).await;
|
||||
let mut acquire_request = AccountAcquireRequest::new(request.model(), now);
|
||||
if let Some(preferred_account_id) = preferred_account_id.as_deref() {
|
||||
if let Some(preferred_account_id) = account_affinity.preferred_account_id() {
|
||||
acquire_request = acquire_request.with_preferred_account_id(preferred_account_id);
|
||||
}
|
||||
let mut excluded_account_ids = Vec::new();
|
||||
@@ -493,11 +567,17 @@ impl ResponseDispatchService {
|
||||
self.apply_cascading_ban_defense(
|
||||
&mut request,
|
||||
&mut implicit_resume,
|
||||
preferred_account_id.as_deref(),
|
||||
account_affinity.preferred_account_id(),
|
||||
&acquired.account.id,
|
||||
explicit_previous_response_id.as_deref(),
|
||||
)
|
||||
.await;
|
||||
Self::strip_history_if_account_changed(
|
||||
&mut request,
|
||||
&mut implicit_resume,
|
||||
&account_affinity,
|
||||
&acquired.account.id,
|
||||
);
|
||||
self.account_pool.wait_for_request_interval(&acquired).await;
|
||||
let account = acquired.account;
|
||||
let release_account_id = account.id.clone();
|
||||
@@ -758,7 +838,7 @@ impl ResponseDispatchService {
|
||||
}
|
||||
Err(error) => {
|
||||
record_response_upstream_error_event(ResponseUpstreamErrorEventRecord {
|
||||
usage_records: &self.usage_records,
|
||||
ops_errors: &self.ops_errors,
|
||||
request_id,
|
||||
account_id: &release_account_id,
|
||||
account_email: account.email.as_deref(),
|
||||
@@ -904,7 +984,7 @@ impl ResponseDispatchService {
|
||||
error: &ResponseDispatchError,
|
||||
) {
|
||||
record_response_dispatch_error_event(ResponseDispatchErrorEventRecord {
|
||||
usage_records: &self.usage_records,
|
||||
ops_errors: &self.ops_errors,
|
||||
request_id,
|
||||
account_id: details.account_id,
|
||||
route,
|
||||
@@ -935,9 +1015,9 @@ impl ResponseDispatchService {
|
||||
let now = Utc::now();
|
||||
let explicit_previous_response_id = request.previous_response_id().map(ToString::to_string);
|
||||
let mut implicit_resume = self.prepare_response_session(&mut request).await;
|
||||
let preferred_account_id = self.preferred_account_id_for_request(&request, now).await;
|
||||
let account_affinity = self.account_affinity_for_request(&request, now).await;
|
||||
let mut acquire_request = AccountAcquireRequest::new(request.model(), now);
|
||||
if let Some(preferred_account_id) = preferred_account_id.as_deref() {
|
||||
if let Some(preferred_account_id) = account_affinity.preferred_account_id() {
|
||||
acquire_request = acquire_request.with_preferred_account_id(preferred_account_id);
|
||||
}
|
||||
let mut excluded_account_ids = Vec::new();
|
||||
@@ -1053,11 +1133,17 @@ impl ResponseDispatchService {
|
||||
self.apply_cascading_ban_defense(
|
||||
&mut request,
|
||||
&mut implicit_resume,
|
||||
preferred_account_id.as_deref(),
|
||||
account_affinity.preferred_account_id(),
|
||||
&acquired.account.id,
|
||||
explicit_previous_response_id.as_deref(),
|
||||
)
|
||||
.await;
|
||||
Self::strip_history_if_account_changed(
|
||||
&mut request,
|
||||
&mut implicit_resume,
|
||||
&account_affinity,
|
||||
&acquired.account.id,
|
||||
);
|
||||
|
||||
self.account_pool.wait_for_request_interval(&acquired).await;
|
||||
let account = acquired.account;
|
||||
@@ -1236,7 +1322,7 @@ impl ResponseDispatchService {
|
||||
if let ResponseDispatchError::Upstream(upstream_error) = &error {
|
||||
record_response_upstream_error_event(
|
||||
ResponseUpstreamErrorEventRecord {
|
||||
usage_records: &self.usage_records,
|
||||
ops_errors: &self.ops_errors,
|
||||
request_id,
|
||||
account_id: &release_account_id,
|
||||
account_email: account.email.as_deref(),
|
||||
@@ -1341,7 +1427,7 @@ impl ResponseDispatchService {
|
||||
self.account_pool.release(&release_account_id).await;
|
||||
record_prefetched_response_stream_failure_event(
|
||||
ResponseStreamFailureEventRecord {
|
||||
usage_records: &self.usage_records,
|
||||
ops_errors: &self.ops_errors,
|
||||
request_id,
|
||||
account_id: &release_account_id,
|
||||
route,
|
||||
@@ -1367,6 +1453,7 @@ impl ResponseDispatchService {
|
||||
session_affinity: Arc::clone(&self.session_affinity),
|
||||
reasoning_replay: Arc::clone(&self.reasoning_replay),
|
||||
usage_records: Arc::clone(&self.usage_records),
|
||||
ops_errors: Arc::clone(&self.ops_errors),
|
||||
cloudflare: self.cloudflare.clone(),
|
||||
account_id: account.id,
|
||||
account_plan_type: account.plan_type,
|
||||
@@ -1503,7 +1590,7 @@ impl ResponseDispatchService {
|
||||
Err(error) => {
|
||||
self.account_pool.release(&release_account_id).await;
|
||||
record_response_upstream_error_event(ResponseUpstreamErrorEventRecord {
|
||||
usage_records: &self.usage_records,
|
||||
ops_errors: &self.ops_errors,
|
||||
request_id,
|
||||
account_id: &release_account_id,
|
||||
account_email: account.email.as_deref(),
|
||||
@@ -1795,7 +1882,7 @@ impl ResponseDispatchService {
|
||||
error: &ResponseDispatchError,
|
||||
) {
|
||||
record_response_dispatch_error_event(ResponseDispatchErrorEventRecord {
|
||||
usage_records: &self.usage_records,
|
||||
ops_errors: &self.ops_errors,
|
||||
request_id,
|
||||
account_id,
|
||||
route: "/v1/responses/compact",
|
||||
@@ -1814,3 +1901,13 @@ fn strip_request_history(request: &mut CodexResponsesRequest) {
|
||||
request.set_previous_response_id(None);
|
||||
request.turn_state = None;
|
||||
}
|
||||
|
||||
fn restore_request_without_history(
|
||||
request: &mut CodexResponsesRequest,
|
||||
implicit_resume: &mut Option<ImplicitResumeSnapshot>,
|
||||
) {
|
||||
if let Some(snapshot) = implicit_resume.take() {
|
||||
snapshot.restore(request);
|
||||
}
|
||||
strip_request_history(request);
|
||||
}
|
||||
|
||||
@@ -5,7 +5,8 @@ use tokio::sync::Mutex;
|
||||
|
||||
use crate::{
|
||||
admin::monitoring::{
|
||||
usage_record_model::UsageRecordLevel, usage_record_service::AdminUsageRecordService,
|
||||
ops_error_service::AdminOpsErrorLogService, usage_record_model::UsageRecordLevel,
|
||||
usage_record_service::AdminUsageRecordService,
|
||||
},
|
||||
infra::time::elapsed_millis_i64,
|
||||
proxy::{
|
||||
@@ -49,6 +50,7 @@ pub(super) struct LiveResponseStreamContext {
|
||||
pub(super) session_affinity: Arc<RuntimeSessionAffinityService>,
|
||||
pub(super) reasoning_replay: Arc<Mutex<ReasoningReplayCache>>,
|
||||
pub(super) usage_records: Arc<AdminUsageRecordService>,
|
||||
pub(super) ops_errors: Arc<AdminOpsErrorLogService>,
|
||||
pub(super) cloudflare: CloudflareRecovery,
|
||||
pub(super) account_id: String,
|
||||
pub(super) account_plan_type: Option<String>,
|
||||
|
||||
@@ -106,6 +106,38 @@ impl SessionAffinityMap {
|
||||
variant_hash: Option<&str>,
|
||||
now: DateTime<Utc>,
|
||||
) -> Option<String> {
|
||||
self.latest_entry_by_conversation(conversation_id, max_age, variant_hash, now)
|
||||
.map(|(response_id, _)| response_id.clone())
|
||||
}
|
||||
|
||||
pub fn lookup_latest_account_by_conversation(
|
||||
&self,
|
||||
conversation_id: &str,
|
||||
max_age: Option<Duration>,
|
||||
variant_hash: Option<&str>,
|
||||
now: DateTime<Utc>,
|
||||
) -> Option<String> {
|
||||
self.latest_entry_by_conversation(conversation_id, max_age, variant_hash, now)
|
||||
.map(|(_, entry)| entry.account_id.clone())
|
||||
}
|
||||
|
||||
pub fn forget(&mut self, response_id: &str) -> bool {
|
||||
self.entries.remove(response_id).is_some()
|
||||
}
|
||||
|
||||
fn active_entry(&self, response_id: &str, now: DateTime<Utc>) -> Option<&SessionAffinityEntry> {
|
||||
self.entries
|
||||
.get(response_id)
|
||||
.filter(|entry| self.entry_is_active(entry, now))
|
||||
}
|
||||
|
||||
fn latest_entry_by_conversation(
|
||||
&self,
|
||||
conversation_id: &str,
|
||||
max_age: Option<Duration>,
|
||||
variant_hash: Option<&str>,
|
||||
now: DateTime<Utc>,
|
||||
) -> Option<(&String, &SessionAffinityEntry)> {
|
||||
self.entries
|
||||
.iter()
|
||||
.filter(|(_, entry)| entry.conversation_id == conversation_id)
|
||||
@@ -118,17 +150,6 @@ impl SessionAffinityMap {
|
||||
max_age.is_none_or(|max_age| now.signed_duration_since(entry.created_at) <= max_age)
|
||||
})
|
||||
.max_by_key(|(_, entry)| entry.created_at)
|
||||
.map(|(response_id, _)| response_id.clone())
|
||||
}
|
||||
|
||||
pub fn forget(&mut self, response_id: &str) -> bool {
|
||||
self.entries.remove(response_id).is_some()
|
||||
}
|
||||
|
||||
fn active_entry(&self, response_id: &str, now: DateTime<Utc>) -> Option<&SessionAffinityEntry> {
|
||||
self.entries
|
||||
.get(response_id)
|
||||
.filter(|entry| self.entry_is_active(entry, now))
|
||||
}
|
||||
|
||||
fn stored_record_is_active(&self, record: &StoredSessionAffinity, now: DateTime<Utc>) -> bool {
|
||||
@@ -630,6 +651,21 @@ impl RuntimeSessionAffinityService {
|
||||
.lookup_latest_response_by_conversation(conversation_id, max_age, variant_hash, now)
|
||||
}
|
||||
|
||||
pub async fn lookup_latest_account_by_conversation(
|
||||
&self,
|
||||
conversation_id: &str,
|
||||
max_age: Option<Duration>,
|
||||
variant_hash: Option<&str>,
|
||||
now: DateTime<Utc>,
|
||||
) -> Option<String> {
|
||||
self.map.read().await.lookup_latest_account_by_conversation(
|
||||
conversation_id,
|
||||
max_age,
|
||||
variant_hash,
|
||||
now,
|
||||
)
|
||||
}
|
||||
|
||||
pub async fn forget(&self, response_id: &str) -> bool {
|
||||
let in_memory = self.map.write().await.forget(response_id);
|
||||
match self.store.forget(response_id).await {
|
||||
|
||||
@@ -103,7 +103,10 @@ pub(crate) async fn verify_acquired_quota_if_required(
|
||||
.apply_quota_snapshot(&account_id, "a)
|
||||
.await;
|
||||
if quota_snapshot_limit_reached("a) {
|
||||
context.account_pool.release(&account_id).await;
|
||||
context
|
||||
.account_pool
|
||||
.release_without_request_usage(&account_id)
|
||||
.await;
|
||||
context.excluded_account_ids.push(account_id.clone());
|
||||
*context.verify_attempts += 1;
|
||||
let max_attempts_reached = *context.verify_attempts >= MAX_QUOTA_VERIFY_ATTEMPTS;
|
||||
|
||||
@@ -6,15 +6,16 @@ use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
admin::monitoring::{
|
||||
usage_record_model::{ResponseUsageRecord, UsageRecord, UsageRecordLevel},
|
||||
usage_record_service::AdminUsageRecordService,
|
||||
ops_error_model::OpsErrorLog,
|
||||
ops_error_service::AdminOpsErrorLogService,
|
||||
usage_record_model::{ResponseUsageRecord, UsageRecord},
|
||||
},
|
||||
infra::time::elapsed_millis_i64,
|
||||
upstream::protocol::responses::{CodexCompactRequest, CodexResponsesRequest},
|
||||
};
|
||||
|
||||
pub(super) struct DispatchErrorUsageRecord<'a> {
|
||||
pub usage_records: &'a AdminUsageRecordService,
|
||||
pub(super) struct DispatchErrorLogRecord<'a> {
|
||||
pub ops_errors: &'a AdminOpsErrorLogService,
|
||||
pub request_id: &'a str,
|
||||
pub account_id: Option<&'a str>,
|
||||
pub route: &'a str,
|
||||
@@ -25,15 +26,11 @@ pub(super) struct DispatchErrorUsageRecord<'a> {
|
||||
pub metadata: Value,
|
||||
}
|
||||
|
||||
pub(super) async fn record_dispatch_error_event(record: DispatchErrorUsageRecord<'_>) {
|
||||
pub(super) async fn record_dispatch_error_event(record: DispatchErrorLogRecord<'_>) {
|
||||
let mut metadata = record.metadata;
|
||||
enrich_event_route_metadata(&mut metadata, record.route);
|
||||
|
||||
let mut event = UsageRecord::new(
|
||||
event_kind(record.route),
|
||||
UsageRecordLevel::Error,
|
||||
record.message,
|
||||
);
|
||||
let mut event = OpsErrorLog::new(event_kind(record.route), record.message);
|
||||
event.request_id = Some(record.request_id.to_string());
|
||||
event.account_id = record.account_id.map(ToString::to_string);
|
||||
event.route = Some(record.route.to_string());
|
||||
@@ -42,7 +39,7 @@ pub(super) async fn record_dispatch_error_event(record: DispatchErrorUsageRecord
|
||||
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
|
||||
event.metadata = metadata;
|
||||
|
||||
if let Err(error) = record.usage_records.record(event).await {
|
||||
if let Err(error) = record.ops_errors.record(event).await {
|
||||
tracing::warn!(
|
||||
account_id = record.account_id.unwrap_or(""),
|
||||
error = %error,
|
||||
|
||||
@@ -5,7 +5,6 @@ use chrono::Utc;
|
||||
|
||||
use crate::config::schema::AppConfig;
|
||||
use crate::config::settings::RuntimeSettingsService;
|
||||
use crate::http::middleware::request_id::TrustedProxyConfig;
|
||||
use crate::infra::database::connect_sqlite;
|
||||
use crate::infra::logging::{init_tracing, LogGuard, RotationConfig};
|
||||
use crate::infra::paths::{ensure_data_dir, load_or_create_installation_id};
|
||||
@@ -113,11 +112,10 @@ async fn build_application(
|
||||
);
|
||||
|
||||
let task_coordinator = TaskCoordinator::start(&runtime_config, &services);
|
||||
let trusted_proxies = TrustedProxyConfig::from_entries(&config.server.trusted_proxies)?;
|
||||
let app_state = AppState { services };
|
||||
|
||||
Ok((
|
||||
crate::http::router::router_with_trusted_proxies(trusted_proxies).with_state(app_state),
|
||||
crate::http::router::router().with_state(app_state),
|
||||
task_coordinator,
|
||||
))
|
||||
}
|
||||
|
||||
@@ -6,6 +6,8 @@ use crate::{
|
||||
admin::monitoring::{
|
||||
account_usage_service::AdminUsageService,
|
||||
account_usage_store::SqliteUsageStore,
|
||||
ops_error_service::AdminOpsErrorLogService,
|
||||
ops_error_store::SqliteOpsErrorLogStore,
|
||||
usage_record_service::AdminUsageRecordService,
|
||||
usage_record_store::{SqliteUsageRecordStore, DEFAULT_USAGE_RECORD_CAPTURE_BODY},
|
||||
},
|
||||
@@ -91,6 +93,7 @@ pub struct Services {
|
||||
pub refresh_policy: RuntimeRefreshPolicy,
|
||||
pub admin_accounts: StdArc<AdminAccountService>,
|
||||
pub usage_records: StdArc<AdminUsageRecordService>,
|
||||
pub ops_errors: StdArc<AdminOpsErrorLogService>,
|
||||
pub usage: StdArc<AdminUsageService>,
|
||||
pub account_pool: StdArc<RuntimeAccountPoolService>,
|
||||
pub token_refresh: StdArc<RuntimeTokenRefreshService<OpenAiTokenClient>>,
|
||||
@@ -193,6 +196,10 @@ impl Services {
|
||||
usage_record_options.enabled,
|
||||
usage_record_options.capture_body,
|
||||
));
|
||||
let ops_errors = StdArc::new(AdminOpsErrorLogService::new(
|
||||
SqliteOpsErrorLogStore::new(stores.usage_records.pool().clone()),
|
||||
usage_record_options.capture_body,
|
||||
));
|
||||
let account_pool = StdArc::new(RuntimeAccountPoolService::new(
|
||||
account_store_trait.clone(),
|
||||
account_pool_options_from_config(config),
|
||||
@@ -247,6 +254,7 @@ impl Services {
|
||||
models: models.clone(),
|
||||
account_pool: account_pool.clone(),
|
||||
token_refresher: StdArc::new(token_client),
|
||||
refresh_leases: stores.refresh_leases.clone(),
|
||||
oauth: crate::admin::accounts::service::oauth::AccountOAuthService::new(
|
||||
reqwest::Client::new(),
|
||||
config.auth.oauth_client_id.clone(),
|
||||
@@ -269,6 +277,7 @@ impl Services {
|
||||
codex: codex.clone(),
|
||||
session_affinity: session_affinity.clone(),
|
||||
usage_records: usage_records.clone(),
|
||||
ops_errors: ops_errors.clone(),
|
||||
token_refresh: token_refresh.clone(),
|
||||
installation_id: installation_id.clone(),
|
||||
cloudflare: cloudflare_recovery,
|
||||
@@ -284,6 +293,7 @@ impl Services {
|
||||
refresh_policy,
|
||||
admin_accounts,
|
||||
usage_records,
|
||||
ops_errors,
|
||||
usage,
|
||||
account_pool,
|
||||
token_refresh,
|
||||
|
||||
@@ -88,7 +88,9 @@ impl ModelRefreshTask {
|
||||
)
|
||||
.await;
|
||||
for selected in &selected_accounts {
|
||||
self.account_pool.release(&selected.account.id).await;
|
||||
self.account_pool
|
||||
.release_without_request_usage(&selected.account.id)
|
||||
.await;
|
||||
}
|
||||
|
||||
let result = result?;
|
||||
|
||||
@@ -613,7 +613,6 @@ pub fn parse_account_status(status: &str) -> Result<AccountStatus, &'static str>
|
||||
"disabled" => Ok(AccountStatus::Disabled),
|
||||
"expired" => Ok(AccountStatus::Expired),
|
||||
"quota_exhausted" => Ok(AccountStatus::QuotaExhausted),
|
||||
"refreshing" => Ok(AccountStatus::Refreshing),
|
||||
"banned" => Ok(AccountStatus::Banned),
|
||||
_ => Err("invalid account status"),
|
||||
}
|
||||
|
||||
@@ -12,8 +12,6 @@ pub enum AccountStatus {
|
||||
Expired,
|
||||
/// 当前配额已经耗尽。
|
||||
QuotaExhausted,
|
||||
/// 账号正处于刷新流程中。
|
||||
Refreshing,
|
||||
/// 账号被显式禁用。
|
||||
Disabled,
|
||||
/// 账号被上游封禁。
|
||||
@@ -27,7 +25,6 @@ impl AccountStatus {
|
||||
Self::Active => "active",
|
||||
Self::Expired => "expired",
|
||||
Self::QuotaExhausted => "quota_exhausted",
|
||||
Self::Refreshing => "refreshing",
|
||||
Self::Disabled => "disabled",
|
||||
Self::Banned => "banned",
|
||||
}
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
//! 账号池调度策略。
|
||||
|
||||
use std::{
|
||||
cmp::Ordering,
|
||||
collections::{BTreeMap, BTreeSet, VecDeque},
|
||||
sync::Arc as StdArc,
|
||||
time::Duration as StdDuration,
|
||||
@@ -11,7 +10,6 @@ use chrono::{DateTime, Duration, Utc};
|
||||
use indexmap::IndexMap;
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::infra::time::parse_rfc3339_utc;
|
||||
use crate::upstream::accounts::quota::{quota_snapshot_limit_reached, quota_snapshot_reset_at};
|
||||
use crate::upstream::accounts::store::AccountStore;
|
||||
use crate::upstream::accounts::token_refresh::{jwt_expiry, JwtExpiry};
|
||||
@@ -23,30 +21,12 @@ use crate::upstream::accounts::{
|
||||
use crate::upstream::protocol::events::{
|
||||
parse_rate_limit_headers, rate_limit_quota, TokenUsage as CodexTokenUsage,
|
||||
};
|
||||
use crate::upstream::scheduler::{
|
||||
candidates::{self, CandidateFilter, CandidateRequest},
|
||||
AccountScheduler, ScoreWeights,
|
||||
};
|
||||
|
||||
/// 账号轮转策略。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum RotationStrategy {
|
||||
/// 综合并发压力、窗口用量和历史用量进行调度。
|
||||
Smart,
|
||||
/// 优先消耗即将重置的额度窗口。
|
||||
QuotaResetPriority,
|
||||
/// 按候选账号顺序循环选择。
|
||||
RoundRobin,
|
||||
/// 优先选择最近使用过的账号。
|
||||
Sticky,
|
||||
}
|
||||
|
||||
impl RotationStrategy {
|
||||
fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Smart => "smart",
|
||||
Self::QuotaResetPriority => "quota_reset_priority",
|
||||
Self::RoundRobin => "round_robin",
|
||||
Self::Sticky => "sticky",
|
||||
}
|
||||
}
|
||||
}
|
||||
pub use crate::upstream::scheduler::RotationStrategy;
|
||||
|
||||
/// 账号池配置。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
@@ -263,13 +243,17 @@ pub struct AccountWindowUsageDelta {
|
||||
pub image_request_failed: bool,
|
||||
}
|
||||
|
||||
/// 纯内存账号池,负责账号调度和运行时状态维护。
|
||||
/// 纯内存账号池:负责账号存储、在途槽位与运行时状态维护。
|
||||
///
|
||||
/// 账号**选择**已剥离到 [`crate::upstream::scheduler`]:候选过滤走
|
||||
/// [`candidates::filter`],策略选择走 [`AccountScheduler::select`]。账号池只提供选择
|
||||
/// 所需的只读视图(候选切片 + `slot_count`),并持有调度器实例。
|
||||
#[derive(Debug)]
|
||||
pub struct AccountPool {
|
||||
accounts: IndexMap<String, Account>,
|
||||
slots: BTreeMap<String, VecDeque<AccountSlot>>,
|
||||
options: AccountPoolOptions,
|
||||
rotation_cursor: usize,
|
||||
scheduler: AccountScheduler,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
@@ -278,16 +262,6 @@ struct AccountSlot {
|
||||
model: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
||||
struct SmartAccountLoadScore {
|
||||
quota_limited: bool,
|
||||
slot_count: usize,
|
||||
window_request_count: u64,
|
||||
window_token_count: u64,
|
||||
window_cached_tokens: u64,
|
||||
request_count: u64,
|
||||
}
|
||||
|
||||
/// 账号释放结果。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ReleasedAccountUsage {
|
||||
@@ -308,7 +282,7 @@ impl AccountPool {
|
||||
accounts: IndexMap::new(),
|
||||
slots: BTreeMap::new(),
|
||||
options,
|
||||
rotation_cursor: 0,
|
||||
scheduler: AccountScheduler::new(ScoreWeights::default()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -322,10 +296,11 @@ impl AccountPool {
|
||||
self.accounts.get(account_id).cloned()
|
||||
}
|
||||
|
||||
/// 移除账号及其在途槽位。
|
||||
/// 移除账号及其在途槽位与运行时反馈。
|
||||
pub fn remove(&mut self, account_id: &str) -> bool {
|
||||
let removed = self.accounts.shift_remove(account_id).is_some();
|
||||
self.slots.remove(account_id);
|
||||
self.scheduler.forget_feedback(account_id);
|
||||
removed
|
||||
}
|
||||
|
||||
@@ -333,22 +308,30 @@ impl AccountPool {
|
||||
pub fn clear(&mut self) {
|
||||
self.accounts.clear();
|
||||
self.slots.clear();
|
||||
self.rotation_cursor = 0;
|
||||
self.scheduler.reset_cursor();
|
||||
self.scheduler.clear_feedback();
|
||||
}
|
||||
|
||||
/// 替换账号池运行参数。
|
||||
pub fn set_options(&mut self, options: AccountPoolOptions) {
|
||||
if self.options.rotation_strategy != options.rotation_strategy {
|
||||
self.rotation_cursor = 0;
|
||||
self.scheduler.reset_cursor();
|
||||
}
|
||||
self.options = options;
|
||||
}
|
||||
|
||||
/// 回灌一次请求结果到调度器的运行时反馈(供 Smart 打分)。
|
||||
pub fn report_feedback(&self, account_id: &str, success: bool, first_token_ms: Option<u64>) {
|
||||
self.scheduler
|
||||
.report_feedback(account_id, success, first_token_ms);
|
||||
}
|
||||
|
||||
/// 更新账号状态。
|
||||
pub fn set_status(&mut self, account_id: &str, status: AccountStatus) -> bool {
|
||||
let Some(account) = self.accounts.get_mut(account_id) else {
|
||||
return false;
|
||||
};
|
||||
let status = status_after_quota_limit(status, account.quota_limit_reached);
|
||||
account.status = status;
|
||||
if status != AccountStatus::Active {
|
||||
self.slots.remove(account_id);
|
||||
@@ -371,6 +354,7 @@ impl AccountPool {
|
||||
.unwrap_or(cooldown_until);
|
||||
account.quota_limit_reached = true;
|
||||
account.quota_cooldown_until = Some(final_cooldown_until);
|
||||
account.status = status_after_quota_limit(account.status, true);
|
||||
account.window_reset_at = account
|
||||
.window_reset_at
|
||||
.filter(|existing| *existing > final_cooldown_until)
|
||||
@@ -498,7 +482,9 @@ impl AccountPool {
|
||||
) -> AccountAcquireWithStatusRefresh {
|
||||
self.cleanup_stale_slots(request.now);
|
||||
let refreshed_accounts = self.refresh_account_statuses(request.now);
|
||||
let candidates = self.candidates(request);
|
||||
let candidates = self.filter_candidates(request);
|
||||
// 亲和优先:required 硬锁定该账号;preferred 若在候选内则直接选中;
|
||||
// 二者都不命中时,回退到调度器按当前策略选择。
|
||||
let selected = if let Some(required_account_id) = &request.required_account_id {
|
||||
candidates.first().cloned().map(|account| {
|
||||
(
|
||||
@@ -508,43 +494,24 @@ impl AccountPool {
|
||||
)
|
||||
})
|
||||
} else if let Some(preferred_account_id) = &request.preferred_account_id {
|
||||
let preferred = candidates
|
||||
candidates
|
||||
.iter()
|
||||
.find(|account| account.id == *preferred_account_id)
|
||||
.cloned();
|
||||
preferred.map(|account| {
|
||||
(
|
||||
account,
|
||||
self.previous_slot_at(preferred_account_id),
|
||||
"preferred",
|
||||
)
|
||||
})
|
||||
.cloned()
|
||||
.map(|account| {
|
||||
(
|
||||
account,
|
||||
self.previous_slot_at(preferred_account_id),
|
||||
"preferred",
|
||||
)
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
.or_else(|| {
|
||||
Some(match self.options.rotation_strategy {
|
||||
RotationStrategy::Smart => {
|
||||
let account = self.select_smart(&candidates)?;
|
||||
let previous = self.previous_slot_at(&account.id);
|
||||
(account, previous, "smart")
|
||||
}
|
||||
RotationStrategy::QuotaResetPriority => {
|
||||
let account = self.select_quota_reset_priority(&candidates)?;
|
||||
let previous = self.previous_slot_at(&account.id);
|
||||
(account, previous, "quota_reset_priority")
|
||||
}
|
||||
RotationStrategy::RoundRobin => {
|
||||
let account = self.select_round_robin(&candidates)?;
|
||||
let previous = self.previous_slot_at(&account.id);
|
||||
(account, previous, "round_robin")
|
||||
}
|
||||
RotationStrategy::Sticky => {
|
||||
let account = Self::select_sticky(&candidates)?;
|
||||
let previous = self.previous_slot_at(&account.id);
|
||||
(account, previous, "sticky")
|
||||
}
|
||||
})
|
||||
let account = self.select_candidate(&candidates, request.now)?;
|
||||
let previous = self.previous_slot_at(&account.id);
|
||||
Some((account, previous, self.options.rotation_strategy.as_str()))
|
||||
});
|
||||
let Some((selected, previous_slot_at, selection_source)) = selected else {
|
||||
return AccountAcquireWithStatusRefresh {
|
||||
@@ -552,26 +519,20 @@ impl AccountPool {
|
||||
refreshed_accounts,
|
||||
};
|
||||
};
|
||||
let selected_slot_count = self.slot_count(&selected.id);
|
||||
let selected_window_token_count = account_window_token_count(&selected);
|
||||
let selected_smart_load_score = self.smart_load_score(&selected);
|
||||
tracing::info!(
|
||||
model = %request.model,
|
||||
rotation_strategy = self.options.rotation_strategy.as_str(),
|
||||
selection_source,
|
||||
account_id = %selected.id,
|
||||
candidate_count = candidates.len(),
|
||||
slot_count = selected_slot_count,
|
||||
slot_count = self.slot_count(&selected.id),
|
||||
request_count = selected.request_count,
|
||||
window_request_count = selected.window_request_count,
|
||||
window_token_count = selected_window_token_count,
|
||||
window_reset_at = ?selected.window_reset_at,
|
||||
last_used_at = selected.last_used_at.as_deref().unwrap_or_default(),
|
||||
quota_limit_reached = selected.quota_limit_reached,
|
||||
quota_cooldown_until = ?selected.quota_cooldown_until,
|
||||
previous_slot_at = ?previous_slot_at,
|
||||
smart_load_score = ?selected_smart_load_score,
|
||||
rotation_cursor = self.rotation_cursor,
|
||||
"account selected for upstream request"
|
||||
);
|
||||
|
||||
@@ -615,6 +576,22 @@ impl AccountPool {
|
||||
|
||||
/// 释放指定账号的一个在途槽位。
|
||||
pub fn release(&mut self, account_id: &str) -> Option<ReleasedAccountUsage> {
|
||||
self.release_slot(account_id, true)
|
||||
}
|
||||
|
||||
/// 释放指定账号的一个在途槽位,不累计请求用量。
|
||||
pub fn release_without_request_usage(
|
||||
&mut self,
|
||||
account_id: &str,
|
||||
) -> Option<ReleasedAccountUsage> {
|
||||
self.release_slot(account_id, false)
|
||||
}
|
||||
|
||||
fn release_slot(
|
||||
&mut self,
|
||||
account_id: &str,
|
||||
record_request_usage: bool,
|
||||
) -> Option<ReleasedAccountUsage> {
|
||||
self.accounts.get(account_id)?;
|
||||
|
||||
let mut remove_slots = false;
|
||||
@@ -626,7 +603,9 @@ impl AccountPool {
|
||||
if remove_slots {
|
||||
self.slots.remove(account_id);
|
||||
}
|
||||
self.mark_request_usage(account_id, Utc::now());
|
||||
if record_request_usage {
|
||||
self.mark_request_usage(account_id, Utc::now());
|
||||
}
|
||||
Some(ReleasedAccountUsage {
|
||||
model: slot.and_then(|slot| slot.model),
|
||||
})
|
||||
@@ -689,48 +668,48 @@ impl AccountPool {
|
||||
let refreshed_accounts = self.refresh_account_statuses(now);
|
||||
let mut by_plan = IndexMap::<String, Vec<Account>>::new();
|
||||
|
||||
for account in self.accounts.values() {
|
||||
if !self.is_model_refresh_available(account, now) {
|
||||
continue;
|
||||
let max_concurrent = self.options.max_concurrent_per_account;
|
||||
let skip_quota_limited = self.options.skip_quota_limited;
|
||||
{
|
||||
let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len);
|
||||
for account in self.accounts.values() {
|
||||
if !candidates::is_model_refresh_available(
|
||||
account,
|
||||
max_concurrent,
|
||||
skip_quota_limited,
|
||||
&slot_count,
|
||||
now,
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
let Some(plan_type) = account.plan_type.as_ref() else {
|
||||
continue;
|
||||
};
|
||||
by_plan
|
||||
.entry(plan_type.clone())
|
||||
.or_default()
|
||||
.push(account.clone());
|
||||
}
|
||||
let Some(plan_type) = account.plan_type.as_ref() else {
|
||||
continue;
|
||||
};
|
||||
by_plan
|
||||
.entry(plan_type.clone())
|
||||
.or_default()
|
||||
.push(account.clone());
|
||||
}
|
||||
|
||||
let strategy = self.options.rotation_strategy;
|
||||
let mut accounts = Vec::new();
|
||||
for (plan_type, group) in by_plan {
|
||||
let selected = match self.options.rotation_strategy {
|
||||
RotationStrategy::Smart => self.select_smart(&group),
|
||||
RotationStrategy::QuotaResetPriority => self.select_quota_reset_priority(&group),
|
||||
RotationStrategy::RoundRobin => self.select_round_robin(&group),
|
||||
RotationStrategy::Sticky => Self::select_sticky(&group),
|
||||
};
|
||||
let Some(account) = selected else {
|
||||
let Some(account) = self.select_candidate(&group, now) else {
|
||||
continue;
|
||||
};
|
||||
let selected_slot_count = self.slot_count(&account.id);
|
||||
let selected_window_token_count = account_window_token_count(&account);
|
||||
let selected_smart_load_score = self.smart_load_score(&account);
|
||||
tracing::info!(
|
||||
plan_type,
|
||||
rotation_strategy = self.options.rotation_strategy.as_str(),
|
||||
rotation_strategy = strategy.as_str(),
|
||||
account_id = %account.id,
|
||||
candidate_count = group.len(),
|
||||
slot_count = selected_slot_count,
|
||||
slot_count = self.slot_count(&account.id),
|
||||
request_count = account.request_count,
|
||||
window_request_count = account.window_request_count,
|
||||
window_token_count = selected_window_token_count,
|
||||
window_reset_at = ?account.window_reset_at,
|
||||
last_used_at = account.last_used_at.as_deref().unwrap_or_default(),
|
||||
quota_limit_reached = account.quota_limit_reached,
|
||||
quota_cooldown_until = ?account.quota_cooldown_until,
|
||||
smart_load_score = ?selected_smart_load_score,
|
||||
rotation_cursor = self.rotation_cursor,
|
||||
"account selected for model refresh"
|
||||
);
|
||||
self.push_slot(&account.id, now, None);
|
||||
@@ -743,106 +722,6 @@ impl AccountPool {
|
||||
}
|
||||
}
|
||||
|
||||
fn select_smart(&mut self, candidates: &[Account]) -> Option<Account> {
|
||||
self.select_by(candidates, |pool, a, b| pool.compare_smart(a, b))
|
||||
}
|
||||
|
||||
fn select_quota_reset_priority(&mut self, candidates: &[Account]) -> Option<Account> {
|
||||
self.select_by(candidates, |_pool, a, b| compare_quota_reset_priority(a, b))
|
||||
}
|
||||
|
||||
fn select_by(
|
||||
&mut self,
|
||||
candidates: &[Account],
|
||||
compare: impl Fn(&Self, &Account, &Account) -> Ordering,
|
||||
) -> Option<Account> {
|
||||
let mut sorted = candidates.iter().collect::<Vec<_>>();
|
||||
sorted.sort_by(|a, b| compare(self, a, b));
|
||||
let best = *sorted.first()?;
|
||||
let tied_count = sorted
|
||||
.iter()
|
||||
.take_while(|account| compare(self, best, account) == Ordering::Equal)
|
||||
.count();
|
||||
let index = self.rotation_cursor % tied_count;
|
||||
self.rotation_cursor = self.rotation_cursor.wrapping_add(1);
|
||||
Some((*sorted[index]).clone())
|
||||
}
|
||||
|
||||
fn compare_smart(&self, a: &Account, b: &Account) -> Ordering {
|
||||
self.smart_load_score(a)
|
||||
.cmp(&self.smart_load_score(b))
|
||||
.then_with(|| compare_window_reset(a.window_reset_at, b.window_reset_at))
|
||||
.then_with(|| compare_last_used(a.last_used_at.as_deref(), b.last_used_at.as_deref()))
|
||||
}
|
||||
|
||||
fn smart_load_score(&self, account: &Account) -> SmartAccountLoadScore {
|
||||
SmartAccountLoadScore {
|
||||
quota_limited: account.quota_limit_reached,
|
||||
slot_count: self.slot_count(&account.id),
|
||||
window_request_count: account.window_request_count,
|
||||
window_token_count: account_window_token_count(account),
|
||||
window_cached_tokens: account.window_cached_tokens,
|
||||
request_count: account.request_count,
|
||||
}
|
||||
}
|
||||
|
||||
fn select_sticky(candidates: &[Account]) -> Option<Account> {
|
||||
let mut selected = candidates.first()?;
|
||||
for candidate in &candidates[1..] {
|
||||
if compare_last_used(
|
||||
candidate.last_used_at.as_deref(),
|
||||
selected.last_used_at.as_deref(),
|
||||
) == Ordering::Greater
|
||||
{
|
||||
selected = candidate;
|
||||
}
|
||||
}
|
||||
Some((*selected).clone())
|
||||
}
|
||||
|
||||
fn select_round_robin(&mut self, candidates: &[Account]) -> Option<Account> {
|
||||
if candidates.is_empty() {
|
||||
return None;
|
||||
}
|
||||
self.rotation_cursor %= candidates.len();
|
||||
let index = self.rotation_cursor;
|
||||
self.rotation_cursor = self.rotation_cursor.wrapping_add(1);
|
||||
Some(candidates[index].clone())
|
||||
}
|
||||
|
||||
fn candidates(&self, request: &AccountAcquireRequest) -> Vec<Account> {
|
||||
let mut candidates = self
|
||||
.accounts
|
||||
.values()
|
||||
.filter(|account| self.is_base_available(account, request))
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
if let Some(required_account_id) = &request.required_account_id {
|
||||
candidates.retain(|account| account.id == *required_account_id);
|
||||
return candidates;
|
||||
}
|
||||
if let Some(best_tier) = self.best_available_tier(&candidates) {
|
||||
candidates.retain(|account| account.plan_type.as_deref() == Some(best_tier.as_str()));
|
||||
}
|
||||
candidates
|
||||
}
|
||||
|
||||
fn is_base_available(&self, account: &Account, request: &AccountAcquireRequest) -> bool {
|
||||
account.status == AccountStatus::Active
|
||||
&& self.slot_count(&account.id) < self.options.max_concurrent_per_account
|
||||
&& AccountService::quota_available_at(
|
||||
account,
|
||||
request.now,
|
||||
self.options.skip_quota_limited,
|
||||
)
|
||||
&& self.is_model_allowed(account, &request.model)
|
||||
&& !request
|
||||
.exclude_account_ids
|
||||
.iter()
|
||||
.any(|account_id| account_id == &account.id)
|
||||
&& AccountService::cloudflare_available_at(account, request.now)
|
||||
}
|
||||
|
||||
fn slot_count(&self, account_id: &str) -> usize {
|
||||
self.slots.get(account_id).map_or(0, VecDeque::len)
|
||||
}
|
||||
@@ -864,13 +743,6 @@ impl AccountPool {
|
||||
});
|
||||
}
|
||||
|
||||
fn is_model_refresh_available(&self, account: &Account, now: DateTime<Utc>) -> bool {
|
||||
account.status == AccountStatus::Active
|
||||
&& self.slot_count(&account.id) < self.options.max_concurrent_per_account
|
||||
&& AccountService::quota_available_at(account, now, self.options.skip_quota_limited)
|
||||
&& AccountService::cloudflare_available_at(account, now)
|
||||
}
|
||||
|
||||
fn mark_request_usage(&mut self, account_id: &str, now: DateTime<Utc>) -> Option<Account> {
|
||||
let account = self.accounts.get_mut(account_id)?;
|
||||
account.last_used_at = Some(now.to_rfc3339());
|
||||
@@ -920,24 +792,33 @@ impl AccountPool {
|
||||
});
|
||||
}
|
||||
|
||||
fn is_model_allowed(&self, account: &Account, model: &str) -> bool {
|
||||
let Some(allowed_plans) = self.options.model_plan_allowlist.get(model) else {
|
||||
return true;
|
||||
};
|
||||
let Some(plan_type) = account.plan_type.as_deref() else {
|
||||
return false;
|
||||
};
|
||||
allowed_plans.iter().any(|plan| plan == plan_type)
|
||||
|| !self.options.fetched_model_plan_types.contains(plan_type)
|
||||
/// 过滤出本次请求的可用候选集合(委托 [`candidates::filter`])。
|
||||
fn filter_candidates(&self, request: &AccountAcquireRequest) -> Vec<Account> {
|
||||
let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len);
|
||||
candidates::filter(
|
||||
self.accounts.values(),
|
||||
&CandidateFilter {
|
||||
max_concurrent_per_account: self.options.max_concurrent_per_account,
|
||||
skip_quota_limited: self.options.skip_quota_limited,
|
||||
tier_priority: &self.options.tier_priority,
|
||||
model_plan_allowlist: &self.options.model_plan_allowlist,
|
||||
fetched_model_plan_types: &self.options.fetched_model_plan_types,
|
||||
slot_count: &slot_count,
|
||||
},
|
||||
&CandidateRequest {
|
||||
model: &request.model,
|
||||
exclude_account_ids: &request.exclude_account_ids,
|
||||
required_account_id: request.required_account_id.as_deref(),
|
||||
now: request.now,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
fn best_available_tier(&self, candidates: &[Account]) -> Option<String> {
|
||||
self.options.tier_priority.iter().find_map(|tier| {
|
||||
candidates
|
||||
.iter()
|
||||
.find(|account| account.plan_type.as_deref() == Some(tier.as_str()))
|
||||
.and_then(|account| account.plan_type.clone())
|
||||
})
|
||||
/// 按当前策略从候选集合中选择一个账号(委托 [`AccountScheduler::select`])。
|
||||
fn select_candidate(&self, candidates: &[Account], now: DateTime<Utc>) -> Option<Account> {
|
||||
let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len);
|
||||
self.scheduler
|
||||
.select(self.options.rotation_strategy, candidates, &slot_count, now)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1039,10 +920,29 @@ impl RuntimeAccountPoolService {
|
||||
|
||||
/// 释放账号在途槽位。
|
||||
pub async fn release(&self, account_id: &str) {
|
||||
let released = self.pool.lock().await.release(account_id);
|
||||
self.release_slot(account_id, true).await;
|
||||
}
|
||||
|
||||
/// 释放账号在途槽位,不累计请求用量。
|
||||
pub async fn release_without_request_usage(&self, account_id: &str) {
|
||||
self.release_slot(account_id, false).await;
|
||||
}
|
||||
|
||||
async fn release_slot(&self, account_id: &str, record_request_usage: bool) {
|
||||
let released = if record_request_usage {
|
||||
self.pool.lock().await.release(account_id)
|
||||
} else {
|
||||
self.pool
|
||||
.lock()
|
||||
.await
|
||||
.release_without_request_usage(account_id)
|
||||
};
|
||||
let Some(released) = released else {
|
||||
return;
|
||||
};
|
||||
if !record_request_usage {
|
||||
return;
|
||||
}
|
||||
if let Err(error) = self.store.record_request(account_id).await {
|
||||
tracing::warn!(
|
||||
account_id,
|
||||
@@ -1478,6 +1378,13 @@ fn refresh_quota_window(account: &mut Account, now: DateTime<Utc>) {
|
||||
}
|
||||
}
|
||||
|
||||
fn status_after_quota_limit(status: AccountStatus, quota_limit_reached: bool) -> AccountStatus {
|
||||
match (status, quota_limit_reached) {
|
||||
(AccountStatus::Active, true) => AccountStatus::QuotaExhausted,
|
||||
(status, _) => status,
|
||||
}
|
||||
}
|
||||
|
||||
fn reset_window_counters(account: &mut Account) {
|
||||
account.window_request_count = 0;
|
||||
account.window_input_tokens = 0;
|
||||
@@ -1520,37 +1427,3 @@ fn refresh_cloudflare_cooldown(account: &mut Account, now: DateTime<Utc>) {
|
||||
account.cloudflare_cooldown_until = None;
|
||||
}
|
||||
}
|
||||
|
||||
fn compare_quota_reset_priority(a: &Account, b: &Account) -> Ordering {
|
||||
a.quota_limit_reached
|
||||
.cmp(&b.quota_limit_reached)
|
||||
.then_with(|| compare_window_reset(a.window_reset_at, b.window_reset_at))
|
||||
.then_with(|| a.request_count.cmp(&b.request_count))
|
||||
.then_with(|| compare_last_used(a.last_used_at.as_deref(), b.last_used_at.as_deref()))
|
||||
}
|
||||
|
||||
fn account_window_token_count(account: &Account) -> u64 {
|
||||
account
|
||||
.window_input_tokens
|
||||
.saturating_add(account.window_output_tokens)
|
||||
.saturating_add(account.window_image_input_tokens)
|
||||
.saturating_add(account.window_image_output_tokens)
|
||||
}
|
||||
|
||||
fn compare_window_reset(a: Option<DateTime<Utc>>, b: Option<DateTime<Utc>>) -> Ordering {
|
||||
match (a, b) {
|
||||
(Some(a), Some(b)) => a.cmp(&b),
|
||||
(Some(_) | None, None) | (None, Some(_)) => Ordering::Equal,
|
||||
}
|
||||
}
|
||||
|
||||
fn compare_last_used(a: Option<&str>, b: Option<&str>) -> Ordering {
|
||||
last_used_millis(a).cmp(&last_used_millis(b))
|
||||
}
|
||||
|
||||
fn last_used_millis(value: Option<&str>) -> i64 {
|
||||
value
|
||||
.and_then(|value| parse_rfc3339_utc(value).ok())
|
||||
.map(|datetime| datetime.timestamp_millis())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
@@ -540,6 +540,7 @@ impl SqliteAccountStore {
|
||||
.bind(optional_update_value(&update.plan_type))
|
||||
.bind(update.status.is_some())
|
||||
.bind(status)
|
||||
.bind(status)
|
||||
.bind(&now)
|
||||
.bind(account_id)
|
||||
.execute(&self.pool)
|
||||
@@ -703,6 +704,9 @@ impl SqliteAccountStore {
|
||||
let cloudflare_cooldown_until = account.cloudflare_cooldown_until.map(|dt| dt.to_rfc3339());
|
||||
let mut tx = self.pool.begin().await?;
|
||||
let result = sqlx::query(SYNC_RUNTIME_ACCOUNT_STATE_SQL)
|
||||
.bind(quota_limit_reached)
|
||||
.bind(&now)
|
||||
.bind(quota_limit_reached)
|
||||
.bind(status_to_db(account.status))
|
||||
.bind(quota_limit_reached)
|
||||
.bind(&now)
|
||||
@@ -773,12 +777,24 @@ impl SqliteAccountStore {
|
||||
status: AccountStatus,
|
||||
) -> SqliteAccountStoreResult<bool> {
|
||||
let now = Utc::now().to_rfc3339();
|
||||
let result = sqlx::query("update accounts set status = ?, updated_at = ? where id = ?")
|
||||
.bind(status_to_db(status))
|
||||
.bind(&now)
|
||||
.bind(account_id)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
let status = status_to_db(status);
|
||||
let result = sqlx::query(
|
||||
r"
|
||||
update accounts
|
||||
set
|
||||
status = case
|
||||
when ? = 'active' and quota_limit_reached = 1 then 'quota_exhausted'
|
||||
else ?
|
||||
end,
|
||||
updated_at = ?
|
||||
where id = ?",
|
||||
)
|
||||
.bind(status)
|
||||
.bind(status)
|
||||
.bind(&now)
|
||||
.bind(account_id)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
|
||||
|
||||
@@ -165,7 +165,13 @@ set
|
||||
chatgpt_user_id = case when ? then ? else chatgpt_user_id end,
|
||||
label = case when ? then ? else label end,
|
||||
plan_type = case when ? then ? else plan_type end,
|
||||
status = case when ? then ? else status end,
|
||||
status = case
|
||||
when ? then case
|
||||
when ? = 'active' and quota_limit_reached = 1 then 'quota_exhausted'
|
||||
else ?
|
||||
end
|
||||
else status
|
||||
end,
|
||||
updated_at = ?
|
||||
where id = ?";
|
||||
|
||||
@@ -354,7 +360,7 @@ set
|
||||
quota_fetched_at = ?,
|
||||
plan_type = coalesce(?, plan_type),
|
||||
status = case
|
||||
when status in ('disabled', 'banned', 'expired', 'refreshing') then status
|
||||
when status in ('disabled', 'banned', 'expired') then status
|
||||
when ? = 1 then 'quota_exhausted'
|
||||
when status = 'quota_exhausted' then 'active'
|
||||
else status
|
||||
@@ -413,6 +419,10 @@ on conflict(account_id) do update set
|
||||
pub(super) const MARK_QUOTA_LIMITED_UNTIL_SQL: &str = r"
|
||||
update accounts
|
||||
set
|
||||
status = case
|
||||
when status in ('disabled', 'banned', 'expired') then status
|
||||
else 'quota_exhausted'
|
||||
end,
|
||||
quota_limit_reached = 1,
|
||||
quota_verify_required = 0,
|
||||
quota_cooldown_until = ?,
|
||||
@@ -423,7 +433,13 @@ pub(super) const SYNC_RUNTIME_ACCOUNT_STATE_SQL: &str = r"
|
||||
update accounts
|
||||
set
|
||||
status = case
|
||||
when status in ('disabled', 'banned', 'refreshing') then status
|
||||
when status in ('disabled', 'banned') then status
|
||||
when (
|
||||
case
|
||||
when ? = 0 and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_limit_reached
|
||||
else ?
|
||||
end
|
||||
) = 1 then 'quota_exhausted'
|
||||
else ?
|
||||
end,
|
||||
quota_limit_reached = case
|
||||
@@ -496,6 +512,7 @@ set
|
||||
next_refresh_at = ?,
|
||||
status = case
|
||||
when status in ('disabled', 'banned') then status
|
||||
when quota_limit_reached = 1 then 'quota_exhausted'
|
||||
else ?
|
||||
end,
|
||||
updated_at = ?
|
||||
@@ -513,6 +530,7 @@ set
|
||||
next_refresh_at = ?,
|
||||
status = case
|
||||
when status in ('disabled', 'banned') then status
|
||||
when quota_limit_reached = 1 then 'quota_exhausted'
|
||||
else ?
|
||||
end,
|
||||
updated_at = ?
|
||||
|
||||
@@ -255,7 +255,6 @@ pub(super) fn status_from_db(value: &str) -> SqliteAccountStoreResult<AccountSta
|
||||
"active" => Ok(AccountStatus::Active),
|
||||
"expired" => Ok(AccountStatus::Expired),
|
||||
"quota_exhausted" => Ok(AccountStatus::QuotaExhausted),
|
||||
"refreshing" => Ok(AccountStatus::Refreshing),
|
||||
"disabled" => Ok(AccountStatus::Disabled),
|
||||
"banned" => Ok(AccountStatus::Banned),
|
||||
other => Err(SqliteAccountStoreError::InvalidStatus(other.to_string())),
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
//! 刷新租约存储、JWT 解码与账号 claims 验证。
|
||||
|
||||
use std::{
|
||||
collections::hash_map::DefaultHasher,
|
||||
collections::{hash_map::DefaultHasher, HashSet},
|
||||
hash::{Hash, Hasher},
|
||||
sync::{Arc, RwLock},
|
||||
};
|
||||
@@ -369,22 +369,53 @@ where
|
||||
trigger: RefreshTrigger,
|
||||
now: DateTime<Utc>,
|
||||
) -> bool {
|
||||
if account.status != AccountStatus::Active {
|
||||
return false;
|
||||
}
|
||||
|
||||
match trigger {
|
||||
RefreshTrigger::Unauthorized => true,
|
||||
RefreshTrigger::BeforeExpiry => account
|
||||
.access_token_expires_at
|
||||
.is_some_and(|expires_at| expires_at <= now + self.refresh_margin()),
|
||||
RefreshTrigger::Unauthorized => can_refresh_after_auth_failure(account.status),
|
||||
RefreshTrigger::BeforeExpiry => {
|
||||
can_refresh_before_expiry(account.status)
|
||||
&& self.should_refresh_before_expiry(account, now)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn refresh_margin(&self) -> Duration {
|
||||
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();
|
||||
let seconds = policy.refresh_margin_seconds.min(86_400 * 7) as i64;
|
||||
Duration::seconds(seconds)
|
||||
jittered_refresh_at(
|
||||
account.id.as_str(),
|
||||
expires_at,
|
||||
policy.refresh_margin_seconds,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn can_refresh_after_auth_failure(status: AccountStatus) -> bool {
|
||||
matches!(
|
||||
status,
|
||||
AccountStatus::Active | AccountStatus::Expired | AccountStatus::QuotaExhausted
|
||||
)
|
||||
}
|
||||
|
||||
fn can_refresh_before_expiry(status: AccountStatus) -> bool {
|
||||
matches!(
|
||||
status,
|
||||
AccountStatus::Active | AccountStatus::QuotaExhausted
|
||||
)
|
||||
}
|
||||
|
||||
fn status_after_successful_token_refresh(status: AccountStatus) -> AccountStatus {
|
||||
match status {
|
||||
AccountStatus::Expired => AccountStatus::Active,
|
||||
status => status,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -398,7 +429,7 @@ pub fn apply_token_pair(account: &Account, token_pair: TokenPair) -> Account {
|
||||
refreshed.refresh_token = Some(refresh_token);
|
||||
}
|
||||
|
||||
refreshed.status = AccountStatus::Active;
|
||||
refreshed.status = status_after_successful_token_refresh(account.status);
|
||||
refreshed
|
||||
}
|
||||
|
||||
@@ -418,7 +449,7 @@ pub fn apply_refresh_failure(account: &Account, failure: RefreshFailure) -> Acco
|
||||
// 刷新租约存储(SQLite)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
use sqlx::SqlitePool;
|
||||
use sqlx::{QueryBuilder, Sqlite, SqlitePool};
|
||||
use thiserror::Error;
|
||||
|
||||
/// SQLite 刷新租约存储结果。
|
||||
@@ -485,6 +516,35 @@ where account_refresh_leases.expires_at <= ?
|
||||
.await?;
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
|
||||
/// 返回给定账号集合中仍有效的刷新租约。
|
||||
pub async fn active_account_ids(
|
||||
&self,
|
||||
account_ids: &[String],
|
||||
now: DateTime<Utc>,
|
||||
) -> RefreshLeaseStoreResult<HashSet<String>> {
|
||||
if account_ids.is_empty() {
|
||||
return Ok(HashSet::new());
|
||||
}
|
||||
|
||||
let mut builder = QueryBuilder::<Sqlite>::new(
|
||||
"select account_id from account_refresh_leases where datetime(expires_at) > datetime(",
|
||||
);
|
||||
builder.push_bind(now.to_rfc3339());
|
||||
builder.push(") and 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_query_as::<(String,)>()
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(rows.into_iter().map(|(account_id,)| account_id).collect())
|
||||
}
|
||||
}
|
||||
|
||||
mod runtime;
|
||||
|
||||
@@ -127,6 +127,41 @@ where
|
||||
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,
|
||||
@@ -165,7 +200,6 @@ where
|
||||
}
|
||||
|
||||
let timer = match account.status {
|
||||
AccountStatus::Refreshing => Some((Duration::ZERO, RefreshTrigger::Unauthorized)),
|
||||
AccountStatus::Expired => {
|
||||
let delay = match account.next_refresh_at {
|
||||
Some(next_refresh_at) if next_refresh_at > now => {
|
||||
@@ -543,7 +577,7 @@ where
|
||||
account: &Account,
|
||||
now: DateTime<Utc>,
|
||||
) -> Option<RefreshTrigger> {
|
||||
if account.status == AccountStatus::Refreshing {
|
||||
if account.status == AccountStatus::Expired {
|
||||
return Some(RefreshTrigger::Unauthorized);
|
||||
}
|
||||
|
||||
@@ -562,7 +596,7 @@ where
|
||||
let mut permanent_failures = 0;
|
||||
|
||||
for attempt_index in 0..max_attempts {
|
||||
let Some(attempt) = self.prepare_refresh_attempt(account, now).await? else {
|
||||
let Some(attempt) = self.prepare_refresh_attempt(account, trigger, now).await? else {
|
||||
return Ok(TokenRefreshOutcome::Skipped);
|
||||
};
|
||||
|
||||
@@ -577,32 +611,34 @@ where
|
||||
.refresh_account_at(&attempt, trigger, now)
|
||||
.await
|
||||
{
|
||||
Ok(mut updated) if account_requires_token_write(&attempt, &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);
|
||||
persist_token_update(&self.store, &updated).await?;
|
||||
return Ok(TokenRefreshOutcome::Refreshed(Box::new(updated)));
|
||||
}
|
||||
Ok(updated) if attempt.status != updated.status => {
|
||||
if is_permanent_refresh_status(updated.status) {
|
||||
permanent_failures += 1;
|
||||
if permanent_failures < PERMANENT_FAILURE_CONFIRMATION_THRESHOLD
|
||||
&& self.sleep_before_retry(&attempt.id, attempt_index).await
|
||||
{
|
||||
continue;
|
||||
}
|
||||
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(_) => return Ok(TokenRefreshOutcome::Skipped),
|
||||
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);
|
||||
persist_token_update(&self.store, &updated).await?;
|
||||
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, AccountStatus::Active).await?;
|
||||
persist_status(
|
||||
&self.store,
|
||||
&attempt.id,
|
||||
status_after_temporary_refresh_failure(attempt.status),
|
||||
)
|
||||
.await?;
|
||||
let recovery_delay = stable_jittered_duration(
|
||||
&attempt.id,
|
||||
Duration::from_secs(RECOVERY_DELAY_SECONDS as u64),
|
||||
@@ -624,6 +660,7 @@ where
|
||||
async fn prepare_refresh_attempt(
|
||||
&self,
|
||||
account: &Account,
|
||||
trigger: RefreshTrigger,
|
||||
now: DateTime<Utc>,
|
||||
) -> TokenRefreshServiceResult<Option<Account>> {
|
||||
let Some(stored) = self.store.get(&account.id).await? else {
|
||||
@@ -631,7 +668,7 @@ where
|
||||
account.id.clone(),
|
||||
));
|
||||
};
|
||||
let mut attempt = stored_account_to_refresh_account(stored);
|
||||
let attempt = stored_account_to_refresh_account(stored);
|
||||
if matches!(
|
||||
attempt.status,
|
||||
AccountStatus::Disabled | AccountStatus::Banned
|
||||
@@ -641,11 +678,13 @@ where
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if attempt.status != AccountStatus::Refreshing {
|
||||
persist_status(&self.store, &attempt.id, AccountStatus::Refreshing).await?;
|
||||
if !self
|
||||
.scheduler
|
||||
.should_refresh_account_at(&attempt, trigger, now)
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
attempt.status = AccountStatus::Active;
|
||||
Ok(Some(attempt))
|
||||
}
|
||||
|
||||
@@ -839,20 +878,16 @@ async fn persist_next_refresh_at(
|
||||
}
|
||||
}
|
||||
|
||||
fn account_requires_token_write(previous: &Account, updated: &Account) -> bool {
|
||||
previous.access_token != updated.access_token
|
||||
|| previous.refresh_token != updated.refresh_token
|
||||
|| (updated.status == AccountStatus::Active && previous.status != updated.status)
|
||||
fn is_permanent_refresh_failure_status(status: AccountStatus) -> bool {
|
||||
matches!(status, AccountStatus::Expired | AccountStatus::Banned)
|
||||
}
|
||||
|
||||
fn is_permanent_refresh_status(status: AccountStatus) -> bool {
|
||||
matches!(
|
||||
status,
|
||||
AccountStatus::Expired
|
||||
| AccountStatus::QuotaExhausted
|
||||
| AccountStatus::Disabled
|
||||
| AccountStatus::Banned
|
||||
)
|
||||
fn status_after_temporary_refresh_failure(status: AccountStatus) -> AccountStatus {
|
||||
match status {
|
||||
AccountStatus::Expired | AccountStatus::QuotaExhausted => status,
|
||||
AccountStatus::Active => AccountStatus::Active,
|
||||
AccountStatus::Disabled | AccountStatus::Banned => status,
|
||||
}
|
||||
}
|
||||
|
||||
fn refresh_lease_owner() -> String {
|
||||
|
||||
@@ -4,5 +4,6 @@ pub mod accounts;
|
||||
pub mod fingerprint;
|
||||
pub mod models;
|
||||
pub mod protocol;
|
||||
pub mod scheduler;
|
||||
pub mod token_client;
|
||||
pub mod transport;
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
//! 候选账号过滤。
|
||||
//!
|
||||
//! 所有调度策略共享的前置:从账号池中筛出「本次请求可用」的候选集合。这是选择流程
|
||||
//! 的第一步,与具体策略无关——策略只在过滤后的候选之间做选择。
|
||||
//!
|
||||
//! 过滤规则(全部满足才入选):
|
||||
//!
|
||||
//! 1. 账号状态为 [`AccountStatus::Active`];
|
||||
//! 2. 在途槽位数未达单账号并发上限;
|
||||
//! 3. 配额可用([`AccountService::quota_available_at`],受 `skip_quota_limited` 控制);
|
||||
//! 4. 模型在账号订阅计划的允许列表内;
|
||||
//! 5. 不在本次请求的排除列表内;
|
||||
//! 6. 不处于 Cloudflare 冷却期。
|
||||
//!
|
||||
//! 之后再按订阅层级优先级(`tier_priority`)收窄到最高可用层级;若请求指定了
|
||||
//! `required_account_id`,则只保留该账号(跳过层级收窄)。
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
|
||||
use crate::upstream::accounts::model::{Account, AccountStatus};
|
||||
use crate::upstream::accounts::service::AccountService;
|
||||
|
||||
/// 候选过滤所需的账号池策略参数(只读视图)。
|
||||
pub struct CandidateFilter<'a> {
|
||||
/// 单账号最大并发请求数。
|
||||
pub max_concurrent_per_account: usize,
|
||||
/// 是否跳过已标记配额耗尽的账号。
|
||||
pub skip_quota_limited: bool,
|
||||
/// 订阅层级优先级,越靠前优先级越高。
|
||||
pub tier_priority: &'a [String],
|
||||
/// 模型到允许订阅计划的映射。
|
||||
pub model_plan_allowlist: &'a std::collections::BTreeMap<String, Vec<String>>,
|
||||
/// 已成功拉取过模型列表的订阅计划。
|
||||
pub fetched_model_plan_types: &'a std::collections::BTreeSet<String>,
|
||||
/// 按账号 ID 读取当前在途槽位数。
|
||||
pub slot_count: &'a dyn Fn(&str) -> usize,
|
||||
}
|
||||
|
||||
/// 本次账号获取请求的过滤条件。
|
||||
pub struct CandidateRequest<'a> {
|
||||
/// 本次请求使用的模型。
|
||||
pub model: &'a str,
|
||||
/// 本次请求需要排除的账号 ID。
|
||||
pub exclude_account_ids: &'a [String],
|
||||
/// 本次请求必须使用的账号 ID(存在时跳过层级收窄,只保留该账号)。
|
||||
pub required_account_id: Option<&'a str>,
|
||||
/// 本次调度使用的当前时间。
|
||||
pub now: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// 从账号池筛出本次请求的可用候选集合。
|
||||
pub fn filter<'a>(
|
||||
accounts: impl Iterator<Item = &'a Account>,
|
||||
filter: &CandidateFilter<'_>,
|
||||
request: &CandidateRequest<'_>,
|
||||
) -> Vec<Account> {
|
||||
let mut candidates = accounts
|
||||
.filter(|account| is_base_available(account, filter, request))
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
if let Some(required_account_id) = request.required_account_id {
|
||||
candidates.retain(|account| account.id == required_account_id);
|
||||
return candidates;
|
||||
}
|
||||
|
||||
if let Some(best_tier) = best_available_tier(&candidates, filter.tier_priority) {
|
||||
candidates.retain(|account| account.plan_type.as_deref() == Some(best_tier.as_str()));
|
||||
}
|
||||
candidates
|
||||
}
|
||||
|
||||
/// 判断账号是否满足全部基础可用条件。
|
||||
fn is_base_available(
|
||||
account: &Account,
|
||||
filter: &CandidateFilter<'_>,
|
||||
request: &CandidateRequest<'_>,
|
||||
) -> 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)
|
||||
&& is_model_allowed(
|
||||
account,
|
||||
request.model,
|
||||
filter.model_plan_allowlist,
|
||||
filter.fetched_model_plan_types,
|
||||
)
|
||||
&& !request
|
||||
.exclude_account_ids
|
||||
.iter()
|
||||
.any(|account_id| account_id == &account.id)
|
||||
&& AccountService::cloudflare_available_at(account, request.now)
|
||||
}
|
||||
|
||||
/// 判断模型是否允许该账号的订阅计划使用。
|
||||
fn is_model_allowed(
|
||||
account: &Account,
|
||||
model: &str,
|
||||
model_plan_allowlist: &std::collections::BTreeMap<String, Vec<String>>,
|
||||
fetched_model_plan_types: &std::collections::BTreeSet<String>,
|
||||
) -> bool {
|
||||
let Some(allowed_plans) = model_plan_allowlist.get(model) else {
|
||||
return true;
|
||||
};
|
||||
let Some(plan_type) = account.plan_type.as_deref() else {
|
||||
return false;
|
||||
};
|
||||
allowed_plans.iter().any(|plan| plan == plan_type)
|
||||
|| !fetched_model_plan_types.contains(plan_type)
|
||||
}
|
||||
|
||||
/// 在候选集合中找出订阅层级优先级最高的可用层级。
|
||||
fn best_available_tier(candidates: &[Account], tier_priority: &[String]) -> Option<String> {
|
||||
for tier in tier_priority {
|
||||
let matched = candidates
|
||||
.iter()
|
||||
.any(|account| account.plan_type.as_deref() == Some(tier.as_str()));
|
||||
if matched {
|
||||
return Some(tier.clone());
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// 判断账号是否可用于模型列表刷新(不涉及排除列表 / 模型允许列表 / 层级收窄)。
|
||||
pub fn is_model_refresh_available(
|
||||
account: &Account,
|
||||
max_concurrent_per_account: usize,
|
||||
skip_quota_limited: bool,
|
||||
slot_count: &dyn Fn(&str) -> usize,
|
||||
now: DateTime<Utc>,
|
||||
) -> 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)
|
||||
}
|
||||
@@ -0,0 +1,171 @@
|
||||
//! 账号运行时反馈:错误率与 TTFT 的 EWMA 统计。
|
||||
//!
|
||||
//! 每个账号维护两个指数加权移动平均(EWMA):错误率与首 token 延迟(TTFT,
|
||||
//! 毫秒)。请求结束后经 [`FeedbackStats::report`] 回灌,直接进 [`super::score`]
|
||||
//! 打分并可驱动 sticky escape。统计以 `AtomicU64`(存 `f64` 位模式)+ CAS 无锁更新,
|
||||
//! 对齐 sub2api 的设计,避免选择热路径上的锁与 DB 往返。
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::RwLock;
|
||||
|
||||
/// EWMA 平滑系数(新样本权重)。与 sub2api 一致取 0.2。
|
||||
const EWMA_ALPHA: f64 = 0.2;
|
||||
|
||||
/// 单账号的运行时 EWMA 统计。
|
||||
///
|
||||
/// `*_bits` 存 `f64` 的位模式;`NaN`(全 1 的哨兵)表示"尚无样本"。
|
||||
#[derive(Debug)]
|
||||
struct AccountFeedback {
|
||||
/// 错误率 EWMA,范围 `[0, 1]`。
|
||||
error_rate_bits: AtomicU64,
|
||||
/// TTFT EWMA,单位毫秒。
|
||||
ttft_ms_bits: AtomicU64,
|
||||
}
|
||||
|
||||
/// 表示"尚无样本"的哨兵位模式。
|
||||
const EMPTY_BITS: u64 = f64::NAN.to_bits();
|
||||
|
||||
impl Default for AccountFeedback {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
error_rate_bits: AtomicU64::new(EMPTY_BITS),
|
||||
ttft_ms_bits: AtomicU64::new(EMPTY_BITS),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AccountFeedback {
|
||||
/// 读取当前错误率 EWMA;无样本时返回 `None`。
|
||||
fn error_rate(&self) -> Option<f64> {
|
||||
load_sample(&self.error_rate_bits)
|
||||
}
|
||||
|
||||
/// 读取当前 TTFT EWMA(毫秒);无样本时返回 `None`。
|
||||
fn ttft_ms(&self) -> Option<f64> {
|
||||
load_sample(&self.ttft_ms_bits)
|
||||
}
|
||||
|
||||
/// 回灌一次请求结果。
|
||||
fn report(&self, success: bool, first_token_ms: Option<u64>) {
|
||||
let error_sample = if success { 0.0 } else { 1.0 };
|
||||
update_ewma(&self.error_rate_bits, error_sample);
|
||||
if let Some(ms) = first_token_ms {
|
||||
update_ewma(&self.ttft_ms_bits, ms as f64);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 读取 EWMA 样本;哨兵(NaN)返回 `None`。
|
||||
fn load_sample(bits: &AtomicU64) -> Option<f64> {
|
||||
let value = f64::from_bits(bits.load(Ordering::Relaxed));
|
||||
if value.is_nan() {
|
||||
None
|
||||
} else {
|
||||
Some(value)
|
||||
}
|
||||
}
|
||||
|
||||
/// 无锁 CAS 更新 EWMA:首个样本直接落值,其后按 alpha 混合。
|
||||
fn update_ewma(bits: &AtomicU64, sample: f64) {
|
||||
let mut current = bits.load(Ordering::Relaxed);
|
||||
loop {
|
||||
let previous = f64::from_bits(current);
|
||||
let next = if previous.is_nan() {
|
||||
sample
|
||||
} else {
|
||||
EWMA_ALPHA * sample + (1.0 - EWMA_ALPHA) * previous
|
||||
};
|
||||
match bits.compare_exchange_weak(
|
||||
current,
|
||||
next.to_bits(),
|
||||
Ordering::Relaxed,
|
||||
Ordering::Relaxed,
|
||||
) {
|
||||
Ok(_) => return,
|
||||
Err(observed) => current = observed,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 账号 EWMA 反馈的读数快照,供打分使用。
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq)]
|
||||
pub struct FeedbackSample {
|
||||
/// 错误率 EWMA,范围 `[0, 1]`;无样本时为 `None`。
|
||||
pub error_rate: Option<f64>,
|
||||
/// TTFT EWMA(毫秒);无样本时为 `None`。
|
||||
pub ttft_ms: Option<f64>,
|
||||
}
|
||||
|
||||
/// 账号运行时反馈存储:按账号 ID 聚合 EWMA 统计。
|
||||
///
|
||||
/// 选择热路径只读(`sample`),请求结束后写(`report`)。内部按账号分片为独立的
|
||||
/// 原子对象,`RwLock` 仅保护 map 结构本身(插入新账号时短暂写锁)。
|
||||
#[derive(Debug, Default)]
|
||||
pub struct FeedbackStats {
|
||||
accounts: RwLock<HashMap<String, AccountFeedback>>,
|
||||
}
|
||||
|
||||
impl FeedbackStats {
|
||||
/// 创建空的反馈存储。
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// 读取指定账号的 EWMA 快照;账号无记录时返回全 `None`。
|
||||
pub fn sample(&self, account_id: &str) -> FeedbackSample {
|
||||
let accounts = self
|
||||
.accounts
|
||||
.read()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
match accounts.get(account_id) {
|
||||
Some(feedback) => FeedbackSample {
|
||||
error_rate: feedback.error_rate(),
|
||||
ttft_ms: feedback.ttft_ms(),
|
||||
},
|
||||
None => FeedbackSample::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 回灌一次请求结果:成功/失败 + 首 token 延迟(毫秒)。
|
||||
pub fn report(&self, account_id: &str, success: bool, first_token_ms: Option<u64>) {
|
||||
// 快路径:账号已存在时只需读锁 + 原子更新。
|
||||
{
|
||||
let accounts = self
|
||||
.accounts
|
||||
.read()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if let Some(feedback) = accounts.get(account_id) {
|
||||
feedback.report(success, first_token_ms);
|
||||
return;
|
||||
}
|
||||
}
|
||||
// 慢路径:首次见到该账号,取写锁插入后回灌。
|
||||
let mut accounts = self
|
||||
.accounts
|
||||
.write()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
accounts
|
||||
.entry(account_id.to_string())
|
||||
.or_default()
|
||||
.report(success, first_token_ms);
|
||||
}
|
||||
|
||||
/// 移除指定账号的反馈记录(账号被删除时调用)。
|
||||
pub fn remove(&self, account_id: &str) {
|
||||
let mut accounts = self
|
||||
.accounts
|
||||
.write()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
accounts.remove(account_id);
|
||||
}
|
||||
|
||||
/// 清空全部反馈记录(账号池清空时调用)。
|
||||
pub fn clear(&self) {
|
||||
let mut accounts = self
|
||||
.accounts
|
||||
.write()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
accounts.clear();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
//! 智能账号调度模块。
|
||||
//!
|
||||
//! 账号选择的唯一归属地。把「选择」从账号池存储中剥离,集中于此:
|
||||
//!
|
||||
//! - [`candidates`] —— 候选过滤(可用性 / 层级 / 模型允许列表),所有策略共享的前置。
|
||||
//! - [`strategy`] —— 全部选择策略,每个策略一个文件([`strategy::smart`] /
|
||||
//! [`strategy::quota_reset`] / [`strategy::round_robin`] / [`strategy::sticky`])。
|
||||
//! - [`feedback`] —— 跨策略的运行时 EWMA 反馈(错误率 / TTFT),喂给 Smart 打分。
|
||||
//!
|
||||
//! [`AccountScheduler`] 是对外门面:持有轮转游标、打分权重与 EWMA 反馈存储,接收
|
||||
//! **已过滤的候选切片**并返回选中账号。它不 lock 账号池、不刷状态、不写库——存储与
|
||||
//! 状态维护仍归 `AccountPool`,二者经候选切片与 `slot_count` 闭包解耦。
|
||||
|
||||
pub mod candidates;
|
||||
pub mod feedback;
|
||||
pub mod strategy;
|
||||
|
||||
use std::sync::Mutex;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
|
||||
use crate::upstream::accounts::model::Account;
|
||||
|
||||
pub use candidates::{CandidateFilter, CandidateRequest};
|
||||
pub use feedback::{FeedbackSample, FeedbackStats};
|
||||
pub use strategy::smart::rank_candidates;
|
||||
pub use strategy::{RotationStrategy, ScoreBreakdown, ScoreWeights};
|
||||
|
||||
/// 账号调度器:账号选择的对外门面。
|
||||
///
|
||||
/// 持有轮转游标(并列破并列用)、打分权重(Smart 用)与 EWMA 反馈存储。选择热路径
|
||||
/// 只读反馈与权重,请求结束后经 [`AccountScheduler::report_feedback`] 回灌。游标用
|
||||
/// [`Mutex`] 保护(选择需可变,但 `AccountScheduler` 以共享引用参与并发请求)。
|
||||
#[derive(Debug)]
|
||||
pub struct AccountScheduler {
|
||||
/// 轮转游标:在并列最优账号之间做公平轮转。
|
||||
cursor: Mutex<usize>,
|
||||
/// Smart 策略的打分权重。
|
||||
weights: ScoreWeights,
|
||||
/// 账号运行时 EWMA 反馈(错误率 / TTFT)。
|
||||
feedback: FeedbackStats,
|
||||
}
|
||||
|
||||
impl Default for AccountScheduler {
|
||||
fn default() -> Self {
|
||||
Self::new(ScoreWeights::default())
|
||||
}
|
||||
}
|
||||
|
||||
impl AccountScheduler {
|
||||
/// 使用指定打分权重创建调度器。
|
||||
pub fn new(weights: ScoreWeights) -> Self {
|
||||
Self {
|
||||
cursor: Mutex::new(0),
|
||||
weights,
|
||||
feedback: FeedbackStats::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 替换打分权重(策略切换或热更时调用)。
|
||||
pub fn set_weights(&mut self, weights: ScoreWeights) {
|
||||
self.weights = weights;
|
||||
}
|
||||
|
||||
/// 重置轮转游标(策略切换时调用,保证轮转从头开始)。
|
||||
pub fn reset_cursor(&self) {
|
||||
*self.lock_cursor() = 0;
|
||||
}
|
||||
|
||||
/// 从已过滤的候选集合中按指定策略选择一个账号。
|
||||
///
|
||||
/// `candidates` 必须是 [`candidates::filter`] 的输出;`slot_count` 按账号 ID 读取
|
||||
/// 当前在途槽位数(Smart 打分与候选无关的运行时状态由调用方提供)。
|
||||
pub fn select(
|
||||
&self,
|
||||
strategy: RotationStrategy,
|
||||
candidates: &[Account],
|
||||
slot_count: &dyn Fn(&str) -> usize,
|
||||
now: DateTime<Utc>,
|
||||
) -> Option<Account> {
|
||||
let input = strategy::SelectionInput {
|
||||
candidates,
|
||||
slot_count,
|
||||
feedback: &self.feedback,
|
||||
weights: &self.weights,
|
||||
now,
|
||||
};
|
||||
let mut cursor = self.lock_cursor();
|
||||
strategy::select(strategy, &input, &mut cursor)
|
||||
}
|
||||
|
||||
/// 回灌一次请求结果到运行时 EWMA 反馈(错误率 / TTFT),供 Smart 打分。
|
||||
pub fn report_feedback(&self, account_id: &str, success: bool, first_token_ms: Option<u64>) {
|
||||
self.feedback.report(account_id, success, first_token_ms);
|
||||
}
|
||||
|
||||
/// 移除指定账号的运行时反馈(账号被删除时调用)。
|
||||
pub fn forget_feedback(&self, account_id: &str) {
|
||||
self.feedback.remove(account_id);
|
||||
}
|
||||
|
||||
/// 清空全部运行时反馈(账号池清空时调用)。
|
||||
pub fn clear_feedback(&self) {
|
||||
self.feedback.clear();
|
||||
}
|
||||
|
||||
fn lock_cursor(&self) -> std::sync::MutexGuard<'_, usize> {
|
||||
self.cursor
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
//! 账号调度策略层。
|
||||
//!
|
||||
//! 所有账号选择策略集中于此,每个策略一个文件,互不依赖:
|
||||
//!
|
||||
//! - [`smart`] —— 加权线性打分 + EWMA 反馈(“智能调度”)。
|
||||
//! - [`quota_reset`] —— 优先消耗即将重置的额度窗口。
|
||||
//! - [`round_robin`] —— 按候选顺序循环。
|
||||
//! - [`sticky`] —— 优先最近使用过的账号。
|
||||
//!
|
||||
//! 本模块只承担三件事:定义 [`RotationStrategy`] 枚举、把选择请求分派到对应策略、
|
||||
//! 提供各策略共享的 tie-break 比较原语([`select_by`] / [`compare_window_reset`] /
|
||||
//! [`compare_last_used`])。策略本身不触碰账号池存储:在途槽位数经 `slot_count`
|
||||
//! 闭包按需读取,EWMA 反馈与打分权重经 [`SelectionInput`] 传入。
|
||||
|
||||
use std::cmp::Ordering;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
|
||||
use crate::infra::time::parse_rfc3339_utc;
|
||||
use crate::upstream::accounts::model::Account;
|
||||
use crate::upstream::scheduler::feedback::FeedbackStats;
|
||||
|
||||
pub mod quota_reset;
|
||||
pub mod round_robin;
|
||||
pub mod smart;
|
||||
pub mod sticky;
|
||||
|
||||
pub use smart::{ScoreBreakdown, ScoreWeights};
|
||||
|
||||
/// 账号轮转策略。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum RotationStrategy {
|
||||
/// 加权线性打分调度:多因子(负载/窗口/错误率/TTFT/重置)归一化后加权求和,
|
||||
/// 配合 EWMA 反馈。
|
||||
Smart,
|
||||
/// 优先消耗即将重置的额度窗口。
|
||||
QuotaResetPriority,
|
||||
/// 按候选账号顺序循环选择。
|
||||
RoundRobin,
|
||||
/// 优先选择最近使用过的账号。
|
||||
Sticky,
|
||||
}
|
||||
|
||||
impl RotationStrategy {
|
||||
/// 返回策略的持久化/日志字符串。
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Smart => "smart",
|
||||
Self::QuotaResetPriority => "quota_reset_priority",
|
||||
Self::RoundRobin => "round_robin",
|
||||
Self::Sticky => "sticky",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 策略选择的运行时输入。
|
||||
///
|
||||
/// `slot_count` 按账号 ID 读取当前在途槽位数;`feedback` / `weights` 仅 [`smart`]
|
||||
/// 使用,其余策略忽略;`now` 用于窗口重置相关的打分。
|
||||
pub struct SelectionInput<'a> {
|
||||
/// 已过滤的候选账号集合。
|
||||
pub candidates: &'a [Account],
|
||||
/// 按账号 ID 读取当前在途槽位数。
|
||||
pub slot_count: &'a dyn Fn(&str) -> usize,
|
||||
/// 账号 EWMA 反馈存储(Smart 打分用)。
|
||||
pub feedback: &'a FeedbackStats,
|
||||
/// 打分权重(Smart 打分用)。
|
||||
pub weights: &'a ScoreWeights,
|
||||
/// 当前时刻(窗口重置因子用)。
|
||||
pub now: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// 按指定策略从候选集合中选择一个账号。
|
||||
///
|
||||
/// `cursor` 是账号池持有的轮转游标,用于在并列最优账号之间做公平轮转;不同策略对它
|
||||
/// 的读写方式不同(详见各策略模块),本函数只负责透传可变引用。
|
||||
pub fn select(
|
||||
strategy: RotationStrategy,
|
||||
input: &SelectionInput<'_>,
|
||||
cursor: &mut usize,
|
||||
) -> Option<Account> {
|
||||
match strategy {
|
||||
RotationStrategy::Smart => smart::select(input, cursor),
|
||||
RotationStrategy::QuotaResetPriority => quota_reset::select(input, cursor),
|
||||
RotationStrategy::RoundRobin => round_robin::select(input, cursor),
|
||||
RotationStrategy::Sticky => sticky::select(input, cursor),
|
||||
}
|
||||
}
|
||||
|
||||
// ====================================================================
|
||||
// 共享 tie-break 原语
|
||||
// ====================================================================
|
||||
|
||||
/// 按比较函数升序排序后,在并列最优(`Ordering::Equal`)的账号之间用轮转游标破并列。
|
||||
///
|
||||
/// 这是 [`smart`] 与 [`quota_reset`] 共享的选择骨架:先按各自的比较规则排序,再在得分
|
||||
/// 相同的账号之间轮转,保证并列账号被公平地依次选中。
|
||||
pub(crate) fn select_by(
|
||||
candidates: &[Account],
|
||||
cursor: &mut usize,
|
||||
compare: impl Fn(&Account, &Account) -> Ordering,
|
||||
) -> Option<Account> {
|
||||
let mut sorted = candidates.iter().collect::<Vec<_>>();
|
||||
sorted.sort_by(|a, b| compare(a, b));
|
||||
let best = *sorted.first()?;
|
||||
let tied_count = sorted
|
||||
.iter()
|
||||
.take_while(|account| compare(best, account) == Ordering::Equal)
|
||||
.count();
|
||||
let index = *cursor % tied_count;
|
||||
*cursor = cursor.wrapping_add(1);
|
||||
Some((*sorted[index]).clone())
|
||||
}
|
||||
|
||||
/// 计算账号当前窗口的总 token 数(输入 + 输出 + 图片输入 + 图片输出)。
|
||||
pub(crate) fn account_window_token_count(account: &Account) -> u64 {
|
||||
account
|
||||
.window_input_tokens
|
||||
.saturating_add(account.window_output_tokens)
|
||||
.saturating_add(account.window_image_input_tokens)
|
||||
.saturating_add(account.window_image_output_tokens)
|
||||
}
|
||||
|
||||
/// 比较两个窗口重置时间:越早重置越优先;缺失重置时间视为并列(不惩罚)。
|
||||
pub(crate) fn compare_window_reset(a: Option<DateTime<Utc>>, b: Option<DateTime<Utc>>) -> Ordering {
|
||||
match (a, b) {
|
||||
(Some(a), Some(b)) => a.cmp(&b),
|
||||
(Some(_) | None, None) | (None, Some(_)) => Ordering::Equal,
|
||||
}
|
||||
}
|
||||
|
||||
/// 比较两个账号的最近使用时间:越久未用越优先(升序,最小者排前)。
|
||||
pub(crate) fn compare_last_used(a: Option<&str>, b: Option<&str>) -> Ordering {
|
||||
last_used_millis(a).cmp(&last_used_millis(b))
|
||||
}
|
||||
|
||||
fn last_used_millis(value: Option<&str>) -> i64 {
|
||||
value
|
||||
.and_then(|value| parse_rfc3339_utc(value).ok())
|
||||
.map(|datetime| datetime.timestamp_millis())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
//! QuotaResetPriority 策略:优先消耗即将重置的额度窗口。
|
||||
//!
|
||||
//! 排序规则(升序,越靠前越优先):
|
||||
//! 1. 未封顶配额优先于已封顶(`quota_limit_reached` false < true);
|
||||
//! 2. 窗口重置时间越早越优先;
|
||||
//! 3. 历史请求数越少越优先;
|
||||
//! 4. 最近使用时间越久越优先(LRU)。
|
||||
//!
|
||||
//! 并列最优的账号之间由 [`super::select_by`] 用轮转游标破并列。
|
||||
|
||||
use std::cmp::Ordering;
|
||||
|
||||
use crate::upstream::accounts::model::Account;
|
||||
|
||||
use super::{compare_last_used, compare_window_reset, select_by, SelectionInput};
|
||||
|
||||
/// QuotaResetPriority 策略选择。
|
||||
pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option<Account> {
|
||||
select_by(input.candidates, cursor, compare_quota_reset_priority)
|
||||
}
|
||||
|
||||
/// 配额重置优先级比较:升序,最优(最应先用)排在最前。
|
||||
fn compare_quota_reset_priority(a: &Account, b: &Account) -> Ordering {
|
||||
a.quota_limit_reached
|
||||
.cmp(&b.quota_limit_reached)
|
||||
.then_with(|| compare_window_reset(a.window_reset_at, b.window_reset_at))
|
||||
.then_with(|| a.request_count.cmp(&b.request_count))
|
||||
.then_with(|| compare_last_used(a.last_used_at.as_deref(), b.last_used_at.as_deref()))
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
//! RoundRobin 策略:按候选顺序循环选择。
|
||||
//!
|
||||
//! 最简单的策略:忽略账号负载与用量,纯粹按候选集合的顺序用轮转游标依次选中,
|
||||
//! 保证请求在候选账号间均匀轮转。
|
||||
|
||||
use crate::upstream::accounts::model::Account;
|
||||
|
||||
use super::SelectionInput;
|
||||
|
||||
/// 按候选顺序做纯轮转选择。
|
||||
pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option<Account> {
|
||||
let candidates = input.candidates;
|
||||
if candidates.is_empty() {
|
||||
return None;
|
||||
}
|
||||
*cursor %= candidates.len();
|
||||
let index = *cursor;
|
||||
*cursor = cursor.wrapping_add(1);
|
||||
Some(candidates[index].clone())
|
||||
}
|
||||
@@ -0,0 +1,291 @@
|
||||
//! Smart 策略:加权线性打分 + EWMA 反馈。
|
||||
//!
|
||||
//! 参考 sub2api 的调度设计:把多个归一化因子(负载、窗口 token 用量、窗口缓存、窗口
|
||||
//! 请求数、错误率、TTFT、窗口重置)按可配置权重线性求和,得分越高越优先。相较字典序
|
||||
//! 元组比较,加权打分是更平滑的负载均衡,不会让单一维度无条件压制其余维度,同时通过
|
||||
//! EWMA 反馈感知账号健康度与延迟。
|
||||
//!
|
||||
//! 各归一化因子都落在 `[0, 1]` 且方向统一(越大越好),因此权重之间可直接比较。配额
|
||||
//! 已封顶(`quota_limit_reached`)的账号会被施加一个强惩罚,确保只要还有未封顶的候选就
|
||||
//! 不会选中它。并列最高分的账号之间由 [`super::select_by`] 用轮转游标破并列。
|
||||
|
||||
use crate::upstream::accounts::model::Account;
|
||||
use crate::upstream::scheduler::feedback::{FeedbackSample, FeedbackStats};
|
||||
|
||||
use super::{account_window_token_count, SelectionInput};
|
||||
|
||||
/// 配额封顶账号的惩罚分。远大于各因子加权和的理论上限(各因子 ∈ [0,1],权重和通常为
|
||||
/// 个位数),保证未封顶账号总是优先于封顶账号。
|
||||
const QUOTA_LIMITED_PENALTY: f64 = 1_000.0;
|
||||
|
||||
/// 打分权重。每个字段对应一个归一化因子的权重,权重为 0 即关闭该维度。
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub struct ScoreWeights {
|
||||
/// 在途槽位压力(越空闲越优先)。
|
||||
pub load: f64,
|
||||
/// 当前窗口实际 token 用量(用得越少越优先)。
|
||||
pub window: f64,
|
||||
/// 当前窗口缓存 token 用量(缓存越少越优先,权重小,仅在实际用量相当时区分)。
|
||||
pub cached: f64,
|
||||
/// 当前窗口请求数(请求越少越优先)。
|
||||
pub window_requests: f64,
|
||||
/// 错误率 EWMA(越健康越优先)。
|
||||
pub error_rate: f64,
|
||||
/// 首 token 延迟 EWMA(越快越优先)。
|
||||
pub ttft: f64,
|
||||
/// 窗口重置临近度(越快重置越优先,use-it-or-lose-it,默认关闭)。
|
||||
pub reset: f64,
|
||||
}
|
||||
|
||||
impl Default for ScoreWeights {
|
||||
fn default() -> Self {
|
||||
// 默认权重对齐 sub2api 的取向:负载与健康度为主,窗口用量次之,TTFT 再次,
|
||||
// 窗口重置默认关闭(避免过早榨干即将重置的账号)。cached 权重远小于 window,
|
||||
// 确保实际 token 用量主导、缓存仅作细粒度区分。
|
||||
Self {
|
||||
load: 1.0,
|
||||
window: 0.8,
|
||||
cached: 0.05,
|
||||
window_requests: 0.6,
|
||||
error_rate: 1.0,
|
||||
ttft: 0.5,
|
||||
reset: 0.0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 单个候选账号的打分明细,供日志与可观测使用。
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub struct ScoreBreakdown {
|
||||
/// 各因子加权后的总分(已含配额封顶惩罚)。
|
||||
pub total: f64,
|
||||
/// 负载因子(归一化后,未乘权重)。
|
||||
pub load: f64,
|
||||
/// 窗口实际 token 用量因子。
|
||||
pub window: f64,
|
||||
/// 窗口缓存 token 用量因子。
|
||||
pub cached: f64,
|
||||
/// 窗口请求数因子。
|
||||
pub window_requests: f64,
|
||||
/// 错误率因子。
|
||||
pub error_rate: f64,
|
||||
/// TTFT 因子。
|
||||
pub ttft: f64,
|
||||
/// 窗口重置因子。
|
||||
pub reset: f64,
|
||||
/// 是否被施加配额封顶惩罚。
|
||||
pub quota_limited: bool,
|
||||
}
|
||||
|
||||
/// Smart 策略选择:按加权得分降序排序,在并列最高分之间用轮转游标破并列。
|
||||
pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option<Account> {
|
||||
let ranked = rank_candidates(
|
||||
input.candidates,
|
||||
input.weights,
|
||||
input.now.timestamp(),
|
||||
input.slot_count,
|
||||
input.feedback,
|
||||
);
|
||||
let best_total = ranked.first()?.1.total;
|
||||
// 并列最高分的账号数量(浮点相等用位相等判定,得分由相同输入算出,稳定)。
|
||||
let tied_count = ranked
|
||||
.iter()
|
||||
.take_while(|(_, breakdown)| breakdown.total.to_bits() == best_total.to_bits())
|
||||
.count();
|
||||
let index = *cursor % tied_count;
|
||||
*cursor = cursor.wrapping_add(1);
|
||||
let account_index = ranked[index].0;
|
||||
Some(input.candidates[account_index].clone())
|
||||
}
|
||||
|
||||
/// 对候选集合按加权得分排序(降序),返回 `(账号索引, 打分明细)` 列表。
|
||||
///
|
||||
/// 只做打分与排序,最终的并列破并列交给 [`select`]。公开供测试直接断言排序结果。
|
||||
pub fn rank_candidates(
|
||||
candidates: &[Account],
|
||||
weights: &ScoreWeights,
|
||||
now_secs: i64,
|
||||
slot_count: &dyn Fn(&str) -> usize,
|
||||
feedback: &FeedbackStats,
|
||||
) -> Vec<(usize, ScoreBreakdown)> {
|
||||
let samples = candidates
|
||||
.iter()
|
||||
.map(|account| feedback.sample(&account.id))
|
||||
.collect::<Vec<_>>();
|
||||
let ctx = NormalizationContext::build(candidates, &samples, now_secs, slot_count);
|
||||
let mut ranked = candidates
|
||||
.iter()
|
||||
.zip(&samples)
|
||||
.enumerate()
|
||||
.map(|(index, (account, sample))| {
|
||||
(
|
||||
index,
|
||||
score_account(account, sample, weights, &ctx, now_secs, slot_count),
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
// 降序:得分高在前。NaN 不应出现(因子均有界),保守用 total_cmp。
|
||||
ranked.sort_by(|a, b| b.1.total.total_cmp(&a.1.total));
|
||||
ranked
|
||||
}
|
||||
|
||||
/// 候选集合的归一化上下文:预先算好各维度的 min/max,供逐账号归一化。
|
||||
struct NormalizationContext {
|
||||
max_slot_count: f64,
|
||||
max_window_tokens: f64,
|
||||
max_window_cached: f64,
|
||||
max_window_requests: f64,
|
||||
min_ttft_ms: f64,
|
||||
max_ttft_ms: f64,
|
||||
min_reset_secs: f64,
|
||||
max_reset_secs: f64,
|
||||
}
|
||||
|
||||
impl NormalizationContext {
|
||||
fn build(
|
||||
candidates: &[Account],
|
||||
samples: &[FeedbackSample],
|
||||
now_secs: i64,
|
||||
slot_count: &dyn Fn(&str) -> usize,
|
||||
) -> Self {
|
||||
let mut max_slot_count = 0.0_f64;
|
||||
let mut max_window_tokens = 0.0_f64;
|
||||
let mut max_window_cached = 0.0_f64;
|
||||
let mut max_window_requests = 0.0_f64;
|
||||
let mut min_ttft_ms = f64::INFINITY;
|
||||
let mut max_ttft_ms = 0.0_f64;
|
||||
let mut min_reset_secs = f64::INFINITY;
|
||||
let mut max_reset_secs = 0.0_f64;
|
||||
|
||||
for (account, sample) in candidates.iter().zip(samples) {
|
||||
max_slot_count = max_slot_count.max(slot_count(&account.id) as f64);
|
||||
max_window_tokens = max_window_tokens.max(account_window_token_count(account) as f64);
|
||||
max_window_cached = max_window_cached.max(account.window_cached_tokens as f64);
|
||||
max_window_requests = max_window_requests.max(account.window_request_count as f64);
|
||||
|
||||
if let Some(ttft) = sample.ttft_ms {
|
||||
if ttft > 0.0 {
|
||||
min_ttft_ms = min_ttft_ms.min(ttft);
|
||||
max_ttft_ms = max_ttft_ms.max(ttft);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(reset_at) = account.window_reset_at {
|
||||
let secs = (reset_at.timestamp() - now_secs).max(0) as f64;
|
||||
min_reset_secs = min_reset_secs.min(secs);
|
||||
max_reset_secs = max_reset_secs.max(secs);
|
||||
}
|
||||
}
|
||||
|
||||
Self {
|
||||
max_slot_count,
|
||||
max_window_tokens,
|
||||
max_window_cached,
|
||||
max_window_requests,
|
||||
min_ttft_ms: if min_ttft_ms.is_finite() {
|
||||
min_ttft_ms
|
||||
} else {
|
||||
0.0
|
||||
},
|
||||
max_ttft_ms,
|
||||
min_reset_secs: if min_reset_secs.is_finite() {
|
||||
min_reset_secs
|
||||
} else {
|
||||
0.0
|
||||
},
|
||||
max_reset_secs,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 把 `value` 从 `[min, max]` 线性映射到 `[0, 1]`;区间退化时返回中性值 `1.0`
|
||||
/// (表示该维度对所有候选无区分度,不应影响排序)。
|
||||
fn normalize_ascending_good(value: f64, min: f64, max: f64) -> f64 {
|
||||
if max <= min {
|
||||
return 1.0;
|
||||
}
|
||||
((value - min) / (max - min)).clamp(0.0, 1.0)
|
||||
}
|
||||
|
||||
/// 计算指定账号在候选集合中的加权得分明细。
|
||||
fn score_account(
|
||||
account: &Account,
|
||||
sample: &FeedbackSample,
|
||||
weights: &ScoreWeights,
|
||||
ctx: &NormalizationContext,
|
||||
now_secs: i64,
|
||||
slot_count: &dyn Fn(&str) -> usize,
|
||||
) -> ScoreBreakdown {
|
||||
// 负载:槽位越少越好。归一化到 [0,1] 后取反(1 - x),使“越空闲得分越高”。
|
||||
let load =
|
||||
1.0 - normalize_ascending_good(slot_count(&account.id) as f64, 0.0, ctx.max_slot_count);
|
||||
|
||||
// 窗口实际 token 用量:token 越少越好。
|
||||
let window = 1.0
|
||||
- normalize_ascending_good(
|
||||
account_window_token_count(account) as f64,
|
||||
0.0,
|
||||
ctx.max_window_tokens,
|
||||
);
|
||||
|
||||
// 窗口缓存 token 用量:缓存越少越好(权重小,仅在实际用量相当时细分)。
|
||||
let cached = 1.0
|
||||
- normalize_ascending_good(
|
||||
account.window_cached_tokens as f64,
|
||||
0.0,
|
||||
ctx.max_window_cached,
|
||||
);
|
||||
|
||||
// 窗口请求数:请求越少越好。
|
||||
let window_requests = 1.0
|
||||
- normalize_ascending_good(
|
||||
account.window_request_count as f64,
|
||||
0.0,
|
||||
ctx.max_window_requests,
|
||||
);
|
||||
|
||||
// 错误率 EWMA:本身即 [0,1],越低越健康。无样本时视为 0(满分健康)。
|
||||
let error_rate = 1.0 - sample.error_rate.unwrap_or(0.0).clamp(0.0, 1.0);
|
||||
|
||||
// TTFT:延迟越低越好。无样本的账号视为中性满分,避免冷启动被压。
|
||||
let ttft = match sample.ttft_ms {
|
||||
Some(ms) if ms > 0.0 => {
|
||||
1.0 - normalize_ascending_good(ms, ctx.min_ttft_ms, ctx.max_ttft_ms)
|
||||
}
|
||||
_ => 1.0,
|
||||
};
|
||||
|
||||
// 窗口重置:越快重置越优先(use-it-or-lose-it)。无 reset 时间视为中性满分。
|
||||
let reset = match account.window_reset_at {
|
||||
Some(reset_at) => {
|
||||
let secs = (reset_at.timestamp() - now_secs).max(0) as f64;
|
||||
1.0 - normalize_ascending_good(secs, ctx.min_reset_secs, ctx.max_reset_secs)
|
||||
}
|
||||
None => 1.0,
|
||||
};
|
||||
|
||||
let mut total = weights.load * load
|
||||
+ weights.window * window
|
||||
+ weights.cached * cached
|
||||
+ weights.window_requests * window_requests
|
||||
+ weights.error_rate * error_rate
|
||||
+ weights.ttft * ttft
|
||||
+ weights.reset * reset;
|
||||
|
||||
// 配额封顶硬惩罚:确保只要还有未封顶候选就不会选中封顶账号。
|
||||
if account.quota_limit_reached {
|
||||
total -= QUOTA_LIMITED_PENALTY;
|
||||
}
|
||||
|
||||
ScoreBreakdown {
|
||||
total,
|
||||
load,
|
||||
window,
|
||||
cached,
|
||||
window_requests,
|
||||
error_rate,
|
||||
ttft,
|
||||
reset,
|
||||
quota_limited: account.quota_limit_reached,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
//! Sticky 策略:优先选择最近使用过的账号。
|
||||
//!
|
||||
//! 选中 `last_used_at` 最大(最近使用)的账号,倾向于持续复用同一账号(利于上游的
|
||||
//! prompt 缓存命中)。该策略不使用轮转游标——最近使用时间本身即确定性顺序。
|
||||
|
||||
use std::cmp::Ordering;
|
||||
|
||||
use crate::upstream::accounts::model::Account;
|
||||
|
||||
use super::{compare_last_used, SelectionInput};
|
||||
|
||||
/// 选中最近使用过的账号(`last_used_at` 最大者)。
|
||||
///
|
||||
/// `_cursor` 未使用:Sticky 由 `last_used_at` 决定顺序,不参与轮转破并列。
|
||||
pub fn select(input: &SelectionInput<'_>, _cursor: &mut usize) -> Option<Account> {
|
||||
let candidates = input.candidates;
|
||||
let mut selected = candidates.first()?;
|
||||
for candidate in &candidates[1..] {
|
||||
if compare_last_used(
|
||||
candidate.last_used_at.as_deref(),
|
||||
selected.last_used_at.as_deref(),
|
||||
) == Ordering::Greater
|
||||
{
|
||||
selected = candidate;
|
||||
}
|
||||
}
|
||||
Some(selected.clone())
|
||||
}
|
||||
@@ -755,7 +755,7 @@ async fn execute_fresh_pooled_response_create_request(
|
||||
key,
|
||||
PooledWebSocketConnection {
|
||||
websocket,
|
||||
metadata,
|
||||
metadata: reusable_websocket_metadata(metadata),
|
||||
created_at: now,
|
||||
},
|
||||
)
|
||||
@@ -827,7 +827,7 @@ async fn execute_pooled_response_create_request(
|
||||
let created_at = connection.created_at;
|
||||
match collect_websocket_response(
|
||||
connection.websocket,
|
||||
connection.metadata,
|
||||
reusable_websocket_metadata(connection.metadata),
|
||||
true,
|
||||
started_at,
|
||||
first_token_started_at,
|
||||
@@ -840,7 +840,7 @@ async fn execute_pooled_response_create_request(
|
||||
key,
|
||||
PooledWebSocketConnection {
|
||||
websocket,
|
||||
metadata,
|
||||
metadata: reusable_websocket_metadata(metadata),
|
||||
created_at,
|
||||
},
|
||||
)
|
||||
@@ -861,7 +861,7 @@ async fn execute_pooled_response_create_request_stream(
|
||||
connection: PooledWebSocketConnection,
|
||||
) -> Result<CodexWebSocketStreamingExchange, CodexWebSocketExchangeError> {
|
||||
let mut websocket = connection.websocket;
|
||||
let mut metadata = connection.metadata;
|
||||
let mut metadata = reusable_websocket_metadata(connection.metadata);
|
||||
let created_at = connection.created_at;
|
||||
let first_token_started_at = Instant::now();
|
||||
if let Err(error) = websocket
|
||||
@@ -1144,6 +1144,13 @@ fn websocket_connection_metadata(
|
||||
}
|
||||
}
|
||||
|
||||
fn reusable_websocket_metadata(
|
||||
mut metadata: CodexWebSocketConnectionMetadata,
|
||||
) -> CodexWebSocketConnectionMetadata {
|
||||
metadata.rate_limit_headers.clear();
|
||||
metadata
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Streaming
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -1406,7 +1413,7 @@ async fn finish_stream_websocket(
|
||||
pool_return.key,
|
||||
PooledWebSocketConnection {
|
||||
websocket,
|
||||
metadata,
|
||||
metadata: reusable_websocket_metadata(metadata),
|
||||
created_at: pool_return.created_at,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -73,6 +73,51 @@ async fn admin_accounts_list_should_not_expose_account_tokens() {
|
||||
assert!(!response_body.contains("refresh-secret-token"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_accounts_list_should_derive_refreshing_display_status_from_active_lease() {
|
||||
let (app, _state, pool, _dir) =
|
||||
admin_accounts_test_app("admin-accounts-refreshing-display.sqlite", 72).await;
|
||||
sqlx::query("insert into accounts (id, email, access_token, refresh_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?)")
|
||||
.bind("acct_refreshing")
|
||||
.bind("refreshing@example.com")
|
||||
.bind("access-token")
|
||||
.bind("refresh-token")
|
||||
.bind("active")
|
||||
.bind("2026-07-09T00:00:00Z")
|
||||
.bind("2026-07-09T00:00:00Z")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query("insert into account_refresh_leases (account_id, owner, expires_at, updated_at) values (?, ?, ?, ?)")
|
||||
.bind("acct_refreshing")
|
||||
.bind("test-owner")
|
||||
.bind("2099-01-01T00:00:00Z")
|
||||
.bind("2026-07-09T00:00:00Z")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/api/admin/accounts?page=1&pageSize=10")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_accounts_refreshing_display")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = response_json(response).await;
|
||||
let item = &body["data"]["items"][0];
|
||||
assert_eq!(item["status"], "active");
|
||||
assert_eq!(item["displayStatus"], "refreshing");
|
||||
assert_eq!(item["tokenRefreshing"], true);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_accounts_list_should_include_usage_quota_and_model_stats() {
|
||||
let (app, _state, pool, _dir) =
|
||||
@@ -271,13 +316,22 @@ async fn admin_accounts_list_should_include_usage_quota_and_model_stats() {
|
||||
assert_eq!(item["quota"]["windows"].as_array().unwrap().len(), 3);
|
||||
assert_eq!(item["quota"]["windows"][0]["labelDisplay"], "月限额");
|
||||
assert_eq!(item["quota"]["windows"][0]["group"], "monthly");
|
||||
assert_eq!(item["quota"]["windows"][0]["tokenUsageDisplay"], "1.12M");
|
||||
assert_eq!(
|
||||
item["quota"]["windows"][0]["localUsage"]["totalTokensDisplay"],
|
||||
"1.12M"
|
||||
);
|
||||
assert_eq!(item["quota"]["windows"][1]["labelDisplay"], "5小时限额");
|
||||
assert_eq!(item["quota"]["windows"][1]["group"], "shortTerm");
|
||||
assert_eq!(item["quota"]["windows"][1]["tokenUsageDisplay"], "2.5K");
|
||||
assert_eq!(
|
||||
item["quota"]["windows"][1]["localUsage"]["totalTokensDisplay"],
|
||||
"2.5K"
|
||||
);
|
||||
assert_eq!(item["quota"]["windows"][2]["labelDisplay"], "周限额");
|
||||
assert_eq!(item["quota"]["windows"][2]["group"], "shortTerm");
|
||||
assert_eq!(item["quota"]["windows"][2]["tokenUsageDisplay"], "124K");
|
||||
assert_eq!(
|
||||
item["quota"]["windows"][2]["localUsage"]["totalTokensDisplay"],
|
||||
"124K"
|
||||
);
|
||||
assert!(item["quota"]["windows"][2]["windowUsedDisplay"]
|
||||
.as_str()
|
||||
.is_some_and(|value| value.contains(" / 7.0d")));
|
||||
|
||||
@@ -175,7 +175,7 @@ async fn admin_account_quota_should_send_usage_cookie() {
|
||||
"shortTerm"
|
||||
);
|
||||
assert_eq!(
|
||||
body["data"]["quotaData"]["windows"][0]["tokenUsageDisplay"],
|
||||
body["data"]["quotaData"]["windows"][0]["localUsage"]["totalTokensDisplay"],
|
||||
"2K"
|
||||
);
|
||||
assert_eq!(
|
||||
@@ -183,7 +183,7 @@ async fn admin_account_quota_should_send_usage_cookie() {
|
||||
"周限额"
|
||||
);
|
||||
assert_eq!(
|
||||
body["data"]["quotaData"]["windows"][1]["tokenUsageDisplay"],
|
||||
body["data"]["quotaData"]["windows"][1]["localUsage"]["totalTokensDisplay"],
|
||||
"2K"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -70,7 +70,7 @@ async fn dashboard_summary_should_render_dash_when_fingerprint_updated_at_is_mis
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_calculate_today_traffic_and_average_first_token_latency_from_time_buckets(
|
||||
) {
|
||||
let (app, store, pool, _dir) = dashboard_test_app(
|
||||
let (app, store, _pool, _dir) = dashboard_test_app(
|
||||
"dashboard-traffic-latency.sqlite",
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
)
|
||||
@@ -96,23 +96,22 @@ async fn dashboard_summary_should_calculate_today_traffic_and_average_first_toke
|
||||
.append(&usage_record_with_tokens(yesterday_record_at, 30))
|
||||
.await
|
||||
.unwrap();
|
||||
seed_usage_summary_with_cached_tokens(&pool, 120, 30, 24).await;
|
||||
|
||||
let body = dashboard_summary(app).await;
|
||||
|
||||
assert_eq!(body["data"]["cards"]["traffic"]["todayRequests"], "2");
|
||||
assert_eq!(body["data"]["cards"]["traffic"]["todayRequestsValue"], 2);
|
||||
assert_eq!(body["data"]["cards"]["traffic"]["totalRequests"], "42");
|
||||
assert_eq!(body["data"]["cards"]["traffic"]["totalRequests"], "3");
|
||||
assert_eq!(
|
||||
body["data"]["cards"]["traffic"]["yesterdayRequestsValue"],
|
||||
1
|
||||
);
|
||||
assert_eq!(body["data"]["cards"]["tokens"]["todayTokens"], "30");
|
||||
assert_eq!(body["data"]["cards"]["tokens"]["todayTokensValue"], 30);
|
||||
assert_eq!(body["data"]["cards"]["tokens"]["totalTokens"], "150");
|
||||
assert_eq!(body["data"]["cards"]["tokens"]["totalTokens"], "60");
|
||||
assert_eq!(body["data"]["cards"]["tokens"]["yesterdayTokensValue"], 30);
|
||||
assert_eq!(body["data"]["cards"]["cache"]["totalHitRate"], "20.0%");
|
||||
assert_eq!(body["data"]["cards"]["cache"]["totalCachedTokens"], "24");
|
||||
assert_eq!(body["data"]["cards"]["cache"]["totalHitRate"], "0.0%");
|
||||
assert_eq!(body["data"]["cards"]["cache"]["totalCachedTokens"], "0");
|
||||
assert_eq!(
|
||||
body["data"]["cards"]["cache"]["averageFirstTokenLatencyMs"],
|
||||
"300 ms"
|
||||
@@ -719,39 +718,6 @@ async fn seed_usage_summary(pool: &SqlitePool, input_tokens: u64, output_tokens:
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn seed_usage_summary_with_cached_tokens(
|
||||
pool: &SqlitePool,
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cached_tokens: u64,
|
||||
) {
|
||||
sqlx::query(
|
||||
"insert into accounts (id, email, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind("acct_usage_cached")
|
||||
.bind("acct-usage-cached@example.com")
|
||||
.bind("cipher")
|
||||
.bind("active")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"insert into account_usage (account_id, request_count, input_tokens, output_tokens, cached_tokens, total_tokens, last_used_at) values (?, ?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind("acct_usage_cached")
|
||||
.bind(42_i64)
|
||||
.bind(input_tokens as i64)
|
||||
.bind(output_tokens as i64)
|
||||
.bind(cached_tokens as i64)
|
||||
.bind((input_tokens + output_tokens) as i64)
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn seed_usage_summary_with_last_used(
|
||||
pool: &SqlitePool,
|
||||
input_tokens: u64,
|
||||
@@ -857,6 +823,7 @@ fn usage_record_with_account_tokens(
|
||||
fn usage_record(created_at: chrono::DateTime<Utc>, model: &str, metadata: Value) -> UsageRecord {
|
||||
let mut record = UsageRecord::new("v1.response", UsageRecordLevel::Info, "completed");
|
||||
record.model = Some(model.to_string());
|
||||
record.status_code = Some(200);
|
||||
record.created_at = created_at;
|
||||
record.metadata = metadata;
|
||||
record
|
||||
|
||||
@@ -54,10 +54,16 @@ async fn admin_usage_records_should_cursor_page_events_and_include_request_id()
|
||||
let now = Utc::now();
|
||||
let mut older = UsageRecord::new("request", UsageRecordLevel::Info, "older");
|
||||
older.id = "usage_older".to_string();
|
||||
older.account_id = Some("acct_usage".to_string());
|
||||
older.model = Some("gpt-5.5".to_string());
|
||||
older.status_code = Some(200);
|
||||
older.created_at = now;
|
||||
store.append(&older).await.unwrap();
|
||||
let mut newer = UsageRecord::new("request", UsageRecordLevel::Info, "newer");
|
||||
newer.id = "usage_newer".to_string();
|
||||
newer.account_id = Some("acct_usage".to_string());
|
||||
newer.model = Some("gpt-5.5".to_string());
|
||||
newer.status_code = Some(200);
|
||||
newer.created_at = now + Duration::seconds(1);
|
||||
store.append(&newer).await.unwrap();
|
||||
|
||||
@@ -88,8 +94,11 @@ async fn admin_usage_records_should_return_numbered_page_metadata() {
|
||||
("usage_old", "older timeout", 0),
|
||||
("usage_new", "newer timeout", 1),
|
||||
] {
|
||||
let mut event = UsageRecord::new("request", UsageRecordLevel::Error, message);
|
||||
let mut event = UsageRecord::new("request", UsageRecordLevel::Info, message);
|
||||
event.id = id.to_string();
|
||||
event.account_id = Some("acct_usage".to_string());
|
||||
event.model = Some("gpt-5.5".to_string());
|
||||
event.status_code = Some(200);
|
||||
event.created_at = now + Duration::seconds(offset);
|
||||
store.append(&event).await.unwrap();
|
||||
}
|
||||
@@ -98,7 +107,7 @@ async fn admin_usage_records_should_return_numbered_page_metadata() {
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/api/admin/usage/records?page=1&pageSize=1&level=error&search=timeout")
|
||||
.uri("/api/admin/usage/records?page=1&pageSize=1&search=timeout")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_usage_records_numbered")
|
||||
.body(Body::empty())
|
||||
@@ -203,12 +212,13 @@ async fn admin_usage_records_should_return_table_display_fields() {
|
||||
async fn admin_usage_records_should_use_record_account_email_snapshot_when_account_is_missing() {
|
||||
let (app, store, _dir) =
|
||||
admin_usage_records_test_app("admin-usage-records-email-snapshot.sqlite").await;
|
||||
let mut event = UsageRecord::new("v1.response", UsageRecordLevel::Error, "upstream failed");
|
||||
let mut event = UsageRecord::new("v1.response", UsageRecordLevel::Info, "completed");
|
||||
event.id = "usage_email_snapshot".to_string();
|
||||
event.account_id = Some("acct_deleted".to_string());
|
||||
event.model = Some("gpt-5.5".to_string());
|
||||
event.status_code = Some(200);
|
||||
event.metadata = json!({
|
||||
"accountEmail": "deleted@example.com",
|
||||
"failed": true
|
||||
"accountEmail": "deleted@example.com"
|
||||
});
|
||||
store.append(&event).await.unwrap();
|
||||
|
||||
@@ -237,7 +247,9 @@ async fn admin_usage_records_summary_should_aggregate_filtered_usage() {
|
||||
admin_usage_records_test_app("admin-usage-records-summary.sqlite").await;
|
||||
let mut success = UsageRecord::new("request", UsageRecordLevel::Info, "summary success");
|
||||
success.id = "usage_summary_success".to_string();
|
||||
success.account_id = Some("acct_summary".to_string());
|
||||
success.model = Some("gpt-5.5".to_string());
|
||||
success.status_code = Some(200);
|
||||
success.latency_ms = Some(1200);
|
||||
success.metadata = json!({
|
||||
"usage": {
|
||||
@@ -250,6 +262,7 @@ async fn admin_usage_records_summary_should_aggregate_filtered_usage() {
|
||||
|
||||
let mut failure = UsageRecord::new("request", UsageRecordLevel::Error, "summary failure");
|
||||
failure.id = "usage_summary_failure".to_string();
|
||||
failure.account_id = Some("acct_summary".to_string());
|
||||
failure.model = Some("gpt-5.5-mini".to_string());
|
||||
failure.status_code = Some(429);
|
||||
failure.latency_ms = Some(600);
|
||||
@@ -291,8 +304,10 @@ async fn admin_usage_record_insight_cards_should_aggregate_filtered_usage_dimens
|
||||
admin_usage_records_test_app_with_pool("admin-usage-records-insights.sqlite").await;
|
||||
let mut first = UsageRecord::new("request", UsageRecordLevel::Info, "insight success");
|
||||
first.id = "usage_insight_success".to_string();
|
||||
first.account_id = Some("acct_insight".to_string());
|
||||
first.route = Some("/v1/responses".to_string());
|
||||
first.model = Some("gpt-5.5".to_string());
|
||||
first.status_code = Some(200);
|
||||
first.latency_ms = Some(240);
|
||||
first.metadata = json!({
|
||||
"stream": true,
|
||||
@@ -310,6 +325,7 @@ async fn admin_usage_record_insight_cards_should_aggregate_filtered_usage_dimens
|
||||
|
||||
let mut second = UsageRecord::new("request", UsageRecordLevel::Error, "insight failure");
|
||||
second.id = "usage_insight_failure".to_string();
|
||||
second.account_id = Some("acct_insight".to_string());
|
||||
second.route = Some("/v1/responses/review".to_string());
|
||||
second.model = Some("gpt-5.5-mini".to_string());
|
||||
second.status_code = Some(500);
|
||||
@@ -451,6 +467,9 @@ async fn admin_usage_record_trends_should_bucket_by_china_calendar_day() {
|
||||
admin_usage_records_test_app_with_pool("admin-usage-records-china-day-trend.sqlite").await;
|
||||
let mut june_29 = UsageRecord::new("request", UsageRecordLevel::Info, "china-day trend");
|
||||
june_29.id = "usage_china_day_29".to_string();
|
||||
june_29.account_id = Some("acct_trend".to_string());
|
||||
june_29.model = Some("gpt-5.5".to_string());
|
||||
june_29.status_code = Some(200);
|
||||
june_29.created_at = chrono::DateTime::parse_from_rfc3339("2026-06-29T15:30:00Z")
|
||||
.unwrap()
|
||||
.with_timezone(&Utc);
|
||||
@@ -465,6 +484,9 @@ async fn admin_usage_record_trends_should_bucket_by_china_calendar_day() {
|
||||
|
||||
let mut june_30 = UsageRecord::new("request", UsageRecordLevel::Info, "china-day trend");
|
||||
june_30.id = "usage_china_day_30".to_string();
|
||||
june_30.account_id = Some("acct_trend".to_string());
|
||||
june_30.model = Some("gpt-5.5".to_string());
|
||||
june_30.status_code = Some(200);
|
||||
june_30.created_at = chrono::DateTime::parse_from_rfc3339("2026-06-29T16:30:00Z")
|
||||
.unwrap()
|
||||
.with_timezone(&Utc);
|
||||
@@ -517,8 +539,10 @@ async fn admin_usage_records_insight_card_endpoints_should_return_source_specifi
|
||||
admin_usage_records_test_app_with_pool("admin-usage-records-insight-cards.sqlite").await;
|
||||
let mut first = UsageRecord::new("request", UsageRecordLevel::Info, "card success");
|
||||
first.id = "usage_card_success".to_string();
|
||||
first.account_id = Some("acct_card".to_string());
|
||||
first.route = Some("/v1/responses".to_string());
|
||||
first.model = Some("gpt-5.5".to_string());
|
||||
first.status_code = Some(200);
|
||||
first.latency_ms = Some(240);
|
||||
first.metadata = json!({
|
||||
"stream": true,
|
||||
@@ -535,8 +559,10 @@ async fn admin_usage_records_insight_card_endpoints_should_return_source_specifi
|
||||
|
||||
let mut second = UsageRecord::new("request", UsageRecordLevel::Info, "card mini");
|
||||
second.id = "usage_card_mini".to_string();
|
||||
second.account_id = Some("acct_card".to_string());
|
||||
second.route = Some("/v1/responses/compact".to_string());
|
||||
second.model = Some("gpt-5.5-mini".to_string());
|
||||
second.status_code = Some(200);
|
||||
second.latency_ms = Some(120);
|
||||
second.metadata = json!({
|
||||
"stream": false,
|
||||
@@ -662,11 +688,17 @@ async fn admin_usage_records_should_filter_summary_and_table_by_time_range() {
|
||||
let now = Utc::now();
|
||||
let mut older = UsageRecord::new("request", UsageRecordLevel::Info, "outside range");
|
||||
older.id = "usage_outside_range".to_string();
|
||||
older.account_id = Some("acct_range".to_string());
|
||||
older.model = Some("gpt-5.5".to_string());
|
||||
older.status_code = Some(200);
|
||||
older.created_at = now - Duration::days(3);
|
||||
store.append(&older).await.unwrap();
|
||||
|
||||
let mut current = UsageRecord::new("request", UsageRecordLevel::Info, "inside range");
|
||||
current.id = "usage_inside_range".to_string();
|
||||
current.account_id = Some("acct_range".to_string());
|
||||
current.model = Some("gpt-5.5".to_string());
|
||||
current.status_code = Some(200);
|
||||
current.created_at = now;
|
||||
store.append(¤t).await.unwrap();
|
||||
|
||||
@@ -710,18 +742,18 @@ async fn admin_usage_records_should_filter_summary_and_table_by_time_range() {
|
||||
#[tokio::test]
|
||||
async fn admin_usage_records_should_filter_and_cursor_page_events() {
|
||||
let (app, store, _dir) = admin_usage_records_test_app("admin-usage-records.sqlite").await;
|
||||
let mut matching = UsageRecord::new("request", UsageRecordLevel::Error, "upstream timeout");
|
||||
let mut matching = UsageRecord::new("request", UsageRecordLevel::Info, "upstream timeout");
|
||||
matching.id = "usage_matching".to_string();
|
||||
matching.account_id = Some("acct_filter".to_string());
|
||||
matching.route = Some("/v1/responses".to_string());
|
||||
matching.model = Some("gpt-5.5".to_string());
|
||||
matching.status_code = Some(200);
|
||||
store.append(&matching).await.unwrap();
|
||||
store
|
||||
.append(&UsageRecord::new(
|
||||
"request",
|
||||
UsageRecordLevel::Info,
|
||||
"upstream timeout",
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let mut non_matching_message = UsageRecord::new("request", UsageRecordLevel::Info, "completed");
|
||||
non_matching_message.account_id = Some("acct_filter".to_string());
|
||||
non_matching_message.model = Some("gpt-5.5".to_string());
|
||||
non_matching_message.status_code = Some(200);
|
||||
store.append(&non_matching_message).await.unwrap();
|
||||
store
|
||||
.append(&UsageRecord::new(
|
||||
"account",
|
||||
@@ -735,7 +767,7 @@ async fn admin_usage_records_should_filter_and_cursor_page_events() {
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/api/admin/usage/records?kind=request&level=error&search=timeout&limit=1")
|
||||
.uri("/api/admin/usage/records?kind=request&level=info&search=timeout&limit=1")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_usage_records_filter")
|
||||
.body(Body::empty())
|
||||
|
||||
@@ -860,6 +860,70 @@ async fn update_events_should_close_on_shutdown() {
|
||||
assert!(body.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn update_events_should_close_after_terminal_update_log() {
|
||||
let _guard = SYSTEM_ENV_LOCK.lock().await;
|
||||
let repository = "zyycn/codex-proxy-rs-events-terminal-route";
|
||||
let github = MockServer::start().await;
|
||||
let deploy = tempfile::tempdir().unwrap();
|
||||
let exe_path = deploy.path().join("codex-proxy-rs");
|
||||
let web_dist = deploy.path().join("web/dist");
|
||||
std::fs::create_dir_all(&web_dist).unwrap();
|
||||
std::fs::write(&exe_path, "old-binary").unwrap();
|
||||
std::fs::write(web_dist.join("index.html"), "old-web").unwrap();
|
||||
mount_latest_release_with_archive(&github, repository, "0.4.0", "new-binary", "new-web").await;
|
||||
set_system_update_env(repository, &github.uri());
|
||||
set_system_update_paths(deploy.path(), &exe_path, &web_dist);
|
||||
let (app, _dir) = admin_system_test_app("system-events-terminal.sqlite").await;
|
||||
|
||||
let sse_response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/api/admin/system/update-events")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_system_update_events_terminal")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let update_app = app.clone();
|
||||
let update_task = tokio::spawn(async move {
|
||||
update_app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/api/admin/system/update")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_system_update_events_terminal_update")
|
||||
.header(CONTENT_TYPE, "application/json")
|
||||
.body(confirmed_update_body("0.4.0"))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
let body = tokio::time::timeout(
|
||||
Duration::from_secs(2),
|
||||
to_bytes(sse_response.into_body(), 128 * 1024),
|
||||
)
|
||||
.await
|
||||
.expect("update event stream should end after terminal update log")
|
||||
.unwrap();
|
||||
let update_response = update_task.await.unwrap();
|
||||
let text = String::from_utf8(body.to_vec()).unwrap();
|
||||
|
||||
assert!(
|
||||
update_response.status() == StatusCode::OK
|
||||
&& text.contains("\"step\":\"done\"")
|
||||
&& text.contains("\"terminal\":true")
|
||||
&& text.contains("更新文件已替换,等待服务重启生效")
|
||||
);
|
||||
}
|
||||
|
||||
async fn get_update_detail(app: &axum::Router, refresh: bool, request_id: &str) -> Value {
|
||||
let response = app
|
||||
.clone()
|
||||
|
||||
@@ -8,7 +8,6 @@ const DEFAULT_CONFIG_YAML: &str = r#"
|
||||
server:
|
||||
host: 0.0.0.0
|
||||
port: 8080
|
||||
trusted_proxies: []
|
||||
api:
|
||||
base_url: https://chatgpt.com/backend-api
|
||||
database:
|
||||
@@ -80,7 +79,6 @@ logging:
|
||||
fn default_config_keeps_only_codex_backend() {
|
||||
let cfg: AppConfig = parse_yaml_config(DEFAULT_CONFIG_YAML).unwrap();
|
||||
assert_eq!(cfg.server.host, "0.0.0.0");
|
||||
assert!(cfg.server.trusted_proxies.is_empty());
|
||||
assert_eq!(cfg.api.base_url, "https://chatgpt.com/backend-api");
|
||||
assert!(cfg.model_aliases.is_empty());
|
||||
assert_eq!(cfg.auth.refresh_margin_seconds, 300);
|
||||
@@ -259,7 +257,6 @@ ws_pool:
|
||||
|
||||
assert_eq!(cfg.server.host, "127.0.0.1");
|
||||
assert_eq!(cfg.server.port, 8080);
|
||||
assert!(cfg.server.trusted_proxies.is_empty());
|
||||
assert_eq!(
|
||||
cfg.database.url,
|
||||
"sqlite://.runtime/data/codex-proxy-rs.sqlite"
|
||||
@@ -280,30 +277,15 @@ ws_pool:
|
||||
assert_eq!(cfg.logging.retention_days, 14);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_should_parse_trusted_proxy_entries() {
|
||||
let cfg: AppConfig = parse_yaml_config(&DEFAULT_CONFIG_YAML.replace(
|
||||
"trusted_proxies: []",
|
||||
"trusted_proxies:\n - 127.0.0.1\n - 10.0.0.0/8",
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
cfg.server.trusted_proxies,
|
||||
vec!["127.0.0.1".to_string(), "10.0.0.0/8".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn server_config_should_reject_unknown_fields() {
|
||||
let err = parse_yaml_config::<AppConfig>(&DEFAULT_CONFIG_YAML.replace(
|
||||
"trusted_proxies: []",
|
||||
"trusted_proxies: []\n trusted_proxy: 127.0.0.1",
|
||||
))
|
||||
let err = parse_yaml_config::<AppConfig>(
|
||||
&DEFAULT_CONFIG_YAML.replace(" port: 8080", " port: 8080\n trusted_proxies: []"),
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(
|
||||
err.to_string().contains("trusted_proxy"),
|
||||
err.to_string().contains("trusted_proxies"),
|
||||
"expected unknown server field to be rejected, got {err}"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -9,14 +9,12 @@ use axum::{
|
||||
extract::connect_info::ConnectInfo,
|
||||
http::HeaderMap,
|
||||
http::{Request, StatusCode},
|
||||
middleware::{from_fn, from_fn_with_state},
|
||||
middleware::from_fn,
|
||||
routing::get,
|
||||
Extension, Router,
|
||||
};
|
||||
use codex_proxy_rs::http::middleware::{
|
||||
request_id::{
|
||||
attach_request_id, attach_request_id_with_proxy_config, ClientIp, TrustedProxyConfig,
|
||||
},
|
||||
request_id::{attach_request_id, ClientIp},
|
||||
trace::http_trace_layer,
|
||||
};
|
||||
use tokio::time::{sleep, Duration};
|
||||
@@ -285,7 +283,7 @@ async fn request_context_should_auto_detect_forwarded_ip_without_trusted_proxy_c
|
||||
let response = app.oneshot(request).await.expect("response");
|
||||
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
|
||||
|
||||
// 未配置可信代理时进入自动模式:CF-Connecting-IP 优先于 X-Real-IP / X-Forwarded-For。
|
||||
// CF-Connecting-IP 优先于 X-Real-IP / X-Forwarded-For。
|
||||
assert_eq!(
|
||||
std::str::from_utf8(&bytes).unwrap(),
|
||||
"203.0.113.43|cf=203.0.113.43|xff=203.0.113.42, 10.0.0.2|real=203.0.113.44"
|
||||
@@ -319,45 +317,7 @@ async fn request_context_should_skip_private_forwarded_hops_in_auto_mode() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn request_context_should_accept_forwarded_ip_from_trusted_proxy() {
|
||||
let trusted_proxies =
|
||||
TrustedProxyConfig::from_entries(&["127.0.0.0/8".to_string()]).expect("trusted proxy");
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/ip",
|
||||
get(
|
||||
|Extension(client_ip): Extension<ClientIp>, headers: HeaderMap| async move {
|
||||
let forwarded_present = headers.contains_key("x-forwarded-for");
|
||||
format!("{}|forwarded={forwarded_present}", client_ip.as_str())
|
||||
},
|
||||
),
|
||||
)
|
||||
.layer(from_fn_with_state(
|
||||
trusted_proxies,
|
||||
attach_request_id_with_proxy_config,
|
||||
));
|
||||
let mut request = Request::builder()
|
||||
.uri("/ip")
|
||||
.header("x-forwarded-for", "203.0.113.42, 10.0.0.2")
|
||||
.body(Body::empty())
|
||||
.expect("request should build");
|
||||
request
|
||||
.extensions_mut()
|
||||
.insert(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 61234))));
|
||||
|
||||
let response = app.oneshot(request).await.expect("response");
|
||||
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
std::str::from_utf8(&bytes).unwrap(),
|
||||
"203.0.113.42|forwarded=true"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn request_context_should_accept_real_ip_from_trusted_proxy_when_forwarded_for_missing() {
|
||||
let trusted_proxies =
|
||||
TrustedProxyConfig::from_entries(&["127.0.0.1".to_string()]).expect("trusted proxy");
|
||||
async fn request_context_should_accept_real_ip_when_forwarded_for_missing() {
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/ip",
|
||||
@@ -365,10 +325,7 @@ async fn request_context_should_accept_real_ip_from_trusted_proxy_when_forwarded
|
||||
client_ip.as_str().to_string()
|
||||
}),
|
||||
)
|
||||
.layer(from_fn_with_state(
|
||||
trusted_proxies,
|
||||
attach_request_id_with_proxy_config,
|
||||
));
|
||||
.layer(from_fn(attach_request_id));
|
||||
let mut request = Request::builder()
|
||||
.uri("/ip")
|
||||
.header("x-real-ip", "198.51.100.23")
|
||||
@@ -385,9 +342,7 @@ async fn request_context_should_accept_real_ip_from_trusted_proxy_when_forwarded
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn request_context_should_prefer_real_ip_over_forwarded_for_from_trusted_proxy() {
|
||||
let trusted_proxies =
|
||||
TrustedProxyConfig::from_entries(&["127.0.0.1".to_string()]).expect("trusted proxy");
|
||||
async fn request_context_should_prefer_real_ip_over_forwarded_for() {
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/ip",
|
||||
@@ -395,10 +350,7 @@ async fn request_context_should_prefer_real_ip_over_forwarded_for_from_trusted_p
|
||||
client_ip.as_str().to_string()
|
||||
}),
|
||||
)
|
||||
.layer(from_fn_with_state(
|
||||
trusted_proxies,
|
||||
attach_request_id_with_proxy_config,
|
||||
));
|
||||
.layer(from_fn(attach_request_id));
|
||||
let mut request = Request::builder()
|
||||
.uri("/ip")
|
||||
.header("x-forwarded-for", "198.51.100.99, 203.0.113.42")
|
||||
@@ -415,36 +367,6 @@ async fn request_context_should_prefer_real_ip_over_forwarded_for_from_trusted_p
|
||||
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "203.0.113.42");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn request_context_should_ignore_forwarded_ip_from_untrusted_peer() {
|
||||
let trusted_proxies =
|
||||
TrustedProxyConfig::from_entries(&["10.0.0.0/8".to_string()]).expect("trusted proxy");
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/ip",
|
||||
get(|Extension(client_ip): Extension<ClientIp>| async move {
|
||||
client_ip.as_str().to_string()
|
||||
}),
|
||||
)
|
||||
.layer(from_fn_with_state(
|
||||
trusted_proxies,
|
||||
attach_request_id_with_proxy_config,
|
||||
));
|
||||
let mut request = Request::builder()
|
||||
.uri("/ip")
|
||||
.header("x-forwarded-for", "203.0.113.42")
|
||||
.body(Body::empty())
|
||||
.expect("request should build");
|
||||
request
|
||||
.extensions_mut()
|
||||
.insert(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 61234))));
|
||||
|
||||
let response = app.oneshot(request).await.expect("response");
|
||||
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
|
||||
|
||||
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "127.0.0.1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn request_context_should_not_inject_real_ip_header() {
|
||||
let app = Router::new()
|
||||
@@ -474,13 +396,3 @@ async fn request_context_should_not_inject_real_ip_header() {
|
||||
|
||||
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "missing");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trusted_proxy_config_should_reject_invalid_entries() {
|
||||
let error = TrustedProxyConfig::from_entries(&["not-an-ip".to_string()]).unwrap_err();
|
||||
|
||||
assert!(
|
||||
error.to_string().contains("not-an-ip"),
|
||||
"unexpected error: {error}"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -10,7 +10,7 @@ async fn sqlite_schema_should_create_current_tables() {
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
|
||||
let rows: Vec<(String,)> = sqlx::query_as(
|
||||
"select name from sqlite_master where type = 'table' and name in ('schema_migrations', 'admin_users', 'admin_sessions', 'client_api_keys', 'runtime_settings', 'accounts', 'account_refresh_leases', 'account_usage', 'account_model_usage', 'usage_time_buckets', 'account_cookies', 'fingerprints', 'fingerprint_update_history', 'usage_records', 'model_plan_snapshots', 'session_affinities') order by name",
|
||||
"select name from sqlite_master where type = 'table' and name in ('schema_migrations', 'admin_users', 'admin_sessions', 'client_api_keys', 'runtime_settings', 'accounts', 'account_refresh_leases', 'account_usage', 'account_model_usage', 'usage_time_buckets', 'account_cookies', 'fingerprints', 'fingerprint_update_history', 'usage_records', 'ops_error_logs', 'model_plan_snapshots', 'session_affinities') order by name",
|
||||
)
|
||||
.fetch_all(&pool)
|
||||
.await
|
||||
@@ -29,6 +29,7 @@ async fn sqlite_schema_should_create_current_tables() {
|
||||
"fingerprint_update_history",
|
||||
"fingerprints",
|
||||
"model_plan_snapshots",
|
||||
"ops_error_logs",
|
||||
"runtime_settings",
|
||||
"schema_migrations",
|
||||
"session_affinities",
|
||||
@@ -56,10 +57,29 @@ async fn sqlite_schema_should_record_applied_migrations() {
|
||||
[
|
||||
(1, "initial".to_string()),
|
||||
(2, "rename_rotation_strategies".to_string()),
|
||||
(3, "ops_error_logs".to_string()),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sqlite_schema_should_reject_refreshing_account_status() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("account-status.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
|
||||
let error = sqlx::query(
|
||||
"insert into accounts (id, access_token, status, added_at, updated_at)
|
||||
values ('acct_refreshing', 'access', 'refreshing', '2026-07-09T00:00:00Z', '2026-07-09T00:00:00Z')",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(error.to_string().contains("CHECK constraint failed"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn initial_migration_should_not_use_idempotent_business_schema_ddl() {
|
||||
let migration_path = Path::new(env!("CARGO_MANIFEST_DIR"))
|
||||
@@ -161,7 +181,7 @@ async fn sqlite_schema_should_persist_diagnostic_filter_indexes() {
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
|
||||
let rows: Vec<(String,)> = sqlx::query_as(
|
||||
"select name from sqlite_master where type = 'index' and name in ('idx_account_cookies_expires', 'idx_usage_records_level_created', 'idx_usage_records_model_created', 'idx_usage_records_route_created', 'idx_usage_records_status_created', 'idx_usage_records_upstream_status_created', 'idx_session_affinities_active_order', 'idx_usage_time_buckets_bucket', 'idx_usage_time_buckets_model_bucket') order by name",
|
||||
"select name from sqlite_master where type = 'index' and name in ('idx_account_cookies_expires', 'idx_usage_records_level_created', 'idx_usage_records_model_created', 'idx_usage_records_route_created', 'idx_usage_records_status_created', 'idx_usage_records_upstream_status_created', 'idx_ops_error_logs_created_id', 'idx_ops_error_logs_status_created', 'idx_ops_error_logs_failure_class', 'idx_session_affinities_active_order', 'idx_usage_time_buckets_bucket', 'idx_usage_time_buckets_model_bucket') order by name",
|
||||
)
|
||||
.fetch_all(&pool)
|
||||
.await
|
||||
@@ -171,6 +191,9 @@ async fn sqlite_schema_should_persist_diagnostic_filter_indexes() {
|
||||
rows.into_iter().map(|row| row.0).collect::<Vec<_>>(),
|
||||
[
|
||||
"idx_account_cookies_expires",
|
||||
"idx_ops_error_logs_created_id",
|
||||
"idx_ops_error_logs_failure_class",
|
||||
"idx_ops_error_logs_status_created",
|
||||
"idx_session_affinities_active_order",
|
||||
"idx_usage_records_level_created",
|
||||
"idx_usage_records_model_created",
|
||||
@@ -475,6 +498,41 @@ async fn sqlite_schema_should_persist_structured_event_diagnostic_columns() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sqlite_schema_should_persist_ops_error_log_columns() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("ops-error-columns.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
|
||||
let rows: Vec<(String,)> = sqlx::query_as(
|
||||
"select name from pragma_table_info('ops_error_logs') where name in ('request_id', 'account_id', 'route', 'model', 'status_code', 'client_status_code', 'upstream_status_code', 'transport', 'attempt_index', 'failure_class', 'response_id', 'upstream_request_id', 'latency_ms', 'metadata_json') order by name",
|
||||
)
|
||||
.fetch_all(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
rows.into_iter().map(|row| row.0).collect::<Vec<_>>(),
|
||||
[
|
||||
"account_id",
|
||||
"attempt_index",
|
||||
"client_status_code",
|
||||
"failure_class",
|
||||
"latency_ms",
|
||||
"metadata_json",
|
||||
"model",
|
||||
"request_id",
|
||||
"response_id",
|
||||
"route",
|
||||
"status_code",
|
||||
"transport",
|
||||
"upstream_request_id",
|
||||
"upstream_status_code",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sqlite_schema_should_reject_invalid_account_status() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
@@ -507,6 +565,22 @@ async fn sqlite_schema_should_reject_invalid_usage_record_level() {
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sqlite_schema_should_reject_invalid_ops_error_status_code() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("ops-error-status-check.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
|
||||
let result = sqlx::query(
|
||||
"insert into ops_error_logs (id, kind, status_code, message, metadata_json, created_at) values ('ops_bad', 'request', 99, 'bad', '{}', '2026-06-14T00:00:00Z')",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sqlite_schema_should_reject_non_boolean_flags() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
@@ -274,7 +274,7 @@ async fn responses_compact_should_return_rate_limit_error_when_fallback_is_exhau
|
||||
assert_eq!(body["error"]["code"], "rate_limit_exceeded");
|
||||
assert!(message.contains("All accounts exhausted (1 rate-limited)"));
|
||||
assert!(message.contains("compact quota reached"));
|
||||
let event = latest_response_usage_record(&pool).await;
|
||||
let event = latest_response_ops_error_log(&pool).await;
|
||||
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
|
||||
assert_eq!(
|
||||
event.request_id.as_deref(),
|
||||
@@ -339,7 +339,7 @@ async fn responses_compact_should_preserve_upstream_client_error_status() {
|
||||
assert_eq!(status, StatusCode::BAD_REQUEST);
|
||||
assert_eq!(body["error"]["type"], "invalid_request_error");
|
||||
assert_eq!(body["error"]["code"], "codex_api_error");
|
||||
let event = latest_response_usage_record(&pool).await;
|
||||
let event = latest_response_ops_error_log(&pool).await;
|
||||
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
|
||||
assert_eq!(
|
||||
event.request_id.as_deref(),
|
||||
@@ -361,7 +361,7 @@ async fn responses_compact_should_preserve_upstream_client_error_status() {
|
||||
#[tokio::test]
|
||||
async fn responses_compact_should_return_responses_error_when_no_accounts_are_available() {
|
||||
let server = MockServer::start().await;
|
||||
let (app, api_key, _dir) = test_app_without_accounts(server.uri()).await;
|
||||
let (app, api_key, _pool, _dir) = test_app_without_accounts(server.uri()).await;
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
|
||||
@@ -22,7 +22,8 @@ use codex_proxy_rs::{
|
||||
http::router,
|
||||
infra::database::connect_sqlite,
|
||||
proxy::dispatch::session_affinity::{
|
||||
build_conversation_identity, SessionAffinityEntry, SqliteSessionAffinityStore,
|
||||
build_conversation_identity, compute_variant_hash, prepare_variant_identity,
|
||||
SessionAffinityEntry, SqliteSessionAffinityStore,
|
||||
},
|
||||
runtime::services::{BackgroundTaskStores, Services, UsageRecordOptions},
|
||||
runtime::state::AppState,
|
||||
@@ -31,6 +32,7 @@ use codex_proxy_rs::{
|
||||
store::SqliteAccountStore, token_refresh::RefreshLeaseStore,
|
||||
},
|
||||
upstream::fingerprint::FingerprintRepository,
|
||||
upstream::protocol::responses::CodexResponsesRequest,
|
||||
};
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use serde_json::{json, Value};
|
||||
@@ -318,7 +320,16 @@ async fn test_app_with_account_and_pool(
|
||||
test_app_with_account_pool_config(base_url, |_| {}).await
|
||||
}
|
||||
|
||||
async fn test_app_without_accounts(base_url: String) -> (axum::Router, String, tempfile::TempDir) {
|
||||
async fn test_app_with_account_pool_and_affinity(
|
||||
base_url: String,
|
||||
response_id: &str,
|
||||
) -> (axum::Router, String, SqlitePool, tempfile::TempDir) {
|
||||
test_app_with_account_pool_config_and_affinity(base_url, |_| {}, Some(response_id)).await
|
||||
}
|
||||
|
||||
async fn test_app_without_accounts(
|
||||
base_url: String,
|
||||
) -> (axum::Router, String, SqlitePool, tempfile::TempDir) {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("openai-responses-no-accounts.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
@@ -327,7 +338,7 @@ async fn test_app_without_accounts(base_url: String) -> (axum::Router, String, t
|
||||
insert_model_snapshot(&pool).await;
|
||||
let state = test_app_state_with_pool_and_installation_id(
|
||||
&test_config(url, base_url),
|
||||
pool,
|
||||
pool.clone(),
|
||||
TEST_INSTALLATION_ID.to_string(),
|
||||
)
|
||||
.await;
|
||||
@@ -337,7 +348,7 @@ async fn test_app_without_accounts(base_url: String) -> (axum::Router, String, t
|
||||
.restore_from_repository()
|
||||
.await
|
||||
.expect("empty account pool should restore");
|
||||
(router::router().with_state(state), api_key, dir)
|
||||
(router::router().with_state(state), api_key, pool, dir)
|
||||
}
|
||||
|
||||
async fn test_app_with_account_pool_and_logging(
|
||||
@@ -407,6 +418,14 @@ async fn test_app_with_account_pool_and_logging_capture_body(
|
||||
async fn test_app_with_account_pool_config(
|
||||
base_url: String,
|
||||
configure: impl FnOnce(&mut AppConfig),
|
||||
) -> (axum::Router, String, SqlitePool, tempfile::TempDir) {
|
||||
test_app_with_account_pool_config_and_affinity(base_url, configure, None).await
|
||||
}
|
||||
|
||||
async fn test_app_with_account_pool_config_and_affinity(
|
||||
base_url: String,
|
||||
configure: impl FnOnce(&mut AppConfig),
|
||||
affinity_response_id: Option<&str>,
|
||||
) -> (axum::Router, String, SqlitePool, tempfile::TempDir) {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("openai-record-affinity.sqlite");
|
||||
@@ -414,6 +433,9 @@ async fn test_app_with_account_pool_config(
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
let api_key = insert_client_api_key(&pool).await;
|
||||
insert_account(&pool).await;
|
||||
if let Some(response_id) = affinity_response_id {
|
||||
insert_session_affinity(&pool, response_id, "acct_chat").await;
|
||||
}
|
||||
let mut config = test_config(url, base_url);
|
||||
configure(&mut config);
|
||||
let state = test_app_state_with_pool_and_installation_id(
|
||||
@@ -428,6 +450,14 @@ async fn test_app_with_account_pool_config(
|
||||
.restore_from_repository()
|
||||
.await
|
||||
.expect("account pool should restore");
|
||||
if affinity_response_id.is_some() {
|
||||
state
|
||||
.services
|
||||
.session_affinity
|
||||
.restore_from_repository(Utc::now())
|
||||
.await
|
||||
.expect("session affinity should restore");
|
||||
}
|
||||
(router::router().with_state(state), api_key, pool, dir)
|
||||
}
|
||||
|
||||
@@ -713,6 +743,26 @@ async fn insert_model_snapshot(pool: &SqlitePool) {
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn insert_session_affinity(pool: &SqlitePool, response_id: &str, account_id: &str) {
|
||||
SqliteSessionAffinityStore::new(pool.clone())
|
||||
.upsert(
|
||||
response_id,
|
||||
&SessionAffinityEntry {
|
||||
account_id: account_id.to_string(),
|
||||
conversation_id: format!("conv_{response_id}"),
|
||||
turn_state: Some("turn-stale".to_string()),
|
||||
instructions_hash: None,
|
||||
input_tokens: None,
|
||||
function_call_ids: Vec::new(),
|
||||
variant_hash: None,
|
||||
created_at: Utc::now(),
|
||||
},
|
||||
Duration::hours(4),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn insert_named_account(
|
||||
pool: &SqlitePool,
|
||||
id: &str,
|
||||
@@ -1393,10 +1443,52 @@ struct ResponseUsageRecordSnapshot {
|
||||
metadata_json: String,
|
||||
}
|
||||
|
||||
#[derive(sqlx::FromRow)]
|
||||
struct OpsErrorLogSnapshot {
|
||||
request_id: Option<String>,
|
||||
account_id: Option<String>,
|
||||
route: Option<String>,
|
||||
model: Option<String>,
|
||||
status_code: Option<i64>,
|
||||
response_id: Option<String>,
|
||||
metadata_json: String,
|
||||
}
|
||||
|
||||
impl From<OpsErrorLogSnapshot> for ResponseUsageRecordSnapshot {
|
||||
fn from(snapshot: OpsErrorLogSnapshot) -> Self {
|
||||
Self {
|
||||
level: "error".to_string(),
|
||||
request_id: snapshot.request_id,
|
||||
account_id: snapshot.account_id,
|
||||
route: snapshot.route,
|
||||
model: snapshot.model,
|
||||
status_code: snapshot.status_code,
|
||||
response_id: snapshot.response_id,
|
||||
metadata_json: snapshot.metadata_json,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn latest_response_usage_record(pool: &SqlitePool) -> ResponseUsageRecordSnapshot {
|
||||
latest_usage_record(pool, "v1.response").await
|
||||
}
|
||||
|
||||
async fn latest_response_ops_error_log(pool: &SqlitePool) -> ResponseUsageRecordSnapshot {
|
||||
let row = sqlx::query_as::<_, OpsErrorLogSnapshot>(
|
||||
"select request_id, account_id, route, model, status_code, response_id, metadata_json
|
||||
from ops_error_logs
|
||||
where kind = ?
|
||||
order by created_at desc, id desc
|
||||
limit 1",
|
||||
)
|
||||
.bind("v1.response")
|
||||
.fetch_optional(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
row.unwrap_or_else(|| panic!("expected a v1.response ops error log"))
|
||||
.into()
|
||||
}
|
||||
|
||||
async fn latest_usage_record(pool: &SqlitePool, kind: &str) -> ResponseUsageRecordSnapshot {
|
||||
let page = SqliteUsageRecordStore::new(pool.clone())
|
||||
.list(
|
||||
|
||||
@@ -1249,7 +1249,7 @@ async fn responses_should_strip_banned_affinity_history_when_switching_to_active
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_should_keep_history_when_affinity_account_is_quota_exhausted() {
|
||||
async fn responses_should_strip_quota_history_when_switching_to_active_account() {
|
||||
let (base_url, upstream) =
|
||||
spawn_single_websocket_completed_upstream("resp_after_quota_affinity").await;
|
||||
let (app, api_key, pool, _dir) =
|
||||
@@ -1287,13 +1287,7 @@ async fn responses_should_keep_history_when_affinity_account_is_quota_exhausted(
|
||||
captured_header(&captured.headers, "authorization"),
|
||||
Some("Bearer access-primary")
|
||||
);
|
||||
assert_eq!(
|
||||
captured_header(&captured.headers, "x-codex-turn-state"),
|
||||
Some("turn_affinity_risk")
|
||||
);
|
||||
assert_eq!(
|
||||
captured.payload["previous_response_id"],
|
||||
"resp_affinity_risk"
|
||||
);
|
||||
assert!(captured_header(&captured.headers, "x-codex-turn-state").is_none());
|
||||
assert!(captured.payload.get("previous_response_id").is_none());
|
||||
assert_eq!(affinity_count.0, 1);
|
||||
}
|
||||
|
||||
@@ -93,7 +93,7 @@ async fn responses_should_classify_sse_cyber_policy_failure_as_bad_request() {
|
||||
#[tokio::test]
|
||||
async fn responses_should_return_no_available_accounts_error_when_no_accounts_are_available() {
|
||||
let server = MockServer::start().await;
|
||||
let (app, api_key, _dir) = test_app_without_accounts(server.uri()).await;
|
||||
let (app, api_key, pool, _dir) = test_app_without_accounts(server.uri()).await;
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
@@ -121,6 +121,35 @@ async fn responses_should_return_no_available_accounts_error_when_no_accounts_ar
|
||||
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(body["error"]["type"], "server_error");
|
||||
assert_eq!(body["error"]["code"], "no_available_accounts");
|
||||
|
||||
let (usage_count,): (i64,) = sqlx::query_as("select count(*) from usage_records")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let (ops_count, account_id, model, failure_class): (
|
||||
i64,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
) = sqlx::query_as(
|
||||
"select count(*), max(account_id), max(model), max(failure_class) from ops_error_logs",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let (bucket_requests, bucket_errors): (i64, i64) = sqlx::query_as(
|
||||
"select coalesce(sum(request_count), 0), coalesce(sum(error_count), 0) from usage_time_buckets",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(usage_count, 0);
|
||||
assert_eq!(ops_count, 1);
|
||||
assert_eq!(account_id, None);
|
||||
assert_eq!(model.as_deref(), Some("gpt-5.5"));
|
||||
assert_eq!(failure_class.as_deref(), Some("no_available_accounts"));
|
||||
assert_eq!((bucket_requests, bucket_errors), (1, 1));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -1021,7 +1050,8 @@ async fn responses_should_strip_history_after_previous_response_not_found() {
|
||||
"Previous response with id resp_stale was not found",
|
||||
))
|
||||
.await;
|
||||
let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await;
|
||||
let (app, api_key, _pool, _dir) =
|
||||
test_app_with_account_pool_and_affinity(base_url, "resp_stale").await;
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
@@ -1073,7 +1103,8 @@ async fn responses_should_strip_history_after_unanswered_function_call() {
|
||||
"No tool output found for function call call_1",
|
||||
))
|
||||
.await;
|
||||
let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await;
|
||||
let (app, api_key, _pool, _dir) =
|
||||
test_app_with_account_pool_and_affinity(base_url, "resp_stale").await;
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
@@ -1117,7 +1148,8 @@ async fn responses_should_strip_history_after_websocket_invalid_encrypted_reason
|
||||
"Invalid encrypted content in reasoning replay",
|
||||
))
|
||||
.await;
|
||||
let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await;
|
||||
let (app, api_key, _pool, _dir) =
|
||||
test_app_with_account_pool_and_affinity(base_url, "resp_stale").await;
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
@@ -1161,7 +1193,8 @@ async fn responses_should_strip_history_after_sse_invalid_encrypted_reasoning_re
|
||||
"Invalid encrypted content in reasoning replay",
|
||||
))
|
||||
.await;
|
||||
let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await;
|
||||
let (app, api_key, _pool, _dir) =
|
||||
test_app_with_account_pool_and_affinity(base_url, "resp_stale").await;
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
@@ -1205,7 +1238,8 @@ async fn responses_stream_should_strip_history_after_previous_response_not_found
|
||||
"Previous response with id resp_stale was not found",
|
||||
))
|
||||
.await;
|
||||
let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await;
|
||||
let (app, api_key, _pool, _dir) =
|
||||
test_app_with_account_pool_and_affinity(base_url, "resp_stale").await;
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
@@ -1258,7 +1292,8 @@ async fn responses_stream_should_strip_history_after_websocket_invalid_encrypted
|
||||
"Invalid encrypted content in reasoning replay",
|
||||
))
|
||||
.await;
|
||||
let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await;
|
||||
let (app, api_key, _pool, _dir) =
|
||||
test_app_with_account_pool_and_affinity(base_url, "resp_stale").await;
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
@@ -2122,7 +2157,7 @@ async fn responses_should_return_rate_limit_error_when_429_fallback_is_exhausted
|
||||
assert_eq!(body["error"]["code"], "rate_limit_exceeded");
|
||||
assert_eq!(quota_state.0, 1);
|
||||
assert!(quota_state.1.is_some());
|
||||
let event = latest_response_usage_record(&pool).await;
|
||||
let event = latest_response_ops_error_log(&pool).await;
|
||||
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
|
||||
assert_eq!(event.level, "error");
|
||||
assert_eq!(
|
||||
@@ -2206,7 +2241,7 @@ async fn responses_stream_should_return_rate_limit_error_when_429_fallback_is_ex
|
||||
assert_eq!(quota_state.0, 1);
|
||||
assert!(quota_state.1.is_some());
|
||||
assert!(!response_request_id.is_empty());
|
||||
let event = latest_response_usage_record(&pool).await;
|
||||
let event = latest_response_ops_error_log(&pool).await;
|
||||
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
|
||||
assert_eq!(event.level, "error");
|
||||
assert_eq!(
|
||||
|
||||
@@ -1539,24 +1539,76 @@ async fn responses_websocket_should_route_previous_response_id_to_recorded_accou
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_non_stream_previous_response_not_found_should_strip_history_and_retry_same_account(
|
||||
async fn responses_websocket_should_prefer_conversation_account_without_previous_response_id() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base_url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let upstream = tokio::spawn(async move {
|
||||
accept_successful_websocket_response_with_authorization(
|
||||
&listener,
|
||||
"Bearer access-secondary",
|
||||
"resp_conversation_affinity",
|
||||
)
|
||||
.await
|
||||
});
|
||||
let (app, state, api_key, _pool, _dir) = test_app_with_two_accounts_and_state(base_url).await;
|
||||
let request_body = json!({
|
||||
"model": "gpt-5.5",
|
||||
"prompt_cache_key": "conv_conversation_affinity",
|
||||
"input": [],
|
||||
"stream": false,
|
||||
"use_websocket": true
|
||||
});
|
||||
let mut affinity_body = request_body.as_object().unwrap().clone();
|
||||
affinity_body.remove("use_websocket");
|
||||
let mut affinity_request = CodexResponsesRequest::from_body(affinity_body);
|
||||
affinity_request.explicit_prompt_cache_key = true;
|
||||
prepare_variant_identity(&mut affinity_request);
|
||||
let variant_hash = compute_variant_hash(&affinity_request);
|
||||
state
|
||||
.services
|
||||
.session_affinity
|
||||
.record(
|
||||
"resp_conversation_latest".to_string(),
|
||||
SessionAffinityEntry {
|
||||
account_id: "acct_secondary".to_string(),
|
||||
conversation_id: "conv_conversation_affinity".to_string(),
|
||||
turn_state: None,
|
||||
instructions_hash: None,
|
||||
input_tokens: None,
|
||||
function_call_ids: Vec::new(),
|
||||
variant_hash: Some(variant_hash),
|
||||
created_at: Utc::now(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let response = app
|
||||
.oneshot(responses_json_request(&api_key, &request_body))
|
||||
.await
|
||||
.unwrap();
|
||||
let body = response_text(response).await;
|
||||
let payload = upstream.await.unwrap();
|
||||
|
||||
assert!(body.contains("\"id\":\"resp_conversation_affinity\""));
|
||||
assert!(payload["prompt_cache_key"]
|
||||
.as_str()
|
||||
.is_some_and(|value| value.starts_with("cp_")));
|
||||
assert!(payload.get("previous_response_id").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_non_stream_unknown_previous_response_should_strip_history_before_upstream(
|
||||
) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base_url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let upstream = tokio::spawn(async move {
|
||||
let first_payload = accept_websocket_response_with_authorization_and_message(
|
||||
&listener,
|
||||
"Bearer access-secret",
|
||||
WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND.trim().to_string(),
|
||||
)
|
||||
.await;
|
||||
let second_payload = accept_websocket_response_with_authorization_and_message(
|
||||
accept_websocket_response_with_authorization_and_message(
|
||||
&listener,
|
||||
"Bearer access-secret",
|
||||
websocket_completed_response("resp_after_history_strip", 3, 1),
|
||||
)
|
||||
.await;
|
||||
(first_payload, second_payload)
|
||||
.await
|
||||
});
|
||||
let (app, api_key, _dir) = test_app_with_account(base_url).await;
|
||||
|
||||
@@ -1573,32 +1625,23 @@ async fn responses_websocket_non_stream_previous_response_not_found_should_strip
|
||||
.await
|
||||
.unwrap();
|
||||
let body = response_json(response).await;
|
||||
let (first_payload, second_payload) = upstream.await.unwrap();
|
||||
let payload = upstream.await.unwrap();
|
||||
|
||||
assert_eq!(body["id"], "resp_after_history_strip");
|
||||
assert_eq!(first_payload["previous_response_id"], "resp_missing");
|
||||
assert!(second_payload.get("previous_response_id").is_none());
|
||||
assert!(payload.get("previous_response_id").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_history_recovery_should_stay_on_fallback_account_when_preferred_is_full(
|
||||
) {
|
||||
async fn responses_websocket_should_strip_known_history_when_preferred_account_is_full() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base_url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let upstream = tokio::spawn(async move {
|
||||
let first_payload = accept_websocket_response_with_authorization_and_message(
|
||||
&listener,
|
||||
"Bearer access-secondary",
|
||||
WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND.trim().to_string(),
|
||||
)
|
||||
.await;
|
||||
let second_payload = accept_websocket_response_with_authorization_and_message(
|
||||
accept_websocket_response_with_authorization_and_message(
|
||||
&listener,
|
||||
"Bearer access-secondary",
|
||||
websocket_completed_response("resp_after_fallback_history_strip", 3, 1),
|
||||
)
|
||||
.await;
|
||||
(first_payload, second_payload)
|
||||
.await
|
||||
});
|
||||
let (app, state, api_key, _pool, _dir) =
|
||||
test_app_with_two_accounts_and_state_config(base_url, |config| {
|
||||
@@ -1620,35 +1663,23 @@ async fn responses_websocket_history_recovery_should_stay_on_fallback_account_wh
|
||||
.await
|
||||
.unwrap();
|
||||
let body = response_json(response).await;
|
||||
let (first_payload, second_payload) = upstream.await.unwrap();
|
||||
let payload = upstream.await.unwrap();
|
||||
|
||||
assert_eq!(body["id"], "resp_after_fallback_history_strip");
|
||||
assert_eq!(
|
||||
first_payload["previous_response_id"],
|
||||
"resp_owned_by_primary"
|
||||
);
|
||||
assert!(second_payload.get("previous_response_id").is_none());
|
||||
assert!(payload.get("previous_response_id").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_stream_history_recovery_should_stay_on_fallback_account_when_preferred_is_full(
|
||||
) {
|
||||
async fn responses_websocket_stream_should_strip_known_history_when_preferred_account_is_full() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base_url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let upstream = tokio::spawn(async move {
|
||||
let first_payload = accept_websocket_response_with_authorization_and_message(
|
||||
&listener,
|
||||
"Bearer access-secondary",
|
||||
WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND.trim().to_string(),
|
||||
)
|
||||
.await;
|
||||
let second_payload = accept_websocket_response_with_authorization_and_message(
|
||||
accept_websocket_response_with_authorization_and_message(
|
||||
&listener,
|
||||
"Bearer access-secondary",
|
||||
websocket_completed_response("resp_stream_after_fallback_history_strip", 3, 1),
|
||||
)
|
||||
.await;
|
||||
(first_payload, second_payload)
|
||||
.await
|
||||
});
|
||||
let (app, state, api_key, _pool, _dir) =
|
||||
test_app_with_two_accounts_and_state_config(base_url, |config| {
|
||||
@@ -1670,35 +1701,24 @@ async fn responses_websocket_stream_history_recovery_should_stay_on_fallback_acc
|
||||
.await
|
||||
.unwrap();
|
||||
let body = response_text(response).await;
|
||||
let (first_payload, second_payload) = upstream.await.unwrap();
|
||||
let payload = upstream.await.unwrap();
|
||||
|
||||
assert!(body.contains("resp_stream_after_fallback_history_strip"));
|
||||
assert_eq!(
|
||||
first_payload["previous_response_id"],
|
||||
"resp_owned_by_primary"
|
||||
);
|
||||
assert!(second_payload.get("previous_response_id").is_none());
|
||||
assert!(payload.get("previous_response_id").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_stream_previous_response_not_found_should_strip_history_and_retry_same_account(
|
||||
) {
|
||||
async fn responses_websocket_stream_unknown_previous_response_should_strip_history_before_upstream()
|
||||
{
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base_url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let upstream = tokio::spawn(async move {
|
||||
let first_payload = accept_websocket_response_with_authorization_and_message(
|
||||
&listener,
|
||||
"Bearer access-secret",
|
||||
WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND.trim().to_string(),
|
||||
)
|
||||
.await;
|
||||
let second_payload = accept_websocket_response_with_authorization_and_message(
|
||||
accept_websocket_response_with_authorization_and_message(
|
||||
&listener,
|
||||
"Bearer access-secret",
|
||||
websocket_completed_response("resp_stream_after_history_strip", 3, 1),
|
||||
)
|
||||
.await;
|
||||
(first_payload, second_payload)
|
||||
.await
|
||||
});
|
||||
let (app, api_key, _dir) = test_app_with_account(base_url).await;
|
||||
|
||||
@@ -1715,11 +1735,10 @@ async fn responses_websocket_stream_previous_response_not_found_should_strip_his
|
||||
.await
|
||||
.unwrap();
|
||||
let body = response_text(response).await;
|
||||
let (first_payload, second_payload) = upstream.await.unwrap();
|
||||
let payload = upstream.await.unwrap();
|
||||
|
||||
assert!(body.contains("\"id\":\"resp_stream_after_history_strip\""));
|
||||
assert_eq!(first_payload["previous_response_id"], "resp_missing");
|
||||
assert!(second_payload.get("previous_response_id").is_none());
|
||||
assert!(payload.get("previous_response_id").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -1742,7 +1761,8 @@ async fn responses_websocket_non_stream_unanswered_function_call_should_strip_hi
|
||||
.await;
|
||||
(first_payload, second_payload)
|
||||
});
|
||||
let (app, api_key, _dir) = test_app_with_account(base_url).await;
|
||||
let (app, api_key, _pool, _dir) =
|
||||
test_app_with_account_pool_and_affinity(base_url, "resp_with_call").await;
|
||||
|
||||
let response = app
|
||||
.oneshot(responses_json_request(
|
||||
@@ -1808,7 +1828,9 @@ async fn responses_websocket_previous_response_id_should_retry_fallback_account_
|
||||
.unwrap();
|
||||
|
||||
assert!(body.contains("\"id\":\"resp_history_fallback\""));
|
||||
assert_eq!(fallback_payload["previous_response_id"], "resp_prev");
|
||||
// 亲和缓存无 `resp_prev` 记录(Unknown owner),换账号前历史即被剥离,
|
||||
// 由 `input` 让备用账号重建上下文,不把上一账号作用域的 id 透传给它。
|
||||
assert_eq!(fallback_payload.get("previous_response_id"), None);
|
||||
assert_eq!(secondary_usage.0, 1);
|
||||
}
|
||||
|
||||
@@ -1858,7 +1880,9 @@ async fn responses_websocket_non_stream_previous_response_id_should_retry_fallba
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(body["id"], "resp_history_fallback_non_stream");
|
||||
assert_eq!(fallback_payload["previous_response_id"], "resp_prev");
|
||||
// 亲和缓存无 `resp_prev` 记录(Unknown owner),换账号前历史即被剥离,
|
||||
// 由 `input` 让备用账号重建上下文,不把上一账号作用域的 id 透传给它。
|
||||
assert_eq!(fallback_payload.get("previous_response_id"), None);
|
||||
assert_eq!(secondary_usage.0, 1);
|
||||
}
|
||||
|
||||
@@ -2304,10 +2328,9 @@ async fn responses_stream_with_previous_response_id_should_forward_websocket_chu
|
||||
captured_header(&captured.headers, "authorization"),
|
||||
Some("Bearer access-secret")
|
||||
);
|
||||
assert_eq!(
|
||||
captured.payload["previous_response_id"],
|
||||
"resp_runtime_pool_previous"
|
||||
);
|
||||
// 亲和缓存无该 previous_response_id 的记录(Unknown),故上游 payload 里应被剥离,
|
||||
// 靠已展开的 input 让选中账号自行重建续链——与 sub2api/codex2api 的默认行为一致。
|
||||
assert_eq!(captured.payload["previous_response_id"], Value::Null);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -2335,10 +2358,8 @@ async fn responses_stream_with_previous_response_id_should_record_websocket_audi
|
||||
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
|
||||
|
||||
assert!(body.contains("resp_live_websocket_stream"));
|
||||
assert_eq!(
|
||||
captured.payload["previous_response_id"],
|
||||
"resp_runtime_pool_previous"
|
||||
);
|
||||
// 亲和缓存无记录(Unknown)→ 上游 payload 里 previous_response_id 被剥离。
|
||||
assert_eq!(captured.payload["previous_response_id"], Value::Null);
|
||||
assert_eq!(event.level, "info");
|
||||
assert_eq!(metadata["stream"], true);
|
||||
assert_eq!(metadata["transport"], "websocket");
|
||||
|
||||
@@ -678,7 +678,7 @@ async fn responses_stream_should_record_usage_record_after_late_disconnect() {
|
||||
let response = response_task.await.unwrap();
|
||||
close_upstream.send(()).unwrap();
|
||||
let body = response_text(response).await;
|
||||
let event = latest_response_usage_record(&pool).await;
|
||||
let event = latest_response_ops_error_log(&pool).await;
|
||||
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
|
||||
|
||||
assert!(body.contains("stream_disconnected"));
|
||||
@@ -744,7 +744,7 @@ async fn responses_stream_should_record_failure_detail_after_transport_error() {
|
||||
let response = response_task.await.unwrap();
|
||||
close_upstream.send(()).unwrap();
|
||||
let _body = response_text(response).await;
|
||||
let event = latest_response_usage_record(&pool).await;
|
||||
let event = latest_response_ops_error_log(&pool).await;
|
||||
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
|
||||
|
||||
assert_eq!(metadata["failureSource"], "proxy");
|
||||
@@ -788,7 +788,7 @@ async fn responses_stream_should_not_record_first_token_when_metadata_only_prefi
|
||||
let response = response_task.await.unwrap();
|
||||
close_upstream.send(()).unwrap();
|
||||
let body = response_text(response).await;
|
||||
let event = latest_response_usage_record(&pool).await;
|
||||
let event = latest_response_ops_error_log(&pool).await;
|
||||
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
|
||||
let first_token_bucket: (i64, i64) = sqlx::query_as(
|
||||
"select coalesce(sum(first_token_latency_sum), 0), coalesce(sum(first_token_latency_count), 0) from usage_time_buckets",
|
||||
@@ -855,7 +855,7 @@ async fn responses_should_record_request_count_when_5xx_retries_are_exhausted()
|
||||
assert_eq!(status, StatusCode::BAD_GATEWAY);
|
||||
assert_eq!(authorizations.len(), 3, "requests: {authorizations:?}");
|
||||
assert_eq!(failed_usage.map(|row| row.0).unwrap_or_default(), 1);
|
||||
let event = latest_response_usage_record(&pool).await;
|
||||
let event = latest_response_ops_error_log(&pool).await;
|
||||
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
|
||||
assert_eq!(event.account_id.as_deref(), Some("acct_chat"));
|
||||
assert_eq!(metadata["accountEmail"], "user@example.com");
|
||||
@@ -906,7 +906,7 @@ async fn responses_should_record_attempt_trace_after_retrying_another_account()
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let event = latest_response_usage_record(&pool).await;
|
||||
let event = latest_response_ops_error_log(&pool).await;
|
||||
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
|
||||
|
||||
@@ -261,6 +261,14 @@ fn session_affinity_map_should_find_latest_response_by_conversation_and_variant(
|
||||
map.lookup_latest_response_by_conversation("conv", None, Some("variant_a"), now),
|
||||
Some("resp_new".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
map.lookup_latest_account_by_conversation("conv", None, Some("variant_a"), now),
|
||||
Some("acct_b".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
map.lookup_latest_account_by_conversation("conv", None, Some("variant_b"), now),
|
||||
Some("acct_c".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
map.lookup_turn_state("resp_new", now),
|
||||
Some("turn_new".to_string())
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
use super::*;
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
sync::atomic::{AtomicUsize, Ordering},
|
||||
};
|
||||
|
||||
use codex_proxy_rs::upstream::{
|
||||
accounts::pool::{AccountAcquireRequest, AccountPoolOptions, RuntimeAccountPoolService},
|
||||
@@ -50,6 +53,8 @@ async fn model_refresh_task_should_sync_model_plan_allowlist_to_account_pool() {
|
||||
plan_account("acct-plus", "plus"),
|
||||
plan_account("acct-team", "team"),
|
||||
],
|
||||
request_usage_records: Arc::new(AtomicUsize::new(0)),
|
||||
model_request_usage_records: Arc::new(AtomicUsize::new(0)),
|
||||
});
|
||||
let account_pool = Arc::new(RuntimeAccountPoolService::new(
|
||||
account_store.clone(),
|
||||
@@ -74,6 +79,16 @@ async fn model_refresh_task_should_sync_model_plan_allowlist_to_account_pool() {
|
||||
.await
|
||||
.expect("team model should acquire an account");
|
||||
assert_eq!(acquired.account.id, "acct-team");
|
||||
assert_eq!(
|
||||
account_store.request_usage_records.load(Ordering::SeqCst),
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
account_store
|
||||
.model_request_usage_records
|
||||
.load(Ordering::SeqCst),
|
||||
0
|
||||
);
|
||||
}
|
||||
|
||||
fn empty_model_config() -> ModelConfig {
|
||||
@@ -139,6 +154,8 @@ impl CodexModelCatalogClient for FakeModelCatalogClient {
|
||||
|
||||
struct PlanAccountStore {
|
||||
accounts: Vec<Account>,
|
||||
request_usage_records: Arc<AtomicUsize>,
|
||||
model_request_usage_records: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -174,8 +191,11 @@ impl AccountStore for PlanAccountStore {
|
||||
async fn record_usage_delta(
|
||||
&self,
|
||||
_account_id: &str,
|
||||
_usage: AccountUsageDelta,
|
||||
usage: AccountUsageDelta,
|
||||
) -> AccountStoreResult<()> {
|
||||
if usage.requests > 0 {
|
||||
self.request_usage_records.fetch_add(1, Ordering::SeqCst);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -183,8 +203,12 @@ impl AccountStore for PlanAccountStore {
|
||||
&self,
|
||||
_account_id: &str,
|
||||
_model: &str,
|
||||
_usage: AccountModelUsageDelta,
|
||||
usage: AccountModelUsageDelta,
|
||||
) -> AccountStoreResult<()> {
|
||||
if usage.requests > 0 {
|
||||
self.model_request_usage_records
|
||||
.fetch_add(1, Ordering::SeqCst);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@@ -8,7 +8,6 @@ pub(crate) fn test_config(database_url: String) -> AppConfig {
|
||||
server: ServerConfig {
|
||||
host: "127.0.0.1".to_string(),
|
||||
port: 0,
|
||||
trusted_proxies: Vec::new(),
|
||||
},
|
||||
api: ApiConfig {
|
||||
base_url: "https://chatgpt.com/backend-api".to_string(),
|
||||
|
||||
@@ -39,6 +39,22 @@ fn account_pool_should_mark_quota_state_as_quota_exhausted_status() {
|
||||
assert!(!restored.quota_limit_reached);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn account_pool_should_mark_quota_cooldown_as_quota_exhausted_status() {
|
||||
let now = fixed_time();
|
||||
let mut pool = AccountPool::default();
|
||||
pool.insert(crate::support::accounts::test_account(
|
||||
"limited",
|
||||
AccountStatus::Active,
|
||||
));
|
||||
|
||||
pool.mark_quota_limited_until("limited", now + Duration::seconds(60));
|
||||
|
||||
let limited = pool.get("limited").unwrap();
|
||||
assert_eq!(limited.status, AccountStatus::QuotaExhausted);
|
||||
assert!(limited.quota_limit_reached);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn account_pool_should_reuse_quota_limited_accounts_after_cooldown() {
|
||||
let now = fixed_time();
|
||||
|
||||
@@ -562,6 +562,34 @@ async fn account_repository_should_mark_quota_exhausted_for_any_exhausted_window
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_repository_should_mark_quota_cooldown_as_quota_exhausted_status() {
|
||||
let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 48).await;
|
||||
let repo = SqliteAccountStore::new(pool.clone());
|
||||
seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await;
|
||||
let future_cooldown = Utc::now() + Duration::hours(1);
|
||||
|
||||
repo.mark_quota_limited_until("acct_a", future_cooldown)
|
||||
.await
|
||||
.unwrap();
|
||||
let stored: (String, i64, Option<String>) = sqlx::query_as(
|
||||
"select status, quota_limit_reached, quota_cooldown_until from accounts where id = ?",
|
||||
)
|
||||
.bind("acct_a")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
(stored.0, stored.1, stored.2),
|
||||
(
|
||||
String::from("quota_exhausted"),
|
||||
1,
|
||||
Some(future_cooldown.to_rfc3339())
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_repository_runtime_state_sync_should_not_clear_newer_future_quota_cooldown() {
|
||||
let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 42).await;
|
||||
@@ -569,7 +597,7 @@ async fn account_repository_runtime_state_sync_should_not_clear_newer_future_quo
|
||||
seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await;
|
||||
let future_cooldown = Utc::now() + Duration::hours(1);
|
||||
sqlx::query(
|
||||
"update accounts set quota_limit_reached = 1, quota_verify_required = 0, quota_cooldown_until = ? where id = ?",
|
||||
"update accounts set status = 'quota_exhausted', quota_limit_reached = 1, quota_verify_required = 0, quota_cooldown_until = ? where id = ?",
|
||||
)
|
||||
.bind(future_cooldown.to_rfc3339())
|
||||
.bind("acct_a")
|
||||
@@ -582,6 +610,7 @@ async fn account_repository_runtime_state_sync_should_not_clear_newer_future_quo
|
||||
.into_iter()
|
||||
.find(|account| account.id == "acct_a")
|
||||
.unwrap();
|
||||
stale_account.status = AccountStatus::Active;
|
||||
stale_account.quota_limit_reached = false;
|
||||
stale_account.quota_verify_required = true;
|
||||
stale_account.quota_cooldown_until = None;
|
||||
@@ -589,17 +618,18 @@ async fn account_repository_runtime_state_sync_should_not_clear_newer_future_quo
|
||||
repo.sync_runtime_account_state(&stale_account, false)
|
||||
.await
|
||||
.unwrap();
|
||||
let stored: (i64, i64, Option<String>) = sqlx::query_as(
|
||||
"select quota_limit_reached, quota_verify_required, quota_cooldown_until from accounts where id = ?",
|
||||
let stored: (String, i64, i64, Option<String>) = sqlx::query_as(
|
||||
"select status, quota_limit_reached, quota_verify_required, quota_cooldown_until from accounts where id = ?",
|
||||
)
|
||||
.bind("acct_a")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(stored.0, 1);
|
||||
assert_eq!(stored.1, 0);
|
||||
assert_eq!(stored.2, Some(future_cooldown.to_rfc3339()));
|
||||
assert_eq!(stored.0, "quota_exhausted");
|
||||
assert_eq!(stored.1, 1);
|
||||
assert_eq!(stored.2, 0);
|
||||
assert_eq!(stored.3, Some(future_cooldown.to_rfc3339()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -68,3 +68,66 @@ async fn refresh_lease_store_should_acquire_release_and_respect_expiry() {
|
||||
.await
|
||||
.expect("lease should be acquirable after release"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_lease_store_should_list_active_account_ids() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let db = dir.path().join("refresh-lease-active-ids.sqlite");
|
||||
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
|
||||
.await
|
||||
.expect("sqlite pool");
|
||||
let accounts = SqliteAccountStore::new(pool.clone());
|
||||
for id in ["acct-active-lease", "acct-expired-lease"] {
|
||||
accounts
|
||||
.insert(NewAccount {
|
||||
id: id.to_string(),
|
||||
email: None,
|
||||
account_id: None,
|
||||
user_id: None,
|
||||
label: None,
|
||||
plan_type: None,
|
||||
access_token: SecretString::new("access".to_string().into()),
|
||||
refresh_token: Some(SecretString::new("refresh".to_string().into())),
|
||||
access_token_expires_at: None,
|
||||
status: AccountStatus::Active,
|
||||
added_at: None,
|
||||
})
|
||||
.await
|
||||
.expect("account should be inserted");
|
||||
}
|
||||
let leases = RefreshLeaseStore::new(pool);
|
||||
let now = Utc.with_ymd_and_hms(2026, 7, 9, 8, 0, 0).unwrap();
|
||||
leases
|
||||
.try_acquire(
|
||||
"acct-active-lease",
|
||||
"owner-a",
|
||||
now + Duration::minutes(5),
|
||||
now,
|
||||
)
|
||||
.await
|
||||
.expect("active lease should be acquired");
|
||||
leases
|
||||
.try_acquire(
|
||||
"acct-expired-lease",
|
||||
"owner-b",
|
||||
now - Duration::minutes(5),
|
||||
now - Duration::minutes(10),
|
||||
)
|
||||
.await
|
||||
.expect("expired lease should be acquired");
|
||||
|
||||
let account_ids = vec![
|
||||
"acct-active-lease".to_string(),
|
||||
"acct-expired-lease".to_string(),
|
||||
"acct-missing".to_string(),
|
||||
];
|
||||
let active = leases
|
||||
.active_account_ids(&account_ids, now)
|
||||
.await
|
||||
.expect("active leases should load");
|
||||
|
||||
assert_eq!(
|
||||
active,
|
||||
std::collections::HashSet::from(["acct-active-lease".to_string()])
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn token_refresh_task_should_mark_due_account_refreshing_before_refresher_call() {
|
||||
async fn token_refresh_task_should_keep_business_status_before_refresher_call() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let db = dir.path().join("token-refresh-status.sqlite");
|
||||
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
|
||||
@@ -13,7 +13,7 @@ async fn token_refresh_task_should_mark_due_account_refreshing_before_refresher_
|
||||
let new_access_token = test_jwt(new_expires_at.timestamp());
|
||||
store
|
||||
.insert(NewAccount {
|
||||
id: "acct-refreshing-marker".to_string(),
|
||||
id: "acct-status-marker".to_string(),
|
||||
email: Some("marker@example.com".to_string()),
|
||||
account_id: Some("chatgpt-marker".to_string()),
|
||||
user_id: Some("user-marker".to_string()),
|
||||
@@ -32,7 +32,7 @@ async fn token_refresh_task_should_mark_due_account_refreshing_before_refresher_
|
||||
let observed_statuses = Arc::new(Mutex::new(Vec::new()));
|
||||
let refresher = StoreObservingTokenRefresher {
|
||||
store: store.clone(),
|
||||
account_id: "acct-refreshing-marker".to_string(),
|
||||
account_id: "acct-status-marker".to_string(),
|
||||
observed_statuses: observed_statuses.clone(),
|
||||
response: Ok(TokenPair {
|
||||
access_token: new_access_token.clone(),
|
||||
@@ -53,14 +53,14 @@ async fn token_refresh_task_should_mark_due_account_refreshing_before_refresher_
|
||||
.await
|
||||
.expect("refresh should succeed");
|
||||
let stored = store
|
||||
.get("acct-refreshing-marker")
|
||||
.get("acct-status-marker")
|
||||
.await
|
||||
.expect("account should load")
|
||||
.expect("account should exist");
|
||||
let observed_statuses = observed_statuses.lock().await.clone();
|
||||
|
||||
assert_eq!(summary.refreshed, 1);
|
||||
assert_eq!(observed_statuses, [AccountStatus::Refreshing]);
|
||||
assert_eq!(observed_statuses, [AccountStatus::Active]);
|
||||
assert_eq!(stored.status, AccountStatus::Active);
|
||||
assert_eq!(
|
||||
stored.access_token.expose_secret(),
|
||||
@@ -69,7 +69,7 @@ async fn token_refresh_task_should_mark_due_account_refreshing_before_refresher_
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn token_refresh_task_should_recover_refreshing_account_after_restart() {
|
||||
async fn token_refresh_task_should_recover_expired_account_after_restart() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let db = dir.path().join("token-refresh-recovery.sqlite");
|
||||
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
|
||||
@@ -81,7 +81,7 @@ async fn token_refresh_task_should_recover_refreshing_account_after_restart() {
|
||||
let new_access_token = test_jwt((now + Duration::hours(2)).timestamp());
|
||||
store
|
||||
.insert(NewAccount {
|
||||
id: "acct-refreshing-recovery".to_string(),
|
||||
id: "acct-expired-recovery".to_string(),
|
||||
email: Some("recovery@example.com".to_string()),
|
||||
account_id: Some("chatgpt-recovery".to_string()),
|
||||
user_id: Some("user-recovery".to_string()),
|
||||
@@ -91,7 +91,7 @@ async fn token_refresh_task_should_recover_refreshing_account_after_restart() {
|
||||
refresh_token: Some(SecretString::new("refresh-recovery".to_string().into())),
|
||||
added_at: None,
|
||||
access_token_expires_at: Some(now + Duration::hours(1)),
|
||||
status: AccountStatus::Refreshing,
|
||||
status: AccountStatus::Expired,
|
||||
})
|
||||
.await
|
||||
.expect("account should be inserted");
|
||||
@@ -113,9 +113,9 @@ async fn token_refresh_task_should_recover_refreshing_account_after_restart() {
|
||||
let summary = task
|
||||
.refresh_due_accounts_once_at(now)
|
||||
.await
|
||||
.expect("refreshing account should recover");
|
||||
.expect("expired account should recover");
|
||||
let stored = store
|
||||
.get("acct-refreshing-recovery")
|
||||
.get("acct-expired-recovery")
|
||||
.await
|
||||
.expect("account should load")
|
||||
.expect("account should exist");
|
||||
@@ -183,7 +183,7 @@ async fn token_refresh_task_should_retry_transport_failure_before_recovery() {
|
||||
let observed_statuses = observed_statuses.lock().await.clone();
|
||||
|
||||
assert_eq!(summary.failed, 1);
|
||||
assert_eq!(observed_statuses, [AccountStatus::Refreshing; 5]);
|
||||
assert_eq!(observed_statuses, [AccountStatus::Active; 5]);
|
||||
assert_eq!(stored.status, AccountStatus::Active);
|
||||
assert!(stored.next_refresh_at.is_some_and(|next| next > now));
|
||||
assert_eq!(
|
||||
@@ -275,7 +275,7 @@ async fn token_refresh_task_should_delay_recovery_after_retry_exhaustion() {
|
||||
assert!(recovery_at <= now + Duration::minutes(12));
|
||||
assert_eq!(delayed.skipped, 1);
|
||||
assert_eq!(refreshed.refreshed, 1);
|
||||
assert_eq!(observed_statuses, [AccountStatus::Refreshing; 6]);
|
||||
assert_eq!(observed_statuses, [AccountStatus::Active; 6]);
|
||||
assert_eq!(stored.status, AccountStatus::Active);
|
||||
assert_eq!(
|
||||
stored.access_token.expose_secret(),
|
||||
@@ -412,7 +412,7 @@ async fn token_refresh_task_should_confirm_invalid_grant_before_expiring_account
|
||||
let observed_statuses = observed_statuses.lock().await.clone();
|
||||
|
||||
assert_eq!(summary.status_updated, 1);
|
||||
assert_eq!(observed_statuses, [AccountStatus::Refreshing; 2]);
|
||||
assert_eq!(observed_statuses, [AccountStatus::Active; 2]);
|
||||
assert_eq!(stored.status, AccountStatus::Expired);
|
||||
assert_eq!(
|
||||
stored.access_token.expose_secret(),
|
||||
|
||||
@@ -494,6 +494,92 @@ async fn token_refresh_task_should_fire_per_account_timer_at_refresh_time() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn token_refresh_task_should_fire_quota_exhausted_timer_without_clearing_status() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let db = dir
|
||||
.path()
|
||||
.join("token-refresh-quota-exhausted-timer.sqlite");
|
||||
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
|
||||
.await
|
||||
.expect("sqlite pool");
|
||||
let store = SqliteAccountStore::new(pool);
|
||||
let now = Utc::now();
|
||||
let old_expires_at = now + Duration::milliseconds(50);
|
||||
let new_access_token = test_jwt((now + Duration::hours(1)).timestamp());
|
||||
store
|
||||
.insert(NewAccount {
|
||||
id: "acct-refresh-quota-exhausted-timer".to_string(),
|
||||
email: Some("quota-exhausted-timer@example.com".to_string()),
|
||||
account_id: Some("chatgpt-quota-exhausted-timer".to_string()),
|
||||
user_id: Some("user-quota-exhausted-timer".to_string()),
|
||||
label: None,
|
||||
plan_type: Some("plus".to_string()),
|
||||
access_token: SecretString::new(test_jwt(old_expires_at.timestamp()).into()),
|
||||
refresh_token: Some(SecretString::new(
|
||||
"refresh-quota-exhausted-timer".to_string().into(),
|
||||
)),
|
||||
added_at: None,
|
||||
access_token_expires_at: Some(old_expires_at),
|
||||
status: AccountStatus::QuotaExhausted,
|
||||
})
|
||||
.await
|
||||
.expect("account should be inserted");
|
||||
let refresher = NotifyingTokenRefresher {
|
||||
calls: Arc::new(AtomicUsize::new(0)),
|
||||
started: Arc::new(Notify::new()),
|
||||
response: Ok(TokenPair {
|
||||
access_token: new_access_token.clone(),
|
||||
refresh_token: None,
|
||||
}),
|
||||
};
|
||||
let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new(
|
||||
store.clone(),
|
||||
RefreshPolicy {
|
||||
refresh_margin_seconds: 0,
|
||||
refresh_concurrency: 1,
|
||||
},
|
||||
refresher.clone(),
|
||||
)
|
||||
.with_retry_delays(vec![StdDuration::ZERO; 4]);
|
||||
|
||||
let started = refresher.started.notified();
|
||||
tokio::pin!(started);
|
||||
let summary = task
|
||||
.schedule_account_timers_once_at(now)
|
||||
.await
|
||||
.expect("quota-exhausted timer should be scheduled");
|
||||
timeout(StdDuration::from_secs(2), &mut started)
|
||||
.await
|
||||
.expect("quota-exhausted timer should trigger refresh");
|
||||
|
||||
let mut stored = store
|
||||
.get("acct-refresh-quota-exhausted-timer")
|
||||
.await
|
||||
.expect("account should load")
|
||||
.expect("account should exist");
|
||||
for _ in 0..50 {
|
||||
if stored.access_token.expose_secret() == new_access_token.as_str() {
|
||||
break;
|
||||
}
|
||||
sleep(StdDuration::from_millis(20)).await;
|
||||
stored = store
|
||||
.get("acct-refresh-quota-exhausted-timer")
|
||||
.await
|
||||
.expect("account should load")
|
||||
.expect("account should exist");
|
||||
}
|
||||
|
||||
assert_eq!(summary.scheduled, 1);
|
||||
assert_eq!(refresher.calls.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(stored.status, AccountStatus::QuotaExhausted);
|
||||
assert_eq!(
|
||||
stored.access_token.expose_secret(),
|
||||
new_access_token.as_str()
|
||||
);
|
||||
assert!(stored.next_refresh_at.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn token_refresh_task_should_reschedule_next_timer_after_scheduled_refresh() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
|
||||
@@ -75,6 +75,160 @@ async fn token_refresh_task_should_persist_refreshed_access_token_and_keep_refre
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn token_refresh_task_should_refresh_quota_exhausted_account_without_clearing_status() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let db = dir.path().join("token-refresh-quota-exhausted.sqlite");
|
||||
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
|
||||
.await
|
||||
.expect("sqlite pool");
|
||||
let store = SqliteAccountStore::new(pool);
|
||||
let now = Utc.with_ymd_and_hms(2026, 6, 19, 11, 0, 0).unwrap();
|
||||
let new_expires_at = now + Duration::hours(1);
|
||||
let new_access_token = test_jwt(new_expires_at.timestamp());
|
||||
store
|
||||
.insert(NewAccount {
|
||||
id: "acct-refresh-quota-exhausted".to_string(),
|
||||
email: Some("quota-exhausted@example.com".to_string()),
|
||||
account_id: Some("chatgpt-quota-exhausted".to_string()),
|
||||
user_id: Some("user-quota-exhausted".to_string()),
|
||||
label: None,
|
||||
plan_type: Some("plus".to_string()),
|
||||
access_token: SecretString::new(
|
||||
test_jwt((now + Duration::seconds(30)).timestamp()).into(),
|
||||
),
|
||||
refresh_token: Some(SecretString::new(
|
||||
"refresh-quota-exhausted".to_string().into(),
|
||||
)),
|
||||
added_at: None,
|
||||
access_token_expires_at: Some(now + Duration::seconds(30)),
|
||||
status: AccountStatus::QuotaExhausted,
|
||||
})
|
||||
.await
|
||||
.expect("account should be inserted");
|
||||
let observed_statuses = Arc::new(Mutex::new(Vec::new()));
|
||||
let refresher = StoreObservingTokenRefresher {
|
||||
store: store.clone(),
|
||||
account_id: "acct-refresh-quota-exhausted".to_string(),
|
||||
observed_statuses: observed_statuses.clone(),
|
||||
response: Ok(TokenPair {
|
||||
access_token: new_access_token.clone(),
|
||||
refresh_token: None,
|
||||
}),
|
||||
};
|
||||
let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new(
|
||||
store.clone(),
|
||||
RefreshPolicy {
|
||||
refresh_margin_seconds: 300,
|
||||
refresh_concurrency: 1,
|
||||
},
|
||||
refresher,
|
||||
);
|
||||
|
||||
let summary = task
|
||||
.refresh_due_accounts_once_at(now)
|
||||
.await
|
||||
.expect("quota-exhausted account should refresh token");
|
||||
let stored = store
|
||||
.get("acct-refresh-quota-exhausted")
|
||||
.await
|
||||
.expect("account should load")
|
||||
.expect("account should exist");
|
||||
let observed_statuses = observed_statuses.lock().await.clone();
|
||||
|
||||
assert_eq!(summary.refreshed, 1);
|
||||
assert_eq!(observed_statuses, [AccountStatus::QuotaExhausted]);
|
||||
assert_eq!(stored.status, AccountStatus::QuotaExhausted);
|
||||
assert_eq!(
|
||||
stored.access_token.expose_secret(),
|
||||
new_access_token.as_str()
|
||||
);
|
||||
assert_eq!(stored.access_token_expires_at, Some(new_expires_at));
|
||||
assert!(stored.next_refresh_at.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn token_refresh_task_should_not_clear_quota_limit_when_expired_token_refreshes() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let db = dir.path().join("token-refresh-expired-quota.sqlite");
|
||||
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
|
||||
.await
|
||||
.expect("sqlite pool");
|
||||
let store = SqliteAccountStore::new(pool.clone());
|
||||
let now = Utc.with_ymd_and_hms(2026, 6, 19, 11, 30, 0).unwrap();
|
||||
let cooldown_until = now + Duration::hours(1);
|
||||
let new_expires_at = now + Duration::hours(2);
|
||||
let new_access_token = test_jwt(new_expires_at.timestamp());
|
||||
store
|
||||
.insert(NewAccount {
|
||||
id: "acct-expired-quota".to_string(),
|
||||
email: Some("expired-quota@example.com".to_string()),
|
||||
account_id: Some("chatgpt-expired-quota".to_string()),
|
||||
user_id: Some("user-expired-quota".to_string()),
|
||||
label: None,
|
||||
plan_type: Some("plus".to_string()),
|
||||
access_token: SecretString::new(
|
||||
test_jwt((now - Duration::seconds(30)).timestamp()).into(),
|
||||
),
|
||||
refresh_token: Some(SecretString::new(
|
||||
"refresh-expired-quota".to_string().into(),
|
||||
)),
|
||||
added_at: None,
|
||||
access_token_expires_at: Some(now - Duration::seconds(30)),
|
||||
status: AccountStatus::Expired,
|
||||
})
|
||||
.await
|
||||
.expect("account should be inserted");
|
||||
sqlx::query(
|
||||
"update accounts set quota_limit_reached = 1, quota_cooldown_until = ? where id = ?",
|
||||
)
|
||||
.bind(cooldown_until.to_rfc3339())
|
||||
.bind("acct-expired-quota")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let refresher = StaticTokenRefresher {
|
||||
response: Arc::new(Mutex::new(Ok(TokenPair {
|
||||
access_token: new_access_token.clone(),
|
||||
refresh_token: None,
|
||||
}))),
|
||||
};
|
||||
let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new(
|
||||
store.clone(),
|
||||
RefreshPolicy {
|
||||
refresh_margin_seconds: 300,
|
||||
refresh_concurrency: 1,
|
||||
},
|
||||
refresher,
|
||||
);
|
||||
|
||||
let summary = task
|
||||
.refresh_due_accounts_once_at(now)
|
||||
.await
|
||||
.expect("expired quota-limited account should refresh token");
|
||||
let stored = store
|
||||
.get("acct-expired-quota")
|
||||
.await
|
||||
.expect("account should load")
|
||||
.expect("account should exist");
|
||||
let quota_state: (i64, Option<String>) = sqlx::query_as(
|
||||
"select quota_limit_reached, quota_cooldown_until from accounts where id = ?",
|
||||
)
|
||||
.bind("acct-expired-quota")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(summary.refreshed, 1);
|
||||
assert_eq!(stored.status, AccountStatus::QuotaExhausted);
|
||||
assert_eq!(
|
||||
stored.access_token.expose_secret(),
|
||||
new_access_token.as_str()
|
||||
);
|
||||
assert_eq!(quota_state.0, 1);
|
||||
assert_eq!(quota_state.1, Some(cooldown_until.to_rfc3339()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn token_refresh_task_should_persist_success_after_transient_transport_retry() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
@@ -139,7 +293,7 @@ async fn token_refresh_task_should_persist_success_after_transient_transport_ret
|
||||
|
||||
assert_eq!(summary.refreshed, 1);
|
||||
assert_eq!(summary.failed, 0);
|
||||
assert_eq!(observed_statuses, [AccountStatus::Refreshing; 2]);
|
||||
assert_eq!(observed_statuses, [AccountStatus::Active; 2]);
|
||||
assert_eq!(stored.status, AccountStatus::Active);
|
||||
assert_eq!(
|
||||
stored.access_token.expose_secret(),
|
||||
|
||||
@@ -2,5 +2,6 @@ pub(crate) mod accounts;
|
||||
pub(crate) mod fingerprint;
|
||||
pub(crate) mod models;
|
||||
pub(crate) mod protocol;
|
||||
pub(crate) mod scheduler;
|
||||
pub(crate) mod token_client;
|
||||
pub(crate) mod transport;
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
use codex_proxy_rs::upstream::scheduler::{FeedbackSample, FeedbackStats};
|
||||
|
||||
#[test]
|
||||
fn sample_of_unknown_account_is_empty() {
|
||||
let stats = FeedbackStats::new();
|
||||
assert_eq!(stats.sample("missing"), FeedbackSample::default());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn first_sample_lands_directly() {
|
||||
let stats = FeedbackStats::new();
|
||||
stats.report("acct", false, Some(400));
|
||||
let sample = stats.sample("acct");
|
||||
assert_eq!(sample.error_rate, Some(1.0));
|
||||
assert_eq!(sample.ttft_ms, Some(400.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn subsequent_samples_blend_by_alpha() {
|
||||
let stats = FeedbackStats::new();
|
||||
stats.report("acct", false, Some(1000)); // error_rate=1.0, ttft=1000
|
||||
stats.report("acct", true, Some(0)); // error_rate=0.8, ttft=800
|
||||
let sample = stats.sample("acct");
|
||||
assert!((sample.error_rate.unwrap() - 0.8).abs() < 1e-9);
|
||||
assert!((sample.ttft_ms.unwrap() - 800.0).abs() < 1e-9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn success_only_reports_decay_error_rate() {
|
||||
let stats = FeedbackStats::new();
|
||||
stats.report("acct", false, None); // error_rate=1.0, ttft untouched
|
||||
stats.report("acct", true, None); // error_rate=0.8
|
||||
let sample = stats.sample("acct");
|
||||
assert!((sample.error_rate.unwrap() - 0.8).abs() < 1e-9);
|
||||
assert_eq!(sample.ttft_ms, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_clears_account() {
|
||||
let stats = FeedbackStats::new();
|
||||
stats.report("acct", false, Some(100));
|
||||
stats.remove("acct");
|
||||
assert_eq!(stats.sample("acct"), FeedbackSample::default());
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub(crate) mod feedback;
|
||||
pub(crate) mod score;
|
||||
@@ -0,0 +1,144 @@
|
||||
use chrono::Utc;
|
||||
use codex_proxy_rs::upstream::accounts::model::AccountStatus;
|
||||
use codex_proxy_rs::upstream::scheduler::{rank_candidates, FeedbackStats, ScoreWeights};
|
||||
|
||||
use crate::support::accounts::test_account;
|
||||
|
||||
#[test]
|
||||
fn ranks_idle_account_above_loaded_one() {
|
||||
let mut loaded = test_account("loaded", AccountStatus::Active);
|
||||
loaded.window_input_tokens = 1_000_000;
|
||||
let idle = test_account("idle", AccountStatus::Active);
|
||||
let candidates = vec![loaded, idle];
|
||||
let feedback = FeedbackStats::new();
|
||||
let slot_count = |_: &str| 0;
|
||||
|
||||
let ranked = rank_candidates(
|
||||
&candidates,
|
||||
&ScoreWeights::default(),
|
||||
Utc::now().timestamp(),
|
||||
&slot_count,
|
||||
&feedback,
|
||||
);
|
||||
// idle 账号(索引 1)窗口用量为 0,应排在负载重的账号之前。
|
||||
assert_eq!(ranked.first().unwrap().0, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn high_error_rate_account_is_penalized() {
|
||||
let healthy = test_account("healthy", AccountStatus::Active);
|
||||
let flaky = test_account("flaky", AccountStatus::Active);
|
||||
let candidates = vec![healthy, flaky];
|
||||
let feedback = FeedbackStats::new();
|
||||
// flaky 连续失败,错误率 EWMA 拉高。
|
||||
for _ in 0..5 {
|
||||
feedback.report("flaky", false, None);
|
||||
}
|
||||
let slot_count = |_: &str| 0;
|
||||
|
||||
let ranked = rank_candidates(
|
||||
&candidates,
|
||||
&ScoreWeights::default(),
|
||||
Utc::now().timestamp(),
|
||||
&slot_count,
|
||||
&feedback,
|
||||
);
|
||||
// healthy(索引 0)应排在 flaky 之前。
|
||||
assert_eq!(ranked.first().unwrap().0, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quota_limited_account_ranks_last_despite_being_idle() {
|
||||
let mut limited = test_account("limited", AccountStatus::Active);
|
||||
limited.quota_limit_reached = true;
|
||||
let mut busy = test_account("busy", AccountStatus::Active);
|
||||
busy.window_input_tokens = 1_000_000;
|
||||
busy.window_request_count = 100;
|
||||
let candidates = vec![limited, busy];
|
||||
let feedback = FeedbackStats::new();
|
||||
let slot_count = |_: &str| 0;
|
||||
|
||||
let ranked = rank_candidates(
|
||||
&candidates,
|
||||
&ScoreWeights::default(),
|
||||
Utc::now().timestamp(),
|
||||
&slot_count,
|
||||
&feedback,
|
||||
);
|
||||
// 即使 limited 完全空闲,配额封顶惩罚也应让繁忙但未封顶的 busy(索引 1)排前。
|
||||
assert_eq!(ranked.first().unwrap().0, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn actual_token_load_dominates_cached() {
|
||||
// cache_heavy 实际 token 更少(但缓存高),heavier_actual 实际 token 更多。
|
||||
// 实际用量权重远大于缓存,cache_heavy(索引 0)应排前。
|
||||
let mut cache_heavy = test_account("cache_heavy", AccountStatus::Active);
|
||||
cache_heavy.window_input_tokens = 10_000;
|
||||
cache_heavy.window_cached_tokens = 100_000;
|
||||
let mut heavier_actual = test_account("heavier_actual", AccountStatus::Active);
|
||||
heavier_actual.window_input_tokens = 20_000;
|
||||
let candidates = vec![cache_heavy, heavier_actual];
|
||||
let feedback = FeedbackStats::new();
|
||||
let slot_count = |_: &str| 0;
|
||||
|
||||
let ranked = rank_candidates(
|
||||
&candidates,
|
||||
&ScoreWeights::default(),
|
||||
Utc::now().timestamp(),
|
||||
&slot_count,
|
||||
&feedback,
|
||||
);
|
||||
assert_eq!(ranked.first().unwrap().0, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cached_breaks_ties_when_actual_load_matches() {
|
||||
// 实际 token 相同,缓存少者(cache_light,索引 1)应排前。
|
||||
let mut cache_heavy = test_account("cache_heavy", AccountStatus::Active);
|
||||
cache_heavy.window_input_tokens = 10_000;
|
||||
cache_heavy.window_cached_tokens = 100_000;
|
||||
let mut cache_light = test_account("cache_light", AccountStatus::Active);
|
||||
cache_light.window_input_tokens = 10_000;
|
||||
cache_light.window_cached_tokens = 1_000;
|
||||
let candidates = vec![cache_heavy, cache_light];
|
||||
let feedback = FeedbackStats::new();
|
||||
let slot_count = |_: &str| 0;
|
||||
|
||||
let ranked = rank_candidates(
|
||||
&candidates,
|
||||
&ScoreWeights::default(),
|
||||
Utc::now().timestamp(),
|
||||
&slot_count,
|
||||
&feedback,
|
||||
);
|
||||
assert_eq!(ranked.first().unwrap().0, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_weights_produce_flat_scores() {
|
||||
let candidates = vec![
|
||||
test_account("a", AccountStatus::Active),
|
||||
test_account("b", AccountStatus::Active),
|
||||
];
|
||||
let feedback = FeedbackStats::new();
|
||||
let weights = ScoreWeights {
|
||||
load: 0.0,
|
||||
window: 0.0,
|
||||
cached: 0.0,
|
||||
window_requests: 0.0,
|
||||
error_rate: 0.0,
|
||||
ttft: 0.0,
|
||||
reset: 0.0,
|
||||
};
|
||||
let slot_count = |_: &str| 0;
|
||||
|
||||
let ranked = rank_candidates(
|
||||
&candidates,
|
||||
&weights,
|
||||
Utc::now().timestamp(),
|
||||
&slot_count,
|
||||
&feedback,
|
||||
);
|
||||
assert!(ranked.iter().all(|(_, b)| b.total == 0.0));
|
||||
}
|
||||
@@ -4,6 +4,29 @@ use codex_proxy_rs::upstream::{
|
||||
transport::websocket::write_websocket_audit_artifact_for_dir,
|
||||
};
|
||||
|
||||
fn rate_limit_event(used_percent: u64) -> String {
|
||||
json!({
|
||||
"type": "codex.rate_limits",
|
||||
"rate_limits": {
|
||||
"primary": {
|
||||
"used_percent": used_percent,
|
||||
"window_minutes": 43200,
|
||||
"reset_at": 1893456000 + used_percent as i64,
|
||||
}
|
||||
}
|
||||
})
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn primary_used_percent_values(headers: &[(String, String)]) -> Vec<&str> {
|
||||
headers
|
||||
.iter()
|
||||
.filter_map(|(name, value)| {
|
||||
(name == "x-codex-primary-used-percent").then_some(value.as_str())
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_audit_artifact_should_require_explicit_directory() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
@@ -663,6 +686,74 @@ async fn websocket_execute_response_create_request_should_capture_internal_metad
|
||||
.any(|(name, value)| name == "x-codex-primary-reset-at" && value == "1893456300"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn codex_backend_client_should_not_reuse_websocket_rate_limit_headers() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let server = tokio::spawn(async move {
|
||||
let (stream, _) = listener.accept().await.unwrap();
|
||||
let mut websocket = accept_codex_test_websocket(stream).await;
|
||||
|
||||
let _first_message = websocket.next().await.unwrap().unwrap();
|
||||
websocket
|
||||
.send(Message::Text(rate_limit_event(10).into()))
|
||||
.await
|
||||
.unwrap();
|
||||
websocket
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_rate_limit_first", 3, 1).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let _second_message = websocket.next().await.unwrap().unwrap();
|
||||
websocket
|
||||
.send(Message::Text(rate_limit_event(20).into()))
|
||||
.await
|
||||
.unwrap();
|
||||
websocket
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_rate_limit_second", 5, 2).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
websocket.close(None).await.unwrap();
|
||||
});
|
||||
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(pool);
|
||||
let request = pooled_websocket_request("conversation-rate-limit-reuse");
|
||||
|
||||
let first = backend
|
||||
.create_response(
|
||||
&request,
|
||||
request_context("req_rate_limit_first", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("first pooled websocket response should succeed");
|
||||
let second = backend
|
||||
.create_response(
|
||||
&request,
|
||||
request_context("req_rate_limit_second", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("second pooled websocket response should reuse connection");
|
||||
server.await.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
primary_used_percent_values(&first.rate_limit_headers),
|
||||
vec!["10"]
|
||||
);
|
||||
assert_eq!(
|
||||
primary_used_percent_values(&second.rate_limit_headers),
|
||||
vec!["20"]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_execute_response_create_request_should_surface_incomplete_response() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
|
||||
+1
-3
@@ -19,9 +19,7 @@ cp deploy/config.example.yaml .runtime/config.yaml
|
||||
|
||||
复制后必须设置 `admin.default_password`。空值、长度不足 12 位或常见弱口令会被后端拒绝。
|
||||
|
||||
如果前面还有 Nginx、Caddy、Cloudflare Tunnel 等反向代理,把后端实际看到的代理 peer IP 或 CIDR 写入 `server.trusted_proxies`。经宿主机反代访问 Docker 端口时,这个 peer 通常是 Compose 网络网关,例如 `172.18.0.1/32`,不是宿主 `127.0.0.1/32`。
|
||||
|
||||
`server.trusted_proxies` 非空时,只有这些 peer 的 `CF-Connecting-IP` / `X-Real-IP` / `X-Forwarded-For` 会被采信;为空时进入自动解析模式,按同一组头解析真实 IP。
|
||||
如果前面还有 Nginx、Caddy、Cloudflare Tunnel 等反向代理,确保反代正常传递 `CF-Connecting-IP`、`X-Real-IP` 或 `X-Forwarded-For`。后端会按这组头自动解析真实客户端 IP;没有这些头时回落到直连 peer IP。
|
||||
|
||||
构建并启动:
|
||||
|
||||
|
||||
@@ -3,11 +3,6 @@ server:
|
||||
host: '0.0.0.0'
|
||||
# 容器内服务端口,需与 docker-compose.yml 的 target port 一致。
|
||||
port: 8080
|
||||
# 可信反向代理 IP 或 CIDR。非空时只有这些 peer 的转发 IP 头会被采信。
|
||||
# Docker 端口映射经宿主反代访问时,后端看到的 peer 通常是 Compose 网络网关
|
||||
# (例如 172.18.0.1),不是宿主 127.0.0.1。
|
||||
# 为空时进入自动解析模式,按 CF-Connecting-IP / X-Real-IP / X-Forwarded-For 解析。
|
||||
trusted_proxies: []
|
||||
|
||||
api:
|
||||
# ChatGPT / Codex 上游 API 地址。
|
||||
|
||||
+426
-392
@@ -1,440 +1,474 @@
|
||||
# Codex Proxy RS 架构
|
||||
# Architecture —— 后端架构权威文档
|
||||
|
||||
## 总览
|
||||
本文档是 Codex Proxy RS 的**唯一**架构权威文档:分层模型、依赖规则、目录与文件命名规范、模块约束、现状审计结论、旧→新完整映射与分阶段执行方案。
|
||||
|
||||
Codex Proxy RS 是一个单进程代理和管理系统。Rust 后端同时承载 OpenAI Responses 兼容代理、管理端 API、SQLite 持久化、后台任务、静态前端托管和在线更新;Vue 管理端作为 SPA 由同一个进程托管。
|
||||
- 数据库 schema、口径与迁移见 [database.md](database.md);本文不重复。
|
||||
- 基线:2026-07-08,`backend/src` 共 160 个 `.rs` 文件、约 41k 行。
|
||||
- 本文的目标结构与 database.md 的 v4 改造**共用同一个改造窗口**:遥测域(usage/ops/buckets)的文件本来就要因 schema v4 重写,搬家与重写一次完成,不做两遍。
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
Client["客户端 / OpenAI SDK"] -->|"Authorization: Bearer<br/>/v1/responses"| Http["Axum Router"]
|
||||
Browser["浏览器"] -->|"管理端 SPA<br/>/api/admin/*"| Http
|
||||
---
|
||||
|
||||
subgraph App["codex-proxy-rs"]
|
||||
Http --> Proxy["proxy<br/>OpenAI compatible routes"]
|
||||
Http --> Admin["admin<br/>accounts / keys / monitoring / system"]
|
||||
Http --> Web["web<br/>static SPA"]
|
||||
## 1. 系统总览
|
||||
|
||||
Proxy --> Dispatch["dispatch<br/>account selection / affinity / retry"]
|
||||
Dispatch --> Protocol["upstream protocol<br/>request / SSE / error mapping"]
|
||||
Protocol --> Transport["upstream transport<br/>HTTP / SSE / WebSocket"]
|
||||
|
||||
Admin --> Services["runtime services"]
|
||||
Dispatch --> Services
|
||||
Services --> Store[("SQLite")]
|
||||
end
|
||||
|
||||
Transport --> Upstream["ChatGPT / Codex backend"]
|
||||
Services --> Tasks["background tasks"]
|
||||
Tasks --> Store
|
||||
Updater["self update"] --> Store
|
||||
```
|
||||
|
||||
核心原则:
|
||||
|
||||
- 启动配置只负责进程启动必需项。
|
||||
- 账号、Key、模型别名和运行参数保存在 SQLite,并通过管理端热更新。
|
||||
- 代理热路径只依赖运行时内存服务,持久化写入在请求完成、后台任务或设置更新时同步。
|
||||
- 上游请求统一经过账号池、session affinity、quota、fingerprint 和 transport 层。
|
||||
|
||||
## 仓库边界
|
||||
|
||||
- `backend/`:Rust/Axum 后端、SQLite migration、构建脚本和集成测试。
|
||||
- `frontend/`:Vue 3 管理端,使用 Vite、Pinia、Vue Router、ECharts 和 lucide 图标。
|
||||
- `deploy/`:Dockerfile、Compose 和部署配置模板。
|
||||
- `docs/`:长期维护文档。
|
||||
- `release/`:版本号、平台矩阵和发布脚本。
|
||||
- `skills/`:项目本地 Codex skill。
|
||||
|
||||
测试统一放在 `backend/tests/`,不在 `backend/src/` 新增测试模块。
|
||||
|
||||
## 后端结构
|
||||
|
||||
入口和组合层:
|
||||
|
||||
- `backend/src/main.rs`:二进制入口。
|
||||
- `backend/src/lib.rs`:二进制和集成测试共享入口。
|
||||
- `runtime/bootstrap.rs`:加载配置、初始化日志、连接 SQLite、构造服务、恢复运行时状态并启动 HTTP。
|
||||
- `runtime/services.rs`:集中构造账号、模型、监控、更新、fingerprint、WebSocket 池等服务。
|
||||
- `runtime/tasks/`:后台任务协调和周期任务。
|
||||
|
||||
主要业务模块:
|
||||
|
||||
- `admin/`:管理端 API,包括登录 session、账号、API Key、监控、设置和系统更新。
|
||||
- `proxy/`:OpenAI 兼容路由、客户端 API Key 鉴权和请求分派。
|
||||
- `proxy/openai/`:Responses、review、compact、Models 的 HTTP 入口。
|
||||
- `proxy/dispatch/`:账号选择、session affinity、quota 验证、历史恢复、错误映射、用量记录。
|
||||
- `upstream/accounts/`:账号模型、账号池、导入、Cookie、quota、token refresh、SQLite store。
|
||||
- `upstream/models/`:模型目录、模型别名、模型计划快照。
|
||||
- `upstream/protocol/`:Responses 请求、SSE、WebSocket 事件和错误转换。
|
||||
- `upstream/transport/`:Reqwest、TLS、WebSocket、WebSocket pool、fingerprint 请求头。
|
||||
- `config/`:启动配置 schema、运行时设置和数据库设置加载。
|
||||
- `infra/`:SQLite、migration、日志、时间、JSON、格式化和身份哈希。
|
||||
- `http/`:router、中间件、可信代理和请求上下文。
|
||||
- `web/`:前端静态资源和 SPA fallback。
|
||||
|
||||
## 配置
|
||||
|
||||
启动配置由 `CPR_CONFIG_FILE` 指定;未指定时读取当前工作目录的 `config.yaml`。
|
||||
|
||||
配置文件字段:
|
||||
|
||||
- `server`:监听地址、端口、可信代理。
|
||||
- `api`:上游 base URL。
|
||||
- `database`:SQLite URL。
|
||||
- `tls`:HTTP/TLS 偏好。
|
||||
- `ws_pool`:WebSocket 连接池和首内容帧超时。
|
||||
- `quota`:quota 刷新周期和是否跳过耗尽账号。
|
||||
- `fingerprint`:默认 Codex Desktop 指纹和请求头。
|
||||
- `admin`:管理员首次初始化和 session 清理配置。
|
||||
- `logging`:日志目录、保留期和使用记录开关。
|
||||
|
||||
运行时设置保存在 `runtime_settings`:
|
||||
|
||||
- `modelAliases`
|
||||
- `refreshMarginSeconds`
|
||||
- `refreshConcurrency`
|
||||
- `maxConcurrentPerAccount`
|
||||
- `requestIntervalMs`
|
||||
- `rotationStrategy`
|
||||
|
||||
这些设置由管理端更新,并同步到账号池、token refresh 策略和模型服务。
|
||||
|
||||
## 运行目录
|
||||
|
||||
`.runtime` 是默认运行目录:
|
||||
单进程:Rust/Axum 同时承载 OpenAI Responses 兼容代理、管理端 API、SQLite 持久化、后台任务、静态 SPA 托管与在线更新。Vue 3 管理端由同一进程托管。
|
||||
|
||||
```text
|
||||
.runtime/config.yaml
|
||||
.runtime/data/codex-proxy-rs.sqlite
|
||||
.runtime/data/installation_id
|
||||
.runtime/data/update-state.json
|
||||
.runtime/data/update.lock
|
||||
.runtime/logs/
|
||||
客户端 SDK ──/v1/*──┐ ┌── ChatGPT / Codex backend
|
||||
├─ api ─ dispatch ─ upstream
|
||||
浏览器 ──/api/admin─┘ │ │
|
||||
│ accounts / models
|
||||
│ │
|
||||
telemetry / keys / auth / settings / update
|
||||
│
|
||||
infra(SQLite / 日志 / 时间)
|
||||
```
|
||||
|
||||
Docker 默认映射:
|
||||
核心原则(不变):
|
||||
|
||||
1. 启动配置只含进程启动必需项;其余全部在 SQLite 热更新。
|
||||
2. 代理热路径只依赖内存态服务;持久化写入在请求完成、后台任务或设置更新时发生。
|
||||
3. 上游请求统一经过账号池、session affinity、quota、fingerprint 与 transport。
|
||||
|
||||
---
|
||||
|
||||
## 2. 分层模型与依赖规则(权威)
|
||||
|
||||
### 2.1 层定义
|
||||
|
||||
| 层 | 模块 | 职责 | 禁止 |
|
||||
| --- | --- | --- | --- |
|
||||
| L0 | `infra` | SQLite 连接/迁移、时间、JSON、日志、身份生成、路径、格式化 | 出现任何业务概念 |
|
||||
| L1 | `upstream` | 纯上游客户端:协议编解码、HTTP/SSE/WebSocket 传输、指纹、OAuth token client。**每家 provider 一个子目录**(`openai/`,规划中 `cloudflare/`) | 依赖任何领域模块;出现账号池/调度/遥测概念 |
|
||||
| L1 | `telemetry` | 遥测域:usage_records / ops_error_logs / request_time_buckets / account_usage 的**写入与查询**、billing、dashboard 聚合、rebuild | 被冠以 admin 语义;HTTP handler |
|
||||
| L1 | `keys` | 客户端 API key 域:store、内存缓存、鉴权校验、管理操作 | — |
|
||||
| L1 | `auth` | 管理端认证域:admin 用户、会话(token_hash) | 客户端 key 概念(那是 `keys`) |
|
||||
| L1 | `settings` | 运行时设置域:持久化、快照、变更广播(watch)、admin API key 校验 | 直接调用上层服务推送变更(§2.3) |
|
||||
| L1 | `update` | 自更新域:release 查询、下载、校验、解压、替换、状态 | — |
|
||||
| L2 | `accounts` | 账号域:实体、池、调度器、store、quota、token 刷新、Cookie、导入导出、管理操作 | — |
|
||||
| L2 | `models` | 模型目录域:catalog、别名、plan 快照 | — |
|
||||
| L3 | `dispatch` | 请求编排域:账号获取、affinity、上游调用、流转发、失败恢复、重试、遥测记录 | HTTP 类型(axum)出现在签名 |
|
||||
| L4 | `api` | 全部入站 HTTP:`api/client`(/v1)、`api/admin`(/api/admin)、路由组合、中间件、SPA 资源 | 写 SQL;承载领域逻辑 |
|
||||
| L5 | `bootstrap` | 进程装配:启动配置、服务构造、AppState、后台任务、关闭 | 被任何层依赖 |
|
||||
|
||||
### 2.2 依赖方向
|
||||
|
||||
只允许**高层 → 低层**。同层依赖默认禁止,例外必须登记在下表(新增例外 = 改本文档 + code review):
|
||||
|
||||
| 允许的同层/特殊依赖 | 理由 |
|
||||
| --- | --- |
|
||||
| `accounts` → `models` | 候选过滤需要模型允许列表;单向 |
|
||||
| `accounts` → `settings` | 池参数/刷新策略订阅设置快照(watch 类型);单向 |
|
||||
| `auth` → `settings` | admin API key 哈希存放在 runtime_settings 表,经 settings 的 store 方法读取;单向 |
|
||||
|
||||
**全库硬禁**(CI 脚本断言,§9):
|
||||
|
||||
- L1–L3 任何模块 import `crate::api` / `crate::bootstrap`。
|
||||
- `dispatch`、`accounts`、`telemetry` 等领域模块 import `axum`。
|
||||
- `api` 下出现 `sqlx`(SQL 只允许在各域的 `store` 文件里)。
|
||||
- `upstream` import `crate::{accounts, dispatch, telemetry, keys, auth, settings, update, models}`。
|
||||
- 任何模块 import 旧路径 `crate::admin` / `crate::proxy` / `crate::runtime` / `crate::web` / `crate::http` / `crate::config`(改造完成后这些模块不存在)。
|
||||
|
||||
### 2.3 控制反转两处(消除向上依赖)
|
||||
|
||||
1. **设置传播**:现状 `RuntimeSettingsService` 持有账号池与刷新策略的引用、主动推送更新(低层持有高层引用)。终态:`settings` 只暴露 `tokio::sync::watch::Receiver<SettingsSnapshot>`;`accounts`(池、刷新策略)、`models`(别名)各自订阅。装配在 `bootstrap` 完成。
|
||||
2. **遥测写入**:现状 proxy 热路径 import `admin::monitoring` 的 Admin* 服务。终态:`dispatch` 依赖 `telemetry::Recorder`(L3 → L1,方向正确),admin 侧只是同一域的查询消费者。
|
||||
|
||||
---
|
||||
|
||||
## 3. 目标目录树(终态,权威)
|
||||
|
||||
```text
|
||||
.runtime/config.yaml -> /app/config.yaml
|
||||
.runtime/data -> /app/data
|
||||
.runtime/logs -> /app/logs
|
||||
backend/src/
|
||||
├── main.rs
|
||||
├── lib.rs
|
||||
├── infra/ # L0
|
||||
│ ├── mod.rs
|
||||
│ ├── database.rs # 连接、pragma、迁移框架(含 Rust post 钩子,database.md §6.1)
|
||||
│ ├── migrations/ # *.sql
|
||||
│ ├── identity.rs time.rs json.rs format.rs logging.rs paths.rs
|
||||
├── upstream/ # L1 纯上游客户端(每家 provider 一个子目录)
|
||||
│ ├── mod.rs # 第二家 provider 出现时,UpstreamClient trait 落此
|
||||
│ └── openai/
|
||||
│ ├── mod.rs
|
||||
│ ├── token_client.rs
|
||||
│ ├── protocol/
|
||||
│ │ ├── mod.rs schema.rs responses.rs events.rs sse.rs websocket.rs
|
||||
│ ├── transport/
|
||||
│ │ ├── mod.rs client.rs endpoints.rs headers.rs tls.rs
|
||||
│ │ ├── response_meta.rs usage.rs diagnostics.rs
|
||||
│ │ ├── websocket.rs websocket_pool.rs websocket_pump.rs
|
||||
│ └── fingerprint/
|
||||
│ ├── mod.rs types.rs store.rs runtime.rs updater.rs
|
||||
├── telemetry/ # L1 遥测域(原 admin/monitoring + proxy/dispatch/usage_events)
|
||||
│ ├── mod.rs
|
||||
│ ├── recorder.rs # 唯一写入口:成功/失败两条事务(database.md §5.2)
|
||||
│ ├── billing.rs
|
||||
│ ├── rebuild.rs # rebuild-derived(database.md §7.6)
|
||||
│ ├── dashboard.rs # Dashboard 聚合查询(无 HTTP)
|
||||
│ ├── usage/ # usage_records
|
||||
│ │ ├── mod.rs types.rs store.rs query.rs
|
||||
│ ├── ops/ # ops_error_logs
|
||||
│ │ ├── mod.rs types.rs store.rs query.rs
|
||||
│ ├── buckets/ # request_time_buckets
|
||||
│ │ ├── mod.rs store.rs query.rs
|
||||
│ └── account_usage/ # account_usage / account_model_usage
|
||||
│ ├── mod.rs store.rs query.rs
|
||||
├── keys/ # L1 客户端 key 域(原 admin/keys/service.rs 拆分)
|
||||
│ ├── mod.rs types.rs store.rs cache.rs service.rs manage.rs
|
||||
├── auth/ # L1 管理端认证域(原 admin/auth/service.rs)
|
||||
│ ├── mod.rs types.rs store.rs service.rs
|
||||
├── settings/ # L1 运行时设置域(原 config/settings.rs)
|
||||
│ ├── mod.rs types.rs store.rs service.rs # service 持 watch 广播
|
||||
├── update/ # L1 自更新域(原 admin/system/*)
|
||||
│ ├── mod.rs types.rs state.rs release.rs download.rs archive.rs service.rs
|
||||
├── accounts/ # L2 账号域(原 upstream/accounts + upstream/scheduler + admin/accounts/service)
|
||||
│ ├── mod.rs
|
||||
│ ├── account.rs # 实体(原 model.rs)
|
||||
│ ├── window.rs
|
||||
│ ├── pool.rs # 池:存储态/槽位/过滤(§5.4 行数预算内需拆出 pool/ 子文件)
|
||||
│ ├── scheduler.rs # 调度门面(原 scheduler/mod.rs 的类型)
|
||||
│ ├── scheduler/
|
||||
│ │ ├── candidates.rs feedback.rs
|
||||
│ │ └── strategy/
|
||||
│ │ ├── mod.rs smart.rs quota_reset.rs round_robin.rs sticky.rs
|
||||
│ ├── store.rs # AccountStore trait + Sqlite 门面
|
||||
│ ├── store/
|
||||
│ │ ├── queries.rs rows.rs
|
||||
│ ├── quota/
|
||||
│ │ ├── mod.rs service.rs runtime.rs
|
||||
│ ├── refresh/ # 原 token_refresh
|
||||
│ │ ├── mod.rs policy.rs lease.rs service.rs
|
||||
│ ├── cookies/
|
||||
│ │ ├── mod.rs store.rs
|
||||
│ ├── import.rs # 导入 payload 解析(原 upstream/accounts/importing.rs)
|
||||
│ └── manage/ # 管理操作(原 admin/accounts/service/*)
|
||||
│ ├── mod.rs types.rs service.rs lifecycle.rs
|
||||
│ ├── import.rs export.rs oauth.rs probe.rs quota.rs quota_view.rs
|
||||
├── models/ # L2 模型目录域
|
||||
│ ├── mod.rs types.rs catalog.rs service.rs store.rs
|
||||
├── dispatch/ # L3 请求编排域(原 proxy/dispatch)
|
||||
│ ├── mod.rs
|
||||
│ ├── service.rs # 编排入口(原 responses/service.rs,按 §5.4 拆分)
|
||||
│ ├── upstream_call.rs # 原 dispatch/upstream.rs
|
||||
│ ├── errors.rs # 原 dispatch/errors.rs + responses/errors.rs 合并
|
||||
│ ├── recording.rs # 原 responses/event_recording.rs,调 telemetry::recorder
|
||||
│ ├── affinity/ # 原 session_affinity.rs(683 行三合一)+ responses/affinity.rs
|
||||
│ │ ├── mod.rs types.rs store.rs service.rs resolve.rs
|
||||
│ ├── recovery/
|
||||
│ │ ├── mod.rs auth.rs cloudflare.rs exhaustion.rs
|
||||
│ │ ├── implicit_resume.rs reasoning_replay.rs
|
||||
│ └── stream/
|
||||
│ ├── mod.rs live.rs lifecycle.rs prefetch.rs sse_failure.rs trace.rs
|
||||
├── api/ # L4 全部入站 HTTP(原 http + web + proxy 入口 + admin routes)
|
||||
│ ├── mod.rs
|
||||
│ ├── router.rs # 总路由 + healthz(经 infra 的 db ping,不穿透到具体 store 的 pool)
|
||||
│ ├── assets.rs # SPA 静态资源(原 web/assets.rs)
|
||||
│ ├── middleware/
|
||||
│ │ ├── mod.rs request_id.rs trace.rs
|
||||
│ ├── client/ # /v1(原 proxy/auth + proxy/openai)
|
||||
│ │ ├── mod.rs router.rs auth.rs
|
||||
│ │ ├── responses.rs models.rs errors.rs sse.rs
|
||||
│ └── admin/ # /api/admin(原 admin 的全部 routes 与 DTO)
|
||||
│ ├── mod.rs router.rs response.rs session.rs # session = 鉴权提取器/中间件
|
||||
│ ├── auth_routes.rs settings_routes.rs system_routes.rs
|
||||
│ ├── accounts_routes.rs keys_routes.rs
|
||||
│ ├── usage_routes.rs dashboard_routes.rs
|
||||
└── bootstrap/ # L5 进程装配(原 runtime + config/schema,loader)
|
||||
├── mod.rs
|
||||
├── config.rs # 启动配置 schema + loader(原 config/schema.rs + loader.rs)
|
||||
├── services.rs state.rs shutdown.rs
|
||||
└── tasks/
|
||||
├── mod.rs coordinator.rs periodic.rs cleanup.rs
|
||||
├── cookie_cleanup.rs session_cleanup.rs session_affinity_cleanup.rs
|
||||
├── model_refresh.rs token_refresh.rs quota_refresh.rs
|
||||
├── fingerprint_update.rs retention_trim.rs # retention_trim 新增(database.md §7.7)
|
||||
```
|
||||
|
||||
容器以非 root 用户 `10001:10001` 运行。Compose 设置 `HOME=/app` 和 `XDG_DATA_HOME=/app/data`,确保 installation id、SQLite、更新状态和日志写入持久化目录。
|
||||
`backend/tests/` 目录结构**镜像 `src/`**:模块搬家时测试目录同步搬家,`fixtures/`、`support/` 不动。
|
||||
|
||||
## SQLite
|
||||
---
|
||||
|
||||
当前 schema 由 `backend/src/infra/migrations` 管理。
|
||||
## 4. 命名规范
|
||||
|
||||
核心表:
|
||||
### 4.1 目录名
|
||||
|
||||
- `admin_users`、`admin_sessions`:管理端账号和 session。
|
||||
- `client_api_keys`:客户端 API Key。
|
||||
- `runtime_settings`:管理端运行设置。
|
||||
- `accounts`:账号身份、token、状态、quota、fingerprint 关联信息。
|
||||
- `account_usage`:账号累计用量和当前 quota 窗口用量。
|
||||
- `account_model_usage`:账号和模型维度的请求、错误和 token 用量。
|
||||
- `account_refresh_leases`:token refresh 分布式租约。
|
||||
- `account_cookies`:账号 Cookie。
|
||||
- `fingerprints`、`fingerprint_update_history`:当前指纹和更新历史。
|
||||
- `usage_records`:网关和上游事件明细。
|
||||
- `usage_time_buckets`:15 分钟聚合桶,供 Dashboard 和趋势接口使用。
|
||||
- `model_plan_snapshots`:模型到账号计划的可用性快照。
|
||||
- `session_affinities`:response ID、conversation、账号、turn state 和 function call 关联。
|
||||
1. snake_case,全小写,禁止连字符与点。
|
||||
2. **实体集合域用复数**(`accounts`、`models`、`keys`、`buckets`);**能力/抽象域用单数**(`dispatch`、`telemetry`、`auth`、`settings`、`update`、`upstream`、`api`、`infra`、`bootstrap`)。
|
||||
3. 子目录表达子能力(`scheduler/strategy/`),不表达文件类型——禁止 `services/`、`stores/`、`utils/`、`helpers/`、`common/`、`misc/` 这类"按技术角色装桶"的目录。
|
||||
|
||||
## 代理入口
|
||||
### 4.2 文件名
|
||||
|
||||
客户端请求入口:
|
||||
1. snake_case。**主类型 = 文件名的 PascalCase**(`account.rs` ↔ `Account`;`scheduler.rs` ↔ `AccountScheduler` 属可接受的域前缀加长)。
|
||||
2. **角色后缀词表**(一个文件一个角色):
|
||||
|
||||
- `POST /v1/responses`
|
||||
- `POST /v1/responses/review`
|
||||
- `POST /v1/responses/compact`
|
||||
- `GET /v1/models`
|
||||
- `GET /v1/models/{model}`
|
||||
| 文件名 | 角色 | 约束 |
|
||||
| --- | --- | --- |
|
||||
| `router.rs` | 路由组合 | 只 merge/nest,不写 handler |
|
||||
| `*_routes.rs` / `routes.rs` | HTTP handler + 请求/响应 DTO | 不写 SQL、不含领域规则 |
|
||||
| `service.rs` | 领域服务(该域对外门面) | 不含 SQL 文本 |
|
||||
| `store.rs` | SQLite 持久化 | SQL 只允许出现在这里;类型名 `Sqlite*Store` |
|
||||
| `cache.rs` | 内存缓存/运行时快照 | 类型名 `Runtime*` |
|
||||
| `query.rs` | 复杂只读查询(列表/聚合/分页) | 与写路径分文件 |
|
||||
| `types.rs` | 该域跨文件共享的实体/值类型 | 无 IO |
|
||||
| `rows.rs` | DB 行结构(store 私有) | `pub(crate)` 以内 |
|
||||
| `errors.rs` | 该域错误类型 | — |
|
||||
| `mod.rs` | 声明 + re-export | **≤ 50 行,零逻辑** |
|
||||
|
||||
请求链路:
|
||||
3. **禁止动名词文件名**:`importing.rs` → `import.rs`、`exporting.rs` → `export.rs`、`testing.rs` → `probe.rs`("testing" 会被误读为测试代码,永久禁用)。
|
||||
4. 禁止 `util.rs` / `helper.rs` / `common.rs` / `misc.rs`——放不进角色词表的代码说明职责没想清楚。
|
||||
5. 缩写白名单:`sse`、`tls`、`ws`(仅类型名内)、`db`、`id`、`json`、`oauth`。其余单词写全。
|
||||
6. 一个概念一个词,全库同义同名:持久化一律 **Store**(废除 Repository——现 `FingerprintRepository` 改 `FingerprintStore`);行类型一律 rows;DTO 一律在 routes/types,废除 `*_model.rs` 与 `contracts.rs` 两种叫法。
|
||||
|
||||
1. HTTP 中间件生成 request id、记录访问日志、解析可信代理 IP。
|
||||
2. `proxy/auth` 校验客户端 API Key。
|
||||
3. `proxy/openai` 解析 OpenAI compatible 请求。
|
||||
4. `proxy/dispatch` 构造账号获取请求,处理 session affinity 和隐式续写。
|
||||
5. 账号池按模型、账号状态、quota、Cloudflare、并发槽和策略选择账号。
|
||||
6. `upstream/protocol` 转为 Codex backend 请求。
|
||||
7. `upstream/transport` 选择 HTTP、SSE 或 WebSocket 访问上游。
|
||||
8. 响应被转换回 OpenAI Responses/Models 格式。
|
||||
9. 请求记录、账号用量、模型用量、quota header、Cookie、turn state 和 session affinity 写回运行时和 SQLite。
|
||||
### 4.3 类型名
|
||||
|
||||
## 账号调度
|
||||
| 前缀/形态 | 保留给 | 示例 |
|
||||
| --- | --- | --- |
|
||||
| `Sqlite*Store` | store 实现 | `SqliteAccountStore` |
|
||||
| `Runtime*` | 内存缓存/快照服务 | `RuntimeClientKeyCache` |
|
||||
| `Admin*` | **仅** `api/admin` 下的 DTO/提取器 | `AdminAccountPayload` |
|
||||
| 无前缀 | 领域服务与实体 | `Recorder`、`AccountScheduler`、`SettingsService` |
|
||||
|
||||
账号池位于 `upstream/accounts/pool.rs`。调度前先做硬过滤:
|
||||
由此更名(现状 → 终态):`AdminUsageRecordService` → `telemetry::usage::UsageQueryService`(读)+ `telemetry::Recorder`(写);`AdminOpsErrorLogService` → `telemetry::ops` 同拆;`AdminUsageService`(实际管 account_usage)→ `telemetry::account_usage::AccountUsageQueryService`;`AdminSessionService` → `auth::SessionService`;`AdminClientKeyService` → `keys::KeyManageService`;`ClientKeyService` → `keys::KeyVerifier`。**类型名与文件名对不上的历史包袱全部在搬家时清算,不留别名 re-export。**
|
||||
|
||||
- 账号状态必须是 `active`。
|
||||
- 模型必须被账号计划允许。
|
||||
- quota 必须可用。
|
||||
- Cloudflare 冷却必须结束。
|
||||
- 账号不在本次请求排除列表。
|
||||
- 当前在途槽位小于 `maxConcurrentPerAccount`。
|
||||
禁止无意义别名导入:`use std::sync::Arc as StdArc` 这类一律 `use std::sync::Arc`。
|
||||
|
||||
策略:
|
||||
---
|
||||
|
||||
- `smart`:默认策略。按账号实时负载选择账号,依次比较当前在途槽位数、窗口请求数、窗口 token 数和历史请求数;负载相同,再参考额度窗口重置时间和最近使用时间。
|
||||
- `quota_reset_priority`:按额度窗口重置时间选择账号,优先使用更早重置的窗口;重置时间相同,再比较历史请求数和最近使用时间。
|
||||
- `round_robin`:按候选账号顺序循环选择。
|
||||
- `sticky`:按最近使用时间优先选择。
|
||||
## 5. 模块约束
|
||||
|
||||
请求带有 `required_account_id` 时只尝试该账号。请求带有 session affinity 的 `preferred_account_id` 时,只有该账号仍在硬过滤结果中才会命中,否则回到当前调度策略。选中账号后立即占用在途槽位,释放时记录请求用量并更新窗口计数。
|
||||
### 5.1 mod.rs 纪律
|
||||
|
||||
## Session Affinity
|
||||
`mod.rs` 只做 `pub mod` 声明与 `pub use` 门面导出,≤ 50 行。现状违规(`accounts/store/mod.rs` 1126 行、`token_refresh/mod.rs` 495 行、`quota/mod.rs` 419 行、`updater/mod.rs` 811 行、`scheduler/mod.rs` 含门面类型)全部在搬家时把实现移入具名文件。
|
||||
|
||||
Session affinity 由 `proxy/dispatch/session_affinity.rs` 和 `proxy/dispatch/responses/affinity.rs` 维护。
|
||||
### 5.2 SQL 边界
|
||||
|
||||
存储内容:
|
||||
SQL 文本只允许出现在 `store.rs` / `store/` / `query.rs`。`api` 层与 `service.rs` 出现 `sqlx::query` 即为违规(现状 `http/router.rs` 的 healthz 直接拿某个 store 的连接池执行 SQL——终态 `infra::database` 提供 `ping(pool)`,router 调它)。
|
||||
|
||||
- response id。
|
||||
- conversation id。
|
||||
- account id。
|
||||
- prompt/instructions 变体。
|
||||
- turn state。
|
||||
- function call ids。
|
||||
- 输入 token。
|
||||
- 过期时间。
|
||||
### 5.3 HTTP 边界
|
||||
|
||||
行为:
|
||||
axum/tower 类型只允许出现在 `api` 与 `bootstrap`。领域函数签名出现 `HeaderMap`、`StatusCode` 等 HTTP 类型即为违规(用域内枚举/结构表达,api 层转换)。
|
||||
|
||||
- 请求显式携带 `previous_response_id` 时,优先查找该 response 对应账号。
|
||||
- 请求没有 `previous_response_id` 但满足隐式续写条件时,按 conversation 查找最新 response 并补充上下文。
|
||||
- 当亲和账号不可用并切换到其它账号时,按请求类型和失败类型处理历史字段。
|
||||
- 上游返回 `previous_response_not_found`、未闭合 function call 等历史错误时,恢复请求历史并重试。
|
||||
- 成功响应会记录新的 response affinity,供后续请求继续使用。
|
||||
### 5.4 文件行数预算
|
||||
|
||||
## Quota 和账号状态
|
||||
- 软上限 **400 行**(超过即在 review 中说明理由);硬上限 **800 行**(超过必须拆,CI 警告,§9)。
|
||||
- 现状超硬线清单(11 个,全部有既定拆法):
|
||||
|
||||
账号状态包括:
|
||||
|
||||
- `active`
|
||||
- `expired`
|
||||
- `quota_exhausted`
|
||||
- `refreshing`
|
||||
- `disabled`
|
||||
- `banned`
|
||||
|
||||
quota 数据来源:
|
||||
|
||||
- 管理端主动刷新账号 quota。
|
||||
- 后台 quota refresh 周期任务。
|
||||
- 上游 rate-limit header 被动同步。
|
||||
- 请求中遇到 429/限额错误时标记冷却。
|
||||
|
||||
quota 运行态字段包括 `quota_json`、`quota_fetched_at`、`quota_limit_reached`、`quota_verify_required`、`quota_cooldown_until`、当前窗口计数、窗口开始时间、窗口重置时间和窗口长度。
|
||||
|
||||
当 quota 冷却或窗口到期时,账号会恢复为可验证状态,并在下次使用前执行 quota verify。验证通过后继续使用;验证仍为耗尽时继续冷却或切换账号。
|
||||
|
||||
## Token Refresh
|
||||
|
||||
token refresh 由 `RuntimeTokenRefreshService` 和 `runtime/tasks/token_refresh.rs` 执行。
|
||||
|
||||
行为:
|
||||
|
||||
- 按 `refreshMarginSeconds` 判断 access token 是否需要提前刷新。
|
||||
- 用 `refreshConcurrency` 控制刷新并发。
|
||||
- 使用 `account_refresh_leases` 避免多任务重复刷新同一账号。
|
||||
- refresh token 缺失、失效、复用或账号被封禁时更新账号状态。
|
||||
- 设置更新后刷新策略热生效。
|
||||
|
||||
## WebSocket Transport
|
||||
|
||||
Responses 请求可通过 WebSocket 访问上游 `/codex/responses`。WebSocket 层负责:
|
||||
|
||||
- 构造 opening handshake 和首个 `response.create` 文本帧。
|
||||
- 将上游 WebSocket 事件转换为 SSE。
|
||||
- 聚合非流式响应或转发流式响应。
|
||||
- 解析上游错误帧、`response.incomplete`、rate-limit 事件和 metadata turn state。
|
||||
- 记录首内容帧延迟、WebSocket pool 决策和上游诊断信息。
|
||||
|
||||
WebSocket pool 配置:
|
||||
|
||||
- `enabled`:是否启用连接池。
|
||||
- `max_age_ms`:连接最大生命周期。
|
||||
- `max_per_account`:单账号池 slot 上限。
|
||||
- `first_token_timeout_ms`:发出 `response.create` 后等待首个内容帧的绝对超时;`0` 表示禁用。
|
||||
|
||||
池行为:
|
||||
|
||||
- pool key 由 base URL、账号 ID 和 conversation ID 组成。
|
||||
- idle 连接按 key 复用,busy 连接会让请求绕过池新建连接。
|
||||
- idle 连接由后台 pump 做 ping/pong 和关闭检测。
|
||||
- 连接超过 max age、被 pump 标记关闭、账号被驱逐或请求失败时丢弃。
|
||||
- 复用连接在首内容帧前失效时,丢弃连接并换新连接。
|
||||
- fresh 连接首内容帧超时时,丢弃当前连接并按配置重试。
|
||||
- 流式请求在返回下游前会预取到首个内容帧;这样首内容帧超时可以在响应开始前完成换连接。
|
||||
|
||||
## Fingerprint
|
||||
|
||||
fingerprint 位于 `upstream/fingerprint.rs`。运行时持有一个 `RuntimeFingerprint` 快照,上游请求从该快照生成 Codex Desktop 风格请求头。
|
||||
|
||||
数据来源:
|
||||
|
||||
- `config.yaml` 中的默认 fingerprint。
|
||||
- SQLite 中的当前 fingerprint。
|
||||
- 后台 fingerprint update 任务。
|
||||
|
||||
启动时会用配置默认值初始化缺失的当前 fingerprint。后台任务从 Codex Desktop appcast 获取版本信息,更新 SQLite 中的 current fingerprint,并替换运行时快照。更新后的请求使用最新 fingerprint。
|
||||
|
||||
## 后台任务
|
||||
|
||||
后台任务由 `runtime/tasks/coordinator.rs` 统一启动和关闭:
|
||||
|
||||
- `cookie_cleanup`:清理过期 Cookie。
|
||||
- `session_cleanup`:清理过期管理端 session。
|
||||
- `session_affinity_cleanup`:清理过期 session affinity。
|
||||
- `model_refresh`:按账号计划刷新模型列表和模型计划快照。
|
||||
- `token_refresh`:刷新即将过期的 access token。
|
||||
- `quota_refresh`:主动刷新账号 quota,并同步账号池状态。
|
||||
- `fingerprint_update`:更新 Codex Desktop fingerprint。
|
||||
- `websocket_pool`:关闭阶段统一 shutdown WebSocket pool。
|
||||
|
||||
周期任务通过统一 runner 管理 shutdown 信号和超时退出。
|
||||
|
||||
## 监控和用量
|
||||
|
||||
请求完成后写入:
|
||||
|
||||
- `usage_records`:请求、上游尝试、状态码、失败分类、延迟、metadata。
|
||||
- `usage_time_buckets`:按中国时区 15 分钟槽聚合请求数、错误数、token、缓存、首内容帧延迟和总延迟。
|
||||
- `account_usage`:账号累计用量和当前 quota 窗口用量。
|
||||
- `account_model_usage`:账号和模型维度用量。
|
||||
|
||||
Dashboard:
|
||||
|
||||
- `/api/admin/dashboard/summary`
|
||||
- `/api/admin/dashboard/trend?kind=usage|latency|errors`
|
||||
|
||||
Dashboard 的今日指标、趋势和健康时间线按中国自然日计算。账号容量、账号状态、服务状态和调度策略取当前运行时状态。累计请求、累计 token、累计缓存和累计命中率来自账号用量汇总。
|
||||
|
||||
请求明细:
|
||||
|
||||
- `/api/admin/usage/records`
|
||||
- `/api/admin/usage/records/summary`
|
||||
- `/api/admin/usage/records/detail`
|
||||
- `/api/admin/usage/records/insights/models`
|
||||
- `/api/admin/usage/records/insights/endpoints`
|
||||
- `/api/admin/usage/records/insights/token-trend`
|
||||
- `/api/admin/usage/records/insights/latency-trend`
|
||||
- `/api/admin/usage/records/delete`
|
||||
|
||||
## 管理端 API
|
||||
|
||||
管理端 API 位于 `/api/admin/*`,使用管理员 session cookie 保护。
|
||||
|
||||
主要路由:
|
||||
|
||||
- Auth:`/api/admin/login`、`/api/admin/auth/status`、`/api/admin/logout`。
|
||||
- Settings:`/api/admin/settings`、`/api/admin/settings/admin-api-key`。
|
||||
- System:`/api/admin/system/version`、`/api/admin/system/update-detail`、`/api/admin/system/update-events`、`/api/admin/system/update-status`、`/api/admin/system/update`、`/api/admin/system/restart`、`/api/admin/system/rollback`。
|
||||
- Dashboard:`/api/admin/dashboard/summary`、`/api/admin/dashboard/trend`。
|
||||
- Accounts:`/api/admin/accounts`、`/api/admin/accounts/import`、`/api/admin/accounts/export`、`/api/admin/accounts/oauth/*`、`/api/admin/accounts/test`、`/api/admin/accounts/models`、`/api/admin/accounts/refresh`、`/api/admin/accounts/health-check`、`/api/admin/accounts/quota`。
|
||||
- Usage Records:`/api/admin/usage/records*`。
|
||||
- API Keys:`/api/admin/keys`、`/api/admin/keys/update`、`/api/admin/keys/delete`。
|
||||
|
||||
## 前端
|
||||
|
||||
前端由后端托管,不单独部署 Node 服务。
|
||||
|
||||
结构:
|
||||
|
||||
- `frontend/src/api/modules`:管理端 API 客户端。
|
||||
- `frontend/src/components/base`:基础 UI 组件。
|
||||
- `frontend/src/layout`:主布局、侧边栏和系统更新弹窗。
|
||||
- `frontend/src/views`:业务页面。
|
||||
- `frontend/src/stores`:Pinia 状态。
|
||||
- `frontend/src/styles/tokens.css`:亮色/暗色主题 token。
|
||||
|
||||
主要页面:
|
||||
|
||||
- 首页 Dashboard。
|
||||
- 账号管理。
|
||||
- API Key 管理。
|
||||
- 请求明细。
|
||||
- 系统设置。
|
||||
- 系统更新。
|
||||
|
||||
## 在线更新
|
||||
|
||||
在线更新由主服务处理,没有独立 updater sidecar。
|
||||
|
||||
流程:
|
||||
|
||||
1. 查询 GitHub Release。
|
||||
2. 按当前 OS/arch 选择 asset。
|
||||
3. 下载归档和 checksum 文件。
|
||||
4. 校验 checksum。
|
||||
5. 解压后端二进制和 `web/dist`。
|
||||
6. 替换当前二进制和前端静态资源。
|
||||
7. 写入 `.runtime/data/update-state.json`。
|
||||
8. 通过 SSE 推送中文更新日志。
|
||||
9. 管理端触发重启。
|
||||
|
||||
Docker 模式依靠 `restart: unless-stopped` 拉起新进程。非 Docker 模式会安排新进程延迟启动,再关闭当前进程。
|
||||
|
||||
## Docker
|
||||
|
||||
`deploy/Dockerfile` 的构建上下文是仓库根目录。
|
||||
|
||||
阶段:
|
||||
|
||||
- `frontend-builder`:安装前端依赖并构建 Vue SPA。
|
||||
- `backend-builder`:构建 Rust release 二进制。
|
||||
- `release-asset-builder` / `release-asset`:组合后端二进制和前端 dist,生成 GitHub Release 归档。
|
||||
- `runtime-base`:Debian slim、CA、curl、非 root 用户和运行目录。
|
||||
- `runtime`:源码构建运行镜像。
|
||||
- `runtime-prebuilt`:发布 workflow 使用的预构建二进制镜像阶段。
|
||||
|
||||
Compose 默认:
|
||||
|
||||
- 镜像:`ghcr.io/zyycn/codex-proxy-rs:latest`,可用 `CPR_IMAGE` 覆盖。
|
||||
- 端口:`127.0.0.1:8080:8080`。
|
||||
- 数据:`.runtime/data:/app/data`。
|
||||
- 日志:`.runtime/logs:/app/logs`。
|
||||
- 配置:`.runtime/config.yaml:/app/config.yaml:ro`。
|
||||
- 更新:`CPR_UPDATE_REPOSITORY=zyycn/codex-proxy-rs`,`CPR_ENABLE_SELF_RESTART=true`。
|
||||
|
||||
## 发布
|
||||
|
||||
版本文件:
|
||||
|
||||
- `release/version.yaml`
|
||||
- `release/platforms.yaml`
|
||||
|
||||
发布命令:
|
||||
| 文件(现路径) | 行数 | 拆法 |
|
||||
| --- | --- | --- |
|
||||
| `proxy/dispatch/responses/service.rs` | 1913 | → `dispatch/service.rs`(编排骨架)+ `upstream_call.rs` + `stream/*` + `recovery/*` 已分走的逻辑归位 |
|
||||
| `upstream/transport/websocket.rs` | 1470 | → `websocket.rs`(连接/握手)+ `websocket_frames.rs`(帧编解码) |
|
||||
| `admin/accounts/routes.rs` | 1455 | → `api/admin/accounts_routes.rs` 按子资源拆(accounts / oauth / quota / import-export) |
|
||||
| `upstream/accounts/pool.rs` | 1420 | → `accounts/pool.rs`(门面)+ `pool/state.rs`(槽位/在途)+ `pool/filters.rs` |
|
||||
| `upstream/transport/client.rs` | 1251 | → `client.rs` + `client_sse.rs`(SSE 请求路径) |
|
||||
| `upstream/accounts/store/mod.rs` | 1126 | → `store.rs` 门面 + `store/queries.rs` 按读写再分 |
|
||||
| `admin/monitoring/usage_record_store.rs` | 1114 | v4 重写时天然拆为 `telemetry/usage/{store,query}.rs` |
|
||||
| `admin/monitoring/dashboard.rs` | 1080 | → `telemetry/dashboard.rs`(聚合)+ `api/admin/dashboard_routes.rs`(HTTP) |
|
||||
| `upstream/protocol/websocket.rs` | 1033 | → `websocket.rs`(事件转换)+ `websocket_errors.rs`(错误帧解析) |
|
||||
| `upstream/accounts/token_refresh/runtime.rs` | 953 | → `accounts/refresh/{service,lease,policy}.rs` |
|
||||
| `admin/system/updater/mod.rs` | 811 | → `update/{service,state}.rs` |
|
||||
|
||||
### 5.5 一域一账本
|
||||
|
||||
每张数据库表恰好属于一个域,只有该域的 store 可以写它:
|
||||
|
||||
| 表(database.md v4) | 属主域 |
|
||||
| --- | --- |
|
||||
| usage_records / ops_error_logs / request_time_buckets / account_usage / account_model_usage | `telemetry` |
|
||||
| accounts / account_refresh_leases / account_cookies | `accounts` |
|
||||
| model_plan_snapshots | `models` |
|
||||
| client_api_keys | `keys` |
|
||||
| admin_users / admin_sessions | `auth` |
|
||||
| runtime_settings | `settings` |
|
||||
| session_affinities | `dispatch`(affinity/store.rs) |
|
||||
| fingerprints / fingerprint_update_history | `upstream`(fingerprint/store.rs) |
|
||||
| schema_migrations | `infra` |
|
||||
|
||||
跨域读通过属主域的 query/service 方法,不直接跨域写 SQL。
|
||||
|
||||
---
|
||||
|
||||
## 6. 现状审计结论(改造动机存档)
|
||||
|
||||
### 6.1 分层违规(P1,本次改造的核心动机)
|
||||
|
||||
1. **热路径反向依赖管理端**:`proxy/dispatch/usage_events.rs`、`event_recording.rs` import `admin::monitoring` 的 `Admin*Service` 写账本;`runtime/services.rs` 用 Admin 命名的服务装配核心链路。遥测是领域层,管理端只是它的一个读者。→ 建 `telemetry` 域(PR1)。
|
||||
2. **核心编排寄居 proxy**:`proxy/dispatch/`(约 5600 行,含 1913 行 service)是系统心脏,却挂在客户端 HTTP 入口下;`SqliteSessionAffinityStore`(DB store)住在 proxy 里。→ 建 `dispatch` 域(PR4)。
|
||||
3. **config 模块越权**:`config/settings.rs`(601 行)= 持久化 + 热更新推送 + admin API key 校验,还持有账号池引用向上推送。config 应只剩启动 schema + loader。→ 建 `settings` 域 + watch 反转(PR2)。
|
||||
4. **admin/keys 四合一**:`admin/keys/service.rs`(698 行)同时装着 proxy 鉴权用的 `ClientKeyService`、运行时缓存、SQLite store、管理端 service——客户端鉴权路径穿过 admin 模块。→ 建 `keys` 域(PR2)。
|
||||
5. **web/http 假分家**:`web/` 只有 69 行 assets,是路由的 fallback 组成部分。→ 并入 `api`(PR5)。
|
||||
6. healthz 穿透三层拿 store 连接池执行裸 SQL(§5.2)。
|
||||
|
||||
### 6.2 命名不规范(P2)
|
||||
|
||||
1. 持久化四种形态并存:目录 `store/`、后缀 `_store.rs`、藏在功能文件(`cookies.rs`、`session_affinity.rs`)、`Repository`(fingerprint)。
|
||||
2. 行/DTO 三种叫法:`rows.rs` / `*_model.rs` / `contracts.rs`。
|
||||
3. 动名词文件:`importing.rs` ×2、`exporting.rs`、`testing.rs`(后者必然被误读为测试)。
|
||||
4. 类型名与文件名脱节:`account_usage_service.rs` 里是 `AdminUsageService`;`SqliteUsageStore` 存的是 account_usage;"usage" 一词横跨三个概念。
|
||||
5. `model.rs`(账号实体)与 `models/`(LLM 目录域)一词双义——实体文件改 `account.rs` 后消解。
|
||||
6. `admin/system/routes.rs` 是 5 行 re-export 假门面;`admin/response.rs`、`admin/update_payload.rs`、`upstream/token_client.rs`、`upstream/fingerprint.rs` 游离在模块根。
|
||||
7. `Arc as StdArc` 式无意义别名。
|
||||
|
||||
### 6.3 模块肥大(P2)
|
||||
|
||||
超 800 行文件 11 个(§5.4 已列);巨型 `mod.rs` 5 个(§5.1)。
|
||||
|
||||
### 6.4 仓库级(P3)
|
||||
|
||||
1. `backend/build/build.rs`:Cargo 惯例是包根 `build.rs`;非标准位置徒增一行配置与一次困惑。→ 移回 `backend/build.rs`。
|
||||
2. 根目录 `server-pulls.runtime/`:点分隔命名不合仓库其余惯例(gitignored 运行产物)。→ 更名 `.server-pulls/` 或移出仓库根。
|
||||
3. `frontend/src` 结构合格(api/components/views/stores 分层清晰),仅一条规范:`utils/` 内文件必须按 §4.2 角色词表命名,禁止继续堆积泛用文件。
|
||||
|
||||
**结论**:模块内聚质量整体不差(scheduler 的文档注释、tests 镜像结构都是亮点),债务集中在**域归属错位**与**命名不一致**,属于"搬家 + 更名"能解决的范畴,不需要重写逻辑。
|
||||
|
||||
---
|
||||
|
||||
## 7. 旧 → 新映射总表(可执行)
|
||||
|
||||
按目录粒度;`†` = 搬家同时拆分/更名,见 §5.4 与 §4.2。
|
||||
|
||||
| 现路径 | 终态路径 |
|
||||
| --- | --- |
|
||||
| `config/schema.rs` + `config/loader.rs` | `bootstrap/config.rs` |
|
||||
| `config/settings.rs` | `settings/{types,store,service}.rs` † |
|
||||
| `runtime/{bootstrap,services,state,shutdown}.rs` | `bootstrap/` 同名 |
|
||||
| `runtime/tasks/*` | `bootstrap/tasks/*`(新增 `retention_trim.rs`) |
|
||||
| `http/router.rs` | `api/router.rs` |
|
||||
| `http/middleware/*` | `api/middleware/*` |
|
||||
| `web/assets.rs` | `api/assets.rs` |
|
||||
| `proxy/{auth,router}.rs`、`proxy/openai/*` | `api/client/*`(`routes.rs` → `router.rs` 与各资源文件) |
|
||||
| `proxy/dispatch/usage_events.rs` | `telemetry/recorder.rs` †(Admin* 类型消失) |
|
||||
| `proxy/dispatch/session_affinity.rs` | `dispatch/affinity/{types,store,service}.rs` † |
|
||||
| `proxy/dispatch/responses/affinity.rs` | `dispatch/affinity/resolve.rs` |
|
||||
| `proxy/dispatch/responses/service.rs` | `dispatch/service.rs` † |
|
||||
| `proxy/dispatch/responses/{live_stream,stream_lifecycle,prefetch,sse_failure,trace}.rs` | `dispatch/stream/{live,lifecycle,prefetch,sse_failure,trace}.rs` |
|
||||
| `proxy/dispatch/responses/event_recording.rs` | `dispatch/recording.rs` |
|
||||
| `proxy/dispatch/responses/errors.rs` + `proxy/dispatch/errors.rs` | `dispatch/errors.rs` † |
|
||||
| `proxy/dispatch/{auth_recovery,cloudflare,exhaustion,implicit_resume,reasoning_replay}.rs` | `dispatch/recovery/{auth,cloudflare,exhaustion,implicit_resume,reasoning_replay}.rs` |
|
||||
| `proxy/dispatch/upstream.rs` | `dispatch/upstream_call.rs` |
|
||||
| `admin/monitoring/usage_record_{model,store,service}.rs` | `telemetry/usage/{types,store,query}.rs` †(v4 重写同窗口) |
|
||||
| `admin/monitoring/ops_error_{model,store,service}.rs` | `telemetry/ops/{types,store,query}.rs` † |
|
||||
| `admin/monitoring/account_usage_{store,service}.rs` | `telemetry/account_usage/{store,query}.rs` † |
|
||||
| `admin/monitoring/billing.rs` | `telemetry/billing.rs` |
|
||||
| `admin/monitoring/dashboard.rs` | `telemetry/dashboard.rs` + `api/admin/dashboard_routes.rs` † |
|
||||
| `admin/monitoring/diagnostics.rs` | `telemetry/diagnostics.rs` |
|
||||
| `admin/monitoring/usage_record_routes.rs` | `api/admin/usage_routes.rs` |
|
||||
| `admin/keys/service.rs` | `keys/{types,store,cache,service,manage}.rs` † |
|
||||
| `admin/keys/routes.rs` | `api/admin/keys_routes.rs` |
|
||||
| `admin/auth/service.rs` | `auth/{types,store,service}.rs` † |
|
||||
| `admin/auth/session.rs` | `api/admin/session.rs` |
|
||||
| `admin/accounts/routes.rs` | `api/admin/accounts_routes.rs` †(按子资源拆) |
|
||||
| `admin/accounts/quota_view.rs` | `accounts/manage/quota_view.rs` |
|
||||
| `admin/accounts/service/*` | `accounts/manage/*`(`importing`→`import`、`exporting`→`export`、`testing`→`probe`、`contracts`→`types`) |
|
||||
| `admin/{response,update_payload}.rs` | `api/admin/response.rs`、`update/types.rs` |
|
||||
| `admin/router.rs` | `api/admin/router.rs` |
|
||||
| `admin/settings/routes.rs` | `api/admin/settings_routes.rs` |
|
||||
| `admin/system/{routes,state}.rs`、`admin/system/updater/*` | `api/admin/system_routes.rs`、`update/{state,service,release,download,archive}.rs` † |
|
||||
| `upstream/accounts/model.rs` | `accounts/account.rs` |
|
||||
| `upstream/accounts/{pool,window}.rs` | `accounts/{pool,window}.rs`(pool 拆 †) |
|
||||
| `upstream/accounts/store/*` | `accounts/store.rs` + `accounts/store/{queries,rows}.rs` † |
|
||||
| `upstream/accounts/quota/*` | `accounts/quota/{service,runtime}.rs` † |
|
||||
| `upstream/accounts/token_refresh/*` | `accounts/refresh/{policy,lease,service}.rs` † |
|
||||
| `upstream/accounts/cookies.rs` | `accounts/cookies/store.rs` |
|
||||
| `upstream/accounts/importing.rs` | `accounts/import.rs` |
|
||||
| `upstream/accounts/service.rs` | 并入 `accounts/mod.rs` 门面导出(32 行) |
|
||||
| `upstream/scheduler/*` | `accounts/scheduler.rs` + `accounts/scheduler/*`(mod.rs 门面类型移出) |
|
||||
| `upstream/models/*` | `models/{types,catalog,service,store}.rs`(`backend_entry`+`info`+`snapshot`+`config` 并入 `types.rs`/`service.rs`) |
|
||||
| `upstream/fingerprint.rs` | `upstream/openai/fingerprint/{types,store,runtime,updater}.rs` †(Repository→Store) |
|
||||
| `upstream/{protocol,transport}/*`、`upstream/token_client.rs` | `upstream/openai/` 下同名路径(websocket/client 按 §5.4 拆) |
|
||||
| `infra/*` | 原地保留 |
|
||||
| `backend/build/build.rs` | `backend/build.rs` |
|
||||
|
||||
---
|
||||
|
||||
## 8. 执行计划(六个顺序 PR,每个独立编译绿、测试绿)
|
||||
|
||||
每个 PR 的固定验收:`cargo fmt --check`、`cargo clippy -- -D warnings`、`cargo test`、`backend/scripts/check_architecture.sh`(§9,随 PR 逐步启用对应规则)、`backend/tests` 镜像目录同步搬移。**搬家 PR 原则上零行为变更**(更名与移动,不改逻辑);唯一例外是 PR1 与 database.md v4 的联动重写。
|
||||
|
||||
| PR | 内容 | 关键动作 |
|
||||
| --- | --- | --- |
|
||||
| **PR1** | `telemetry` 域成立(与 database.md 0004 + §7 代码改造同窗口) | monitoring 全部迁出 admin;`usage_events.rs` → `recorder.rs`;Admin* 遥测类型按 §4.3 更名;dispatch/api 改依赖 telemetry;启用"禁止 `crate::admin::monitoring`"检查 |
|
||||
| **PR2** | `keys` / `auth` / `settings` / `update` 域成立 | 四个域从 admin/config 拆出;settings 改 watch 广播(§2.3);admin API key / session token 哈希化(database.md)落在此处;config 只剩 schema+loader |
|
||||
| **PR3** | `accounts` / `models` 域成立 | `upstream/accounts`、`upstream/scheduler`、`admin/accounts/service` 三处合并为 `accounts`;`upstream` 缩为纯客户端并落成 `openai/` 子目录;动名词文件更名;Repository→Store |
|
||||
| **PR4** | `dispatch` 域成立 | proxy/dispatch 整体迁出并按目标树重组;affinity store 独立成文件;`proxy` 只剩 api/client 的料 |
|
||||
| **PR5** | `api` / `bootstrap` 壳成立 | http + web + proxy 入口 + admin routes 合并为 `api`;runtime + config 装配部分改名 `bootstrap`;删除 `crate::{admin,proxy,runtime,web,http,config}` 六个旧根模块;启用全部依赖检查 |
|
||||
| **PR6** | 行数预算清账 | §5.4 清单内 11 个文件拆到硬线以下;启用行数预算 CI 警告 |
|
||||
|
||||
顺序依据:PR1 与数据库改造互锁必须最先;PR2–4 自底向上消除旧模块的存在必要;PR5 收壳;PR6 纯内部拆分无跨域影响。
|
||||
|
||||
---
|
||||
|
||||
## 9. CI 强制
|
||||
|
||||
新增 `backend/scripts/check_architecture.sh`,进 `.github/workflows/ci.yml`(在 clippy 之后):
|
||||
|
||||
```bash
|
||||
release/publish <version>
|
||||
#!/usr/bin/env bash
|
||||
# 架构依赖与命名断言。违规即失败。
|
||||
set -euo pipefail
|
||||
cd "$(dirname "$0")/.."
|
||||
fail=0
|
||||
deny() { # deny <描述> <grep-pattern> <路径...>
|
||||
if grep -rn --include='*.rs' -E "$2" "${@:3}" ; then
|
||||
echo "ARCH VIOLATION: $1" >&2; fail=1
|
||||
fi
|
||||
}
|
||||
|
||||
# 分层:领域层不得依赖 api / bootstrap
|
||||
deny "domain -> api/bootstrap" 'crate::(api|bootstrap)\b' \
|
||||
src/infra src/upstream src/telemetry src/keys src/auth src/settings src/update \
|
||||
src/accounts src/models src/dispatch
|
||||
# upstream 是纯客户端
|
||||
deny "upstream -> domain" 'crate::(accounts|dispatch|telemetry|keys|auth|settings|update|models)\b' src/upstream
|
||||
# HTTP 类型只在 api / bootstrap
|
||||
deny "axum outside api" '\baxum(::|\b)' \
|
||||
src/infra src/upstream src/telemetry src/keys src/auth src/settings src/update \
|
||||
src/accounts src/models src/dispatch
|
||||
# SQL 只在 store / query
|
||||
deny "sqlx outside store/query" 'sqlx::query' \
|
||||
$(find src/api -name '*.rs')
|
||||
# 旧模块根不得复活
|
||||
deny "legacy module path" 'crate::(admin|proxy|runtime|web|http|config)::' src
|
||||
# 禁用文件名
|
||||
banned=$(find src -name 'importing.rs' -o -name 'exporting.rs' -o -name 'testing.rs' \
|
||||
-o -name 'util.rs' -o -name 'helper.rs' -o -name 'common.rs' -o -name 'misc.rs')
|
||||
[ -n "$banned" ] && { echo "ARCH VIOLATION: banned file names: $banned" >&2; fail=1; }
|
||||
# mod.rs 纪律
|
||||
while read -r f; do
|
||||
lines=$(wc -l < "$f")
|
||||
[ "$lines" -gt 50 ] && { echo "ARCH VIOLATION: $f has $lines lines (mod.rs <= 50)" >&2; fail=1; }
|
||||
done < <(find src -name mod.rs)
|
||||
# 行数硬上限(警告期:仅报告;PR6 后改 fail)
|
||||
find src -name '*.rs' ! -name mod.rs | while read -r f; do
|
||||
lines=$(wc -l < "$f")
|
||||
[ "$lines" -gt 800 ] && echo "WARN: $f has $lines lines (> 800)"
|
||||
done
|
||||
exit $fail
|
||||
```
|
||||
|
||||
发布 workflow 校验触发 tag 必须等于 `v` + `release/version.yaml`。
|
||||
规则随 PR 分阶段启用(PR1 起启用 legacy-path 的 monitoring 子集,PR5 后全量);`deny` 的路径参数在过渡期只指向已完成搬家的目录。
|
||||
|
||||
构建时注入:
|
||||
---
|
||||
|
||||
- `CPR_VERSION`
|
||||
- `CPR_GIT_SHA`
|
||||
- `CPR_BUILD_TIME`
|
||||
- `CPR_BUILD_TYPE=release`
|
||||
## 10. 仓库级约定
|
||||
|
||||
Release 产物按 `release/platforms.yaml` 生成 Linux、Darwin 和 Windows 包。GHCR 镜像发布 `linux/amd64` 和 `linux/arm64` 多平台 manifest。
|
||||
- `backend/tests/` 镜像 `src/`,不在 `src/` 写测试模块。
|
||||
- `docs/` 只保留权威文档:`architecture.md`(本文)、`database.md`。行为细节(调度策略、WS 池、更新流程等)以模块文档注释为准(`scheduler/mod.rs` 是范本),不再在 docs 里维护第二份会腐化的副本。
|
||||
- `deploy/`、`release/`、`skills/` 现状合格,不动。
|
||||
- 前端改动仅一处联动:keys 列表不再返回完整 key(database.md §7.8)。
|
||||
|
||||
## CI
|
||||
---
|
||||
|
||||
常规 CI 位于 `.github/workflows/ci.yml`:
|
||||
## 11. 扩展路径:第二上游 provider(已规划:Cloudflare Workers AI)
|
||||
|
||||
- Rust fmt、clippy 和集成测试。
|
||||
- 前端 install、typecheck/build。
|
||||
- Docker Compose 配置校验。
|
||||
- Runtime image build。
|
||||
分层已为多 provider 留好隔离区,接入时**不新增层、不动依赖规则**,动作清单:
|
||||
|
||||
安全扫描位于 `.github/workflows/security-scan.yml`:
|
||||
| 位置 | 动作 |
|
||||
| --- | --- |
|
||||
| `upstream/` | 新增 `upstream/cloudflare/`:client(按账号拼 `accounts/{account_id}/ai/v1` base URL)+ protocol 适配器(**Responses ↔ chat/completions 双向转换**,接入的主要工作量)。无 fingerprint、无 token_client(静态 API key,无刷新)。`upstream/mod.rs` 引入 `UpstreamClient` trait,`dispatch` 从具体 `CodexBackendClient` 改依赖 trait——在只有一家实现时不预建该 trait(投机抽象),第二家动工即引入 |
|
||||
| `models` | catalog 收录 `@cf/...` 命名空间;别名映射升级为"客户端名 → (provider, 上游名)";**models 域是"模型 → provider"的路由枢纽**,dispatch 据此选池 |
|
||||
| `accounts` | `scheduler/candidates` 加 provider 硬过滤(单池 + 过滤,不按 provider 拆池);`manage/{oauth,import}` 按 provider 分实现(CF 导入格式 = `(account_id, api_token)`);refresh/租约对 CF 账号整体跳过(`refresh_token` NULL 语义已覆盖) |
|
||||
| `dispatch` | `affinity` / `recovery` 按 provider 策略分发:CF 无会话语义 → 跳过 affinity;恢复只剩朴素 429/5xx 重试 |
|
||||
| `telemetry` | 无结构改动——provider 归因维度已随 schema v4 就位(database.md §4.10–4.12);billing 单价查找加 provider 参数 |
|
||||
| 不动 | `api/client`(入站永远是 OpenAI 兼容协议)、`keys`、`auth`、`settings`、`update`、`infra` |
|
||||
|
||||
- `cargo audit`。
|
||||
- `pnpm audit --prod --audit-level=high`。
|
||||
- Docker runtime image security scan。
|
||||
|
||||
发布流程位于 `.github/workflows/release.yml`,由 `v*` tag 或手动指定 tag 触发。
|
||||
DB 侧一次性迁移(accounts 加 provider、身份列更名、快照表 PK 加维度)见 database.md §8。动工当天需实测验证三点:`cfk_` 老式 authkey 对 OpenAI 兼容端点的鉴权方式(`X-Auth-Email`/`X-Auth-Key` 还是 Bearer)、GLM 类模型是否支持上游 Responses 端点(决定适配器能否对部分模型直通)、SSE 分块与 OpenAI delta 事件的逐字段兼容性。
|
||||
|
||||
@@ -0,0 +1,880 @@
|
||||
# Database —— 唯一权威文档
|
||||
|
||||
本文档是 Codex Proxy RS 数据库的**唯一**权威文档:终态 Schema 全量定义、全局纪律、统计口径、一次到位的迁移方案(v3 → v4)与代码改造清单。
|
||||
|
||||
- 取代并合并旧的三份文档(database-design.md / database-audit.md / database-schema.md),旧文档一律作废删除。
|
||||
- 基线:2026-07-08,生产库 schema v3(`backend/src/infra/database.rs`)。
|
||||
- 策略:**一次破坏性迁移直达终态,不做兼容层**。没有双写期、没有回填过渡期、没有查询层兼容过滤;不满足终态契约的历史数据按 §6.3 的规则迁移或丢弃,丢弃项全部显式列出。
|
||||
|
||||
---
|
||||
|
||||
## 1. 运行形态
|
||||
|
||||
单进程 SQLite。这是刻意选择:自托管代理、单实例部署,SQLite 免去外部依赖且性能远超需求。多实例是明确的未来触发条件而非当前需求(§8)。
|
||||
|
||||
连接参数(`database.rs` 连接期设置,全部显式):
|
||||
|
||||
| Pragma | 值 | 理由 |
|
||||
| --- | --- | --- |
|
||||
| `journal_mode` | WAL | 读写不互斥,单写多读的正确形态 |
|
||||
| `synchronous` | NORMAL | WAL 下的标准取舍:掉电最多丢最后一次 checkpoint 后的事务,不损坏库。FULL 的额外 fsync 对本系统无意义 |
|
||||
| `foreign_keys` | ON | 级联删除是业务语义的一部分 |
|
||||
| `busy_timeout` | 5s | 单写者下的排队上限 |
|
||||
| `auto_vacuum` | INCREMENTAL | 三张增长表按保留期批量删除,必须有空间回收路径;trim 任务后执行 `pragma incremental_vacuum`(§7.7)。对已存在的库,该设置在 0004 迁移后的一次 `vacuum` 时生效 |
|
||||
|
||||
---
|
||||
|
||||
## 2. 全局纪律
|
||||
|
||||
### 2.1 命名
|
||||
|
||||
| 对象 | 规则 | 示例 |
|
||||
| --- | --- | --- |
|
||||
| 表名 | snake_case;实体集合复数;子表 `<父实体单数>_<内容>` | `accounts`、`account_cookies` |
|
||||
| 表名后缀 | 事实 `_records`/`_logs`;聚合 `_buckets`;缓存 `_snapshots`;运行态 `_leases`/`_affinities`;审计 `_history` | `usage_records` |
|
||||
| 列名 | snake_case,全库同义同名 | `account_id` 处处一致 |
|
||||
| 时间戳 | `*_at`,TEXT,RFC3339 UTC | `created_at` |
|
||||
| 时长 | `*_seconds` / `*_ms`,INTEGER | `latency_ms` |
|
||||
| 计数 | `*_count`,INTEGER,`check >= 0` | `request_count` |
|
||||
| JSON | `*_json`,TEXT | `metadata_json` |
|
||||
| 布尔 | 裸形容词/过去式,INTEGER,`check in (0,1)` | `enabled` |
|
||||
| 外键/逻辑引用列 | `<被引用实体单数>_id` | `client_api_key_id` |
|
||||
| 凭据哈希 | `*_hash`,明文永不落库(§2.6) | `key_hash`、`token_hash` |
|
||||
| 主键 | `id`,除非自然复合键或哈希键 | — |
|
||||
| 普通/唯一索引 | `idx_<表>_<列摘要>` / `ux_<表>_<列摘要>` | `ux_accounts_chatgpt_identity` |
|
||||
| 双端状态码 | 视角前缀 | `status_code` / `client_status_code` / `upstream_status_code` |
|
||||
| 模型归因 | 来源前缀 | `model` / `requested_model` / `upstream_model` |
|
||||
|
||||
### 2.2 类型
|
||||
|
||||
只用两种存储类型表达业务数据:
|
||||
|
||||
| 用途 | 类型 | 约定 |
|
||||
| --- | --- | --- |
|
||||
| ID、名称、枚举、时间戳、JSON、哈希 | TEXT | 枚举加 `check in (...)` |
|
||||
| 计数、时长、布尔、状态码 | INTEGER | 计数 `>= 0`;布尔 `in (0,1)` |
|
||||
|
||||
**不使用 REAL,不存储金额**。成本一律查询期由 token × 单价计算(`admin/monitoring/billing.rs`):单价会变、需追溯重算,落库的浮点成本既漂移又过期。
|
||||
|
||||
### 2.3 时间戳
|
||||
|
||||
- 库内一律 UTC,统一经 `chrono::DateTime::to_rfc3339()` 写入(`+00:00` 后缀)。RFC3339 UTC 文本的字典序 = 时间序,TEXT 时间戳因此可直接做范围查询与排序索引。
|
||||
- 时区换算(中国时区展示、15 分钟槽对齐)只在代码层(`infra/time.rs` 的 `china_*` 系列),数据库不感知时区。
|
||||
- **禁止**在同一张做范围比较的表里混用 `strftime('…Z')` 与 `to_rfc3339()`。唯一例外:`schema_migrations.applied_at` 用 `strftime`(不做范围比较)。
|
||||
|
||||
### 2.4 ID
|
||||
|
||||
主键一律 TEXT 随机 ID(`infra/identity.rs` 生成):无中心分配、写入前可生成、跨表无碰撞。例外:`schema_migrations.version`(天然有序整数)、`runtime_settings.id`(单例锚点)、凭据表以哈希为键(§2.6)。
|
||||
|
||||
### 2.5 NULL 与 sentinel
|
||||
|
||||
- "无值 / 未观测"一律 NULL;**禁止** 0、空字符串做哨兵。NULL 与 0 语义不同(`input_tokens` NULL = 上游未报告,0 = 确认为零)。
|
||||
- 聚合维度列不可空时,未知值用 `__unknown__`;预聚合汇总行(预留)用 `__all__`。
|
||||
- 唯一例外:`ux_accounts_chatgpt_identity` 的索引表达式 `coalesce(chatgpt_user_id, '')`——空串只存在于索引表达式,不是存储值,不构成先例。
|
||||
|
||||
### 2.6 Secret 纪律(无例外)
|
||||
|
||||
凡是"出示即获得权限"的值都是凭据,**明文一律不落库**:
|
||||
|
||||
| 凭据 | 存储 | 哈希类型 | 理由 |
|
||||
| --- | --- | --- | --- |
|
||||
| 管理员密码 | `admin_users.password_hash` | 慢哈希(argon2/bcrypt) | 低熵人类输入,必须抗字典攻击 |
|
||||
| 客户端 API key | `client_api_keys.key_hash` | SHA-256 | 高熵随机串无字典攻击面,快哈希保住 unique 索引 O(log n) 查找;明文只在创建响应出现一次,`prefix` 供事后定位 |
|
||||
| 管理端 session token | `admin_sessions.token_hash` | SHA-256 | 同上。库文件/备份泄露不再等于会话劫持 |
|
||||
| 管理 API key | `runtime_settings.admin_api_key_hash` | SHA-256 | 同上;校验 = 哈希后 `ct_eq` |
|
||||
| 上游 access/refresh token、Cookie value | 明文列,标记 secret | —(功能上必须可逆) | 不进日志、不进列表 API;静态加密的触发条件见 §8 |
|
||||
|
||||
### 2.7 JSON 边界
|
||||
|
||||
`*_json` 列只允许三类内容:① 上游原样快照(`quota_json`、`models_json`、`manifest_json`);② 结构化模板(`default_headers_json`、`header_order_json`);③ 调试细节(`metadata_json`)。**凡进入 WHERE / GROUP BY / 告警的字段必须是一等列**;列一旦提升即为唯一可查询真相,JSON 内同名字段仅供人工排查,写入时不再往 metadata 双写已提升字段。
|
||||
|
||||
### 2.8 索引
|
||||
|
||||
1. 每个索引对应一条真实查询;unique 约束已覆盖的等值查找不再建索引。
|
||||
2. 列表分页统一 keyset 模式,配 `(排序时间 desc, id desc)` 复合索引。
|
||||
3. 可空列上的筛选索引一律部分索引 `where <col> is not null`。
|
||||
4. FK 子表必须有以 FK 列开头的索引(SQLite 级联删除在子表无索引时对父表每行全扫)。免除条件必须注明(如 `admin_sessions`:行数个位数)。
|
||||
5. 高频写入表的索引总数进 code review,新增需说明查询来源。
|
||||
|
||||
### 2.9 保留与可重建性(诚实边界)
|
||||
|
||||
| 表 | 保留 | 配置 |
|
||||
| --- | --- | --- |
|
||||
| `usage_records` | 30 天(默认) | `runtime_settings.usage_retention_days` |
|
||||
| `ops_error_logs` | 30 天(默认) | `runtime_settings.ops_error_retention_days` |
|
||||
| `request_time_buckets` | 90 天(默认) | `runtime_settings.bucket_retention_days` |
|
||||
| `admin_sessions` / `account_cookies` / `session_affinities` | 按 `expires_at` 清理 | 定时任务 |
|
||||
| 其余 | 常量级增长 | — |
|
||||
|
||||
**可重建性以事实保留期为界,不夸大**:
|
||||
|
||||
- `request_time_buckets` 在事实保留期内可由两张事实表重算(`rebuild-derived` 命令,§7.6);超出事实保留期的桶是**因事实过期而升格的一级数据**,写入纪律 + 同事务写入(§5.2)是它们的唯一保障。
|
||||
- `account_usage` / `account_model_usage` 的累计列**不可重建**(事实早已过期),同样只受事务保护。文档不再声称"一切派生皆可重算"。
|
||||
- `window_*` 列是运行态(窗口边界来自上游),不可重建,可丢弃。
|
||||
|
||||
---
|
||||
|
||||
## 3. 架构分层与关系
|
||||
|
||||
```text
|
||||
┌─ 元数据 ──────────────────────────────────────────────┐
|
||||
│ schema_migrations │
|
||||
├─ 配置 ────────────────────────────────────────────────┤
|
||||
│ runtime_settings(单例热配置) │
|
||||
├─ 接入 ────────────────────────────────────────────────┤
|
||||
│ admin_users · admin_sessions · client_api_keys │
|
||||
├─ 账号 ────────────────────────────────────────────────┤
|
||||
│ accounts · account_cookies · fingerprints │
|
||||
│ fingerprint_update_history · model_plan_snapshots │
|
||||
├─ 事实(不可替代的业务事件)────────────────────────────┤
|
||||
│ usage_records(成功) · ops_error_logs(失败/运维) │
|
||||
├─ 派生(口径见 §2.9)──────────────────────────────────┤
|
||||
│ request_time_buckets · account_usage │
|
||||
│ account_model_usage │
|
||||
└─ 运行态(临时协调,可丢弃)───────────────────────────┤
|
||||
session_affinities · account_refresh_leases │
|
||||
```
|
||||
|
||||
```text
|
||||
admin_users 1──< admin_sessions FK cascade(索引免除:行数个位数)
|
||||
accounts 1──1 account_usage FK cascade(PK=FK)
|
||||
accounts 1──1 account_refresh_leases FK cascade(PK=FK)
|
||||
accounts 1──< account_model_usage FK cascade(复合 PK 前缀)
|
||||
accounts 1──< account_cookies FK cascade(account_domain 前缀)
|
||||
accounts 1──< session_affinities FK cascade(idx_…_account)
|
||||
fingerprints 1──< fingerprint_update_history FK cascade(idx_…_fingerprint)
|
||||
|
||||
usage_records.{account_id, client_api_key_id} → 逻辑引用,无 FK
|
||||
ops_error_logs.{account_id, client_api_key_id} → 逻辑引用,无 FK
|
||||
request_time_buckets.account_id → 逻辑维度,无 FK
|
||||
```
|
||||
|
||||
**事实表/聚合表不建 FK 是有意为之**:审计事实的生命周期必须独立于账号与 key——删除账号或 key 不得抹掉历史用量和错误记录。运行态表相反:账号没了亲和/租约就失去意义,级联删除是正确语义。
|
||||
|
||||
---
|
||||
|
||||
## 4. 表设计(终态 DDL)
|
||||
|
||||
### 4.1 schema_migrations
|
||||
|
||||
迁移版本的唯一事实来源。由 `database.rs` 建库时创建,不走迁移文件。只增不改;已发布迁移永不修改;`CURRENT_SCHEMA_VERSION` 与迁移列表同步提交。
|
||||
|
||||
```sql
|
||||
create table schema_migrations (
|
||||
version integer not null check (version > 0),
|
||||
name text not null,
|
||||
applied_at text not null,
|
||||
primary key (version)
|
||||
);
|
||||
```
|
||||
|
||||
### 4.2 admin_users
|
||||
|
||||
单管理员模型的登录凭据。不预建 `username`/`email`/`role` 死列;多管理员的扩展路径见 §8。
|
||||
|
||||
```sql
|
||||
create table admin_users (
|
||||
id text primary key,
|
||||
password_hash text not null,
|
||||
created_at text not null,
|
||||
updated_at text not null
|
||||
);
|
||||
```
|
||||
|
||||
### 4.3 admin_sessions
|
||||
|
||||
管理端会话的服务端状态。session 存库(而非 JWT)使"登出即失效、改密码踢会话、重启不掉线"都是一条 SQL。
|
||||
|
||||
```sql
|
||||
create table admin_sessions (
|
||||
token_hash text primary key,
|
||||
user_id text not null references admin_users(id) on delete cascade,
|
||||
expires_at text not null,
|
||||
created_at text not null
|
||||
);
|
||||
|
||||
create index idx_admin_sessions_expires on admin_sessions(expires_at);
|
||||
```
|
||||
|
||||
| 字段 | 说明 |
|
||||
| --- | --- |
|
||||
| `token_hash` | `SHA-256(session token)`。客户端持有明文 token,鉴权时哈希后按 PK 等值查找。**明文 token 不落库**(§2.6) |
|
||||
| `user_id` | FK cascade,删用户即踢会话 |
|
||||
| `expires_at` | 鉴权比较 + 清理任务扫描(`idx_admin_sessions_expires`) |
|
||||
|
||||
`user_id` 不建索引是注明的免除:单管理员下行数个位数。IP/UA 审计不进本表;需要时新建 `admin_login_events`。
|
||||
|
||||
### 4.4 client_api_keys
|
||||
|
||||
客户端调用 `/v1/*` 的凭据。key 与上游账号解耦,客户端不感知账号池。
|
||||
|
||||
```sql
|
||||
create table client_api_keys (
|
||||
id text primary key,
|
||||
name text not null,
|
||||
prefix text not null,
|
||||
key_hash text not null unique,
|
||||
label text,
|
||||
enabled integer not null default 1 check (enabled in (0, 1)),
|
||||
created_at text not null,
|
||||
last_used_at text
|
||||
);
|
||||
|
||||
create index idx_client_api_keys_created_id on client_api_keys(created_at desc, id desc);
|
||||
```
|
||||
|
||||
| 字段 | 说明 |
|
||||
| --- | --- |
|
||||
| `id` | 稳定标识;管理操作与事实归因(`usage_records.client_api_key_id`)都引用 id,key 轮换不影响引用 |
|
||||
| `name` / `label` | `name` 必填主展示,`label` 可选注记 |
|
||||
| `prefix` | 明文前几位(`sk-abc1…`),明文销毁后定位 key 的唯一手段 |
|
||||
| `key_hash` | `SHA-256(key)`,unique 索引即鉴权查找路径。明文只在创建响应出现一次,**管理端列表 API 永不回显完整 key** |
|
||||
| `enabled` | 软开关;刻意不建索引(低基数布尔) |
|
||||
| `last_used_at` | 鉴权成功后异步更新,不在请求关键路径强一致 |
|
||||
|
||||
每 key 配额/限流不进本表——未来 `client_api_key_policies` 附属表(§8),归因数据已由事实表的 `client_api_key_id` 备好。
|
||||
|
||||
### 4.5 runtime_settings
|
||||
|
||||
热更新全局配置,单行表,`check (id = 1)` 约束层面保证单例。选宽表不选 KV:列级类型与 check 是 KV 给不了的;加配置要加列,恰好强迫走迁移评审。只放小型标量配置,禁塞大 JSON 状态。
|
||||
|
||||
```sql
|
||||
create table runtime_settings (
|
||||
id integer primary key check (id = 1),
|
||||
model_aliases_json text not null default '{}',
|
||||
refresh_margin_seconds integer not null check (refresh_margin_seconds > 0),
|
||||
refresh_concurrency integer not null check (refresh_concurrency > 0),
|
||||
max_concurrent_per_account integer not null check (max_concurrent_per_account > 0),
|
||||
request_interval_ms integer not null check (request_interval_ms >= 0),
|
||||
rotation_strategy text not null check (rotation_strategy in ('smart', 'quota_reset_priority', 'round_robin', 'sticky')),
|
||||
admin_api_key_hash text,
|
||||
usage_retention_days integer not null default 30 check (usage_retention_days > 0),
|
||||
ops_error_retention_days integer not null default 30 check (ops_error_retention_days > 0),
|
||||
bucket_retention_days integer not null default 90 check (bucket_retention_days > 0),
|
||||
updated_at text not null
|
||||
);
|
||||
```
|
||||
|
||||
| 字段 | 说明 |
|
||||
| --- | --- |
|
||||
| `model_aliases_json` | 客户端模型名 → 上游模型名映射。JSON 合规:整体读写、不按别名查询 |
|
||||
| `refresh_*` / `max_concurrent_per_account` / `request_interval_ms` | 刷新与调度参数 |
|
||||
| `rotation_strategy` | 枚举 check;枚举值变更必须走迁移(0002 先例),库内无死值 |
|
||||
| `admin_api_key_hash` | 管理 API 机器凭据的 SHA-256。NULL = 未启用。校验 = 对来键哈希后 `ct_eq` 比较 |
|
||||
| `*_retention_days` 三列 | 三张增长表各自独立的保留天数。bucket 默认更长:趋势比明细有更长查询价值 |
|
||||
|
||||
### 4.6 accounts
|
||||
|
||||
上游账号的身份、凭据、调度状态与 quota 快照。**保持单表是当前正确决策**:调度器每次决策同时读身份+状态+冷却,拆表只增加 join 与事务面;拆分触发条件明确(§8)。
|
||||
|
||||
```sql
|
||||
create table accounts (
|
||||
id text primary key,
|
||||
email text,
|
||||
chatgpt_account_id text,
|
||||
chatgpt_user_id text,
|
||||
label text,
|
||||
plan_type text,
|
||||
access_token text not null,
|
||||
refresh_token text,
|
||||
access_token_expires_at text,
|
||||
next_refresh_at text,
|
||||
status text not null check (status in ('active', 'expired', 'quota_exhausted', 'disabled', 'banned')),
|
||||
quota_json text,
|
||||
quota_fetched_at text,
|
||||
quota_limit_reached integer not null default 0 check (quota_limit_reached in (0, 1)),
|
||||
quota_verify_required integer not null default 0 check (quota_verify_required in (0, 1)),
|
||||
quota_cooldown_until text,
|
||||
cloudflare_cooldown_until text,
|
||||
added_at text not null,
|
||||
updated_at text not null
|
||||
);
|
||||
|
||||
create index idx_accounts_status on accounts(status);
|
||||
create index idx_accounts_added_id on accounts(added_at desc, id desc);
|
||||
create unique index ux_accounts_chatgpt_identity
|
||||
on accounts(chatgpt_account_id, coalesce(chatgpt_user_id, ''))
|
||||
where chatgpt_account_id is not null;
|
||||
```
|
||||
|
||||
要点:
|
||||
|
||||
- **身份**:`id` 与上游 ID 解耦(导入半成品账号仍可运转);`chatgpt_account_id` + `chatgpt_user_id` 复合唯一防重复导入,身份未知的账号不参与去重(部分索引)。`plan_type` 不加枚举 check:取值由上游定义。
|
||||
- **凭据(secret)**:`access_token` not null 是账号存在的最低要求;`refresh_token` NULL = 不可续期,到期即 `expired`。不进日志、不进列表 API。
|
||||
- **状态**:单列业务状态机 + 布尔/冷却补充列,避免状态爆炸。`refreshing` 不落 `accounts.status`,由 `account_refresh_leases` / 进程内 in-flight 派生为运行时展示态。冷却用时间而非布尔:自然过期无需回写。quota 与 CF 冷却分列:成因、时长、解除策略都不同。
|
||||
- **快照**:`quota_json` 上游原样快照(JSON 合规);`added_at` 语义化命名——账号是"被添加进池子"的。
|
||||
|
||||
### 4.7 account_refresh_leases
|
||||
|
||||
token 刷新互斥租约。租约(owner + 到期)是最小充分的互斥原语:持有者崩溃后自然失效,无需显式释放;天然支持未来多实例竞争。运行态表,不承载历史(需要时另建 `account_refresh_events`)。
|
||||
|
||||
```sql
|
||||
create table account_refresh_leases (
|
||||
account_id text primary key references accounts(id) on delete cascade,
|
||||
owner text not null,
|
||||
expires_at text not null,
|
||||
updated_at text not null
|
||||
);
|
||||
|
||||
create index idx_account_refresh_leases_expires on account_refresh_leases(expires_at);
|
||||
```
|
||||
|
||||
PK=FK 的 1:1 建模把"每账号至多一个持有者"变成主键约束。`owner` 在续约/释放时校验,防误释放他人租约。
|
||||
|
||||
### 4.8 account_usage
|
||||
|
||||
账号维度两组计数:① 生命周期累计(列表展示,**不可重建**,§2.9);② 当前 quota 窗口计数(调度限流输入,运行态)。1:1 派生表,不与 accounts 合并:更新频率差多个数量级,合表会让高频计数写污染低频实体行。
|
||||
|
||||
```sql
|
||||
create table account_usage (
|
||||
account_id text primary key references accounts(id) on delete cascade,
|
||||
request_count integer not null default 0 check (request_count >= 0),
|
||||
empty_response_count integer not null default 0 check (empty_response_count >= 0),
|
||||
input_tokens integer not null default 0 check (input_tokens >= 0),
|
||||
output_tokens integer not null default 0 check (output_tokens >= 0),
|
||||
cached_tokens integer not null default 0 check (cached_tokens >= 0),
|
||||
reasoning_tokens integer not null default 0 check (reasoning_tokens >= 0),
|
||||
total_tokens integer not null default 0 check (total_tokens >= 0),
|
||||
image_input_tokens integer not null default 0 check (image_input_tokens >= 0),
|
||||
image_output_tokens integer not null default 0 check (image_output_tokens >= 0),
|
||||
image_request_count integer not null default 0 check (image_request_count >= 0),
|
||||
image_request_failed_count integer not null default 0 check (image_request_failed_count >= 0),
|
||||
window_request_count integer not null default 0 check (window_request_count >= 0),
|
||||
window_input_tokens integer not null default 0 check (window_input_tokens >= 0),
|
||||
window_output_tokens integer not null default 0 check (window_output_tokens >= 0),
|
||||
window_cached_tokens integer not null default 0 check (window_cached_tokens >= 0),
|
||||
window_image_input_tokens integer not null default 0 check (window_image_input_tokens >= 0),
|
||||
window_image_output_tokens integer not null default 0 check (window_image_output_tokens >= 0),
|
||||
window_image_request_count integer not null default 0 check (window_image_request_count >= 0),
|
||||
window_image_request_failed_count integer not null default 0 check (window_image_request_failed_count >= 0),
|
||||
window_started_at text,
|
||||
window_reset_at text,
|
||||
limit_window_seconds integer check (limit_window_seconds is null or limit_window_seconds > 0),
|
||||
last_used_at text
|
||||
);
|
||||
|
||||
create index idx_account_usage_last_used_account
|
||||
on account_usage(last_used_at desc, account_id desc);
|
||||
```
|
||||
|
||||
| 组 | 说明 |
|
||||
| --- | --- |
|
||||
| 累计 | 口径 = usage_records 成功事实;失败率不得由本表推断。`total_tokens` 是已知的"有偿冗余"(列表排序高频),与分量列同一 upsert 内更新,事务保证一致 |
|
||||
| 窗口 | 当前 quota 窗口内计数,窗口重置清零。**窗口组刻意不含 `reasoning_tokens`/`total_tokens`**:窗口列只镜像上游 quota 核算维度,reasoning 不参与上游窗口计数;累计组是展示口径,两组不对称是设计而非遗漏 |
|
||||
| 边界 | `window_started_at` / `window_reset_at` / `limit_window_seconds` 来自上游 rate-limit 响应,NULL = 未知 |
|
||||
|
||||
宽表(20 计数列)而非 EAV `(account_id, metric, value)`:计数种类是代码常量,宽表一次 upsert 全更新、check 逐列可写;EAV 只换来动态 SQL 和类型丢失。
|
||||
|
||||
### 4.9 account_model_usage
|
||||
|
||||
账号 × 模型分布,服务账号详情页与调度辅助。自然复合键,不引入代理 id。
|
||||
|
||||
```sql
|
||||
create table account_model_usage (
|
||||
account_id text not null references accounts(id) on delete cascade,
|
||||
model text not null,
|
||||
request_count integer not null default 0 check (request_count >= 0),
|
||||
error_count integer not null default 0 check (error_count >= 0),
|
||||
input_tokens integer not null default 0 check (input_tokens >= 0),
|
||||
output_tokens integer not null default 0 check (output_tokens >= 0),
|
||||
cached_tokens integer not null default 0 check (cached_tokens >= 0),
|
||||
last_used_at text,
|
||||
primary key (account_id, model)
|
||||
);
|
||||
|
||||
create index idx_account_model_usage_last_used
|
||||
on account_model_usage(last_used_at desc, account_id, model);
|
||||
```
|
||||
|
||||
`error_count` 是**收窄口径**:仅统计已归属到账号+模型之后的失败,用作调度回避信号(频繁失败 → 降权)。这是全库唯一允许的口径重叠(与 ops_error_logs),因为用途不同(调度局部信号 vs 全局错误事实);Dashboard 错误率禁止取自本列。
|
||||
|
||||
### 4.10 usage_records ★
|
||||
|
||||
**每个成功的客户端请求恰好一行**。用量、token、成本、账号与调用方归因的唯一事实来源。这是系统的账本;"只收成功"的边界由 DB 约束兜底,不再只靠 service 层。
|
||||
|
||||
```sql
|
||||
create table usage_records (
|
||||
id text primary key,
|
||||
request_id text,
|
||||
client_api_key_id text,
|
||||
kind text not null,
|
||||
route text,
|
||||
provider text not null,
|
||||
account_id text not null,
|
||||
model text not null,
|
||||
requested_model text,
|
||||
upstream_model text,
|
||||
service_tier text,
|
||||
status_code integer not null check (status_code between 200 and 399),
|
||||
transport text,
|
||||
attempt_index integer check (attempt_index is null or attempt_index >= 0),
|
||||
response_id text,
|
||||
upstream_request_id text,
|
||||
latency_ms integer check (latency_ms is null or latency_ms >= 0),
|
||||
first_token_ms integer check (first_token_ms is null or first_token_ms >= 0),
|
||||
input_tokens integer check (input_tokens is null or input_tokens >= 0),
|
||||
output_tokens integer check (output_tokens is null or output_tokens >= 0),
|
||||
cached_tokens integer check (cached_tokens is null or cached_tokens >= 0),
|
||||
reasoning_tokens integer check (reasoning_tokens is null or reasoning_tokens >= 0),
|
||||
message text not null,
|
||||
metadata_json text not null,
|
||||
created_at text not null
|
||||
);
|
||||
|
||||
create index idx_usage_records_created_id on usage_records(created_at desc, id desc);
|
||||
create index idx_usage_records_request_id on usage_records(request_id) where request_id is not null;
|
||||
create index idx_usage_records_kind_created on usage_records(kind, created_at desc);
|
||||
create index idx_usage_records_account_created on usage_records(account_id, created_at desc);
|
||||
create index idx_usage_records_model_created on usage_records(model, created_at desc);
|
||||
create index idx_usage_records_key_created on usage_records(client_api_key_id, created_at desc) where client_api_key_id is not null;
|
||||
create index idx_usage_records_response_id on usage_records(response_id) where response_id is not null;
|
||||
create index idx_usage_records_upstream_request_id on usage_records(upstream_request_id) where upstream_request_id is not null;
|
||||
```
|
||||
|
||||
| 字段 | 说明 |
|
||||
| --- | --- |
|
||||
| `request_id` | 请求链路 ID,与 ops 表同域——跨两表串起一次请求的完整故事(成功前的失败尝试在 ops 表) |
|
||||
| `client_api_key_id` | **调用方归因**(逻辑引用 client_api_keys.id,无 FK)。可空:管理端代发/内部探测请求无 key。归因列是事实字段,错过不可回补——按 key 用量、未来按 key 配额核算都依赖它 |
|
||||
| `kind` / `route` | kind 是稳定事件族(`v1.response`),route 是原始 HTTP 事实;route 变体收敛到同一 kind,历史统计不断裂。映射只允许存在于 `event_kind()` 一处 |
|
||||
| `provider` | **上游 provider 归属**(`'openai'`,规划中 `'cloudflare'`)。not null:成功请求必知上游。事实必须自包含——账号可删除、模型别名可重映射,provider 不能靠 join 推导。开放取值不加 check(SQLite 改 check 要重建大表,新 provider 只应是新增取值)。低基数列刻意不建索引 |
|
||||
| `account_id` | **not null**:成功请求必有归属。无 FK(§3) |
|
||||
| `model` / `requested_model` / `upstream_model` | 展示归因(计费默认)/ 客户端原始请求(别名映射前)/ 上游实际执行。三列语义互不混用 |
|
||||
| `service_tier` | 计费单价第二维度 |
|
||||
| `status_code` | **`between 200 and 399`**:终态响应码 + 只收成功,数据库层兜底。1xx 非终态,4xx/5xx 属 ops 表 |
|
||||
| `transport` / `attempt_index` | 传输方式与成功发生在第几次尝试(>0 意味着此前有失败,记录在 ops 表) |
|
||||
| `latency_ms` / `first_token_ms` | 端到端完成延迟 / 首 token 延迟。NULL = 未测得,区别于 0 |
|
||||
| `input/output/cached/reasoning_tokens` | 一等列(§2.7)。NULL = 上游未报告,与 0 语义不同——因此不用 `not null default 0` |
|
||||
| `message` | 一行人类可读摘要,必须脱敏 |
|
||||
| `metadata_json` | 仅调试细节;已提升为列的字段不再写入 metadata |
|
||||
|
||||
8 个索引全部对应真实查询:keyset 分页 / 链路排查 / 事件族过滤 / 账号维度 / 模型维度 / 调用方维度 / response 续写排查 / 上游报障对账。每行 8 索引是全库最高写放大,是明细可查性的直接代价;量级失控的出路是缩短保留期(配置已就位),不是删索引。
|
||||
|
||||
**禁止**:写入失败事件;token 写回 metadata 当真相;加任何兼容过滤查询历史数据。
|
||||
|
||||
### 4.11 ops_error_logs ★
|
||||
|
||||
**每个失败请求/运维错误恰好一行**。与 usage_records 互斥:一次请求终态只进一张表。失败与成功的查询维度完全不同(阶段/分类/归责 vs token/成本),拆表让两边 schema 各自朝自己的查询模型演进。
|
||||
|
||||
```sql
|
||||
create table ops_error_logs (
|
||||
id text primary key,
|
||||
request_id text,
|
||||
client_api_key_id text,
|
||||
kind text not null,
|
||||
provider text,
|
||||
account_id text,
|
||||
route text,
|
||||
model text,
|
||||
status_code integer check (status_code is null or (status_code between 100 and 599)),
|
||||
client_status_code integer check (client_status_code is null or (client_status_code between 100 and 599)),
|
||||
upstream_status_code integer check (upstream_status_code is null or (upstream_status_code between 100 and 599)),
|
||||
transport text,
|
||||
attempt_index integer check (attempt_index is null or attempt_index >= 0),
|
||||
failure_class text,
|
||||
response_id text,
|
||||
upstream_request_id text,
|
||||
latency_ms integer check (latency_ms is null or latency_ms >= 0),
|
||||
message text not null,
|
||||
metadata_json text not null,
|
||||
created_at text not null
|
||||
);
|
||||
|
||||
create index idx_ops_error_logs_created_id on ops_error_logs(created_at desc, id desc);
|
||||
create index idx_ops_error_logs_request_id on ops_error_logs(request_id) where request_id is not null;
|
||||
create index idx_ops_error_logs_key_created on ops_error_logs(client_api_key_id, created_at desc) where client_api_key_id is not null;
|
||||
create index idx_ops_error_logs_account on ops_error_logs(account_id, created_at desc) where account_id is not null;
|
||||
create index idx_ops_error_logs_route_created on ops_error_logs(route, created_at desc) where route is not null;
|
||||
create index idx_ops_error_logs_model_created on ops_error_logs(model, created_at desc) where model is not null;
|
||||
create index idx_ops_error_logs_status_created on ops_error_logs(status_code, created_at desc) where status_code is not null;
|
||||
create index idx_ops_error_logs_transport_created on ops_error_logs(transport, created_at desc) where transport is not null;
|
||||
create index idx_ops_error_logs_failure_class on ops_error_logs(failure_class, created_at desc) where failure_class is not null;
|
||||
create index idx_ops_error_logs_response_id on ops_error_logs(response_id) where response_id is not null;
|
||||
create index idx_ops_error_logs_upstream_request_id on ops_error_logs(upstream_request_id) where upstream_request_id is not null;
|
||||
```
|
||||
|
||||
| 字段 | 说明 |
|
||||
| --- | --- |
|
||||
| `client_api_key_id` | 调用方归因,可空(鉴权前失败、后台任务无 key) |
|
||||
| `provider` | 上游 provider 归属,可空:模型解析前失败(鉴权失败、路由失败)时上游未定。与 `account_id` 同为"失败可以不知道,成功必须知道"的口径差异 |
|
||||
| `account_id` | **可空是本表的关键差异**:调度前失败(无可用账号、模型不可用)天然无归属——这正是这类事件进不了 usage_records 的原因 |
|
||||
| 三视角状态码 | `status_code`(事件主视角)/ `client_status_code`(用户看到什么)/ `upstream_status_code`(上游返回什么)。错误排查的第一个问题永远是"哪一层出的错" |
|
||||
| `failure_class` | 开放取值不加 check:分类法仍在演进 |
|
||||
| `metadata_json` | **禁止**存长请求/响应体(capture body 开启也必须截断脱敏);**禁止**存计费 token——失败不是计费事实,将来统计失败中观测到的 token 必须用 `observed_` 前缀新列 |
|
||||
|
||||
11 个索引全部对应错误排查的真实切片;可空维度全部部分索引;错误表写入频率远低于成功表,写放大可承受。
|
||||
|
||||
**预留演进列(有 UI/告警需求时增列,刻意不预建,本清单为唯一权威版本)**:`surface`(client_api/admin/system)、`error_phase`(auth/dispatch/upstream/stream/parse/internal)、`error_type`(稳定错误类型)、`severity`(P1–P3)、`requested_model` / `upstream_model` / `service_tier`、`client_ip` / `user_agent`、`error_source` / `error_owner`、`is_retryable` / `retry_count`。原则:字段跟着真实查询需求走,不照搬外部项目全集。
|
||||
|
||||
### 4.12 request_time_buckets ★
|
||||
|
||||
15 分钟 × provider × 账号 × 模型 × service_tier 预聚合,Dashboard 趋势与流量卡片的唯一数据源。**单表同时服务两种口径**:traffic(success + error)与 usage(token,仅成功)。曾有双表方案,已否决:每请求两次 upsert 的写放大不值得;分列 + 写入纪律已同时提供两种口径,rebuild 命令兜底(保留期内,§2.9)。
|
||||
|
||||
```sql
|
||||
create table request_time_buckets (
|
||||
bucket_start text not null,
|
||||
provider text not null default '__unknown__',
|
||||
account_id text not null default '__unknown__',
|
||||
model text not null default '__unknown__',
|
||||
service_tier text not null default '__unknown__',
|
||||
success_count integer not null default 0 check (success_count >= 0),
|
||||
error_count integer not null default 0 check (error_count >= 0),
|
||||
input_tokens integer not null default 0 check (input_tokens >= 0),
|
||||
output_tokens integer not null default 0 check (output_tokens >= 0),
|
||||
cached_tokens integer not null default 0 check (cached_tokens >= 0),
|
||||
first_token_latency_sum integer not null default 0 check (first_token_latency_sum >= 0),
|
||||
first_token_latency_count integer not null default 0 check (first_token_latency_count >= 0),
|
||||
latency_sum integer not null default 0 check (latency_sum >= 0),
|
||||
latency_count integer not null default 0 check (latency_count >= 0),
|
||||
max_latency_ms integer not null default 0 check (max_latency_ms >= 0),
|
||||
min_latency_ms integer check (min_latency_ms is null or min_latency_ms >= 0),
|
||||
updated_at text not null,
|
||||
primary key (bucket_start, provider, account_id, model, service_tier)
|
||||
);
|
||||
|
||||
create index idx_request_time_buckets_model on request_time_buckets(model, bucket_start);
|
||||
```
|
||||
|
||||
| 字段 | 说明 |
|
||||
| --- | --- |
|
||||
| `bucket_start` | UTC 时刻,按中国时区 15 分钟槽对齐(`china_quarter_hour_start`)。对齐规则只活在代码层 |
|
||||
| 维度四列(provider / account_id / model / service_tier) | 不可空 + `__unknown__` sentinel(§2.5)。provider 由两条写入路径显式给值(成功/失败路径都知道自己在打哪家),`__unknown__` 仅兜底 |
|
||||
| `success_count` / `error_count` | 分别由 usage / ops 写入路径 +1。**`request_count` 不落库**,恒等于两者之和,查询期推导——存储恒等式必然漂移 |
|
||||
| token 三列 | **仅成功路径写入**;错误路径只 +1 `error_count`,不得触碰 token 与延迟列 |
|
||||
| latency sum/count 对 | 平均值不可再聚合,sum/count 才能跨桶合并、跨维度上卷。**口径 = 仅成功**(失败延迟混入会让"变慢"和"在报错"互相污染) |
|
||||
| `max_latency_ms` / `min_latency_ms` | 可单调合并的极值。min 可空:NULL = 无样本,`min()` 天然忽略 NULL。**check 是 `>= 0`**——0ms 是合法样本(缓存命中/本地短路),与 `latency_ms >= 0` 的合法域一致;`> 0` 会让 0ms 样本在同事务写入中炸掉整笔成功事实 |
|
||||
|
||||
**写入事务纪律**:事实 insert 与桶 upsert 必须同事务(§5.2)。
|
||||
|
||||
**禁止**:错误路径写 token/延迟列;落库 `request_count`;空字符串维度。
|
||||
|
||||
### 4.13 account_cookies
|
||||
|
||||
账号维度上游 Cookie(部分上游流程需要浏览器态凭据)。语义直接采用 RFC 6265 模型;unique 约束 = RFC 替换语义,upsert 依据。
|
||||
|
||||
```sql
|
||||
create table account_cookies (
|
||||
id text primary key,
|
||||
account_id text not null references accounts(id) on delete cascade,
|
||||
domain text not null,
|
||||
name text not null,
|
||||
value text not null,
|
||||
path text not null default '/',
|
||||
expires_at text,
|
||||
updated_at text not null,
|
||||
unique(account_id, domain, name, path)
|
||||
);
|
||||
|
||||
create index idx_account_cookies_account_domain on account_cookies(account_id, domain);
|
||||
create index idx_account_cookies_expires on account_cookies(expires_at) where expires_at is not null;
|
||||
```
|
||||
|
||||
`value` 是 secret(§2.6)。`expires_at` NULL = 会话 Cookie,不参与清理。v3 单独的 `idx_account_cookies_account` 被 `account_domain` 前缀覆盖,终态不保留。
|
||||
|
||||
### 4.14 fingerprints
|
||||
|
||||
设备指纹模板(UA、header 集合与顺序——顺序本身是指纹的一部分)。行数常量级、按 PK 取用,**无二级索引**。版本三元组用 TEXT:版本号不是数字(`1.2.3-beta`)。
|
||||
|
||||
```sql
|
||||
create table fingerprints (
|
||||
id text primary key,
|
||||
originator text not null,
|
||||
app_version text not null,
|
||||
build_number text not null,
|
||||
platform text not null,
|
||||
arch text not null,
|
||||
chromium_version text not null,
|
||||
user_agent_template text not null,
|
||||
default_headers_json text not null,
|
||||
header_order_json text not null,
|
||||
source text not null,
|
||||
created_at text not null,
|
||||
updated_at text not null
|
||||
);
|
||||
```
|
||||
|
||||
### 4.15 fingerprint_update_history
|
||||
|
||||
指纹更新审计。版本快照列固化"当时更新到了什么"——不能只存 FK,指纹行会被后续更新覆写。级联删除可接受:指纹本体删除后其历史无独立价值。
|
||||
|
||||
```sql
|
||||
create table fingerprint_update_history (
|
||||
id text primary key,
|
||||
current_fingerprint_id text not null references fingerprints(id) on delete cascade,
|
||||
app_version text not null,
|
||||
build_number text not null,
|
||||
chromium_version text,
|
||||
source text not null,
|
||||
manifest_json text,
|
||||
created_at text not null
|
||||
);
|
||||
|
||||
create index idx_fingerprint_update_history_created_id
|
||||
on fingerprint_update_history(created_at desc, id desc);
|
||||
create index idx_fingerprint_update_history_fingerprint
|
||||
on fingerprint_update_history(current_fingerprint_id);
|
||||
```
|
||||
|
||||
### 4.16 model_plan_snapshots
|
||||
|
||||
套餐 → 可用模型清单缓存。刻意"最新态覆盖"而非追加历史——缓存表的正确形态,可随时清空重建。**禁止**当模型目录事实表用;需要 diff/审计时新建历史表。
|
||||
|
||||
```sql
|
||||
create table model_plan_snapshots (
|
||||
plan_type text primary key,
|
||||
models_json text not null,
|
||||
fetched_at text not null
|
||||
);
|
||||
```
|
||||
|
||||
### 4.17 session_affinities
|
||||
|
||||
`response_id` → 账号/会话上下文映射,保证续写请求路由回同一账号(上游 response 链只在原账号内有效)。纯运行态:全表清空只影响路由优化,不损失事实。**禁止**承载会话历史/审计。
|
||||
|
||||
```sql
|
||||
create table session_affinities (
|
||||
response_id text primary key,
|
||||
account_id text not null references accounts(id) on delete cascade,
|
||||
conversation_id text not null,
|
||||
turn_state text,
|
||||
instructions_hash text,
|
||||
input_tokens integer check (input_tokens is null or input_tokens >= 0),
|
||||
function_call_ids_json text not null default '[]',
|
||||
variant_hash text,
|
||||
expires_at text not null,
|
||||
created_at text not null
|
||||
);
|
||||
|
||||
create index idx_session_affinities_conversation on session_affinities(conversation_id, created_at desc);
|
||||
create index idx_session_affinities_active_order on session_affinities(expires_at, created_at, response_id);
|
||||
create index idx_session_affinities_account on session_affinities(account_id);
|
||||
```
|
||||
|
||||
`instructions_hash` / `variant_hash` 共同约束隐式续写安全(指令或参数变了就不能隐式续写);存哈希不存原文。`active_order` 同时服务过期清理(前缀)与活跃排序;v3 单独的 `expires` 索引是其前缀,终态删除。`account` 索引是 FK 子表索引(级联删除路径)。
|
||||
|
||||
---
|
||||
|
||||
## 5. 统计口径(全系统唯一定义)
|
||||
|
||||
### 5.1 口径恒等式
|
||||
|
||||
任何 UI 卡片不得自创口径:
|
||||
|
||||
```text
|
||||
traffic: request_count = success_count + error_count ← request_time_buckets(查询期推导)
|
||||
usage: token / 成本 ← usage_records 列(明细)、bucket token 列(趋势,仅成功)
|
||||
errors: 错误分布 / 排查 ← ops_error_logs
|
||||
按 key: 调用方用量 / 失败 ← usage_records / ops_error_logs 的 client_api_key_id
|
||||
调度负载: account_usage / account_model_usage ← 派生累计,事务保护,不可重建(§2.9)
|
||||
```
|
||||
|
||||
### 5.2 写入路径与事务
|
||||
|
||||
```text
|
||||
请求成功 ─┬─ usage_records insert
|
||||
├─ request_time_buckets upsert(success_count+1,token、延迟列)
|
||||
├─ account_usage upsert(累计 + 窗口)
|
||||
└─ account_model_usage upsert
|
||||
全部在同一个事务(pool.begin() … commit)
|
||||
|
||||
请求失败 ─┬─ ops_error_logs insert
|
||||
└─ request_time_buckets upsert(error_count+1,仅此一列)
|
||||
同一个事务
|
||||
```
|
||||
|
||||
### 5.3 Dashboard 口径决策(钉死)
|
||||
|
||||
- 所有 traffic 卡片(今日请求、区间请求、错误率、QPS 趋势)**一律来自 request_time_buckets**,语义为"bucket 保留期内"。卡片文案如实标注区间,**不提供也不暗示"历史总量"**——bucket 有 90 天保留期,任何"total"字样都是口径谎言。
|
||||
- 生命周期累计仅有一处合法来源:`account_usage` 累计列(成功口径),使用时必须标注"成功累计"。
|
||||
- 同一张卡片禁止混排两种口径的数字(v3 的 `todayRequests` 来自 bucket、`totalRequests` 来自 usage summary 就是这个病)。
|
||||
- 真正的长期趋势需求 → 日粒度归档表 `request_day_buckets`(§8),不是延长 15 分钟桶保留期。
|
||||
|
||||
---
|
||||
|
||||
## 6. 迁移 0004:一次到位(v3 → v4)
|
||||
|
||||
### 6.1 前置:迁移框架支持 Rust 步骤
|
||||
|
||||
SQL 无法计算 SHA-256,0004 需要一个事务内的 Rust 钩子。改造 `database.rs`:
|
||||
|
||||
```rust
|
||||
struct Migration {
|
||||
version: i64,
|
||||
name: &'static str,
|
||||
sql: &'static str,
|
||||
/// SQL 之后、同一事务内执行的 Rust 步骤(哈希回填等 SQL 做不到的事)。
|
||||
post: Option<fn(&mut Transaction<'_, Sqlite>) -> BoxFuture<'_, Result<(), sqlx::Error>>>,
|
||||
}
|
||||
```
|
||||
|
||||
应用顺序不变:`raw_sql(migration.sql)` → `post`(若有)→ `record_migration`,全程同一事务。当前 `CURRENT_SCHEMA_VERSION = 3`;本节 0004 为目标迁移草案,落地时再提升版本。
|
||||
|
||||
### 6.2 0004_final_schema.sql(SQL 部分)
|
||||
|
||||
```sql
|
||||
-- ===== admin_sessions:清空重建(哈希键) =====
|
||||
-- 数据决策:现存会话全部作废,管理员重新登录。会话是运行态,无保留价值。
|
||||
drop table admin_sessions;
|
||||
create table admin_sessions ( … §4.3 DDL … );
|
||||
|
||||
-- ===== client_api_keys:改名暂存,Rust 步骤哈希回填 =====
|
||||
alter table client_api_keys rename to client_api_keys_v3;
|
||||
drop index if exists idx_client_api_keys_key_enabled;
|
||||
drop index if exists idx_client_api_keys_created_id;
|
||||
create table client_api_keys ( … §4.4 DDL … );
|
||||
create index idx_client_api_keys_created_id on client_api_keys(created_at desc, id desc);
|
||||
|
||||
-- ===== runtime_settings:重建(admin_api_key → 暂存,Rust 哈希) =====
|
||||
create table _migration_admin_key as select admin_api_key as value from runtime_settings;
|
||||
alter table runtime_settings rename to runtime_settings_v3;
|
||||
create table runtime_settings ( … §4.5 DDL … );
|
||||
insert into runtime_settings (id, model_aliases_json, refresh_margin_seconds, refresh_concurrency,
|
||||
max_concurrent_per_account, request_interval_ms, rotation_strategy, admin_api_key_hash, updated_at)
|
||||
select id, model_aliases_json, refresh_margin_seconds, refresh_concurrency,
|
||||
max_concurrent_per_account, request_interval_ms, rotation_strategy, null, updated_at
|
||||
from runtime_settings_v3;
|
||||
drop table runtime_settings_v3;
|
||||
|
||||
-- ===== ops_error_logs:先增列(下方 usage_records 迁移要写入这些列) =====
|
||||
alter table ops_error_logs add column client_api_key_id text;
|
||||
alter table ops_error_logs add column provider text;
|
||||
update ops_error_logs set provider = 'openai';
|
||||
create index idx_ops_error_logs_key_created on ops_error_logs(client_api_key_id, created_at desc) where client_api_key_id is not null;
|
||||
|
||||
-- ===== usage_records:重建为成功事实表 =====
|
||||
alter table usage_records rename to usage_records_v3;
|
||||
create table usage_records ( … §4.10 DDL 与全部索引 … );
|
||||
|
||||
-- 数据决策 1:满足成功契约的行迁入,token 等列从 metadata 提升(此后不再回读 metadata)。
|
||||
insert into usage_records (
|
||||
id, request_id, client_api_key_id, kind, route, provider, account_id, model,
|
||||
requested_model, upstream_model, service_tier, status_code, transport, attempt_index,
|
||||
response_id, upstream_request_id, latency_ms, first_token_ms,
|
||||
input_tokens, output_tokens, cached_tokens, reasoning_tokens,
|
||||
message, metadata_json, created_at)
|
||||
select
|
||||
id, request_id, null, kind, route, 'openai', account_id, model,
|
||||
nullif(trim(json_extract(metadata_json, '$.requestedModel')), ''),
|
||||
nullif(trim(json_extract(metadata_json, '$.upstreamModel')), ''),
|
||||
nullif(trim(json_extract(metadata_json, '$.serviceTier')), ''),
|
||||
status_code, transport, attempt_index,
|
||||
response_id, upstream_request_id, latency_ms,
|
||||
coalesce(cast(json_extract(metadata_json, '$.firstTokenMs') as integer),
|
||||
cast(json_extract(metadata_json, '$.usage.firstTokenMs') as integer)),
|
||||
coalesce(cast(json_extract(metadata_json, '$.usage.inputTokens') as integer),
|
||||
cast(json_extract(metadata_json, '$.inputTokens') as integer)),
|
||||
coalesce(cast(json_extract(metadata_json, '$.usage.outputTokens') as integer),
|
||||
cast(json_extract(metadata_json, '$.outputTokens') as integer)),
|
||||
coalesce(cast(json_extract(metadata_json, '$.usage.cachedTokens') as integer),
|
||||
cast(json_extract(metadata_json, '$.cachedTokens') as integer)),
|
||||
coalesce(cast(json_extract(metadata_json, '$.usage.reasoningTokens') as integer),
|
||||
cast(json_extract(metadata_json, '$.reasoningTokens') as integer)),
|
||||
message, metadata_json, created_at
|
||||
from usage_records_v3
|
||||
where level <> 'error'
|
||||
and account_id is not null and trim(account_id) <> ''
|
||||
and model is not null and trim(model) <> ''
|
||||
and status_code between 200 and 399;
|
||||
|
||||
-- 数据决策 2:历史错误行迁入 ops_error_logs(列同构)。
|
||||
insert into ops_error_logs (
|
||||
id, request_id, client_api_key_id, kind, provider, account_id, route, model,
|
||||
status_code, client_status_code, upstream_status_code, transport, attempt_index,
|
||||
failure_class, response_id, upstream_request_id, latency_ms, message, metadata_json, created_at)
|
||||
select
|
||||
id, request_id, null, kind, 'openai', account_id, route, model,
|
||||
status_code, null, upstream_status_code, transport, attempt_index,
|
||||
failure_class, response_id, upstream_request_id, latency_ms, message, metadata_json, created_at
|
||||
from usage_records_v3
|
||||
where level = 'error';
|
||||
|
||||
-- 数据决策 3:既非成功契约也非 error 的行(无账号/无模型/无终态状态码的历史噪音)不迁移,随 v3 表一起删除。
|
||||
drop table usage_records_v3;
|
||||
|
||||
-- ===== usage_time_buckets → request_time_buckets =====
|
||||
create table request_time_buckets ( … §4.12 DDL … );
|
||||
create index idx_request_time_buckets_model on request_time_buckets(model, bucket_start);
|
||||
|
||||
-- 数据决策 4:历史桶迁入(sentinel 归一、success 推导、0-sentinel 转 NULL)。
|
||||
-- 历史桶的 token 列含 v3 错误路径污染,按旧口径保留展示;事实保留期内的窗口
|
||||
-- 随后由 rebuild-derived 用新口径重算覆盖(§6.4 验收第 4 步)。
|
||||
insert into request_time_buckets (
|
||||
bucket_start, provider, account_id, model, service_tier,
|
||||
success_count, error_count, input_tokens, output_tokens, cached_tokens,
|
||||
first_token_latency_sum, first_token_latency_count,
|
||||
latency_sum, latency_count, max_latency_ms, min_latency_ms, updated_at)
|
||||
select
|
||||
bucket_start,
|
||||
'openai',
|
||||
case when account_id = '' then '__unknown__' else account_id end,
|
||||
case when model = '' then '__unknown__' else model end,
|
||||
case when service_tier = '' then '__unknown__' else service_tier end,
|
||||
max(request_count - error_count, 0), error_count,
|
||||
input_tokens, output_tokens, cached_tokens,
|
||||
first_token_latency_sum, first_token_latency_count,
|
||||
latency_sum, latency_count, max_latency_ms,
|
||||
nullif(min_latency_ms, 0), updated_at
|
||||
from usage_time_buckets;
|
||||
|
||||
drop table usage_time_buckets;
|
||||
|
||||
-- ===== 索引修正(其余表) =====
|
||||
drop index if exists idx_session_affinities_expires;
|
||||
create index idx_session_affinities_account on session_affinities(account_id);
|
||||
create index idx_fingerprint_update_history_fingerprint on fingerprint_update_history(current_fingerprint_id);
|
||||
drop index if exists idx_account_cookies_account;
|
||||
```
|
||||
|
||||
### 6.3 0004 的 Rust 步骤(`post` 钩子,同一事务)
|
||||
|
||||
1. **client key 哈希搬迁**:读 `client_api_keys_v3` 全部行 → 每行计算 `SHA-256(key)` → insert 进新表(`key_hash` = 哈希,其余列原样)→ `drop table client_api_keys_v3`。明文自此不存在于库中。
|
||||
2. **admin API key 哈希**:读 `_migration_admin_key.value`,非空则 `update runtime_settings set admin_api_key_hash = <sha256>`;`drop table _migration_admin_key`。
|
||||
|
||||
### 6.4 迁移后动作与验收
|
||||
|
||||
| # | 动作 | 验收 |
|
||||
| --- | --- | --- |
|
||||
| 1 | 事务提交后执行一次 `vacuum`(使 `auto_vacuum = incremental` 对旧库生效;vacuum 不能在事务内) | `pragma auto_vacuum` 返回 2 |
|
||||
| 2 | 管理员重新登录(会话已清空) | 登录成功,新会话 `token_hash` 长度 64 |
|
||||
| 3 | 抽查 `select count(*) from client_api_keys where length(key_hash) <> 64` = 0 | 明文 key 消失,鉴权仍通过 |
|
||||
| 4 | 执行 `rebuild-derived`(§7.6),用新口径重算事实保留期内的桶 | 保留期内桶 token 列 = usage_records 聚合值 |
|
||||
| 5 | 核对 Dashboard 卡片全部走 bucket 口径 | 同卡片无混合口径数字 |
|
||||
|
||||
**显式数据损失清单**(接受即执行,不做挽留方案):全部管理端会话;v3 usage_records 中既非成功契约也非 error 的噪音行;历史桶中错误路径污染的 token 按旧口径保留、不追溯清洗(保留期内被 rebuild 覆盖为新口径)。
|
||||
|
||||
---
|
||||
|
||||
## 7. 代码改造清单
|
||||
|
||||
与 0004 同一 PR 交付,按依赖顺序:
|
||||
|
||||
1. **迁移框架**(`infra/database.rs`):`Migration.post` 钩子;连接参数补 `synchronous(Normal)`、`auto_vacuum(Incremental)`;目标迁移落地时再提升 `CURRENT_SCHEMA_VERSION`。
|
||||
2. **鉴权归因**(`proxy/auth.rs`):`authorize_client_api_key_result` 改为返回 key 的 `id`(现在验完即弃);SHA-256 后查 `key_hash`。调用链把 `client_api_key_id` 装进请求上下文。
|
||||
3. **事件结构**(`proxy/dispatch/usage_events.rs`、`proxy/dispatch/responses/event_recording.rs`):成功/失败事件携带 `client_api_key_id`;token、requested/upstream_model、service_tier、first_token_ms 从 metadata 移到一等字段;metadata 不再写已提升字段。
|
||||
4. **写入路径**(`admin/monitoring/usage_record_store.rs`、ops error store):§5.2 的两条事务;错误路径对桶只 `error_count + 1`;`__unknown__` sentinel;min/max 延迟直接 `min()`/`max()`(无 0-sentinel 分支)。
|
||||
5. **查询路径**(`usage_record_service.rs`、`dashboard.rs`、`account_usage_service.rs`):summary/分布/趋势全部走列,删除所有 `json_extract` 聚合与 `like '%…%'` 搜索;Dashboard 按 §5.3 口径改卡片来源与文案。
|
||||
6. **维护命令**:管理端 `rebuild-derived`——删除事实保留期内的桶并从两张事实表重算;保留期外只读不动。范围明确**不含** `account_usage` / `account_model_usage`(§2.9,不可重建)。
|
||||
7. **清理任务**:三个 trim 任务改读 `runtime_settings` 三列;新增 bucket trim;每轮 trim 后 `pragma incremental_vacuum`。
|
||||
8. **管理端 API**(`admin/auth/service.rs`、`config/settings.rs`、keys 路由):session 创建/校验走 `token_hash`;`verify_admin_api_key` 对来键哈希后 `ct_eq`;key 列表响应删除完整 key 字段(只余 `prefix`)——**前端同步移除完整 key 展示**。
|
||||
9. **测试**(`backend/tests/infra/storage_schema/`):断言终态 DDL、迁移数据规则(成功行迁入 / error 行转 ops / 噪音行丢弃)、两条写入事务的原子性、0ms 延迟样本可入桶。
|
||||
|
||||
---
|
||||
|
||||
## 8. 扩展路径与明确不做
|
||||
|
||||
每个可预见的增长方向都有"只加不改"的路径:
|
||||
|
||||
| 未来需求 | 扩展方式 |
|
||||
| --- | --- |
|
||||
| 多管理员 | `admin_users` 加 `username`/`role` + 新表 `admin_login_events` |
|
||||
| 按 key 限流/配额 | 新表 `client_api_key_policies (key_id FK, …)`;归因数据已在事实表就位 |
|
||||
| 错误运维大盘 | `ops_error_logs` 按 §4.11 预留清单增列 |
|
||||
| 长期趋势/年度账务 | 新表 `request_day_buckets`(由 15 分钟桶降采样归档),Dashboard 才允许出现"历史总量" |
|
||||
| 刷新历史分析 | 新表 `account_refresh_events` |
|
||||
| 模型别名审计 | `model_aliases_json` 拆 `model_aliases` 表 |
|
||||
| 凭据静态加密 / 合规 | `accounts` 拆 `account_credentials`(1:1,PK=FK),集中一个破坏窗口 |
|
||||
| 接入第二上游 provider(已规划:Cloudflare Workers AI) | 事实表与桶的 `provider` 维度已在 v4 就位(§4.10–4.12),CF 首条请求落库前无需再动事实表。届时一个迁移完成:`accounts` 加 `provider` + `chatgpt_account_id/user_id` 更名 `upstream_account_id/upstream_user_id` + 唯一索引加 provider 维度;`model_plan_snapshots` PK 改 `(provider, plan_type)`;billing 单价查找加 provider 参数。回填全部确定为 `'openai'` |
|
||||
| 多实例部署 | SQLite → Postgres:TEXT 时间戳 → timestamptz、`check in` → enum、upsert 语法兼容;租约表天然支持竞争 |
|
||||
|
||||
Cloudflare 账号对现有 schema 的适配(无需新列):静态 API key = `access_token`,`refresh_token`/`access_token_expires_at` NULL(永不过期,语义已覆盖);无套餐配额窗口 → `account_usage.window_*` 恒为初始值、窗口边界 NULL;`account_cookies` 对 CF 账号空置。**硬规则**:任何新 provider 的第一条请求落库之前,其事实行必须已携带 provider 归因——v4 之后这条自动满足。
|
||||
|
||||
**明确不做**:无多实例需求前不引入外部数据库/缓存;不为历史数据加任何兼容过滤;不预建"可能有用"的列——每一列都必须有当前的真实读者;不照搬外部项目的字段全集。
|
||||
@@ -19,6 +19,7 @@ export interface SystemUpdateLog {
|
||||
level: SystemUpdateLogLevel
|
||||
step?: string
|
||||
message: string
|
||||
terminal?: boolean
|
||||
at: string
|
||||
}
|
||||
|
||||
@@ -84,8 +85,12 @@ function clearUpdateLogs() {
|
||||
updateStreamError.value = ''
|
||||
}
|
||||
|
||||
function connectUpdateEvents() {
|
||||
if (updateEventSource.value) return
|
||||
function connectUpdateEvents(options: { force?: boolean } = {}) {
|
||||
if (options.force) {
|
||||
disconnectUpdateEvents()
|
||||
} else if (updateEventSource.value) {
|
||||
return
|
||||
}
|
||||
|
||||
if (typeof EventSource === 'undefined') {
|
||||
updateStreamError.value = '当前浏览器不支持实时更新日志'
|
||||
@@ -97,6 +102,8 @@ function connectUpdateEvents() {
|
||||
updateEventSource.value = source
|
||||
|
||||
source.onopen = () => {
|
||||
if (updateEventSource.value !== source) return
|
||||
|
||||
updateStreaming.value = true
|
||||
updateStreamError.value = ''
|
||||
}
|
||||
@@ -106,6 +113,9 @@ function connectUpdateEvents() {
|
||||
const log = JSON.parse((event as MessageEvent<string>).data)
|
||||
if (isUpdateLog(log)) {
|
||||
appendUpdateLog(log)
|
||||
if (log.terminal) {
|
||||
disconnectUpdateEvents(source)
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
updateStreamError.value = '更新日志解析失败'
|
||||
@@ -113,6 +123,8 @@ function connectUpdateEvents() {
|
||||
})
|
||||
|
||||
source.onerror = () => {
|
||||
if (updateEventSource.value !== source) return
|
||||
|
||||
updateStreaming.value = false
|
||||
updateStreamError.value = '更新日志连接中断'
|
||||
if (source.readyState === EventSource.CLOSED) {
|
||||
@@ -121,8 +133,10 @@ function connectUpdateEvents() {
|
||||
}
|
||||
}
|
||||
|
||||
function disconnectUpdateEvents() {
|
||||
updateEventSource.value?.close()
|
||||
function disconnectUpdateEvents(source: EventSource | null = updateEventSource.value) {
|
||||
source?.close()
|
||||
if (source && updateEventSource.value !== source) return
|
||||
|
||||
updateEventSource.value = null
|
||||
updateStreaming.value = false
|
||||
}
|
||||
@@ -184,7 +198,7 @@ async function updateNow(targetVersion: string) {
|
||||
if (!canUpdate.value || !currentInfo || updating.value || !confirmedTargetVersion) return null
|
||||
|
||||
clearUpdateLogs()
|
||||
connectUpdateEvents()
|
||||
connectUpdateEvents({ force: true })
|
||||
updating.value = true
|
||||
updateError.value = ''
|
||||
updateSuccess.value = false
|
||||
|
||||
@@ -286,7 +286,7 @@ watch(
|
||||
<button
|
||||
v-if="hasVersionLabel"
|
||||
type="button"
|
||||
class="inline-flex h-4.5 min-w-0 items-center gap-1 rounded-(--cp-input-radius-small) border-0 px-1.5 font-mono text-[10px] leading-none font-[720] cursor-pointer transition-colors"
|
||||
class="inline-flex h-4.5 min-w-0 cursor-pointer items-center gap-1 rounded-(--cp-input-radius-small) border-0 px-1.5 font-mono text-[10px] leading-none font-[720] transition-colors outline-none focus-visible:ring-2 focus-visible:ring-(--cp-info-border) focus-visible:ring-offset-2 focus-visible:ring-offset-(--cp-bg-surface)"
|
||||
:class="[
|
||||
hasUpdate
|
||||
? 'bg-(--cp-success-bg) text-(--cp-success-text) hover:bg-(--cp-success-bg-hover)'
|
||||
@@ -309,7 +309,7 @@ watch(
|
||||
v-for="item in navItems"
|
||||
:key="item.label"
|
||||
type="button"
|
||||
class="inline-flex h-11.5 items-center rounded-(--cp-icon-button-radius) text-sm leading-[1.15] border-0 cursor-pointer transition-colors duration-200"
|
||||
class="inline-flex h-11.5 cursor-pointer items-center rounded-(--cp-icon-button-radius) border-0 text-sm leading-[1.15] transition-colors duration-200 outline-none focus-visible:ring-2 focus-visible:ring-(--cp-info-border) focus-visible:ring-offset-2 focus-visible:ring-offset-(--cp-bg-surface)"
|
||||
:class="[
|
||||
isCollapsed ? 'w-11.5 justify-center' : 'w-full gap-3 px-4',
|
||||
isActive(item.path)
|
||||
|
||||
@@ -45,7 +45,6 @@ const {
|
||||
updateNow,
|
||||
restartNow,
|
||||
clearRestartTimer,
|
||||
connectUpdateEvents,
|
||||
disconnectUpdateEvents,
|
||||
} = useSystemUpdate()
|
||||
|
||||
@@ -54,7 +53,6 @@ const updateConfirmOpen = shallowRef(false)
|
||||
const updateConfirmInfo = shallowRef<SystemUpdateInfo | null>(null)
|
||||
const updateConfirmPreviousTarget = shallowRef('')
|
||||
const preparingUpdate = shallowRef(false)
|
||||
let updateLogPinTimer: number | undefined
|
||||
|
||||
const statusView = computed(() => {
|
||||
if (updating.value) {
|
||||
@@ -137,7 +135,6 @@ const renderedReleaseNotes = computed(() => renderMarkdown(updateInfo.value?.not
|
||||
const showUpdateProgress = computed(
|
||||
() => hasUpdate.value || updating.value || restarting.value || updateLogRows.value.length > 0,
|
||||
)
|
||||
const shouldConnectUpdateEvents = computed(() => (open.value && hasUpdate.value) || updating.value)
|
||||
|
||||
const updateConfirmRows = computed(() => [
|
||||
{
|
||||
@@ -212,20 +209,6 @@ function pinUpdateLogsToBottom() {
|
||||
})
|
||||
}
|
||||
|
||||
function startUpdateLogPinning() {
|
||||
if (updateLogPinTimer !== undefined) return
|
||||
|
||||
pinUpdateLogsToBottom()
|
||||
updateLogPinTimer = window.setInterval(pinUpdateLogsToBottom, 250)
|
||||
}
|
||||
|
||||
function stopUpdateLogPinning() {
|
||||
if (updateLogPinTimer === undefined) return
|
||||
|
||||
window.clearInterval(updateLogPinTimer)
|
||||
updateLogPinTimer = undefined
|
||||
}
|
||||
|
||||
async function handleCheckUpdates(force = true) {
|
||||
try {
|
||||
const data = await checkUpdates(force)
|
||||
@@ -304,45 +287,17 @@ watch(open, (visible) => {
|
||||
}
|
||||
})
|
||||
|
||||
watch(
|
||||
shouldConnectUpdateEvents,
|
||||
(shouldConnect) => {
|
||||
if (shouldConnect) {
|
||||
connectUpdateEvents()
|
||||
} else {
|
||||
disconnectUpdateEvents()
|
||||
}
|
||||
},
|
||||
{ immediate: true },
|
||||
)
|
||||
|
||||
watch(
|
||||
() => updateLogs.value.at(-1)?.id,
|
||||
() => {
|
||||
(logId, previousLogId) => {
|
||||
if (!logId || logId === previousLogId) return
|
||||
|
||||
pinUpdateLogsToBottom()
|
||||
},
|
||||
)
|
||||
|
||||
watch(
|
||||
() => updateLogRows.value.length,
|
||||
() => {
|
||||
pinUpdateLogsToBottom()
|
||||
},
|
||||
)
|
||||
|
||||
watch(
|
||||
() => updating.value || restarting.value || updateStreaming.value,
|
||||
(active) => {
|
||||
if (active) {
|
||||
startUpdateLogPinning()
|
||||
return
|
||||
}
|
||||
stopUpdateLogPinning()
|
||||
},
|
||||
{ flush: 'post' },
|
||||
)
|
||||
|
||||
onUnmounted(() => {
|
||||
stopUpdateLogPinning()
|
||||
clearRestartTimer()
|
||||
disconnectUpdateEvents()
|
||||
})
|
||||
|
||||
@@ -42,6 +42,29 @@ function quotaWindowBarClass(window?: any) {
|
||||
}
|
||||
return 'bg-(--cp-success)'
|
||||
}
|
||||
|
||||
function quotaWindowPercentTextClass(window?: any) {
|
||||
if (window?.usedPercent === null || window?.usedPercent === undefined) {
|
||||
return 'text-(--cp-text-muted)'
|
||||
}
|
||||
if (window.usedPercent >= 95) {
|
||||
return 'text-(--cp-danger-text)'
|
||||
}
|
||||
if (window.usedPercent >= 80) {
|
||||
return 'text-(--cp-warning-text)'
|
||||
}
|
||||
return 'text-(--cp-success-text)'
|
||||
}
|
||||
|
||||
function quotaWindowLocalUsageDisplay(window?: any) {
|
||||
const display = window?.localUsage?.totalTokensDisplay
|
||||
return typeof display === 'string' && display.trim() ? display : '-'
|
||||
}
|
||||
|
||||
function shouldShowQuotaWindowLocalUsage(window?: any) {
|
||||
const totalTokens = window?.localUsage?.totalTokens
|
||||
return typeof totalTokens === 'number' && totalTokens > 0
|
||||
}
|
||||
</script>
|
||||
|
||||
<template>
|
||||
@@ -80,7 +103,17 @@ function quotaWindowBarClass(window?: any) {
|
||||
>
|
||||
<div class="flex items-center justify-between gap-3 text-[12px] font-[720]">
|
||||
<span class="text-(--cp-text-secondary)">{{ window.labelDisplay }}</span>
|
||||
<span class="text-(--cp-text-primary)">{{ window.usedPercentDisplay }}</span>
|
||||
<span class="flex shrink-0 items-baseline justify-end gap-1.5 font-mono tabular-nums">
|
||||
<span
|
||||
v-if="shouldShowQuotaWindowLocalUsage(window)"
|
||||
class="text-[12px] font-[680] text-(--cp-text-muted)"
|
||||
>
|
||||
{{ quotaWindowLocalUsageDisplay(window) }}
|
||||
</span>
|
||||
<span class="text-[12px] font-[780]" :class="quotaWindowPercentTextClass(window)">
|
||||
{{ window.usedPercentDisplay }}
|
||||
</span>
|
||||
</span>
|
||||
</div>
|
||||
<div class="mt-2 h-2 overflow-hidden rounded-full bg-(--cp-default-border)">
|
||||
<div
|
||||
@@ -105,7 +138,17 @@ function quotaWindowBarClass(window?: any) {
|
||||
>
|
||||
<div class="flex items-center justify-between gap-3 text-[12px] font-[720]">
|
||||
<span class="text-(--cp-text-secondary)">{{ window.labelDisplay }}</span>
|
||||
<span class="text-(--cp-text-primary)">{{ window.usedPercentDisplay }}</span>
|
||||
<span class="flex shrink-0 items-baseline justify-end gap-1.5 font-mono tabular-nums">
|
||||
<span
|
||||
v-if="shouldShowQuotaWindowLocalUsage(window)"
|
||||
class="text-[12px] font-[680] text-(--cp-text-muted)"
|
||||
>
|
||||
{{ quotaWindowLocalUsageDisplay(window) }}
|
||||
</span>
|
||||
<span class="text-[12px] font-[780]" :class="quotaWindowPercentTextClass(window)">
|
||||
{{ window.usedPercentDisplay }}
|
||||
</span>
|
||||
</span>
|
||||
</div>
|
||||
<div class="mt-2 h-2 overflow-hidden rounded-full bg-(--cp-default-border)">
|
||||
<div
|
||||
@@ -130,7 +173,17 @@ function quotaWindowBarClass(window?: any) {
|
||||
>
|
||||
<div class="flex items-center justify-between gap-3 text-[12px] font-[720]">
|
||||
<span class="text-(--cp-text-secondary)">{{ window.labelDisplay }}</span>
|
||||
<span class="text-(--cp-text-primary)">{{ window.usedPercentDisplay }}</span>
|
||||
<span class="flex shrink-0 items-baseline justify-end gap-1.5 font-mono tabular-nums">
|
||||
<span
|
||||
v-if="shouldShowQuotaWindowLocalUsage(window)"
|
||||
class="text-[12px] font-[680] text-(--cp-text-muted)"
|
||||
>
|
||||
{{ quotaWindowLocalUsageDisplay(window) }}
|
||||
</span>
|
||||
<span class="text-[12px] font-[780]" :class="quotaWindowPercentTextClass(window)">
|
||||
{{ window.usedPercentDisplay }}
|
||||
</span>
|
||||
</span>
|
||||
</div>
|
||||
<div class="mt-2 h-2 overflow-hidden rounded-full bg-(--cp-default-border)">
|
||||
<div
|
||||
|
||||
@@ -50,16 +50,21 @@ function quotaWindowBarClass(window?: any) {
|
||||
return 'bg-(--cp-success)'
|
||||
}
|
||||
|
||||
function windowUsageDisplay(window?: any) {
|
||||
const display = window?.tokenUsageDisplay
|
||||
return typeof display === 'string' && display.trim() ? display : '-'
|
||||
}
|
||||
|
||||
function windowPercentDisplay(window?: any) {
|
||||
const display = window?.usedPercentDisplay
|
||||
return typeof display === 'string' && display.trim() ? display : '-'
|
||||
}
|
||||
|
||||
function windowLocalUsageDisplay(window?: any) {
|
||||
const display = window?.localUsage?.totalTokensDisplay
|
||||
return typeof display === 'string' && display.trim() ? display : '-'
|
||||
}
|
||||
|
||||
function shouldShowWindowLocalUsage(window?: any) {
|
||||
const totalTokens = window?.localUsage?.totalTokens
|
||||
return typeof totalTokens === 'number' && totalTokens > 0
|
||||
}
|
||||
|
||||
function quotaWindowPercentTextClass(window?: any) {
|
||||
if (window?.usedPercent === null || window?.usedPercent === undefined) {
|
||||
return 'text-(--cp-text-muted)'
|
||||
@@ -99,15 +104,18 @@ function quotaWindowMatches(actual: unknown, expected: number) {
|
||||
{{ window.labelDisplay }}
|
||||
</span>
|
||||
<span class="flex shrink-0 items-baseline justify-end gap-1.5 font-mono tabular-nums">
|
||||
<span
|
||||
v-if="shouldShowWindowLocalUsage(window)"
|
||||
class="text-[10px] leading-none font-[680] text-(--cp-text-muted)"
|
||||
>
|
||||
{{ windowLocalUsageDisplay(window) }}
|
||||
</span>
|
||||
<span
|
||||
class="text-[10px] leading-none font-[780]"
|
||||
:class="quotaWindowPercentTextClass(window)"
|
||||
>
|
||||
{{ windowPercentDisplay(window) }}
|
||||
</span>
|
||||
<span class="text-[10px] leading-none font-[680] text-(--cp-text-muted)">
|
||||
{{ windowUsageDisplay(window) }}
|
||||
</span>
|
||||
</span>
|
||||
</div>
|
||||
<div class="h-1 w-full overflow-hidden rounded-full bg-(--cp-default-border)">
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import { clamp } from 'es-toolkit'
|
||||
import { computed, onMounted, ref, type Ref } from 'vue'
|
||||
import { computed, onMounted, onUnmounted, ref, watch, type Ref } from 'vue'
|
||||
import dayjs from 'dayjs'
|
||||
|
||||
import {
|
||||
@@ -23,6 +23,8 @@ type TokenImportAccount = { token: string } | { refreshToken: string }
|
||||
type AccountRow = {
|
||||
id: string
|
||||
status: string
|
||||
displayStatus: string
|
||||
tokenRefreshing: boolean
|
||||
planType?: string | null
|
||||
quota: any
|
||||
usage: any
|
||||
@@ -33,6 +35,11 @@ type AccountQuotaRefreshResult = {
|
||||
account: AccountRow
|
||||
}
|
||||
|
||||
type AccountRefreshResult = {
|
||||
result: 'alive' | 'dead' | 'skipped'
|
||||
error?: string | null
|
||||
}
|
||||
|
||||
export function useAccountMutations(options: {
|
||||
page: Ref<number>
|
||||
pageSize: Ref<number>
|
||||
@@ -70,6 +77,7 @@ export function useAccountMutations(options: {
|
||||
const authorizingOAuth = authorizingOAuthAction.loading
|
||||
const batchDeleting = batchDeletingAction.loading
|
||||
const exportingAccounts = exportingAccountsAction.loading
|
||||
let tokenRefreshPollTimer: ReturnType<typeof window.setInterval> | undefined
|
||||
|
||||
const createForm = ref({
|
||||
mode: 'oauth',
|
||||
@@ -338,11 +346,18 @@ export function useAccountMutations(options: {
|
||||
async function handleRefresh(accountId: string) {
|
||||
await refreshingAccounts.run(accountId, async () => {
|
||||
try {
|
||||
await withMinimumDuration(async () => {
|
||||
await refreshAccount({ id: accountId })
|
||||
const result = await withMinimumDuration(async () => {
|
||||
const result = (await refreshAccount({ id: accountId })) as AccountRefreshResult
|
||||
await loadAccounts()
|
||||
return result
|
||||
})
|
||||
toast.success('Token 已刷新')
|
||||
if (result.result === 'alive') {
|
||||
toast.success('Token 已刷新')
|
||||
} else if (result.result === 'skipped') {
|
||||
toast.warning(result.error || 'Token 正在刷新中')
|
||||
} else {
|
||||
toast.error(result.error || '刷新失败')
|
||||
}
|
||||
} catch (error: any) {
|
||||
toast.error(error.message || '刷新失败')
|
||||
}
|
||||
@@ -389,10 +404,44 @@ export function useAccountMutations(options: {
|
||||
return `cpr-accounts-selected-${selectedCount}-${dayjs().format('YYYY-MM-DD')}.json`
|
||||
}
|
||||
|
||||
function hasTokenRefreshingAccount() {
|
||||
return accounts.value.some((account) => account.tokenRefreshing)
|
||||
}
|
||||
|
||||
function startTokenRefreshPolling() {
|
||||
if (tokenRefreshPollTimer !== undefined) return
|
||||
tokenRefreshPollTimer = window.setInterval(() => {
|
||||
if (!loading.value && hasTokenRefreshingAccount()) {
|
||||
void loadAccounts()
|
||||
}
|
||||
}, 2000)
|
||||
}
|
||||
|
||||
function stopTokenRefreshPolling() {
|
||||
if (tokenRefreshPollTimer === undefined) return
|
||||
window.clearInterval(tokenRefreshPollTimer)
|
||||
tokenRefreshPollTimer = undefined
|
||||
}
|
||||
|
||||
watch(
|
||||
() => hasTokenRefreshingAccount(),
|
||||
(hasRefreshing) => {
|
||||
if (hasRefreshing) {
|
||||
startTokenRefreshPolling()
|
||||
} else {
|
||||
stopTokenRefreshPolling()
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
onMounted(() => {
|
||||
loadAccounts()
|
||||
})
|
||||
|
||||
onUnmounted(() => {
|
||||
stopTokenRefreshPolling()
|
||||
})
|
||||
|
||||
return {
|
||||
loading,
|
||||
accounts,
|
||||
|
||||
@@ -222,7 +222,7 @@ bindAccountLoader(loadAccounts)
|
||||
</template>
|
||||
|
||||
<template #status="{ row }">
|
||||
<AccountStatusBadge :status="row.status" />
|
||||
<AccountStatusBadge :status="row.displayStatus" />
|
||||
</template>
|
||||
|
||||
<template #planType="{ row }">
|
||||
@@ -237,7 +237,7 @@ bindAccountLoader(loadAccounts)
|
||||
<AccountTableActions
|
||||
:account="row"
|
||||
:deleting="deletingAccount"
|
||||
:refreshing="refreshingAccountIds.has(row.id)"
|
||||
:refreshing="refreshingAccountIds.has(row.id) || row.tokenRefreshing"
|
||||
:schedule-label="scheduleActionLabel(row)"
|
||||
:testing="testingConnectionIds.has(row.id)"
|
||||
:updating-status="updatingStatusAccountIds.has(row.id)"
|
||||
|
||||
@@ -19,8 +19,8 @@ const model = defineModel<RotationStrategy>({ required: true })
|
||||
<template>
|
||||
<BaseCard
|
||||
:padded="false"
|
||||
title="账号选择"
|
||||
description="决定每次请求如何使用账号池。"
|
||||
title="调度策略"
|
||||
description="决定每次请求如何调度账号池。"
|
||||
header-class="px-5 pt-4"
|
||||
body-class="px-5 py-5"
|
||||
>
|
||||
|
||||
@@ -76,22 +76,22 @@ const rotationOptions: RotationOption[] = [
|
||||
{
|
||||
label: '智能调度(推荐)',
|
||||
value: 'smart',
|
||||
description: '优先避开高并发压力账号,并结合窗口用量和历史用量打散请求。',
|
||||
description: '按负载、窗口用量、请求数和健康反馈评分,优先选择更空闲的账号。',
|
||||
},
|
||||
{
|
||||
label: '额度重置优先',
|
||||
value: 'quota_reset_priority',
|
||||
description: '优先使用即将刷新额度的账号,最大化总使用量。',
|
||||
description: '优先选择额度窗口更快重置的账号,适合在重置前消耗剩余额度。',
|
||||
},
|
||||
{
|
||||
label: '轮询',
|
||||
value: 'round_robin',
|
||||
description: '按顺序轮流使用各账号。',
|
||||
description: '在可用候选账号间按顺序轮转,分配结果最可预测。',
|
||||
},
|
||||
{
|
||||
label: '粘滞',
|
||||
value: 'sticky',
|
||||
description: '持续使用同一账号,直到限速或额度耗尽。',
|
||||
description: '优先复用最近使用的账号,直到不可用后再切换。',
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@@ -1 +1 @@
|
||||
version: "1.0.8"
|
||||
version: "1.0.11"
|
||||
|
||||
Reference in New Issue
Block a user