mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-10 00:06:00 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6e24037c4d | ||
|
|
95daf7093e | ||
|
|
2e121c0670 | ||
|
|
66916ab8b2 | ||
|
|
2945523140 | ||
|
|
f375d2d77c | ||
|
|
c6e1b347be | ||
|
|
1ac9803c34 | ||
|
|
352b781340 | ||
|
|
62a6dfc489 | ||
|
|
e91f91b542 | ||
|
|
380556369b | ||
|
|
595cda541d | ||
|
|
90ed8bf6bc | ||
|
|
1101eee628 | ||
|
|
0e952c9d83 | ||
|
|
875d610812 | ||
|
|
9c37d262d1 | ||
|
|
bd95adadf9 | ||
|
|
6e6627b0c9 | ||
|
|
3579c28b75 | ||
|
|
2143de2791 | ||
|
|
298357f97c |
@@ -15,11 +15,11 @@
|
||||
|
||||
## 概览
|
||||
|
||||
Codex Proxy RS 是一个单进程服务:
|
||||
Codex Proxy RS 是一个单进程网关:
|
||||
|
||||
- Rust/Axum 后端提供 OpenAI 兼容 `/v1/*` 代理、管理端 API、SQLite 持久化和静态前端托管。
|
||||
- Vue 管理端用于账号导入、API Key、用量统计、请求记录、运行参数、模型映射和系统更新。
|
||||
- 运行数据默认收敛到仓库 `.runtime/`,Docker 也使用同一套宿主机目录做 bind mount。
|
||||
- Rust/Axum 后端提供 Responses 兼容代理、管理端 API、SQLite 持久化和静态前端托管。
|
||||
- Vue 管理端负责账号导入、API Key、用量统计、请求记录、运行参数、模型映射和系统更新。
|
||||
- 运行数据默认写入仓库 `.runtime/`;Docker 部署也使用这套宿主机目录挂载。
|
||||
|
||||
## 环境
|
||||
|
||||
@@ -55,7 +55,7 @@ mkdir -p .runtime/data .runtime/logs
|
||||
CPR_CONFIG_FILE=.runtime/config.yaml cargo run --manifest-path backend/Cargo.toml
|
||||
```
|
||||
|
||||
服务默认监听 `http://0.0.0.0:8080`。
|
||||
服务默认监听 `0.0.0.0:8080`。本机访问通常使用 `http://127.0.0.1:8080`。
|
||||
|
||||
## Docker 部署
|
||||
|
||||
@@ -106,8 +106,11 @@ docker compose -f deploy/docker-compose.yml up -d
|
||||
| 端点 | 说明 |
|
||||
| --- | --- |
|
||||
| `POST /v1/responses` | OpenAI Responses 兼容接口 |
|
||||
| `POST /v1/responses/review` | Review 模型请求入口 |
|
||||
| `POST /v1/responses/compact` | Compact 请求入口 |
|
||||
| `GET /v1/models` | 模型列表 |
|
||||
| `GET /v1/models/{id}` | 模型详情 |
|
||||
| `GET /v1/models/{id}/info` | 模型运行信息 |
|
||||
| `GET /v1/models/catalog` | 管理端可见模型目录 |
|
||||
| `/api/admin/*` | 管理端 API |
|
||||
|
||||
@@ -136,11 +139,11 @@ pnpm --dir frontend build
|
||||
CPR_CONFIG_FILE=.runtime/config.yaml cargo run --manifest-path backend/Cargo.toml
|
||||
```
|
||||
|
||||
管理端主要功能:
|
||||
管理端功能:
|
||||
|
||||
- 账号导入、OAuth 授权、连接测试、quota 刷新、token 刷新。
|
||||
- 客户端 API Key 管理。
|
||||
- 请求记录、Token 用量、模型用量和账号额度。
|
||||
- Dashboard、请求记录、Token 用量、模型用量和账号额度。
|
||||
- 运行参数、模型别名、账号选择策略。
|
||||
- 版本检查、在线更新、更新日志、重启和回滚。
|
||||
|
||||
@@ -151,7 +154,7 @@ CPR_CONFIG_FILE=.runtime/config.yaml cargo run --manifest-path backend/Cargo.tom
|
||||
`.runtime` 约定:
|
||||
|
||||
```text
|
||||
.runtime/config.yaml # 本地或 Docker 启动配置,按实际运行环境写路径
|
||||
.runtime/config.yaml # 本地或 Docker 启动配置,路径按运行环境写
|
||||
.runtime/data/codex-proxy-rs.sqlite # SQLite 数据库
|
||||
.runtime/data/installation_id # 上游 installation id
|
||||
.runtime/data/update-state.json # 在线更新状态
|
||||
@@ -161,22 +164,21 @@ CPR_CONFIG_FILE=.runtime/config.yaml cargo run --manifest-path backend/Cargo.tom
|
||||
|
||||
## 发布与更新
|
||||
|
||||
版本信息在 `release/version.yaml`,发布平台在 `release/platforms.yaml`。
|
||||
当前版本记录在 `release/version.yaml`,发布平台记录在 `release/platforms.yaml`。
|
||||
|
||||
发布命令:
|
||||
|
||||
```bash
|
||||
release/publish 0.1.5
|
||||
release/publish 1.0.4
|
||||
```
|
||||
|
||||
该命令会更新版本文件、生成版本提交、创建 `v<version>` tag 并推送,随后 GitHub Actions 构建 Release 归档和 GHCR 多平台镜像。
|
||||
该命令会更新版本文件、生成版本提交、创建 `v<version>` tag 并推送。GitHub Actions 随后构建 Release 归档和 GHCR 多平台镜像。
|
||||
|
||||
在线更新由主服务处理,管理端调用 `/api/admin/system/*`:
|
||||
|
||||
- 检查 GitHub Release。
|
||||
- 下载匹配当前平台的归档和 `checksums.txt`。
|
||||
- 校验 checksum。
|
||||
- 替换二进制和 `web/dist`。
|
||||
- 查询 GitHub Release。
|
||||
- 下载当前平台对应的归档和 `checksums.txt`。
|
||||
- 校验 checksum 后替换二进制和 `web/dist`。
|
||||
- 重启时,Docker 模式依赖 `restart: unless-stopped` 拉起新容器进程;非 Docker 模式会先安排新进程延迟启动,再关闭当前进程。
|
||||
|
||||
## 项目结构
|
||||
@@ -185,7 +187,7 @@ release/publish 0.1.5
|
||||
backend/ Rust/Axum 后端、SQLite migration、集成测试
|
||||
frontend/ Vue 3 管理端
|
||||
deploy/ Dockerfile、Compose、部署配置模板
|
||||
docs/ 架构、审计和维护文档
|
||||
docs/ 架构和维护文档
|
||||
release/ 版本、平台和发布脚本
|
||||
skills/ 项目本地 Codex skill
|
||||
```
|
||||
|
||||
@@ -43,6 +43,16 @@ pub(crate) struct AdminAccountQuotaUsageWindow {
|
||||
pub(crate) key: String,
|
||||
pub(crate) start: DateTime<Utc>,
|
||||
pub(crate) end: DateTime<Utc>,
|
||||
pub(crate) window_seconds: u64,
|
||||
}
|
||||
|
||||
impl AdminAccountQuotaUsageWindow {
|
||||
pub(crate) fn duration_seconds(&self) -> u64 {
|
||||
self.end
|
||||
.signed_duration_since(self.start)
|
||||
.num_seconds()
|
||||
.max(0) as u64
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
@@ -101,12 +111,14 @@ impl AdminAccountQuotaData {
|
||||
impl AdminAccountQuotaWindowData {
|
||||
fn usage_window(&self) -> Option<AdminAccountQuotaUsageWindow> {
|
||||
let end = self.reset_at?;
|
||||
let seconds = i64::try_from(self.window_seconds?).ok()?;
|
||||
let window_seconds = self.window_seconds?;
|
||||
let seconds = i64::try_from(window_seconds).ok()?;
|
||||
let start = end.checked_sub_signed(Duration::seconds(seconds))?;
|
||||
(start <= end).then(|| AdminAccountQuotaUsageWindow {
|
||||
key: self.key.clone(),
|
||||
start,
|
||||
end,
|
||||
window_seconds,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -303,6 +303,14 @@ 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,
|
||||
@@ -416,6 +424,34 @@ impl From<AdminAccountRefreshResult> for AccountRefreshData {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountQuotaData {
|
||||
quota: Value,
|
||||
raw: Value,
|
||||
quota_data: AdminAccountQuotaData,
|
||||
plan_type: Option<String>,
|
||||
account: AdminAccountData,
|
||||
}
|
||||
|
||||
impl AccountQuotaData {
|
||||
fn from_account(
|
||||
quota: Value,
|
||||
raw: Value,
|
||||
quota_data: AdminAccountQuotaData,
|
||||
account: AdminAccountData,
|
||||
) -> Self {
|
||||
let plan_type = account.plan_type.clone();
|
||||
Self {
|
||||
quota,
|
||||
raw,
|
||||
quota_data,
|
||||
plan_type,
|
||||
account,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Default)]
|
||||
#[serde(default)]
|
||||
pub(crate) struct AccountHealthCheckRequest {
|
||||
@@ -656,18 +692,13 @@ pub(crate) async fn account_quota(
|
||||
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;
|
||||
let plan_type = quota
|
||||
.get("plan_type")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToString::to_string);
|
||||
let account =
|
||||
account_data_for_quota_refresh(&state, &account_id, quota_data.clone()).await?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(serde_json::json!({
|
||||
"quota": quota,
|
||||
"raw": raw,
|
||||
"quotaData": quota_data,
|
||||
"planType": plan_type,
|
||||
})),
|
||||
AdminEnvelope::ok(AccountQuotaData::from_account(
|
||||
quota, raw, quota_data, account,
|
||||
)),
|
||||
))
|
||||
}
|
||||
Err(AdminAccountError::NotFound) => Err(account_not_found()),
|
||||
@@ -682,6 +713,25 @@ pub(crate) async fn account_quota(
|
||||
}
|
||||
}
|
||||
|
||||
async fn account_data_for_quota_refresh(
|
||||
state: &AppState,
|
||||
account_id: &str,
|
||||
quota_data: AdminAccountQuotaData,
|
||||
) -> Result<AdminAccountData, AdminError> {
|
||||
let Some(account) = state
|
||||
.services
|
||||
.admin_accounts
|
||||
.get(account_id)
|
||||
.await
|
||||
.map_err(|error| account_error(&error))?
|
||||
else {
|
||||
return Err(account_not_found());
|
||||
};
|
||||
let quota_by_account = HashMap::from([(account_id.to_string(), quota_data)]);
|
||||
let stats = account_list_stats(state, std::slice::from_ref(&account), "a_by_account).await;
|
||||
Ok(stats.data_for(account))
|
||||
}
|
||||
|
||||
/// `POST /api/admin/accounts/import`
|
||||
pub(crate) async fn import_accounts(
|
||||
State(state): State<AppState>,
|
||||
@@ -873,6 +923,19 @@ async fn account_list_stats(
|
||||
.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 quota_window_usage_stats =
|
||||
account_usage_stats_by_quota_window(state, &selected_quota_windows).await;
|
||||
let usage_records = usage_records
|
||||
.into_iter()
|
||||
.map(|usage| {
|
||||
apply_selected_quota_window_usage(
|
||||
usage,
|
||||
&selected_quota_windows,
|
||||
"a_window_usage_stats,
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let models_by_account = list_current_window_model_usage(state, &usage_records).await;
|
||||
|
||||
AccountListStats {
|
||||
@@ -885,6 +948,105 @@ async fn account_list_stats(
|
||||
}
|
||||
}
|
||||
|
||||
fn selected_quota_windows_by_account(
|
||||
accounts: &[AdminAccountMetadata],
|
||||
quota_by_account: &HashMap<String, AdminAccountQuotaData>,
|
||||
) -> HashMap<String, AdminAccountQuotaUsageWindow> {
|
||||
accounts
|
||||
.iter()
|
||||
.filter_map(|account| {
|
||||
quota_by_account
|
||||
.get(&account.id)
|
||||
.and_then(selected_quota_window)
|
||||
.map(|window| (account.id.clone(), window))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn selected_quota_window(quota: &AdminAccountQuotaData) -> Option<AdminAccountQuotaUsageWindow> {
|
||||
quota.usage_windows().into_iter().max_by(|left, right| {
|
||||
left.duration_seconds()
|
||||
.cmp(&right.duration_seconds())
|
||||
.then_with(|| left.end.cmp(&right.end))
|
||||
.then_with(|| left.key.cmp(&right.key))
|
||||
})
|
||||
}
|
||||
|
||||
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,
|
||||
|
||||
@@ -38,15 +38,11 @@ impl AdminAccountService {
|
||||
.await
|
||||
.map_err(|e| AdminAccountError::FetchQuota(e.to_string()))?;
|
||||
let normalized = crate::upstream::accounts::quota::quota_from_usage(&raw);
|
||||
if let Ok(json_str) = serde_json::to_string(&normalized) {
|
||||
if matches!(
|
||||
self.store.update_quota_json(account_id, &json_str).await,
|
||||
Ok(true)
|
||||
) {
|
||||
self.sync_account_pool_best_effort(account_id, "account quota refresh")
|
||||
.await;
|
||||
}
|
||||
}
|
||||
self.account_pool
|
||||
.apply_quota_snapshot(account_id, &normalized)
|
||||
.await;
|
||||
self.sync_account_pool_best_effort(account_id, "account quota refresh")
|
||||
.await;
|
||||
Ok(serde_json::json!({ "quota": normalized, "raw": raw }))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
//! 管理端 Dashboard 聚合视图。
|
||||
|
||||
use std::{cmp::Reverse, collections::HashMap};
|
||||
use std::collections::HashMap;
|
||||
|
||||
use axum::{
|
||||
extract::{Query, State},
|
||||
@@ -15,14 +15,19 @@ use crate::{
|
||||
admin::{
|
||||
auth::session::AdminAuth,
|
||||
monitoring::{
|
||||
account_usage_service::{
|
||||
AdminUsageRecord, AdminUsageSummary, AdminUsageTimeBucketRecord,
|
||||
},
|
||||
account_usage_service::{AdminUsageSummary, AdminUsageTimeBucketRecord},
|
||||
diagnostics::{
|
||||
fingerprint_diagnostics, AccountCapacityDiagnostics, AccountPoolDiagnostics,
|
||||
},
|
||||
usage_record_routes::{usage_record_items, UsageRecordData},
|
||||
usage_record_model::{
|
||||
metadata_nonnegative_i64, metadata_string, UsageRecord, UsageRecordLevel,
|
||||
},
|
||||
usage_record_routes::{
|
||||
usage_cost_details, usage_record_models, usage_token_details,
|
||||
UsageRecordCostDetailsData, UsageRecordTokenDetailsData,
|
||||
},
|
||||
usage_record_service::AdminUsageRecordFilter,
|
||||
usage_record_store::UsageRecordAccountUsage,
|
||||
},
|
||||
response::{AdminEnvelope, AdminError, AdminResponse},
|
||||
},
|
||||
@@ -40,11 +45,15 @@ use crate::{
|
||||
upstream::accounts::model::{Account, AccountStatus},
|
||||
};
|
||||
|
||||
const DASHBOARD_USAGE_RECORD_LIMIT: u32 = 1000;
|
||||
const DASHBOARD_ACCOUNT_USAGE_LIMIT: u32 = 4;
|
||||
const DASHBOARD_USAGE_RECORD_LIMIT: u32 = 10;
|
||||
const HEALTH_TIMELINE_SLOT_MINUTES: i64 = 15;
|
||||
const DASHBOARD_TIME_BUCKET_SLOTS: i64 = 7 * 24 * 4;
|
||||
const HEALTH_TIMELINE_SLOTS: i64 = 24 * 4;
|
||||
const HEALTH_TIMELINE_STABLE_SUCCESS_THRESHOLD: u64 = 3;
|
||||
const FIVE_HOUR_WINDOW_SECONDS: u64 = 18_000;
|
||||
const WEEK_WINDOW_SECONDS: u64 = 604_800;
|
||||
const MONTH_WINDOW_SECONDS: u64 = 2_592_000;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
@@ -69,7 +78,7 @@ struct DashboardSummaryData {
|
||||
health_timeline: DashboardHealthTimelineData,
|
||||
account_usage: Vec<DashboardAccountUsageData>,
|
||||
service_statuses: Vec<DashboardServiceStatusData>,
|
||||
usage_records: Vec<UsageRecordData>,
|
||||
usage_records: Vec<DashboardUsageRecordData>,
|
||||
pool_summary: AccountPoolDiagnostics,
|
||||
capacity_info: AccountCapacityDiagnostics,
|
||||
rotation_strategy: Option<String>,
|
||||
@@ -80,10 +89,7 @@ struct DashboardSummaryData {
|
||||
struct DashboardHealthTimelineData {
|
||||
title: String,
|
||||
description: String,
|
||||
range_display: String,
|
||||
reliability_display: String,
|
||||
oldest_label: String,
|
||||
newest_label: String,
|
||||
points: String,
|
||||
}
|
||||
|
||||
@@ -112,11 +118,8 @@ struct DashboardAccountsCardData {
|
||||
struct DashboardTrafficCardData {
|
||||
today_requests: String,
|
||||
today_requests_value: u64,
|
||||
yesterday_requests: String,
|
||||
yesterday_requests_value: u64,
|
||||
total_requests: String,
|
||||
rpm: String,
|
||||
tpm: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
@@ -124,13 +127,9 @@ struct DashboardTrafficCardData {
|
||||
struct DashboardTokenCardData {
|
||||
today_tokens: String,
|
||||
today_tokens_value: u64,
|
||||
yesterday_tokens: String,
|
||||
yesterday_tokens_value: u64,
|
||||
total_tokens: String,
|
||||
today_cost_usd: String,
|
||||
today_cost_usd_value: f64,
|
||||
total_cost_usd: String,
|
||||
total_cost_usd_value: f64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
@@ -138,12 +137,10 @@ struct DashboardTokenCardData {
|
||||
struct DashboardCacheCardData {
|
||||
today_hit_rate: String,
|
||||
today_hit_rate_value: Option<f64>,
|
||||
yesterday_hit_rate: String,
|
||||
yesterday_hit_rate_value: Option<f64>,
|
||||
total_hit_rate: String,
|
||||
total_cached_tokens: String,
|
||||
first_token_latency_ms: String,
|
||||
completion_latency_ms: String,
|
||||
average_first_token_latency_ms: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
@@ -166,7 +163,7 @@ struct DashboardTrendPointData {
|
||||
output_tokens_value: u64,
|
||||
cached_tokens: String,
|
||||
cached_tokens_value: u64,
|
||||
tokens: String,
|
||||
cache_hit_rate_value: f64,
|
||||
tokens_value: u64,
|
||||
errors: String,
|
||||
errors_value: u64,
|
||||
@@ -191,14 +188,34 @@ struct DashboardTrendSummaryData {
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct DashboardAccountUsageData {
|
||||
name: String,
|
||||
id: String,
|
||||
email: String,
|
||||
plan: String,
|
||||
requests: String,
|
||||
plan_type: Option<String>,
|
||||
tokens: String,
|
||||
quota_used_percent: Option<f64>,
|
||||
last_used: String,
|
||||
status: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct DashboardUsageRecordData {
|
||||
id: String,
|
||||
level: UsageRecordLevel,
|
||||
route: Option<String>,
|
||||
model: Option<String>,
|
||||
status_code: Option<i64>,
|
||||
transport: Option<String>,
|
||||
created_at_display: String,
|
||||
account_email: Option<String>,
|
||||
requested_model: Option<String>,
|
||||
upstream_model: Option<String>,
|
||||
client_ip: Option<String>,
|
||||
user_agent: Option<String>,
|
||||
reasoning_effort: Option<String>,
|
||||
token_details: UsageRecordTokenDetailsData,
|
||||
cost_details: Option<UsageRecordCostDetailsData>,
|
||||
first_token_latency_ms_display: String,
|
||||
latency_ms_display: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
@@ -221,8 +238,14 @@ struct UsageWindow {
|
||||
first_token_latency_count: u64,
|
||||
max_first_token_bucket_latency: u64,
|
||||
min_first_token_bucket_latency: u64,
|
||||
latency_sum: u64,
|
||||
latency_count: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
||||
enum QuotaWindowPriority {
|
||||
FiveHour,
|
||||
Weekly,
|
||||
Monthly,
|
||||
Other,
|
||||
}
|
||||
|
||||
impl UsageWindow {
|
||||
@@ -250,10 +273,6 @@ impl UsageWindow {
|
||||
fn min_first_token_bucket_latency(self) -> Option<u64> {
|
||||
(self.min_first_token_bucket_latency > 0).then_some(self.min_first_token_bucket_latency)
|
||||
}
|
||||
|
||||
fn avg_latency(self) -> Option<u64> {
|
||||
self.latency_sum.checked_div(self.latency_count)
|
||||
}
|
||||
}
|
||||
|
||||
/// `GET /api/admin/dashboard/summary`
|
||||
@@ -269,53 +288,56 @@ pub(crate) async fn dashboard_summary(
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
let capacity = state.services.account_pool.capacity_summary_now().await;
|
||||
let summary = state
|
||||
let now = Utc::now();
|
||||
let today_filter = today_usage_record_filter(now);
|
||||
let usage_summary = state
|
||||
.services
|
||||
.usage
|
||||
.summary()
|
||||
.await
|
||||
.unwrap_or_else(|_| empty_usage_summary());
|
||||
let usage_records = state
|
||||
let account_usage_records = state
|
||||
.services
|
||||
.usage
|
||||
.list(None, 200)
|
||||
.usage_records
|
||||
.account_usage(today_filter.clone(), DASHBOARD_ACCOUNT_USAGE_LIMIT)
|
||||
.await
|
||||
.map(|page| page.items)
|
||||
.unwrap_or_default();
|
||||
let recent_events = state
|
||||
.services
|
||||
.usage_records
|
||||
.list(
|
||||
None,
|
||||
DASHBOARD_USAGE_RECORD_LIMIT,
|
||||
AdminUsageRecordFilter::default(),
|
||||
)
|
||||
.list(None, DASHBOARD_USAGE_RECORD_LIMIT, today_filter)
|
||||
.await
|
||||
.map(|page| page.items)
|
||||
.unwrap_or_default();
|
||||
|
||||
let now = Utc::now();
|
||||
let time_buckets = dashboard_time_buckets(&state, now).await;
|
||||
let pool_summary = account_pool_summary(&accounts);
|
||||
let trend = dashboard_trend_data(&time_buckets, query.kind.unwrap_or_default());
|
||||
let quota_used_by_account = account_quota_used_percent_by_id(&state, &usage_records).await;
|
||||
let quota_used_by_account =
|
||||
account_quota_used_percent_by_id(&state, &account_usage_records).await;
|
||||
let settings = state.services.settings.current();
|
||||
let recent_usage_records = recent_events.into_iter().take(10).collect::<Vec<_>>();
|
||||
let recent_usage_records = recent_events;
|
||||
let account_emails = state
|
||||
.services
|
||||
.usage_records
|
||||
.account_email_map(&recent_usage_records)
|
||||
.await
|
||||
.map_err(|error| AdminError::usage_record_accounts_failed(error.to_string()))?;
|
||||
let dashboard_usage_records = usage_record_items(recent_usage_records, &account_emails);
|
||||
let dashboard_usage_records =
|
||||
dashboard_usage_record_items(recent_usage_records, &account_emails);
|
||||
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(DashboardSummaryData {
|
||||
cards: dashboard_cards(&accounts, &summary, &time_buckets),
|
||||
cards: dashboard_cards(&accounts, &time_buckets, &usage_summary),
|
||||
trend,
|
||||
health_timeline: dashboard_health_timeline_data(&time_buckets),
|
||||
account_usage: account_usage_data(&accounts, &usage_records, "a_used_by_account),
|
||||
account_usage: account_usage_data(
|
||||
&accounts,
|
||||
&account_usage_records,
|
||||
"a_used_by_account,
|
||||
now,
|
||||
),
|
||||
service_statuses: service_status_data(&state),
|
||||
usage_records: dashboard_usage_records,
|
||||
pool_summary,
|
||||
@@ -325,6 +347,14 @@ pub(crate) async fn dashboard_summary(
|
||||
))
|
||||
}
|
||||
|
||||
fn today_usage_record_filter(now: DateTime<Utc>) -> AdminUsageRecordFilter {
|
||||
AdminUsageRecordFilter {
|
||||
start_time: Some(china_day_start(now)),
|
||||
end_time: Some(now),
|
||||
..AdminUsageRecordFilter::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// `GET /api/admin/dashboard/trend?kind=usage|latency|errors`
|
||||
pub(crate) async fn dashboard_trend(
|
||||
State(state): State<AppState>,
|
||||
@@ -344,16 +374,16 @@ pub(crate) async fn dashboard_trend(
|
||||
|
||||
fn dashboard_cards(
|
||||
accounts: &[Account],
|
||||
summary: &AdminUsageSummary,
|
||||
buckets: &[AdminUsageTimeBucketRecord],
|
||||
usage_summary: &AdminUsageSummary,
|
||||
) -> DashboardCardsData {
|
||||
dashboard_cards_at(accounts, summary, buckets, Utc::now())
|
||||
dashboard_cards_at(accounts, buckets, usage_summary, Utc::now())
|
||||
}
|
||||
|
||||
fn dashboard_cards_at(
|
||||
accounts: &[Account],
|
||||
summary: &AdminUsageSummary,
|
||||
buckets: &[AdminUsageTimeBucketRecord],
|
||||
usage_summary: &AdminUsageSummary,
|
||||
now: DateTime<Utc>,
|
||||
) -> DashboardCardsData {
|
||||
let today_start = china_day_start(now);
|
||||
@@ -361,10 +391,11 @@ 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_cost = total_cost(buckets);
|
||||
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_hit_rate = summary_cache_hit_rate(usage_summary);
|
||||
|
||||
let total_input = nonnegative_i64_to_u64(summary.input_tokens);
|
||||
let total_cached = nonnegative_i64_to_u64(summary.cached_tokens);
|
||||
let total_accounts = accounts.len() as u64;
|
||||
let enabled_accounts = accounts
|
||||
.iter()
|
||||
@@ -382,13 +413,6 @@ fn dashboard_cards_at(
|
||||
)
|
||||
})
|
||||
.count() as u64;
|
||||
let total_requests = nonnegative_i64_to_u64(summary.request_count);
|
||||
let total_tokens = nonnegative_i64_to_u64(summary.input_tokens + summary.output_tokens);
|
||||
let total_hit_rate = if total_input > 0 {
|
||||
Some(total_cached as f64 / total_input as f64)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
DashboardCardsData {
|
||||
accounts: DashboardAccountsCardData {
|
||||
@@ -402,38 +426,53 @@ fn dashboard_cards_at(
|
||||
traffic: DashboardTrafficCardData {
|
||||
today_requests: format_compact_number(today.requests),
|
||||
today_requests_value: today.requests,
|
||||
yesterday_requests: format_compact_number(yesterday.requests),
|
||||
yesterday_requests_value: yesterday.requests,
|
||||
total_requests: format_compact_number(total_requests),
|
||||
rpm: format_compact_number(today.requests),
|
||||
tpm: format_tokens(today.tokens()),
|
||||
},
|
||||
tokens: DashboardTokenCardData {
|
||||
today_tokens: format_tokens(today.tokens()),
|
||||
today_tokens_value: today.tokens(),
|
||||
yesterday_tokens: format_tokens(yesterday.tokens()),
|
||||
yesterday_tokens_value: yesterday.tokens(),
|
||||
total_tokens: format_tokens(total_tokens),
|
||||
today_cost_usd: format_optional_cost(Some(today_cost)),
|
||||
today_cost_usd_value: today_cost,
|
||||
total_cost_usd: format_optional_cost(Some(total_cost)),
|
||||
total_cost_usd_value: total_cost,
|
||||
today_cost_usd: format_cost(today_cost),
|
||||
},
|
||||
cache: DashboardCacheCardData {
|
||||
today_hit_rate: format_rate(today.cache_hit_rate()),
|
||||
today_hit_rate_value: today.cache_hit_rate(),
|
||||
yesterday_hit_rate: format_rate(yesterday.cache_hit_rate()),
|
||||
yesterday_hit_rate_value: yesterday.cache_hit_rate(),
|
||||
total_hit_rate: format_rate(total_hit_rate),
|
||||
total_cached_tokens: format_tokens(total_cached),
|
||||
first_token_latency_ms: format_optional_duration_ms(today.avg_first_token_latency()),
|
||||
completion_latency_ms: format_optional_duration_ms(today.avg_latency()),
|
||||
total_cached_tokens: format_tokens(total_cached_tokens),
|
||||
average_first_token_latency_ms: format_optional_duration_ms(
|
||||
today.avg_first_token_latency(),
|
||||
),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn format_optional_cost(value: Option<f64>) -> String {
|
||||
value.map_or_else(|| "-".to_string(), format_cost)
|
||||
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);
|
||||
if input_tokens > 0 {
|
||||
Some(nonnegative_i64_to_u64(summary.cached_tokens) as f64 / input_tokens as f64)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn format_optional_duration_ms(value: Option<u64>) -> String {
|
||||
@@ -468,8 +507,9 @@ fn dashboard_trend_data_at(
|
||||
now: DateTime<Utc>,
|
||||
) -> DashboardTrendData {
|
||||
let current_hour = china_hour_start(now);
|
||||
let start = current_hour - Duration::hours(23);
|
||||
let mut buckets = (0..24)
|
||||
let start = china_day_start(now);
|
||||
let elapsed_hours = current_hour.signed_duration_since(start).num_hours();
|
||||
let mut buckets = (0..=elapsed_hours)
|
||||
.map(|index| {
|
||||
let bucket_start = start + Duration::hours(index);
|
||||
(bucket_start, UsageWindow::default())
|
||||
@@ -496,6 +536,7 @@ fn dashboard_trend_data_at(
|
||||
let max_latency = bucket.max_first_token_bucket_latency().unwrap_or(0);
|
||||
let min_latency = bucket.min_first_token_bucket_latency().unwrap_or(0);
|
||||
let tokens = bucket.tokens();
|
||||
let cache_hit_rate = bucket.cache_hit_rate().unwrap_or(0.0);
|
||||
let success_rate = if bucket.requests > 0 {
|
||||
((bucket.requests - bucket.errors) as f64 / bucket.requests as f64 * 1000.0).round()
|
||||
/ 10.0
|
||||
@@ -512,7 +553,7 @@ fn dashboard_trend_data_at(
|
||||
output_tokens_value: bucket.output_tokens,
|
||||
cached_tokens: format_tokens(bucket.cached_tokens),
|
||||
cached_tokens_value: bucket.cached_tokens,
|
||||
tokens: format_tokens(tokens),
|
||||
cache_hit_rate_value: cache_hit_rate,
|
||||
tokens_value: tokens,
|
||||
errors: format_compact_number(bucket.errors),
|
||||
errors_value: bucket.errors,
|
||||
@@ -585,15 +626,12 @@ fn dashboard_health_timeline_data_at(
|
||||
};
|
||||
DashboardHealthTimelineData {
|
||||
title: "请求健康时间线".to_string(),
|
||||
description: "今日请求可靠性".to_string(),
|
||||
range_display: format!("{} - {}", china_datetime(&start), china_datetime(&now)),
|
||||
description: "请求可靠性".to_string(),
|
||||
reliability_display: if requests > 0 {
|
||||
format!("{reliability:.1}%")
|
||||
} else {
|
||||
"-".to_string()
|
||||
},
|
||||
oldest_label: "00:00".to_string(),
|
||||
newest_label: "24:00".to_string(),
|
||||
points: buckets
|
||||
.into_iter()
|
||||
.map(|(bucket_start, bucket)| {
|
||||
@@ -759,22 +797,6 @@ fn cost_window(
|
||||
has_usage.then_some(total)
|
||||
}
|
||||
|
||||
fn total_cost(records: &[AdminUsageTimeBucketRecord]) -> f64 {
|
||||
let mut total = 0.0;
|
||||
let mut has_usage = false;
|
||||
for record in records {
|
||||
if let Some(cost) = record.cost_usd {
|
||||
total += cost;
|
||||
has_usage = true;
|
||||
}
|
||||
}
|
||||
if has_usage {
|
||||
return total;
|
||||
}
|
||||
|
||||
0.0
|
||||
}
|
||||
|
||||
fn apply_bucket(window: &mut UsageWindow, record: &AdminUsageTimeBucketRecord) {
|
||||
window.requests += nonnegative_i64_to_u64(record.request_count);
|
||||
window.input_tokens += nonnegative_i64_to_u64(record.input_tokens);
|
||||
@@ -792,8 +814,6 @@ fn apply_bucket(window: &mut UsageWindow, record: &AdminUsageTimeBucketRecord) {
|
||||
window.min_first_token_bucket_latency.min(latency)
|
||||
};
|
||||
}
|
||||
window.latency_sum += nonnegative_i64_to_u64(record.latency_sum);
|
||||
window.latency_count += nonnegative_i64_to_u64(record.latency_count);
|
||||
}
|
||||
|
||||
fn first_token_bucket_latency(record: &AdminUsageTimeBucketRecord) -> Option<u64> {
|
||||
@@ -822,68 +842,51 @@ fn account_pool_summary(accounts: &[Account]) -> AccountPoolDiagnostics {
|
||||
|
||||
fn account_usage_data(
|
||||
accounts: &[Account],
|
||||
usage_records: &[AdminUsageRecord],
|
||||
usage_records: &[UsageRecordAccountUsage],
|
||||
quota_used_by_account: &HashMap<String, f64>,
|
||||
now: DateTime<Utc>,
|
||||
) -> Vec<DashboardAccountUsageData> {
|
||||
let account_by_id = accounts
|
||||
.iter()
|
||||
.map(|account| (account.id.as_str(), account))
|
||||
.collect::<HashMap<_, _>>();
|
||||
let mut records = usage_records.iter().collect::<Vec<_>>();
|
||||
records.sort_by_key(|record| Reverse((record.last_used_at, record.account_id.clone())));
|
||||
let mut rows = records
|
||||
usage_records
|
||||
.iter()
|
||||
.map(|usage| {
|
||||
let account = account_by_id.get(usage.account_id.as_str()).copied();
|
||||
let name = account
|
||||
.and_then(|account| account.label.clone())
|
||||
.or_else(|| usage.label.clone())
|
||||
.or_else(|| {
|
||||
usage
|
||||
.email
|
||||
.as_ref()
|
||||
.and_then(|email| email.split('@').next().map(ToString::to_string))
|
||||
})
|
||||
.unwrap_or_else(|| usage.account_id.chars().take(8).collect());
|
||||
let status = account.map_or_else(
|
||||
|| "active".to_string(),
|
||||
|account| account.status.to_string(),
|
||||
);
|
||||
let quota_used_percent = quota_used_by_account
|
||||
.get(&usage.account_id)
|
||||
.copied()
|
||||
.or_else(|| (status == "quota_exhausted").then_some(100.0));
|
||||
let requests = nonnegative_i64_to_u64(usage.request_count);
|
||||
let tokens = nonnegative_i64_to_u64(usage.input_tokens + usage.output_tokens);
|
||||
.or_else(|| {
|
||||
matches!(
|
||||
account.map(|account| account.status),
|
||||
Some(AccountStatus::QuotaExhausted)
|
||||
)
|
||||
.then_some(100.0)
|
||||
});
|
||||
DashboardAccountUsageData {
|
||||
name,
|
||||
email: usage.email.clone().unwrap_or_else(|| "-".to_string()),
|
||||
plan: usage
|
||||
.plan_type
|
||||
.clone()
|
||||
.unwrap_or_else(|| "free".to_string()),
|
||||
requests: format_compact_number(requests),
|
||||
tokens: format_tokens(tokens),
|
||||
id: usage.account_id.clone(),
|
||||
email: account
|
||||
.and_then(|account| account.email.as_ref())
|
||||
.cloned()
|
||||
.expect("dashboard account usage requires email"),
|
||||
plan_type: account
|
||||
.and_then(|account| account.plan_type.as_ref())
|
||||
.cloned(),
|
||||
tokens: format_tokens(usage.total_tokens),
|
||||
quota_used_percent,
|
||||
last_used: china_relative_time(usage.last_used_at, Utc::now()),
|
||||
status,
|
||||
last_used: china_relative_time(Some(usage.last_used_at), now),
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
rows.truncate(4);
|
||||
rows
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn account_quota_used_percent_by_id(
|
||||
state: &AppState,
|
||||
usage_records: &[AdminUsageRecord],
|
||||
usage_records: &[UsageRecordAccountUsage],
|
||||
) -> HashMap<String, f64> {
|
||||
let mut records = usage_records.iter().collect::<Vec<_>>();
|
||||
records.sort_by_key(|record| Reverse((record.last_used_at, record.account_id.clone())));
|
||||
records.truncate(4);
|
||||
|
||||
let mut quota_used_by_account = HashMap::with_capacity(records.len());
|
||||
for usage in records {
|
||||
let mut quota_used_by_account = HashMap::with_capacity(usage_records.len());
|
||||
for usage in usage_records {
|
||||
let Ok(Some(quota_json)) = state
|
||||
.services
|
||||
.accounts
|
||||
@@ -901,27 +904,67 @@ async fn account_quota_used_percent_by_id(
|
||||
|
||||
fn quota_used_percent(quota_json: &str) -> Option<f64> {
|
||||
let quota = serde_json::from_str::<Value>(quota_json).ok()?;
|
||||
let mut values = Vec::new();
|
||||
if let Some(used_percent) = quota
|
||||
.pointer("/monthly_limit/used_percent")
|
||||
.and_then(percent_value)
|
||||
{
|
||||
values.push(used_percent);
|
||||
let mut selected = None;
|
||||
if let Some(monthly_limit) = quota.get("monthly_limit") {
|
||||
select_quota_used_percent(
|
||||
&mut selected,
|
||||
quota_window_priority(monthly_limit, QuotaWindowPriority::Monthly),
|
||||
monthly_limit.get("used_percent").and_then(percent_value),
|
||||
);
|
||||
}
|
||||
if let Some(snapshots) = quota.get("snapshots").and_then(Value::as_array) {
|
||||
for snapshot in snapshots {
|
||||
for role in ["primary", "secondary"] {
|
||||
if let Some(used_percent) = snapshot
|
||||
.get(role)
|
||||
.and_then(|window| window.get("used_percent"))
|
||||
.and_then(percent_value)
|
||||
{
|
||||
values.push(used_percent);
|
||||
if let Some(window) = snapshot.get(role) {
|
||||
select_quota_used_percent(
|
||||
&mut selected,
|
||||
quota_window_priority(window, QuotaWindowPriority::Other),
|
||||
window.get("used_percent").and_then(percent_value),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
values.into_iter().max_by(f64::total_cmp)
|
||||
selected.map(|(_, used_percent)| used_percent)
|
||||
}
|
||||
|
||||
fn select_quota_used_percent(
|
||||
selected: &mut Option<(QuotaWindowPriority, f64)>,
|
||||
priority: QuotaWindowPriority,
|
||||
used_percent: Option<f64>,
|
||||
) {
|
||||
let Some(used_percent) = used_percent else {
|
||||
return;
|
||||
};
|
||||
match selected {
|
||||
Some((selected_priority, selected_percent))
|
||||
if priority > *selected_priority
|
||||
|| (priority == *selected_priority && used_percent <= *selected_percent) => {}
|
||||
_ => *selected = Some((priority, used_percent)),
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_window_priority(window: &Value, fallback: QuotaWindowPriority) -> QuotaWindowPriority {
|
||||
let window_seconds = window
|
||||
.get("window_minutes")
|
||||
.and_then(Value::as_u64)
|
||||
.and_then(|minutes| minutes.checked_mul(60));
|
||||
match window_seconds {
|
||||
Some(seconds) if quota_window_matches(seconds, FIVE_HOUR_WINDOW_SECONDS) => {
|
||||
QuotaWindowPriority::FiveHour
|
||||
}
|
||||
Some(seconds) if quota_window_matches(seconds, WEEK_WINDOW_SECONDS) => {
|
||||
QuotaWindowPriority::Weekly
|
||||
}
|
||||
Some(seconds) if quota_window_matches(seconds, MONTH_WINDOW_SECONDS) => {
|
||||
QuotaWindowPriority::Monthly
|
||||
}
|
||||
_ => fallback,
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_window_matches(actual: u64, expected: u64) -> bool {
|
||||
actual > 0 && actual.abs_diff(expected) <= expected / 20
|
||||
}
|
||||
|
||||
fn percent_value(value: &Value) -> Option<f64> {
|
||||
@@ -931,8 +974,67 @@ fn percent_value(value: &Value) -> Option<f64> {
|
||||
percent.is_finite().then_some(percent.clamp(0.0, 100.0))
|
||||
}
|
||||
|
||||
fn dashboard_usage_record_items(
|
||||
items: Vec<UsageRecord>,
|
||||
account_emails: &HashMap<String, String>,
|
||||
) -> Vec<DashboardUsageRecordData> {
|
||||
items
|
||||
.into_iter()
|
||||
.map(|record| dashboard_usage_record_data(record, account_emails))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn dashboard_usage_record_data(
|
||||
record: UsageRecord,
|
||||
account_emails: &HashMap<String, String>,
|
||||
) -> DashboardUsageRecordData {
|
||||
let account_email = record
|
||||
.account_id
|
||||
.as_deref()
|
||||
.and_then(|account_id| account_emails.get(account_id))
|
||||
.cloned()
|
||||
.or_else(|| metadata_string(&record.metadata, &["accountEmail", "account_email"]));
|
||||
let (requested_model, upstream_model) = usage_record_models(&record);
|
||||
let client_ip = metadata_string(&record.metadata, &["clientIp", "ipAddress", "ip_address"]);
|
||||
let user_agent = metadata_string(&record.metadata, &["userAgent", "user_agent"]);
|
||||
let reasoning_effort =
|
||||
metadata_string(&record.metadata, &["reasoningEffort", "reasoning_effort"]);
|
||||
let token_details = usage_token_details(&record.metadata);
|
||||
let cost_details = usage_cost_details(&record, upstream_model.as_deref(), &token_details);
|
||||
let first_token_latency_ms = metadata_nonnegative_i64(
|
||||
&record.metadata,
|
||||
&[
|
||||
"firstTokenMs",
|
||||
"first_token_ms",
|
||||
"firstTokenLatencyMs",
|
||||
"first_token_latency_ms",
|
||||
],
|
||||
);
|
||||
|
||||
DashboardUsageRecordData {
|
||||
id: record.id,
|
||||
level: record.level,
|
||||
route: record.route,
|
||||
model: record.model,
|
||||
status_code: record.status_code,
|
||||
transport: record.transport,
|
||||
created_at_display: china_datetime(&record.created_at),
|
||||
account_email,
|
||||
requested_model,
|
||||
upstream_model,
|
||||
client_ip,
|
||||
user_agent,
|
||||
reasoning_effort,
|
||||
token_details,
|
||||
cost_details,
|
||||
first_token_latency_ms_display: format_duration_ms(first_token_latency_ms),
|
||||
latency_ms_display: format_duration_ms(record.latency_ms),
|
||||
}
|
||||
}
|
||||
|
||||
fn service_status_data(state: &AppState) -> Vec<DashboardServiceStatusData> {
|
||||
let fingerprint = fingerprint_diagnostics(&state.services.fingerprint);
|
||||
let current_fingerprint = state.services.fingerprint.snapshot();
|
||||
let fingerprint = fingerprint_diagnostics(¤t_fingerprint);
|
||||
vec![
|
||||
DashboardServiceStatusData {
|
||||
label: "客户端版本".to_string(),
|
||||
@@ -993,20 +1095,3 @@ fn format_fingerprint_updated_at(value: Option<String>) -> String {
|
||||
}
|
||||
china_datetime_rfc3339_str(value)
|
||||
}
|
||||
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -463,7 +463,7 @@ fn log_error(error: &AdminUsageRecordError) -> AdminError {
|
||||
| AdminUsageRecordError::Get
|
||||
| AdminUsageRecordError::Clear
|
||||
| AdminUsageRecordError::Append
|
||||
| AdminUsageRecordError::Trim => AdminError::internal(error.to_string()),
|
||||
| AdminUsageRecordError::Retention => AdminError::internal(error.to_string()),
|
||||
AdminUsageRecordError::Accounts => {
|
||||
AdminError::usage_record_accounts_failed(error.to_string())
|
||||
}
|
||||
@@ -545,7 +545,7 @@ fn usage_record_data(
|
||||
}
|
||||
}
|
||||
|
||||
fn usage_record_models(record: &UsageRecord) -> (Option<String>, Option<String>) {
|
||||
pub(crate) fn usage_record_models(record: &UsageRecord) -> (Option<String>, Option<String>) {
|
||||
let stored_model = record
|
||||
.model
|
||||
.as_deref()
|
||||
@@ -560,7 +560,7 @@ fn usage_record_models(record: &UsageRecord) -> (Option<String>, Option<String>)
|
||||
(requested_model, upstream_model)
|
||||
}
|
||||
|
||||
fn usage_token_details(metadata: &Value) -> UsageRecordTokenDetailsData {
|
||||
pub(crate) fn usage_token_details(metadata: &Value) -> UsageRecordTokenDetailsData {
|
||||
let input_tokens = metadata_usage_u64(metadata, "inputTokens");
|
||||
let output_tokens = metadata_usage_u64(metadata, "outputTokens");
|
||||
let cached_tokens = metadata_usage_u64(metadata, "cachedTokens");
|
||||
@@ -582,7 +582,7 @@ fn usage_token_details(metadata: &Value) -> UsageRecordTokenDetailsData {
|
||||
}
|
||||
}
|
||||
|
||||
fn usage_cost_details(
|
||||
pub(crate) fn usage_cost_details(
|
||||
record: &UsageRecord,
|
||||
upstream_model: Option<&str>,
|
||||
tokens: &UsageRecordTokenDetailsData,
|
||||
|
||||
@@ -11,8 +11,9 @@ use crate::{
|
||||
usage_record_model::metadata_service_tier,
|
||||
usage_record_model::{UsageRecord, UsageRecordLevel},
|
||||
usage_record_store::{
|
||||
SqliteUsageRecordStore, UsageRecordBreakdown, UsageRecordEndpointSource,
|
||||
UsageRecordFilter, UsageRecordModelSource, UsageRecordSummary, UsageRecordTrendPoint,
|
||||
SqliteUsageRecordStore, UsageRecordAccountUsage, UsageRecordBreakdown,
|
||||
UsageRecordEndpointSource, UsageRecordFilter, UsageRecordModelSource,
|
||||
UsageRecordSummary, UsageRecordTrendPoint,
|
||||
},
|
||||
},
|
||||
infra::json::{NumberedPage, Page},
|
||||
@@ -116,9 +117,9 @@ pub enum AdminUsageRecordError {
|
||||
/// 写入失败。
|
||||
#[error("failed to append usage record")]
|
||||
Append,
|
||||
/// 裁剪失败。
|
||||
#[error("failed to trim usage records")]
|
||||
Trim,
|
||||
/// 保留期清理失败。
|
||||
#[error("failed to trim expired usage records")]
|
||||
Retention,
|
||||
/// 账号关联信息读取失败。
|
||||
#[error("failed to load usage record accounts")]
|
||||
Accounts,
|
||||
@@ -134,23 +135,16 @@ pub struct AdminUsageRecordService {
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct AdminUsageRecordSettings {
|
||||
enabled: bool,
|
||||
capacity: u32,
|
||||
capture_body: bool,
|
||||
}
|
||||
|
||||
impl AdminUsageRecordService {
|
||||
/// 构造管理端使用记录服务。
|
||||
pub fn new(
|
||||
store: SqliteUsageRecordStore,
|
||||
enabled: bool,
|
||||
capacity: u32,
|
||||
capture_body: bool,
|
||||
) -> Self {
|
||||
pub fn new(store: SqliteUsageRecordStore, enabled: bool, capture_body: bool) -> Self {
|
||||
Self {
|
||||
store,
|
||||
settings: Arc::new(tokio::sync::RwLock::new(AdminUsageRecordSettings {
|
||||
enabled,
|
||||
capacity,
|
||||
capture_body,
|
||||
})),
|
||||
}
|
||||
@@ -212,6 +206,18 @@ impl AdminUsageRecordService {
|
||||
.map_err(|_| AdminUsageRecordError::List)
|
||||
}
|
||||
|
||||
/// 按账号聚合使用记录。
|
||||
pub async fn account_usage(
|
||||
&self,
|
||||
filter: AdminUsageRecordFilter,
|
||||
limit: u32,
|
||||
) -> Result<Vec<UsageRecordAccountUsage>, AdminUsageRecordError> {
|
||||
self.store
|
||||
.account_usage(filter.into(), limit)
|
||||
.await
|
||||
.map_err(|_| AdminUsageRecordError::List)
|
||||
}
|
||||
|
||||
/// 按模型来源聚合使用记录分布。
|
||||
pub async fn model_distribution(
|
||||
&self,
|
||||
@@ -276,9 +282,9 @@ impl AdminUsageRecordService {
|
||||
.await
|
||||
.map_err(|_| AdminUsageRecordError::Append)?;
|
||||
self.store
|
||||
.trim_to_capacity(settings.capacity)
|
||||
.trim_to_retention(Utc::now())
|
||||
.await
|
||||
.map_err(|_| AdminUsageRecordError::Trim)?;
|
||||
.map_err(|_| AdminUsageRecordError::Retention)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
use std::collections::{BTreeMap, HashMap};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use serde_json::Value;
|
||||
use sqlx::{QueryBuilder, Row, Sqlite, SqlitePool};
|
||||
use thiserror::Error;
|
||||
@@ -19,8 +19,8 @@ use crate::infra::{
|
||||
time::{china_quarter_hour_start, parse_rfc3339_utc as parse_rfc3339},
|
||||
};
|
||||
|
||||
/// 默认保留最新使用记录数量。
|
||||
pub const DEFAULT_USAGE_RECORD_CAPACITY: u32 = 2_000;
|
||||
/// 使用记录保留天数。
|
||||
pub const USAGE_RECORD_RETENTION_DAYS: i64 = 30;
|
||||
/// 默认不保存请求/响应体,避免把用户内容写入诊断日志。
|
||||
pub const DEFAULT_USAGE_RECORD_CAPTURE_BODY: bool = false;
|
||||
|
||||
@@ -105,9 +105,9 @@ impl SqliteUsageRecordStore {
|
||||
append_event(&self.pool, event).await
|
||||
}
|
||||
|
||||
/// 保留最新的指定数量使用记录。
|
||||
pub async fn trim_to_capacity(&self, capacity: u32) -> SqliteUsageRecordStoreResult<u64> {
|
||||
trim_to_capacity(&self.pool, capacity).await
|
||||
/// 按保留期清理使用记录。
|
||||
pub async fn trim_to_retention(&self, now: DateTime<Utc>) -> SqliteUsageRecordStoreResult<u64> {
|
||||
trim_to_retention(&self.pool, now).await
|
||||
}
|
||||
|
||||
/// 分页查询使用记录。
|
||||
@@ -198,6 +198,15 @@ impl SqliteUsageRecordStore {
|
||||
usage_summary(&self.pool, &filter).await
|
||||
}
|
||||
|
||||
/// 按账号聚合使用记录。
|
||||
pub async fn account_usage(
|
||||
&self,
|
||||
filter: UsageRecordFilter,
|
||||
limit: u32,
|
||||
) -> SqliteUsageRecordStoreResult<Vec<UsageRecordAccountUsage>> {
|
||||
usage_account_usage(&self.pool, &filter, limit).await
|
||||
}
|
||||
|
||||
/// 按模型来源聚合使用记录分布。
|
||||
pub async fn model_distribution(
|
||||
&self,
|
||||
@@ -277,7 +286,6 @@ fn usage_record_summary_fields(event: &UsageRecord) -> UsageRecordSummaryFields
|
||||
"upstreamRequestId",
|
||||
"upstream_request_id",
|
||||
"openaiRequestId",
|
||||
"cfRay",
|
||||
],
|
||||
)
|
||||
}),
|
||||
@@ -417,17 +425,15 @@ fn metadata_first_token_latency(metadata: &Value) -> Option<i64> {
|
||||
})
|
||||
}
|
||||
|
||||
async fn trim_to_capacity(pool: &SqlitePool, capacity: u32) -> SqliteUsageRecordStoreResult<u64> {
|
||||
let result = sqlx::query(
|
||||
"delete from usage_records where id in (
|
||||
select id from usage_records
|
||||
order by created_at desc, id desc
|
||||
limit -1 offset ?
|
||||
)",
|
||||
)
|
||||
.bind(i64::from(capacity))
|
||||
.execute(pool)
|
||||
.await?;
|
||||
async fn trim_to_retention(
|
||||
pool: &SqlitePool,
|
||||
now: DateTime<Utc>,
|
||||
) -> SqliteUsageRecordStoreResult<u64> {
|
||||
let cutoff = now - Duration::days(USAGE_RECORD_RETENTION_DAYS);
|
||||
let result = sqlx::query("delete from usage_records where created_at < ?")
|
||||
.bind(cutoff.to_rfc3339())
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected())
|
||||
}
|
||||
|
||||
@@ -477,6 +483,45 @@ async fn usage_summary(
|
||||
})
|
||||
}
|
||||
|
||||
async fn usage_account_usage(
|
||||
pool: &SqlitePool,
|
||||
filter: &UsageRecordFilter,
|
||||
limit: u32,
|
||||
) -> SqliteUsageRecordStoreResult<Vec<UsageRecordAccountUsage>> {
|
||||
let mut builder = QueryBuilder::<Sqlite>::new(
|
||||
"select
|
||||
account_id,
|
||||
sum(coalesce(cast(json_extract(metadata_json, '$.usage.inputTokens') as integer), cast(json_extract(metadata_json, '$.inputTokens') as integer), 0)) as input_tokens,
|
||||
sum(coalesce(cast(json_extract(metadata_json, '$.usage.outputTokens') as integer), cast(json_extract(metadata_json, '$.outputTokens') as integer), 0)) as output_tokens,
|
||||
max(created_at) as last_used_at
|
||||
from usage_records",
|
||||
);
|
||||
push_filter(&mut builder, filter, None)?;
|
||||
builder.push(" and account_id is not null and trim(account_id) <> ''");
|
||||
builder.push(" group by account_id order by last_used_at desc, account_id asc limit ");
|
||||
builder.push_bind(i64::from(limit.clamp(1, 50)));
|
||||
|
||||
builder
|
||||
.build()
|
||||
.fetch_all(pool)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
let input_tokens =
|
||||
optional_nonnegative_i64_to_u64(row.get::<Option<i64>, _>("input_tokens"));
|
||||
let output_tokens =
|
||||
optional_nonnegative_i64_to_u64(row.get::<Option<i64>, _>("output_tokens"));
|
||||
Ok(UsageRecordAccountUsage {
|
||||
account_id: row.get("account_id"),
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens: input_tokens + output_tokens,
|
||||
last_used_at: parse_rfc3339(&row.get::<String, _>("last_used_at"))?,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn usage_record_account_email_map(
|
||||
pool: &SqlitePool,
|
||||
items: &[UsageRecord],
|
||||
@@ -978,6 +1023,21 @@ pub struct UsageRecordSummary {
|
||||
pub average_latency_ms: Option<f64>,
|
||||
}
|
||||
|
||||
/// 账号使用记录聚合。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UsageRecordAccountUsage {
|
||||
/// 账号 ID。
|
||||
pub account_id: String,
|
||||
/// 输入 Token。
|
||||
pub input_tokens: u64,
|
||||
/// 输出 Token。
|
||||
pub output_tokens: u64,
|
||||
/// 总 Token。
|
||||
pub total_tokens: u64,
|
||||
/// 最近使用时间。
|
||||
pub last_used_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// 模型分布来源。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum UsageRecordModelSource {
|
||||
|
||||
@@ -16,9 +16,8 @@ use crate::{
|
||||
upstream::models::config::ModelConfig,
|
||||
};
|
||||
|
||||
const ALLOWED_SETTINGS_KEYS: [&str; 7] = [
|
||||
const ALLOWED_SETTINGS_KEYS: [&str; 6] = [
|
||||
"modelAliases",
|
||||
"modelAccountRoutes",
|
||||
"refreshMarginSeconds",
|
||||
"refreshConcurrency",
|
||||
"maxConcurrentPerAccount",
|
||||
@@ -30,7 +29,6 @@ const ALLOWED_SETTINGS_KEYS: [&str; 7] = [
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AdminSettingsData {
|
||||
model_aliases: BTreeMap<String, String>,
|
||||
model_account_routes: BTreeMap<String, Vec<String>>,
|
||||
refresh_margin_seconds: u64,
|
||||
refresh_concurrency: u32,
|
||||
max_concurrent_per_account: usize,
|
||||
@@ -42,7 +40,6 @@ impl AdminSettingsData {
|
||||
fn from_config(config: &AppConfig) -> Self {
|
||||
Self {
|
||||
model_aliases: config.model_aliases.clone(),
|
||||
model_account_routes: config.model_account_routes.clone(),
|
||||
refresh_margin_seconds: config.auth.refresh_margin_seconds,
|
||||
refresh_concurrency: config.auth.refresh_concurrency,
|
||||
max_concurrent_per_account: config.auth.max_concurrent_per_account,
|
||||
|
||||
@@ -24,9 +24,6 @@ pub struct AppConfig {
|
||||
/// 运行时模型别名,由数据库设置加载,配置文件不承载该字段。
|
||||
#[serde(skip)]
|
||||
pub model_aliases: BTreeMap<String, String>,
|
||||
/// 运行时模型账号路由,由 `model_account_routes` 表加载,配置文件不承载该字段。
|
||||
#[serde(skip)]
|
||||
pub model_account_routes: BTreeMap<String, Vec<String>>,
|
||||
/// 认证与刷新配置,由固定默认值和数据库运行时设置承载,配置文件不承载该字段。
|
||||
#[serde(skip)]
|
||||
pub auth: AuthConfig,
|
||||
@@ -60,7 +57,7 @@ pub struct ServerConfig {
|
||||
pub host: String,
|
||||
/// 监听端口。
|
||||
pub port: u16,
|
||||
/// 可信反向代理 IP 或 CIDR。只有这些 peer 的转发 IP 头会被采信。
|
||||
/// 可信反向代理 IP 或 CIDR。非空时只有这些 peer 的转发 IP 头会被采信。
|
||||
#[serde(default)]
|
||||
pub trusted_proxies: Vec<String>,
|
||||
}
|
||||
@@ -103,7 +100,7 @@ impl Default for AuthConfig {
|
||||
refresh_concurrency: 2,
|
||||
max_concurrent_per_account: 3,
|
||||
request_interval_ms: 50,
|
||||
rotation_strategy: "least_used".to_string(),
|
||||
rotation_strategy: "smart".to_string(),
|
||||
tier_priority: Vec::new(),
|
||||
oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(),
|
||||
oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(),
|
||||
@@ -164,6 +161,17 @@ pub struct WebSocketPoolConfig {
|
||||
pub max_age_ms: u64,
|
||||
/// 单账号最大连接数。
|
||||
pub max_per_account: usize,
|
||||
/// 首个内容帧(首 token)到达前的绝对超时;`0` 表示禁用。
|
||||
///
|
||||
/// 从发出 `response.create` 起算,覆盖建连/发送/上游排队/思考的全程,直到收到
|
||||
/// 首个真实内容帧(`response.created`/`response.in_progress` 等纯生命周期帧不计)。
|
||||
/// 超时判定为连接落到病态上游后端,丢弃并换新连接重试。
|
||||
#[serde(default = "default_ws_first_token_timeout_ms")]
|
||||
pub first_token_timeout_ms: u64,
|
||||
}
|
||||
|
||||
fn default_ws_first_token_timeout_ms() -> u64 {
|
||||
20_000
|
||||
}
|
||||
|
||||
impl Default for WebSocketPoolConfig {
|
||||
@@ -172,6 +180,7 @@ impl Default for WebSocketPoolConfig {
|
||||
enabled: true,
|
||||
max_age_ms: 3_300_000,
|
||||
max_per_account: 8,
|
||||
first_token_timeout_ms: default_ws_first_token_timeout_ms(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+10
-113
@@ -18,7 +18,7 @@ use crate::upstream::accounts::{
|
||||
token_refresh::{RefreshPolicy, RuntimeRefreshPolicy},
|
||||
};
|
||||
|
||||
const ROTATION_STRATEGIES: [&str; 3] = ["least_used", "round_robin", "sticky"];
|
||||
const ROTATION_STRATEGIES: [&str; 4] = ["smart", "quota_reset_priority", "round_robin", "sticky"];
|
||||
const RUNTIME_SETTINGS_ID: i64 = 1;
|
||||
|
||||
/// 管理端可变设置。
|
||||
@@ -26,8 +26,6 @@ const RUNTIME_SETTINGS_ID: i64 = 1;
|
||||
pub struct AdminSettings {
|
||||
/// 模型别名映射。
|
||||
pub model_aliases: BTreeMap<String, String>,
|
||||
/// 模型到账号 ID 的显式路由。
|
||||
pub model_account_routes: BTreeMap<String, Vec<String>>,
|
||||
/// 访问令牌过期前多少秒开始刷新。
|
||||
pub refresh_margin_seconds: u64,
|
||||
/// 访问令牌刷新并发数。
|
||||
@@ -44,12 +42,11 @@ impl Default for AdminSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
model_aliases: BTreeMap::new(),
|
||||
model_account_routes: BTreeMap::new(),
|
||||
refresh_margin_seconds: 300,
|
||||
refresh_concurrency: 2,
|
||||
max_concurrent_per_account: 3,
|
||||
request_interval_ms: 50,
|
||||
rotation_strategy: "least_used".to_string(),
|
||||
rotation_strategy: "smart".to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -60,8 +57,6 @@ impl Default for AdminSettings {
|
||||
pub struct AdminSettingsPatch {
|
||||
/// 模型别名映射。
|
||||
pub model_aliases: Option<BTreeMap<String, String>>,
|
||||
/// 模型到账号 ID 的显式路由。
|
||||
pub model_account_routes: Option<BTreeMap<String, Vec<String>>>,
|
||||
/// 访问令牌过期前多少秒开始刷新。
|
||||
pub refresh_margin_seconds: Option<u64>,
|
||||
/// 访问令牌刷新并发数。
|
||||
@@ -87,9 +82,6 @@ impl SettingsService {
|
||||
if let Some(model_aliases) = patch.model_aliases {
|
||||
current.model_aliases = validate_model_aliases(model_aliases)?;
|
||||
}
|
||||
if let Some(model_account_routes) = patch.model_account_routes {
|
||||
current.model_account_routes = validate_model_account_routes(model_account_routes)?;
|
||||
}
|
||||
if let Some(refresh_margin_seconds) = patch.refresh_margin_seconds {
|
||||
current.refresh_margin_seconds =
|
||||
positive_u64("refreshMarginSeconds", refresh_margin_seconds)?;
|
||||
@@ -151,34 +143,6 @@ fn validate_model_aliases(
|
||||
Ok(normalized)
|
||||
}
|
||||
|
||||
fn validate_model_account_routes(
|
||||
routes: BTreeMap<String, Vec<String>>,
|
||||
) -> Result<BTreeMap<String, Vec<String>>, SettingsServiceError> {
|
||||
let mut normalized = BTreeMap::new();
|
||||
for (model, account_ids) in routes {
|
||||
let model = non_empty_string("modelAccountRoutes", &model)?;
|
||||
let mut normalized_account_ids = Vec::new();
|
||||
for account_id in account_ids {
|
||||
let account_id = non_empty_string("modelAccountRoutes", &account_id)?;
|
||||
if normalized_account_ids.contains(&account_id) {
|
||||
return Err(invalid_field(
|
||||
"modelAccountRoutes",
|
||||
format!("duplicate account id `{account_id}`"),
|
||||
));
|
||||
}
|
||||
normalized_account_ids.push(account_id);
|
||||
}
|
||||
if normalized_account_ids.is_empty() {
|
||||
return Err(invalid_field(
|
||||
"modelAccountRoutes",
|
||||
"account list must not be empty",
|
||||
));
|
||||
}
|
||||
normalized.insert(model, normalized_account_ids);
|
||||
}
|
||||
Ok(normalized)
|
||||
}
|
||||
|
||||
fn validate_rotation_strategy(strategy: &str) -> Result<String, SettingsServiceError> {
|
||||
let strategy = non_empty_string("rotationStrategy", strategy)?;
|
||||
if ROTATION_STRATEGIES.contains(&strategy.as_str()) {
|
||||
@@ -186,7 +150,7 @@ fn validate_rotation_strategy(strategy: &str) -> Result<String, SettingsServiceE
|
||||
} else {
|
||||
Err(invalid_field(
|
||||
"rotationStrategy",
|
||||
"must be one of least_used, round_robin, sticky",
|
||||
"must be one of smart, quota_reset_priority, round_robin, sticky",
|
||||
))
|
||||
}
|
||||
}
|
||||
@@ -394,7 +358,6 @@ pub enum RuntimeSettingsError {
|
||||
fn admin_settings_from_config(config: &AppConfig) -> AdminSettings {
|
||||
AdminSettings {
|
||||
model_aliases: config.model_aliases.clone(),
|
||||
model_account_routes: config.model_account_routes.clone(),
|
||||
refresh_margin_seconds: config.auth.refresh_margin_seconds,
|
||||
refresh_concurrency: config.auth.refresh_concurrency,
|
||||
max_concurrent_per_account: config.auth.max_concurrent_per_account,
|
||||
@@ -405,7 +368,6 @@ fn admin_settings_from_config(config: &AppConfig) -> AdminSettings {
|
||||
|
||||
fn apply_admin_settings_to_config(config: &mut AppConfig, settings: AdminSettings) {
|
||||
config.model_aliases = settings.model_aliases;
|
||||
config.model_account_routes = settings.model_account_routes;
|
||||
config.auth.refresh_margin_seconds = settings.refresh_margin_seconds;
|
||||
config.auth.refresh_concurrency = settings.refresh_concurrency;
|
||||
config.auth.max_concurrent_per_account = settings.max_concurrent_per_account;
|
||||
@@ -418,8 +380,7 @@ async fn insert_runtime_settings_if_missing(
|
||||
settings: &AdminSettings,
|
||||
) -> Result<(), RuntimeSettingsError> {
|
||||
let model_aliases_json = serde_json::to_string(&settings.model_aliases)?;
|
||||
let mut tx = pool.begin().await?;
|
||||
let result = sqlx::query(
|
||||
sqlx::query(
|
||||
r"
|
||||
insert or ignore into runtime_settings (
|
||||
id,
|
||||
@@ -449,12 +410,8 @@ insert or ignore into runtime_settings (
|
||||
)
|
||||
.bind(&settings.rotation_strategy)
|
||||
.bind(Utc::now().to_rfc3339())
|
||||
.execute(&mut *tx)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
if result.rows_affected() > 0 {
|
||||
replace_model_account_routes(&mut tx, &settings.model_account_routes).await?;
|
||||
}
|
||||
tx.commit().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -463,7 +420,6 @@ async fn save_runtime_settings(
|
||||
settings: &AdminSettings,
|
||||
) -> Result<(), RuntimeSettingsError> {
|
||||
let model_aliases_json = serde_json::to_string(&settings.model_aliases)?;
|
||||
let mut tx = pool.begin().await?;
|
||||
sqlx::query(
|
||||
r"
|
||||
insert into runtime_settings (
|
||||
@@ -502,10 +458,8 @@ on conflict(id) do update set
|
||||
)
|
||||
.bind(&settings.rotation_strategy)
|
||||
.bind(Utc::now().to_rfc3339())
|
||||
.execute(&mut *tx)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
replace_model_account_routes(&mut tx, &settings.model_account_routes).await?;
|
||||
tx.commit().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -525,9 +479,7 @@ where id = ?",
|
||||
.bind(RUNTIME_SETTINGS_ID)
|
||||
.fetch_one(pool)
|
||||
.await?;
|
||||
let mut settings = runtime_settings_from_row(&row)?;
|
||||
settings.model_account_routes = load_model_account_routes(pool).await?;
|
||||
Ok(settings)
|
||||
runtime_settings_from_row(&row)
|
||||
}
|
||||
|
||||
async fn admin_api_key_status(
|
||||
@@ -566,7 +518,6 @@ fn runtime_settings_from_row(row: &SqliteRow) -> Result<AdminSettings, RuntimeSe
|
||||
let model_aliases_json: String = row.get("model_aliases_json");
|
||||
Ok(AdminSettings {
|
||||
model_aliases: serde_json::from_str(&model_aliases_json)?,
|
||||
model_account_routes: BTreeMap::new(),
|
||||
refresh_margin_seconds: positive_i64_to_u64(
|
||||
"refreshMarginSeconds",
|
||||
row.get("refresh_margin_seconds"),
|
||||
@@ -587,61 +538,6 @@ fn runtime_settings_from_row(row: &SqliteRow) -> Result<AdminSettings, RuntimeSe
|
||||
})
|
||||
}
|
||||
|
||||
async fn replace_model_account_routes(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
|
||||
routes: &BTreeMap<String, Vec<String>>,
|
||||
) -> Result<(), RuntimeSettingsError> {
|
||||
sqlx::query("delete from model_account_routes")
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
let now = Utc::now().to_rfc3339();
|
||||
for (model, account_ids) in routes {
|
||||
for (priority, account_id) in account_ids.iter().enumerate() {
|
||||
sqlx::query(
|
||||
r"
|
||||
insert into model_account_routes (
|
||||
model,
|
||||
account_id,
|
||||
priority,
|
||||
enabled,
|
||||
created_at,
|
||||
updated_at
|
||||
) values (?, ?, ?, 1, ?, ?)",
|
||||
)
|
||||
.bind(model)
|
||||
.bind(account_id)
|
||||
.bind(
|
||||
i64::try_from(priority)
|
||||
.map_err(|_| stored_field_error("modelAccountRoutes", "out of range"))?,
|
||||
)
|
||||
.bind(&now)
|
||||
.bind(&now)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn load_model_account_routes(
|
||||
pool: &SqlitePool,
|
||||
) -> Result<BTreeMap<String, Vec<String>>, RuntimeSettingsError> {
|
||||
let rows: Vec<(String, String)> = sqlx::query_as(
|
||||
r"
|
||||
select model, account_id
|
||||
from model_account_routes
|
||||
where enabled = 1
|
||||
order by model asc, priority asc, account_id asc",
|
||||
)
|
||||
.fetch_all(pool)
|
||||
.await?;
|
||||
let mut routes = BTreeMap::<String, Vec<String>>::new();
|
||||
for (model, account_id) in rows {
|
||||
routes.entry(model).or_default().push(account_id);
|
||||
}
|
||||
Ok(routes)
|
||||
}
|
||||
|
||||
/// 从当前运行配置生成账号池调度参数。
|
||||
pub fn account_pool_options_from_config(config: &AppConfig) -> AccountPoolOptions {
|
||||
AccountPoolOptions {
|
||||
@@ -649,16 +545,17 @@ pub fn account_pool_options_from_config(config: &AppConfig) -> AccountPoolOption
|
||||
rotation_strategy: rotation_strategy_from_config(&config.auth.rotation_strategy),
|
||||
skip_quota_limited: config.quota.skip_exhausted,
|
||||
tier_priority: config.auth.tier_priority.clone(),
|
||||
model_account_routes: config.model_account_routes.clone(),
|
||||
..AccountPoolOptions::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn rotation_strategy_from_config(strategy: &str) -> RotationStrategy {
|
||||
match strategy {
|
||||
"smart" => RotationStrategy::Smart,
|
||||
"quota_reset_priority" => RotationStrategy::QuotaResetPriority,
|
||||
"round_robin" => RotationStrategy::RoundRobin,
|
||||
"sticky" => RotationStrategy::Sticky,
|
||||
_ => RotationStrategy::LeastUsed,
|
||||
_ => RotationStrategy::Smart,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -214,8 +214,9 @@ fn forwarded_for_public_ip(headers: &HeaderMap) -> Option<String> {
|
||||
|
||||
fn trusted_forwarded_client_ip(headers: &HeaderMap) -> Option<String> {
|
||||
header_string(headers, CF_CONNECTING_IP_HEADER)
|
||||
.or_else(|| first_forwarded_for_value(headers))
|
||||
.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> {
|
||||
|
||||
@@ -4,12 +4,19 @@ use std::{fs, path::Path, str::FromStr, time::Duration};
|
||||
|
||||
use sqlx::{sqlite::SqliteConnectOptions, SqlitePool};
|
||||
|
||||
const CURRENT_SCHEMA_VERSION: i64 = 1;
|
||||
const MIGRATIONS: &[Migration] = &[Migration {
|
||||
version: 1,
|
||||
name: "initial",
|
||||
sql: include_str!("migrations/0001_initial.sql"),
|
||||
}];
|
||||
const CURRENT_SCHEMA_VERSION: i64 = 2;
|
||||
const MIGRATIONS: &[Migration] = &[
|
||||
Migration {
|
||||
version: 1,
|
||||
name: "initial",
|
||||
sql: include_str!("migrations/0001_initial.sql"),
|
||||
},
|
||||
Migration {
|
||||
version: 2,
|
||||
name: "rename_rotation_strategies",
|
||||
sql: include_str!("migrations/0002_rename_rotation_strategies.sql"),
|
||||
},
|
||||
];
|
||||
|
||||
struct Migration {
|
||||
version: i64,
|
||||
|
||||
@@ -49,7 +49,7 @@ create table runtime_settings (
|
||||
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 ('least_used', 'round_robin', 'sticky')),
|
||||
rotation_strategy text not null check (rotation_strategy in ('smart', 'quota_reset_priority', 'round_robin', 'sticky')),
|
||||
admin_api_key text,
|
||||
updated_at text not null
|
||||
);
|
||||
@@ -176,24 +176,6 @@ on usage_time_buckets(bucket_start);
|
||||
create index idx_usage_time_buckets_model_bucket
|
||||
on usage_time_buckets(model, bucket_start);
|
||||
|
||||
-- ============================================
|
||||
-- Model Account Routes
|
||||
-- ============================================
|
||||
create table model_account_routes (
|
||||
model text not null,
|
||||
account_id text not null references accounts(id) on delete cascade,
|
||||
priority integer not null default 0 check (priority >= 0),
|
||||
enabled integer not null default 1 check (enabled in (0, 1)),
|
||||
created_at text not null,
|
||||
updated_at text not null,
|
||||
primary key (model, account_id)
|
||||
);
|
||||
|
||||
create index idx_model_account_routes_account
|
||||
on model_account_routes(account_id);
|
||||
create index idx_model_account_routes_enabled_model
|
||||
on model_account_routes(enabled, model, priority asc, account_id);
|
||||
|
||||
-- ============================================
|
||||
-- Account Cookies
|
||||
-- ============================================
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
create table runtime_settings_new (
|
||||
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 text,
|
||||
updated_at text not null
|
||||
);
|
||||
|
||||
insert into runtime_settings_new (
|
||||
id,
|
||||
model_aliases_json,
|
||||
refresh_margin_seconds,
|
||||
refresh_concurrency,
|
||||
max_concurrent_per_account,
|
||||
request_interval_ms,
|
||||
rotation_strategy,
|
||||
admin_api_key,
|
||||
updated_at
|
||||
)
|
||||
select
|
||||
id,
|
||||
model_aliases_json,
|
||||
refresh_margin_seconds,
|
||||
refresh_concurrency,
|
||||
max_concurrent_per_account,
|
||||
request_interval_ms,
|
||||
case rotation_strategy
|
||||
when 'least_used' then 'quota_reset_priority'
|
||||
else rotation_strategy
|
||||
end,
|
||||
admin_api_key,
|
||||
updated_at
|
||||
from runtime_settings;
|
||||
|
||||
drop table runtime_settings;
|
||||
alter table runtime_settings_new rename to runtime_settings;
|
||||
@@ -90,60 +90,3 @@ fn parse_installation_id(raw: &str) -> Option<String> {
|
||||
fn generate_installation_id() -> String {
|
||||
Uuid::new_v4().to_string()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Mutex;
|
||||
|
||||
use super::*;
|
||||
|
||||
static ENV_LOCK: Mutex<()> = Mutex::new(());
|
||||
|
||||
#[test]
|
||||
fn data_dir_should_use_xdg_data_home_directly() {
|
||||
let _guard = ENV_LOCK.lock().unwrap();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let previous = std::env::var_os("XDG_DATA_HOME");
|
||||
std::env::set_var("XDG_DATA_HOME", dir.path());
|
||||
|
||||
let resolved = data_dir();
|
||||
|
||||
restore_env("XDG_DATA_HOME", previous);
|
||||
assert_eq!(resolved, dir.path());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_or_create_installation_id_should_persist_under_data_dir_root() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let data_path = dir.path().join(INSTALLATION_ID_FILE_NAME);
|
||||
|
||||
let installation_id = load_or_create_installation_id_from_paths(None, Some(&data_path))
|
||||
.expect("installation id should be generated");
|
||||
|
||||
assert_eq!(fs::read_to_string(data_path).unwrap(), installation_id);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_or_create_installation_id_should_seed_data_dir_from_codex_desktop_id() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let codex_path = dir.path().join(".codex").join(INSTALLATION_ID_FILE_NAME);
|
||||
let data_path = dir.path().join("data").join(INSTALLATION_ID_FILE_NAME);
|
||||
let desktop_id = "018f8f6b-1d7b-7b7c-b9c8-8c9f4c6d0e1a";
|
||||
persist_installation_id(&codex_path, desktop_id).unwrap();
|
||||
|
||||
let installation_id =
|
||||
load_or_create_installation_id_from_paths(Some(&codex_path), Some(&data_path))
|
||||
.expect("installation id should be seeded");
|
||||
|
||||
assert_eq!(installation_id, desktop_id);
|
||||
assert_eq!(fs::read_to_string(data_path).unwrap(), desktop_id);
|
||||
}
|
||||
|
||||
fn restore_env(key: &str, value: Option<std::ffi::OsString>) {
|
||||
if let Some(value) = value {
|
||||
std::env::set_var(key, value);
|
||||
} else {
|
||||
std::env::remove_var(key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,7 +6,9 @@ use serde_json::{json, Map, Value};
|
||||
use crate::{
|
||||
proxy::dispatch::exhaustion::{ExhaustedAccountKind, ExhaustedAccountRef},
|
||||
upstream::accounts::model::AccountStatus,
|
||||
upstream::transport::{is_banned_auth_signal, CodexBackendTransport, CodexClientError},
|
||||
upstream::transport::{
|
||||
is_banned_auth_signal, CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics,
|
||||
},
|
||||
};
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
@@ -15,6 +17,7 @@ pub(crate) struct DispatchErrorMetadata {
|
||||
pub exhausted_count: Option<usize>,
|
||||
pub upstream_error: Option<String>,
|
||||
pub upstream_status: Option<u16>,
|
||||
pub diagnostics: Option<CodexUpstreamDiagnostics>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
@@ -63,6 +66,7 @@ impl DispatchErrorMetadata {
|
||||
exhausted_count: None,
|
||||
upstream_error: None,
|
||||
upstream_status: None,
|
||||
diagnostics: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -72,6 +76,7 @@ impl DispatchErrorMetadata {
|
||||
exhausted_count: Some(exhausted.count),
|
||||
upstream_error: Some(exhausted.upstream_error.to_string()),
|
||||
upstream_status: None,
|
||||
diagnostics: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -84,6 +89,7 @@ impl DispatchErrorMetadata {
|
||||
CodexClientError::Upstream { status, .. } => Some(status.as_u16()),
|
||||
_ => None,
|
||||
},
|
||||
diagnostics: upstream_error_diagnostics(error).cloned(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -93,6 +99,7 @@ impl DispatchErrorMetadata {
|
||||
exhausted_count: None,
|
||||
upstream_error: None,
|
||||
upstream_status: None,
|
||||
diagnostics: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -129,6 +136,45 @@ pub(crate) fn insert_dispatch_error_metadata(
|
||||
if let Some(status) = metadata.upstream_status {
|
||||
object.insert("upstreamStatus".to_string(), json!(status));
|
||||
}
|
||||
if let Some(diagnostics) = metadata.diagnostics {
|
||||
insert_upstream_diagnostics_metadata(object, &diagnostics);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn insert_upstream_diagnostics_metadata(
|
||||
object: &mut Map<String, Value>,
|
||||
diagnostics: &CodexUpstreamDiagnostics,
|
||||
) {
|
||||
if diagnostics.is_empty() {
|
||||
return;
|
||||
}
|
||||
if let Some(status_code) = diagnostics.status_code {
|
||||
object.insert("upstreamStatus".to_string(), json!(status_code));
|
||||
}
|
||||
if let Some(request_id) = &diagnostics.request_id {
|
||||
object.insert(
|
||||
"upstreamRequestId".to_string(),
|
||||
Value::String(request_id.clone()),
|
||||
);
|
||||
}
|
||||
if !diagnostics.trace_headers.is_empty() {
|
||||
object.insert(
|
||||
"upstreamTraceHeaders".to_string(),
|
||||
json!(diagnostics.trace_headers),
|
||||
);
|
||||
}
|
||||
if let Some(cf_ray) = diagnostics.cf_ray() {
|
||||
object.insert("cfRay".to_string(), Value::String(cf_ray.to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn upstream_error_diagnostics(
|
||||
error: &CodexClientError,
|
||||
) -> Option<&CodexUpstreamDiagnostics> {
|
||||
match error {
|
||||
CodexClientError::Upstream { diagnostics, .. } => Some(diagnostics),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn is_rate_limit_upstream_error(error: &CodexClientError) -> bool {
|
||||
|
||||
@@ -199,6 +199,7 @@ impl ResponseDispatchError {
|
||||
exhausted_count: None,
|
||||
upstream_error: Some(sse_failure_error_body(failure)),
|
||||
upstream_status: None,
|
||||
diagnostics: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use std::time::Instant;
|
||||
|
||||
use serde_json::{json, Value};
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use crate::{
|
||||
admin::monitoring::{
|
||||
@@ -9,7 +9,11 @@ use crate::{
|
||||
},
|
||||
infra::time::elapsed_millis_i64,
|
||||
proxy::dispatch::{
|
||||
errors::{backend_transport_name, upstream_error_http_status},
|
||||
errors::{
|
||||
backend_transport_name, insert_dispatch_error_metadata,
|
||||
insert_upstream_diagnostics_metadata, upstream_error_http_status,
|
||||
DispatchErrorMetadata,
|
||||
},
|
||||
usage_events::{
|
||||
enrich_event_route_metadata, enrich_usage_record_identity,
|
||||
event_kind as response_event_kind, reasoning_effort_from_request,
|
||||
@@ -18,7 +22,10 @@ use crate::{
|
||||
},
|
||||
upstream::{
|
||||
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
|
||||
transport::{CodexBackendTransport, CodexClientError, WebSocketPoolDecision},
|
||||
transport::{
|
||||
CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics,
|
||||
WebSocketPoolDecision,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
@@ -28,6 +35,7 @@ use super::{
|
||||
},
|
||||
sse_failure::{status_code_for_stream_failure, stream_failure_metadata},
|
||||
stream_lifecycle::{latest_response_id, LiveResponseStreamContext},
|
||||
trace::{ResponseDispatchAttempt, ResponseDispatchTrace},
|
||||
};
|
||||
|
||||
pub(super) struct ResponseUpstreamErrorEventRecord<'a> {
|
||||
@@ -40,7 +48,10 @@ pub(super) struct ResponseUpstreamErrorEventRecord<'a> {
|
||||
pub(super) started_at: Instant,
|
||||
pub(super) stream: bool,
|
||||
pub(super) transport: CodexBackendTransport,
|
||||
pub(super) request: &'a CodexResponsesRequest,
|
||||
pub(super) error: &'a CodexClientError,
|
||||
pub(super) trace: &'a ResponseDispatchTrace,
|
||||
pub(super) attempt: Option<&'a ResponseDispatchAttempt>,
|
||||
}
|
||||
|
||||
pub(super) struct ResponseStreamFailureEventRecord<'a> {
|
||||
@@ -54,8 +65,11 @@ pub(super) struct ResponseStreamFailureEventRecord<'a> {
|
||||
pub(super) transport: CodexBackendTransport,
|
||||
pub(super) request: &'a CodexResponsesRequest,
|
||||
pub(super) failure: &'a ResponsesSseFailure,
|
||||
pub(super) diagnostics: &'a CodexUpstreamDiagnostics,
|
||||
pub(super) rate_limit_headers: &'a [(String, String)],
|
||||
pub(super) prefetched: &'a [u8],
|
||||
pub(super) trace: &'a ResponseDispatchTrace,
|
||||
pub(super) attempt: &'a ResponseDispatchAttempt,
|
||||
}
|
||||
|
||||
pub(super) struct ResponseDispatchErrorEventRecord<'a> {
|
||||
@@ -88,6 +102,14 @@ pub(super) async fn record_response_dispatch_error_event(
|
||||
record.transport,
|
||||
);
|
||||
enrich_response_dispatch_error_metadata(&mut metadata, record.error);
|
||||
if let Some(object) = metadata.as_object_mut() {
|
||||
insert_response_status_metadata_object(
|
||||
object,
|
||||
i64::from(record.error.http_status_code()),
|
||||
dispatch_error_client_status_code(record.error, record.stream, record.compact),
|
||||
dispatch_error_upstream_status_code(record.error),
|
||||
);
|
||||
}
|
||||
record_dispatch_error_event(DispatchErrorUsageRecord {
|
||||
usage_records: record.usage_records,
|
||||
request_id: record.request_id,
|
||||
@@ -105,6 +127,7 @@ pub(super) async fn record_response_dispatch_error_event(
|
||||
pub(super) async fn record_response_upstream_error_event(
|
||||
record: ResponseUpstreamErrorEventRecord<'_>,
|
||||
) {
|
||||
let event_status_code = i64::from(upstream_error_http_status(record.error));
|
||||
let mut metadata = dispatch_error_metadata(
|
||||
record.error,
|
||||
record.stream,
|
||||
@@ -121,9 +144,18 @@ pub(super) async fn record_response_upstream_error_event(
|
||||
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(i64::from(upstream_error_http_status(record.error)));
|
||||
event.status_code = Some(event_status_code);
|
||||
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
|
||||
if let Some(object) = metadata.as_object_mut() {
|
||||
insert_dispatch_error_metadata(object, DispatchErrorMetadata::upstream(record.error));
|
||||
insert_response_status_metadata_object(
|
||||
object,
|
||||
event_status_code,
|
||||
upstream_failure_client_status_code(record.stream),
|
||||
Some(event_status_code),
|
||||
);
|
||||
insert_response_request_summary_object(object, record.request, record.transport);
|
||||
insert_response_trace_metadata_object(object, record.trace.attempts(), record.attempt);
|
||||
if let Some(account_email) = record
|
||||
.account_email
|
||||
.map(str::trim)
|
||||
@@ -144,6 +176,7 @@ pub(super) async fn record_response_upstream_error_event(
|
||||
pub(super) async fn record_prefetched_response_stream_failure_event(
|
||||
record: ResponseStreamFailureEventRecord<'_>,
|
||||
) {
|
||||
let event_status_code = status_code_for_stream_failure(record.failure);
|
||||
let mut metadata = stream_failure_metadata(record.failure, None);
|
||||
if let Some(object) = metadata.as_object_mut() {
|
||||
if record.transport == CodexBackendTransport::WebSocket {
|
||||
@@ -152,11 +185,24 @@ pub(super) async fn record_prefetched_response_stream_failure_event(
|
||||
Value::String("websocket".to_string()),
|
||||
);
|
||||
}
|
||||
insert_response_status_metadata_object(
|
||||
object,
|
||||
event_status_code,
|
||||
200,
|
||||
diagnostics_status_code(record.diagnostics),
|
||||
);
|
||||
insert_response_request_summary_object(object, record.request, record.transport);
|
||||
object.insert("requestBody".to_string(), json!(record.request));
|
||||
object.insert(
|
||||
"responseBody".to_string(),
|
||||
Value::String(String::from_utf8_lossy(record.prefetched).to_string()),
|
||||
);
|
||||
insert_upstream_diagnostics_metadata(object, record.diagnostics);
|
||||
insert_response_trace_metadata_object(
|
||||
object,
|
||||
record.trace.attempts(),
|
||||
Some(record.attempt),
|
||||
);
|
||||
}
|
||||
record_response_event(ResponseUsageRecord {
|
||||
usage_records: record.usage_records,
|
||||
@@ -170,7 +216,7 @@ pub(super) async fn record_prefetched_response_stream_failure_event(
|
||||
reasoning_effort: reasoning_effort_from_request(record.request),
|
||||
service_tier: record.request.service_tier(),
|
||||
started_at: record.started_at,
|
||||
status_code: status_code_for_stream_failure(record.failure),
|
||||
status_code: event_status_code,
|
||||
level: UsageRecordLevel::Error,
|
||||
message: "v1 responses stream failed",
|
||||
metadata,
|
||||
@@ -193,6 +239,7 @@ fn enrich_live_response_stream_metadata(
|
||||
context: &LiveResponseStreamContext,
|
||||
rate_limit_headers: &[(String, String)],
|
||||
metadata: &mut Value,
|
||||
status_code: i64,
|
||||
body: &str,
|
||||
) {
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
@@ -211,6 +258,15 @@ fn enrich_live_response_stream_metadata(
|
||||
.entry("websocketPool".to_string())
|
||||
.or_insert_with(|| decision.metadata_value());
|
||||
}
|
||||
insert_response_status_metadata_object(
|
||||
object,
|
||||
status_code,
|
||||
200,
|
||||
diagnostics_status_code(&context.diagnostics),
|
||||
);
|
||||
insert_response_request_summary_object(object, &context.request, context.transport);
|
||||
insert_upstream_diagnostics_metadata(object, &context.diagnostics);
|
||||
insert_response_trace_metadata_object(object, &context.attempts, Some(&context.attempt));
|
||||
object
|
||||
.entry("requestBody".to_string())
|
||||
.or_insert_with(|| serde_json::json!(context.request));
|
||||
@@ -219,6 +275,181 @@ fn enrich_live_response_stream_metadata(
|
||||
.or_insert_with(|| Value::String(body.to_string()));
|
||||
}
|
||||
|
||||
pub(super) fn insert_response_upstream_diagnostics(
|
||||
metadata: &mut Value,
|
||||
diagnostics: &CodexUpstreamDiagnostics,
|
||||
) {
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
insert_upstream_diagnostics_metadata(object, diagnostics);
|
||||
}
|
||||
|
||||
pub(super) fn insert_response_status_metadata(
|
||||
metadata: &mut Value,
|
||||
event_status_code: i64,
|
||||
client_status_code: i64,
|
||||
upstream_status_code: Option<i64>,
|
||||
) {
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
insert_response_status_metadata_object(
|
||||
object,
|
||||
event_status_code,
|
||||
client_status_code,
|
||||
upstream_status_code,
|
||||
);
|
||||
}
|
||||
|
||||
pub(super) fn insert_response_trace_metadata(
|
||||
metadata: &mut Value,
|
||||
trace: &ResponseDispatchTrace,
|
||||
current_attempt: Option<&ResponseDispatchAttempt>,
|
||||
) {
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
insert_response_trace_metadata_object(object, trace.attempts(), current_attempt);
|
||||
}
|
||||
|
||||
fn insert_response_attempt_metadata_object(
|
||||
object: &mut serde_json::Map<String, Value>,
|
||||
attempt: Option<&ResponseDispatchAttempt>,
|
||||
) {
|
||||
let Some(attempt) = attempt else {
|
||||
return;
|
||||
};
|
||||
object.insert("attemptIndex".to_string(), json!(attempt.index()));
|
||||
object.insert(
|
||||
"attemptAccountId".to_string(),
|
||||
Value::String(attempt.account_id().to_string()),
|
||||
);
|
||||
}
|
||||
|
||||
fn insert_response_trace_metadata_object(
|
||||
object: &mut Map<String, Value>,
|
||||
attempts: &[ResponseDispatchAttempt],
|
||||
current_attempt: Option<&ResponseDispatchAttempt>,
|
||||
) {
|
||||
insert_response_attempt_metadata_object(object, current_attempt);
|
||||
object.insert("attemptCount".to_string(), json!(attempts.len()));
|
||||
object.insert(
|
||||
"attempts".to_string(),
|
||||
Value::Array(
|
||||
attempts
|
||||
.iter()
|
||||
.map(|attempt| {
|
||||
json!({
|
||||
"index": attempt.index(),
|
||||
"accountId": attempt.account_id(),
|
||||
})
|
||||
})
|
||||
.collect(),
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
fn insert_response_status_metadata_object(
|
||||
object: &mut Map<String, Value>,
|
||||
event_status_code: i64,
|
||||
client_status_code: i64,
|
||||
upstream_status_code: Option<i64>,
|
||||
) {
|
||||
object.insert("eventStatusCode".to_string(), json!(event_status_code));
|
||||
object.insert("clientStatusCode".to_string(), json!(client_status_code));
|
||||
if let Some(upstream_status_code) = upstream_status_code {
|
||||
object.insert(
|
||||
"upstreamStatusCode".to_string(),
|
||||
json!(upstream_status_code),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn insert_response_request_summary_object(
|
||||
object: &mut Map<String, Value>,
|
||||
request: &CodexResponsesRequest,
|
||||
transport: CodexBackendTransport,
|
||||
) {
|
||||
object.insert(
|
||||
"requestSummary".to_string(),
|
||||
response_request_summary(request, transport),
|
||||
);
|
||||
}
|
||||
|
||||
fn response_request_summary(
|
||||
request: &CodexResponsesRequest,
|
||||
transport: CodexBackendTransport,
|
||||
) -> Value {
|
||||
let body = request.body();
|
||||
let input = body.get("input");
|
||||
let tools = body.get("tools");
|
||||
json!({
|
||||
"model": request.model(),
|
||||
"stream": request.stream(),
|
||||
"store": request.store(),
|
||||
"transport": backend_transport_name(transport),
|
||||
"inputType": json_value_kind(input),
|
||||
"inputItemsCount": input.and_then(Value::as_array).map(Vec::len),
|
||||
"toolsType": json_value_kind(tools),
|
||||
"toolsCount": tools.and_then(Value::as_array).map(Vec::len),
|
||||
"topLevelFields": body.keys().cloned().collect::<Vec<_>>(),
|
||||
"previousResponseIdPresent": request.previous_response_id().is_some(),
|
||||
"serviceTier": request.service_tier(),
|
||||
"localTransport": {
|
||||
"useWebsocket": request.use_websocket,
|
||||
"forceHttpSse": request.force_http_sse,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
fn json_value_kind(value: Option<&Value>) -> &'static str {
|
||||
match value {
|
||||
None => "missing",
|
||||
Some(Value::Null) => "null",
|
||||
Some(Value::Bool(_)) => "boolean",
|
||||
Some(Value::Number(_)) => "number",
|
||||
Some(Value::String(_)) => "string",
|
||||
Some(Value::Array(_)) => "array",
|
||||
Some(Value::Object(_)) => "object",
|
||||
}
|
||||
}
|
||||
|
||||
fn diagnostics_status_code(diagnostics: &CodexUpstreamDiagnostics) -> Option<i64> {
|
||||
diagnostics.status_code.map(i64::from)
|
||||
}
|
||||
|
||||
fn dispatch_error_upstream_status_code(error: &ResponseDispatchError) -> Option<i64> {
|
||||
match error {
|
||||
ResponseDispatchError::Upstream(error) => {
|
||||
Some(i64::from(upstream_error_http_status(error)))
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch_error_client_status_code(
|
||||
error: &ResponseDispatchError,
|
||||
stream: bool,
|
||||
compact: bool,
|
||||
) -> i64 {
|
||||
if stream {
|
||||
return 200;
|
||||
}
|
||||
if matches!(error, ResponseDispatchError::Upstream(_)) && !compact {
|
||||
return 502;
|
||||
}
|
||||
i64::from(error.http_status_code())
|
||||
}
|
||||
|
||||
fn upstream_failure_client_status_code(stream: bool) -> i64 {
|
||||
if stream {
|
||||
200
|
||||
} else {
|
||||
502
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn record_live_response_stream_event(
|
||||
context: &LiveResponseStreamContext,
|
||||
status_code: i64,
|
||||
@@ -230,7 +461,13 @@ pub(super) async fn record_live_response_stream_event(
|
||||
) {
|
||||
ensure_stream_metadata_flag(&mut metadata);
|
||||
enrich_event_route_metadata(&mut metadata, &context.route);
|
||||
enrich_live_response_stream_metadata(context, rate_limit_headers, &mut metadata, body);
|
||||
enrich_live_response_stream_metadata(
|
||||
context,
|
||||
rate_limit_headers,
|
||||
&mut metadata,
|
||||
status_code,
|
||||
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());
|
||||
|
||||
@@ -8,3 +8,4 @@ mod prefetch;
|
||||
pub(crate) mod service;
|
||||
mod sse_failure;
|
||||
mod stream_lifecycle;
|
||||
mod trace;
|
||||
|
||||
@@ -82,10 +82,12 @@ use super::{
|
||||
affinity::{evict_reasoning_replay, record_response_affinity},
|
||||
errors::{ResponseDispatchError, ResponseDispatchStreamError},
|
||||
event_recording::{
|
||||
insert_websocket_pool_decision, record_prefetched_response_stream_failure_event,
|
||||
record_response_dispatch_error_event, record_response_upstream_error_event,
|
||||
ResponseDispatchErrorDetails, ResponseDispatchErrorEventRecord,
|
||||
ResponseStreamFailureEventRecord, ResponseUpstreamErrorEventRecord,
|
||||
insert_response_status_metadata, insert_response_trace_metadata,
|
||||
insert_response_upstream_diagnostics, insert_websocket_pool_decision,
|
||||
record_prefetched_response_stream_failure_event, record_response_dispatch_error_event,
|
||||
record_response_upstream_error_event, ResponseDispatchErrorDetails,
|
||||
ResponseDispatchErrorEventRecord, ResponseStreamFailureEventRecord,
|
||||
ResponseUpstreamErrorEventRecord,
|
||||
},
|
||||
live_stream::spawn_live_response_stream,
|
||||
prefetch::prefetch_first_sse_chunk,
|
||||
@@ -96,6 +98,7 @@ use super::{
|
||||
sse_failure_invalid_reasoning_replay, stream_failure_http_status,
|
||||
},
|
||||
stream_lifecycle::LiveResponseStreamContext,
|
||||
trace::{ResponseDispatchAttempt, ResponseDispatchTrace},
|
||||
};
|
||||
|
||||
/// OpenAI Responses 调度服务。
|
||||
@@ -385,17 +388,26 @@ impl ResponseDispatchService {
|
||||
let mut excluded_account_ids = Vec::new();
|
||||
let mut exhausted_accounts = AccountExhaustionTracker::default();
|
||||
let mut history_recovery_used = false;
|
||||
let mut next_required_account_id: Option<String> = None;
|
||||
let mut empty_response_retries = 0u8;
|
||||
let mut quota_verify_attempts = 0usize;
|
||||
let mut trace = ResponseDispatchTrace::default();
|
||||
const MAX_EMPTY_RESPONSE_RETRIES: u8 = 2;
|
||||
let (account, response, collected_response): (
|
||||
let (account, response, collected_response, attempt): (
|
||||
Account,
|
||||
CodexBackendResponse,
|
||||
CollectedResponse,
|
||||
ResponseDispatchAttempt,
|
||||
) = loop {
|
||||
let mut attempt_acquire_request = acquire_request
|
||||
.clone()
|
||||
.with_exclude_account_ids(excluded_account_ids.iter().cloned());
|
||||
if let Some(account_id) = next_required_account_id.take() {
|
||||
attempt_acquire_request =
|
||||
attempt_acquire_request.with_required_account_id(account_id);
|
||||
}
|
||||
let allow_quota_verify_account_retry =
|
||||
attempt_acquire_request.required_account_id.is_none();
|
||||
attempt_acquire_request.now = Utc::now();
|
||||
let Some(acquired) = self
|
||||
.account_pool
|
||||
@@ -436,6 +448,7 @@ impl ResponseDispatchService {
|
||||
request_id,
|
||||
excluded_account_ids: &mut excluded_account_ids,
|
||||
verify_attempts: &mut quota_verify_attempts,
|
||||
allow_retry_with_another_account: allow_quota_verify_account_retry,
|
||||
},
|
||||
acquired,
|
||||
)
|
||||
@@ -447,7 +460,8 @@ impl ResponseDispatchService {
|
||||
.record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE);
|
||||
continue;
|
||||
}
|
||||
QuotaVerificationDecision::MaxAttemptsReached => {
|
||||
QuotaVerificationDecision::MaxAttemptsReached
|
||||
| QuotaVerificationDecision::RequiredAccountUnavailable => {
|
||||
exhausted_accounts.record_rate_limited(
|
||||
Some(&acquired_account_id),
|
||||
QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
|
||||
@@ -487,6 +501,7 @@ impl ResponseDispatchService {
|
||||
self.account_pool.wait_for_request_interval(&acquired).await;
|
||||
let account = acquired.account;
|
||||
let release_account_id = account.id.clone();
|
||||
let attempt = trace.start_attempt(&release_account_id);
|
||||
let response_result = create_response_with_account_retrying_5xx(
|
||||
&self.codex,
|
||||
self.installation_id.as_deref(),
|
||||
@@ -562,6 +577,7 @@ impl ResponseDispatchService {
|
||||
self.recover_request_history(&mut request, &mut implicit_resume)
|
||||
.await;
|
||||
history_recovery_used = true;
|
||||
next_required_account_id = Some(release_account_id);
|
||||
continue;
|
||||
}
|
||||
if is_model_unsupported_sse_failure(failure) {
|
||||
@@ -626,7 +642,7 @@ impl ResponseDispatchService {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
break (account, response, collected_response);
|
||||
break (account, response, collected_response, attempt);
|
||||
}
|
||||
Err(error) if is_rate_limit_upstream_error(&error) => {
|
||||
exhausted_accounts.record_rate_limited(
|
||||
@@ -659,6 +675,7 @@ impl ResponseDispatchService {
|
||||
self.recover_request_history(&mut request, &mut implicit_resume)
|
||||
.await;
|
||||
history_recovery_used = true;
|
||||
next_required_account_id = Some(release_account_id);
|
||||
}
|
||||
Err(error) if is_auth_upstream_error(&error) => {
|
||||
let upstream_error = upstream_error_body(&error);
|
||||
@@ -750,7 +767,10 @@ impl ResponseDispatchService {
|
||||
started_at,
|
||||
stream: false,
|
||||
transport: backend_transport_for_response_request(&request),
|
||||
request: &request,
|
||||
error: &error,
|
||||
trace: &trace,
|
||||
attempt: Some(&attempt),
|
||||
})
|
||||
.await;
|
||||
return Err(ResponseDispatchError::Upstream(error));
|
||||
@@ -787,6 +807,14 @@ impl ResponseDispatchService {
|
||||
"firstTokenMs": response.first_token_ms,
|
||||
"usage": response.usage,
|
||||
});
|
||||
insert_response_status_metadata(
|
||||
&mut metadata,
|
||||
200,
|
||||
200,
|
||||
response.diagnostics.status_code.map(i64::from),
|
||||
);
|
||||
insert_response_upstream_diagnostics(&mut metadata, &response.diagnostics);
|
||||
insert_response_trace_metadata(&mut metadata, &trace, Some(&attempt));
|
||||
insert_websocket_pool_decision(&mut metadata, response.websocket_pool_decision);
|
||||
record_response_event(ResponseUsageRecord {
|
||||
usage_records: &self.usage_records,
|
||||
@@ -915,7 +943,9 @@ impl ResponseDispatchService {
|
||||
let mut excluded_account_ids = Vec::new();
|
||||
let mut exhausted_accounts = AccountExhaustionTracker::default();
|
||||
let mut history_recovery_used = false;
|
||||
let mut next_required_account_id: Option<String> = None;
|
||||
let mut quota_verify_attempts = 0usize;
|
||||
let mut trace = ResponseDispatchTrace::default();
|
||||
macro_rules! return_stream_dispatch_error {
|
||||
($error:expr) => {{
|
||||
let error = $error;
|
||||
@@ -960,6 +990,12 @@ impl ResponseDispatchService {
|
||||
let mut attempt_acquire_request = acquire_request
|
||||
.clone()
|
||||
.with_exclude_account_ids(excluded_account_ids.iter().cloned());
|
||||
if let Some(account_id) = next_required_account_id.take() {
|
||||
attempt_acquire_request =
|
||||
attempt_acquire_request.with_required_account_id(account_id);
|
||||
}
|
||||
let allow_quota_verify_account_retry =
|
||||
attempt_acquire_request.required_account_id.is_none();
|
||||
attempt_acquire_request.now = Utc::now();
|
||||
let Some(acquired) = self
|
||||
.account_pool
|
||||
@@ -982,6 +1018,7 @@ impl ResponseDispatchService {
|
||||
request_id,
|
||||
excluded_account_ids: &mut excluded_account_ids,
|
||||
verify_attempts: &mut quota_verify_attempts,
|
||||
allow_retry_with_another_account: allow_quota_verify_account_retry,
|
||||
},
|
||||
acquired,
|
||||
)
|
||||
@@ -993,7 +1030,8 @@ impl ResponseDispatchService {
|
||||
.record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE);
|
||||
continue;
|
||||
}
|
||||
QuotaVerificationDecision::MaxAttemptsReached => {
|
||||
QuotaVerificationDecision::MaxAttemptsReached
|
||||
| QuotaVerificationDecision::RequiredAccountUnavailable => {
|
||||
exhausted_accounts.record_rate_limited(
|
||||
Some(&acquired_account_id),
|
||||
QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
|
||||
@@ -1024,6 +1062,7 @@ impl ResponseDispatchService {
|
||||
self.account_pool.wait_for_request_interval(&acquired).await;
|
||||
let account = acquired.account;
|
||||
let release_account_id = account.id.clone();
|
||||
let attempt = trace.start_attempt(&release_account_id);
|
||||
let response_result = create_response_stream_with_account_retrying_5xx(
|
||||
&self.codex,
|
||||
self.installation_id.as_deref(),
|
||||
@@ -1051,6 +1090,7 @@ impl ResponseDispatchService {
|
||||
let turn_state_update = response.turn_state_update;
|
||||
let websocket_pool_decision = response.websocket_pool_decision;
|
||||
let turn_state = response.turn_state;
|
||||
let diagnostics = response.diagnostics;
|
||||
self.cloudflare
|
||||
.capture_set_cookie_headers(&release_account_id, &set_cookie_headers)
|
||||
.await;
|
||||
@@ -1071,6 +1111,7 @@ impl ResponseDispatchService {
|
||||
self.recover_request_history(&mut request, &mut implicit_resume)
|
||||
.await;
|
||||
history_recovery_used = true;
|
||||
next_required_account_id = Some(release_account_id);
|
||||
continue;
|
||||
}
|
||||
Err(ResponseDispatchError::Upstream(error))
|
||||
@@ -1204,7 +1245,10 @@ impl ResponseDispatchService {
|
||||
started_at,
|
||||
stream: true,
|
||||
transport,
|
||||
request: &request,
|
||||
error: upstream_error,
|
||||
trace: &trace,
|
||||
attempt: Some(&attempt),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
@@ -1238,6 +1282,7 @@ impl ResponseDispatchService {
|
||||
self.recover_request_history(&mut request, &mut implicit_resume)
|
||||
.await;
|
||||
history_recovery_used = true;
|
||||
next_required_account_id = Some(release_account_id);
|
||||
continue;
|
||||
}
|
||||
if is_model_unsupported_sse_failure(&failure) {
|
||||
@@ -1306,8 +1351,11 @@ impl ResponseDispatchService {
|
||||
transport,
|
||||
request: &request,
|
||||
failure: &failure,
|
||||
diagnostics: &diagnostics,
|
||||
rate_limit_headers: &rate_limit_headers,
|
||||
prefetched: &prefetched,
|
||||
trace: &trace,
|
||||
attempt: &attempt,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
@@ -1336,6 +1384,9 @@ impl ResponseDispatchService {
|
||||
turn_state_update,
|
||||
websocket_pool_decision,
|
||||
turn_state,
|
||||
diagnostics,
|
||||
attempt: attempt.clone(),
|
||||
attempts: trace.attempts().to_vec(),
|
||||
started_at,
|
||||
};
|
||||
return Ok(spawn_live_response_stream(context, prefetched, body));
|
||||
@@ -1374,6 +1425,7 @@ impl ResponseDispatchService {
|
||||
self.recover_request_history(&mut request, &mut implicit_resume)
|
||||
.await;
|
||||
history_recovery_used = true;
|
||||
next_required_account_id = Some(release_account_id);
|
||||
}
|
||||
Err(error) if is_auth_upstream_error(&error) => {
|
||||
self.account_pool.release(&release_account_id).await;
|
||||
@@ -1460,7 +1512,10 @@ impl ResponseDispatchService {
|
||||
started_at,
|
||||
stream: true,
|
||||
transport: backend_transport_for_response_request(&request),
|
||||
request: &request,
|
||||
error: &error,
|
||||
trace: &trace,
|
||||
attempt: Some(&attempt),
|
||||
})
|
||||
.await;
|
||||
return Err(ResponseDispatchError::Upstream(error));
|
||||
@@ -1483,6 +1538,7 @@ impl ResponseDispatchService {
|
||||
let mut excluded_account_ids = Vec::new();
|
||||
let mut exhausted_accounts = AccountExhaustionTracker::default();
|
||||
let mut quota_verify_attempts = 0usize;
|
||||
let mut trace = ResponseDispatchTrace::default();
|
||||
|
||||
loop {
|
||||
let acquire_request = AccountAcquireRequest::new(request.model(), Utc::now())
|
||||
@@ -1515,6 +1571,7 @@ impl ResponseDispatchService {
|
||||
request_id,
|
||||
excluded_account_ids: &mut excluded_account_ids,
|
||||
verify_attempts: &mut quota_verify_attempts,
|
||||
allow_retry_with_another_account: true,
|
||||
},
|
||||
acquired,
|
||||
)
|
||||
@@ -1526,7 +1583,8 @@ impl ResponseDispatchService {
|
||||
.record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE);
|
||||
continue;
|
||||
}
|
||||
QuotaVerificationDecision::MaxAttemptsReached => {
|
||||
QuotaVerificationDecision::MaxAttemptsReached
|
||||
| QuotaVerificationDecision::RequiredAccountUnavailable => {
|
||||
exhausted_accounts.record_rate_limited(
|
||||
Some(&acquired_account_id),
|
||||
QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
|
||||
@@ -1548,6 +1606,7 @@ impl ResponseDispatchService {
|
||||
};
|
||||
let account = acquired.account;
|
||||
let release_account_id = account.id.clone();
|
||||
let attempt = trace.start_attempt(&release_account_id);
|
||||
let response_result = create_compact_response_with_account_retrying_5xx(
|
||||
&self.codex,
|
||||
self.installation_id.as_deref(),
|
||||
@@ -1585,6 +1644,19 @@ impl ResponseDispatchService {
|
||||
.record_token_usage(&account.id, request.model(), &usage)
|
||||
.await;
|
||||
}
|
||||
let mut metadata = json!({
|
||||
"stream": false,
|
||||
"compact": true,
|
||||
"usage": usage,
|
||||
});
|
||||
insert_response_status_metadata(
|
||||
&mut metadata,
|
||||
200,
|
||||
200,
|
||||
response.diagnostics.status_code.map(i64::from),
|
||||
);
|
||||
insert_response_upstream_diagnostics(&mut metadata, &response.diagnostics);
|
||||
insert_response_trace_metadata(&mut metadata, &trace, Some(&attempt));
|
||||
record_response_event(ResponseUsageRecord {
|
||||
usage_records: &self.usage_records,
|
||||
request_id,
|
||||
@@ -1600,11 +1672,7 @@ impl ResponseDispatchService {
|
||||
status_code: 200,
|
||||
level: UsageRecordLevel::Info,
|
||||
message: "v1 responses compact completed",
|
||||
metadata: json!({
|
||||
"stream": false,
|
||||
"compact": true,
|
||||
"usage": usage,
|
||||
}),
|
||||
metadata,
|
||||
rate_limit_headers: &response.rate_limit_headers,
|
||||
})
|
||||
.await;
|
||||
|
||||
@@ -24,7 +24,7 @@ use crate::{
|
||||
},
|
||||
transport::{
|
||||
CodexBackendTransport, CodexRateLimitHeaderUpdates, CodexTurnStateUpdate,
|
||||
WebSocketPoolDecision,
|
||||
CodexUpstreamDiagnostics, WebSocketPoolDecision,
|
||||
},
|
||||
},
|
||||
};
|
||||
@@ -40,6 +40,7 @@ use super::{
|
||||
stream_failure_metadata, stream_failure_source, synthetic_stream_disconnected_detail,
|
||||
STREAM_DISCONNECTED_CODE, STREAM_DISCONNECTED_MESSAGE,
|
||||
},
|
||||
trace::ResponseDispatchAttempt,
|
||||
};
|
||||
use crate::proxy::dispatch::cloudflare::CloudflareRecovery;
|
||||
|
||||
@@ -65,6 +66,9 @@ pub(super) struct LiveResponseStreamContext {
|
||||
pub(super) turn_state_update: Option<CodexTurnStateUpdate>,
|
||||
pub(super) websocket_pool_decision: Option<WebSocketPoolDecision>,
|
||||
pub(super) turn_state: Option<String>,
|
||||
pub(super) diagnostics: CodexUpstreamDiagnostics,
|
||||
pub(super) attempt: ResponseDispatchAttempt,
|
||||
pub(super) attempts: Vec<ResponseDispatchAttempt>,
|
||||
pub(super) started_at: Instant,
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
//! Request-level trace state for Responses dispatch.
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub(super) struct ResponseDispatchTrace {
|
||||
next_attempt_index: i64,
|
||||
attempts: Vec<ResponseDispatchAttempt>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(super) struct ResponseDispatchAttempt {
|
||||
index: i64,
|
||||
account_id: String,
|
||||
}
|
||||
|
||||
impl ResponseDispatchTrace {
|
||||
pub(super) fn start_attempt(&mut self, account_id: &str) -> ResponseDispatchAttempt {
|
||||
let attempt = ResponseDispatchAttempt {
|
||||
index: self.next_attempt_index,
|
||||
account_id: account_id.to_string(),
|
||||
};
|
||||
self.next_attempt_index += 1;
|
||||
self.attempts.push(attempt.clone());
|
||||
attempt
|
||||
}
|
||||
|
||||
pub(super) fn attempts(&self) -> &[ResponseDispatchAttempt] {
|
||||
&self.attempts
|
||||
}
|
||||
}
|
||||
|
||||
impl ResponseDispatchAttempt {
|
||||
pub(super) fn index(&self) -> i64 {
|
||||
self.index
|
||||
}
|
||||
|
||||
pub(super) fn account_id(&self) -> &str {
|
||||
&self.account_id
|
||||
}
|
||||
}
|
||||
@@ -34,6 +34,7 @@ pub(crate) const QUOTA_VERIFY_LIMIT_REACHED_MESSAGE: &str =
|
||||
pub(crate) enum QuotaVerificationDecision {
|
||||
Ready(Box<AcquiredAccount>),
|
||||
RetryWithAnotherAccount,
|
||||
RequiredAccountUnavailable,
|
||||
MaxAttemptsReached,
|
||||
}
|
||||
|
||||
@@ -45,6 +46,7 @@ pub(crate) struct QuotaVerificationContext<'a> {
|
||||
pub request_id: &'a str,
|
||||
pub excluded_account_ids: &'a mut Vec<String>,
|
||||
pub verify_attempts: &'a mut usize,
|
||||
pub allow_retry_with_another_account: bool,
|
||||
}
|
||||
|
||||
pub(crate) async fn verify_acquired_quota_if_required(
|
||||
@@ -111,9 +113,12 @@ pub(crate) async fn verify_acquired_quota_if_required(
|
||||
quota_verify_required = true,
|
||||
quota_verify_result = "limit_reached",
|
||||
verify_attempts = *context.verify_attempts,
|
||||
retry_with_another_account = !max_attempts_reached,
|
||||
retry_with_another_account = context.allow_retry_with_another_account && !max_attempts_reached,
|
||||
"quota verification reported exhausted account before upstream request"
|
||||
);
|
||||
if !context.allow_retry_with_another_account {
|
||||
return QuotaVerificationDecision::RequiredAccountUnavailable;
|
||||
}
|
||||
if max_attempts_reached {
|
||||
return QuotaVerificationDecision::MaxAttemptsReached;
|
||||
}
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
//! OpenAI 兼容 API 路由。
|
||||
|
||||
use axum::{
|
||||
extract::DefaultBodyLimit,
|
||||
routing::{get, post},
|
||||
Router,
|
||||
};
|
||||
@@ -22,4 +23,5 @@ pub fn router() -> Router<AppState> {
|
||||
.route("/v1/models/catalog", get(model_catalog))
|
||||
.route("/v1/models/{model_id}", get(model_detail))
|
||||
.route("/v1/models/{model_id}/info", get(model_info))
|
||||
.layer(DefaultBodyLimit::disable())
|
||||
}
|
||||
|
||||
@@ -13,7 +13,7 @@ use crate::runtime::services::{BackgroundTaskStores, Services};
|
||||
use crate::runtime::shutdown::{restart_executable_path, shutdown_signal};
|
||||
use crate::runtime::state::{AppState, RuntimeConfig};
|
||||
use crate::runtime::tasks::coordinator::TaskCoordinator;
|
||||
use crate::upstream::fingerprint::{Fingerprint, FingerprintRepository};
|
||||
use crate::upstream::fingerprint::{Fingerprint, FingerprintRepository, RuntimeFingerprint};
|
||||
|
||||
const RESTART_DELAY_ENV: &str = "CPR_RESTART_DELAY_MS";
|
||||
|
||||
@@ -81,6 +81,7 @@ async fn build_application(
|
||||
let runtime_fingerprint = fingerprint_repository
|
||||
.ensure_current_seed(&default_fingerprint)
|
||||
.await?;
|
||||
let runtime_fingerprint = RuntimeFingerprint::new(runtime_fingerprint);
|
||||
let runtime_config = RuntimeConfig::from(config.clone());
|
||||
let services = Services::try_with_installation_id(
|
||||
&config,
|
||||
|
||||
@@ -7,10 +7,7 @@ use crate::{
|
||||
account_usage_service::AdminUsageService,
|
||||
account_usage_store::SqliteUsageStore,
|
||||
usage_record_service::AdminUsageRecordService,
|
||||
usage_record_store::{
|
||||
SqliteUsageRecordStore, DEFAULT_USAGE_RECORD_CAPACITY,
|
||||
DEFAULT_USAGE_RECORD_CAPTURE_BODY,
|
||||
},
|
||||
usage_record_store::{SqliteUsageRecordStore, DEFAULT_USAGE_RECORD_CAPTURE_BODY},
|
||||
},
|
||||
admin::{
|
||||
accounts::service::{AdminAccountService, AdminAccountServiceParts},
|
||||
@@ -38,7 +35,7 @@ use crate::{
|
||||
},
|
||||
},
|
||||
upstream::{
|
||||
fingerprint::{Fingerprint, FingerprintRepository},
|
||||
fingerprint::{FingerprintRepository, RuntimeFingerprint},
|
||||
models::{
|
||||
config::ModelConfig,
|
||||
service::{ModelService, ModelServiceError},
|
||||
@@ -101,7 +98,7 @@ pub struct Services {
|
||||
pub session_affinity: StdArc<RuntimeSessionAffinityService>,
|
||||
pub codex: StdArc<CodexBackendClient>,
|
||||
pub websocket_pool: Option<StdArc<CodexWebSocketPool>>,
|
||||
pub fingerprint: Fingerprint,
|
||||
pub fingerprint: RuntimeFingerprint,
|
||||
pub installation_id: Option<String>,
|
||||
pub background_tasks: BackgroundTaskStores,
|
||||
}
|
||||
@@ -110,7 +107,6 @@ pub struct Services {
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct UsageRecordOptions {
|
||||
pub enabled: bool,
|
||||
pub capacity: u32,
|
||||
pub capture_body: bool,
|
||||
}
|
||||
|
||||
@@ -118,14 +114,17 @@ impl UsageRecordOptions {
|
||||
pub fn from_config(config: &AppConfig) -> Self {
|
||||
Self {
|
||||
enabled: config.logging.enabled,
|
||||
capacity: DEFAULT_USAGE_RECORD_CAPACITY,
|
||||
capture_body: DEFAULT_USAGE_RECORD_CAPTURE_BODY,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Services {
|
||||
pub fn new(config: &AppConfig, stores: BackgroundTaskStores, fingerprint: Fingerprint) -> Self {
|
||||
pub fn new(
|
||||
config: &AppConfig,
|
||||
stores: BackgroundTaskStores,
|
||||
fingerprint: RuntimeFingerprint,
|
||||
) -> Self {
|
||||
Self::try_new(config, stores, fingerprint)
|
||||
.expect("failed to build runtime services with configured TLS transport")
|
||||
}
|
||||
@@ -133,7 +132,7 @@ impl Services {
|
||||
pub fn try_new(
|
||||
config: &AppConfig,
|
||||
stores: BackgroundTaskStores,
|
||||
fingerprint: Fingerprint,
|
||||
fingerprint: RuntimeFingerprint,
|
||||
) -> Result<Self, CustomCaError> {
|
||||
Self::try_with_installation_id(config, stores, fingerprint, None)
|
||||
}
|
||||
@@ -141,7 +140,7 @@ impl Services {
|
||||
pub(crate) fn try_with_installation_id(
|
||||
config: &AppConfig,
|
||||
stores: BackgroundTaskStores,
|
||||
fingerprint: Fingerprint,
|
||||
fingerprint: RuntimeFingerprint,
|
||||
installation_id: Option<String>,
|
||||
) -> Result<Self, CustomCaError> {
|
||||
Self::try_with_installation_id_and_usage_record_options(
|
||||
@@ -156,7 +155,7 @@ impl Services {
|
||||
pub fn try_with_installation_id_and_usage_record_options(
|
||||
config: &AppConfig,
|
||||
stores: BackgroundTaskStores,
|
||||
fingerprint: Fingerprint,
|
||||
fingerprint: RuntimeFingerprint,
|
||||
installation_id: Option<String>,
|
||||
usage_record_options: UsageRecordOptions,
|
||||
) -> Result<Self, CustomCaError> {
|
||||
@@ -168,21 +167,32 @@ impl Services {
|
||||
enabled: true,
|
||||
max_age: std::time::Duration::from_millis(config.ws_pool.max_age_ms),
|
||||
max_per_account: config.ws_pool.max_per_account,
|
||||
first_token_timeout: (config.ws_pool.first_token_timeout_ms > 0).then(|| {
|
||||
std::time::Duration::from_millis(config.ws_pool.first_token_timeout_ms)
|
||||
}),
|
||||
..CodexWebSocketPoolConfig::default()
|
||||
}))
|
||||
});
|
||||
let codex = {
|
||||
let websocket_first_token_timeout = (config.ws_pool.first_token_timeout_ms > 0)
|
||||
.then(|| std::time::Duration::from_millis(config.ws_pool.first_token_timeout_ms));
|
||||
let client = CodexBackendClient::new(
|
||||
build_reqwest_client(config.tls.force_http11)?,
|
||||
config.api.base_url.clone(),
|
||||
fingerprint.clone(),
|
||||
);
|
||||
)
|
||||
.with_websocket_first_token_timeout(websocket_first_token_timeout);
|
||||
if let Some(pool) = &websocket_pool {
|
||||
StdArc::new(client.with_websocket_pool(pool.clone()))
|
||||
} else {
|
||||
StdArc::new(client)
|
||||
}
|
||||
};
|
||||
let usage_records = StdArc::new(AdminUsageRecordService::new(
|
||||
stores.usage_records.clone(),
|
||||
usage_record_options.enabled,
|
||||
usage_record_options.capture_body,
|
||||
));
|
||||
let account_pool = StdArc::new(RuntimeAccountPoolService::new(
|
||||
account_store_trait.clone(),
|
||||
account_pool_options_from_config(config),
|
||||
@@ -245,12 +255,6 @@ impl Services {
|
||||
refresh_policy: refresh_policy.clone(),
|
||||
installation_id: installation_id.clone(),
|
||||
}));
|
||||
let usage_records = StdArc::new(AdminUsageRecordService::new(
|
||||
stores.usage_records.clone(),
|
||||
usage_record_options.enabled,
|
||||
usage_record_options.capacity,
|
||||
usage_record_options.capture_body,
|
||||
));
|
||||
let usage_store = SqliteUsageStore::new(stores.accounts.pool().clone());
|
||||
let usage = StdArc::new(AdminUsageService::new(usage_store));
|
||||
let session_affinity = StdArc::new(RuntimeSessionAffinityService::new(
|
||||
|
||||
@@ -86,16 +86,17 @@ impl TaskCoordinator {
|
||||
.with_account_pool(services.account_pool.clone())
|
||||
.start(),
|
||||
);
|
||||
coordinator.push(
|
||||
"fingerprint_update",
|
||||
coordinator.push("fingerprint_update", {
|
||||
let fingerprint = services.fingerprint.snapshot();
|
||||
start_fingerprint_update_task(
|
||||
Some(stores.fingerprints.clone()),
|
||||
stores.fingerprints.clone(),
|
||||
services.fingerprint.clone(),
|
||||
CODEX_DESKTOP_APPCAST_URL.to_string(),
|
||||
PathBuf::from(DEFAULT_EXTRACTED_FINGERPRINT_PATH),
|
||||
services.fingerprint.app_version.clone(),
|
||||
services.fingerprint.build_number.clone(),
|
||||
),
|
||||
);
|
||||
fingerprint.app_version,
|
||||
fingerprint.build_number,
|
||||
)
|
||||
});
|
||||
if let Some(pool) = &services.websocket_pool {
|
||||
coordinator.push(
|
||||
"websocket_pool",
|
||||
|
||||
@@ -2,9 +2,11 @@
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use tracing::warn;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::upstream::fingerprint::{FingerprintRepository, UpdateChecker, APPCAST_POLL_INTERVAL};
|
||||
use crate::upstream::fingerprint::{
|
||||
FingerprintRepository, RuntimeFingerprint, UpdateChecker, APPCAST_POLL_INTERVAL,
|
||||
};
|
||||
|
||||
use super::{
|
||||
coordinator::SchedulerHandle,
|
||||
@@ -19,7 +21,8 @@ pub const DEFAULT_EXTRACTED_FINGERPRINT_PATH: &str = "data/extracted-fingerprint
|
||||
|
||||
/// 启动指纹自动更新后台任务。
|
||||
pub fn start_fingerprint_update_task(
|
||||
repository: Option<FingerprintRepository>,
|
||||
repository: FingerprintRepository,
|
||||
runtime_fingerprint: RuntimeFingerprint,
|
||||
appcast_url: String,
|
||||
extracted_fingerprint_path: PathBuf,
|
||||
current_version: String,
|
||||
@@ -36,6 +39,7 @@ pub fn start_fingerprint_update_task(
|
||||
spawn_periodic_task(
|
||||
FingerprintUpdateTask {
|
||||
checker: update_checker,
|
||||
runtime_fingerprint,
|
||||
first_tick: true,
|
||||
},
|
||||
PeriodicTaskConfig::new(
|
||||
@@ -48,6 +52,7 @@ pub fn start_fingerprint_update_task(
|
||||
|
||||
struct FingerprintUpdateTask {
|
||||
checker: UpdateChecker,
|
||||
runtime_fingerprint: RuntimeFingerprint,
|
||||
first_tick: bool,
|
||||
}
|
||||
|
||||
@@ -55,11 +60,24 @@ impl PeriodicTaskRunner for FingerprintUpdateTask {
|
||||
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
|
||||
Box::pin(async move {
|
||||
let first_tick = std::mem::take(&mut self.first_tick);
|
||||
if let Err(error) = self.checker.check_and_apply_update().await {
|
||||
if first_tick {
|
||||
warn!(error = %error, "fingerprint 首次版本检查失败");
|
||||
} else {
|
||||
warn!(error = %error, "fingerprint 定期版本检查失败");
|
||||
match self.checker.check_and_apply_update().await {
|
||||
Ok(Some(fingerprint)) => {
|
||||
let app_version = fingerprint.app_version.clone();
|
||||
let build_number = fingerprint.build_number.clone();
|
||||
self.runtime_fingerprint.replace(fingerprint);
|
||||
info!(
|
||||
app_version = %app_version,
|
||||
build_number = %build_number,
|
||||
"fingerprint 运行时版本已更新"
|
||||
);
|
||||
}
|
||||
Ok(None) => {}
|
||||
Err(error) => {
|
||||
if first_tick {
|
||||
warn!(error = %error, "fingerprint 首次版本检查失败");
|
||||
} else {
|
||||
warn!(error = %error, "fingerprint 定期版本检查失败");
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
@@ -89,6 +89,7 @@ impl PeriodicTaskRunner for QuotaRefreshTask {
|
||||
Ok(summary) => {
|
||||
debug!(
|
||||
candidates = summary.candidates,
|
||||
skipped_cooldown = summary.skipped_cooldown,
|
||||
skipped_recent = summary.skipped_recent,
|
||||
failed = summary.failed,
|
||||
"没有需要刷新的 quota 锁定或待验证账号"
|
||||
|
||||
@@ -27,8 +27,10 @@ use crate::upstream::protocol::events::{
|
||||
/// 账号轮转策略。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum RotationStrategy {
|
||||
/// 优先选择当前窗口用量最低的账号。
|
||||
LeastUsed,
|
||||
/// 综合并发压力、窗口用量和历史用量进行调度。
|
||||
Smart,
|
||||
/// 优先消耗即将重置的额度窗口。
|
||||
QuotaResetPriority,
|
||||
/// 按候选账号顺序循环选择。
|
||||
RoundRobin,
|
||||
/// 优先选择最近使用过的账号。
|
||||
@@ -38,7 +40,8 @@ pub enum RotationStrategy {
|
||||
impl RotationStrategy {
|
||||
fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::LeastUsed => "least_used",
|
||||
Self::Smart => "smart",
|
||||
Self::QuotaResetPriority => "quota_reset_priority",
|
||||
Self::RoundRobin => "round_robin",
|
||||
Self::Sticky => "sticky",
|
||||
}
|
||||
@@ -62,8 +65,6 @@ pub struct AccountPoolOptions {
|
||||
pub model_plan_allowlist: BTreeMap<String, Vec<String>>,
|
||||
/// 已成功拉取过模型列表的订阅计划。
|
||||
pub fetched_model_plan_types: BTreeSet<String>,
|
||||
/// 模型到指定账号 ID 的显式路由。
|
||||
pub model_account_routes: BTreeMap<String, Vec<String>>,
|
||||
}
|
||||
|
||||
impl Default for AccountPoolOptions {
|
||||
@@ -71,12 +72,11 @@ impl Default for AccountPoolOptions {
|
||||
Self {
|
||||
max_concurrent_per_account: 3,
|
||||
stale_slot_ttl: Duration::minutes(5),
|
||||
rotation_strategy: RotationStrategy::LeastUsed,
|
||||
rotation_strategy: RotationStrategy::Smart,
|
||||
skip_quota_limited: true,
|
||||
tier_priority: Vec::new(),
|
||||
model_plan_allowlist: BTreeMap::new(),
|
||||
fetched_model_plan_types: BTreeSet::new(),
|
||||
model_account_routes: BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -88,6 +88,8 @@ pub struct AccountAcquireRequest {
|
||||
pub model: String,
|
||||
/// 本次请求需要排除的账号 ID。
|
||||
pub exclude_account_ids: Vec<String>,
|
||||
/// 本次请求必须使用的账号 ID,用于同请求内的同账号重试。
|
||||
pub required_account_id: Option<String>,
|
||||
/// 会话亲和性建议的优先账号 ID。
|
||||
pub preferred_account_id: Option<String>,
|
||||
/// 本次调度使用的当前时间。
|
||||
@@ -100,6 +102,7 @@ impl AccountAcquireRequest {
|
||||
Self {
|
||||
model: model.into(),
|
||||
exclude_account_ids: Vec::new(),
|
||||
required_account_id: None,
|
||||
preferred_account_id: None,
|
||||
now,
|
||||
}
|
||||
@@ -119,6 +122,12 @@ impl AccountAcquireRequest {
|
||||
self.preferred_account_id = Some(account_id.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置本次调度必须使用的账号 ID。
|
||||
pub fn with_required_account_id(mut self, account_id: impl Into<String>) -> Self {
|
||||
self.required_account_id = Some(account_id.into());
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// 成功获取到的账号及调度元数据。
|
||||
@@ -269,6 +278,16 @@ 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 {
|
||||
@@ -480,7 +499,15 @@ impl AccountPool {
|
||||
self.cleanup_stale_slots(request.now);
|
||||
let refreshed_accounts = self.refresh_account_statuses(request.now);
|
||||
let candidates = self.candidates(request);
|
||||
let selected = if let Some(preferred_account_id) = &request.preferred_account_id {
|
||||
let selected = if let Some(required_account_id) = &request.required_account_id {
|
||||
candidates.first().cloned().map(|account| {
|
||||
(
|
||||
account,
|
||||
self.previous_slot_at(required_account_id),
|
||||
"required",
|
||||
)
|
||||
})
|
||||
} else if let Some(preferred_account_id) = &request.preferred_account_id {
|
||||
let preferred = candidates
|
||||
.iter()
|
||||
.find(|account| account.id == *preferred_account_id)
|
||||
@@ -497,10 +524,15 @@ impl AccountPool {
|
||||
}
|
||||
.or_else(|| {
|
||||
Some(match self.options.rotation_strategy {
|
||||
RotationStrategy::LeastUsed => {
|
||||
let account = self.select_least_used(&candidates)?;
|
||||
RotationStrategy::Smart => {
|
||||
let account = self.select_smart(&candidates)?;
|
||||
let previous = self.previous_slot_at(&account.id);
|
||||
(account, previous, "least_used")
|
||||
(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)?;
|
||||
@@ -520,19 +552,25 @@ 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,
|
||||
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"
|
||||
);
|
||||
@@ -562,9 +600,10 @@ impl AccountPool {
|
||||
}
|
||||
let after = RuntimeAccountStateSnapshot::from(&*account);
|
||||
if before != after {
|
||||
let sync_usage_window = before.usage_window_changed(&after);
|
||||
refreshed_accounts.push(RefreshedAccountState {
|
||||
account: account.clone(),
|
||||
sync_usage_window: before.usage_window_changed(&after),
|
||||
sync_usage_window,
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -666,24 +705,31 @@ impl AccountPool {
|
||||
let mut accounts = Vec::new();
|
||||
for (plan_type, group) in by_plan {
|
||||
let selected = match self.options.rotation_strategy {
|
||||
RotationStrategy::LeastUsed => self.select_least_used(&group),
|
||||
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 {
|
||||
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(),
|
||||
account_id = %account.id,
|
||||
candidate_count = group.len(),
|
||||
slot_count = selected_slot_count,
|
||||
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"
|
||||
);
|
||||
@@ -697,19 +743,49 @@ impl AccountPool {
|
||||
}
|
||||
}
|
||||
|
||||
fn select_least_used(&mut self, candidates: &[Account]) -> Option<Account> {
|
||||
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_least_used(a, b));
|
||||
sorted.sort_by(|a, b| compare(self, a, b));
|
||||
let best = *sorted.first()?;
|
||||
let tied_count = sorted
|
||||
.iter()
|
||||
.take_while(|account| compare_least_used(best, account) == Ordering::Equal)
|
||||
.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..] {
|
||||
@@ -741,6 +817,10 @@ impl AccountPool {
|
||||
.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()));
|
||||
}
|
||||
@@ -755,7 +835,6 @@ impl AccountPool {
|
||||
request.now,
|
||||
self.options.skip_quota_limited,
|
||||
)
|
||||
&& self.is_model_account_allowed(account, &request.model)
|
||||
&& self.is_model_allowed(account, &request.model)
|
||||
&& !request
|
||||
.exclude_account_ids
|
||||
@@ -852,15 +931,6 @@ impl AccountPool {
|
||||
|| !self.options.fetched_model_plan_types.contains(plan_type)
|
||||
}
|
||||
|
||||
fn is_model_account_allowed(&self, account: &Account, model: &str) -> bool {
|
||||
let Some(account_ids) = self.options.model_account_routes.get(model) else {
|
||||
return true;
|
||||
};
|
||||
account_ids
|
||||
.iter()
|
||||
.any(|account_id| account_id == &account.id)
|
||||
}
|
||||
|
||||
fn best_available_tier(&self, candidates: &[Account]) -> Option<String> {
|
||||
self.options.tier_priority.iter().find_map(|tier| {
|
||||
candidates
|
||||
@@ -1310,11 +1380,14 @@ impl RuntimeAccountPoolService {
|
||||
}
|
||||
}
|
||||
|
||||
let mut pool = self.pool.lock().await;
|
||||
let in_memory = pool.apply_quota_state(account_id, limit_reached, cooldown_until);
|
||||
if let Some(reset_at) = reset_at {
|
||||
pool.sync_rate_limit_window(account_id, reset_at, limit_window_seconds);
|
||||
}
|
||||
let in_memory = {
|
||||
let mut pool = self.pool.lock().await;
|
||||
let in_memory = pool.apply_quota_state(account_id, limit_reached, cooldown_until);
|
||||
if let Some(reset_at) = reset_at {
|
||||
pool.sync_rate_limit_window(account_id, reset_at, limit_window_seconds);
|
||||
}
|
||||
in_memory
|
||||
};
|
||||
|
||||
persisted || in_memory
|
||||
}
|
||||
@@ -1448,7 +1521,7 @@ fn refresh_cloudflare_cooldown(account: &mut Account, now: DateTime<Utc>) {
|
||||
}
|
||||
}
|
||||
|
||||
fn compare_least_used(a: &Account, b: &Account) -> Ordering {
|
||||
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))
|
||||
@@ -1456,6 +1529,14 @@ fn compare_least_used(a: &Account, b: &Account) -> Ordering {
|
||||
.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),
|
||||
|
||||
@@ -10,7 +10,7 @@ const MONTH_WINDOW_SECONDS: u64 = 2_592_000;
|
||||
|
||||
/// 判断配额快照是否已经触顶。
|
||||
pub fn quota_snapshot_limit_reached(quota: &Value) -> bool {
|
||||
core_snapshot_blocked(quota)
|
||||
snapshot_blocked(quota)
|
||||
|| spend_control_limit_reached(quota)
|
||||
|| monthly_limit_reached(quota)
|
||||
|| credits_overage_limit_reached(quota)
|
||||
@@ -20,8 +20,10 @@ pub fn quota_snapshot_limit_reached(quota: &Value) -> bool {
|
||||
pub fn quota_snapshot_reset_at(quota: &Value) -> Option<DateTime<Utc>> {
|
||||
let reset_at = if monthly_limit_reached(quota) || spend_control_limit_reached(quota) {
|
||||
monthly_limit_reset_at(quota)
|
||||
.or_else(|| blocking_core_bucket(quota).and_then(bucket_reset_at))
|
||||
.or_else(|| blocking_bucket(quota).and_then(bucket_reset_at))
|
||||
.or_else(|| core_bucket_with_reset_at(quota).and_then(bucket_reset_at))
|
||||
} else if let Some(reset_at) = blocking_bucket(quota).and_then(bucket_reset_at) {
|
||||
Some(reset_at)
|
||||
} else {
|
||||
core_bucket_with_reset_at(quota)
|
||||
.and_then(bucket_reset_at)
|
||||
@@ -35,9 +37,14 @@ pub fn quota_snapshot_limit_window_seconds(quota: &Value) -> Option<u64> {
|
||||
if monthly_limit_reached(quota) || spend_control_limit_reached(quota) {
|
||||
return monthly_limit_window_seconds(quota).or(Some(MONTH_WINDOW_SECONDS));
|
||||
}
|
||||
core_bucket_with_window_minutes(quota)
|
||||
blocking_bucket(quota)
|
||||
.and_then(bucket_window_minutes)
|
||||
.and_then(|minutes| minutes.checked_mul(60))
|
||||
.or_else(|| {
|
||||
core_bucket_with_window_minutes(quota)
|
||||
.and_then(bucket_window_minutes)
|
||||
.and_then(|minutes| minutes.checked_mul(60))
|
||||
})
|
||||
.or_else(|| monthly_limit_window_seconds(quota))
|
||||
}
|
||||
|
||||
@@ -266,41 +273,66 @@ fn monthly_limit_reached(quota: &Value) -> bool {
|
||||
.pointer("/monthly_limit/limit_reached")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
|| quota
|
||||
.pointer("/monthly_limit/used_percent")
|
||||
.and_then(number_value)
|
||||
.is_some_and(|used_percent| used_percent >= 100.0)
|
||||
}
|
||||
|
||||
fn core_snapshot_blocked(quota: &Value) -> bool {
|
||||
fn quota_bucket_limit_reached(bucket: &Value) -> bool {
|
||||
!bucket.is_null()
|
||||
&& (bucket
|
||||
.get("limit_reached")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
|| bucket
|
||||
.get("used_percent")
|
||||
.and_then(number_value)
|
||||
.is_some_and(|used_percent| used_percent >= 100.0))
|
||||
}
|
||||
|
||||
fn quota_snapshot_blocked(snapshot: &Value) -> bool {
|
||||
snapshot
|
||||
.get("blocked")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
|| snapshot
|
||||
.get("allowed")
|
||||
.and_then(Value::as_bool)
|
||||
.is_some_and(|allowed| !allowed)
|
||||
|| snapshot_buckets(snapshot).any(quota_bucket_limit_reached)
|
||||
}
|
||||
|
||||
fn snapshot_blocked(quota: &Value) -> bool {
|
||||
quota
|
||||
.get("snapshots")
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.any(|snapshot| {
|
||||
snapshot.get("source").and_then(Value::as_str) == Some("core")
|
||||
&& snapshot
|
||||
.get("blocked")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
.any(quota_snapshot_blocked)
|
||||
}
|
||||
|
||||
fn blocking_bucket(quota: &Value) -> Option<&Value> {
|
||||
let snapshots = quota.get("snapshots").and_then(Value::as_array)?;
|
||||
snapshots
|
||||
.iter()
|
||||
.filter(|snapshot| quota_snapshot_blocked(snapshot))
|
||||
.flat_map(snapshot_buckets)
|
||||
.find(|bucket| quota_bucket_limit_reached(bucket))
|
||||
.or_else(|| {
|
||||
snapshots
|
||||
.iter()
|
||||
.filter(|snapshot| quota_snapshot_blocked(snapshot))
|
||||
.flat_map(snapshot_buckets)
|
||||
.find(|bucket| bucket_reset_at(bucket).is_some())
|
||||
})
|
||||
}
|
||||
|
||||
fn blocking_core_bucket(quota: &Value) -> Option<&Value> {
|
||||
quota
|
||||
.get("snapshots")
|
||||
.and_then(Value::as_array)?
|
||||
.iter()
|
||||
.filter(|snapshot| snapshot.get("source").and_then(Value::as_str) == Some("core"))
|
||||
.flat_map(|snapshot| {
|
||||
["primary", "secondary"]
|
||||
.into_iter()
|
||||
.filter_map(|key| snapshot.get(key))
|
||||
})
|
||||
.find(|bucket| {
|
||||
!bucket.is_null()
|
||||
&& bucket
|
||||
.get("limit_reached")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
})
|
||||
fn snapshot_buckets(snapshot: &Value) -> impl Iterator<Item = &Value> {
|
||||
["primary", "secondary"]
|
||||
.into_iter()
|
||||
.filter_map(|key| snapshot.get(key))
|
||||
.filter(|bucket| !bucket.is_null())
|
||||
}
|
||||
|
||||
fn core_bucket_with_reset_at(quota: &Value) -> Option<&Value> {
|
||||
@@ -318,12 +350,7 @@ fn core_buckets<'a>(quota: &'a Value) -> impl Iterator<Item = &'a Value> + 'a {
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter(|snapshot| snapshot.get("source").and_then(Value::as_str) == Some("core"))
|
||||
.flat_map(|snapshot| {
|
||||
["primary", "secondary"]
|
||||
.into_iter()
|
||||
.filter_map(|key| snapshot.get(key))
|
||||
})
|
||||
.filter(|bucket| !bucket.is_null())
|
||||
.flat_map(snapshot_buckets)
|
||||
}
|
||||
|
||||
fn bucket_reset_at(bucket: &Value) -> Option<i64> {
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
use std::{collections::HashMap, sync::Arc, time::Duration};
|
||||
|
||||
use chrono::Utc;
|
||||
use thiserror::Error;
|
||||
use tokio::time::{sleep, Instant};
|
||||
use tracing::warn;
|
||||
@@ -22,6 +23,7 @@ use crate::upstream::{
|
||||
|
||||
const MIN_REFRESH_INTERVAL_SECS: u64 = 30 * 60;
|
||||
const DEFAULT_REQUEST_SPACING_SECS: u64 = 3;
|
||||
const COOLDOWN_REFRESH_GRACE_SECS: i64 = 30;
|
||||
|
||||
/// 运行时 quota refresh 服务。
|
||||
pub struct RuntimeQuotaRefreshService {
|
||||
@@ -109,6 +111,7 @@ impl RuntimeQuotaRefreshService {
|
||||
};
|
||||
let min_interval = Duration::from_secs(self.min_refresh_interval_secs);
|
||||
let now = Instant::now();
|
||||
let now_utc = Utc::now();
|
||||
|
||||
let mut candidates = accounts
|
||||
.into_iter()
|
||||
@@ -121,9 +124,18 @@ impl RuntimeQuotaRefreshService {
|
||||
|
||||
while let Some(account) = candidates.next() {
|
||||
summary.candidates += 1;
|
||||
if last_refreshed
|
||||
.get(&account.id)
|
||||
.is_some_and(|last| now.duration_since(*last) < min_interval)
|
||||
let cooldown_ready = match account.quota_cooldown_until {
|
||||
Some(cooldown_until) if !quota_cooldown_ready(cooldown_until, now_utc) => {
|
||||
summary.skipped_cooldown += 1;
|
||||
continue;
|
||||
}
|
||||
Some(_) => true,
|
||||
None => false,
|
||||
};
|
||||
if !cooldown_ready
|
||||
&& last_refreshed
|
||||
.get(&account.id)
|
||||
.is_some_and(|last| now.duration_since(*last) < min_interval)
|
||||
{
|
||||
summary.skipped_recent += 1;
|
||||
continue;
|
||||
@@ -221,6 +233,12 @@ impl RuntimeQuotaRefreshService {
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_cooldown_ready(cooldown_until: chrono::DateTime<Utc>, now: chrono::DateTime<Utc>) -> bool {
|
||||
cooldown_until
|
||||
.checked_add_signed(chrono::Duration::seconds(COOLDOWN_REFRESH_GRACE_SECS))
|
||||
.is_some_and(|ready_at| now >= ready_at)
|
||||
}
|
||||
|
||||
/// 单次配额刷新摘要。
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct QuotaRefreshSummary {
|
||||
@@ -228,6 +246,8 @@ pub struct QuotaRefreshSummary {
|
||||
pub scanned: usize,
|
||||
/// 配额锁定或待验证候选账号数。
|
||||
pub candidates: usize,
|
||||
/// 配额冷却宽限期内跳过的账号数。
|
||||
pub skipped_cooldown: usize,
|
||||
/// 成功刷新账号数。
|
||||
pub refreshed: usize,
|
||||
/// 最近刷新过而跳过的账号数。
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
//! 指纹类型、更新与持久化。
|
||||
|
||||
use std::{path::PathBuf, sync::Arc, time::Duration};
|
||||
use std::{
|
||||
path::PathBuf,
|
||||
sync::{Arc, RwLock},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use indexmap::IndexMap;
|
||||
@@ -89,6 +93,38 @@ impl Fingerprint {
|
||||
}
|
||||
}
|
||||
|
||||
/// 运行时共享指纹快照。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RuntimeFingerprint {
|
||||
current: Arc<RwLock<Arc<Fingerprint>>>,
|
||||
}
|
||||
|
||||
impl RuntimeFingerprint {
|
||||
pub fn new(fingerprint: Fingerprint) -> Self {
|
||||
Self {
|
||||
current: Arc::new(RwLock::new(Arc::new(fingerprint))),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn current(&self) -> Arc<Fingerprint> {
|
||||
self.current
|
||||
.read()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.clone()
|
||||
}
|
||||
|
||||
pub fn snapshot(&self) -> Fingerprint {
|
||||
self.current().as_ref().clone()
|
||||
}
|
||||
|
||||
pub fn replace(&self, fingerprint: Fingerprint) {
|
||||
*self
|
||||
.current
|
||||
.write()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner) = Arc::new(fingerprint);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 指纹历史记录与仓储
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -345,7 +381,7 @@ fn json_error(error: serde_json::Error) -> sqlx::Error {
|
||||
/// Appcast 指纹更新检查器。
|
||||
#[derive(Clone)]
|
||||
pub struct UpdateChecker {
|
||||
repository: Option<FingerprintRepository>,
|
||||
repository: FingerprintRepository,
|
||||
client: reqwest::Client,
|
||||
appcast_url: String,
|
||||
extracted_fingerprint_path: PathBuf,
|
||||
@@ -360,7 +396,7 @@ struct InternalState {
|
||||
impl UpdateChecker {
|
||||
/// 使用自定义 client 和路径构造检查器。
|
||||
pub fn with_client(
|
||||
repository: Option<FingerprintRepository>,
|
||||
repository: FingerprintRepository,
|
||||
client: reqwest::Client,
|
||||
appcast_url: impl Into<String>,
|
||||
extracted_fingerprint_path: PathBuf,
|
||||
@@ -406,18 +442,18 @@ impl UpdateChecker {
|
||||
}
|
||||
|
||||
/// 检查并在需要时应用更新。
|
||||
pub async fn check_and_apply_update(&self) -> Result<bool, UpdateError> {
|
||||
pub async fn check_and_apply_update(&self) -> Result<Option<Fingerprint>, UpdateError> {
|
||||
let update_state = self.check_for_update().await?;
|
||||
|
||||
if !update_state.update_available {
|
||||
return Ok(false);
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let (Some(version), Some(build)) = (
|
||||
update_state.latest_version.as_deref(),
|
||||
update_state.latest_build.as_deref(),
|
||||
) else {
|
||||
return Ok(false);
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let state = self.state.lock().await;
|
||||
@@ -433,13 +469,13 @@ impl UpdateChecker {
|
||||
"发现新的 fingerprint 版本"
|
||||
);
|
||||
|
||||
self.apply_version_update(version, build).await?;
|
||||
let updated_fingerprint = self.apply_version_update(version, build).await?;
|
||||
|
||||
let mut state = self.state.lock().await;
|
||||
state.current_version = version.to_string();
|
||||
state.current_build = build.to_string();
|
||||
|
||||
Ok(true)
|
||||
Ok(Some(updated_fingerprint))
|
||||
}
|
||||
|
||||
async fn fetch_appcast(&self) -> Result<String, UpdateError> {
|
||||
@@ -457,20 +493,26 @@ impl UpdateChecker {
|
||||
response.text().await.map_err(UpdateError::from)
|
||||
}
|
||||
|
||||
async fn apply_version_update(&self, version: &str, build: &str) -> Result<(), UpdateError> {
|
||||
async fn apply_version_update(
|
||||
&self,
|
||||
version: &str,
|
||||
build: &str,
|
||||
) -> Result<Fingerprint, UpdateError> {
|
||||
let chromium_version =
|
||||
load_matching_chromium_version(&self.extracted_fingerprint_path, version, build);
|
||||
|
||||
if let Some(repository) = &self.repository {
|
||||
repository
|
||||
.update_current_version(version, build, chromium_version.as_deref())
|
||||
.await?;
|
||||
repository
|
||||
.insert_update_history(version, build, chromium_version.as_deref(), None)
|
||||
.await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
self.repository
|
||||
.update_current_version(version, build, chromium_version.as_deref())
|
||||
.await?;
|
||||
self.repository
|
||||
.insert_update_history(version, build, chromium_version.as_deref(), None)
|
||||
.await?;
|
||||
let updated = self
|
||||
.repository
|
||||
.load_current()
|
||||
.await?
|
||||
.ok_or(sqlx::Error::RowNotFound)?;
|
||||
Ok(updated)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -43,11 +43,17 @@ pub enum ModelServiceError {
|
||||
#[error("no active accounts available for model refresh")]
|
||||
NoAccounts,
|
||||
/// 快照写入失败。
|
||||
#[error("failed to store model snapshot")]
|
||||
StoreSnapshot,
|
||||
#[error("failed to store model snapshot: {source}")]
|
||||
StoreSnapshot {
|
||||
#[source]
|
||||
source: ModelSnapshotStoreError,
|
||||
},
|
||||
/// 刷新后重新读取快照失败。
|
||||
#[error("failed to load model snapshots")]
|
||||
LoadSnapshots,
|
||||
#[error("failed to load model snapshots: {source}")]
|
||||
LoadSnapshots {
|
||||
#[source]
|
||||
source: ModelSnapshotStoreError,
|
||||
},
|
||||
/// 所有计划都刷新失败。
|
||||
#[error("all model refresh plans failed")]
|
||||
AllPlansFailed(ModelRefreshResult),
|
||||
@@ -251,10 +257,10 @@ impl ModelService {
|
||||
}
|
||||
}
|
||||
|
||||
fn map_store_snapshot_error(_: ModelSnapshotStoreError) -> ModelServiceError {
|
||||
ModelServiceError::StoreSnapshot
|
||||
fn map_store_snapshot_error(source: ModelSnapshotStoreError) -> ModelServiceError {
|
||||
ModelServiceError::StoreSnapshot { source }
|
||||
}
|
||||
|
||||
fn map_load_snapshots_error(_: ModelSnapshotStoreError) -> ModelServiceError {
|
||||
ModelServiceError::LoadSnapshots
|
||||
fn map_load_snapshots_error(source: ModelSnapshotStoreError) -> ModelServiceError {
|
||||
ModelServiceError::LoadSnapshots { source }
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ use thiserror::Error;
|
||||
use super::{info::CodexModelInfo, snapshot::ModelPlanSnapshot};
|
||||
|
||||
/// 模型快照存储错误。
|
||||
#[derive(Debug, Error)]
|
||||
#[derive(Debug, Error, PartialEq, Eq)]
|
||||
pub enum ModelSnapshotStoreError {
|
||||
/// 底层存储失败。
|
||||
#[error("model snapshot store operation failed: {message}")]
|
||||
|
||||
@@ -540,11 +540,10 @@ fn quota_snapshot_from_windows(
|
||||
if primary.is_none() && secondary.is_none() {
|
||||
return None;
|
||||
}
|
||||
let blocked = limit_reached.unwrap_or_else(|| {
|
||||
allowed.is_some_and(|allowed| !allowed)
|
||||
|| primary.is_some_and(|window| window.used_percent >= 100.0)
|
||||
|| secondary.is_some_and(|window| window.used_percent >= 100.0)
|
||||
});
|
||||
let blocked = limit_reached.unwrap_or(false)
|
||||
|| allowed.is_some_and(|allowed| !allowed)
|
||||
|| primary.is_some_and(|window| window.used_percent >= 100.0)
|
||||
|| secondary.is_some_and(|window| window.used_percent >= 100.0);
|
||||
Some(json!({
|
||||
"source": source,
|
||||
"limit_name": limit_name,
|
||||
|
||||
@@ -640,18 +640,6 @@ impl CodexResponsesRequest {
|
||||
self.body.get("reasoning")
|
||||
}
|
||||
|
||||
/// 设置 reasoning 配置(模型后缀路由)。
|
||||
pub fn set_reasoning(&mut self, reasoning: Option<Value>) {
|
||||
match reasoning {
|
||||
Some(value) => {
|
||||
self.body.insert("reasoning".to_string(), value);
|
||||
}
|
||||
None => {
|
||||
self.body.remove("reasoning");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 工具定义数组(非数组或空时 None)。
|
||||
pub fn tools(&self) -> Option<&[Value]> {
|
||||
self.body
|
||||
|
||||
@@ -834,6 +834,17 @@ pub fn is_terminal_websocket_event(event: &str) -> bool {
|
||||
matches!(event, "response.completed" | "response.failed" | "error")
|
||||
}
|
||||
|
||||
/// 判断是否为「首个内容帧之前的纯生命周期帧」。
|
||||
///
|
||||
/// `response.created` / `response.in_progress` 仅表示上游已受理请求,尚未开始产出
|
||||
/// 真实内容(含 reasoning)。首 token 熔断计时不应被它们解除,否则连接一收到
|
||||
/// `response.created`(通常 1s 内)就被误判为「已开始产出」,使真正的首内容帧
|
||||
/// 可以无限拖延而不触发保护。其余任何帧(reasoning/delta/output_item 等)到达即视为
|
||||
/// 上游已开始真实产出,解除熔断——对慢思考模型友好。
|
||||
pub fn is_pre_content_lifecycle_event(event: &str) -> bool {
|
||||
matches!(event, "response.created" | "response.in_progress")
|
||||
}
|
||||
|
||||
/// 将 WebSocket 错误帧映射为上游 HTTP 状态分类。
|
||||
pub fn classify_websocket_error_frame(raw: &str) -> Option<ClassifiedWebSocketError> {
|
||||
let value = serde_json::from_str::<Value>(raw).ok()?;
|
||||
|
||||
@@ -22,7 +22,7 @@ use thiserror::Error;
|
||||
use tokio_tungstenite::tungstenite::handshake::client::generate_key;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::upstream::fingerprint::Fingerprint;
|
||||
use crate::upstream::fingerprint::{Fingerprint, RuntimeFingerprint};
|
||||
use crate::upstream::models::backend_entry::BackendModelEntry;
|
||||
use crate::upstream::protocol::events::{extract_sse_usage, retry_after_seconds_from_body};
|
||||
use crate::upstream::protocol::responses::{
|
||||
@@ -34,6 +34,7 @@ use crate::upstream::protocol::websocket::{
|
||||
websocket_audit_artifact_from_attempt, websocket_payload_audit_snapshot,
|
||||
};
|
||||
|
||||
use super::diagnostics::CodexUpstreamDiagnostics;
|
||||
use super::endpoints::{endpoint_url, CODEX_RESPONSES_COMPACT_PATH, CODEX_RESPONSES_PATH};
|
||||
use super::headers::{
|
||||
build_ordered_codex_base_headers, insert_optional_header, insert_ordered_headers,
|
||||
@@ -46,7 +47,9 @@ use super::websocket::{
|
||||
write_websocket_audit_artifact_from_env, CodexWebSocketConnection, CodexWebSocketExchangeError,
|
||||
CodexWebSocketRateLimitHeaderUpdates, CodexWebSocketTurnStateUpdate,
|
||||
};
|
||||
use super::websocket_pool::{CodexWebSocketPool, CodexWebSocketPoolKey, WebSocketPoolDecision};
|
||||
use super::websocket_pool::{
|
||||
CodexWebSocketPool, CodexWebSocketPoolKey, WebSocketPoolDecision, DEFAULT_FIRST_TOKEN_TIMEOUT,
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Constants
|
||||
@@ -127,6 +130,8 @@ pub enum CodexClientError {
|
||||
body: String,
|
||||
/// 推导出的重试秒数。
|
||||
retry_after_seconds: Option<u64>,
|
||||
/// 上游诊断元数据。
|
||||
diagnostics: CodexUpstreamDiagnostics,
|
||||
/// 上游透传的 `set-cookie` 列表。
|
||||
set_cookie_headers: Vec<String>,
|
||||
},
|
||||
@@ -253,6 +258,8 @@ pub struct CodexBackendResponse {
|
||||
pub first_token_ms: Option<i64>,
|
||||
/// WebSocket 连接池决策。
|
||||
pub websocket_pool_decision: Option<WebSocketPoolDecision>,
|
||||
/// 上游诊断元数据。
|
||||
pub diagnostics: CodexUpstreamDiagnostics,
|
||||
}
|
||||
|
||||
/// Codex Responses 实际使用的上游传输。
|
||||
@@ -300,6 +307,8 @@ pub struct CodexBackendStreamingResponse {
|
||||
pub turn_state_update: Option<CodexTurnStateUpdate>,
|
||||
/// WebSocket 连接池决策。
|
||||
pub websocket_pool_decision: Option<WebSocketPoolDecision>,
|
||||
/// 上游诊断元数据。
|
||||
pub diagnostics: CodexUpstreamDiagnostics,
|
||||
}
|
||||
|
||||
/// Codex compact 端点响应。
|
||||
@@ -311,6 +320,8 @@ pub struct CodexCompactResponse {
|
||||
pub set_cookie_headers: Vec<String>,
|
||||
/// 上游透传的限流头。
|
||||
pub rate_limit_headers: Vec<(String, String)>,
|
||||
/// 上游诊断元数据。
|
||||
pub diagnostics: CodexUpstreamDiagnostics,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -322,21 +333,32 @@ pub struct CodexCompactResponse {
|
||||
pub struct CodexBackendClient {
|
||||
pub(super) client: Client,
|
||||
pub(super) base_url: String,
|
||||
fingerprint: Fingerprint,
|
||||
fingerprint: RuntimeFingerprint,
|
||||
websocket_pool: Option<Arc<CodexWebSocketPool>>,
|
||||
websocket_first_token_timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
impl CodexBackendClient {
|
||||
/// 构造客户端。
|
||||
pub fn new(client: Client, base_url: impl Into<String>, fingerprint: Fingerprint) -> Self {
|
||||
pub fn new(
|
||||
client: Client,
|
||||
base_url: impl Into<String>,
|
||||
fingerprint: RuntimeFingerprint,
|
||||
) -> Self {
|
||||
Self {
|
||||
client,
|
||||
base_url: base_url.into().trim_end_matches('/').to_string(),
|
||||
fingerprint,
|
||||
websocket_pool: None,
|
||||
websocket_first_token_timeout: Some(DEFAULT_FIRST_TOKEN_TIMEOUT),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_websocket_first_token_timeout(mut self, timeout: Option<Duration>) -> Self {
|
||||
self.websocket_first_token_timeout = timeout.filter(|timeout| !timeout.is_zero());
|
||||
self
|
||||
}
|
||||
|
||||
/// 为 Responses WebSocket 请求启用连接池。
|
||||
pub fn with_websocket_pool(mut self, pool: Arc<CodexWebSocketPool>) -> Self {
|
||||
self.websocket_pool = Some(pool);
|
||||
@@ -367,6 +389,7 @@ impl CodexBackendClient {
|
||||
.send()
|
||||
.await?;
|
||||
let status = response.status();
|
||||
let diagnostics = response_meta::diagnostics(Some(status.as_u16()), response.headers());
|
||||
let turn_state = response_meta::turn_state(response.headers());
|
||||
let set_cookie_headers = response_meta::set_cookie_headers(response.headers());
|
||||
let rate_limit_headers = response_meta::rate_limit_headers(response.headers());
|
||||
@@ -379,6 +402,7 @@ impl CodexBackendClient {
|
||||
retry_after_seconds: retry_after_seconds
|
||||
.or_else(|| retry_after_seconds_from_body(&body)),
|
||||
body,
|
||||
diagnostics,
|
||||
set_cookie_headers,
|
||||
});
|
||||
}
|
||||
@@ -401,6 +425,7 @@ impl CodexBackendClient {
|
||||
rate_limit_headers,
|
||||
first_token_ms,
|
||||
websocket_pool_decision: None,
|
||||
diagnostics,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -419,6 +444,7 @@ impl CodexBackendClient {
|
||||
.send()
|
||||
.await?;
|
||||
let status = response.status();
|
||||
let diagnostics = response_meta::diagnostics(Some(status.as_u16()), response.headers());
|
||||
let turn_state = response_meta::turn_state(response.headers());
|
||||
let set_cookie_headers = response_meta::set_cookie_headers(response.headers());
|
||||
let rate_limit_headers = response_meta::rate_limit_headers(response.headers());
|
||||
@@ -431,6 +457,7 @@ impl CodexBackendClient {
|
||||
retry_after_seconds: retry_after_seconds
|
||||
.or_else(|| retry_after_seconds_from_body(&body)),
|
||||
body,
|
||||
diagnostics,
|
||||
set_cookie_headers,
|
||||
});
|
||||
}
|
||||
@@ -444,6 +471,7 @@ impl CodexBackendClient {
|
||||
rate_limit_header_updates: None,
|
||||
turn_state_update: None,
|
||||
websocket_pool_decision: None,
|
||||
diagnostics,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -590,10 +618,23 @@ impl CodexBackendClient {
|
||||
let pool_log_context = pool_key.as_ref().map(WebSocketPoolLogContext::from_key);
|
||||
let exchange = match (self.websocket_pool.as_deref(), pool_key) {
|
||||
(Some(pool), Some(key)) => {
|
||||
execute_response_create_request_with_pool(&prepared, Some((pool, key)), started_at)
|
||||
.await
|
||||
execute_response_create_request_with_pool(
|
||||
&prepared,
|
||||
Some((pool, key)),
|
||||
started_at,
|
||||
self.websocket_first_token_timeout,
|
||||
)
|
||||
.await
|
||||
}
|
||||
_ => {
|
||||
execute_response_create_request_with_pool(
|
||||
&prepared,
|
||||
None,
|
||||
started_at,
|
||||
self.websocket_first_token_timeout,
|
||||
)
|
||||
.await
|
||||
}
|
||||
_ => execute_response_create_request_with_pool(&prepared, None, started_at).await,
|
||||
}
|
||||
.map_err(websocket_exchange_error_to_client_error)?;
|
||||
log_websocket_pool_decision(
|
||||
@@ -612,6 +653,7 @@ impl CodexBackendClient {
|
||||
rate_limit_headers: exchange.rate_limit_headers,
|
||||
first_token_ms: exchange.first_token_ms,
|
||||
websocket_pool_decision: exchange.pool_decision,
|
||||
diagnostics: exchange.diagnostics,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -642,9 +684,21 @@ impl CodexBackendClient {
|
||||
let pool_log_context = pool_key.as_ref().map(WebSocketPoolLogContext::from_key);
|
||||
let exchange = match (self.websocket_pool.as_deref(), pool_key) {
|
||||
(Some(pool), Some(key)) => {
|
||||
execute_response_create_request_stream_with_pool(&prepared, Some((pool, key))).await
|
||||
execute_response_create_request_stream_with_pool(
|
||||
&prepared,
|
||||
Some((pool, key)),
|
||||
self.websocket_first_token_timeout,
|
||||
)
|
||||
.await
|
||||
}
|
||||
_ => {
|
||||
execute_response_create_request_stream_with_pool(
|
||||
&prepared,
|
||||
None,
|
||||
self.websocket_first_token_timeout,
|
||||
)
|
||||
.await
|
||||
}
|
||||
_ => execute_response_create_request_stream_with_pool(&prepared, None).await,
|
||||
}
|
||||
.map_err(websocket_exchange_error_to_client_error)?;
|
||||
log_websocket_pool_decision(
|
||||
@@ -667,6 +721,7 @@ impl CodexBackendClient {
|
||||
rate_limit_header_updates: Some(exchange.rate_limit_header_updates),
|
||||
turn_state_update: Some(exchange.turn_state_update),
|
||||
websocket_pool_decision: exchange.pool_decision,
|
||||
diagnostics: exchange.diagnostics,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -704,6 +759,7 @@ impl CodexBackendClient {
|
||||
.await?;
|
||||
|
||||
let status = response.status();
|
||||
let diagnostics = response_meta::diagnostics(Some(status.as_u16()), response.headers());
|
||||
let set_cookie_headers = response_meta::set_cookie_headers(response.headers());
|
||||
let rate_limit_headers = response_meta::rate_limit_headers(response.headers());
|
||||
let retry_after_seconds = retry_after_seconds(response.headers(), None);
|
||||
@@ -715,6 +771,7 @@ impl CodexBackendClient {
|
||||
retry_after_seconds: retry_after_seconds
|
||||
.or_else(|| retry_after_seconds_from_body(&body)),
|
||||
body,
|
||||
diagnostics,
|
||||
set_cookie_headers,
|
||||
});
|
||||
}
|
||||
@@ -727,12 +784,16 @@ impl CodexBackendClient {
|
||||
"Compact response is not valid JSON: {}",
|
||||
truncate_for_error(&body)
|
||||
),
|
||||
diagnostics: CodexUpstreamDiagnostics::with_status(
|
||||
StatusCode::BAD_GATEWAY.as_u16(),
|
||||
),
|
||||
set_cookie_headers: set_cookie_headers.clone(),
|
||||
})?;
|
||||
Ok(CodexCompactResponse {
|
||||
body: parsed,
|
||||
set_cookie_headers,
|
||||
rate_limit_headers,
|
||||
diagnostics,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -741,17 +802,18 @@ impl CodexBackendClient {
|
||||
&self,
|
||||
context: CodexRequestContext<'_>,
|
||||
) -> CodexClientResult<Vec<BackendModelEntry>> {
|
||||
let fingerprint = self.fingerprint.current();
|
||||
let endpoints = [
|
||||
format!(
|
||||
"{}/codex/models?client_version={}",
|
||||
self.base_url, self.fingerprint.app_version
|
||||
self.base_url, fingerprint.app_version
|
||||
),
|
||||
format!("{}/models", self.base_url),
|
||||
format!("{}/sentinel/chat-requirements", self.base_url),
|
||||
];
|
||||
|
||||
for endpoint in endpoints {
|
||||
let headers = self.auxiliary_request_headers(context)?;
|
||||
let headers = self.auxiliary_request_headers(&fingerprint, context)?;
|
||||
let response = self.client.get(endpoint).headers(headers).send().await?;
|
||||
if !response.status().is_success() {
|
||||
continue;
|
||||
@@ -767,6 +829,7 @@ impl CodexBackendClient {
|
||||
status: StatusCode::BAD_GATEWAY,
|
||||
retry_after_seconds: None,
|
||||
body: "backend model catalog is unavailable".to_string(),
|
||||
diagnostics: CodexUpstreamDiagnostics::with_status(StatusCode::BAD_GATEWAY.as_u16()),
|
||||
set_cookie_headers: Vec::new(),
|
||||
})
|
||||
}
|
||||
@@ -787,9 +850,10 @@ impl CodexBackendClient {
|
||||
}
|
||||
|
||||
fn request_headers(&self, context: CodexRequestContext<'_>) -> CodexClientResult<HeaderMap> {
|
||||
let fingerprint = self.fingerprint.current();
|
||||
let request_id = context.session_id.unwrap_or(context.request_id);
|
||||
let ordered_headers = build_ordered_codex_base_headers(
|
||||
&self.fingerprint,
|
||||
&fingerprint,
|
||||
context.access_token,
|
||||
context.account_id,
|
||||
);
|
||||
@@ -839,13 +903,11 @@ impl CodexBackendClient {
|
||||
|
||||
fn auxiliary_request_headers(
|
||||
&self,
|
||||
fingerprint: &Fingerprint,
|
||||
context: CodexRequestContext<'_>,
|
||||
) -> CodexClientResult<HeaderMap> {
|
||||
let ordered_headers = build_ordered_codex_base_headers(
|
||||
&self.fingerprint,
|
||||
context.access_token,
|
||||
context.account_id,
|
||||
);
|
||||
let ordered_headers =
|
||||
build_ordered_codex_base_headers(fingerprint, context.access_token, context.account_id);
|
||||
|
||||
let mut headers = HeaderMap::new();
|
||||
insert_ordered_headers(&mut headers, &ordered_headers)?;
|
||||
@@ -866,10 +928,11 @@ impl CodexBackendClient {
|
||||
&self,
|
||||
context: CodexRequestContext<'_>,
|
||||
) -> CodexClientResult<HeaderMap> {
|
||||
let fingerprint = self.fingerprint.current();
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(
|
||||
USER_AGENT,
|
||||
HeaderValue::from_str(&self.fingerprint.user_agent())?,
|
||||
HeaderValue::from_str(&fingerprint.user_agent())?,
|
||||
);
|
||||
headers.insert(
|
||||
AUTHORIZATION,
|
||||
@@ -877,7 +940,7 @@ impl CodexBackendClient {
|
||||
);
|
||||
headers.insert(
|
||||
HeaderName::from_static("originator"),
|
||||
HeaderValue::from_str(&self.fingerprint.originator)?,
|
||||
HeaderValue::from_str(&fingerprint.originator)?,
|
||||
);
|
||||
insert_optional_header(&mut headers, "chatgpt-account-id", context.account_id)?;
|
||||
insert_optional_header(&mut headers, "cookie", context.cookie_header)?;
|
||||
@@ -889,8 +952,9 @@ impl CodexBackendClient {
|
||||
&self,
|
||||
context: CodexRequestContext<'_>,
|
||||
) -> CodexClientResult<HeaderMap> {
|
||||
let fingerprint = self.fingerprint.current();
|
||||
let ordered_headers = build_ordered_codex_base_headers(
|
||||
&self.fingerprint,
|
||||
&fingerprint,
|
||||
context.access_token,
|
||||
context.account_id,
|
||||
);
|
||||
@@ -1167,17 +1231,17 @@ fn websocket_exchange_error_to_client_error(
|
||||
error: CodexWebSocketExchangeError,
|
||||
) -> CodexClientError {
|
||||
match error {
|
||||
CodexWebSocketExchangeError::Upstream {
|
||||
status_code,
|
||||
retry_after_seconds,
|
||||
body,
|
||||
set_cookie_headers,
|
||||
} => CodexClientError::Upstream {
|
||||
status: StatusCode::from_u16(status_code).unwrap_or(StatusCode::BAD_GATEWAY),
|
||||
body,
|
||||
retry_after_seconds,
|
||||
set_cookie_headers,
|
||||
},
|
||||
CodexWebSocketExchangeError::Upstream(upstream) => {
|
||||
let upstream = *upstream;
|
||||
CodexClientError::Upstream {
|
||||
status: StatusCode::from_u16(upstream.status_code)
|
||||
.unwrap_or(StatusCode::BAD_GATEWAY),
|
||||
body: upstream.body,
|
||||
retry_after_seconds: upstream.retry_after_seconds,
|
||||
diagnostics: upstream.diagnostics,
|
||||
set_cookie_headers: upstream.set_cookie_headers,
|
||||
}
|
||||
}
|
||||
error => CodexClientError::WebSocket(error),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
//! Upstream response diagnostics captured at the transport boundary.
|
||||
|
||||
use reqwest::header::HeaderMap;
|
||||
|
||||
const UPSTREAM_REQUEST_ID_HEADERS: &[&str] = &[
|
||||
"x-request-id",
|
||||
"x-openai-request-id",
|
||||
"openai-request-id",
|
||||
"request-id",
|
||||
];
|
||||
|
||||
const UPSTREAM_TRACE_HEADERS: &[&str] = &[
|
||||
"x-request-id",
|
||||
"x-openai-request-id",
|
||||
"openai-request-id",
|
||||
"request-id",
|
||||
"cf-ray",
|
||||
];
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct CodexUpstreamDiagnostics {
|
||||
pub status_code: Option<u16>,
|
||||
pub request_id: Option<String>,
|
||||
pub trace_headers: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
impl CodexUpstreamDiagnostics {
|
||||
pub fn from_headers(status_code: Option<u16>, headers: &HeaderMap) -> Self {
|
||||
Self {
|
||||
status_code,
|
||||
request_id: first_header(headers, UPSTREAM_REQUEST_ID_HEADERS),
|
||||
trace_headers: trace_headers(headers),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_status(status_code: u16) -> Self {
|
||||
Self {
|
||||
status_code: Some(status_code),
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.status_code.is_none() && self.request_id.is_none() && self.trace_headers.is_empty()
|
||||
}
|
||||
|
||||
pub fn cf_ray(&self) -> Option<&str> {
|
||||
self.trace_headers
|
||||
.iter()
|
||||
.find(|(name, _)| name.eq_ignore_ascii_case("cf-ray"))
|
||||
.map(|(_, value)| value.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
fn first_header(headers: &HeaderMap, names: &[&str]) -> Option<String> {
|
||||
names.iter().find_map(|name| header_value(headers, name))
|
||||
}
|
||||
|
||||
fn trace_headers(headers: &HeaderMap) -> Vec<(String, String)> {
|
||||
UPSTREAM_TRACE_HEADERS
|
||||
.iter()
|
||||
.filter_map(|name| header_value(headers, name).map(|value| ((*name).to_string(), value)))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn header_value(headers: &HeaderMap, name: &str) -> Option<String> {
|
||||
headers
|
||||
.get(name)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToString::to_string)
|
||||
}
|
||||
@@ -1,6 +1,7 @@
|
||||
//! Codex HTTP/SSE/WebSocket 上游 transport。
|
||||
|
||||
pub mod client;
|
||||
pub mod diagnostics;
|
||||
pub mod endpoints;
|
||||
pub mod headers;
|
||||
mod response_meta;
|
||||
@@ -19,6 +20,7 @@ pub use self::{
|
||||
CodexModelCatalogRequest, CodexRateLimitHeaderUpdates, CodexRequestContext,
|
||||
CodexTurnStateUpdate,
|
||||
},
|
||||
diagnostics::CodexUpstreamDiagnostics,
|
||||
endpoints::{
|
||||
endpoint_request_path, endpoint_url, usage_endpoint_urls, CODEX_RESPONSES_COMPACT_PATH,
|
||||
CODEX_RESPONSES_PATH, CODEX_USAGE_API_PATH, CODEX_USAGE_PATH, WHAM_USAGE_PATH,
|
||||
|
||||
@@ -6,6 +6,15 @@ use reqwest::header::{HeaderMap, SET_COOKIE};
|
||||
|
||||
use crate::upstream::protocol::responses::update_first_response_event_ms;
|
||||
|
||||
use super::diagnostics::CodexUpstreamDiagnostics;
|
||||
|
||||
pub(super) fn diagnostics(
|
||||
status_code: Option<u16>,
|
||||
headers: &HeaderMap,
|
||||
) -> CodexUpstreamDiagnostics {
|
||||
CodexUpstreamDiagnostics::from_headers(status_code, headers)
|
||||
}
|
||||
|
||||
pub(super) fn turn_state(headers: &HeaderMap) -> Option<String> {
|
||||
headers
|
||||
.get("x-codex-turn-state")
|
||||
|
||||
@@ -6,7 +6,8 @@ use crate::upstream::protocol::events::retry_after_seconds_from_body;
|
||||
use super::{
|
||||
client::{read_capped_error_body, retry_after_seconds, truncate_for_error},
|
||||
endpoints::usage_endpoint_urls,
|
||||
CodexBackendClient, CodexClientError, CodexClientResult, CodexRequestContext,
|
||||
response_meta, CodexBackendClient, CodexClientError, CodexClientResult, CodexRequestContext,
|
||||
CodexUpstreamDiagnostics,
|
||||
};
|
||||
|
||||
impl CodexBackendClient {
|
||||
@@ -18,6 +19,7 @@ impl CodexBackendClient {
|
||||
let headers = self.usage_request_headers(context)?;
|
||||
let response = self.client.get(endpoint).headers(headers).send().await?;
|
||||
let status = response.status();
|
||||
let diagnostics = response_meta::diagnostics(Some(status.as_u16()), response.headers());
|
||||
let retry_after_seconds = retry_after_seconds(response.headers(), None);
|
||||
|
||||
if status == StatusCode::NOT_FOUND {
|
||||
@@ -31,6 +33,7 @@ impl CodexBackendClient {
|
||||
retry_after_seconds: retry_after_seconds
|
||||
.or_else(|| retry_after_seconds_from_body(&body)),
|
||||
body,
|
||||
diagnostics,
|
||||
set_cookie_headers: Vec::new(),
|
||||
});
|
||||
}
|
||||
@@ -49,6 +52,7 @@ impl CodexBackendClient {
|
||||
|| "usage endpoint is unavailable".to_string(),
|
||||
|body| format!("invalid usage response: {}", truncate_for_error(&body)),
|
||||
),
|
||||
diagnostics: CodexUpstreamDiagnostics::with_status(StatusCode::BAD_GATEWAY.as_u16()),
|
||||
set_cookie_headers: Vec::new(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
//! Codex WebSocket 连接建立(关键函数)。
|
||||
|
||||
use std::{
|
||||
io,
|
||||
collections::VecDeque,
|
||||
fmt, io,
|
||||
path::{Path, PathBuf},
|
||||
pin::Pin,
|
||||
sync::Arc,
|
||||
@@ -29,19 +30,19 @@ use crate::upstream::protocol::events::{self, TokenUsage};
|
||||
use crate::upstream::protocol::responses::CodexResponsesRequest;
|
||||
use crate::upstream::protocol::sse::SseError;
|
||||
use crate::upstream::protocol::websocket::{
|
||||
classify_websocket_error_frame, is_terminal_websocket_event,
|
||||
classify_websocket_error_frame, is_pre_content_lifecycle_event, is_terminal_websocket_event,
|
||||
retry_after_seconds_from_wrapped_error_headers, websocket_event_to_sse_frame,
|
||||
websocket_event_type, websocket_incomplete_response_reason, websocket_metadata_turn_state,
|
||||
websocket_response_completed_parse_error, websocket_response_create_payload_text,
|
||||
OpeningAuditHeader, OpeningAuditSnapshot, WebSocketAuditArtifact,
|
||||
};
|
||||
|
||||
use super::response_meta;
|
||||
use super::websocket_pool::{
|
||||
CodexWebSocketConnectionMetadata, CodexWebSocketPool, CodexWebSocketPoolKey,
|
||||
PooledWebSocketConnection, WebSocketPoolAcquire, WebSocketPoolDecision,
|
||||
};
|
||||
use super::websocket_pump::{PumpKeepalive, PumpedWebSocket, RawWsStream};
|
||||
use super::{diagnostics::CodexUpstreamDiagnostics, response_meta};
|
||||
use uuid::Uuid;
|
||||
|
||||
const REDACTED_HEADER_VALUE: &str = "<redacted>";
|
||||
@@ -50,6 +51,7 @@ const WEBSOCKET_EXTENSIONS: &str = "permessage-deflate; client_max_window_bits";
|
||||
const WEBSOCKET_RECEIVE_IDLE_TIMEOUT: Duration = Duration::from_secs(20);
|
||||
const WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
|
||||
const WEBSOCKET_STREAM_BUFFER: usize = 16;
|
||||
const WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS: usize = 2;
|
||||
/// WebSocket audit artifact 输出目录环境变量。
|
||||
pub const WS_AUDIT_DIR_ENV: &str = "CODEX_PROXY_WS_AUDIT_DIR";
|
||||
|
||||
@@ -118,6 +120,8 @@ pub struct CodexWebSocketExchange {
|
||||
pub first_token_ms: Option<i64>,
|
||||
/// WebSocket 连接池决策。
|
||||
pub pool_decision: Option<WebSocketPoolDecision>,
|
||||
/// 上游诊断元数据。
|
||||
pub diagnostics: CodexUpstreamDiagnostics,
|
||||
}
|
||||
|
||||
/// Responses WebSocket live SSE exchange result.
|
||||
@@ -136,6 +140,8 @@ pub struct CodexWebSocketStreamingExchange {
|
||||
pub turn_state_update: CodexWebSocketTurnStateUpdate,
|
||||
/// WebSocket 连接池决策。
|
||||
pub pool_decision: Option<WebSocketPoolDecision>,
|
||||
/// 上游诊断元数据。
|
||||
pub diagnostics: CodexUpstreamDiagnostics,
|
||||
}
|
||||
|
||||
/// Responses WebSocket live SSE byte stream.
|
||||
@@ -159,17 +165,8 @@ pub enum CodexWebSocketExchangeError {
|
||||
#[error("invalid websocket SSE response: {0}")]
|
||||
InvalidSse(#[from] SseError),
|
||||
/// 上游 WebSocket 错误帧。
|
||||
#[error("websocket upstream error {status_code}: {body}")]
|
||||
Upstream {
|
||||
/// HTTP-style upstream status code.
|
||||
status_code: u16,
|
||||
/// 推导出的重试秒数。
|
||||
retry_after_seconds: Option<u64>,
|
||||
/// 原始错误帧。
|
||||
body: String,
|
||||
/// 上游透传的 `set-cookie` 列表。
|
||||
set_cookie_headers: Vec<String>,
|
||||
},
|
||||
#[error("{0}")]
|
||||
Upstream(Box<CodexWebSocketUpstreamError>),
|
||||
/// 上游返回 `response.incomplete`。
|
||||
#[error("Incomplete response returned, reason: {reason}")]
|
||||
IncompleteResponse {
|
||||
@@ -200,6 +197,55 @@ pub enum CodexWebSocketExchangeError {
|
||||
/// 底层失效原因。
|
||||
message: String,
|
||||
},
|
||||
/// 首个内容帧在配置的绝对超时内未到达(连接落到病态上游后端)。
|
||||
#[error("websocket first content frame not received within {timeout:?}")]
|
||||
FirstTokenTimeout {
|
||||
/// 首 token 超时时长。
|
||||
timeout: Duration,
|
||||
},
|
||||
}
|
||||
|
||||
/// WebSocket 上游错误帧载荷。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CodexWebSocketUpstreamError {
|
||||
/// HTTP-style upstream status code.
|
||||
pub status_code: u16,
|
||||
/// 推导出的重试秒数。
|
||||
pub retry_after_seconds: Option<u64>,
|
||||
/// 原始错误帧。
|
||||
pub body: String,
|
||||
/// 上游透传的 `set-cookie` 列表。
|
||||
pub set_cookie_headers: Vec<String>,
|
||||
/// 上游诊断元数据。
|
||||
pub diagnostics: CodexUpstreamDiagnostics,
|
||||
}
|
||||
|
||||
impl fmt::Display for CodexWebSocketUpstreamError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"websocket upstream error {}: {}",
|
||||
self.status_code, self.body
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl CodexWebSocketExchangeError {
|
||||
fn upstream(
|
||||
status_code: u16,
|
||||
retry_after_seconds: Option<u64>,
|
||||
body: String,
|
||||
set_cookie_headers: Vec<String>,
|
||||
diagnostics: CodexUpstreamDiagnostics,
|
||||
) -> Self {
|
||||
Self::Upstream(Box::new(CodexWebSocketUpstreamError {
|
||||
status_code,
|
||||
retry_after_seconds,
|
||||
body,
|
||||
set_cookie_headers,
|
||||
diagnostics,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
impl CodexWebSocketRequest {
|
||||
@@ -381,12 +427,13 @@ fn websocket_opening_error(response: &WsResponse<Option<Vec<u8>>>) -> CodexWebSo
|
||||
.and_then(|value| value.parse::<u64>().ok())
|
||||
.filter(|seconds| *seconds > 0)
|
||||
.or_else(|| events::retry_after_seconds_from_body(&body));
|
||||
CodexWebSocketExchangeError::Upstream {
|
||||
CodexWebSocketExchangeError::upstream(
|
||||
status_code,
|
||||
retry_after_seconds,
|
||||
body,
|
||||
set_cookie_headers: response_meta::set_cookie_headers(response.headers()),
|
||||
}
|
||||
response_meta::set_cookie_headers(response.headers()),
|
||||
response_meta::diagnostics(Some(status_code), response.headers()),
|
||||
)
|
||||
}
|
||||
|
||||
fn request_line_for_endpoint(endpoint: &str) -> String {
|
||||
@@ -412,30 +459,94 @@ fn request_line_for_endpoint(endpoint: &str) -> String {
|
||||
async fn execute_fresh_response_create_request(
|
||||
request: &CodexWebSocketRequest,
|
||||
started_at: Instant,
|
||||
first_token_timeout: Option<Duration>,
|
||||
) -> Result<CodexWebSocketExchange, CodexWebSocketExchangeError> {
|
||||
let (websocket, response) =
|
||||
connect_pumped_websocket(request.connection(), PumpKeepalive::disabled()).await?;
|
||||
let first_token_started_at = Instant::now();
|
||||
websocket
|
||||
.send(Message::Text(request.payload_text().to_string().into()))
|
||||
.await?;
|
||||
|
||||
let metadata = websocket_connection_metadata(&response);
|
||||
let (exchange, _websocket, _metadata) =
|
||||
collect_websocket_response(websocket, metadata, false, started_at).await?;
|
||||
let (exchange, _websocket, _metadata) = collect_websocket_response(
|
||||
websocket,
|
||||
metadata,
|
||||
false,
|
||||
started_at,
|
||||
first_token_started_at,
|
||||
first_token_timeout,
|
||||
)
|
||||
.await?;
|
||||
Ok(exchange)
|
||||
}
|
||||
|
||||
async fn execute_fresh_response_create_request_stream(
|
||||
request: &CodexWebSocketRequest,
|
||||
first_token_timeout: Option<Duration>,
|
||||
) -> Result<CodexWebSocketStreamingExchange, CodexWebSocketExchangeError> {
|
||||
let (websocket, response) =
|
||||
let (mut websocket, response) =
|
||||
connect_pumped_websocket(request.connection(), PumpKeepalive::disabled()).await?;
|
||||
let first_token_started_at = Instant::now();
|
||||
websocket
|
||||
.send(Message::Text(request.payload_text().to_string().into()))
|
||||
.await?;
|
||||
|
||||
let metadata = websocket_connection_metadata(&response);
|
||||
Ok(stream_websocket_response(websocket, metadata, None, None))
|
||||
let mut metadata = websocket_connection_metadata(&response);
|
||||
let prefetched_frames = match prefetch_first_content_stream_frames(
|
||||
&mut websocket,
|
||||
&mut metadata,
|
||||
first_token_started_at,
|
||||
first_token_timeout,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(prefetched_frames) => prefetched_frames,
|
||||
Err(error) => {
|
||||
websocket.close().await;
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
Ok(stream_websocket_response(
|
||||
websocket,
|
||||
metadata,
|
||||
None,
|
||||
prefetched_frames,
|
||||
))
|
||||
}
|
||||
|
||||
async fn execute_fresh_response_create_request_with_retries(
|
||||
request: &CodexWebSocketRequest,
|
||||
started_at: Instant,
|
||||
first_token_timeout: Option<Duration>,
|
||||
attempts: usize,
|
||||
) -> Result<CodexWebSocketExchange, CodexWebSocketExchangeError> {
|
||||
let attempts = attempts.max(1);
|
||||
for attempt in 1..=attempts {
|
||||
match execute_fresh_response_create_request(request, started_at, first_token_timeout).await
|
||||
{
|
||||
Ok(exchange) => return Ok(exchange),
|
||||
Err(error) if attempt < attempts && is_first_token_timeout(&error) => continue,
|
||||
Err(error) => return Err(error),
|
||||
}
|
||||
}
|
||||
unreachable!("fresh websocket retry loop always returns");
|
||||
}
|
||||
|
||||
async fn execute_fresh_response_create_request_stream_with_retries(
|
||||
request: &CodexWebSocketRequest,
|
||||
first_token_timeout: Option<Duration>,
|
||||
attempts: usize,
|
||||
) -> Result<CodexWebSocketStreamingExchange, CodexWebSocketExchangeError> {
|
||||
let attempts = attempts.max(1);
|
||||
for attempt in 1..=attempts {
|
||||
match execute_fresh_response_create_request_stream(request, first_token_timeout).await {
|
||||
Ok(exchange) => return Ok(exchange),
|
||||
Err(error) if attempt < attempts && is_first_token_timeout(&error) => continue,
|
||||
Err(error) => return Err(error),
|
||||
}
|
||||
}
|
||||
unreachable!("fresh websocket stream retry loop always returns");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -446,9 +557,16 @@ pub(crate) async fn execute_response_create_request_with_pool(
|
||||
request: &CodexWebSocketRequest,
|
||||
pool: Option<(&CodexWebSocketPool, CodexWebSocketPoolKey)>,
|
||||
started_at: Instant,
|
||||
fallback_first_token_timeout: Option<Duration>,
|
||||
) -> Result<CodexWebSocketExchange, CodexWebSocketExchangeError> {
|
||||
let Some((pool, key)) = pool else {
|
||||
return execute_fresh_response_create_request(request, started_at).await;
|
||||
return execute_fresh_response_create_request_with_retries(
|
||||
request,
|
||||
started_at,
|
||||
fallback_first_token_timeout,
|
||||
WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS,
|
||||
)
|
||||
.await;
|
||||
};
|
||||
|
||||
match pool.acquire(&key).await {
|
||||
@@ -459,9 +577,15 @@ pub(crate) async fn execute_response_create_request_with_pool(
|
||||
match result {
|
||||
Err(CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame {
|
||||
..
|
||||
}) => {
|
||||
let mut exchange =
|
||||
execute_fresh_response_create_request(request, started_at).await?;
|
||||
})
|
||||
| Err(CodexWebSocketExchangeError::FirstTokenTimeout { .. }) => {
|
||||
let mut exchange = execute_fresh_response_create_request_with_retries(
|
||||
request,
|
||||
started_at,
|
||||
pool.first_token_timeout(),
|
||||
WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS,
|
||||
)
|
||||
.await?;
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::retry_after_stale_reuse());
|
||||
Ok(exchange)
|
||||
}
|
||||
@@ -480,16 +604,37 @@ pub(crate) async fn execute_response_create_request_with_pool(
|
||||
started_at,
|
||||
)
|
||||
.await;
|
||||
if result.is_err() {
|
||||
pool.discard(&key).await;
|
||||
match result {
|
||||
Ok(mut exchange) => {
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::new());
|
||||
Ok(exchange)
|
||||
}
|
||||
Err(error) if is_first_token_timeout(&error) => {
|
||||
pool.discard(&key).await;
|
||||
let mut exchange = execute_fresh_response_create_request_with_retries(
|
||||
request,
|
||||
started_at,
|
||||
pool.first_token_timeout(),
|
||||
1,
|
||||
)
|
||||
.await?;
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::new());
|
||||
Ok(exchange)
|
||||
}
|
||||
Err(error) => {
|
||||
pool.discard(&key).await;
|
||||
Err(error)
|
||||
}
|
||||
}
|
||||
result.map(|mut exchange| {
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::new());
|
||||
exchange
|
||||
})
|
||||
}
|
||||
WebSocketPoolAcquire::Bypass(reason) => {
|
||||
let mut exchange = execute_fresh_response_create_request(request, started_at).await?;
|
||||
let mut exchange = execute_fresh_response_create_request_with_retries(
|
||||
request,
|
||||
started_at,
|
||||
pool.first_token_timeout(),
|
||||
WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS,
|
||||
)
|
||||
.await?;
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::bypass(reason));
|
||||
Ok(exchange)
|
||||
}
|
||||
@@ -499,9 +644,15 @@ pub(crate) async fn execute_response_create_request_with_pool(
|
||||
pub(crate) async fn execute_response_create_request_stream_with_pool(
|
||||
request: &CodexWebSocketRequest,
|
||||
pool: Option<(&CodexWebSocketPool, CodexWebSocketPoolKey)>,
|
||||
fallback_first_token_timeout: Option<Duration>,
|
||||
) -> Result<CodexWebSocketStreamingExchange, CodexWebSocketExchangeError> {
|
||||
let Some((pool, key)) = pool else {
|
||||
return execute_fresh_response_create_request_stream(request).await;
|
||||
return execute_fresh_response_create_request_stream_with_retries(
|
||||
request,
|
||||
fallback_first_token_timeout,
|
||||
WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS,
|
||||
)
|
||||
.await;
|
||||
};
|
||||
|
||||
match pool.acquire(&key).await {
|
||||
@@ -516,9 +667,14 @@ pub(crate) async fn execute_response_create_request_stream_with_pool(
|
||||
match result {
|
||||
Err(CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame {
|
||||
..
|
||||
}) => {
|
||||
let mut exchange =
|
||||
execute_fresh_response_create_request_stream(request).await?;
|
||||
})
|
||||
| Err(CodexWebSocketExchangeError::FirstTokenTimeout { .. }) => {
|
||||
let mut exchange = execute_fresh_response_create_request_stream_with_retries(
|
||||
request,
|
||||
pool.first_token_timeout(),
|
||||
WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS,
|
||||
)
|
||||
.await?;
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::retry_after_stale_reuse());
|
||||
Ok(exchange)
|
||||
}
|
||||
@@ -536,16 +692,35 @@ pub(crate) async fn execute_response_create_request_stream_with_pool(
|
||||
key.clone(),
|
||||
)
|
||||
.await;
|
||||
if result.is_err() {
|
||||
pool.discard(&key).await;
|
||||
match result {
|
||||
Ok(mut exchange) => {
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::new());
|
||||
Ok(exchange)
|
||||
}
|
||||
Err(error) if is_first_token_timeout(&error) => {
|
||||
pool.discard(&key).await;
|
||||
let mut exchange = execute_fresh_response_create_request_stream_with_retries(
|
||||
request,
|
||||
pool.first_token_timeout(),
|
||||
1,
|
||||
)
|
||||
.await?;
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::new());
|
||||
Ok(exchange)
|
||||
}
|
||||
Err(error) => {
|
||||
pool.discard(&key).await;
|
||||
Err(error)
|
||||
}
|
||||
}
|
||||
result.map(|mut exchange| {
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::new());
|
||||
exchange
|
||||
})
|
||||
}
|
||||
WebSocketPoolAcquire::Bypass(reason) => {
|
||||
let mut exchange = execute_fresh_response_create_request_stream(request).await?;
|
||||
let mut exchange = execute_fresh_response_create_request_stream_with_retries(
|
||||
request,
|
||||
pool.first_token_timeout(),
|
||||
WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS,
|
||||
)
|
||||
.await?;
|
||||
exchange.pool_decision = Some(WebSocketPoolDecision::bypass(reason));
|
||||
Ok(exchange)
|
||||
}
|
||||
@@ -560,14 +735,22 @@ async fn execute_fresh_pooled_response_create_request(
|
||||
) -> Result<CodexWebSocketExchange, CodexWebSocketExchangeError> {
|
||||
let (websocket, response) =
|
||||
connect_pumped_websocket(request.connection(), pool.keepalive()).await?;
|
||||
let first_token_started_at = Instant::now();
|
||||
websocket
|
||||
.send(Message::Text(request.payload_text().to_string().into()))
|
||||
.await?;
|
||||
|
||||
let now = Instant::now();
|
||||
let metadata = websocket_connection_metadata(&response);
|
||||
let (exchange, websocket, metadata) =
|
||||
collect_websocket_response(websocket, metadata, false, started_at).await?;
|
||||
let (exchange, websocket, metadata) = collect_websocket_response(
|
||||
websocket,
|
||||
metadata,
|
||||
false,
|
||||
started_at,
|
||||
first_token_started_at,
|
||||
pool.first_token_timeout(),
|
||||
)
|
||||
.await?;
|
||||
pool.put(
|
||||
key,
|
||||
PooledWebSocketConnection {
|
||||
@@ -585,14 +768,29 @@ async fn execute_fresh_pooled_response_create_request_stream(
|
||||
pool: CodexWebSocketPool,
|
||||
key: CodexWebSocketPoolKey,
|
||||
) -> Result<CodexWebSocketStreamingExchange, CodexWebSocketExchangeError> {
|
||||
let (websocket, response) =
|
||||
let (mut websocket, response) =
|
||||
connect_pumped_websocket(request.connection(), pool.keepalive()).await?;
|
||||
let first_token_started_at = Instant::now();
|
||||
websocket
|
||||
.send(Message::Text(request.payload_text().to_string().into()))
|
||||
.await?;
|
||||
|
||||
let now = Instant::now();
|
||||
let metadata = websocket_connection_metadata(&response);
|
||||
let mut metadata = websocket_connection_metadata(&response);
|
||||
let prefetched_frames = match prefetch_first_content_stream_frames(
|
||||
&mut websocket,
|
||||
&mut metadata,
|
||||
first_token_started_at,
|
||||
pool.first_token_timeout(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(prefetched_frames) => prefetched_frames,
|
||||
Err(error) => {
|
||||
websocket.close().await;
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
Ok(stream_websocket_response(
|
||||
websocket,
|
||||
metadata,
|
||||
@@ -601,7 +799,7 @@ async fn execute_fresh_pooled_response_create_request_stream(
|
||||
key,
|
||||
created_at: now,
|
||||
}),
|
||||
None,
|
||||
prefetched_frames,
|
||||
))
|
||||
}
|
||||
|
||||
@@ -612,6 +810,7 @@ async fn execute_pooled_response_create_request(
|
||||
connection: PooledWebSocketConnection,
|
||||
started_at: Instant,
|
||||
) -> Result<CodexWebSocketExchange, CodexWebSocketExchangeError> {
|
||||
let first_token_started_at = Instant::now();
|
||||
if let Err(error) = connection
|
||||
.websocket
|
||||
.send(Message::Text(request.payload_text().to_string().into()))
|
||||
@@ -626,8 +825,15 @@ async fn execute_pooled_response_create_request(
|
||||
}
|
||||
|
||||
let created_at = connection.created_at;
|
||||
match collect_websocket_response(connection.websocket, connection.metadata, true, started_at)
|
||||
.await
|
||||
match collect_websocket_response(
|
||||
connection.websocket,
|
||||
connection.metadata,
|
||||
true,
|
||||
started_at,
|
||||
first_token_started_at,
|
||||
pool.first_token_timeout(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok((exchange, websocket, metadata)) => {
|
||||
pool.put(
|
||||
@@ -652,10 +858,13 @@ async fn execute_pooled_response_create_request_stream(
|
||||
request: &CodexWebSocketRequest,
|
||||
pool: CodexWebSocketPool,
|
||||
key: CodexWebSocketPoolKey,
|
||||
mut connection: PooledWebSocketConnection,
|
||||
connection: PooledWebSocketConnection,
|
||||
) -> Result<CodexWebSocketStreamingExchange, CodexWebSocketExchangeError> {
|
||||
if let Err(error) = connection
|
||||
.websocket
|
||||
let mut websocket = connection.websocket;
|
||||
let mut metadata = connection.metadata;
|
||||
let created_at = connection.created_at;
|
||||
let first_token_started_at = Instant::now();
|
||||
if let Err(error) = websocket
|
||||
.send(Message::Text(request.payload_text().to_string().into()))
|
||||
.await
|
||||
{
|
||||
@@ -667,51 +876,34 @@ async fn execute_pooled_response_create_request_stream(
|
||||
);
|
||||
}
|
||||
|
||||
let prefetched_text = first_reused_stream_text(&mut connection.websocket).await?;
|
||||
let prefetched_frames = match prefetch_first_content_stream_frames(
|
||||
&mut websocket,
|
||||
&mut metadata,
|
||||
first_token_started_at,
|
||||
pool.first_token_timeout(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(prefetched_frames) => prefetched_frames,
|
||||
Err(error) => {
|
||||
pool.discard(&key).await;
|
||||
websocket.close().await;
|
||||
return Err(reused_stream_prefetch_error(error));
|
||||
}
|
||||
};
|
||||
|
||||
Ok(stream_websocket_response(
|
||||
connection.websocket,
|
||||
connection.metadata,
|
||||
websocket,
|
||||
metadata,
|
||||
Some(WebSocketStreamPoolReturn {
|
||||
pool,
|
||||
key,
|
||||
created_at: connection.created_at,
|
||||
created_at,
|
||||
}),
|
||||
Some(prefetched_text),
|
||||
prefetched_frames,
|
||||
))
|
||||
}
|
||||
|
||||
async fn first_reused_stream_text(
|
||||
websocket: &mut PumpedWebSocket,
|
||||
) -> Result<String, CodexWebSocketExchangeError> {
|
||||
loop {
|
||||
let message = match next_websocket_message(websocket, WEBSOCKET_RECEIVE_IDLE_TIMEOUT).await
|
||||
{
|
||||
Ok(Some(message)) => message,
|
||||
Ok(None) => {
|
||||
return Err(
|
||||
CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame {
|
||||
message: "websocket closed".to_string(),
|
||||
},
|
||||
);
|
||||
}
|
||||
Err(error) => return Err(reused_connection_died_before_first_frame(&error)),
|
||||
};
|
||||
match message {
|
||||
Message::Text(text) => return Ok(text.to_string()),
|
||||
Message::Close(_) => {
|
||||
return Err(
|
||||
CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame {
|
||||
message: "websocket closed".to_string(),
|
||||
},
|
||||
);
|
||||
}
|
||||
Message::Binary(_) => return Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Public entry points
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -720,7 +912,7 @@ async fn first_reused_stream_text(
|
||||
pub async fn execute_response_create_request(
|
||||
request: &CodexWebSocketRequest,
|
||||
) -> Result<CodexWebSocketExchange, CodexWebSocketExchangeError> {
|
||||
execute_response_create_request_with_pool(request, None, Instant::now()).await
|
||||
execute_response_create_request_with_pool(request, None, Instant::now(), None).await
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -732,6 +924,8 @@ async fn collect_websocket_response(
|
||||
mut metadata: CodexWebSocketConnectionMetadata,
|
||||
reused_connection: bool,
|
||||
started_at: Instant,
|
||||
first_token_started_at: Instant,
|
||||
first_token_timeout: Option<Duration>,
|
||||
) -> Result<
|
||||
(
|
||||
CodexWebSocketExchange,
|
||||
@@ -741,27 +935,37 @@ async fn collect_websocket_response(
|
||||
CodexWebSocketExchangeError,
|
||||
> {
|
||||
let mut body = String::new();
|
||||
let mut saw_response_frame = false;
|
||||
let mut saw_content_frame = false;
|
||||
let mut first_token_ms = None;
|
||||
|
||||
loop {
|
||||
let message =
|
||||
match next_websocket_message(&mut websocket, receive_idle_timeout(saw_response_frame))
|
||||
.await
|
||||
{
|
||||
Ok(message) => message,
|
||||
Err(error) if reused_connection && !saw_response_frame => {
|
||||
let receive_timeout = receive_idle_timeout_with_first_token_deadline(
|
||||
saw_content_frame,
|
||||
first_token_started_at,
|
||||
first_token_timeout,
|
||||
)?;
|
||||
let message = match next_websocket_message(&mut websocket, receive_timeout).await {
|
||||
Ok(message) => message,
|
||||
Err(error) if !saw_content_frame => {
|
||||
if let Some(timeout) =
|
||||
elapsed_first_token_timeout(first_token_started_at, first_token_timeout)
|
||||
{
|
||||
return Err(CodexWebSocketExchangeError::FirstTokenTimeout { timeout });
|
||||
}
|
||||
if reused_connection {
|
||||
return Err(reused_connection_died_before_first_frame(&error));
|
||||
}
|
||||
Err(error) => return Err(error),
|
||||
};
|
||||
return Err(error);
|
||||
}
|
||||
Err(error) => return Err(error),
|
||||
};
|
||||
let Some(message) = message else {
|
||||
break;
|
||||
};
|
||||
let text = match message {
|
||||
Message::Text(text) => text,
|
||||
Message::Binary(_) => return Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent),
|
||||
Message::Close(_) if reused_connection && !saw_response_frame => {
|
||||
Message::Close(_) if reused_connection && !saw_content_frame => {
|
||||
return Err(
|
||||
CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame {
|
||||
message: "websocket closed".to_string(),
|
||||
@@ -775,12 +979,13 @@ async fn collect_websocket_response(
|
||||
if let Some(classified) = classify_websocket_error_frame(&raw) {
|
||||
let retry_after_seconds = retry_after_seconds_from_wrapped_error_headers(&raw)
|
||||
.or_else(|| events::retry_after_seconds_from_body(&raw));
|
||||
return Err(CodexWebSocketExchangeError::Upstream {
|
||||
status_code: classified.status_code,
|
||||
return Err(CodexWebSocketExchangeError::upstream(
|
||||
classified.status_code,
|
||||
retry_after_seconds,
|
||||
body: raw,
|
||||
set_cookie_headers: Vec::new(),
|
||||
});
|
||||
raw,
|
||||
Vec::new(),
|
||||
CodexUpstreamDiagnostics::with_status(classified.status_code),
|
||||
));
|
||||
}
|
||||
if let Some(reason) = websocket_incomplete_response_reason(&raw) {
|
||||
return Err(CodexWebSocketExchangeError::IncompleteResponse { reason });
|
||||
@@ -799,8 +1004,14 @@ async fn collect_websocket_response(
|
||||
let event = websocket_event_type(&raw);
|
||||
let forwarded = if let Some(frame) = websocket_event_to_sse_frame(&raw) {
|
||||
body.push_str(&frame);
|
||||
saw_response_frame = true;
|
||||
response_meta::update_first_token_ms(started_at, body.as_bytes(), &mut first_token_ms);
|
||||
if event_is_first_content(event.as_deref()) {
|
||||
saw_content_frame = true;
|
||||
response_meta::update_first_token_ms(
|
||||
started_at,
|
||||
body.as_bytes(),
|
||||
&mut first_token_ms,
|
||||
);
|
||||
}
|
||||
true
|
||||
} else {
|
||||
false
|
||||
@@ -815,12 +1026,13 @@ async fn collect_websocket_response(
|
||||
rate_limit_headers: metadata.rate_limit_headers.clone(),
|
||||
first_token_ms,
|
||||
pool_decision: None,
|
||||
diagnostics: metadata.diagnostics.clone(),
|
||||
};
|
||||
return Ok((exchange, websocket, metadata));
|
||||
}
|
||||
}
|
||||
|
||||
if reused_connection && !saw_response_frame {
|
||||
if reused_connection && !saw_content_frame {
|
||||
return Err(
|
||||
CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame {
|
||||
message: "websocket closed before terminal event".to_string(),
|
||||
@@ -831,6 +1043,85 @@ async fn collect_websocket_response(
|
||||
Err(CodexWebSocketExchangeError::ClosedBeforeTerminal)
|
||||
}
|
||||
|
||||
async fn prefetch_first_content_stream_frames(
|
||||
websocket: &mut PumpedWebSocket,
|
||||
metadata: &mut CodexWebSocketConnectionMetadata,
|
||||
first_token_started_at: Instant,
|
||||
first_token_timeout: Option<Duration>,
|
||||
) -> Result<Vec<String>, CodexWebSocketExchangeError> {
|
||||
let mut prefetched_frames = Vec::new();
|
||||
loop {
|
||||
let receive_timeout =
|
||||
receive_timeout_before_first_content(first_token_started_at, first_token_timeout)?;
|
||||
let message = match next_websocket_message(websocket, receive_timeout).await {
|
||||
Ok(Some(message)) => message,
|
||||
Ok(None) => return Err(CodexWebSocketExchangeError::ClosedBeforeTerminal),
|
||||
Err(error) => {
|
||||
if let Some(timeout) =
|
||||
elapsed_first_token_timeout(first_token_started_at, first_token_timeout)
|
||||
{
|
||||
return Err(CodexWebSocketExchangeError::FirstTokenTimeout { timeout });
|
||||
}
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
|
||||
let raw = match message {
|
||||
Message::Text(text) => text.to_string(),
|
||||
Message::Binary(_) => return Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent),
|
||||
Message::Close(_) => return Err(CodexWebSocketExchangeError::ClosedBeforeTerminal),
|
||||
_ => continue,
|
||||
};
|
||||
|
||||
if let Some(classified) = classify_websocket_error_frame(&raw) {
|
||||
let retry_after_seconds = retry_after_seconds_from_wrapped_error_headers(&raw)
|
||||
.or_else(|| events::retry_after_seconds_from_body(&raw));
|
||||
return Err(CodexWebSocketExchangeError::upstream(
|
||||
classified.status_code,
|
||||
retry_after_seconds,
|
||||
raw,
|
||||
Vec::new(),
|
||||
CodexUpstreamDiagnostics::with_status(classified.status_code),
|
||||
));
|
||||
}
|
||||
if let Some(reason) = websocket_incomplete_response_reason(&raw) {
|
||||
return Err(CodexWebSocketExchangeError::IncompleteResponse { reason });
|
||||
}
|
||||
if let Some(message) = websocket_response_completed_parse_error(&raw) {
|
||||
return Err(CodexWebSocketExchangeError::InvalidCompletedResponse { message });
|
||||
}
|
||||
if let Some(headers) = websocket_rate_limit_event_headers(&raw) {
|
||||
metadata.rate_limit_headers.extend(headers);
|
||||
continue;
|
||||
}
|
||||
if let Some(metadata_turn_state) = websocket_metadata_turn_state(&raw) {
|
||||
metadata.turn_state = Some(metadata_turn_state);
|
||||
continue;
|
||||
}
|
||||
|
||||
let event = websocket_event_type(&raw);
|
||||
if websocket_event_to_sse_frame(&raw).is_none() {
|
||||
continue;
|
||||
}
|
||||
let is_first_content = event_is_first_content(event.as_deref());
|
||||
prefetched_frames.push(raw);
|
||||
if is_first_content {
|
||||
return Ok(prefetched_frames);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn reused_stream_prefetch_error(error: CodexWebSocketExchangeError) -> CodexWebSocketExchangeError {
|
||||
match error {
|
||||
CodexWebSocketExchangeError::ClosedBeforeTerminal
|
||||
| CodexWebSocketExchangeError::ReceiveIdleTimeout { .. }
|
||||
| CodexWebSocketExchangeError::Transport(_) => {
|
||||
reused_connection_died_before_first_frame(&error)
|
||||
}
|
||||
error => error,
|
||||
}
|
||||
}
|
||||
|
||||
fn reused_connection_died_before_first_frame(
|
||||
error: &CodexWebSocketExchangeError,
|
||||
) -> CodexWebSocketExchangeError {
|
||||
@@ -846,6 +1137,10 @@ fn websocket_connection_metadata(
|
||||
turn_state: response_meta::turn_state(response.headers()),
|
||||
set_cookie_headers: response_meta::set_cookie_headers(response.headers()),
|
||||
rate_limit_headers: response_meta::rate_limit_headers(response.headers()),
|
||||
diagnostics: response_meta::diagnostics(
|
||||
Some(response.status().as_u16()),
|
||||
response.headers(),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -863,7 +1158,7 @@ fn stream_websocket_response(
|
||||
websocket: PumpedWebSocket,
|
||||
metadata: CodexWebSocketConnectionMetadata,
|
||||
pool_return: Option<WebSocketStreamPoolReturn>,
|
||||
prefetched_text: Option<String>,
|
||||
prefetched_frames: Vec<String>,
|
||||
) -> CodexWebSocketStreamingExchange {
|
||||
let response_metadata = metadata.clone();
|
||||
let rate_limit_header_updates = Arc::new(Mutex::new(Vec::new()));
|
||||
@@ -876,7 +1171,7 @@ fn stream_websocket_response(
|
||||
websocket,
|
||||
metadata,
|
||||
pool_return,
|
||||
prefetched_text,
|
||||
prefetched_frames,
|
||||
rate_limit_header_updates_for_task,
|
||||
turn_state_update_for_task,
|
||||
tx,
|
||||
@@ -896,6 +1191,7 @@ fn stream_websocket_response(
|
||||
rate_limit_header_updates,
|
||||
turn_state_update,
|
||||
pool_decision: None,
|
||||
diagnostics: response_metadata.diagnostics,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -903,16 +1199,16 @@ async fn forward_websocket_response_stream(
|
||||
mut websocket: PumpedWebSocket,
|
||||
mut metadata: CodexWebSocketConnectionMetadata,
|
||||
pool_return: Option<WebSocketStreamPoolReturn>,
|
||||
prefetched_text: Option<String>,
|
||||
prefetched_frames: Vec<String>,
|
||||
rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates,
|
||||
turn_state_update: CodexWebSocketTurnStateUpdate,
|
||||
tx: mpsc::Sender<Result<Bytes, CodexWebSocketExchangeError>>,
|
||||
) {
|
||||
let mut pool_return = pool_return;
|
||||
let mut prefetched_text = prefetched_text;
|
||||
let mut saw_response_frame = false;
|
||||
let mut prefetched_frames: VecDeque<String> = prefetched_frames.into();
|
||||
let mut saw_content_frame = false;
|
||||
loop {
|
||||
let raw = if let Some(text) = prefetched_text.take() {
|
||||
let raw = if let Some(text) = prefetched_frames.pop_front() {
|
||||
text
|
||||
} else {
|
||||
let message = tokio::select! {
|
||||
@@ -925,7 +1221,7 @@ async fn forward_websocket_response_stream(
|
||||
}
|
||||
message = next_websocket_message(
|
||||
&mut websocket,
|
||||
receive_idle_timeout(saw_response_frame),
|
||||
receive_idle_timeout(saw_content_frame),
|
||||
) => message,
|
||||
};
|
||||
let message = match message {
|
||||
@@ -963,12 +1259,13 @@ async fn forward_websocket_response_stream(
|
||||
let retry_after_seconds = retry_after_seconds_from_wrapped_error_headers(&raw)
|
||||
.or_else(|| events::retry_after_seconds_from_body(&raw));
|
||||
let _ = tx
|
||||
.send(Err(CodexWebSocketExchangeError::Upstream {
|
||||
status_code: classified.status_code,
|
||||
.send(Err(CodexWebSocketExchangeError::upstream(
|
||||
classified.status_code,
|
||||
retry_after_seconds,
|
||||
body: raw,
|
||||
set_cookie_headers: Vec::new(),
|
||||
}))
|
||||
raw,
|
||||
Vec::new(),
|
||||
CodexUpstreamDiagnostics::with_status(classified.status_code),
|
||||
)))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
@@ -1009,7 +1306,9 @@ async fn forward_websocket_response_stream(
|
||||
discard_stream_websocket(websocket, pool_return).await;
|
||||
return;
|
||||
}
|
||||
saw_response_frame = true;
|
||||
if event_is_first_content(event.as_deref()) {
|
||||
saw_content_frame = true;
|
||||
}
|
||||
if terminal {
|
||||
finish_stream_websocket(websocket, metadata, pool_return.take()).await;
|
||||
return;
|
||||
@@ -1030,6 +1329,56 @@ fn receive_idle_timeout(saw_response_frame: bool) -> Duration {
|
||||
}
|
||||
}
|
||||
|
||||
fn receive_idle_timeout_with_first_token_deadline(
|
||||
saw_content_frame: bool,
|
||||
first_token_started_at: Instant,
|
||||
first_token_timeout: Option<Duration>,
|
||||
) -> Result<Duration, CodexWebSocketExchangeError> {
|
||||
if saw_content_frame {
|
||||
return Ok(WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT);
|
||||
}
|
||||
receive_timeout_before_first_content(first_token_started_at, first_token_timeout)
|
||||
}
|
||||
|
||||
fn receive_timeout_before_first_content(
|
||||
first_token_started_at: Instant,
|
||||
first_token_timeout: Option<Duration>,
|
||||
) -> Result<Duration, CodexWebSocketExchangeError> {
|
||||
let Some(first_token_timeout) = first_token_timeout.filter(|timeout| !timeout.is_zero()) else {
|
||||
return Ok(WEBSOCKET_RECEIVE_IDLE_TIMEOUT);
|
||||
};
|
||||
let Some(remaining) = first_token_timeout.checked_sub(first_token_started_at.elapsed()) else {
|
||||
return Err(CodexWebSocketExchangeError::FirstTokenTimeout {
|
||||
timeout: first_token_timeout,
|
||||
});
|
||||
};
|
||||
if remaining.is_zero() {
|
||||
return Err(CodexWebSocketExchangeError::FirstTokenTimeout {
|
||||
timeout: first_token_timeout,
|
||||
});
|
||||
}
|
||||
Ok(WEBSOCKET_RECEIVE_IDLE_TIMEOUT.min(remaining))
|
||||
}
|
||||
|
||||
fn elapsed_first_token_timeout(
|
||||
first_token_started_at: Instant,
|
||||
first_token_timeout: Option<Duration>,
|
||||
) -> Option<Duration> {
|
||||
let first_token_timeout = first_token_timeout.filter(|timeout| !timeout.is_zero())?;
|
||||
(first_token_started_at.elapsed() >= first_token_timeout).then_some(first_token_timeout)
|
||||
}
|
||||
|
||||
fn event_is_first_content(event: Option<&str>) -> bool {
|
||||
match event {
|
||||
Some(event) => !is_pre_content_lifecycle_event(event),
|
||||
None => true,
|
||||
}
|
||||
}
|
||||
|
||||
fn is_first_token_timeout(error: &CodexWebSocketExchangeError) -> bool {
|
||||
matches!(error, CodexWebSocketExchangeError::FirstTokenTimeout { .. })
|
||||
}
|
||||
|
||||
async fn next_websocket_message(
|
||||
websocket: &mut PumpedWebSocket,
|
||||
receive_timeout: Duration,
|
||||
|
||||
@@ -17,6 +17,7 @@ const DEFAULT_MAX_AGE: Duration = Duration::from_mins(55);
|
||||
const DEFAULT_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(25);
|
||||
const DEFAULT_PING_INTERVAL: Duration = Duration::from_secs(25);
|
||||
const DEFAULT_PING_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
pub(crate) const DEFAULT_FIRST_TOKEN_TIMEOUT: Duration = Duration::from_secs(20);
|
||||
|
||||
/// WebSocket 连接池 key。
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
@@ -87,6 +88,8 @@ pub struct CodexWebSocketPoolConfig {
|
||||
pub ping_timeout: Duration,
|
||||
/// idle socket 无活动多久后视为失活。
|
||||
pub liveness_timeout: Option<Duration>,
|
||||
/// 首个内容帧到达前的绝对超时;`None` 表示禁用首 token 熔断。
|
||||
pub first_token_timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
impl Default for CodexWebSocketPoolConfig {
|
||||
@@ -101,6 +104,7 @@ impl Default for CodexWebSocketPoolConfig {
|
||||
// idle 连接不设失活截断:靠 ping/pong 保活,只在 max_age(55 分钟)
|
||||
// 或 ping 失败时关闭,最大化跨轮复用(对齐 Codex CLI 的长连接策略)。
|
||||
liveness_timeout: None,
|
||||
first_token_timeout: Some(DEFAULT_FIRST_TOKEN_TIMEOUT),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -143,6 +147,11 @@ impl CodexWebSocketPool {
|
||||
self.config.keepalive()
|
||||
}
|
||||
|
||||
/// 首个内容帧到达前的绝对超时;`None` 表示禁用首 token 熔断。
|
||||
pub(crate) fn first_token_timeout(&self) -> Option<Duration> {
|
||||
self.config.first_token_timeout
|
||||
}
|
||||
|
||||
pub(crate) async fn acquire(&self, key: &CodexWebSocketPoolKey) -> WebSocketPoolAcquire {
|
||||
let mut expired_connection = None;
|
||||
let acquire = {
|
||||
@@ -343,6 +352,7 @@ pub(crate) struct CodexWebSocketConnectionMetadata {
|
||||
pub(crate) turn_state: Option<String>,
|
||||
pub(crate) set_cookie_headers: Vec<String>,
|
||||
pub(crate) rate_limit_headers: Vec<(String, String)>,
|
||||
pub(crate) diagnostics: super::diagnostics::CodexUpstreamDiagnostics,
|
||||
}
|
||||
|
||||
pub(crate) struct PooledWebSocketConnection {
|
||||
|
||||
@@ -30,7 +30,11 @@ async fn admin_accounts_import_should_store_cpr_account_tokens() {
|
||||
usage_records: SqliteUsageRecordStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let services = Arc::new(Services::new(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
let state = AppState {
|
||||
services: (*services).clone(),
|
||||
};
|
||||
@@ -639,7 +643,11 @@ async fn admin_accounts_import_should_fetch_usage_to_complete_missing_plan_and_q
|
||||
usage_records: SqliteUsageRecordStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let services = Arc::new(Services::new(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
let state = AppState {
|
||||
services: (*services).clone(),
|
||||
};
|
||||
@@ -855,7 +863,11 @@ async fn admin_accounts_import_should_require_admin_session_cookie() {
|
||||
usage_records: SqliteUsageRecordStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let services = Arc::new(Services::new(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
let state = AppState {
|
||||
services: (*services).clone(),
|
||||
};
|
||||
|
||||
@@ -23,7 +23,11 @@ async fn admin_accounts_lifecycle_should_update_and_delete_accounts() {
|
||||
usage_records: SqliteUsageRecordStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let services = Arc::new(Services::new(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
let state = AppState {
|
||||
services: (*services).clone(),
|
||||
};
|
||||
@@ -118,16 +122,6 @@ async fn seed_account_related_rows(pool: &SqlitePool, account_id: &str) {
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"insert into model_account_routes (model, account_id, priority, created_at, updated_at) values (?, ?, 0, ?, ?)",
|
||||
)
|
||||
.bind("gpt-5.5")
|
||||
.bind(account_id)
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query("insert into account_cookies (id, account_id, domain, name, value, updated_at) values (?, ?, ?, ?, ?, ?)")
|
||||
.bind("cookie_lifecycle")
|
||||
.bind(account_id)
|
||||
@@ -180,12 +174,6 @@ async fn assert_account_related_rows_deleted(pool: &SqlitePool, account_id: &str
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let model_routes: i64 =
|
||||
sqlx::query_scalar("select count(*) from model_account_routes where account_id = ?")
|
||||
.bind(account_id)
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let leases: i64 =
|
||||
sqlx::query_scalar("select count(*) from account_refresh_leases where account_id = ?")
|
||||
.bind(account_id)
|
||||
@@ -202,7 +190,6 @@ async fn assert_account_related_rows_deleted(pool: &SqlitePool, account_id: &str
|
||||
assert_eq!(accounts, 0);
|
||||
assert_eq!(usage, 0);
|
||||
assert_eq!(model_usage, 0);
|
||||
assert_eq!(model_routes, 0);
|
||||
assert_eq!(cookies, 0);
|
||||
assert_eq!(leases, 0);
|
||||
assert_eq!(affinities, 0);
|
||||
@@ -765,7 +752,11 @@ async fn admin_account_create_should_derive_claims_and_store_plain_tokens() {
|
||||
usage_records: SqliteUsageRecordStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let services = Arc::new(Services::new(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
let state = AppState {
|
||||
services: (*services).clone(),
|
||||
};
|
||||
@@ -808,7 +799,11 @@ async fn admin_account_manual_create_should_reject_missing_invalid_expired_or_un
|
||||
usage_records: SqliteUsageRecordStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let services = Arc::new(Services::new(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
let state = AppState {
|
||||
services: (*services).clone(),
|
||||
};
|
||||
|
||||
@@ -27,7 +27,11 @@ async fn admin_accounts_list_should_not_expose_account_tokens() {
|
||||
usage_records: SqliteUsageRecordStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let services = Arc::new(Services::new(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
let state = AppState {
|
||||
services: (*services).clone(),
|
||||
};
|
||||
@@ -259,11 +263,11 @@ async fn admin_accounts_list_should_include_usage_quota_and_model_stats() {
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = response_json(response).await;
|
||||
let item = &body["data"]["items"][0];
|
||||
assert_eq!(item["usage"]["requestCount"], 7);
|
||||
assert_eq!(item["usage"]["inputTokens"], 120_000);
|
||||
assert_eq!(item["usage"]["outputTokens"], 4_000);
|
||||
assert_eq!(item["usage"]["cachedTokens"], 40_000);
|
||||
assert_eq!(item["usage"]["totalTokens"], 124_000);
|
||||
assert_eq!(item["usage"]["requestCount"], 107);
|
||||
assert_eq!(item["usage"]["inputTokens"], 1_022_000);
|
||||
assert_eq!(item["usage"]["outputTokens"], 94_500);
|
||||
assert_eq!(item["usage"]["cachedTokens"], 49_100);
|
||||
assert_eq!(item["usage"]["totalTokens"], 1_116_500);
|
||||
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");
|
||||
@@ -277,12 +281,128 @@ async fn admin_accounts_list_should_include_usage_quota_and_model_stats() {
|
||||
assert!(item["quota"]["windows"][2]["windowUsedDisplay"]
|
||||
.as_str()
|
||||
.is_some_and(|value| value.contains(" / 7.0d")));
|
||||
assert_eq!(item["usage"]["createdTokens"], 80_000);
|
||||
assert_eq!(item["usage"]["readTokens"], 40_000);
|
||||
assert_eq!(item["usage"]["models"].as_array().unwrap().len(), 2);
|
||||
assert_eq!(item["usage"]["models"][0]["model"], "gpt-5.5");
|
||||
assert_eq!(item["usage"]["models"][0]["requestCount"], 4);
|
||||
assert_eq!(item["usage"]["models"][1]["model"], "gpt-5");
|
||||
assert_eq!(item["usage"]["createdTokens"], 972_900);
|
||||
assert_eq!(item["usage"]["readTokens"], 49_100);
|
||||
assert_eq!(item["usage"]["models"].as_array().unwrap().len(), 4);
|
||||
assert_eq!(item["usage"]["models"][0]["model"], "gpt-outside-window");
|
||||
assert_eq!(item["usage"]["models"][0]["requestCount"], 99);
|
||||
assert_eq!(item["usage"]["models"][1]["model"], "gpt-5.5");
|
||||
assert_eq!(item["usage"]["models"][2]["model"], "gpt-5");
|
||||
assert_eq!(item["usage"]["models"][3]["model"], "gpt-five-hour-window");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_accounts_list_should_use_largest_quota_reset_interval_for_usage_stats() {
|
||||
let (app, _state, pool, _dir) =
|
||||
admin_accounts_test_app("admin-accounts-largest-quota-window.sqlite", 71).await;
|
||||
seed_usage_account(
|
||||
&pool,
|
||||
UsageAccountSeed {
|
||||
id: "acct_custom_window",
|
||||
email: "custom-window@example.com",
|
||||
label: "custom",
|
||||
plan_type: "plus",
|
||||
request_count: 10,
|
||||
empty_response_count: 0,
|
||||
input_tokens: 500,
|
||||
output_tokens: 100,
|
||||
cached_tokens: 50,
|
||||
window_request_count: 1,
|
||||
window_input_tokens: 100,
|
||||
window_output_tokens: 10,
|
||||
window_cached_tokens: 5,
|
||||
window_started_at: "2026-07-13T19:00:00Z",
|
||||
window_reset_at: "2026-07-14T00:00:00Z",
|
||||
limit_window_seconds: 18_000,
|
||||
last_used_at: "2026-06-30T01:10:00Z",
|
||||
},
|
||||
)
|
||||
.await;
|
||||
sqlx::query("update accounts set quota_json = ?, quota_fetched_at = ? where id = ?")
|
||||
.bind(
|
||||
json!({
|
||||
"plan_type": "plus",
|
||||
"snapshots": [{
|
||||
"source": "custom",
|
||||
"limit_name": "Codex",
|
||||
"metered_feature": "codex",
|
||||
"allowed": true,
|
||||
"limit_reached": false,
|
||||
"blocked": false,
|
||||
"primary": {
|
||||
"used_percent": 50,
|
||||
"remaining_percent": 50,
|
||||
"reset_at": 1783987200,
|
||||
"window_minutes": 20160,
|
||||
"limit_reached": false
|
||||
},
|
||||
"secondary": {
|
||||
"used_percent": 10,
|
||||
"remaining_percent": 90,
|
||||
"reset_at": 1783987200,
|
||||
"window_minutes": 10080,
|
||||
"limit_reached": false
|
||||
}
|
||||
}]
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.bind("2026-06-30T01:10:00Z")
|
||||
.bind("acct_custom_window")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"insert into usage_time_buckets (bucket_start, account_id, model, request_count, input_tokens, output_tokens, cached_tokens, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind("2026-06-30T01:00:00+00:00")
|
||||
.bind("acct_custom_window")
|
||||
.bind("gpt-custom-14d")
|
||||
.bind(5)
|
||||
.bind(10_000)
|
||||
.bind(2_000)
|
||||
.bind(1_000)
|
||||
.bind("2026-06-30T01:10:00Z")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"insert into usage_time_buckets (bucket_start, account_id, model, request_count, input_tokens, output_tokens, cached_tokens, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind("2026-06-29T23:45:00+00:00")
|
||||
.bind("acct_custom_window")
|
||||
.bind("gpt-before-14d")
|
||||
.bind(99)
|
||||
.bind(900_000)
|
||||
.bind(90_000)
|
||||
.bind(9_000)
|
||||
.bind("2026-06-29T23:59: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_largest_quota_window")
|
||||
.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["usage"]["requestCount"], 5);
|
||||
assert_eq!(item["usage"]["inputTokens"], 10_000);
|
||||
assert_eq!(item["usage"]["outputTokens"], 2_000);
|
||||
assert_eq!(item["usage"]["cachedTokens"], 1_000);
|
||||
assert_eq!(item["usage"]["models"].as_array().unwrap().len(), 1);
|
||||
assert_eq!(item["usage"]["models"][0]["model"], "gpt-custom-14d");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -26,7 +26,10 @@ use sqlx::SqlitePool;
|
||||
use tower::util::ServiceExt;
|
||||
|
||||
use crate::support::jwt::unsigned_jwt;
|
||||
use crate::support::{admin::seed_admin_session, config::test_config, http::response_json};
|
||||
use crate::support::{
|
||||
admin::seed_admin_session, config::test_config, fingerprint::runtime_fingerprint,
|
||||
http::response_json,
|
||||
};
|
||||
|
||||
mod exporting;
|
||||
mod importing;
|
||||
@@ -170,7 +173,11 @@ async fn admin_accounts_test_app_with_overrides(
|
||||
usage_records: SqliteUsageRecordStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let services = Arc::new(Services::new(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
let state = AppState {
|
||||
services: (*services).clone(),
|
||||
};
|
||||
|
||||
@@ -46,11 +46,38 @@ async fn admin_account_quota_should_send_usage_cookie() {
|
||||
access_token: SecretString::new("access-quota-cookie".to_string().into()),
|
||||
refresh_token: None,
|
||||
access_token_expires_at: None,
|
||||
status: AccountStatus::Active,
|
||||
status: AccountStatus::QuotaExhausted,
|
||||
added_at: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
sqlx::query(
|
||||
"insert into account_usage (
|
||||
account_id,
|
||||
request_count,
|
||||
window_request_count,
|
||||
window_input_tokens,
|
||||
window_output_tokens,
|
||||
window_cached_tokens,
|
||||
window_started_at,
|
||||
window_reset_at,
|
||||
limit_window_seconds,
|
||||
last_used_at
|
||||
) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind("acct_quota_cookie")
|
||||
.bind(9)
|
||||
.bind(9)
|
||||
.bind(9_000)
|
||||
.bind(900)
|
||||
.bind(90)
|
||||
.bind("2027-01-01T00:00:00Z")
|
||||
.bind("2027-01-01T05:00:00Z")
|
||||
.bind(18_000)
|
||||
.bind("2027-01-15T04:00:00Z")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"insert into usage_time_buckets (
|
||||
bucket_start,
|
||||
@@ -74,6 +101,29 @@ async fn admin_account_quota_should_send_usage_cookie() {
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"insert into usage_time_buckets (
|
||||
bucket_start,
|
||||
account_id,
|
||||
model,
|
||||
request_count,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cached_tokens,
|
||||
updated_at
|
||||
) values (?, ?, ?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind("2027-01-15T09:00:00+00:00")
|
||||
.bind("acct_quota_cookie")
|
||||
.bind("gpt-5")
|
||||
.bind(1)
|
||||
.bind(1_500)
|
||||
.bind(500)
|
||||
.bind(300)
|
||||
.bind("2027-01-15T09:00:00Z")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
SqliteCookieStore::new(pool.clone())
|
||||
.capture_set_cookie(
|
||||
"acct_quota_cookie",
|
||||
@@ -107,6 +157,15 @@ async fn admin_account_quota_should_send_usage_cookie() {
|
||||
);
|
||||
let body = response_json(response).await;
|
||||
assert_eq!(body["data"]["planType"], "plus");
|
||||
assert_eq!(body["data"]["account"]["status"], "active");
|
||||
assert_eq!(body["data"]["account"]["usage"]["requestCount"], 1);
|
||||
assert_eq!(body["data"]["account"]["usage"]["inputTokens"], 1_500);
|
||||
assert_eq!(body["data"]["account"]["usage"]["outputTokens"], 500);
|
||||
assert_eq!(body["data"]["account"]["usage"]["cachedTokens"], 300);
|
||||
assert_eq!(
|
||||
body["data"]["account"]["usage"]["models"][0]["model"],
|
||||
"gpt-5"
|
||||
);
|
||||
assert_eq!(
|
||||
body["data"]["quotaData"]["windows"][0]["labelDisplay"],
|
||||
"5小时限额"
|
||||
@@ -125,6 +184,6 @@ async fn admin_account_quota_should_send_usage_cookie() {
|
||||
);
|
||||
assert_eq!(
|
||||
body["data"]["quotaData"]["windows"][1]["tokenUsageDisplay"],
|
||||
"0"
|
||||
"2K"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ use codex_proxy_rs::infra::{database::connect_sqlite, identity::hash_admin_passw
|
||||
use sqlx::SqlitePool;
|
||||
use tower::util::ServiceExt;
|
||||
|
||||
use crate::support::{config::test_config, http::response_json};
|
||||
use crate::support::{config::test_config, fingerprint::runtime_fingerprint, http::response_json};
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_login_should_issue_http_only_session_cookie() {
|
||||
@@ -202,7 +202,7 @@ async fn admin_login_test_app(
|
||||
let services = std::sync::Arc::new(codex_proxy_rs::runtime::services::Services::new(
|
||||
&config,
|
||||
stores,
|
||||
fingerprint,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
let state = codex_proxy_rs::runtime::state::AppState {
|
||||
services: (*services).clone(),
|
||||
|
||||
@@ -17,7 +17,10 @@ use codex_proxy_rs::{
|
||||
use serde_json::json;
|
||||
use tower::util::ServiceExt;
|
||||
|
||||
use crate::support::{admin::seed_admin_session, config::test_config, http::response_json};
|
||||
use crate::support::{
|
||||
admin::seed_admin_session, config::test_config, fingerprint::runtime_fingerprint,
|
||||
http::response_json,
|
||||
};
|
||||
|
||||
mod authorization;
|
||||
mod lifecycle;
|
||||
@@ -41,7 +44,11 @@ async fn admin_client_key_test_app(db_name: &str) -> (axum::Router, tempfile::Te
|
||||
usage_records: SqliteUsageRecordStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let services = std::sync::Arc::new(Services::new(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
let state = AppState {
|
||||
services: (*services).clone(),
|
||||
};
|
||||
|
||||
@@ -14,7 +14,7 @@ use codex_proxy_rs::{
|
||||
},
|
||||
infra::{
|
||||
database::connect_sqlite,
|
||||
time::{china_day_start, china_hour, china_hour_start},
|
||||
time::{china_datetime, china_day_start, china_hour},
|
||||
},
|
||||
proxy::dispatch::session_affinity::SqliteSessionAffinityStore,
|
||||
runtime::{
|
||||
@@ -32,7 +32,9 @@ use serde_json::{json, Value};
|
||||
use sqlx::SqlitePool;
|
||||
use tower::util::ServiceExt;
|
||||
|
||||
use crate::support::{admin::seed_admin_session, http::response_json};
|
||||
use crate::support::{
|
||||
admin::seed_admin_session, fingerprint::runtime_fingerprint, http::response_json,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_render_fingerprint_updated_at_as_beijing_time() {
|
||||
@@ -66,8 +68,9 @@ async fn dashboard_summary_should_render_dash_when_fingerprint_updated_at_is_mis
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_calculate_today_traffic_and_latency_from_time_buckets() {
|
||||
let (app, store, _pool, _dir) = dashboard_test_app(
|
||||
async fn dashboard_summary_should_calculate_today_traffic_and_average_first_token_latency_from_time_buckets(
|
||||
) {
|
||||
let (app, store, pool, _dir) = dashboard_test_app(
|
||||
"dashboard-traffic-latency.sqlite",
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
)
|
||||
@@ -93,25 +96,32 @@ async fn dashboard_summary_should_calculate_today_traffic_and_latency_from_time_
|
||||
.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"]["rpm"], "2");
|
||||
assert_eq!(body["data"]["cards"]["traffic"]["tpm"], "30");
|
||||
assert_eq!(body["data"]["cards"]["traffic"]["todayRequests"], "2");
|
||||
assert_eq!(body["data"]["cards"]["traffic"]["todayRequestsValue"], 2);
|
||||
assert_eq!(body["data"]["cards"]["traffic"]["yesterdayRequests"], "1");
|
||||
assert_eq!(body["data"]["cards"]["traffic"]["totalRequests"], "42");
|
||||
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"]["yesterdayTokens"], "30");
|
||||
assert_eq!(body["data"]["cards"]["tokens"]["totalTokens"], "150");
|
||||
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"]["firstTokenLatencyMs"],
|
||||
body["data"]["cards"]["cache"]["averageFirstTokenLatencyMs"],
|
||||
"300 ms"
|
||||
);
|
||||
assert_eq!(
|
||||
body["data"]["cards"]["cache"]["completionLatencyMs"],
|
||||
"1.50 s"
|
||||
);
|
||||
assert!(body["data"]["cards"]["cache"]
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.get("averageLatencyMs")
|
||||
.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -151,7 +161,12 @@ async fn dashboard_trend_should_bucket_usage_by_china_hour() {
|
||||
.await;
|
||||
let now = Utc::now();
|
||||
store
|
||||
.append(&usage_record_with_tokens(now - Duration::seconds(1), 10))
|
||||
.append(&usage_record_with_token_parts(
|
||||
now - Duration::seconds(1),
|
||||
40,
|
||||
60,
|
||||
10,
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -164,7 +179,41 @@ async fn dashboard_trend_should_bucket_usage_by_china_hour() {
|
||||
.expect("trend should include the inserted usage bucket");
|
||||
|
||||
assert_eq!(point["time"], format!("{:02}", china_hour(&now)));
|
||||
assert_eq!(point["tokensValue"], 10);
|
||||
assert_eq!(point["tokensValue"], 100);
|
||||
assert_eq!(point["cachedTokensValue"], 10);
|
||||
assert_f64_eq(point["cacheHitRateValue"].as_f64().unwrap(), 0.25);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_trend_should_only_include_china_today() {
|
||||
let (app, store, _pool, _dir) = dashboard_test_app(
|
||||
"dashboard-trend-china-today.sqlite",
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
)
|
||||
.await;
|
||||
let now = Utc::now();
|
||||
let today_start = china_day_start(now);
|
||||
store
|
||||
.append(&usage_record_with_tokens(
|
||||
today_start - Duration::hours(1),
|
||||
99,
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.append(&usage_record_with_tokens(now, 7))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let body = dashboard_trend(app, "usage").await;
|
||||
let points = body["data"]["points"].as_array().unwrap();
|
||||
let token_total = points
|
||||
.iter()
|
||||
.map(|point| point["tokensValue"].as_u64().unwrap())
|
||||
.sum::<u64>();
|
||||
|
||||
assert_eq!(points.first().unwrap()["time"], "00");
|
||||
assert_eq!(token_total, 7);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -175,20 +224,14 @@ async fn dashboard_latency_trend_should_use_first_token_latency() {
|
||||
)
|
||||
.await;
|
||||
let now = Utc::now();
|
||||
let current_hour_start = china_hour_start(now);
|
||||
let sample_hour_start = if now - current_hour_start > Duration::minutes(31) {
|
||||
current_hour_start
|
||||
} else {
|
||||
current_hour_start - Duration::hours(1)
|
||||
};
|
||||
let mut slow_completion =
|
||||
usage_record_with_tokens(sample_hour_start + Duration::minutes(1), 10);
|
||||
let mut slow_completion = usage_record_with_tokens(now, 10);
|
||||
slow_completion.latency_ms = Some(10_000);
|
||||
slow_completion.metadata["firstTokenMs"] = json!(100_i64);
|
||||
let mut fast_completion =
|
||||
usage_record_with_tokens(sample_hour_start + Duration::minutes(16), 10);
|
||||
slow_completion.metadata["serviceTier"] = json!("default");
|
||||
let mut fast_completion = usage_record_with_tokens(now, 10);
|
||||
fast_completion.latency_ms = Some(1_000);
|
||||
fast_completion.metadata["firstTokenMs"] = json!(500_i64);
|
||||
fast_completion.metadata["serviceTier"] = json!("priority");
|
||||
store.append(&slow_completion).await.unwrap();
|
||||
store.append(&fast_completion).await.unwrap();
|
||||
|
||||
@@ -241,7 +284,7 @@ async fn dashboard_summary_should_return_requested_trend_kind() {
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_calculate_priority_cost_from_event_service_tier() {
|
||||
let cost = dashboard_summary_total_cost_for_usage_record(
|
||||
let cost = dashboard_summary_today_cost_for_usage_record(
|
||||
"dashboard-cost-priority.sqlite",
|
||||
"gpt-5.5",
|
||||
json!({
|
||||
@@ -255,12 +298,12 @@ async fn dashboard_summary_should_calculate_priority_cost_from_event_service_tie
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_f64_eq(cost, 8.75);
|
||||
assert_eq!(cost, "$8.75");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_charge_cached_input_at_cache_read_price() {
|
||||
let cost = dashboard_summary_total_cost_for_usage_record(
|
||||
let cost = dashboard_summary_today_cost_for_usage_record(
|
||||
"dashboard-cost-cached.sqlite",
|
||||
"gpt-5.5",
|
||||
json!({
|
||||
@@ -273,12 +316,12 @@ async fn dashboard_summary_should_charge_cached_input_at_cache_read_price() {
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_f64_eq(cost, 0.03275);
|
||||
assert_eq!(cost, "$0.03");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_apply_long_context_prices() {
|
||||
let cost = dashboard_summary_total_cost_for_usage_record(
|
||||
let cost = dashboard_summary_today_cost_for_usage_record(
|
||||
"dashboard-cost-long-context.sqlite",
|
||||
"gpt-5.5",
|
||||
json!({
|
||||
@@ -291,12 +334,12 @@ async fn dashboard_summary_should_apply_long_context_prices() {
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_f64_eq(cost, 3.045);
|
||||
assert_eq!(cost, "$3.04");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_normalize_codex_model_names_for_cost() {
|
||||
let cost = dashboard_summary_total_cost_for_usage_record(
|
||||
let cost = dashboard_summary_today_cost_for_usage_record(
|
||||
"dashboard-cost-codex-model.sqlite",
|
||||
"codex-auto-review",
|
||||
json!({
|
||||
@@ -309,7 +352,7 @@ async fn dashboard_summary_should_normalize_codex_model_names_for_cost() {
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_f64_eq(cost, 1.75);
|
||||
assert_eq!(cost, "$1.75");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -323,62 +366,106 @@ async fn dashboard_summary_should_not_price_usage_summary_without_model_dimensio
|
||||
|
||||
let body = dashboard_summary(app).await;
|
||||
|
||||
assert_f64_eq(
|
||||
body["data"]["cards"]["tokens"]["totalCostUsdValue"]
|
||||
.as_f64()
|
||||
.unwrap(),
|
||||
0.0,
|
||||
);
|
||||
assert_eq!(body["data"]["cards"]["tokens"]["todayCostUsd"], "$0.00");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_return_backend_formatted_time_fields() {
|
||||
let (app, store, pool, _dir) = dashboard_test_app(
|
||||
let (app, store, _pool, _dir) = dashboard_test_app(
|
||||
"dashboard-backend-formatted-time.sqlite",
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
)
|
||||
.await;
|
||||
let record_at = Utc::now() - Duration::seconds(1);
|
||||
store
|
||||
.append(&usage_record_with_tokens(
|
||||
"2026-06-18T12:34:56Z".parse().unwrap(),
|
||||
10,
|
||||
))
|
||||
.append(&usage_record_with_tokens(record_at, 10))
|
||||
.await
|
||||
.unwrap();
|
||||
seed_usage_summary_with_last_used(&pool, 1, 1, "2000-01-01T00:00:00Z").await;
|
||||
|
||||
let body = dashboard_summary(app).await;
|
||||
|
||||
assert_eq!(
|
||||
body["data"]["usageRecords"][0]["createdAtDisplay"],
|
||||
"2026-06-18 20:34:56"
|
||||
china_datetime(&record_at)
|
||||
);
|
||||
assert_eq!(body["data"]["accountUsage"][0]["lastUsed"], "2000-01-01");
|
||||
assert!(body["data"]["usageRecords"][0].get("metadata").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_only_return_today_account_usage_and_usage_records() {
|
||||
let (app, store, pool, _dir) = dashboard_test_app(
|
||||
"dashboard-summary-today-records.sqlite",
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
)
|
||||
.await;
|
||||
seed_dashboard_account(&pool, "acct_today", "today@example.com").await;
|
||||
seed_dashboard_account(&pool, "acct_yesterday", "yesterday@example.com").await;
|
||||
let now = Utc::now();
|
||||
let today_start = china_day_start(now);
|
||||
let today_record_at = if now - today_start > Duration::minutes(1) {
|
||||
today_start + Duration::minutes(1)
|
||||
} else {
|
||||
now - Duration::seconds(1)
|
||||
};
|
||||
let yesterday_record_at = today_start - Duration::minutes(1);
|
||||
store
|
||||
.append(&usage_record_with_account_tokens(
|
||||
yesterday_record_at,
|
||||
"acct_yesterday",
|
||||
90,
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.append(&usage_record_with_account_tokens(
|
||||
today_record_at,
|
||||
"acct_today",
|
||||
7,
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let body = dashboard_summary(app).await;
|
||||
|
||||
assert_eq!(body["data"]["usageRecords"].as_array().unwrap().len(), 1);
|
||||
assert_eq!(
|
||||
body["data"]["usageRecords"][0]["accountEmail"],
|
||||
"today@example.com"
|
||||
);
|
||||
assert_eq!(body["data"]["accountUsage"].as_array().unwrap().len(), 1);
|
||||
assert_eq!(
|
||||
body["data"]["accountUsage"][0]["email"],
|
||||
"today@example.com"
|
||||
);
|
||||
assert_eq!(body["data"]["accountUsage"][0]["tokens"], "7");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_order_account_usage_by_recent_usage() {
|
||||
let (app, _store, pool, _dir) = dashboard_test_app(
|
||||
let (app, store, pool, _dir) = dashboard_test_app(
|
||||
"dashboard-account-usage-recent-order.sqlite",
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
)
|
||||
.await;
|
||||
seed_dashboard_account_usage(
|
||||
&pool,
|
||||
"acct_old_heavy",
|
||||
"old-heavy@example.com",
|
||||
2_600_000,
|
||||
"2026-06-29T09:00:00Z",
|
||||
)
|
||||
.await;
|
||||
seed_dashboard_account_usage(
|
||||
&pool,
|
||||
"acct_recent_light",
|
||||
"recent-light@example.com",
|
||||
660_000,
|
||||
"2026-06-30T08:00:00Z",
|
||||
)
|
||||
.await;
|
||||
seed_dashboard_account(&pool, "acct_old_heavy", "old-heavy@example.com").await;
|
||||
seed_dashboard_account(&pool, "acct_recent_light", "recent-light@example.com").await;
|
||||
let now = Utc::now();
|
||||
store
|
||||
.append(&usage_record_with_account_tokens(
|
||||
now - Duration::minutes(2),
|
||||
"acct_old_heavy",
|
||||
2_600_000,
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.append(&usage_record_with_account_tokens(
|
||||
now - Duration::minutes(1),
|
||||
"acct_recent_light",
|
||||
660_000,
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let body = dashboard_summary(app).await;
|
||||
|
||||
@@ -386,13 +473,117 @@ async fn dashboard_summary_should_order_account_usage_by_recent_usage() {
|
||||
body["data"]["accountUsage"][0]["email"],
|
||||
"recent-light@example.com"
|
||||
);
|
||||
assert_eq!(body["data"]["accountUsage"][0]["requests"], "660K");
|
||||
assert_eq!(body["data"]["accountUsage"][0]["id"], "acct_recent_light");
|
||||
assert_eq!(body["data"]["accountUsage"][0]["planType"], "K12");
|
||||
assert_eq!(body["data"]["accountUsage"][0]["tokens"], "660K");
|
||||
assert_eq!(
|
||||
body["data"]["accountUsage"][1]["email"],
|
||||
"old-heavy@example.com"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_use_five_hour_quota_window_for_account_usage() {
|
||||
let (app, store, pool, _dir) = dashboard_test_app(
|
||||
"dashboard-account-usage-five-hour-quota.sqlite",
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
)
|
||||
.await;
|
||||
seed_dashboard_account(&pool, "acct_quota_five_hour", "quota-five-hour@example.com").await;
|
||||
store
|
||||
.append(&usage_record_with_account_tokens(
|
||||
Utc::now() - Duration::seconds(1),
|
||||
"acct_quota_five_hour",
|
||||
1,
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
set_account_quota_json(
|
||||
&pool,
|
||||
"acct_quota_five_hour",
|
||||
json!({
|
||||
"monthly_limit": {
|
||||
"used_percent": 92.0,
|
||||
"window_minutes": 43_200
|
||||
},
|
||||
"snapshots": [
|
||||
{
|
||||
"source": "core",
|
||||
"primary": {
|
||||
"used_percent": 31.0,
|
||||
"window_minutes": 300
|
||||
},
|
||||
"secondary": {
|
||||
"used_percent": 88.0,
|
||||
"window_minutes": 10_080
|
||||
}
|
||||
}
|
||||
]
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let body = dashboard_summary(app).await;
|
||||
|
||||
assert_f64_eq(
|
||||
body["data"]["accountUsage"][0]["quotaUsedPercent"]
|
||||
.as_f64()
|
||||
.unwrap(),
|
||||
31.0,
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_fall_back_to_weekly_quota_window_when_five_hour_missing() {
|
||||
let (app, store, pool, _dir) = dashboard_test_app(
|
||||
"dashboard-account-usage-weekly-quota.sqlite",
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
)
|
||||
.await;
|
||||
seed_dashboard_account(&pool, "acct_quota_weekly", "quota-weekly@example.com").await;
|
||||
store
|
||||
.append(&usage_record_with_account_tokens(
|
||||
Utc::now() - Duration::seconds(1),
|
||||
"acct_quota_weekly",
|
||||
1,
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
set_account_quota_json(
|
||||
&pool,
|
||||
"acct_quota_weekly",
|
||||
json!({
|
||||
"monthly_limit": {
|
||||
"used_percent": 92.0,
|
||||
"window_minutes": 43_200
|
||||
},
|
||||
"snapshots": [
|
||||
{
|
||||
"source": "experimental",
|
||||
"primary": {
|
||||
"used_percent": 99.0,
|
||||
"window_minutes": 120
|
||||
},
|
||||
"secondary": {
|
||||
"used_percent": 47.0,
|
||||
"window_minutes": 10_080
|
||||
}
|
||||
}
|
||||
]
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let body = dashboard_summary(app).await;
|
||||
|
||||
assert_f64_eq(
|
||||
body["data"]["accountUsage"][0]["quotaUsedPercent"]
|
||||
.as_f64()
|
||||
.unwrap(),
|
||||
47.0,
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_return_today_health_timeline() {
|
||||
let (app, store, _pool, _dir) = dashboard_test_app(
|
||||
@@ -416,19 +607,16 @@ async fn dashboard_summary_should_return_today_health_timeline() {
|
||||
let points = body["data"]["healthTimeline"]["points"].as_str().unwrap();
|
||||
|
||||
assert_eq!(body["data"]["healthTimeline"]["title"], "请求健康时间线");
|
||||
assert_eq!(
|
||||
body["data"]["healthTimeline"]["description"],
|
||||
"今日请求可靠性"
|
||||
);
|
||||
assert_eq!(body["data"]["healthTimeline"]["description"], "请求可靠性");
|
||||
assert_eq!(points.len(), 96);
|
||||
assert_eq!(points.matches('4').count(), 1);
|
||||
}
|
||||
|
||||
async fn dashboard_summary_total_cost_for_usage_record(
|
||||
async fn dashboard_summary_today_cost_for_usage_record(
|
||||
db_name: &str,
|
||||
model: &str,
|
||||
metadata: Value,
|
||||
) -> f64 {
|
||||
) -> String {
|
||||
let (app, store, _pool, _dir) =
|
||||
dashboard_test_app(db_name, crate::support::fingerprint::test_fingerprint()).await;
|
||||
store
|
||||
@@ -437,9 +625,10 @@ async fn dashboard_summary_total_cost_for_usage_record(
|
||||
.unwrap();
|
||||
|
||||
let body = dashboard_summary(app).await;
|
||||
body["data"]["cards"]["tokens"]["totalCostUsdValue"]
|
||||
.as_f64()
|
||||
body["data"]["cards"]["tokens"]["todayCostUsd"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.to_string()
|
||||
}
|
||||
|
||||
async fn dashboard_test_app(
|
||||
@@ -467,7 +656,11 @@ async fn dashboard_test_app(
|
||||
client_keys: SqliteClientKeyStore::new(pool.clone()),
|
||||
usage_records: SqliteUsageRecordStore::new(pool.clone()),
|
||||
};
|
||||
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let services = std::sync::Arc::new(Services::new(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
let state = AppState {
|
||||
services: (*services).clone(),
|
||||
};
|
||||
@@ -526,6 +719,39 @@ 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,
|
||||
@@ -533,9 +759,10 @@ async fn seed_usage_summary_with_last_used(
|
||||
last_used_at: &str,
|
||||
) {
|
||||
sqlx::query(
|
||||
"insert into accounts (id, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?)",
|
||||
"insert into accounts (id, email, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind("acct_usage")
|
||||
.bind("acct-usage@example.com")
|
||||
.bind("cipher")
|
||||
.bind("active")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
@@ -557,18 +784,14 @@ async fn seed_usage_summary_with_last_used(
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn seed_dashboard_account_usage(
|
||||
pool: &SqlitePool,
|
||||
account_id: &str,
|
||||
email: &str,
|
||||
request_count: i64,
|
||||
last_used_at: &str,
|
||||
) {
|
||||
async fn seed_dashboard_account(pool: &SqlitePool, account_id: &str, email: &str) {
|
||||
sqlx::query(
|
||||
"insert into accounts (id, email, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)",
|
||||
"insert into accounts (id, email, chatgpt_account_id, plan_type, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(account_id)
|
||||
.bind(email)
|
||||
.bind(format!("upstream_{account_id}"))
|
||||
.bind("K12")
|
||||
.bind("cipher")
|
||||
.bind("active")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
@@ -576,18 +799,16 @@ async fn seed_dashboard_account_usage(
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"insert into account_usage (account_id, request_count, input_tokens, output_tokens, total_tokens, last_used_at) values (?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(account_id)
|
||||
.bind(request_count)
|
||||
.bind(0_i64)
|
||||
.bind(0_i64)
|
||||
.bind(0_i64)
|
||||
.bind(last_used_at)
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn set_account_quota_json(pool: &SqlitePool, account_id: &str, quota_json: Value) {
|
||||
sqlx::query("update accounts set quota_json = ?, quota_fetched_at = ? where id = ?")
|
||||
.bind(quota_json.to_string())
|
||||
.bind("2026-06-30T08:01:00Z")
|
||||
.bind(account_id)
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
fn service_status_value<'a>(body: &'a Value, label: &str) -> &'a str {
|
||||
@@ -601,19 +822,38 @@ fn service_status_value<'a>(body: &'a Value, label: &str) -> &'a str {
|
||||
}
|
||||
|
||||
fn usage_record_with_tokens(created_at: chrono::DateTime<Utc>, total_tokens: u64) -> UsageRecord {
|
||||
usage_record_with_token_parts(created_at, total_tokens, 0, 0)
|
||||
}
|
||||
|
||||
fn usage_record_with_token_parts(
|
||||
created_at: chrono::DateTime<Utc>,
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cached_tokens: u64,
|
||||
) -> UsageRecord {
|
||||
usage_record(
|
||||
created_at,
|
||||
"gpt-5.5",
|
||||
json!({
|
||||
"usage": {
|
||||
"inputTokens": total_tokens,
|
||||
"outputTokens": 0u64,
|
||||
"cachedTokens": 0u64
|
||||
"inputTokens": input_tokens,
|
||||
"outputTokens": output_tokens,
|
||||
"cachedTokens": cached_tokens
|
||||
}
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
fn usage_record_with_account_tokens(
|
||||
created_at: chrono::DateTime<Utc>,
|
||||
account_id: &str,
|
||||
total_tokens: u64,
|
||||
) -> UsageRecord {
|
||||
let mut record = usage_record_with_tokens(created_at, total_tokens);
|
||||
record.account_id = Some(account_id.to_string());
|
||||
record
|
||||
}
|
||||
|
||||
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());
|
||||
|
||||
@@ -23,7 +23,10 @@ use serde_json::json;
|
||||
use sqlx::SqlitePool;
|
||||
use tower::util::ServiceExt;
|
||||
|
||||
use crate::support::{admin::seed_admin_session, config::test_config, http::response_json};
|
||||
use crate::support::{
|
||||
admin::seed_admin_session, config::test_config, fingerprint::runtime_fingerprint,
|
||||
http::response_json,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_usage_records_should_require_admin_session_cookie() {
|
||||
@@ -883,7 +886,11 @@ async fn admin_usage_records_test_app_with_config(
|
||||
usage_records: SqliteUsageRecordStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let services = std::sync::Arc::new(Services::new(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
let state = AppState {
|
||||
services: (*services).clone(),
|
||||
};
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
use chrono::{Duration, Utc};
|
||||
use codex_proxy_rs::{
|
||||
admin::monitoring::{
|
||||
usage_record_model::{UsageRecord, UsageRecordLevel},
|
||||
@@ -91,6 +92,22 @@ async fn usage_record_store_should_promote_diagnostic_metadata_fields() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_record_store_should_not_promote_cf_ray_as_upstream_request_id() {
|
||||
let (store, _dir) = usage_record_store("usage-records-cf-ray.sqlite").await;
|
||||
let mut event = UsageRecord::new("request", UsageRecordLevel::Error, "upstream failed");
|
||||
event.id = "usage_cf_ray".to_string();
|
||||
event.metadata = json!({
|
||||
"upstreamStatus": 403,
|
||||
"cfRay": "ray-1"
|
||||
});
|
||||
|
||||
store.append(&event).await.unwrap();
|
||||
let loaded = store.get("usage_cf_ray").await.unwrap().unwrap();
|
||||
|
||||
assert_eq!(loaded.upstream_request_id, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_record_store_should_get_and_clear_events() {
|
||||
let (store, _dir) = usage_record_store("usage-records-clear.sqlite").await;
|
||||
@@ -108,6 +125,28 @@ async fn usage_record_store_should_get_and_clear_events() {
|
||||
assert!(page.items.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_record_store_should_trim_records_older_than_thirty_days() {
|
||||
let (store, _dir) = usage_record_store("usage-records-retention.sqlite").await;
|
||||
let now = Utc::now();
|
||||
let mut expired = UsageRecord::new("v1.response", UsageRecordLevel::Info, "expired");
|
||||
expired.id = "usage_expired".to_string();
|
||||
expired.route = Some("/v1/responses".to_string());
|
||||
expired.created_at = now - Duration::days(31);
|
||||
store.append(&expired).await.unwrap();
|
||||
|
||||
let mut retained = UsageRecord::new("v1.response", UsageRecordLevel::Info, "retained");
|
||||
retained.id = "usage_retained".to_string();
|
||||
retained.route = Some("/v1/responses".to_string());
|
||||
retained.created_at = now - Duration::days(29);
|
||||
store.append(&retained).await.unwrap();
|
||||
|
||||
assert_eq!(store.trim_to_retention(now).await.unwrap(), 1);
|
||||
|
||||
assert!(store.get("usage_expired").await.unwrap().is_none());
|
||||
assert!(store.get("usage_retained").await.unwrap().is_some());
|
||||
}
|
||||
|
||||
async fn usage_record_store(db_name: &str) -> (SqliteUsageRecordStore, tempfile::TempDir) {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join(db_name);
|
||||
|
||||
@@ -7,7 +7,7 @@ fn settings_service_should_apply_retained_settings_patch() {
|
||||
let mut current = AdminSettings {
|
||||
max_concurrent_per_account: 4,
|
||||
request_interval_ms: 50,
|
||||
rotation_strategy: "least_used".to_string(),
|
||||
rotation_strategy: "smart".to_string(),
|
||||
..AdminSettings::default()
|
||||
};
|
||||
|
||||
@@ -21,13 +21,6 @@ fn settings_service_should_apply_retained_settings_patch() {
|
||||
]
|
||||
.into(),
|
||||
),
|
||||
model_account_routes: Some(
|
||||
[(
|
||||
" gpt-5.5 ".to_string(),
|
||||
vec![" acct_a ".to_string(), "acct_b".to_string()],
|
||||
)]
|
||||
.into(),
|
||||
),
|
||||
refresh_margin_seconds: Some(600),
|
||||
refresh_concurrency: Some(3),
|
||||
max_concurrent_per_account: Some(5),
|
||||
@@ -39,10 +32,6 @@ fn settings_service_should_apply_retained_settings_patch() {
|
||||
|
||||
assert_eq!(current.model_aliases["claude-sonnet"], "gpt-5.5");
|
||||
assert_eq!(current.model_aliases["openai-fast"], "openai:gpt-4o");
|
||||
assert_eq!(
|
||||
current.model_account_routes["gpt-5.5"],
|
||||
vec!["acct_a".to_string(), "acct_b".to_string()]
|
||||
);
|
||||
assert_eq!(current.refresh_margin_seconds, 600);
|
||||
assert_eq!(current.refresh_concurrency, 3);
|
||||
assert_eq!(current.rotation_strategy, "round_robin");
|
||||
@@ -50,27 +39,6 @@ fn settings_service_should_apply_retained_settings_patch() {
|
||||
assert_eq!(current.request_interval_ms, 125);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settings_service_should_reject_invalid_model_account_routes() {
|
||||
let mut current = AdminSettings::default();
|
||||
let error = SettingsService::apply_patch(
|
||||
&mut current,
|
||||
AdminSettingsPatch {
|
||||
model_account_routes: Some(
|
||||
[(
|
||||
"gpt-5.5".to_string(),
|
||||
vec!["acct_a".to_string(), "acct_a".to_string()],
|
||||
)]
|
||||
.into(),
|
||||
),
|
||||
..AdminSettingsPatch::default()
|
||||
},
|
||||
)
|
||||
.expect_err("duplicate model route account should be rejected");
|
||||
|
||||
assert_eq!(error.field(), "modelAccountRoutes");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settings_service_should_reject_invalid_patch_values() {
|
||||
let mut current = AdminSettings::default();
|
||||
|
||||
@@ -28,7 +28,8 @@ use sqlx::SqlitePool;
|
||||
use tower::util::ServiceExt;
|
||||
|
||||
use crate::support::{
|
||||
admin::seed_admin_session, config::test_config as base_test_config, http::response_json,
|
||||
admin::seed_admin_session, config::test_config as base_test_config,
|
||||
fingerprint::runtime_fingerprint, http::response_json,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
@@ -66,10 +67,9 @@ async fn admin_settings_should_return_runtime_fields() {
|
||||
.unwrap();
|
||||
let body = response_json(response).await;
|
||||
assert_eq!(body["data"]["modelAliases"]["codex-fast"], "gpt-5.5");
|
||||
assert_eq!(body["data"]["modelAccountRoutes"], json!({}));
|
||||
assert_eq!(body["data"]["refreshMarginSeconds"], 240);
|
||||
assert_eq!(body["data"]["refreshConcurrency"], 2);
|
||||
assert_eq!(body["data"]["rotationStrategy"], "least_used");
|
||||
assert_eq!(body["data"]["rotationStrategy"], "smart");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -267,8 +267,6 @@ async fn admin_settings_update_should_persist_runtime_settings_to_database() {
|
||||
let config = test_config(url);
|
||||
let pool = connect_sqlite(&config.database.url).await.unwrap();
|
||||
seed_admin_session(&pool, "session_1").await;
|
||||
seed_account(&pool, "acct_route_a").await;
|
||||
seed_account(&pool, "acct_route_b").await;
|
||||
let stores = BackgroundTaskStores {
|
||||
accounts: SqliteAccountStore::new(pool.clone()),
|
||||
admin_sessions: SqliteAdminSessionStore::new(pool.clone()),
|
||||
@@ -280,7 +278,7 @@ async fn admin_settings_update_should_persist_runtime_settings_to_database() {
|
||||
usage_records: SqliteUsageRecordStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = Services::new(&config, stores, fingerprint);
|
||||
let services = Services::new(&config, stores, runtime_fingerprint(fingerprint));
|
||||
let services = std::sync::Arc::new(services);
|
||||
let state = AppState {
|
||||
services: (*services).clone(),
|
||||
@@ -303,9 +301,6 @@ async fn admin_settings_update_should_persist_runtime_settings_to_database() {
|
||||
"gpt-5.2": "gpt-5.5",
|
||||
"claude-sonnet": "gpt-5.5"
|
||||
},
|
||||
"modelAccountRoutes": {
|
||||
"gpt-5.5": ["acct_route_a", "acct_route_b"]
|
||||
},
|
||||
"refreshMarginSeconds": 900,
|
||||
"refreshConcurrency": 4,
|
||||
"maxConcurrentPerAccount": 7,
|
||||
@@ -320,10 +315,6 @@ async fn admin_settings_update_should_persist_runtime_settings_to_database() {
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = response_json(response).await;
|
||||
assert_eq!(body["data"]["modelAliases"]["gpt-5.2"], "gpt-5.5");
|
||||
assert_eq!(
|
||||
body["data"]["modelAccountRoutes"]["gpt-5.5"][0],
|
||||
"acct_route_a"
|
||||
);
|
||||
assert_eq!(body["data"]["refreshMarginSeconds"], 900);
|
||||
assert_eq!(body["data"]["refreshConcurrency"], 4);
|
||||
|
||||
@@ -357,29 +348,12 @@ async fn admin_settings_update_should_persist_runtime_settings_to_database() {
|
||||
assert_eq!(row.3, 7);
|
||||
assert_eq!(row.4, 80);
|
||||
assert_eq!(row.5, "round_robin");
|
||||
let route_rows: Vec<(String, String, i64)> = sqlx::query_as(
|
||||
"select model, account_id, priority from model_account_routes order by model, priority",
|
||||
)
|
||||
.fetch_all(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
route_rows,
|
||||
vec![
|
||||
("gpt-5.5".to_string(), "acct_route_a".to_string(), 0),
|
||||
("gpt-5.5".to_string(), "acct_route_b".to_string(), 1),
|
||||
]
|
||||
);
|
||||
assert!(!dir.path().join("config.yaml").exists());
|
||||
|
||||
let restarted_config = RuntimeSettingsService::load_or_initialize_config(config.clone(), &pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(restarted_config.model_aliases["gpt-5.2"], "gpt-5.5");
|
||||
assert_eq!(
|
||||
restarted_config.model_account_routes["gpt-5.5"],
|
||||
vec!["acct_route_a".to_string(), "acct_route_b".to_string()]
|
||||
);
|
||||
assert_eq!(restarted_config.auth.refresh_margin_seconds, 900);
|
||||
assert_eq!(restarted_config.auth.refresh_concurrency, 4);
|
||||
assert_eq!(restarted_config.auth.rotation_strategy, "round_robin");
|
||||
@@ -409,7 +383,11 @@ async fn admin_settings_update_should_apply_runtime_services() {
|
||||
usage_records: SqliteUsageRecordStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let services = std::sync::Arc::new(Services::new(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
services
|
||||
.account_pool
|
||||
.restore_from_repository()
|
||||
@@ -433,10 +411,6 @@ async fn admin_settings_update_should_apply_runtime_services() {
|
||||
"modelAliases": {
|
||||
"runtime-alias": "gpt-5.5"
|
||||
},
|
||||
"modelAccountRoutes": {
|
||||
"gpt-5.5": ["acct_runtime_b"],
|
||||
"gpt-5.6": ["acct_runtime_a", "acct_runtime_b"]
|
||||
},
|
||||
"refreshMarginSeconds": 120,
|
||||
"refreshConcurrency": 3,
|
||||
"maxConcurrentPerAccount": 2,
|
||||
@@ -466,15 +440,6 @@ async fn admin_settings_update_should_apply_runtime_services() {
|
||||
assert_eq!(capacity.total_slots, 4);
|
||||
assert_eq!(capacity.max_concurrent_per_account, 2);
|
||||
|
||||
let route_request = AccountAcquireRequest::new("gpt-5.5", Utc::now());
|
||||
let routed = services
|
||||
.account_pool
|
||||
.acquire_with(&route_request)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(routed.account.id, "acct_runtime_b");
|
||||
services.account_pool.release(&routed.account.id).await;
|
||||
|
||||
let rotation_request = AccountAcquireRequest::new("gpt-5.6", Utc::now());
|
||||
let first = services
|
||||
.account_pool
|
||||
@@ -486,15 +451,13 @@ async fn admin_settings_update_should_apply_runtime_services() {
|
||||
.acquire_with(&rotation_request)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(first.account.id, "acct_runtime_b");
|
||||
assert_eq!(second.account.id, "acct_runtime_a");
|
||||
assert_ne!(first.account.id, second.account.id);
|
||||
|
||||
let repeated = services
|
||||
.account_pool
|
||||
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", Utc::now()))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(repeated.account.id, "acct_runtime_b");
|
||||
let started_at = Instant::now();
|
||||
services
|
||||
.account_pool
|
||||
@@ -522,7 +485,11 @@ async fn admin_settings_update_should_reject_unsupported_or_invalid_fields() {
|
||||
usage_records: SqliteUsageRecordStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let services = std::sync::Arc::new(Services::new(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
let state = AppState {
|
||||
services: (*services).clone(),
|
||||
};
|
||||
@@ -564,7 +531,11 @@ async fn admin_settings_test_app(db_name: &str) -> (axum::Router, tempfile::Temp
|
||||
usage_records: SqliteUsageRecordStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let services = std::sync::Arc::new(Services::new(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
let state = AppState {
|
||||
services: (*services).clone(),
|
||||
};
|
||||
@@ -574,26 +545,6 @@ async fn admin_settings_test_app(db_name: &str) -> (axum::Router, tempfile::Temp
|
||||
)
|
||||
}
|
||||
|
||||
async fn seed_account(pool: &SqlitePool, account_id: &str) {
|
||||
sqlx::query(
|
||||
r"
|
||||
insert into accounts (
|
||||
id,
|
||||
access_token,
|
||||
status,
|
||||
added_at,
|
||||
updated_at
|
||||
) values (?, ?, 'active', ?, ?)",
|
||||
)
|
||||
.bind(account_id)
|
||||
.bind("access-token")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn seed_schedulable_account(pool: &SqlitePool, account_id: &str) {
|
||||
sqlx::query(
|
||||
r"
|
||||
|
||||
@@ -35,7 +35,9 @@ use wiremock::{
|
||||
};
|
||||
|
||||
use crate::support::{
|
||||
admin::seed_admin_session, config::test_config, fingerprint::test_fingerprint,
|
||||
admin::seed_admin_session,
|
||||
config::test_config,
|
||||
fingerprint::{runtime_fingerprint, test_fingerprint},
|
||||
http::response_json,
|
||||
};
|
||||
|
||||
@@ -1218,7 +1220,11 @@ async fn admin_system_test_app(db_name: &str) -> (axum::Router, tempfile::TempDi
|
||||
seed_admin_session(&pool, "session_1").await;
|
||||
let config = test_config(url);
|
||||
let stores = stores(pool);
|
||||
let services = Arc::new(Services::new(&config, stores, test_fingerprint()));
|
||||
let services = Arc::new(Services::new(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint(test_fingerprint()),
|
||||
));
|
||||
let state = AppState {
|
||||
services: (*services).clone(),
|
||||
};
|
||||
|
||||
@@ -84,7 +84,7 @@ fn default_config_keeps_only_codex_backend() {
|
||||
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);
|
||||
assert_eq!(cfg.auth.rotation_strategy, "least_used");
|
||||
assert_eq!(cfg.auth.rotation_strategy, "smart");
|
||||
assert_eq!(cfg.auth.oauth_client_id, "app_EMoamEEZ73f0CkXaXp7hrann");
|
||||
assert_eq!(
|
||||
cfg.auth.oauth_token_endpoint,
|
||||
|
||||
@@ -384,6 +384,37 @@ async fn request_context_should_accept_real_ip_from_trusted_proxy_when_forwarded
|
||||
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "198.51.100.23");
|
||||
}
|
||||
|
||||
#[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");
|
||||
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", "198.51.100.99, 203.0.113.42")
|
||||
.header("x-real-ip", "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(), "203.0.113.42");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn request_context_should_ignore_forwarded_ip_from_untrusted_peer() {
|
||||
let trusted_proxies =
|
||||
|
||||
@@ -10,7 +10,9 @@ use serde_json::Value;
|
||||
use sqlx::SqlitePool;
|
||||
use tower::ServiceExt;
|
||||
|
||||
use crate::support::{config::test_config as base_test_config, http::response_json};
|
||||
use crate::support::{
|
||||
config::test_config as base_test_config, fingerprint::runtime_fingerprint, http::response_json,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn server_router_should_serve_frontend_assets_without_shadowing_api_routes() {
|
||||
@@ -200,7 +202,7 @@ async fn router_with_dist_and_database(dist: &Path, database_url: &str) -> (Rout
|
||||
let services = std::sync::Arc::new(codex_proxy_rs::runtime::services::Services::new(
|
||||
&config,
|
||||
stores,
|
||||
fingerprint,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
let state = codex_proxy_rs::runtime::state::AppState {
|
||||
services: (*services).clone(),
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
pub(crate) mod format;
|
||||
pub(crate) mod log_rotation;
|
||||
pub(crate) mod paths;
|
||||
pub(crate) mod storage_schema;
|
||||
pub(crate) mod time;
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
use std::{ffi::OsString, fs, sync::Mutex};
|
||||
|
||||
use codex_proxy_rs::infra::paths::{ensure_data_dir, load_or_create_installation_id};
|
||||
|
||||
const INSTALLATION_ID_FILE_NAME: &str = "installation_id";
|
||||
|
||||
static ENV_LOCK: Mutex<()> = Mutex::new(());
|
||||
|
||||
#[test]
|
||||
fn data_dir_should_use_xdg_data_home_directly() {
|
||||
let _guard = ENV_LOCK.lock().unwrap();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let _xdg_data_home = EnvVarGuard::set("XDG_DATA_HOME", dir.path().as_os_str().to_owned());
|
||||
|
||||
let resolved = ensure_data_dir().unwrap();
|
||||
|
||||
assert_eq!(resolved, dir.path());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_or_create_installation_id_should_persist_under_data_dir_root() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
let installation_id =
|
||||
load_or_create_installation_id(Some(dir.path())).expect("installation id should generate");
|
||||
|
||||
assert_eq!(
|
||||
fs::read_to_string(dir.path().join(INSTALLATION_ID_FILE_NAME)).unwrap(),
|
||||
installation_id
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_or_create_installation_id_should_seed_data_dir_from_codex_desktop_id() {
|
||||
let _guard = ENV_LOCK.lock().unwrap();
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let data = tempfile::tempdir().unwrap();
|
||||
let codex_path = home.path().join(".codex").join(INSTALLATION_ID_FILE_NAME);
|
||||
let desktop_id = "018f8f6b-1d7b-7b7c-b9c8-8c9f4c6d0e1a";
|
||||
fs::create_dir_all(codex_path.parent().unwrap()).unwrap();
|
||||
fs::write(&codex_path, desktop_id).unwrap();
|
||||
let _home = EnvVarGuard::set("HOME", home.path().as_os_str().to_owned());
|
||||
|
||||
let installation_id =
|
||||
load_or_create_installation_id(Some(data.path())).expect("installation id should seed");
|
||||
|
||||
assert_eq!(installation_id, desktop_id);
|
||||
assert_eq!(
|
||||
fs::read_to_string(data.path().join(INSTALLATION_ID_FILE_NAME)).unwrap(),
|
||||
desktop_id
|
||||
);
|
||||
}
|
||||
|
||||
struct EnvVarGuard {
|
||||
key: &'static str,
|
||||
previous: Option<OsString>,
|
||||
}
|
||||
|
||||
impl EnvVarGuard {
|
||||
fn set(key: &'static str, value: OsString) -> Self {
|
||||
let previous = std::env::var_os(key);
|
||||
std::env::set_var(key, value);
|
||||
Self { key, previous }
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for EnvVarGuard {
|
||||
fn drop(&mut self) {
|
||||
match &self.previous {
|
||||
Some(value) => std::env::set_var(self.key, value),
|
||||
None => std::env::remove_var(self.key),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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', 'model_account_routes', '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', 'model_plan_snapshots', 'session_affinities') order by name",
|
||||
)
|
||||
.fetch_all(&pool)
|
||||
.await
|
||||
@@ -28,7 +28,6 @@ async fn sqlite_schema_should_create_current_tables() {
|
||||
"client_api_keys",
|
||||
"fingerprint_update_history",
|
||||
"fingerprints",
|
||||
"model_account_routes",
|
||||
"model_plan_snapshots",
|
||||
"runtime_settings",
|
||||
"schema_migrations",
|
||||
@@ -52,7 +51,13 @@ async fn sqlite_schema_should_record_applied_migrations() {
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(rows, [(1, "initial".to_string())]);
|
||||
assert_eq!(
|
||||
rows,
|
||||
[
|
||||
(1, "initial".to_string()),
|
||||
(2, "rename_rotation_strategies".to_string()),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -156,7 +161,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_model_account_routes_account', 'idx_model_account_routes_enabled_model', '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_session_affinities_active_order', 'idx_usage_time_buckets_bucket', 'idx_usage_time_buckets_model_bucket') order by name",
|
||||
)
|
||||
.fetch_all(&pool)
|
||||
.await
|
||||
@@ -166,8 +171,6 @@ async fn sqlite_schema_should_persist_diagnostic_filter_indexes() {
|
||||
rows.into_iter().map(|row| row.0).collect::<Vec<_>>(),
|
||||
[
|
||||
"idx_account_cookies_expires",
|
||||
"idx_model_account_routes_account",
|
||||
"idx_model_account_routes_enabled_model",
|
||||
"idx_session_affinities_active_order",
|
||||
"idx_usage_records_level_created",
|
||||
"idx_usage_records_model_created",
|
||||
@@ -246,33 +249,6 @@ async fn sqlite_schema_should_persist_usage_time_bucket_columns() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sqlite_schema_should_persist_model_account_routes_columns() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("model-account-routes.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('model_account_routes') where name in ('model', 'account_id', 'priority', 'enabled', 'created_at', 'updated_at') order by name",
|
||||
)
|
||||
.fetch_all(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
rows.into_iter().map(|row| row.0).collect::<Vec<_>>(),
|
||||
[
|
||||
"account_id",
|
||||
"created_at",
|
||||
"enabled",
|
||||
"model",
|
||||
"priority",
|
||||
"updated_at",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sqlite_schema_should_persist_plain_account_cookie_value_column() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
@@ -553,22 +529,9 @@ async fn sqlite_schema_should_reject_non_boolean_flags() {
|
||||
)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
sqlx::query(
|
||||
"insert into accounts (id, access_token, status, added_at, updated_at) values ('acct_route', 'access-token', 'active', '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let model_route_result = sqlx::query(
|
||||
"insert into model_account_routes (model, account_id, enabled, created_at, updated_at) values ('gpt-5.5', 'acct_route', 2, '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
|
||||
assert!(api_key_result.is_err());
|
||||
assert!(account_result.is_err());
|
||||
assert!(quota_verify_result.is_err());
|
||||
assert!(model_route_result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -579,12 +542,12 @@ async fn sqlite_schema_should_reject_invalid_runtime_settings() {
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
|
||||
let invalid_id_result = sqlx::query(
|
||||
"insert into runtime_settings (id, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, updated_at) values (2, 300, 2, 3, 50, 'least_used', '2026-06-14T00:00:00Z')",
|
||||
"insert into runtime_settings (id, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, updated_at) values (2, 300, 2, 3, 50, 'smart', '2026-06-14T00:00:00Z')",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
let invalid_refresh_result = sqlx::query(
|
||||
"insert into runtime_settings (id, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, updated_at) values (1, 0, 2, 3, 50, 'least_used', '2026-06-14T00:00:00Z')",
|
||||
"insert into runtime_settings (id, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, updated_at) values (1, 0, 2, 3, 50, 'smart', '2026-06-14T00:00:00Z')",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
@@ -599,28 +562,6 @@ async fn sqlite_schema_should_reject_invalid_runtime_settings() {
|
||||
assert!(invalid_strategy_result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sqlite_schema_should_reject_invalid_model_account_route_priority() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("model-route-check.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
|
||||
sqlx::query(
|
||||
"insert into accounts (id, access_token, status, added_at, updated_at) values ('acct_route', 'access-token', 'active', '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let result = sqlx::query(
|
||||
"insert into model_account_routes (model, account_id, priority, created_at, updated_at) values ('gpt-5.5', 'acct_route', -1, '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await;
|
||||
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sqlite_schema_should_reject_negative_usage_time_bucket_counts() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
@@ -16,9 +16,7 @@ use codex_proxy_rs::{
|
||||
admin::keys::service::SqliteClientKeyStore,
|
||||
admin::monitoring::{
|
||||
usage_record_model::UsageRecordLevel,
|
||||
usage_record_store::{
|
||||
SqliteUsageRecordStore, UsageRecordFilter, DEFAULT_USAGE_RECORD_CAPACITY,
|
||||
},
|
||||
usage_record_store::{SqliteUsageRecordStore, UsageRecordFilter},
|
||||
},
|
||||
config::schema::AppConfig,
|
||||
http::router,
|
||||
@@ -29,7 +27,8 @@ use codex_proxy_rs::{
|
||||
runtime::services::{BackgroundTaskStores, Services, UsageRecordOptions},
|
||||
runtime::state::AppState,
|
||||
upstream::accounts::{
|
||||
cookies::SqliteCookieStore, store::SqliteAccountStore, token_refresh::RefreshLeaseStore,
|
||||
cookies::SqliteCookieStore, model::AccountStatus, pool::AccountAcquireRequest,
|
||||
store::SqliteAccountStore, token_refresh::RefreshLeaseStore,
|
||||
},
|
||||
upstream::fingerprint::FingerprintRepository,
|
||||
};
|
||||
@@ -272,7 +271,7 @@ async fn test_app_state_with_pool_installation_id_and_usage_record_options(
|
||||
let services = Services::try_with_installation_id_and_usage_record_options(
|
||||
&config,
|
||||
stores,
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
Some(installation_id),
|
||||
usage_record_options,
|
||||
)
|
||||
@@ -387,7 +386,6 @@ async fn test_app_with_account_pool_and_logging_capture_body(
|
||||
config.logging.enabled = true;
|
||||
let usage_record_options = UsageRecordOptions {
|
||||
enabled: true,
|
||||
capacity: DEFAULT_USAGE_RECORD_CAPACITY,
|
||||
capture_body: true,
|
||||
};
|
||||
let state = test_app_state_with_pool_installation_id_and_usage_record_options(
|
||||
@@ -549,6 +547,43 @@ async fn test_app_with_two_accounts_and_affinity_status(
|
||||
async fn test_app_with_two_accounts(
|
||||
base_url: String,
|
||||
) -> (axum::Router, String, SqlitePool, tempfile::TempDir) {
|
||||
let (app, _state, api_key, pool, dir) = test_app_with_two_accounts_and_state(base_url).await;
|
||||
(app, api_key, pool, dir)
|
||||
}
|
||||
|
||||
async fn test_app_with_two_accounts_and_logging(
|
||||
base_url: String,
|
||||
) -> (axum::Router, String, SqlitePool, tempfile::TempDir) {
|
||||
let (app, _state, api_key, pool, dir) =
|
||||
test_app_with_two_accounts_and_state_config(base_url, |config| {
|
||||
config.logging.enabled = true;
|
||||
})
|
||||
.await;
|
||||
(app, api_key, pool, dir)
|
||||
}
|
||||
|
||||
async fn test_app_with_two_accounts_and_state(
|
||||
base_url: String,
|
||||
) -> (
|
||||
axum::Router,
|
||||
AppState,
|
||||
String,
|
||||
SqlitePool,
|
||||
tempfile::TempDir,
|
||||
) {
|
||||
test_app_with_two_accounts_and_state_config(base_url, |_| {}).await
|
||||
}
|
||||
|
||||
async fn test_app_with_two_accounts_and_state_config(
|
||||
base_url: String,
|
||||
configure: impl FnOnce(&mut AppConfig),
|
||||
) -> (
|
||||
axum::Router,
|
||||
AppState,
|
||||
String,
|
||||
SqlitePool,
|
||||
tempfile::TempDir,
|
||||
) {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("openai-responses-fallback.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
@@ -562,8 +597,10 @@ async fn test_app_with_two_accounts(
|
||||
"chatgpt-secondary",
|
||||
)
|
||||
.await;
|
||||
let mut config = test_config(url, base_url);
|
||||
configure(&mut config);
|
||||
let state = test_app_state_with_pool_and_installation_id(
|
||||
&test_config(url, base_url),
|
||||
&config,
|
||||
pool.clone(),
|
||||
TEST_INSTALLATION_ID.to_string(),
|
||||
)
|
||||
@@ -574,7 +611,13 @@ async fn test_app_with_two_accounts(
|
||||
.restore_from_repository()
|
||||
.await
|
||||
.expect("account pool should restore");
|
||||
(router::router().with_state(state), api_key, pool, dir)
|
||||
(
|
||||
router::router().with_state(state.clone()),
|
||||
state,
|
||||
api_key,
|
||||
pool,
|
||||
dir,
|
||||
)
|
||||
}
|
||||
|
||||
async fn seed_openai_admin_session(pool: &SqlitePool, session_id: &str) {
|
||||
|
||||
@@ -781,15 +781,37 @@ async fn responses_websocket_should_implicitly_resume_full_history_with_reasonin
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_should_strip_implicit_disabled_affinity_when_switching_accounts() {
|
||||
responses_websocket_should_strip_implicit_affinity_when_switching_accounts(
|
||||
AccountStatus::Disabled,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_should_strip_implicit_banned_affinity_when_switching_accounts() {
|
||||
responses_websocket_should_strip_implicit_affinity_when_switching_accounts(
|
||||
AccountStatus::Banned,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn responses_websocket_should_strip_implicit_affinity_when_switching_accounts(
|
||||
status: AccountStatus,
|
||||
) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base_url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let status_name = status.as_str();
|
||||
let first_response_id = format!("resp_implicit_{status_name}_first");
|
||||
let second_response_id = format!("resp_implicit_{status_name}_second");
|
||||
let upstream_first_response_id = first_response_id.clone();
|
||||
let upstream_second_response_id = second_response_id.clone();
|
||||
let upstream = tokio::spawn(async move {
|
||||
let (first_stream, _) = listener.accept().await.unwrap();
|
||||
let mut first_websocket =
|
||||
accept_websocket_with_authorization(first_stream, "Bearer access-primary").await;
|
||||
let first_payload = send_websocket_response_and_capture_payload(
|
||||
&mut first_websocket,
|
||||
websocket_completed_response("resp_implicit_ban_first", 4, 1),
|
||||
websocket_completed_response(&upstream_first_response_id, 4, 1),
|
||||
)
|
||||
.await;
|
||||
first_websocket.close(None).await.unwrap();
|
||||
@@ -804,7 +826,7 @@ async fn responses_websocket_should_strip_implicit_disabled_affinity_when_switch
|
||||
.unwrap();
|
||||
let second_payload = send_websocket_response_and_capture_payload(
|
||||
&mut second_websocket,
|
||||
websocket_completed_response("resp_implicit_ban_second", 3, 1),
|
||||
websocket_completed_response(&upstream_second_response_id, 3, 1),
|
||||
)
|
||||
.await;
|
||||
second_websocket.close(None).await.unwrap();
|
||||
@@ -812,8 +834,7 @@ async fn responses_websocket_should_strip_implicit_disabled_affinity_when_switch
|
||||
|
||||
(first_payload, second_payload, captured_second_headers)
|
||||
});
|
||||
let (app, api_key, pool, _dir) = test_app_with_two_accounts(base_url).await;
|
||||
seed_openai_admin_session(&pool, "session_status_cycle").await;
|
||||
let (app, state, api_key, pool, _dir) = test_app_with_two_accounts_and_state(base_url).await;
|
||||
|
||||
let first_response = app
|
||||
.clone()
|
||||
@@ -831,9 +852,15 @@ async fn responses_websocket_should_strip_implicit_disabled_affinity_when_switch
|
||||
assert_eq!(first_response.status(), StatusCode::OK);
|
||||
assert!(response_text(first_response)
|
||||
.await
|
||||
.contains("\"id\":\"resp_implicit_ban_first\""));
|
||||
.contains(&format!("\"id\":\"{first_response_id}\"")));
|
||||
|
||||
update_admin_account_status(&app, "acct_primary", "disabled").await;
|
||||
assert!(
|
||||
state
|
||||
.services
|
||||
.account_pool
|
||||
.set_status("acct_primary", status)
|
||||
.await
|
||||
);
|
||||
|
||||
let second_response = app
|
||||
.oneshot(responses_json_request(
|
||||
@@ -854,11 +881,11 @@ async fn responses_websocket_should_strip_implicit_disabled_affinity_when_switch
|
||||
assert_eq!(second_response.status(), StatusCode::OK);
|
||||
assert!(response_text(second_response)
|
||||
.await
|
||||
.contains("\"id\":\"resp_implicit_ban_second\""));
|
||||
.contains(&format!("\"id\":\"{second_response_id}\"")));
|
||||
let (first_payload, second_payload, second_headers) = upstream.await.unwrap();
|
||||
let affinity_count: (i64,) =
|
||||
sqlx::query_as("select count(*) from session_affinities where response_id = ?")
|
||||
.bind("resp_implicit_ban_first")
|
||||
.bind(&first_response_id)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -1553,6 +1580,106 @@ async fn responses_websocket_non_stream_previous_response_not_found_should_strip
|
||||
assert!(second_payload.get("previous_response_id").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_history_recovery_should_stay_on_fallback_account_when_preferred_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(
|
||||
&listener,
|
||||
"Bearer access-secondary",
|
||||
websocket_completed_response("resp_after_fallback_history_strip", 3, 1),
|
||||
)
|
||||
.await;
|
||||
(first_payload, second_payload)
|
||||
});
|
||||
let (app, state, api_key, _pool, _dir) =
|
||||
test_app_with_two_accounts_and_state_config(base_url, |config| {
|
||||
config.auth.max_concurrent_per_account = 1;
|
||||
})
|
||||
.await;
|
||||
seed_primary_affinity_and_hold_primary_slot(&state).await;
|
||||
|
||||
let response = app
|
||||
.oneshot(responses_json_request(
|
||||
&api_key,
|
||||
&json!({
|
||||
"model": "gpt-5.5",
|
||||
"input": [],
|
||||
"stream": false,
|
||||
"previous_response_id": "resp_owned_by_primary"
|
||||
}),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let body = response_json(response).await;
|
||||
let (first_payload, second_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());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_stream_history_recovery_should_stay_on_fallback_account_when_preferred_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(
|
||||
&listener,
|
||||
"Bearer access-secondary",
|
||||
websocket_completed_response("resp_stream_after_fallback_history_strip", 3, 1),
|
||||
)
|
||||
.await;
|
||||
(first_payload, second_payload)
|
||||
});
|
||||
let (app, state, api_key, _pool, _dir) =
|
||||
test_app_with_two_accounts_and_state_config(base_url, |config| {
|
||||
config.auth.max_concurrent_per_account = 1;
|
||||
})
|
||||
.await;
|
||||
seed_primary_affinity_and_hold_primary_slot(&state).await;
|
||||
|
||||
let response = app
|
||||
.oneshot(responses_json_request(
|
||||
&api_key,
|
||||
&json!({
|
||||
"model": "gpt-5.5",
|
||||
"input": [],
|
||||
"stream": true,
|
||||
"previous_response_id": "resp_owned_by_primary"
|
||||
}),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let body = response_text(response).await;
|
||||
let (first_payload, second_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());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_websocket_stream_previous_response_not_found_should_strip_history_and_retry_same_account(
|
||||
) {
|
||||
@@ -2237,3 +2364,34 @@ async fn responses_stream_with_previous_response_id_should_record_websocket_audi
|
||||
Some(44.0)
|
||||
);
|
||||
}
|
||||
|
||||
async fn seed_primary_affinity_and_hold_primary_slot(state: &AppState) {
|
||||
state
|
||||
.services
|
||||
.session_affinity
|
||||
.record(
|
||||
"resp_owned_by_primary".to_string(),
|
||||
SessionAffinityEntry {
|
||||
account_id: "acct_primary".to_string(),
|
||||
conversation_id: "conv_primary".to_string(),
|
||||
turn_state: Some("turn-primary".to_string()),
|
||||
instructions_hash: None,
|
||||
input_tokens: None,
|
||||
function_call_ids: Vec::new(),
|
||||
variant_hash: None,
|
||||
created_at: Utc::now(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let held_primary_slot = state
|
||||
.services
|
||||
.account_pool
|
||||
.acquire_with(
|
||||
&AccountAcquireRequest::new("gpt-5.5", Utc::now())
|
||||
.with_required_account_id("acct_primary"),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(held_primary_slot.account.id, "acct_primary");
|
||||
}
|
||||
|
||||
@@ -861,6 +861,82 @@ async fn responses_should_record_request_count_when_5xx_retries_are_exhausted()
|
||||
assert_eq!(metadata["accountEmail"], "user@example.com");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_should_record_attempt_trace_after_retrying_another_account() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.and(header("authorization", "Bearer access-primary"))
|
||||
.respond_with(ResponseTemplate::new(429).set_body_json(json!({
|
||||
"error": {"message": "primary rate limited"}
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.and(header("authorization", "Bearer access-secondary"))
|
||||
.respond_with(ResponseTemplate::new(500).set_body_json(json!({
|
||||
"error": {"message": "secondary failed"}
|
||||
})))
|
||||
.expect(3)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let secret_input = "do not persist this input";
|
||||
let (app, api_key, pool, _dir) = test_app_with_two_accounts_and_logging(server.uri()).await;
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/responses")
|
||||
.header("authorization", format!("Bearer {api_key}"))
|
||||
.header("content-type", "application/json")
|
||||
.header("x-request-id", "req_attempt_trace_retry")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "gpt-5.5",
|
||||
"input": secret_input,
|
||||
"stream": false,
|
||||
"use_websocket": false
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let event = latest_response_usage_record(&pool).await;
|
||||
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
|
||||
assert_eq!(event.level, "error");
|
||||
assert_eq!(event.request_id.as_deref(), Some("req_attempt_trace_retry"));
|
||||
assert_eq!(event.account_id.as_deref(), Some("acct_secondary"));
|
||||
assert_eq!(event.status_code, Some(500));
|
||||
assert_eq!(metadata["eventStatusCode"], 500);
|
||||
assert_eq!(metadata["clientStatusCode"], 502);
|
||||
assert_eq!(metadata["upstreamStatusCode"], 500);
|
||||
assert_eq!(metadata["attemptIndex"], 1);
|
||||
assert_eq!(metadata["attemptAccountId"], "acct_secondary");
|
||||
assert_eq!(metadata["attemptCount"], 2);
|
||||
assert_eq!(metadata["attempts"][0]["index"], 0);
|
||||
assert_eq!(metadata["attempts"][0]["accountId"], "acct_primary");
|
||||
assert_eq!(metadata["attempts"][1]["index"], 1);
|
||||
assert_eq!(metadata["attempts"][1]["accountId"], "acct_secondary");
|
||||
assert_eq!(metadata["requestSummary"]["inputType"], "string");
|
||||
assert!(metadata["requestSummary"]["inputItemsCount"].is_null());
|
||||
assert_eq!(metadata["requestSummary"]["transport"], "http_sse");
|
||||
assert_eq!(
|
||||
metadata["requestSummary"]["localTransport"]["useWebsocket"],
|
||||
false
|
||||
);
|
||||
assert!(
|
||||
!metadata.to_string().contains(secret_input),
|
||||
"request summary must not persist input content: {metadata:?}",
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_should_record_session_affinity_for_completed_response() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
@@ -20,7 +20,8 @@ use codex_proxy_rs::{
|
||||
use tower::util::ServiceExt;
|
||||
|
||||
use crate::support::{
|
||||
client_keys::insert_client_api_key, config::test_config, http::response_json,
|
||||
client_keys::insert_client_api_key, config::test_config, fingerprint::runtime_fingerprint,
|
||||
http::response_json,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
@@ -164,7 +165,11 @@ async fn build_test_app(
|
||||
usage_records: SqliteUsageRecordStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let services = std::sync::Arc::new(Services::new(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
services
|
||||
.initialize_hot_path_state()
|
||||
.await
|
||||
|
||||
@@ -18,7 +18,8 @@ use sqlx::SqlitePool;
|
||||
use tower::util::ServiceExt;
|
||||
|
||||
use crate::support::{
|
||||
client_keys::insert_client_api_key, config::test_config, http::response_json,
|
||||
client_keys::insert_client_api_key, config::test_config, fingerprint::runtime_fingerprint,
|
||||
http::response_json,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
@@ -59,6 +60,30 @@ async fn responses_route_should_reject_unknown_models_with_openai_error() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_route_should_accept_body_larger_than_axum_default_limit() {
|
||||
let (app, api_key, _dir) = test_app_with_client_api_key().await;
|
||||
let large_input = "x".repeat(3 * 1024 * 1024);
|
||||
let body = serde_json::json!({
|
||||
"model": "unknown-model",
|
||||
"input": large_input,
|
||||
})
|
||||
.to_string();
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/responses")
|
||||
.header("authorization", format!("Bearer {api_key}"))
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(body))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::NOT_FOUND);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn responses_route_should_terminate_generated_stream_errors_with_done_marker() {
|
||||
let (app, api_key, _dir) = test_app_with_client_api_key().await;
|
||||
@@ -130,7 +155,11 @@ async fn test_app_with_client_api_key() -> (axum::Router, String, tempfile::Temp
|
||||
usage_records: SqliteUsageRecordStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let services = std::sync::Arc::new(Services::new(
|
||||
&config,
|
||||
stores,
|
||||
runtime_fingerprint(fingerprint),
|
||||
));
|
||||
services
|
||||
.initialize_hot_path_state()
|
||||
.await
|
||||
|
||||
@@ -49,7 +49,7 @@ async fn services_try_new_should_use_configured_tls_transport_builder() {
|
||||
let Err(error) = Services::try_new(
|
||||
&config,
|
||||
stores,
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
) else {
|
||||
panic!("invalid custom CA should fail service transport construction");
|
||||
};
|
||||
@@ -90,7 +90,7 @@ async fn services_try_new_should_expose_websocket_pool_when_enabled() {
|
||||
let services = Services::try_new(
|
||||
&config,
|
||||
stores(pool),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.expect("services should build");
|
||||
|
||||
@@ -111,7 +111,7 @@ async fn services_try_new_should_not_expose_websocket_pool_when_disabled() {
|
||||
let services = Services::try_new(
|
||||
&config,
|
||||
stores(pool),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.expect("services should build");
|
||||
|
||||
|
||||
@@ -21,6 +21,7 @@ async fn scheduler_handle_should_shutdown_websocket_pool() {
|
||||
max_per_account: 1,
|
||||
enabled: true,
|
||||
ping_timeout: Duration::ZERO,
|
||||
first_token_timeout: None,
|
||||
}));
|
||||
|
||||
SchedulerHandle::from_websocket_pool(pool.clone())
|
||||
|
||||
@@ -11,9 +11,11 @@ async fn fingerprint_update_task_should_start_background_checker() {
|
||||
repo.ensure_current_seed(&crate::support::fingerprint::test_fingerprint())
|
||||
.await
|
||||
.expect("seed current fingerprint");
|
||||
let runtime_fingerprint = runtime_fingerprint_from_repo(&repo).await;
|
||||
|
||||
let handle = codex_proxy_rs::runtime::tasks::fingerprint_update::start_fingerprint_update_task(
|
||||
Some(repo),
|
||||
repo,
|
||||
runtime_fingerprint,
|
||||
"https://example.invalid/appcast.xml".to_string(),
|
||||
dir.path().join("extracted-fingerprint.json"),
|
||||
"26.800.1".to_string(),
|
||||
@@ -49,6 +51,7 @@ async fn fingerprint_update_task_should_apply_initial_appcast_update_to_reposito
|
||||
repo.ensure_current_seed(&crate::support::fingerprint::test_fingerprint())
|
||||
.await
|
||||
.expect("seed current fingerprint");
|
||||
let runtime_fingerprint = runtime_fingerprint_from_repo(&repo).await;
|
||||
let extracted_path = dir.path().join("extracted-fingerprint.json");
|
||||
std::fs::write(
|
||||
&extracted_path,
|
||||
@@ -57,7 +60,8 @@ async fn fingerprint_update_task_should_apply_initial_appcast_update_to_reposito
|
||||
.expect("extracted fingerprint should be written");
|
||||
|
||||
let handle = codex_proxy_rs::runtime::tasks::fingerprint_update::start_fingerprint_update_task(
|
||||
Some(repo.clone()),
|
||||
repo.clone(),
|
||||
runtime_fingerprint.clone(),
|
||||
format!("{}/appcast.xml", server.uri()),
|
||||
extracted_path,
|
||||
"26.800.1".to_string(),
|
||||
@@ -75,6 +79,15 @@ async fn fingerprint_update_task_should_apply_initial_appcast_update_to_reposito
|
||||
),
|
||||
("26.900.1", "7001", "147")
|
||||
);
|
||||
let runtime_snapshot = runtime_fingerprint.snapshot();
|
||||
assert_eq!(
|
||||
(
|
||||
runtime_snapshot.app_version.as_str(),
|
||||
runtime_snapshot.build_number.as_str(),
|
||||
runtime_snapshot.chromium_version.as_str()
|
||||
),
|
||||
("26.900.1", "7001", "147")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -103,8 +116,10 @@ async fn fingerprint_update_task_should_not_persist_when_appcast_matches_current
|
||||
repo.ensure_current_seed(&crate::support::fingerprint::test_fingerprint())
|
||||
.await
|
||||
.expect("seed current fingerprint");
|
||||
let runtime_fingerprint = runtime_fingerprint_from_repo(&repo).await;
|
||||
let handle = codex_proxy_rs::runtime::tasks::fingerprint_update::start_fingerprint_update_task(
|
||||
Some(repo.clone()),
|
||||
repo.clone(),
|
||||
runtime_fingerprint,
|
||||
format!("{}/appcast.xml", server.uri()),
|
||||
dir.path().join("missing-extracted-fingerprint.json"),
|
||||
"26.900.1".to_string(),
|
||||
@@ -120,33 +135,3 @@ async fn fingerprint_update_task_should_not_persist_when_appcast_matches_current
|
||||
|
||||
assert_eq!(history_count, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fingerprint_update_task_should_check_appcast_without_repository() {
|
||||
let server = MockServer::start().await;
|
||||
mount_appcast(
|
||||
&server,
|
||||
r#"
|
||||
<rss>
|
||||
<channel>
|
||||
<item>
|
||||
<enclosure url="https://example.invalid/download" sparkle:shortVersionString="26.900.1" sparkle:version="7001" />
|
||||
</item>
|
||||
</channel>
|
||||
</rss>
|
||||
"#,
|
||||
)
|
||||
.await;
|
||||
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let handle = codex_proxy_rs::runtime::tasks::fingerprint_update::start_fingerprint_update_task(
|
||||
None,
|
||||
format!("{}/appcast.xml", server.uri()),
|
||||
dir.path().join("missing-extracted-fingerprint.json"),
|
||||
"26.800.1".to_string(),
|
||||
"6001".to_string(),
|
||||
);
|
||||
|
||||
wait_for_appcast_requests(&server, 1).await;
|
||||
handle.shutdown().await;
|
||||
}
|
||||
|
||||
@@ -10,7 +10,7 @@ use codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore;
|
||||
use codex_proxy_rs::upstream::accounts::model::{Account, AccountStatus};
|
||||
use codex_proxy_rs::upstream::accounts::model::{AccountModelUsageDelta, AccountUsageDelta};
|
||||
use codex_proxy_rs::upstream::accounts::store::{AccountStore, AccountStoreResult};
|
||||
use codex_proxy_rs::upstream::fingerprint::FingerprintRepository;
|
||||
use codex_proxy_rs::upstream::fingerprint::{FingerprintRepository, RuntimeFingerprint};
|
||||
use codex_proxy_rs::upstream::models::config::ModelConfig;
|
||||
use codex_proxy_rs::upstream::models::service::ModelService;
|
||||
use wiremock::{
|
||||
@@ -137,6 +137,15 @@ async fn wait_for_current_fingerprint_version(
|
||||
}
|
||||
}
|
||||
|
||||
async fn runtime_fingerprint_from_repo(repo: &FingerprintRepository) -> RuntimeFingerprint {
|
||||
RuntimeFingerprint::new(
|
||||
repo.load_current()
|
||||
.await
|
||||
.expect("current fingerprint should load")
|
||||
.expect("current fingerprint should exist"),
|
||||
)
|
||||
}
|
||||
|
||||
async fn wait_for_appcast_requests(server: &MockServer, expected_count: usize) {
|
||||
let deadline = tokio::time::Instant::now() + StdDuration::from_secs(2);
|
||||
loop {
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use codex_proxy_rs::config::schema::{
|
||||
AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, FingerprintConfig,
|
||||
LoggingConfig, QuotaConfig, ServerConfig, TlsConfig, WebSocketPoolConfig,
|
||||
@@ -15,15 +13,14 @@ pub(crate) fn test_config(database_url: String) -> AppConfig {
|
||||
api: ApiConfig {
|
||||
base_url: "https://chatgpt.com/backend-api".to_string(),
|
||||
},
|
||||
model_aliases: BTreeMap::new(),
|
||||
model_account_routes: BTreeMap::new(),
|
||||
model_aliases: std::collections::BTreeMap::new(),
|
||||
auth: AuthConfig {
|
||||
refresh_margin_seconds: 300,
|
||||
refresh_enabled: true,
|
||||
refresh_concurrency: 2,
|
||||
max_concurrent_per_account: 3,
|
||||
request_interval_ms: 50,
|
||||
rotation_strategy: "least_used".to_string(),
|
||||
rotation_strategy: "smart".to_string(),
|
||||
tier_priority: Vec::new(),
|
||||
oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(),
|
||||
oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(),
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
use codex_proxy_rs::{config::schema::FingerprintConfig, upstream::fingerprint::Fingerprint};
|
||||
use codex_proxy_rs::{
|
||||
config::schema::FingerprintConfig,
|
||||
upstream::fingerprint::{Fingerprint, RuntimeFingerprint},
|
||||
};
|
||||
|
||||
pub(crate) fn test_fingerprint() -> Fingerprint {
|
||||
Fingerprint::from_config(&FingerprintConfig::default())
|
||||
@@ -9,3 +12,11 @@ pub(crate) fn test_fingerprint_with_updated_at(updated_at: Option<&str>) -> Fing
|
||||
fingerprint.updated_at = updated_at.map(ToString::to_string);
|
||||
fingerprint
|
||||
}
|
||||
|
||||
pub(crate) fn runtime_test_fingerprint() -> RuntimeFingerprint {
|
||||
RuntimeFingerprint::new(test_fingerprint())
|
||||
}
|
||||
|
||||
pub(crate) fn runtime_fingerprint(fingerprint: Fingerprint) -> RuntimeFingerprint {
|
||||
RuntimeFingerprint::new(fingerprint)
|
||||
}
|
||||
|
||||
@@ -153,10 +153,10 @@ fn account_pool_should_skip_accounts_in_cloudflare_cooldown() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn least_used_should_deprioritize_quota_limited_accounts_when_skip_is_disabled() {
|
||||
fn quota_reset_priority_should_deprioritize_quota_limited_accounts_when_skip_is_disabled() {
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
skip_quota_limited: false,
|
||||
rotation_strategy: RotationStrategy::LeastUsed,
|
||||
rotation_strategy: RotationStrategy::QuotaResetPriority,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
let mut limited = crate::support::accounts::test_account("limited", AccountStatus::Active);
|
||||
@@ -172,10 +172,10 @@ fn least_used_should_deprioritize_quota_limited_accounts_when_skip_is_disabled()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn least_used_should_prefer_earlier_rate_limit_window_reset() {
|
||||
fn quota_reset_priority_should_prefer_earlier_rate_limit_window_reset() {
|
||||
let now = fixed_time();
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
rotation_strategy: RotationStrategy::LeastUsed,
|
||||
rotation_strategy: RotationStrategy::QuotaResetPriority,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
let mut soon = crate::support::accounts::test_account("soon", AccountStatus::Active);
|
||||
|
||||
@@ -352,6 +352,14 @@ async fn runtime_account_pool_should_persist_quota_window_reset_when_cooldown_ex
|
||||
);
|
||||
assert!(persisted_reset_at.is_some_and(|reset_at| reset_at > now));
|
||||
assert_eq!(stored_limit_window_seconds, Some(60));
|
||||
let usage_record_count: i64 =
|
||||
sqlx::query_scalar("select count(*) from usage_records where account_id = ?")
|
||||
.bind("acct_quota_reset")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(usage_record_count, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -437,28 +445,6 @@ fn account_pool_distinct_plan_accounts_should_filter_like_model_refresh() {
|
||||
assert_eq!(pool.release("plus-ok").unwrap().model, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn account_pool_should_filter_by_model_account_routes() {
|
||||
let mut model_routes = BTreeMap::new();
|
||||
model_routes.insert("gpt-5.5".to_string(), vec!["acct_b".to_string()]);
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
max_concurrent_per_account: 1,
|
||||
model_account_routes: model_routes,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
pool.insert(crate::support::accounts::test_account(
|
||||
"acct_a",
|
||||
AccountStatus::Active,
|
||||
));
|
||||
pool.insert(crate::support::accounts::test_account(
|
||||
"acct_b",
|
||||
AccountStatus::Active,
|
||||
));
|
||||
|
||||
assert_eq!(acquire_account(&mut pool, "gpt-5.5").unwrap().id, "acct_b");
|
||||
assert!(acquire_account(&mut pool, "gpt-5.5").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn account_pool_should_exclude_requested_account_ids() {
|
||||
let mut pool = AccountPool::default();
|
||||
@@ -503,6 +489,33 @@ fn account_pool_should_prefer_session_affinity_account_when_available() {
|
||||
assert_eq!(acquired.account.id, "acct_b");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn account_pool_should_not_fallback_when_required_account_is_unavailable() {
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
max_concurrent_per_account: 1,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
pool.insert(crate::support::accounts::test_account(
|
||||
"acct_a",
|
||||
AccountStatus::Active,
|
||||
));
|
||||
pool.insert(crate::support::accounts::test_account(
|
||||
"acct_b",
|
||||
AccountStatus::Active,
|
||||
));
|
||||
|
||||
assert!(pool
|
||||
.acquire_with(
|
||||
&AccountAcquireRequest::new("gpt-5.5", fixed_time()).with_required_account_id("acct_a"),
|
||||
)
|
||||
.is_some());
|
||||
let acquired = pool.acquire_with(
|
||||
&AccountAcquireRequest::new("gpt-5.5", fixed_time()).with_required_account_id("acct_a"),
|
||||
);
|
||||
|
||||
assert!(acquired.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn account_pool_should_cleanup_stale_slots_before_acquire() {
|
||||
let now = fixed_time();
|
||||
@@ -530,10 +543,244 @@ fn account_pool_should_cleanup_stale_slots_before_acquire() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn account_pool_should_rotate_tied_least_used_accounts() {
|
||||
fn smart_should_spread_concurrent_requests_by_slot_pressure() {
|
||||
let now = fixed_time();
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
max_concurrent_per_account: 3,
|
||||
rotation_strategy: RotationStrategy::Smart,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
pool.insert(crate::support::accounts::test_account(
|
||||
"acct_a",
|
||||
AccountStatus::Active,
|
||||
));
|
||||
pool.insert(crate::support::accounts::test_account(
|
||||
"acct_b",
|
||||
AccountStatus::Active,
|
||||
));
|
||||
pool.insert(crate::support::accounts::test_account(
|
||||
"acct_c",
|
||||
AccountStatus::Active,
|
||||
));
|
||||
|
||||
let first = pool
|
||||
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now))
|
||||
.unwrap();
|
||||
let second = pool
|
||||
.acquire_with(&AccountAcquireRequest::new(
|
||||
"gpt-5.5",
|
||||
now + Duration::milliseconds(1),
|
||||
))
|
||||
.unwrap();
|
||||
let third = pool
|
||||
.acquire_with(&AccountAcquireRequest::new(
|
||||
"gpt-5.5",
|
||||
now + Duration::milliseconds(2),
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
[first.account.id, second.account.id, third.account.id]
|
||||
.into_iter()
|
||||
.collect::<BTreeSet<_>>()
|
||||
.len(),
|
||||
3
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn smart_should_balance_released_requests_by_runtime_usage() {
|
||||
let now = fixed_time();
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
rotation_strategy: RotationStrategy::Smart,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
pool.insert(crate::support::accounts::test_account(
|
||||
"acct_a",
|
||||
AccountStatus::Active,
|
||||
));
|
||||
pool.insert(crate::support::accounts::test_account(
|
||||
"acct_b",
|
||||
AccountStatus::Active,
|
||||
));
|
||||
|
||||
let first = pool
|
||||
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now))
|
||||
.unwrap();
|
||||
pool.release(&first.account.id);
|
||||
let second = pool
|
||||
.acquire_with(&AccountAcquireRequest::new(
|
||||
"gpt-5.5",
|
||||
now + Duration::seconds(1),
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_ne!(first.account.id, second.account.id);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn smart_should_not_let_reset_window_override_lower_request_count() {
|
||||
let now = fixed_time();
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
rotation_strategy: RotationStrategy::Smart,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
let mut reset_soon =
|
||||
crate::support::accounts::test_account("reset_soon", AccountStatus::Active);
|
||||
reset_soon.request_count = 10;
|
||||
reset_soon.window_request_count = 10;
|
||||
reset_soon.window_reset_at = Some(now + Duration::seconds(30));
|
||||
let mut quieter = crate::support::accounts::test_account("quieter", AccountStatus::Active);
|
||||
quieter.request_count = 2;
|
||||
quieter.window_request_count = 2;
|
||||
quieter.window_reset_at = Some(now + Duration::seconds(300));
|
||||
pool.insert(reset_soon);
|
||||
pool.insert(quieter);
|
||||
|
||||
let acquired = pool
|
||||
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(acquired.account.id, "quieter");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn smart_should_not_drain_earliest_reset_account_when_other_accounts_are_idle() {
|
||||
let now = fixed_time();
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
rotation_strategy: RotationStrategy::Smart,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
let mut reset_first =
|
||||
crate::support::accounts::test_account("reset_first", AccountStatus::Active);
|
||||
reset_first.window_reset_at = Some(now + Duration::seconds(30));
|
||||
let mut reset_later =
|
||||
crate::support::accounts::test_account("reset_later", AccountStatus::Active);
|
||||
reset_later.window_reset_at = Some(now + Duration::minutes(10));
|
||||
let mut reset_latest =
|
||||
crate::support::accounts::test_account("reset_latest", AccountStatus::Active);
|
||||
reset_latest.window_reset_at = Some(now + Duration::minutes(20));
|
||||
pool.insert(reset_first);
|
||||
pool.insert(reset_later);
|
||||
pool.insert(reset_latest);
|
||||
|
||||
let mut selected_ids = Vec::new();
|
||||
for offset in 0..4 {
|
||||
let acquired = pool
|
||||
.acquire_with(&AccountAcquireRequest::new(
|
||||
"gpt-5.5",
|
||||
now + Duration::seconds(offset),
|
||||
))
|
||||
.unwrap();
|
||||
selected_ids.push(acquired.account.id.clone());
|
||||
pool.release(&acquired.account.id);
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
selected_ids
|
||||
.iter()
|
||||
.take(3)
|
||||
.cloned()
|
||||
.collect::<BTreeSet<_>>()
|
||||
.len(),
|
||||
3
|
||||
);
|
||||
assert!(selected_ids.windows(2).all(|pair| pair[0] != pair[1]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn smart_should_use_window_token_load_when_request_counts_tie() {
|
||||
let now = fixed_time();
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
rotation_strategy: RotationStrategy::Smart,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
let mut heavy = crate::support::accounts::test_account("heavy", AccountStatus::Active);
|
||||
heavy.request_count = 5;
|
||||
heavy.window_request_count = 3;
|
||||
heavy.window_input_tokens = 1_000_000;
|
||||
heavy.window_output_tokens = 500_000;
|
||||
let mut light = crate::support::accounts::test_account("light", AccountStatus::Active);
|
||||
light.request_count = 5;
|
||||
light.window_request_count = 3;
|
||||
light.window_input_tokens = 8_000;
|
||||
light.window_output_tokens = 2_000;
|
||||
pool.insert(heavy);
|
||||
pool.insert(light);
|
||||
|
||||
let acquired = pool
|
||||
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(acquired.account.id, "light");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn smart_should_not_double_count_cached_tokens_when_request_counts_tie() {
|
||||
let now = fixed_time();
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
rotation_strategy: RotationStrategy::Smart,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
let mut cache_heavy =
|
||||
crate::support::accounts::test_account("cache_heavy", AccountStatus::Active);
|
||||
cache_heavy.request_count = 5;
|
||||
cache_heavy.window_request_count = 3;
|
||||
cache_heavy.window_input_tokens = 10_000;
|
||||
cache_heavy.window_output_tokens = 1_000;
|
||||
cache_heavy.window_cached_tokens = 100_000;
|
||||
let mut heavier_actual =
|
||||
crate::support::accounts::test_account("heavier_actual", AccountStatus::Active);
|
||||
heavier_actual.request_count = 5;
|
||||
heavier_actual.window_request_count = 3;
|
||||
heavier_actual.window_input_tokens = 20_000;
|
||||
heavier_actual.window_output_tokens = 1_000;
|
||||
pool.insert(cache_heavy);
|
||||
pool.insert(heavier_actual);
|
||||
|
||||
let acquired = pool
|
||||
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(acquired.account.id, "cache_heavy");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn smart_should_use_cached_tokens_as_tie_break_when_actual_token_load_ties() {
|
||||
let now = fixed_time();
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
rotation_strategy: RotationStrategy::Smart,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
let mut cache_heavy =
|
||||
crate::support::accounts::test_account("cache_heavy", AccountStatus::Active);
|
||||
cache_heavy.request_count = 5;
|
||||
cache_heavy.window_request_count = 3;
|
||||
cache_heavy.window_input_tokens = 10_000;
|
||||
cache_heavy.window_output_tokens = 1_000;
|
||||
cache_heavy.window_cached_tokens = 100_000;
|
||||
let mut cache_light =
|
||||
crate::support::accounts::test_account("cache_light", AccountStatus::Active);
|
||||
cache_light.request_count = 5;
|
||||
cache_light.window_request_count = 3;
|
||||
cache_light.window_input_tokens = 10_000;
|
||||
cache_light.window_output_tokens = 1_000;
|
||||
cache_light.window_cached_tokens = 1_000;
|
||||
pool.insert(cache_heavy);
|
||||
pool.insert(cache_light);
|
||||
|
||||
let acquired = pool
|
||||
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(acquired.account.id, "cache_light");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn account_pool_should_rotate_tied_quota_reset_priority_accounts() {
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
max_concurrent_per_account: 1,
|
||||
rotation_strategy: RotationStrategy::LeastUsed,
|
||||
rotation_strategy: RotationStrategy::QuotaResetPriority,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
pool.insert(crate::support::accounts::test_account(
|
||||
@@ -553,11 +800,11 @@ fn account_pool_should_rotate_tied_least_used_accounts() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn least_used_should_prefer_lru_before_tie_rotation() {
|
||||
fn quota_reset_priority_should_prefer_lru_before_tie_rotation() {
|
||||
let now = fixed_time();
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
max_concurrent_per_account: 1,
|
||||
rotation_strategy: RotationStrategy::LeastUsed,
|
||||
rotation_strategy: RotationStrategy::QuotaResetPriority,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
pool.insert(crate::support::accounts::test_account(
|
||||
@@ -590,10 +837,10 @@ fn least_used_should_prefer_lru_before_tie_rotation() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn least_used_should_prefer_lower_runtime_request_count() {
|
||||
fn quota_reset_priority_should_prefer_lower_runtime_request_count() {
|
||||
let now = fixed_time();
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
rotation_strategy: RotationStrategy::LeastUsed,
|
||||
rotation_strategy: RotationStrategy::QuotaResetPriority,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
let mut busy = crate::support::accounts::test_account("busy", AccountStatus::Active);
|
||||
@@ -611,10 +858,10 @@ fn least_used_should_prefer_lower_runtime_request_count() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn least_used_should_not_penalize_missing_window_reset() {
|
||||
fn quota_reset_priority_should_not_penalize_missing_window_reset() {
|
||||
let now = fixed_time();
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
rotation_strategy: RotationStrategy::LeastUsed,
|
||||
rotation_strategy: RotationStrategy::QuotaResetPriority,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
let mut fresh = crate::support::accounts::test_account("fresh", AccountStatus::Active);
|
||||
@@ -634,10 +881,10 @@ fn least_used_should_not_penalize_missing_window_reset() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn least_used_should_fall_through_to_request_count_without_window_resets() {
|
||||
fn quota_reset_priority_should_fall_through_to_request_count_without_window_resets() {
|
||||
let now = fixed_time();
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
rotation_strategy: RotationStrategy::LeastUsed,
|
||||
rotation_strategy: RotationStrategy::QuotaResetPriority,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
let mut busier = crate::support::accounts::test_account("busier", AccountStatus::Active);
|
||||
@@ -655,11 +902,11 @@ fn least_used_should_fall_through_to_request_count_without_window_resets() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn least_used_should_sort_quota_limited_accounts_by_reset_when_not_skipping() {
|
||||
fn quota_reset_priority_should_sort_quota_limited_accounts_by_reset_when_not_skipping() {
|
||||
let now = fixed_time();
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
skip_quota_limited: false,
|
||||
rotation_strategy: RotationStrategy::LeastUsed,
|
||||
rotation_strategy: RotationStrategy::QuotaResetPriority,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
let mut later = crate::support::accounts::test_account("later", AccountStatus::Active);
|
||||
@@ -679,10 +926,10 @@ fn least_used_should_sort_quota_limited_accounts_by_reset_when_not_skipping() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn least_used_should_not_mutate_candidate_order() {
|
||||
fn quota_reset_priority_should_not_mutate_candidate_order() {
|
||||
let now = fixed_time();
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
rotation_strategy: RotationStrategy::LeastUsed,
|
||||
rotation_strategy: RotationStrategy::QuotaResetPriority,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
let mut first = crate::support::accounts::test_account("first", AccountStatus::Active);
|
||||
|
||||
@@ -68,10 +68,10 @@ fn capacity_summary_should_not_exclude_cloudflare_cooldown_accounts() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn least_used_should_compare_request_count_when_only_one_account_has_window_reset() {
|
||||
fn quota_reset_priority_should_compare_request_count_when_only_one_account_has_window_reset() {
|
||||
let now = fixed_time();
|
||||
let mut pool = AccountPool::with_options(AccountPoolOptions {
|
||||
rotation_strategy: RotationStrategy::LeastUsed,
|
||||
rotation_strategy: RotationStrategy::QuotaResetPriority,
|
||||
..AccountPoolOptions::default()
|
||||
});
|
||||
let mut reset_known =
|
||||
|
||||
@@ -3,10 +3,12 @@ use codex_proxy_rs::{
|
||||
infra::database::connect_sqlite,
|
||||
upstream::accounts::{
|
||||
model::{AccountStatus, AccountUsageDelta},
|
||||
quota::{quota_snapshot_limit_reached, quota_snapshot_reset_at},
|
||||
store::{AccountClaimsUpdate, AccountStore, NewAccount, SqliteAccountStore},
|
||||
},
|
||||
};
|
||||
use secrecy::{ExposeSecret, SecretString};
|
||||
use serde_json::json;
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_repository_should_store_plain_tokens_and_load_secret_wrappers() {
|
||||
@@ -518,6 +520,48 @@ async fn account_repository_should_map_quota_snapshot_to_account_status() {
|
||||
assert!(restored_status.2.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_repository_should_mark_quota_exhausted_for_any_exhausted_window() {
|
||||
let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 47).await;
|
||||
let repo = SqliteAccountStore::new(pool.clone());
|
||||
seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await;
|
||||
let quota = json!({
|
||||
"plan_type": "plus",
|
||||
"snapshots": [{
|
||||
"source": "additional",
|
||||
"blocked": false,
|
||||
"allowed": true,
|
||||
"primary": {
|
||||
"used_percent": 100,
|
||||
"remaining_percent": 0,
|
||||
"reset_at": 1_893_456_300,
|
||||
"window_minutes": 300,
|
||||
"limit_reached": false
|
||||
},
|
||||
"secondary": null
|
||||
}]
|
||||
});
|
||||
let limit_reached = quota_snapshot_limit_reached("a);
|
||||
let cooldown_until = quota_snapshot_reset_at("a);
|
||||
|
||||
let updated = repo
|
||||
.apply_quota_snapshot("acct_a", "a.to_string(), limit_reached, cooldown_until)
|
||||
.await
|
||||
.unwrap();
|
||||
let status: (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!(
|
||||
(updated, status.0, status.1, status.2.is_some()),
|
||||
(true, "quota_exhausted".to_string(), 1, true)
|
||||
);
|
||||
}
|
||||
|
||||
#[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;
|
||||
|
||||
@@ -227,6 +227,53 @@ fn quota_snapshot_limit_reached_should_block_explicit_credit_overage_limit() {
|
||||
assert!(quota_snapshot_limit_reached("a));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quota_snapshot_limit_reached_should_block_any_exhausted_window() {
|
||||
let quota = json!({
|
||||
"snapshots": [{
|
||||
"source": "additional",
|
||||
"blocked": false,
|
||||
"allowed": true,
|
||||
"primary": {
|
||||
"used_percent": 100,
|
||||
"remaining_percent": 0,
|
||||
"reset_at": 1_893_456_300,
|
||||
"window_minutes": 300,
|
||||
"limit_reached": false
|
||||
},
|
||||
"secondary": null
|
||||
}]
|
||||
});
|
||||
|
||||
assert!(quota_snapshot_limit_reached("a));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quota_snapshot_window_should_follow_blocking_window() {
|
||||
let quota = json!({
|
||||
"snapshots": [{
|
||||
"source": "additional",
|
||||
"blocked": false,
|
||||
"primary": {
|
||||
"used_percent": 100,
|
||||
"remaining_percent": 0,
|
||||
"reset_at": 1_893_456_300,
|
||||
"window_minutes": 10_080,
|
||||
"limit_reached": false
|
||||
},
|
||||
"secondary": null
|
||||
}]
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
(
|
||||
quota_snapshot_reset_at("a).map(|reset_at| reset_at.timestamp()),
|
||||
quota_snapshot_limit_window_seconds("a),
|
||||
),
|
||||
(Some(1_893_456_300), Some(604_800))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_refresh_service_should_send_usage_cookie_when_cookie_store_is_configured() {
|
||||
let server = MockServer::start().await;
|
||||
@@ -265,7 +312,7 @@ async fn quota_refresh_service_should_send_usage_cookie_when_cookie_store_is_con
|
||||
let codex = CodexBackendClient::new(
|
||||
reqwest::Client::new(),
|
||||
server.uri(),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
let service =
|
||||
RuntimeQuotaRefreshService::new(store, Arc::new(codex)).with_cookie_store(cookies);
|
||||
@@ -346,7 +393,7 @@ async fn quota_refresh_service_should_fetch_usage_for_quota_locked_accounts_and_
|
||||
let codex = CodexBackendClient::new(
|
||||
reqwest::Client::new(),
|
||||
server.uri(),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
let service = RuntimeQuotaRefreshService::new(store.clone(), Arc::new(codex));
|
||||
|
||||
@@ -421,7 +468,7 @@ async fn quota_refresh_service_should_refresh_quota_exhausted_accounts() {
|
||||
let codex = CodexBackendClient::new(
|
||||
reqwest::Client::new(),
|
||||
server.uri(),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
let service = RuntimeQuotaRefreshService::new(store, Arc::new(codex));
|
||||
|
||||
@@ -441,6 +488,145 @@ async fn quota_refresh_service_should_refresh_quota_exhausted_accounts() {
|
||||
assert_eq!(stored, ("active".to_string(), 0));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_refresh_service_should_skip_recent_locked_account_before_cooldown_expires() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let db = dir
|
||||
.path()
|
||||
.join("quota-refresh-recent-future-cooldown.sqlite");
|
||||
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
|
||||
.await
|
||||
.expect("sqlite pool");
|
||||
let store = SqliteAccountStore::new(pool.clone());
|
||||
insert_quota_locked_account(&store, &pool, "acct-quota", "access-token").await;
|
||||
set_quota_cooldown_until(&pool, "acct-quota", Utc::now() + Duration::minutes(5)).await;
|
||||
let codex = CodexBackendClient::new(
|
||||
reqwest::Client::new(),
|
||||
server.uri(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
let service =
|
||||
RuntimeQuotaRefreshService::with_min_refresh_interval_secs(store, Arc::new(codex), 1_800);
|
||||
|
||||
let mut last_refreshed = HashMap::new();
|
||||
last_refreshed.insert("acct-quota".to_string(), tokio::time::Instant::now());
|
||||
let summary = service
|
||||
.refresh_locked_accounts(&mut last_refreshed)
|
||||
.await
|
||||
.expect("quota refresh should succeed");
|
||||
|
||||
assert_eq!(
|
||||
(
|
||||
summary.candidates,
|
||||
summary.skipped_cooldown,
|
||||
summary.skipped_recent,
|
||||
summary.refreshed,
|
||||
usage_request_count(&server).await,
|
||||
),
|
||||
(1, 1, 0, 0, 0)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_refresh_service_should_skip_recent_locked_account_inside_cooldown_grace() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let db = dir
|
||||
.path()
|
||||
.join("quota-refresh-recent-cooldown-grace.sqlite");
|
||||
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
|
||||
.await
|
||||
.expect("sqlite pool");
|
||||
let store = SqliteAccountStore::new(pool.clone());
|
||||
insert_quota_locked_account(&store, &pool, "acct-quota", "access-token").await;
|
||||
set_quota_cooldown_until(&pool, "acct-quota", Utc::now() - Duration::seconds(20)).await;
|
||||
let codex = CodexBackendClient::new(
|
||||
reqwest::Client::new(),
|
||||
server.uri(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
let service =
|
||||
RuntimeQuotaRefreshService::with_min_refresh_interval_secs(store, Arc::new(codex), 1_800);
|
||||
|
||||
let mut last_refreshed = HashMap::new();
|
||||
last_refreshed.insert("acct-quota".to_string(), tokio::time::Instant::now());
|
||||
let summary = service
|
||||
.refresh_locked_accounts(&mut last_refreshed)
|
||||
.await
|
||||
.expect("quota refresh should succeed");
|
||||
|
||||
assert_eq!(
|
||||
(
|
||||
summary.candidates,
|
||||
summary.skipped_cooldown,
|
||||
summary.skipped_recent,
|
||||
summary.refreshed,
|
||||
usage_request_count(&server).await,
|
||||
),
|
||||
(1, 1, 0, 0, 0)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_refresh_service_should_bypass_recent_skip_after_cooldown_grace() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/codex/usage"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"plan_type": "plus",
|
||||
"rate_limit": {
|
||||
"allowed": true,
|
||||
"limit_reached": false,
|
||||
"primary_window": {
|
||||
"used_percent": 8,
|
||||
"reset_at": 1_800_000_000,
|
||||
"limit_window_seconds": 18_000
|
||||
}
|
||||
}
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let db = dir
|
||||
.path()
|
||||
.join("quota-refresh-recent-expired-cooldown.sqlite");
|
||||
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
|
||||
.await
|
||||
.expect("sqlite pool");
|
||||
let store = SqliteAccountStore::new(pool.clone());
|
||||
insert_quota_locked_account(&store, &pool, "acct-quota", "access-token").await;
|
||||
set_quota_cooldown_until(&pool, "acct-quota", Utc::now() - Duration::seconds(31)).await;
|
||||
let codex = CodexBackendClient::new(
|
||||
reqwest::Client::new(),
|
||||
server.uri(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
let service =
|
||||
RuntimeQuotaRefreshService::with_min_refresh_interval_secs(store, Arc::new(codex), 1_800);
|
||||
|
||||
let mut last_refreshed = HashMap::new();
|
||||
last_refreshed.insert("acct-quota".to_string(), tokio::time::Instant::now());
|
||||
let summary = service
|
||||
.refresh_locked_accounts(&mut last_refreshed)
|
||||
.await
|
||||
.expect("quota refresh should succeed");
|
||||
|
||||
assert_eq!(
|
||||
(
|
||||
summary.skipped_recent,
|
||||
summary.skipped_cooldown,
|
||||
summary.refreshed,
|
||||
usage_request_count(&server).await,
|
||||
),
|
||||
(0, 0, 1, 1)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_refresh_service_should_stagger_multiple_locked_account_requests() {
|
||||
let server = MockServer::start().await;
|
||||
@@ -472,7 +658,7 @@ async fn quota_refresh_service_should_stagger_multiple_locked_account_requests()
|
||||
let codex = CodexBackendClient::new(
|
||||
reqwest::Client::new(),
|
||||
server.uri(),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
let service = RuntimeQuotaRefreshService::new(store, Arc::new(codex))
|
||||
.with_request_spacing(StdDuration::from_millis(200));
|
||||
@@ -526,6 +712,19 @@ async fn insert_quota_locked_account(
|
||||
.expect("quota lock should be set");
|
||||
}
|
||||
|
||||
async fn set_quota_cooldown_until(
|
||||
pool: &sqlx::SqlitePool,
|
||||
id: &str,
|
||||
cooldown_until: chrono::DateTime<Utc>,
|
||||
) {
|
||||
sqlx::query("update accounts set quota_cooldown_until = ? where id = ?")
|
||||
.bind(cooldown_until.to_rfc3339())
|
||||
.bind(id)
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("quota cooldown should be set");
|
||||
}
|
||||
|
||||
async fn wait_for_usage_requests(server: &MockServer, expected_count: usize) {
|
||||
let deadline = tokio::time::Instant::now() + StdDuration::from_secs(2);
|
||||
loop {
|
||||
|
||||
@@ -57,14 +57,21 @@ async fn update_checker_should_report_available_update_from_appcast() {
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let db = dir.path().join("fingerprints-check.sqlite");
|
||||
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
|
||||
.await
|
||||
.expect("sqlite pool");
|
||||
let repo = FingerprintRepository::new(pool);
|
||||
repo.ensure_current_seed(&crate::support::fingerprint::test_fingerprint())
|
||||
.await
|
||||
.expect("seed current fingerprint");
|
||||
|
||||
let checker = UpdateChecker::with_client(
|
||||
None,
|
||||
repo,
|
||||
reqwest::Client::new(),
|
||||
format!("{}/appcast.xml", server.uri()),
|
||||
tempfile::tempdir()
|
||||
.expect("temp dir")
|
||||
.path()
|
||||
.join("extracted-fingerprint.json"),
|
||||
dir.path().join("extracted-fingerprint.json"),
|
||||
"26.800.1",
|
||||
"6001",
|
||||
);
|
||||
@@ -113,7 +120,7 @@ async fn update_checker_should_apply_available_update_to_repository() {
|
||||
.expect("write extracted fingerprint");
|
||||
|
||||
let checker = UpdateChecker::with_client(
|
||||
Some(repo.clone()),
|
||||
repo.clone(),
|
||||
reqwest::Client::new(),
|
||||
format!("{}/appcast.xml", server.uri()),
|
||||
extracted_path,
|
||||
@@ -121,17 +128,20 @@ async fn update_checker_should_apply_available_update_to_repository() {
|
||||
"6001",
|
||||
);
|
||||
|
||||
let applied = checker
|
||||
let updated = checker
|
||||
.check_and_apply_update()
|
||||
.await
|
||||
.expect("apply update");
|
||||
.expect("apply update")
|
||||
.expect("available update should return current fingerprint");
|
||||
let stored = repo
|
||||
.load_current()
|
||||
.await
|
||||
.expect("load current")
|
||||
.expect("stored fingerprint");
|
||||
|
||||
assert!(applied);
|
||||
assert_eq!(updated.app_version, "26.900.1");
|
||||
assert_eq!(updated.build_number, "7001");
|
||||
assert_eq!(updated.chromium_version, "147");
|
||||
assert_eq!(stored.app_version, "26.900.1");
|
||||
assert_eq!(stored.build_number, "7001");
|
||||
assert_eq!(stored.chromium_version, "147");
|
||||
|
||||
@@ -11,7 +11,6 @@ fn codex_websocket_payload_audit_snapshot_should_redact_user_content() {
|
||||
})],
|
||||
);
|
||||
request.set_previous_response_id(Some("resp_secret".to_string()));
|
||||
request.set_reasoning(Some(json!({"effort": "medium"})));
|
||||
request.set_service_tier(Some("flex".to_string()));
|
||||
request.set_prompt_cache_key(Some("cache-secret".to_string()));
|
||||
request.set_client_metadata(Some(json!({
|
||||
@@ -33,7 +32,6 @@ fn codex_websocket_payload_audit_snapshot_should_redact_user_content() {
|
||||
"stream",
|
||||
"store",
|
||||
"previous_response_id",
|
||||
"reasoning",
|
||||
"service_tier",
|
||||
"prompt_cache_key",
|
||||
"client_metadata",
|
||||
|
||||
@@ -344,6 +344,26 @@ fn rate_limit_quota_should_preserve_existing_monthly_and_credits_when_passive_da
|
||||
assert_eq!(quota["credits"]["balance"], 12);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rate_limit_quota_should_block_when_window_is_exhausted_even_if_flag_is_false() {
|
||||
let event = json!({
|
||||
"type": "codex.rate_limits",
|
||||
"rate_limits": {
|
||||
"limit_reached": false,
|
||||
"primary": {
|
||||
"used_percent": 100,
|
||||
"window_minutes": 300,
|
||||
"reset_at": 1893456300
|
||||
}
|
||||
}
|
||||
});
|
||||
let parsed = parse_rate_limits_event(&event).expect("event should parse");
|
||||
|
||||
let quota = rate_limit_quota(&parsed, Some("plus"), None);
|
||||
|
||||
assert_eq!(quota["snapshots"][0]["blocked"], true);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_scheduler_should_refresh_before_expiry_and_preserve_refresh_token() {
|
||||
use async_trait::async_trait;
|
||||
|
||||
@@ -59,7 +59,7 @@ async fn codex_backend_client_should_apply_configured_websocket_pool() {
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::new(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::new(pool));
|
||||
let mut request = CodexResponsesRequest::new_http_sse("gpt-5.5", "be brief", Vec::new());
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
use codex_proxy_rs::upstream::transport::CodexUpstreamDiagnostics;
|
||||
use reqwest::header::{HeaderMap, HeaderValue};
|
||||
|
||||
#[test]
|
||||
fn diagnostics_should_extract_request_id_and_trace_headers() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert("x-request-id", HeaderValue::from_static("req_1"));
|
||||
headers.insert("cf-ray", HeaderValue::from_static("ray_1"));
|
||||
|
||||
let diagnostics = CodexUpstreamDiagnostics::from_headers(Some(429), &headers);
|
||||
|
||||
assert_eq!(diagnostics.status_code, Some(429));
|
||||
assert_eq!(diagnostics.request_id.as_deref(), Some("req_1"));
|
||||
assert_eq!(diagnostics.cf_ray(), Some("ray_1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn diagnostics_should_not_treat_cf_ray_as_request_id() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert("cf-ray", HeaderValue::from_static("ray_1"));
|
||||
|
||||
let diagnostics = CodexUpstreamDiagnostics::from_headers(Some(403), &headers);
|
||||
|
||||
assert_eq!(diagnostics.request_id, None);
|
||||
assert_eq!(diagnostics.cf_ray(), Some("ray_1"));
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
use super::*;
|
||||
use codex_proxy_rs::upstream::fingerprint::RuntimeFingerprint;
|
||||
use codex_proxy_rs::upstream::transport::build_ordered_codex_headers;
|
||||
use codex_proxy_rs::upstream::transport::websocket::CodexWebSocketConnection;
|
||||
use serde_json::Value;
|
||||
@@ -279,7 +280,7 @@ async fn codex_backend_client_websocket_should_forward_security_chain_headers_an
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(pool);
|
||||
|
||||
@@ -420,7 +421,7 @@ async fn codex_backend_client_should_send_desktop_headers_and_capture_response_m
|
||||
let client = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
server.uri(),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
let mut request =
|
||||
codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse(
|
||||
@@ -471,6 +472,91 @@ async fn codex_backend_client_should_send_desktop_headers_and_capture_response_m
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn codex_backend_client_should_use_latest_runtime_fingerprint_for_new_requests() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
let sse_body = include_str!("../../fixtures/responses/http_sse/completed_usage_basic.sse");
|
||||
wiremock::Mock::given(wiremock::matchers::method("POST"))
|
||||
.and(wiremock::matchers::path("/codex/responses"))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
.set_body_string(sse_body),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let mut initial = crate::support::fingerprint::test_fingerprint();
|
||||
initial.app_version = "26.800.1".to_string();
|
||||
let runtime_fingerprint = RuntimeFingerprint::new(initial);
|
||||
let client = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
server.uri(),
|
||||
runtime_fingerprint.clone(),
|
||||
);
|
||||
let mut request =
|
||||
codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse(
|
||||
"gpt-5.5",
|
||||
"",
|
||||
Vec::new(),
|
||||
);
|
||||
request.force_http_sse = true;
|
||||
let request_context = |request_id| CodexRequestContext {
|
||||
access_token: "access-token",
|
||||
account_id: Some("chatgpt-account"),
|
||||
request_id,
|
||||
turn_state: None,
|
||||
turn_metadata: None,
|
||||
beta_features: None,
|
||||
include_timing_metrics: None,
|
||||
version: None,
|
||||
codex_window_id: None,
|
||||
parent_thread_id: None,
|
||||
cookie_header: None,
|
||||
installation_id: None,
|
||||
session_id: None,
|
||||
};
|
||||
|
||||
client
|
||||
.create_response(&request, request_context("req_fingerprint_1"))
|
||||
.await
|
||||
.expect("first response should succeed");
|
||||
|
||||
let mut updated = runtime_fingerprint.snapshot();
|
||||
updated.app_version = "26.900.1".to_string();
|
||||
updated.build_number = "7001".to_string();
|
||||
runtime_fingerprint.replace(updated);
|
||||
|
||||
client
|
||||
.create_response(&request, request_context("req_fingerprint_2"))
|
||||
.await
|
||||
.expect("second response should succeed");
|
||||
|
||||
let user_agents = server
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("received requests should load")
|
||||
.iter()
|
||||
.map(|request| {
|
||||
request
|
||||
.headers
|
||||
.get("user-agent")
|
||||
.expect("user-agent should be sent")
|
||||
.to_str()
|
||||
.expect("user-agent should be valid")
|
||||
.to_string()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(
|
||||
user_agents,
|
||||
vec![
|
||||
"Codex Desktop/26.800.1 (darwin; arm64)".to_string(),
|
||||
"Codex Desktop/26.900.1 (darwin; arm64)".to_string(),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn codex_backend_client_usage_should_use_wham_usage_headers() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
@@ -486,7 +572,7 @@ async fn codex_backend_client_usage_should_use_wham_usage_headers() {
|
||||
let client = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
server.uri(),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
|
||||
let usage = client
|
||||
@@ -571,7 +657,7 @@ async fn codex_backend_client_models_should_use_original_auxiliary_headers() {
|
||||
let client = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
server.uri(),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
|
||||
let models = client
|
||||
@@ -638,7 +724,7 @@ async fn codex_backend_client_should_send_http_sse_headers_in_fingerprint_order(
|
||||
let client = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
|
||||
client
|
||||
@@ -713,7 +799,7 @@ async fn codex_backend_client_should_send_compact_headers_in_fingerprint_order()
|
||||
let client = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
|
||||
client
|
||||
|
||||
@@ -107,7 +107,7 @@ async fn codex_backend_client_should_cap_non_success_error_body_at_one_mib() {
|
||||
let client = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
server.uri(),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
let mut request =
|
||||
codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse(
|
||||
@@ -161,7 +161,7 @@ async fn codex_backend_client_should_parse_retry_after_from_rate_limit_error_bod
|
||||
let client = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
server.uri(),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
let mut request =
|
||||
codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse(
|
||||
|
||||
@@ -37,6 +37,7 @@ use tokio_tungstenite::{
|
||||
use crate::support::assertions::assert_substrings_appear_in_order;
|
||||
|
||||
pub(crate) mod client;
|
||||
pub(crate) mod diagnostics;
|
||||
pub(crate) mod headers;
|
||||
pub(crate) mod http_client;
|
||||
pub(crate) mod websocket;
|
||||
@@ -246,6 +247,7 @@ pub(crate) fn websocket_pool_config_for_tests(
|
||||
ping_interval,
|
||||
ping_timeout: Duration::from_secs(1),
|
||||
liveness_timeout,
|
||||
first_token_timeout: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -107,7 +107,7 @@ async fn codex_backend_client_should_decode_permessage_deflate_context_takeover_
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(pool);
|
||||
let mut request =
|
||||
@@ -322,18 +322,12 @@ async fn websocket_execute_response_create_request_should_surface_response_faile
|
||||
.expect_err("response.failed should be surfaced as upstream error");
|
||||
server.await.unwrap();
|
||||
|
||||
let CodexWebSocketExchangeError::Upstream {
|
||||
status_code,
|
||||
retry_after_seconds,
|
||||
body,
|
||||
..
|
||||
} = error
|
||||
else {
|
||||
let CodexWebSocketExchangeError::Upstream(error) = error else {
|
||||
panic!("expected upstream websocket error");
|
||||
};
|
||||
assert_eq!(status_code, 429);
|
||||
assert_eq!(retry_after_seconds, Some(12));
|
||||
assert!(body.contains("rate_limit_exceeded"));
|
||||
assert_eq!(error.status_code, 429);
|
||||
assert_eq!(error.retry_after_seconds, Some(12));
|
||||
assert!(error.body.contains("rate_limit_exceeded"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -405,18 +399,12 @@ async fn websocket_execute_response_create_request_should_preserve_opening_error
|
||||
.expect_err("failed opening should surface upstream status");
|
||||
server.await.unwrap();
|
||||
|
||||
let CodexWebSocketExchangeError::Upstream {
|
||||
status_code,
|
||||
retry_after_seconds,
|
||||
body,
|
||||
..
|
||||
} = error
|
||||
else {
|
||||
let CodexWebSocketExchangeError::Upstream(error) = error else {
|
||||
panic!("expected upstream opening error");
|
||||
};
|
||||
assert_eq!(status_code, 429);
|
||||
assert_eq!(retry_after_seconds, Some(33));
|
||||
assert!(body.contains("rate limited"));
|
||||
assert_eq!(error.status_code, 429);
|
||||
assert_eq!(error.retry_after_seconds, Some(33));
|
||||
assert!(error.body.contains("rate limited"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -482,18 +470,12 @@ async fn websocket_execute_response_create_request_should_surface_wrapped_error_
|
||||
.expect_err("wrapped error should surface as upstream error");
|
||||
server.await.unwrap();
|
||||
|
||||
let CodexWebSocketExchangeError::Upstream {
|
||||
status_code,
|
||||
retry_after_seconds,
|
||||
body,
|
||||
..
|
||||
} = error
|
||||
else {
|
||||
let CodexWebSocketExchangeError::Upstream(error) = error else {
|
||||
panic!("expected wrapped upstream error");
|
||||
};
|
||||
assert_eq!(status_code, 409);
|
||||
assert_eq!(retry_after_seconds, Some(17));
|
||||
assert!(body.contains("wrapped conflict"));
|
||||
assert_eq!(error.status_code, 409);
|
||||
assert_eq!(error.retry_after_seconds, Some(17));
|
||||
assert!(error.body.contains("wrapped conflict"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -530,14 +512,11 @@ async fn websocket_execute_response_create_request_should_surface_connection_lim
|
||||
.expect_err("connection limit should surface as upstream error");
|
||||
server.await.unwrap();
|
||||
|
||||
let CodexWebSocketExchangeError::Upstream {
|
||||
status_code, body, ..
|
||||
} = error
|
||||
else {
|
||||
let CodexWebSocketExchangeError::Upstream(error) = error else {
|
||||
panic!("expected connection limit upstream error");
|
||||
};
|
||||
assert_eq!(status_code, 503);
|
||||
assert!(body.contains("websocket_connection_limit_reached"));
|
||||
assert_eq!(error.status_code, 503);
|
||||
assert!(error.body.contains("websocket_connection_limit_reached"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -1021,22 +1000,18 @@ async fn codex_backend_client_stream_should_reject_binary_websocket_event() {
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
|
||||
let mut response = backend
|
||||
let Err(error) = backend
|
||||
.create_response_stream(
|
||||
&request,
|
||||
request_context("req_stream_binary", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("websocket stream should open");
|
||||
let error = response
|
||||
.body
|
||||
.next()
|
||||
.await
|
||||
.expect("stream should yield binary error")
|
||||
.expect_err("binary event should be rejected");
|
||||
else {
|
||||
panic!("binary event before first content should reject stream open");
|
||||
};
|
||||
server.await.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
@@ -1045,6 +1020,87 @@ async fn codex_backend_client_stream_should_reject_binary_websocket_event() {
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn codex_backend_client_stream_should_retry_without_pool_when_first_content_times_out() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let accepted_connections = Arc::new(AtomicUsize::new(0));
|
||||
let accepted_connections_for_server = Arc::clone(&accepted_connections);
|
||||
let server = tokio::spawn(async move {
|
||||
let (first_stream, _) = listener.accept().await.unwrap();
|
||||
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
|
||||
let mut first_websocket = accept_codex_test_websocket(first_stream).await;
|
||||
let _first_message = first_websocket.next().await.unwrap().unwrap();
|
||||
first_websocket
|
||||
.send(Message::Text(
|
||||
json!({
|
||||
"type": "response.created",
|
||||
"response": {
|
||||
"id": "resp_no_pool_first_token_stalled",
|
||||
"object": "response"
|
||||
}
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let (second_stream, _) = listener.accept().await.unwrap();
|
||||
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
|
||||
let mut second_websocket = accept_codex_test_websocket(second_stream).await;
|
||||
let _second_message = second_websocket.next().await.unwrap().unwrap();
|
||||
second_websocket
|
||||
.send(Message::Text(
|
||||
json!({
|
||||
"type": "response.output_text.delta",
|
||||
"delta": "no-pool retry recovered"
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
second_websocket
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_no_pool_retry_recovered", 3, 1).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
second_websocket.close(None).await.unwrap();
|
||||
drop(first_websocket);
|
||||
});
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_first_token_timeout(Some(Duration::from_millis(30)));
|
||||
let request = pooled_websocket_request("conversation-first-token-no-pool-timeout");
|
||||
|
||||
let response = backend
|
||||
.create_response_stream(
|
||||
&request,
|
||||
request_context("req_first_token_no_pool_timeout", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("no-pool first-token timeout should retry before returning stream");
|
||||
let decision = response.websocket_pool_decision;
|
||||
let mut stream = response.body;
|
||||
let mut body = String::new();
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk.expect("retried no-pool stream chunk should be valid");
|
||||
body.push_str(std::str::from_utf8(&chunk).unwrap());
|
||||
}
|
||||
server.await.unwrap();
|
||||
|
||||
assert!(body.contains("no-pool retry recovered"));
|
||||
assert!(body.contains("resp_no_pool_retry_recovered"));
|
||||
assert!(!body.contains("resp_no_pool_first_token_stalled"));
|
||||
assert!(decision.is_none());
|
||||
assert_eq!(accepted_connections.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn codex_backend_client_stream_should_error_when_websocket_closes_before_terminal() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
@@ -1077,7 +1133,7 @@ async fn codex_backend_client_stream_should_error_when_websocket_closes_before_t
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
|
||||
let mut response = backend
|
||||
@@ -1147,7 +1203,7 @@ async fn codex_backend_client_stream_should_wait_for_terminal_after_active_webso
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
|
||||
let mut response = backend
|
||||
@@ -1211,7 +1267,7 @@ async fn codex_backend_client_stream_should_timeout_when_active_websocket_stalls
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
);
|
||||
|
||||
let mut response = backend
|
||||
@@ -1307,7 +1363,7 @@ async fn codex_backend_client_should_use_websocket_when_previous_response_id_is_
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(pool);
|
||||
|
||||
|
||||
@@ -39,7 +39,7 @@ async fn codex_backend_client_should_reuse_pooled_websocket_for_same_account_and
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool));
|
||||
let mut request =
|
||||
@@ -117,7 +117,7 @@ async fn codex_backend_client_should_open_fresh_socket_when_idle_pooled_websocke
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool));
|
||||
let mut request =
|
||||
@@ -155,6 +155,280 @@ async fn codex_backend_client_should_open_fresh_socket_when_idle_pooled_websocke
|
||||
assert_eq!(accepted_connections.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn codex_backend_client_stream_should_retry_fresh_socket_when_first_content_times_out() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let accepted_connections = Arc::new(AtomicUsize::new(0));
|
||||
let accepted_connections_for_server = Arc::clone(&accepted_connections);
|
||||
let server = tokio::spawn(async move {
|
||||
let (first_stream, _) = listener.accept().await.unwrap();
|
||||
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
|
||||
let mut first_websocket = accept_codex_test_websocket(first_stream).await;
|
||||
let _first_message = first_websocket.next().await.unwrap().unwrap();
|
||||
first_websocket
|
||||
.send(Message::Text(
|
||||
json!({
|
||||
"type": "response.created",
|
||||
"response": {
|
||||
"id": "resp_first_token_fresh_stalled",
|
||||
"object": "response"
|
||||
}
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let (second_stream, _) = listener.accept().await.unwrap();
|
||||
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
|
||||
let mut second_websocket = accept_codex_test_websocket(second_stream).await;
|
||||
let _second_message = second_websocket.next().await.unwrap().unwrap();
|
||||
second_websocket
|
||||
.send(Message::Text(
|
||||
json!({
|
||||
"type": "response.output_text.delta",
|
||||
"delta": "fresh retry recovered"
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
second_websocket
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_first_token_fresh_recovered", 3, 1).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
second_websocket.close(None).await.unwrap();
|
||||
drop(first_websocket);
|
||||
});
|
||||
let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig {
|
||||
first_token_timeout: Some(Duration::from_millis(30)),
|
||||
..websocket_pool_config_for_tests(None, None, None)
|
||||
}));
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool));
|
||||
let request = pooled_websocket_request("conversation-first-token-fresh-timeout");
|
||||
|
||||
let response = backend
|
||||
.create_response_stream(
|
||||
&request,
|
||||
request_context("req_first_token_fresh_timeout", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("fresh first-token timeout should retry before returning stream");
|
||||
let decision = response.websocket_pool_decision;
|
||||
let mut stream = response.body;
|
||||
let mut body = String::new();
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk.expect("retried stream chunk should be valid");
|
||||
body.push_str(std::str::from_utf8(&chunk).unwrap());
|
||||
}
|
||||
server.await.unwrap();
|
||||
|
||||
assert!(body.contains("fresh retry recovered"));
|
||||
assert!(body.contains("resp_first_token_fresh_recovered"));
|
||||
assert!(!body.contains("resp_first_token_fresh_stalled"));
|
||||
assert_eq!(decision.unwrap().kind(), "new");
|
||||
assert_eq!(accepted_connections.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn codex_backend_client_stream_should_retry_reused_socket_when_first_content_times_out() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let accepted_connections = Arc::new(AtomicUsize::new(0));
|
||||
let accepted_connections_for_server = Arc::clone(&accepted_connections);
|
||||
let server = tokio::spawn(async move {
|
||||
let (first_stream, _) = listener.accept().await.unwrap();
|
||||
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
|
||||
let mut first_websocket = accept_codex_test_websocket(first_stream).await;
|
||||
let _seed_message = first_websocket.next().await.unwrap().unwrap();
|
||||
first_websocket
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_first_token_reuse_seed", 2, 1).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let _reused_message = first_websocket.next().await.unwrap().unwrap();
|
||||
first_websocket
|
||||
.send(Message::Text(
|
||||
json!({
|
||||
"type": "response.in_progress",
|
||||
"response": {
|
||||
"id": "resp_first_token_reuse_stalled",
|
||||
"object": "response"
|
||||
}
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let (second_stream, _) = listener.accept().await.unwrap();
|
||||
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
|
||||
let mut second_websocket = accept_codex_test_websocket(second_stream).await;
|
||||
let _retry_message = second_websocket.next().await.unwrap().unwrap();
|
||||
second_websocket
|
||||
.send(Message::Text(
|
||||
json!({
|
||||
"type": "response.output_text.delta",
|
||||
"delta": "reuse retry recovered"
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
second_websocket
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_first_token_reuse_recovered", 3, 1).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
second_websocket.close(None).await.unwrap();
|
||||
drop(first_websocket);
|
||||
});
|
||||
let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig {
|
||||
first_token_timeout: Some(Duration::from_millis(30)),
|
||||
..websocket_pool_config_for_tests(None, None, None)
|
||||
}));
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool));
|
||||
let request = pooled_websocket_request("conversation-first-token-reuse-timeout");
|
||||
|
||||
let seed = backend
|
||||
.create_response(
|
||||
&request,
|
||||
request_context("req_first_token_reuse_seed", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("seed response should populate pool");
|
||||
let response = backend
|
||||
.create_response_stream(
|
||||
&request,
|
||||
request_context("req_first_token_reuse_timeout", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("reused first-token timeout should retry before returning stream");
|
||||
let decision = response.websocket_pool_decision;
|
||||
let mut stream = response.body;
|
||||
let mut body = String::new();
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk.expect("retried reused stream chunk should be valid");
|
||||
body.push_str(std::str::from_utf8(&chunk).unwrap());
|
||||
}
|
||||
server.await.unwrap();
|
||||
|
||||
assert!(seed.body.contains("resp_first_token_reuse_seed"));
|
||||
assert!(body.contains("reuse retry recovered"));
|
||||
assert!(body.contains("resp_first_token_reuse_recovered"));
|
||||
assert!(!body.contains("resp_first_token_reuse_stalled"));
|
||||
assert_eq!(decision.unwrap().kind(), "retry_after_stale_reuse");
|
||||
assert_eq!(accepted_connections.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn codex_backend_client_stream_should_retry_bypass_socket_when_first_content_times_out() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let accepted_connections = Arc::new(AtomicUsize::new(0));
|
||||
let accepted_connections_for_server = Arc::clone(&accepted_connections);
|
||||
let server = tokio::spawn(async move {
|
||||
let (first_stream, _) = listener.accept().await.unwrap();
|
||||
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
|
||||
let mut first_websocket = accept_codex_test_websocket(first_stream).await;
|
||||
let _first_message = first_websocket.next().await.unwrap().unwrap();
|
||||
first_websocket
|
||||
.send(Message::Text(
|
||||
json!({
|
||||
"type": "response.created",
|
||||
"response": {
|
||||
"id": "resp_bypass_first_token_stalled",
|
||||
"object": "response"
|
||||
}
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let (second_stream, _) = listener.accept().await.unwrap();
|
||||
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
|
||||
let mut second_websocket = accept_codex_test_websocket(second_stream).await;
|
||||
let _second_message = second_websocket.next().await.unwrap().unwrap();
|
||||
second_websocket
|
||||
.send(Message::Text(
|
||||
json!({
|
||||
"type": "response.output_text.delta",
|
||||
"delta": "bypass retry recovered"
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
second_websocket
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_bypass_retry_recovered", 3, 1).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
second_websocket.close(None).await.unwrap();
|
||||
drop(first_websocket);
|
||||
});
|
||||
let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig {
|
||||
max_per_account: 0,
|
||||
first_token_timeout: Some(Duration::from_millis(30)),
|
||||
..websocket_pool_config_for_tests(None, None, None)
|
||||
}));
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool));
|
||||
let request = pooled_websocket_request("conversation-first-token-bypass-timeout");
|
||||
|
||||
let response = backend
|
||||
.create_response_stream(
|
||||
&request,
|
||||
request_context("req_first_token_bypass_timeout", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("bypass first-token timeout should retry before returning stream");
|
||||
let decision = response.websocket_pool_decision;
|
||||
let mut stream = response.body;
|
||||
let mut body = String::new();
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk.expect("retried bypass stream chunk should be valid");
|
||||
body.push_str(std::str::from_utf8(&chunk).unwrap());
|
||||
}
|
||||
server.await.unwrap();
|
||||
|
||||
assert!(body.contains("bypass retry recovered"));
|
||||
assert!(body.contains("resp_bypass_retry_recovered"));
|
||||
assert!(!body.contains("resp_bypass_first_token_stalled"));
|
||||
let decision = decision.expect("bypass decision");
|
||||
assert_eq!(decision.kind(), "bypass");
|
||||
assert_eq!(decision.reason(), Some("cap"));
|
||||
assert_eq!(accepted_connections.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn codex_backend_client_should_not_reuse_pooled_websocket_across_local_accounts() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
@@ -180,7 +454,7 @@ async fn codex_backend_client_should_not_reuse_pooled_websocket_across_local_acc
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool));
|
||||
let mut request =
|
||||
@@ -273,7 +547,7 @@ async fn websocket_pool_should_bypass_busy_key_with_one_shot_connections() {
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(pool);
|
||||
let request = pooled_websocket_request("conversation-busy");
|
||||
@@ -366,7 +640,7 @@ async fn websocket_pool_should_release_slot_when_client_drops_stream() {
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(pool);
|
||||
let request = pooled_websocket_request("conversation-drop");
|
||||
@@ -440,7 +714,7 @@ async fn websocket_pool_should_bypass_new_keys_after_account_cap() {
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(pool);
|
||||
let first_request = pooled_websocket_request("conversation-cap-one");
|
||||
@@ -526,7 +800,7 @@ async fn websocket_pool_should_evict_idle_connection_when_liveness_lapses_despit
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool));
|
||||
let request = pooled_websocket_request("conversation-no-pong");
|
||||
@@ -600,7 +874,7 @@ async fn websocket_pool_should_gc_expired_idle_connections() {
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool));
|
||||
let request = pooled_websocket_request("conversation-gc");
|
||||
@@ -704,7 +978,7 @@ async fn codex_backend_client_should_keep_idle_pooled_websocket_alive_across_rep
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool));
|
||||
let mut request =
|
||||
@@ -810,7 +1084,7 @@ async fn codex_backend_client_should_close_idle_pooled_websocket_when_account_is
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool));
|
||||
let mut request =
|
||||
@@ -914,7 +1188,7 @@ async fn codex_backend_client_should_stop_reusing_pooled_websockets_after_shutdo
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool));
|
||||
let mut request =
|
||||
@@ -1018,7 +1292,7 @@ async fn codex_backend_client_should_close_idle_pooled_websocket_after_liveness_
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(Arc::clone(&pool));
|
||||
let mut request =
|
||||
@@ -1113,7 +1387,7 @@ async fn codex_backend_client_should_discard_pooled_websocket_after_upstream_err
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(pool);
|
||||
let mut request =
|
||||
@@ -1193,7 +1467,7 @@ async fn codex_backend_client_should_reuse_pooled_websocket_after_unmapped_respo
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(pool);
|
||||
let request = pooled_websocket_request("conversation-pool-unmapped-failed");
|
||||
|
||||
+12
-9
@@ -1,6 +1,6 @@
|
||||
# Deploy
|
||||
# 部署
|
||||
|
||||
Docker 部署入口集中在本目录。
|
||||
本目录包含 Dockerfile、Compose 文件和 Docker 配置模板。
|
||||
|
||||
## Compose
|
||||
|
||||
@@ -11,14 +11,17 @@ mkdir -p .runtime/data .runtime/logs
|
||||
cp deploy/config.example.yaml .runtime/config.yaml
|
||||
```
|
||||
|
||||
`deploy/config.example.yaml` 是 Docker 部署模板。Compose 默认把配置、数据和日志都放在仓库 `.runtime/` 下:
|
||||
`deploy/config.example.yaml` 是 Docker 部署模板。Compose 默认使用仓库根目录的 `.runtime/`:
|
||||
|
||||
- `.runtime/config.yaml` -> `/app/config.yaml`
|
||||
- `.runtime/data` -> `/app/data`
|
||||
- `.runtime/logs` -> `/app/logs`
|
||||
|
||||
复制后必须设置 `admin.default_password`,空值或常见弱口令会被后端拒绝启动。
|
||||
如果前面还有 Nginx、Caddy、Cloudflare Tunnel 等反向代理,只在该代理会覆盖或清理 `X-Forwarded-For` / `X-Real-IP` / `CF-Connecting-IP` 时,把代理的 IP 或 CIDR 写入 `server.trusted_proxies`;默认空列表会剥离这些客户端可伪造的头并使用 TCP peer IP。
|
||||
复制后必须设置 `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。
|
||||
|
||||
构建并启动:
|
||||
|
||||
@@ -27,7 +30,7 @@ docker compose -f deploy/docker-compose.yml build
|
||||
docker compose -f deploy/docker-compose.yml up -d
|
||||
```
|
||||
|
||||
可以通过环境变量覆盖宿主机路径;容器内会通过 `CPR_CONFIG_FILE=/app/config.yaml` 显式读取挂载后的配置:
|
||||
可以通过环境变量覆盖宿主机路径。容器内始终通过 `CPR_CONFIG_FILE=/app/config.yaml` 读取挂载后的配置:
|
||||
|
||||
```bash
|
||||
CPR_CONFIG_FILE=/path/to/config.yaml \
|
||||
@@ -36,7 +39,7 @@ CPR_LOG_DIR=/path/to/logs \
|
||||
docker compose -f deploy/docker-compose.yml up -d
|
||||
```
|
||||
|
||||
## Build
|
||||
## 构建
|
||||
|
||||
Docker 构建上下文保持仓库根目录:
|
||||
|
||||
@@ -44,13 +47,13 @@ Docker 构建上下文保持仓库根目录:
|
||||
docker build -f deploy/Dockerfile -t codex-proxy-rs:latest .
|
||||
```
|
||||
|
||||
也可以通过 Compose 构建:
|
||||
也可以用 Compose 构建:
|
||||
|
||||
```bash
|
||||
docker compose -f deploy/docker-compose.yml build
|
||||
```
|
||||
|
||||
需要注入版本信息时:
|
||||
需要写入版本元数据时:
|
||||
|
||||
```bash
|
||||
CPR_VERSION="$(ruby -ryaml -e 'puts YAML.load_file("release/version.yaml").fetch("version").delete_prefix("v")')" \
|
||||
|
||||
@@ -3,7 +3,10 @@ server:
|
||||
host: '0.0.0.0'
|
||||
# 容器内服务端口,需与 docker-compose.yml 的 target port 一致。
|
||||
port: 8080
|
||||
# 可信反向代理 IP 或 CIDR。为空时会忽略并剥离客户端传入的转发 IP 头。
|
||||
# 可信反向代理 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:
|
||||
|
||||
+322
-122
@@ -1,86 +1,107 @@
|
||||
# Codex Proxy RS 架构说明
|
||||
# Codex Proxy RS 架构
|
||||
|
||||
## 总览
|
||||
|
||||
Codex Proxy RS 是一个单主进程应用。Rust 后端同时承担 OpenAI 兼容代理、管理端 API、SQLite 持久化、后台任务、静态前端托管和在线更新;Vue 管理端作为 SPA 由同一个后端进程服务。
|
||||
Codex Proxy RS 是一个单进程代理和管理系统。Rust 后端同时承载 OpenAI Responses 兼容代理、管理端 API、SQLite 持久化、后台任务、静态前端托管和在线更新;Vue 管理端作为 SPA 由同一个进程托管。
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
Client["客户端 / OpenAI SDK"] -->|"Authorization: Bearer<br/>/v1/*"| Http["Axum Router"]
|
||||
Browser["浏览器"] -->|"管理端 SPA"| Web["web/dist"]
|
||||
Browser -->|"/api/admin/*"| Http
|
||||
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 兼容入口<br/>客户端 API Key 鉴权"]
|
||||
Http --> Admin["admin<br/>账号 / Key / 设置 / 监控 / 系统更新"]
|
||||
Http --> Web
|
||||
Http --> Proxy["proxy<br/>OpenAI compatible routes"]
|
||||
Http --> Admin["admin<br/>accounts / keys / monitoring / system"]
|
||||
Http --> Web["web<br/>static SPA"]
|
||||
|
||||
Proxy --> Dispatch["dispatch<br/>账号选择<br/>限流 / quota / 亲和性 / 回退"]
|
||||
Dispatch --> Protocol["upstream protocol<br/>请求转换<br/>SSE / 错误映射"]
|
||||
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 --> Runtime["runtime services<br/>账号池 / 后台任务 / 更新状态"]
|
||||
Dispatch --> Runtime
|
||||
Runtime --> Store[("SQLite<br/>账号 / Key / session / 用量 / 设置")]
|
||||
Admin --> Store
|
||||
Admin --> Services["runtime services"]
|
||||
Dispatch --> Services
|
||||
Services --> Store[("SQLite")]
|
||||
end
|
||||
|
||||
Transport --> Upstream["ChatGPT / Codex Upstream"]
|
||||
Release["GitHub Release<br/>archives + checksums"] --> Updater["self update<br/>下载 / 校验 / 替换 / 重启"]
|
||||
Updater --> Admin
|
||||
GHCR["GHCR image<br/>linux/amd64 + linux/arm64"] --> Docker["Docker Compose<br/>.runtime bind mounts"]
|
||||
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 8、Tailwind v4、Pinia、Vue Router、ECharts 和 lucide 图标。
|
||||
- `deploy/`:Dockerfile、Compose 和 Docker 部署配置模板。
|
||||
- `docs/`:长期维护文档和审计记录,只记录当前有效结构与决策。
|
||||
- `frontend/`:Vue 3 管理端,使用 Vite、Pinia、Vue Router、ECharts 和 lucide 图标。
|
||||
- `deploy/`:Dockerfile、Compose 和部署配置模板。
|
||||
- `docs/`:长期维护文档。
|
||||
- `release/`:版本号、平台矩阵和发布脚本。
|
||||
- `skills/`:项目本地 Codex skill。
|
||||
|
||||
当前后端是 `backend/` 单 crate。不要在文档里保留旧迁移布局或兼容性重导出约束,除非代码重新引入对应结构。
|
||||
测试统一放在 `backend/tests/`,不在 `backend/src/` 新增测试模块。
|
||||
|
||||
## 后端结构
|
||||
|
||||
入口:
|
||||
入口和组合层:
|
||||
|
||||
- `backend/src/main.rs`:二进制入口。
|
||||
- `backend/src/lib.rs`:集成测试和二进制共享的库入口。
|
||||
- `backend/src/runtime/bootstrap.rs`:加载配置、初始化日志、连接数据库、构造服务、启动 HTTP。
|
||||
- `backend/src/runtime/services.rs`:集中构造账号、模型、监控、更新等运行时服务。
|
||||
- `backend/src/lib.rs`:二进制和集成测试共享入口。
|
||||
- `runtime/bootstrap.rs`:加载配置、初始化日志、连接 SQLite、构造服务、恢复运行时状态并启动 HTTP。
|
||||
- `runtime/services.rs`:集中构造账号、模型、监控、更新、fingerprint、WebSocket 池等服务。
|
||||
- `runtime/tasks/`:后台任务协调和周期任务。
|
||||
|
||||
主要模块:
|
||||
主要业务模块:
|
||||
|
||||
- `admin/`:管理端 API,包括登录 session、账号、API Key、监控、运行设置和系统更新。
|
||||
- `proxy/`:OpenAI 兼容 HTTP 入口、客户端 API Key 鉴权、请求分派。
|
||||
- `proxy/openai/`:Responses 和 Models 的协议入口。
|
||||
- `proxy/dispatch/`:账号选择、quota/限流处理、错误映射、session affinity 和失败回退。
|
||||
- `upstream/`:ChatGPT/Codex 上游账号、模型、协议、传输、token 刷新和 fingerprint。
|
||||
- `config/`:启动配置、运行时设置和管理端设置落库。
|
||||
- `infra/`:SQLite、migration、日志、时间、格式化、JSON 和身份哈希。
|
||||
- `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。
|
||||
|
||||
测试统一放在 `backend/tests/`,不在 `backend/src/` 新增测试模块。
|
||||
|
||||
## 配置和运行时目录
|
||||
## 配置
|
||||
|
||||
启动配置由 `CPR_CONFIG_FILE` 指定;未指定时读取当前工作目录的 `config.yaml`。
|
||||
|
||||
配置文件只负责启动必需项:
|
||||
配置文件字段:
|
||||
|
||||
- HTTP 监听。
|
||||
- 上游 base URL。
|
||||
- SQLite URL。
|
||||
- 日志目录。
|
||||
- 管理员首次初始化密码。
|
||||
- TLS、WebSocket 池、quota 周期和 fingerprint 默认值。
|
||||
- `server`:监听地址、端口、可信代理。
|
||||
- `api`:上游 base URL。
|
||||
- `database`:SQLite URL。
|
||||
- `tls`:HTTP/TLS 偏好。
|
||||
- `ws_pool`:WebSocket 连接池和首内容帧超时。
|
||||
- `quota`:quota 刷新周期和是否跳过耗尽账号。
|
||||
- `fingerprint`:默认 Codex Desktop 指纹和请求头。
|
||||
- `admin`:管理员首次初始化和 session 清理配置。
|
||||
- `logging`:日志目录、保留期和使用记录开关。
|
||||
|
||||
账号、客户端 API Key、模型别名、账号选择策略和多数运行参数由管理端写入 SQLite。
|
||||
运行时设置保存在 `runtime_settings`:
|
||||
|
||||
`.runtime` 是本项目统一运行目录:
|
||||
- `modelAliases`
|
||||
- `refreshMarginSeconds`
|
||||
- `refreshConcurrency`
|
||||
- `maxConcurrentPerAccount`
|
||||
- `requestIntervalMs`
|
||||
- `rotationStrategy`
|
||||
|
||||
这些设置由管理端更新,并同步到账号池、token refresh 策略和模型服务。
|
||||
|
||||
## 运行目录
|
||||
|
||||
`.runtime` 是默认运行目录:
|
||||
|
||||
```text
|
||||
.runtime/config.yaml
|
||||
@@ -91,7 +112,7 @@ flowchart LR
|
||||
.runtime/logs/
|
||||
```
|
||||
|
||||
Docker 也使用宿主机 `.runtime`,但容器内配置路径仍然写 `/app/data` 和 `/app/logs`。Compose 默认映射:
|
||||
Docker 默认映射:
|
||||
|
||||
```text
|
||||
.runtime/config.yaml -> /app/config.yaml
|
||||
@@ -99,27 +120,227 @@ Docker 也使用宿主机 `.runtime`,但容器内配置路径仍然写 `/app/d
|
||||
.runtime/logs -> /app/logs
|
||||
```
|
||||
|
||||
容器以非 root 用户 `10001:10001` 运行。Compose 设置 `HOME=/app` 和 `XDG_DATA_HOME=/app/data`,确保 installation id、SQLite、更新状态和日志都写入可持久化目录。
|
||||
容器以非 root 用户 `10001:10001` 运行。Compose 设置 `HOME=/app` 和 `XDG_DATA_HOME=/app/data`,确保 installation id、SQLite、更新状态和日志写入持久化目录。
|
||||
|
||||
## 代理请求链路
|
||||
## SQLite
|
||||
|
||||
客户端请求进入 `/v1/*`:
|
||||
当前 schema 由 `backend/src/infra/migrations` 管理。
|
||||
|
||||
1. `proxy/auth` 校验客户端 API Key,并记录调用上下文。
|
||||
2. `proxy/openai` 解析 OpenAI 兼容请求。
|
||||
3. `proxy/dispatch` 按模型、quota、账号状态和选择策略挑选账号。
|
||||
4. `upstream/protocol` 转换为 ChatGPT/Codex 上游协议。
|
||||
5. `upstream/transport` 使用 HTTP、SSE 或 WebSocket 访问上游。
|
||||
6. 响应和错误被转换回 OpenAI 兼容格式。
|
||||
7. 用量、模型维度统计、请求记录、账号窗口和失败状态写回 SQLite/账号池。
|
||||
核心表:
|
||||
|
||||
账号状态由 SQLite 恢复到内存账号池,并在请求、token 刷新、quota 刷新、Cloudflare 处理和失败回退时同步更新。
|
||||
- `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 关联。
|
||||
|
||||
## 管理端
|
||||
## 代理入口
|
||||
|
||||
管理端 API 位于 `/api/admin/*`,使用 session cookie 保护。前端由后端托管,不单独部署 Node 服务。
|
||||
客户端请求入口:
|
||||
|
||||
前端结构:
|
||||
- `POST /v1/responses`
|
||||
- `POST /v1/responses/review`
|
||||
- `POST /v1/responses/compact`
|
||||
- `GET /v1/models`
|
||||
- `GET /v1/models/{model}`
|
||||
|
||||
请求链路:
|
||||
|
||||
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。
|
||||
|
||||
## 账号调度
|
||||
|
||||
账号池位于 `upstream/accounts/pool.rs`。调度前先做硬过滤:
|
||||
|
||||
- 账号状态必须是 `active`。
|
||||
- 模型必须被账号计划允许。
|
||||
- quota 必须可用。
|
||||
- Cloudflare 冷却必须结束。
|
||||
- 账号不在本次请求排除列表。
|
||||
- 当前在途槽位小于 `maxConcurrentPerAccount`。
|
||||
|
||||
策略:
|
||||
|
||||
- `smart`:默认策略。按账号实时负载选择账号,依次比较当前在途槽位数、窗口请求数、窗口 token 数和历史请求数;负载相同,再参考额度窗口重置时间和最近使用时间。
|
||||
- `quota_reset_priority`:按额度窗口重置时间选择账号,优先使用更早重置的窗口;重置时间相同,再比较历史请求数和最近使用时间。
|
||||
- `round_robin`:按候选账号顺序循环选择。
|
||||
- `sticky`:按最近使用时间优先选择。
|
||||
|
||||
请求带有 `required_account_id` 时只尝试该账号。请求带有 session affinity 的 `preferred_account_id` 时,只有该账号仍在硬过滤结果中才会命中,否则回到当前调度策略。选中账号后立即占用在途槽位,释放时记录请求用量并更新窗口计数。
|
||||
|
||||
## Session Affinity
|
||||
|
||||
Session affinity 由 `proxy/dispatch/session_affinity.rs` 和 `proxy/dispatch/responses/affinity.rs` 维护。
|
||||
|
||||
存储内容:
|
||||
|
||||
- response id。
|
||||
- conversation id。
|
||||
- account id。
|
||||
- prompt/instructions 变体。
|
||||
- turn state。
|
||||
- function call ids。
|
||||
- 输入 token。
|
||||
- 过期时间。
|
||||
|
||||
行为:
|
||||
|
||||
- 请求显式携带 `previous_response_id` 时,优先查找该 response 对应账号。
|
||||
- 请求没有 `previous_response_id` 但满足隐式续写条件时,按 conversation 查找最新 response 并补充上下文。
|
||||
- 当亲和账号不可用并切换到其它账号时,按请求类型和失败类型处理历史字段。
|
||||
- 上游返回 `previous_response_not_found`、未闭合 function call 等历史错误时,恢复请求历史并重试。
|
||||
- 成功响应会记录新的 response affinity,供后续请求继续使用。
|
||||
|
||||
## Quota 和账号状态
|
||||
|
||||
账号状态包括:
|
||||
|
||||
- `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 组件。
|
||||
@@ -128,77 +349,32 @@ Docker 也使用宿主机 `.runtime`,但容器内配置路径仍然写 `/app/d
|
||||
- `frontend/src/stores`:Pinia 状态。
|
||||
- `frontend/src/styles/tokens.css`:亮色/暗色主题 token。
|
||||
|
||||
管理端页面包括:
|
||||
主要页面:
|
||||
|
||||
- 账号管理:导入、OAuth、连接测试、额度、模型、调度状态、刷新能力。
|
||||
- API Key 管理:客户端 Key 和管理员 API Key。
|
||||
- Dashboard:请求趋势、账号概览、服务状态。
|
||||
- 请求明细:网关调用、上游状态、Token 消耗、错误和详情。
|
||||
- 设置:运行参数、模型映射、账号选择策略。
|
||||
- 系统更新:版本、检查更新、SSE 更新日志、更新、重启、回滚。
|
||||
|
||||
## 发布
|
||||
|
||||
版本契约集中在 `release/`:
|
||||
|
||||
- `release/version.yaml`:当前产品版本。
|
||||
- `release/platforms.yaml`:Release asset 和 Docker 平台矩阵。
|
||||
- `release/publish`:更新版本文件、生成版本提交、创建 tag 并 push。
|
||||
|
||||
发布命令:
|
||||
|
||||
```bash
|
||||
release/publish 0.1.5
|
||||
```
|
||||
|
||||
发布 workflow 校验触发 tag 必须等于 `v` + `release/version.yaml`。构建时注入:
|
||||
|
||||
- `CPR_VERSION`
|
||||
- `CPR_GIT_SHA`
|
||||
- `CPR_BUILD_TIME`
|
||||
- `CPR_BUILD_TYPE=release`
|
||||
|
||||
`backend/build/build.rs` 负责把这些值写入二进制。没有 `CPR_VERSION` 时,版本回落到 `release/version.yaml`,最后才回落到 Cargo 包版本。
|
||||
|
||||
Release 产物:
|
||||
|
||||
```text
|
||||
codex-proxy-rs_<version>_linux_amd64.tar.gz
|
||||
codex-proxy-rs_<version>_linux_arm64.tar.gz
|
||||
codex-proxy-rs_<version>_darwin_arm64.tar.gz
|
||||
codex-proxy-rs_<version>_windows_amd64.zip
|
||||
checksums.txt
|
||||
```
|
||||
|
||||
GHCR 镜像发布 `linux/amd64` 和 `linux/arm64` 多平台 manifest。发布 workflow 先在对应原生 runner 上构建平台镜像,再合并正式 manifest。Docker 在线更新使用的 Linux 二进制在 `rust:1.95-bookworm` 中构建,保持与 runtime 镜像的 glibc ABI 兼容。
|
||||
- 首页 Dashboard。
|
||||
- 账号管理。
|
||||
- API Key 管理。
|
||||
- 请求明细。
|
||||
- 系统设置。
|
||||
- 系统更新。
|
||||
|
||||
## 在线更新
|
||||
|
||||
在线更新由主服务处理,没有独立 updater sidecar。
|
||||
|
||||
核心接口:
|
||||
|
||||
- `GET /api/admin/system/version`
|
||||
- `GET /api/admin/system/update-detail`
|
||||
- `GET /api/admin/system/update-events`
|
||||
- `GET /api/admin/system/update-status`
|
||||
- `POST /api/admin/system/update`
|
||||
- `POST /api/admin/system/restart`
|
||||
- `POST /api/admin/system/rollback`
|
||||
|
||||
更新流程:
|
||||
流程:
|
||||
|
||||
1. 查询 GitHub Release。
|
||||
2. 按当前 OS/arch 选择 asset,支持 `amd64/x86_64` 与 `arm64/aarch64` 命名。
|
||||
3. 下载归档和 `checksums.txt`。
|
||||
2. 按当前 OS/arch 选择 asset。
|
||||
3. 下载归档和 checksum 文件。
|
||||
4. 校验 checksum。
|
||||
5. 解压二进制和 `web/dist`。
|
||||
5. 解压后端二进制和 `web/dist`。
|
||||
6. 替换当前二进制和前端静态资源。
|
||||
7. 写入 `.runtime/data/update-state.json`。
|
||||
8. 通过 SSE 推送中文更新日志。
|
||||
9. 管理端触发重启;Docker 模式依靠 `restart: unless-stopped` 拉起新容器进程,非 Docker 模式先安排新进程延迟启动,再关闭当前进程。
|
||||
9. 管理端触发重启。
|
||||
|
||||
替换逻辑必须支持跨文件系统 fallback,因为容器环境中 `rename` 可能遇到 `Invalid cross-device link`。
|
||||
Docker 模式依靠 `restart: unless-stopped` 拉起新进程。非 Docker 模式会安排新进程延迟启动,再关闭当前进程。
|
||||
|
||||
## Docker
|
||||
|
||||
@@ -207,10 +383,10 @@ GHCR 镜像发布 `linux/amd64` 和 `linux/arm64` 多平台 manifest。发布 wo
|
||||
阶段:
|
||||
|
||||
- `frontend-builder`:安装前端依赖并构建 Vue SPA。
|
||||
- `backend-builder`:构建 Rust release 二进制,不依赖前端 dist。
|
||||
- `backend-builder`:构建 Rust release 二进制。
|
||||
- `release-asset-builder` / `release-asset`:组合后端二进制和前端 dist,生成 GitHub Release 归档。
|
||||
- `runtime-base`:Debian slim、CA、curl、非 root 用户、运行目录。
|
||||
- `runtime`:从源码构建出的运行镜像。
|
||||
- `runtime-base`:Debian slim、CA、curl、非 root 用户和运行目录。
|
||||
- `runtime`:源码构建运行镜像。
|
||||
- `runtime-prebuilt`:发布 workflow 使用的预构建二进制镜像阶段。
|
||||
|
||||
Compose 默认:
|
||||
@@ -222,6 +398,30 @@ Compose 默认:
|
||||
- 配置:`.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`
|
||||
|
||||
发布命令:
|
||||
|
||||
```bash
|
||||
release/publish <version>
|
||||
```
|
||||
|
||||
发布 workflow 校验触发 tag 必须等于 `v` + `release/version.yaml`。
|
||||
|
||||
构建时注入:
|
||||
|
||||
- `CPR_VERSION`
|
||||
- `CPR_GIT_SHA`
|
||||
- `CPR_BUILD_TIME`
|
||||
- `CPR_BUILD_TYPE=release`
|
||||
|
||||
Release 产物按 `release/platforms.yaml` 生成 Linux、Darwin 和 Windows 包。GHCR 镜像发布 `linux/amd64` 和 `linux/arm64` 多平台 manifest。
|
||||
|
||||
## CI
|
||||
|
||||
常规 CI 位于 `.github/workflows/ci.yml`:
|
||||
|
||||
+4
-1
@@ -2,7 +2,10 @@
|
||||
<html lang="zh-CN">
|
||||
<head>
|
||||
<meta charset="UTF-8" />
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
|
||||
<meta
|
||||
name="viewport"
|
||||
content="width=device-width, initial-scale=1.0, maximum-scale=1.0, user-scalable=no, viewport-fit=cover"
|
||||
/>
|
||||
<link rel="icon" type="image/svg+xml" href="/favicon.svg" />
|
||||
<title>Codex Proxy RS</title>
|
||||
</head>
|
||||
|
||||
@@ -16,6 +16,7 @@ const props = defineProps<{
|
||||
scrollable?: boolean
|
||||
bodyMaxHeight?: string
|
||||
bodyViewClass?: string
|
||||
hideFooter?: boolean
|
||||
}>()
|
||||
|
||||
const open = defineModel<boolean>({ default: false })
|
||||
@@ -123,7 +124,7 @@ function closeModal() {
|
||||
</BaseScrollbar>
|
||||
<slot v-else />
|
||||
</div>
|
||||
<footer class="flex justify-end gap-3 px-7 pb-7">
|
||||
<footer v-if="!hideFooter" class="flex justify-end gap-3 px-7 pb-7">
|
||||
<slot name="footer">
|
||||
<BaseButton variant="default" @click="open = false">取消</BaseButton>
|
||||
<BaseButton variant="primary" @click="open = false">确认</BaseButton>
|
||||
|
||||
@@ -1,10 +1,12 @@
|
||||
<script setup lang="ts">
|
||||
import { clamp } from 'es-toolkit'
|
||||
import { computed } from 'vue'
|
||||
import type { Component } from 'vue'
|
||||
|
||||
interface SegmentedOption {
|
||||
label: string
|
||||
value: string
|
||||
icon?: Component
|
||||
}
|
||||
|
||||
const props = withDefaults(
|
||||
@@ -54,7 +56,7 @@ function selectOption(value: string) {
|
||||
<button
|
||||
v-for="option in options"
|
||||
:key="option.value"
|
||||
class="relative z-10 h-7 min-w-0 rounded-(--cp-input-radius-base) border-0 bg-transparent px-3 text-xs leading-none font-[650] transition-colors duration-150 outline-none focus-visible:ring-2 focus-visible:ring-(--cp-info-border)"
|
||||
class="relative z-10 inline-flex h-7 min-w-0 items-center justify-center gap-1.5 rounded-(--cp-input-radius-base) border-0 bg-transparent px-3 text-xs leading-none font-[650] transition-colors duration-150 outline-none focus-visible:ring-2 focus-visible:ring-(--cp-info-border)"
|
||||
:class="[
|
||||
model === option.value
|
||||
? 'text-(--cp-text-primary)'
|
||||
@@ -67,6 +69,7 @@ function selectOption(value: string) {
|
||||
:disabled="disabled"
|
||||
@click="selectOption(option.value)"
|
||||
>
|
||||
<component :is="option.icon" v-if="option.icon" class="size-3.5 shrink-0" />
|
||||
{{ option.label }}
|
||||
</button>
|
||||
</div>
|
||||
|
||||
@@ -32,7 +32,6 @@ export function resolveColumns(columns: BaseTableColumn[], minWidth?: number | s
|
||||
const minWidthPixels = minWidth === undefined ? 0 : numericPixelWidth(minWidth)
|
||||
const flexibleColumns = columns.filter((column) => column.width === undefined)
|
||||
const flexTotal = sumBy(flexibleColumns, (column) => column.flex ?? 1)
|
||||
const flexibleColumnCount = clamp(flexibleColumns.length, 1, Number.POSITIVE_INFINITY)
|
||||
const available = clamp(100 - fixedPercentTotal, 0, Number.POSITIVE_INFINITY)
|
||||
const availablePixels = clamp(minWidthPixels - fixedPixelTotal, 0, Number.POSITIVE_INFINITY)
|
||||
|
||||
@@ -41,10 +40,12 @@ export function resolveColumns(columns: BaseTableColumn[], minWidth?: number | s
|
||||
const width =
|
||||
column.width === undefined
|
||||
? minWidthPixels > 0 && fixedPercentTotal === 0
|
||||
? `${availablePixels / flexibleColumnCount}px`
|
||||
? flexTotal > 0
|
||||
? `${(availablePixels * flex) / flexTotal}px`
|
||||
: `${availablePixels / clamp(flexibleColumns.length, 1, Number.POSITIVE_INFINITY)}px`
|
||||
: flexTotal > 0
|
||||
? `${(available * flex) / flexTotal}%`
|
||||
: `${available / flexibleColumnCount}%`
|
||||
: `${available / clamp(flexibleColumns.length, 1, Number.POSITIVE_INFINITY)}%`
|
||||
: normalizeWidth(column.width)
|
||||
|
||||
return {
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
<script setup lang="ts">
|
||||
import { ExternalLink, GitBranch } from '@lucide/vue'
|
||||
import { computed } from 'vue'
|
||||
|
||||
import BaseModal from '@/components/base/BaseModal.vue'
|
||||
import { useSystemUpdate } from '@/composables/useSystemUpdate'
|
||||
|
||||
const open = defineModel<boolean>({ default: false })
|
||||
|
||||
const { version } = useSystemUpdate()
|
||||
|
||||
const author = 'Zyy'
|
||||
const githubUrl = 'https://github.com/zyycn/codex-proxy-rs'
|
||||
|
||||
function normalizeBuildValue(value: string | undefined) {
|
||||
const normalized = value?.trim()
|
||||
|
||||
if (!normalized) {
|
||||
return ''
|
||||
}
|
||||
|
||||
return ['unknown', 'null'].includes(normalized.toLowerCase()) ? '' : normalized
|
||||
}
|
||||
|
||||
const versionDisplay = computed(() => {
|
||||
const versionText = normalizeBuildValue(version.value?.version).replace(/^v/i, '')
|
||||
return versionText ? `v${versionText}` : ''
|
||||
})
|
||||
const gitShaDisplay = computed(() => {
|
||||
const gitSha = normalizeBuildValue(version.value?.gitSha)
|
||||
return gitSha ? gitSha.slice(0, 8) : ''
|
||||
})
|
||||
|
||||
const versionLine = computed(() => {
|
||||
const parts = [versionDisplay.value, gitShaDisplay.value].filter(Boolean)
|
||||
return parts.length ? `版本 ${parts.join(' · ')}` : '版本信息不可用'
|
||||
})
|
||||
|
||||
const linkItems = [
|
||||
{
|
||||
label: 'GitHub',
|
||||
value: 'codex-proxy-rs',
|
||||
href: githubUrl,
|
||||
icon: GitBranch,
|
||||
},
|
||||
]
|
||||
</script>
|
||||
|
||||
<template>
|
||||
<BaseModal v-model="open" title="关于" width="420px" hide-footer>
|
||||
<div class="grid gap-5">
|
||||
<section class="flex min-w-0 items-center gap-3">
|
||||
<span
|
||||
class="inline-flex size-9 shrink-0 items-center justify-center rounded-(--cp-icon-button-radius) bg-(--cp-bg-subtle) font-mono text-[15px] leading-none font-[820] text-(--cp-text-primary)"
|
||||
>
|
||||
Z
|
||||
</span>
|
||||
<div class="min-w-0">
|
||||
<p class="m-0 truncate text-[15px] leading-none font-[780] text-(--cp-text-primary)">
|
||||
{{ author }}
|
||||
</p>
|
||||
<p class="mt-1.5 mb-0 text-[12px] leading-none font-[650] text-(--cp-text-secondary)">
|
||||
Built by Zyy · Codex Proxy RS
|
||||
</p>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
<section class="grid gap-1">
|
||||
<div
|
||||
v-for="item in linkItems"
|
||||
:key="item.label"
|
||||
class="group flex min-w-0 items-center justify-between gap-3 rounded-(--cp-input-radius-base) px-1 py-2.5 transition-colors hover:bg-(--cp-bg-subtle)"
|
||||
>
|
||||
<div class="flex min-w-0 items-center gap-3">
|
||||
<span
|
||||
class="inline-flex size-8 shrink-0 items-center justify-center rounded-(--cp-icon-button-radius) bg-(--cp-bg-muted) text-(--cp-text-secondary)"
|
||||
>
|
||||
<component :is="item.icon" class="size-4" />
|
||||
</span>
|
||||
<div class="min-w-0">
|
||||
<p class="m-0 text-[11px] leading-none font-[760] text-(--cp-text-muted)">
|
||||
{{ item.label }}
|
||||
</p>
|
||||
<p
|
||||
class="mt-2 mb-0 truncate font-mono text-[12px] leading-none font-[720] text-(--cp-text-primary)"
|
||||
:title="item.value"
|
||||
>
|
||||
{{ item.value }}
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<a
|
||||
:href="item.href"
|
||||
target="_blank"
|
||||
rel="noreferrer"
|
||||
class="inline-flex size-7 shrink-0 items-center justify-center rounded-(--cp-icon-button-radius) text-(--cp-text-muted) transition-colors hover:bg-(--cp-bg-muted) hover:text-(--cp-info)"
|
||||
:aria-label="`打开 ${item.label}`"
|
||||
>
|
||||
<ExternalLink class="size-3.5" />
|
||||
</a>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
<p class="m-0 font-mono text-[11px] leading-none font-[650] text-(--cp-text-muted)">
|
||||
{{ versionLine }}
|
||||
</p>
|
||||
</div>
|
||||
</BaseModal>
|
||||
</template>
|
||||
@@ -8,6 +8,7 @@ import {
|
||||
Cat,
|
||||
ChartNoAxesColumn,
|
||||
ArrowUpCircle,
|
||||
Info,
|
||||
KeyRound,
|
||||
LayoutDashboard,
|
||||
LogOut,
|
||||
@@ -24,6 +25,7 @@ import { useSystemUpdate } from '@/composables/useSystemUpdate'
|
||||
import { useAuthStore } from '@/stores/modules/auth'
|
||||
import { useUiStore } from '@/stores/modules/ui'
|
||||
|
||||
import AppAboutModal from './AppAboutModal.vue'
|
||||
import SystemUpdateModal from './SystemUpdateModal.vue'
|
||||
|
||||
const route = useRoute()
|
||||
@@ -35,6 +37,23 @@ const { toggleTheme } = uiStore
|
||||
const preferredMotion = usePreferredReducedMotion()
|
||||
const { version, hasUpdate, loadedOnce, loadVersion, loadSystem } = useSystemUpdate()
|
||||
|
||||
const props = withDefaults(
|
||||
defineProps<{
|
||||
collapsed?: boolean
|
||||
mobile?: boolean
|
||||
}>(),
|
||||
{
|
||||
collapsed: false,
|
||||
mobile: false,
|
||||
},
|
||||
)
|
||||
|
||||
const emit = defineEmits<{
|
||||
close: []
|
||||
navigate: []
|
||||
toggle: []
|
||||
}>()
|
||||
|
||||
const navItems = [
|
||||
{ label: '概览', icon: LayoutDashboard, path: '/' },
|
||||
{ label: '账号管理', icon: Users, path: '/accounts' },
|
||||
@@ -49,7 +68,8 @@ const isActive = (path: string) => {
|
||||
}
|
||||
|
||||
function navigate(path: string) {
|
||||
router.push(path)
|
||||
void router.push(path)
|
||||
emit('navigate')
|
||||
}
|
||||
|
||||
const systemUpdateOpening = shallowRef(false)
|
||||
@@ -72,26 +92,17 @@ async function openSystemUpdate() {
|
||||
|
||||
async function handleLogout() {
|
||||
await authStore.logout()
|
||||
router.push('/login')
|
||||
await router.push('/login')
|
||||
emit('navigate')
|
||||
}
|
||||
|
||||
const props = withDefaults(
|
||||
defineProps<{
|
||||
collapsed?: boolean
|
||||
}>(),
|
||||
{
|
||||
collapsed: false,
|
||||
},
|
||||
)
|
||||
|
||||
defineEmits<{
|
||||
toggle: []
|
||||
}>()
|
||||
|
||||
const sidebarEl = ref<HTMLElement | null>(null)
|
||||
const brandLabelEl = ref<HTMLElement | null>(null)
|
||||
const brandLabelVisible = shallowRef(!props.collapsed)
|
||||
const systemUpdateOpen = shallowRef(false)
|
||||
const aboutOpen = shallowRef(false)
|
||||
const isCollapsed = computed(() => !props.mobile && Boolean(props.collapsed))
|
||||
const sidebarWidth = computed(() => (isCollapsed.value ? 88 : 256))
|
||||
const brandLabelVisible = shallowRef(!isCollapsed.value)
|
||||
const themeToggleLabel = computed(() =>
|
||||
effectiveTheme.value === 'dark' ? '切换浅色模式' : '切换暗黑模式',
|
||||
)
|
||||
@@ -208,10 +219,10 @@ function animateSidebarWidth(collapsed: boolean) {
|
||||
|
||||
onMounted(() => {
|
||||
gsap.set(sidebarEl.value, {
|
||||
width: props.collapsed ? 88 : 256,
|
||||
flexBasis: props.collapsed ? 88 : 256,
|
||||
width: sidebarWidth.value,
|
||||
flexBasis: sidebarWidth.value,
|
||||
})
|
||||
if (props.collapsed) {
|
||||
if (isCollapsed.value) {
|
||||
hideBrandLabel()
|
||||
} else {
|
||||
gsap.set(brandLabelEl.value, {
|
||||
@@ -219,12 +230,12 @@ onMounted(() => {
|
||||
x: 0,
|
||||
})
|
||||
}
|
||||
animateSidebarLabels(Boolean(props.collapsed))
|
||||
animateSidebarLabels(isCollapsed.value)
|
||||
void loadVersion().catch(() => undefined)
|
||||
})
|
||||
|
||||
watch(
|
||||
() => props.collapsed,
|
||||
() => isCollapsed.value,
|
||||
async (collapsed) => {
|
||||
if (collapsed) {
|
||||
hideBrandLabel()
|
||||
@@ -245,12 +256,15 @@ watch(
|
||||
<template>
|
||||
<aside
|
||||
ref="sidebarEl"
|
||||
class="z-20 hidden h-screen shrink-0 flex-col overflow-hidden bg-(--cp-bg-surface) px-4 shadow-(--cp-shadow-sidebar) min-[961px]:flex"
|
||||
:class="collapsed ? 'w-22 basis-22 items-center' : 'w-64 basis-64'"
|
||||
class="z-20 h-dvh shrink-0 flex-col overflow-hidden bg-(--cp-bg-surface) px-4 shadow-(--cp-shadow-sidebar)"
|
||||
:class="[
|
||||
mobile ? 'flex' : 'hidden min-[961px]:flex',
|
||||
isCollapsed ? 'w-22 basis-22 items-center' : 'w-64 basis-64',
|
||||
]"
|
||||
>
|
||||
<div
|
||||
class="mt-6 grid h-12 grid-cols-[44px_minmax(0,1fr)] items-center"
|
||||
:class="collapsed ? 'w-11 justify-start' : 'w-full gap-3'"
|
||||
:class="isCollapsed ? 'w-11 justify-start' : 'w-full gap-3'"
|
||||
>
|
||||
<span
|
||||
class="inline-flex size-11 items-center justify-center relative -top-0.5 rounded-(--cp-icon-button-radius) bg-(--cp-bg-muted) text-(--cp-text-primary)"
|
||||
@@ -290,14 +304,14 @@ watch(
|
||||
</span>
|
||||
</div>
|
||||
|
||||
<nav class="mt-7 grid gap-3" :class="collapsed ? '' : 'w-full'" aria-label="主导航">
|
||||
<nav class="mt-7 grid gap-3" :class="isCollapsed ? '' : 'w-full'" aria-label="主导航">
|
||||
<button
|
||||
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="[
|
||||
collapsed ? 'w-11.5 justify-center' : 'w-full gap-3 px-4',
|
||||
isCollapsed ? 'w-11.5 justify-center' : 'w-full gap-3 px-4',
|
||||
isActive(item.path)
|
||||
? 'bg-(--cp-bg-nav-active) font-bold text-(--cp-text-primary)'
|
||||
: 'bg-transparent font-semibold text-(--cp-text-secondary) hover:bg-(--cp-bg-subtle) hover:text-(--cp-text-primary)',
|
||||
@@ -307,32 +321,32 @@ watch(
|
||||
<component :is="item.icon" class="shrink-0" :size="20" />
|
||||
<span
|
||||
class="sidebar-label overflow-hidden whitespace-nowrap transition-[opacity,transform] duration-200"
|
||||
:class="collapsed ? 'pointer-events-none w-0' : 'w-auto'"
|
||||
:class="isCollapsed ? 'pointer-events-none w-0' : 'w-auto'"
|
||||
>{{ item.label }}</span
|
||||
>
|
||||
</button>
|
||||
</nav>
|
||||
|
||||
<div class="mt-auto mb-6" :class="collapsed ? 'w-11' : 'w-full'">
|
||||
<div class="mt-auto mb-6" :class="isCollapsed ? 'w-11' : 'w-full'">
|
||||
<div
|
||||
class="bg-(--cp-bg-subtle)"
|
||||
:class="
|
||||
collapsed
|
||||
isCollapsed
|
||||
? 'grid gap-1 rounded-(--cp-icon-button-radius) p-1'
|
||||
: 'flex h-11 items-center justify-between rounded-(--cp-panel-radius) px-2'
|
||||
"
|
||||
>
|
||||
<span
|
||||
v-if="!collapsed"
|
||||
v-if="!isCollapsed"
|
||||
class="inline-flex h-7 items-center gap-1.5 rounded-lg bg-(--cp-success-bg) px-2.5 text-xs leading-none font-[650] text-(--cp-success-text)"
|
||||
>
|
||||
<i class="size-1.5 rounded-full bg-(--cp-success)" />
|
||||
在线
|
||||
</span>
|
||||
|
||||
<div class="flex items-center" :class="collapsed ? 'grid gap-1' : 'gap-1'">
|
||||
<div class="flex items-center" :class="isCollapsed ? 'grid gap-1' : 'gap-1'">
|
||||
<BaseButton
|
||||
v-if="collapsed"
|
||||
v-if="isCollapsed && hasUpdate"
|
||||
icon-only
|
||||
:variant="hasUpdate ? 'success' : 'ghost'"
|
||||
size="default"
|
||||
@@ -346,34 +360,56 @@ watch(
|
||||
<BaseButton
|
||||
icon-only
|
||||
variant="ghost"
|
||||
:size="collapsed ? 'default' : 'sm'"
|
||||
:size="isCollapsed ? 'default' : 'sm'"
|
||||
label="退出登录"
|
||||
class="hover:bg-(--cp-danger-bg) hover:text-(--cp-danger)"
|
||||
@click="handleLogout"
|
||||
>
|
||||
<LogOut :size="collapsed ? 19 : 18" />
|
||||
<LogOut :size="isCollapsed ? 19 : 18" />
|
||||
</BaseButton>
|
||||
|
||||
<BaseButton
|
||||
icon-only
|
||||
variant="ghost"
|
||||
:size="collapsed ? 'default' : 'sm'"
|
||||
:size="isCollapsed ? 'default' : 'sm'"
|
||||
:label="themeToggleLabel"
|
||||
@click="toggleTheme($event)"
|
||||
>
|
||||
<Sun v-if="effectiveTheme === 'dark'" :size="collapsed ? 19 : 18" />
|
||||
<Moon v-else :size="collapsed ? 19 : 18" />
|
||||
<Sun v-if="effectiveTheme === 'dark'" :size="isCollapsed ? 19 : 18" />
|
||||
<Moon v-else :size="isCollapsed ? 19 : 18" />
|
||||
</BaseButton>
|
||||
|
||||
<BaseButton
|
||||
icon-only
|
||||
variant="ghost"
|
||||
:size="collapsed ? 'default' : 'sm'"
|
||||
data-sidebar-toggle
|
||||
:label="collapsed ? '展开侧边栏' : '收缩侧边栏'"
|
||||
@click="$emit('toggle')"
|
||||
:size="isCollapsed ? 'default' : 'sm'"
|
||||
label="关于"
|
||||
@click="aboutOpen = true"
|
||||
>
|
||||
<PanelLeftOpen v-if="collapsed" :size="19" />
|
||||
<Info :size="isCollapsed ? 19 : 18" />
|
||||
</BaseButton>
|
||||
|
||||
<BaseButton
|
||||
v-if="mobile"
|
||||
icon-only
|
||||
variant="ghost"
|
||||
size="sm"
|
||||
label="关闭侧边栏"
|
||||
@click="emit('close')"
|
||||
>
|
||||
<PanelLeftClose :size="18" />
|
||||
</BaseButton>
|
||||
|
||||
<BaseButton
|
||||
v-else
|
||||
icon-only
|
||||
variant="ghost"
|
||||
:size="isCollapsed ? 'default' : 'sm'"
|
||||
data-sidebar-toggle
|
||||
:label="isCollapsed ? '展开侧边栏' : '收缩侧边栏'"
|
||||
@click="emit('toggle')"
|
||||
>
|
||||
<PanelLeftOpen v-if="isCollapsed" :size="19" />
|
||||
<PanelLeftClose v-else :size="18" />
|
||||
</BaseButton>
|
||||
</div>
|
||||
@@ -381,5 +417,6 @@ watch(
|
||||
</div>
|
||||
</aside>
|
||||
|
||||
<AppAboutModal v-model="aboutOpen" />
|
||||
<SystemUpdateModal v-model="systemUpdateOpen" />
|
||||
</template>
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user