mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 16:06:08 +00:00
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:
@@ -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()),
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -13,23 +13,18 @@ use crate::models::service::{
|
||||
ModelRefreshPlanAccount, ModelRefreshResult, ModelService, ModelServiceError,
|
||||
};
|
||||
|
||||
use super::{
|
||||
coordinator::SchedulerHandle,
|
||||
periodic::{PeriodicTaskConfig, PeriodicTaskRunner, PeriodicTaskStartup, spawn_periodic_task},
|
||||
};
|
||||
use super::coordinator::SchedulerHandle;
|
||||
|
||||
/// 模型刷新任务接线器。
|
||||
pub struct ModelRefreshTask {
|
||||
model_service: Arc<ModelService>,
|
||||
account_pool: Arc<AccountPoolService>,
|
||||
refresh_interval_secs: u64,
|
||||
initial_delay_ms: u64,
|
||||
retry_delay_ms: u64,
|
||||
max_retries: u32,
|
||||
account_pseudonymizer: Arc<AccountPseudonymizer>,
|
||||
}
|
||||
|
||||
const DEFAULT_REFRESH_INTERVAL_SECS: u64 = 3600;
|
||||
const INITIAL_DELAY_MS: u64 = 1000;
|
||||
const RETRY_DELAY_MS: u64 = 10_000;
|
||||
const MAX_RETRIES: u32 = 12;
|
||||
@@ -44,7 +39,6 @@ impl ModelRefreshTask {
|
||||
Self {
|
||||
model_service,
|
||||
account_pool,
|
||||
refresh_interval_secs: DEFAULT_REFRESH_INTERVAL_SECS,
|
||||
initial_delay_ms: INITIAL_DELAY_MS,
|
||||
retry_delay_ms: RETRY_DELAY_MS,
|
||||
max_retries: MAX_RETRIES,
|
||||
@@ -52,15 +46,13 @@ impl ModelRefreshTask {
|
||||
}
|
||||
}
|
||||
|
||||
/// 启动后台刷新任务。
|
||||
/// 启动一次性首次刷新任务。
|
||||
pub fn start(self) -> SchedulerHandle {
|
||||
let config = PeriodicTaskConfig::new(
|
||||
self.refresh_interval_secs,
|
||||
"模型刷新任务已启动",
|
||||
"模型刷新任务正在关闭",
|
||||
)
|
||||
.wait_first_interval();
|
||||
spawn_periodic_task(self, config)
|
||||
let (shutdown_tx, mut shutdown_rx) = mpsc::channel(1);
|
||||
let handle = tokio::spawn(async move {
|
||||
self.run_initial_refresh(&mut shutdown_rx).await;
|
||||
});
|
||||
SchedulerHandle::new(shutdown_tx, handle)
|
||||
}
|
||||
|
||||
/// 在 Responses 响应声明模型目录 ETag 变化时异步刷新。
|
||||
@@ -115,84 +107,52 @@ impl ModelRefreshTask {
|
||||
.await;
|
||||
Ok(result)
|
||||
}
|
||||
}
|
||||
|
||||
impl PeriodicTaskRunner for ModelRefreshTask {
|
||||
fn before_loop<'a>(
|
||||
&'a mut self,
|
||||
shutdown_rx: &'a mut mpsc::Receiver<()>,
|
||||
) -> super::periodic::TaskFuture<'a, PeriodicTaskStartup> {
|
||||
Box::pin(async move {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
info!("模型刷新任务在首次拉取前关闭");
|
||||
return PeriodicTaskStartup::Stop;
|
||||
}
|
||||
() = sleep(Duration::from_millis(self.initial_delay_ms)) => {}
|
||||
async fn run_initial_refresh(&self, shutdown_rx: &mut mpsc::Receiver<()>) {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
info!("模型刷新任务在首次拉取前关闭");
|
||||
return;
|
||||
}
|
||||
() = sleep(Duration::from_millis(self.initial_delay_ms)) => {}
|
||||
}
|
||||
|
||||
let mut attempt = 0;
|
||||
let mut has_fetched_once = false;
|
||||
let mut deferred_until_accounts_available = false;
|
||||
|
||||
while attempt < self.max_retries {
|
||||
match self.refresh_once().await {
|
||||
Ok(result) => {
|
||||
info!(
|
||||
refreshed_plans = result.refreshed_plans,
|
||||
model_count = result.model_count,
|
||||
"首次模型刷新成功"
|
||||
);
|
||||
has_fetched_once = true;
|
||||
break;
|
||||
}
|
||||
Err(error) => {
|
||||
attempt += 1;
|
||||
if matches!(error, ModelServiceError::NoAccounts) {
|
||||
info!("首次模型刷新跳过:没有可用账户,将进入周期重试");
|
||||
deferred_until_accounts_available = true;
|
||||
break;
|
||||
}
|
||||
|
||||
warn!(
|
||||
attempt,
|
||||
max_retries = self.max_retries,
|
||||
error = ?error,
|
||||
"首次模型刷新失败"
|
||||
);
|
||||
|
||||
if attempt < self.max_retries {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
info!("模型刷新任务在重试等待期间关闭");
|
||||
return PeriodicTaskStartup::Stop;
|
||||
}
|
||||
() = sleep(Duration::from_millis(self.retry_delay_ms)) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !has_fetched_once && !deferred_until_accounts_available {
|
||||
warn!("模型刷新任务首次尝试全部失败,切换到周期刷新");
|
||||
}
|
||||
|
||||
PeriodicTaskStartup::Continue
|
||||
})
|
||||
}
|
||||
|
||||
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
|
||||
Box::pin(async move {
|
||||
let mut attempt = 0;
|
||||
while attempt < self.max_retries {
|
||||
match self.refresh_once().await {
|
||||
Ok(_) => {}
|
||||
Ok(result) => {
|
||||
info!(
|
||||
refreshed_plans = result.refreshed_plans,
|
||||
model_count = result.model_count,
|
||||
"首次模型刷新成功"
|
||||
);
|
||||
return;
|
||||
}
|
||||
Err(ModelServiceError::NoAccounts) => {
|
||||
info!("模型刷新跳过:没有可用账户");
|
||||
info!("首次模型刷新跳过:没有可用账户,后续由 ETag 或手动操作触发");
|
||||
return;
|
||||
}
|
||||
Err(error) => {
|
||||
warn!(error = ?error, "刷新模型列表失败");
|
||||
attempt += 1;
|
||||
warn!(
|
||||
attempt,
|
||||
max_retries = self.max_retries,
|
||||
error = ?error,
|
||||
"首次模型刷新失败"
|
||||
);
|
||||
|
||||
if attempt < self.max_retries {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
info!("模型刷新任务在重试等待期间关闭");
|
||||
return;
|
||||
}
|
||||
() = sleep(Duration::from_millis(self.retry_delay_ms)) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
warn!("模型刷新任务首次尝试全部失败,等待 ETag 或手动操作再次触发");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,20 +10,11 @@ use super::coordinator::SchedulerHandle;
|
||||
|
||||
pub(crate) type TaskFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
|
||||
|
||||
/// 周期任务进入循环前的决策。
|
||||
pub(crate) enum PeriodicTaskStartup {
|
||||
/// 继续进入周期循环。
|
||||
Continue,
|
||||
/// 任务已自行结束,不再进入周期循环。
|
||||
Stop,
|
||||
}
|
||||
|
||||
/// 周期任务调度配置。
|
||||
pub(crate) struct PeriodicTaskConfig {
|
||||
interval_secs: u64,
|
||||
started_message: &'static str,
|
||||
stopped_message: &'static str,
|
||||
run_immediately: bool,
|
||||
}
|
||||
|
||||
impl PeriodicTaskConfig {
|
||||
@@ -36,25 +27,12 @@ impl PeriodicTaskConfig {
|
||||
interval_secs,
|
||||
started_message,
|
||||
stopped_message,
|
||||
run_immediately: true,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn wait_first_interval(mut self) -> Self {
|
||||
self.run_immediately = false;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// 周期任务行为。
|
||||
pub(crate) trait PeriodicTaskRunner: Send + 'static {
|
||||
fn before_loop<'a>(
|
||||
&'a mut self,
|
||||
_shutdown_rx: &'a mut mpsc::Receiver<()>,
|
||||
) -> TaskFuture<'a, PeriodicTaskStartup> {
|
||||
Box::pin(std::future::ready(PeriodicTaskStartup::Continue))
|
||||
}
|
||||
|
||||
fn tick(&mut self) -> TaskFuture<'_, ()>;
|
||||
|
||||
fn shutdown(&mut self) -> TaskFuture<'_, ()> {
|
||||
@@ -75,17 +53,7 @@ where
|
||||
"{}", config.started_message
|
||||
);
|
||||
|
||||
if matches!(
|
||||
task.before_loop(&mut shutdown_rx).await,
|
||||
PeriodicTaskStartup::Stop
|
||||
) {
|
||||
return;
|
||||
}
|
||||
|
||||
let mut ticker = interval(Duration::from_secs(config.interval_secs));
|
||||
if !config.run_immediately {
|
||||
ticker.tick().await;
|
||||
}
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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, "a)
|
||||
.await;
|
||||
if quota_snapshot_limit_reached("a) {
|
||||
acquired.release_without_usage().await;
|
||||
tracing::info!(
|
||||
request_id = %context.request_id,
|
||||
account_id = %account_id,
|
||||
quota_verify_required = true,
|
||||
quota_verify_result = "limit_reached",
|
||||
retry_with_another_account = true,
|
||||
"Quota verification reported exhausted account before upstream request"
|
||||
);
|
||||
return QuotaEnterOutcome::LimitReached;
|
||||
}
|
||||
|
||||
QuotaEnterOutcome::Ready(Box::new(acquired_with_verified_quota(acquired, "a)))
|
||||
}
|
||||
|
||||
pub(super) fn classify(observation: &AttemptObservation) -> Option<ClassifiedQuotaFailure> {
|
||||
if matches!(
|
||||
observation.kind,
|
||||
AttemptObservationKind::CandidatePreparationRejected
|
||||
) {
|
||||
return observation
|
||||
.account
|
||||
.as_ref()
|
||||
.map(|account| ClassifiedQuotaFailure {
|
||||
exhaustion: AccountExhaustionRecord::new(
|
||||
account.id.clone(),
|
||||
ExhaustedAccountKind::RateLimited,
|
||||
LIMIT_REACHED_MESSAGE,
|
||||
),
|
||||
effect: None,
|
||||
});
|
||||
}
|
||||
classify(observation)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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")]
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"}}}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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()"
|
||||
/>
|
||||
|
||||
|
||||
Reference in New Issue
Block a user