Compare commits

..
9 Commits
Author SHA1 Message Date
zyycnandCodex 40a2b70bf2 chore: bump version to 2.3.2
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 12:07:00 +08:00
zyycnandCodex 1e6ac77d3c fix(frontend): scope usage record refresh
Refresh only the request detail table with an independent current-time snapshot, while preserving latest-request-wins ordering and leaving analytics cards untouched.

Co-authored-by: Codex <noreply@openai.com>
2026-07-15 12:06:38 +08:00
zyycnandCodex 3340934944 fix(transport): preserve websocket openings after fallback
Continue pool-owned WebSocket handshakes after the foreground budget expires, retain late successes for reuse, and cancel them on account eviction or shutdown. Keep breaker accounting and transport metrics aligned with the foreground decision.

Co-authored-by: Codex <noreply@openai.com>
2026-07-15 12:06:26 +08:00
zyycnandCodex 37a368de32 fix(ci): avoid duplicate pull request checks
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 11:33:11 +08:00
zyycnandCodex 79cd8ac965 chore(deps): delay fresh npm updates
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 11:13:50 +08:00
zyycnandCodex 57c6b95596 fix: report total dashboard billing
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 10:28:07 +08:00
zyycnandCodex a4f1843876 fix: unify account failures and usage diagnostics
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 10:15:10 +08:00
zyycnandCodex 096e48621a chore: bump version to 2.3.1
Co-authored-by: Codex <noreply@openai.com>
2026-07-15 03:09:52 +08:00
zyycnandCodex 36828b1a7e fix: harden websocket liveness and usage reporting
Enforce pump-owned Pong deadlines without adding request-path probes, and pause watchdogs during inbound backpressure.

Refresh usage snapshots against the current time, expose account identities, preserve full plan badges, and source dashboard totals from global request buckets.

Co-authored-by: Codex <noreply@openai.com>
2026-07-15 03:09:39 +08:00
58 changed files with 2050 additions and 1024 deletions
+2
View File
@@ -34,6 +34,8 @@ updates:
day: monday
time: '09:20'
timezone: Asia/Shanghai
cooldown:
default-days: 1
groups:
minor-and-patch:
update-types:
+3 -1
View File
@@ -3,8 +3,10 @@ name: CI
on:
push:
branches:
- '**'
- main
pull_request:
branches:
- main
permissions:
contents: read
+3 -1
View File
@@ -3,8 +3,10 @@ name: Security Scan
on:
push:
branches:
- '**'
- main
pull_request:
branches:
- main
schedule:
- cron: '0 3 * * 1'
+4 -4
View File
@@ -205,12 +205,12 @@ pub(crate) async fn dashboard_summary(
let capacity = state.services.account_pool.capacity_summary_now().await;
let now = Utc::now();
let today_filter = today_usage_record_filter(now);
let lifetime_usage = state
let retained_usage = state
.services
.usage
.summary()
.retained_summary()
.await
.map_err(|error| dashboard_data_error("account usage summary", &error))?;
.map_err(|error| dashboard_data_error("retained usage summary", &error))?;
let account_usage_records = state
.services
.usage_records
@@ -256,7 +256,7 @@ pub(crate) async fn dashboard_summary(
cards: dashboard_cards(
dashboard_account_counts(&accounts),
&time_buckets,
&lifetime_usage,
&retained_usage,
),
trend,
health_timeline: dashboard_health_timeline_data(&time_buckets),
+1 -1
View File
@@ -82,6 +82,7 @@ impl TaskCoordinator {
stores.accounts.clone(),
stores.account_usage.clone(),
services.codex.clone(),
services.account_pool.clone(),
services.account_pseudonymizer.clone(),
config
.quota
@@ -91,7 +92,6 @@ impl TaskCoordinator {
30 * 60,
)
.with_cookie_store(stores.cookies.clone())
.with_account_pool(services.account_pool.clone())
.start(),
);
coordinator.push("fingerprint_update", {
+2 -6
View File
@@ -35,6 +35,7 @@ impl QuotaRefreshTask {
store: PgAccountStore,
usage_store: PgAccountUsageStore,
codex: Arc<CodexBackendClient>,
account_pool: Arc<AccountPoolService>,
account_pseudonymizer: Arc<AccountPseudonymizer>,
interval_secs: u64,
min_refresh_interval_secs: u64,
@@ -44,6 +45,7 @@ impl QuotaRefreshTask {
store,
usage_store,
codex,
account_pool,
account_pseudonymizer,
min_refresh_interval_secs,
),
@@ -58,12 +60,6 @@ impl QuotaRefreshTask {
self
}
/// 设置运行时账号池,用于刷新后同步内存调度状态。
pub fn with_account_pool(mut self, account_pool: Arc<AccountPoolService>) -> Self {
self.service = self.service.with_account_pool(account_pool);
self
}
/// 启动后台刷新任务。
pub fn start(self) -> SchedulerHandle {
let config = PeriodicTaskConfig::new(
@@ -2,19 +2,25 @@
use crate::{
dispatch::{
controllers::{
ControllerFailureFact,
account_state::{AccountStateEffect, AccountStateEffects},
},
controllers::ControllerFailureFact,
errors::ClientFailure,
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind},
lifecycle::contract::{
AttemptDecision, AttemptObservation, AttemptObservationKind, CompleteResponseFacts,
},
transport::observation::{UpstreamFailureFacts, UpstreamFailureKind},
},
fleet::{account::AccountStatus, pool::AccountPoolService},
upstream::openai::{protocol::responses::ResponsesSseFailure, transport::CodexBackendClient},
fleet::{
account_failure::{
AccountFailureKind, AccountStateEffect, apply_account_state_effect_immediately,
classify_response_failure, classify_upstream_failure,
},
pool::AccountPoolService,
},
upstream::openai::{
failure::{UpstreamFailureFacts, UpstreamFailureKind},
protocol::responses::ResponsesSseFailure,
transport::CodexBackendClient,
},
};
pub(super) struct AccountFailureController;
@@ -46,7 +52,7 @@ impl AccountFailureController {
return;
};
if let Some(effect) = &classified.effect {
AccountStateEffects::apply(account_pool, codex, account_id, effect).await;
apply_account_state_effect_immediately(account_pool, codex, account_id, effect).await;
}
}
@@ -110,15 +116,7 @@ impl AccountFailureController {
.as_deref()
.unwrap_or_default()
.to_ascii_lowercase();
let error_type = failure
.upstream_type
.as_deref()
.unwrap_or_default()
.to_ascii_lowercase();
let message = failure.message.to_ascii_lowercase();
if is_model_unsupported(&code)
|| is_model_unsupported(&message)
|| code.contains("invalid_request")
if code.contains("invalid_request")
|| code.contains("not_found")
|| code.contains("context_window")
|| code.contains("invalid_prompt")
@@ -126,17 +124,18 @@ impl AccountFailureController {
{
return Some(ClientFailure::new(failure.clone(), 400, false));
}
if is_auth_failure(&code, &error_type, &message) {
return Some(ClientFailure::new(failure.clone(), 401, false));
}
if code.contains("forbidden") || code.contains("banned") || error_type == "permission_error"
{
return Some(ClientFailure::new(failure.clone(), 403, false));
}
if code.contains("server_overloaded") {
return Some(ClientFailure::new(failure.clone(), 503, true));
}
None
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)
})
}
}
@@ -168,167 +167,42 @@ fn classify_failure(
}
fn classify_upstream(account_id: &str, facts: &UpstreamFailureFacts) -> Option<ClassifiedFailure> {
let status = facts.status_code;
let body = facts.body.as_str();
let code = facts.code.as_deref().unwrap_or_default();
let error_type = facts.error_type.as_deref().unwrap_or_default();
let message = facts.message.as_str();
if is_model_unsupported(code) || is_model_unsupported(message) || is_model_unsupported(body) {
return Some(model_unsupported_failure(account_id, body));
}
if let Some(account_status) = explicit_account_status(code, error_type, body) {
return Some(account_status_failure(
account_id,
account_status,
body,
status,
));
}
if status == Some(403) && !is_html(body) {
return Some(account_status_failure(
account_id,
AccountStatus::Banned,
body,
status,
));
}
if status == Some(401) {
return Some(account_status_failure(
account_id,
AccountStatus::Expired,
body,
status,
));
}
None
classified_failure(
account_id,
&facts.body,
facts.status_code,
classify_upstream_failure(facts)?,
)
}
fn classify_sse(account_id: &str, failure: &ResponsesSseFailure) -> Option<ClassifiedFailure> {
let code = failure.upstream_code.as_deref().unwrap_or_default();
let error_type = failure.upstream_type.as_deref().unwrap_or_default();
if is_model_unsupported(code) || is_model_unsupported(&failure.message) {
return Some(model_unsupported_failure(
account_id,
&crate::dispatch::failure::sse::sse_failure_error_body(failure),
));
}
let account_status =
explicit_account_status(code, error_type, &failure.message).or_else(|| {
(code.eq_ignore_ascii_case("forbidden") || error_type == "permission_error")
.then_some(AccountStatus::Banned)
})?;
Some(account_status_failure(
classified_failure(
account_id,
account_status,
&crate::dispatch::failure::sse::sse_failure_error_body(failure),
failure.explicit_status_code,
))
classify_response_failure(failure)?,
)
}
fn model_unsupported_failure(account_id: &str, message: &str) -> ClassifiedFailure {
ClassifiedFailure {
exhaustion: AccountExhaustionRecord::new(
account_id,
ExhaustedAccountKind::ModelUnsupported,
message,
),
effect: None,
}
}
fn account_status_failure(
fn classified_failure(
account_id: &str,
status: AccountStatus,
message: &str,
status_code: Option<u16>,
) -> ClassifiedFailure {
let kind = match status {
AccountStatus::Expired => ExhaustedAccountKind::Expired,
AccountStatus::Disabled => ExhaustedAccountKind::Disabled,
AccountStatus::Banned => ExhaustedAccountKind::Banned,
AccountStatus::Active | AccountStatus::QuotaExhausted => {
unreachable!("account failure controller only invalidates unavailable statuses")
}
classified: crate::fleet::account_failure::ClassifiedAccountFailure,
) -> Option<ClassifiedFailure> {
let kind = match classified.kind {
AccountFailureKind::ModelUnsupported => ExhaustedAccountKind::ModelUnsupported,
AccountFailureKind::Expired => ExhaustedAccountKind::Expired,
AccountFailureKind::Disabled => ExhaustedAccountKind::Disabled,
AccountFailureKind::Banned => ExhaustedAccountKind::Banned,
AccountFailureKind::QuotaExhausted | AccountFailureKind::RateLimited => return None,
};
let mut exhaustion = AccountExhaustionRecord::new(account_id, kind, message);
if let Some(status_code) = status_code {
exhaustion = exhaustion.with_status_code(status_code);
}
ClassifiedFailure {
Some(ClassifiedFailure {
exhaustion,
effect: Some(AccountStateEffect::SetStatus(status)),
}
}
fn explicit_account_status(code: &str, error_type: &str, message: &str) -> Option<AccountStatus> {
let code = code.trim().to_ascii_lowercase();
let error_type = error_type.trim().to_ascii_lowercase();
let message = message.to_ascii_lowercase();
if matches!(
code.as_str(),
"identity_verification_required" | "verification_required"
) || message.contains("identity verification is required")
{
return Some(AccountStatus::Disabled);
}
if matches!(
code.as_str(),
"account_banned"
| "account_deactivated"
| "account_disabled"
| "account_suspended"
| "deactivated_workspace"
| "organization_disabled"
| "workspace_deactivated"
) || message.contains("account is banned")
|| message.contains("account has been banned")
|| message.contains("account deactivated")
|| message.contains("account has been deactivated")
|| message.contains("account disabled")
|| message.contains("account has been disabled")
|| message.contains("account suspended")
|| message.contains("organization has been disabled")
|| message.contains("workspace has been deactivated")
|| message.contains("deactivated_workspace")
{
return Some(AccountStatus::Banned);
}
is_auth_failure(&code, &error_type, &message).then_some(AccountStatus::Expired)
}
fn is_auth_failure(code: &str, error_type: &str, message: &str) -> bool {
matches!(
code,
"token_invalid"
| "token_invalidated"
| "token_expired"
| "token_revoked"
| "refresh_token_invalidated"
| "unauthorized"
| "invalid_api_key"
| "authentication_error"
) || error_type == "authentication_error"
|| message.contains("token revoked")
|| message.contains("token invalidated")
|| message.contains("token invalid")
|| message.contains("token expired")
|| message.contains("unauthorized")
|| message.contains("invalid api key")
}
fn is_model_unsupported(value: &str) -> bool {
let value = value.to_ascii_lowercase();
value.contains("model_not_supported")
|| value.contains("model_not_available")
|| (value.contains("model")
&& (value.contains("not supported")
|| value.contains("not available")
|| 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")
effect: classified.effect,
})
}
@@ -1,38 +0,0 @@
//! 明确账号事实触发的运行时失效与连接驱逐。
use chrono::{DateTime, Utc};
use crate::{
fleet::{account::AccountStatus, pool::AccountPoolService},
upstream::openai::transport::CodexBackendClient,
};
pub(super) enum AccountStateEffect {
SetStatus(AccountStatus),
MarkQuotaLimitedUntil(DateTime<Utc>),
}
pub(super) struct AccountStateEffects;
impl AccountStateEffects {
pub(super) async fn apply(
account_pool: &AccountPoolService,
codex: &CodexBackendClient,
account_id: &str,
effect: &AccountStateEffect,
) {
match effect {
AccountStateEffect::SetStatus(status) => {
account_pool
.set_status_immediately(account_id, *status)
.await;
}
AccountStateEffect::MarkQuotaLimitedUntil(until) => {
account_pool
.mark_quota_limited_until_immediately(account_id, *until)
.await;
}
}
codex.evict_websocket_account(account_id).await;
}
}
@@ -9,13 +9,13 @@ use crate::{
AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind,
},
lifecycle::stream::StreamTerminal,
transport::observation::UpstreamFailureFacts,
},
fleet::{
account::AccountStatus,
cookies::{CloudflareChallengeCooldownTracker, CloudflarePathBlockTracker, PgCookieStore},
pool::AccountPoolService,
},
upstream::openai::failure::UpstreamFailureFacts,
};
pub(super) struct CloudflareController;
@@ -15,9 +15,11 @@ use crate::{
controllers::ControllerFailureFact,
errors::ClientFailure,
lifecycle::contract::{AttemptDecision, AttemptObservation, AttemptReturnKind},
transport::observation::UpstreamFailureFacts,
},
upstream::openai::protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
upstream::openai::{
failure::UpstreamFailureFacts,
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
},
};
use self::{store::Store, types::SessionState};
+1 -1
View File
@@ -13,9 +13,9 @@ use crate::{
AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind,
AttemptRoutingFacts, CompleteResponseFacts, PinnedCandidateAcquireFailureKind,
},
transport::observation::{UpstreamFailureFacts, UpstreamFailureKind},
},
upstream::openai::{
failure::{UpstreamFailureFacts, UpstreamFailureKind},
protocol::responses::{CodexResponsesRequest, PreviousResponseScope, StreamCommitPolicy},
transport::CodexBackendTransport,
transport::websocket::PreviousResponseUnavailableReason,
+4 -3
View File
@@ -49,14 +49,15 @@ use crate::{
},
stream::StreamTerminal,
},
transport::observation::UpstreamFailureFacts,
},
fleet::pool::AccountPoolService,
upstream::openai::protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
upstream::openai::{
failure::UpstreamFailureFacts,
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
},
};
mod account_failure;
mod account_state;
mod affinity;
pub(crate) mod cloudflare;
pub(crate) mod cyber_policy;
+25 -123
View File
@@ -1,23 +1,22 @@
//! 请求前配额验证规则 owner。
use chrono::{Duration, Utc};
use serde_json::Value;
use crate::{
dispatch::{
affinity::AccountIdentityService,
controllers::{
ControllerFailureFact,
account_state::{AccountStateEffect, AccountStateEffects},
},
controllers::ControllerFailureFact,
errors::ClientFailure,
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind},
lifecycle::contract::{
AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind,
},
transport::observation::UpstreamFailureFacts,
},
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,
@@ -25,13 +24,13 @@ use crate::{
},
},
upstream::openai::{
failure::UpstreamFailureFacts,
protocol::responses::ResponsesSseFailure,
transport::{CodexBackendClient, CodexRequestContext},
},
};
const LIMIT_REACHED_MESSAGE: &str = "Upstream usage quota still reports limit_reached";
const DEFAULT_RATE_LIMIT_RETRY_SECONDS: u64 = 60;
pub(super) enum QuotaEnterOutcome {
Ready(Box<AccountLease>),
@@ -167,7 +166,7 @@ impl QuotaController {
return;
};
if let Some(effect) = &classified.effect {
AccountStateEffects::apply(account_pool, codex, account_id, effect).await;
apply_account_state_effect_immediately(account_pool, codex, account_id, effect).await;
}
}
@@ -201,131 +200,34 @@ fn classify_upstream(
account_id: &str,
facts: &UpstreamFailureFacts,
) -> Option<ClassifiedQuotaFailure> {
let signaled_kind = facts
.code
.as_deref()
.and_then(classify_quota_signal)
.or_else(|| facts.error_type.as_deref().and_then(classify_quota_signal))
.or_else(|| classify_quota_message(&facts.message));
match signaled_kind.or(match facts.status_code {
Some(429) => Some(StreamQuotaKind::RateLimited),
Some(402) => Some(StreamQuotaKind::QuotaExhausted),
_ => None,
})? {
StreamQuotaKind::RateLimited => Some(rate_limited_failure(
account_id,
facts.body.clone(),
facts.retry_after_seconds,
)),
StreamQuotaKind::QuotaExhausted => {
Some(quota_exhausted_failure(account_id, facts.body.clone()))
}
}
classified_quota_failure(account_id, &facts.body, classify_upstream_failure(facts)?)
}
fn classify_stream(
account_id: &str,
failure: &ResponsesSseFailure,
) -> Option<ClassifiedQuotaFailure> {
let kind = classify_stream_signal(failure).or(match failure.explicit_status_code {
Some(429) => Some(StreamQuotaKind::RateLimited),
Some(402) => Some(StreamQuotaKind::QuotaExhausted),
_ => None,
})?;
let body = crate::dispatch::failure::sse::sse_failure_error_body(failure);
Some(match kind {
StreamQuotaKind::RateLimited => {
rate_limited_failure(account_id, body, failure.retry_after_seconds)
}
StreamQuotaKind::QuotaExhausted => quota_exhausted_failure(account_id, body),
})
classified_quota_failure(account_id, &body, classify_response_failure(failure)?)
}
#[derive(Clone, Copy)]
enum StreamQuotaKind {
RateLimited,
QuotaExhausted,
}
fn classify_stream_signal(failure: &ResponsesSseFailure) -> Option<StreamQuotaKind> {
failure
.upstream_code
.as_deref()
.and_then(classify_quota_signal)
.or_else(|| {
failure
.upstream_type
.as_deref()
.and_then(classify_quota_signal)
})
.or_else(|| classify_quota_message(&failure.message))
}
fn classify_quota_signal(signal: &str) -> Option<StreamQuotaKind> {
let signal = signal.trim().to_ascii_lowercase();
match signal.as_str() {
"usage_limit_reached"
| "rate_limit_exceeded"
| "rate_limit_reached"
| "rate_limit_error"
| "workspace_owner_usage_limit_reached"
| "workspace_member_usage_limit_reached" => Some(StreamQuotaKind::RateLimited),
"quota_exhausted"
| "quota_exceeded"
| "payment_required"
| "insufficient_quota"
| "workspace_owner_credits_depleted"
| "workspace_member_credits_depleted" => Some(StreamQuotaKind::QuotaExhausted),
signal if signal.starts_with("billing_limit") => Some(StreamQuotaKind::QuotaExhausted),
_ => None,
}
}
fn classify_quota_message(message: &str) -> Option<StreamQuotaKind> {
let message = message.to_ascii_lowercase();
if message.contains("rate limit") || message.contains("usage limit") {
return Some(StreamQuotaKind::RateLimited);
}
if message.contains("quota")
|| message.contains("payment required")
|| message.contains("billing limit")
{
return Some(StreamQuotaKind::QuotaExhausted);
}
None
}
fn rate_limited_failure(
fn classified_quota_failure(
account_id: &str,
message: String,
retry_after_seconds: Option<u64>,
) -> ClassifiedQuotaFailure {
let seconds = retry_after_seconds
.unwrap_or(DEFAULT_RATE_LIMIT_RETRY_SECONDS)
.min(i64::MAX as u64) as i64;
ClassifiedQuotaFailure {
exhaustion: AccountExhaustionRecord::new(
account_id,
ExhaustedAccountKind::RateLimited,
message,
),
effect: Some(AccountStateEffect::MarkQuotaLimitedUntil(
Utc::now() + Duration::seconds(seconds),
)),
}
}
fn quota_exhausted_failure(account_id: &str, message: String) -> ClassifiedQuotaFailure {
ClassifiedQuotaFailure {
exhaustion: AccountExhaustionRecord::new(
account_id,
ExhaustedAccountKind::QuotaExhausted,
message,
),
effect: Some(AccountStateEffect::SetStatus(
crate::fleet::account::AccountStatus::QuotaExhausted,
)),
}
message: &str,
classified: crate::fleet::account_failure::ClassifiedAccountFailure,
) -> Option<ClassifiedQuotaFailure> {
let kind = match classified.kind {
AccountFailureKind::RateLimited => ExhaustedAccountKind::RateLimited,
AccountFailureKind::QuotaExhausted => ExhaustedAccountKind::QuotaExhausted,
AccountFailureKind::ModelUnsupported
| AccountFailureKind::Expired
| AccountFailureKind::Disabled
| AccountFailureKind::Banned => return None,
};
Some(ClassifiedQuotaFailure {
exhaustion: AccountExhaustionRecord::new(account_id, kind, message),
effect: classified.effect,
})
}
fn acquired_with_verified_quota(mut acquired: AccountLease, quota: &Value) -> AccountLease {
@@ -546,6 +546,9 @@ pub(super) async fn record_live_response_stream_event(
status_code,
body,
);
if !is_success_usage_event(level, status_code) {
ensure_stream_failure_class(&mut metadata);
}
log_live_response_stream_finalized(context, status_code, level, message, &metadata);
enrich_usage_record_identity(
&mut metadata,
@@ -601,6 +604,19 @@ fn is_success_usage_event(level: UsageRecordLevel, status_code: i64) -> bool {
level != UsageRecordLevel::Error && status_code < 400
}
fn ensure_stream_failure_class(metadata: &mut Value) {
let failure_class = ["failureClass", "upstreamCode", "terminal"]
.into_iter()
.find_map(|field| metadata_string_field(metadata, field))
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string);
let (Some(failure_class), Some(object)) = (failure_class, metadata.as_object_mut()) else {
return;
};
object.insert("failureClass".to_string(), Value::String(failure_class));
}
fn log_live_response_stream_finalized(
context: &StreamContext<'_>,
status_code: i64,
+1 -1
View File
@@ -21,7 +21,6 @@ use crate::{
create_response_with_account, prepare_response_transport_with_account,
},
canonical::{CanonicalStreamDecoder, normalize_complete_response},
observation::upstream_failure_facts,
prefetch::{PrefetchedStream, StreamPrefetchError, prefetch_until_commit},
},
},
@@ -30,6 +29,7 @@ use crate::{
pool::{AccountCandidateLease, AccountPoolService},
},
upstream::openai::{
failure::upstream_failure_facts,
protocol::responses::{
CodexResponsesRequest, CollectedResponse, TransportRequirement, transport_requirement,
},
+2 -4
View File
@@ -6,13 +6,11 @@ use crate::{
dispatch::{
errors::{ClientFailure, ResponseDispatchError},
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccount},
transport::{
canonical::{CanonicalStreamBatch, CanonicalStreamDecoder},
observation::UpstreamFailureFacts,
},
transport::canonical::{CanonicalStreamBatch, CanonicalStreamDecoder},
},
fleet::{account::Account, pool::AccountLease},
upstream::openai::{
failure::UpstreamFailureFacts,
protocol::{
responses::{CodexResponsesRequest, CollectedResponse, ResponsesSseFailure},
sse::SseError,
+1 -2
View File
@@ -1,6 +1,5 @@
//! Attempt 阶段的账号作用域上游调用与事实规范化。
//! Attempt 阶段的账号作用域上游调用与流规范化。
pub(in crate::dispatch) mod account;
pub(crate) mod canonical;
pub(in crate::dispatch) mod observation;
pub(in crate::dispatch) mod prefetch;
+317
View File
@@ -0,0 +1,317 @@
//! 上游失败到账号状态 effect 的唯一领域分类。
use chrono::{Duration, Utc};
use crate::{
fleet::{account::AccountStatus, pool::AccountPoolService},
upstream::openai::{
failure::{UpstreamFailureFacts, upstream_failure_facts},
protocol::responses::ResponsesSseFailure,
transport::{CodexBackendClient, CodexClientError},
},
};
const DEFAULT_RATE_LIMIT_RETRY_SECONDS: u64 = 60;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountFailureKind {
ModelUnsupported,
Expired,
Disabled,
Banned,
QuotaExhausted,
RateLimited,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AccountStateEffect {
SetStatus(AccountStatus),
MarkQuotaLimitedUntil(chrono::DateTime<Utc>),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClassifiedAccountFailure {
pub kind: AccountFailureKind,
pub effect: Option<AccountStateEffect>,
}
impl ClassifiedAccountFailure {
pub fn account_status(&self) -> Option<AccountStatus> {
match self.effect.as_ref()? {
AccountStateEffect::SetStatus(status) => Some(*status),
AccountStateEffect::MarkQuotaLimitedUntil(_) => None,
}
}
}
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> {
classify_failure(FailureFields {
status_code: facts.status_code,
code: facts.code.as_deref(),
error_type: facts.error_type.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,
account_id: &str,
effect: &AccountStateEffect,
) {
match effect {
AccountStateEffect::SetStatus(status) => {
account_pool
.set_status_immediately(account_id, *status)
.await;
}
AccountStateEffect::MarkQuotaLimitedUntil(until) => {
account_pool
.mark_quota_limited_until_immediately(account_id, *until)
.await;
}
}
codex.evict_websocket_account(account_id).await;
}
pub async fn apply_account_state_effect(
account_pool: &AccountPoolService,
codex: &CodexBackendClient,
account_id: &str,
effect: &AccountStateEffect,
) -> bool {
let updated = match effect {
AccountStateEffect::SetStatus(status) => account_pool.set_status(account_id, *status).await,
AccountStateEffect::MarkQuotaLimitedUntil(until) => {
account_pool
.mark_quota_limited_until(account_id, *until)
.await
}
};
codex.evict_websocket_account(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>,
message: &'a str,
body: &'a str,
retry_after_seconds: Option<u64>,
}
fn classify_failure(fields: FailureFields<'_>) -> Option<ClassifiedAccountFailure> {
let code = fields.code.unwrap_or_default();
let error_type = fields.error_type.unwrap_or_default();
if [code, fields.message, fields.body]
.into_iter()
.any(is_model_unsupported)
{
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))
{
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));
}
let quota_kind = [code, error_type]
.into_iter()
.find_map(classify_quota_signal)
.or_else(|| classify_quota_message(fields.message))
.or_else(|| classify_quota_message(fields.body))
.or(match fields.status_code {
Some(429) => Some(AccountFailureKind::RateLimited),
Some(402) => Some(AccountFailureKind::QuotaExhausted),
_ => None,
})?;
Some(match quota_kind {
AccountFailureKind::RateLimited => {
let seconds = fields
.retry_after_seconds
.unwrap_or(DEFAULT_RATE_LIMIT_RETRY_SECONDS)
.min(i64::MAX as u64) as i64;
classified(
AccountFailureKind::RateLimited,
Some(AccountStateEffect::MarkQuotaLimitedUntil(
Utc::now() + Duration::seconds(seconds),
)),
)
}
AccountFailureKind::QuotaExhausted => classified(
AccountFailureKind::QuotaExhausted,
Some(AccountStateEffect::SetStatus(AccountStatus::QuotaExhausted)),
),
AccountFailureKind::ModelUnsupported
| AccountFailureKind::Expired
| AccountFailureKind::Disabled
| AccountFailureKind::Banned => unreachable!("quota classifier returned account status"),
})
}
fn status_failure(status: AccountStatus) -> ClassifiedAccountFailure {
let kind = match status {
AccountStatus::Expired => AccountFailureKind::Expired,
AccountStatus::Disabled => AccountFailureKind::Disabled,
AccountStatus::Banned => AccountFailureKind::Banned,
AccountStatus::Active | AccountStatus::QuotaExhausted => {
unreachable!("explicit account failure must make the account unavailable")
}
};
classified(kind, Some(AccountStateEffect::SetStatus(status)))
}
fn classified(
kind: AccountFailureKind,
effect: Option<AccountStateEffect>,
) -> ClassifiedAccountFailure {
ClassifiedAccountFailure { kind, effect }
}
fn explicit_account_status(code: &str, error_type: &str, message: &str) -> Option<AccountStatus> {
let code = code.trim().to_ascii_lowercase();
let error_type = error_type.trim().to_ascii_lowercase();
let message = message.to_ascii_lowercase();
if matches!(
code.as_str(),
"identity_verification_required" | "verification_required"
) || message.contains("identity verification is required")
{
return Some(AccountStatus::Disabled);
}
if matches!(
code.as_str(),
"account_banned"
| "account_deactivated"
| "account_disabled"
| "account_suspended"
| "deactivated_workspace"
| "organization_disabled"
| "workspace_deactivated"
) || message.contains("account is banned")
|| message.contains("account has been banned")
|| message.contains("account deactivated")
|| message.contains("account has been deactivated")
|| message.contains("account disabled")
|| message.contains("account has been disabled")
|| message.contains("account suspended")
|| 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);
}
is_auth_failure(&code, &error_type, &message).then_some(AccountStatus::Expired)
}
fn is_auth_failure(code: &str, error_type: &str, message: &str) -> bool {
matches!(
code,
"token_invalid"
| "token_invalidated"
| "token_expired"
| "token_revoked"
| "refresh_token_invalidated"
| "unauthorized"
| "invalid_api_key"
| "authentication_error"
) || error_type == "authentication_error"
|| message.contains("token revoked")
|| message.contains("token invalidated")
|| message.contains("token invalid")
|| message.contains("token expired")
|| message.contains("unauthorized")
|| message.contains("invalid api key")
}
fn classify_quota_signal(signal: &str) -> Option<AccountFailureKind> {
let signal = signal.trim().to_ascii_lowercase();
match signal.as_str() {
"usage_limit_reached"
| "rate_limit_exceeded"
| "rate_limit_reached"
| "rate_limit_error"
| "workspace_owner_usage_limit_reached"
| "workspace_member_usage_limit_reached" => Some(AccountFailureKind::RateLimited),
"quota_exhausted"
| "quota_exceeded"
| "payment_required"
| "insufficient_quota"
| "workspace_owner_credits_depleted"
| "workspace_member_credits_depleted" => Some(AccountFailureKind::QuotaExhausted),
signal if signal.starts_with("billing_limit") => Some(AccountFailureKind::QuotaExhausted),
_ => None,
}
}
fn classify_quota_message(message: &str) -> Option<AccountFailureKind> {
let message = message.to_ascii_lowercase();
if message.contains("rate limit") || message.contains("usage limit") {
return Some(AccountFailureKind::RateLimited);
}
if message.contains("quota")
|| message.contains("payment required")
|| message.contains("billing limit")
{
return Some(AccountFailureKind::QuotaExhausted);
}
None
}
fn is_model_unsupported(value: &str) -> bool {
let value = value.to_ascii_lowercase();
value.contains("model_not_supported")
|| value.contains("model_not_available")
|| (value.contains("model")
&& (value.contains("not supported")
|| value.contains("not available")
|| 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")
}
+8 -49
View File
@@ -7,16 +7,18 @@ use serde_json::{Value, json};
use tokio::sync::mpsc;
use crate::{
fleet::{account::AccountStatus, store::StoredAccount},
fleet::{
account::AccountStatus,
account_failure::{classify_client_failure, classify_response_failure},
store::StoredAccount,
},
models::{service::ModelRefreshPlanAccount, types::CodexModelInfo},
upstream::openai::{
protocol::{
responses::{CodexResponsesRequest, ResponsesSseFailure},
sse::{SseEvent, encode_sse_event, parse_sse_events},
},
transport::{
CodexClientError, CodexRequestContext, is_banned_auth_signal, is_banned_upstream_error,
},
transport::{CodexClientError, CodexRequestContext},
},
};
@@ -366,54 +368,11 @@ async fn process_sse_event(
}
fn account_status_from_client_error(error: &CodexClientError) -> Option<AccountStatus> {
if is_banned_upstream_error(error) {
Some(AccountStatus::Banned)
} else {
match error {
CodexClientError::Upstream { status, .. } if status.as_u16() == 402 => {
Some(AccountStatus::QuotaExhausted)
}
CodexClientError::Upstream { status, body, .. } if status.as_u16() == 401 => {
Some(if is_banned_auth_signal(body) {
AccountStatus::Banned
} else {
AccountStatus::Expired
})
}
_ => None,
}
}
classify_client_failure(error)?.account_status()
}
fn sse_failure_account_status(failure: &ResponsesSseFailure) -> Option<AccountStatus> {
let code = failure
.upstream_code
.as_deref()
.unwrap_or_default()
.to_ascii_lowercase();
let message = failure.message.to_ascii_lowercase();
if matches!(code.as_str(), "quota_exceeded" | "insufficient_quota") || message.contains("quota")
{
return Some(AccountStatus::QuotaExhausted);
}
let authentication_failed = matches!(
code.as_str(),
"token_invalid"
| "token_expired"
| "token_revoked"
| "account_deactivated"
| "unauthorized"
| "invalid_api_key"
) || message.contains("token revoked")
|| message.contains("token invalid")
|| message.contains("token expired");
authentication_failed.then(|| {
if is_banned_auth_signal(&failure.message) {
AccountStatus::Banned
} else {
AccountStatus::Expired
}
})
classify_response_failure(failure)?.account_status()
}
fn take_sse_frame(buffer: &mut String) -> Option<String> {
+9 -5
View File
@@ -1,6 +1,7 @@
use secrecy::ExposeSecret;
use super::{AccountManageService, types::AccountManageError};
use crate::fleet::account_failure::apply_client_failure_effect;
impl AccountManageService {
async fn usage_cookie_header(&self, account_id: &str) -> Option<String> {
@@ -53,11 +54,14 @@ impl AccountManageService {
client_request_id: None,
turn_id: None,
};
let raw = self
.codex
.fetch_usage(context)
.await
.map_err(|e| AccountManageError::FetchQuota(e.to_string()))?;
let raw = match self.codex.fetch_usage(context).await {
Ok(raw) => raw,
Err(error) => {
apply_client_failure_effect(&self.account_pool, &self.codex, account_id, &error)
.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)
+4 -5
View File
@@ -7,6 +7,7 @@ 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},
@@ -323,11 +324,9 @@ fn normalized_plan_type(value: &str) -> Option<String> {
}
pub(super) fn import_status_from_usage_error(error: &CodexClientError) -> Option<AccountStatus> {
if crate::upstream::openai::transport::is_banned_upstream_error(error) {
Some(AccountStatus::Banned)
} else {
None
}
classify_client_failure(error)?
.account_status()
.filter(|status| *status == AccountStatus::Banned)
}
pub(super) fn parse_account_status(status: &str) -> Result<AcctStatus, AccountManageError> {
+1
View File
@@ -1,6 +1,7 @@
//! 账号、调度、配额与刷新领域。
pub mod account;
pub mod account_failure;
pub mod cookies;
pub mod import;
pub mod manage;
+14 -13
View File
@@ -13,6 +13,7 @@ use tracing::warn;
use crate::fleet::{
account::{Account, AccountStatus},
account_failure::apply_client_failure_effect,
pool::AccountPoolService,
quota::{
quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds,
@@ -41,7 +42,7 @@ pub struct QuotaRefreshService {
request_spacing: Duration,
account_pseudonymizer: Arc<crate::infra::identity::AccountPseudonymizer>,
cookie_store: Option<PgCookieStore>,
account_pool: Option<Arc<AccountPoolService>>,
account_pool: Arc<AccountPoolService>,
}
impl QuotaRefreshService {
@@ -50,6 +51,7 @@ impl QuotaRefreshService {
store: PgAccountStore,
usage_store: PgAccountUsageStore,
codex: Arc<CodexBackendClient>,
account_pool: Arc<AccountPoolService>,
account_pseudonymizer: Arc<crate::infra::identity::AccountPseudonymizer>,
) -> Self {
Self {
@@ -60,7 +62,7 @@ impl QuotaRefreshService {
request_spacing: Self::default_request_spacing(),
account_pseudonymizer,
cookie_store: None,
account_pool: None,
account_pool,
}
}
@@ -73,6 +75,7 @@ impl QuotaRefreshService {
store: PgAccountStore,
usage_store: PgAccountUsageStore,
codex: Arc<CodexBackendClient>,
account_pool: Arc<AccountPoolService>,
account_pseudonymizer: Arc<crate::infra::identity::AccountPseudonymizer>,
min_refresh_interval_secs: u64,
) -> Self {
@@ -84,7 +87,7 @@ impl QuotaRefreshService {
request_spacing: Self::default_request_spacing(),
account_pseudonymizer,
cookie_store: None,
account_pool: None,
account_pool,
}
}
@@ -100,12 +103,6 @@ impl QuotaRefreshService {
self
}
/// 设置运行时账号池,用于刷新后同步内存调度状态。
pub fn with_account_pool(mut self, account_pool: Arc<AccountPoolService>) -> Self {
self.account_pool = Some(account_pool);
self
}
/// 执行一次配额锁定或待验证账号刷新,复用调用方传入的最近刷新记录。
pub async fn refresh_locked_accounts(
&self,
@@ -211,6 +208,13 @@ 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;
warn!(
account_id = %account.id,
error = %error,
@@ -238,10 +242,7 @@ impl QuotaRefreshService {
}
async fn sync_runtime_pool_account(&self, account_id: &str) {
let Some(account_pool) = &self.account_pool else {
return;
};
if let Err(error) = account_pool.sync_account_from_store(account_id).await {
if let Err(error) = self.account_pool.sync_account_from_store(account_id).await {
warn!(
account_id = %account_id,
error = %error,
+28 -70
View File
@@ -6,7 +6,7 @@ use thiserror::Error;
use crate::infra::format::nonnegative_i64_to_u64;
use crate::telemetry::{
account_usage::store::{PgAccountUsageStore, UsageListRecord, UsageSummary},
account_usage::store::{PgAccountUsageStore, UsageListRecord},
billing,
buckets::query::{
ModelBucketUsage, ModelUsageWindow, PgRequestBucketQuery, UsageBucketTotals,
@@ -42,13 +42,25 @@ impl AccountUsageQueryService {
.map_err(|_| AccountUsageQueryError::List)
}
/// 汇总账号用量。
pub async fn summary(&self) -> Result<AccountUsageSummary, AccountUsageQueryError> {
self.store
.usage_summary()
/// 汇总请求时间桶保留期内的全局用量与计费。
pub async fn retained_summary(&self) -> Result<RetainedUsageSummary, AccountUsageQueryError> {
let buckets = self
.buckets
.retained_usage_buckets()
.await
.map(AccountUsageSummary::from)
.map_err(|_| AccountUsageQueryError::Summary)
.map_err(|_| AccountUsageQueryError::RetainedSummary)?;
let mut summary = RetainedUsageSummary::default();
for bucket in buckets {
summary.billing_amount_usd += billing::calculate_billing_amount(
nonnegative_i64_to_u64(bucket.totals.input_tokens),
nonnegative_i64_to_u64(bucket.totals.output_tokens),
nonnegative_i64_to_u64(bucket.totals.cached_tokens),
&bucket.model,
bucket.service_tier.as_deref(),
);
summary.totals.add(bucket.totals);
}
Ok(summary)
}
/// 列出指定时间范围内的聚合时间桶。
@@ -95,12 +107,19 @@ impl AccountUsageQueryService {
pub enum AccountUsageQueryError {
#[error("failed to list account usage")]
List,
#[error("failed to summarize account usage")]
Summary,
#[error("failed to summarize retained usage")]
RetainedSummary,
#[error("failed to list usage time buckets")]
TimeBuckets,
}
/// 请求时间桶保留期内的全局用量与计费汇总。
#[derive(Debug, Clone, Default, PartialEq)]
pub struct RetainedUsageSummary {
pub totals: UsageBucketTotals,
pub billing_amount_usd: f64,
}
/// 管理端账号用量记录。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountUsageRecord {
@@ -129,23 +148,6 @@ pub struct AccountUsageRecord {
pub last_used_at: Option<DateTime<Utc>>,
}
/// 管理端账号用量汇总。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountUsageSummary {
pub account_count: i64,
pub request_count: i64,
pub empty_response_count: i64,
pub input_tokens: i64,
pub output_tokens: i64,
pub cached_tokens: i64,
pub reasoning_tokens: i64,
pub total_tokens: i64,
pub image_input_tokens: i64,
pub image_output_tokens: i64,
pub image_request_count: i64,
pub image_request_failed_count: i64,
}
/// 管理端时间桶用量记录。
#[derive(Debug, Clone, PartialEq)]
pub struct AccountUsageTimeBucket {
@@ -163,31 +165,10 @@ pub struct AccountUsageTimeBucket {
pub latency_count: i64,
pub max_latency_ms: i64,
pub min_latency_ms: i64,
pub billing_amount_usd: Option<f64>,
}
impl From<UsageSummary> for AccountUsageSummary {
fn from(s: UsageSummary) -> Self {
Self {
account_count: s.account_count,
request_count: s.request_count,
empty_response_count: s.empty_response_count,
input_tokens: s.input_tokens,
output_tokens: s.output_tokens,
cached_tokens: s.cached_tokens,
reasoning_tokens: s.reasoning_tokens,
total_tokens: s.total_tokens,
image_input_tokens: s.image_input_tokens,
image_output_tokens: s.image_output_tokens,
image_request_count: s.image_request_count,
image_request_failed_count: s.image_request_failed_count,
}
}
}
impl From<UsageTimeBucketRecord> for AccountUsageTimeBucket {
fn from(record: UsageTimeBucketRecord) -> Self {
let billing_amount_usd = time_bucket_billing_amount_usd(&record);
Self {
bucket_start: record.bucket_start,
model: record.model,
@@ -203,33 +184,10 @@ impl From<UsageTimeBucketRecord> for AccountUsageTimeBucket {
latency_count: record.latency_count,
max_latency_ms: record.max_latency_ms,
min_latency_ms: record.min_latency_ms,
billing_amount_usd,
}
}
}
fn time_bucket_billing_amount_usd(record: &UsageTimeBucketRecord) -> Option<f64> {
let input_tokens = nonnegative_i64_to_u64(record.input_tokens);
let output_tokens = nonnegative_i64_to_u64(record.output_tokens);
let cached_tokens = nonnegative_i64_to_u64(record.cached_tokens);
if input_tokens == 0 && output_tokens == 0 && cached_tokens == 0 {
return None;
}
let model = record.model.trim();
if model.is_empty() {
return None;
}
Some(billing::calculate_billing_amount(
input_tokens,
output_tokens,
cached_tokens,
model,
record.service_tier.as_deref(),
))
}
impl From<UsageListRecord> for AccountUsageRecord {
fn from(usage: UsageListRecord) -> Self {
Self {
@@ -129,22 +129,6 @@ select
from account_usage au
left join accounts a on a.id = au.account_id";
const USAGE_SUMMARY_SQL: &str = r"
select
count(*) as account_count,
coalesce(sum(request_count), 0)::bigint as request_count,
coalesce(sum(empty_response_count), 0)::bigint as empty_response_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(reasoning_tokens), 0)::bigint as reasoning_tokens,
coalesce(sum(total_tokens), 0)::bigint as total_tokens,
coalesce(sum(image_input_tokens), 0)::bigint as image_input_tokens,
coalesce(sum(image_output_tokens), 0)::bigint as image_output_tokens,
coalesce(sum(image_request_count), 0)::bigint as image_request_count,
coalesce(sum(image_request_failed_count), 0)::bigint as image_request_failed_count
from account_usage";
/// PostgreSQL 用量存储错误。
#[derive(Debug, Error)]
pub enum PgAccountUsageStoreError {
@@ -312,35 +296,6 @@ pub struct UsageListRecord {
pub last_used_at: Option<DateTime<Utc>>,
}
/// 账号用量汇总。
#[derive(Debug, Clone)]
pub struct UsageSummary {
/// 账号数。
pub account_count: i64,
/// 总请求数。
pub request_count: i64,
/// 总空响应数。
pub empty_response_count: i64,
/// 总输入 token。
pub input_tokens: i64,
/// 总输出 token。
pub output_tokens: i64,
/// 总缓存 token。
pub cached_tokens: i64,
/// 总 reasoning token。
pub reasoning_tokens: i64,
/// 总 token。
pub total_tokens: i64,
/// 总图片输入 token。
pub image_input_tokens: i64,
/// 总图片输出 token。
pub image_output_tokens: i64,
/// 总图片请求成功数。
pub image_request_count: i64,
/// 总图片请求失败数。
pub image_request_failed_count: i64,
}
/// PostgreSQL 用量存储。
#[derive(Clone)]
pub struct PgAccountUsageStore {
@@ -492,25 +447,6 @@ where account_id in (",
let rows = builder.build().fetch_all(&self.pool).await?;
rows.iter().map(usage_list_from_row).collect()
}
/// 汇总账号用量。
pub async fn usage_summary(&self) -> PgAccountUsageStoreResult<UsageSummary> {
let row = sqlx::query(USAGE_SUMMARY_SQL).fetch_one(&self.pool).await?;
Ok(UsageSummary {
account_count: row.get("account_count"),
request_count: row.get("request_count"),
empty_response_count: row.get("empty_response_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
reasoning_tokens: row.get("reasoning_tokens"),
total_tokens: row.get("total_tokens"),
image_input_tokens: row.get("image_input_tokens"),
image_output_tokens: row.get("image_output_tokens"),
image_request_count: row.get("image_request_count"),
image_request_failed_count: row.get("image_request_failed_count"),
})
}
}
#[async_trait]
+42 -1
View File
@@ -29,6 +29,17 @@ where bucket_start >= $1 and bucket_start <= $2
group by bucket_start, model, service_tier
order by bucket_start asc, model asc, service_tier asc";
const RETAINED_USAGE_BUCKETS_SQL: &str = r"
select
model,
nullif(service_tier, '__unknown__') as service_tier,
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
from request_time_buckets
group by bucket_start, model, service_tier";
/// 请求时间桶查询错误。
#[derive(Debug, Error)]
pub enum PgRequestBucketQueryError {
@@ -72,8 +83,16 @@ pub struct UsageBucketTotals {
pub cached_tokens: i64,
}
/// 保留期内单个计费桶的用量事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RetainedUsageBucket {
pub model: String,
pub service_tier: Option<String>,
pub totals: UsageBucketTotals,
}
impl UsageBucketTotals {
fn add(&mut self, other: Self) {
pub(crate) fn add(&mut self, other: Self) {
self.request_count = self.request_count.saturating_add(other.request_count);
self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
@@ -127,6 +146,28 @@ impl PgRequestBucketQuery {
Ok(rows.iter().map(usage_time_bucket_from_row).collect())
}
/// 列出请求时间桶保留期内的全局计费桶,不受当前账号集合影响。
pub async fn retained_usage_buckets(
&self,
) -> Result<Vec<RetainedUsageBucket>, PgRequestBucketQueryError> {
let rows = sqlx::query(RETAINED_USAGE_BUCKETS_SQL)
.fetch_all(&self.pool)
.await?;
Ok(rows
.iter()
.map(|row| RetainedUsageBucket {
model: row.get("model"),
service_tier: row.get("service_tier"),
totals: UsageBucketTotals {
request_count: row.get("request_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
},
})
.collect())
}
pub async fn usage_by_windows(
&self,
windows: &[UsageBucketWindow],
+14 -27
View File
@@ -11,7 +11,7 @@ use crate::infra::{
time::{china_day_start, china_quarter_hour_start},
};
use super::account_usage::query::{AccountUsageSummary, AccountUsageTimeBucket};
use super::account_usage::query::{AccountUsageTimeBucket, RetainedUsageSummary};
const HEALTH_TIMELINE_SLOT_MINUTES: i64 = 15;
const HEALTH_TIMELINE_SLOTS: i64 = 24 * 4;
@@ -79,7 +79,7 @@ struct DashboardTokenCardData {
today_tokens_value: u64,
yesterday_tokens_value: u64,
total_tokens: String,
today_billing_amount_usd: String,
total_billing_amount_usd: String,
}
#[derive(Debug, Serialize)]
@@ -174,25 +174,25 @@ impl UsageWindow {
pub fn dashboard_cards(
accounts: DashboardAccountCounts,
buckets: &[AccountUsageTimeBucket],
lifetime_usage: &AccountUsageSummary,
retained_usage: &RetainedUsageSummary,
) -> DashboardCardsData {
dashboard_cards_at(accounts, buckets, lifetime_usage, Utc::now())
dashboard_cards_at(accounts, buckets, retained_usage, Utc::now())
}
pub fn dashboard_cards_at(
accounts: DashboardAccountCounts,
buckets: &[AccountUsageTimeBucket],
lifetime_usage: &AccountUsageSummary,
retained_usage: &RetainedUsageSummary,
now: DateTime<Utc>,
) -> DashboardCardsData {
let today_start = china_day_start(now);
let yesterday_start = today_start - Duration::days(1);
let today = usage_window(buckets, today_start, now);
let yesterday = usage_window(buckets, yesterday_start, today_start);
let today_billing_amount = billing_window(buckets, today_start, now).unwrap_or(0.0);
let total_requests = nonnegative_i64_to_u64(lifetime_usage.request_count);
let total_tokens = nonnegative_i64_to_u64(lifetime_usage.total_tokens);
let total_cached_tokens = nonnegative_i64_to_u64(lifetime_usage.cached_tokens);
let total_requests = nonnegative_i64_to_u64(retained_usage.totals.request_count);
let total_tokens = nonnegative_i64_to_u64(retained_usage.totals.input_tokens)
.saturating_add(nonnegative_i64_to_u64(retained_usage.totals.output_tokens));
let total_cached_tokens = nonnegative_i64_to_u64(retained_usage.totals.cached_tokens);
DashboardCardsData {
accounts: DashboardAccountsCardData {
@@ -214,13 +214,13 @@ pub fn dashboard_cards_at(
today_tokens_value: today.tokens(),
yesterday_tokens_value: yesterday.tokens(),
total_tokens: format_tokens(total_tokens),
today_billing_amount_usd: format_billing_amount(today_billing_amount),
total_billing_amount_usd: format_billing_amount(retained_usage.billing_amount_usd),
},
cache: DashboardCacheCardData {
today_hit_rate: format_rate(today.cache_hit_rate()),
today_hit_rate_value: today.cache_hit_rate(),
yesterday_hit_rate_value: yesterday.cache_hit_rate(),
total_hit_rate: format_rate(summary_cache_hit_rate(lifetime_usage)),
total_hit_rate: format_rate(summary_cache_hit_rate(retained_usage)),
total_cached_tokens: format_tokens(total_cached_tokens),
average_first_token_latency_ms: format_optional_duration_ms(
today.avg_first_token_latency(),
@@ -486,19 +486,6 @@ fn usage_window(
})
}
fn billing_window(
records: &[AccountUsageTimeBucket],
start: DateTime<Utc>,
end: DateTime<Utc>,
) -> Option<f64> {
let billing_amounts = records
.iter()
.filter(|record| record.bucket_start >= start && record.bucket_start < end)
.filter_map(|record| record.billing_amount_usd)
.collect::<Vec<_>>();
(!billing_amounts.is_empty()).then(|| billing_amounts.into_iter().sum())
}
fn apply_bucket(window: &mut UsageWindow, record: &AccountUsageTimeBucket) {
window.requests += nonnegative_i64_to_u64(record.request_count);
window.input_tokens += nonnegative_i64_to_u64(record.input_tokens);
@@ -540,10 +527,10 @@ fn health_tone(bucket: UsageWindow, success_rate: f64, is_future: bool) -> char
}
}
fn summary_cache_hit_rate(summary: &AccountUsageSummary) -> Option<f64> {
let input_tokens = nonnegative_i64_to_u64(summary.input_tokens);
fn summary_cache_hit_rate(summary: &RetainedUsageSummary) -> Option<f64> {
let input_tokens = nonnegative_i64_to_u64(summary.totals.input_tokens);
(input_tokens > 0)
.then(|| nonnegative_i64_to_u64(summary.cached_tokens) as f64 / input_tokens as f64)
.then(|| nonnegative_i64_to_u64(summary.totals.cached_tokens) as f64 / input_tokens as f64)
}
fn format_optional_duration_ms(value: Option<u64>) -> String {
+35 -11
View File
@@ -49,7 +49,13 @@ error_terminals as (
case when request_id is null then 'event:' || id else 'request:' || request_id end as request_key,
id, request_id, created_at as terminal_at, client_api_key_id, provider,
account_id, transport, coalesce(nullif(model, ''), '未知模型') as model_name,
coalesce(nullif(failure_class, ''), '未分类错误') as failure_class,
coalesce(
nullif(failure_class, ''),
nullif(metadata_json->>'failureClass', ''),
nullif(metadata_json->>'upstreamCode', ''),
nullif(metadata_json->>'terminal', ''),
'未分类错误'
) as failure_class,
latency_ms
from ops_error_logs
where created_at >= $1 and created_at < $2
@@ -162,12 +168,28 @@ impl UsageDiagnosticsDimension {
fn expression(self) -> &'static str {
match self {
Self::Model => "model_name",
Self::Account => "coalesce(nullif(account_id, ''), '未知账号')",
Self::ApiKey => "coalesce(nullif(client_api_key_id, ''), '未知 API Key')",
Self::Provider => "coalesce(nullif(provider, ''), '未知 Provider')",
Self::Transport => "coalesce(nullif(transport, ''), '未知传输')",
Self::FailureClass => "failure_class",
Self::Model => "terminal_requests.model_name",
Self::Account => {
"coalesce(nullif(concat_ws(' → ', nullif(diagnostic_accounts.email, ''), nullif(terminal_requests.account_id, '')), ''), '未知账号')"
}
Self::ApiKey => "coalesce(nullif(diagnostic_api_keys.name, ''), '未知密钥')",
Self::Provider => "coalesce(nullif(terminal_requests.provider, ''), '未知 Provider')",
Self::Transport => "coalesce(nullif(terminal_requests.transport, ''), '未知传输')",
Self::FailureClass => "terminal_requests.failure_class",
}
}
fn source(self) -> &'static str {
match self {
Self::Account => {
"terminal_requests left join accounts diagnostic_accounts on diagnostic_accounts.id = terminal_requests.account_id"
}
Self::ApiKey => {
"terminal_requests left join client_api_keys diagnostic_api_keys on diagnostic_api_keys.id = terminal_requests.client_api_key_id"
}
Self::Model | Self::Provider | Self::Transport | Self::FailureClass => {
"terminal_requests"
}
}
}
}
@@ -307,8 +329,9 @@ pub async fn diagnostics(
dimension: UsageDiagnosticsDimension,
) -> PgUsageRecordStoreResult<UsageDiagnosticsInsights> {
let expression = dimension.expression();
let source = dimension.source();
let failure_filter = if dimension == UsageDiagnosticsDimension::FailureClass {
"where not is_success"
"where not terminal_requests.is_success"
} else {
""
};
@@ -321,7 +344,7 @@ pub async fn diagnostics(
percentile_cont(0.95) within group (order by latency_ms)\n\
filter (where is_success and latency_ms is not null) as latency_p95_ms,\n\
(select count(*) from terminal_requests)::bigint as total_requests\n\
from terminal_requests {failure_filter}\n\
from {source} {failure_filter}\n\
group by name\n\
order by request_count desc, name asc\n\
limit 8"
@@ -610,6 +633,7 @@ async fn load_diagnostic_costs(
dimension: UsageDiagnosticsDimension,
) -> PgUsageRecordStoreResult<HashMap<String, f64>> {
let expression = dimension.expression();
let source = dimension.source();
let sql = format!(
"with {TERMINAL_REQUESTS_CTE}\n\
select {expression} as name, billing_model, service_tier,\n\
@@ -622,8 +646,8 @@ async fn load_diagnostic_costs(
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\
from terminal_requests\n\
where is_success\n\
from {source}\n\
where terminal_requests.is_success\n\
group by name, billing_model, service_tier, is_long"
);
// SAFETY: the grouping expression comes from the closed diagnostics dimension enum.
@@ -1,4 +1,4 @@
//! 上游原始错误到 attempt 通用事实的单次规范化。
//! Codex 上游原始错误到稳定失败事实的规范化。
use serde_json::Value;
@@ -8,7 +8,7 @@ use crate::upstream::openai::transport::{
};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(in crate::dispatch) enum UpstreamFailureKind {
pub enum UpstreamFailureKind {
HttpStatus,
HttpConnect,
HttpTimeout,
@@ -25,7 +25,7 @@ pub(in crate::dispatch) enum UpstreamFailureKind {
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub(in crate::dispatch) struct UpstreamFailureFacts {
pub struct UpstreamFailureFacts {
pub kind: UpstreamFailureKind,
pub status_code: Option<u16>,
pub code: Option<String>,
@@ -35,9 +35,7 @@ pub(in crate::dispatch) struct UpstreamFailureFacts {
pub retry_after_seconds: Option<u64>,
}
pub(in crate::dispatch) fn upstream_failure_facts(
error: &CodexClientError,
) -> UpstreamFailureFacts {
pub fn upstream_failure_facts(error: &CodexClientError) -> UpstreamFailureFacts {
match error {
CodexClientError::Http(error) => UpstreamFailureFacts {
kind: if error.is_connect() {
@@ -156,6 +154,7 @@ fn error_fields(body: &str) -> (Option<String>, Option<String>, Option<String>)
let error = value
.pointer("/response/error")
.or_else(|| value.get("error"))
.or_else(|| value.get("detail"))
.unwrap_or(&value);
(
error
+1
View File
@@ -1,5 +1,6 @@
//! OpenAI/Codex 上游客户端。
pub mod failure;
pub mod fingerprint;
pub mod protocol;
pub mod token_client;
@@ -166,39 +166,6 @@ impl CodexClientError {
}
}
/// 判断上游错误正文是否表示账号已封禁或停用。
pub fn is_banned_auth_signal(value: &str) -> bool {
let value = value.to_ascii_lowercase();
value.contains("account_deactivated")
|| value.contains("account deactivated")
|| value.contains("account has been deactivated")
|| value.contains("deactivated")
|| value.contains("banned")
}
/// 判断 402 错误正文是否带有 OpenAI 工作区停用标记。
pub fn is_deactivated_workspace_error_body(value: &str) -> bool {
serde_json::from_str::<Value>(value).is_ok_and(|value| {
value.pointer("/detail/code").and_then(Value::as_str) == Some("deactivated_workspace")
})
}
/// 判断 Codex 上游错误是否表示账号已封禁或停用。
pub fn is_banned_upstream_error(error: &CodexClientError) -> bool {
matches!(
error,
CodexClientError::Upstream { status, body, .. }
if (status.as_u16() == 403
&& !is_html_error_body(body))
|| (status.as_u16() == 402 && is_deactivated_workspace_error_body(body))
)
}
fn is_html_error_body(value: &str) -> bool {
let value = value.trim_start().to_ascii_lowercase();
value.starts_with("<!doctype") || value.starts_with("<html") || value.contains("<html")
}
/// Codex 客户端结果类型。
pub type CodexClientResult<T> = Result<T, CodexClientError>;
@@ -342,7 +309,7 @@ pub enum CodexTransportDecision {
ConnectedWebSocket,
ExactWebSocket,
RequiredWebSocket,
Http2WebSocketSlow,
Http2WebSocketBudgetExhausted,
Http2BreakerOpen,
Http2PoolUnavailable,
Http2PreSendFailure,
@@ -356,7 +323,7 @@ impl CodexTransportDecision {
Self::ConnectedWebSocket => "ws_connected_fast",
Self::ExactWebSocket => "ws_exact_required",
Self::RequiredWebSocket => "ws_required",
Self::Http2WebSocketSlow => "http2_ws_slow",
Self::Http2WebSocketBudgetExhausted => "http2_ws_budget_exhausted",
Self::Http2BreakerOpen => "http2_breaker_open",
Self::Http2PoolUnavailable => "http2_pool_unavailable",
Self::Http2PreSendFailure => "http2_ws_pre_send_failure",
@@ -660,7 +627,7 @@ fn websocket_success_decision(
fn http_fallback_decision(error: &CodexWebSocketExchangeError) -> CodexTransportDecision {
match error {
CodexWebSocketExchangeError::FastPathTimeout { .. } => {
CodexTransportDecision::Http2WebSocketSlow
CodexTransportDecision::Http2WebSocketBudgetExhausted
}
CodexWebSocketExchangeError::OriginCircuitOpen
| CodexWebSocketExchangeError::OriginHalfOpenBusy => {
@@ -314,11 +314,7 @@ impl CodexBackendClient {
route: PreparedResponseRoute::Http,
metrics: CodexTransportMetrics {
decision: Some(decision),
ws_connect_ms: matches!(
error,
CodexWebSocketExchangeError::FastPathTimeout { .. }
)
.then_some(wait_ms),
ws_connect_ms: None,
transport_decision_wait_ms: Some(wait_ms),
..CodexTransportMetrics::default()
},
@@ -20,7 +20,6 @@ pub use self::{
CodexModelCatalogClient, CodexModelCatalogClientError, CodexModelCatalogRequest,
CodexRateLimitHeaderUpdates, CodexRequestContext, CodexResponseMetadata,
CodexTransportDecision, CodexTransportMetrics, CodexTurnStateUpdate, build_reqwest_client,
is_banned_auth_signal, is_banned_upstream_error, is_deactivated_workspace_error_body,
},
diagnostics::CodexUpstreamDiagnostics,
endpoints::{
@@ -12,7 +12,7 @@ use bytes::Bytes;
use chrono::Utc;
use futures::Stream;
use thiserror::Error;
use tokio::sync::mpsc;
use tokio::sync::{mpsc, oneshot};
use tokio::{sync::Mutex, time::timeout};
use tokio_tungstenite::{Connector, connect_async_tls_with_config};
use tungstenite::{
@@ -39,12 +39,13 @@ use crate::upstream::openai::protocol::websocket::{
use super::websocket_breaker::{
WebSocketOriginBreaker, WebSocketOriginBreakerDecision, WebSocketOriginBreakerPermit,
WebSocketOriginFastPathReporter,
};
use super::websocket_pool::{
CodexWebSocketConnectionMetadata, CodexWebSocketPool, CodexWebSocketPoolKey,
PooledWebSocketConnection, WebSocketContinuationState, WebSocketPoolAcquire,
WebSocketPoolBypassReason, WebSocketPoolConnectOutcome, WebSocketPoolDecision,
WebSocketPoolLease,
WebSocketPoolBypassReason, WebSocketPoolConnectLease, WebSocketPoolConnectOutcome,
WebSocketPoolDecision, WebSocketPoolLease,
};
use super::websocket_pump::{PumpExitReason, PumpKeepalive, PumpedWebSocket, RawWsStream};
use super::{client::CodexResponseMetadata, diagnostics::CodexUpstreamDiagnostics, response_meta};
@@ -576,7 +577,7 @@ async fn prepare_unpooled_websocket(
async fn prepare_pooled_websocket(
request: &CodexWebSocketRequest,
pool: &CodexWebSocketPool,
connect_lease: super::websocket_pool::WebSocketPoolConnectLease,
connect_lease: WebSocketPoolConnectLease,
breaker: &WebSocketOriginBreaker,
origin_key: &str,
fast_path_budget: Option<Duration>,
@@ -589,27 +590,21 @@ async fn prepare_pooled_websocket(
return Err(error);
}
};
let connected =
connect_with_budget(request.connection(), pool.keepalive(), fast_path_budget).await;
let (connection, connect_elapsed) = match finish_breaker_attempt(permit, connected) {
Ok(connected) => connected,
Err(error) => {
connect_lease.failed().await;
return Err(error);
}
};
let connected = match connect_lease.connected(connection).await {
Ok(connected) => connected,
Err(connection) => {
connection.websocket.close().await;
return Err(continuation_unavailable(
PreviousResponseUnavailableReason::PoolUnavailable,
));
}
let mut waiter = start_pooled_websocket_connect(
request.connection().clone(),
pool.clone(),
connect_lease,
permit,
);
let connect_elapsed = waiter.wait(fast_path_budget).await?;
let Some((connection, lease)) = pool.take_idle(waiter.key()).await else {
return Err(continuation_unavailable(
PreviousResponseUnavailableReason::PoolUnavailable,
));
};
Ok(PreparedWebSocket {
connection: *connected.connection,
lease: Some(connected.lease),
connection: *connection,
lease: Some(lease),
reused: false,
pool_decision: Some(WebSocketPoolDecision::new()),
connect_elapsed: Some(connect_elapsed),
@@ -618,6 +613,131 @@ async fn prepare_pooled_websocket(
})
}
struct PooledWebSocketConnectWaiter {
key: CodexWebSocketPoolKey,
started_at: tokio::time::Instant,
receiver: oneshot::Receiver<Result<Duration, CodexWebSocketExchangeError>>,
fast_path_reporter: WebSocketOriginFastPathReporter,
}
impl PooledWebSocketConnectWaiter {
fn key(&self) -> &CodexWebSocketPoolKey {
&self.key
}
async fn wait(
&mut self,
fast_path_budget: Option<Duration>,
) -> Result<Duration, CodexWebSocketExchangeError> {
let received = match fast_path_budget {
Some(budget) => {
let remaining = budget.saturating_sub(self.started_at.elapsed());
match timeout(remaining, &mut self.receiver).await {
Ok(received) => received,
Err(_) => {
self.fast_path_reporter.missed();
return Err(CodexWebSocketExchangeError::FastPathTimeout {
timeout: budget,
});
}
}
}
None => (&mut self.receiver).await,
};
received.map_err(|_| CodexWebSocketExchangeError::SharedConnectFailed)?
}
}
fn start_pooled_websocket_connect(
connection: CodexWebSocketConnection,
pool: CodexWebSocketPool,
connect_lease: WebSocketPoolConnectLease,
permit: WebSocketOriginBreakerPermit,
) -> PooledWebSocketConnectWaiter {
let key = connect_lease.key().clone();
let started_at = connect_lease.started_at();
let cancellation = connect_lease.cancellation_token();
let fast_path_reporter = permit.fast_path_reporter();
let keepalive = pool.keepalive();
let task_key = key.clone();
let (sender, receiver) = oneshot::channel();
pool.spawn_connect_task(async move {
let connected = tokio::select! {
biased;
_ = cancellation.cancelled() => {
permit.cancel();
connect_lease.failed().await;
let _ = sender.send(Err(CodexWebSocketExchangeError::SharedConnectFailed));
tracing::info!(
account_id = task_key.account_id(),
conversation_id_hash = task_key.conversation_id_hash(),
ws_preconnect_duration_ms = duration_millis_u64(started_at.elapsed()),
ws_preconnect_outcome = "cancelled",
"WebSocket pool connect finished"
);
return;
}
result = connect_with_budget(&connection, keepalive, None) => result,
};
match finish_breaker_attempt(permit, connected) {
Ok((connection, connect_elapsed)) => {
match connect_lease.connected_idle(connection).await {
Ok(()) => {
let foreground_waiting = sender.send(Ok(connect_elapsed)).is_ok();
tracing::info!(
account_id = task_key.account_id(),
conversation_id_hash = task_key.conversation_id_hash(),
ws_preconnect_duration_ms = duration_millis_u64(connect_elapsed),
foreground_waiting,
ws_preconnect_outcome = "ready",
"WebSocket pool connect finished"
);
}
Err(connection) => {
connection.websocket.close().await;
let _ = sender.send(Err(continuation_unavailable(
PreviousResponseUnavailableReason::PoolUnavailable,
)));
tracing::info!(
account_id = task_key.account_id(),
conversation_id_hash = task_key.conversation_id_hash(),
ws_preconnect_duration_ms = duration_millis_u64(connect_elapsed),
ws_preconnect_outcome = "rejected",
"WebSocket pool connect finished"
);
}
}
}
Err(error) => {
let error_message = error.to_string();
connect_lease.failed().await;
let foreground_waiting = sender.send(Err(error)).is_ok();
tracing::warn!(
account_id = task_key.account_id(),
conversation_id_hash = task_key.conversation_id_hash(),
ws_preconnect_duration_ms = duration_millis_u64(started_at.elapsed()),
foreground_waiting,
error = %error_message,
ws_preconnect_outcome = "failed",
"WebSocket pool connect finished"
);
}
}
});
PooledWebSocketConnectWaiter {
key,
started_at,
receiver,
fast_path_reporter,
}
}
fn duration_millis_u64(duration: Duration) -> u64 {
u64::try_from(duration.as_millis())
.unwrap_or(u64::MAX)
.max(1)
}
async fn connect_with_budget(
connection: &CodexWebSocketConnection,
keepalive: PumpKeepalive,
@@ -683,11 +803,11 @@ fn finish_breaker_attempt(
async fn wait_for_shared_connect(
waiter: super::websocket_pool::WebSocketPoolConnectWaiter,
fast_path_budget: Option<Duration>,
decision_started_at: Instant,
_decision_started_at: Instant,
) -> Result<WebSocketPoolConnectOutcome, CodexWebSocketExchangeError> {
match fast_path_budget {
Some(budget) => {
let remaining = budget.saturating_sub(decision_started_at.elapsed());
let remaining = waiter.remaining_budget(budget);
timeout(remaining, waiter.wait())
.await
.map_err(|_| CodexWebSocketExchangeError::FastPathTimeout { timeout: budget })
@@ -2,7 +2,10 @@
use std::{
collections::{HashMap, VecDeque},
sync::{Arc, Mutex},
sync::{
Arc, Mutex,
atomic::{AtomicBool, Ordering},
},
time::{Duration, Instant},
};
@@ -73,8 +76,8 @@ impl WebSocketOriginBreaker {
.or_insert_with(CircuitState::closed);
match state {
CircuitState::Closed { fast_timeouts } => {
retain_window(fast_timeouts, now, self.config.failure_window);
CircuitState::Closed { failures } => {
retain_window(failures, now, self.config.failure_window);
WebSocketOriginBreakerDecision::Allowed(WebSocketOriginBreakerPermit::new(
self.clone(),
origin_key.to_string(),
@@ -105,7 +108,7 @@ impl WebSocketOriginBreaker {
}
}
fn record_fast_timeout(&self, origin_key: &str, probe_id: Option<Uuid>) {
fn record_failure(&self, origin_key: &str, probe_id: Option<Uuid>) {
let now = Instant::now();
let mut circuits = self
.inner
@@ -127,35 +130,17 @@ impl WebSocketOriginBreaker {
let state = circuits
.entry(origin_key.to_string())
.or_insert_with(CircuitState::closed);
let CircuitState::Closed { fast_timeouts } = state else {
let CircuitState::Closed { failures } = state else {
return;
};
retain_window(fast_timeouts, now, self.config.failure_window);
fast_timeouts.push_back(now);
if fast_timeouts.len() >= self.config.failure_threshold {
retain_window(failures, now, self.config.failure_window);
failures.push_back(now);
if failures.len() >= self.config.failure_threshold {
*state = CircuitState::Open {
until: now + self.config.open_duration,
};
}
}
fn record_failure(&self, origin_key: &str, probe_id: Option<Uuid>) {
let Some(probe_id) = probe_id else {
return;
};
let mut circuits = self
.inner
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if permit_still_owns_state(circuits.get(origin_key), Some(probe_id)) {
circuits.insert(
origin_key.to_string(),
CircuitState::Open {
until: Instant::now() + self.config.open_duration,
},
);
}
}
}
/// 一次新的 WebSocket opening 许可;消费式完成保证 half-open 探针不会泄漏。
@@ -163,6 +148,7 @@ pub struct WebSocketOriginBreakerPermit {
breaker: WebSocketOriginBreaker,
origin_key: String,
probe_id: Option<Uuid>,
degradation_recorded: Arc<AtomicBool>,
armed: bool,
}
@@ -172,6 +158,7 @@ impl WebSocketOriginBreakerPermit {
breaker,
origin_key,
probe_id,
degradation_recorded: Arc::new(AtomicBool::new(false)),
armed: true,
}
}
@@ -180,33 +167,71 @@ impl WebSocketOriginBreakerPermit {
self.probe_id.is_some()
}
/// 创建可跨前台预算与后台 opening 生命周期共享的快路径观察器。
pub(crate) fn fast_path_reporter(&self) -> WebSocketOriginFastPathReporter {
WebSocketOriginFastPathReporter {
breaker: self.breaker.clone(),
origin_key: self.origin_key.clone(),
probe_id: self.probe_id,
degradation_recorded: Arc::clone(&self.degradation_recorded),
}
}
pub fn succeed(mut self) {
self.breaker.record_success(&self.origin_key, self.probe_id);
if !self.degradation_recorded.load(Ordering::Acquire) {
self.breaker.record_success(&self.origin_key, self.probe_id);
}
self.armed = false;
}
pub fn fast_timeout(mut self) {
self.breaker
.record_fast_timeout(&self.origin_key, self.probe_id);
self.fast_path_reporter().missed();
self.armed = false;
}
pub fn fail(mut self) {
self.breaker.record_failure(&self.origin_key, self.probe_id);
self.record_failure_once();
self.armed = false;
}
/// opening 因账号驱逐或服务关闭而取消,不污染 origin 可用性。
pub(crate) fn cancel(mut self) {
self.armed = false;
}
fn record_failure_once(&self) {
if !self.degradation_recorded.swap(true, Ordering::AcqRel) {
self.breaker.record_failure(&self.origin_key, self.probe_id);
}
}
}
impl Drop for WebSocketOriginBreakerPermit {
fn drop(&mut self) {
if self.armed {
self.record_failure_once();
}
}
}
/// 前台等待结束后仍可报告同一次后台 opening 错过快路径预算。
pub(crate) struct WebSocketOriginFastPathReporter {
breaker: WebSocketOriginBreaker,
origin_key: String,
probe_id: Option<Uuid>,
degradation_recorded: Arc<AtomicBool>,
}
impl WebSocketOriginFastPathReporter {
pub(crate) fn missed(&self) {
if !self.degradation_recorded.swap(true, Ordering::AcqRel) {
self.breaker.record_failure(&self.origin_key, self.probe_id);
}
}
}
enum CircuitState {
Closed { fast_timeouts: VecDeque<Instant> },
Closed { failures: VecDeque<Instant> },
Open { until: Instant },
HalfOpen { probe_id: Uuid },
}
@@ -214,7 +239,7 @@ enum CircuitState {
impl CircuitState {
fn closed() -> Self {
Self::Closed {
fast_timeouts: VecDeque::new(),
failures: VecDeque::new(),
}
}
}
@@ -1,6 +1,6 @@
//! Codex WebSocket 连接池。
use std::{collections::HashMap, sync::Arc, time::Duration};
use std::{collections::HashMap, future::Future, sync::Arc, time::Duration};
use serde_json::{Value, json};
use sha2::{Digest, Sha256};
@@ -8,6 +8,10 @@ use tokio::{
sync::{Mutex, watch},
time::Instant,
};
use tokio_util::{
sync::CancellationToken,
task::{TaskTracker, task_tracker::TaskTrackerToken},
};
use uuid::Uuid;
use super::websocket_pump::{PumpKeepalive, PumpedWebSocket};
@@ -63,6 +67,8 @@ impl CodexWebSocketPoolKey {
pub struct CodexWebSocketPool {
inner: Arc<Mutex<WebSocketPoolState>>,
config: CodexWebSocketPoolConfig,
tasks: TaskTracker,
shutdown: CancellationToken,
}
impl Default for CodexWebSocketPool {
@@ -84,7 +90,7 @@ pub struct CodexWebSocketPoolConfig {
pub maintenance_interval: Option<Duration>,
/// idle socket 探活 ping 间隔;`None` 表示维护时只做过期清理。
pub ping_interval: Option<Duration>,
/// 发送 ping 后等待上游响应的超时时间。
/// 发送 ping 后等待上游响应的超时时间;零值表示不校验 Pong deadline。
pub ping_timeout: Duration,
/// idle socket 无活动多久后视为失活。
pub liveness_timeout: Option<Duration>,
@@ -110,10 +116,11 @@ impl Default for CodexWebSocketPoolConfig {
}
impl CodexWebSocketPoolConfig {
/// pump 后台任务的保活策略:从连接池配置派生出 ping 间隔与 liveness 超时。
/// pump 后台任务的保活策略:从连接池配置派生出 ping/pong 与 liveness 策略。
pub(crate) fn keepalive(&self) -> PumpKeepalive {
PumpKeepalive {
ping_interval: self.ping_interval,
ping_timeout: (!self.ping_timeout.is_zero()).then_some(self.ping_timeout),
liveness_timeout: self.liveness_timeout,
}
}
@@ -137,6 +144,8 @@ impl CodexWebSocketPool {
let pool = Self {
inner: Arc::new(Mutex::new(WebSocketPoolState::default())),
config,
tasks: TaskTracker::new(),
shutdown: CancellationToken::new(),
};
pool.spawn_maintenance_task();
pool
@@ -152,6 +161,11 @@ impl CodexWebSocketPool {
self.config.initial_event_timeout
}
/// 注册由连接池生命周期托管的 opening 任务。
pub(crate) fn spawn_connect_task(&self, future: impl Future<Output = ()> + Send + 'static) {
drop(self.tasks.spawn(future));
}
pub(crate) async fn acquire(&self, key: &CodexWebSocketPoolKey) -> WebSocketPoolAcquire {
let mut expired_connection = None;
let acquire = {
@@ -166,6 +180,7 @@ impl CodexWebSocketPool {
Some(WebSocketPoolSlot::Connecting(connecting)) => {
return WebSocketPoolAcquire::Wait(WebSocketPoolConnectWaiter {
receiver: connecting.outcome.subscribe(),
started_at: connecting.started_at,
});
}
Some(WebSocketPoolSlot::Idle(_)) => {
@@ -198,8 +213,9 @@ impl CodexWebSocketPool {
key.clone(),
WebSocketPoolSlot::Connecting(WebSocketPoolConnecting {
id: lease.id,
started_at: Instant::now(),
started_at: lease.started_at,
outcome: lease.outcome.clone(),
cancellation: lease.cancellation.clone(),
}),
);
WebSocketPoolAcquire::Connect(lease)
@@ -215,6 +231,39 @@ impl CodexWebSocketPool {
acquire
}
/// 只租用已经处于 Idle 的连接,不创建新的 Connecting slot。
pub(crate) async fn take_idle(
&self,
key: &CodexWebSocketPoolKey,
) -> Option<(Box<PooledWebSocketConnection>, WebSocketPoolLease)> {
let mut expired_connection = None;
let acquired = {
let mut state = self.inner.lock().await;
if !self.config.enabled || state.shutting_down {
return None;
}
let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(key) else {
return None;
};
if connection.created_at.elapsed() < self.config.max_age
&& !connection.websocket.is_closed()
{
let lease = WebSocketPoolLease::reserve(self.clone(), key.clone());
state
.slots
.insert(key.clone(), WebSocketPoolSlot::Busy(lease.reservation));
Some((connection, lease))
} else {
expired_connection = Some(*connection);
None
}
};
if let Some(connection) = expired_connection {
close_pooled_connection(connection).await;
}
acquired
}
async fn put_reserved(
&self,
key: &CodexWebSocketPoolKey,
@@ -255,26 +304,22 @@ impl CodexWebSocketPool {
}
}
async fn finish_connect(
async fn finish_connect_idle(
&self,
key: &CodexWebSocketPoolKey,
connect_id: Uuid,
connection: PooledWebSocketConnection,
) -> Result<WebSocketPoolConnected, Box<PooledWebSocketConnection>> {
) -> Result<(), Box<PooledWebSocketConnection>> {
let mut state = self.inner.lock().await;
let owns_connect = matches!(
state.slots.get(key),
Some(WebSocketPoolSlot::Connecting(connecting)) if connecting.id == connect_id
);
if owns_connect && !state.shutting_down && self.config.enabled {
let lease = WebSocketPoolLease::reserve(self.clone(), key.clone());
state
.slots
.insert(key.clone(), WebSocketPoolSlot::Busy(lease.reservation));
Ok(WebSocketPoolConnected {
connection: Box::new(connection),
lease,
})
.insert(key.clone(), WebSocketPoolSlot::Idle(Box::new(connection)));
Ok(())
} else {
if owns_connect {
state.slots.remove(key);
@@ -293,7 +338,7 @@ impl CodexWebSocketPool {
}
}
/// 驱逐指定账号的 idle 连接,并阻止已占用 slot 回收到池中。
/// 驱逐指定账号的全部 slot,取消 opening,并阻止 busy 连接回收到池中。
pub async fn evict_account(&self, account_id: &str) {
let mut idle_connections = Vec::new();
{
@@ -310,6 +355,7 @@ impl CodexWebSocketPool {
idle_connections.push(*connection);
}
Some(WebSocketPoolSlot::Connecting(connecting)) => {
connecting.cancellation.cancel();
connecting
.outcome
.send_replace(WebSocketPoolConnectOutcome::Failed);
@@ -321,8 +367,10 @@ impl CodexWebSocketPool {
close_pooled_connections(idle_connections).await;
}
/// 关闭连接池,关闭所有 idle 连接,并让后续 acquire 直接绕过池。
/// 关闭连接池,取消受管任务、关闭 idle 连接,并让后续 acquire 直接绕过池。
pub async fn shutdown(&self) {
self.tasks.close();
self.shutdown.cancel();
let idle_connections = {
let mut state = self.inner.lock().await;
state.shutting_down = true;
@@ -332,6 +380,7 @@ impl CodexWebSocketPool {
.filter_map(|(_, slot)| match slot {
WebSocketPoolSlot::Idle(connection) => Some(*connection),
WebSocketPoolSlot::Connecting(connecting) => {
connecting.cancellation.cancel();
connecting
.outcome
.send_replace(WebSocketPoolConnectOutcome::Failed);
@@ -342,6 +391,7 @@ impl CodexWebSocketPool {
.collect::<Vec<_>>()
};
close_pooled_connections(idle_connections).await;
self.tasks.wait().await;
}
/// 维护池 slot:清扫已死亡或超龄的 idle 连接,以及异常残留的 Busy reservation。
@@ -363,20 +413,30 @@ impl CodexWebSocketPool {
};
let inner = Arc::downgrade(&self.inner);
let config = self.config;
tokio::spawn(async move {
let tasks = self.tasks.clone();
let shutdown = self.shutdown.clone();
drop(self.tasks.spawn(async move {
let mut interval = tokio::time::interval(interval_duration);
loop {
interval.tick().await;
tokio::select! {
_ = shutdown.cancelled() => break,
_ = interval.tick() => {}
}
let Some(inner) = inner.upgrade() else {
break;
};
let pool = CodexWebSocketPool { inner, config };
let pool = CodexWebSocketPool {
inner,
config,
tasks: tasks.clone(),
shutdown: shutdown.clone(),
};
if pool.is_shutdown().await {
break;
}
pool.maintain_idle_connections().await;
}
});
}));
}
/// 返回连接池是否已进入关闭状态。
@@ -427,6 +487,7 @@ impl CodexWebSocketPool {
);
}
Some(WebSocketPoolSlot::Connecting(connecting)) => {
connecting.cancellation.cancel();
connecting
.outcome
.send_replace(WebSocketPoolConnectOutcome::Failed);
@@ -502,6 +563,7 @@ struct WebSocketPoolConnecting {
id: Uuid,
started_at: Instant,
outcome: watch::Sender<WebSocketPoolConnectOutcome>,
cancellation: CancellationToken,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -513,9 +575,14 @@ pub(crate) enum WebSocketPoolConnectOutcome {
pub(crate) struct WebSocketPoolConnectWaiter {
receiver: watch::Receiver<WebSocketPoolConnectOutcome>,
started_at: Instant,
}
impl WebSocketPoolConnectWaiter {
pub(crate) fn remaining_budget(&self, budget: Duration) -> Duration {
budget.saturating_sub(self.started_at.elapsed())
}
pub(crate) async fn wait(mut self) -> WebSocketPoolConnectOutcome {
loop {
let outcome = *self.receiver.borrow_and_update();
@@ -533,34 +600,50 @@ pub(crate) struct WebSocketPoolConnectLease {
pool: CodexWebSocketPool,
key: CodexWebSocketPoolKey,
id: Uuid,
started_at: Instant,
outcome: watch::Sender<WebSocketPoolConnectOutcome>,
cancellation: CancellationToken,
// slot 分配时即注册,封闭 acquire 与后台 task spawn 之间的 shutdown 竞态。
_task_registration: TaskTrackerToken,
armed: bool,
}
pub(crate) struct WebSocketPoolConnected {
pub(crate) connection: Box<PooledWebSocketConnection>,
pub(crate) lease: WebSocketPoolLease,
}
impl WebSocketPoolConnectLease {
fn reserve(pool: CodexWebSocketPool, key: CodexWebSocketPoolKey) -> Self {
let (outcome, _) = watch::channel(WebSocketPoolConnectOutcome::Pending);
let cancellation = pool.shutdown.child_token();
let task_registration = pool.tasks.token();
Self {
pool,
key,
id: Uuid::new_v4(),
started_at: Instant::now(),
outcome,
cancellation,
_task_registration: task_registration,
armed: true,
}
}
pub(crate) async fn connected(
pub(crate) fn started_at(&self) -> Instant {
self.started_at
}
pub(crate) fn key(&self) -> &CodexWebSocketPoolKey {
&self.key
}
pub(crate) fn cancellation_token(&self) -> CancellationToken {
self.cancellation.clone()
}
pub(crate) async fn connected_idle(
mut self,
connection: PooledWebSocketConnection,
) -> Result<WebSocketPoolConnected, Box<PooledWebSocketConnection>> {
) -> Result<(), Box<PooledWebSocketConnection>> {
let result = self
.pool
.finish_connect(&self.key, self.id, connection)
.finish_connect_idle(&self.key, self.id, connection)
.await;
let outcome = if result.is_ok() {
WebSocketPoolConnectOutcome::Ready
@@ -573,6 +656,7 @@ impl WebSocketPoolConnectLease {
}
pub(crate) async fn failed(mut self) {
self.cancellation.cancel();
self.pool.fail_connect(&self.key, self.id).await;
self.outcome
.send_replace(WebSocketPoolConnectOutcome::Failed);
@@ -585,6 +669,7 @@ impl Drop for WebSocketPoolConnectLease {
if !self.armed {
return;
}
self.cancellation.cancel();
self.outcome
.send_replace(WebSocketPoolConnectOutcome::Failed);
let Ok(runtime) = tokio::runtime::Handle::try_current() else {
@@ -2,7 +2,7 @@
//!
//! 每条上游 WebSocket 连接都由一个后台 pump 任务独占:
//! - 持续读取 socket:一旦观察到 `Close` / EOF / 传输错误,立即把连接标记为 `closed`。
//! - 自动回应上游 `Ping`(保活),并按 `ping_interval` 主动 `Ping`(穿透 NAT/中间盒空闲计时器)。
//! - 自动回应上游 `Ping`,并按 `ping_interval` 主动 `Ping`;匹配的 `Pong` 必须在 deadline 内返回。
//! - 可选 `liveness_timeout`:长时间无任何入站活动时判定连接失活并退出。
//!
//! 因此空闲连接的“死没死”在后台被实时感知;复用方只需零成本读取 [`PumpedWebSocket::is_closed`],
@@ -11,7 +11,7 @@
//! 收发都通过 channel 与 pump 任务交互:
//! - 发送:`send` 走 command channel,等待 pump 回执。
//! - 接收:`next` 从 message channel 取出 pump 转发的入站帧(`Ping`/`Pong` 已被 pump 吞掉)。
//! - 入站缓冲满时暂停读取 socket,并继续处理关闭命令与保活;消费恢复后按原顺序继续转发。
//! - 入站缓冲满时暂停读取 socket 和主动探活,仅继续处理发送/关闭命令;消费恢复后按原顺序继续转发。
use std::sync::{
Arc, Mutex,
@@ -39,6 +39,8 @@ const PUMP_COMMAND_BUFFER: usize = 32;
pub(crate) struct PumpKeepalive {
/// 主动 `Ping` 间隔;`None` 表示 pump 不主动 ping(仅被动读取 + 回应上游 ping)。
pub(crate) ping_interval: Option<Duration>,
/// 主动 `Ping` 发出后等待匹配 `Pong` 的 deadline;`None` 表示不校验响应。
pub(crate) ping_timeout: Option<Duration>,
/// 无入站活动多久后判定失活;`None` 表示只靠显式 close/传输错误发现死亡。
pub(crate) liveness_timeout: Option<Duration>,
}
@@ -48,6 +50,7 @@ impl PumpKeepalive {
pub(crate) fn disabled() -> Self {
Self {
ping_interval: None,
ping_timeout: None,
liveness_timeout: None,
}
}
@@ -70,6 +73,7 @@ pub(crate) enum PumpExitReason {
UpstreamCloseFrame { frame: Option<String> },
UpstreamEof,
InboundTransportError { message: String },
PongTimeout { timeout: Duration },
LivenessTimeout { timeout: Duration },
KeepaliveTransportError { message: String },
MessageReceiverClosed,
@@ -84,6 +88,7 @@ impl PumpExitReason {
Self::UpstreamCloseFrame { .. } => "upstream_close_frame",
Self::UpstreamEof => "upstream_eof",
Self::InboundTransportError { .. } => "inbound_transport_error",
Self::PongTimeout { .. } => "pong_timeout",
Self::LivenessTimeout { .. } => "liveness_timeout",
Self::KeepaliveTransportError { .. } => "keepalive_transport_error",
Self::MessageReceiverClosed => "message_receiver_closed",
@@ -96,7 +101,9 @@ impl PumpExitReason {
| Self::InboundTransportError { message }
| Self::KeepaliveTransportError { message } => Some(message.clone()),
Self::UpstreamCloseFrame { frame } => frame.clone(),
Self::LivenessTimeout { timeout } => Some(format!("{timeout:?}")),
Self::PongTimeout { timeout } | Self::LivenessTimeout { timeout } => {
Some(format!("{timeout:?}"))
}
Self::CommandChannelClosed
| Self::LocalClose
| Self::UpstreamEof
@@ -110,6 +117,7 @@ impl PumpExitReason {
Self::OutboundTransportError { .. }
| Self::UpstreamEof
| Self::InboundTransportError { .. }
| Self::PongTimeout { .. }
| Self::LivenessTimeout { .. }
| Self::KeepaliveTransportError { .. }
)
@@ -120,6 +128,7 @@ impl PumpExitReason {
self,
Self::CommandChannelClosed
| Self::UpstreamCloseFrame { .. }
| Self::PongTimeout { .. }
| Self::LivenessTimeout { .. }
| Self::MessageReceiverClosed
)
@@ -131,6 +140,12 @@ struct PendingInbound {
exit_reason: Option<PumpExitReason>,
}
#[derive(Debug, Clone, Copy)]
struct PendingPing {
payload: [u8; 8],
deadline: Instant,
}
/// 由后台 pump 任务托管的 WebSocket 连接句柄。
pub(crate) struct PumpedWebSocket {
connection_id: Uuid,
@@ -248,16 +263,9 @@ async fn pump_loop(
) {
let mut last_activity = Instant::now();
let ping_interval = keepalive.ping_interval.filter(|d| !d.is_zero());
let ping_timeout = keepalive.ping_timeout.filter(|d| !d.is_zero());
let liveness_timeout = keepalive.liveness_timeout.filter(|d| !d.is_zero());
// ticker 由 ping 间隔或 liveness 检查间隔中较小者驱动:即使不主动 ping,
// 也要周期性醒来检查 liveness。两者都未配置时不启动 ticker。
let tick_interval = match (ping_interval, liveness_timeout) {
(Some(ping), Some(live)) => Some(ping.min(live)),
(Some(ping), None) => Some(ping),
(None, Some(live)) => Some(live),
(None, None) => None,
};
let mut ticker = tick_interval.map(|d| {
let mut ping_ticker = ping_interval.map(|d| {
// 首个 tick 推迟一整个间隔:tokio::time::interval 默认会让首个 tick 立即就绪,
// 否则连接一建立就会立刻发一帧 Ping,与首个请求 Text 抢跑、打乱帧序。
let mut ticker = tokio::time::interval_at(tokio::time::Instant::now() + d, d);
@@ -266,6 +274,8 @@ async fn pump_loop(
});
let mut pending_inbound: Option<PendingInbound> = None;
let mut pending_ping: Option<PendingPing> = None;
let mut ping_sequence = 0_u64;
let mut backpressure_events = 0_u64;
let reason = 'pump: loop {
@@ -282,6 +292,8 @@ async fn pump_loop(
if let Some(reason) = pending.exit_reason {
break 'pump reason;
}
// 背压期间无法读取 socket,不能把这段本地下游阻塞误算成上游静默。
last_activity = Instant::now();
}
command = rx_command.recv() => {
let Some(command) = command else {
@@ -291,21 +303,12 @@ async fn pump_loop(
break 'pump reason;
}
}
_ = tick(&mut ticker) => {
if let Some(reason) = handle_tick(
&mut inner,
last_activity,
ping_interval,
liveness_timeout,
false,
).await {
break 'pump reason;
}
}
}
continue;
}
let pong_deadline = pending_ping.map(|ping| ping.deadline);
let liveness_deadline = liveness_timeout.map(|timeout| last_activity + timeout);
tokio::select! {
command = rx_command.recv() => {
let Some(command) = command else {
@@ -318,36 +321,38 @@ async fn pump_loop(
message = inner.next() => {
match message {
None => break 'pump PumpExitReason::UpstreamEof,
Some(Ok(Message::Ping(payload))) => {
last_activity = Instant::now();
if let Err(error) = inner.send(Message::Pong(payload)).await {
break 'pump PumpExitReason::KeepaliveTransportError {
message: error.to_string(),
};
}
}
Some(Ok(Message::Pong(_))) => {
last_activity = Instant::now();
}
Some(Ok(message)) => {
last_activity = Instant::now();
let terminal_reason = match &message {
Message::Close(frame) => Some(PumpExitReason::UpstreamCloseFrame {
frame: frame.as_ref().map(|frame| format!("{frame:?}")),
}),
_ => None,
};
if let Some(reason) = enqueue_inbound(
&tx_message,
&mut pending_inbound,
connection_id,
&mut backpressure_events,
PendingInbound {
item: Ok(message),
exit_reason: terminal_reason,
},
) {
break 'pump reason;
observe_inbound(connection_id, &mut pending_ping, &message);
match message {
Message::Ping(payload) => {
if let Err(error) = inner.send(Message::Pong(payload)).await {
break 'pump PumpExitReason::KeepaliveTransportError {
message: error.to_string(),
};
}
}
Message::Pong(_) => {}
message => {
let terminal_reason = match &message {
Message::Close(frame) => Some(PumpExitReason::UpstreamCloseFrame {
frame: frame.as_ref().map(|frame| format!("{frame:?}")),
}),
_ => None,
};
if let Some(reason) = enqueue_inbound(
&tx_message,
&mut pending_inbound,
connection_id,
&mut backpressure_events,
PendingInbound {
item: Ok(message),
exit_reason: terminal_reason,
},
) {
break 'pump reason;
}
}
}
}
Some(Err(err)) => {
@@ -369,17 +374,28 @@ async fn pump_loop(
}
}
}
_ = tick(&mut ticker) => {
if let Some(reason) = handle_tick(
&mut inner,
last_activity,
ping_interval,
liveness_timeout,
true,
).await {
break 'pump reason;
_ = tick(&mut ping_ticker) => {
if pending_ping.is_none() {
match send_keepalive_ping(
&mut inner,
&mut ping_sequence,
ping_timeout,
).await {
Ok(pending) => pending_ping = pending,
Err(reason) => break 'pump reason,
}
}
}
_ = wait_until(pong_deadline) => {
break 'pump PumpExitReason::PongTimeout {
timeout: ping_timeout.expect("pending ping must have a timeout"),
};
}
_ = wait_until(liveness_deadline) => {
break 'pump PumpExitReason::LivenessTimeout {
timeout: liveness_timeout.expect("liveness deadline must have a timeout"),
};
}
}
};
@@ -433,27 +449,39 @@ async fn handle_command(inner: &mut RawWsStream, command: PumpCommand) -> Option
}
}
async fn handle_tick(
fn observe_inbound(connection_id: Uuid, pending_ping: &mut Option<PendingPing>, message: &Message) {
let Some(pending) = pending_ping else {
return;
};
if let Message::Pong(payload) = message
&& payload.as_ref() != pending.payload
{
tracing::debug!(
websocket_connection_id = %connection_id,
"Responses WebSocket pump received a non-matching Pong"
);
}
// 匹配 Pong 是探针的直接回执;其他入站帧同样证明上游链路仍可达。
*pending_ping = None;
}
async fn send_keepalive_ping(
inner: &mut RawWsStream,
last_activity: Instant,
ping_interval: Option<Duration>,
liveness_timeout: Option<Duration>,
check_liveness: bool,
) -> Option<PumpExitReason> {
if check_liveness
&& let Some(timeout) = liveness_timeout
&& last_activity.elapsed() >= timeout
{
return Some(PumpExitReason::LivenessTimeout { timeout });
}
if ping_interval.is_some()
&& let Err(error) = inner.send(Message::Ping(Vec::new().into())).await
{
return Some(PumpExitReason::KeepaliveTransportError {
sequence: &mut u64,
ping_timeout: Option<Duration>,
) -> Result<Option<PendingPing>, PumpExitReason> {
let payload = sequence.to_be_bytes();
*sequence = sequence.wrapping_add(1);
inner
.send(Message::Ping(payload.to_vec().into()))
.await
.map_err(|error| PumpExitReason::KeepaliveTransportError {
message: error.to_string(),
});
}
None
})?;
Ok(ping_timeout.map(|timeout| PendingPing {
payload,
deadline: Instant::now() + timeout,
}))
}
fn log_pump_exit(connection_id: Uuid, reason: &PumpExitReason, backpressure_events: u64) {
@@ -494,3 +522,11 @@ async fn tick(ticker: &mut Option<tokio::time::Interval>) {
None => std::future::pending::<()>().await,
}
}
/// 等待可选 deadline;`None` 时永远挂起,让 `select!` 分支实际禁用。
async fn wait_until(deadline: Option<Instant>) {
match deadline {
Some(deadline) => tokio::time::sleep_until(deadline).await,
None => std::future::pending::<()>().await,
}
}
@@ -195,3 +195,60 @@ async fn admin_account_quota_should_send_usage_cookie() {
"2K"
);
}
#[tokio::test]
async fn admin_account_quota_should_ban_deactivated_workspace() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/backend-api/wham/usage"))
.respond_with(
ResponseTemplate::new(402)
.set_body_json(json!({ "detail": { "code": "deactivated_workspace" } })),
)
.expect(1)
.mount(&server)
.await;
let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url(
"admin-account-quota-deactivated",
89,
format!("{}/backend-api", server.uri()),
)
.await;
seed_account(
&pool,
NewAccount {
id: "acct_quota_deactivated".to_string(),
email: Some("quota-deactivated@example.com".to_string()),
account_id: Some("chatgpt-quota-deactivated".to_string()),
user_id: None,
label: None,
plan_type: Some("plus".to_string()),
access_token: SecretString::new("access-quota-deactivated".to_string().into()),
refresh_token: None,
access_token_expires_at: None,
status: AccountStatus::Active,
added_at: None,
},
)
.await;
let response = app
.oneshot(
Request::builder()
.method("GET")
.uri("/api/admin/accounts/quota?id=acct_quota_deactivated")
.header("cookie", "cpr_admin_session=session_1")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let status: String = sqlx::query_scalar("select status from accounts where id = $1")
.bind("acct_quota_deactivated")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
assert_eq!(status, "banned");
}
+56 -68
View File
@@ -117,7 +117,7 @@ async fn dashboard_summary_should_calculate_today_traffic_and_average_first_toke
.append(&usage_record_with_tokens(yesterday_record_at, 30))
.await
.unwrap();
seed_lifetime_usage_summary(&pool, 7, 400, 500, 100).await;
seed_retained_usage_bucket(&pool, today_start - Duration::days(8), 4, 340, 500, 100).await;
let body = dashboard_summary(app).await;
@@ -370,7 +370,7 @@ async fn dashboard_summary_should_return_requested_trend_kind() {
#[tokio::test]
async fn dashboard_summary_should_calculate_priority_billing_from_event_service_tier() {
let billing_amount = dashboard_summary_today_billing_for_usage_record(
let billing_amount = dashboard_summary_total_billing_for_usage_record(
"dashboard-billing-priority",
"gpt-5.5",
json!({
@@ -389,7 +389,7 @@ async fn dashboard_summary_should_calculate_priority_billing_from_event_service_
#[tokio::test]
async fn dashboard_summary_should_charge_cached_input_at_cache_read_price() {
let billing_amount = dashboard_summary_today_billing_for_usage_record(
let billing_amount = dashboard_summary_total_billing_for_usage_record(
"dashboard-billing-cached",
"gpt-5.5",
json!({
@@ -407,7 +407,7 @@ async fn dashboard_summary_should_charge_cached_input_at_cache_read_price() {
#[tokio::test]
async fn dashboard_summary_should_apply_long_context_prices() {
let billing_amount = dashboard_summary_today_billing_for_usage_record(
let billing_amount = dashboard_summary_total_billing_for_usage_record(
"dashboard-billing-long-context",
"gpt-5.5",
json!({
@@ -425,7 +425,7 @@ async fn dashboard_summary_should_apply_long_context_prices() {
#[tokio::test]
async fn dashboard_summary_should_match_versioned_model_names_for_billing() {
let billing_amount = dashboard_summary_today_billing_for_usage_record(
let billing_amount = dashboard_summary_total_billing_for_usage_record(
"dashboard-billing-codex-model",
"gpt-5.6-terra-2026-07-01",
json!({
@@ -441,23 +441,6 @@ async fn dashboard_summary_should_match_versioned_model_names_for_billing() {
assert_eq!(billing_amount, "$1.75");
}
#[tokio::test]
async fn dashboard_summary_should_not_price_usage_summary_without_model_dimension() {
let (app, _store, pool, _dir) = dashboard_test_app(
"dashboard-billing-summary-fallback",
crate::support::fingerprint::test_fingerprint(),
)
.await;
seed_usage_summary(&pool, 100_000, 100_000).await;
let body = dashboard_summary(app).await;
assert_eq!(
body["data"]["cards"]["tokens"]["todayBillingAmountUsd"],
"$0.00"
);
}
#[tokio::test]
async fn dashboard_summary_should_return_backend_formatted_time_fields() {
let (app, store, _pool, _dir) = dashboard_test_app(
@@ -706,7 +689,44 @@ async fn dashboard_summary_should_return_today_health_timeline() {
assert_eq!(points.matches('4').count(), 1);
}
async fn dashboard_summary_today_billing_for_usage_record(
#[tokio::test]
async fn dashboard_summary_should_include_retained_history_in_total_billing() {
let (app, store, _pool, _dir) = dashboard_test_app(
"dashboard-total-billing-retained-history",
crate::support::fingerprint::test_fingerprint(),
)
.await;
let metadata = json!({
"usage": {
"inputTokens": 100_000u64,
"outputTokens": 100_000u64,
"cachedTokens": 0u64
},
"serviceTier": "priority"
});
for created_at in [Utc::now(), Utc::now() - Duration::days(8)] {
let mut record = usage_record(created_at, "gpt-5.5", metadata.clone());
record.input_tokens = Some(100_000);
record.output_tokens = Some(100_000);
record.cached_tokens = Some(0);
record.service_tier = Some("priority".to_string());
store.append(&record).await.unwrap();
}
let body = dashboard_summary(app).await;
assert_eq!(
body["data"]["cards"]["tokens"]["totalBillingAmountUsd"],
"$17.50"
);
assert!(
body["data"]["cards"]["tokens"]
.get("todayBillingAmountUsd")
.is_none()
);
}
async fn dashboard_summary_total_billing_for_usage_record(
db_name: &str,
model: &str,
metadata: Value,
@@ -733,7 +753,7 @@ async fn dashboard_summary_today_billing_for_usage_record(
store.append(&record).await.unwrap();
let body = dashboard_summary(app).await;
body["data"]["cards"]["tokens"]["todayBillingAmountUsd"]
body["data"]["cards"]["tokens"]["totalBillingAmountUsd"]
.as_str()
.unwrap()
.to_string()
@@ -812,62 +832,30 @@ async fn dashboard_trend(app: axum::Router, kind: &str) -> Value {
response_json(response).await
}
async fn seed_usage_summary(pool: &PgPool, input_tokens: u64, output_tokens: u64) {
seed_usage_summary_with_last_used(pool, input_tokens, output_tokens, "2026-06-18T00:00:00Z")
.await;
}
async fn seed_lifetime_usage_summary(
async fn seed_retained_usage_bucket(
pool: &PgPool,
bucket_start: chrono::DateTime<Utc>,
request_count: u64,
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
) {
seed_dashboard_account(pool, "acct_lifetime", "lifetime@example.com").await;
sqlx::query(
"insert into account_usage
(account_id, request_count, input_tokens, output_tokens, cached_tokens, total_tokens)
values ($1, $2, $3, $4, $5, $6)",
"insert into request_time_buckets
(bucket_start, provider, account_id, model, service_tier, success_count,
input_tokens, output_tokens, cached_tokens, updated_at)
values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)",
)
.bind("acct_lifetime")
.bind(china_quarter_hour_start(bucket_start))
.bind("openai")
.bind("acct_retained")
.bind("gpt-5.5")
.bind("__unknown__")
.bind(request_count as i64)
.bind(input_tokens as i64)
.bind(output_tokens as i64)
.bind(cached_tokens as i64)
.bind((input_tokens + output_tokens) as i64)
.execute(pool)
.await
.unwrap();
}
async fn seed_usage_summary_with_last_used(
pool: &PgPool,
input_tokens: u64,
output_tokens: u64,
last_used_at: &str,
) {
sqlx::query(
"insert into accounts (id, email, access_token, status, added_at, updated_at) values ($1, $2, $3, $4, $5, $6)",
)
.bind("acct_usage")
.bind("acct-usage@example.com")
.bind("cipher")
.bind("active")
.bind(crate::support::storage::timestamp("2026-06-18T00:00:00Z"))
.bind(crate::support::storage::timestamp("2026-06-18T00:00:00Z"))
.execute(pool)
.await
.unwrap();
sqlx::query(
"insert into account_usage (account_id, request_count, input_tokens, output_tokens, total_tokens, last_used_at) values ($1, $2, $3, $4, $5, $6)",
)
.bind("acct_usage")
.bind(1_i64)
.bind(input_tokens as i64)
.bind(output_tokens as i64)
.bind((input_tokens + output_tokens) as i64)
.bind(crate::support::storage::timestamp(last_used_at))
.bind(Utc::now())
.execute(pool)
.await
.unwrap();
+62 -4
View File
@@ -124,10 +124,38 @@ async fn usage_summary_and_trend_expose_only_real_success_metrics() {
async fn usage_insights_should_collapse_retry_errors_to_terminal_requests() {
let (app, usage, ops, _guard) = monitoring_test_app("monitoring-insights", true).await;
let now = chrono::Utc::now();
sqlx::query(
"insert into accounts
(id, email, access_token, status, added_at, updated_at)
values ($1, $2, $3, $4, $5, $6)",
)
.bind("acct_usage")
.bind("usage-account@example.com")
.bind("cipher")
.bind("active")
.bind(now)
.bind(now)
.execute(usage.pool())
.await
.unwrap();
sqlx::query(
"insert into client_api_keys
(id, name, prefix, key, enabled, created_at, last_used_at)
values ($1, $2, $3, $4, true, $5, null)",
)
.bind("key_diagnostic")
.bind("生产密钥")
.bind("sk-diagnostic")
.bind("sk-diagnostic-secret")
.bind(now)
.execute(usage.pool())
.await
.unwrap();
let mut success = success_record("eventually succeeded");
success.id = "usage_insights_success".to_string();
success.request_id = Some("req_retried".to_string());
success.client_api_key_id = Some("key_diagnostic".to_string());
success.input_tokens = Some(100);
success.output_tokens = Some(20);
success.cached_tokens = Some(20);
@@ -143,14 +171,18 @@ async fn usage_insights_should_collapse_retry_errors_to_terminal_requests() {
retried_error.created_at = now - chrono::Duration::seconds(2);
ops.append(&retried_error).await.unwrap();
for (id, failure_class, seconds) in [
("ops_insights_failed_older", "upstream", 2),
("ops_insights_failed_latest", "response_failed", 1),
for (id, failure_class, seconds, persist_class) in [
("ops_insights_failed_older", "upstream", 2, true),
("ops_insights_failed_latest", "response_failed", 1, false),
] {
let mut error = OpsErrorLog::new("upstream", failure_class);
error.id = id.to_string();
error.request_id = Some("req_failed".to_string());
error.failure_class = Some(failure_class.to_string());
if persist_class {
error.failure_class = Some(failure_class.to_string());
} else {
error.metadata = json!({ "upstreamCode": failure_class });
}
error.created_at = now - chrono::Duration::seconds(seconds);
ops.append(&error).await.unwrap();
}
@@ -184,6 +216,32 @@ async fn usage_insights_should_collapse_retry_errors_to_terminal_requests() {
.unwrap()
.is_empty()
);
let account_diagnostics = admin_get(
&app,
"/api/admin/usage/records/insights/diagnostics?dimension=account",
)
.await;
assert!(
account_diagnostics["data"]["items"]
.as_array()
.unwrap()
.iter()
.any(|item| item["name"] == "usage-account@example.com → acct_usage")
);
let api_key_diagnostics = admin_get(
&app,
"/api/admin/usage/records/insights/diagnostics?dimension=apiKey",
)
.await;
assert!(
api_key_diagnostics["data"]["items"]
.as_array()
.unwrap()
.iter()
.any(|item| item["name"] == "生产密钥")
);
}
#[tokio::test]
@@ -796,6 +796,7 @@ async fn responses_stream_should_record_usage_record_after_late_disconnect() {
assert_eq!(metadata["stream"], true);
assert_eq!(metadata["failed"], true);
assert_eq!(metadata["upstreamCode"], "stream_disconnected");
assert_eq!(event.failure_class.as_deref(), Some("stream_disconnected"));
assert_eq!(metadata["failureSource"], "proxy");
assert_eq!(metadata["synthetic"], true);
assert_rate_limit_header(&metadata, "retry-after", "11");
+79
View File
@@ -0,0 +1,79 @@
use codex_proxy_rs::{
fleet::{
account::AccountStatus,
account_failure::{AccountFailureKind, AccountStateEffect, classify_client_failure},
},
upstream::openai::transport::{
CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics,
},
};
use reqwest::StatusCode;
#[test]
fn deactivated_workspace_should_be_banned_before_generic_payment_classification() {
let failure = classify(402, r#"{"detail":{"code":"deactivated_workspace"}}"#);
assert_eq!(failure.kind, AccountFailureKind::Banned);
assert_eq!(
failure.effect,
Some(AccountStateEffect::SetStatus(AccountStatus::Banned))
);
}
#[test]
fn generic_payment_required_should_exhaust_quota() {
let failure = classify(402, r#"{"detail":{"code":"payment_required"}}"#);
assert_eq!(failure.kind, AccountFailureKind::QuotaExhausted);
assert_eq!(
failure.effect,
Some(AccountStateEffect::SetStatus(AccountStatus::QuotaExhausted))
);
}
#[test]
fn unauthorized_should_expire_account() {
let failure = classify(401, r#"{"error":{"message":"unauthorized"}}"#);
assert_eq!(failure.kind, AccountFailureKind::Expired);
assert_eq!(
failure.effect,
Some(AccountStateEffect::SetStatus(AccountStatus::Expired))
);
}
#[test]
fn rate_limit_should_create_temporary_quota_effect() {
let failure = classify(429, r#"{"error":{"code":"rate_limit_exceeded"}}"#);
assert_eq!(failure.kind, AccountFailureKind::RateLimited);
assert!(matches!(
failure.effect,
Some(AccountStateEffect::MarkQuotaLimitedUntil(_))
));
}
#[test]
fn html_forbidden_page_should_not_ban_account() {
let error = upstream_error(403, "<!doctype html><html>temporary edge failure</html>");
assert!(classify_client_failure(&error).is_none());
}
fn classify(
status: u16,
body: &str,
) -> codex_proxy_rs::fleet::account_failure::ClassifiedAccountFailure {
classify_client_failure(&upstream_error(status, body)).expect("failure should be classified")
}
fn upstream_error(status: u16, body: &str) -> CodexClientError {
CodexClientError::Upstream {
status: StatusCode::from_u16(status).unwrap(),
body: body.to_string(),
retry_after_seconds: None,
diagnostics: CodexUpstreamDiagnostics::with_status(status),
set_cookie_headers: Vec::new(),
transport: CodexBackendTransport::HttpSse,
}
}
+1
View File
@@ -1,3 +1,4 @@
pub(crate) mod account_failure;
pub(crate) mod cookies;
pub(crate) mod pool;
pub(crate) mod quota;
+68 -6
View File
@@ -1,13 +1,16 @@
use std::{collections::HashMap, sync::Arc, time::Duration as StdDuration};
use chrono::{Duration, Utc};
use codex_proxy_rs::fleet::account::AccountStatus;
use codex_proxy_rs::fleet::cookies::PgCookieStore;
use codex_proxy_rs::fleet::quota::{
QuotaRefreshService, quota_from_usage, quota_snapshot_limit_reached,
quota_snapshot_limit_window_seconds, quota_snapshot_reset_at,
use codex_proxy_rs::fleet::{
account::AccountStatus,
cookies::PgCookieStore,
pool::{AccountPoolOptions, AccountPoolService},
quota::{
QuotaRefreshService, quota_from_usage, quota_snapshot_limit_reached,
quota_snapshot_limit_window_seconds, quota_snapshot_reset_at,
},
store::{AccountStore, NewAccount, PgAccountStore},
};
use codex_proxy_rs::fleet::store::{NewAccount, PgAccountStore};
use codex_proxy_rs::infra::identity::AccountPseudonymizer;
use codex_proxy_rs::telemetry::account_usage::store::PgAccountUsageStore;
use codex_proxy_rs::upstream::openai::transport::CodexBackendClient;
@@ -24,6 +27,15 @@ fn test_account_pseudonymizer() -> Arc<AccountPseudonymizer> {
Arc::new(AccountPseudonymizer::new([7; 32]))
}
fn test_account_pool(pool: &sqlx::PgPool) -> Arc<AccountPoolService> {
Arc::new(AccountPoolService::new(
Arc::new(PgAccountStore::new(pool.clone())) as Arc<dyn AccountStore>,
Arc::new(PgAccountUsageStore::new(pool.clone())),
AccountPoolOptions::default(),
0,
))
}
#[test]
fn quota_from_usage_should_preserve_spend_control_individual_limit() {
let usage = json!({
@@ -321,6 +333,7 @@ async fn quota_refresh_service_should_send_usage_cookie_when_cookie_store_is_con
store,
PgAccountUsageStore::new(pool.clone()),
Arc::new(codex),
test_account_pool(&pool),
test_account_pseudonymizer(),
)
.with_cookie_store(cookies);
@@ -403,6 +416,7 @@ async fn quota_refresh_service_should_fetch_usage_for_quota_locked_accounts_and_
store.clone(),
PgAccountUsageStore::new(pool.clone()),
Arc::new(codex),
test_account_pool(&pool),
test_account_pseudonymizer(),
);
@@ -435,6 +449,49 @@ async fn quota_refresh_service_should_fetch_usage_for_quota_locked_accounts_and_
);
}
#[tokio::test]
async fn quota_refresh_service_should_ban_deactivated_workspace() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(
ResponseTemplate::new(402)
.set_body_json(json!({ "detail": { "code": "deactivated_workspace" } })),
)
.expect(1)
.mount(&server)
.await;
let (pool, _guard) = init_test_db("quota-refresh-deactivated-workspace").await;
let store = PgAccountStore::new(pool.clone());
insert_quota_locked_account(&store, &pool, "acct-deactivated", "access-token").await;
let codex = Arc::new(CodexBackendClient::new(
reqwest::Client::new(),
server.uri(),
crate::support::fingerprint::runtime_test_fingerprint(),
));
let service = QuotaRefreshService::new(
store,
PgAccountUsageStore::new(pool.clone()),
codex,
test_account_pool(&pool),
test_account_pseudonymizer(),
);
let summary = service
.refresh_locked_accounts(&mut HashMap::new())
.await
.expect("quota refresh should complete");
let status: String = sqlx::query_scalar("select status from accounts where id = $1")
.bind("acct-deactivated")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(summary.failed, 1);
assert_eq!(status, "banned");
}
#[tokio::test]
async fn quota_refresh_service_should_refresh_quota_exhausted_accounts() {
let server = MockServer::start().await;
@@ -479,6 +536,7 @@ async fn quota_refresh_service_should_refresh_quota_exhausted_accounts() {
store,
PgAccountUsageStore::new(pool.clone()),
Arc::new(codex),
test_account_pool(&pool),
test_account_pseudonymizer(),
);
@@ -515,6 +573,7 @@ async fn quota_refresh_service_should_skip_recent_locked_account_before_cooldown
store,
PgAccountUsageStore::new(pool.clone()),
Arc::new(codex),
test_account_pool(&pool),
test_account_pseudonymizer(),
1_800,
);
@@ -555,6 +614,7 @@ async fn quota_refresh_service_should_skip_recent_locked_account_inside_cooldown
store,
PgAccountUsageStore::new(pool.clone()),
Arc::new(codex),
test_account_pool(&pool),
test_account_pseudonymizer(),
1_800,
);
@@ -612,6 +672,7 @@ async fn quota_refresh_service_should_bypass_recent_skip_after_cooldown_grace()
store,
PgAccountUsageStore::new(pool.clone()),
Arc::new(codex),
test_account_pool(&pool),
test_account_pseudonymizer(),
1_800,
);
@@ -667,6 +728,7 @@ async fn quota_refresh_service_should_stagger_multiple_locked_account_requests()
store,
PgAccountUsageStore::new(pool.clone()),
Arc::new(codex),
test_account_pool(&pool),
test_account_pseudonymizer(),
)
.with_request_spacing(StdDuration::from_millis(200));
+27 -3
View File
@@ -3,7 +3,7 @@ use std::{env, sync::OnceLock, thread};
use chrono::{DateTime, Utc};
use codex_proxy_rs::{
auth::store::{PgAdminUserStore, RedisAdminSessionStore},
bootstrap::services::BackgroundTaskStores,
bootstrap::{config::BootstrapConfig, services::BackgroundTaskStores},
dispatch::affinity::RedisSessionAffinityStore,
fleet::{cookies::PgCookieStore, refresh::RedisRefreshLeaseStore, store::PgAccountStore},
infra::{database::migrate, redis::RedisConnection},
@@ -24,6 +24,7 @@ const DEFAULT_TEST_REDIS_URL: &str = "redis://127.0.0.1:6379";
static DATABASE_INIT_LIMIT: tokio::sync::Semaphore = tokio::sync::Semaphore::const_new(2);
static REDIS_TEST_CLEANUP: tokio::sync::OnceCell<()> = tokio::sync::OnceCell::const_new();
static REDIS_TEST_RUN_ID: OnceLock<String> = OnceLock::new();
static LOCAL_BOOTSTRAP_CONFIG: OnceLock<Option<BootstrapConfig>> = OnceLock::new();
pub(crate) struct TestDatabaseGuard {
database_name: Option<String>,
@@ -186,11 +187,34 @@ pub(crate) fn background_task_stores(pool: PgPool, redis: RedisConnection) -> Ba
}
pub(crate) fn test_database_url() -> String {
env::var("CPR_TEST_DATABASE_URL").unwrap_or_else(|_| DEFAULT_TEST_DATABASE_URL.to_string())
// 显式测试拓扑优先;本地开发默认跟随 Compose 配置;无本地配置的 CI 使用固定测试值。
env::var("CPR_TEST_DATABASE_URL")
.ok()
.or_else(|| local_bootstrap_config().map(|config| config.database_url().to_string()))
.unwrap_or_else(|| DEFAULT_TEST_DATABASE_URL.to_string())
}
pub(crate) fn test_redis_url() -> String {
env::var("CPR_TEST_REDIS_URL").unwrap_or_else(|_| DEFAULT_TEST_REDIS_URL.to_string())
env::var("CPR_TEST_REDIS_URL")
.ok()
.or_else(|| local_bootstrap_config().map(|config| config.redis_url().to_string()))
.unwrap_or_else(|| DEFAULT_TEST_REDIS_URL.to_string())
}
fn local_bootstrap_config() -> Option<&'static BootstrapConfig> {
LOCAL_BOOTSTRAP_CONFIG
.get_or_init(|| {
let path =
std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../deploy/config.yaml");
if !path.is_file() {
return None;
}
Some(
BootstrapConfig::load_from_path(path)
.unwrap_or_else(|error| panic!("load local test service config: {error}")),
)
})
.as_ref()
}
pub(crate) fn timestamp(value: &str) -> DateTime<Utc> {
@@ -5,27 +5,11 @@ use std::sync::{
use super::{accept_codex_test_websocket, completed_websocket_response, request_context};
use codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest;
use codex_proxy_rs::upstream::openai::transport::{
CodexBackendClient, CodexWebSocketPool, is_deactivated_workspace_error_body,
};
use codex_proxy_rs::upstream::openai::transport::{CodexBackendClient, CodexWebSocketPool};
use futures::{SinkExt, StreamExt};
use tokio::net::TcpListener;
use tokio_tungstenite::tungstenite::Message;
#[test]
fn deactivated_workspace_error_body_should_match_detail_code() {
let body = r#"{"detail":{"code":"deactivated_workspace"}}"#;
assert!(is_deactivated_workspace_error_body(body));
}
#[test]
fn deactivated_workspace_error_body_should_not_match_other_json_paths() {
let body = r#"{"error":{"code":"deactivated_workspace"}}"#;
assert!(!is_deactivated_workspace_error_body(body));
}
#[tokio::test]
async fn codex_backend_client_should_apply_configured_websocket_pool() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -19,7 +19,7 @@ fn new_chain_request(conversation_id: &str) -> CodexResponsesRequest {
}
#[tokio::test]
async fn cold_websocket_should_fall_back_to_same_account_http_after_fast_budget() {
async fn cold_websocket_should_fall_back_without_recording_a_successful_connect() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
@@ -32,12 +32,13 @@ async fn cold_websocket_should_fall_back_to_same_account_http_after_fast_budget(
assert!(request.starts_with("POST /codex/responses HTTP/1.1"));
write_completed_sse_response(&mut http).await;
});
let pool = Arc::new(CodexWebSocketPool::default());
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::runtime_test_fingerprint(),
)
.with_websocket_pool(Arc::new(CodexWebSocketPool::default()))
.with_websocket_pool(Arc::clone(&pool))
.with_websocket_fast_path_budget(Duration::from_millis(30));
let response = backend
@@ -52,10 +53,251 @@ async fn cold_websocket_should_fall_back_to_same_account_http_after_fast_budget(
assert_eq!(response.transport, CodexBackendTransport::HttpSse);
assert_eq!(
response.transport_metrics.decision,
Some(CodexTransportDecision::Http2WebSocketSlow)
Some(CodexTransportDecision::Http2WebSocketBudgetExhausted)
);
assert_eq!(response.transport_metrics.ws_connect_ms, None);
assert!(response.transport_metrics.first_event_ms.is_some());
assert!(response.body.contains("response.completed"));
pool.shutdown().await;
}
#[tokio::test]
async fn timed_out_websocket_should_finish_in_background_and_serve_the_next_request() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let (websocket_ready_tx, websocket_ready_rx) = tokio::sync::oneshot::channel();
let server = tokio::spawn(async move {
let (websocket_stream, _) = listener.accept().await.unwrap();
let websocket_server = tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(60)).await;
let mut websocket = accept_codex_test_websocket(websocket_stream).await;
websocket_ready_tx.send(()).unwrap();
let _payload = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
completed_websocket_response("resp_background_ready", 2, 1).into(),
))
.await
.unwrap();
});
let (mut http, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut http).await;
assert!(request.starts_with("POST /codex/responses HTTP/1.1"));
write_completed_sse_response(&mut http).await;
websocket_server.await.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::default());
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::runtime_test_fingerprint(),
)
.with_websocket_pool(Arc::clone(&pool))
.with_websocket_fast_path_budget(Duration::from_millis(30));
let request = new_chain_request("conversation-background-ready");
let first = backend
.create_response(
&request,
request_context("req_background_http", Some("chatgpt-account")),
)
.await
.expect("foreground request should use HTTP after its fast-path budget");
websocket_ready_rx.await.unwrap();
tokio::time::sleep(Duration::from_millis(20)).await;
let second = backend
.create_response(
&request,
request_context("req_background_reuse", Some("chatgpt-account")),
)
.await
.expect("next request should reuse the background websocket");
server.await.unwrap();
assert_eq!(first.transport, CodexBackendTransport::HttpSse);
assert_eq!(first.transport_metrics.ws_connect_ms, None);
assert_eq!(second.transport, CodexBackendTransport::WebSocket);
assert_eq!(
second.transport_metrics.decision,
Some(CodexTransportDecision::ReusedWebSocket)
);
assert!(second.body.contains("resp_background_ready"));
pool.shutdown().await;
}
#[tokio::test]
async fn shared_websocket_opening_should_keep_the_original_fast_path_deadline() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (mut stalled_websocket, _) = listener.accept().await.unwrap();
let opening = read_http_request(&mut stalled_websocket).await;
assert!(opening.starts_with("GET /codex/responses HTTP/1.1"));
for _ in 0..2 {
let (mut http, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut http).await;
assert!(request.starts_with("POST /codex/responses HTTP/1.1"));
write_completed_sse_response(&mut http).await;
}
});
let pool = Arc::new(CodexWebSocketPool::default());
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::runtime_test_fingerprint(),
)
.with_websocket_pool(Arc::clone(&pool))
.with_websocket_fast_path_budget(Duration::from_millis(200));
let request = new_chain_request("conversation-original-deadline");
let first = backend
.create_response(
&request,
request_context("req_original_deadline_first", Some("chatgpt-account")),
)
.await
.expect("first request should use HTTP after the opening budget");
let second = timeout(
Duration::from_millis(100),
backend.create_response(
&request,
request_context("req_original_deadline_second", Some("chatgpt-account")),
),
)
.await
.expect("shared opening must not grant the second request another full budget")
.expect("second request should use HTTP");
server.await.unwrap();
assert_eq!(first.transport, CodexBackendTransport::HttpSse);
assert_eq!(second.transport, CodexBackendTransport::HttpSse);
assert_eq!(
second.transport_metrics.decision,
Some(CodexTransportDecision::Http2WebSocketBudgetExhausted)
);
pool.shutdown().await;
}
#[tokio::test]
async fn account_eviction_should_cancel_a_background_websocket_opening() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (mut stale_opening, _) = listener.accept().await.unwrap();
let opening = read_http_request(&mut stale_opening).await;
assert!(opening.starts_with("GET /codex/responses HTTP/1.1"));
let (mut http, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut http).await;
assert!(request.starts_with("POST /codex/responses HTTP/1.1"));
write_completed_sse_response(&mut http).await;
let mut byte = [0_u8; 1];
let read = timeout(Duration::from_secs(1), stale_opening.read(&mut byte))
.await
.expect("account eviction should close the opening socket")
.unwrap();
assert_eq!(read, 0);
let (fresh_stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_codex_test_websocket(fresh_stream).await;
let _payload = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
completed_websocket_response("resp_after_eviction", 2, 1).into(),
))
.await
.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::runtime_test_fingerprint(),
)
.with_websocket_pool(Arc::clone(&pool))
.with_websocket_fast_path_budget(Duration::from_millis(30));
let request = new_chain_request("conversation-evict-opening");
let first = backend
.create_response(
&request,
request_context("req_evict_opening_first", Some("chatgpt-account")),
)
.await
.expect("first request should use HTTP after the opening budget");
pool.evict_account("chatgpt-account").await;
let second = timeout(
Duration::from_secs(2),
backend
.clone()
.with_websocket_fast_path_budget(Duration::from_millis(500))
.create_response(
&request,
request_context("req_evict_opening_second", Some("chatgpt-account")),
),
)
.await
.expect("fresh websocket request should finish")
.expect("the next request should build a fresh websocket");
timeout(Duration::from_secs(2), server)
.await
.expect("test server should finish")
.unwrap();
assert_eq!(first.transport, CodexBackendTransport::HttpSse);
assert_eq!(second.transport, CodexBackendTransport::WebSocket);
assert!(second.body.contains("resp_after_eviction"));
timeout(Duration::from_secs(2), pool.shutdown())
.await
.expect("pool shutdown should finish");
}
#[tokio::test]
async fn pool_shutdown_should_cancel_and_join_a_background_websocket_opening() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (mut stalled_websocket, _) = listener.accept().await.unwrap();
let opening = read_http_request(&mut stalled_websocket).await;
assert!(opening.starts_with("GET /codex/responses HTTP/1.1"));
let (mut http, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut http).await;
assert!(request.starts_with("POST /codex/responses HTTP/1.1"));
write_completed_sse_response(&mut http).await;
let mut byte = [0_u8; 1];
let read = timeout(Duration::from_secs(1), stalled_websocket.read(&mut byte))
.await
.expect("pool shutdown should close the opening socket")
.unwrap();
assert_eq!(read, 0);
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::runtime_test_fingerprint(),
)
.with_websocket_pool(Arc::clone(&pool))
.with_websocket_fast_path_budget(Duration::from_millis(30));
backend
.create_response(
&new_chain_request("conversation-shutdown-opening"),
request_context("req_shutdown_opening", Some("chatgpt-account")),
)
.await
.expect("foreground request should use HTTP after the opening budget");
timeout(Duration::from_secs(1), pool.shutdown())
.await
.expect("shutdown should join the cancelled opening task");
server.await.unwrap();
assert!(pool.is_shutdown().await);
}
#[tokio::test]
@@ -399,6 +641,145 @@ async fn origin_breaker_should_open_then_allow_only_one_half_open_probe() {
));
}
#[test]
fn origin_breaker_should_count_hard_opening_failures() {
let breaker = WebSocketOriginBreaker::with_config(WebSocketOriginBreakerConfig {
failure_threshold: 1,
failure_window: Duration::from_secs(1),
open_duration: Duration::from_secs(1),
});
let WebSocketOriginBreakerDecision::Allowed(permit) =
breaker.try_acquire("https://example.test:443")
else {
panic!("closed breaker should allow an opening");
};
permit.fail();
assert!(matches!(
breaker.try_acquire("https://example.test:443"),
WebSocketOriginBreakerDecision::Open
));
}
#[tokio::test]
async fn fast_path_miss_and_late_failure_should_count_as_one_breaker_failure() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let (late_failure_tx, late_failure_rx) = tokio::sync::oneshot::channel();
let (second_used_websocket_tx, second_used_websocket_rx) = tokio::sync::oneshot::channel();
let server = tokio::spawn(async move {
let (mut delayed_opening, _) = listener.accept().await.unwrap();
let opening = read_http_request(&mut delayed_opening).await;
assert!(opening.starts_with("GET /codex/responses HTTP/1.1"));
let delayed_failure = tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(80)).await;
delayed_opening
.write_all(
b"HTTP/1.1 503 Service Unavailable\r\ncontent-length: 0\r\nconnection: close\r\n\r\n",
)
.await
.unwrap();
late_failure_tx.send(()).unwrap();
});
let (mut first_http, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut first_http).await;
assert!(request.starts_with("POST /codex/responses HTTP/1.1"));
write_completed_sse_response(&mut first_http).await;
delayed_failure.await.unwrap();
let (second_stream, _) = listener.accept().await.unwrap();
let second_is_websocket = timeout(Duration::from_secs(1), async {
let mut prefix = [0_u8; 4];
loop {
let read = second_stream.peek(&mut prefix).await.unwrap();
if read == 0 {
return false;
}
if read == prefix.len() {
return prefix == *b"GET ";
}
tokio::task::yield_now().await;
}
})
.await
.expect("second request should reach the upstream");
if second_is_websocket {
second_used_websocket_tx.send(true).unwrap();
let mut websocket = accept_codex_test_websocket(second_stream).await;
let _payload = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
completed_websocket_response("resp_after_late_failure", 2, 1).into(),
))
.await
.unwrap();
} else {
second_used_websocket_tx.send(false).unwrap();
let mut second_http = second_stream;
let request = read_http_request(&mut second_http).await;
assert!(request.starts_with("POST /codex/responses HTTP/1.1"));
write_completed_sse_response(&mut second_http).await;
}
});
let breaker = WebSocketOriginBreaker::with_config(WebSocketOriginBreakerConfig {
failure_threshold: 2,
failure_window: Duration::from_secs(1),
open_duration: Duration::from_secs(1),
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::runtime_test_fingerprint(),
)
.with_websocket_pool(Arc::clone(&pool))
.with_websocket_fast_path_budget(Duration::from_millis(30))
.with_websocket_origin_breaker(breaker);
let first = backend
.create_response(
&new_chain_request("conversation-late-failure-first"),
request_context("req_late_failure_first", Some("chatgpt-account")),
)
.await
.expect("first request should use HTTP after the opening budget");
timeout(Duration::from_secs(2), late_failure_rx)
.await
.expect("delayed opening should receive its 503 response")
.unwrap();
tokio::time::sleep(Duration::from_millis(20)).await;
let second = timeout(
Duration::from_secs(2),
backend
.clone()
.with_websocket_fast_path_budget(Duration::from_millis(500))
.create_response(
&new_chain_request("conversation-late-failure-second"),
request_context("req_late_failure_second", Some("chatgpt-account")),
),
)
.await
.expect("second request should finish")
.expect("one degraded opening must not open a threshold-two breaker");
let second_used_websocket = timeout(Duration::from_secs(2), second_used_websocket_rx)
.await
.expect("test server should observe the second transport")
.unwrap();
timeout(Duration::from_secs(2), server)
.await
.expect("test server should finish")
.unwrap();
assert_eq!(first.transport, CodexBackendTransport::HttpSse);
assert!(second_used_websocket);
assert_eq!(second.transport, CodexBackendTransport::WebSocket);
assert!(second.body.contains("resp_after_late_failure"));
timeout(Duration::from_secs(2), pool.shutdown())
.await
.expect("pool shutdown should finish");
}
#[tokio::test]
async fn half_open_upstream_response_should_close_origin_breaker() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -1327,21 +1327,33 @@ async fn codex_backend_client_stream_should_preserve_burst_during_downstream_bac
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let (terminal_drained_tx, terminal_drained_rx) = tokio::sync::oneshot::channel();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
json!({
"type": "response.output_text.delta",
"delta": "initial;"
})
.to_string()
.into(),
))
.await
.unwrap();
let (mut sink, mut source) = websocket.split();
// 服务端持续读取控制帧,让 tungstenite 正常回 Pong。若客户端在本地入站背压
// 期间仍执行 Pong deadline,它仍会因为暂时读不到已返回的 Pong 而误杀连接。
let control_frames = tokio::spawn(async move {
while let Some(message) = source.next().await {
match message.unwrap() {
Message::Ping(_) | Message::Pong(_) => {}
Message::Close(_) => break,
other => panic!("unexpected client frame during response: {other:?}"),
}
}
});
sink.send(Message::Text(
json!({
"type": "response.output_text.delta",
"delta": "initial;"
})
.to_string()
.into(),
))
.await
.unwrap();
for index in 0..BURST_FRAMES {
let frame = json!({
@@ -1349,22 +1361,34 @@ async fn codex_backend_client_stream_should_preserve_burst_during_downstream_bac
"delta": format!("chunk-{index};")
})
.to_string();
if websocket.send(Message::Text(frame.into())).await.is_err() {
if sink.send(Message::Text(frame.into())).await.is_err() {
control_frames.abort();
return;
}
}
let _ = websocket
let _ = sink
.send(Message::Text(
completed_websocket_response("resp_backpressure", 3, 257).into(),
))
.await;
let _ = websocket.close(None).await;
terminal_drained_rx
.await
.expect("client should drain the terminal response before server close");
let _ = sink.send(Message::Close(None)).await;
control_frames.abort();
});
let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig {
ping_interval: Some(Duration::from_millis(5)),
ping_timeout: Duration::from_millis(20),
liveness_timeout: None,
..websocket_pool_config_for_tests(None, None, None)
}));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::runtime_test_fingerprint(),
);
)
.with_websocket_pool(pool);
let request = pooled_websocket_request("conversation-backpressure");
let mut response = backend
@@ -1386,6 +1410,7 @@ async fn codex_backend_client_stream_should_preserve_burst_during_downstream_bac
})
.await
.expect("backpressured websocket stream should finish");
terminal_drained_tx.send(()).unwrap();
server.await.unwrap();
assert_eq!(
@@ -162,7 +162,7 @@ async fn codex_backend_client_should_open_fresh_socket_when_idle_pooled_websocke
assert_eq!(accepted_connections.load(Ordering::SeqCst), 2);
}
#[tokio::test(start_paused = true)]
#[tokio::test]
async fn codex_backend_client_stream_should_keep_fresh_socket_after_structural_activity() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
@@ -550,7 +550,7 @@ async fn websocket_pool_should_use_http_while_exact_key_is_busy() {
);
}
#[tokio::test(start_paused = true)]
#[tokio::test]
async fn websocket_pool_should_release_slot_when_client_drops_stream() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
@@ -714,29 +714,48 @@ async fn websocket_pool_should_use_http_for_new_keys_after_account_cap() {
}
#[tokio::test]
async fn websocket_pool_should_evict_idle_connection_when_liveness_lapses_despite_pings() {
// 不变量:pump 自己发出的 keepalive ping 不算“入站活动”,因此一个只收 ping、
// 从不回 pong / 从不发帧的静默连接,仍会被 liveness watchdog 判定失活并驱逐;
// 复用前 acquire 读到 is_closed 即开新连接。
async fn websocket_pool_should_replace_idle_connection_after_pong_deadline() {
// 服务端读取 Ping 后暂不继续 poll,避免 tungstenite 自动回 Pong;pump 必须在独立
// deadline 到期时主动关闭连接,acquire 随后只读 closed 状态并直接新建连接。
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let accepted_connections = Arc::new(AtomicUsize::new(0));
let accepted_connections_for_server = Arc::clone(&accepted_connections);
let (release_first_tx, release_first_rx) = tokio::sync::oneshot::channel();
let (ping_seen_tx, ping_seen_rx) = tokio::sync::oneshot::channel();
let (inspect_close_tx, inspect_close_rx) = tokio::sync::oneshot::channel();
let (pong_timeout_closed_tx, pong_timeout_closed_rx) = tokio::sync::oneshot::channel();
let server = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let first_connection = tokio::spawn(async move {
let mut first_websocket = accept_codex_test_websocket(first_stream).await;
let _first_message = first_websocket.next().await.unwrap().unwrap();
first_websocket
.send(Message::Text(
completed_websocket_response("resp_no_pong_first", 2, 1).into(),
))
.await
.unwrap();
let _ = release_first_rx.await;
});
let mut first_websocket = accept_codex_test_websocket(first_stream).await;
let _first_message = first_websocket.next().await.unwrap().unwrap();
first_websocket
.send(Message::Text(
completed_websocket_response("resp_no_pong_first", 2, 1).into(),
))
.await
.unwrap();
let ping = timeout(Duration::from_secs(1), first_websocket.next())
.await
.expect("pump should send a keepalive ping")
.expect("keepalive ping should be present")
.expect("keepalive ping should be valid");
let Message::Ping(payload) = ping else {
panic!("expected keepalive ping, got {ping:?}");
};
assert_eq!(payload.len(), 8);
ping_seen_tx.send(()).unwrap();
inspect_close_rx
.await
.expect("test should release the server after the Pong deadline");
let close = timeout(Duration::from_secs(1), first_websocket.next())
.await
.expect("Pong deadline should close the idle websocket")
.expect("Pong deadline should send a close frame")
.expect("close frame should be valid");
std::assert_matches!(close, Message::Close(_));
pong_timeout_closed_tx.send(()).unwrap();
let (second_stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
@@ -749,11 +768,11 @@ async fn websocket_pool_should_evict_idle_connection_when_liveness_lapses_despit
.await
.unwrap();
second_websocket.close(None).await.unwrap();
first_connection.await.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig {
ping_interval: Some(Duration::from_millis(1)),
liveness_timeout: Some(Duration::from_millis(20)),
ping_interval: Some(Duration::from_millis(10)),
ping_timeout: Duration::from_millis(30),
liveness_timeout: None,
maintenance_interval: None,
..websocket_pool_config_for_tests(None, None, None)
}));
@@ -772,12 +791,16 @@ async fn websocket_pool_should_evict_idle_connection_when_liveness_lapses_despit
)
.await
.expect("first websocket response should succeed");
tokio::time::pause();
tokio::time::advance(Duration::from_millis(25)).await;
tokio::task::yield_now().await;
// 只加速首条 idle 连接的失活窗口;恢复实时时钟,避免新连接在发送
// response.create 前被测试时钟自动推进到下一次 liveness timeout。
tokio::time::resume();
timeout(Duration::from_secs(1), ping_seen_rx)
.await
.expect("server should observe the keepalive ping")
.expect("server should report the keepalive ping");
tokio::time::sleep(Duration::from_millis(80)).await;
inspect_close_tx.send(()).unwrap();
timeout(Duration::from_secs(1), pong_timeout_closed_rx)
.await
.expect("server should observe the Pong-timeout close")
.expect("server should report the Pong-timeout close");
let second = backend
.create_response(
&request,
@@ -785,7 +808,6 @@ async fn websocket_pool_should_evict_idle_connection_when_liveness_lapses_despit
)
.await
.expect("second websocket response should use a fresh connection");
release_first_tx.send(()).unwrap();
server.await.unwrap();
assert!(first.body.contains("resp_no_pong_first"));
@@ -794,7 +816,7 @@ async fn websocket_pool_should_evict_idle_connection_when_liveness_lapses_despit
assert_eq!(accepted_connections.load(Ordering::SeqCst), 2);
}
#[tokio::test(start_paused = true)]
#[tokio::test]
async fn websocket_pool_should_gc_expired_idle_connections() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
@@ -868,7 +890,7 @@ async fn websocket_pool_should_gc_expired_idle_connections() {
assert_eq!(accepted_connections.load(Ordering::SeqCst), 2);
}
#[tokio::test(start_paused = true)]
#[tokio::test]
async fn codex_backend_client_should_keep_idle_pooled_websocket_alive_across_repeated_pings() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
@@ -901,6 +923,7 @@ async fn codex_backend_client_should_keep_idle_pooled_websocket_alive_across_rep
// pump 会在 idle 期间反复发送 keepalive ping;服务端计数并回 pong,
// 直到下一个业务请求(response.create)到达为止。
let mut previous_ping = None;
loop {
let message = timeout(Duration::from_secs(1), websocket.next())
.await
@@ -909,6 +932,11 @@ async fn codex_backend_client_should_keep_idle_pooled_websocket_alive_across_rep
.expect("frame should be valid");
match message {
Message::Ping(payload) => {
assert_eq!(payload.len(), 8);
if let Some(previous) = &previous_ping {
assert_ne!(previous, &payload, "keepalive ping sequence must be unique");
}
previous_ping = Some(payload.clone());
ping_count_for_server.fetch_add(1, Ordering::SeqCst);
websocket.send(Message::Pong(payload)).await.unwrap();
}
@@ -978,7 +1006,66 @@ async fn codex_backend_client_should_keep_idle_pooled_websocket_alive_across_rep
assert!(first.body.contains("resp_pool_background_first"));
assert!(second.body.contains("resp_pool_background_second"));
assert!(ping_count.load(Ordering::SeqCst) >= 1);
assert!(ping_count.load(Ordering::SeqCst) >= 2);
}
#[tokio::test]
async fn codex_backend_client_should_treat_active_business_frames_as_ping_liveness() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _request = websocket.next().await.unwrap().unwrap();
for index in 0..8 {
tokio::time::sleep(Duration::from_millis(20)).await;
websocket
.send(Message::Text(
json!({
"type": "response.output_text.delta",
"delta": format!("active-{index};")
})
.to_string()
.into(),
))
.await
.unwrap();
}
websocket
.send(Message::Text(
completed_websocket_response("resp_active_without_pong", 3, 8).into(),
))
.await
.unwrap();
websocket.close(None).await.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig {
ping_interval: Some(Duration::from_millis(10)),
ping_timeout: Duration::from_millis(60),
liveness_timeout: None,
..websocket_pool_config_for_tests(None, None, None)
}));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::runtime_test_fingerprint(),
)
.with_websocket_pool(Arc::clone(&pool));
let request = pooled_websocket_request("conversation-active-without-pong");
let response = backend
.create_response(
&request,
request_context("req_active_without_pong", Some("chatgpt-account")),
)
.await
.expect("active business frames should satisfy the Ping liveness probe");
server.await.unwrap();
assert!(response.body.contains("active-0;"));
assert!(response.body.contains("active-7;"));
assert!(response.body.contains("resp_active_without_pong"));
}
#[tokio::test]
+34 -19
View File
@@ -28,11 +28,11 @@ PostgreSQL 和 Redis 都是启动必需依赖。运行中任一健康检查失
## 架构不变量
1. `v1/*` 使用 Request → Attempt → Stream 三层静态洋葱生命周期;管理端、健康检查和静态资源不进入业务洋葱。
2. 一个功能只有一个规则 owner。阈值、分类、状态机和 decision 只能存在于对应 controller。
2. 一个功能只有一个规则 owner。跨入口的账号失败语义属于 `fleet/account_failure.rs`;`v1/*` 的 retry、exhaustion 和 commit decision 属于对应 controller。
3. API 只解析和编码;lifecycle 只编排顺序、retry、commit 和 finalize;transport 只建立连接并产出 typed facts。
4. 上游响应一旦可能收到 payload,transport 或账号都不得自动重放同一请求。
5. 账号认证身份按 attempt 重建;客户端会话拓扑按原值透传;两者不能共用一个“身份”抽象。
6. 账号明确失效先在线性化的内存路径生效并驱逐 WebSocket,再异步持久化 PostgreSQL。
6. `v1/*` 热路径中的账号明确失效先在线性化的内存路径生效并驱逐 WebSocket,再异步持久化 PostgreSQL;管理端和后台任务可等待同一状态 effect 持久化完成。
7. 每项生产能力只有一条执行路径,不设置兼容 wrapper、双状态机或运行时切换的新旧实现。
8. `backend/src/` 只放生产代码;所有测试和 test-only helper 位于 `backend/tests/`。
@@ -104,6 +104,7 @@ backend/src/
```text
upstream/openai/
├── token_client.rs
├── failure.rs
├── protocol/
│ ├── responses.rs
│ ├── events.rs
@@ -127,6 +128,7 @@ upstream/openai/
```
- `protocol` 定义 Responses body、canonical SSE/WS 事实和序列化,不包含账号策略。
- `failure.rs` 将各 transport 错误规范化为稳定的 `UpstreamFailureFacts`,不解释账号状态。
- `transport` 拥有 HTTP/SSE、WebSocket、TLS、header、连接池、熔断和 transport metrics。
- `fingerprint` 拥有 Codex Desktop 指纹的存储、运行时快照和更新。
- `token_client.rs` 只实现 OAuth token 刷新协议。
@@ -137,6 +139,7 @@ upstream/openai/
- `store`:账号与 quota 的 PostgreSQL 映射和事务写入。
- `pool`:进程内账号快照、并发 slot、请求间隔和候选 lease。
- `account_failure.rs`:HTTP、SSE、手动额度查询和后台额度刷新共享的账号失败分类与状态 effect。
- `scheduler`:候选过滤与排序的唯一 owner,支持 `smart`、`quota_reset_priority`、`round_robin`、`sticky`。
- `scheduler/feedback`:进程内 EWMA 错误率和首字延迟。
- `quota`:配额查询、窗口与运行时限制。
@@ -225,19 +228,20 @@ Stream terminal 包括 completed、incomplete、failed、上游断开、下游
| --- | --- |
| previous response 所有权、scope、full replay | `controllers/history.rs` |
| cyber policy 会话排除和 CAS 清理 | `controllers/cyber_policy/` |
| quota 与 rate limit | `controllers/quota.rs` |
| 封禁、过期、模型不支持等账号事实 | `controllers/account_failure.rs` |
| 内存失效、异步持久化和 WS 驱逐 effect | `controllers/account_state.rs` |
| v1 quota 与 rate limit attempt decision | `controllers/quota.rs` |
| v1 封禁、过期、模型不支持的 exhaustion decision | `controllers/account_failure.rs` |
| transport 错误到稳定上游事实 | `upstream/openai/failure.rs` |
| 跨入口账号失败分类、状态 effect 与 WS 驱逐 | `fleet/account_failure.rs` |
| Cloudflare Cookie 与冷却 | `controllers/cloudflare.rs` |
| 会话亲和写入 | `controllers/affinity.rs` |
| usage 结算 | `controllers/usage.rs` |
| 代理事件与 transport metrics | `controllers/telemetry/` |
Controller 之间不直接调用。跨功能优先级只在 `ControllerSet` 声明;API、transport、store 和 lifecycle 不重复识别功能语义。
Controller 之间不直接调用。跨功能优先级只在 `ControllerSet` 声明;API、transport、store 和 lifecycle 不重复识别功能语义。Controller 只把共享账号分类翻译为 v1 生命周期 decision;手动额度查询和后台额度刷新执行同一状态 effect,但各自保留自己的返回与调度职责。
## 上游双热传输
Responses transport 使用精确 WebSocket 与共享 HTTP/2 两条常热通道。WebSocket 保存 connection-local response 状态;HTTP/2 提供新链和可持久化续链的低延迟后备。
Responses transport 使用会话级 WebSocket 池与 origin 级 HTTP/2 池两条常热通道。WebSocket 保存 connection-local response 状态;HTTP/2 提供新链和可持久化续链的低延迟后备。transport 只竞速连接准备,不并行发送业务 payload。
### TransportRequirement
@@ -258,11 +262,14 @@ Responses transport 使用精确 WebSocket 与共享 HTTP/2 两条常热通道
`PreparedResponseTransport` 只表示“账号、header 和连接已准备,但 payload 尚未发送”。WebSocket opening 与 `response.create` 发送是单向类型边界:
- 普通新链、persisted continuation 和 external unknown 的冷 WS opening 预算为 800ms。
- 预算内连接成功或复用热 socket 时发送 WS payload。
- 预算耗尽、DNS/TCP/TLS 失败、pool 不可用或 breaker 打开时,同账号使用 HTTP/2。
- 热 WS 可用时立即租用并发送 WS payload。
- 普通新链、persisted continuation 和 external unknown 没有热 WS 时创建一个 `Connecting` slot;前台从该 slot 的统一起点最多等待 800ms。
- 800ms 内握手成功时,连接先原子进入 `Idle`,当前请求再租用为 `Busy` 并发送 WS payload。
- 800ms 耗尽时,当前请求不发送 WS payload,立即在同账号使用 HTTP/2;未完成的握手由 pool 在后台继续,仍受 15 秒硬连接超时约束。
- 后台握手成功后只进入 `Idle`,供后续请求复用;最终失败则删除 `Connecting` slot。相同 key 的后续请求沿用原 opening 的绝对 deadline,不会重新获得一段 800ms。
- DNS/TCP/TLS 立即失败、pool 不可用或 breaker 打开时,同账号使用 HTTP/2。
- opening 已返回明确上游响应时,事实直接进入 controller,不追加 HTTP 请求。
- warmup 与 exact continuation 禁止 HTTP fallback。
- warmup 与 exact continuation 必须等待所需 WS,不使用 HTTP fallback。
- payload 发送后的连接断开、发送超时或首事件超时统一为 `PostSendAmbiguous`,不得换 transport 或账号。
AttemptRunner 将 prepared transport 与账号请求间隔并行,800ms、熔断阈值和连接状态只存在于 `upstream/openai/transport`。
@@ -273,21 +280,26 @@ Pool key 为 `(base_url, account_id, local_conversation_id)`,slot 状态为:
- `Idle`:可租用的热 socket。
- `Busy`:单 socket 上已有一个 in-flight response。
- `Connecting`:同 key opening 单飞,其他请求不会再建第二条 socket。
- `Connecting`:同 key opening 单飞,包含统一 started-at、结果通知和取消令牌。
每条连接记录最新 completed response ID。Exact continuation 只有 ID 完全匹配时才能租用;socket 缺失、busy、失活或 ID 不匹配时不新建连接碰运气,由 HistoryController 决定 full replay 或失败。
池化 socket 默认 25 秒 ping、5 秒探活、55 分钟最大寿命;后台 pump 处理 ping/pong、EOF、close 和 liveness,复用前只读取连接状态。
池化 socket 默认每 25 秒发送带唯一序列的 Ping,匹配 Pong 的 deadline 为 5 秒,最大寿命为 55 分钟。任意业务入站帧同样证明链路存活;下游背压导致 pump 无法读取 socket 时暂停主动探活。后台 pump 统一处理 ping/pong、EOF、close 和 liveness,复用前只读取原子连接状态,不增加网络往返。
`Idle`、`Busy` 和 `Connecting` 都计入 `ws_pool.max_per_account`,默认单账号最多 8 个 slot;后台续建不设置另一套队列或全局连接上限。账号驱逐会删除该账号全部 slot、关闭 idle socket 并取消 connecting;busy socket 在归还时因 reservation 已失效而关闭。进程关闭先取消所有 pool-owned opening,再等待维护任务和建连任务退出,迟到连接不能在驱逐或 shutdown 后重新入池。
### Origin breaker
Breaker key 为 origin 与 TLS profile,不按账号区分:
- 30 秒窗口内 3 次 WS 快路径超时后打开 30 秒。
- 30 秒窗口内 3 次退化 opening 后打开 30 秒;退化包括 800ms 快路径预算耗尽、transport 失败和 5xx opening 响应。
- 同一次 opening 的前台预算耗尽与后台最终失败只计一次;迟到成功也不抹去已经发生的快路径退化。
- 预算内成功 opening 会清空该 origin 的历史退化记录。
- 打开期间热 WS 仍可使用;普通冷 opening 直接走 HTTP/2。
- 到期后只允许一个 half-open probe。
- 账号或请求级 4xx opening 响应证明 origin 可达并关闭 half-open。
- 5xx 或 transport probe 失败重新打开 breaker。
- 账号驱逐和进程 shutdown 触发的主动取消不计为 origin 故障。
Breaker 不更新账号状态;账号失效只由 controller 的明确 typed fact 触发。
@@ -318,7 +330,9 @@ HTTP/2 连接按 origin/TLS profile 共享,账号 Authorization 只属于单
- `httpVersion`
- `websocketPool.kind`
稳定 decision 包括 `ws_reused`、`ws_connected_fast`、`ws_exact_required`、`ws_required`、`http2_ws_slow`、`http2_breaker_open`、`http2_pool_unavailable` 和 `http2_ws_pre_send_failure`。
`wsConnectMs` 只表示已经成功的 WS 握手耗时;复用、预算耗尽和其他 HTTP fallback 不填写该字段。`transportDecisionWaitMs` 表示当前请求为 transport 选择实际等待的时间。后台 opening 另外记录 `ws_preconnect_duration_ms`、`ws_preconnect_outcome` 和 `foreground_waiting`,不把未成功握手的 800ms 伪装成连接耗时。
稳定 decision 包括 `ws_reused`、`ws_connected_fast`、`ws_exact_required`、`ws_required`、`http2_ws_budget_exhausted`、`http2_breaker_open`、`http2_pool_unavailable` 和 `http2_ws_pre_send_failure`。
## 账号调度与失效
@@ -327,10 +341,11 @@ Scheduler 先过滤不可用账号,再按运行时 `rotation_strategy` 排序
明确上游事实的处理顺序:
1. transport 提取 status、code、type、message、retry-after 等事实。
2. controller 分类为 quota、rate limit、expired、disabled、banned、model unsupported、Cloudflare 或 cyber policy。
3. 对账号全局事实,账号池 mutex 内立即失效 runtime state 并移除 slot。
4. 驱逐该账号的 idle/busy/connecting WebSocket entry。
5. lifecycle 从冻结候选中选择下一账号。
2. `fleet/account_failure.rs` 将跨入口账号事实统一分类为 quota、rate limit、expired、disabled、banned 或 model unsupported;Cloudflare 与 cyber policy 保持各自 owner。
3. v1 controller 把共享分类翻译为 retry/exhaustion decision。
4. 对账号全局事实,账号池 mutex 内立即失效 runtime state 并移除 slot。
5. 驱逐该账号的 idle/busy/connecting WebSocket entry。
6. lifecycle 从冻结候选中选择下一账号。
6. PostgreSQL 状态写入后台有序队列。
数据库延迟不会延长当前换号路径,也不会使并发请求重新租到已失效账号。
@@ -61,9 +61,9 @@ const avatarClass = computed(() => {
>
{{ initial }}
</span>
<div class="min-w-0">
<div class="min-w-0 flex-1">
<div class="flex min-w-0 items-center gap-2">
<span class="truncate text-[14px] font-[760] text-(--cp-text-primary)">
<span class="min-w-0 flex-1 truncate text-[14px] font-[760] text-(--cp-text-primary)">
{{ displayTitle }}
</span>
<AccountPlanBadge v-if="showPlan" :plan-type="account.planType" size="sm" />
@@ -23,8 +23,8 @@ const label = computed(() => props.planType?.trim() || 'Free')
const sizeClass = computed(() =>
props.size === 'sm'
? 'h-5 max-w-24 rounded-full px-1.75 text-[11px] font-[720]'
: 'h-5.5 max-w-full rounded-full px-2 text-[11px] font-[760]',
? 'h-5 rounded-full px-1.75 text-[11px] font-[720]'
: 'h-5.5 rounded-full px-2 text-[11px] font-[760]',
)
const paletteClass = computed(() => {
@@ -39,9 +39,9 @@ const paletteClass = computed(() => {
<template>
<span
class="inline-flex min-w-0 items-center justify-center leading-none capitalize"
class="inline-flex shrink-0 items-center justify-center whitespace-nowrap leading-none capitalize"
:class="[sizeClass, paletteClass]"
>
<span class="min-w-0 truncate">{{ label }}</span>
<span>{{ label }}</span>
</span>
</template>
@@ -139,7 +139,7 @@ export function useDashboard(): any {
todayTokensValue: 0,
yesterdayTokensValue: 0,
totalTokens: '0',
todayBillingAmountUsd: '—',
totalBillingAmountUsd: '—',
},
cache: {
todayHitRate: '—',
@@ -222,7 +222,7 @@ export function useDashboard(): any {
details: [
{ label: '总请求', value: traffic.totalRequests, tone: 'info' },
{
label: '首字延迟',
label: '今日首字',
value: cache.averageFirstTokenLatencyMs,
tone: 'info',
},
@@ -243,8 +243,8 @@ export function useDashboard(): any {
details: [
{ label: '总 Token', value: tokens.totalTokens, tone: 'success' },
{
label: '计费',
value: tokens.todayBillingAmountUsd,
label: '总计费',
value: tokens.totalBillingAmountUsd,
tone: 'success',
},
],
@@ -266,9 +266,9 @@ export function useDashboard(): any {
'warning',
),
details: [
{ label: '平均缓存命中', value: cache.totalHitRate, tone: 'warning' },
{ label: '总缓存命中', value: cache.totalHitRate, tone: 'warning' },
{
label: '缓存',
label: '总缓存',
value: cache.totalCachedTokens,
tone: 'warning',
},
@@ -79,8 +79,10 @@ const displayItems = computed(() =>
const tableRows = computed(() => (isStale.value ? [] : displayItems.value))
function diagnosticNameDisplay(name: string) {
const full = name.trim() || '未知'
if (dimension.value !== 'model') {
const raw = name.trim() || '未知'
const full =
dimension.value === 'transport' ? ({ websocket: 'WS', http_sse: 'SSE' }[raw] ?? raw) : raw
if (dimension.value !== 'model' && dimension.value !== 'account') {
return { primary: full, secondary: '', full }
}
@@ -141,7 +143,11 @@ function diagnosticNameDisplay(name: string) {
v-if="row.nameDisplay.secondary"
class="flex min-w-0 items-center gap-1.25 text-(--cp-text-secondary)"
>
<CornerDownRight class="size-3.25 shrink-0 text-(--cp-info)" stroke-width="2.4" />
<CornerDownRight
v-if="dimension !== 'account'"
class="size-3.25 shrink-0 text-(--cp-info)"
stroke-width="2.4"
/>
<code class="block truncate font-mono text-[11px] leading-none font-bold">
{{ row.nameDisplay.secondary }}
</code>
@@ -1,5 +1,6 @@
import { clamp } from 'es-toolkit'
import { onMounted, shallowRef, watch } from 'vue'
import type { Ref } from 'vue'
import {
getUsageRecordInsightsDiagnostics,
@@ -10,12 +11,28 @@ import {
import { toast } from '@/components/base/BaseToast'
import { withMinimumDuration } from '@/utils/async'
export function useUsageRecordsTable(options: any) {
import type { UsageTimeRangeParams } from './useUsageTimeRange'
interface UseUsageRecordsTableOptions {
page: Ref<number>
pageSize: Ref<number>
searchQuery: Ref<string>
timeRangeParams: Readonly<Ref<UsageTimeRangeParams>>
totalRecords: Ref<number>
latestTimeRangeParams: () => UsageTimeRangeParams
}
type UsageLoadScope = 'all' | 'table'
export function useUsageRecordsTable(options: UseUsageRecordsTableOptions) {
const loading = shallowRef(true)
const analyticsLoading = shallowRef(true)
const records = shallowRef<any[]>([])
const summary = shallowRef(emptySummary())
const insights = shallowRef(emptyInsights())
const tableTimeRangeParams = shallowRef<UsageTimeRangeParams>({
...options.timeRangeParams.value,
})
const refreshingList = shallowRef(false)
const diagnosticDimension = shallowRef('model')
const diagnosticLoading = shallowRef(false)
@@ -26,7 +43,7 @@ export function useUsageRecordsTable(options: any) {
search: options.searchQuery.value || undefined,
})
async function loadUsageRecords(scope = 'all') {
async function loadUsageRecords(scope: UsageLoadScope = 'all') {
const requestId = ++loadRequestId
try {
loading.value = true
@@ -37,8 +54,11 @@ export function useUsageRecordsTable(options: any) {
}
const globalParams = scopedParams()
if (scope === 'all') {
tableTimeRangeParams.value = { ...globalParams }
}
const tableParams = {
...globalParams,
...tableTimeRangeParams.value,
...filterParams(),
}
const resultPromise = getUsageRecords({
@@ -135,6 +155,7 @@ export function useUsageRecordsTable(options: any) {
if (refreshingList.value || loading.value) return
refreshingList.value = true
try {
tableTimeRangeParams.value = options.latestTimeRangeParams()
await withMinimumDuration(() => loadUsageRecords('table'))
} finally {
refreshingList.value = false
@@ -0,0 +1,45 @@
import dayjs, { type Dayjs } from 'dayjs'
import { computed, shallowRef } from 'vue'
export type UsageTimeRange = 'today' | '7d' | '30d'
export interface UsageTimeRangeParams extends Record<string, string> {
startTime: string
endTime: string
}
export function useUsageTimeRange(initialRange: UsageTimeRange = 'today') {
const timeRange = shallowRef<UsageTimeRange>(initialRange)
const rangeEnd = shallowRef(dayjs())
const timeRangeParams = computed(() => buildTimeRange(timeRange.value, rangeEnd.value))
function refreshTimeRangeEnd() {
rangeEnd.value = dayjs()
}
function latestTimeRangeParams() {
return buildTimeRange(timeRange.value, dayjs())
}
return {
timeRange,
timeRangeParams,
refreshTimeRangeEnd,
latestTimeRangeParams,
}
}
function buildTimeRange(range: UsageTimeRange, end: Dayjs): UsageTimeRangeParams {
if (range === 'today') {
return {
startTime: end.startOf('day').toISOString(),
endTime: end.toISOString(),
}
}
const days = range === '30d' ? 29 : 6
return {
startTime: end.subtract(days, 'day').startOf('day').toISOString(),
endTime: end.toISOString(),
}
}
+6 -25
View File
@@ -1,7 +1,6 @@
<script setup lang="ts">
import { CalendarDays, Eye, Minimize2 } from '@lucide/vue'
import dayjs from 'dayjs'
import { computed, shallowRef, watch } from 'vue'
import { shallowRef, watch } from 'vue'
import BaseButton from '@/components/base/BaseButton.vue'
import BaseCard from '@/components/base/BaseCard.vue'
@@ -19,6 +18,7 @@ import {
import { useUsageRecordDetail } from './composables/useUsageRecordDetail'
import { useUsageFilters } from './composables/useUsageFilters'
import { useUsageRecordsTable } from './composables/useUsageRecordsTable'
import { useUsageTimeRange } from './composables/useUsageTimeRange'
import UsageClientIpCell from './components/UsageClientIpCell.vue'
import UsageBillingCell from './components/UsageBillingCell.vue'
import UsageFilters from './components/UsageFilters.vue'
@@ -31,34 +31,13 @@ import UsageSummaryCards from './components/UsageSummaryCards.vue'
import UsageTokenCell from './components/UsageTokenCell.vue'
const totalRecords = shallowRef(0)
const timeRange = shallowRef('today')
const recordView = shallowRef('success')
const recordViewOptions = [
{ label: '成功记录', value: 'success' },
{ label: '错误排查', value: 'errors' },
]
const timeRangeParams = computed(() => {
const now = dayjs()
if (timeRange.value === 'today') {
return {
startTime: now.startOf('day').toISOString(),
endTime: now.toISOString(),
}
}
if (timeRange.value === '30d') {
return {
startTime: now.subtract(29, 'day').startOf('day').toISOString(),
endTime: now.toISOString(),
}
}
return {
startTime: now.subtract(6, 'day').startOf('day').toISOString(),
endTime: now.toISOString(),
}
})
const { timeRange, timeRangeParams, refreshTimeRangeEnd, latestTimeRangeParams } =
useUsageTimeRange()
const {
page,
@@ -87,6 +66,7 @@ const {
searchQuery,
timeRangeParams,
totalRecords,
latestTimeRangeParams,
})
const { showDetailModal, selectedUsageRecord, handleViewDetail } = useUsageRecordDetail({
@@ -96,6 +76,7 @@ const { showDetailModal, selectedUsageRecord, handleViewDetail } = useUsageRecor
bindUsageRecordLoader(loadUsageRecords)
watch(timeRange, () => {
refreshTimeRangeEnd()
page.value = 1
void loadUsageRecords('all')
})
+1 -1
View File
@@ -1 +1 @@
version: "2.3.0"
version: "2.3.2"