mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 08:05:32 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b11e1ecf4b | ||
|
|
a34d3ca7bc | ||
|
|
28f2879a05 | ||
|
|
2474259c22 | ||
|
|
79c2c5948f | ||
|
|
3c4bd8757f | ||
|
|
06f3443e56 | ||
|
|
88fd21f169 | ||
|
|
521201c926 | ||
|
|
be2123ea30 | ||
|
|
4b6d369225 | ||
|
|
76b81a6ceb | ||
|
|
29c8b679ba | ||
|
|
4091a0c2d8 | ||
|
|
f0ccc981fb | ||
|
|
d01ea7f775 | ||
|
|
fd749ea8f8 | ||
|
|
f0543b637e | ||
|
|
549bb9f62b | ||
|
|
9e57f78bab | ||
|
|
73c0ecdc4c |
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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) }}
|
||||
|
||||
Generated
+1
@@ -382,6 +382,7 @@ dependencies = [
|
||||
"sha2 0.11.0",
|
||||
"sqlx",
|
||||
"subtle",
|
||||
"syn",
|
||||
"tar",
|
||||
"tempfile",
|
||||
"thiserror",
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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, "a_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("a_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), "a_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;
|
||||
@@ -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,
|
||||
"a_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("a_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, "a_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()
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(())
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 首次发布检查失败");
|
||||
|
||||
@@ -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 或手动操作再次触发");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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! {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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, "a)
|
||||
.await;
|
||||
if quota_snapshot_limit_reached("a) {
|
||||
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, "a)))
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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>,
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
@@ -586,11 +586,22 @@ pub fn normalized_imported_account_status(
|
||||
status: AccountStatus,
|
||||
source: AccountImportSource,
|
||||
access_token: &str,
|
||||
access_token_expires_at: Option<&DateTime<Utc>>,
|
||||
) -> AccountStatus {
|
||||
if source == AccountImportSource::Sub2Api
|
||||
&& status == AccountStatus::Active
|
||||
&& jwt_expiry(access_token, Utc::now()) != JwtExpiry::Valid
|
||||
{
|
||||
if source != AccountImportSource::Sub2Api || status != AccountStatus::Active {
|
||||
return status;
|
||||
}
|
||||
|
||||
let now = Utc::now();
|
||||
if let Some(expires_at) = access_token_expires_at {
|
||||
return if now >= *expires_at {
|
||||
AccountStatus::Expired
|
||||
} else {
|
||||
status
|
||||
};
|
||||
}
|
||||
|
||||
if jwt_expiry(access_token, now) != JwtExpiry::Valid {
|
||||
AccountStatus::Expired
|
||||
} else {
|
||||
status
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
};
|
||||
|
||||
@@ -81,16 +83,17 @@ impl AccountManageService {
|
||||
let parsed_status =
|
||||
crate::fleet::import::parse_account_import_status(entry.status.as_deref())
|
||||
.map_err(|error| AccountManageError::InvalidStatus(error.to_string()))?;
|
||||
let mut status = crate::fleet::import::normalized_imported_account_status(
|
||||
parsed_status,
|
||||
source,
|
||||
&resolved_tokens.access_token,
|
||||
);
|
||||
let access_token_expires_at = resolved_tokens
|
||||
.claims
|
||||
.as_ref()
|
||||
.map(|claims| claims.expires_at)
|
||||
.or(access_token_expires_at);
|
||||
let mut status = crate::fleet::import::normalized_imported_account_status(
|
||||
parsed_status,
|
||||
source,
|
||||
&resolved_tokens.access_token,
|
||||
access_token_expires_at.as_ref(),
|
||||
);
|
||||
let claims = resolved_tokens.claims.as_ref();
|
||||
let mut plan_type = claims
|
||||
.and_then(|claims| claims.plan_type.clone())
|
||||
@@ -190,41 +193,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,
|
||||
|
||||
@@ -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::{
|
||||
|
||||
@@ -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
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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("a);
|
||||
|
||||
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")
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,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;
|
||||
|
||||
@@ -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>("a_json).ok(),
|
||||
Ok(Some(quota_json)) => QuotaSnapshot::from_json("a_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, "a).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, "a_json, limit_reached, cooldown_until)
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
/// 更新账号状态。
|
||||
|
||||
@@ -1,7 +1,11 @@
|
||||
//! 账号配额解析与刷新。
|
||||
|
||||
mod model;
|
||||
mod observation;
|
||||
mod runtime;
|
||||
mod service;
|
||||
|
||||
pub use model::*;
|
||||
pub use observation::*;
|
||||
pub use runtime::*;
|
||||
pub use service::*;
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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"
|
||||
)
|
||||
}
|
||||
@@ -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("a);
|
||||
let reset_at = quota_snapshot_reset_at("a);
|
||||
let quota_json = quota.to_json()?;
|
||||
if self
|
||||
.store
|
||||
.apply_quota_snapshot(
|
||||
&account.id,
|
||||
"a.to_string(),
|
||||
"a_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 服务结果。
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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>;
|
||||
}
|
||||
@@ -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::*;
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
/// 移除指定账号的反馈记录(账号被删除时调用)。
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
@@ -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,6 +1,7 @@
|
||||
//! 模型目录领域。
|
||||
|
||||
pub mod catalog;
|
||||
pub mod gateway;
|
||||
pub mod service;
|
||||
pub mod store;
|
||||
pub mod types;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -3,8 +3,8 @@
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::telemetry::{
|
||||
ops::store::{OpsErrorFilter, PgOpsErrorLogStore},
|
||||
ops::types::OpsErrorLog,
|
||||
ops::store::PgOpsErrorLogStore,
|
||||
ops::types::{OpsErrorFilter, OpsErrorLog},
|
||||
};
|
||||
|
||||
/// 管理端运维错误事件错误。
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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)],
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user