feat(proxy): prefer passive upstream state refresh

Remove request-path quota probes and consume account, request, and rate-limit facts from normal upstream responses, including 429 active-limit metadata.

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

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-17 02:12:16 +08:00
co-authored by Codex
parent f0ccc981fb
commit 4091a0c2d8
34 changed files with 673 additions and 438 deletions
@@ -60,3 +60,29 @@ pub async fn account_models(
Err(error) => Err(account_error(&error)),
}
}
/// `POST /api/admin/accounts/models`
pub async fn refresh_account_models(
State(state): State<AppState>,
_auth: AdminAuth,
Json(payload): Json<AccountActionRequest>,
) -> Result<impl IntoResponse, AdminError> {
let models = state
.services
.admin_accounts
.refresh_account_models(&payload.id)
.await
.map_err(|error| account_error(&error))?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(AccountModelsData {
models: models
.into_iter()
.map(|model| AccountModelData {
id: model.id,
label: model.label,
})
.collect(),
}),
))
}
@@ -513,7 +513,9 @@ pub(super) fn account_error(error: &AccountManageError) -> AdminError {
| AccountManageError::OAuthCallbackInvalid
| AccountManageError::OAuthStateMismatch
| AccountManageError::NoValidCookies => AdminError::bad_request(error.to_string()),
AccountManageError::OAuthCodeExchange(_) => AdminError::bad_gateway(error.to_string()),
AccountManageError::OAuthCodeExchange(_) | AccountManageError::RefreshModels(_) => {
AdminError::bad_gateway(error.to_string())
}
AccountManageError::NotFound => account_not_found(),
AccountManageError::Inactive(_) => AdminError::conflict(error.to_string()),
_ => AdminError::internal(error.to_string()),
+6 -2
View File
@@ -14,7 +14,8 @@ use super::{
accounts_routes::{
account_models, account_quota, accounts, batch_delete_accounts, create_account,
export_accounts, health_check_accounts, import_accounts, oauth_authorize_account,
oauth_exchange_account, refresh_account, test_account_connection, update_account,
oauth_exchange_account, refresh_account, refresh_account_models, test_account_connection,
update_account,
},
auth_routes::{login, logout, session_status},
dashboard_routes::{dashboard_summary, dashboard_trend},
@@ -72,7 +73,10 @@ pub fn router() -> Router<AppState> {
post(oauth_exchange_account),
)
.route("/api/admin/accounts/test", get(test_account_connection))
.route("/api/admin/accounts/models", get(account_models))
.route(
"/api/admin/accounts/models",
get(account_models).post(refresh_account_models),
)
.route("/api/admin/accounts/delete", post(batch_delete_accounts))
.route("/api/admin/accounts/update", post(update_account))
.route("/api/admin/accounts/refresh", post(refresh_account))
+45 -85
View File
@@ -13,23 +13,18 @@ use crate::models::service::{
ModelRefreshPlanAccount, ModelRefreshResult, ModelService, ModelServiceError,
};
use super::{
coordinator::SchedulerHandle,
periodic::{PeriodicTaskConfig, PeriodicTaskRunner, PeriodicTaskStartup, spawn_periodic_task},
};
use super::coordinator::SchedulerHandle;
/// 模型刷新任务接线器。
pub struct ModelRefreshTask {
model_service: Arc<ModelService>,
account_pool: Arc<AccountPoolService>,
refresh_interval_secs: u64,
initial_delay_ms: u64,
retry_delay_ms: u64,
max_retries: u32,
account_pseudonymizer: Arc<AccountPseudonymizer>,
}
const DEFAULT_REFRESH_INTERVAL_SECS: u64 = 3600;
const INITIAL_DELAY_MS: u64 = 1000;
const RETRY_DELAY_MS: u64 = 10_000;
const MAX_RETRIES: u32 = 12;
@@ -44,7 +39,6 @@ impl ModelRefreshTask {
Self {
model_service,
account_pool,
refresh_interval_secs: DEFAULT_REFRESH_INTERVAL_SECS,
initial_delay_ms: INITIAL_DELAY_MS,
retry_delay_ms: RETRY_DELAY_MS,
max_retries: MAX_RETRIES,
@@ -52,15 +46,13 @@ impl ModelRefreshTask {
}
}
/// 启动后台刷新任务。
/// 启动一次性首次刷新任务。
pub fn start(self) -> SchedulerHandle {
let config = PeriodicTaskConfig::new(
self.refresh_interval_secs,
"模型刷新任务已启动",
"模型刷新任务正在关闭",
)
.wait_first_interval();
spawn_periodic_task(self, config)
let (shutdown_tx, mut shutdown_rx) = mpsc::channel(1);
let handle = tokio::spawn(async move {
self.run_initial_refresh(&mut shutdown_rx).await;
});
SchedulerHandle::new(shutdown_tx, handle)
}
/// 在 Responses 响应声明模型目录 ETag 变化时异步刷新。
@@ -115,84 +107,52 @@ impl ModelRefreshTask {
.await;
Ok(result)
}
}
impl PeriodicTaskRunner for ModelRefreshTask {
fn before_loop<'a>(
&'a mut self,
shutdown_rx: &'a mut mpsc::Receiver<()>,
) -> super::periodic::TaskFuture<'a, PeriodicTaskStartup> {
Box::pin(async move {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在首次拉取前关闭");
return PeriodicTaskStartup::Stop;
}
() = sleep(Duration::from_millis(self.initial_delay_ms)) => {}
async fn run_initial_refresh(&self, shutdown_rx: &mut mpsc::Receiver<()>) {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在首次拉取前关闭");
return;
}
() = sleep(Duration::from_millis(self.initial_delay_ms)) => {}
}
let mut attempt = 0;
let mut has_fetched_once = false;
let mut deferred_until_accounts_available = false;
while attempt < self.max_retries {
match self.refresh_once().await {
Ok(result) => {
info!(
refreshed_plans = result.refreshed_plans,
model_count = result.model_count,
"首次模型刷新成功"
);
has_fetched_once = true;
break;
}
Err(error) => {
attempt += 1;
if matches!(error, ModelServiceError::NoAccounts) {
info!("首次模型刷新跳过:没有可用账户,将进入周期重试");
deferred_until_accounts_available = true;
break;
}
warn!(
attempt,
max_retries = self.max_retries,
error = ?error,
"首次模型刷新失败"
);
if attempt < self.max_retries {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在重试等待期间关闭");
return PeriodicTaskStartup::Stop;
}
() = sleep(Duration::from_millis(self.retry_delay_ms)) => {}
}
}
}
}
}
if !has_fetched_once && !deferred_until_accounts_available {
warn!("模型刷新任务首次尝试全部失败,切换到周期刷新");
}
PeriodicTaskStartup::Continue
})
}
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
let mut attempt = 0;
while attempt < self.max_retries {
match self.refresh_once().await {
Ok(_) => {}
Ok(result) => {
info!(
refreshed_plans = result.refreshed_plans,
model_count = result.model_count,
"首次模型刷新成功"
);
return;
}
Err(ModelServiceError::NoAccounts) => {
info!("模型刷新跳过:没有可用账户");
info!("首次模型刷新跳过:没有可用账户,后续由 ETag 或手动操作触发");
return;
}
Err(error) => {
warn!(error = ?error, "刷新模型列表失败");
attempt += 1;
warn!(
attempt,
max_retries = self.max_retries,
error = ?error,
"首次模型刷新失败"
);
if attempt < self.max_retries {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在重试等待期间关闭");
return;
}
() = sleep(Duration::from_millis(self.retry_delay_ms)) => {}
}
}
}
}
})
}
warn!("模型刷新任务首次尝试全部失败,等待 ETag 或手动操作再次触发");
}
}
-32
View File
@@ -10,20 +10,11 @@ use super::coordinator::SchedulerHandle;
pub(crate) type TaskFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
/// 周期任务进入循环前的决策。
pub(crate) enum PeriodicTaskStartup {
/// 继续进入周期循环。
Continue,
/// 任务已自行结束,不再进入周期循环。
Stop,
}
/// 周期任务调度配置。
pub(crate) struct PeriodicTaskConfig {
interval_secs: u64,
started_message: &'static str,
stopped_message: &'static str,
run_immediately: bool,
}
impl PeriodicTaskConfig {
@@ -36,25 +27,12 @@ impl PeriodicTaskConfig {
interval_secs,
started_message,
stopped_message,
run_immediately: true,
}
}
pub(crate) fn wait_first_interval(mut self) -> Self {
self.run_immediately = false;
self
}
}
/// 周期任务行为。
pub(crate) trait PeriodicTaskRunner: Send + 'static {
fn before_loop<'a>(
&'a mut self,
_shutdown_rx: &'a mut mpsc::Receiver<()>,
) -> TaskFuture<'a, PeriodicTaskStartup> {
Box::pin(std::future::ready(PeriodicTaskStartup::Continue))
}
fn tick(&mut self) -> TaskFuture<'_, ()>;
fn shutdown(&mut self) -> TaskFuture<'_, ()> {
@@ -75,17 +53,7 @@ where
"{}", config.started_message
);
if matches!(
task.before_loop(&mut shutdown_rx).await,
PeriodicTaskStartup::Stop
) {
return;
}
let mut ticker = interval(Duration::from_secs(config.interval_secs));
if !config.run_immediately {
ticker.tick().await;
}
loop {
tokio::select! {
@@ -191,12 +191,6 @@ impl AccountIdentityService {
AccountScopedRequest { request, identity }
}
pub(in crate::dispatch) fn scope_auxiliary(&self, account_id: &str) -> AccountScopedIdentity {
AccountScopedIdentity {
installation_id: self.pseudonymizer.installation_id(account_id),
}
}
fn scope_turn_metadata(
raw: &str,
identity: &AccountScopedIdentity,
@@ -119,14 +119,20 @@ impl AccountFailureController {
if code.contains("invalid_request")
|| code.contains("not_found")
|| code.contains("context_window")
|| code == "context_length_exceeded"
|| code.contains("invalid_prompt")
|| code.contains("bad_request")
{
return Some(ClientFailure::new(failure.clone(), 400, false));
}
if code.contains("server_overloaded") {
if matches!(code.as_str(), "server_is_overloaded" | "slow_down")
|| code.contains("server_overloaded")
{
return Some(ClientFailure::new(failure.clone(), 503, true));
}
if code == "usage_not_included" {
return Some(ClientFailure::new(failure.clone(), 403, false));
}
classify_response_failure(failure).map(|classified| {
let status = match classified.kind {
AccountFailureKind::ModelUnsupported => 400,
+18 -46
View File
@@ -19,7 +19,10 @@ use crate::{
AccountIdentityScope, AccountIdentityService, AccountScopedRequest,
SessionAffinityService,
},
errors::{ClientFailure, ResponseDispatchError, upstream_error_set_cookie_headers},
errors::{
ClientFailure, ResponseDispatchError, upstream_error_rate_limit_headers,
upstream_error_set_cookie_headers,
},
lifecycle::trace::{ResponseDispatchAttempt, ResponseDispatchTrace},
},
fleet::{account::Account, pool::AccountLease},
@@ -271,22 +274,11 @@ pub(in crate::dispatch) struct ControllerEnter {
pub excluded_account_ids: BTreeSet<String>,
}
pub(in crate::dispatch) struct AttemptAccountPreparationContext<'a> {
pub codex: &'a CodexBackendClient,
pub account_identity: &'a AccountIdentityService,
pub request_id: &'a str,
}
pub(in crate::dispatch) struct PreparedAttemptAccount {
pub lease: AccountLease,
pub cookie_header: Option<String>,
}
pub(in crate::dispatch) enum AttemptAccountPreparation {
Ready(Box<PreparedAttemptAccount>),
Rejected,
}
pub(in crate::dispatch) struct AttemptRoutePreparationContext<'a> {
pub account_identity: &'a AccountIdentityService,
pub account_id: &'a str,
@@ -300,38 +292,16 @@ pub(in crate::dispatch) enum AttemptRoutePreparation {
impl ControllerSet {
pub(in crate::dispatch) async fn prepare_attempt_account(
&self,
context: AttemptAccountPreparationContext<'_>,
acquired: AccountLease,
) -> AttemptAccountPreparation {
let usage_cookie = self
.cloudflare
.cookie_header_for_request(&acquired.account.id, "/codex/usage")
.await;
let acquired = match QuotaController::enter(
quota::QuotaEnterContext {
account_pool: &self.account_pool,
codex: context.codex,
cookie_header: usage_cookie.as_deref(),
account_identity: context.account_identity,
request_id: context.request_id,
},
acquired,
)
.await
{
quota::QuotaEnterOutcome::Ready(acquired) => *acquired,
quota::QuotaEnterOutcome::LimitReached => {
return AttemptAccountPreparation::Rejected;
}
};
) -> PreparedAttemptAccount {
let cookie_header = self
.cloudflare
.cookie_header_for_request(&acquired.account.id, "/codex/responses")
.await;
AttemptAccountPreparation::Ready(Box::new(PreparedAttemptAccount {
PreparedAttemptAccount {
lease: acquired,
cookie_header,
}))
}
}
pub(in crate::dispatch) fn prepare_attempt_route(
@@ -396,19 +366,24 @@ impl ControllerSet {
&self,
context: AttemptUpstreamErrorContext<'_>,
) {
let _ = best_effort_controller_io(
let cookie = best_effort_controller_io(
"cookie.attempt_error",
self.cloudflare.capture_set_cookie_headers(
&context.account.id,
upstream_error_set_cookie_headers(context.error),
),
)
.await;
let _ = best_effort_controller_io(
);
let rate_limit_headers = upstream_error_rate_limit_headers(context.error);
let quota = best_effort_controller_io(
"quota.attempt_error_headers",
self.account_pool
.sync_passive_rate_limit_headers(context.account, &rate_limit_headers),
);
let telemetry = best_effort_controller_io(
"telemetry.attempt_error",
telemetry::TelemetryController::observe_upstream_error(&self.recorder, &context),
)
.await;
);
let _ = tokio::join!(cookie, quota, telemetry);
}
pub(in crate::dispatch) async fn observe_dispatch_error(
@@ -944,9 +919,6 @@ impl ControllerSet {
AttemptObservationKind::NoCandidate { .. } => {
AttemptDecision::Return(AttemptReturnKind::Observed)
}
AttemptObservationKind::CandidatePreparationRejected => {
unreachable!("candidate preparation controller owns it")
}
AttemptObservationKind::RoutePreparationRejected { message } => {
AttemptDecision::Return(AttemptReturnKind::RouteUnavailable {
message: message.clone(),
+5 -131
View File
@@ -1,139 +1,29 @@
//! 请求前配额验证规则 owner。
use serde_json::Value;
//! 配额失败规则 owner。
use crate::{
dispatch::{
affinity::AccountIdentityService,
controllers::ControllerFailureFact,
errors::ClientFailure,
failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind},
lifecycle::contract::{
AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind,
},
lifecycle::contract::{AttemptDecision, AttemptObservation, AttemptReturnKind},
},
fleet::{
account_failure::{
AccountFailureKind, AccountStateEffect, apply_account_state_effect_immediately,
classify_response_failure, classify_upstream_failure,
},
pool::{AccountLease, AccountPoolService},
quota::{
quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds,
quota_snapshot_reset_at,
},
pool::AccountPoolService,
},
upstream::openai::{
failure::UpstreamFailureFacts,
protocol::responses::ResponsesSseFailure,
transport::{CodexBackendClient, CodexRequestContext},
failure::UpstreamFailureFacts, protocol::responses::ResponsesSseFailure,
transport::CodexBackendClient,
},
};
const LIMIT_REACHED_MESSAGE: &str = "Upstream usage quota still reports limit_reached";
pub(super) enum QuotaEnterOutcome {
Ready(Box<AccountLease>),
LimitReached,
}
pub(super) struct QuotaEnterContext<'a> {
pub account_pool: &'a AccountPoolService,
pub codex: &'a CodexBackendClient,
pub account_identity: &'a AccountIdentityService,
pub request_id: &'a str,
pub cookie_header: Option<&'a str>,
}
pub(super) struct QuotaController;
impl QuotaController {
pub(super) async fn enter(
context: QuotaEnterContext<'_>,
acquired: AccountLease,
) -> QuotaEnterOutcome {
if !acquired.account.quota_verify_required {
return QuotaEnterOutcome::Ready(Box::new(acquired));
}
let account_id = acquired.account.id.clone();
let identity = context.account_identity.scope_auxiliary(&account_id);
let usage = context
.codex
.fetch_usage(CodexRequestContext {
access_token: &acquired.account.access_token,
account_id: acquired.account.account_id.as_deref(),
request_id: context.request_id,
turn_state: None,
turn_metadata: None,
beta_features: None,
include_timing_metrics: None,
version: None,
codex_window_id: None,
parent_thread_id: None,
cookie_header: context.cookie_header,
installation_id: Some(&identity.installation_id),
session_id: None,
thread_id: None,
client_request_id: None,
turn_id: None,
})
.await;
let raw = match usage {
Ok(raw) => raw,
Err(error) => {
tracing::warn!(
request_id = %context.request_id,
account_id = %account_id,
quota_verify_required = true,
quota_verify_result = "upstream_error",
retry_with_another_account = false,
error = %error,
"Failed to verify stale quota state before upstream request"
);
return QuotaEnterOutcome::Ready(Box::new(acquired));
}
};
let quota = quota_from_usage(&raw);
context
.account_pool
.apply_quota_snapshot(&account_id, &quota)
.await;
if quota_snapshot_limit_reached(&quota) {
acquired.release_without_usage().await;
tracing::info!(
request_id = %context.request_id,
account_id = %account_id,
quota_verify_required = true,
quota_verify_result = "limit_reached",
retry_with_another_account = true,
"Quota verification reported exhausted account before upstream request"
);
return QuotaEnterOutcome::LimitReached;
}
QuotaEnterOutcome::Ready(Box::new(acquired_with_verified_quota(acquired, &quota)))
}
pub(super) fn classify(observation: &AttemptObservation) -> Option<ClassifiedQuotaFailure> {
if matches!(
observation.kind,
AttemptObservationKind::CandidatePreparationRejected
) {
return observation
.account
.as_ref()
.map(|account| ClassifiedQuotaFailure {
exhaustion: AccountExhaustionRecord::new(
account.id.clone(),
ExhaustedAccountKind::RateLimited,
LIMIT_REACHED_MESSAGE,
),
effect: None,
});
}
classify(observation)
}
@@ -229,19 +119,3 @@ fn classified_quota_failure(
effect: classified.effect,
})
}
fn acquired_with_verified_quota(mut acquired: AccountLease, quota: &Value) -> AccountLease {
let limit_reached = quota_snapshot_limit_reached(quota);
acquired.account.quota_verify_required = false;
acquired.account.quota_limit_reached = limit_reached;
acquired.account.quota_cooldown_until = limit_reached
.then_some(quota_snapshot_reset_at(quota))
.flatten();
if let Some(reset_at) = quota_snapshot_reset_at(quota) {
acquired.account.window_reset_at = Some(reset_at);
if let Some(limit_window_seconds) = quota_snapshot_limit_window_seconds(quota) {
acquired.account.limit_window_seconds = Some(limit_window_seconds);
}
}
acquired
}
+22 -2
View File
@@ -9,8 +9,13 @@ use crate::{
sse::sse_failure_error_body,
},
upstream::openai::{
protocol::{responses::ResponsesSseFailure, sse::SseError},
transport::{CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics},
protocol::{
events::is_rate_limit_header_name, responses::ResponsesSseFailure, sse::SseError,
},
transport::{
CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics,
websocket::CodexWebSocketExchangeError,
},
},
};
@@ -238,6 +243,21 @@ pub(crate) fn upstream_error_set_cookie_headers(error: &CodexClientError) -> &[S
}
}
pub(crate) fn upstream_error_rate_limit_headers(error: &CodexClientError) -> Vec<(String, String)> {
match error {
CodexClientError::Upstream {
rate_limit_headers, ..
} => rate_limit_headers.clone(),
CodexClientError::WebSocket(CodexWebSocketExchangeError::Upstream(upstream)) => upstream
.headers
.iter()
.filter(|(name, _)| is_rate_limit_header_name(name))
.cloned()
.collect(),
_ => Vec::new(),
}
}
pub(crate) fn upstream_error_http_status(error: &CodexClientError) -> u16 {
match error {
CodexClientError::Upstream { status, .. } => status.as_u16(),
+15 -40
View File
@@ -8,9 +8,8 @@ use crate::{
dispatch::{
affinity::AccountIdentityService,
controllers::{
AttemptAccountPreparation, AttemptAccountPreparationContext, AttemptRoutePreparation,
AttemptRoutePreparationContext, AttemptUpstreamErrorContext, ControllerRequestScope,
ControllerSet,
AttemptRoutePreparation, AttemptRoutePreparationContext, AttemptUpstreamErrorContext,
ControllerRequestScope, ControllerSet,
},
errors::{ClientFailure, ResponseDispatchError},
failure::exhaustion::AccountExhaustionTracker,
@@ -191,25 +190,11 @@ impl<'a> AttemptRunner<'a> {
.controller_scope
.as_mut()
.ok_or(AttemptContractError::Terminal)?;
let prepared_account = match self
let prepared_account = self
.dependencies
.controllers
.prepare_attempt_account(
AttemptAccountPreparationContext {
codex: self.dependencies.codex,
account_identity: self.dependencies.account_identity,
request_id: self.dependencies.request_id,
},
acquired,
)
.await
{
AttemptAccountPreparation::Ready(prepared) => *prepared,
AttemptAccountPreparation::Rejected => {
self.pending = Some(PendingAttempt::CandidatePreparationRejected { account });
return self.step();
}
};
.prepare_attempt_account(acquired)
.await;
let acquired = prepared_account.lease;
let cookie_header = prepared_account.cookie_header;
@@ -769,9 +754,6 @@ fn observation_kind(
kind: *kind,
}
}
PendingAttempt::CandidatePreparationRejected { .. } => {
AttemptObservationKind::CandidatePreparationRejected
}
PendingAttempt::RoutePreparationRejected { message, .. } => {
AttemptObservationKind::RoutePreparationRejected {
message: message.clone(),
@@ -840,17 +822,15 @@ async fn rejection_error(
})
}
AttemptReturnKind::Observed => match pending {
PendingAttempt::NoCandidate | PendingAttempt::CandidatePreparationRejected { .. } => {
Ok(on_exhaustion.map_or_else(
|| {
exhaustion
.last_exhausted()
.map(ResponseDispatchError::from_exhausted_account)
.unwrap_or(ResponseDispatchError::NoActiveAccount)
},
ResponseDispatchError::Failed,
))
}
PendingAttempt::NoCandidate => Ok(on_exhaustion.map_or_else(
|| {
exhaustion
.last_exhausted()
.map(ResponseDispatchError::from_exhausted_account)
.unwrap_or(ResponseDispatchError::NoActiveAccount)
},
ResponseDispatchError::Failed,
)),
PendingAttempt::RoutePreparationRejected { message, .. } => {
Ok(ResponseDispatchError::HistoryUnavailable {
upstream_error: message,
@@ -902,8 +882,7 @@ fn pending_account(pending: &PendingAttempt) -> Option<&str> {
fn pending_account_value(pending: &PendingAttempt) -> Option<&Account> {
match pending {
PendingAttempt::NoCandidate | PendingAttempt::PinnedCandidateUnavailable { .. } => None,
PendingAttempt::CandidatePreparationRejected { account }
| PendingAttempt::RoutePreparationRejected { account, .. }
PendingAttempt::RoutePreparationRejected { account, .. }
| PendingAttempt::UpstreamFailure { account, .. }
| PendingAttempt::ProtocolFailure { account, .. } => Some(account),
PendingAttempt::CompleteResponse(complete) => Some(&complete.account),
@@ -919,7 +898,6 @@ fn pending_attempt(pending: &PendingAttempt) -> Option<&super::trace::ResponseDi
PendingAttempt::StreamResponse(stream) => Some(&stream.attempt),
PendingAttempt::NoCandidate
| PendingAttempt::PinnedCandidateUnavailable { .. }
| PendingAttempt::CandidatePreparationRejected { .. }
| PendingAttempt::RoutePreparationRejected { .. } => None,
}
}
@@ -936,7 +914,6 @@ fn pending_attempt_request(pending: &PendingAttempt) -> Option<&CodexResponsesRe
PendingAttempt::StreamResponse(stream) => Some(&stream.attempt_request),
PendingAttempt::NoCandidate
| PendingAttempt::PinnedCandidateUnavailable { .. }
| PendingAttempt::CandidatePreparationRejected { .. }
| PendingAttempt::RoutePreparationRejected { .. } => None,
}
}
@@ -951,7 +928,6 @@ fn pending_transport(
PendingAttempt::StreamResponse(stream) => Some(stream.response.transport),
PendingAttempt::NoCandidate
| PendingAttempt::PinnedCandidateUnavailable { .. }
| PendingAttempt::CandidatePreparationRejected { .. }
| PendingAttempt::RoutePreparationRejected { .. } => None,
}
}
@@ -973,7 +949,6 @@ fn invalid(decision: &'static str, pending: &PendingAttempt) -> AttemptContractE
observation: match pending {
PendingAttempt::NoCandidate => "no_candidate",
PendingAttempt::PinnedCandidateUnavailable { .. } => "pinned_candidate_unavailable",
PendingAttempt::CandidatePreparationRejected { .. } => "candidate_preparation_rejected",
PendingAttempt::RoutePreparationRejected { .. } => "route_preparation_rejected",
PendingAttempt::UpstreamFailure { .. } => "upstream_failure",
PendingAttempt::ProtocolFailure { .. } => "protocol_failure",
@@ -93,7 +93,6 @@ pub(in crate::dispatch) enum AttemptObservationKind {
ledger: CandidateLedgerFacts,
last_exhausted: Option<ExhaustedAccount>,
},
CandidatePreparationRejected,
RoutePreparationRejected {
message: String,
},
@@ -257,9 +256,6 @@ pub(super) enum PendingAttempt {
account_id: String,
kind: PinnedCandidateAcquireFailureKind,
},
CandidatePreparationRejected {
account: Account,
},
RoutePreparationRejected {
account: Account,
message: String,
+12 -12
View File
@@ -53,6 +53,8 @@ pub fn classify_upstream_failure(facts: &UpstreamFailureFacts) -> Option<Classif
status_code: facts.status_code,
code: facts.code.as_deref(),
error_type: facts.error_type.as_deref(),
identity_authorization_error: facts.identity_authorization_error.as_deref(),
identity_error_code: facts.identity_error_code.as_deref(),
message: &facts.message,
body: &facts.body,
retry_after_seconds: facts.retry_after_seconds,
@@ -66,6 +68,8 @@ pub fn classify_response_failure(
status_code: failure.explicit_status_code,
code: failure.upstream_code.as_deref(),
error_type: failure.upstream_type.as_deref(),
identity_authorization_error: None,
identity_error_code: None,
message: &failure.message,
body: &failure.message,
retry_after_seconds: failure.retry_after_seconds,
@@ -128,6 +132,8 @@ struct FailureFields<'a> {
status_code: Option<u16>,
code: Option<&'a str>,
error_type: Option<&'a str>,
identity_authorization_error: Option<&'a str>,
identity_error_code: Option<&'a str>,
message: &'a str,
body: &'a str,
retry_after_seconds: Option<u64>,
@@ -136,6 +142,8 @@ struct FailureFields<'a> {
fn classify_failure(fields: FailureFields<'_>) -> Option<ClassifiedAccountFailure> {
let code = fields.code.unwrap_or_default();
let error_type = fields.error_type.unwrap_or_default();
let identity_error_code = fields.identity_error_code.unwrap_or_default();
let identity_authorization_error = fields.identity_authorization_error.unwrap_or_default();
if [code, fields.message, fields.body]
.into_iter()
@@ -143,14 +151,13 @@ fn classify_failure(fields: FailureFields<'_>) -> Option<ClassifiedAccountFailur
{
return Some(classified(AccountFailureKind::ModelUnsupported, None));
}
if let Some(status) = explicit_account_status(code, error_type, fields.message)
.or_else(|| explicit_account_status(code, error_type, fields.body))
if let Some(status) =
explicit_account_status(identity_error_code, "", identity_authorization_error)
.or_else(|| explicit_account_status(code, error_type, fields.message))
.or_else(|| explicit_account_status(code, error_type, fields.body))
{
return Some(status_failure(status));
}
if fields.status_code == Some(403) && !is_html(fields.body) {
return Some(status_failure(AccountStatus::Banned));
}
if fields.status_code == Some(401) {
return Some(status_failure(AccountStatus::Expired));
}
@@ -238,8 +245,6 @@ fn explicit_account_status(code: &str, error_type: &str, message: &str) -> Optio
|| message.contains("organization has been disabled")
|| message.contains("workspace has been deactivated")
|| message.contains("deactivated_workspace")
|| code == "forbidden"
|| error_type == "permission_error"
{
return Some(AccountStatus::Banned);
}
@@ -310,8 +315,3 @@ fn is_model_unsupported(value: &str) -> bool {
|| value.contains("not_supported")
|| value.contains("not_available")))
}
fn is_html(value: &str) -> bool {
let value = value.trim_start().to_ascii_lowercase();
value.starts_with("<!doctype") || value.starts_with("<html") || value.contains("<html")
}
+43 -17
View File
@@ -12,7 +12,10 @@ use crate::{
account_failure::{classify_client_failure, classify_response_failure},
store::StoredAccount,
},
models::{service::ModelRefreshPlanAccount, types::CodexModelInfo},
models::{
service::{ModelRefreshPlanAccount, ModelRefreshResult},
types::CodexModelInfo,
},
upstream::openai::{
protocol::{
responses::{CodexResponsesRequest, ResponsesSseFailure},
@@ -76,11 +79,7 @@ impl AccountManageService {
.map_err(|_| AccountManageError::Inspect)?
.ok_or(AccountManageError::NotFound)?;
let plan_type = account_plan_type(&account);
let mut models = self.models.catalog().await.models_for_plan(&plan_type);
if models.is_empty() {
self.refresh_account_plan_models(&account, &plan_type).await;
models = self.models.catalog().await.models_for_plan(&plan_type);
}
let models = self.models.catalog().await.models_for_plan(&plan_type);
let models = models.iter().map(account_model_option).collect::<Vec<_>>();
if models.is_empty() {
return Err(AccountManageError::NoModels);
@@ -88,7 +87,32 @@ impl AccountManageService {
Ok(models)
}
async fn refresh_account_plan_models(&self, account: &StoredAccount, plan_type: &str) {
pub async fn refresh_account_models(
&self,
account_id: &str,
) -> Result<Vec<AccountModelOption>, AccountManageError> {
let account = self
.store
.get(account_id)
.await
.map_err(|_| AccountManageError::Inspect)?
.ok_or(AccountManageError::NotFound)?;
let plan_type = account_plan_type(&account);
self.refresh_account_plan_models(&account, &plan_type)
.await?;
let models = self.models.catalog().await.models_for_plan(&plan_type);
let models = models.iter().map(account_model_option).collect::<Vec<_>>();
if models.is_empty() {
return Err(AccountManageError::NoModels);
}
Ok(models)
}
async fn refresh_account_plan_models(
&self,
account: &StoredAccount,
plan_type: &str,
) -> Result<ModelRefreshResult, AccountManageError> {
let request_id = uuid::Uuid::new_v4().to_string();
let plan_account = ModelRefreshPlanAccount {
plan_type: plan_type.to_string(),
@@ -96,18 +120,20 @@ impl AccountManageService {
account_id: account.account_id.clone(),
installation_id: self.account_pseudonymizer.installation_id(&account.id),
};
if let Err(error) = self
let result = self
.models
.refresh_backend_models(&[plan_account], &request_id)
.refresh_selected_plan_models(&[plan_account], &request_id)
.await
{
tracing::warn!(
account_id = %account.id,
plan_type,
error = %error,
"Failed to refresh account plan models"
);
}
.map_err(|error| AccountManageError::RefreshModels(error.to_string()))?;
let routing = self
.models
.model_plan_routing()
.await
.map_err(|error| AccountManageError::RefreshModels(error.to_string()))?;
self.account_pool
.apply_model_plan_routing(routing.allowlist, routing.fetched_plan_types)
.await;
Ok(result)
}
pub async fn test_connection_stream(
+2
View File
@@ -86,6 +86,8 @@ pub enum AccountManageError {
FetchQuota(String),
#[error("no available models returned")]
NoModels,
#[error("failed to refresh account models: {0}")]
RefreshModels(String),
#[error("invalid status: {0}")]
InvalidStatus(String),
#[error("label must be 64 characters or fewer")]
+25 -1
View File
@@ -189,6 +189,26 @@ impl ModelService {
&self,
plan_accounts: &[ModelRefreshPlanAccount],
request_id: &str,
) -> Result<ModelRefreshResult, ModelServiceError> {
self.refresh_backend_models_inner(plan_accounts, request_id, false)
.await
}
/// 刷新指定计划,并保留未参与本次刷新的计划快照。
pub async fn refresh_selected_plan_models(
&self,
plan_accounts: &[ModelRefreshPlanAccount],
request_id: &str,
) -> Result<ModelRefreshResult, ModelServiceError> {
self.refresh_backend_models_inner(plan_accounts, request_id, true)
.await
}
async fn refresh_backend_models_inner(
&self,
plan_accounts: &[ModelRefreshPlanAccount],
request_id: &str,
preserve_other_plans: bool,
) -> Result<ModelRefreshResult, ModelServiceError> {
let _refresh_guard = self.refresh_lock.lock().await;
let store = self
@@ -211,7 +231,11 @@ impl ModelService {
.into_iter()
.map(|snapshot| (snapshot.plan_type.clone(), snapshot))
.collect::<BTreeMap<_, _>>();
let mut next_snapshots = BTreeMap::new();
let mut next_snapshots = if preserve_other_plans {
previous_snapshots.clone()
} else {
BTreeMap::new()
};
let mut result = ModelRefreshResult {
refreshed_plans: 0,
+14
View File
@@ -30,6 +30,8 @@ pub struct UpstreamFailureFacts {
pub status_code: Option<u16>,
pub code: Option<String>,
pub error_type: Option<String>,
pub identity_authorization_error: Option<String>,
pub identity_error_code: Option<String>,
pub message: String,
pub body: String,
pub retry_after_seconds: Option<u64>,
@@ -48,6 +50,8 @@ pub fn upstream_failure_facts(error: &CodexClientError) -> UpstreamFailureFacts
status_code: error.status().map(|status| status.as_u16()),
code: None,
error_type: None,
identity_authorization_error: None,
identity_error_code: None,
message: error.to_string(),
body: error.to_string(),
retry_after_seconds: None,
@@ -57,6 +61,7 @@ pub fn upstream_failure_facts(error: &CodexClientError) -> UpstreamFailureFacts
body,
retry_after_seconds,
transport,
diagnostics,
..
} => {
let (code, error_type, message) = error_fields(body);
@@ -72,6 +77,8 @@ pub fn upstream_failure_facts(error: &CodexClientError) -> UpstreamFailureFacts
status_code: Some(status.as_u16()),
code,
error_type,
identity_authorization_error: diagnostics.identity_authorization_error.clone(),
identity_error_code: diagnostics.identity_error_code.clone(),
message: message.unwrap_or_else(|| body.clone()),
body: body.clone(),
retry_after_seconds: *retry_after_seconds,
@@ -87,6 +94,11 @@ pub fn upstream_failure_facts(error: &CodexClientError) -> UpstreamFailureFacts
status_code: Some(upstream.status_code),
code: upstream.code.clone().or(body_code),
error_type: upstream.error_type.clone().or(body_type),
identity_authorization_error: upstream
.diagnostics
.identity_authorization_error
.clone(),
identity_error_code: upstream.diagnostics.identity_error_code.clone(),
message: upstream
.message
.clone()
@@ -141,6 +153,8 @@ fn simple_facts(kind: UpstreamFailureKind, error: &CodexClientError) -> Upstream
status_code: None,
code: None,
error_type: None,
identity_authorization_error: None,
identity_error_code: None,
body: message.clone(),
message,
retry_after_seconds: None,
+25 -1
View File
@@ -238,6 +238,8 @@ pub struct CreditsSnapshot {
pub struct ParsedRateLimits {
/// 以标准化 limit ID 为键的全部计量项。
pub limits: BTreeMap<String, RateLimitDetails>,
/// 当前响应对应的计量项 ID。
pub active_limit: Option<String>,
pub credits: Option<CreditsSnapshot>,
pub plan_type: Option<String>,
pub promo_message: Option<String>,
@@ -251,10 +253,15 @@ pub fn parse_rate_limit_headers(headers: &[(String, String)]) -> Option<ParsedRa
normalized.insert(name.to_ascii_lowercase(), value.trim());
}
let active_limit = lookup_non_empty(&normalized, "x-codex-active-limit")
.map(|limit_id| normalize_limit_id(&limit_id));
let mut limit_ids = normalized
.keys()
.filter_map(|name| rate_limit_id_from_header_name(name))
.collect::<std::collections::BTreeSet<_>>();
if let Some(active_limit) = &active_limit {
limit_ids.insert(active_limit.clone());
}
limit_ids.insert("codex".to_string());
let limits = limit_ids
.into_iter()
@@ -270,6 +277,7 @@ pub fn parse_rate_limit_headers(headers: &[(String, String)]) -> Option<ParsedRa
let rate_limit_reached_type = lookup_non_empty(&normalized, "x-codex-rate-limit-reached-type");
if limits.is_empty()
&& active_limit.is_none()
&& credits.is_none()
&& plan_type.is_none()
&& promo_message.is_none()
@@ -280,6 +288,7 @@ pub fn parse_rate_limit_headers(headers: &[(String, String)]) -> Option<ParsedRa
Some(ParsedRateLimits {
limits,
active_limit,
credits,
plan_type,
promo_message,
@@ -298,6 +307,7 @@ pub(crate) fn is_rate_limit_header_name(name: &str) -> bool {
"x-codex-credits-has-credits"
| "x-codex-credits-unlimited"
| "x-codex-credits-balance"
| "x-codex-active-limit"
| "x-codex-plan-type"
| "x-codex-promo-message"
| "x-codex-rate-limit-reached-type"
@@ -343,10 +353,11 @@ pub fn parse_rate_limits_event(value: &Value) -> Option<ParsedRateLimits> {
let mut limits = BTreeMap::new();
if let Some(details) = details {
limits.insert(limit_id, details);
limits.insert(limit_id.clone(), details);
}
Some(ParsedRateLimits {
limits,
active_limit: Some(limit_id),
credits,
plan_type,
promo_message: value
@@ -430,9 +441,19 @@ pub fn rate_limit_quota(
let spend_control = existing_quota
.and_then(|quota| quota.get("spend_control").cloned())
.unwrap_or(Value::Null);
let active_limit = rate_limits.active_limit.as_ref().map_or_else(
|| {
existing_quota
.and_then(|quota| quota.get("active_limit"))
.cloned()
.unwrap_or(Value::Null)
},
|limit| Value::String(limit.clone()),
);
json!({
"plan_type": rate_limits.plan_type.as_deref().or(plan_type).unwrap_or("unknown"),
"active_limit": active_limit,
"snapshots": snapshots,
"monthly_limit": monthly_limit,
"credits": credits,
@@ -445,6 +466,9 @@ pub fn rate_limit_quota(
/// 将限流状态转换回内部传输头键值对。
pub fn rate_limits_to_header_pairs(rate_limits: &ParsedRateLimits) -> Vec<(String, String)> {
let mut headers = Vec::new();
if let Some(active_limit) = &rate_limits.active_limit {
headers.push(("x-codex-active-limit".to_string(), active_limit.clone()));
}
for details in rate_limits.limits.values() {
let prefix = format!("x-{}", details.limit_id.replace('_', "-"));
push_window_headers(&mut headers, &format!("{prefix}-primary"), details.primary);
@@ -139,9 +139,11 @@ pub enum CodexClientError {
/// 推导出的重试秒数。
retry_after_seconds: Option<u64>,
/// 上游诊断元数据。
diagnostics: CodexUpstreamDiagnostics,
diagnostics: Box<CodexUpstreamDiagnostics>,
/// 上游透传的 `set-cookie` 列表。
set_cookie_headers: Vec<String>,
/// 上游错误响应携带的限流头。
rate_limit_headers: Vec<(String, String)>,
/// 实际收到该上游响应的 transport。
transport: CodexBackendTransport,
},
@@ -616,13 +618,22 @@ fn websocket_exchange_error_to_client_error(
match error {
CodexWebSocketExchangeError::Upstream(upstream) => {
let upstream = *upstream;
let rate_limit_headers = upstream
.headers
.iter()
.filter(|(name, _)| {
crate::upstream::openai::protocol::events::is_rate_limit_header_name(name)
})
.cloned()
.collect();
CodexClientError::Upstream {
status: StatusCode::from_u16(upstream.status_code)
.unwrap_or(StatusCode::BAD_GATEWAY),
body: upstream.body,
retry_after_seconds: upstream.retry_after_seconds,
diagnostics: upstream.diagnostics,
diagnostics: Box::new(upstream.diagnostics),
set_cookie_headers: upstream.set_cookie_headers,
rate_limit_headers,
transport: CodexBackendTransport::WebSocket,
}
}
@@ -84,8 +84,9 @@ impl CodexBackendClient {
retry_after_seconds: retry_after_seconds
.or_else(|| retry_after_seconds_from_body(&body)),
body,
diagnostics,
diagnostics: Box::new(diagnostics),
set_cookie_headers,
rate_limit_headers,
transport: CodexBackendTransport::HttpSse,
});
}
@@ -166,8 +167,9 @@ impl CodexBackendClient {
retry_after_seconds: retry_after_seconds
.or_else(|| retry_after_seconds_from_body(&body)),
body,
diagnostics,
diagnostics: Box::new(diagnostics),
set_cookie_headers,
rate_limit_headers,
transport: CodexBackendTransport::HttpSse,
});
}
@@ -539,8 +541,9 @@ impl CodexBackendClient {
retry_after_seconds: retry_after_seconds
.or_else(|| retry_after_seconds_from_body(&body)),
body,
diagnostics,
diagnostics: Box::new(diagnostics),
set_cookie_headers,
rate_limit_headers: Vec::new(),
transport: CodexBackendTransport::HttpSse,
});
}
@@ -552,10 +555,11 @@ impl CodexBackendClient {
"model catalog response is not valid JSON: {error}; body: {}",
truncate_for_error(&body)
),
diagnostics: CodexUpstreamDiagnostics::with_status(
diagnostics: Box::new(CodexUpstreamDiagnostics::with_status(
StatusCode::BAD_GATEWAY.as_u16(),
),
)),
set_cookie_headers,
rate_limit_headers: Vec::new(),
transport: CodexBackendTransport::HttpSse,
})?;
let models = extract_model_entries(&parsed);
@@ -567,8 +571,11 @@ impl CodexBackendClient {
status: StatusCode::BAD_GATEWAY,
retry_after_seconds: None,
body: "model catalog response contains no models".to_string(),
diagnostics: CodexUpstreamDiagnostics::with_status(StatusCode::BAD_GATEWAY.as_u16()),
diagnostics: Box::new(CodexUpstreamDiagnostics::with_status(
StatusCode::BAD_GATEWAY.as_u16(),
)),
set_cookie_headers: Vec::new(),
rate_limit_headers: Vec::new(),
transport: CodexBackendTransport::HttpSse,
})
}
@@ -1,9 +1,12 @@
//! Upstream response diagnostics captured at the transport boundary.
use base64::{Engine as _, engine::general_purpose::STANDARD};
use reqwest::header::HeaderMap;
use serde_json::Value;
const UPSTREAM_REQUEST_ID_HEADERS: &[&str] = &[
"x-request-id",
"x-oai-request-id",
"x-openai-request-id",
"openai-request-id",
"request-id",
@@ -11,16 +14,21 @@ const UPSTREAM_REQUEST_ID_HEADERS: &[&str] = &[
const UPSTREAM_TRACE_HEADERS: &[&str] = &[
"x-request-id",
"x-oai-request-id",
"x-openai-request-id",
"openai-request-id",
"request-id",
"cf-ray",
];
const IDENTITY_AUTHORIZATION_ERROR_HEADER: &str = "x-openai-authorization-error";
const IDENTITY_ERROR_JSON_HEADER: &str = "x-error-json";
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct CodexUpstreamDiagnostics {
pub status_code: Option<u16>,
pub request_id: Option<String>,
pub identity_authorization_error: Option<String>,
pub identity_error_code: Option<String>,
pub trace_headers: Vec<(String, String)>,
}
@@ -29,6 +37,12 @@ impl CodexUpstreamDiagnostics {
Self {
status_code,
request_id: first_header(headers, UPSTREAM_REQUEST_ID_HEADERS),
identity_authorization_error: header_value(
headers,
IDENTITY_AUTHORIZATION_ERROR_HEADER,
),
identity_error_code: header_value(headers, IDENTITY_ERROR_JSON_HEADER)
.and_then(|encoded| decode_identity_error_code(&encoded)),
trace_headers: trace_headers(headers),
}
}
@@ -48,15 +62,24 @@ impl CodexUpstreamDiagnostics {
.iter()
.filter_map(|name| pair_value(headers, name).map(|value| ((*name).to_string(), value)))
.collect();
let identity_authorization_error = pair_value(headers, IDENTITY_AUTHORIZATION_ERROR_HEADER);
let identity_error_code = pair_value(headers, IDENTITY_ERROR_JSON_HEADER)
.and_then(|encoded| decode_identity_error_code(&encoded));
Self {
status_code,
request_id,
identity_authorization_error,
identity_error_code,
trace_headers,
}
}
pub fn is_empty(&self) -> bool {
self.status_code.is_none() && self.request_id.is_none() && self.trace_headers.is_empty()
self.status_code.is_none()
&& self.request_id.is_none()
&& self.identity_authorization_error.is_none()
&& self.identity_error_code.is_none()
&& self.trace_headers.is_empty()
}
pub fn cf_ray(&self) -> Option<&str> {
@@ -95,3 +118,14 @@ fn header_value(headers: &HeaderMap, name: &str) -> Option<String> {
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
fn decode_identity_error_code(encoded: &str) -> Option<String> {
let decoded = STANDARD.decode(encoded).ok()?;
serde_json::from_slice::<Value>(&decoded)
.ok()?
.pointer("/error/code")?
.as_str()
.map(str::trim)
.filter(|code| !code.is_empty())
.map(ToString::to_string)
}
@@ -34,8 +34,9 @@ impl CodexBackendClient {
retry_after_seconds: retry_after_seconds
.or_else(|| retry_after_seconds_from_body(&body)),
body,
diagnostics,
diagnostics: Box::new(diagnostics),
set_cookie_headers: Vec::new(),
rate_limit_headers: Vec::new(),
transport: super::client::CodexBackendTransport::HttpSse,
});
}
@@ -54,8 +55,11 @@ impl CodexBackendClient {
|| "usage endpoint is unavailable".to_string(),
|body| format!("invalid usage response: {}", truncate_for_error(&body)),
),
diagnostics: CodexUpstreamDiagnostics::with_status(StatusCode::BAD_GATEWAY.as_u16()),
diagnostics: Box::new(CodexUpstreamDiagnostics::with_status(
StatusCode::BAD_GATEWAY.as_u16(),
)),
set_cookie_headers: Vec::new(),
rate_limit_headers: Vec::new(),
transport: super::client::CodexBackendTransport::HttpSse,
})
}
@@ -85,27 +85,94 @@ async fn account_models_should_return_redis_snapshot_models() {
}
#[tokio::test]
async fn account_models_should_fetch_missing_plan_snapshot_with_account_token() {
async fn refresh_account_models_should_fetch_persist_and_preserve_other_plans() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/backend-api/codex/models"))
.and(header("authorization", "Bearer access-test"))
.and(header("chatgpt-account-id", "chatgpt-test"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"models": [{
"slug": "gpt-plus-refreshed",
"display_name": "GPT Plus Refreshed"
}]
})))
.expect(1)
.mount(&server)
.await;
let (app, state, pool, _dir) = admin_accounts_test_app_with_api_base_url(
"admin-account-models-refresh",
92,
format!("{}/backend-api", server.uri()),
)
.await;
seed_test_account(&pool).await;
let model_snapshots =
codex_proxy_rs::models::store::RedisModelSnapshotStore::new(state.redis.clone());
model_snapshots
.replace_plan_snapshots(&[
ModelPlanSnapshot::from_backend_entries(
"plus",
vec![BackendModelEntry {
id: Some("gpt-plus-old".to_string()),
name: Some("GPT Plus Old".to_string()),
..BackendModelEntry::default()
}],
),
ModelPlanSnapshot::from_backend_entries(
"team",
vec![BackendModelEntry {
id: Some("gpt-team-existing".to_string()),
name: Some("GPT Team Existing".to_string()),
..BackendModelEntry::default()
}],
),
])
.await
.unwrap();
state
.services
.models
.reload_from_store()
.await
.expect("model catalog should reload from seeded snapshots");
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/api/admin/accounts/models")
.header("cookie", "cpr_admin_session=session_1")
.header("content-type", "application/json")
.body(Body::from(json!({ "id": "acct_test" }).to_string()))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let body = response_json(response).await;
assert_eq!(body["data"]["models"][0]["id"], "gpt-plus-refreshed");
let stored = model_snapshots.list_plan_snapshots().await.unwrap();
assert!(stored.iter().any(|snapshot| {
snapshot.plan_type == "plus"
&& snapshot
.models
.iter()
.any(|model| model.id == "gpt-plus-refreshed")
}));
assert!(stored.iter().any(|snapshot| snapshot.plan_type == "team"));
}
#[tokio::test]
async fn account_models_should_not_fetch_missing_plan_snapshot() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/backend-api/codex/models"))
.and(header("authorization", "Bearer access-plan-b"))
.and(header("chatgpt-account-id", "chatgpt-plan-b"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"models": [
{
"slug": "gpt-plan-b-live",
"display_name": "GPT Plan B Live",
"description": "Live model",
"is_default": true,
"supported_reasoning_efforts": [{"reasoning_effort": "medium", "description": "medium"}],
"default_reasoning_effort": "medium",
"input_modalities": ["text"],
"output_modalities": ["text"],
"supports_personality": false
}
]
})))
.respond_with(ResponseTemplate::new(500))
.expect(0)
.mount(&server)
.await;
let (app, state, pool, _dir) = admin_accounts_test_app_with_api_base_url(
@@ -153,15 +220,12 @@ async fn account_models_should_fetch_missing_plan_snapshot_with_account_token()
.unwrap();
let requests = server.received_requests().await.unwrap();
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(requests.len(), 1, "missing plan should be fetched upstream");
let body = response_json(response).await;
assert_eq!(body["data"]["models"][0]["id"], "gpt-plan-b-live");
assert_eq!(body["data"]["models"][0]["label"], "GPT Plan B Live");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
assert!(requests.is_empty(), "model snapshot lookup must stay local");
let model_snapshots =
codex_proxy_rs::models::store::RedisModelSnapshotStore::new(state.redis.clone());
let stored = model_snapshots.list_plan_snapshots().await.unwrap();
assert!(stored.iter().any(|snapshot| snapshot.plan_type == "plan-b"));
assert!(!stored.iter().any(|snapshot| snapshot.plan_type == "plan-b"));
}
#[tokio::test]
@@ -44,7 +44,7 @@ async fn responses_json_should_preserve_upstream_client_failure() {
let body = response_json(response).await;
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(body["error"]["code"], "server_overloaded");
assert_eq!(body["error"]["code"], "server_is_overloaded");
assert_eq!(
body["error"]["message"],
"upstream is temporarily overloaded"
@@ -97,11 +97,25 @@ async fn responses_sse_should_preserve_upstream_client_failure() {
assert_eq!(status, StatusCode::OK);
assert!(content_type.starts_with("text/event-stream"));
assert!(body.contains(r#""code":"server_overloaded""#));
assert!(body.contains(r#""code":"server_is_overloaded""#));
assert!(body.contains("upstream is temporarily overloaded"));
assert!(body.ends_with("data: [DONE]\n\n"));
}
#[tokio::test]
async fn responses_context_length_exceeded_should_return_bad_request() {
let status = complete_response_failure_status("context_length_exceeded").await;
assert_eq!(status, StatusCode::BAD_REQUEST);
}
#[tokio::test]
async fn responses_usage_not_included_should_return_forbidden() {
let status = complete_response_failure_status("usage_not_included").await;
assert_eq!(status, StatusCode::FORBIDDEN);
}
#[tokio::test]
async fn responses_should_reject_invalid_json_without_upstream_request() {
let server = MockServer::start().await;
@@ -128,6 +142,53 @@ async fn responses_should_reject_invalid_json_without_upstream_request() {
assert!(requests.is_empty());
}
async fn complete_response_failure_status(code: &str) -> StatusCode {
let server = MockServer::start().await;
let body = format!(
"event: response.failed\ndata: {}\n\n",
json!({
"response": {
"id": "resp_failed",
"object": "response",
"status": "failed",
"error": {"code": code, "message": "upstream rejected request"}
}
})
);
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(body),
)
.expect(1)
.mount(&server)
.await;
let (app, api_key, _dir) = test_app_with_account(server.uri()).await;
app.oneshot(
Request::builder()
.method("POST")
.uri("/v1/responses")
.header("authorization", format!("Bearer {api_key}"))
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "gpt-5.5",
"input": [],
"stream": false,
"use_websocket": false
})
.to_string(),
))
.unwrap(),
)
.await
.unwrap()
.status()
}
#[tokio::test]
async fn responses_should_reject_non_object_json_without_upstream_request() {
let server = MockServer::start().await;
@@ -414,12 +414,17 @@ async fn responses_should_return_no_available_accounts_error_when_no_accounts_ar
#[tokio::test]
async fn responses_should_fallback_to_next_account_after_rate_limit() {
let server = MockServer::start().await;
let reset_at = Utc::now().timestamp() + 120;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.and(header("authorization", "Bearer access-primary"))
.respond_with(
ResponseTemplate::new(429)
.insert_header("retry-after", "120")
.insert_header("x-codex-active-limit", "codex")
.insert_header("x-codex-primary-used-percent", "100")
.insert_header("x-codex-primary-window-minutes", "5")
.insert_header("x-codex-primary-reset-at", reset_at.to_string())
.set_body_json(json!({
"error": {
"message": "rate limited",
@@ -484,6 +489,13 @@ async fn responses_should_fallback_to_next_account_after_rate_limit() {
.unwrap();
assert_eq!(primary_usage, (1, 0, 0));
assert_eq!(secondary_usage, (1, 5, 2));
let quota: (Value,) = sqlx::query_as("select quota_json from accounts where id = $1")
.bind("acct_primary")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(quota.0["active_limit"], "codex");
assert_eq!(quota.0["snapshots"][0]["primary"]["reset_at"], reset_at);
}
#[tokio::test]
@@ -856,7 +868,7 @@ async fn responses_stream_should_mark_banned_after_402_deactivated_workspace_and
}
#[tokio::test]
async fn responses_should_mark_banned_after_403_and_fallback() {
async fn responses_should_return_generic_403_without_banning_or_fallback() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
@@ -875,7 +887,7 @@ async fn responses_should_mark_banned_after_403_and_fallback() {
.insert_header("content-type", "text/event-stream")
.set_body_string(RESPONSES_AFTER_403_SSE),
)
.expect(1)
.expect(0)
.mount(&server)
.await;
@@ -908,9 +920,9 @@ async fn responses_should_mark_banned_after_403_and_fallback() {
.await
.unwrap();
assert_eq!(status, StatusCode::OK);
assert_eq!(body["id"], "resp_after_403");
assert_eq!(account_status.0, "banned");
assert_eq!(status, StatusCode::FORBIDDEN);
assert_eq!(body["error"]["message"], "request forbidden");
assert_eq!(account_status.0, "active");
}
#[tokio::test]
@@ -2152,7 +2164,7 @@ async fn responses_stream_should_return_auth_error_when_401_fallback_is_exhauste
}
#[tokio::test]
async fn responses_stream_should_mark_banned_after_403_and_fallback() {
async fn responses_stream_should_return_generic_403_without_banning_or_fallback() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
@@ -2171,7 +2183,7 @@ async fn responses_stream_should_mark_banned_after_403_and_fallback() {
.insert_header("content-type", "text/event-stream")
.set_body_string(RESPONSES_STREAM_AFTER_403_SSE),
)
.expect(1)
.expect(0)
.mount(&server)
.await;
@@ -2210,10 +2222,10 @@ async fn responses_stream_should_mark_banned_after_403_and_fallback() {
.await
.unwrap();
assert_eq!(status, StatusCode::OK);
assert!(content_type.starts_with("text/event-stream"));
assert!(body.contains("resp_stream_after_403"));
assert_eq!(account_status.0, "banned");
assert_eq!(status, StatusCode::FORBIDDEN);
assert!(content_type.starts_with("application/json"));
assert!(body.contains("request forbidden"));
assert_eq!(account_status.0, "active");
}
#[tokio::test]
@@ -399,6 +399,7 @@ async fn responses_should_passively_cache_rate_limit_headers() {
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.insert_header("x-codex-active-limit", "codex")
.insert_header("x-codex-primary-used-percent", "100")
.insert_header("x-codex-primary-window-minutes", "5")
.insert_header("x-codex-primary-reset-at", reset_at.to_string())
@@ -440,6 +441,10 @@ async fn responses_should_passively_cache_rate_limit_headers() {
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let upstream_requests = server.received_requests().await.unwrap();
assert_eq!(upstream_requests.len(), 1);
assert_eq!(upstream_requests[0].method.as_str(), "POST");
assert_eq!(upstream_requests[0].url.path(), "/codex/responses");
let stored: (
Value,
String,
@@ -454,6 +459,7 @@ async fn responses_should_passively_cache_rate_limit_headers() {
.await
.unwrap();
let quota = stored.0;
assert_eq!(quota["active_limit"], "codex");
assert_eq!(quota["snapshots"][0]["source"], "core");
assert_eq!(quota["snapshots"][0]["primary"]["limit_reached"], true);
assert_eq!(quota["snapshots"][0]["primary"]["reset_at"], reset_at);
@@ -1,2 +1,2 @@
event: response.failed
data: {"response":{"id":"resp_server_overloaded","object":"response","status":"failed","error":{"code":"server_overloaded","message":"upstream is temporarily overloaded"}}}
data: {"response":{"id":"resp_server_overloaded","object":"response","status":"failed","error":{"code":"server_is_overloaded","message":"upstream is temporarily overloaded"}}}
+42 -1
View File
@@ -60,6 +60,38 @@ fn html_forbidden_page_should_not_ban_account() {
assert!(classify_client_failure(&error).is_none());
}
#[test]
fn generic_forbidden_should_not_ban_account() {
let error = upstream_error(403, r#"{"error":{"message":"request forbidden"}}"#);
assert!(classify_client_failure(&error).is_none());
}
#[test]
fn generic_permission_error_should_not_ban_account() {
let error = upstream_error(
403,
r#"{"error":{"code":"forbidden","type":"permission_error"}}"#,
);
assert!(classify_client_failure(&error).is_none());
}
#[test]
fn identity_error_header_should_expire_account() {
let mut diagnostics = CodexUpstreamDiagnostics::with_status(403);
diagnostics.identity_error_code = Some("token_expired".to_string());
let error = upstream_error_with_diagnostics(
403,
r#"{"error":{"message":"request forbidden"}}"#,
diagnostics,
);
let failure = classify_client_failure(&error).expect("identity failure should be classified");
assert_eq!(failure.kind, AccountFailureKind::Expired);
}
fn classify(
status: u16,
body: &str,
@@ -68,12 +100,21 @@ fn classify(
}
fn upstream_error(status: u16, body: &str) -> CodexClientError {
upstream_error_with_diagnostics(status, body, CodexUpstreamDiagnostics::with_status(status))
}
fn upstream_error_with_diagnostics(
status: u16,
body: &str,
diagnostics: CodexUpstreamDiagnostics,
) -> CodexClientError {
CodexClientError::Upstream {
status: StatusCode::from_u16(status).unwrap(),
body: body.to_string(),
retry_after_seconds: None,
diagnostics: CodexUpstreamDiagnostics::with_status(status),
diagnostics: Box::new(diagnostics),
set_cookie_headers: Vec::new(),
rate_limit_headers: Vec::new(),
transport: CodexBackendTransport::HttpSse,
}
}
@@ -348,6 +348,10 @@ fn parse_rate_limit_headers_should_extract_arbitrary_limits_and_account_metadata
"x-codex-other-primary-used-percent".to_string(),
"63.5".to_string(),
),
(
"x-codex-active-limit".to_string(),
"codex-other".to_string(),
),
(
"x-codex-other-primary-window-minutes".to_string(),
"1440".to_string(),
@@ -371,6 +375,7 @@ fn parse_rate_limit_headers_should_extract_arbitrary_limits_and_account_metadata
];
let parsed = parse_rate_limit_headers(&headers).expect("rate limits should parse");
assert_eq!(parsed.active_limit.as_deref(), Some("codex_other"));
let other = parsed.limits.get("codex_other").expect("dynamic limit");
assert_eq!(other.limit_name.as_deref(), Some("Codex Other"));
assert_eq!(
@@ -1,3 +1,4 @@
use base64::{Engine as _, engine::general_purpose::STANDARD};
use codex_proxy_rs::upstream::openai::transport::CodexUpstreamDiagnostics;
use reqwest::header::{HeaderMap, HeaderValue};
@@ -24,3 +25,45 @@ fn diagnostics_should_not_treat_cf_ray_as_request_id() {
assert_eq!(diagnostics.request_id, None);
assert_eq!(diagnostics.cf_ray(), Some("ray_1"));
}
#[test]
fn diagnostics_should_extract_oai_request_id() {
let mut headers = HeaderMap::new();
headers.insert("x-oai-request-id", HeaderValue::from_static("oai_req_1"));
let diagnostics = CodexUpstreamDiagnostics::from_headers(Some(500), &headers);
assert_eq!(diagnostics.request_id.as_deref(), Some("oai_req_1"));
}
#[test]
fn diagnostics_should_extract_identity_error_facts() {
let encoded = STANDARD.encode(r#"{"error":{"code":"token_expired"}}"#);
let mut headers = HeaderMap::new();
headers.insert(
"x-openai-authorization-error",
HeaderValue::from_static("authorization failed"),
);
headers.insert("x-error-json", HeaderValue::from_str(&encoded).unwrap());
let diagnostics = CodexUpstreamDiagnostics::from_headers(Some(403), &headers);
assert_eq!(
diagnostics.identity_authorization_error.as_deref(),
Some("authorization failed")
);
assert_eq!(
diagnostics.identity_error_code.as_deref(),
Some("token_expired")
);
}
#[test]
fn diagnostics_should_ignore_malformed_identity_error_json() {
let mut headers = HeaderMap::new();
headers.insert("x-error-json", HeaderValue::from_static("not-base64"));
let diagnostics = CodexUpstreamDiagnostics::from_headers(Some(403), &headers);
assert_eq!(diagnostics.identity_error_code, None);
}
+8
View File
@@ -86,3 +86,11 @@ export function getAccountModels(data: any) {
},
})
}
export function refreshAccountModels(data: any) {
return request({
url: '/api/admin/accounts/models',
method: 'POST',
data,
})
}
@@ -1,4 +1,6 @@
<script setup lang="ts">
import { RefreshCw } from '@lucide/vue'
import BaseButton from '@/components/base/BaseButton.vue'
import BaseModal from '@/components/base/BaseModal.vue'
import BaseScrollbar from '@/components/base/BaseScrollbar.vue'
@@ -16,6 +18,7 @@ defineProps<{
finishedAt: string
durationMs: number | null
loadingModels: boolean
refreshingModels: boolean
modelOptions: any[]
statusView: any
}>()
@@ -25,6 +28,7 @@ const selectedModel = defineModel<string>('selectedModel', { required: true })
const emit = defineEmits<{
test: []
refreshModels: []
}>()
function connectionLogClass(tone: string) {
@@ -53,11 +57,26 @@ function connectionLogClass(tone: string) {
<section class="rounded-(--cp-card-radius) bg-(--cp-bg-subtle) px-4 py-3">
<div class="grid gap-2">
<label class="text-[12px] font-[760] text-(--cp-text-muted)">测试模型</label>
<div class="flex min-h-8 items-center justify-between gap-3">
<label class="text-[12px] font-[760] text-(--cp-text-muted)">测试模型</label>
<BaseButton
variant="ghost"
size="sm"
icon-only
label="刷新上游模型"
:loading="refreshingModels"
:disabled="status === 'running' || loadingModels"
@click="emit('refreshModels')"
>
<template #icon>
<RefreshCw class="size-3.5" />
</template>
</BaseButton>
</div>
<BaseSelect
v-model="selectedModel"
:options="modelOptions"
:disabled="status === 'running' || loadingModels"
:disabled="status === 'running' || loadingModels || refreshingModels"
:placeholder="loadingModels ? '加载模型中...' : '选择上游模型'"
empty-text="上游没有返回模型"
/>
@@ -175,7 +194,7 @@ function connectionLogClass(tone: string) {
<BaseButton
variant="primary"
:loading="status === 'running'"
:disabled="!account || loadingModels || !selectedModel"
:disabled="!account || loadingModels || refreshingModels || !selectedModel"
@click="emit('test')"
>
{{ logs.length > 0 || error ? '重新测试' : '开始测试' }}
@@ -3,8 +3,9 @@ import { CheckCircle2, Clock3, Wifi, XCircle } from '@lucide/vue'
import { useEventSource } from '@vueuse/core'
import { clamp } from 'es-toolkit'
import { getAccountModels } from '@/api'
import { getAccountModels, refreshAccountModels } from '@/api'
import { API_BASE_URL } from '@/api/constants'
import { toast } from '@/components/base/BaseToast'
import { withMinimumDuration } from '@/utils/async'
import { formatDateTime, formatTime } from '@/utils/date'
@@ -77,6 +78,7 @@ export function useAccountConnectionTest(options: AccountConnectionTestOptions)
const connectionTestDurationMs = shallowRef<number | null>(null)
const testingConnectionIds = ref<Set<string>>(new Set())
const loadingConnectionTestModels = shallowRef(false)
const refreshingConnectionTestModels = shallowRef(false)
const connectionTestSelectedModel = shallowRef('')
const connectionTestModelOptions = ref<ConnectionTestModelOption[]>([])
const connectionTestUrl = shallowRef<string | undefined>()
@@ -323,11 +325,7 @@ export function useAccountConnectionTest(options: AccountConnectionTestOptions)
connectionTestError.value = ''
try {
const result = (await getAccountModels({ id: account.id })) as AccountModelsResponse
connectionTestModelOptions.value = (result.models ?? []).map((model) => ({
label: model.label || model.id,
value: model.id,
}))
connectionTestSelectedModel.value = connectionTestModelOptions.value[0]?.value || ''
applyConnectionTestModels(result)
if (!connectionTestSelectedModel.value) {
connectionTestError.value = '没有可测试模型'
}
@@ -340,6 +338,35 @@ export function useAccountConnectionTest(options: AccountConnectionTestOptions)
}
}
function applyConnectionTestModels(result: AccountModelsResponse, preserveSelection = false) {
const previousSelection = preserveSelection ? connectionTestSelectedModel.value : ''
connectionTestModelOptions.value = (result.models ?? []).map((model) => ({
label: model.label || model.id,
value: model.id,
}))
connectionTestSelectedModel.value = connectionTestModelOptions.value.some(
(model) => model.value === previousSelection,
)
? previousSelection
: connectionTestModelOptions.value[0]?.value || ''
}
async function handleRefreshConnectionTestModels(account = testingAccount.value) {
if (!account?.id || refreshingConnectionTestModels.value) return
refreshingConnectionTestModels.value = true
connectionTestError.value = ''
try {
const result = (await refreshAccountModels({ id: account.id })) as AccountModelsResponse
applyConnectionTestModels(result, true)
toast.success(`已刷新 ${connectionTestModelOptions.value.length} 个上游模型`)
} catch (error: unknown) {
connectionTestError.value = errorMessage(error, '刷新上游模型失败')
toast.error(connectionTestError.value)
} finally {
refreshingConnectionTestModels.value = false
}
}
async function handleTestConnection(account = testingAccount.value) {
if (!account?.id) return
if (!connectionTestSelectedModel.value) {
@@ -429,10 +456,12 @@ export function useAccountConnectionTest(options: AccountConnectionTestOptions)
connectionTestDurationMs,
testingConnectionIds,
loadingConnectionTestModels,
refreshingConnectionTestModels,
connectionTestSelectedModel,
connectionTestModelOptions,
connectionTestStatusView,
openConnectionTest,
handleRefreshConnectionTestModels,
handleTestConnection,
}
}
+4
View File
@@ -95,10 +95,12 @@ const {
connectionTestDurationMs,
testingConnectionIds,
loadingConnectionTestModels,
refreshingConnectionTestModels,
connectionTestSelectedModel,
connectionTestModelOptions,
connectionTestStatusView,
openConnectionTest,
handleRefreshConnectionTestModels,
handleTestConnection,
} = useAccountConnectionTest({
onAccountStatus: patchAccountStatus,
@@ -292,12 +294,14 @@ bindAccountLoader(loadAccounts)
:error="connectionTestError"
:finished-at="connectionTestFinishedAt"
:loading-models="loadingConnectionTestModels"
:refreshing-models="refreshingConnectionTestModels"
:logs="connectionTestLogs"
:model="connectionTestModel"
:model-options="connectionTestModelOptions"
:started-at="connectionTestStartedAt"
:status="connectionTestStatus"
:status-view="connectionTestStatusView"
@refresh-models="handleRefreshConnectionTestModels()"
@test="handleTestConnection()"
/>