mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 00:05:49 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
40a2b70bf2 | ||
|
|
1e6ac77d3c | ||
|
|
3340934944 | ||
|
|
37a368de32 | ||
|
|
79cd8ac965 | ||
|
|
57c6b95596 | ||
|
|
a4f1843876 | ||
|
|
096e48621a | ||
|
|
36828b1a7e |
@@ -34,6 +34,8 @@ updates:
|
||||
day: monday
|
||||
time: '09:20'
|
||||
timezone: Asia/Shanghai
|
||||
cooldown:
|
||||
default-days: 1
|
||||
groups:
|
||||
minor-and-patch:
|
||||
update-types:
|
||||
|
||||
@@ -3,8 +3,10 @@ name: CI
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- '**'
|
||||
- main
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
@@ -3,8 +3,10 @@ name: Security Scan
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- '**'
|
||||
- main
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
schedule:
|
||||
- cron: '0 3 * * 1'
|
||||
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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", {
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
|
||||
@@ -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,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;
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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> {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,6 +1,7 @@
|
||||
//! 账号、调度、配额与刷新领域。
|
||||
|
||||
pub mod account;
|
||||
pub mod account_failure;
|
||||
pub mod cookies;
|
||||
pub mod import;
|
||||
pub mod manage;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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],
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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.
|
||||
|
||||
+5
-6
@@ -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,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");
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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");
|
||||
|
||||
@@ -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,3 +1,4 @@
|
||||
pub(crate) mod account_failure;
|
||||
pub(crate) mod cookies;
|
||||
pub(crate) mod pool;
|
||||
pub(crate) mod quota;
|
||||
|
||||
@@ -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));
|
||||
|
||||
@@ -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
@@ -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(),
|
||||
}
|
||||
}
|
||||
@@ -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 @@
|
||||
version: "2.3.0"
|
||||
version: "2.3.2"
|
||||
|
||||
Reference in New Issue
Block a user