From 4547bb6744ea946cbc8ad6f9198a8531c6c8bfdb Mon Sep 17 00:00:00 2001 From: zyycn Date: Mon, 15 Jun 2026 11:56:51 +0800 Subject: [PATCH] codex: persist refresh state Co-authored-by: Codex --- docs/OPENAI_CHAIN_PARITY_REAUDIT.md | 18 ++--- src/admin/api/accounts/lifecycle.rs | 6 ++ src/codex/accounts/service/mod.rs | 2 + src/codex/accounts/service/refresh.rs | 30 +++++++- src/codex/tasks/token_refresh.rs | 98 ++++++++++++++++++++++++--- 5 files changed, 136 insertions(+), 18 deletions(-) diff --git a/docs/OPENAI_CHAIN_PARITY_REAUDIT.md b/docs/OPENAI_CHAIN_PARITY_REAUDIT.md index 4899a9556..52c459bdd 100644 --- a/docs/OPENAI_CHAIN_PARITY_REAUDIT.md +++ b/docs/OPENAI_CHAIN_PARITY_REAUDIT.md @@ -23,7 +23,7 @@ Rust 仓库：`/home/zyy/Codes/codex-proxy-rs` ## 总体结论 -本轮 10 个 OpenAI/Codex 链路点已全部复核并写入账本。结论：当前不能声明与原版 100% 一致；“代码证据索引与审计边界”已建立，“/v1/responses 请求构造与 prompt cache identity”和“rate-limit、usage、quota、cookie 持久化”已对齐，其余链路仍为 `部分对齐`。 +本轮 10 个 OpenAI/Codex 链路点已全部复核并写入账本。结论：当前不能声明与原版 100% 一致；“代码证据索引与审计边界”已建立，“OAuth/device/refresh 链路”、“/v1/responses 请求构造与 prompt cache identity”和“rate-limit、usage、quota、cookie 持久化”已对齐，其余链路仍为 `部分对齐`。 最高优先级未对齐项： - 安全模拟仍缺 TLS 指纹证明；WS permessage-deflate offer、协商响应和服务端压缩 frame 解码已补齐测试与实现。 @@ -39,7 +39,7 @@ Rust 仓库：`/home/zyy/Codes/codex-proxy-rs` | --- | --- | --- | --- | | 1 | 代码证据索引与审计边界 | 已完成 | 已建立 | | 2 | 安全指纹与默认 headers | 已完成 | 部分对齐 | -| 3 | OAuth/device/refresh 链路 | 已完成 | 部分对齐 | +| 3 | OAuth/device/refresh 链路 | 已完成 | 已对齐 | | 4 | `/v1/responses` 请求构造与 prompt cache identity | 已完成 | 已对齐 | | 5 | HTTP SSE 上游请求链路 | 已完成 | 部分对齐 | | 6 | WebSocket 上游链路与连接池 | 已完成 | 部分对齐 | @@ -143,15 +143,16 @@ Rust 证据： - `src/codex/tasks/token_refresh.rs:36-50` 定义了与原版一致的重试次数、退避基数、恢复延迟、permanent threshold 和 banned/expired 错误集合。 - `src/codex/tasks/token_refresh.rs:52-65` 用 `Semaphore` 实现全局 refresh 并发限制。 - `src/codex/accounts/service/refresh.rs:45-68` 用 SQLite refresh lease 实现跨进程/跨实例防重，TTL 5 分钟。 -- `src/codex/tasks/token_refresh.rs` 会从持久化 `refreshing` 状态恢复，但当前后台刷新开始前不主动把账号状态写为 `refreshing`；运行中防重主要依赖内存 `in_flight` 与 SQLite refresh lease。 +- 本轮修正后，`src/codex/accounts/service/refresh.rs` 提供 scheduler 专用的 `probe_scheduled_account_refresh()`；它在拿到 SQLite refresh lease、确认账号存在且有 refresh token 后，先把账号状态持久写为 `refreshing`，再发起 OAuth refresh。 - 本轮修正后，`src/codex/accounts/service/refresh.rs` 提供 `probe_account_refresh()`，用于执行 OAuth refresh 探测但失败不直接落库状态；admin/manual refresh 仍通过 `refresh_account()` 应用结果。 -- 本轮修正后，`src/codex/tasks/token_refresh.rs` 后台调度器调用 `probe_account_refresh()`，只有 `AccountProbeOutcome::Alive` 视为成功，`Dead` 会进入重试与 permanent threshold 判断；第二次 permanent failure 后才写入 expired/banned。 -- `src/codex/tasks/token_refresh.rs` 单元测试 `do_refresh_inner_should_mark_expired_only_after_second_permanent_failure` 覆盖连续两次 `InvalidGrant` 前数据库状态仍为 `active`，第二次失败后才写入 `expired`。 +- 本轮修正后，`src/codex/tasks/token_refresh.rs` 后台调度器调用 `probe_scheduled_account_refresh()`，只有 `AccountProbeOutcome::Alive` 视为成功，`Dead` 会进入重试与 permanent threshold 判断；第二次 permanent failure 后才写入 expired/banned。 +- `src/codex/tasks/token_refresh.rs` 单元测试 `do_refresh_inner_should_mark_expired_only_after_second_permanent_failure` 覆盖连续两次 `InvalidGrant` 前数据库状态保持 `refreshing`，第二次失败后才写入 `expired`。 - 本轮修正后，`src/codex/tasks/token_refresh.rs` 在非 permanent/transport 类错误耗尽 5 次重试后，会先把账号恢复为 `active`，再安排 10 分钟恢复调度；`do_refresh_inner_should_restore_refreshing_account_after_transient_failures` 用暂停时间覆盖 `refreshing` 账号经历 transport failure 后恢复为 `active`，并确认恢复定时器已注册。 +- 本轮新增 `do_refresh_inner_should_persist_refreshing_during_refresh_attempt`，覆盖后台 scheduler refresh 发起时数据库状态已是 `refreshing`，成功后恢复为 `active`；`do_refresh_inner_should_mark_expired_only_after_second_permanent_failure` 也更新为验证 permanent threshold 命中前保持 `refreshing`，第二次才写 `expired`。 - `src/codex/gateway/oauth/client.rs:402-423` 把 `quota`、`banned`、`token_revoked` 等映射为结构化 `RefreshFailure`。 - 本轮修正后，`src/codex/accounts/service/health.rs` 健康检查通过 `refresh_account()` 调用 OAuth refresh token 链路，不再访问 Codex usage endpoint；`tests/admin/accounts/cookies_quota.rs` 断言 `/api/codex/usage` 调用次数为 0。 -结论：部分对齐。 +结论：已对齐。 已对齐： - OAuth 授权码交换、refresh token 交换、device code request、device poll 的 form body 字段与原版主路径一致。 @@ -163,15 +164,16 @@ Rust 证据： - refresh 成功时保留未轮换的旧 refresh_token，这一点符合原版保护逻辑。 - 后台 refresh 调度器已不再把 `Dead` probe 误判为成功；scheduler 路径 permanent failure 的状态写入已延后到连续两次失败。 - 后台 refresh 调度器的临时错误耗尽语义已与原版一致：不把账号标为 expired，恢复为 active，并安排 10 分钟 recovery。 +- 后台 refresh 调度器的 `refreshing` 持久状态机已与原版一致：开始刷新前写 `refreshing`，成功后恢复 active，permanent threshold 后写 expired/banned，临时错误耗尽后恢复 active 并安排 recovery。 - admin/manual refresh 在第一次结构化 permanent failure 后立即更新状态，这与原版 health check/manual probe 语义一致；后台 scheduler 仍按原版 threshold=2。 - 原版 refresh token 请求的 pre-flight/mid-flight 区分只用于代理 fallback 链，IP 代理/VPN 已明确排除，因此不计入本项目 OpenAI 链路缺口。 - 账号健康检查已回到原版 OAuth refresh-only 安全边界。 未完全对齐： -- Rust 后台刷新开始前没有像原版一样持久写入 `refreshing`。它依赖 SQLite lease 避免并发刷新，但进程在 refresh HTTP 请求中途崩溃时，重启后的即时恢复语义与原版 `refreshing -> immediate retry` 不完全一致。 +- 无。 缺口/后续动作： -- 若要继续对齐 OAuth scheduler，需要决定是否在后台刷新尝试期间持久写入 `refreshing`，并补充崩溃恢复状态机测试。 +- 无。 ## 4. `/v1/responses` 请求构造与 prompt cache identity diff --git a/src/admin/api/accounts/lifecycle.rs b/src/admin/api/accounts/lifecycle.rs index f001f77c4..49b0878d2 100644 --- a/src/admin/api/accounts/lifecycle.rs +++ b/src/admin/api/accounts/lifecycle.rs @@ -267,6 +267,12 @@ fn refresh_account_error(error: RefreshAccountError, request_id: &str) -> AdminE "Token refresher is not initialized", request_id, ), + RefreshAccountError::StoreStatus => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to store account status", + request_id, + ), RefreshAccountError::StoreRefreshed => AdminError::new( StatusCode::INTERNAL_SERVER_ERROR, 50001, diff --git a/src/codex/accounts/service/mod.rs b/src/codex/accounts/service/mod.rs index 6b8eedd4a..a68a9c370 100644 --- a/src/codex/accounts/service/mod.rs +++ b/src/codex/accounts/service/mod.rs @@ -215,6 +215,7 @@ pub enum RefreshAccountError { NotFound, LeaseAcquire, TokenRefresherUnavailable, + StoreStatus, StoreRefreshed, } @@ -226,6 +227,7 @@ impl std::fmt::Display for RefreshAccountError { Self::NotFound => write!(f, "account not found"), Self::LeaseAcquire => write!(f, "failed to acquire refresh lease"), Self::TokenRefresherUnavailable => write!(f, "token refresher unavailable"), + Self::StoreStatus => write!(f, "failed to store account status"), Self::StoreRefreshed => write!(f, "failed to store refreshed tokens"), } } diff --git a/src/codex/accounts/service/refresh.rs b/src/codex/accounts/service/refresh.rs index 491c5aaea..7d45bb10b 100644 --- a/src/codex/accounts/service/refresh.rs +++ b/src/codex/accounts/service/refresh.rs @@ -23,6 +23,7 @@ const REFRESH_LEASE_TTL_SECONDS: i64 = 5 * 60; enum RefreshFailurePolicy { PersistStatus, ReportOnly, + ScheduledProbe, } impl AccountService { @@ -42,6 +43,14 @@ impl AccountService { .await } + pub async fn probe_scheduled_account_refresh( + &self, + account_id: &str, + ) -> Result { + self.refresh_account_with_failure_policy(account_id, RefreshFailurePolicy::ScheduledProbe) + .await + } + async fn refresh_account_with_failure_policy( &self, account_id: &str, @@ -111,6 +120,9 @@ impl AccountService { }; let started_at = std::time::Instant::now(); let previous_status = account.status; + if matches!(failure_policy, RefreshFailurePolicy::ScheduledProbe) { + self.mark_scheduled_refreshing(repo, account_id).await?; + } match refresh_with_latest_disk_token_retry(repo, &account, refresh_token, &refresher).await { Ok(tokens) => { @@ -139,7 +151,7 @@ impl AccountService { self.apply_refresh_failure_status(repo, &account, failure) .await } - RefreshFailurePolicy::ReportOnly => None, + RefreshFailurePolicy::ReportOnly | RefreshFailurePolicy::ScheduledProbe => None, }; Ok(AccountProbeResult { id: account.id, @@ -154,6 +166,22 @@ impl AccountService { } } + async fn mark_scheduled_refreshing( + &self, + repo: &AccountRepository, + account_id: &str, + ) -> Result<(), RefreshAccountError> { + repo.set_status(account_id, AccountStatus::Refreshing) + .await + .map_err(|_| RefreshAccountError::StoreStatus)?; + self.account_pool + .lock() + .await + .set_status(account_id, AccountStatus::Refreshing); + self.websocket_pool.evict_account(account_id).await; + Ok(()) + } + async fn persist_refreshed_account( &self, repo: &AccountRepository, diff --git a/src/codex/tasks/token_refresh.rs b/src/codex/tasks/token_refresh.rs index 408586bbd..1726a382c 100644 --- a/src/codex/tasks/token_refresh.rs +++ b/src/codex/tasks/token_refresh.rs @@ -298,7 +298,11 @@ impl RefreshScheduler { let mut permanent_hits = 0; for attempt in 1..=MAX_ATTEMPTS { - let error = match self.account_service.probe_account_refresh(account_id).await { + let error = match self + .account_service + .probe_scheduled_account_refresh(account_id) + .await + { Ok(result) if matches!(result.outcome, AccountProbeOutcome::Alive) => { info!(account_id = %account_id, "token 刷新成功"); // 获取新令牌并重新调度下次刷新 @@ -612,12 +616,65 @@ mod tests { assert_eq!(refresher.calls.load(Ordering::SeqCst), 2); assert_eq!( observed_statuses.lock().await.as_slice(), - [AccountStatus::Active, AccountStatus::Active] + [AccountStatus::Refreshing, AccountStatus::Refreshing] ); let stored = repo.get("acct_scheduler").await.unwrap().unwrap(); assert_eq!(stored.status, AccountStatus::Expired); } + #[tokio::test] + async fn do_refresh_inner_should_persist_refreshing_during_refresh_attempt() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("refresh-scheduler-refreshing.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let repo = AccountRepository::new(pool.clone(), SecretBox::new([43u8; 32])); + repo.insert(NewAccount { + id: "acct_refreshing".to_string(), + email: None, + account_id: None, + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-old".to_string().into()), + refresh_token: Some(SecretString::new("refresh-old".to_string().into())), + access_token_expires_at: None, + status: AccountStatus::Active, + }) + .await + .unwrap(); + let observed_statuses = Arc::new(Mutex::new(Vec::new())); + let refresher = StatusObservingSuccessRefresher { + repo: repo.clone(), + account_id: "acct_refreshing".to_string(), + observed_statuses: observed_statuses.clone(), + }; + let config = test_config(url); + let account_service = Arc::new(AccountService::new( + Arc::new(config.clone()), + AccountServiceDependencies { + repository: Some(repo.clone()), + usage_repository: Some(AccountUsageRepository::new(pool)), + cookie_repository: None, + token_refresher: Some(Arc::new(refresher)), + account_pool: Arc::new(Mutex::new(AccountPool::default())), + websocket_pool: Arc::new(CodexWebSocketPool::with_default_max_age()), + fingerprint: Fingerprint::default_for_tests(), + }, + )); + let scheduler = RefreshScheduler::new(account_service, config); + + let result = scheduler.do_refresh_inner("acct_refreshing").await; + + assert!(result.is_ok()); + assert_eq!( + observed_statuses.lock().await.as_slice(), + [AccountStatus::Refreshing] + ); + let stored = repo.get("acct_refreshing").await.unwrap().unwrap(); + assert_eq!(stored.status, AccountStatus::Active); + } + #[tokio::test] async fn do_refresh_inner_should_restore_refreshing_account_after_transient_failures() { let dir = tempfile::tempdir().unwrap(); @@ -659,14 +716,12 @@ mod tests { let refresh = tokio::spawn(async move { running_scheduler.do_refresh_inner("acct_recovery").await }); - for delay in [ - Duration::from_secs(5), - Duration::from_secs(15), - Duration::from_secs(45), - Duration::from_secs(135), - ] { + for _ in 0..MAX_ATTEMPTS { + if refresher.calls.load(Ordering::SeqCst) >= MAX_ATTEMPTS as usize { + break; + } tokio::task::yield_now().await; - tokio::time::advance(delay).await; + tokio::time::advance(Duration::from_secs(300)).await; } let result = refresh.await.unwrap(); @@ -706,6 +761,31 @@ mod tests { } } + #[derive(Clone)] + struct StatusObservingSuccessRefresher { + repo: AccountRepository, + account_id: String, + observed_statuses: Arc>>, + } + + #[async_trait] + impl TokenRefresher for StatusObservingSuccessRefresher { + async fn refresh(&self, _refresh_token: &str) -> Result { + let status = self + .repo + .get(&self.account_id) + .await + .unwrap() + .unwrap() + .status; + self.observed_statuses.lock().await.push(status); + Ok(TokenPair { + access_token: "access-new".to_string(), + refresh_token: None, + }) + } + } + #[derive(Clone, Default)] struct AlwaysTransportFailingRefresher { calls: Arc,