starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-06 16:06:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(openai): align wire profile os_type and harden quota 5xx backoff

Browse Source 
- Default wire profile os_type to macOS to match the audited official
  desktop user-agent
- Retry quota usage 5xx rejections with 1s/2s exponential backoff,
  leaving 4xx to the existing quota state transitions
- Classify 503 with Retry-After into a resettable cooldown instead of
  an unclassified transient warning
- Add account_id and conversation hash to WebSocket pump exit logs so
  idle-connection resets can be attributed to an account

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-08-02 14:27:04 +08:00
co-authored by Codex
parent 31d6ef9585
commit f314ebb3a2
7 changed files with 155 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/providers/openai/src/config.rs
+1 -1
View File
@@ -295,7 +295,7 @@ impl Default for CodexWireProfileConfig {
codex_version: "0.145.0".to_owned(),
desktop_version: "26.721.31836".to_owned(),
desktop_build: "5828".to_owned(),
os_type: "Mac OS".to_owned(),
os_type: "macOS".to_owned(),
os_version: "15.7.1".to_owned(),
arch: "arm64".to_owned(),
terminal: "unknown".to_owned(),
backend/crates/providers/openai/src/credential/quota.rs
+40 -3
View File
@@ -30,6 +30,9 @@ const MAX_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60 * 60);
const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60);
const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5);
const QUOTA_PERIODIC_REFRESH_MIN_INTERVAL: Duration = Duration::from_secs(30 * 60);
// 5xx 上游拒绝的短退避重试预算；指数退避 1s/2s，吞掉瞬时抖动（对齐 TS withRetry）。
const QUOTA_FETCH_5XX_MAX_RETRIES: u32 = 2;
const QUOTA_FETCH_5XX_BASE_DELAY: Duration = Duration::from_secs(1);
/// OpenAI Provider 主动额度刷新的调度策略。
///
@@ -1089,7 +1092,7 @@ impl CodexCredentialQuotaService {
.prepare(account)
.await
.map_err(|_| CodexQuotaFetchError::InvalidCredential)?;
let mut result = fetch_usage_once(client, &prepared).await;
let mut result = fetch_usage_with_5xx_retry(client, &prepared).await;
if let Err(CodexQuotaFetchAttemptError::Upstream(error)) = &result
&& let Some(recovered) = self
.agent_identity
@@ -1102,7 +1105,7 @@ impl CodexCredentialQuotaService {
.map_err(|_| CodexQuotaFetchError::Recovery)?
{
prepared = recovered;
result = fetch_usage_once(client, &prepared).await;
result = fetch_usage_with_5xx_retry(client, &prepared).await;
}
match result {
Ok(value) => Ok(FetchedCodexQuota {
@@ -1143,6 +1146,39 @@ async fn fetch_usage_once(
.map_err(CodexQuotaFetchAttemptError::Upstream)
}
/// 对 5xx 上游拒绝做有限次指数退避重试（1s/2s），吞掉瞬时抖动。
///
/// 4xx（含 402/429）不重试：它们已经走额度状态转换，重试只会放大上游负载。
async fn fetch_usage_with_5xx_retry(
client: &CodexBackendClient,
prepared: &PreparedCodexRuntimeCredential,
) -> Result<Value, CodexQuotaFetchAttemptError> {
let mut attempt = 0_u32;
loop {
let result = fetch_usage_once(client, prepared).await;
let retryable = match &result {
Ok(_) => false,
Err(CodexQuotaFetchAttemptError::Upstream(CodexClientError::Upstream {
status,
..
})) => status.is_server_error(),
Err(_) => false,
};
if !retryable || attempt >= QUOTA_FETCH_5XX_MAX_RETRIES {
return result;
}
attempt += 1;
let delay = QUOTA_FETCH_5XX_BASE_DELAY.saturating_mul(attempt);
tracing::warn!(
account_id = %prepared.account.id(),
retry_attempt = attempt,
retry_delay_ms = delay.as_millis(),
"OpenAI quota usage 5xx upstream rejection; retrying with backoff"
);
tokio::time::sleep(delay).await;
}
}
fn eligible_periodic_quota_refresh(account: &ProviderAccount, now: SystemTime) -> bool {
account.enabled() && access_token_is_current(account, now)
}
@@ -1354,7 +1390,8 @@ fn quota_state_transition(
}
match *status {
StatusCode::PAYMENT_REQUIRED => Some(CodexQuotaStateTransition::Exhausted),
StatusCode::TOO_MANY_REQUESTS => {
// 429 与 503 都是"稍后重试"语义：优先尊重 Retry-After，缺失时用默认冷却。
StatusCode::TOO_MANY_REQUESTS | StatusCode::SERVICE_UNAVAILABLE => {
let duration = Duration::from_secs(
retry_after_seconds.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN.as_secs()),
)
backend/crates/providers/openai/src/transport/websocket/coordinator.rs
+21 -4
View File
@@ -24,7 +24,7 @@ use super::{
WebSocketPoolConnectLease, WebSocketPoolConnectOutcome, WebSocketPoolDecision,
WebSocketPoolLease,
},
pump::{PumpKeepalive, PumpedWebSocket},
pump::{PumpKeepalive, PumpLogContext, PumpedWebSocket},
};
pub(crate) const WEBSOCKET_FAST_PATH_BUDGET: Duration = Duration::from_millis(800);
@@ -203,10 +203,12 @@ async fn prepare_unpooled_websocket(
decision_started_at: Instant,
) -> Result<PreparedWebSocket, CodexWebSocketExchangeError> {
let permit = acquire_breaker_permit(breaker, origin_key)?;
let context = pump_log_context_from_connection(request.connection());
let connected = connect_with_budget(
request.connection(),
PumpKeepalive::disabled(),
fast_path_budget,
context,
)
.await;
let (connection, connect_elapsed) = finish_breaker_attempt(permit, connected)?;
@@ -302,6 +304,10 @@ fn start_pooled_websocket_connect(
let cancellation = connect_lease.cancellation_token();
let fast_path_reporter = permit.fast_path_reporter();
let keepalive = pool.keepalive();
let context = PumpLogContext::new(
Some(task_key.account_id().to_owned()),
Some(task_key.conversation_id_hash()),
);
let (sender, receiver) = oneshot::channel();
pool.spawn_connect_task(async move {
let connected = tokio::select! {
@@ -319,7 +325,7 @@ fn start_pooled_websocket_connect(
);
return;
}
result = connect_with_budget(&connection, keepalive, None) => result,
result = connect_with_budget(&connection, keepalive, None, context) => result,
};
match finish_breaker_attempt(permit, connected) {
Ok((connection, connect_elapsed)) => {
@@ -392,17 +398,28 @@ fn duration_millis_u64(duration: Duration) -> u64 {
.max(1)
}
/// 非池化建连的 pump 日志上下文：从业务头提取账号归属（会话 hash 只有池 key 才有）。
fn pump_log_context_from_connection(connection: &CodexWebSocketConnection) -> PumpLogContext {
let account_id = connection
.headers()
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("chatgpt-account-id"))
.map(|(_, value)| value.clone());
PumpLogContext::new(account_id, None)
}
async fn connect_with_budget(
connection: &CodexWebSocketConnection,
keepalive: PumpKeepalive,
fast_path_budget: Option<Duration>,
context: PumpLogContext,
) -> Result<(PooledWebSocketConnection, Duration), CodexWebSocketExchangeError> {
let started_at = Instant::now();
let connected = match fast_path_budget {
Some(budget) => timeout(budget, connect_pumped_websocket(connection, keepalive))
Some(budget) => timeout(budget, connect_pumped_websocket(connection, keepalive, context))
.await
.map_err(|_| CodexWebSocketExchangeError::FastPathTimeout { timeout: budget })?,
None => connect_pumped_websocket(connection, keepalive).await,
None => connect_pumped_websocket(connection, keepalive, context).await,
}?;
let (websocket, response) = connected;
Ok((
backend/crates/providers/openai/src/transport/websocket/handshake.rs
+3 -2
View File
@@ -30,7 +30,7 @@ use super::{
error::CodexWebSocketExchangeError,
model::{CodexWebSocketConnection, CodexWebSocketRequest, WebSocketContinuationRequirement},
pool::CodexWebSocketConnectionMetadata,
pump::{PumpKeepalive, PumpedWebSocket, RawWsStream},
pump::{PumpKeepalive, PumpLogContext, PumpedWebSocket, RawWsStream},
};
const WEBSOCKET_EXTENSIONS: &str = "permessage-deflate; client_max_window_bits";
@@ -93,9 +93,10 @@ pub fn responses_websocket_endpoint(base_url: &str) -> String {
pub(super) async fn connect_pumped_websocket(
connection: &CodexWebSocketConnection,
keepalive: PumpKeepalive,
context: PumpLogContext,
) -> Result<(PumpedWebSocket, WsResponse<Option<Vec<u8>>>), CodexWebSocketExchangeError> {
let (raw, response) = connect_websocket(connection).await?;
Ok((PumpedWebSocket::new(raw, keepalive), response))
Ok((PumpedWebSocket::new(raw, keepalive, context), response))
}
pub(super) async fn send_websocket_request(
backend/crates/providers/openai/src/transport/websocket/pump.rs
+38 -3
View File
@@ -36,6 +36,22 @@ pub(crate) type RawWsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
const PUMP_COMMAND_BUFFER: usize = 32;
/// pump 日志的账号级关联上下文；空闲连接被上游重置等场景用于回溯归属。
#[derive(Debug, Clone, Default)]
pub(crate) struct PumpLogContext {
pub(crate) account_id: Option<String>,
pub(crate) conversation_id_hash: Option<String>,
}
impl PumpLogContext {
pub(crate) fn new(account_id: Option<String>, conversation_id_hash: Option<String>) -> Self {
Self {
account_id,
conversation_id_hash,
}
}
}
/// pump 保活策略。
#[derive(Debug, Clone, Copy)]
pub(crate) struct PumpKeepalive {
@@ -188,7 +204,7 @@ impl std::fmt::Debug for PumpedWebSocket {
impl PumpedWebSocket {
/// 用底层流启动一个 pump 任务并返回句柄。
pub(crate) fn new(inner: RawWsStream, keepalive: PumpKeepalive) -> Self {
pub(crate) fn new(inner: RawWsStream, keepalive: PumpKeepalive, context: PumpLogContext) -> Self {
let (tx_command, rx_command) = mpsc::channel::<PumpCommand>(PUMP_COMMAND_BUFFER);
let (tx_message, rx_message) = mpsc::channel(PUMP_MESSAGE_BUFFER);
let connection_id = Uuid::new_v4();
@@ -203,6 +219,7 @@ impl PumpedWebSocket {
rx_command,
tx_message,
keepalive,
context,
closed_for_task,
exit_reason_for_task,
)
@@ -273,12 +290,17 @@ impl Drop for PumpedWebSocket {
}
}
#[expect(
clippy::too_many_arguments,
reason = "pump owns one flat set of channel/state handles and log context"
)]
async fn pump_loop(
connection_id: Uuid,
mut inner: RawWsStream,
mut rx_command: mpsc::Receiver<PumpCommand>,
tx_message: mpsc::Sender<Result<Message, tungstenite::Error>>,
keepalive: PumpKeepalive,
context: PumpLogContext,
closed: Arc<AtomicBool>,
exit_reason: Arc<Mutex<Option<PumpExitReason>>>,
) {
@@ -429,7 +451,7 @@ async fn pump_loop(
*exit_reason
.lock()
.expect("WebSocket pump exit reason lock poisoned") = Some(reason.clone());
log_pump_exit(connection_id, &reason, backpressure_events);
log_pump_exit(connection_id, &context, &reason, backpressure_events);
if reason.should_send_close() {
let _ = inner.send(Message::Close(None)).await;
@@ -510,11 +532,20 @@ async fn send_keepalive_ping(
}))
}
fn log_pump_exit(connection_id: Uuid, reason: &PumpExitReason, backpressure_events: u64) {
fn log_pump_exit(
connection_id: Uuid,
context: &PumpLogContext,
reason: &PumpExitReason,
backpressure_events: u64,
) {
let detail = reason.detail().unwrap_or_default();
let account_id = context.account_id.as_deref().unwrap_or("unknown");
let conversation_id_hash = context.conversation_id_hash.as_deref().unwrap_or("unknown");
if reason.is_unexpected() {
tracing::warn!(
websocket_connection_id = %connection_id,
account_id = %account_id,
conversation_id_hash = %conversation_id_hash,
pump_exit_reason = reason.as_str(),
pump_exit_detail = detail,
backpressure_events,
@@ -523,6 +554,8 @@ fn log_pump_exit(connection_id: Uuid, reason: &PumpExitReason, backpressure_even
} else if matches!(reason, PumpExitReason::UpstreamCloseFrame { .. }) {
tracing::info!(
websocket_connection_id = %connection_id,
account_id = %account_id,
conversation_id_hash = %conversation_id_hash,
pump_exit_reason = reason.as_str(),
pump_exit_detail = detail,
backpressure_events,
@@ -531,6 +564,8 @@ fn log_pump_exit(connection_id: Uuid, reason: &PumpExitReason, backpressure_even
} else {
tracing::debug!(
websocket_connection_id = %connection_id,
account_id = %account_id,
conversation_id_hash = %conversation_id_hash,
pump_exit_reason = reason.as_str(),
pump_exit_detail = detail,
backpressure_events,
backend/crates/providers/openai/tests/credential/quota.rs
+49 -1
View File
@@ -945,6 +945,53 @@ async fn periodic_quota_synchronization_does_not_downgrade_exhaustion_on_rate_li
);
}
#[tokio::test]
async fn periodic_quota_synchronization_respects_retry_after_on_service_unavailable() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_periodic_503_retry_after").await;
let account = store
.account("acct_periodic_503_retry_after")
.expect("created account");
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota_exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
// 5xx 短退避重试后仍失败，Retry-After 用于写冷却，不再落入"未分类"。
.respond_with(ResponseTemplate::new(503).insert_header("Retry-After", "120"))
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
service.synchronize().await.expect("quota synchronization");
let account = store
.account("acct_periodic_503_retry_after")
.expect("persisted account");
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
}
#[tokio::test]
async fn periodic_quota_synchronization_skips_ordinary_cooldown_without_exhausted_quota() {
let store = Arc::new(MemoryAccountStore::default());
@@ -1028,7 +1075,8 @@ async fn periodic_quota_synchronization_rechecks_exhausted_accounts_without_wait
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(500))
.expect(1)
// 5xx 短退避重试 1s/2s 后仍失败，最终计入 transient。
.expect(3)
.mount(&server)
.await;
let service = quota_service_with_base_url(
backend/crates/providers/openai/tests/transport/profile.rs
+3 -3
View File
@@ -66,7 +66,7 @@ fn wire_profile_should_generate_codex_core_user_agent() {
codex_version: "0.144.2".to_owned(),
desktop_version: "26.707.72221".to_owned(),
desktop_build: "72221".to_owned(),
os_type: "Mac OS".to_owned(),
os_type: "macOS".to_owned(),
os_version: "15.7.1".to_owned(),
arch: "arm64".to_owned(),
terminal: "unknown".to_owned(),
@@ -78,7 +78,7 @@ fn wire_profile_should_generate_codex_core_user_agent() {
assert_eq!(
profile.user_agent(),
"Codex Desktop/0.144.2 (Mac OS 15.7.1; arm64) unknown"
"Codex Desktop/0.144.2 (macOS 15.7.1; arm64) unknown"
);
}
@@ -219,7 +219,7 @@ fn wire_profile() -> CodexWireProfile {
codex_version: "0.144.2".to_owned(),
desktop_version: "26.707.72221".to_owned(),
desktop_build: "72221".to_owned(),
os_type: "Mac OS".to_owned(),
os_type: "macOS".to_owned(),
os_version: "15.7.1".to_owned(),
arch: "arm64".to_owned(),
terminal: "unknown".to_owned(),