Compare commits

...
19 Commits
Author SHA1 Message Date
zyycnandCodex 28f2879a05 chore: bump version to 2.5.2
Co-authored-by: Codex <noreply@openai.com>
2026-07-19 20:40:47 +08:00
zyycnandCodex 2474259c22 fix(auth): synchronize refreshed tokens with account pool
Co-authored-by: Codex <noreply@openai.com>
2026-07-19 20:30:43 +08:00
ZyyandCodex 79c2c5948f fix(backend): skip drain for self-restart
Route in-process update restarts through an explicit shutdown action so the runtime releases the listener and logging guard before exec, without waiting for the single instance's graceful drain.

Co-authored-by: Codex <noreply@openai.com>
2026-07-17 18:54:33 +08:00
ZyyandCodex 3c4bd8757f chore: bump version to 2.5.1
Co-authored-by: Codex <noreply@openai.com>
2026-07-17 17:38:38 +08:00
ZyyandCodex 06f3443e56 docs(architecture): define multi-provider gateway target
Document the provider-neutral gateway boundary and its terminal PostgreSQL and Redis data model, including configuration, credentials, routing, attempts, continuation, telemetry, retention, and migration contracts.

Co-authored-by: Codex <noreply@openai.com>
2026-07-17 17:38:11 +08:00
Zyy 88fd21f169 fix(frontend): restore table scroll and loading states 2026-07-17 16:23:35 +08:00
Zyy 521201c926 fix(telemetry): preserve timed-out record writes 2026-07-17 14:52:19 +08:00
ZyyandCodex be2123ea30 chore: bump version to 2.5.0
Co-authored-by: Codex <noreply@openai.com>
2026-07-17 13:40:14 +08:00
ZyyandCodex 4b6d369225 refactor(backend): simplify admin query structure
Move account and dashboard read orchestration into their existing admin API modules, and remove the top-level admin_queries package and service wrapper.

Update composition wiring, dependency guards, and architecture documentation to reflect feature-local query ownership.

Co-authored-by: Codex <noreply@openai.com>
2026-07-17 13:09:17 +08:00
ZyyandCodex 76b81a6ceb refactor(frontend): simplify client architecture
Remove frontend schemas and tests, keep API modules as thin GET/POST adapters, and replace manual SSE and browser helpers with existing VueUse and utility primitives.

Reorganize feature helpers, move unauthorized handling into an app plugin, align the admin key removal route, and satisfy the enabled accessibility rules.

Co-authored-by: Codex <noreply@openai.com>
2026-07-17 11:44:04 +08:00
zyycnandCodex 29c8b679ba refactor(architecture): complete dependency boundary audit
Harden WebSocket ownership and capacity semantics, split the transport into focused modules, and move admin queries, quota observations, and upstream adapters behind explicit consumer-owned ports.

Parse frontend payloads at the API boundary, centralize SSE and store ownership, add latest-request-wins queries with inferred view types, and enforce dependency direction through lint and AST tests.

Co-authored-by: Codex <noreply@openai.com>
2026-07-17 09:03:33 +08:00
zyycnandCodex 4091a0c2d8 feat(proxy): prefer passive upstream state refresh
Remove request-path quota probes and consume account, request, and rate-limit facts from normal upstream responses, including 429 active-limit metadata.

Refresh model catalogs only at startup, on x-models-etag changes, or through the explicit account model refresh control in the connection test dialog.

Co-authored-by: Codex <noreply@openai.com>
2026-07-17 02:12:16 +08:00
zyycnandCodex f0ccc981fb feat(proxy): align latest Responses semantics
Preserve Responses Lite and Memgen controls across HTTP and WebSocket transports without rewriting request bodies, and standardize continuation recovery signals.

Track cache-write usage and billing while applying appcast profile updates to new WebSocket generations without breaking exact continuations.

Co-authored-by: Codex <noreply@openai.com>
2026-07-17 00:50:27 +08:00
zyycnandCodex d01ea7f775 fix(dispatch): generalize cross-account replay sanitization
Sanitize every replay item envelope without maintaining an output-type allowlist, while preserving nested client and tool semantics.

Build connection snapshots from the actual successful attempt input so cleaned source-account state cannot return during target-account recovery.

Co-authored-by: Codex <noreply@openai.com>
2026-07-16 23:59:10 +08:00
zyycnandCodex fd749ea8f8 fix(dispatch): isolate cross-account replay state
Track connection transcripts by the actual response account and sanitize typed Responses output state before replaying across accounts.

Rebuild target auth, Cookie, installation metadata, and turn state boundaries, fail closed on malformed turn metadata, and persist only target-account state after replay.

Co-authored-by: Codex <noreply@openai.com>
2026-07-16 23:18:26 +08:00
zyycnandCodex f0543b637e fix(frontend): refine analytics interactions
Align usage chart bounds and diagnostic loading states, simplify dashboard hover feedback, and improve User-Agent layout.

Restore the login light palette to the primary accent color.

Co-authored-by: Codex <noreply@openai.com>
2026-07-16 20:03:51 +08:00
ZyyandCodex 549bb9f62b feat(telemetry): expose latency breakdowns
Capture first semantic output and transport timing, then surface concise latency values with safe phase details in dashboard and usage records.

Keep the Codex Desktop request profile synchronized with successful appcast checks at runtime, without rewriting configuration or evicting existing WebSocket connections.

Co-authored-by: Codex <noreply@openai.com>
2026-07-16 18:07:08 +08:00
ZyyandCodex 9e57f78bab ci: speed up validation pipelines
Run affected CI checks in parallel and skip unrelated backend or frontend work while retaining workflow lint.

Start release binary builds after release metadata is available, while keeping image assembly, packaging, and publication gated by quality checks.

Co-authored-by: Codex <noreply@openai.com>
2026-07-16 16:07:55 +08:00
ZyyandCodex 73c0ecdc4c feat(telemetry): add first-token latency observability
Measure TTFT from the first semantic response event across streaming, retry, scheduler, usage telemetry, and dashboard paths.

Expose first-token percentiles and concise UI labels, and isolate concurrent Redis-backed test runs.

Co-authored-by: Codex <noreply@openai.com>
2026-07-16 15:57:19 +08:00
333 changed files with 21046 additions and 10891 deletions
+12
View File
@@ -12,12 +12,23 @@ on:
required: false
default: false
type: boolean
run_backend:
description: Run backend quality checks
required: false
default: true
type: boolean
run_frontend:
description: Run frontend quality checks
required: false
default: true
type: boolean
permissions:
contents: read
jobs:
backend:
if: inputs.run_backend
runs-on: ubuntu-latest
timeout-minutes: 30
env:
@@ -77,6 +88,7 @@ jobs:
run: cargo test --test main --locked
frontend:
if: inputs.run_frontend
runs-on: ubuntu-latest
timeout-minutes: 15
env:
+14 -2
View File
@@ -20,6 +20,8 @@ jobs:
runs-on: ubuntu-latest
timeout-minutes: 5
outputs:
backend: ${{ steps.filter.outputs.backend }}
frontend: ${{ steps.filter.outputs.frontend }}
container: ${{ steps.filter.outputs.container }}
steps:
- name: Checkout
@@ -28,11 +30,18 @@ jobs:
ref: ${{ github.event.pull_request.head.sha || github.sha }}
fetch-depth: 0
- name: Detect container changes
- name: Detect affected checks
id: filter
uses: dorny/paths-filter@7b450fff21473bca461d4b92ce414b9d0420d706 # v4.0.2
with:
filters: |
backend:
- 'backend/**'
- 'release/version.yaml'
- '.github/workflows/**'
frontend:
- 'frontend/**'
- '.github/workflows/**'
container:
- 'backend/**'
- 'frontend/**'
@@ -43,12 +52,15 @@ jobs:
- '.github/workflows/ci.yml'
quality:
needs: changes
uses: ./.github/workflows/_quality.yml
with:
ref: ${{ github.event.pull_request.head.sha || github.sha }}
run_backend: ${{ needs.changes.outputs.backend == 'true' }}
run_frontend: ${{ needs.changes.outputs.frontend == 'true' }}
container:
needs: [changes, quality]
needs: changes
if: needs.changes.outputs.container == 'true'
permissions:
contents: read
+2 -2
View File
@@ -165,7 +165,7 @@ jobs:
build-docker-binary:
name: build-docker-binary (${{ matrix.os }}-${{ matrix.arch }})
needs: [release-info, quality]
needs: release-info
strategy:
fail-fast: false
matrix: ${{ fromJson(needs.release-info.outputs.docker_matrix) }}
@@ -225,7 +225,7 @@ jobs:
build-binary:
name: build-binary (${{ matrix.os }}-${{ matrix.arch }})
needs: [release-info, quality]
needs: release-info
strategy:
fail-fast: false
matrix: ${{ fromJson(needs.release-info.outputs.non_docker_matrix) }}
+1
View File
@@ -382,6 +382,7 @@ dependencies = [
"sha2 0.11.0",
"sqlx",
"subtle",
"syn",
"tar",
"tempfile",
"thiserror",
+1
View File
@@ -54,6 +54,7 @@ uuid = { version = "1.23.3", features = ["v4", "v7", "serde"] }
[dev-dependencies]
sqlx = { version = "0.9.0", default-features = false, features = ["derive"] }
syn = { version = "2.0.118", features = ["full", "visit"] }
tempfile = "3.27.0"
tokio = { version = "1.52.3", default-features = false, features = ["test-util"] }
wiremock = "0.6.5"
+41 -106
View File
@@ -10,9 +10,7 @@ use axum::{
},
response::{IntoResponse, Response},
};
use std::collections::{HashMap, HashSet};
use chrono::{DateTime, Utc};
use chrono::Utc;
use serde::{Deserialize, Serialize};
use serde_json::Value;
@@ -30,11 +28,7 @@ use crate::{
manage::{
AccountHealthCheck, AccountManageError, AccountRefreshOutcome, AccountRefreshResult,
AccountUpdate, ManagedAccount, OAuthExchangeInput,
quota_view::{
AccountQuotaData, AccountQuotaUsageWindow, AccountQuotaWindowLocalUsage, quota_data,
},
},
refresh::token_refresh_status_eligible,
},
infra::{
format::{
@@ -44,14 +38,9 @@ use crate::{
json::{SortDirection, clamp_limit, clamp_page, total_pages},
time::{china_datetime, china_relative_time, china_rfc3339},
},
telemetry::{
account_usage::query::AccountUsageRecord,
billing,
buckets::query::{ModelUsageWindow, UsageBucketWindow},
},
telemetry::account_usage::query::AccountUsageRecord,
};
const ACCOUNT_STATS_PAGE_LIMIT: u32 = 200;
const ACCOUNT_EXPORT_CONFIRMATION: &str = "export_sensitive_accounts";
#[derive(Debug, Deserialize)]
@@ -195,21 +184,14 @@ impl AdminAccountData {
struct AdminAccountUsageData {
request_count: i64,
request_count_display: String,
empty_response_count: i64,
input_tokens: i64,
input_tokens_display: String,
output_tokens: i64,
output_tokens_display: String,
cached_tokens: i64,
cached_tokens_display: String,
reasoning_tokens: i64,
total_tokens: i64,
total_tokens_display: String,
image_input_tokens: i64,
image_output_tokens: i64,
image_tokens_display: String,
image_request_count: i64,
image_request_failed_count: i64,
created_tokens: i64,
created_tokens_display: String,
read_tokens: i64,
@@ -225,16 +207,10 @@ impl AdminAccountUsageData {
models: Vec<AdminAccountModelUsageData>,
) -> Self {
let request_count = usage.map_or(0, |usage| usage.window_request_count);
let empty_response_count = 0;
let input_tokens = usage.map_or(0, |usage| usage.window_input_tokens);
let output_tokens = usage.map_or(0, |usage| usage.window_output_tokens);
let cached_tokens = usage.map_or(0, |usage| usage.window_cached_tokens);
let reasoning_tokens = 0;
let total_tokens = input_tokens.saturating_add(output_tokens);
let image_input_tokens = 0;
let image_output_tokens = 0;
let image_request_count = 0;
let image_request_failed_count = 0;
let created_tokens = input_tokens.saturating_sub(cached_tokens);
let read_tokens = cached_tokens;
let last_used_at = usage.and_then(|usage| usage.last_used_at);
@@ -242,23 +218,14 @@ impl AdminAccountUsageData {
Self {
request_count,
request_count_display: format_plain_number(nonnegative_i64_to_u64(request_count)),
empty_response_count,
input_tokens,
input_tokens_display: format_tokens(nonnegative_i64_to_u64(input_tokens)),
output_tokens,
output_tokens_display: format_tokens(nonnegative_i64_to_u64(output_tokens)),
cached_tokens,
cached_tokens_display: format_tokens(nonnegative_i64_to_u64(cached_tokens)),
reasoning_tokens,
total_tokens,
total_tokens_display: format_tokens(nonnegative_i64_to_u64(total_tokens)),
image_input_tokens,
image_output_tokens,
image_tokens_display: format_tokens(nonnegative_i64_to_u64(
image_input_tokens + image_output_tokens,
)),
image_request_count,
image_request_failed_count,
created_tokens,
created_tokens_display: format_tokens(nonnegative_i64_to_u64(created_tokens)),
read_tokens,
@@ -292,50 +259,6 @@ struct AdminAccountModelUsageData {
last_used_at_display: String,
}
#[derive(Debug, Clone, Default)]
struct AccountListStats {
usage_by_account: HashMap<String, AccountUsageRecord>,
quota_by_account: HashMap<String, AccountQuotaData>,
models_by_account: HashMap<String, Vec<AdminAccountModelUsageData>>,
refreshing_account_ids: HashSet<String>,
}
struct AccountModelUsageRecord {
account_id: String,
model: String,
request_count: i64,
error_count: i64,
input_tokens: i64,
output_tokens: i64,
cached_tokens: i64,
billing_amount_usd: f64,
last_used_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone)]
struct AccountQuotaWindowSelection {
account_id: String,
window: AccountQuotaUsageWindow,
}
impl AccountListStats {
fn data_for(&self, account: ManagedAccount) -> AdminAccountData {
let account_id = account.id.clone();
let token_refreshing = token_refresh_status_eligible(account.status)
&& self.refreshing_account_ids.contains(&account_id);
AdminAccountData::from_parts(
account,
self.usage_by_account.get(&account_id),
self.quota_by_account.get(&account_id).cloned(),
self.models_by_account
.get(&account_id)
.cloned()
.unwrap_or_default(),
token_refreshing,
)
}
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AdminAccountSummaryData {
@@ -345,6 +268,17 @@ struct AdminAccountSummaryData {
attention: u64,
}
impl From<AccountSummary> for AdminAccountSummaryData {
fn from(summary: AccountSummary) -> Self {
Self {
total: summary.total,
active: summary.active,
quota_exhausted: summary.quota_exhausted,
attention: summary.attention,
}
}
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AdminAccountPageData {
@@ -511,34 +445,33 @@ pub(crate) async fn accounts(
let page_size = clamp_limit(params.page_size.unwrap_or(50));
let status = account_status_filter(params.status)?;
let sort = account_list_sort(params.sort_by, params.sort_direction)?;
let quota_by_account = quota_snapshots_by_account(&state).await;
let summary = account_summary_data(&state).await;
match state
let result = state
.services
.admin_accounts
.list_page(page, page_size, params.search, status, sort)
.account_list
.query(AccountListQuery {
page,
page_size,
search: params.search,
status,
sort,
})
.await
{
Ok(page) => {
let stats = account_list_stats(&state, &page.items, &quota_by_account).await?;
let page = crate::infra::json::NumberedPage {
items: page
.items
.into_iter()
.map(|item| stats.data_for(item))
.collect(),
total: page.total,
page: page.page,
page_size: page.page_size,
};
Ok(AdminResponse::new(
StatusCode::OK,
AdminAccountPageEnvelope::ok(page, summary),
))
}
Err(error) => Err(account_error(&error)),
}
.map_err(|error| AdminError::internal(error.to_string()))?;
let page = crate::infra::json::NumberedPage {
items: result
.page
.items
.into_iter()
.map(account_list_item_data)
.collect(),
total: result.page.total,
page: result.page.page,
page_size: result.page.page_size,
};
Ok(AdminResponse::new(
StatusCode::OK,
AdminAccountPageEnvelope::ok(page, result.summary.into()),
))
}
/// `POST /api/admin/accounts`
@@ -566,7 +499,8 @@ mod import_routes;
mod lifecycle_routes;
mod oauth_routes;
mod probe_routes;
mod query;
pub(crate) mod query;
mod quota_presenter;
mod quota_routes;
pub(crate) use export_routes::*;
@@ -575,4 +509,5 @@ pub(crate) use lifecycle_routes::*;
pub(crate) use oauth_routes::*;
pub(crate) use probe_routes::*;
use query::*;
use quota_presenter::*;
pub(crate) use quota_routes::*;
@@ -1,4 +1,11 @@
use super::*;
use std::convert::Infallible;
use bytes::Bytes;
use futures::StreamExt;
use serde_json::json;
use crate::fleet::manage::AccountTestEvent;
/// `GET /api/admin/accounts/test?id=...&modelId=...`
pub async fn test_account_connection(
@@ -21,6 +28,12 @@ pub async fn test_account_connection(
.test_connection_stream(&account_id, model)
.await
.map_err(|error| account_error(&error))?;
let stream = stream.map(|event| {
Ok::<Bytes, Infallible>(Bytes::from(format!(
"data: {}\n\n",
account_test_event_json(event)
)))
});
Response::builder()
.status(StatusCode::OK)
@@ -32,6 +45,43 @@ pub async fn test_account_connection(
.map_err(|_| AdminError::internal("Failed to build account test stream"))
}
fn account_test_event_json(event: AccountTestEvent) -> Value {
match event {
AccountTestEvent::Started { model } => json!({
"type": "test_start",
"model": model,
"text": "正在连接 Codex Responses"
}),
AccountTestEvent::Request { payload } => json!({
"type": "request",
"payload": payload
}),
AccountTestEvent::Content { text } => json!({ "type": "content", "text": text }),
AccountTestEvent::Complete { account_status } => account_test_terminal_event(
json!({ "type": "test_complete", "success": true }),
account_status,
),
AccountTestEvent::Error {
error,
account_status,
} => {
account_test_terminal_event(json!({ "type": "error", "error": error }), account_status)
}
}
}
fn account_test_terminal_event(mut event: Value, account_status: Option<AccountStatus>) -> Value {
if let Some(account_status) = account_status
&& let Some(event) = event.as_object_mut()
{
event.insert(
"accountStatus".to_string(),
Value::String(account_status.as_str().to_string()),
);
}
event
}
/// `GET /api/admin/accounts/models?id=...`
pub async fn account_models(
State(state): State<AppState>,
@@ -60,3 +110,29 @@ pub async fn account_models(
Err(error) => Err(account_error(&error)),
}
}
/// `POST /api/admin/accounts/models`
pub async fn refresh_account_models(
State(state): State<AppState>,
_auth: AdminAuth,
Json(payload): Json<AccountActionRequest>,
) -> Result<impl IntoResponse, AdminError> {
let models = state
.services
.admin_accounts
.refresh_account_models(&payload.id)
.await
.map_err(|error| account_error(&error))?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(AccountModelsData {
models: models
.into_iter()
.map(|model| AccountModelData {
id: model.id,
label: model.label,
})
.collect(),
}),
))
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,211 @@
//! typed quota read model 的管理端 presenter。
use chrono::{DateTime, Utc};
use serde::Serialize;
use crate::{
api::admin::accounts_routes::query::{AccountQuotaReadModel, AccountQuotaWindowLocalUsage},
fleet::quota::{QuotaWindowKind, QuotaWindowRole},
infra::{
format::{
format_compact_percent, format_plain_number, format_tokens, nonnegative_i64_to_u64,
},
time::{china_datetime, china_relative_time},
},
};
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct AccountQuotaData {
refreshed_at_display: String,
windows: Vec<AccountQuotaWindowData>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountQuotaWindowData {
key: String,
group: &'static str,
window_seconds: Option<u64>,
label_display: String,
used_percent: Option<f64>,
used_percent_display: String,
#[serde(skip_serializing_if = "Option::is_none")]
local_usage: Option<AccountQuotaWindowLocalUsageData>,
reset_at_display: String,
window_used_display: String,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountQuotaWindowLocalUsageData {
request_count: i64,
request_count_display: String,
input_tokens: i64,
input_tokens_display: String,
output_tokens: i64,
output_tokens_display: String,
cached_tokens: i64,
cached_tokens_display: String,
total_tokens: i64,
total_tokens_display: String,
}
impl From<AccountQuotaWindowLocalUsage> for AccountQuotaWindowLocalUsageData {
fn from(usage: AccountQuotaWindowLocalUsage) -> Self {
let total_tokens = usage.input_tokens.saturating_add(usage.output_tokens);
Self {
request_count: usage.request_count,
request_count_display: format_plain_number(nonnegative_i64_to_u64(usage.request_count)),
input_tokens: usage.input_tokens,
input_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.input_tokens)),
output_tokens: usage.output_tokens,
output_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.output_tokens)),
cached_tokens: usage.cached_tokens,
cached_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.cached_tokens)),
total_tokens,
total_tokens_display: format_tokens(nonnegative_i64_to_u64(total_tokens)),
}
}
}
impl Default for AccountQuotaData {
fn default() -> Self {
Self {
refreshed_at_display: "-".to_string(),
windows: Vec::new(),
}
}
}
pub(super) fn quota_data(read: AccountQuotaReadModel) -> AccountQuotaData {
let now = Utc::now();
AccountQuotaData {
refreshed_at_display: china_relative_time(read.fetched_at, now),
windows: read
.windows
.into_iter()
.map(|read| {
let window = read.quota.window;
let window_seconds = window.window_seconds();
let reset_at = window.reset_datetime();
let used_percent = window.used_percent;
AccountQuotaWindowData {
key: read.quota.key,
group: quota_window_group(window.kind),
window_seconds,
label_display: quota_window_label(
read.quota.source.as_str(),
read.quota.limit_name.as_deref(),
read.quota.metered_feature.as_deref(),
read.quota.role,
window.kind,
window_seconds,
),
used_percent,
used_percent_display: used_percent
.map_or_else(|| "-".to_string(), format_compact_percent),
local_usage: read.local_usage.map(Into::into),
reset_at_display: reset_at
.as_ref()
.map_or_else(|| "-".to_string(), china_datetime),
window_used_display: quota_window_used_display(reset_at, window_seconds, now),
}
})
.collect(),
}
}
fn quota_window_group(kind: QuotaWindowKind) -> &'static str {
match kind {
QuotaWindowKind::Monthly => "monthly",
QuotaWindowKind::ShortTerm | QuotaWindowKind::Weekly => "shortTerm",
QuotaWindowKind::Other => "other",
}
}
fn quota_window_label(
source: &str,
limit_name: Option<&str>,
metered_feature: Option<&str>,
role: QuotaWindowRole,
kind: QuotaWindowKind,
window_seconds: Option<u64>,
) -> String {
if role == QuotaWindowRole::Monthly {
return "月限额".to_string();
}
let base = match kind {
QuotaWindowKind::ShortTerm => "5小时限额".to_string(),
QuotaWindowKind::Weekly => "周限额".to_string(),
QuotaWindowKind::Monthly => "月限额".to_string(),
QuotaWindowKind::Other => custom_window_label(window_seconds),
};
let label = limit_name
.or(metered_feature)
.map(str::trim)
.filter(|value| !value.is_empty() && !value.eq_ignore_ascii_case("codex"));
if source == "core" || label.is_none() {
return base;
}
let label = label.unwrap_or_default();
let label = if is_review_limit_label(label) {
"代码审查"
} else {
label
};
format!("{label} · {base}")
}
fn custom_window_label(window_seconds: Option<u64>) -> String {
let Some(seconds) = window_seconds.filter(|seconds| *seconds > 0) else {
return "额度".to_string();
};
if seconds % 86_400 == 0 {
format!("{}天限额", seconds / 86_400)
} else if seconds % 3_600 == 0 {
format!("{}小时限额", seconds / 3_600)
} else {
format!("{}分钟限额", seconds.div_ceil(60))
}
}
fn quota_window_used_display(
reset_at: Option<DateTime<Utc>>,
window_seconds: Option<u64>,
now: DateTime<Utc>,
) -> String {
let (Some(reset_at), Some(window_seconds)) = (reset_at, window_seconds) else {
return "-".to_string();
};
let remaining = reset_at
.signed_duration_since(now)
.num_seconds()
.max(0)
.cast_unsigned();
let used = window_seconds.saturating_sub(remaining);
format!(
"{} / {}",
format_duration_days(used),
format_duration_days(window_seconds)
)
}
fn format_duration_days(seconds: u64) -> String {
let days = seconds as f64 / 86_400.0;
if days >= 1.0 {
format!("{days:.1}d")
} else {
format!("{:.1}h", seconds as f64 / 3_600.0)
}
}
fn is_review_limit_label(value: &str) -> bool {
let normalized = value.trim().to_ascii_lowercase().replace(['-', ' '], "_");
normalized == "review"
|| normalized == "code_review"
|| normalized == "codex_review"
|| normalized == "codex_code_review"
|| normalized.contains("code_review")
|| normalized.contains("codex_review")
}
@@ -14,13 +14,25 @@ pub async fn account_quota(
.await
{
Ok(data) => {
let quota = data.get("quota").cloned().unwrap_or(Value::Null);
let raw = data.get("raw").cloned().unwrap_or(Value::Null);
let quota_json = quota.to_string();
let mut quota_data = quota_data(&quota_json, Some(Utc::now()));
apply_account_quota_window_local_usage(&state, &account_id, &mut quota_data).await;
let account =
account_data_for_quota_refresh(&state, &account_id, quota_data.clone()).await?;
let quota = serde_json::to_value(&data.quota)
.map_err(|error| AdminError::internal(error.to_string()))?;
let raw = data.raw;
let quota_read = AccountQuotaReadModel::from_snapshot(data.quota, Some(Utc::now()));
let account = state
.services
.admin_accounts
.get(&account_id)
.await
.map_err(|error| account_error(&error))?
.ok_or_else(account_not_found)?;
let item = state
.services
.account_list
.enrich_account(account, quota_read)
.await
.map_err(|error| AdminError::internal(error.to_string()))?;
let quota_data = item.quota.clone().map(quota_data).unwrap_or_default();
let account = account_list_item_data(item);
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(AdminAccountQuotaResponseData::from_account(
@@ -39,23 +51,3 @@ pub async fn account_quota(
Err(e) => Err(account_error(&e)),
}
}
async fn account_data_for_quota_refresh(
state: &AppState,
account_id: &str,
quota_data: AccountQuotaData,
) -> Result<AdminAccountData, AdminError> {
let Some(account) = state
.services
.admin_accounts
.get(account_id)
.await
.map_err(|error| account_error(&error))?
else {
return Err(account_not_found());
};
let quota_by_account = HashMap::from([(account_id.to_string(), quota_data)]);
let stats =
account_list_stats(state, std::slice::from_ref(&account), &quota_by_account).await?;
Ok(stats.data_for(account))
}
@@ -1,4 +1,4 @@
//! Dashboard 遥测聚合与展示模型。
//! Dashboard 遥测事实的 HTTP presenter。
use chrono::{DateTime, Duration, Utc};
use serde::{Deserialize, Serialize};
@@ -11,8 +11,8 @@ use crate::infra::{
time::{china_day_start, china_quarter_hour_start},
};
use super::account_usage::query::{AccountUsageTimeBucket, RetainedUsageSummary};
use super::usage::insights::RequestHealthTimeBucket;
use crate::telemetry::account_usage::query::{AccountUsageTimeBucket, RetainedUsageSummary};
use crate::telemetry::usage::insights::RequestHealthTimeBucket;
const HEALTH_TIMELINE_SLOT_MINUTES: i64 = 15;
const HEALTH_TIMELINE_SLOTS: i64 = 24 * 4;
+421 -299
View File
@@ -1,6 +1,6 @@
//! 管理端 Dashboard 聚合视图。
use std::collections::{HashMap, HashSet};
use std::{collections::HashMap, sync::Arc};
use axum::{
extract::{Query, State},
@@ -10,10 +10,17 @@ use axum::{
use chrono::{DateTime, Duration, Utc};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;
use crate::{
api::AppState,
api::admin::{
accounts_routes::query::RefreshActivityQuery,
dashboard_presenter::{
DashboardAccountCounts, DashboardCardsData, DashboardHealthTimelineData,
DashboardTrendData, DashboardTrendKind, dashboard_cards,
dashboard_health_timeline_data, dashboard_trend_data,
},
response::{AdminEnvelope, AdminError, AdminResponse},
session::AdminAuth,
usage_routes::{
@@ -23,36 +30,28 @@ use crate::{
},
fleet::{
account::{Account, AccountStatus},
pool::AccountCapacitySummary,
pool::{AccountCapacitySummary, AccountPoolService},
quota::QuotaSnapshot,
refresh::token_refresh_status_eligible,
store::AccountStore,
},
infra::{
format::{format_duration_ms, format_tokens},
time::{china_datetime, china_day_start, china_quarter_hour_start, china_relative_time},
},
settings::service::SettingsService,
telemetry::{
account_usage::query::AccountUsageTimeBucket,
dashboard::{
DashboardAccountCounts, DashboardCardsData, DashboardHealthTimelineData,
DashboardTrendData, DashboardTrendKind, dashboard_cards,
dashboard_health_timeline_data, dashboard_trend_data,
account_usage::query::{
AccountUsageQueryService, AccountUsageTimeBucket, RetainedUsageSummary,
},
usage::{
insights::RequestHealthTimeBucket,
query::{UsageQueryFilter, UsageRecordAccountUsage},
types::{UsageRecord, metadata_string},
query::{UsageQueryFilter, UsageQueryService},
types::{UsageRecord, metadata_i64, metadata_string},
},
},
};
const DASHBOARD_ACCOUNT_USAGE_LIMIT: u32 = 4;
const DASHBOARD_USAGE_RECORD_LIMIT: u32 = 10;
const DASHBOARD_TIME_BUCKET_MINUTES: i64 = 15;
const DASHBOARD_TIME_BUCKET_SLOTS: i64 = 7 * 24 * 4;
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")]
pub(crate) struct DashboardTrendQuery {
@@ -105,10 +104,25 @@ struct DashboardUsageRecordData {
subagent_kind: Option<String>,
token_details: UsageRecordTokenDetailsData,
billing: Option<UsageRecordBillingData>,
first_token_latency_ms: Option<i64>,
latency_ms: Option<i64>,
latency_details: DashboardUsageLatencyDetailsData,
first_token_latency_ms_display: String,
latency_ms_display: String,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardUsageLatencyDetailsData {
first_reasoning_ms: Option<i64>,
first_text_ms: Option<i64>,
transport_decision_wait_ms: Option<i64>,
ws_connect_ms: Option<i64>,
upstream_headers_ms: Option<i64>,
first_event_ms: Option<i64>,
openai_processing_ms: Option<i64>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardWireProfileData {
@@ -178,6 +192,20 @@ struct AccountCapacityDiagnostics {
available_slots: usize,
}
impl From<DashboardPoolSummary> for AccountPoolDiagnostics {
fn from(summary: DashboardPoolSummary) -> Self {
Self {
total: summary.total,
active: summary.active,
expired: summary.expired,
quota_exhausted: summary.quota_exhausted,
refreshing: summary.refreshing,
disabled: summary.disabled,
banned: summary.banned,
}
}
}
impl From<AccountCapacitySummary> for AccountCapacityDiagnostics {
fn from(summary: AccountCapacitySummary) -> Self {
Self {
@@ -189,117 +217,59 @@ impl From<AccountCapacitySummary> for AccountCapacityDiagnostics {
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
enum QuotaWindowPriority {
FiveHour,
Weekly,
Monthly,
Other,
}
/// `GET /api/admin/dashboard/summary`
pub(crate) async fn dashboard_summary(
State(state): State<AppState>,
_auth: AdminAuth,
Query(query): Query<DashboardTrendQuery>,
) -> Result<impl IntoResponse, AdminError> {
let accounts = state
.services
.accounts
.list_pool_accounts()
.await
.map_err(|error| dashboard_data_error("accounts", &error))?;
let capacity = state.services.account_pool.capacity_summary_now().await;
let now = Utc::now();
let today_filter = today_usage_record_filter(now);
let retained_usage = state
let dashboard = state
.services
.usage
.retained_summary()
.dashboard
.summary(now)
.await
.map_err(|error| dashboard_data_error("retained usage summary", &error))?;
let account_usage_records = state
.services
.usage_records
.account_usage(today_filter.clone(), DASHBOARD_ACCOUNT_USAGE_LIMIT)
.await
.map_err(|error| dashboard_data_error("account usage ranking", &error))?;
let recent_events = state
.services
.usage_records
.list_recent(DASHBOARD_USAGE_RECORD_LIMIT, today_filter)
.await
.map_err(|error| dashboard_data_error("recent usage records", &error))?;
let (time_buckets, health_buckets) = tokio::try_join!(
dashboard_time_buckets(&state, now),
dashboard_health_buckets(&state, now),
)?;
let account_ids = accounts
.iter()
.map(|account| account.id.clone())
.collect::<Vec<_>>();
let refreshing_account_ids = state
.services
.token_refresh
.refreshing_account_ids(&account_ids, now)
.await
.map_err(|error| dashboard_data_error("refreshing accounts", &error))?;
let pool_summary = account_pool_summary(&accounts, &refreshing_account_ids);
let trend = dashboard_trend_data(&time_buckets, query.kind.unwrap_or_default());
let quota_used_by_account =
account_quota_used_percent_by_id(&state, &account_usage_records).await?;
let settings = state.services.settings.current();
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()))?;
.map_err(dashboard_query_error)?;
let trend = dashboard_trend_data(&dashboard.time_buckets, query.kind.unwrap_or_default());
let dashboard_usage_records =
dashboard_usage_record_items(recent_usage_records, &account_emails);
dashboard_usage_record_items(dashboard.usage_records, &dashboard.account_emails);
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(DashboardSummaryData {
cards: dashboard_cards(
dashboard_account_counts(&accounts),
&time_buckets,
&retained_usage,
DashboardAccountCounts {
total: dashboard.account_counts.total,
enabled: dashboard.account_counts.enabled,
abnormal: dashboard.account_counts.abnormal,
},
&dashboard.time_buckets,
&dashboard.retained_usage,
),
trend,
health_timeline: dashboard_health_timeline_data(&health_buckets),
account_usage: account_usage_data(
&accounts,
&account_usage_records,
&quota_used_by_account,
now,
),
wire_profile: wire_profile_data(&state),
health_timeline: dashboard_health_timeline_data(&dashboard.health_buckets),
account_usage: account_usage_data(&dashboard.account_usage, now),
wire_profile: wire_profile_data(dashboard.wire_profile, dashboard.desktop_release),
usage_records: dashboard_usage_records,
pool_summary,
capacity_info: AccountCapacityDiagnostics::from(capacity),
rotation_strategy: Some(settings.rotation_strategy.clone()),
pool_summary: AccountPoolDiagnostics::from(dashboard.pool),
capacity_info: AccountCapacityDiagnostics::from(dashboard.capacity),
rotation_strategy: Some(dashboard.rotation_strategy),
}),
))
}
fn today_usage_record_filter(now: DateTime<Utc>) -> UsageQueryFilter {
UsageQueryFilter {
start_time: Some(china_day_start(now)),
end_time: Some(now),
..UsageQueryFilter::default()
}
}
/// `GET /api/admin/dashboard/trend?kind=usage|latency|errors`
pub(crate) async fn dashboard_trend(
State(state): State<AppState>,
_auth: AdminAuth,
Query(query): Query<DashboardTrendQuery>,
) -> Result<impl IntoResponse, AdminError> {
let time_buckets = dashboard_time_buckets(&state, Utc::now()).await?;
let time_buckets = state
.services
.dashboard
.time_buckets(Utc::now())
.await
.map_err(dashboard_query_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
@@ -310,213 +280,28 @@ pub(crate) async fn dashboard_trend(
))
}
async fn dashboard_time_buckets(
state: &AppState,
now: DateTime<Utc>,
) -> Result<Vec<AccountUsageTimeBucket>, AdminError> {
let current_slot = china_quarter_hour_start(now);
let start = current_slot
- Duration::minutes(DASHBOARD_TIME_BUCKET_MINUTES * (DASHBOARD_TIME_BUCKET_SLOTS - 1));
state
.services
.usage
.time_buckets(start, now)
.await
.map_err(|error| dashboard_data_error("time buckets", &error))
}
async fn dashboard_health_buckets(
state: &AppState,
now: DateTime<Utc>,
) -> Result<Vec<RequestHealthTimeBucket>, AdminError> {
state
.services
.usage_records
.health_timeline(china_day_start(now), now)
.await
.map_err(|error| dashboard_data_error("health timeline", &error))
}
fn dashboard_account_counts(accounts: &[Account]) -> DashboardAccountCounts {
DashboardAccountCounts {
total: accounts.len() as u64,
enabled: accounts
.iter()
.filter(|account| account.status == AccountStatus::Active)
.count() as u64,
abnormal: accounts
.iter()
.filter(|account| account.status != AccountStatus::Active)
.count() as u64,
}
}
fn account_pool_summary(
accounts: &[Account],
refreshing_account_ids: &HashSet<String>,
) -> AccountPoolDiagnostics {
let mut summary = AccountPoolDiagnostics {
total: accounts.len(),
..AccountPoolDiagnostics::default()
};
for account in accounts {
match account.status {
status
if token_refresh_status_eligible(status)
&& refreshing_account_ids.contains(&account.id) =>
{
summary.refreshing += 1;
}
AccountStatus::Active => summary.active += 1,
AccountStatus::Expired => summary.expired += 1,
AccountStatus::QuotaExhausted => summary.quota_exhausted += 1,
AccountStatus::Disabled => summary.disabled += 1,
AccountStatus::Banned => summary.banned += 1,
}
}
summary
}
fn account_usage_data(
accounts: &[Account],
usage_records: &[UsageRecordAccountUsage],
quota_used_by_account: &HashMap<String, f64>,
usage_records: &[DashboardAccountUsage],
now: DateTime<Utc>,
) -> Vec<DashboardAccountUsageData> {
let account_by_id = accounts
.iter()
.map(|account| (account.id.as_str(), account))
.collect::<HashMap<_, _>>();
usage_records
.iter()
.map(|usage| {
let account = account_by_id.get(usage.account_id.as_str()).copied();
let quota_used_percent = quota_used_by_account
.get(&usage.account_id)
.copied()
.or_else(|| {
matches!(
account.map(|account| account.status),
Some(AccountStatus::QuotaExhausted)
)
.then_some(100.0)
});
DashboardAccountUsageData {
id: usage.account_id.clone(),
email: account
.and_then(|account| account.email.as_ref())
.cloned()
.unwrap_or_else(|| usage.account_id.clone()),
plan_type: account
.and_then(|account| account.plan_type.as_ref())
.cloned(),
tokens: format_tokens(usage.total_tokens),
quota_used_percent,
last_used: china_relative_time(Some(usage.last_used_at), now),
}
.map(|usage| DashboardAccountUsageData {
id: usage.id.clone(),
email: usage.email.clone(),
plan_type: usage.plan_type.clone(),
tokens: format_tokens(usage.total_tokens),
quota_used_percent: usage.quota_used_percent,
last_used: china_relative_time(Some(usage.last_used_at), now),
})
.collect()
}
async fn account_quota_used_percent_by_id(
state: &AppState,
usage_records: &[UsageRecordAccountUsage],
) -> Result<HashMap<String, f64>, AdminError> {
let mut quota_used_by_account = HashMap::with_capacity(usage_records.len());
for usage in usage_records {
let quota_json = state
.services
.accounts
.get_quota_json(&usage.account_id)
.await
.map_err(|error| dashboard_data_error("account quota", &error))?;
let Some(quota_json) = quota_json else {
continue;
};
if let Some(used_percent) = quota_used_percent(&quota_json) {
quota_used_by_account.insert(usage.account_id.clone(), used_percent);
}
}
Ok(quota_used_by_account)
}
fn dashboard_data_error(source: &'static str, error: &impl std::fmt::Display) -> AdminError {
tracing::error!(source, error = %error, "Failed to load dashboard data");
fn dashboard_query_error(error: DashboardQueryError) -> AdminError {
tracing::error!(error = %error, "Failed to load dashboard data");
AdminError::internal("Failed to load dashboard data")
}
fn quota_used_percent(quota_json: &str) -> Option<f64> {
let quota = serde_json::from_str::<Value>(quota_json).ok()?;
let mut selected = None;
if let Some(monthly_limit) = quota.get("monthly_limit") {
select_quota_used_percent(
&mut selected,
quota_window_priority(monthly_limit, QuotaWindowPriority::Monthly),
monthly_limit.get("used_percent").and_then(percent_value),
);
}
if let Some(snapshots) = quota.get("snapshots").and_then(Value::as_array) {
for snapshot in snapshots {
for role in ["primary", "secondary"] {
if let Some(window) = snapshot.get(role) {
select_quota_used_percent(
&mut selected,
quota_window_priority(window, QuotaWindowPriority::Other),
window.get("used_percent").and_then(percent_value),
);
}
}
}
}
selected.map(|(_, used_percent)| used_percent)
}
fn select_quota_used_percent(
selected: &mut Option<(QuotaWindowPriority, f64)>,
priority: QuotaWindowPriority,
used_percent: Option<f64>,
) {
let Some(used_percent) = used_percent else {
return;
};
match selected {
Some((selected_priority, selected_percent))
if priority > *selected_priority
|| (priority == *selected_priority && used_percent <= *selected_percent) => {}
_ => *selected = Some((priority, used_percent)),
}
}
fn quota_window_priority(window: &Value, fallback: QuotaWindowPriority) -> QuotaWindowPriority {
let window_seconds = window
.get("window_minutes")
.and_then(Value::as_u64)
.and_then(|minutes| minutes.checked_mul(60));
match window_seconds {
Some(seconds) if quota_window_matches(seconds, FIVE_HOUR_WINDOW_SECONDS) => {
QuotaWindowPriority::FiveHour
}
Some(seconds) if quota_window_matches(seconds, WEEK_WINDOW_SECONDS) => {
QuotaWindowPriority::Weekly
}
Some(seconds) if quota_window_matches(seconds, MONTH_WINDOW_SECONDS) => {
QuotaWindowPriority::Monthly
}
_ => fallback,
}
}
fn quota_window_matches(actual: u64, expected: u64) -> bool {
actual > 0 && actual.abs_diff(expected) <= expected / 20
}
fn percent_value(value: &Value) -> Option<f64> {
let percent = value
.as_f64()
.or_else(|| value.as_str().and_then(|value| value.parse::<f64>().ok()))?;
percent.is_finite().then_some(percent.clamp(0.0, 100.0))
}
fn dashboard_usage_record_items(
items: Vec<UsageRecord>,
account_emails: &HashMap<String, String>,
@@ -551,6 +336,7 @@ fn dashboard_usage_record_data(
let token_details = usage_token_details(&record);
let billing = usage_billing(&record, upstream_model.as_deref(), &token_details);
let first_token_latency_ms = record.first_token_ms;
let latency_details = dashboard_usage_latency_details(&record.metadata);
DashboardUsageRecordData {
id: record.id,
@@ -571,14 +357,34 @@ fn dashboard_usage_record_data(
subagent_kind,
token_details,
billing,
first_token_latency_ms,
latency_ms: record.latency_ms,
latency_details,
first_token_latency_ms_display: format_duration_ms(first_token_latency_ms),
latency_ms_display: format_duration_ms(record.latency_ms),
}
}
fn wire_profile_data(state: &AppState) -> DashboardWireProfileData {
let profile = &state.services.wire_profile;
let release_snapshot = state.services.desktop_release.snapshot();
fn dashboard_usage_latency_details(metadata: &Value) -> DashboardUsageLatencyDetailsData {
DashboardUsageLatencyDetailsData {
first_reasoning_ms: metadata_duration_ms(metadata, "firstReasoningMs"),
first_text_ms: metadata_duration_ms(metadata, "firstTextMs"),
transport_decision_wait_ms: metadata_duration_ms(metadata, "transportDecisionWaitMs"),
ws_connect_ms: metadata_duration_ms(metadata, "wsConnectMs"),
upstream_headers_ms: metadata_duration_ms(metadata, "upstreamHeadersMs"),
first_event_ms: metadata_duration_ms(metadata, "firstEventMs"),
openai_processing_ms: metadata_duration_ms(metadata, "openaiProcessingMs"),
}
}
fn metadata_duration_ms(metadata: &Value, key: &str) -> Option<i64> {
metadata_i64(metadata, &[key]).filter(|value| *value >= 0)
}
fn wire_profile_data(
profile: DashboardWireProfile,
release_snapshot: DashboardDesktopReleaseSnapshot,
) -> DashboardWireProfileData {
let release_status = if release_snapshot.last_error.is_some() {
"check_failed"
} else if let Some(latest) = &release_snapshot.latest {
@@ -603,7 +409,7 @@ fn wire_profile_data(state: &AppState) -> DashboardWireProfileData {
arch: profile.arch.clone(),
terminal: profile.terminal.clone(),
},
user_agent: profile.user_agent(),
user_agent: profile.user_agent,
verified_at: profile.verified_at,
release: DashboardDesktopReleaseData {
status: release_status,
@@ -626,3 +432,319 @@ fn wire_profile_data(state: &AppState) -> DashboardWireProfileData {
},
}
}
const ACCOUNT_USAGE_LIMIT: u32 = 4;
const USAGE_RECORD_LIMIT: u32 = 10;
const TIME_BUCKET_MINUTES: i64 = 15;
const TIME_BUCKET_SLOTS: i64 = 7 * 24 * 4;
/// Dashboard 消费的请求画像快照端口。
pub(crate) trait DashboardWireProfileQuery: Send + Sync {
fn snapshot(&self) -> DashboardWireProfile;
}
/// Dashboard 消费的 Desktop 发布观测端口。
pub(crate) trait DashboardDesktopReleaseQuery: Send + Sync {
fn snapshot(&self) -> DashboardDesktopReleaseSnapshot;
}
#[derive(Debug, Clone)]
pub(crate) struct DashboardWireProfile {
pub(crate) originator: String,
pub(crate) codex_version: String,
pub(crate) desktop_version: String,
pub(crate) desktop_build: String,
pub(crate) os_type: String,
pub(crate) os_version: String,
pub(crate) arch: String,
pub(crate) terminal: String,
pub(crate) user_agent: String,
pub(crate) verified_at: DateTime<Utc>,
}
#[derive(Debug, Clone, Default)]
pub(crate) struct DashboardDesktopReleaseSnapshot {
pub(crate) checked_at: Option<DateTime<Utc>>,
pub(crate) latest: Option<DashboardDesktopRelease>,
pub(crate) last_error: Option<String>,
}
#[derive(Debug, Clone)]
pub(crate) struct DashboardDesktopRelease {
pub(crate) version: String,
pub(crate) build: String,
pub(crate) published_at: Option<DateTime<Utc>>,
pub(crate) minimum_system_version: Option<String>,
pub(crate) hardware_requirements: Option<String>,
pub(crate) download_url: Option<String>,
pub(crate) download_size: Option<u64>,
pub(crate) signature_present: bool,
}
#[derive(Debug, Clone)]
pub(crate) struct DashboardReadModel {
pub(crate) account_counts: DashboardQueryAccountCounts,
pub(crate) retained_usage: RetainedUsageSummary,
pub(crate) time_buckets: Vec<AccountUsageTimeBucket>,
pub(crate) health_buckets: Vec<RequestHealthTimeBucket>,
pub(crate) account_usage: Vec<DashboardAccountUsage>,
pub(crate) usage_records: Vec<UsageRecord>,
pub(crate) account_emails: HashMap<String, String>,
pub(crate) pool: DashboardPoolSummary,
pub(crate) capacity: AccountCapacitySummary,
pub(crate) rotation_strategy: String,
pub(crate) wire_profile: DashboardWireProfile,
pub(crate) desktop_release: DashboardDesktopReleaseSnapshot,
}
#[derive(Debug, Clone, Copy, Default)]
pub(crate) struct DashboardQueryAccountCounts {
pub(crate) total: u64,
pub(crate) enabled: u64,
pub(crate) abnormal: u64,
}
#[derive(Debug, Clone)]
pub(crate) struct DashboardAccountUsage {
pub(crate) id: String,
pub(crate) email: String,
pub(crate) plan_type: Option<String>,
pub(crate) total_tokens: u64,
pub(crate) quota_used_percent: Option<f64>,
pub(crate) last_used_at: DateTime<Utc>,
}
#[derive(Debug, Clone, Copy, Default)]
pub(crate) struct DashboardPoolSummary {
pub(crate) total: usize,
pub(crate) active: usize,
pub(crate) expired: usize,
pub(crate) quota_exhausted: usize,
pub(crate) refreshing: usize,
pub(crate) disabled: usize,
pub(crate) banned: usize,
}
#[derive(Debug, Error)]
pub(crate) enum DashboardQueryError {
#[error("failed to load dashboard {0}")]
Data(&'static str),
}
#[derive(Clone)]
pub(crate) struct DashboardQueryService {
accounts: Arc<dyn AccountStore>,
usage: Arc<AccountUsageQueryService>,
usage_records: Arc<UsageQueryService>,
account_pool: Arc<AccountPoolService>,
refresh_activity: Arc<dyn RefreshActivityQuery>,
settings: Arc<SettingsService>,
wire_profile: Arc<dyn DashboardWireProfileQuery>,
desktop_release: Arc<dyn DashboardDesktopReleaseQuery>,
}
pub(crate) struct DashboardQueryServiceParts {
pub(crate) accounts: Arc<dyn AccountStore>,
pub(crate) usage: Arc<AccountUsageQueryService>,
pub(crate) usage_records: Arc<UsageQueryService>,
pub(crate) account_pool: Arc<AccountPoolService>,
pub(crate) refresh_activity: Arc<dyn RefreshActivityQuery>,
pub(crate) settings: Arc<SettingsService>,
pub(crate) wire_profile: Arc<dyn DashboardWireProfileQuery>,
pub(crate) desktop_release: Arc<dyn DashboardDesktopReleaseQuery>,
}
impl DashboardQueryService {
pub(crate) fn new(parts: DashboardQueryServiceParts) -> Self {
Self {
accounts: parts.accounts,
usage: parts.usage,
usage_records: parts.usage_records,
account_pool: parts.account_pool,
refresh_activity: parts.refresh_activity,
settings: parts.settings,
wire_profile: parts.wire_profile,
desktop_release: parts.desktop_release,
}
}
pub(crate) async fn summary(
&self,
now: DateTime<Utc>,
) -> Result<DashboardReadModel, DashboardQueryError> {
let accounts = self
.accounts
.list_pool_accounts()
.await
.map_err(|_| DashboardQueryError::Data("accounts"))?;
let capacity = self.account_pool.capacity_summary_now().await;
let filter = UsageQueryFilter {
start_time: Some(china_day_start(now)),
end_time: Some(now),
..UsageQueryFilter::default()
};
let retained_usage = self
.usage
.retained_summary()
.await
.map_err(|_| DashboardQueryError::Data("retained usage summary"))?;
let account_usage_records = self
.usage_records
.account_usage(filter.clone(), ACCOUNT_USAGE_LIMIT)
.await
.map_err(|_| DashboardQueryError::Data("account usage ranking"))?;
let usage_records = self
.usage_records
.list_recent(USAGE_RECORD_LIMIT, filter)
.await
.map_err(|_| DashboardQueryError::Data("recent usage records"))?;
let time_buckets = self.time_buckets(now).await?;
let health_buckets = self.health_buckets(now).await?;
let account_ids = accounts
.iter()
.map(|account| account.id.clone())
.collect::<Vec<_>>();
let refreshing = self
.refresh_activity
.refreshing_account_ids(&account_ids, now)
.await
.map_err(|_| DashboardQueryError::Data("refreshing accounts"))?;
let quota_used = self.quota_used_by_account(&account_usage_records).await?;
let account_emails = self
.usage_records
.account_email_map(&usage_records)
.await
.map_err(|_| DashboardQueryError::Data("usage record accounts"))?;
Ok(DashboardReadModel {
account_counts: account_counts(&accounts),
retained_usage,
time_buckets,
health_buckets,
account_usage: account_usage(&accounts, &account_usage_records, &quota_used),
usage_records,
account_emails,
pool: pool_summary(&accounts, &refreshing),
capacity,
rotation_strategy: self.settings.current().rotation_strategy,
wire_profile: self.wire_profile.snapshot(),
desktop_release: self.desktop_release.snapshot(),
})
}
pub(crate) async fn time_buckets(
&self,
now: DateTime<Utc>,
) -> Result<Vec<AccountUsageTimeBucket>, DashboardQueryError> {
let current_slot = china_quarter_hour_start(now);
let start = current_slot - Duration::minutes(TIME_BUCKET_MINUTES * (TIME_BUCKET_SLOTS - 1));
self.usage
.time_buckets(start, now)
.await
.map_err(|_| DashboardQueryError::Data("time buckets"))
}
async fn health_buckets(
&self,
now: DateTime<Utc>,
) -> Result<Vec<RequestHealthTimeBucket>, DashboardQueryError> {
self.usage_records
.health_timeline(china_day_start(now), now)
.await
.map_err(|_| DashboardQueryError::Data("health timeline"))
}
async fn quota_used_by_account(
&self,
usage: &[crate::telemetry::usage::query::UsageRecordAccountUsage],
) -> Result<HashMap<String, f64>, DashboardQueryError> {
let mut quota_used = HashMap::with_capacity(usage.len());
for usage in usage {
let quota_json = self
.accounts
.get_quota_json(&usage.account_id)
.await
.map_err(|_| DashboardQueryError::Data("account quota"))?;
let Some(quota) = quota_json
.as_deref()
.and_then(|value| QuotaSnapshot::from_json(value).ok())
else {
continue;
};
if let Some(used_percent) = quota.representative_used_percent() {
quota_used.insert(usage.account_id.clone(), used_percent);
}
}
Ok(quota_used)
}
}
fn account_counts(accounts: &[Account]) -> DashboardQueryAccountCounts {
DashboardQueryAccountCounts {
total: accounts.len() as u64,
enabled: accounts
.iter()
.filter(|account| account.status == AccountStatus::Active)
.count() as u64,
abnormal: accounts
.iter()
.filter(|account| account.status != AccountStatus::Active)
.count() as u64,
}
}
fn pool_summary(
accounts: &[Account],
refreshing: &std::collections::HashSet<String>,
) -> DashboardPoolSummary {
let mut summary = DashboardPoolSummary {
total: accounts.len(),
..DashboardPoolSummary::default()
};
for account in accounts {
match account.status {
status if token_refresh_status_eligible(status) && refreshing.contains(&account.id) => {
summary.refreshing += 1;
}
AccountStatus::Active => summary.active += 1,
AccountStatus::Expired => summary.expired += 1,
AccountStatus::QuotaExhausted => summary.quota_exhausted += 1,
AccountStatus::Disabled => summary.disabled += 1,
AccountStatus::Banned => summary.banned += 1,
}
}
summary
}
fn account_usage(
accounts: &[Account],
usage_records: &[crate::telemetry::usage::query::UsageRecordAccountUsage],
quota_used: &HashMap<String, f64>,
) -> Vec<DashboardAccountUsage> {
let accounts = accounts
.iter()
.map(|account| (account.id.as_str(), account))
.collect::<HashMap<_, _>>();
usage_records
.iter()
.map(|usage| {
let account = accounts.get(usage.account_id.as_str()).copied();
DashboardAccountUsage {
id: usage.account_id.clone(),
email: account
.and_then(|account| account.email.clone())
.unwrap_or_else(|| usage.account_id.clone()),
plan_type: account.and_then(|account| account.plan_type.clone()),
total_tokens: usage.total_tokens,
quota_used_percent: quota_used.get(&usage.account_id).copied().or_else(|| {
matches!(
account.map(|account| account.status),
Some(AccountStatus::QuotaExhausted)
)
.then_some(100.0)
}),
last_used_at: usage.last_used_at,
}
})
.collect()
}
+1
View File
@@ -2,6 +2,7 @@
pub mod accounts_routes;
pub mod auth_routes;
mod dashboard_presenter;
pub mod dashboard_routes;
pub mod keys_routes;
pub mod ops_routes;
+4 -1
View File
@@ -17,7 +17,10 @@ use crate::{
json::{NumberedPage, clamp_limit, clamp_page},
time::china_datetime,
},
telemetry::{ops::query::OpsQueryError, ops::store::OpsErrorFilter, ops::types::OpsErrorLog},
telemetry::{
ops::query::OpsQueryError,
ops::types::{OpsErrorFilter, OpsErrorLog},
},
};
#[derive(Debug, Deserialize)]
+7 -3
View File
@@ -14,7 +14,8 @@ use super::{
accounts_routes::{
account_models, account_quota, accounts, batch_delete_accounts, create_account,
export_accounts, health_check_accounts, import_accounts, oauth_authorize_account,
oauth_exchange_account, refresh_account, test_account_connection, update_account,
oauth_exchange_account, refresh_account, refresh_account_models, test_account_connection,
update_account,
},
auth_routes::{login, logout, session_status},
dashboard_routes::{dashboard_summary, dashboard_trend},
@@ -45,7 +46,7 @@ pub fn router() -> Router<AppState> {
.route("/api/admin/settings", get(settings).post(update_settings))
.route(
"/api/admin/settings/admin-api-key",
get(admin_api_key_status).delete(delete_admin_api_key),
get(admin_api_key_status).post(delete_admin_api_key),
)
.route(
"/api/admin/settings/admin-api-key/regenerate",
@@ -72,7 +73,10 @@ pub fn router() -> Router<AppState> {
post(oauth_exchange_account),
)
.route("/api/admin/accounts/test", get(test_account_connection))
.route("/api/admin/accounts/models", get(account_models))
.route(
"/api/admin/accounts/models",
get(account_models).post(refresh_account_models),
)
.route("/api/admin/accounts/delete", post(batch_delete_accounts))
.route("/api/admin/accounts/update", post(update_account))
.route("/api/admin/accounts/refresh", post(refresh_account))
+1 -1
View File
@@ -115,7 +115,7 @@ pub(crate) async fn regenerate_admin_api_key(
))
}
/// `DELETE /api/admin/settings/admin-api-key`
/// `POST /api/admin/settings/admin-api-key`
pub(crate) async fn delete_admin_api_key(
State(state): State<AppState>,
_auth: AdminAuth,
+25
View File
@@ -121,11 +121,13 @@ pub(crate) struct UsageRecordTokenDetailsData {
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
cache_write_tokens: u64,
reasoning_tokens: u64,
total_tokens: u64,
input_tokens_display: String,
output_tokens_display: String,
cached_tokens_display: String,
cache_write_tokens_display: String,
reasoning_tokens_display: String,
total_tokens_display: String,
}
@@ -136,22 +138,26 @@ pub(crate) struct UsageRecordBillingData {
input_amount: f64,
output_amount: f64,
cache_read_amount: f64,
cache_write_amount: f64,
standard_amount: f64,
total_amount: f64,
input_price_per_mtoken: f64,
output_price_per_mtoken: f64,
cache_read_price_per_mtoken: f64,
cache_write_price_per_mtoken: f64,
service_tier: Option<String>,
service_tier_display: String,
multiplier: f64,
input_amount_display: String,
output_amount_display: String,
cache_read_amount_display: String,
cache_write_amount_display: String,
standard_amount_display: String,
total_amount_display: String,
input_price_display: String,
output_price_display: String,
cache_read_price_display: String,
cache_write_price_display: String,
multiplier_display: String,
}
@@ -162,6 +168,7 @@ struct UsageRecordSummaryData {
input_tokens: String,
output_tokens: String,
cached_tokens: String,
cache_write_tokens: String,
total_tokens: String,
average_latency_ms: String,
}
@@ -175,6 +182,7 @@ struct UsageRecordBreakdownData {
input_tokens: String,
output_tokens: String,
cached_tokens: String,
cache_write_tokens: String,
total_tokens: String,
total_tokens_value: u64,
total_tokens_total: String,
@@ -195,6 +203,8 @@ struct UsageRecordTrendPointData {
output_tokens_value: u64,
cached_tokens: String,
cached_tokens_value: u64,
cache_write_tokens: String,
cache_write_tokens_value: u64,
total_tokens: String,
total_tokens_value: u64,
cache_hit_rate: String,
@@ -574,6 +584,10 @@ pub(crate) fn usage_token_details(record: &UsageRecord) -> UsageRecordTokenDetai
.cached_tokens
.and_then(|value| value.try_into().ok())
.unwrap_or(0);
let cache_write_tokens: u64 = record
.cache_write_tokens
.and_then(|value| value.try_into().ok())
.unwrap_or(0);
let reasoning_tokens: u64 = record
.reasoning_tokens
.and_then(|value| value.try_into().ok())
@@ -584,11 +598,13 @@ pub(crate) fn usage_token_details(record: &UsageRecord) -> UsageRecordTokenDetai
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
reasoning_tokens,
total_tokens,
input_tokens_display: format_number(input_tokens),
output_tokens_display: format_number(output_tokens),
cached_tokens_display: format_compact_number(cached_tokens),
cache_write_tokens_display: format_compact_number(cache_write_tokens),
reasoning_tokens_display: format_number(reasoning_tokens),
total_tokens_display: format_number(total_tokens),
}
@@ -605,28 +621,33 @@ pub(crate) fn usage_billing(
tokens.input_tokens,
tokens.output_tokens,
tokens.cached_tokens,
tokens.cache_write_tokens,
)?;
Some(UsageRecordBillingData {
input_amount: billing.input_amount,
output_amount: billing.output_amount,
cache_read_amount: billing.cache_read_amount,
cache_write_amount: billing.cache_write_amount,
standard_amount: billing.standard_amount,
total_amount: billing.total_amount,
input_price_per_mtoken: billing.input_price_per_mtoken,
output_price_per_mtoken: billing.output_price_per_mtoken,
cache_read_price_per_mtoken: billing.cache_read_price_per_mtoken,
cache_write_price_per_mtoken: billing.cache_write_price_per_mtoken,
service_tier: billing.service_tier,
service_tier_display: billing.service_tier_display,
multiplier: billing.multiplier,
input_amount_display: format_billing_amount(billing.input_amount),
output_amount_display: format_billing_amount(billing.output_amount),
cache_read_amount_display: format_billing_amount(billing.cache_read_amount),
cache_write_amount_display: format_billing_amount(billing.cache_write_amount),
standard_amount_display: format_billing_amount(billing.standard_amount),
total_amount_display: format_billing_amount(billing.total_amount),
input_price_display: format_token_price(billing.input_price_per_mtoken),
output_price_display: format_token_price(billing.output_price_per_mtoken),
cache_read_price_display: format_token_price(billing.cache_read_price_per_mtoken),
cache_write_price_display: format_token_price(billing.cache_write_price_per_mtoken),
multiplier_display: format_multiplier(billing.multiplier),
})
}
@@ -638,6 +659,7 @@ impl From<UsageRecordSummary> for UsageRecordSummaryData {
input_tokens: format_tokens(summary.input_tokens),
output_tokens: format_tokens(summary.output_tokens),
cached_tokens: format_tokens(summary.cached_tokens),
cache_write_tokens: format_tokens(summary.cache_write_tokens),
total_tokens: format_tokens(summary.total_tokens),
average_latency_ms: format_duration_ms_f64(summary.average_latency_ms),
}
@@ -659,6 +681,7 @@ impl UsageRecordBreakdownData {
input_tokens: format_tokens(item.input_tokens),
output_tokens: format_tokens(item.output_tokens),
cached_tokens: format_tokens(item.cached_tokens),
cache_write_tokens: format_tokens(item.cache_write_tokens),
total_tokens: format_tokens(item.total_tokens),
total_tokens_value: item.total_tokens,
total_tokens_total: format_tokens(total_tokens_total),
@@ -698,6 +721,8 @@ impl From<UsageRecordTrendPoint> for UsageRecordTrendPointData {
output_tokens_value: point.output_tokens,
cached_tokens: format_tokens(point.cached_tokens),
cached_tokens_value: point.cached_tokens,
cache_write_tokens: format_tokens(point.cache_write_tokens),
cache_write_tokens_value: point.cache_write_tokens,
total_tokens: format_tokens(point.total_tokens),
total_tokens_value: point.total_tokens,
cache_hit_rate: format_percent(cache_hit_rate_value),
+1 -1
View File
@@ -83,7 +83,7 @@ const MODEL_NOT_FOUND_ERROR: OpenAiErrorKind = OpenAiErrorKind {
};
const HISTORY_UNAVAILABLE_ERROR: OpenAiErrorKind = OpenAiErrorKind {
error_type: "invalid_request_error",
code: "previous_response_unavailable",
code: "previous_response_not_found",
};
const CONTINUATION_BUSY_ERROR: OpenAiErrorKind = OpenAiErrorKind {
error_type: "server_error",
+25 -1
View File
@@ -19,7 +19,13 @@ use crate::{
api::AppState,
api::middleware::request_id::{ClientIp, RequestId},
dispatch::{errors::ResponseDispatchError, service::ResponseDispatchStream},
upstream::openai::protocol::{responses::CodexResponsesRequest, sse::sse_body_has_done},
upstream::openai::protocol::{
responses::{
CodexResponsesRequest, WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY,
X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER, X_OPENAI_MEMGEN_REQUEST_HEADER,
},
sse::sse_body_has_done,
},
};
use self::sse::{done_sse_frame, event_stream_response as sse_event_stream_response};
@@ -103,11 +109,29 @@ pub fn build_codex_request(
None => {}
}
apply_context_header_fields(&mut request, headers);
apply_request_semantic_headers(&mut request, headers);
apply_session_fields(&mut request, headers);
request.client_user_agent = user_agent_from_headers(headers);
request
}
fn apply_request_semantic_headers(request: &mut CodexResponsesRequest, headers: &HeaderMap) {
request.responses_lite = header_string(headers, X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER)
.or_else(|| {
request
.client_metadata()
.and_then(Value::as_object)
.and_then(|metadata| {
metadata.get(WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY)
})
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
});
request.memgen_request = header_string(headers, X_OPENAI_MEMGEN_REQUEST_HEADER);
}
fn apply_session_fields(request: &mut CodexResponsesRequest, headers: &HeaderMap) {
request.client_conversation_id = session_context_string(request.body(), "conversation_id")
.or_else(|| header_string(headers, "conversation-id"))
+33 -4
View File
@@ -131,8 +131,10 @@ struct ResponsesWebSocketSession {
connection_id: String,
}
#[derive(Default)]
struct ConnectionTranscriptCollector {
account_id: String,
request_input: Vec<Value>,
continued_from_previous_response: bool,
completed_items: Vec<Value>,
completed: Option<ConnectionTranscriptFacts>,
}
@@ -143,6 +145,20 @@ enum ForwardDispatchOutcome {
}
impl ConnectionTranscriptCollector {
fn new(
account_id: String,
request_input: Vec<Value>,
continued_from_previous_response: bool,
) -> Self {
Self {
account_id,
request_input,
continued_from_previous_response,
completed_items: Vec::new(),
completed: None,
}
}
fn observe(&mut self, event: &CanonicalResponseEvent) {
match event.event_type() {
Some("response.output_item.done") => {
@@ -166,7 +182,13 @@ impl ConnectionTranscriptCollector {
.filter(|output| !output.is_empty())
.cloned()
.unwrap_or_else(|| self.completed_items.clone());
self.completed = Some(ConnectionTranscriptFacts::new(response_id, output));
self.completed = Some(ConnectionTranscriptFacts::new(
self.account_id.clone(),
self.request_input.clone(),
self.continued_from_previous_response,
response_id,
output,
));
}
_ => {}
}
@@ -346,6 +368,9 @@ async fn forward_dispatch_stream(
request_id: &str,
) -> ForwardDispatchOutcome {
let ResponseDispatchStream {
account_id,
request_input,
continued_from_previous_response,
body,
mut canonical_events,
response_headers,
@@ -361,7 +386,11 @@ async fn forward_dispatch_stream(
let mut body = body;
let mut terminal_sent = false;
let mut completed = ConnectionTranscriptCollector::default();
let mut completed = ConnectionTranscriptCollector::new(
account_id,
request_input,
continued_from_previous_response,
);
loop {
let next = if terminal_sent {
ActiveResponseInput::Upstream(body.next().await)
@@ -516,7 +545,7 @@ fn parse_response_create(payload: &str) -> Result<Map<String, Value>, ResponseCr
Some(Value::String(message_type)) if message_type == "response.create" => {}
_ => return Err(ResponseCreateFrameError::UnsupportedType),
}
if body.get("stream").and_then(Value::as_bool) != Some(true) {
if matches!(body.get("stream"), Some(value) if value.as_bool() != Some(true)) {
return Err(ResponseCreateFrameError::StreamingRequired);
}
Ok(body)
+9 -18
View File
@@ -10,35 +10,30 @@ use axum::{Router, extract::State, http::StatusCode, middleware, routing::get};
use crate::{
api::{
admin, assets, client,
admin::{
self, accounts_routes::query::AccountListQueryService,
dashboard_routes::DashboardQueryService,
},
assets, client,
middleware::{request_id::attach_request_id, trace::http_trace_layer},
},
auth::service::SessionService,
dispatch::{affinity::SessionAffinityService, service::ResponseDispatchService},
fleet::{
manage::AccountManageService, pool::AccountPoolService, refresh::TokenRefreshService,
store::AccountStore,
},
fleet::{manage::AccountManageService, pool::AccountPoolService},
keys::{manage::KeyManageService, service::KeyVerifier},
models::service::ModelService,
settings::service::SettingsService,
telemetry::{
account_usage::query::AccountUsageQueryService, ops::query::OpsQueryService,
usage::query::UsageQueryService,
},
telemetry::{ops::query::OpsQueryService, usage::query::UsageQueryService},
update::service::SystemUpdateService,
upstream::openai::{
desktop_release::DesktopReleaseStatus, profile::CodexWireProfile,
token_client::OpenAiTokenClient,
},
};
/// HTTP API 所需的领域服务集合。
#[derive(Clone)]
pub struct ApiServices {
pub(crate) account_list: Arc<AccountListQueryService>,
pub(crate) dashboard: Arc<DashboardQueryService>,
pub health_probe: Arc<dyn HealthProbe>,
pub models: Arc<ModelService>,
pub accounts: Arc<dyn AccountStore>,
pub client_keys: Arc<KeyVerifier>,
pub admin_client_keys: Arc<KeyManageService>,
pub admin_sessions: Arc<SessionService>,
@@ -46,13 +41,9 @@ pub struct ApiServices {
pub admin_accounts: Arc<AccountManageService>,
pub usage_records: Arc<UsageQueryService>,
pub ops_errors: Arc<OpsQueryService>,
pub usage: Arc<AccountUsageQueryService>,
pub account_pool: Arc<AccountPoolService>,
pub token_refresh: Arc<TokenRefreshService<OpenAiTokenClient>>,
pub responses: Arc<ResponseDispatchService>,
pub session_affinity: Arc<SessionAffinityService>,
pub wire_profile: Arc<CodexWireProfile>,
pub desktop_release: DesktopReleaseStatus,
pub process_control: Arc<dyn ProcessControl>,
pub system_update: Arc<SystemUpdateService>,
pub connection_drain: crate::api::middleware::connection_drain::ConnectionDrain,
+12
View File
@@ -350,6 +350,10 @@ pub struct WebSocketPoolSettings {
pub max_age_ms: u64,
/// 单账号最大连接数。
pub max_per_account: usize,
/// 所有账号合计最大连接数。
pub max_total: usize,
/// 所有账号合计最大并发建连数。
pub max_connecting: usize,
/// 首个真实输出(首事件)到达前的绝对超时;`0` 表示禁用。
pub initial_event_timeout_ms: u64,
}
@@ -539,6 +543,14 @@ fn validate_file_config(config: &FileAppConfig) -> Result<(), ConfigError> {
if config.ws_pool.max_per_account == 0 {
return Err(ConfigError::InvalidField("ws_pool.max_per_account"));
}
if config.ws_pool.max_total == 0 {
return Err(ConfigError::InvalidField("ws_pool.max_total"));
}
if config.ws_pool.max_connecting == 0
|| config.ws_pool.max_connecting > config.ws_pool.max_total
{
return Err(ConfigError::InvalidField("ws_pool.max_connecting"));
}
validate_wire_profile(&config.wire_profile)?;
validate_logging(&config.logging)?;
Ok(())
+161 -62
View File
@@ -6,6 +6,14 @@ use axum::Router;
use secrecy::ExposeSecret;
use crate::{
api::admin::{
accounts_routes::query::{AccountListQueryService, RefreshActivityQuery},
dashboard_routes::{
DashboardDesktopRelease, DashboardDesktopReleaseQuery, DashboardDesktopReleaseSnapshot,
DashboardQueryService, DashboardQueryServiceParts, DashboardWireProfile,
DashboardWireProfileQuery,
},
},
auth::{
service::SessionService,
store::{PgAdminUserStore, RedisAdminSessionStore},
@@ -18,11 +26,15 @@ use crate::{
fleet::{
cookies::PgCookieStore,
manage::{AccountManageService, AccountManageServiceParts},
pool::{AccountPoolOptions, AccountPoolService, AccountPoolStaticSettings},
pool::{
AccountPoolOptions, AccountPoolRuntimeOptions, AccountPoolService,
AccountPoolStaticSettings, RotationStrategy,
},
refresh::{
RedisRefreshLeaseStore, RefreshPolicy, RuntimeRefreshPolicy, TokenRefreshService,
},
store::{AccountStore, PgAccountStore},
usage::AccountUsageStore,
},
infra::{
database::connect,
@@ -33,6 +45,7 @@ use crate::{
},
keys::{manage::KeyManageService, service::KeyVerifier, store::PgClientKeyStore},
models::{
gateway::ModelCatalogSource,
service::{ModelService, ModelServiceError},
store::{ModelSnapshotStore, RedisModelSnapshotStore},
types::ModelConfig,
@@ -40,7 +53,7 @@ use crate::{
settings::{SettingsSnapshot, service::SettingsService},
telemetry::{
account_usage::query::AccountUsageQueryService,
account_usage::store::{AccountUsageStore, PgAccountUsageStore},
account_usage::store::PgAccountUsageStore,
buckets::store::PgRequestBucketStore,
ops::query::OpsQueryService,
ops::store::PgOpsErrorLogStore,
@@ -51,11 +64,11 @@ use crate::{
update::service::SystemUpdateService,
upstream::openai::{
desktop_release::DesktopReleaseStatus,
profile::CodexWireProfile,
profile::{CodexWireProfile, CodexWireProfileState},
token_client::{OpenAiTokenClient, TokenClientConfig, default_openai_token_client},
transport::{
CodexBackendClient, CodexModelCatalogClient, CodexWebSocketPool,
CodexWebSocketPoolConfig, build_reqwest_client, tls::CustomCaError,
CodexBackendClient, CodexWebSocketPool, CodexWebSocketPoolConfig, build_reqwest_client,
tls::CustomCaError,
},
},
};
@@ -117,7 +130,7 @@ pub struct Services {
pub session_affinity: Arc<SessionAffinityService>,
pub codex: Arc<CodexBackendClient>,
pub websocket_pool: Option<Arc<CodexWebSocketPool>>,
pub wire_profile: Arc<CodexWireProfile>,
pub wire_profile: CodexWireProfileState,
pub desktop_release: DesktopReleaseStatus,
pub account_pseudonymizer: Arc<AccountPseudonymizer>,
pub system_update: Arc<SystemUpdateService>,
@@ -142,11 +155,50 @@ impl UsageRecordOptions {
}
}
impl DashboardWireProfileQuery for CodexWireProfileState {
fn snapshot(&self) -> DashboardWireProfile {
let profile = CodexWireProfileState::snapshot(self);
let user_agent = profile.user_agent();
DashboardWireProfile {
originator: profile.originator,
codex_version: profile.codex_version,
desktop_version: profile.desktop_version,
desktop_build: profile.desktop_build,
os_type: profile.os_type,
os_version: profile.os_version,
arch: profile.arch,
terminal: profile.terminal,
user_agent,
verified_at: profile.verified_at,
}
}
}
impl DashboardDesktopReleaseQuery for DesktopReleaseStatus {
fn snapshot(&self) -> DashboardDesktopReleaseSnapshot {
let snapshot = DesktopReleaseStatus::snapshot(self);
DashboardDesktopReleaseSnapshot {
checked_at: snapshot.checked_at,
latest: snapshot.latest.map(|release| DashboardDesktopRelease {
version: release.version,
build: release.build,
published_at: release.published_at,
minimum_system_version: release.minimum_system_version,
hardware_requirements: release.hardware_requirements,
download_url: release.download_url,
download_size: release.download_size,
signature_present: release.signature_present,
}),
last_error: snapshot.last_error,
}
}
}
impl Services {
pub fn new(
config: &AppConfig,
stores: BackgroundTaskStores,
wire_profile: Arc<CodexWireProfile>,
wire_profile: CodexWireProfileState,
) -> Self {
Self::try_new(config, stores, wire_profile)
.expect("failed to build runtime services with configured TLS transport")
@@ -155,7 +207,7 @@ impl Services {
pub fn try_new(
config: &AppConfig,
stores: BackgroundTaskStores,
wire_profile: Arc<CodexWireProfile>,
wire_profile: CodexWireProfileState,
) -> Result<Self, CustomCaError> {
Self::try_with_usage_record_options(
config,
@@ -168,7 +220,7 @@ impl Services {
pub fn try_with_usage_record_options(
config: &AppConfig,
stores: BackgroundTaskStores,
wire_profile: Arc<CodexWireProfile>,
wire_profile: CodexWireProfileState,
usage_record_options: UsageRecordOptions,
) -> Result<Self, CustomCaError> {
let mut identity_secret = [0u8; 32];
@@ -185,7 +237,7 @@ impl Services {
fn try_with_identity_secret_and_usage_record_options(
config: &AppConfig,
stores: BackgroundTaskStores,
wire_profile: Arc<CodexWireProfile>,
wire_profile: CodexWireProfileState,
identity_secret: [u8; 32],
usage_record_options: UsageRecordOptions,
) -> Result<Self, CustomCaError> {
@@ -203,6 +255,8 @@ 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,
max_total: config.ws_pool.max_total,
max_connecting: config.ws_pool.max_connecting,
initial_event_timeout: (config.ws_pool.initial_event_timeout_ms > 0).then(|| {
std::time::Duration::from_millis(config.ws_pool.initial_event_timeout_ms)
}),
@@ -267,9 +321,10 @@ impl Services {
refresh_policy.clone(),
token_client.clone(),
)
.with_refresh_lease_store(stores.refresh_leases.clone()),
.with_refresh_lease_store(stores.refresh_leases.clone())
.with_account_pool_sync(account_pool.clone()),
);
let upstream_client: Arc<dyn CodexModelCatalogClient> = codex.clone();
let upstream_client: Arc<dyn ModelCatalogSource> = codex.clone();
let store: Arc<dyn ModelSnapshotStore> = Arc::new(stores.model_snapshots.clone());
let models = Arc::new(ModelService::new(
ModelConfig {
@@ -281,7 +336,7 @@ impl Services {
let admin_accounts = Arc::new(AccountManageService::new(AccountManageServiceParts {
store: stores.accounts.clone(),
cookies: stores.cookies.clone(),
codex: codex.clone(),
upstream: codex.clone(),
models: models.clone(),
account_pool: account_pool.clone(),
token_refresher: Arc::new(token_client),
@@ -355,13 +410,30 @@ impl Services {
impl From<&Services> for crate::api::router::ApiServices {
fn from(services: &Services) -> Self {
let refresh_activity = services.token_refresh.clone() as Arc<dyn RefreshActivityQuery>;
let account_list = Arc::new(AccountListQueryService::new(
services.admin_accounts.clone(),
services.usage.clone(),
refresh_activity.clone(),
));
let dashboard = Arc::new(DashboardQueryService::new(DashboardQueryServiceParts {
accounts: services.accounts.clone(),
usage: services.usage.clone(),
usage_records: services.usage_records.clone(),
account_pool: services.account_pool.clone(),
refresh_activity,
settings: services.settings.clone(),
wire_profile: Arc::new(services.wire_profile.clone()),
desktop_release: Arc::new(services.desktop_release.clone()),
}));
Self {
account_list,
dashboard,
health_probe: Arc::new(crate::bootstrap::state::RuntimeHealthProbe::new(
services.database.clone(),
services.redis.clone(),
)),
models: services.models.clone(),
accounts: services.accounts.clone(),
client_keys: services.client_keys.clone(),
admin_client_keys: services.admin_client_keys.clone(),
admin_sessions: services.admin_sessions.clone(),
@@ -369,13 +441,9 @@ impl From<&Services> for crate::api::router::ApiServices {
admin_accounts: services.admin_accounts.clone(),
usage_records: services.usage_records.clone(),
ops_errors: services.ops_errors.clone(),
usage: services.usage.clone(),
account_pool: services.account_pool.clone(),
token_refresh: services.token_refresh.clone(),
responses: services.responses.clone(),
session_affinity: services.session_affinity.clone(),
wire_profile: services.wire_profile.clone(),
desktop_release: services.desktop_release.clone(),
process_control: Arc::new(crate::bootstrap::shutdown::RuntimeProcessControl),
system_update: services.system_update.clone(),
connection_drain: services.connection_drain.clone(),
@@ -441,11 +509,33 @@ pub fn wire_profile_from_config(
}
pub fn account_pool_options_from_config(config: &AppConfig) -> AccountPoolOptions {
AccountPoolStaticSettings {
let static_settings = AccountPoolStaticSettings {
skip_quota_limited: config.quota.skip_exhausted,
tier_priority: config.auth.tier_priority.clone(),
};
account_pool_runtime_options(&settings_snapshot_from_config(config), &static_settings).options
}
pub fn account_pool_runtime_options(
settings: &SettingsSnapshot,
static_settings: &AccountPoolStaticSettings,
) -> AccountPoolRuntimeOptions {
AccountPoolRuntimeOptions {
options: AccountPoolOptions {
max_concurrent_per_account: settings.max_concurrent_per_account,
rotation_strategy: match settings.rotation_strategy.as_str() {
"smart" => RotationStrategy::Smart,
"quota_reset_priority" => RotationStrategy::QuotaResetPriority,
"round_robin" => RotationStrategy::RoundRobin,
"sticky" => RotationStrategy::Sticky,
_ => RotationStrategy::Smart,
},
skip_quota_limited: static_settings.skip_quota_limited,
tier_priority: static_settings.tier_priority.clone(),
..AccountPoolOptions::default()
},
request_interval_ms: settings.request_interval_ms,
}
.pool_options(&settings_snapshot_from_config(config))
}
pub async fn run(config: BootstrapConfig) -> Result<(), Box<dyn Error + Send + Sync>> {
@@ -457,6 +547,7 @@ pub async fn run(config: BootstrapConfig) -> Result<(), Box<dyn Error + Send + S
let listener = tokio::net::TcpListener::bind((host.as_str(), port)).await?;
let (router, task_coordinator, connection_drain) = build_router(config).await?;
let mut task_coordinator = Some(task_coordinator);
let mut restart_path = None;
let serve_result = {
let (graceful_shutdown_tx, graceful_shutdown_rx) = tokio::sync::oneshot::channel();
@@ -471,46 +562,63 @@ pub async fn run(config: BootstrapConfig) -> Result<(), Box<dyn Error + Send + S
tokio::pin!(serve);
tokio::select! {
result = &mut serve => result,
() = crate::bootstrap::shutdown::shutdown_signal() => {
let _ = graceful_shutdown_tx.send(());
let active_websocket_connections = connection_drain.begin_shutdown();
tracing::info!(
active_websocket_connections,
"正在排空入站连接并关闭运行时任务"
);
let coordinator = task_coordinator.take();
let connection_drain = connection_drain.clone();
let runtime_shutdown = async move {
let shutdown_tasks = async move {
if let Some(coordinator) = coordinator {
coordinator.shutdown().await;
}
};
tokio::join!(shutdown_tasks, connection_drain.wait());
};
let drain_runtime = async {
let (serve_result, ()) = tokio::join!(&mut serve, runtime_shutdown);
serve_result
};
tokio::pin!(drain_runtime);
tokio::select! {
result = &mut drain_runtime => result,
() = tokio::time::sleep(HTTP_DRAIN_TIMEOUT) => {
tracing::warn!(
timeout_secs = HTTP_DRAIN_TIMEOUT.as_secs(),
"HTTP graceful shutdown timed out; cancelling remaining connections"
);
action = crate::bootstrap::shutdown::shutdown_signal() => {
match action {
crate::bootstrap::shutdown::ShutdownAction::Restart(executable_path) => {
restart_path = Some(executable_path);
Ok(())
}
() = crate::bootstrap::shutdown::shutdown_signal() => {
tracing::warn!("Received a second shutdown signal; cancelling remaining connections");
Ok(())
crate::bootstrap::shutdown::ShutdownAction::Graceful => {
let _ = graceful_shutdown_tx.send(());
let active_websocket_connections = connection_drain.begin_shutdown();
tracing::info!(
active_websocket_connections,
"正在排空入站连接并关闭运行时任务"
);
let coordinator = task_coordinator.take();
let connection_drain = connection_drain.clone();
let runtime_shutdown = async move {
let shutdown_tasks = async move {
if let Some(coordinator) = coordinator {
coordinator.shutdown().await;
}
};
tokio::join!(shutdown_tasks, connection_drain.wait());
};
let drain_runtime = async {
let (serve_result, ()) = tokio::join!(&mut serve, runtime_shutdown);
serve_result
};
tokio::pin!(drain_runtime);
tokio::select! {
result = &mut drain_runtime => result,
() = tokio::time::sleep(HTTP_DRAIN_TIMEOUT) => {
tracing::warn!(
timeout_secs = HTTP_DRAIN_TIMEOUT.as_secs(),
"HTTP graceful shutdown timed out; cancelling remaining connections"
);
Ok(())
}
_ = crate::bootstrap::shutdown::shutdown_signal() => {
tracing::warn!("Received a second shutdown signal; cancelling remaining connections");
Ok(())
}
}
}
}
}
}
};
if let Some(executable_path) = restart_path {
tracing::info!(
executable = %executable_path.display(),
"正在以更新后的二进制替换当前进程"
);
drop(log_guard);
return exec_replacement_process(executable_path);
}
if let Some(coordinator) = task_coordinator.take() {
let active_websocket_connections = connection_drain.begin_shutdown();
tracing::info!(
@@ -521,15 +629,6 @@ pub async fn run(config: BootstrapConfig) -> Result<(), Box<dyn Error + Send + S
}
serve_result?;
if let Some(executable_path) = crate::bootstrap::shutdown::restart_executable_path() {
tracing::info!(
executable = %executable_path.display(),
"正在以更新后的二进制替换当前进程"
);
drop(log_guard);
exec_replacement_process(executable_path)?;
}
Ok(())
}
@@ -578,7 +677,7 @@ async fn build_router(
let data_dir = ensure_data_dir(&config.runtime.data_directory)?;
let identity_secret = load_or_create_identity_secret(&data_dir)?;
let wire_profile = Arc::new(wire_profile_from_config(&config.wire_profile));
let wire_profile = CodexWireProfileState::new(wire_profile_from_config(&config.wire_profile));
let runtime_config = crate::bootstrap::state::RuntimeConfig::from(&config);
let services = Services::try_with_identity_secret_and_usage_record_options(
&config,
+22 -32
View File
@@ -1,19 +1,17 @@
//! 优雅关闭信号处理。
//! 进程关闭与重启信号处理。
use std::{
path::PathBuf,
sync::{
OnceLock,
atomic::{AtomicBool, Ordering},
},
};
use std::{path::PathBuf, sync::OnceLock};
use tokio::signal;
use tokio::sync::broadcast;
static SHUTDOWN_REQUESTS: OnceLock<broadcast::Sender<()>> = OnceLock::new();
static RESTART_REQUESTED: AtomicBool = AtomicBool::new(false);
static RESTART_EXECUTABLE_PATH: OnceLock<PathBuf> = OnceLock::new();
static SHUTDOWN_REQUESTS: OnceLock<broadcast::Sender<ShutdownAction>> = OnceLock::new();
#[derive(Clone)]
pub(crate) enum ShutdownAction {
Graceful,
Restart(PathBuf),
}
#[derive(Debug, Clone, Copy, Default)]
pub struct RuntimeProcessControl;
@@ -24,12 +22,12 @@ impl crate::api::router::ProcessControl for RuntimeProcessControl {
}
fn request_restart(&self, executable_path: PathBuf) {
request_process_restart(executable_path);
request_immediate_restart(executable_path);
}
}
/// 等待进程关闭信号(Ctrl+C 或 SIGTERM)。
pub async fn shutdown_signal() {
/// 等待进程关闭或重启信号。
pub(crate) async fn shutdown_signal() -> ShutdownAction {
let mut requested = shutdown_sender().subscribe();
let ctrl_c = async {
signal::ctrl_c()
@@ -49,33 +47,25 @@ pub async fn shutdown_signal() {
let terminate = std::future::pending::<()>();
tokio::select! {
_ = requested.recv() => {},
() = ctrl_c => {},
() = terminate => {},
action = requested.recv() => match action {
Ok(action) => action,
Err(_) => ShutdownAction::Graceful,
},
() = ctrl_c => ShutdownAction::Graceful,
() = terminate => ShutdownAction::Graceful,
}
}
/// 请求进程按优雅关闭路径退出。
pub fn request_shutdown() {
let _ = shutdown_sender().send(());
let _ = shutdown_sender().send(ShutdownAction::Graceful);
}
/// 请求进程完成优雅关闭后以新二进制替换当前进程。
pub fn request_process_restart(executable_path: PathBuf) {
let _ = RESTART_EXECUTABLE_PATH.set(executable_path);
RESTART_REQUESTED.store(true, Ordering::SeqCst);
request_shutdown();
fn request_immediate_restart(executable_path: PathBuf) {
let _ = shutdown_sender().send(ShutdownAction::Restart(executable_path));
}
/// 获取待 exec 的重启目标。
pub fn restart_executable_path() -> Option<PathBuf> {
RESTART_REQUESTED
.load(Ordering::SeqCst)
.then(|| RESTART_EXECUTABLE_PATH.get().cloned())
.flatten()
}
fn shutdown_sender() -> &'static broadcast::Sender<()> {
fn shutdown_sender() -> &'static broadcast::Sender<ShutdownAction> {
SHUTDOWN_REQUESTS.get_or_init(|| {
let (sender, _receiver) = broadcast::channel(16);
sender
+39 -16
View File
@@ -91,6 +91,7 @@ impl TaskCoordinator {
"desktop_release_update",
start_desktop_release_update_task(
services.desktop_release.clone(),
services.wire_profile.clone(),
crate::upstream::openai::desktop_release::CODEX_DESKTOP_APPCAST_URL.to_string(),
),
);
@@ -107,30 +108,52 @@ impl TaskCoordinator {
/// 启动各领域的设置订阅任务。
pub fn start_settings_subscriptions(services: &crate::bootstrap::services::Services) -> Self {
let mut coordinator = TaskCoordinator::default();
let mut settings = services.settings.subscribe();
let initial = settings.borrow().clone();
let (model_tx, model_rx) = tokio::sync::watch::channel(crate::models::types::ModelConfig {
model_aliases: initial.model_aliases.clone(),
});
let (pool_tx, pool_rx) =
tokio::sync::watch::channel(crate::bootstrap::services::account_pool_runtime_options(
&initial,
&services.account_pool_static,
));
let (refresh_tx, refresh_rx) =
tokio::sync::watch::channel(crate::fleet::refresh::RefreshPolicy {
refresh_margin_seconds: initial.refresh_margin_seconds,
refresh_concurrency: initial.refresh_concurrency,
});
coordinator.push(
"model_settings",
SchedulerHandle::spawn(
services
.models
.clone()
.subscribe_settings(services.settings.subscribe()),
),
SchedulerHandle::spawn(services.models.clone().subscribe_config(model_rx)),
);
coordinator.push(
"account_pool_settings",
SchedulerHandle::spawn(services.account_pool.clone().subscribe_settings(
services.settings.subscribe(),
services.account_pool_static.clone(),
)),
SchedulerHandle::spawn(services.account_pool.clone().subscribe_options(pool_rx)),
);
coordinator.push(
"refresh_policy_settings",
SchedulerHandle::spawn(
services
.refresh_policy
.clone()
.subscribe_settings(services.settings.subscribe()),
),
SchedulerHandle::spawn(services.refresh_policy.clone().subscribe_policy(refresh_rx)),
);
let static_settings = services.account_pool_static.clone();
coordinator.push(
"settings_fanout",
SchedulerHandle::spawn(async move {
while settings.changed().await.is_ok() {
let snapshot = settings.borrow_and_update().clone();
model_tx.send_replace(crate::models::types::ModelConfig {
model_aliases: snapshot.model_aliases.clone(),
});
pool_tx.send_replace(crate::bootstrap::services::account_pool_runtime_options(
&snapshot,
&static_settings,
));
refresh_tx.send_replace(crate::fleet::refresh::RefreshPolicy {
refresh_margin_seconds: snapshot.refresh_margin_seconds,
refresh_concurrency: snapshot.refresh_concurrency,
});
}
}),
);
coordinator
}
@@ -5,6 +5,7 @@ use tracing::{info, warn};
use crate::upstream::openai::desktop_release::{
APPCAST_POLL_INTERVAL, DesktopReleaseChecker, DesktopReleaseStatus,
};
use crate::upstream::openai::profile::CodexWireProfileState;
use super::{
coordinator::SchedulerHandle,
@@ -14,9 +15,15 @@ use super::{
/// 启动 Desktop 官方发布检查后台任务。
pub fn start_desktop_release_update_task(
status: DesktopReleaseStatus,
wire_profile: CodexWireProfileState,
appcast_url: String,
) -> SchedulerHandle {
let checker = DesktopReleaseChecker::with_client(reqwest::Client::new(), appcast_url, status);
let checker = DesktopReleaseChecker::with_client(
reqwest::Client::new(),
appcast_url,
status,
wire_profile,
);
spawn_periodic_task(
DesktopReleaseUpdateTask {
checker,
@@ -43,7 +50,7 @@ impl PeriodicTaskRunner for DesktopReleaseUpdateTask {
Ok(release) => info!(
version = %release.version,
build = %release.build,
"Codex Desktop 最新发布信息已刷新"
"Codex Desktop 最新发布信息已刷新并已同步请求身份"
),
Err(error) if first_tick => {
warn!(error = %error, "Codex Desktop 首次发布检查失败");
+45 -85
View File
@@ -13,23 +13,18 @@ use crate::models::service::{
ModelRefreshPlanAccount, ModelRefreshResult, ModelService, ModelServiceError,
};
use super::{
coordinator::SchedulerHandle,
periodic::{PeriodicTaskConfig, PeriodicTaskRunner, PeriodicTaskStartup, spawn_periodic_task},
};
use super::coordinator::SchedulerHandle;
/// 模型刷新任务接线器。
pub struct ModelRefreshTask {
model_service: Arc<ModelService>,
account_pool: Arc<AccountPoolService>,
refresh_interval_secs: u64,
initial_delay_ms: u64,
retry_delay_ms: u64,
max_retries: u32,
account_pseudonymizer: Arc<AccountPseudonymizer>,
}
const DEFAULT_REFRESH_INTERVAL_SECS: u64 = 3600;
const INITIAL_DELAY_MS: u64 = 1000;
const RETRY_DELAY_MS: u64 = 10_000;
const MAX_RETRIES: u32 = 12;
@@ -44,7 +39,6 @@ impl ModelRefreshTask {
Self {
model_service,
account_pool,
refresh_interval_secs: DEFAULT_REFRESH_INTERVAL_SECS,
initial_delay_ms: INITIAL_DELAY_MS,
retry_delay_ms: RETRY_DELAY_MS,
max_retries: MAX_RETRIES,
@@ -52,15 +46,13 @@ impl ModelRefreshTask {
}
}
/// 启动后台刷新任务。
/// 启动一次性首次刷新任务。
pub fn start(self) -> SchedulerHandle {
let config = PeriodicTaskConfig::new(
self.refresh_interval_secs,
"模型刷新任务已启动",
"模型刷新任务正在关闭",
)
.wait_first_interval();
spawn_periodic_task(self, config)
let (shutdown_tx, mut shutdown_rx) = mpsc::channel(1);
let handle = tokio::spawn(async move {
self.run_initial_refresh(&mut shutdown_rx).await;
});
SchedulerHandle::new(shutdown_tx, handle)
}
/// 在 Responses 响应声明模型目录 ETag 变化时异步刷新。
@@ -115,84 +107,52 @@ impl ModelRefreshTask {
.await;
Ok(result)
}
}
impl PeriodicTaskRunner for ModelRefreshTask {
fn before_loop<'a>(
&'a mut self,
shutdown_rx: &'a mut mpsc::Receiver<()>,
) -> super::periodic::TaskFuture<'a, PeriodicTaskStartup> {
Box::pin(async move {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在首次拉取前关闭");
return PeriodicTaskStartup::Stop;
}
() = sleep(Duration::from_millis(self.initial_delay_ms)) => {}
async fn run_initial_refresh(&self, shutdown_rx: &mut mpsc::Receiver<()>) {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在首次拉取前关闭");
return;
}
() = sleep(Duration::from_millis(self.initial_delay_ms)) => {}
}
let mut attempt = 0;
let mut has_fetched_once = false;
let mut deferred_until_accounts_available = false;
while attempt < self.max_retries {
match self.refresh_once().await {
Ok(result) => {
info!(
refreshed_plans = result.refreshed_plans,
model_count = result.model_count,
"首次模型刷新成功"
);
has_fetched_once = true;
break;
}
Err(error) => {
attempt += 1;
if matches!(error, ModelServiceError::NoAccounts) {
info!("首次模型刷新跳过:没有可用账户,将进入周期重试");
deferred_until_accounts_available = true;
break;
}
warn!(
attempt,
max_retries = self.max_retries,
error = ?error,
"首次模型刷新失败"
);
if attempt < self.max_retries {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在重试等待期间关闭");
return PeriodicTaskStartup::Stop;
}
() = sleep(Duration::from_millis(self.retry_delay_ms)) => {}
}
}
}
}
}
if !has_fetched_once && !deferred_until_accounts_available {
warn!("模型刷新任务首次尝试全部失败,切换到周期刷新");
}
PeriodicTaskStartup::Continue
})
}
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
let mut attempt = 0;
while attempt < self.max_retries {
match self.refresh_once().await {
Ok(_) => {}
Ok(result) => {
info!(
refreshed_plans = result.refreshed_plans,
model_count = result.model_count,
"首次模型刷新成功"
);
return;
}
Err(ModelServiceError::NoAccounts) => {
info!("模型刷新跳过:没有可用账户");
info!("首次模型刷新跳过:没有可用账户,后续由 ETag 或手动操作触发");
return;
}
Err(error) => {
warn!(error = ?error, "刷新模型列表失败");
attempt += 1;
warn!(
attempt,
max_retries = self.max_retries,
error = ?error,
"首次模型刷新失败"
);
if attempt < self.max_retries {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在重试等待期间关闭");
return;
}
() = sleep(Duration::from_millis(self.retry_delay_ms)) => {}
}
}
}
}
})
}
warn!("模型刷新任务首次尝试全部失败,等待 ETag 或手动操作再次触发");
}
}
-32
View File
@@ -10,20 +10,11 @@ use super::coordinator::SchedulerHandle;
pub(crate) type TaskFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
/// 周期任务进入循环前的决策。
pub(crate) enum PeriodicTaskStartup {
/// 继续进入周期循环。
Continue,
/// 任务已自行结束,不再进入周期循环。
Stop,
}
/// 周期任务调度配置。
pub(crate) struct PeriodicTaskConfig {
interval_secs: u64,
started_message: &'static str,
stopped_message: &'static str,
run_immediately: bool,
}
impl PeriodicTaskConfig {
@@ -36,25 +27,12 @@ impl PeriodicTaskConfig {
interval_secs,
started_message,
stopped_message,
run_immediately: true,
}
}
pub(crate) fn wait_first_interval(mut self) -> Self {
self.run_immediately = false;
self
}
}
/// 周期任务行为。
pub(crate) trait PeriodicTaskRunner: Send + 'static {
fn before_loop<'a>(
&'a mut self,
_shutdown_rx: &'a mut mpsc::Receiver<()>,
) -> TaskFuture<'a, PeriodicTaskStartup> {
Box::pin(std::future::ready(PeriodicTaskStartup::Continue))
}
fn tick(&mut self) -> TaskFuture<'_, ()>;
fn shutdown(&mut self) -> TaskFuture<'_, ()> {
@@ -75,17 +53,7 @@ where
"{}", config.started_message
);
if matches!(
task.before_loop(&mut shutdown_rx).await,
PeriodicTaskStartup::Stop
) {
return;
}
let mut ticker = interval(Duration::from_secs(config.interval_secs));
if !config.run_immediately {
ticker.tick().await;
}
loop {
tokio::select! {
+4 -3
View File
@@ -10,7 +10,8 @@ use crate::{
cookies::PgCookieStore,
pool::AccountPoolService,
quota::{QuotaRefreshService, QuotaRefreshSummary},
store::PgAccountStore,
store::{AccountStore, PgAccountStore},
usage::AccountUsageStore,
},
infra::identity::AccountPseudonymizer,
telemetry::account_usage::store::PgAccountUsageStore,
@@ -42,8 +43,8 @@ impl QuotaRefreshTask {
) -> Self {
Self {
service: QuotaRefreshService::with_min_refresh_interval_secs(
store,
usage_store,
Arc::new(store) as Arc<dyn AccountStore>,
Arc::new(usage_store) as Arc<dyn AccountUsageStore>,
codex,
account_pool,
account_pseudonymizer,
+1 -1
View File
@@ -2,11 +2,11 @@
use tracing::{debug, info, warn};
use crate::fleet::refresh::TokenRefresher;
use crate::fleet::{
refresh::{RedisRefreshLeaseStore, RuntimeRefreshPolicy, TokenRefreshService},
store::PgAccountStore,
};
use crate::upstream::openai::token_client::TokenRefresher;
use super::{
coordinator::SchedulerHandle,
+114 -38
View File
@@ -3,12 +3,56 @@
use serde_json::{Map, Value};
use crate::{
dispatch::controllers::history::{CrossAccountReplay, sanitize_cross_account_input},
infra::identity::AccountPseudonymizer,
upstream::openai::protocol::responses::CodexResponsesRequest,
};
const TURN_STATE_METADATA_KEY: &str = "x-codex-turn-state";
const TURN_METADATA_KEY: &str = "x-codex-turn-metadata";
const INSTALLATION_IDENTITY_KEYS: [&str; 3] = [
"installation_id",
"installationId",
"x-codex-installation-id",
];
const CROSS_ACCOUNT_IDENTITY_KEYS: [&str; 26] = [
"authorization",
"Authorization",
"cookie",
"Cookie",
"chatgpt-account-id",
"chatgpt_account_id",
"chatgptAccountId",
"account_id",
"accountId",
"user_id",
"userId",
"chatgpt_user_id",
"chatgptUserId",
"access_token",
"accessToken",
"session_token",
"sessionToken",
"refresh_token",
"refreshToken",
"id_token",
"idToken",
"token",
"cookies",
"cookie_header",
"cookieHeader",
"cf_clearance",
];
const CROSS_ACCOUNT_OPAQUE_STATE_KEYS: [&str; 8] = [
"turnState",
"turn_state",
"x-codex-turn-state",
"previous_response_id",
"previousResponseId",
"response_id",
"responseId",
"conversation",
];
/// 持久化密钥驱动的账号身份隔离服务。
#[derive(Clone)]
@@ -32,24 +76,21 @@ pub(in crate::dispatch) struct AccountScopedIdentity {
/// 单个请求内账号绑定 opaque 状态的 owner。
pub(in crate::dispatch) struct AccountIdentityScope {
turn_state_account_id: Option<String>,
account_bound_state_owner_id: Option<String>,
}
impl AccountIdentityScope {
pub(in crate::dispatch) fn new(preferred_account_id: Option<String>) -> Self {
Self {
turn_state_account_id: preferred_account_id,
account_bound_state_owner_id: preferred_account_id,
}
}
fn preserves_turn_state(&mut self, account_id: &str, has_turn_state: bool) -> bool {
if !has_turn_state {
return false;
}
fn bind_and_is_cross_account(&mut self, account_id: &str) -> bool {
let owner = self
.turn_state_account_id
.account_bound_state_owner_id
.get_or_insert_with(|| account_id.to_string());
owner == account_id
owner != account_id
}
}
@@ -100,15 +141,19 @@ impl AccountIdentityService {
scope: &mut AccountIdentityScope,
mut request: CodexResponsesRequest,
account_id: &str,
cross_account_replay: Option<&CrossAccountReplay>,
) -> AccountScopedRequest {
if let Some(replay) = cross_account_replay {
replay.validate_target(account_id);
}
let cross_account =
cross_account_replay.is_some() || scope.bind_and_is_cross_account(account_id);
let identity = AccountScopedIdentity {
installation_id: self.pseudonymizer.installation_id(account_id),
};
let client_metadata_turn_state = client_metadata_string(&request, TURN_STATE_METADATA_KEY);
let preserve_turn_state = scope.preserves_turn_state(
account_id,
request.turn_state.is_some() || client_metadata_turn_state.is_some(),
);
let preserve_turn_state = !cross_account
&& (request.turn_state.is_some() || client_metadata_turn_state.is_some());
let turn_state = if preserve_turn_state {
request.turn_state.clone()
} else {
@@ -122,9 +167,14 @@ impl AccountIdentityService {
let turn_metadata = request
.turn_metadata
.as_deref()
.map(|raw| Self::scope_turn_metadata(raw, &identity));
.and_then(|raw| Self::scope_turn_metadata(raw, &identity, cross_account));
let client_metadata_turn_metadata = client_metadata_string(&request, TURN_METADATA_KEY)
.map(|raw| Self::scope_turn_metadata(&raw, &identity));
.and_then(|raw| Self::scope_turn_metadata(&raw, &identity, cross_account));
if cross_account {
sanitize_cross_account_input(&mut request);
remove_cross_account_request_state(&mut request);
}
apply_identity_to_request(
&mut request,
@@ -133,6 +183,7 @@ impl AccountIdentityService {
turn_metadata.as_deref(),
client_metadata_turn_state.as_deref(),
client_metadata_turn_metadata.as_deref(),
cross_account,
);
request.turn_state = turn_state;
request.turn_metadata = turn_metadata;
@@ -140,37 +191,33 @@ impl AccountIdentityService {
AccountScopedRequest { request, identity }
}
pub(in crate::dispatch) fn scope_auxiliary(&self, account_id: &str) -> AccountScopedIdentity {
AccountScopedIdentity {
installation_id: self.pseudonymizer.installation_id(account_id),
}
}
fn scope_turn_metadata(raw: &str, identity: &AccountScopedIdentity) -> String {
fn scope_turn_metadata(
raw: &str,
identity: &AccountScopedIdentity,
cross_account: bool,
) -> Option<String> {
let Ok(Value::Object(mut metadata)) = serde_json::from_str::<Value>(raw) else {
return raw.to_string();
return (!cross_account).then(|| raw.to_string());
};
let contains_installation_identity = [
"installation_id",
"installationId",
"x-codex-installation-id",
]
.into_iter()
.any(|key| metadata.contains_key(key));
if !contains_installation_identity {
return raw.to_string();
if cross_account {
remove_cross_account_metadata_state(&mut metadata);
} else if !INSTALLATION_IDENTITY_KEYS
.into_iter()
.any(|key| metadata.contains_key(key))
{
return Some(raw.to_string());
}
for key in [
"installation_id",
"installationId",
"x-codex-installation-id",
] {
for key in INSTALLATION_IDENTITY_KEYS {
replace_existing_metadata_field(&mut metadata, key, Some(&identity.installation_id));
}
serde_json::to_string(&metadata).unwrap_or_else(|_| raw.to_string())
match serde_json::to_string(&metadata) {
Ok(metadata) => Some(metadata),
Err(_) if cross_account => None,
Err(_) => Some(raw.to_string()),
}
}
fn pseudonym(&self, domain: &str, account_id: Option<&str>, value: &str) -> String {
@@ -194,6 +241,7 @@ fn apply_identity_to_request(
turn_metadata: Option<&str>,
client_metadata_turn_state: Option<&str>,
client_metadata_turn_metadata: Option<&str>,
cross_account: bool,
) {
for (key, value) in [
(
@@ -215,6 +263,7 @@ fn apply_identity_to_request(
identity,
client_metadata_turn_state,
client_metadata_turn_metadata,
cross_account,
);
request.set_client_metadata(client_metadata);
}
@@ -224,12 +273,17 @@ fn scoped_client_metadata(
identity: &AccountScopedIdentity,
turn_state: Option<&str>,
turn_metadata: Option<&str>,
cross_account: bool,
) -> Option<Value> {
let client_metadata = client_metadata?;
let Value::Object(mut metadata) = client_metadata.clone() else {
return Some(client_metadata.clone());
return (!cross_account).then(|| client_metadata.clone());
};
if cross_account {
remove_cross_account_metadata_state(&mut metadata);
}
for (key, value) in [
(
"x-codex-installation-id",
@@ -254,6 +308,28 @@ fn scoped_client_metadata(
(!metadata.is_empty()).then_some(Value::Object(metadata))
}
fn remove_cross_account_request_state(request: &mut CodexResponsesRequest) {
for key in CROSS_ACCOUNT_IDENTITY_KEYS
.into_iter()
.chain(CROSS_ACCOUNT_OPAQUE_STATE_KEYS)
{
request.replace_existing_identity_field(key, None);
}
for key in ["turn_metadata", TURN_METADATA_KEY] {
request.replace_existing_identity_field(key, None);
}
}
fn remove_cross_account_metadata_state(metadata: &mut Map<String, Value>) {
for key in CROSS_ACCOUNT_IDENTITY_KEYS
.into_iter()
.chain(CROSS_ACCOUNT_OPAQUE_STATE_KEYS)
.chain(["turnMetadata", "turn_metadata", TURN_METADATA_KEY])
{
metadata.remove(key);
}
}
fn replace_existing_metadata_field(
metadata: &mut Map<String, Value>,
key: &str,
@@ -12,7 +12,7 @@ use crate::{
fleet::{
account_failure::{
AccountFailureKind, AccountStateEffect, apply_account_state_effect_immediately,
classify_response_failure, classify_upstream_failure,
classify_account_failure,
},
pool::AccountPoolService,
},
@@ -119,23 +119,31 @@ impl AccountFailureController {
if code.contains("invalid_request")
|| code.contains("not_found")
|| code.contains("context_window")
|| code == "context_length_exceeded"
|| code.contains("invalid_prompt")
|| code.contains("bad_request")
{
return Some(ClientFailure::new(failure.clone(), 400, false));
}
if code.contains("server_overloaded") {
if matches!(code.as_str(), "server_is_overloaded" | "slow_down")
|| code.contains("server_overloaded")
{
return Some(ClientFailure::new(failure.clone(), 503, true));
}
classify_response_failure(failure).map(|classified| {
let status = match classified.kind {
AccountFailureKind::ModelUnsupported => 400,
AccountFailureKind::Expired => 401,
AccountFailureKind::Disabled | AccountFailureKind::Banned => 403,
AccountFailureKind::QuotaExhausted | AccountFailureKind::RateLimited => 429,
};
ClientFailure::new(failure.clone(), status, false)
})
if code == "usage_not_included" {
return Some(ClientFailure::new(failure.clone(), 403, false));
}
classify_account_failure(&super::failure_observation::from_response(failure)).map(
|classified| {
let status = match classified.kind {
AccountFailureKind::ModelUnsupported => 400,
AccountFailureKind::Expired => 401,
AccountFailureKind::Disabled | AccountFailureKind::Banned => 403,
AccountFailureKind::QuotaExhausted | AccountFailureKind::RateLimited => 429,
};
ClientFailure::new(failure.clone(), status, false)
},
)
}
}
@@ -171,7 +179,7 @@ fn classify_upstream(account_id: &str, facts: &UpstreamFailureFacts) -> Option<C
account_id,
&facts.body,
facts.status_code,
classify_upstream_failure(facts)?,
classify_account_failure(&super::failure_observation::from_upstream(facts))?,
)
}
@@ -180,7 +188,7 @@ fn classify_sse(account_id: &str, failure: &ResponsesSseFailure) -> Option<Class
account_id,
&crate::dispatch::failure::sse::sse_failure_error_body(failure),
failure.explicit_status_code,
classify_response_failure(failure)?,
classify_account_failure(&super::failure_observation::from_response(failure))?,
)
}
@@ -0,0 +1,31 @@
//! 上游 typed failure 到 fleet 账号事实的边界映射。
use crate::{
fleet::account_gateway::AccountFailureObservation,
upstream::openai::{failure::UpstreamFailureFacts, protocol::responses::ResponsesSseFailure},
};
pub(super) fn from_upstream(facts: &UpstreamFailureFacts) -> AccountFailureObservation {
AccountFailureObservation {
status_code: facts.status_code,
code: facts.code.clone(),
error_type: facts.error_type.clone(),
identity_authorization_error: facts.identity_authorization_error.clone(),
identity_error_code: facts.identity_error_code.clone(),
message: facts.message.clone(),
body: facts.body.clone(),
retry_after_seconds: facts.retry_after_seconds,
}
}
pub(super) fn from_response(failure: &ResponsesSseFailure) -> AccountFailureObservation {
AccountFailureObservation {
status_code: failure.explicit_status_code,
code: failure.upstream_code.clone(),
error_type: failure.upstream_type.clone(),
message: failure.message.clone(),
body: failure.message.clone(),
retry_after_seconds: failure.retry_after_seconds,
..AccountFailureObservation::default()
}
}
+266 -78
View File
@@ -16,7 +16,10 @@ use crate::{
},
upstream::openai::{
failure::{UpstreamFailureFacts, UpstreamFailureKind},
protocol::responses::{CodexResponsesRequest, PreviousResponseScope, StreamCommitPolicy},
protocol::responses::{
CodexResponsesRequest, LocalReplayItem, LocalReplayTranscript, PreviousResponseScope,
StreamCommitPolicy,
},
transport::CodexBackendTransport,
transport::websocket::PreviousResponseUnavailableReason,
},
@@ -38,7 +41,8 @@ enum HistorySource {
pub(super) struct HistoryState {
source: HistorySource,
managed_entry: Option<SessionAffinityEntry>,
local_replay_input: Option<Arc<Vec<Value>>>,
local_replay_transcript: Option<Arc<LocalReplayTranscript>>,
local_replay_account_id: Option<String>,
external_attempted_account_id: Option<String>,
replay_owner_account: bool,
}
@@ -46,7 +50,8 @@ pub(super) struct HistoryState {
#[derive(Default)]
pub(in crate::dispatch) struct ConnectionReplaySnapshot {
last_response_id: Option<String>,
input: Option<Arc<Vec<Value>>>,
account_id: Option<String>,
transcript: Option<Arc<LocalReplayTranscript>>,
}
pub(in crate::dispatch) struct ConnectionReplayPlan {
@@ -54,16 +59,51 @@ pub(in crate::dispatch) struct ConnectionReplayPlan {
}
enum ConnectionReplayUpdate {
Replace(Vec<Value>),
Append {
history: Arc<Vec<Value>>,
input: Vec<Value>,
},
Replace,
Append { history: Arc<LocalReplayTranscript> },
Unavailable,
}
/// 已确认来源和目标账号不同的完整 transcript 重放边界。
pub(in crate::dispatch) struct CrossAccountReplay {
source_account_id: String,
target_account_id: String,
}
impl CrossAccountReplay {
fn between(source_account_id: &str, target_account_id: &str) -> Option<Self> {
(source_account_id != target_account_id).then(|| Self {
source_account_id: source_account_id.to_string(),
target_account_id: target_account_id.to_string(),
})
}
pub(in crate::dispatch) fn validate_target(&self, account_id: &str) {
debug_assert_ne!(self.source_account_id, self.target_account_id);
debug_assert_eq!(self.target_account_id, account_id);
}
}
pub(super) struct PreparedHistoryRequest {
pub(super) request: CodexResponsesRequest,
pub(super) cross_account_replay: Option<CrossAccountReplay>,
}
pub(in crate::dispatch) struct HistoryController;
pub(in crate::dispatch) fn sanitize_cross_account_input(request: &mut CodexResponsesRequest) {
if request.input().is_empty() {
return;
}
let input = request
.input()
.iter()
.cloned()
.filter_map(sanitize_cross_account_item)
.collect();
request.set_input(input);
}
impl HistoryController {
pub(in crate::dispatch) fn new_connection_replay_snapshot() -> ConnectionReplaySnapshot {
ConnectionReplaySnapshot::default()
@@ -73,28 +113,29 @@ impl HistoryController {
snapshot: &ConnectionReplaySnapshot,
request: &mut CodexResponsesRequest,
) -> ConnectionReplayPlan {
let turn_input = request.input().to_vec();
let update = match request.previous_response_id() {
None => {
request.local_replay_input = None;
ConnectionReplayUpdate::Replace(turn_input)
clear_local_replay(request);
ConnectionReplayUpdate::Replace
}
Some(previous_response_id)
if snapshot.last_response_id.as_deref() == Some(previous_response_id) =>
{
if let Some(history) = snapshot.input.as_ref() {
request.local_replay_input = Some(Arc::clone(history));
if let (Some(account_id), Some(history)) =
(snapshot.account_id.as_ref(), snapshot.transcript.as_ref())
{
request.local_replay_transcript = Some(Arc::clone(history));
request.local_replay_account_id = Some(account_id.clone());
ConnectionReplayUpdate::Append {
history: Arc::clone(history),
input: turn_input,
}
} else {
request.local_replay_input = None;
clear_local_replay(request);
ConnectionReplayUpdate::Unavailable
}
}
Some(_) => {
request.local_replay_input = None;
clear_local_replay(request);
ConnectionReplayUpdate::Unavailable
}
};
@@ -107,18 +148,38 @@ impl HistoryController {
plan: ConnectionReplayPlan,
response_id: String,
output: Vec<Value>,
account_id: String,
request_input: Vec<Value>,
continued_from_previous_response: bool,
) {
snapshot.last_response_id = Some(response_id);
snapshot.account_id = Some(account_id.clone());
match plan.update {
ConnectionReplayUpdate::Replace(input) => {
snapshot.input = Some(Arc::new(sanitized_transcript(input, output)));
ConnectionReplayUpdate::Replace => {
snapshot.transcript = Some(Arc::new(transcript_for_turn(
request_input,
output,
&account_id,
)));
}
ConnectionReplayUpdate::Append { mut history, input } => {
let mut turn = sanitized_transcript(input, output);
Arc::make_mut(&mut history).append(&mut turn);
snapshot.input = Some(history);
ConnectionReplayUpdate::Append { mut history } => {
if continued_from_previous_response {
append_turn(
Arc::make_mut(&mut history),
request_input,
output,
&account_id,
);
snapshot.transcript = Some(history);
} else {
snapshot.transcript = Some(Arc::new(transcript_for_turn(
request_input,
output,
&account_id,
)));
}
}
ConnectionReplayUpdate::Unavailable => snapshot.input = None,
ConnectionReplayUpdate::Unavailable => snapshot.transcript = None,
}
}
@@ -145,51 +206,67 @@ impl HistoryController {
None
}
};
let local_replay_input = request.local_replay_input.clone();
let source = match (&managed_entry, &local_replay_input) {
let (local_replay_transcript, local_replay_account_id) = match (
request.local_replay_transcript.clone(),
request.local_replay_account_id.clone(),
) {
(Some(transcript), Some(account_id)) => (Some(transcript), Some(account_id)),
(None, None) => (None, None),
_ => {
tracing::warn!(
previous_response_id,
"Ignoring incomplete connection-local replay provenance"
);
(None, None)
}
};
let has_complete_local_replay = local_replay_transcript.is_some();
let source = match (&managed_entry, has_complete_local_replay) {
(Some(_), _) => HistorySource::Managed,
(None, Some(_)) => HistorySource::ConnectionLocal,
(None, None) => HistorySource::ExternalUnknown,
(None, true) => HistorySource::ConnectionLocal,
(None, false) => HistorySource::ExternalUnknown,
};
HistoryState {
source,
managed_entry,
local_replay_input,
local_replay_transcript,
local_replay_account_id,
external_attempted_account_id: None,
replay_owner_account: false,
}
}
pub(super) fn preferred_account_id(state: &HistoryState) -> Option<&str> {
state
.managed_entry
.as_ref()
.map(|entry| entry.account_id.as_str())
state.local_replay_account_id.as_deref().or_else(|| {
state
.managed_entry
.as_ref()
.map(|entry| entry.account_id.as_str())
})
}
pub(super) fn prepare_attempt(
state: &mut HistoryState,
request: &CodexResponsesRequest,
account_id: &str,
) -> Result<CodexResponsesRequest, String> {
) -> Result<PreparedHistoryRequest, String> {
let prepared = match state.source {
HistorySource::None => Some(request.clone()),
HistorySource::None => Some(PreparedHistoryRequest::plain(request.clone())),
HistorySource::ExternalUnknown => {
match state.external_attempted_account_id.as_deref() {
None => {
state.external_attempted_account_id = Some(account_id.to_string());
Some(external_history_request(request))
}
Some(first_account_id) if first_account_id == account_id => {
Some(external_history_request(request))
Some(PreparedHistoryRequest::plain(external_history_request(
request,
)))
}
Some(first_account_id) if first_account_id == account_id => Some(
PreparedHistoryRequest::plain(external_history_request(request)),
),
Some(_) => None,
}
}
HistorySource::ConnectionLocal => state
.local_replay_input
.as_deref()
.map(|replay_input| full_replay_request(request, replay_input)),
HistorySource::ConnectionLocal => full_replay_from_state(state, request, account_id),
HistorySource::Managed => prepare_managed_attempt(state, request, account_id),
};
prepared.ok_or_else(|| CROSS_ACCOUNT_HISTORY_UNAVAILABLE_MESSAGE.to_string())
@@ -213,7 +290,7 @@ impl HistoryController {
.is_some_and(|account_id| Self::can_retry_same_account(state, account_id)),
can_retry_next_candidate: match state.source {
HistorySource::None | HistorySource::ConnectionLocal => true,
HistorySource::Managed => state.local_replay_input.is_some(),
HistorySource::Managed => state.local_replay_transcript.is_some(),
HistorySource::ExternalUnknown => false,
},
}
@@ -294,7 +371,7 @@ impl HistoryController {
return false;
};
entry.account_id == account_id
&& state.local_replay_input.is_some()
&& state.local_replay_transcript.is_some()
&& !state.replay_owner_account
}
}
@@ -303,10 +380,10 @@ fn prepare_managed_attempt(
state: &HistoryState,
request: &CodexResponsesRequest,
account_id: &str,
) -> Option<CodexResponsesRequest> {
) -> Option<PreparedHistoryRequest> {
let entry = state.managed_entry.as_ref()?;
if entry.account_id != account_id || state.replay_owner_account {
return replay_for_conversation(state, request, &entry.conversation_id);
return replay_for_conversation(state, request, account_id, &entry.conversation_id);
}
match entry.continuation_scope {
@@ -320,10 +397,10 @@ fn prepare_managed_attempt(
if prepared.turn_state.as_deref().is_none_or(str::is_empty) {
prepared.turn_state.clone_from(&entry.turn_state);
}
Some(prepared)
Some(PreparedHistoryRequest::plain(prepared))
}
PreviousResponseScope::ReplayRequired => {
replay_for_conversation(state, request, &entry.conversation_id)
replay_for_conversation(state, request, account_id, &entry.conversation_id)
}
PreviousResponseScope::Unavailable => None,
}
@@ -332,21 +409,37 @@ fn prepare_managed_attempt(
fn replay_for_conversation(
state: &HistoryState,
request: &CodexResponsesRequest,
account_id: &str,
conversation_id: &str,
) -> Option<CodexResponsesRequest> {
state.local_replay_input.as_deref().map(|replay_input| {
let mut prepared = full_replay_request(request, replay_input);
prepared.local_conversation_id = Some(conversation_id.to_string());
) -> Option<PreparedHistoryRequest> {
full_replay_from_state(state, request, account_id).map(|mut prepared| {
prepared.request.local_conversation_id = Some(conversation_id.to_string());
prepared
})
}
fn full_replay_from_state(
state: &HistoryState,
request: &CodexResponsesRequest,
account_id: &str,
) -> Option<PreparedHistoryRequest> {
let transcript = state.local_replay_transcript.as_deref()?;
let source_account_id = state.local_replay_account_id.as_deref()?;
Some(full_replay_request(
request,
transcript,
source_account_id,
account_id,
))
}
impl HistoryState {
fn empty() -> Self {
Self {
source: HistorySource::None,
managed_entry: None,
local_replay_input: None,
local_replay_transcript: None,
local_replay_account_id: None,
external_attempted_account_id: None,
replay_owner_account: false,
}
@@ -363,43 +456,138 @@ fn external_history_request(original: &CodexResponsesRequest) -> CodexResponsesR
fn full_replay_request(
original: &CodexResponsesRequest,
replay_input: &[Value],
) -> CodexResponsesRequest {
transcript: &LocalReplayTranscript,
source_account_id: &str,
target_account_id: &str,
) -> PreparedHistoryRequest {
let cross_account_replay = CrossAccountReplay::between(source_account_id, target_account_id);
let mut request = original.clone();
let mut input = Vec::with_capacity(replay_input.len() + original.input().len());
input.extend_from_slice(replay_input);
let mut input = replay_input_for_account(transcript, target_account_id);
input.reserve(original.input().len());
input.extend_from_slice(original.input());
sanitize_replay_items(&mut input);
request.set_input(input);
request.set_previous_response_id(None);
request.turn_state = None;
request.stream_commit_policy = StreamCommitPolicy::UntilOutputOrTerminal;
request
PreparedHistoryRequest {
request,
cross_account_replay,
}
}
fn sanitized_transcript(mut input: Vec<Value>, mut output: Vec<Value>) -> Vec<Value> {
sanitize_replay_items(&mut input);
sanitize_replay_items(&mut output);
input.append(&mut output);
input
}
fn sanitize_replay_items(values: &mut [Value]) {
for value in values {
match value {
Value::Array(values) => sanitize_replay_items(values),
Value::Object(object) => {
object.remove("id");
object.remove("encrypted_content");
for value in object.values_mut() {
sanitize_replay_items(std::slice::from_mut(value));
}
}
Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {}
impl PreparedHistoryRequest {
fn plain(request: CodexResponsesRequest) -> Self {
Self {
request,
cross_account_replay: None,
}
}
}
fn clear_local_replay(request: &mut CodexResponsesRequest) {
request.local_replay_transcript = None;
request.local_replay_account_id = None;
}
fn transcript_for_turn(
input: Vec<Value>,
output: Vec<Value>,
account_id: &str,
) -> LocalReplayTranscript {
let mut transcript = LocalReplayTranscript::default();
append_client_input(&mut transcript, input);
append_account_output(&mut transcript, output, account_id);
transcript
}
fn append_turn(
transcript: &mut LocalReplayTranscript,
input: Vec<Value>,
output: Vec<Value>,
account_id: &str,
) {
project_transcript_to_account(transcript, account_id);
append_client_input(transcript, input);
append_account_output(transcript, output, account_id);
}
fn append_client_input(transcript: &mut LocalReplayTranscript, input: Vec<Value>) {
transcript
.items
.extend(input.into_iter().map(LocalReplayItem::ClientInput));
}
fn append_account_output(
transcript: &mut LocalReplayTranscript,
output: Vec<Value>,
account_id: &str,
) {
transcript.items.extend(
output
.into_iter()
.map(|item| LocalReplayItem::AccountOutput {
account_id: account_id.to_string(),
item,
}),
);
}
fn project_transcript_to_account(transcript: &mut LocalReplayTranscript, account_id: &str) {
transcript.items = transcript
.items
.drain(..)
.filter_map(|item| match item {
LocalReplayItem::AccountOutput {
account_id: owner,
item,
} if owner != account_id => {
sanitize_cross_account_item(item).map(LocalReplayItem::SanitizedOutput)
}
item => Some(item),
})
.collect();
}
fn replay_input_for_account(transcript: &LocalReplayTranscript, account_id: &str) -> Vec<Value> {
transcript
.items
.iter()
.filter_map(|item| match item {
LocalReplayItem::ClientInput(value) | LocalReplayItem::SanitizedOutput(value) => {
Some(value.clone())
}
LocalReplayItem::AccountOutput {
account_id: owner,
item,
} if owner == account_id => Some(without_output_id(item.clone())),
LocalReplayItem::AccountOutput { item, .. } => {
sanitize_cross_account_item(item.clone())
}
})
.collect()
}
fn without_output_id(mut item: Value) -> Value {
if let Value::Object(object) = &mut item {
object.remove("id");
}
item
}
fn sanitize_cross_account_item(mut item: Value) -> Option<Value> {
if let Value::Object(object) = &mut item {
if matches!(
object.get("type").and_then(Value::as_str),
Some("compaction" | "compaction_summary" | "context_compaction")
) {
return None;
}
object.remove("id");
object.remove("encrypted_content");
}
Some(item)
}
fn is_continuation_busy(facts: &UpstreamFailureFacts) -> bool {
matches!(
facts.kind,
+124 -52
View File
@@ -19,7 +19,10 @@ use crate::{
AccountIdentityScope, AccountIdentityService, AccountScopedRequest,
SessionAffinityService,
},
errors::{ClientFailure, ResponseDispatchError, upstream_error_set_cookie_headers},
errors::{
ClientFailure, ResponseDispatchError, upstream_error_rate_limit_headers,
upstream_error_set_cookie_headers,
},
lifecycle::trace::{ResponseDispatchAttempt, ResponseDispatchTrace},
},
fleet::{account::Account, pool::AccountLease},
@@ -51,6 +54,7 @@ use crate::{
},
},
fleet::pool::AccountPoolService,
fleet::scheduler::AttemptFeedback,
upstream::openai::{
failure::UpstreamFailureFacts,
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
@@ -61,6 +65,7 @@ mod account_failure;
mod affinity;
pub(crate) mod cloudflare;
pub(crate) mod cyber_policy;
mod failure_observation;
pub(in crate::dispatch) mod history;
mod quota;
pub(in crate::dispatch) mod telemetry;
@@ -270,22 +275,11 @@ pub(in crate::dispatch) struct ControllerEnter {
pub excluded_account_ids: BTreeSet<String>,
}
pub(in crate::dispatch) struct AttemptAccountPreparationContext<'a> {
pub codex: &'a CodexBackendClient,
pub account_identity: &'a AccountIdentityService,
pub request_id: &'a str,
}
pub(in crate::dispatch) struct PreparedAttemptAccount {
pub lease: AccountLease,
pub cookie_header: Option<String>,
}
pub(in crate::dispatch) enum AttemptAccountPreparation {
Ready(Box<PreparedAttemptAccount>),
Rejected,
}
pub(in crate::dispatch) struct AttemptRoutePreparationContext<'a> {
pub account_identity: &'a AccountIdentityService,
pub account_id: &'a str,
@@ -299,38 +293,16 @@ pub(in crate::dispatch) enum AttemptRoutePreparation {
impl ControllerSet {
pub(in crate::dispatch) async fn prepare_attempt_account(
&self,
context: AttemptAccountPreparationContext<'_>,
acquired: AccountLease,
) -> AttemptAccountPreparation {
let usage_cookie = self
.cloudflare
.cookie_header_for_request(&acquired.account.id, "/codex/usage")
.await;
let acquired = match QuotaController::enter(
quota::QuotaEnterContext {
account_pool: &self.account_pool,
codex: context.codex,
cookie_header: usage_cookie.as_deref(),
account_identity: context.account_identity,
request_id: context.request_id,
},
acquired,
)
.await
{
quota::QuotaEnterOutcome::Ready(acquired) => *acquired,
quota::QuotaEnterOutcome::LimitReached => {
return AttemptAccountPreparation::Rejected;
}
};
) -> PreparedAttemptAccount {
let cookie_header = self
.cloudflare
.cookie_header_for_request(&acquired.account.id, "/codex/responses")
.await;
AttemptAccountPreparation::Ready(Box::new(PreparedAttemptAccount {
PreparedAttemptAccount {
lease: acquired,
cookie_header,
}))
}
}
pub(in crate::dispatch) fn prepare_attempt_route(
@@ -339,18 +311,19 @@ impl ControllerSet {
request: &CodexResponsesRequest,
context: AttemptRoutePreparationContext<'_>,
) -> AttemptRoutePreparation {
let request = match HistoryController::prepare_attempt(
let prepared = match HistoryController::prepare_attempt(
&mut scope.history,
request,
context.account_id,
) {
Ok(request) => request,
Ok(prepared) => prepared,
Err(message) => return AttemptRoutePreparation::Unavailable { message },
};
AttemptRoutePreparation::Ready(Box::new(context.account_identity.scope_request(
&mut scope.identity,
request,
prepared.request,
context.account_id,
prepared.cross_account_replay.as_ref(),
)))
}
@@ -384,8 +357,12 @@ impl ControllerSet {
.await;
let _ = best_effort_controller_io(
"quota.complete_headers",
self.account_pool
.sync_passive_rate_limit_headers(account, rate_limit_headers),
usage::UsageController::sync_passive_quota(
&self.account_pool,
&account.id,
account.plan_type.as_deref(),
rate_limit_headers,
),
)
.await;
}
@@ -394,19 +371,28 @@ impl ControllerSet {
&self,
context: AttemptUpstreamErrorContext<'_>,
) {
let _ = best_effort_controller_io(
let cookie = best_effort_controller_io(
"cookie.attempt_error",
self.cloudflare.capture_set_cookie_headers(
&context.account.id,
upstream_error_set_cookie_headers(context.error),
),
)
.await;
let _ = best_effort_controller_io(
);
let rate_limit_headers = upstream_error_rate_limit_headers(context.error);
let quota = best_effort_controller_io(
"quota.attempt_error_headers",
usage::UsageController::sync_passive_quota(
&self.account_pool,
&context.account.id,
context.account.plan_type.as_deref(),
&rate_limit_headers,
),
);
let telemetry = best_effort_controller_io(
"telemetry.attempt_error",
telemetry::TelemetryController::observe_upstream_error(&self.recorder, &context),
)
.await;
);
let _ = tokio::join!(cookie, quota, telemetry);
}
pub(in crate::dispatch) async fn observe_dispatch_error(
@@ -453,6 +439,18 @@ impl ControllerSet {
)
.await;
}
let feedback = if exit.completed {
AttemptFeedback::Completed {
first_token_ms: nonnegative_u64(exit.response.first_token_ms),
}
} else {
AttemptFeedback::Incomplete {
first_token_ms: nonnegative_u64(exit.response.first_token_ms),
}
};
self.account_pool
.report_attempt(exit.account_id, feedback)
.await;
if successful {
let _ = best_effort_controller_io(
@@ -657,6 +655,12 @@ impl ControllerSet {
}),
)
.await;
self.account_pool
.report_attempt(
&context.account_id,
stream_feedback(summary, context.started_at, context.attempt.started_at()),
)
.await;
let fallback = usage::StreamUsageExit {
rate_limit_headers: context.rate_limit_headers.clone(),
@@ -726,7 +730,6 @@ impl ControllerSet {
}),
)
.await;
if !matches!(outcome, FinalOutcome::Cancelled | FinalOutcome::Shutdown) {
if let Some(classification) = stream_failure.as_ref() {
self.apply_stream_policy_failure(&scope, &context.account_id, classification)
@@ -750,6 +753,13 @@ impl ControllerSet {
usage::UsageController::observe_attempt(&self.account_pool, &scope.usage, observation),
)
.await;
if let (Some(account), Some(feedback)) =
(observation.account.as_ref(), attempt_feedback(observation))
{
self.account_pool
.report_attempt(&account.id, feedback)
.await;
}
if matches!(
observation.kind,
@@ -918,9 +928,6 @@ impl ControllerSet {
AttemptObservationKind::NoCandidate { .. } => {
AttemptDecision::Return(AttemptReturnKind::Observed)
}
AttemptObservationKind::CandidatePreparationRejected => {
unreachable!("candidate preparation controller owns it")
}
AttemptObservationKind::RoutePreparationRejected { message } => {
AttemptDecision::Return(AttemptReturnKind::RouteUnavailable {
message: message.clone(),
@@ -953,6 +960,71 @@ impl ControllerSet {
}
}
fn attempt_feedback(observation: &AttemptObservation) -> Option<AttemptFeedback> {
observation.attempt.as_ref()?;
if matches!(
observation.kind,
AttemptObservationKind::CompleteResponse(
crate::dispatch::lifecycle::contract::CompleteResponseFacts::Completed
| crate::dispatch::lifecycle::contract::CompleteResponseFacts::Incomplete,
)
) {
return None;
}
Some(AttemptFeedback::FailedBeforeFirstToken)
}
fn stream_feedback(
summary: &crate::dispatch::lifecycle::stream::StreamSummary,
request_started_at: Instant,
attempt_started_at: Instant,
) -> AttemptFeedback {
let first_token_ms = attempt_first_token_ms(
summary.first_token_ms,
request_started_at,
attempt_started_at,
);
match summary.terminal {
StreamTerminal::Completed { .. } => AttemptFeedback::Completed { first_token_ms },
StreamTerminal::Incomplete { .. } => AttemptFeedback::Incomplete { first_token_ms },
StreamTerminal::Cancelled | StreamTerminal::DownstreamClosed => AttemptFeedback::Cancelled,
StreamTerminal::Shutdown => AttemptFeedback::Shutdown,
StreamTerminal::Failed { .. }
| StreamTerminal::UpstreamClosed
| StreamTerminal::UpstreamError { .. }
| StreamTerminal::ProtocolError { .. }
| StreamTerminal::CaptureLimitExceeded => {
summary
.first_token_ms
.map_or(AttemptFeedback::FailedBeforeFirstToken, |_| {
AttemptFeedback::FailedAfterFirstToken {
first_token_ms: first_token_ms.unwrap_or(0),
}
})
}
}
}
fn attempt_first_token_ms(
request_first_token_ms: Option<i64>,
request_started_at: Instant,
attempt_started_at: Instant,
) -> Option<u64> {
let offset_ms = i64::try_from(
attempt_started_at
.saturating_duration_since(request_started_at)
.as_millis(),
)
.ok()?;
request_first_token_ms
.and_then(|value| value.checked_sub(offset_ms))
.and_then(|value| u64::try_from(value).ok())
}
fn nonnegative_u64(value: Option<i64>) -> Option<u64> {
value.and_then(|value| u64::try_from(value).ok())
}
fn response_failure(observation: &AttemptObservation) -> Option<&ResponsesSseFailure> {
match &observation.kind {
AttemptObservationKind::CompleteResponse(
+16 -134
View File
@@ -1,139 +1,29 @@
//! 请求前配额验证规则 owner。
use serde_json::Value;
//! 配额失败规则 owner。
use crate::{
dispatch::{
affinity::AccountIdentityService,
controllers::ControllerFailureFact,
errors::ClientFailure,
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind},
lifecycle::contract::{
AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind,
},
lifecycle::contract::{AttemptDecision, AttemptObservation, AttemptReturnKind},
},
fleet::{
account_failure::{
AccountFailureKind, AccountStateEffect, apply_account_state_effect_immediately,
classify_response_failure, classify_upstream_failure,
},
pool::{AccountLease, AccountPoolService},
quota::{
quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds,
quota_snapshot_reset_at,
classify_account_failure,
},
pool::AccountPoolService,
},
upstream::openai::{
failure::UpstreamFailureFacts,
protocol::responses::ResponsesSseFailure,
transport::{CodexBackendClient, CodexRequestContext},
failure::UpstreamFailureFacts, protocol::responses::ResponsesSseFailure,
transport::CodexBackendClient,
},
};
const LIMIT_REACHED_MESSAGE: &str = "Upstream usage quota still reports limit_reached";
pub(super) enum QuotaEnterOutcome {
Ready(Box<AccountLease>),
LimitReached,
}
pub(super) struct QuotaEnterContext<'a> {
pub account_pool: &'a AccountPoolService,
pub codex: &'a CodexBackendClient,
pub account_identity: &'a AccountIdentityService,
pub request_id: &'a str,
pub cookie_header: Option<&'a str>,
}
pub(super) struct QuotaController;
impl QuotaController {
pub(super) async fn enter(
context: QuotaEnterContext<'_>,
acquired: AccountLease,
) -> QuotaEnterOutcome {
if !acquired.account.quota_verify_required {
return QuotaEnterOutcome::Ready(Box::new(acquired));
}
let account_id = acquired.account.id.clone();
let identity = context.account_identity.scope_auxiliary(&account_id);
let usage = context
.codex
.fetch_usage(CodexRequestContext {
access_token: &acquired.account.access_token,
account_id: acquired.account.account_id.as_deref(),
request_id: context.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: context.cookie_header,
installation_id: Some(&identity.installation_id),
session_id: None,
thread_id: None,
client_request_id: None,
turn_id: None,
})
.await;
let raw = match usage {
Ok(raw) => raw,
Err(error) => {
tracing::warn!(
request_id = %context.request_id,
account_id = %account_id,
quota_verify_required = true,
quota_verify_result = "upstream_error",
retry_with_another_account = false,
error = %error,
"Failed to verify stale quota state before upstream request"
);
return QuotaEnterOutcome::Ready(Box::new(acquired));
}
};
let quota = quota_from_usage(&raw);
context
.account_pool
.apply_quota_snapshot(&account_id, &quota)
.await;
if quota_snapshot_limit_reached(&quota) {
acquired.release_without_usage().await;
tracing::info!(
request_id = %context.request_id,
account_id = %account_id,
quota_verify_required = true,
quota_verify_result = "limit_reached",
retry_with_another_account = true,
"Quota verification reported exhausted account before upstream request"
);
return QuotaEnterOutcome::LimitReached;
}
QuotaEnterOutcome::Ready(Box::new(acquired_with_verified_quota(acquired, &quota)))
}
pub(super) fn classify(observation: &AttemptObservation) -> Option<ClassifiedQuotaFailure> {
if matches!(
observation.kind,
AttemptObservationKind::CandidatePreparationRejected
) {
return observation
.account
.as_ref()
.map(|account| ClassifiedQuotaFailure {
exhaustion: AccountExhaustionRecord::new(
account.id.clone(),
ExhaustedAccountKind::RateLimited,
LIMIT_REACHED_MESSAGE,
),
effect: None,
});
}
classify(observation)
}
@@ -200,7 +90,11 @@ fn classify_upstream(
account_id: &str,
facts: &UpstreamFailureFacts,
) -> Option<ClassifiedQuotaFailure> {
classified_quota_failure(account_id, &facts.body, classify_upstream_failure(facts)?)
classified_quota_failure(
account_id,
&facts.body,
classify_account_failure(&super::failure_observation::from_upstream(facts))?,
)
}
fn classify_stream(
@@ -208,7 +102,11 @@ fn classify_stream(
failure: &ResponsesSseFailure,
) -> Option<ClassifiedQuotaFailure> {
let body = crate::dispatch::failure::sse::sse_failure_error_body(failure);
classified_quota_failure(account_id, &body, classify_response_failure(failure)?)
classified_quota_failure(
account_id,
&body,
classify_account_failure(&super::failure_observation::from_response(failure))?,
)
}
fn classified_quota_failure(
@@ -229,19 +127,3 @@ fn classified_quota_failure(
effect: classified.effect,
})
}
fn acquired_with_verified_quota(mut acquired: AccountLease, quota: &Value) -> AccountLease {
let limit_reached = quota_snapshot_limit_reached(quota);
acquired.account.quota_verify_required = false;
acquired.account.quota_limit_reached = limit_reached;
acquired.account.quota_cooldown_until = limit_reached
.then_some(quota_snapshot_reset_at(quota))
.flatten();
if let Some(reset_at) = quota_snapshot_reset_at(quota) {
acquired.account.window_reset_at = Some(reset_at);
if let Some(limit_window_seconds) = quota_snapshot_limit_window_seconds(quota) {
acquired.account.limit_window_seconds = Some(limit_window_seconds);
}
}
acquired
}
@@ -10,23 +10,25 @@ use crate::{
infra::time::elapsed_millis_i64,
telemetry::{
ops::types::OpsErrorLog,
recorder::{
DispatchErrorLogRecord, Recorder, enrich_event_route_metadata,
enrich_response_request_semantics, enrich_usage_record_identity,
event_kind as response_event_kind, reasoning_effort_from_request,
record_dispatch_error_event,
recorder::Recorder,
usage::types::{
UsageRecord, UsageRecordLevel, metadata_i64, metadata_service_tier,
metadata_string as metadata_string_any,
},
usage::types::UsageRecordLevel,
},
upstream::openai::{
protocol::events::TokenUsage,
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
transport::{
CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics,
WebSocketPoolDecision,
CodexBackendTransport, CodexClientError, CodexResponseMetadata,
CodexUpstreamDiagnostics, WebSocketPoolDecision,
},
},
};
const RESPONSE_EVENT_KIND: &str = "v1.response";
const RESPONSES_API_KIND: &str = "responses";
use crate::dispatch::{
errors::{
ResponseDispatchError, dispatch_error_metadata, enrich_response_dispatch_error_metadata,
@@ -86,6 +88,119 @@ pub(super) struct ResponseDispatchErrorEventRecord<'a> {
pub(super) error: &'a ResponseDispatchError,
}
pub(super) struct ResponseUsageEventRecord<'a> {
pub(super) recorder: &'a Recorder,
pub(super) request_id: &'a str,
pub(super) client_api_key_id: Option<&'a str>,
pub(super) account_id: &'a str,
pub(super) route: &'a str,
pub(super) model: &'a str,
pub(super) requested_model: Option<&'a str>,
pub(super) client_ip: Option<&'a str>,
pub(super) client_user_agent: Option<&'a str>,
pub(super) reasoning_effort: Option<&'a str>,
pub(super) service_tier: Option<&'a str>,
pub(super) started_at: Instant,
pub(super) status_code: i64,
pub(super) message: &'a str,
pub(super) usage: Option<TokenUsage>,
pub(super) metadata: Value,
pub(super) rate_limit_headers: &'a [(String, String)],
}
pub(super) struct LiveResponseStreamEventRecord<'a> {
pub(super) context: &'a StreamContext<'a>,
pub(super) status_code: i64,
pub(super) level: UsageRecordLevel,
pub(super) message: &'a str,
pub(super) usage: Option<TokenUsage>,
pub(super) metadata: Value,
pub(super) rate_limit_headers: &'a [(String, String)],
pub(super) body: &'a str,
}
struct DispatchErrorLogRecord<'a> {
recorder: &'a Recorder,
request_id: &'a str,
client_api_key_id: Option<&'a str>,
provider: Option<&'a str>,
account_id: Option<&'a str>,
route: &'a str,
model: &'a str,
started_at: Instant,
status_code: i64,
message: &'a str,
metadata: Value,
}
pub(super) async fn record_response_event(record: ResponseUsageEventRecord<'_>) {
let mut metadata = record.metadata;
enrich_event_route_metadata(&mut metadata, record.route);
let mut event = UsageRecord::new(
RESPONSE_EVENT_KIND,
record.message,
record.account_id,
record.model,
record.status_code,
);
event.request_id = Some(record.request_id.to_string());
event.client_api_key_id = record.client_api_key_id.map(ToString::to_string);
event.route = Some(record.route.to_string());
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
event.requested_model = normalized_string(record.requested_model.or(Some(record.model)));
event.upstream_model = normalized_string(Some(record.model));
event.service_tier = normalized_string(record.service_tier);
enrich_usage_record_identity(
&mut metadata,
record.requested_model,
record.model,
record.client_ip,
record.client_user_agent,
record.reasoning_effort,
record.service_tier,
);
if !record.rate_limit_headers.is_empty()
&& let Some(object) = metadata.as_object_mut()
{
object.insert(
"rateLimitHeaders".to_string(),
serde_json::json!(record.rate_limit_headers),
);
}
lift_success_fact_fields(&mut event, &mut metadata, record.usage);
event.metadata = metadata;
if let Err(error) = record.recorder.record_usage(event).await {
tracing::error!(account_id = %record.account_id, error = %error, "Failed to record response event");
}
}
async fn record_dispatch_error_event(record: DispatchErrorLogRecord<'_>) {
let mut metadata = record.metadata;
enrich_event_route_metadata(&mut metadata, record.route);
let mut event = OpsErrorLog::new(RESPONSE_EVENT_KIND, record.message);
event.request_id = Some(record.request_id.to_string());
event.client_api_key_id = record.client_api_key_id.map(ToString::to_string);
event.provider = record.provider.map(ToString::to_string);
event.account_id = record.account_id.map(ToString::to_string);
event.route = Some(record.route.to_string());
event.model = Some(record.model.to_string());
event.status_code = Some(record.status_code);
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
event.metadata = metadata;
lift_error_fact_fields(&mut event);
if let Err(error) = record.recorder.record_error(event).await {
tracing::error!(
account_id = record.account_id.unwrap_or(""),
error = %error,
"Failed to record dispatch error event"
);
}
}
pub(super) async fn record_response_dispatch_error_event(
record: ResponseDispatchErrorEventRecord<'_>,
) {
@@ -172,6 +287,7 @@ pub(super) async fn record_response_upstream_error_event(
);
enrich_response_request_semantics(&mut metadata, record.request);
event.metadata = metadata;
lift_error_fact_fields(&mut event);
if let Err(error) = record.recorder.record_error(event).await {
tracing::error!(account_id = %record.account_id, error = %error, "Failed to record upstream error event");
}
@@ -221,10 +337,7 @@ pub(super) async fn record_prefetched_response_stream_failure_event(
record.request.service_tier(),
);
enrich_response_request_semantics(&mut metadata, record.request);
let mut event = OpsErrorLog::new(
response_event_kind(record.route),
"v1 responses stream failed",
);
let mut event = OpsErrorLog::new(RESPONSE_EVENT_KIND, "v1 responses stream failed");
event.request_id = Some(record.request_id.to_string());
event.client_api_key_id = record.request.client_api_key_id.clone();
event.provider = Some("openai".to_string());
@@ -242,11 +355,209 @@ pub(super) async fn record_prefetched_response_stream_failure_event(
);
}
event.metadata = metadata;
lift_error_fact_fields(&mut event);
if let Err(error) = record.recorder.record_error(event).await {
tracing::error!(account_id = %record.account_id, error = %error, "Failed to record prefetched stream error event");
}
}
fn lift_error_fact_fields(event: &mut OpsErrorLog) {
event.client_status_code = event
.client_status_code
.or_else(|| metadata_i64(&event.metadata, &["clientStatusCode", "client_status_code"]));
event.upstream_status_code = event
.upstream_status_code
.or_else(|| metadata_i64(&event.metadata, &["upstreamStatusCode", "upstreamStatus"]));
event.transport = event
.transport
.take()
.or_else(|| metadata_string_any(&event.metadata, &["transport"]));
event.attempt_index = event
.attempt_index
.or_else(|| metadata_i64(&event.metadata, &["attemptIndex", "attempt_index"]));
event.failure_class = event
.failure_class
.take()
.or_else(|| metadata_string_any(&event.metadata, &["failureClass", "failure_class"]));
event.response_id = event
.response_id
.take()
.or_else(|| metadata_string_any(&event.metadata, &["responseId", "response_id"]));
event.upstream_request_id = event.upstream_request_id.take().or_else(|| {
metadata_string_any(
&event.metadata,
&[
"upstreamRequestId",
"upstream_request_id",
"openaiRequestId",
],
)
});
event.service_tier = event
.service_tier
.take()
.or_else(|| metadata_service_tier(&event.metadata).map(ToString::to_string));
let Some(metadata) = event.metadata.as_object_mut() else {
return;
};
for key in [
"clientStatusCode",
"client_status_code",
"upstreamStatusCode",
"upstreamStatus",
"transport",
"attemptIndex",
"attempt_index",
"failureClass",
"failure_class",
"responseId",
"response_id",
"upstreamRequestId",
"upstream_request_id",
"openaiRequestId",
"serviceTier",
] {
metadata.remove(key);
}
}
fn lift_success_fact_fields(
event: &mut UsageRecord,
metadata: &mut Value,
usage: Option<TokenUsage>,
) {
event.transport = metadata_string(metadata, "transport");
event.attempt_index = metadata_nonnegative_i64(metadata, "attemptIndex");
event.response_id = metadata_string(metadata, "responseId");
event.upstream_request_id = metadata_string(metadata, "upstreamRequestId")
.or_else(|| metadata_string(metadata, "openaiRequestId"));
event.first_token_ms = metadata_nonnegative_i64(metadata, "firstTokenMs");
if let Some(usage) = usage {
event.input_tokens = i64::try_from(usage.input_tokens).ok();
event.output_tokens = i64::try_from(usage.output_tokens).ok();
event.cached_tokens = i64::try_from(usage.cached_tokens).ok();
event.cache_write_tokens = i64::try_from(usage.cache_write_tokens).ok();
event.reasoning_tokens = i64::try_from(usage.reasoning_tokens).ok();
}
let Some(object) = metadata.as_object_mut() else {
return;
};
for key in [
"usage",
"route",
"apiKind",
"requestedModel",
"upstreamModel",
"serviceTier",
"statusCode",
"transport",
"attemptIndex",
"responseId",
"upstreamRequestId",
"openaiRequestId",
"firstTokenMs",
] {
object.remove(key);
}
}
fn metadata_string(metadata: &Value, field: &str) -> Option<String> {
metadata
.get(field)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
fn metadata_nonnegative_i64(metadata: &Value, field: &str) -> Option<i64> {
metadata
.get(field)
.and_then(Value::as_i64)
.filter(|value| *value >= 0)
}
fn normalized_string(value: Option<&str>) -> Option<String> {
value
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
fn enrich_event_route_metadata(metadata: &mut Value, route: &str) {
let Some(object) = metadata.as_object_mut() else {
return;
};
object
.entry("route".to_string())
.or_insert_with(|| Value::String(route.to_string()));
object
.entry("apiKind".to_string())
.or_insert_with(|| Value::String(RESPONSES_API_KIND.to_string()));
}
fn enrich_usage_record_identity(
metadata: &mut Value,
requested_model: Option<&str>,
upstream_model: &str,
client_ip: Option<&str>,
client_user_agent: Option<&str>,
reasoning_effort: Option<&str>,
service_tier: Option<&str>,
) {
let Some(object) = metadata.as_object_mut() else {
return;
};
let upstream_model = upstream_model.trim();
let requested_model = requested_model
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(upstream_model);
object.insert(
"requestedModel".to_string(),
Value::String(requested_model.to_string()),
);
object.insert(
"upstreamModel".to_string(),
Value::String(upstream_model.to_string()),
);
insert_trimmed_string(object, "clientIp", client_ip);
insert_trimmed_string(object, "userAgent", client_user_agent);
insert_trimmed_string(object, "reasoningEffort", reasoning_effort);
insert_trimmed_string(object, "serviceTier", service_tier);
}
pub(super) fn enrich_response_request_semantics(
metadata: &mut Value,
request: &CodexResponsesRequest,
) {
let Some(object) = metadata.as_object_mut() else {
return;
};
let semantics = request.semantics();
object.insert("compact".to_string(), Value::Bool(semantics.compact));
insert_trimmed_string(object, "requestKind", semantics.request_kind.as_deref());
insert_trimmed_string(object, "subagentKind", semantics.subagent_kind.as_deref());
insert_trimmed_string(object, "reasoningPreset", semantics.reasoning_preset);
}
pub(super) fn reasoning_effort_from_request(request: &CodexResponsesRequest) -> Option<&str> {
request
.reasoning()?
.get("effort")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn insert_trimmed_string(object: &mut Map<String, Value>, key: &str, value: Option<&str>) {
if let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) {
object.insert(key.to_string(), Value::String(value.to_string()));
}
}
fn ensure_stream_metadata_flag(metadata: &mut Value) {
let Some(object) = metadata.as_object_mut() else {
*metadata = serde_json::json!({ "stream": true });
@@ -313,6 +624,11 @@ fn enrich_live_response_stream_metadata(
"firstEventMs",
context.transport_metrics.first_event_ms,
);
insert_optional_i64(
object,
"openaiProcessingMs",
openai_processing_ms(context.response_metadata),
);
if let Some(http_version) = context.transport_metrics.http_version.as_deref() {
object
.entry("httpVersion".to_string())
@@ -505,15 +821,17 @@ fn upstream_failure_client_status_code(upstream_status: i64) -> i64 {
))
}
pub(super) async fn record_live_response_stream_event(
context: &StreamContext<'_>,
status_code: i64,
level: UsageRecordLevel,
message: &str,
mut metadata: Value,
rate_limit_headers: &[(String, String)],
body: &str,
) {
pub(super) async fn record_live_response_stream_event(record: LiveResponseStreamEventRecord<'_>) {
let LiveResponseStreamEventRecord {
context,
status_code,
level,
message,
usage,
mut metadata,
rate_limit_headers,
body,
} = record;
let effective_model = context
.response_metadata
.effective_model
@@ -562,29 +880,28 @@ pub(super) async fn record_live_response_stream_event(
enrich_response_request_semantics(&mut metadata, context.request);
if is_success_usage_event(level, status_code) {
crate::telemetry::recorder::record_response_event(
crate::telemetry::usage::types::ResponseUsageRecord {
recorder: context.recorder,
request_id: context.request_id,
client_api_key_id: context.request.client_api_key_id.as_deref(),
account_id: context.account_id,
route: context.route,
model: effective_model,
requested_model: Some(context.requested_model),
client_ip: context.request.client_ip.as_deref(),
client_user_agent: context.request.client_user_agent.as_deref(),
reasoning_effort: reasoning_effort_from_request(context.request),
service_tier: context.request.service_tier(),
started_at: context.started_at,
status_code,
message,
metadata,
rate_limit_headers,
},
)
record_response_event(ResponseUsageEventRecord {
recorder: context.recorder,
request_id: context.request_id,
client_api_key_id: context.request.client_api_key_id.as_deref(),
account_id: context.account_id,
route: context.route,
model: effective_model,
requested_model: Some(context.requested_model),
client_ip: context.request.client_ip.as_deref(),
client_user_agent: context.request.client_user_agent.as_deref(),
reasoning_effort: reasoning_effort_from_request(context.request),
service_tier: context.request.service_tier(),
started_at: context.started_at,
status_code,
message,
usage,
metadata,
rate_limit_headers,
})
.await;
} else {
let mut event = OpsErrorLog::new(response_event_kind(context.route), message);
let mut event = OpsErrorLog::new(RESPONSE_EVENT_KIND, message);
event.request_id = Some(context.request_id.to_string());
event.client_api_key_id = context.request.client_api_key_id.clone();
event.provider = Some("openai".to_string());
@@ -594,6 +911,7 @@ pub(super) async fn record_live_response_stream_event(
event.status_code = Some(status_code);
event.latency_ms = Some(elapsed_millis_i64(context.started_at));
event.metadata = metadata;
lift_error_fact_fields(&mut event);
if let Err(error) = context.recorder.record_error(event).await {
tracing::error!(account_id = %context.account_id, error = %error, "Failed to record live response stream error event");
}
@@ -683,16 +1001,40 @@ fn metadata_string_field<'a>(metadata: &'a Value, field: &str) -> Option<&'a str
metadata.get(field).and_then(Value::as_str)
}
pub(super) fn insert_first_token_ms(metadata: &mut Value, first_token_ms: Option<i64>) {
let Some(first_token_ms) = first_token_ms else {
pub(super) fn insert_output_timing_ms(
metadata: &mut Value,
first_token_ms: Option<i64>,
first_reasoning_ms: Option<i64>,
first_text_ms: Option<i64>,
) {
if let Some(object) = metadata.as_object_mut() {
insert_optional_i64(object, "firstTokenMs", first_token_ms);
insert_optional_i64(object, "firstReasoningMs", first_reasoning_ms);
insert_optional_i64(object, "firstTextMs", first_text_ms);
}
}
pub(super) fn insert_openai_processing_ms(
metadata: &mut Value,
response_metadata: &CodexResponseMetadata,
) {
let Some(object) = metadata.as_object_mut() else {
return;
};
if let Some(object) = metadata.as_object_mut() {
object.insert(
"firstTokenMs".to_string(),
Value::Number(first_token_ms.into()),
);
}
insert_optional_i64(
object,
"openaiProcessingMs",
openai_processing_ms(response_metadata),
);
}
fn openai_processing_ms(response_metadata: &CodexResponseMetadata) -> Option<i64> {
response_metadata
.client_headers
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("openai-processing-ms"))
.and_then(|(_, value)| value.parse::<i64>().ok())
.filter(|value| *value >= 0)
}
pub(super) fn insert_websocket_pool_decision(
@@ -2,15 +2,20 @@
mod events;
use std::time::{Duration, Instant};
use serde_json::Value;
use crate::{
dispatch::{
controllers::telemetry::events::{
ResponseUpstreamErrorEventRecord, insert_first_token_ms,
insert_response_status_metadata, insert_response_trace_metadata,
insert_response_upstream_diagnostics, insert_websocket_pool_decision,
record_live_response_stream_event, record_response_upstream_error_event,
LiveResponseStreamEventRecord, ResponseUpstreamErrorEventRecord,
ResponseUsageEventRecord, enrich_response_request_semantics,
insert_openai_processing_ms, insert_output_timing_ms, insert_response_status_metadata,
insert_response_trace_metadata, insert_response_upstream_diagnostics,
insert_websocket_pool_decision, reasoning_effort_from_request,
record_live_response_stream_event, record_response_event,
record_response_upstream_error_event,
},
controllers::{
AttemptUpstreamErrorContext, CompleteExit, DispatchErrorObservation,
@@ -27,13 +32,7 @@ use crate::{
},
},
infra::time::elapsed_millis_i64,
telemetry::{
recorder::{
Recorder, enrich_response_request_semantics, reasoning_effort_from_request,
record_response_event,
},
usage::types::{ResponseUsageRecord, UsageRecordLevel},
},
telemetry::{recorder::Recorder, usage::types::UsageRecordLevel},
upstream::openai::{
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
transport::{
@@ -149,13 +148,27 @@ impl TelemetryController {
.effective_model
.as_deref()
.unwrap_or(exit.display_model);
let first_token_ms = request_timing_ms(
exit.started_at,
exit.attempt.started_at(),
exit.response.first_token_ms,
);
let first_reasoning_ms = request_timing_ms(
exit.started_at,
exit.attempt.started_at(),
exit.response.first_reasoning_ms,
);
let first_text_ms = request_timing_ms(
exit.started_at,
exit.attempt.started_at(),
exit.response.first_text_ms,
);
let mut metadata = serde_json::json!({
"responseId": exit.response_id,
"stream": false,
"completed": exit.completed,
"incomplete": !exit.completed,
"transport": backend_transport_name(exit.response.transport),
"firstTokenMs": exit.response.first_token_ms,
"usage": exit.response.usage,
"effectiveModel": effective_model,
"modelsEtag": exit.response.response_metadata.models_etag.as_deref(),
@@ -167,6 +180,13 @@ impl TelemetryController {
"firstEventMs": exit.response.transport_metrics.first_event_ms,
"httpVersion": exit.response.transport_metrics.http_version.as_deref(),
});
insert_output_timing_ms(
&mut metadata,
first_token_ms,
first_reasoning_ms,
first_text_ms,
);
insert_openai_processing_ms(&mut metadata, &exit.response.response_metadata);
insert_response_status_metadata(
&mut metadata,
200,
@@ -177,7 +197,7 @@ impl TelemetryController {
insert_response_trace_metadata(&mut metadata, exit.trace, Some(exit.attempt));
insert_websocket_pool_decision(&mut metadata, exit.response.websocket_pool_decision);
enrich_response_request_semantics(&mut metadata, exit.request);
record_response_event(ResponseUsageRecord {
record_response_event(ResponseUsageEventRecord {
recorder: exit.recorder,
request_id: exit.request_id,
client_api_key_id: exit.request.client_api_key_id.as_deref(),
@@ -196,6 +216,7 @@ impl TelemetryController {
} else {
"v1 responses incomplete"
},
usage: exit.response.usage,
metadata,
rate_limit_headers: &exit.response.rate_limit_headers,
})
@@ -248,22 +269,44 @@ impl TelemetryController {
"firstEventMs": summary.first_event_ms,
"usage": summary.usage,
});
insert_first_token_ms(&mut metadata, summary.first_token_ms);
record_live_response_stream_event(
insert_output_timing_ms(
&mut metadata,
summary.first_token_ms,
summary.first_reasoning_ms,
summary.first_text_ms,
);
record_live_response_stream_event(LiveResponseStreamEventRecord {
context,
status_code,
UsageRecordLevel::Warn,
"v1 responses stream cancelled",
level: UsageRecordLevel::Warn,
message: "v1 responses stream cancelled",
usage: summary.usage,
metadata,
rate_limit_headers,
body,
)
})
.await;
}
}
}
}
fn request_timing_ms(
request_started_at: Instant,
attempt_started_at: Instant,
attempt_first_token_ms: Option<i64>,
) -> Option<i64> {
let attempt_first_token_ms = u64::try_from(attempt_first_token_ms?).ok()?;
let output_at =
attempt_started_at.checked_add(Duration::from_millis(attempt_first_token_ms))?;
i64::try_from(
output_at
.saturating_duration_since(request_started_at)
.as_millis(),
)
.ok()
}
async fn record_success(
context: &StreamContext<'_>,
summary: &StreamSummary,
@@ -272,27 +315,34 @@ async fn record_success(
rate_limit_headers: &[(String, String)],
body: &str,
) {
record_live_response_stream_event(
let mut metadata = serde_json::json!({
"stream": true,
"completed": completed,
"incomplete": !completed,
"responseId": response.get("id").and_then(Value::as_str),
"firstEventMs": summary.first_event_ms,
"usage": summary.usage,
});
insert_output_timing_ms(
&mut metadata,
summary.first_token_ms,
summary.first_reasoning_ms,
summary.first_text_ms,
);
record_live_response_stream_event(LiveResponseStreamEventRecord {
context,
200,
UsageRecordLevel::Info,
if completed {
status_code: 200,
level: UsageRecordLevel::Info,
message: if completed {
"v1 responses stream completed"
} else {
"v1 responses stream incomplete"
},
serde_json::json!({
"stream": true,
"completed": completed,
"incomplete": !completed,
"responseId": response.get("id").and_then(Value::as_str),
"firstEventMs": summary.first_event_ms,
"firstTokenMs": summary.first_token_ms,
"usage": summary.usage,
}),
usage: summary.usage,
metadata,
rate_limit_headers,
body,
)
})
.await;
}
@@ -344,16 +394,22 @@ async fn record_failure(
.map_or(Value::Null, |response_id| {
Value::String(response_id.clone())
});
insert_first_token_ms(&mut metadata, summary.first_token_ms);
insert_output_timing_ms(
&mut metadata,
summary.first_token_ms,
summary.first_reasoning_ms,
summary.first_text_ms,
);
metadata["firstEventMs"] = Value::Number(summary.first_event_ms.into());
record_live_response_stream_event(
record_live_response_stream_event(LiveResponseStreamEventRecord {
context,
status_code,
UsageRecordLevel::Error,
"v1 responses stream failed",
level: UsageRecordLevel::Error,
message: "v1 responses stream failed",
usage: summary.usage,
metadata,
rate_limit_headers,
body,
)
})
.await;
}
+98 -11
View File
@@ -1,12 +1,24 @@
//! 配额响应头与 token 用量写入的唯一 feature owner。
use serde_json::Value;
use crate::{
dispatch::lifecycle::contract::{
AttemptObservation, AttemptObservationKind, CompleteResponseFacts,
},
fleet::pool::AccountPoolService,
fleet::{
pool::AccountPoolService,
quota::{
QuotaCreditsObservation, QuotaLimitObservation, QuotaObservation,
QuotaWindowObservation,
},
usage::ResponseUsage,
},
upstream::openai::{
protocol::{events::TokenUsage, responses::CodexResponsesRequest},
protocol::{
events::{RateLimitWindow, TokenUsage, parse_rate_limit_headers},
responses::CodexResponsesRequest,
},
transport::{CodexRateLimitHeaderUpdates, CodexTurnStateUpdate},
},
};
@@ -71,23 +83,41 @@ impl UsageController {
) {
if let Some(usage) = usage {
account_pool
.record_response_usage(account_id, usage, image_generation_requested)
.record_response_usage(
account_id,
response_usage(usage),
image_generation_requested,
)
.await;
}
}
pub(super) async fn sync_passive_quota(
account_pool: &AccountPoolService,
account_id: &str,
plan_type: Option<&str>,
headers: &[(String, String)],
) {
let Some(observation) = quota_observation(headers) else {
return;
};
account_pool
.sync_passive_quota_observation_for_account(account_id, plan_type, &observation)
.await;
}
pub(super) async fn leave_stream(exit: StreamExit<'_>) -> StreamUsageExit {
let mut rate_limit_headers = exit.rate_limit_headers.to_vec();
if let Some(updates) = exit.rate_limit_header_updates {
rate_limit_headers.extend(updates.lock().await.iter().cloned());
}
exit.account_pool
.sync_passive_rate_limit_headers_for_account(
exit.account_id,
exit.account_plan_type,
&rate_limit_headers,
)
.await;
Self::sync_passive_quota(
exit.account_pool,
exit.account_id,
exit.account_plan_type,
&rate_limit_headers,
)
.await;
let turn_state = if let Some(update) = exit.turn_state_update {
update.lock().await.clone()
} else {
@@ -95,7 +125,7 @@ impl UsageController {
};
if let Some(usage) = exit.usage {
exit.account_pool
.record_token_usage(exit.account_id, &usage)
.record_response_usage(exit.account_id, response_usage(usage), false)
.await;
}
StreamUsageExit {
@@ -104,3 +134,60 @@ impl UsageController {
}
}
}
fn response_usage(usage: TokenUsage) -> ResponseUsage {
ResponseUsage {
input_tokens: usage.input_tokens,
output_tokens: usage.output_tokens,
cached_tokens: usage.cached_tokens,
reasoning_tokens: usage.reasoning_tokens,
total_tokens: usage.total_tokens,
image_input_tokens: usage.image_input_tokens,
image_output_tokens: usage.image_output_tokens,
}
}
fn quota_observation(headers: &[(String, String)]) -> Option<QuotaObservation> {
let parsed = parse_rate_limit_headers(headers)?;
Some(QuotaObservation {
limits: parsed
.limits
.into_iter()
.map(|(limit_id, details)| {
(
limit_id,
QuotaLimitObservation {
limit_id: details.limit_id,
limit_name: details.limit_name,
metered_feature: None,
allowed: details.allowed,
limit_reached: details.limit_reached,
primary: details.primary.map(quota_window_observation),
secondary: details.secondary.map(quota_window_observation),
},
)
})
.collect(),
active_limit: parsed.active_limit,
credits: parsed.credits.map(|credits| QuotaCreditsObservation {
has_credits: credits.has_credits,
unlimited: credits.unlimited,
overage_limit_reached: false,
balance: credits.balance.map(Value::String),
}),
spend_control: None,
plan_type: parsed.plan_type,
promo_message: parsed.promo_message,
rate_limit_reached_type: parsed.rate_limit_reached_type,
})
}
fn quota_window_observation(window: RateLimitWindow) -> QuotaWindowObservation {
QuotaWindowObservation {
used_percent: window.used_percent,
window_minutes: window.window_minutes,
reset_at: window.reset_at,
used: None,
limit: None,
}
}
+9
View File
@@ -238,6 +238,15 @@ pub(crate) fn upstream_error_set_cookie_headers(error: &CodexClientError) -> &[S
}
}
pub(crate) fn upstream_error_rate_limit_headers(error: &CodexClientError) -> Vec<(String, String)> {
match error {
CodexClientError::Upstream {
rate_limit_headers, ..
} => rate_limit_headers.clone(),
_ => Vec::new(),
}
}
pub(crate) fn upstream_error_http_status(error: &CodexClientError) -> u16 {
match error {
CodexClientError::Upstream { status, .. } => status.as_u16(),
+22 -53
View File
@@ -8,9 +8,8 @@ use crate::{
dispatch::{
affinity::AccountIdentityService,
controllers::{
AttemptAccountPreparation, AttemptAccountPreparationContext, AttemptRoutePreparation,
AttemptRoutePreparationContext, AttemptUpstreamErrorContext, ControllerRequestScope,
ControllerSet,
AttemptRoutePreparation, AttemptRoutePreparationContext, AttemptUpstreamErrorContext,
ControllerRequestScope, ControllerSet,
},
errors::{ClientFailure, ResponseDispatchError},
failure::exhaustion::AccountExhaustionTracker,
@@ -50,13 +49,8 @@ use super::{
};
pub(in crate::dispatch) enum AttemptMode {
Complete {
started_at: Instant,
tuple_schema: Option<Value>,
},
Stream {
tuple_schema: Option<Value>,
},
Complete { tuple_schema: Option<Value> },
Stream { tuple_schema: Option<Value> },
}
pub(in crate::dispatch) struct AttemptRunnerDependencies<'a> {
@@ -196,25 +190,11 @@ impl<'a> AttemptRunner<'a> {
.controller_scope
.as_mut()
.ok_or(AttemptContractError::Terminal)?;
let prepared_account = match self
let prepared_account = self
.dependencies
.controllers
.prepare_attempt_account(
AttemptAccountPreparationContext {
codex: self.dependencies.codex,
account_identity: self.dependencies.account_identity,
request_id: self.dependencies.request_id,
},
acquired,
)
.await
{
AttemptAccountPreparation::Ready(prepared) => *prepared,
AttemptAccountPreparation::Rejected => {
self.pending = Some(PendingAttempt::CandidatePreparationRejected { account });
return self.step();
}
};
.prepare_attempt_account(acquired)
.await;
let acquired = prepared_account.lease;
let cookie_header = prepared_account.cookie_header;
@@ -288,15 +268,12 @@ impl<'a> AttemptRunner<'a> {
};
match &self.mode {
AttemptMode::Complete {
started_at,
tuple_schema,
} => {
AttemptMode::Complete { tuple_schema } => {
let result = create_response_with_account(
upstream_context,
&account_scoped_request,
prepared_transport,
*started_at,
attempt.started_at(),
)
.await;
let attempt_request = account_scoped_request.into_request();
@@ -420,7 +397,7 @@ impl<'a> AttemptRunner<'a> {
body,
&mut decoder,
attempt_request.stream_commit_policy,
self.dependencies.started_at,
attempt.started_at(),
)
.await
{
@@ -658,6 +635,7 @@ impl<'a> AttemptRunner<'a> {
PendingAttempt::StreamResponse(stream) => {
let PendingStreamResponse {
account,
attempt_request,
attempt,
lease,
response,
@@ -668,6 +646,7 @@ impl<'a> AttemptRunner<'a> {
} = *stream;
EstablishedResponse::Stream(Box::new(EstablishedStream {
context: self.take_established_context(account, attempt)?,
attempt_request,
lease,
response,
decoder,
@@ -775,9 +754,6 @@ fn observation_kind(
kind: *kind,
}
}
PendingAttempt::CandidatePreparationRejected { .. } => {
AttemptObservationKind::CandidatePreparationRejected
}
PendingAttempt::RoutePreparationRejected { message, .. } => {
AttemptObservationKind::RoutePreparationRejected {
message: message.clone(),
@@ -846,17 +822,15 @@ async fn rejection_error(
})
}
AttemptReturnKind::Observed => match pending {
PendingAttempt::NoCandidate | PendingAttempt::CandidatePreparationRejected { .. } => {
Ok(on_exhaustion.map_or_else(
|| {
exhaustion
.last_exhausted()
.map(ResponseDispatchError::from_exhausted_account)
.unwrap_or(ResponseDispatchError::NoActiveAccount)
},
ResponseDispatchError::Failed,
))
}
PendingAttempt::NoCandidate => Ok(on_exhaustion.map_or_else(
|| {
exhaustion
.last_exhausted()
.map(ResponseDispatchError::from_exhausted_account)
.unwrap_or(ResponseDispatchError::NoActiveAccount)
},
ResponseDispatchError::Failed,
)),
PendingAttempt::RoutePreparationRejected { message, .. } => {
Ok(ResponseDispatchError::HistoryUnavailable {
upstream_error: message,
@@ -908,8 +882,7 @@ fn pending_account(pending: &PendingAttempt) -> Option<&str> {
fn pending_account_value(pending: &PendingAttempt) -> Option<&Account> {
match pending {
PendingAttempt::NoCandidate | PendingAttempt::PinnedCandidateUnavailable { .. } => None,
PendingAttempt::CandidatePreparationRejected { account }
| PendingAttempt::RoutePreparationRejected { account, .. }
PendingAttempt::RoutePreparationRejected { account, .. }
| PendingAttempt::UpstreamFailure { account, .. }
| PendingAttempt::ProtocolFailure { account, .. } => Some(account),
PendingAttempt::CompleteResponse(complete) => Some(&complete.account),
@@ -925,7 +898,6 @@ fn pending_attempt(pending: &PendingAttempt) -> Option<&super::trace::ResponseDi
PendingAttempt::StreamResponse(stream) => Some(&stream.attempt),
PendingAttempt::NoCandidate
| PendingAttempt::PinnedCandidateUnavailable { .. }
| PendingAttempt::CandidatePreparationRejected { .. }
| PendingAttempt::RoutePreparationRejected { .. } => None,
}
}
@@ -942,7 +914,6 @@ fn pending_attempt_request(pending: &PendingAttempt) -> Option<&CodexResponsesRe
PendingAttempt::StreamResponse(stream) => Some(&stream.attempt_request),
PendingAttempt::NoCandidate
| PendingAttempt::PinnedCandidateUnavailable { .. }
| PendingAttempt::CandidatePreparationRejected { .. }
| PendingAttempt::RoutePreparationRejected { .. } => None,
}
}
@@ -957,7 +928,6 @@ fn pending_transport(
PendingAttempt::StreamResponse(stream) => Some(stream.response.transport),
PendingAttempt::NoCandidate
| PendingAttempt::PinnedCandidateUnavailable { .. }
| PendingAttempt::CandidatePreparationRejected { .. }
| PendingAttempt::RoutePreparationRejected { .. } => None,
}
}
@@ -979,7 +949,6 @@ fn invalid(decision: &'static str, pending: &PendingAttempt) -> AttemptContractE
observation: match pending {
PendingAttempt::NoCandidate => "no_candidate",
PendingAttempt::PinnedCandidateUnavailable { .. } => "pinned_candidate_unavailable",
PendingAttempt::CandidatePreparationRejected { .. } => "candidate_preparation_rejected",
PendingAttempt::RoutePreparationRejected { .. } => "route_preparation_rejected",
PendingAttempt::UpstreamFailure { .. } => "upstream_failure",
PendingAttempt::ProtocolFailure { .. } => "protocol_failure",
+1 -4
View File
@@ -93,7 +93,6 @@ pub(in crate::dispatch) enum AttemptObservationKind {
ledger: CandidateLedgerFacts,
last_exhausted: Option<ExhaustedAccount>,
},
CandidatePreparationRejected,
RoutePreparationRejected {
message: String,
},
@@ -173,6 +172,7 @@ impl TryFrom<CollectedResponse> for EstablishedCompleteBody {
pub(in crate::dispatch) struct EstablishedStream {
pub context: EstablishedAttemptContext,
pub attempt_request: CodexResponsesRequest,
pub lease: AccountLease,
pub response: CodexBackendStreamingResponse,
pub decoder: CanonicalStreamDecoder,
@@ -256,9 +256,6 @@ pub(super) enum PendingAttempt {
account_id: String,
kind: PinnedCandidateAcquireFailureKind,
},
CandidatePreparationRejected {
account: Account,
},
RoutePreparationRejected {
account: Account,
message: String,
@@ -27,6 +27,8 @@ pub(in crate::dispatch) struct StreamFinalizer {
controllers: ControllerSet,
controller_scope: ControllerRequestScope,
controller_context: StreamControllerContext,
request_input: Vec<serde_json::Value>,
continued_from_previous_response: bool,
account_lease: AccountLease,
}
@@ -35,12 +37,16 @@ impl StreamFinalizer {
controllers: ControllerSet,
controller_scope: ControllerRequestScope,
controller_context: StreamControllerContext,
request_input: Vec<serde_json::Value>,
continued_from_previous_response: bool,
account_lease: AccountLease,
) -> Self {
Self {
controllers,
controller_scope,
controller_context,
request_input,
continued_from_previous_response,
account_lease,
}
}
@@ -53,6 +59,18 @@ impl StreamFinalizer {
self.controller_context.started_at
}
pub(in crate::dispatch) fn account_id(&self) -> &str {
&self.controller_context.account_id
}
pub(in crate::dispatch) fn request_input(&self) -> &[serde_json::Value] {
&self.request_input
}
pub(in crate::dispatch) fn continued_from_previous_response(&self) -> bool {
self.continued_from_previous_response
}
pub(in crate::dispatch) async fn finalize(
self,
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
@@ -64,6 +82,7 @@ impl StreamFinalizer {
controller_scope,
controller_context,
account_lease,
..
} = self;
let body = String::from_utf8_lossy(&summary.body).into_owned();
controllers
+1 -1
View File
@@ -60,7 +60,7 @@ pub(in crate::dispatch) async fn enter_request(
let catalog = dependencies.models.catalog().await;
let display_model = catalog.resolve_model_id(requested_model);
request.set_model(display_model.clone());
request.set_stream(mode.stream());
debug_assert_eq!(request.stream(), mode.stream());
let compact = request.semantics().compact;
let tuple_schema = request.tuple_schema.clone();
+2
View File
@@ -28,6 +28,8 @@ pub(in crate::dispatch) struct StreamSummary {
pub body: Vec<u8>,
pub terminal_chunks: Vec<CanonicalResponseChunk>,
pub first_token_ms: Option<i64>,
pub first_reasoning_ms: Option<i64>,
pub first_text_ms: Option<i64>,
pub first_event_ms: i64,
pub usage: Option<TokenUsage>,
pub last_response_id: Option<String>,
+8
View File
@@ -1,5 +1,7 @@
//! Responses 请求级账号尝试轨迹。
use std::time::Instant;
#[derive(Debug, Default)]
pub(in crate::dispatch) struct ResponseDispatchTrace {
next_attempt_index: i64,
@@ -10,6 +12,7 @@ pub(in crate::dispatch) struct ResponseDispatchTrace {
pub(in crate::dispatch) struct ResponseDispatchAttempt {
index: i64,
account_id: String,
started_at: Instant,
}
impl ResponseDispatchTrace {
@@ -20,6 +23,7 @@ impl ResponseDispatchTrace {
let attempt = ResponseDispatchAttempt {
index: self.next_attempt_index,
account_id: account_id.to_string(),
started_at: Instant::now(),
};
self.next_attempt_index += 1;
self.attempts.push(attempt.clone());
@@ -39,4 +43,8 @@ impl ResponseDispatchAttempt {
pub(in crate::dispatch) fn account_id(&self) -> &str {
&self.account_id
}
pub(in crate::dispatch) fn started_at(&self) -> Instant {
self.started_at
}
}
+20 -5
View File
@@ -79,6 +79,9 @@ pub type ResponseBodyStream =
/// Responses live SSE 调度结果。
pub struct ResponseDispatchStream {
pub(crate) account_id: String,
pub(crate) request_input: Vec<Value>,
pub(crate) continued_from_previous_response: bool,
pub body: ResponseBodyStream,
pub(crate) canonical_events: tokio::sync::mpsc::UnboundedReceiver<Vec<CanonicalResponseEvent>>,
pub response_headers: Vec<(String, String)>,
@@ -102,13 +105,25 @@ pub(crate) struct ConnectionReplayPlan {
/// API 从 canonical 响应事件采集的原始 transcript 事实。
pub(crate) struct ConnectionTranscriptFacts {
account_id: String,
request_input: Vec<Value>,
continued_from_previous_response: bool,
response_id: String,
output: Vec<Value>,
}
impl ConnectionTranscriptFacts {
pub(crate) fn new(response_id: String, output: Vec<Value>) -> Self {
pub(crate) fn new(
account_id: String,
request_input: Vec<Value>,
continued_from_previous_response: bool,
response_id: String,
output: Vec<Value>,
) -> Self {
Self {
account_id,
request_input,
continued_from_previous_response,
response_id,
output,
}
@@ -168,6 +183,9 @@ impl ResponseDispatchService {
plan.inner,
transcript.response_id,
transcript.output,
transcript.account_id,
transcript.request_input,
transcript.continued_from_previous_response,
);
}
@@ -220,10 +238,7 @@ impl ResponseDispatchService {
requested_model,
started_at,
},
AttemptMode::Complete {
started_at,
tuple_schema,
},
AttemptMode::Complete { tuple_schema },
request,
controller_scope,
candidates,
+5
View File
@@ -141,6 +141,7 @@ impl ResponseDispatchService {
let EstablishedStream {
context,
attempt_request,
lease,
response,
decoder,
@@ -159,6 +160,8 @@ impl ResponseDispatchService {
.observe_models_etag(response.response_metadata.models_etag.as_deref());
let attempts = trace.attempts().to_vec();
let transport = response.transport;
let request_input = attempt_request.input().to_vec();
let continued_from_previous_response = attempt_request.previous_response_id().is_some();
let finalizer = StreamFinalizer::new(
controllers,
controller_scope,
@@ -185,6 +188,8 @@ impl ResponseDispatchService {
attempts,
started_at,
},
request_input,
continued_from_previous_response,
lease,
);
Ok(spawn_live_response_stream(
+54 -9
View File
@@ -57,6 +57,9 @@ pub(in crate::dispatch) fn spawn_live_response_stream(
body: CodexBackendSseStream,
shutdown: CancellationToken,
) -> ResponseDispatchStream {
let account_id = finalizer.account_id().to_string();
let request_input = finalizer.request_input().to_vec();
let continued_from_previous_response = finalizer.continued_from_previous_response();
let response_headers = finalizer.response_headers().to_vec();
let started_at = finalizer.started_at();
let (sender, receiver) = mpsc::channel(8);
@@ -81,6 +84,9 @@ pub(in crate::dispatch) fn spawn_live_response_stream(
});
ResponseDispatchStream {
account_id,
request_input,
continued_from_previous_response,
body: Box::pin(MpscResponseBodyStream {
receiver,
cancel: Some(cancel_sender),
@@ -117,6 +123,8 @@ async fn run_live_response_stream(
let mut body_bytes = Vec::new();
let mut terminal_chunks = Vec::new();
let mut first_token_ms = None;
let mut first_reasoning_ms = None;
let mut first_text_ms = None;
let decoded_terminal = take_stream_terminal(&mut decoder);
let terminal = if decoded_terminal.is_some() {
@@ -128,13 +136,19 @@ async fn run_live_response_stream(
.await
.err()
};
update_first_token_ms(&decoder, started_at, &mut first_token_ms);
update_output_timing_ms(
&decoder,
started_at,
&mut first_token_ms,
&mut first_reasoning_ms,
&mut first_text_ms,
);
if let Some(terminal) = terminal {
return stream_summary(
terminal,
body_bytes,
terminal_chunks,
first_token_ms,
(first_token_ms, first_reasoning_ms, first_text_ms),
first_event_ms,
&decoder,
);
@@ -150,7 +164,13 @@ async fn run_live_response_stream(
match next {
Some(Ok(chunk)) => match decoder.push(chunk) {
Ok(batch) => {
update_first_token_ms(&decoder, started_at, &mut first_token_ms);
update_output_timing_ms(
&decoder,
started_at,
&mut first_token_ms,
&mut first_reasoning_ms,
&mut first_text_ms,
);
if let Some(terminal) = take_stream_terminal(&mut decoder) {
match stage_terminal_batch(&mut body_bytes, &mut terminal_chunks, batch) {
Ok(()) => break terminal,
@@ -196,12 +216,18 @@ async fn run_live_response_stream(
}
};
update_first_token_ms(&decoder, started_at, &mut first_token_ms);
update_output_timing_ms(
&decoder,
started_at,
&mut first_token_ms,
&mut first_reasoning_ms,
&mut first_text_ms,
);
stream_summary(
terminal,
body_bytes,
terminal_chunks,
first_token_ms,
(first_token_ms, first_reasoning_ms, first_text_ms),
first_event_ms,
&decoder,
)
@@ -299,13 +325,29 @@ fn ensure_capture_capacity(
Ok(())
}
fn update_first_token_ms(
fn update_output_timing_ms(
decoder: &CanonicalStreamDecoder,
started_at: Instant,
first_token_ms: &mut Option<i64>,
first_reasoning_ms: &mut Option<i64>,
first_text_ms: &mut Option<i64>,
) {
if first_token_ms.is_none() && decoder.first_output_seen() {
*first_token_ms = Some(elapsed_millis_i64(started_at).max(1));
let first_token_missing = first_token_ms.is_none() && decoder.first_semantic_output_seen();
let first_reasoning_missing =
first_reasoning_ms.is_none() && decoder.first_reasoning_output_seen();
let first_text_missing = first_text_ms.is_none() && decoder.first_text_output_seen();
if !(first_token_missing || first_reasoning_missing || first_text_missing) {
return;
}
let elapsed_ms = elapsed_millis_i64(started_at).max(1);
if first_token_missing {
*first_token_ms = Some(elapsed_ms);
}
if first_reasoning_missing {
*first_reasoning_ms = Some(elapsed_ms);
}
if first_text_missing {
*first_text_ms = Some(elapsed_ms);
}
}
@@ -313,15 +355,18 @@ fn stream_summary(
terminal: StreamTerminal,
body: Vec<u8>,
terminal_chunks: Vec<CanonicalResponseChunk>,
first_token_ms: Option<i64>,
output_timing: (Option<i64>, Option<i64>, Option<i64>),
first_event_ms: i64,
decoder: &CanonicalStreamDecoder,
) -> StreamSummary {
let (first_token_ms, first_reasoning_ms, first_text_ms) = output_timing;
StreamSummary {
terminal,
body,
terminal_chunks,
first_token_ms,
first_reasoning_ms,
first_text_ms,
first_event_ms,
usage: decoder.usage(),
last_response_id: decoder.last_response_id().map(ToString::to_string),
+26 -37
View File
@@ -6,8 +6,8 @@ use serde_json::Value;
use crate::upstream::openai::protocol::{
events::{TokenUsage, extract_usage},
responses::{
CollectedResponse, ResponsesSseFailure, StreamCommitPolicy,
reconvert_responses_sse_event_tuple_values,
CollectedResponse, ResponseEventSignals, ResponsesSseFailure, StreamCommitPolicy,
reconvert_responses_sse_event_tuple_values, response_event_signals,
},
sse::{
MAX_SSE_EVENT_BUFFER_BYTES, SseError, encode_sse_event, parse_sse_events, sse_frame_end,
@@ -123,34 +123,8 @@ impl CanonicalResponseEvent {
.or_else(|| extract_usage(&self.data))
}
pub(in crate::dispatch) fn has_first_output(&self) -> bool {
match self.event_type() {
Some(
"response.output_text.delta"
| "response.reasoning_summary_text.delta"
| "response.reasoning_text.delta"
| "response.function_call_arguments.delta"
| "response.custom_tool_call_input.delta",
) => self
.data
.get("delta")
.and_then(Value::as_str)
.is_some_and(|value| !value.is_empty()),
Some("response.output_text.done") => self
.data
.get("text")
.and_then(Value::as_str)
.is_some_and(|value| !value.is_empty()),
Some("response.function_call_arguments.done") => self
.data
.get("arguments")
.and_then(Value::as_str)
.is_some_and(|value| !value.is_empty()),
Some("response.output_item.done") => {
self.data.get("item").is_some_and(Value::is_object)
}
_ => false,
}
pub(in crate::dispatch) fn signals(&self) -> ResponseEventSignals {
response_event_signals(self.event_type(), &self.data)
}
pub(in crate::dispatch) fn terminal(&self) -> Option<CanonicalStreamTerminal> {
@@ -230,7 +204,7 @@ impl CanonicalStreamBatch {
/// 同一上游 chunk 可能同时包含输出与终态,事件顺序不能被网络分块抹掉。
pub(in crate::dispatch) fn first_failure(&self) -> Option<ResponsesSseFailure> {
for event in self.chunks.iter().flat_map(|chunk| chunk.events()) {
if event.has_first_output() {
if event.signals().semantic_output {
return None;
}
if let Some(terminal) = event.terminal() {
@@ -250,7 +224,9 @@ pub(in crate::dispatch) struct CanonicalStreamDecoder {
tuple_schema: Option<Value>,
pending: Vec<u8>,
forwardable_frame_seen: bool,
first_output_seen: bool,
first_semantic_output_seen: bool,
first_reasoning_output_seen: bool,
first_text_output_seen: bool,
terminal: Option<CanonicalStreamTerminal>,
transport_done_seen: bool,
completed_usage: Option<TokenUsage>,
@@ -264,7 +240,9 @@ impl CanonicalStreamDecoder {
tuple_schema,
pending: Vec::new(),
forwardable_frame_seen: false,
first_output_seen: false,
first_semantic_output_seen: false,
first_reasoning_output_seen: false,
first_text_output_seen: false,
terminal: None,
transport_done_seen: false,
completed_usage: None,
@@ -333,15 +311,23 @@ impl CanonicalStreamDecoder {
self.transport_done_seen
}
pub(in crate::dispatch) fn first_output_seen(&self) -> bool {
self.first_output_seen
pub(in crate::dispatch) fn first_semantic_output_seen(&self) -> bool {
self.first_semantic_output_seen
}
pub(in crate::dispatch) fn first_reasoning_output_seen(&self) -> bool {
self.first_reasoning_output_seen
}
pub(in crate::dispatch) fn first_text_output_seen(&self) -> bool {
self.first_text_output_seen
}
pub(in crate::dispatch) fn commit_boundary_reached(&self, policy: StreamCommitPolicy) -> bool {
match policy {
StreamCommitPolicy::FirstForwardableEvent => self.forwardable_frame_seen,
StreamCommitPolicy::UntilOutputOrTerminal => {
self.first_output_seen || self.terminal.is_some()
self.first_semantic_output_seen || self.terminal.is_some()
}
}
}
@@ -406,7 +392,10 @@ impl CanonicalStreamDecoder {
}
fn observe(&mut self, event: &CanonicalResponseEvent) {
self.first_output_seen |= event.has_first_output();
let signals = event.signals();
self.first_semantic_output_seen |= signals.semantic_output;
self.first_reasoning_output_seen |= signals.reasoning_output;
self.first_text_output_seen |= signals.text_output;
if let Some(response_id) = event.response_id() {
self.last_response_id = Some(response_id.to_string());
}
+21 -54
View File
@@ -2,13 +2,10 @@
use chrono::{Duration, Utc};
use crate::{
fleet::{account::AccountStatus, pool::AccountPoolService},
upstream::openai::{
failure::{UpstreamFailureFacts, upstream_failure_facts},
protocol::responses::ResponsesSseFailure,
transport::{CodexBackendClient, CodexClientError},
},
use crate::fleet::{
account::AccountStatus,
account_gateway::{AccountFailureObservation, AccountUpstreamGateway},
pool::AccountPoolService,
};
const DEFAULT_RATE_LIMIT_RETRY_SECONDS: u64 = 60;
@@ -44,37 +41,24 @@ impl ClassifiedAccountFailure {
}
}
pub fn classify_client_failure(error: &CodexClientError) -> Option<ClassifiedAccountFailure> {
classify_upstream_failure(&upstream_failure_facts(error))
}
pub fn classify_upstream_failure(facts: &UpstreamFailureFacts) -> Option<ClassifiedAccountFailure> {
pub fn classify_account_failure(
facts: &AccountFailureObservation,
) -> Option<ClassifiedAccountFailure> {
classify_failure(FailureFields {
status_code: facts.status_code,
code: facts.code.as_deref(),
error_type: facts.error_type.as_deref(),
identity_authorization_error: facts.identity_authorization_error.as_deref(),
identity_error_code: facts.identity_error_code.as_deref(),
message: &facts.message,
body: &facts.body,
retry_after_seconds: facts.retry_after_seconds,
})
}
pub fn classify_response_failure(
failure: &ResponsesSseFailure,
) -> Option<ClassifiedAccountFailure> {
classify_failure(FailureFields {
status_code: failure.explicit_status_code,
code: failure.upstream_code.as_deref(),
error_type: failure.upstream_type.as_deref(),
message: &failure.message,
body: &failure.message,
retry_after_seconds: failure.retry_after_seconds,
})
}
pub async fn apply_account_state_effect_immediately(
account_pool: &AccountPoolService,
codex: &CodexBackendClient,
upstream: &dyn AccountUpstreamGateway,
account_id: &str,
effect: &AccountStateEffect,
) {
@@ -90,12 +74,12 @@ pub async fn apply_account_state_effect_immediately(
.await;
}
}
codex.evict_websocket_account(account_id).await;
upstream.evict_account_connections(account_id).await;
}
pub async fn apply_account_state_effect(
account_pool: &AccountPoolService,
codex: &CodexBackendClient,
upstream: &dyn AccountUpstreamGateway,
account_id: &str,
effect: &AccountStateEffect,
) -> bool {
@@ -107,27 +91,16 @@ pub async fn apply_account_state_effect(
.await
}
};
codex.evict_websocket_account(account_id).await;
upstream.evict_account_connections(account_id).await;
updated
}
pub async fn apply_client_failure_effect(
account_pool: &AccountPoolService,
codex: &CodexBackendClient,
account_id: &str,
error: &CodexClientError,
) -> Option<ClassifiedAccountFailure> {
let classified = classify_client_failure(error)?;
if let Some(effect) = &classified.effect {
apply_account_state_effect(account_pool, codex, account_id, effect).await;
}
Some(classified)
}
struct FailureFields<'a> {
status_code: Option<u16>,
code: Option<&'a str>,
error_type: Option<&'a str>,
identity_authorization_error: Option<&'a str>,
identity_error_code: Option<&'a str>,
message: &'a str,
body: &'a str,
retry_after_seconds: Option<u64>,
@@ -136,6 +109,8 @@ struct FailureFields<'a> {
fn classify_failure(fields: FailureFields<'_>) -> Option<ClassifiedAccountFailure> {
let code = fields.code.unwrap_or_default();
let error_type = fields.error_type.unwrap_or_default();
let identity_error_code = fields.identity_error_code.unwrap_or_default();
let identity_authorization_error = fields.identity_authorization_error.unwrap_or_default();
if [code, fields.message, fields.body]
.into_iter()
@@ -143,14 +118,13 @@ fn classify_failure(fields: FailureFields<'_>) -> Option<ClassifiedAccountFailur
{
return Some(classified(AccountFailureKind::ModelUnsupported, None));
}
if let Some(status) = explicit_account_status(code, error_type, fields.message)
.or_else(|| explicit_account_status(code, error_type, fields.body))
if let Some(status) =
explicit_account_status(identity_error_code, "", identity_authorization_error)
.or_else(|| explicit_account_status(code, error_type, fields.message))
.or_else(|| explicit_account_status(code, error_type, fields.body))
{
return Some(status_failure(status));
}
if fields.status_code == Some(403) && !is_html(fields.body) {
return Some(status_failure(AccountStatus::Banned));
}
if fields.status_code == Some(401) {
return Some(status_failure(AccountStatus::Expired));
}
@@ -238,8 +212,6 @@ fn explicit_account_status(code: &str, error_type: &str, message: &str) -> Optio
|| message.contains("organization has been disabled")
|| message.contains("workspace has been deactivated")
|| message.contains("deactivated_workspace")
|| code == "forbidden"
|| error_type == "permission_error"
{
return Some(AccountStatus::Banned);
}
@@ -310,8 +282,3 @@ fn is_model_unsupported(value: &str) -> bool {
|| value.contains("not_supported")
|| value.contains("not_available")))
}
fn is_html(value: &str) -> bool {
let value = value.trim_start().to_ascii_lowercase();
value.starts_with("<!doctype") || value.starts_with("<html") || value.contains("<html")
}
+98
View File
@@ -0,0 +1,98 @@
//! 账号管理与配额刷新用例拥有的 OpenAI IO 端口。
use std::pin::Pin;
use async_trait::async_trait;
use futures::Stream;
use secrecy::SecretString;
use serde_json::Value;
use crate::fleet::quota::QuotaSnapshot;
#[derive(Debug, Clone)]
pub struct AccountUpstreamContext {
pub access_token: SecretString,
pub account_id: Option<String>,
pub request_id: String,
pub cookie_header: Option<String>,
pub installation_id: Option<String>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct AccountFailureObservation {
pub status_code: Option<u16>,
pub code: Option<String>,
pub error_type: Option<String>,
pub identity_authorization_error: Option<String>,
pub identity_error_code: Option<String>,
pub message: String,
pub body: String,
pub retry_after_seconds: Option<u64>,
}
#[derive(Debug, Clone, thiserror::Error)]
#[error("{message}")]
pub struct AccountGatewayError {
message: String,
failure: Option<AccountFailureObservation>,
}
impl AccountGatewayError {
pub fn new(message: impl Into<String>, failure: Option<AccountFailureObservation>) -> Self {
Self {
message: message.into(),
failure,
}
}
pub fn failure(&self) -> Option<&AccountFailureObservation> {
self.failure.as_ref()
}
}
#[derive(Debug, Clone)]
pub struct AccountUsageResult {
pub quota: QuotaSnapshot,
pub raw: Value,
pub account_id: Option<String>,
pub user_id: Option<String>,
pub email: Option<String>,
pub plan_type: Option<String>,
}
#[derive(Debug, Clone)]
pub struct AccountProbeRequest {
pub model: String,
pub instructions: String,
pub input_text: String,
}
#[derive(Debug, Clone)]
pub enum AccountProbeEvent {
Content(String),
Complete,
Failed(AccountGatewayError),
}
pub type AccountProbeEventStream = Pin<Box<dyn Stream<Item = AccountProbeEvent> + Send>>;
pub struct AccountProbeSession {
pub request_payload: Value,
pub events: AccountProbeEventStream,
}
#[async_trait]
pub trait AccountUpstreamGateway: Send + Sync + 'static {
async fn fetch_usage(
&self,
context: AccountUpstreamContext,
) -> Result<AccountUsageResult, AccountGatewayError>;
async fn probe_response(
&self,
context: AccountUpstreamContext,
request: AccountProbeRequest,
) -> Result<AccountProbeSession, AccountGatewayError>;
async fn evict_account_connections(&self, account_id: &str);
}
+18 -31
View File
@@ -1,12 +1,14 @@
use chrono::Utc;
use secrecy::SecretString;
use crate::fleet::account_gateway::AccountUpstreamContext;
use super::{
AccountManageService,
types::{
AccountManageError, ImportSupplementalAccountInfo, ImportSupplementalNeeds,
ImportedAccountState, ImportedAccounts, ResolvedImportTokens, import_quota_plan_type,
import_status_from_usage_error, import_usage_plan_type, import_usage_string,
import_status_from_usage_error,
},
};
@@ -190,41 +192,26 @@ impl AccountManageService {
return ImportSupplementalAccountInfo::default();
}
let request_id = uuid::Uuid::new_v4().to_string();
let installation_id =
account_id.map(|account_id| self.account_pseudonymizer.installation_id(account_id));
let context = crate::upstream::openai::transport::CodexRequestContext {
access_token,
account_id,
request_id: &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,
let context = AccountUpstreamContext {
access_token: SecretString::new(access_token.to_string().into()),
account_id: account_id.map(ToString::to_string),
request_id: uuid::Uuid::new_v4().to_string(),
cookie_header: None,
installation_id: installation_id.as_deref(),
session_id: None,
thread_id: None,
client_request_id: None,
turn_id: None,
installation_id,
};
match self.codex.fetch_usage(context).await {
Ok(raw) => {
let normalized = crate::fleet::quota::quota_from_usage(&raw);
ImportSupplementalAccountInfo {
account_id: import_usage_string(&raw, "account_id"),
user_id: import_usage_string(&raw, "user_id"),
email: import_usage_string(&raw, "email"),
plan_type: import_usage_plan_type(&raw),
quota_json: serde_json::to_string(&normalized).ok(),
quota_fetched_at: Some(Utc::now()),
status: None,
}
}
match self.upstream.fetch_usage(context).await {
Ok(result) => ImportSupplementalAccountInfo {
account_id: result.account_id,
user_id: result.user_id,
email: result.email,
plan_type: result.plan_type,
quota_json: result.quota.to_json().ok(),
quota_fetched_at: Some(Utc::now()),
status: None,
},
Err(error) => {
tracing::warn!(
error = %error,
+1 -1
View File
@@ -10,12 +10,12 @@ use tokio::time::sleep;
use uuid::Uuid;
use crate::{
fleet::refresh::RefreshFailure,
fleet::{
account::{AccountListSort, AccountStatus},
store::AccountStore,
},
infra::{json::NumberedPage, time::elapsed_millis_i64},
upstream::openai::token_client::RefreshFailure,
};
use super::{
+1 -1
View File
@@ -6,9 +6,9 @@ mod lifecycle;
pub(crate) mod oauth;
mod probe;
mod quota;
pub mod quota_view;
mod service;
mod types;
pub use probe::{AccountModelOption, AccountTestEvent, AccountTestStream};
pub use service::*;
pub use types::*;
+148 -223
View File
@@ -1,66 +1,86 @@
use std::{convert::Infallible, pin::Pin};
use std::pin::Pin;
use bytes::Bytes;
use futures::{Stream, StreamExt};
use secrecy::ExposeSecret;
use serde_json::{Value, json};
use serde_json::Value;
use tokio::sync::mpsc;
use crate::{
fleet::{
account::AccountStatus,
account_failure::{classify_client_failure, classify_response_failure},
account_failure::classify_account_failure,
account_gateway::{
AccountGatewayError, AccountProbeEvent, AccountProbeRequest, AccountProbeSession,
AccountUpstreamContext,
},
store::StoredAccount,
},
models::{service::ModelRefreshPlanAccount, types::CodexModelInfo},
upstream::openai::{
protocol::{
responses::{CodexResponsesRequest, ResponsesSseFailure},
sse::{SseEvent, encode_sse_event, parse_sse_events},
},
transport::{CodexClientError, CodexRequestContext},
models::{
service::{ModelRefreshPlanAccount, ModelRefreshResult},
types::CodexModelInfo,
},
};
use super::{AccountManageService, types::AccountManageError};
pub(super) type AccountTestStream = Pin<Box<dyn Stream<Item = Result<Bytes, Infallible>> + Send>>;
pub type AccountTestStream = Pin<Box<dyn Stream<Item = AccountTestEvent> + Send>>;
#[derive(Debug, Clone, serde::Serialize)]
#[serde(rename_all = "camelCase")]
#[derive(Debug, Clone)]
pub enum AccountTestEvent {
Started {
model: String,
},
Request {
payload: Value,
},
Content {
text: String,
},
Complete {
account_status: Option<AccountStatus>,
},
Error {
error: String,
account_status: Option<AccountStatus>,
},
}
#[derive(Debug, Clone)]
pub struct AccountModelOption {
pub id: String,
pub label: String,
}
struct AccountTestOutcome {
event: Value,
error: Option<String>,
status: Option<AccountStatus>,
}
impl AccountTestOutcome {
fn success() -> Self {
Self {
event: json!({ "type": "test_complete", "success": true }),
error: None,
status: Some(AccountStatus::Active),
}
}
fn error(error: impl Into<String>, status: Option<AccountStatus>) -> Self {
Self {
event: json!({ "type": "error", "error": error.into() }),
error: Some(error.into()),
status,
}
}
fn attach_account_status(&mut self, status: AccountStatus) {
let Some(event) = self.event.as_object_mut() else {
return;
};
event.insert(
"accountStatus".to_string(),
Value::String(status.as_str().to_string()),
);
fn into_event(self) -> AccountTestEvent {
match self.error {
Some(error) => AccountTestEvent::Error {
error,
account_status: self.status,
},
None => AccountTestEvent::Complete {
account_status: self.status,
},
}
}
}
@@ -76,11 +96,7 @@ impl AccountManageService {
.map_err(|_| AccountManageError::Inspect)?
.ok_or(AccountManageError::NotFound)?;
let plan_type = account_plan_type(&account);
let mut models = self.models.catalog().await.models_for_plan(&plan_type);
if models.is_empty() {
self.refresh_account_plan_models(&account, &plan_type).await;
models = self.models.catalog().await.models_for_plan(&plan_type);
}
let models = self.models.catalog().await.models_for_plan(&plan_type);
let models = models.iter().map(account_model_option).collect::<Vec<_>>();
if models.is_empty() {
return Err(AccountManageError::NoModels);
@@ -88,7 +104,32 @@ impl AccountManageService {
Ok(models)
}
async fn refresh_account_plan_models(&self, account: &StoredAccount, plan_type: &str) {
pub async fn refresh_account_models(
&self,
account_id: &str,
) -> Result<Vec<AccountModelOption>, AccountManageError> {
let account = self
.store
.get(account_id)
.await
.map_err(|_| AccountManageError::Inspect)?
.ok_or(AccountManageError::NotFound)?;
let plan_type = account_plan_type(&account);
self.refresh_account_plan_models(&account, &plan_type)
.await?;
let models = self.models.catalog().await.models_for_plan(&plan_type);
let models = models.iter().map(account_model_option).collect::<Vec<_>>();
if models.is_empty() {
return Err(AccountManageError::NoModels);
}
Ok(models)
}
async fn refresh_account_plan_models(
&self,
account: &StoredAccount,
plan_type: &str,
) -> Result<ModelRefreshResult, AccountManageError> {
let request_id = uuid::Uuid::new_v4().to_string();
let plan_account = ModelRefreshPlanAccount {
plan_type: plan_type.to_string(),
@@ -96,18 +137,20 @@ impl AccountManageService {
account_id: account.account_id.clone(),
installation_id: self.account_pseudonymizer.installation_id(&account.id),
};
if let Err(error) = self
let result = self
.models
.refresh_backend_models(&[plan_account], &request_id)
.refresh_selected_plan_models(&[plan_account], &request_id)
.await
{
tracing::warn!(
account_id = %account.id,
plan_type,
error = %error,
"Failed to refresh account plan models"
);
}
.map_err(|error| AccountManageError::RefreshModels(error.to_string()))?;
let routing = self
.models
.model_plan_routing()
.await
.map_err(|error| AccountManageError::RefreshModels(error.to_string()))?;
self.account_pool
.apply_model_plan_routing(routing.allowlist, routing.fetched_plan_types)
.await;
Ok(result)
}
pub async fn test_connection_stream(
@@ -121,84 +164,55 @@ impl AccountManageService {
.await
.map_err(|_| AccountManageError::Inspect)?
.ok_or(AccountManageError::NotFound)?;
let token = account.access_token.expose_secret().to_string();
let upstream_account_id = account.account_id.clone();
let cookie_header = self
.cookies
.cookie_header_for_request(&account.id, "chatgpt.com", "/codex/responses")
.await
.ok()
.flatten();
let installation_id = self.account_pseudonymizer.installation_id(&account.id);
let codex = self.codex.clone();
let context = AccountUpstreamContext {
access_token: account.access_token.clone(),
account_id: account.account_id.clone(),
request_id: uuid::Uuid::new_v4().to_string(),
cookie_header: self
.cookies
.cookie_header_for_request(&account.id, "chatgpt.com", "/codex/responses")
.await
.ok()
.flatten(),
installation_id: Some(self.account_pseudonymizer.installation_id(&account.id)),
};
let upstream = self.upstream.clone();
let service = self.clone();
let stored_account_id = account.id.clone();
let request_id = uuid::Uuid::new_v4().to_string();
let (tx, rx) = mpsc::channel::<Bytes>(16);
let stored_account_id = account.id;
let (tx, rx) = mpsc::channel(16);
tokio::spawn(async move {
send_test_event(
&tx,
json!({
"type": "test_start",
"model": model.clone(),
"text": "正在连接 Codex Responses"
}),
AccountTestEvent::Started {
model: model.clone(),
},
)
.await;
let request = test_responses_request(model);
send_test_event(
&tx,
json!({
"type": "request",
"payload": serde_json::to_value(&request).unwrap_or_else(|_| json!({}))
}),
)
.await;
let context = CodexRequestContext {
access_token: &token,
account_id: upstream_account_id.as_deref(),
request_id: &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: cookie_header.as_deref(),
installation_id: Some(&installation_id),
session_id: None,
thread_id: None,
client_request_id: None,
turn_id: None,
let request = AccountProbeRequest {
model,
instructions:
"You are checking whether this Codex account can answer. Reply with ok."
.to_string(),
input_text: "hi".to_string(),
};
let mut outcome = match codex.create_response_stream(&request, context).await {
Ok(response) => process_upstream_test_stream(response.body, &tx).await,
let mut outcome = match upstream.probe_response(context, request).await {
Ok(session) => process_probe_session(session, &tx).await,
Err(error) => AccountTestOutcome::error(
error.to_string(),
account_status_from_client_error(&error),
account_status_from_gateway_error(&error),
),
};
if let Some(status) = outcome.status
&& let Some(status) = service
if let Some(status) = outcome.status {
outcome.status = service
.apply_connection_test_status(&stored_account_id, status)
.await
{
outcome.attach_account_status(status);
.await;
}
send_test_event(&tx, outcome.event).await;
send_test_event(&tx, outcome.into_event()).await;
});
let stream = futures::stream::unfold(rx, |mut rx| async {
rx.recv()
.await
.map(|bytes| (Ok::<Bytes, Infallible>(bytes), rx))
rx.recv().await.map(|event| (event, rx))
});
Ok(Box::pin(stream))
}
@@ -223,7 +237,6 @@ impl AccountManageService {
if current.status == AccountStatus::Disabled {
return Some(AccountStatus::Disabled);
}
match self.store.set_status(account_id, status).await {
Ok(true) => {}
Ok(false) => return None,
@@ -237,7 +250,6 @@ impl AccountManageService {
return None;
}
}
if matches!(status, AccountStatus::Expired | AccountStatus::Banned)
&& let Err(error) = self.store.set_next_refresh_at(account_id, None).await
{
@@ -253,6 +265,34 @@ impl AccountManageService {
}
}
async fn process_probe_session(
mut session: AccountProbeSession,
tx: &mpsc::Sender<AccountTestEvent>,
) -> AccountTestOutcome {
send_test_event(
tx,
AccountTestEvent::Request {
payload: session.request_payload,
},
)
.await;
while let Some(event) = session.events.next().await {
match event {
AccountProbeEvent::Content(text) => {
send_test_event(tx, AccountTestEvent::Content { text }).await;
}
AccountProbeEvent::Complete => return AccountTestOutcome::success(),
AccountProbeEvent::Failed(error) => {
return AccountTestOutcome::error(
error.to_string(),
account_status_from_gateway_error(&error),
);
}
}
}
AccountTestOutcome::error("Stream ended before response.completed", None)
}
fn account_model_option(model: &CodexModelInfo) -> AccountModelOption {
AccountModelOption {
id: model.id.clone(),
@@ -270,125 +310,10 @@ fn account_plan_type(account: &StoredAccount) -> String {
.to_string()
}
fn test_responses_request(model: String) -> CodexResponsesRequest {
let mut request = CodexResponsesRequest::new_http_sse(
model,
"You are checking whether this Codex account can answer. Reply with ok.",
vec![json!({
"role": "user",
"content": [
{
"type": "input_text",
"text": "hi"
}
]
})],
);
request.set_stream(true);
request.set_store(false);
request.force_http_sse = true;
request
fn account_status_from_gateway_error(error: &AccountGatewayError) -> Option<AccountStatus> {
classify_account_failure(error.failure()?)?.account_status()
}
async fn process_upstream_test_stream(
mut body: crate::upstream::openai::transport::CodexBackendSseStream,
tx: &mpsc::Sender<Bytes>,
) -> AccountTestOutcome {
let mut buffer = String::new();
while let Some(chunk) = body.next().await {
let chunk = match chunk {
Ok(bytes) => bytes,
Err(error) => {
return AccountTestOutcome::error(
error.to_string(),
account_status_from_client_error(&error),
);
}
};
buffer.push_str(&String::from_utf8_lossy(&chunk));
while let Some(frame) = take_sse_frame(&mut buffer) {
if let Some(outcome) = process_sse_frame(tx, &frame).await {
return outcome;
}
}
}
if !buffer.trim().is_empty()
&& let Some(outcome) = process_sse_frame(tx, &buffer).await
{
return outcome;
}
AccountTestOutcome::error("Stream ended before response.completed", None)
}
async fn process_sse_frame(tx: &mpsc::Sender<Bytes>, frame: &str) -> Option<AccountTestOutcome> {
let events = match parse_sse_events(frame) {
Ok(events) => events,
Err(error) => {
return Some(AccountTestOutcome::error(error.to_string(), None));
}
};
for event in events {
if let Some(outcome) = process_sse_event(tx, &event).await {
return Some(outcome);
}
}
None
}
async fn process_sse_event(
tx: &mpsc::Sender<Bytes>,
event: &SseEvent,
) -> Option<AccountTestOutcome> {
let value: Value = match serde_json::from_str(&event.data) {
Ok(value) => value,
Err(_) => return None,
};
match value.get("type").and_then(Value::as_str) {
Some("response.output_text.delta") => {
if let Some(delta) = value.get("delta").and_then(Value::as_str)
&& !delta.is_empty()
{
send_test_event(tx, json!({ "type": "content", "text": delta })).await;
}
None
}
Some("response.completed" | "response.done") => Some(AccountTestOutcome::success()),
Some(event_name @ ("response.failed" | "error")) => {
let failure = ResponsesSseFailure::from_event(event_name, &value);
let status = sse_failure_account_status(&failure);
Some(AccountTestOutcome::error(failure.message, status))
}
_ => None,
}
}
fn account_status_from_client_error(error: &CodexClientError) -> Option<AccountStatus> {
classify_client_failure(error)?.account_status()
}
fn sse_failure_account_status(failure: &ResponsesSseFailure) -> Option<AccountStatus> {
classify_response_failure(failure)?.account_status()
}
fn take_sse_frame(buffer: &mut String) -> Option<String> {
let index = buffer.find("\n\n").or_else(|| buffer.find("\r\n\r\n"))?;
let delimiter_len = if buffer[index..].starts_with("\r\n\r\n") {
4
} else {
2
};
let frame = buffer[..index + delimiter_len].to_string();
buffer.drain(..index + delimiter_len);
Some(frame)
}
async fn send_test_event(tx: &mpsc::Sender<Bytes>, event: Value) {
let _ = tx
.send(Bytes::from(encode_sse_event("", &event.to_string())))
.await;
async fn send_test_event(tx: &mpsc::Sender<AccountTestEvent>, event: AccountTestEvent) {
let _ = tx.send(event).await;
}
+35 -30
View File
@@ -1,7 +1,16 @@
use secrecy::ExposeSecret;
use serde_json::Value;
use super::{AccountManageService, types::AccountManageError};
use crate::fleet::account_failure::apply_client_failure_effect;
use crate::fleet::{
account_failure::{apply_account_state_effect, classify_account_failure},
account_gateway::AccountUpstreamContext,
quota::QuotaSnapshot,
};
pub struct AccountQuotaRefresh {
pub quota: QuotaSnapshot,
pub raw: Value,
}
impl AccountManageService {
async fn usage_cookie_header(&self, account_id: &str) -> Option<String> {
@@ -24,7 +33,7 @@ impl AccountManageService {
pub async fn account_quota(
&self,
account_id: &str,
) -> Result<serde_json::Value, AccountManageError> {
) -> Result<AccountQuotaRefresh, AccountManageError> {
let stored = self
.store
.get(account_id)
@@ -32,42 +41,38 @@ impl AccountManageService {
.map_err(|_| AccountManageError::NotFound)?
.ok_or(AccountManageError::NotFound)?;
let request_id = uuid::Uuid::new_v4().to_string();
let token = stored.access_token.expose_secret().to_string();
let cookie_header = self.usage_cookie_header(account_id).await;
let installation_id = self.account_pseudonymizer.installation_id(account_id);
let context = crate::upstream::openai::transport::CodexRequestContext {
access_token: &token,
account_id: stored.account_id.as_deref(),
request_id: &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: cookie_header.as_deref(),
installation_id: Some(&installation_id),
session_id: None,
thread_id: None,
client_request_id: None,
turn_id: None,
let context = AccountUpstreamContext {
access_token: stored.access_token.clone(),
account_id: stored.account_id.clone(),
request_id: uuid::Uuid::new_v4().to_string(),
cookie_header: self.usage_cookie_header(account_id).await,
installation_id: Some(self.account_pseudonymizer.installation_id(account_id)),
};
let raw = match self.codex.fetch_usage(context).await {
Ok(raw) => raw,
let result = match self.upstream.fetch_usage(context).await {
Ok(result) => result,
Err(error) => {
apply_client_failure_effect(&self.account_pool, &self.codex, account_id, &error)
if let Some(classified) = error.failure().and_then(classify_account_failure)
&& let Some(effect) = &classified.effect
{
apply_account_state_effect(
&self.account_pool,
self.upstream.as_ref(),
account_id,
effect,
)
.await;
}
return Err(AccountManageError::FetchQuota(error.to_string()));
}
};
let normalized = crate::fleet::quota::quota_from_usage(&raw);
self.account_pool
.apply_quota_snapshot(account_id, &normalized)
.apply_quota_snapshot(account_id, &result.quota)
.await;
self.sync_account_pool_best_effort(account_id, "account quota refresh")
.await;
Ok(serde_json::json!({ "quota": normalized, "raw": raw }))
Ok(AccountQuotaRefresh {
quota: result.quota,
raw: result.raw,
})
}
}
-471
View File
@@ -1,471 +0,0 @@
//! 管理端账号 quota 展示视图。
use std::collections::HashMap;
use chrono::{DateTime, Duration, Utc};
use serde::Serialize;
use serde_json::Value;
use crate::infra::{
format::{format_compact_percent, format_plain_number, format_tokens, nonnegative_i64_to_u64},
time::{china_datetime, china_relative_time},
};
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, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct AccountQuotaData {
refreshed_at_display: String,
windows: Vec<AccountQuotaWindowData>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountQuotaWindowData {
key: String,
group: String,
window_seconds: Option<u64>,
label_display: String,
used_percent: Option<f64>,
used_percent_display: String,
#[serde(skip_serializing_if = "Option::is_none")]
local_usage: Option<AccountQuotaWindowLocalUsageData>,
reset_at_display: String,
window_used_display: String,
#[serde(skip)]
reset_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountQuotaWindowLocalUsageData {
request_count: i64,
request_count_display: String,
input_tokens: i64,
input_tokens_display: String,
output_tokens: i64,
output_tokens_display: String,
cached_tokens: i64,
cached_tokens_display: String,
total_tokens: i64,
total_tokens_display: String,
}
#[derive(Debug, Clone)]
pub(crate) struct AccountQuotaUsageWindow {
pub(crate) key: String,
pub(crate) start: DateTime<Utc>,
pub(crate) end: DateTime<Utc>,
pub(crate) window_seconds: u64,
}
impl AccountQuotaUsageWindow {
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)]
pub(crate) struct AccountQuotaWindowLocalUsage {
pub(crate) request_count: i64,
pub(crate) input_tokens: i64,
pub(crate) output_tokens: i64,
pub(crate) cached_tokens: i64,
}
impl From<AccountQuotaWindowLocalUsage> for AccountQuotaWindowLocalUsageData {
fn from(usage: AccountQuotaWindowLocalUsage) -> Self {
let total_tokens = usage.input_tokens.saturating_add(usage.output_tokens);
Self {
request_count: usage.request_count,
request_count_display: format_plain_number(nonnegative_i64_to_u64(usage.request_count)),
input_tokens: usage.input_tokens,
input_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.input_tokens)),
output_tokens: usage.output_tokens,
output_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.output_tokens)),
cached_tokens: usage.cached_tokens,
cached_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.cached_tokens)),
total_tokens,
total_tokens_display: format_tokens(nonnegative_i64_to_u64(total_tokens)),
}
}
}
impl Default for AccountQuotaData {
fn default() -> Self {
Self {
refreshed_at_display: "-".to_string(),
windows: Vec::new(),
}
}
}
impl AccountQuotaData {
pub(crate) fn usage_windows(&self) -> Vec<AccountQuotaUsageWindow> {
self.windows
.iter()
.filter_map(AccountQuotaWindowData::usage_window)
.collect()
}
pub(crate) fn apply_local_usage(
&mut self,
usage_by_window: &HashMap<String, AccountQuotaWindowLocalUsage>,
) {
for window in &mut self.windows {
window.local_usage = window.usage_window().map(|_| {
usage_by_window
.get(&window.key)
.copied()
.unwrap_or_default()
.into()
});
}
}
}
impl AccountQuotaWindowData {
fn usage_window(&self) -> Option<AccountQuotaUsageWindow> {
let end = self.reset_at?;
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(|| AccountQuotaUsageWindow {
key: self.key.clone(),
start,
end,
window_seconds,
})
}
}
pub(crate) fn quota_data(quota_json: &str, fetched_at: Option<DateTime<Utc>>) -> AccountQuotaData {
let quota = serde_json::from_str::<Value>(quota_json).unwrap_or(Value::Null);
let windows = quota_windows(&quota);
AccountQuotaData {
refreshed_at_display: china_relative_time(fetched_at, Utc::now()),
windows,
}
}
fn quota_windows(quota: &Value) -> Vec<AccountQuotaWindowData> {
let mut windows = Vec::new();
let has_monthly_limit = push_monthly_quota_window(&mut windows, quota.get("monthly_limit"));
if let Some(snapshots) = quota.get("snapshots").and_then(Value::as_array) {
for snapshot in snapshots {
push_snapshot_quota_windows(&mut windows, snapshot, has_monthly_limit);
}
}
windows.sort_by_key(quota_window_sort_key);
windows
}
fn push_monthly_quota_window(
windows: &mut Vec<AccountQuotaWindowData>,
monthly_limit: Option<&Value>,
) -> bool {
let Some(monthly_limit) = monthly_limit.filter(|value| !value.is_null()) else {
return false;
};
let used_percent = monthly_limit
.get("used_percent")
.and_then(number_value)
.map(|value| value.clamp(0.0, 100.0));
let reset_at = monthly_limit
.get("reset_at")
.and_then(Value::as_i64)
.and_then(|value| DateTime::<Utc>::from_timestamp(value, 0));
let window_seconds = monthly_limit
.get("window_minutes")
.and_then(Value::as_u64)
.and_then(|minutes| minutes.checked_mul(60))
.or(Some(MONTH_WINDOW_SECONDS));
if used_percent.is_none() && reset_at.is_none() {
return false;
}
let key = monthly_limit
.get("key")
.and_then(Value::as_str)
.unwrap_or("monthly");
windows.push(AccountQuotaWindowData {
key: quota_key_segment(key),
group: "monthly".to_string(),
window_seconds,
label_display: "月限额".to_string(),
used_percent,
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_compact_percent),
local_usage: None,
reset_at_display: reset_at
.as_ref()
.map_or_else(|| "-".to_string(), china_datetime),
window_used_display: quota_window_used_display(reset_at, window_seconds),
reset_at,
});
true
}
fn push_snapshot_quota_windows(
windows: &mut Vec<AccountQuotaWindowData>,
snapshot: &Value,
skip_core_monthly: bool,
) {
let source = snapshot
.get("source")
.and_then(Value::as_str)
.unwrap_or("quota");
let source_key = snapshot_source_key(source, snapshot);
let label_prefix = snapshot_label(snapshot);
for role in ["primary", "secondary"] {
let Some(window) = snapshot.get(role).filter(|value| !value.is_null()) else {
continue;
};
let window_seconds = window
.get("window_minutes")
.and_then(Value::as_u64)
.and_then(|minutes| minutes.checked_mul(60));
if skip_core_monthly
&& source == "core"
&& window_seconds
.is_some_and(|seconds| quota_window_matches(seconds, MONTH_WINDOW_SECONDS))
{
continue;
}
push_quota_window(
windows,
&source_key,
role,
label_prefix.as_deref(),
Some(window),
);
}
}
fn push_quota_window(
windows: &mut Vec<AccountQuotaWindowData>,
source_key: &str,
role: &str,
label_prefix: Option<&str>,
window: Option<&Value>,
) {
let Some(window) = window.filter(|value| !value.is_null()) else {
return;
};
let used_percent = window
.get("used_percent")
.and_then(number_value)
.map(|value| value.clamp(0.0, 100.0));
let reset_at = window
.get("reset_at")
.and_then(Value::as_i64)
.and_then(|value| DateTime::<Utc>::from_timestamp(value, 0));
let window_seconds = window
.get("window_minutes")
.and_then(Value::as_u64)
.and_then(|minutes| minutes.checked_mul(60));
let label_prefix = label_prefix
.map(str::trim)
.filter(|value| !value.is_empty());
if reset_at.is_none()
&& window_seconds.is_none()
&& (used_percent.is_none() || label_prefix.is_none())
{
return;
}
let base_label = quota_window_label_display(window_seconds);
let label_display = label_prefix
.map(|value| format!("{value} · {base_label}"))
.unwrap_or(base_label);
windows.push(AccountQuotaWindowData {
key: unique_quota_window_key(windows, source_key, role, window_seconds),
group: quota_window_group(window_seconds).to_string(),
window_seconds,
label_display,
used_percent,
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_compact_percent),
local_usage: None,
reset_at_display: reset_at
.as_ref()
.map_or_else(|| "-".to_string(), china_datetime),
window_used_display: quota_window_used_display(reset_at, window_seconds),
reset_at,
});
}
fn snapshot_source_key(source: &str, snapshot: &Value) -> String {
let label = snapshot
.get("limit_name")
.and_then(Value::as_str)
.or_else(|| snapshot.get("metered_feature").and_then(Value::as_str))
.unwrap_or(source);
format!("{}-{}", quota_key_segment(source), quota_key_segment(label))
}
fn snapshot_label(snapshot: &Value) -> Option<String> {
let source = snapshot.get("source").and_then(Value::as_str);
if source == Some("core") {
return None;
}
let label = snapshot
.get("limit_name")
.and_then(Value::as_str)
.or_else(|| snapshot.get("metered_feature").and_then(Value::as_str))
.map(str::trim)
.filter(|value| !value.is_empty())?;
if label.eq_ignore_ascii_case("codex") {
return None;
}
if is_review_limit_label(Some(label)) {
return Some("代码审查".to_string());
}
Some(label.to_string())
}
fn unique_quota_window_key(
windows: &[AccountQuotaWindowData],
source_key: &str,
role: &str,
window_seconds: Option<u64>,
) -> String {
let bucket = quota_window_key_part(window_seconds).unwrap_or(role);
let key = format!("{source_key}-{bucket}");
if windows.iter().any(|window| window.key == key) {
format!("{key}-{role}")
} else {
key
}
}
fn quota_window_key_part(window_seconds: Option<u64>) -> Option<&'static str> {
match window_seconds {
Some(seconds) if quota_window_matches(seconds, FIVE_HOUR_WINDOW_SECONDS) => {
Some("five-hour")
}
Some(seconds) if quota_window_matches(seconds, WEEK_WINDOW_SECONDS) => Some("weekly"),
Some(seconds) if quota_window_matches(seconds, MONTH_WINDOW_SECONDS) => Some("monthly"),
_ => None,
}
}
fn quota_key_segment(value: &str) -> String {
let mut segment = String::with_capacity(value.len());
for ch in value.chars() {
if ch.is_ascii_alphanumeric() {
segment.push(ch.to_ascii_lowercase());
} else if !segment.ends_with('-') {
segment.push('-');
}
}
let segment = segment.trim_matches('-');
if segment.is_empty() {
"quota".to_string()
} else {
segment.to_string()
}
}
fn quota_window_sort_key(window: &AccountQuotaWindowData) -> (u8, u64, String) {
let group_order = match window.group.as_str() {
"monthly" => 0,
"shortTerm" => 1,
_ => 2,
};
(
group_order,
window.window_seconds.unwrap_or(0),
window.key.clone(),
)
}
fn quota_window_group(window_seconds: Option<u64>) -> &'static str {
match window_seconds {
Some(seconds) if quota_window_matches(seconds, MONTH_WINDOW_SECONDS) => "monthly",
Some(seconds)
if quota_window_matches(seconds, FIVE_HOUR_WINDOW_SECONDS)
|| quota_window_matches(seconds, WEEK_WINDOW_SECONDS) =>
{
"shortTerm"
}
_ => "other",
}
}
fn quota_window_matches(actual: u64, expected: u64) -> bool {
actual > 0 && actual.abs_diff(expected) <= expected / 20
}
fn quota_window_label_display(window_seconds: Option<u64>) -> String {
let Some(window_seconds) = window_seconds.filter(|seconds| *seconds > 0) else {
return "额度".to_string();
};
match window_seconds {
seconds if quota_window_matches(seconds, FIVE_HOUR_WINDOW_SECONDS) => {
"5小时限额".to_string()
}
seconds if quota_window_matches(seconds, WEEK_WINDOW_SECONDS) => "周限额".to_string(),
seconds if quota_window_matches(seconds, MONTH_WINDOW_SECONDS) => "月限额".to_string(),
seconds if seconds % 86_400 == 0 => format!("{}天限额", seconds / 86_400),
seconds if seconds % 3_600 == 0 => format!("{}小时限额", seconds / 3_600),
seconds => format!("{}分钟限额", seconds.div_ceil(60)),
}
}
fn quota_window_used_display(
reset_at: Option<DateTime<Utc>>,
window_seconds: Option<u64>,
) -> String {
let (Some(reset_at), Some(window_seconds)) = (reset_at, window_seconds) else {
return "-".to_string();
};
let remaining = reset_at
.signed_duration_since(Utc::now())
.num_seconds()
.max(0)
.cast_unsigned();
let used = window_seconds.saturating_sub(remaining);
format!(
"{} / {}",
format_duration_days(used),
format_duration_days(window_seconds)
)
}
fn format_duration_days(seconds: u64) -> String {
let days = seconds as f64 / 86_400.0;
if days >= 1.0 {
format!("{days:.1}d")
} else {
format!("{:.1}h", seconds as f64 / 3_600.0)
}
}
fn number_value(value: &Value) -> Option<f64> {
value
.as_f64()
.or_else(|| value.as_str().and_then(|value| value.parse::<f64>().ok()))
.filter(|value| value.is_finite())
}
fn is_review_limit_label(value: Option<&str>) -> bool {
let normalized = value
.unwrap_or_default()
.trim()
.to_ascii_lowercase()
.replace(['-', ' '], "_");
normalized == "review"
|| normalized == "code_review"
|| normalized == "codex_review"
|| normalized == "codex_code_review"
|| normalized.contains("code_review")
|| normalized.contains("codex_review")
}
+8 -8
View File
@@ -9,14 +9,14 @@ use super::oauth;
use crate::{
fleet::{
account_gateway::AccountUpstreamGateway,
cookies::PgCookieStore,
pool::AccountPoolService,
refresh::{RedisRefreshLeaseStore, RuntimeRefreshPolicy},
refresh::{RedisRefreshLeaseStore, RuntimeRefreshPolicy, TokenRefresher},
store::PgAccountStore,
},
infra::identity::AccountPseudonymizer,
models::service::ModelService,
upstream::openai::{token_client::TokenRefresher, transport::CodexBackendClient},
};
use super::types::{AccountManageError, RefreshedAccountTokens};
@@ -25,7 +25,7 @@ use super::types::{AccountManageError, RefreshedAccountTokens};
pub struct AccountManageService {
pub store: PgAccountStore,
pub(crate) cookies: PgCookieStore,
pub(crate) codex: Arc<CodexBackendClient>,
pub(crate) upstream: Arc<dyn AccountUpstreamGateway>,
pub(crate) models: Arc<ModelService>,
pub(crate) account_pool: Arc<AccountPoolService>,
pub(crate) token_refresher: Arc<dyn TokenRefresher>,
@@ -39,7 +39,7 @@ pub struct AccountManageService {
pub struct AccountManageServiceParts {
pub store: PgAccountStore,
pub cookies: PgCookieStore,
pub codex: Arc<CodexBackendClient>,
pub upstream: Arc<dyn AccountUpstreamGateway>,
pub models: Arc<ModelService>,
pub account_pool: Arc<AccountPoolService>,
pub token_refresher: Arc<dyn TokenRefresher>,
@@ -54,7 +54,7 @@ impl AccountManageService {
Self {
store: parts.store,
cookies: parts.cookies,
codex: parts.codex,
upstream: parts.upstream,
models: parts.models,
account_pool: parts.account_pool,
token_refresher: parts.token_refresher,
@@ -127,7 +127,7 @@ impl AccountManageService {
}
pub(crate) async fn evict_account_websocket_pool(&self, account_id: &str) {
self.codex.evict_websocket_account(account_id).await;
self.upstream.evict_account_connections(account_id).await;
match self.store.get(account_id).await {
Ok(Some(account)) => {
if let Some(upstream_account_id) = account
@@ -135,8 +135,8 @@ impl AccountManageService {
.as_deref()
.filter(|value| *value != account_id)
{
self.codex
.evict_websocket_account(upstream_account_id)
self.upstream
.evict_account_connections(upstream_account_id)
.await;
}
}
+10 -24
View File
@@ -4,13 +4,12 @@ use chrono::{DateTime, Utc};
use serde::Serialize;
use thiserror::Error;
use crate::{
fleet::{
account::{AccountStatus, AccountStatus as AcctStatus},
account_failure::classify_client_failure,
store::StoredAccountMetadata,
},
upstream::openai::{token_client::RefreshFailure, transport::CodexClientError},
use crate::fleet::{
account::{AccountStatus, AccountStatus as AcctStatus},
account_failure::classify_account_failure,
account_gateway::AccountGatewayError,
refresh::RefreshFailure,
store::StoredAccountMetadata,
};
#[derive(Debug, Clone, serde::Serialize)]
@@ -86,6 +85,8 @@ pub enum AccountManageError {
FetchQuota(String),
#[error("no available models returned")]
NoModels,
#[error("failed to refresh account models: {0}")]
RefreshModels(String),
#[error("invalid status: {0}")]
InvalidStatus(String),
#[error("label must be 64 characters or fewer")]
@@ -296,21 +297,6 @@ pub(super) fn stored_to_admin_metadata(s: crate::fleet::store::StoredAccount) ->
})
}
pub(super) fn import_usage_plan_type(usage: &serde_json::Value) -> Option<String> {
usage
.get("plan_type")
.and_then(serde_json::Value::as_str)
.and_then(normalized_plan_type)
}
pub(super) fn import_usage_string(usage: &serde_json::Value, key: &str) -> Option<String> {
usage
.get(key)
.and_then(serde_json::Value::as_str)
.and_then(|value| crate::fleet::import::normalize_nonempty_str(Some(value)))
.map(ToString::to_string)
}
pub(super) fn import_quota_plan_type(quota: &serde_json::Value) -> Option<String> {
quota
.get("plan_type")
@@ -323,8 +309,8 @@ fn normalized_plan_type(value: &str) -> Option<String> {
(!value.is_empty() && !matches!(value.as_str(), "unknown" | "null")).then_some(value)
}
pub(super) fn import_status_from_usage_error(error: &CodexClientError) -> Option<AccountStatus> {
classify_client_failure(error)?
pub(super) fn import_status_from_usage_error(error: &AccountGatewayError) -> Option<AccountStatus> {
classify_account_failure(error.failure()?)?
.account_status()
.filter(|status| *status == AccountStatus::Banned)
}
+2
View File
@@ -2,6 +2,7 @@
pub mod account;
pub mod account_failure;
pub mod account_gateway;
pub mod cookies;
pub mod import;
pub mod manage;
@@ -10,4 +11,5 @@ pub mod quota;
pub mod refresh;
pub mod scheduler;
pub mod store;
pub mod usage;
mod window;
+46 -55
View File
@@ -12,23 +12,24 @@ use chrono::{DateTime, Duration, Utc};
use indexmap::IndexMap;
use thiserror::Error;
use crate::fleet::quota::{quota_snapshot_limit_reached, quota_snapshot_reset_at};
use crate::fleet::quota::{
QuotaObservation, QuotaSnapshot, quota_from_observation, quota_snapshot_limit_reached,
quota_snapshot_reset_at,
};
use crate::fleet::refresh::{JwtExpiry, jwt_expiry};
use crate::fleet::scheduler::{
AccountScheduler, ScoreWeights,
AccountScheduler, AttemptFeedback, ScoreWeights,
candidates::{self, CandidateFilter, CandidateRequest},
};
use crate::fleet::store::AccountStore;
use crate::fleet::{
account::{Account, AccountStatus},
usage::{
AccountUsageDelta, AccountUsageSnapshot, AccountUsageStore, AccountUsageWindow,
ResponseUsage,
},
window::should_reset_usage_window,
};
use crate::telemetry::account_usage::store::{
AccountUsageDelta, AccountUsageSnapshot, AccountUsageStore, AccountUsageWindow,
};
use crate::upstream::openai::protocol::events::{
TokenUsage as CodexTokenUsage, parse_rate_limit_headers, rate_limit_quota,
};
pub use crate::fleet::scheduler::RotationStrategy;
@@ -48,23 +49,10 @@ pub struct AccountPoolStaticSettings {
pub tier_priority: Vec<String>,
}
impl AccountPoolStaticSettings {
/// 将设置快照转换为账号池选项。
pub fn pool_options(&self, settings: &crate::settings::SettingsSnapshot) -> AccountPoolOptions {
AccountPoolOptions {
max_concurrent_per_account: settings.max_concurrent_per_account,
rotation_strategy: match settings.rotation_strategy.as_str() {
"smart" => RotationStrategy::Smart,
"quota_reset_priority" => RotationStrategy::QuotaResetPriority,
"round_robin" => RotationStrategy::RoundRobin,
"sticky" => RotationStrategy::Sticky,
_ => RotationStrategy::Smart,
},
skip_quota_limited: self.skip_quota_limited,
tier_priority: self.tier_priority.clone(),
..AccountPoolOptions::default()
}
}
#[derive(Debug, Clone)]
pub struct AccountPoolRuntimeOptions {
pub options: AccountPoolOptions,
pub request_interval_ms: u64,
}
#[derive(Clone)]
@@ -223,19 +211,20 @@ impl AccountPoolService {
StdDuration::from_millis(request_interval_ms);
}
/// 回灌一个已经启动的上游 attempt,供 Smart 调度器更新运行时健康反馈。
pub async fn report_attempt(&self, account_id: &str, feedback: AttemptFeedback) {
self.pool.lock().await.report_attempt(account_id, feedback);
}
/// 持续接收运行时设置并更新账号池参数。
pub async fn subscribe_settings(
pub async fn subscribe_options(
self: Arc<Self>,
mut receiver: tokio::sync::watch::Receiver<crate::settings::SettingsSnapshot>,
static_settings: AccountPoolStaticSettings,
mut receiver: tokio::sync::watch::Receiver<AccountPoolRuntimeOptions>,
) {
while receiver.changed().await.is_ok() {
let settings = receiver.borrow_and_update().clone();
self.apply_options(
static_settings.pool_options(&settings),
settings.request_interval_ms,
)
.await;
let runtime = receiver.borrow_and_update().clone();
self.apply_options(runtime.options, runtime.request_interval_ms)
.await;
}
}
@@ -541,16 +530,11 @@ impl AccountPoolService {
}
}
/// 记录上游 token 用量。
pub async fn record_token_usage(&self, account_id: &str, usage: &CodexTokenUsage) {
self.record_response_usage(account_id, *usage, false).await;
}
/// 记录 Responses 请求用量。
pub async fn record_response_usage(
&self,
account_id: &str,
usage: CodexTokenUsage,
usage: ResponseUsage,
image_generation_requested: bool,
) {
let image_request_succeeded = image_generation_requested && usage.image_output_tokens > 0;
@@ -621,32 +605,29 @@ impl AccountPoolService {
}
}
/// 根据上游返回的 rate-limit header 被动同步 quota 状态。
pub async fn sync_passive_rate_limit_headers(
/// 被动同步上游 adapter 已归一化的 quota observation。
pub async fn sync_passive_quota_observation(
&self,
account: &Account,
headers: &[(String, String)],
observation: &QuotaObservation,
) {
self.sync_passive_rate_limit_headers_for_account(
self.sync_passive_quota_observation_for_account(
&account.id,
account.plan_type.as_deref(),
headers,
observation,
)
.await;
}
/// 根据上游返回的 rate-limit header 被动同步指定账号 quota 状态。
pub async fn sync_passive_rate_limit_headers_for_account(
/// 被动同步指定账号的 typed quota observation。
pub async fn sync_passive_quota_observation_for_account(
&self,
account_id: &str,
plan_type: Option<&str>,
headers: &[(String, String)],
observation: &QuotaObservation,
) {
let Some(rate_limits) = parse_rate_limit_headers(headers) else {
return;
};
let existing_quota = match self.store.get_quota_json(account_id).await {
Ok(Some(quota_json)) => serde_json::from_str::<serde_json::Value>(&quota_json).ok(),
Ok(Some(quota_json)) => QuotaSnapshot::from_json(&quota_json).ok(),
Ok(None) => None,
Err(error) => {
tracing::error!(
@@ -657,18 +638,28 @@ impl AccountPoolService {
None
}
};
let quota = rate_limit_quota(&rate_limits, plan_type, existing_quota.as_ref());
let quota = quota_from_observation(observation, plan_type, existing_quota.as_ref());
self.apply_quota_snapshot(account_id, &quota).await;
}
/// 应用 quota 快照到持久化存储和运行时账号池。
pub async fn apply_quota_snapshot(&self, account_id: &str, quota: &serde_json::Value) -> bool {
pub async fn apply_quota_snapshot(&self, account_id: &str, quota: &QuotaSnapshot) -> bool {
let limit_reached = quota_snapshot_limit_reached(quota);
let reset_at = quota_snapshot_reset_at(quota);
let cooldown_until = limit_reached.then_some(reset_at).flatten();
let limit_window_seconds =
reset_at.and_then(|_| crate::fleet::quota::quota_snapshot_limit_window_seconds(quota));
let quota_json = quota.to_string();
let quota_json = match quota.to_json() {
Ok(quota_json) => quota_json,
Err(error) => {
tracing::error!(
account_id,
error = %error,
"Failed to serialize verified quota snapshot"
);
return false;
}
};
let persisted = match self
.store
.apply_quota_snapshot(account_id, &quota_json, limit_reached, cooldown_until)
+7 -4
View File
@@ -206,10 +206,13 @@ impl AccountPool {
self.options = options;
}
/// 回灌一次请求结果到调度器的运行时反馈(供 Smart 打分)。
pub fn report_feedback(&self, account_id: &str, success: bool, first_token_ms: Option<u64>) {
self.scheduler
.report_feedback(account_id, success, first_token_ms);
/// 回灌一个已启动上游 attempt 的运行时反馈(供 Smart 打分)。
pub fn report_attempt(
&self,
account_id: &str,
feedback: crate::fleet::scheduler::AttemptFeedback,
) {
self.scheduler.report_attempt(account_id, feedback);
}
/// 更新账号状态。
+4
View File
@@ -1,7 +1,11 @@
//! 账号配额解析与刷新。
mod model;
mod observation;
mod runtime;
mod service;
pub use model::*;
pub use observation::*;
pub use runtime::*;
pub use service::*;
+673
View File
@@ -0,0 +1,673 @@
//! fleet 唯一拥有的 typed quota 快照。
use std::collections::BTreeMap;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Deserializer, Serialize};
use serde_json::Value;
pub const FIVE_HOUR_WINDOW_MINUTES: u64 = 300;
pub const WEEK_WINDOW_MINUTES: u64 = 10_080;
pub const MONTH_WINDOW_MINUTES: u64 = 43_200;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum QuotaWindowKind {
ShortTerm,
Weekly,
Monthly,
#[default]
Other,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum QuotaWindowRole {
Primary,
Secondary,
Monthly,
}
#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct QuotaWindow {
#[serde(skip)]
pub kind: QuotaWindowKind,
pub used_percent: Option<f64>,
pub remaining_percent: Option<i64>,
pub reset_at: Option<i64>,
pub window_minutes: Option<u64>,
#[serde(default)]
pub limit_reached: bool,
pub used: Option<Value>,
pub limit: Option<Value>,
pub used_credits: Option<Value>,
pub limit_credits: Option<Value>,
}
#[derive(Deserialize)]
struct QuotaWindowWire {
#[serde(default, deserialize_with = "deserialize_optional_finite_number")]
used_percent: Option<f64>,
#[serde(default, deserialize_with = "deserialize_optional_finite_number")]
remaining_percent: Option<f64>,
#[serde(default, deserialize_with = "deserialize_optional_positive_i64")]
reset_at: Option<i64>,
#[serde(default, deserialize_with = "deserialize_optional_positive_u64")]
window_minutes: Option<u64>,
#[serde(default)]
limit_reached: bool,
#[serde(default)]
used: Option<Value>,
#[serde(default)]
limit: Option<Value>,
#[serde(default)]
used_credits: Option<Value>,
#[serde(default)]
limit_credits: Option<Value>,
}
impl<'de> Deserialize<'de> for QuotaWindow {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let wire = QuotaWindowWire::deserialize(deserializer)?;
Ok(Self {
kind: quota_window_kind(wire.window_minutes),
used_percent: wire.used_percent.map(clamp_percent),
remaining_percent: wire
.remaining_percent
.map(clamp_percent)
.map(|value| value.round() as i64),
reset_at: wire.reset_at,
window_minutes: wire.window_minutes,
limit_reached: wire.limit_reached,
used: wire.used,
limit: wire.limit,
used_credits: wire.used_credits,
limit_credits: wire.limit_credits,
})
}
}
impl QuotaWindow {
pub fn from_usage(
used_percent: f64,
reset_at: Option<i64>,
window_minutes: Option<u64>,
used: Option<Value>,
limit: Option<Value>,
) -> Self {
let used_percent = clamp_percent(used_percent);
Self {
kind: quota_window_kind(window_minutes),
used_percent: Some(used_percent),
remaining_percent: Some(remaining_percent(used_percent)),
reset_at,
window_minutes,
limit_reached: used_percent >= 100.0,
used,
limit,
used_credits: None,
limit_credits: None,
}
}
pub fn from_observation(
used_percent: f64,
reset_at: Option<i64>,
window_minutes: Option<u64>,
) -> Self {
Self::from_usage(used_percent, reset_at, window_minutes, None, None)
}
pub fn window_seconds(&self) -> Option<u64> {
self.window_minutes?.checked_mul(60)
}
pub fn reset_datetime(&self) -> Option<DateTime<Utc>> {
DateTime::<Utc>::from_timestamp(self.reset_at?, 0)
}
pub fn is_limit_reached(&self) -> bool {
self.limit_reached || self.used_percent.is_some_and(|used| used >= 100.0)
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct QuotaLimitSnapshot {
pub source: String,
#[serde(default)]
pub limit_name: Option<String>,
#[serde(default)]
pub metered_feature: Option<String>,
#[serde(default)]
pub allowed: Option<bool>,
#[serde(default)]
pub limit_reached: Option<bool>,
#[serde(default)]
pub blocked: bool,
#[serde(default)]
pub primary: Option<QuotaWindow>,
#[serde(default)]
pub secondary: Option<QuotaWindow>,
}
impl QuotaLimitSnapshot {
pub fn is_blocked(&self) -> bool {
self.blocked
|| self.limit_reached.unwrap_or(false)
|| self.allowed.is_some_and(|allowed| !allowed)
|| self.windows().any(QuotaWindow::is_limit_reached)
}
pub fn windows(&self) -> impl Iterator<Item = &QuotaWindow> {
[self.primary.as_ref(), self.secondary.as_ref()]
.into_iter()
.flatten()
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MonthlyQuotaLimit {
#[serde(default)]
pub key: Option<String>,
#[serde(default)]
pub source: Option<String>,
#[serde(default, deserialize_with = "deserialize_optional_finite_number")]
pub used_percent: Option<f64>,
#[serde(default, deserialize_with = "deserialize_optional_finite_number")]
pub remaining_percent: Option<f64>,
#[serde(default, deserialize_with = "deserialize_optional_positive_i64")]
pub reset_at: Option<i64>,
#[serde(default, deserialize_with = "deserialize_optional_positive_u64")]
pub window_minutes: Option<u64>,
#[serde(default)]
pub limit_reached: bool,
#[serde(default)]
pub used: Option<Value>,
#[serde(default)]
pub limit: Option<Value>,
#[serde(default)]
pub used_credits: Option<Value>,
#[serde(default)]
pub limit_credits: Option<Value>,
}
impl MonthlyQuotaLimit {
pub fn as_window(&self) -> QuotaWindow {
QuotaWindow {
kind: QuotaWindowKind::Monthly,
used_percent: self.used_percent.map(clamp_percent),
remaining_percent: self
.remaining_percent
.map(clamp_percent)
.map(|value| value.round() as i64),
reset_at: self.reset_at,
window_minutes: self.window_minutes.or(Some(MONTH_WINDOW_MINUTES)),
limit_reached: self.limit_reached,
used: self.used.clone(),
limit: self.limit.clone(),
used_credits: self.used_credits.clone(),
limit_credits: self.limit_credits.clone(),
}
}
pub fn is_limit_reached(&self) -> bool {
self.limit_reached || self.used_percent.is_some_and(|used| used >= 100.0)
}
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct QuotaCredits {
#[serde(default)]
pub has_credits: bool,
#[serde(default)]
pub unlimited: bool,
#[serde(default)]
pub overage_limit_reached: bool,
#[serde(default)]
pub balance: Option<Value>,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct QuotaSpendLimit {
#[serde(default, deserialize_with = "deserialize_optional_finite_number")]
pub used_percent: Option<f64>,
#[serde(default, deserialize_with = "deserialize_optional_finite_number")]
pub remaining_percent: Option<f64>,
#[serde(default, deserialize_with = "deserialize_optional_positive_i64")]
pub reset_at: Option<i64>,
#[serde(default)]
pub used: Option<Value>,
#[serde(default)]
pub limit: Option<Value>,
#[serde(flatten)]
pub extensions: BTreeMap<String, Value>,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct QuotaSpendControl {
#[serde(default)]
pub reached: bool,
#[serde(default)]
pub individual_limit: Option<QuotaSpendLimit>,
#[serde(flatten)]
pub extensions: BTreeMap<String, Value>,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct QuotaSnapshot {
#[serde(default)]
pub plan_type: Option<String>,
#[serde(default)]
pub active_limit: Option<String>,
#[serde(default, rename = "snapshots")]
pub limits: Vec<QuotaLimitSnapshot>,
#[serde(default)]
pub monthly_limit: Option<MonthlyQuotaLimit>,
#[serde(default)]
pub credits: Option<QuotaCredits>,
#[serde(default)]
pub spend_control: Option<QuotaSpendControl>,
#[serde(default)]
pub promo_message: Option<String>,
#[serde(default)]
pub rate_limit_reached_type: Option<String>,
#[serde(flatten)]
pub extensions: BTreeMap<String, Value>,
}
impl QuotaSnapshot {
pub fn from_json(value: &str) -> serde_json::Result<Self> {
serde_json::from_str(value)
}
pub fn from_value(value: Value) -> serde_json::Result<Self> {
serde_json::from_value(value)
}
pub fn to_json(&self) -> serde_json::Result<String> {
serde_json::to_string(self)
}
pub fn is_limit_reached(&self) -> bool {
self.limits.iter().any(QuotaLimitSnapshot::is_blocked)
|| self
.monthly_limit
.as_ref()
.is_some_and(MonthlyQuotaLimit::is_limit_reached)
|| self
.spend_control
.as_ref()
.is_some_and(quota_spend_control_limit_reached)
|| self
.credits
.as_ref()
.is_some_and(|credits| credits.overage_limit_reached)
}
pub fn reset_at(&self) -> Option<DateTime<Utc>> {
let reset_at = if self.monthly_or_spend_limit_reached() {
self.monthly_reset_at()
.or_else(|| self.blocking_window().and_then(|window| window.reset_at))
.or_else(|| {
self.core_window_with_reset()
.and_then(|window| window.reset_at)
})
} else if let Some(reset_at) = self.blocking_window().and_then(|window| window.reset_at) {
Some(reset_at)
} else {
self.core_window_with_reset()
.and_then(|window| window.reset_at)
.or_else(|| self.monthly_reset_at())
}?;
DateTime::<Utc>::from_timestamp(reset_at, 0)
}
pub fn limit_window_seconds(&self) -> Option<u64> {
if self.monthly_or_spend_limit_reached() {
return self
.monthly_window_seconds()
.or(MONTH_WINDOW_MINUTES.checked_mul(60));
}
self.blocking_window()
.and_then(QuotaWindow::window_seconds)
.or_else(|| self.core_windows().find_map(QuotaWindow::window_seconds))
.or_else(|| self.monthly_window_seconds())
}
pub fn windows(&self) -> impl Iterator<Item = QuotaWindowEntry<'_>> {
let monthly = self.monthly_limit.as_ref().map(|limit| QuotaWindowEntry {
source: limit.source.as_deref().unwrap_or("monthly"),
limit_name: limit.key.as_deref(),
metered_feature: None,
role: QuotaWindowRole::Monthly,
window: QuotaWindowRef::Monthly(limit),
});
monthly
.into_iter()
.chain(self.limits.iter().flat_map(|limit| {
let primary = limit.primary.as_ref().map(|window| QuotaWindowEntry {
source: &limit.source,
limit_name: limit.limit_name.as_deref(),
metered_feature: limit.metered_feature.as_deref(),
role: QuotaWindowRole::Primary,
window: QuotaWindowRef::Regular(window),
});
let secondary = limit.secondary.as_ref().map(|window| QuotaWindowEntry {
source: &limit.source,
limit_name: limit.limit_name.as_deref(),
metered_feature: limit.metered_feature.as_deref(),
role: QuotaWindowRole::Secondary,
window: QuotaWindowRef::Regular(window),
});
primary.into_iter().chain(secondary)
}))
}
fn monthly_or_spend_limit_reached(&self) -> bool {
self.monthly_limit
.as_ref()
.is_some_and(MonthlyQuotaLimit::is_limit_reached)
|| self
.spend_control
.as_ref()
.is_some_and(quota_spend_control_limit_reached)
}
fn monthly_reset_at(&self) -> Option<i64> {
self.monthly_limit
.as_ref()?
.reset_at
.filter(|value| *value > 0)
}
fn monthly_window_seconds(&self) -> Option<u64> {
self.monthly_limit
.as_ref()?
.window_minutes
.filter(|value| *value > 0)?
.checked_mul(60)
}
fn blocking_window(&self) -> Option<&QuotaWindow> {
self.limits
.iter()
.filter(|limit| limit.is_blocked())
.flat_map(QuotaLimitSnapshot::windows)
.find(|window| window.is_limit_reached())
.or_else(|| {
self.limits
.iter()
.filter(|limit| limit.is_blocked())
.flat_map(QuotaLimitSnapshot::windows)
.find(|window| window.reset_at.is_some())
})
}
fn core_window_with_reset(&self) -> Option<&QuotaWindow> {
self.core_windows().find(|window| window.reset_at.is_some())
}
fn core_windows(&self) -> impl Iterator<Item = &QuotaWindow> {
self.limits
.iter()
.filter(|limit| limit.source == "core")
.flat_map(QuotaLimitSnapshot::windows)
}
}
#[derive(Debug, Clone, Copy)]
pub enum QuotaWindowRef<'a> {
Regular(&'a QuotaWindow),
Monthly(&'a MonthlyQuotaLimit),
}
impl QuotaWindowRef<'_> {
pub fn to_owned(self) -> QuotaWindow {
match self {
Self::Regular(window) => window.clone(),
Self::Monthly(limit) => limit.as_window(),
}
}
}
#[derive(Debug, Clone, Copy)]
pub struct QuotaWindowEntry<'a> {
pub source: &'a str,
pub limit_name: Option<&'a str>,
pub metered_feature: Option<&'a str>,
pub role: QuotaWindowRole,
pub window: QuotaWindowRef<'a>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct ClassifiedQuotaWindow {
pub key: String,
pub source: String,
pub limit_name: Option<String>,
pub metered_feature: Option<String>,
pub role: QuotaWindowRole,
pub window: QuotaWindow,
pub order: u8,
}
impl QuotaSnapshot {
pub fn classified_windows(&self) -> Vec<ClassifiedQuotaWindow> {
let has_monthly_limit = self.monthly_limit.is_some();
let mut windows = Vec::new();
for entry in self.windows() {
let window = entry.window.to_owned();
if has_monthly_limit
&& entry.source == "core"
&& entry.role != QuotaWindowRole::Monthly
&& window.kind == QuotaWindowKind::Monthly
{
continue;
}
let has_semantic_label = entry
.limit_name
.or(entry.metered_feature)
.is_some_and(|label| !label.trim().is_empty());
if window.reset_at.is_none()
&& window.window_minutes.is_none()
&& (window.used_percent.is_none() || !has_semantic_label)
{
continue;
}
let source_key =
quota_source_key(entry.source, entry.limit_name, entry.metered_feature);
let bucket = quota_window_key_part(window.kind)
.unwrap_or_else(|| quota_window_role_name(entry.role));
let mut key = format!("{source_key}-{bucket}");
if windows
.iter()
.any(|existing: &ClassifiedQuotaWindow| existing.key == key)
{
key = format!("{key}-{}", quota_window_role_name(entry.role));
}
windows.push(ClassifiedQuotaWindow {
key,
source: entry.source.to_string(),
limit_name: entry.limit_name.map(ToString::to_string),
metered_feature: entry.metered_feature.map(ToString::to_string),
role: entry.role,
order: quota_window_order(window.kind),
window,
});
}
windows.sort_by(|left, right| {
left.order
.cmp(&right.order)
.then_with(|| {
left.window
.window_minutes
.unwrap_or_default()
.cmp(&right.window.window_minutes.unwrap_or_default())
})
.then_with(|| left.key.cmp(&right.key))
});
windows
}
pub fn representative_used_percent(&self) -> Option<f64> {
self.classified_windows()
.into_iter()
.filter_map(|window| {
window
.window
.used_percent
.map(|used| (quota_usage_priority(window.window.kind), used))
})
.fold(None::<(u8, f64)>, |selected, candidate| match selected {
Some(current)
if current.0 < candidate.0
|| (current.0 == candidate.0 && current.1 >= candidate.1) =>
{
Some(current)
}
_ => Some(candidate),
})
.map(|(_, used)| used)
}
}
pub fn quota_window_kind(window_minutes: Option<u64>) -> QuotaWindowKind {
match window_minutes {
Some(minutes) if window_minutes_match(minutes, FIVE_HOUR_WINDOW_MINUTES) => {
QuotaWindowKind::ShortTerm
}
Some(minutes) if window_minutes_match(minutes, WEEK_WINDOW_MINUTES) => {
QuotaWindowKind::Weekly
}
Some(minutes) if window_minutes_match(minutes, MONTH_WINDOW_MINUTES) => {
QuotaWindowKind::Monthly
}
_ => QuotaWindowKind::Other,
}
}
pub fn window_minutes_match(actual: u64, expected: u64) -> bool {
actual > 0 && actual.abs_diff(expected) <= expected / 20
}
fn quota_window_order(kind: QuotaWindowKind) -> u8 {
match kind {
QuotaWindowKind::Monthly => 0,
QuotaWindowKind::ShortTerm => 1,
QuotaWindowKind::Weekly => 2,
QuotaWindowKind::Other => 3,
}
}
fn quota_usage_priority(kind: QuotaWindowKind) -> u8 {
match kind {
QuotaWindowKind::ShortTerm => 0,
QuotaWindowKind::Weekly => 1,
QuotaWindowKind::Monthly => 2,
QuotaWindowKind::Other => 3,
}
}
fn quota_window_key_part(kind: QuotaWindowKind) -> Option<&'static str> {
match kind {
QuotaWindowKind::ShortTerm => Some("five-hour"),
QuotaWindowKind::Weekly => Some("weekly"),
QuotaWindowKind::Monthly => Some("monthly"),
QuotaWindowKind::Other => None,
}
}
fn quota_window_role_name(role: QuotaWindowRole) -> &'static str {
match role {
QuotaWindowRole::Primary => "primary",
QuotaWindowRole::Secondary => "secondary",
QuotaWindowRole::Monthly => "monthly",
}
}
fn quota_source_key(
source: &str,
limit_name: Option<&str>,
metered_feature: Option<&str>,
) -> String {
let label = limit_name.or(metered_feature).unwrap_or(source);
format!("{}-{}", quota_key_segment(source), quota_key_segment(label))
}
fn quota_key_segment(value: &str) -> String {
let mut segment = String::with_capacity(value.len());
for ch in value.chars() {
if ch.is_ascii_alphanumeric() {
segment.push(ch.to_ascii_lowercase());
} else if !segment.ends_with('-') {
segment.push('-');
}
}
let segment = segment.trim_matches('-');
if segment.is_empty() {
"quota".to_string()
} else {
segment.to_string()
}
}
pub fn remaining_percent(used_percent: f64) -> i64 {
(100.0 - clamp_percent(used_percent)).round() as i64
}
fn quota_spend_control_limit_reached(spend_control: &QuotaSpendControl) -> bool {
spend_control.reached
|| spend_control
.individual_limit
.as_ref()
.and_then(|limit| limit.used_percent)
.is_some_and(|used| used >= 100.0)
}
fn clamp_percent(value: f64) -> f64 {
value.clamp(0.0, 100.0)
}
fn deserialize_optional_finite_number<'de, D>(deserializer: D) -> Result<Option<f64>, D::Error>
where
D: Deserializer<'de>,
{
let value = Option::<Value>::deserialize(deserializer)?;
Ok(value.and_then(|value| {
value
.as_f64()
.or_else(|| value.as_str().and_then(|value| value.parse::<f64>().ok()))
.filter(|value| value.is_finite())
}))
}
fn deserialize_optional_positive_u64<'de, D>(deserializer: D) -> Result<Option<u64>, D::Error>
where
D: Deserializer<'de>,
{
let value = Option::<Value>::deserialize(deserializer)?;
Ok(value
.and_then(|value| {
value
.as_u64()
.or_else(|| value.as_i64().and_then(|value| u64::try_from(value).ok()))
.or_else(|| value.as_str().and_then(|value| value.parse::<u64>().ok()))
})
.filter(|value| *value > 0))
}
fn deserialize_optional_positive_i64<'de, D>(deserializer: D) -> Result<Option<i64>, D::Error>
where
D: Deserializer<'de>,
{
let value = Option::<Value>::deserialize(deserializer)?;
Ok(value
.and_then(|value| {
value
.as_i64()
.or_else(|| value.as_u64().and_then(|value| i64::try_from(value).ok()))
.or_else(|| value.as_str().and_then(|value| value.parse::<i64>().ok()))
})
.filter(|value| *value > 0))
}
+238
View File
@@ -0,0 +1,238 @@
//! 上游 adapter 交给 fleet 的 typed quota observation。
use std::collections::BTreeMap;
use serde_json::Value;
use super::{
MONTH_WINDOW_MINUTES, MonthlyQuotaLimit, QuotaCredits, QuotaLimitSnapshot, QuotaSnapshot,
QuotaSpendControl, QuotaWindow, QuotaWindowKind, remaining_percent,
};
#[derive(Debug, Clone, PartialEq)]
pub struct QuotaWindowObservation {
pub used_percent: f64,
pub window_minutes: Option<u64>,
pub reset_at: Option<i64>,
pub used: Option<Value>,
pub limit: Option<Value>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct QuotaLimitObservation {
pub limit_id: String,
pub limit_name: Option<String>,
pub metered_feature: Option<String>,
pub allowed: Option<bool>,
pub limit_reached: Option<bool>,
pub primary: Option<QuotaWindowObservation>,
pub secondary: Option<QuotaWindowObservation>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct QuotaCreditsObservation {
pub has_credits: bool,
pub unlimited: bool,
pub overage_limit_reached: bool,
pub balance: Option<Value>,
}
#[derive(Debug, Clone, Default, PartialEq)]
pub struct QuotaObservation {
pub limits: BTreeMap<String, QuotaLimitObservation>,
pub active_limit: Option<String>,
pub credits: Option<QuotaCreditsObservation>,
pub spend_control: Option<QuotaSpendControl>,
pub plan_type: Option<String>,
pub promo_message: Option<String>,
pub rate_limit_reached_type: Option<String>,
}
/// 将 typed observation 合并为代理内部 quota 快照。
pub fn quota_from_observation(
observation: &QuotaObservation,
plan_type: Option<&str>,
existing: Option<&QuotaSnapshot>,
) -> QuotaSnapshot {
let mut limits = observation
.limits
.values()
.filter_map(quota_limit_from_observation)
.collect::<Vec<_>>();
if let Some(existing) = existing {
for limit in &existing.limits {
if !observation
.limits
.keys()
.any(|limit_id| quota_limit_matches(limit, limit_id))
{
limits.push(limit.clone());
}
}
}
let spend_control = observation
.spend_control
.clone()
.or_else(|| existing.and_then(|quota| quota.spend_control.clone()));
let monthly_limit = spend_control
.as_ref()
.and_then(monthly_limit_from_spend_control)
.or_else(|| monthly_limit_from_limits(&limits))
.or_else(|| existing.and_then(|quota| quota.monthly_limit.clone()));
let credits = observation
.credits
.as_ref()
.map(|credits| QuotaCredits {
has_credits: credits.has_credits,
unlimited: credits.unlimited,
overage_limit_reached: credits.overage_limit_reached,
balance: credits.balance.clone(),
})
.or_else(|| existing.and_then(|quota| quota.credits.clone()));
QuotaSnapshot {
plan_type: observation
.plan_type
.clone()
.or_else(|| plan_type.map(ToString::to_string))
.or_else(|| existing.and_then(|quota| quota.plan_type.clone()))
.or_else(|| Some("unknown".to_string())),
active_limit: observation
.active_limit
.clone()
.or_else(|| existing.and_then(|quota| quota.active_limit.clone())),
limits,
monthly_limit,
credits,
spend_control,
promo_message: observation
.promo_message
.clone()
.or_else(|| existing.and_then(|quota| quota.promo_message.clone())),
rate_limit_reached_type: observation
.rate_limit_reached_type
.clone()
.or_else(|| existing.and_then(|quota| quota.rate_limit_reached_type.clone())),
extensions: existing
.map(|quota| quota.extensions.clone())
.unwrap_or_default(),
}
}
fn quota_limit_from_observation(details: &QuotaLimitObservation) -> Option<QuotaLimitSnapshot> {
let (source, metered_feature) = quota_identity(&details.limit_id);
let limit_name = details
.limit_name
.clone()
.or_else(|| (source != "core").then(|| details.limit_id.clone()));
let primary = details.primary.as_ref().map(quota_window_from_observation);
let secondary = details
.secondary
.as_ref()
.map(quota_window_from_observation);
if primary.is_none() && secondary.is_none() {
return None;
}
let blocked = details.limit_reached.unwrap_or(false)
|| details.allowed.is_some_and(|allowed| !allowed)
|| primary.as_ref().is_some_and(QuotaWindow::is_limit_reached)
|| secondary
.as_ref()
.is_some_and(QuotaWindow::is_limit_reached);
Some(QuotaLimitSnapshot {
source: source.to_string(),
limit_name,
metered_feature: details
.metered_feature
.clone()
.or_else(|| metered_feature.map(ToString::to_string)),
allowed: details.allowed,
limit_reached: details.limit_reached,
blocked,
primary,
secondary,
})
}
fn quota_window_from_observation(window: &QuotaWindowObservation) -> QuotaWindow {
QuotaWindow::from_usage(
window.used_percent,
window.reset_at,
window.window_minutes,
window.used.clone(),
window.limit.clone(),
)
}
fn quota_identity(limit_id: &str) -> (&'static str, Option<&str>) {
if limit_id == "codex" {
("core", None)
} else if is_review_limit_name(limit_id) {
("code_review", Some(limit_id))
} else {
("additional", Some(limit_id))
}
}
fn quota_limit_matches(limit: &QuotaLimitSnapshot, limit_id: &str) -> bool {
let (source, metered_feature) = quota_identity(limit_id);
if limit.source != source {
return false;
}
source == "core"
|| limit.metered_feature.as_deref() == metered_feature
|| limit.limit_name.as_deref() == Some(limit_id)
}
fn monthly_limit_from_limits(limits: &[QuotaLimitSnapshot]) -> Option<MonthlyQuotaLimit> {
let window = limits
.iter()
.filter(|limit| limit.source == "core")
.flat_map(QuotaLimitSnapshot::windows)
.find(|window| window.kind == QuotaWindowKind::Monthly)?;
Some(MonthlyQuotaLimit {
key: Some("core-monthly".to_string()),
source: Some("rate_limit".to_string()),
used_percent: window.used_percent,
remaining_percent: window.remaining_percent.map(|value| value as f64),
reset_at: window.reset_at,
window_minutes: window.window_minutes,
limit_reached: window.is_limit_reached(),
used: window.used.clone(),
limit: window.limit.clone(),
used_credits: window.used_credits.clone(),
limit_credits: window.limit_credits.clone(),
})
}
fn monthly_limit_from_spend_control(
spend_control: &QuotaSpendControl,
) -> Option<MonthlyQuotaLimit> {
let individual = spend_control.individual_limit.as_ref()?;
let used_percent = individual.used_percent.unwrap_or(0.0).clamp(0.0, 100.0);
Some(MonthlyQuotaLimit {
key: Some("spend-control-monthly".to_string()),
source: Some("spend_control".to_string()),
used_percent: Some(used_percent),
remaining_percent: individual
.remaining_percent
.or_else(|| Some(remaining_percent(used_percent) as f64)),
reset_at: individual.reset_at,
window_minutes: Some(MONTH_WINDOW_MINUTES),
limit_reached: spend_control.reached || used_percent >= 100.0,
used: individual.used.clone(),
limit: individual.limit.clone(),
used_credits: individual.used.clone(),
limit_credits: individual.limit.clone(),
})
}
fn is_review_limit_name(value: &str) -> bool {
let normalized = value.trim().to_ascii_lowercase().replace('-', "_");
matches!(
normalized.as_str(),
"review" | "code_review" | "codex_review" | "codex_code_review"
)
}
+44 -55
View File
@@ -7,26 +7,22 @@ use std::{
};
use chrono::Utc;
use secrecy::SecretString;
use thiserror::Error;
use tokio::time::{Instant, sleep};
use tracing::warn;
use crate::fleet::cookies::PgCookieStore;
use crate::fleet::{
account::{Account, AccountStatus},
account_failure::apply_client_failure_effect,
account_failure::{apply_account_state_effect, classify_account_failure},
account_gateway::{AccountGatewayError, AccountUpstreamContext, AccountUpstreamGateway},
pool::AccountPoolService,
quota::{
quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds,
quota_snapshot_reset_at,
quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds, quota_snapshot_reset_at,
},
store::{AccountStore, AccountStoreError, PgAccountStore, PgAccountStoreError},
};
use crate::{
fleet::cookies::PgCookieStore,
telemetry::account_usage::store::{
AccountUsageStore, AccountUsageStoreError, PgAccountUsageStore,
},
upstream::openai::transport::{CodexBackendClient, CodexClientError, CodexRequestContext},
store::{AccountStore, AccountStoreError},
usage::{AccountUsageStore, AccountUsageStoreError},
};
const MIN_REFRESH_INTERVAL_SECS: u64 = 30 * 60;
@@ -35,9 +31,9 @@ const COOLDOWN_REFRESH_GRACE_SECS: i64 = 30;
/// 运行时 quota refresh 服务。
pub struct QuotaRefreshService {
store: PgAccountStore,
usage_store: PgAccountUsageStore,
codex: Arc<CodexBackendClient>,
store: Arc<dyn AccountStore>,
usage_store: Arc<dyn AccountUsageStore>,
upstream: Arc<dyn AccountUpstreamGateway>,
min_refresh_interval_secs: u64,
request_spacing: Duration,
account_pseudonymizer: Arc<crate::infra::identity::AccountPseudonymizer>,
@@ -48,16 +44,16 @@ pub struct QuotaRefreshService {
impl QuotaRefreshService {
/// 构造默认 quota refresh 服务。
pub fn new(
store: PgAccountStore,
usage_store: PgAccountUsageStore,
codex: Arc<CodexBackendClient>,
store: Arc<dyn AccountStore>,
usage_store: Arc<dyn AccountUsageStore>,
upstream: Arc<dyn AccountUpstreamGateway>,
account_pool: Arc<AccountPoolService>,
account_pseudonymizer: Arc<crate::infra::identity::AccountPseudonymizer>,
) -> Self {
Self {
store,
usage_store,
codex,
upstream,
min_refresh_interval_secs: MIN_REFRESH_INTERVAL_SECS,
request_spacing: Self::default_request_spacing(),
account_pseudonymizer,
@@ -72,9 +68,9 @@ impl QuotaRefreshService {
/// 使用自定义最小刷新间隔构造 quota refresh 服务。
pub fn with_min_refresh_interval_secs(
store: PgAccountStore,
usage_store: PgAccountUsageStore,
codex: Arc<CodexBackendClient>,
store: Arc<dyn AccountStore>,
usage_store: Arc<dyn AccountUsageStore>,
upstream: Arc<dyn AccountUpstreamGateway>,
account_pool: Arc<AccountPoolService>,
account_pseudonymizer: Arc<crate::infra::identity::AccountPseudonymizer>,
min_refresh_interval_secs: u64,
@@ -82,7 +78,7 @@ impl QuotaRefreshService {
Self {
store,
usage_store,
codex,
upstream,
min_refresh_interval_secs,
request_spacing: Self::default_request_spacing(),
account_pseudonymizer,
@@ -153,40 +149,29 @@ impl QuotaRefreshService {
}
last_refreshed.insert(account.id.clone(), now);
let request_id = uuid::Uuid::new_v4().to_string();
let cookie_header = self.usage_cookie_header(&account.id).await;
let installation_id = self.account_pseudonymizer.installation_id(&account.id);
match self
.codex
.fetch_usage(CodexRequestContext {
access_token: &account.access_token,
account_id: account.account_id.as_deref(),
request_id: &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: cookie_header.as_deref(),
installation_id: Some(&installation_id),
session_id: None,
thread_id: None,
client_request_id: None,
turn_id: None,
.upstream
.fetch_usage(AccountUpstreamContext {
access_token: SecretString::new(account.access_token.clone().into()),
account_id: account.account_id.clone(),
request_id: uuid::Uuid::new_v4().to_string(),
cookie_header,
installation_id: Some(installation_id),
})
.await
{
Ok(raw) => {
let quota = quota_from_usage(&raw);
Ok(result) => {
let quota = result.quota;
let limit_reached = quota_snapshot_limit_reached(&quota);
let reset_at = quota_snapshot_reset_at(&quota);
let quota_json = quota.to_json()?;
if self
.store
.apply_quota_snapshot(
&account.id,
&quota.to_string(),
&quota_json,
limit_reached,
limit_reached.then_some(reset_at).flatten(),
)
@@ -208,13 +193,17 @@ impl QuotaRefreshService {
self.sync_runtime_pool_account(&account.id).await;
}
Err(error) => {
apply_client_failure_effect(
&self.account_pool,
&self.codex,
&account.id,
&error,
)
.await;
if let Some(classified) = error.failure().and_then(classify_account_failure)
&& let Some(effect) = &classified.effect
{
apply_account_state_effect(
&self.account_pool,
self.upstream.as_ref(),
&account.id,
effect,
)
.await;
}
warn!(
account_id = %account.id,
error = %error,
@@ -289,15 +278,15 @@ pub enum QuotaRefreshServiceError {
/// 账号读取失败。
#[error("failed to list accounts for quota refresh: {0}")]
AccountStore(#[from] AccountStoreError),
/// 账号写入失败。
#[error("failed to persist quota refresh result: {0}")]
Store(#[from] PgAccountStoreError),
/// 用量窗口写入失败。
#[error("failed to persist quota usage window: {0}")]
UsageStore(#[from] AccountUsageStoreError),
/// typed quota 序列化失败。
#[error("failed to serialize quota snapshot: {0}")]
SerializeQuota(#[from] serde_json::Error),
/// 上游 usage 请求失败。
#[error("failed to fetch quota usage: {0}")]
Codex(#[from] CodexClientError),
Upstream(#[from] AccountGatewayError),
}
/// quota refresh 服务结果。
+8 -403
View File
@@ -1,412 +1,17 @@
//! 配额检查策略。
use std::time::Duration;
//! 账号调度读取 typed quota 的稳定策略入口。
use chrono::{DateTime, Utc};
use serde_json::{Value, json};
const MONTH_WINDOW_MINUTES: u64 = 43_200;
const MONTH_WINDOW_SECONDS: u64 = 2_592_000;
use super::QuotaSnapshot;
/// 判断配额快照是否已经触顶。
pub fn quota_snapshot_limit_reached(quota: &Value) -> bool {
snapshot_blocked(quota)
|| spend_control_limit_reached(quota)
|| monthly_limit_reached(quota)
|| credits_overage_limit_reached(quota)
pub fn quota_snapshot_limit_reached(quota: &QuotaSnapshot) -> bool {
quota.is_limit_reached()
}
/// 从配额快照读取 UTC 重置时间。
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_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)
.or_else(|| monthly_limit_reset_at(quota))
}?;
DateTime::<Utc>::from_timestamp(reset_at, 0)
pub fn quota_snapshot_reset_at(quota: &QuotaSnapshot) -> Option<DateTime<Utc>> {
quota.reset_at()
}
/// 从配额快照读取窗口大小。
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));
}
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))
}
/// 从 Codex usage 响应中提取持久化配额快照。
pub fn quota_from_usage(usage: &Value) -> Value {
let mut snapshots = Vec::new();
if let Some(snapshot) = quota_snapshot_from_rate_limit(
"core",
usage.get("rate_limit_name").and_then(trimmed_str),
usage.get("metered_feature").and_then(trimmed_str),
usage.get("rate_limit"),
) {
snapshots.push(snapshot);
}
if let Some(additional) = usage
.get("additional_rate_limits")
.and_then(Value::as_array)
{
for item in additional {
let Some(limit_name) = item.get("limit_name").and_then(trimmed_str) else {
continue;
};
if let Some(snapshot) = quota_snapshot_from_rate_limit(
"additional",
Some(limit_name),
item.get("metered_feature").and_then(trimmed_str),
item.get("rate_limit"),
) {
snapshots.push(snapshot);
}
}
}
let monthly_limit = monthly_limit_from_spend_control(usage.get("spend_control"))
.or_else(|| monthly_limit_from_snapshots(&snapshots))
.unwrap_or(Value::Null);
json!({
"plan_type": usage.get("plan_type").cloned().unwrap_or(Value::Null),
"snapshots": snapshots,
"monthly_limit": monthly_limit,
"credits": normalize_quota_credits(usage.get("credits")),
"spend_control": normalize_spend_control(usage.get("spend_control")),
})
}
fn quota_snapshot_from_rate_limit(
source: &str,
limit_name: Option<&str>,
metered_feature: Option<&str>,
rate_limit: Option<&Value>,
) -> Option<Value> {
let rate_limit = rate_limit.filter(|value| !value.is_null())?;
let primary = window_from_rate_limit(rate_limit.get("primary_window"));
let secondary = window_from_rate_limit(rate_limit.get("secondary_window"));
if primary.is_null() && secondary.is_null() {
return None;
}
Some(json!({
"source": source,
"limit_name": limit_name,
"metered_feature": metered_feature,
"allowed": rate_limit.get("allowed").cloned().unwrap_or(Value::Null),
"limit_reached": rate_limit.get("limit_reached").cloned().unwrap_or(Value::Null),
"blocked": rate_limit_blocked(rate_limit)
|| window_limit_reached(&primary)
|| window_limit_reached(&secondary),
"primary": primary,
"secondary": secondary,
}))
}
fn window_from_rate_limit(window: Option<&Value>) -> Value {
let Some(window) = window.filter(|value| !value.is_null()) else {
return Value::Null;
};
let used_percent = window
.get("used_percent")
.and_then(number_value)
.unwrap_or(0.0)
.clamp(0.0, 100.0);
let window_minutes = window
.get("limit_window_seconds")
.and_then(number_value)
.filter(|seconds| *seconds > 0.0)
.and_then(|seconds| Duration::try_from_secs_f64((seconds / 60.0).round()).ok())
.map(|duration| duration.as_secs());
json!({
"used_percent": used_percent,
"remaining_percent": remaining_percent(used_percent),
"reset_at": window.get("reset_at").and_then(positive_i64),
"window_minutes": window_minutes,
"limit_reached": used_percent >= 100.0,
"used": first_present(window, &["used_tokens", "used", "used_credits", "usage", "consumed_tokens", "consumed"]),
"limit": first_present(window, &["limit_tokens", "limit", "limit_credits", "quota", "total_tokens", "total"]),
})
}
fn normalize_quota_credits(raw: Option<&Value>) -> Value {
let Some(raw) = raw.filter(|value| !value.is_null()) else {
return Value::Null;
};
json!({
"has_credits": raw.get("has_credits").and_then(Value::as_bool).unwrap_or(false),
"unlimited": raw.get("unlimited").and_then(Value::as_bool).unwrap_or(false),
"overage_limit_reached": raw.get("overage_limit_reached").and_then(Value::as_bool).unwrap_or(false),
"balance": raw.get("balance").cloned().unwrap_or(Value::Null),
})
}
fn normalize_spend_control(raw: Option<&Value>) -> Value {
let Some(raw) = raw.filter(|value| !value.is_null()) else {
return Value::Null;
};
raw.clone()
}
fn monthly_limit_from_spend_control(spend_control: Option<&Value>) -> Option<Value> {
let individual_limit = spend_control?
.get("individual_limit")
.filter(|value| !value.is_null())?;
let used_percent = individual_limit
.get("used_percent")
.and_then(number_value)
.unwrap_or(0.0)
.clamp(0.0, 100.0);
Some(json!({
"key": "spend-control-monthly",
"source": "spend_control",
"used_percent": used_percent,
"remaining_percent": individual_limit
.get("remaining_percent")
.and_then(number_value)
.unwrap_or_else(|| remaining_percent(used_percent) as f64),
"reset_at": individual_limit.get("reset_at").and_then(positive_i64),
"window_minutes": MONTH_WINDOW_MINUTES,
"limit_reached": spend_control_limit_reached_value(spend_control.unwrap_or(&Value::Null), used_percent),
"used_credits": individual_limit.get("used").cloned().unwrap_or(Value::Null),
"limit_credits": individual_limit.get("limit").cloned().unwrap_or(Value::Null),
}))
}
fn monthly_limit_from_snapshots(snapshots: &[Value]) -> Option<Value> {
for snapshot in snapshots {
if snapshot.get("source").and_then(Value::as_str) != Some("core") {
continue;
}
for key in ["primary", "secondary"] {
let Some(bucket) = snapshot.get(key).filter(|value| !value.is_null()) else {
continue;
};
if bucket
.get("window_minutes")
.and_then(Value::as_u64)
.is_some_and(|minutes| window_minutes_match(minutes, MONTH_WINDOW_MINUTES))
{
return Some(json!({
"key": "core-monthly",
"source": "rate_limit",
"used_percent": bucket.get("used_percent").cloned().unwrap_or(Value::Null),
"remaining_percent": bucket.get("remaining_percent").cloned().unwrap_or(Value::Null),
"reset_at": bucket.get("reset_at").cloned().unwrap_or(Value::Null),
"window_minutes": bucket.get("window_minutes").cloned().unwrap_or(Value::Null),
"limit_reached": bucket.get("limit_reached").cloned().unwrap_or(Value::Bool(false)),
"used": bucket.get("used").cloned().unwrap_or(Value::Null),
"limit": bucket.get("limit").cloned().unwrap_or(Value::Null),
"used_credits": bucket.get("used_credits").cloned().unwrap_or(Value::Null),
"limit_credits": bucket.get("limit_credits").cloned().unwrap_or(Value::Null),
}));
}
}
}
None
}
fn remaining_percent(used_percent: f64) -> i64 {
(100.0 - used_percent.clamp(0.0, 100.0)).round() as i64
}
fn first_present(value: &Value, keys: &[&str]) -> Value {
keys.iter()
.find_map(|key| value.get(*key).filter(|value| !value.is_null()).cloned())
.unwrap_or(Value::Null)
}
fn rate_limit_blocked(rate_limit: &Value) -> bool {
rate_limit
.get("limit_reached")
.and_then(Value::as_bool)
.unwrap_or(false)
|| rate_limit
.get("allowed")
.and_then(Value::as_bool)
.is_some_and(|allowed| !allowed)
}
fn window_limit_reached(window: &Value) -> bool {
window
.get("limit_reached")
.and_then(Value::as_bool)
.unwrap_or(false)
}
fn spend_control_limit_reached(quota: &Value) -> bool {
quota
.pointer("/spend_control/reached")
.and_then(Value::as_bool)
.unwrap_or(false)
|| quota
.pointer("/spend_control/individual_limit/used_percent")
.and_then(number_value)
.is_some_and(|used_percent| used_percent >= 100.0)
}
fn spend_control_limit_reached_value(spend_control: &Value, used_percent: f64) -> bool {
spend_control
.get("reached")
.and_then(Value::as_bool)
.unwrap_or(false)
|| used_percent >= 100.0
}
fn monthly_limit_reached(quota: &Value) -> bool {
quota
.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 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(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 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> {
core_buckets(quota).find(|bucket| bucket_reset_at(bucket).is_some())
}
fn core_bucket_with_window_minutes(quota: &Value) -> Option<&Value> {
core_buckets(quota).find(|bucket| bucket_window_minutes(bucket).is_some())
}
fn core_buckets<'a>(quota: &'a Value) -> impl Iterator<Item = &'a Value> + 'a {
quota
.get("snapshots")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter(|snapshot| snapshot.get("source").and_then(Value::as_str) == Some("core"))
.flat_map(snapshot_buckets)
}
fn bucket_reset_at(bucket: &Value) -> Option<i64> {
bucket.get("reset_at").and_then(positive_i64)
}
fn bucket_window_minutes(bucket: &Value) -> Option<u64> {
bucket
.get("window_minutes")
.and_then(Value::as_u64)
.filter(|value| *value > 0)
}
fn credits_overage_limit_reached(quota: &Value) -> bool {
quota
.pointer("/credits/overage_limit_reached")
.and_then(Value::as_bool)
.unwrap_or(false)
}
fn monthly_limit_reset_at(quota: &Value) -> Option<i64> {
quota
.pointer("/monthly_limit/reset_at")
.and_then(positive_i64)
}
fn monthly_limit_window_seconds(quota: &Value) -> Option<u64> {
quota
.pointer("/monthly_limit/window_minutes")
.and_then(Value::as_u64)
.filter(|value| *value > 0)
.and_then(|minutes| minutes.checked_mul(60))
}
fn number_value(value: &Value) -> Option<f64> {
value
.as_f64()
.or_else(|| value.as_str().and_then(|value| value.parse::<f64>().ok()))
.filter(|value| value.is_finite())
}
fn positive_i64(value: &Value) -> Option<i64> {
value
.as_i64()
.filter(|value| *value > 0)
.or_else(|| value.as_u64().and_then(|value| i64::try_from(value).ok()))
.filter(|value| *value > 0)
}
fn trimmed_str(value: &Value) -> Option<&str> {
value
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn window_minutes_match(actual: u64, expected: u64) -> bool {
actual > 0 && actual.abs_diff(expected) <= expected / 20
pub fn quota_snapshot_limit_window_seconds(quota: &QuotaSnapshot) -> Option<u64> {
quota.limit_window_seconds()
}
+26
View File
@@ -0,0 +1,26 @@
//! token 刷新用例拥有的上游端口与稳定值。
use async_trait::async_trait;
#[derive(Debug, Clone)]
pub struct TokenPair {
pub access_token: String,
pub refresh_token: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum RefreshFailure {
#[error("refresh token is invalid or expired")]
InvalidGrant,
#[error("account is banned")]
Banned,
#[error("refresh transport failed before server processing")]
RetryableTransport,
#[error("refresh transport failed after possible server processing")]
Transport,
}
#[async_trait]
pub trait TokenRefresher: Send + Sync + 'static {
async fn refresh(&self, refresh_token: &str) -> Result<TokenPair, RefreshFailure>;
}
+2
View File
@@ -1,9 +1,11 @@
//! Token 刷新策略、租约与运行时服务。
mod gateway;
mod lease;
mod policy;
mod service;
pub use gateway::*;
pub use lease::*;
pub use policy::*;
pub use service::*;
+3 -10
View File
@@ -12,8 +12,8 @@ use chrono::{DateTime, Duration, TimeZone, Utc};
use serde_json::{Map, Value};
use tokio::sync::{Semaphore, watch};
use super::{RefreshFailure, TokenPair, TokenRefresher};
use crate::fleet::account::{Account, AccountStatus};
use crate::upstream::openai::token_client::{RefreshFailure, TokenPair, TokenRefresher};
pub(crate) const PERMANENT_FAILURE_CONFIRMATION_THRESHOLD: usize = 2;
pub(crate) const RECOVERY_DELAY_SECONDS: i64 = 10 * 60;
@@ -238,16 +238,9 @@ impl RuntimeRefreshPolicy {
}
/// 持续接收运行时设置并更新刷新策略。
pub async fn subscribe_settings(
self,
mut receiver: watch::Receiver<crate::settings::SettingsSnapshot>,
) {
pub async fn subscribe_policy(self, mut receiver: watch::Receiver<RefreshPolicy>) {
while receiver.changed().await.is_ok() {
let settings = receiver.borrow_and_update();
self.update(RefreshPolicy {
refresh_margin_seconds: settings.refresh_margin_seconds,
refresh_concurrency: settings.refresh_concurrency,
});
self.update(*receiver.borrow_and_update());
}
}
+27 -1
View File
@@ -19,13 +19,14 @@ use tokio_util::task::TaskTracker;
use tracing::{debug, info, warn};
use uuid::Uuid;
use super::TokenRefresher;
use crate::fleet::{
account::{Account, AccountStatus},
pool::AccountPoolService,
store::{
AccountClaimsUpdate, AccountStore, PgAccountStore, PgAccountStoreError, StoredAccount,
},
};
use crate::upstream::openai::token_client::TokenRefresher;
use super::{
lease::{RedisRefreshLeaseStore, RedisRefreshLeaseStoreError},
@@ -48,6 +49,7 @@ where
C: TokenRefresher,
{
store: PgAccountStore,
account_pool: Option<Arc<AccountPoolService>>,
scheduler: Arc<RefreshScheduler<C>>,
policy: RuntimeRefreshPolicy,
refresh_leases: Option<RedisRefreshLeaseStore>,
@@ -193,6 +195,7 @@ where
fn clone(&self) -> Self {
Self {
store: self.store.clone(),
account_pool: self.account_pool.clone(),
scheduler: self.scheduler.clone(),
policy: self.policy.clone(),
refresh_leases: self.refresh_leases.clone(),
@@ -215,6 +218,7 @@ where
let policy = policy.into();
Self {
store,
account_pool: None,
scheduler: Arc::new(RefreshScheduler::new(policy.clone(), client)),
policy,
refresh_leases: None,
@@ -233,6 +237,12 @@ where
self
}
/// 刷新后同步持久化账号到运行时池,避免旧 token 快照覆盖新状态。
pub fn with_account_pool_sync(mut self, account_pool: Arc<AccountPoolService>) -> Self {
self.account_pool = Some(account_pool);
self
}
/// 更新刷新策略。
pub fn update_policy(&self, policy: RefreshPolicy) {
self.policy.update(policy);
@@ -587,6 +597,7 @@ where
if attempt.refresh_token.is_none() {
persist_status(&self.store, &attempt.id, AccountStatus::Expired).await?;
persist_next_refresh_at(&self.store, &attempt.id, None).await?;
self.sync_account_pool_after_update(&attempt.id).await;
return Ok(TokenRefreshOutcome::StatusUpdated);
}
@@ -601,6 +612,7 @@ where
persist_status(&self.store, &attempt.id, updated.status).await?;
persist_next_refresh_at(&self.store, &attempt.id, None).await?;
self.sync_account_pool_after_update(&attempt.id).await;
return Ok(TokenRefreshOutcome::StatusUpdated);
}
Ok(mut updated) => {
@@ -609,6 +621,7 @@ where
updated.next_refresh_at = self.computed_refresh_at(&updated);
let refresh_token_rotated = updated.refresh_token != attempt.refresh_token;
persist_token_update(&self.store, &updated).await?;
self.sync_account_pool_after_update(&updated.id).await;
info!(
account_id = %updated.id,
attempt = attempt_index + 1,
@@ -643,6 +656,19 @@ where
Ok(TokenRefreshOutcome::Skipped)
}
async fn sync_account_pool_after_update(&self, account_id: &str) {
let Some(account_pool) = self.account_pool.as_ref() else {
return;
};
if let Err(error) = account_pool.sync_account_from_store(account_id).await {
warn!(
account_id,
error = %error,
"Failed to sync runtime account pool after token refresh"
);
}
}
async fn prepare_refresh_attempt(
&self,
account: &Account,
+68 -23
View File
@@ -1,7 +1,7 @@
//! 账号运行时反馈:错误率与 TTFT 的 EWMA 统计。
//! 账号运行时 feedback:首字前失败率、首字后中断率与 TTFT 的 EWMA 统计。
//!
//! 每个账号维护两个指数加权移动平均(EWMA):错误率与首个有效输出事件延迟(TTFT,
//! 毫秒)。请求结束后经 [`FeedbackStats::report`] 回灌,直接进 [`super::score`]
//! 每个账号维护三个指数加权移动平均(EWMA)。每个已启动 attempt 结束后经
//! [`FeedbackStats::report_attempt`] 回灌,直接进 [`super::score`]
//! 打分并可驱动 sticky escape。统计以 `AtomicU64`(存 `f64` 位模式)+ CAS 无锁更新,
//! 对齐 sub2api 的设计,避免选择热路径上的锁与 DB 往返。
@@ -12,13 +12,32 @@ use std::sync::atomic::{AtomicU64, Ordering};
/// EWMA 平滑系数(新样本权重)。与 sub2api 一致取 0.2。
const EWMA_ALPHA: f64 = 0.2;
/// 单个上游 attempt 对账号调度器的反馈。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AttemptFeedback {
Completed { first_token_ms: Option<u64> },
Incomplete { first_token_ms: Option<u64> },
FailedBeforeFirstToken,
FailedAfterFirstToken { first_token_ms: u64 },
Cancelled,
Shutdown,
}
impl AttemptFeedback {
fn affects_account_health(self) -> bool {
!matches!(self, Self::Cancelled | Self::Shutdown)
}
}
/// 单账号的运行时 EWMA 统计。
///
/// `*_bits` 存 `f64` 的位模式;`NaN`(全 1 的哨兵)表示"尚无样本"。
#[derive(Debug)]
struct AccountFeedback {
/// 错误率 EWMA,范围 `[0, 1]`。
error_rate_bits: AtomicU64,
/// 首字前失败率 EWMA,范围 `[0, 1]`。
pre_token_failure_rate_bits: AtomicU64,
/// 首字后中断率 EWMA,范围 `[0, 1]`。
post_token_abort_rate_bits: AtomicU64,
/// TTFT EWMA,单位毫秒。
ttft_ms_bits: AtomicU64,
}
@@ -29,16 +48,22 @@ const EMPTY_BITS: u64 = f64::NAN.to_bits();
impl Default for AccountFeedback {
fn default() -> Self {
Self {
error_rate_bits: AtomicU64::new(EMPTY_BITS),
pre_token_failure_rate_bits: AtomicU64::new(EMPTY_BITS),
post_token_abort_rate_bits: AtomicU64::new(EMPTY_BITS),
ttft_ms_bits: AtomicU64::new(EMPTY_BITS),
}
}
}
impl AccountFeedback {
/// 读取当前错误率 EWMA;无样本时返回 `None`。
fn error_rate(&self) -> Option<f64> {
load_sample(&self.error_rate_bits)
/// 读取当前首字前失败率 EWMA;无样本时返回 `None`。
fn pre_token_failure_rate(&self) -> Option<f64> {
load_sample(&self.pre_token_failure_rate_bits)
}
/// 读取当前首字后中断率 EWMA;无样本时返回 `None`。
fn post_token_abort_rate(&self) -> Option<f64> {
load_sample(&self.post_token_abort_rate_bits)
}
/// 读取当前 TTFT EWMA(毫秒);无样本时返回 `None`。
@@ -46,12 +71,26 @@ impl AccountFeedback {
load_sample(&self.ttft_ms_bits)
}
/// 回灌一次请求结果。
fn report(&self, success: bool, first_token_ms: Option<u64>) {
let error_sample = if success { 0.0 } else { 1.0 };
update_ewma(&self.error_rate_bits, error_sample);
if let Some(ms) = first_token_ms {
update_ewma(&self.ttft_ms_bits, ms as f64);
/// 回灌一次已经开始的 attempt。
fn report_attempt(&self, feedback: AttemptFeedback) {
match feedback {
AttemptFeedback::Completed { first_token_ms }
| AttemptFeedback::Incomplete { first_token_ms } => {
update_ewma(&self.pre_token_failure_rate_bits, 0.0);
update_ewma(&self.post_token_abort_rate_bits, 0.0);
if let Some(ms) = first_token_ms {
update_ewma(&self.ttft_ms_bits, ms as f64);
}
}
AttemptFeedback::FailedBeforeFirstToken => {
update_ewma(&self.pre_token_failure_rate_bits, 1.0);
}
AttemptFeedback::FailedAfterFirstToken { first_token_ms } => {
update_ewma(&self.pre_token_failure_rate_bits, 0.0);
update_ewma(&self.post_token_abort_rate_bits, 1.0);
update_ewma(&self.ttft_ms_bits, first_token_ms as f64);
}
AttemptFeedback::Cancelled | AttemptFeedback::Shutdown => {}
}
}
}
@@ -87,8 +126,10 @@ fn update_ewma(bits: &AtomicU64, sample: f64) {
/// 账号 EWMA 反馈的读数快照,供打分使用。
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct FeedbackSample {
/// 错误率 EWMA,范围 `[0, 1]`;无样本时为 `None`。
pub error_rate: Option<f64>,
/// 首字前失败率 EWMA,范围 `[0, 1]`;无样本时为 `None`。
pub pre_token_failure_rate: Option<f64>,
/// 首字后中断率 EWMA,范围 `[0, 1]`;无样本时为 `None`。
pub post_token_abort_rate: Option<f64>,
/// TTFT EWMA(毫秒);无样本时为 `None`。
pub ttft_ms: Option<f64>,
}
@@ -116,23 +157,27 @@ impl FeedbackStats {
.unwrap_or_else(std::sync::PoisonError::into_inner);
match accounts.get(account_id) {
Some(feedback) => FeedbackSample {
error_rate: feedback.error_rate(),
pre_token_failure_rate: feedback.pre_token_failure_rate(),
post_token_abort_rate: feedback.post_token_abort_rate(),
ttft_ms: feedback.ttft_ms(),
},
None => FeedbackSample::default(),
}
}
/// 回灌一次请求结果:成功/失败 + 首个有效输出事件延迟(毫秒)。
pub fn report(&self, account_id: &str, success: bool, first_token_ms: Option<u64>) {
/// 回灌一个已启动上游 attempt 的结果。
pub fn report_attempt(&self, account_id: &str, sample: AttemptFeedback) {
if !sample.affects_account_health() {
return;
}
// 快路径:账号已存在时只需读锁 + 原子更新。
{
let accounts = self
.accounts
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(feedback) = accounts.get(account_id) {
feedback.report(success, first_token_ms);
if let Some(account_feedback) = accounts.get(account_id) {
account_feedback.report_attempt(sample);
return;
}
}
@@ -144,7 +189,7 @@ impl FeedbackStats {
accounts
.entry(account_id.to_string())
.or_default()
.report(success, first_token_ms);
.report_attempt(sample);
}
/// 移除指定账号的反馈记录(账号被删除时调用)。
+5 -5
View File
@@ -23,7 +23,7 @@ use crate::fleet::account::Account;
use crate::infra::time::parse_rfc3339_utc;
pub use candidates::{CandidateFilter, CandidateRequest};
pub use feedback::{FeedbackSample, FeedbackStats};
pub use feedback::{AttemptFeedback, FeedbackSample, FeedbackStats};
pub use strategy::smart::rank_candidates;
pub use strategy::{ScoreBreakdown, ScoreWeights};
@@ -69,7 +69,7 @@ pub struct SelectionInput<'a> {
/// 账号调度器:账号选择的对外门面。
///
/// 持有轮转游标(并列破并列用)、打分权重(Smart 用)与 EWMA 反馈存储。选择热路径
/// 只读反馈与权重,请求结束后经 [`AccountScheduler::report_feedback`] 回灌。游标用
/// 只读反馈与权重,请求结束后经 [`AccountScheduler::report_attempt`] 回灌。游标用
/// [`Mutex`] 保护(选择需可变,但 `AccountScheduler` 以共享引用参与并发请求)。
#[derive(Debug)]
pub struct AccountScheduler {
@@ -150,9 +150,9 @@ impl AccountScheduler {
}
}
/// 回灌一次请求结果到运行时 EWMA 反馈(错误率 / TTFT),供 Smart 打分。
pub fn report_feedback(&self, account_id: &str, success: bool, first_token_ms: Option<u64>) {
self.feedback.report(account_id, success, first_token_ms);
/// 回灌一个已启动上游 attempt 的运行时反馈,供 Smart 打分。
pub fn report_attempt(&self, account_id: &str, feedback: AttemptFeedback) {
self.feedback.report_attempt(account_id, feedback);
}
/// 移除指定账号的运行时反馈(账号被删除时调用)。
@@ -284,7 +284,7 @@ fn score_account(
);
// 错误率 EWMA:本身即 [0,1],越低越健康。无样本时视为 0(满分健康)。
let error_rate = 1.0 - sample.error_rate.unwrap_or(0.0).clamp(0.0, 1.0);
let error_rate = 1.0 - sample.pre_token_failure_rate.unwrap_or(0.0).clamp(0.0, 1.0);
// TTFT:延迟越低越好。无样本的账号视为中性满分,避免冷启动被压。
let ttft = match sample.ttft_ms {
+1
View File
@@ -174,6 +174,7 @@ update accounts
set
status = case
when status in ('disabled', 'banned') then status
when $4 = 'expired' and access_token_expires_at > $2 then status
when (
case
when not $1 and quota_cooldown_until is not null and quota_cooldown_until > $2 then quota_limit_reached
+5 -7
View File
@@ -7,7 +7,7 @@ use sqlx::{PgPool, Postgres, QueryBuilder, Row, postgres::PgRow};
use crate::{
fleet::{
account::{Account, AccountStatus},
quota::{quota_snapshot_limit_window_seconds, quota_snapshot_reset_at},
quota::{QuotaSnapshot, quota_snapshot_limit_window_seconds, quota_snapshot_reset_at},
},
infra::format::nonnegative_i64_to_u64,
};
@@ -45,13 +45,11 @@ pub(super) async fn get_pool_account(
}
pub(super) fn pool_account_from_row(row: &PgRow) -> PgAccountStoreResult<Account> {
let quota_json = row
let quota = row
.get::<Option<sqlx::types::Json<Value>>, _>("quota_json")
.map(|quota_json| quota_json.0);
let quota_window_reset_at = quota_json.as_ref().and_then(quota_snapshot_reset_at);
let quota_limit_window_seconds = quota_json
.as_ref()
.and_then(quota_snapshot_limit_window_seconds);
.and_then(|quota_json| QuotaSnapshot::from_value(quota_json.0).ok());
let quota_window_reset_at = quota.as_ref().and_then(quota_snapshot_reset_at);
let quota_limit_window_seconds = quota.as_ref().and_then(quota_snapshot_limit_window_seconds);
Ok(Account {
id: row.get("id"),
+125
View File
@@ -0,0 +1,125 @@
//! 账号调度消费的用量值与持久化端口。
use std::collections::HashMap;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use thiserror::Error;
/// 账号用量端口错误。
#[derive(Debug, Error)]
#[error("account usage store operation failed: {message}")]
pub struct AccountUsageStoreError {
message: String,
}
impl AccountUsageStoreError {
/// 将 adapter 错误收窄为消费方稳定错误。
pub fn adapter(error: impl std::fmt::Display) -> Self {
Self {
message: error.to_string(),
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountUsageDelta {
pub requests: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub reasoning_tokens: u64,
pub total_tokens: u64,
pub empty_responses: u64,
pub image_input_tokens: u64,
pub image_output_tokens: u64,
pub image_requests: u64,
pub image_request_failures: u64,
}
/// fleet 记录一次 Responses 完成结果所需的稳定用量事实。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ResponseUsage {
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub reasoning_tokens: u64,
pub total_tokens: u64,
pub image_input_tokens: u64,
pub image_output_tokens: u64,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct AccountUsageSnapshot {
pub request_count: u64,
pub empty_response_count: u64,
pub image_input_tokens: u64,
pub image_output_tokens: u64,
pub image_request_count: u64,
pub image_request_failed_count: u64,
pub window_request_count: u64,
pub window_input_tokens: u64,
pub window_output_tokens: u64,
pub window_cached_tokens: u64,
pub window_image_input_tokens: u64,
pub window_image_output_tokens: u64,
pub window_image_request_count: u64,
pub window_image_request_failed_count: u64,
pub window_started_at: Option<DateTime<Utc>>,
pub window_reset_at: Option<DateTime<Utc>>,
pub limit_window_seconds: Option<u64>,
pub last_used_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountUsageWindow {
pub request_count: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub image_input_tokens: u64,
pub image_output_tokens: u64,
pub image_request_count: u64,
pub image_request_failed_count: u64,
pub started_at: Option<DateTime<Utc>>,
pub reset_at: Option<DateTime<Utc>>,
pub limit_window_seconds: Option<u64>,
}
#[async_trait]
pub trait AccountUsageStore: Send + Sync + 'static {
async fn snapshots(
&self,
account_ids: &[String],
) -> Result<HashMap<String, AccountUsageSnapshot>, AccountUsageStoreError>;
async fn record_usage_delta(
&self,
account_id: &str,
usage: AccountUsageDelta,
) -> Result<(), AccountUsageStoreError>;
async fn sync_runtime_window(
&self,
account_id: &str,
window: AccountUsageWindow,
) -> Result<(), AccountUsageStoreError>;
async fn sync_rate_limit_window(
&self,
account_id: &str,
reset_at: DateTime<Utc>,
limit_window_seconds: Option<u64>,
) -> Result<(), AccountUsageStoreError>;
async fn record_request(&self, account_id: &str) -> Result<(), AccountUsageStoreError> {
self.record_usage_delta(
account_id,
AccountUsageDelta {
requests: 1,
..AccountUsageDelta::default()
},
)
.await
}
}
+30 -1
View File
@@ -1,3 +1,32 @@
//! 账号配额窗口规则。
pub(crate) use crate::telemetry::account_usage::should_reset_usage_window;
use chrono::{DateTime, Utc};
/// 根据额度窗口边界漂移判断是否重置本地窗口计数。
pub(crate) fn should_reset_usage_window(
existing_reset_at: Option<DateTime<Utc>>,
existing_limit_window_seconds: Option<u64>,
new_reset_at: DateTime<Utc>,
new_limit_window_seconds: Option<u64>,
) -> bool {
let Some(existing_reset_at) = existing_reset_at else {
return false;
};
if existing_reset_at == new_reset_at {
return false;
}
let drift = existing_reset_at
.signed_duration_since(new_reset_at)
.num_seconds()
.unsigned_abs();
let window_seconds = new_limit_window_seconds
.or(existing_limit_window_seconds)
.unwrap_or(0);
let threshold = if window_seconds > 0 {
window_seconds / 2
} else {
3_600
};
drift >= threshold
}
+12 -5
View File
@@ -8,11 +8,18 @@ use tracing::info;
const MIGRATION_LOCK_KEY: i64 = i64::from_be_bytes(*b"CPR-MIGR");
const MIGRATIONS: &[Migration] = &[Migration {
version: 1,
name: "initial",
sql: include_str!("migrations/0001_initial.sql"),
}];
const MIGRATIONS: &[Migration] = &[
Migration {
version: 1,
name: "initial",
sql: include_str!("migrations/0001_initial.sql"),
},
Migration {
version: 2,
name: "cache_write_tokens",
sql: include_str!("migrations/0002_cache_write_tokens.sql"),
},
];
struct Migration {
version: i64,
@@ -0,0 +1,7 @@
alter table usage_records
add column cache_write_tokens bigint
check (cache_write_tokens is null or cache_write_tokens >= 0);
alter table request_time_buckets
add column cache_write_tokens bigint not null default 0
check (cache_write_tokens >= 0);
+36
View File
@@ -0,0 +1,36 @@
//! 模型目录用例拥有的上游端口。
use async_trait::async_trait;
use super::types::BackendModelEntry;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ModelCatalogRequest<'a> {
pub access_token: &'a str,
pub account_id: Option<&'a str>,
pub request_id: &'a str,
pub installation_id: Option<&'a str>,
pub plan_type: &'a str,
}
#[derive(Debug, thiserror::Error)]
#[error("model catalog request failed: {message}")]
pub struct ModelCatalogSourceError {
message: String,
}
impl ModelCatalogSourceError {
pub fn new(message: impl Into<String>) -> Self {
Self {
message: message.into(),
}
}
}
#[async_trait]
pub trait ModelCatalogSource: Send + Sync + 'static {
async fn fetch_models(
&self,
request: &ModelCatalogRequest<'_>,
) -> Result<Vec<BackendModelEntry>, ModelCatalogSourceError>;
}
+1
View File
@@ -1,6 +1,7 @@
//! 模型目录领域。
pub mod catalog;
pub mod gateway;
pub mod service;
pub mod store;
pub mod types;
+32 -18
View File
@@ -8,13 +8,9 @@ use std::{
use thiserror::Error;
use tokio::sync::watch;
use crate::{
settings::SettingsSnapshot,
upstream::openai::transport::{CodexModelCatalogClient, CodexModelCatalogRequest},
};
use super::{
catalog::ModelCatalog,
gateway::{ModelCatalogRequest, ModelCatalogSource},
store::{ModelSnapshotStore, ModelSnapshotStoreError},
types::{ModelConfig, ModelPlanSnapshot},
};
@@ -91,7 +87,7 @@ pub struct ModelService {
snapshots: Arc<RwLock<Vec<ModelPlanSnapshot>>>,
catalog: Arc<RwLock<ModelCatalog>>,
store: Option<Arc<dyn ModelSnapshotStore>>,
upstream_client: Option<Arc<dyn CodexModelCatalogClient>>,
upstream_client: Option<Arc<dyn ModelCatalogSource>>,
models_etag: Arc<RwLock<Option<String>>>,
etag_change_tx: watch::Sender<u64>,
refresh_lock: Arc<tokio::sync::Mutex<()>>,
@@ -102,7 +98,7 @@ impl ModelService {
pub fn new(
config: ModelConfig,
store: Option<Arc<dyn ModelSnapshotStore>>,
upstream_client: Option<Arc<dyn CodexModelCatalogClient>>,
upstream_client: Option<Arc<dyn ModelCatalogSource>>,
) -> Self {
let catalog = ModelCatalog::from_config(&config);
let (etag_change_tx, _) = watch::channel(0);
@@ -153,15 +149,9 @@ impl ModelService {
}
/// 持续接收运行时设置并更新模型别名。
pub async fn subscribe_settings(
self: Arc<Self>,
mut receiver: watch::Receiver<SettingsSnapshot>,
) {
pub async fn subscribe_config(self: Arc<Self>, mut receiver: watch::Receiver<ModelConfig>) {
while receiver.changed().await.is_ok() {
let settings = receiver.borrow_and_update().clone();
self.update_config(ModelConfig {
model_aliases: settings.model_aliases,
});
self.update_config(receiver.borrow_and_update().clone());
}
}
@@ -189,6 +179,26 @@ impl ModelService {
&self,
plan_accounts: &[ModelRefreshPlanAccount],
request_id: &str,
) -> Result<ModelRefreshResult, ModelServiceError> {
self.refresh_backend_models_inner(plan_accounts, request_id, false)
.await
}
/// 刷新指定计划,并保留未参与本次刷新的计划快照。
pub async fn refresh_selected_plan_models(
&self,
plan_accounts: &[ModelRefreshPlanAccount],
request_id: &str,
) -> Result<ModelRefreshResult, ModelServiceError> {
self.refresh_backend_models_inner(plan_accounts, request_id, true)
.await
}
async fn refresh_backend_models_inner(
&self,
plan_accounts: &[ModelRefreshPlanAccount],
request_id: &str,
preserve_other_plans: bool,
) -> Result<ModelRefreshResult, ModelServiceError> {
let _refresh_guard = self.refresh_lock.lock().await;
let store = self
@@ -211,7 +221,11 @@ impl ModelService {
.into_iter()
.map(|snapshot| (snapshot.plan_type.clone(), snapshot))
.collect::<BTreeMap<_, _>>();
let mut next_snapshots = BTreeMap::new();
let mut next_snapshots = if preserve_other_plans {
previous_snapshots.clone()
} else {
BTreeMap::new()
};
let mut result = ModelRefreshResult {
refreshed_plans: 0,
@@ -220,7 +234,7 @@ impl ModelService {
};
for plan_account in plan_accounts {
let request = CodexModelCatalogRequest {
let request = ModelCatalogRequest {
access_token: &plan_account.access_token,
account_id: plan_account.account_id.as_deref(),
request_id,
@@ -248,7 +262,7 @@ impl ModelService {
};
let snapshot =
ModelPlanSnapshot::from_backend_values(plan_account.plan_type.clone(), entries);
ModelPlanSnapshot::from_backend_entries(plan_account.plan_type.clone(), entries);
if snapshot.models.is_empty() {
result.failed_plans += 1;
if let Some(snapshot) = previous_snapshots.get(&plan_account.plan_type) {
@@ -1,35 +1,4 @@
//! 账号累计用量与时间桶查询。
use chrono::{DateTime, Utc};
pub mod query;
pub mod store;
/// 根据额度窗口边界漂移判断是否重置本地窗口计数。
pub fn should_reset_usage_window(
existing_reset_at: Option<DateTime<Utc>>,
existing_limit_window_seconds: Option<u64>,
new_reset_at: DateTime<Utc>,
new_limit_window_seconds: Option<u64>,
) -> bool {
let Some(existing_reset_at) = existing_reset_at else {
return false;
};
if existing_reset_at == new_reset_at {
return false;
}
let drift = existing_reset_at
.signed_duration_since(new_reset_at)
.num_seconds()
.unsigned_abs();
let window_seconds = new_limit_window_seconds
.or(existing_limit_window_seconds)
.unwrap_or(0);
let threshold = if window_seconds > 0 {
window_seconds / 2
} else {
3_600
};
drift >= threshold
}
@@ -55,6 +55,7 @@ impl AccountUsageQueryService {
nonnegative_i64_to_u64(bucket.totals.input_tokens),
nonnegative_i64_to_u64(bucket.totals.output_tokens),
nonnegative_i64_to_u64(bucket.totals.cached_tokens),
nonnegative_i64_to_u64(bucket.totals.cache_write_tokens),
&bucket.model,
bucket.service_tier.as_deref(),
);
@@ -159,6 +160,7 @@ pub struct AccountUsageTimeBucket {
pub input_tokens: i64,
pub output_tokens: i64,
pub cached_tokens: i64,
pub cache_write_tokens: i64,
pub first_token_latency_sum: i64,
pub first_token_latency_count: i64,
pub latency_sum: i64,
@@ -178,6 +180,7 @@ impl From<UsageTimeBucketRecord> for AccountUsageTimeBucket {
input_tokens: record.input_tokens,
output_tokens: record.output_tokens,
cached_tokens: record.cached_tokens,
cache_write_tokens: record.cache_write_tokens,
first_token_latency_sum: record.first_token_latency_sum,
first_token_latency_count: record.first_token_latency_count,
latency_sum: record.latency_sum,
+8 -101
View File
@@ -7,7 +7,13 @@ use chrono::{DateTime, Utc};
use sqlx::{PgPool, Postgres, QueryBuilder, Row, postgres::PgRow};
use thiserror::Error;
use crate::infra::format::nonnegative_i64_to_u64;
use crate::{
fleet::usage::{
AccountUsageDelta, AccountUsageSnapshot, AccountUsageStore, AccountUsageStoreError,
AccountUsageWindow,
},
infra::format::nonnegative_i64_to_u64,
};
const RECORD_USAGE_SQL: &str = r"
insert into account_usage (
@@ -140,108 +146,9 @@ pub enum PgAccountUsageStoreError {
/// PostgreSQL 用量存储结果。
pub type PgAccountUsageStoreResult<T> = Result<T, PgAccountUsageStoreError>;
/// 账号用量端口错误。
#[derive(Debug, Error)]
#[error("account usage store operation failed: {message}")]
pub struct AccountUsageStoreError {
message: String,
}
impl From<PgAccountUsageStoreError> for AccountUsageStoreError {
fn from(error: PgAccountUsageStoreError) -> Self {
Self {
message: error.to_string(),
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountUsageDelta {
pub requests: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub reasoning_tokens: u64,
pub total_tokens: u64,
pub empty_responses: u64,
pub image_input_tokens: u64,
pub image_output_tokens: u64,
pub image_requests: u64,
pub image_request_failures: u64,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct AccountUsageSnapshot {
pub request_count: u64,
pub empty_response_count: u64,
pub image_input_tokens: u64,
pub image_output_tokens: u64,
pub image_request_count: u64,
pub image_request_failed_count: u64,
pub window_request_count: u64,
pub window_input_tokens: u64,
pub window_output_tokens: u64,
pub window_cached_tokens: u64,
pub window_image_input_tokens: u64,
pub window_image_output_tokens: u64,
pub window_image_request_count: u64,
pub window_image_request_failed_count: u64,
pub window_started_at: Option<DateTime<Utc>>,
pub window_reset_at: Option<DateTime<Utc>>,
pub limit_window_seconds: Option<u64>,
pub last_used_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountUsageWindow {
pub request_count: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub image_input_tokens: u64,
pub image_output_tokens: u64,
pub image_request_count: u64,
pub image_request_failed_count: u64,
pub started_at: Option<DateTime<Utc>>,
pub reset_at: Option<DateTime<Utc>>,
pub limit_window_seconds: Option<u64>,
}
#[async_trait]
pub trait AccountUsageStore: Send + Sync + 'static {
async fn snapshots(
&self,
account_ids: &[String],
) -> Result<HashMap<String, AccountUsageSnapshot>, AccountUsageStoreError>;
async fn record_usage_delta(
&self,
account_id: &str,
usage: AccountUsageDelta,
) -> Result<(), AccountUsageStoreError>;
async fn sync_runtime_window(
&self,
account_id: &str,
window: AccountUsageWindow,
) -> Result<(), AccountUsageStoreError>;
async fn sync_rate_limit_window(
&self,
account_id: &str,
reset_at: DateTime<Utc>,
limit_window_seconds: Option<u64>,
) -> Result<(), AccountUsageStoreError>;
async fn record_request(&self, account_id: &str) -> Result<(), AccountUsageStoreError> {
self.record_usage_delta(
account_id,
AccountUsageDelta {
requests: 1,
..AccountUsageDelta::default()
},
)
.await
Self::adapter(error)
}
}
+55 -6
View File
@@ -10,6 +10,7 @@ struct ModelPricing {
output_price_per_mtoken_priority: f64,
cache_read_price_per_mtoken: f64,
cache_read_price_per_mtoken_priority: f64,
cache_write_multiplier: f64,
long_input_price_per_mtoken: f64,
long_input_price_per_mtoken_priority: f64,
long_output_price_per_mtoken: f64,
@@ -27,6 +28,7 @@ impl ModelPricing {
output_price_per_mtoken_priority: 0.0,
cache_read_price_per_mtoken: 0.0,
cache_read_price_per_mtoken_priority: 0.0,
cache_write_multiplier: 0.0,
long_input_price_per_mtoken: 0.0,
long_input_price_per_mtoken_priority: 0.0,
long_output_price_per_mtoken: 0.0,
@@ -41,6 +43,11 @@ impl ModelPricing {
self
}
const fn with_cache_write(mut self, multiplier: f64) -> Self {
self.cache_write_multiplier = multiplier;
self
}
const fn with_priority(mut self, input: f64, output: f64, cache_read: f64) -> Self {
self.input_price_per_mtoken_priority = input;
self.output_price_per_mtoken_priority = output;
@@ -78,24 +85,28 @@ const MODEL_PRICING_RULES: &[ModelPricingRule] = &[
model: "gpt-5.6-sol",
pricing: ModelPricing::new(5.0, 30.0)
.with_cache(0.5)
.with_cache_write(1.25)
.with_priority(10.0, 60.0, 1.0),
},
ModelPricingRule {
model: "gpt-5.6-terra",
pricing: ModelPricing::new(2.5, 15.0)
.with_cache(0.25)
.with_cache_write(1.25)
.with_priority(5.0, 30.0, 0.5),
},
ModelPricingRule {
model: "gpt-5.6-luna",
pricing: ModelPricing::new(1.0, 6.0)
.with_cache(0.1)
.with_cache_write(1.25)
.with_priority(2.0, 12.0, 0.2),
},
ModelPricingRule {
model: "gpt-5.6",
pricing: ModelPricing::new(5.0, 30.0)
.with_cache(0.5)
.with_cache_write(1.25)
.with_priority(10.0, 60.0, 1.0),
},
ModelPricingRule {
@@ -181,6 +192,8 @@ pub(crate) struct BillingBreakdown {
pub output_amount: f64,
/// 缓存读取金额。
pub cache_read_amount: f64,
/// 缓存写入金额。
pub cache_write_amount: f64,
/// 应用服务档位倍率前的标准金额。
pub standard_amount: f64,
/// 最终计费金额。
@@ -191,6 +204,8 @@ pub(crate) struct BillingBreakdown {
pub output_price_per_mtoken: f64,
/// 缓存读取单价,美元 / 1M token。
pub cache_read_price_per_mtoken: f64,
/// 缓存写入单价,美元 / 1M token。
pub cache_write_price_per_mtoken: f64,
/// 服务档位。
pub service_tier: Option<String>,
/// 服务档位倍率。
@@ -202,16 +217,24 @@ pub(crate) fn calculate_billing(
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
cache_write_tokens: u64,
model: &str,
service_tier: Option<&str>,
) -> BillingBreakdown {
let pricing = model_pricing(model).unwrap_or_else(|| {
highest_pricing_for_request(input_tokens, output_tokens, cached_tokens, service_tier)
highest_pricing_for_request(
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
service_tier,
)
});
calculate_billing_with_pricing(
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
pricing,
service_tier,
)
@@ -221,6 +244,7 @@ fn calculate_billing_with_pricing(
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
cache_write_tokens: u64,
pricing: ModelPricing,
service_tier: Option<&str>,
) -> BillingBreakdown {
@@ -230,6 +254,7 @@ fn calculate_billing_with_pricing(
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
pricing,
service_tier,
is_long,
@@ -240,6 +265,7 @@ fn calculate_billing_with_pricing_for_context(
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
cache_write_tokens: u64,
pricing: ModelPricing,
service_tier: Option<&str>,
is_long: bool,
@@ -277,26 +303,38 @@ fn calculate_billing_with_pricing_for_context(
}
let cached_tokens = cached_tokens.min(input_tokens);
let uncached_input_tokens = if cache_read_price > 0.0 {
input_tokens - cached_tokens
let billed_cache_read_tokens = if cache_read_price > 0.0 {
cached_tokens
} else {
input_tokens
0
};
let cache_write_price = input_price * pricing.cache_write_multiplier;
let billed_cache_write_tokens = if cache_write_price > 0.0 {
cache_write_tokens.min(input_tokens - billed_cache_read_tokens)
} else {
0
};
let uncached_input_tokens = input_tokens
.saturating_sub(billed_cache_read_tokens)
.saturating_sub(billed_cache_write_tokens);
let input_amount = uncached_input_tokens as f64 / 1_000_000.0 * input_price;
let cache_read_amount = cached_tokens as f64 / 1_000_000.0 * cache_read_price;
let cache_read_amount = billed_cache_read_tokens as f64 / 1_000_000.0 * cache_read_price;
let cache_write_amount = billed_cache_write_tokens as f64 / 1_000_000.0 * cache_write_price;
let output_amount = output_tokens as f64 / 1_000_000.0 * output_price;
let standard_amount = input_amount + cache_read_amount + output_amount;
let standard_amount = input_amount + cache_read_amount + cache_write_amount + output_amount;
let total_amount = standard_amount * tier_multiplier;
BillingBreakdown {
input_amount,
output_amount,
cache_read_amount,
cache_write_amount,
standard_amount,
total_amount,
input_price_per_mtoken: input_price,
output_price_per_mtoken: output_price,
cache_read_price_per_mtoken: cache_read_price,
cache_write_price_per_mtoken: cache_write_price,
service_tier: normalize_service_tier(service_tier),
tier_multiplier,
}
@@ -309,6 +347,7 @@ pub(crate) fn calculate_aggregate_billing(
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
cache_write_tokens: u64,
model: &str,
service_tier: Option<&str>,
is_long: bool,
@@ -318,6 +357,7 @@ pub(crate) fn calculate_aggregate_billing(
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
service_tier,
is_long,
)
@@ -326,6 +366,7 @@ pub(crate) fn calculate_aggregate_billing(
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
pricing,
service_tier,
is_long,
@@ -337,6 +378,7 @@ pub(crate) fn calculate_billing_amount(
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
cache_write_tokens: u64,
model: &str,
service_tier: Option<&str>,
) -> f64 {
@@ -344,6 +386,7 @@ pub(crate) fn calculate_billing_amount(
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
model,
service_tier,
)
@@ -365,6 +408,7 @@ fn highest_pricing_for_request(
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
cache_write_tokens: u64,
service_tier: Option<&str>,
) -> ModelPricing {
MODEL_PRICING_RULES
@@ -374,6 +418,7 @@ fn highest_pricing_for_request(
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
highest,
service_tier,
)
@@ -382,6 +427,7 @@ fn highest_pricing_for_request(
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
rule.pricing,
service_tier,
)
@@ -398,6 +444,7 @@ fn highest_pricing_for_context(
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
cache_write_tokens: u64,
service_tier: Option<&str>,
is_long: bool,
) -> ModelPricing {
@@ -408,6 +455,7 @@ fn highest_pricing_for_context(
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
highest,
service_tier,
is_long,
@@ -417,6 +465,7 @@ fn highest_pricing_for_context(
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
rule.pricing,
service_tier,
is_long,
+16 -2
View File
@@ -18,6 +18,7 @@ select
coalesce(sum(input_tokens), 0)::bigint as input_tokens,
coalesce(sum(output_tokens), 0)::bigint as output_tokens,
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens,
coalesce(sum(cache_write_tokens), 0)::bigint as cache_write_tokens,
coalesce(sum(first_token_latency_sum), 0)::bigint as first_token_latency_sum,
coalesce(sum(first_token_latency_count), 0)::bigint as first_token_latency_count,
coalesce(sum(latency_sum), 0)::bigint as latency_sum,
@@ -36,7 +37,8 @@ select
coalesce(sum(success_count + error_count), 0)::bigint as request_count,
coalesce(sum(input_tokens), 0)::bigint as input_tokens,
coalesce(sum(output_tokens), 0)::bigint as output_tokens,
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens,
coalesce(sum(cache_write_tokens), 0)::bigint as cache_write_tokens
from request_time_buckets
group by bucket_start, model, service_tier";
@@ -59,6 +61,7 @@ pub struct UsageTimeBucketRecord {
pub input_tokens: i64,
pub output_tokens: i64,
pub cached_tokens: i64,
pub cache_write_tokens: i64,
pub first_token_latency_sum: i64,
pub first_token_latency_count: i64,
pub latency_sum: i64,
@@ -81,6 +84,7 @@ pub struct UsageBucketTotals {
pub input_tokens: i64,
pub output_tokens: i64,
pub cached_tokens: i64,
pub cache_write_tokens: i64,
}
/// 保留期内单个计费桶的用量事实。
@@ -97,6 +101,9 @@ impl UsageBucketTotals {
self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
self.cached_tokens = self.cached_tokens.saturating_add(other.cached_tokens);
self.cache_write_tokens = self
.cache_write_tokens
.saturating_add(other.cache_write_tokens);
}
}
@@ -117,6 +124,7 @@ pub struct ModelBucketUsage {
pub input_tokens: i64,
pub output_tokens: i64,
pub cached_tokens: i64,
pub cache_write_tokens: i64,
pub last_used_at: Option<DateTime<Utc>>,
}
@@ -163,6 +171,7 @@ impl PgRequestBucketQuery {
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
cache_write_tokens: row.get("cache_write_tokens"),
},
})
.collect())
@@ -198,7 +207,8 @@ impl PgRequestBucketQuery {
coalesce(sum(success_count + error_count), 0)::bigint as request_count,
coalesce(sum(input_tokens), 0)::bigint as input_tokens,
coalesce(sum(output_tokens), 0)::bigint as output_tokens,
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens,
coalesce(sum(cache_write_tokens), 0)::bigint as cache_write_tokens
from request_time_buckets
where account_id in (",
);
@@ -223,6 +233,7 @@ impl PgRequestBucketQuery {
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
cache_write_tokens: row.get("cache_write_tokens"),
};
let Some(account_windows) = windows_by_account.get(account_id.as_str()) else {
continue;
@@ -258,6 +269,7 @@ impl PgRequestBucketQuery {
coalesce(sum(input_tokens), 0)::bigint as input_tokens,
coalesce(sum(output_tokens), 0)::bigint as output_tokens,
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens,
coalesce(sum(cache_write_tokens), 0)::bigint as cache_write_tokens,
max(bucket_start) as last_used_at
from request_time_buckets
where model != '__unknown__' and (",
@@ -290,6 +302,7 @@ impl PgRequestBucketQuery {
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
cache_write_tokens: row.get("cache_write_tokens"),
last_used_at: row.get("last_used_at"),
}
})
@@ -307,6 +320,7 @@ fn usage_time_bucket_from_row(row: &sqlx::postgres::PgRow) -> UsageTimeBucketRec
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
cache_write_tokens: row.get("cache_write_tokens"),
first_token_latency_sum: row.get("first_token_latency_sum"),
first_token_latency_count: row.get("first_token_latency_count"),
latency_sum: row.get("latency_sum"),
+9 -4
View File
@@ -53,15 +53,16 @@ impl PgRequestBucketStore {
r"
insert into request_time_buckets (
bucket_start, provider, account_id, model, service_tier, success_count,
input_tokens, output_tokens, cached_tokens, first_token_latency_sum,
input_tokens, output_tokens, cached_tokens, cache_write_tokens, first_token_latency_sum,
first_token_latency_count, latency_sum, latency_count, max_latency_ms,
min_latency_ms, updated_at
) values ($1, $2, $3, $4, $5, 1, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)
) values ($1, $2, $3, $4, $5, 1, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)
on conflict (bucket_start, provider, account_id, model, service_tier) do update set
success_count = request_time_buckets.success_count + 1,
input_tokens = request_time_buckets.input_tokens + excluded.input_tokens,
output_tokens = request_time_buckets.output_tokens + excluded.output_tokens,
cached_tokens = request_time_buckets.cached_tokens + excluded.cached_tokens,
cache_write_tokens = request_time_buckets.cache_write_tokens + excluded.cache_write_tokens,
first_token_latency_sum = request_time_buckets.first_token_latency_sum + excluded.first_token_latency_sum,
first_token_latency_count = request_time_buckets.first_token_latency_count + excluded.first_token_latency_count,
latency_sum = request_time_buckets.latency_sum + excluded.latency_sum,
@@ -82,6 +83,7 @@ on conflict (bucket_start, provider, account_id, model, service_tier) do update
.bind(event.input_tokens.unwrap_or(0))
.bind(event.output_tokens.unwrap_or(0))
.bind(event.cached_tokens.unwrap_or(0))
.bind(event.cache_write_tokens.unwrap_or(0))
.bind(first_token.unwrap_or(0))
.bind(i64::from(first_token.is_some()))
.bind(latency.unwrap_or(0))
@@ -188,6 +190,7 @@ with facts as (
coalesce(input_tokens, 0)::bigint as input_tokens,
coalesce(output_tokens, 0)::bigint as output_tokens,
coalesce(cached_tokens, 0)::bigint as cached_tokens,
coalesce(cache_write_tokens, 0)::bigint as cache_write_tokens,
coalesce(first_token_ms, 0)::bigint as first_token_latency_sum,
(first_token_ms is not null)::integer::bigint as first_token_latency_count,
coalesce(latency_ms, 0)::bigint as latency_sum,
@@ -215,6 +218,7 @@ with facts as (
0::bigint,
0::bigint,
0::bigint,
0::bigint,
null::bigint
from ops_error_logs
where created_at >= $1
@@ -226,6 +230,7 @@ with facts as (
sum(input_tokens) as input_tokens,
sum(output_tokens) as output_tokens,
sum(cached_tokens) as cached_tokens,
sum(cache_write_tokens) as cache_write_tokens,
sum(first_token_latency_sum) as first_token_latency_sum,
sum(first_token_latency_count) as first_token_latency_count,
sum(latency_sum) as latency_sum,
@@ -237,13 +242,13 @@ with facts as (
)
insert into request_time_buckets (
bucket_start, provider, account_id, model, service_tier,
success_count, error_count, input_tokens, output_tokens, cached_tokens,
success_count, error_count, input_tokens, output_tokens, cached_tokens, cache_write_tokens,
first_token_latency_sum, first_token_latency_count, latency_sum, latency_count,
max_latency_ms, min_latency_ms, updated_at
)
select
bucket_start, provider, account_id, model, service_tier,
success_count, error_count, input_tokens, output_tokens, cached_tokens,
success_count, error_count, input_tokens, output_tokens, cached_tokens, cache_write_tokens,
first_token_latency_sum, first_token_latency_count, latency_sum, latency_count,
max_latency_ms, min_latency_ms, now()
from aggregated",
-1
View File
@@ -3,7 +3,6 @@
pub mod account_usage;
pub(crate) mod billing;
pub mod buckets;
pub mod dashboard;
pub mod ops;
pub mod rebuild;
pub mod recorder;
+2 -2
View File
@@ -3,8 +3,8 @@
use thiserror::Error;
use crate::telemetry::{
ops::store::{OpsErrorFilter, PgOpsErrorLogStore},
ops::types::OpsErrorLog,
ops::store::PgOpsErrorLogStore,
ops::types::{OpsErrorFilter, OpsErrorLog},
};
/// 管理端运维错误事件错误。
+4 -23
View File
@@ -6,7 +6,10 @@ use thiserror::Error;
use crate::{
infra::json::{NumberedPage, page_offset},
telemetry::{buckets::store::PgRequestBucketStore, ops::types::OpsErrorLog},
telemetry::{
buckets::store::PgRequestBucketStore,
ops::types::{OpsErrorFilter, OpsErrorLog},
},
};
const OPS_ERROR_SELECT_SQL: &str = r"select
@@ -143,28 +146,6 @@ insert into ops_error_logs (
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct OpsErrorFilter {
pub kind: Option<String>,
pub client_api_key_id: Option<String>,
pub provider: Option<String>,
pub request_id: Option<String>,
pub account_id: Option<String>,
pub route: Option<String>,
pub model: Option<String>,
pub status_code: Option<i64>,
pub client_status_code: Option<i64>,
pub upstream_status_code: Option<i64>,
pub transport: Option<String>,
pub attempt_index: Option<i64>,
pub failure_class: Option<String>,
pub response_id: Option<String>,
pub upstream_request_id: Option<String>,
pub search: Option<String>,
pub start_time: Option<DateTime<Utc>>,
pub end_time: Option<DateTime<Utc>>,
}
fn push_filter(builder: &mut QueryBuilder<Postgres>, filter: &OpsErrorFilter) {
builder.push(" where true");
push_optional_text(builder, "kind", filter.kind.as_deref());
+23
View File
@@ -55,6 +55,29 @@ pub struct OpsErrorLog {
pub created_at: DateTime<Utc>,
}
/// 运维错误查询条件。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct OpsErrorFilter {
pub kind: Option<String>,
pub client_api_key_id: Option<String>,
pub provider: Option<String>,
pub request_id: Option<String>,
pub account_id: Option<String>,
pub route: Option<String>,
pub model: Option<String>,
pub status_code: Option<i64>,
pub client_status_code: Option<i64>,
pub upstream_status_code: Option<i64>,
pub transport: Option<String>,
pub attempt_index: Option<i64>,
pub failure_class: Option<String>,
pub response_id: Option<String>,
pub upstream_request_id: Option<String>,
pub search: Option<String>,
pub start_time: Option<DateTime<Utc>>,
pub end_time: Option<DateTime<Utc>>,
}
impl OpsErrorLog {
/// 构造运维错误事件。
pub fn new(kind: impl Into<String>, message: impl Into<String>) -> Self {
+20 -334
View File
@@ -1,28 +1,18 @@
//! Dispatch 使用记录事件辅助。
use std::time::Instant;
//! 已完成遥测事实的持久化入口。
use serde_json::Value;
use crate::{
infra::time::elapsed_millis_i64,
telemetry::{
ops::store::{PgOpsErrorLogStore, PgOpsErrorLogStoreError},
ops::types::OpsErrorLog,
usage::{
store::{PgUsageRecordStore, PgUsageRecordStoreError},
types::{
ResponseUsageRecord, UsageRecord, metadata_i64, metadata_service_tier,
metadata_string as metadata_string_any,
},
},
use crate::telemetry::{
ops::store::{PgOpsErrorLogStore, PgOpsErrorLogStoreError},
ops::types::OpsErrorLog,
usage::{
store::{PgUsageRecordStore, PgUsageRecordStoreError},
types::UsageRecord,
},
upstream::openai::protocol::responses::CodexResponsesRequest,
};
const MAX_CAPTURED_BODY_BYTES: usize = 4 * 1024;
/// 成功与失败事实的唯一写入入口。
#[derive(Clone)]
pub struct Recorder {
usage_records: PgUsageRecordStore,
@@ -31,11 +21,12 @@ pub struct Recorder {
capture_body: bool,
}
/// 遥测事实写入错误。
#[derive(Debug, thiserror::Error)]
pub enum RecorderError {
#[error("invalid success usage fact")]
InvalidUsageFact,
#[error("telemetry persistence task failed: {0}")]
Task(#[from] tokio::task::JoinError),
#[error(transparent)]
Usage(#[from] PgUsageRecordStoreError),
#[error(transparent)]
@@ -57,7 +48,6 @@ impl Recorder {
}
}
/// 校验并写入一条成功事实及其时间桶。
pub async fn record_usage(&self, mut event: UsageRecord) -> Result<(), RecorderError> {
if !self.usage_enabled {
return Ok(());
@@ -73,16 +63,17 @@ impl Recorder {
);
return Err(RecorderError::InvalidUsageFact);
}
apply_success_capture_body_policy(&mut event, self.capture_body);
self.usage_records.append(&event).await?;
apply_capture_body_policy(&mut event.metadata, self.capture_body);
let usage_records = self.usage_records.clone();
// Controller timeouts must stop waiting without cancelling an in-flight fact write.
tokio::spawn(async move { usage_records.append(&event).await }).await??;
Ok(())
}
/// 规范化并写入一条失败事实及其时间桶。
pub async fn record_error(&self, mut event: OpsErrorLog) -> Result<(), RecorderError> {
lift_error_fact_fields(&mut event);
apply_error_capture_body_policy(&mut event, self.capture_body);
self.ops_errors.append(&event).await?;
apply_capture_body_policy(&mut event.metadata, self.capture_body);
let ops_errors = self.ops_errors.clone();
tokio::spawn(async move { ops_errors.append(&event).await }).await??;
Ok(())
}
@@ -91,88 +82,6 @@ impl Recorder {
}
}
pub(crate) struct DispatchErrorLogRecord<'a> {
pub recorder: &'a Recorder,
pub request_id: &'a str,
pub client_api_key_id: Option<&'a str>,
pub provider: Option<&'a str>,
pub account_id: Option<&'a str>,
pub route: &'a str,
pub model: &'a str,
pub started_at: Instant,
pub status_code: i64,
pub message: &'a str,
pub metadata: Value,
}
pub(crate) async fn record_dispatch_error_event(record: DispatchErrorLogRecord<'_>) {
let mut metadata = record.metadata;
enrich_event_route_metadata(&mut metadata, record.route);
let mut event = OpsErrorLog::new(event_kind(record.route), record.message);
event.request_id = Some(record.request_id.to_string());
event.client_api_key_id = record.client_api_key_id.map(ToString::to_string);
event.provider = record.provider.map(ToString::to_string);
event.account_id = record.account_id.map(ToString::to_string);
event.route = Some(record.route.to_string());
event.model = Some(record.model.to_string());
event.status_code = Some(record.status_code);
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
event.metadata = metadata;
if let Err(error) = record.recorder.record_error(event).await {
tracing::error!(
account_id = record.account_id.unwrap_or(""),
error = %error,
"Failed to record dispatch error event"
);
}
}
pub(crate) async fn record_response_event(record: ResponseUsageRecord<'_>) {
let mut metadata = record.metadata;
enrich_event_route_metadata(&mut metadata, record.route);
let mut event = UsageRecord::new(
event_kind(record.route),
record.message,
record.account_id,
record.model,
record.status_code,
);
event.request_id = Some(record.request_id.to_string());
event.client_api_key_id = record.client_api_key_id.map(ToString::to_string);
event.route = Some(record.route.to_string());
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
event.requested_model = normalized_string(record.requested_model.or(Some(record.model)));
event.upstream_model = normalized_string(Some(record.model));
event.service_tier = normalized_string(record.service_tier);
enrich_usage_record_identity(
&mut metadata,
record.requested_model,
record.model,
record.client_ip,
record.client_user_agent,
record.reasoning_effort,
record.service_tier,
);
if !record.rate_limit_headers.is_empty()
&& let Some(object) = metadata.as_object_mut()
{
object.insert(
"rateLimitHeaders".to_string(),
serde_json::json!(record.rate_limit_headers),
);
}
lift_success_fact_fields(&mut event, &mut metadata);
event.metadata = metadata;
if let Err(error) = record.recorder.record_usage(event).await {
tracing::error!(account_id = %record.account_id, error = %error, "Failed to record response event");
}
}
fn is_usage_fact(event: &UsageRecord) -> bool {
(200..=399).contains(&event.status_code)
&& !event.provider.trim().is_empty()
@@ -180,19 +89,11 @@ fn is_usage_fact(event: &UsageRecord) -> bool {
&& !event.model.trim().is_empty()
}
fn apply_success_capture_body_policy(event: &mut UsageRecord, capture_body: bool) {
if !capture_body {
remove_body_fields(&mut event.metadata);
fn apply_capture_body_policy(metadata: &mut Value, capture_body: bool) {
if capture_body {
limit_body_fields(metadata);
} else {
limit_body_fields(&mut event.metadata);
}
}
fn apply_error_capture_body_policy(event: &mut OpsErrorLog, capture_body: bool) {
if !capture_body {
remove_body_fields(&mut event.metadata);
} else {
limit_body_fields(&mut event.metadata);
remove_body_fields(metadata);
}
}
@@ -248,218 +149,3 @@ fn body_fields() -> [&'static str; 5] {
"upstreamBody",
]
}
fn lift_error_fact_fields(event: &mut OpsErrorLog) {
event.client_status_code = event
.client_status_code
.or_else(|| metadata_i64(&event.metadata, &["clientStatusCode", "client_status_code"]));
event.upstream_status_code = event
.upstream_status_code
.or_else(|| metadata_i64(&event.metadata, &["upstreamStatusCode", "upstreamStatus"]));
event.transport = event
.transport
.take()
.or_else(|| metadata_string_any(&event.metadata, &["transport"]));
event.attempt_index = event
.attempt_index
.or_else(|| metadata_i64(&event.metadata, &["attemptIndex", "attempt_index"]));
event.failure_class = event
.failure_class
.take()
.or_else(|| metadata_string_any(&event.metadata, &["failureClass", "failure_class"]));
event.response_id = event
.response_id
.take()
.or_else(|| metadata_string_any(&event.metadata, &["responseId", "response_id"]));
event.upstream_request_id = event.upstream_request_id.take().or_else(|| {
metadata_string_any(
&event.metadata,
&[
"upstreamRequestId",
"upstream_request_id",
"openaiRequestId",
],
)
});
event.service_tier = event
.service_tier
.take()
.or_else(|| metadata_service_tier(&event.metadata).map(ToString::to_string));
let Some(metadata) = event.metadata.as_object_mut() else {
return;
};
for key in [
"clientStatusCode",
"client_status_code",
"upstreamStatusCode",
"upstreamStatus",
"transport",
"attemptIndex",
"attempt_index",
"failureClass",
"failure_class",
"responseId",
"response_id",
"upstreamRequestId",
"upstream_request_id",
"openaiRequestId",
"serviceTier",
] {
metadata.remove(key);
}
}
fn lift_success_fact_fields(event: &mut UsageRecord, metadata: &mut Value) {
event.transport = metadata_string(metadata, "transport");
event.attempt_index = metadata_nonnegative_i64(metadata, "attemptIndex");
event.response_id = metadata_string(metadata, "responseId");
event.upstream_request_id = metadata_string(metadata, "upstreamRequestId")
.or_else(|| metadata_string(metadata, "openaiRequestId"));
event.first_token_ms = metadata_nonnegative_i64(metadata, "firstTokenMs");
if let Some(usage) = metadata.get("usage") {
event.input_tokens = value_nonnegative_i64(usage, "inputTokens");
event.output_tokens = value_nonnegative_i64(usage, "outputTokens");
event.cached_tokens = value_nonnegative_i64(usage, "cachedTokens");
event.reasoning_tokens = value_nonnegative_i64(usage, "reasoningTokens");
}
let Some(object) = metadata.as_object_mut() else {
return;
};
for key in [
"usage",
"route",
"apiKind",
"requestedModel",
"upstreamModel",
"serviceTier",
"statusCode",
"transport",
"attemptIndex",
"responseId",
"upstreamRequestId",
"openaiRequestId",
"firstTokenMs",
] {
object.remove(key);
}
}
fn metadata_string(metadata: &Value, field: &str) -> Option<String> {
metadata
.get(field)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
fn metadata_nonnegative_i64(metadata: &Value, field: &str) -> Option<i64> {
value_nonnegative_i64(metadata, field)
}
fn value_nonnegative_i64(value: &Value, field: &str) -> Option<i64> {
value
.get(field)
.and_then(Value::as_i64)
.filter(|value| *value >= 0)
}
fn normalized_string(value: Option<&str>) -> Option<String> {
value
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
pub(crate) fn event_kind(route: &str) -> &'static str {
let _ = route;
"v1.response"
}
pub(crate) fn api_kind(route: &str) -> &'static str {
let _ = route;
"responses"
}
pub(crate) fn enrich_event_route_metadata(metadata: &mut Value, route: &str) {
let Some(object) = metadata.as_object_mut() else {
return;
};
object
.entry("route".to_string())
.or_insert_with(|| Value::String(route.to_string()));
object
.entry("apiKind".to_string())
.or_insert_with(|| Value::String(api_kind(route).to_string()));
}
pub(crate) fn enrich_usage_record_identity(
metadata: &mut Value,
requested_model: Option<&str>,
upstream_model: &str,
client_ip: Option<&str>,
client_user_agent: Option<&str>,
reasoning_effort: Option<&str>,
service_tier: Option<&str>,
) {
let Some(object) = metadata.as_object_mut() else {
return;
};
let upstream_model = upstream_model.trim();
let requested_model = requested_model
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(upstream_model);
object.insert(
"requestedModel".to_string(),
Value::String(requested_model.to_string()),
);
object.insert(
"upstreamModel".to_string(),
Value::String(upstream_model.to_string()),
);
insert_trimmed_string(object, "clientIp", client_ip);
insert_trimmed_string(object, "userAgent", client_user_agent);
insert_trimmed_string(object, "reasoningEffort", reasoning_effort);
insert_trimmed_string(object, "serviceTier", service_tier);
}
pub(crate) fn enrich_response_request_semantics(
metadata: &mut Value,
request: &CodexResponsesRequest,
) {
let Some(object) = metadata.as_object_mut() else {
return;
};
let semantics = request.semantics();
object.insert("compact".to_string(), Value::Bool(semantics.compact));
insert_trimmed_string(object, "requestKind", semantics.request_kind.as_deref());
insert_trimmed_string(object, "subagentKind", semantics.subagent_kind.as_deref());
insert_trimmed_string(object, "reasoningPreset", semantics.reasoning_preset);
}
pub(crate) fn reasoning_effort_from_request(request: &CodexResponsesRequest) -> Option<&str> {
reasoning_effort_from_value(request.reasoning())
}
fn insert_trimmed_string(
object: &mut serde_json::Map<String, Value>,
key: &str,
value: Option<&str>,
) {
if let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) {
object.insert(key.to_string(), Value::String(value.to_string()));
}
}
fn reasoning_effort_from_value(reasoning: Option<&Value>) -> Option<&str> {
reasoning?
.get("effort")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
+36 -26
View File
@@ -21,7 +21,7 @@ success_terminals as (
client_api_key_id, provider, account_id, transport, model_name,
null::text as failure_class, false as is_client_cancelled,
false as is_caller_error, latency_ms, first_token_ms, billing_model,
service_tier, input_tokens, output_tokens, cached_tokens
service_tier, input_tokens, output_tokens, cached_tokens, cache_write_tokens
from (
select
case when request_id is null then 'event:' || id else 'request:' || request_id end as request_key,
@@ -31,7 +31,7 @@ success_terminals as (
coalesce(nullif(upstream_model, ''), model) as model_name,
latency_ms, first_token_ms,
coalesce(nullif(upstream_model, ''), model) as billing_model,
service_tier, input_tokens, output_tokens, cached_tokens
service_tier, input_tokens, output_tokens, cached_tokens, cache_write_tokens
from usage_records
where created_at >= $1 and created_at < $2
) candidates
@@ -45,7 +45,7 @@ error_terminals as (
latency_ms, null::bigint as first_token_ms,
null::text as billing_model, null::text as service_tier,
null::bigint as input_tokens, null::bigint as output_tokens,
null::bigint as cached_tokens
null::bigint as cached_tokens, null::bigint as cache_write_tokens
from (
select
case when request_id is null then 'event:' || id else 'request:' || request_id end as request_key,
@@ -264,11 +264,11 @@ struct UsagePerformanceInsights {
latency_p50_ms: Option<f64>,
latency_p95_ms: Option<f64>,
latency_p99_ms: Option<f64>,
ttft_p50_ms: Option<f64>,
ttft_p95_ms: Option<f64>,
ttft_p99_ms: Option<f64>,
first_token_p50_ms: Option<f64>,
first_token_p95_ms: Option<f64>,
first_token_p99_ms: Option<f64>,
latency_coverage: f64,
ttft_coverage: f64,
first_token_coverage: f64,
points: Vec<UsagePerformancePoint>,
}
@@ -280,9 +280,9 @@ struct UsagePerformancePoint {
latency_p50_ms: Option<f64>,
latency_p95_ms: Option<f64>,
latency_p99_ms: Option<f64>,
ttft_p50_ms: Option<f64>,
ttft_p95_ms: Option<f64>,
ttft_p99_ms: Option<f64>,
first_token_p50_ms: Option<f64>,
first_token_p95_ms: Option<f64>,
first_token_p99_ms: Option<f64>,
}
#[derive(Debug, Serialize)]
@@ -565,13 +565,13 @@ async fn load_performance(
select
count(*)::bigint as request_count,
count(latency_ms)::bigint as latency_count,
count(first_token_ms)::bigint as ttft_count,
count(first_token_ms)::bigint as first_token_count,
percentile_cont(0.50) within group (order by latency_ms) as latency_p50_ms,
percentile_cont(0.95) within group (order by latency_ms) as latency_p95_ms,
percentile_cont(0.99) within group (order by latency_ms) as latency_p99_ms,
percentile_cont(0.50) within group (order by first_token_ms) as ttft_p50_ms,
percentile_cont(0.95) within group (order by first_token_ms) as ttft_p95_ms,
percentile_cont(0.99) within group (order by first_token_ms) as ttft_p99_ms
percentile_cont(0.50) within group (order by first_token_ms) as first_token_p50_ms,
percentile_cont(0.95) within group (order by first_token_ms) as first_token_p95_ms,
percentile_cont(0.99) within group (order by first_token_ms) as first_token_p99_ms
from usage_records
where created_at >= $1 and created_at < $2
"#,
@@ -586,9 +586,9 @@ where created_at >= $1 and created_at < $2
percentile_cont(0.50) within group (order by latency_ms) as latency_p50_ms,\n\
percentile_cont(0.95) within group (order by latency_ms) as latency_p95_ms,\n\
percentile_cont(0.99) within group (order by latency_ms) as latency_p99_ms,\n\
percentile_cont(0.50) within group (order by first_token_ms) as ttft_p50_ms,\n\
percentile_cont(0.95) within group (order by first_token_ms) as ttft_p95_ms,\n\
percentile_cont(0.99) within group (order by first_token_ms) as ttft_p99_ms\n\
percentile_cont(0.50) within group (order by first_token_ms) as first_token_p50_ms,\n\
percentile_cont(0.95) within group (order by first_token_ms) as first_token_p95_ms,\n\
percentile_cont(0.99) within group (order by first_token_ms) as first_token_p99_ms\n\
from usage_records\n\
where created_at >= $1 and created_at < $2\n\
group by bucket_start order by bucket_start"
@@ -610,9 +610,9 @@ where created_at >= $1 and created_at < $2
latency_p50_ms: row.get("latency_p50_ms"),
latency_p95_ms: row.get("latency_p95_ms"),
latency_p99_ms: row.get("latency_p99_ms"),
ttft_p50_ms: row.get("ttft_p50_ms"),
ttft_p95_ms: row.get("ttft_p95_ms"),
ttft_p99_ms: row.get("ttft_p99_ms"),
first_token_p50_ms: row.get("first_token_p50_ms"),
first_token_p95_ms: row.get("first_token_p95_ms"),
first_token_p99_ms: row.get("first_token_p99_ms"),
},
);
}
@@ -633,11 +633,11 @@ where created_at >= $1 and created_at < $2
latency_p50_ms: summary.get("latency_p50_ms"),
latency_p95_ms: summary.get("latency_p95_ms"),
latency_p99_ms: summary.get("latency_p99_ms"),
ttft_p50_ms: summary.get("ttft_p50_ms"),
ttft_p95_ms: summary.get("ttft_p95_ms"),
ttft_p99_ms: summary.get("ttft_p99_ms"),
first_token_p50_ms: summary.get("first_token_p50_ms"),
first_token_p95_ms: summary.get("first_token_p95_ms"),
first_token_p99_ms: summary.get("first_token_p99_ms"),
latency_coverage: rate(nonnegative(summary.get("latency_count")), request_count),
ttft_coverage: rate(nonnegative(summary.get("ttft_count")), request_count),
first_token_coverage: rate(nonnegative(summary.get("first_token_count")), request_count),
points,
})
}
@@ -660,7 +660,8 @@ async fn load_cost(
)::bigint as cache_hit_count,\n\
coalesce(sum(input_tokens), 0)::bigint as input_tokens,\n\
coalesce(sum(output_tokens), 0)::bigint as output_tokens,\n\
coalesce(sum(least(coalesce(cached_tokens, 0), coalesce(input_tokens, 0))), 0)::bigint as cached_tokens\n\
coalesce(sum(least(coalesce(cached_tokens, 0), coalesce(input_tokens, 0))), 0)::bigint as cached_tokens,\n\
coalesce(sum(least(coalesce(cache_write_tokens, 0), coalesce(input_tokens, 0))), 0)::bigint as cache_write_tokens\n\
from usage_records\n\
where created_at >= $1 and created_at < $2\n\
group by bucket_start, billing_model, service_tier, is_long\n\
@@ -735,7 +736,8 @@ async fn load_diagnostic_costs(
)::bigint as cache_hit_count,\n\
coalesce(sum(input_tokens), 0)::bigint as input_tokens,\n\
coalesce(sum(output_tokens), 0)::bigint as output_tokens,\n\
coalesce(sum(least(coalesce(cached_tokens, 0), coalesce(input_tokens, 0))), 0)::bigint as cached_tokens\n\
coalesce(sum(least(coalesce(cached_tokens, 0), coalesce(input_tokens, 0))), 0)::bigint as cached_tokens,\n\
coalesce(sum(least(coalesce(cache_write_tokens, 0), coalesce(input_tokens, 0))), 0)::bigint as cache_write_tokens\n\
from {source}\n\
where terminal_requests.is_success\n\
group by name, billing_model, service_tier, is_long"
@@ -764,6 +766,7 @@ struct CostSample {
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
cache_write_tokens: u64,
estimated_cost: f64,
standard_cost: f64,
}
@@ -776,6 +779,7 @@ struct CostAccumulator {
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
cache_write_tokens: u64,
estimated_cost: f64,
standard_cost: f64,
}
@@ -790,6 +794,9 @@ impl CostAccumulator {
self.input_tokens = self.input_tokens.saturating_add(sample.input_tokens);
self.output_tokens = self.output_tokens.saturating_add(sample.output_tokens);
self.cached_tokens = self.cached_tokens.saturating_add(sample.cached_tokens);
self.cache_write_tokens = self
.cache_write_tokens
.saturating_add(sample.cache_write_tokens);
self.estimated_cost += sample.estimated_cost;
self.standard_cost += sample.standard_cost;
}
@@ -803,12 +810,14 @@ fn cost_sample(row: &sqlx::postgres::PgRow) -> CostSample {
let input_tokens = nonnegative(row.get("input_tokens"));
let output_tokens = nonnegative(row.get("output_tokens"));
let cached_tokens = nonnegative(row.get("cached_tokens")).min(input_tokens);
let cache_write_tokens = nonnegative(row.get("cache_write_tokens")).min(input_tokens);
let billing_model: String = row.get("billing_model");
let service_tier: Option<String> = row.get("service_tier");
let billing = calculate_aggregate_billing(
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
&billing_model,
service_tier.as_deref(),
row.get("is_long"),
@@ -820,6 +829,7 @@ fn cost_sample(row: &sqlx::postgres::PgRow) -> CostSample {
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
estimated_cost: billing.total_amount,
standard_cost: billing.standard_amount,
}
+50 -22
View File
@@ -91,11 +91,13 @@ pub(crate) struct UsageRecordBilling {
pub input_amount: f64,
pub output_amount: f64,
pub cache_read_amount: f64,
pub cache_write_amount: f64,
pub standard_amount: f64,
pub total_amount: f64,
pub input_price_per_mtoken: f64,
pub output_price_per_mtoken: f64,
pub cache_read_price_per_mtoken: f64,
pub cache_write_price_per_mtoken: f64,
pub service_tier: Option<String>,
pub service_tier_display: String,
pub multiplier: f64,
@@ -278,8 +280,9 @@ pub(crate) fn usage_record_billing(
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
cache_write_tokens: u64,
) -> Option<UsageRecordBilling> {
if input_tokens == 0 && output_tokens == 0 && cached_tokens == 0 {
if input_tokens == 0 && output_tokens == 0 && cached_tokens == 0 && cache_write_tokens == 0 {
return None;
}
@@ -293,6 +296,7 @@ pub(crate) fn usage_record_billing(
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
model,
service_tier,
);
@@ -301,11 +305,13 @@ pub(crate) fn usage_record_billing(
input_amount: breakdown.input_amount,
output_amount: breakdown.output_amount,
cache_read_amount: breakdown.cache_read_amount,
cache_write_amount: breakdown.cache_write_amount,
standard_amount: breakdown.standard_amount,
total_amount: breakdown.total_amount,
input_price_per_mtoken: breakdown.input_price_per_mtoken,
output_price_per_mtoken: breakdown.output_price_per_mtoken,
cache_read_price_per_mtoken: breakdown.cache_read_price_per_mtoken,
cache_write_price_per_mtoken: breakdown.cache_write_price_per_mtoken,
service_tier: breakdown.service_tier.clone(),
service_tier_display: breakdown
.service_tier
@@ -346,6 +352,7 @@ select
coalesce(sum(input_tokens), 0)::bigint as input_tokens,
coalesce(sum(output_tokens), 0)::bigint as output_tokens,
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens,
coalesce(sum(cache_write_tokens), 0)::bigint as cache_write_tokens,
avg(latency_ms::double precision) as average_latency_ms
from usage_records",
);
@@ -354,11 +361,13 @@ from usage_records",
let input_tokens = nonnegative(row.get("input_tokens"));
let output_tokens = nonnegative(row.get("output_tokens"));
let cached_tokens = nonnegative(row.get("cached_tokens"));
let cache_write_tokens = nonnegative(row.get("cache_write_tokens"));
Ok(UsageRecordSummary {
total_requests: nonnegative(row.get("total_requests")),
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
total_tokens: input_tokens + output_tokens,
average_latency_ms: row.get("average_latency_ms"),
})
@@ -448,6 +457,7 @@ pub(super) async fn usage_breakdown(
coalesce(sum(input_tokens), 0)::bigint as input_tokens,
coalesce(sum(output_tokens), 0)::bigint as output_tokens,
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens,
coalesce(sum(cache_write_tokens), 0)::bigint as cache_write_tokens,
coalesce(sum(latency_ms), 0)::bigint as latency_sum,
count(latency_ms) as latency_count
from usage_records",
@@ -462,6 +472,7 @@ from usage_records",
let input_tokens = nonnegative(row.get("input_tokens"));
let output_tokens = nonnegative(row.get("output_tokens"));
let cached_tokens = nonnegative(row.get("cached_tokens"));
let cache_write_tokens = nonnegative(row.get("cache_write_tokens"));
let request_count = nonnegative(row.get("request_count"));
let latency_sum = nonnegative(row.get("latency_sum"));
let latency_count = nonnegative(row.get("latency_count"));
@@ -469,6 +480,7 @@ from usage_records",
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
&row.get::<String, _>("billing_model"),
row.get::<Option<String>, _>("service_tier").as_deref(),
);
@@ -480,6 +492,7 @@ from usage_records",
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
billing_amount,
latency_sum,
latency_count,
@@ -512,6 +525,7 @@ select
coalesce(sum(input_tokens), 0)::bigint as input_tokens,
coalesce(sum(output_tokens), 0)::bigint as output_tokens,
coalesce(sum(cached_tokens), 0)::bigint as cached_tokens,
coalesce(sum(cache_write_tokens), 0)::bigint as cache_write_tokens,
coalesce(sum(latency_ms), 0)::bigint as latency_sum,
count(latency_ms) as latency_count
from usage_records",
@@ -525,23 +539,26 @@ from usage_records",
let input_tokens = nonnegative(row.get("input_tokens"));
let output_tokens = nonnegative(row.get("output_tokens"));
let cached_tokens = nonnegative(row.get("cached_tokens"));
let cache_write_tokens = nonnegative(row.get("cache_write_tokens"));
let billing_amount = usage_breakdown_billing_amount(
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
&row.get::<String, _>("billing_model"),
row.get::<Option<String>, _>("service_tier").as_deref(),
);
days.entry(date.clone())
.or_insert_with(|| UsageTrendAccumulator::new(date))
.push(
.push(UsageTrendSample {
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
billing_amount,
nonnegative(row.get("latency_sum")),
nonnegative(row.get("latency_count")),
);
latency_sum: nonnegative(row.get("latency_sum")),
latency_count: nonnegative(row.get("latency_count")),
});
}
let mut points = days
.into_values()
@@ -576,6 +593,7 @@ pub(super) fn usage_record_from_row(row: &sqlx::postgres::PgRow) -> UsageRecord
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
cache_write_tokens: row.get("cache_write_tokens"),
reasoning_tokens: row.get("reasoning_tokens"),
message: row.get("message"),
metadata: row
@@ -593,6 +611,7 @@ fn usage_breakdown_billing_amount(
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
cache_write_tokens: u64,
model: &str,
service_tier: Option<&str>,
) -> f64 {
@@ -600,6 +619,7 @@ fn usage_breakdown_billing_amount(
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
model,
service_tier,
)
@@ -616,6 +636,7 @@ struct BreakdownSample {
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
cache_write_tokens: u64,
billing_amount: f64,
latency_sum: u64,
latency_count: u64,
@@ -638,6 +659,7 @@ impl BreakdownAccumulator {
self.item.input_tokens += sample.input_tokens;
self.item.output_tokens += sample.output_tokens;
self.item.cached_tokens += sample.cached_tokens;
self.item.cache_write_tokens += sample.cache_write_tokens;
self.item.total_tokens += sample.input_tokens + sample.output_tokens;
self.item.standard_billing_amount += sample.billing_amount;
self.item.actual_billing_amount += sample.billing_amount;
@@ -659,6 +681,16 @@ struct UsageTrendAccumulator {
latency_count: u64,
}
struct UsageTrendSample {
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
cache_write_tokens: u64,
billing_amount: f64,
latency_sum: u64,
latency_count: u64,
}
impl UsageTrendAccumulator {
fn new(date: String) -> Self {
Self {
@@ -671,23 +703,16 @@ impl UsageTrendAccumulator {
}
}
fn push(
&mut self,
input: u64,
output: u64,
cached: u64,
billing_amount: f64,
latency_sum: u64,
latency_count: u64,
) {
self.point.input_tokens += input;
self.point.output_tokens += output;
self.point.cached_tokens += cached;
self.point.total_tokens += input + output;
self.point.standard_billing_amount += billing_amount;
self.point.actual_billing_amount += billing_amount;
self.latency_sum += latency_sum;
self.latency_count += latency_count;
fn push(&mut self, sample: UsageTrendSample) {
self.point.input_tokens += sample.input_tokens;
self.point.output_tokens += sample.output_tokens;
self.point.cached_tokens += sample.cached_tokens;
self.point.cache_write_tokens += sample.cache_write_tokens;
self.point.total_tokens += sample.input_tokens + sample.output_tokens;
self.point.standard_billing_amount += sample.billing_amount;
self.point.actual_billing_amount += sample.billing_amount;
self.latency_sum += sample.latency_sum;
self.latency_count += sample.latency_count;
}
fn finish(mut self) -> UsageRecordTrendPoint {
@@ -703,6 +728,7 @@ pub struct UsageRecordSummary {
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub cache_write_tokens: u64,
pub total_tokens: u64,
pub average_latency_ms: Option<f64>,
}
@@ -730,6 +756,7 @@ pub struct UsageRecordBreakdown {
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub cache_write_tokens: u64,
pub total_tokens: u64,
pub standard_billing_amount: f64,
pub actual_billing_amount: f64,
@@ -743,6 +770,7 @@ pub struct UsageRecordTrendPoint {
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub cache_write_tokens: u64,
pub total_tokens: u64,
pub standard_billing_amount: f64,
pub actual_billing_amount: f64,
+6 -3
View File
@@ -24,7 +24,8 @@ select
id, request_id, client_api_key_id, kind, route, provider, account_id, model,
requested_model, upstream_model, service_tier, status_code, transport, attempt_index,
response_id, upstream_request_id, latency_ms, first_token_ms, input_tokens,
output_tokens, cached_tokens, reasoning_tokens, message, metadata_json, created_at
output_tokens, cached_tokens, cache_write_tokens, reasoning_tokens, message, metadata_json,
created_at
from usage_records";
#[derive(Debug, Error)]
@@ -82,10 +83,11 @@ insert into usage_records (
id, request_id, client_api_key_id, kind, route, provider, account_id, model,
requested_model, upstream_model, service_tier, status_code, transport, attempt_index,
response_id, upstream_request_id, latency_ms, first_token_ms, input_tokens,
output_tokens, cached_tokens, reasoning_tokens, message, metadata_json, created_at
output_tokens, cached_tokens, cache_write_tokens, reasoning_tokens, message, metadata_json,
created_at
) values (
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13,
$14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25
$14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25, $26
)",
)
.bind(&event.id)
@@ -109,6 +111,7 @@ insert into usage_records (
.bind(event.input_tokens)
.bind(event.output_tokens)
.bind(event.cached_tokens)
.bind(event.cache_write_tokens)
.bind(event.reasoning_tokens)
.bind(&event.message)
.bind(sqlx::types::Json(&event.metadata))
+3 -20
View File
@@ -65,6 +65,8 @@ pub struct UsageRecord {
pub output_tokens: Option<i64>,
/// 缓存 token。
pub cached_tokens: Option<i64>,
/// 缓存写入 token。
pub cache_write_tokens: Option<i64>,
/// reasoning token。
pub reasoning_tokens: Option<i64>,
/// 展示消息。
@@ -107,6 +109,7 @@ impl UsageRecord {
input_tokens: None,
output_tokens: None,
cached_tokens: None,
cache_write_tokens: None,
reasoning_tokens: None,
message: message.into(),
metadata: Value::Object(Map::default()),
@@ -138,23 +141,3 @@ pub(crate) fn metadata_i64(metadata: &Value, keys: &[&str]) -> Option<i64> {
keys.iter()
.find_map(|key| metadata.get(*key).and_then(Value::as_i64))
}
/// 响应事件记录(供 dispatch 模块使用)。
pub struct ResponseUsageRecord<'a> {
pub recorder: &'a crate::telemetry::recorder::Recorder,
pub request_id: &'a str,
pub client_api_key_id: Option<&'a str>,
pub account_id: &'a str,
pub route: &'a str,
pub model: &'a str,
pub requested_model: Option<&'a str>,
pub client_ip: Option<&'a str>,
pub client_user_agent: Option<&'a str>,
pub reasoning_effort: Option<&'a str>,
pub service_tier: Option<&'a str>,
pub started_at: std::time::Instant,
pub status_code: i64,
pub message: &'a str,
pub metadata: serde_json::Value,
pub rate_limit_headers: &'a [(String, String)],
}
@@ -0,0 +1,436 @@
//! OpenAI 账号 usage/probe adapter。
use std::{collections::VecDeque, time::Duration};
use async_trait::async_trait;
use futures::StreamExt;
use secrecy::ExposeSecret;
use serde_json::{Value, json};
use crate::fleet::{
account_gateway::{
AccountFailureObservation, AccountGatewayError, AccountProbeEvent, AccountProbeRequest,
AccountProbeSession, AccountUpstreamContext, AccountUpstreamGateway, AccountUsageResult,
},
quota::{
QuotaCreditsObservation, QuotaLimitObservation, QuotaObservation, QuotaSpendControl,
QuotaWindowObservation, quota_from_observation,
},
};
use super::{
failure::{UpstreamFailureFacts, upstream_failure_facts},
protocol::{
responses::{CodexResponsesRequest, ResponsesSseFailure},
sse::{SseEvent, parse_sse_events},
},
transport::{CodexBackendClient, CodexBackendSseStream, CodexRequestContext},
};
#[async_trait]
impl AccountUpstreamGateway for CodexBackendClient {
async fn fetch_usage(
&self,
context: AccountUpstreamContext,
) -> Result<AccountUsageResult, AccountGatewayError> {
let raw = CodexBackendClient::fetch_usage(self, codex_context(&context))
.await
.map_err(account_gateway_error)?;
Ok(account_usage_result(raw))
}
async fn probe_response(
&self,
context: AccountUpstreamContext,
request: AccountProbeRequest,
) -> Result<AccountProbeSession, AccountGatewayError> {
let mut codex_request = CodexResponsesRequest::new_http_sse(
request.model,
request.instructions,
vec![json!({
"role": "user",
"content": [{ "type": "input_text", "text": request.input_text }]
})],
);
codex_request.set_stream(true);
codex_request.set_store(false);
codex_request.force_http_sse = true;
let request_payload = serde_json::to_value(&codex_request)
.map_err(|error| AccountGatewayError::new(error.to_string(), None))?;
let response = self
.create_response_stream(&codex_request, codex_context(&context))
.await
.map_err(account_gateway_error)?;
let state = ProbeStreamState::new(response.body);
let events = futures::stream::unfold(state, next_probe_event);
Ok(AccountProbeSession {
request_payload,
events: Box::pin(events),
})
}
async fn evict_account_connections(&self, account_id: &str) {
self.evict_websocket_account(account_id).await;
}
}
fn codex_context(context: &AccountUpstreamContext) -> CodexRequestContext<'_> {
CodexRequestContext {
access_token: context.access_token.expose_secret(),
account_id: context.account_id.as_deref(),
request_id: &context.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: context.cookie_header.as_deref(),
installation_id: context.installation_id.as_deref(),
session_id: None,
thread_id: None,
client_request_id: None,
turn_id: None,
}
}
fn account_usage_result(raw: Value) -> AccountUsageResult {
let observation = usage_quota_observation(&raw);
AccountUsageResult {
quota: quota_from_observation(&observation, None, None),
account_id: string_field(&raw, "account_id"),
user_id: string_field(&raw, "user_id"),
email: string_field(&raw, "email"),
plan_type: string_field(&raw, "plan_type"),
raw,
}
}
fn usage_quota_observation(usage: &Value) -> QuotaObservation {
let mut limits = std::collections::BTreeMap::new();
if let Some(limit) = usage
.get("rate_limit")
.and_then(|rate_limit| usage_limit_observation("codex", None, None, rate_limit))
{
limits.insert("codex".to_string(), limit);
}
if let Some(additional) = usage
.get("additional_rate_limits")
.and_then(Value::as_array)
{
for item in additional {
let limit_name = item.get("limit_name").and_then(trimmed_str);
let metered_feature = item.get("metered_feature").and_then(trimmed_str);
let Some(limit_id) = metered_feature.or(limit_name).map(normalize_limit_id) else {
continue;
};
let Some(limit) = item.get("rate_limit").and_then(|rate_limit| {
usage_limit_observation(&limit_id, limit_name, metered_feature, rate_limit)
}) else {
continue;
};
limits.insert(limit_id, limit);
}
}
QuotaObservation {
limits,
credits: usage
.get("credits")
.filter(|value| !value.is_null())
.map(|credits| QuotaCreditsObservation {
has_credits: credits
.get("has_credits")
.and_then(Value::as_bool)
.unwrap_or(false),
unlimited: credits
.get("unlimited")
.and_then(Value::as_bool)
.unwrap_or(false),
overage_limit_reached: credits
.get("overage_limit_reached")
.and_then(Value::as_bool)
.unwrap_or(false),
balance: credits
.get("balance")
.filter(|value| !value.is_null())
.cloned(),
}),
spend_control: usage
.get("spend_control")
.filter(|value| !value.is_null())
.and_then(|value| serde_json::from_value::<QuotaSpendControl>(value.clone()).ok()),
plan_type: string_field(usage, "plan_type"),
..QuotaObservation::default()
}
}
fn usage_limit_observation(
limit_id: &str,
limit_name: Option<&str>,
metered_feature: Option<&str>,
rate_limit: &Value,
) -> Option<QuotaLimitObservation> {
if rate_limit.is_null() {
return None;
}
let primary = rate_limit
.get("primary_window")
.and_then(usage_window_observation);
let secondary = rate_limit
.get("secondary_window")
.and_then(usage_window_observation);
if primary.is_none() && secondary.is_none() {
return None;
}
Some(QuotaLimitObservation {
limit_id: limit_id.to_string(),
limit_name: limit_name.map(ToString::to_string),
metered_feature: metered_feature.map(ToString::to_string),
allowed: rate_limit.get("allowed").and_then(Value::as_bool),
limit_reached: rate_limit.get("limit_reached").and_then(Value::as_bool),
primary,
secondary,
})
}
fn usage_window_observation(window: &Value) -> Option<QuotaWindowObservation> {
if window.is_null() {
return None;
}
let used_percent = window
.get("used_percent")
.and_then(number_value)
.unwrap_or(0.0);
let window_minutes = window
.get("limit_window_seconds")
.and_then(number_value)
.filter(|seconds| *seconds > 0.0)
.and_then(|seconds| Duration::try_from_secs_f64((seconds / 60.0).round()).ok())
.map(|duration| duration.as_secs());
Some(QuotaWindowObservation {
used_percent,
reset_at: window.get("reset_at").and_then(positive_i64),
window_minutes,
used: first_present(
window,
&[
"used_tokens",
"used",
"used_credits",
"usage",
"consumed_tokens",
"consumed",
],
),
limit: first_present(
window,
&[
"limit_tokens",
"limit",
"limit_credits",
"quota",
"total_tokens",
"total",
],
),
})
}
struct ProbeStreamState {
body: CodexBackendSseStream,
buffer: String,
pending: VecDeque<AccountProbeEvent>,
terminal: bool,
eof: bool,
}
impl ProbeStreamState {
fn new(body: CodexBackendSseStream) -> Self {
Self {
body,
buffer: String::new(),
pending: VecDeque::new(),
terminal: false,
eof: false,
}
}
}
async fn next_probe_event(
mut state: ProbeStreamState,
) -> Option<(AccountProbeEvent, ProbeStreamState)> {
loop {
if let Some(event) = state.pending.pop_front() {
return Some((event, state));
}
if state.terminal {
return None;
}
if let Some(frame) = take_sse_frame(&mut state.buffer) {
process_probe_frame(&mut state, &frame);
continue;
}
if state.eof {
if !state.buffer.trim().is_empty() {
let frame = std::mem::take(&mut state.buffer);
process_probe_frame(&mut state, &frame);
continue;
}
state.terminal = true;
return Some((
AccountProbeEvent::Failed(AccountGatewayError::new(
"Stream ended before response.completed",
None,
)),
state,
));
}
match state.body.next().await {
Some(Ok(chunk)) => state.buffer.push_str(&String::from_utf8_lossy(&chunk)),
Some(Err(error)) => {
state.terminal = true;
return Some((
AccountProbeEvent::Failed(account_gateway_error(error)),
state,
));
}
None => state.eof = true,
}
}
}
fn process_probe_frame(state: &mut ProbeStreamState, frame: &str) {
let events = match parse_sse_events(frame) {
Ok(events) => events,
Err(error) => {
state.terminal = true;
state
.pending
.push_back(AccountProbeEvent::Failed(AccountGatewayError::new(
error.to_string(),
None,
)));
return;
}
};
for event in events {
process_probe_sse_event(state, &event);
if state.terminal {
break;
}
}
}
fn process_probe_sse_event(state: &mut ProbeStreamState, event: &SseEvent) {
let Ok(value) = serde_json::from_str::<Value>(&event.data) else {
return;
};
match value.get("type").and_then(Value::as_str) {
Some("response.output_text.delta") => {
if let Some(delta) = value.get("delta").and_then(Value::as_str)
&& !delta.is_empty()
{
state
.pending
.push_back(AccountProbeEvent::Content(delta.to_string()));
}
}
Some("response.completed" | "response.done") => {
state.terminal = true;
state.pending.push_back(AccountProbeEvent::Complete);
}
Some(event_name @ ("response.failed" | "error")) => {
let failure = ResponsesSseFailure::from_event(event_name, &value);
state.terminal = true;
state
.pending
.push_back(AccountProbeEvent::Failed(response_gateway_error(&failure)));
}
_ => {}
}
}
fn account_gateway_error(error: super::transport::CodexClientError) -> AccountGatewayError {
let message = error.to_string();
let facts = upstream_failure_facts(&error);
AccountGatewayError::new(message, Some(account_failure_observation(&facts)))
}
fn response_gateway_error(failure: &ResponsesSseFailure) -> AccountGatewayError {
AccountGatewayError::new(
failure.message.clone(),
Some(AccountFailureObservation {
status_code: failure.explicit_status_code,
code: failure.upstream_code.clone(),
error_type: failure.upstream_type.clone(),
message: failure.message.clone(),
body: failure.message.clone(),
retry_after_seconds: failure.retry_after_seconds,
..AccountFailureObservation::default()
}),
)
}
fn account_failure_observation(facts: &UpstreamFailureFacts) -> AccountFailureObservation {
AccountFailureObservation {
status_code: facts.status_code,
code: facts.code.clone(),
error_type: facts.error_type.clone(),
identity_authorization_error: facts.identity_authorization_error.clone(),
identity_error_code: facts.identity_error_code.clone(),
message: facts.message.clone(),
body: facts.body.clone(),
retry_after_seconds: facts.retry_after_seconds,
}
}
fn take_sse_frame(buffer: &mut String) -> Option<String> {
let index = buffer.find("\n\n").or_else(|| buffer.find("\r\n\r\n"))?;
let delimiter_len = if buffer[index..].starts_with("\r\n\r\n") {
4
} else {
2
};
let frame = buffer[..index + delimiter_len].to_string();
buffer.drain(..index + delimiter_len);
Some(frame)
}
fn normalize_limit_id(value: &str) -> String {
value.trim().to_ascii_lowercase().replace(['-', ' '], "_")
}
fn string_field(value: &Value, key: &str) -> Option<String> {
value
.get(key)
.and_then(trimmed_str)
.map(ToString::to_string)
}
fn trimmed_str(value: &Value) -> Option<&str> {
value
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn first_present(value: &Value, keys: &[&str]) -> Option<Value> {
keys.iter()
.find_map(|key| value.get(*key).filter(|value| !value.is_null()).cloned())
}
fn number_value(value: &Value) -> Option<f64> {
value
.as_f64()
.or_else(|| value.as_str().and_then(|value| value.parse::<f64>().ok()))
.filter(|value| value.is_finite())
}
fn positive_i64(value: &Value) -> Option<i64> {
value
.as_i64()
.filter(|value| *value > 0)
.or_else(|| value.as_u64().and_then(|value| i64::try_from(value).ok()))
.filter(|value| *value > 0)
}

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