codex: persist refresh state

Co-authored-by: Codex <noreply@openai.com>
2026-11-06 16:06:08 +00:00 · 2026-06-15 11:56:51 +08:00
parent f0e62e4d9c
commit 4547bb6744
5 changed files with 136 additions and 18 deletions
@@ -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

@@ -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,
@@ -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"),
        }
    }
@@ -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<AccountProbeResult, RefreshAccountError> {
+        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,
@@ -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<Mutex<Vec<AccountStatus>>>,
+    }
+
+    #[async_trait]
+    impl TokenRefresher for StatusObservingSuccessRefresher {
+        async fn refresh(&self, _refresh_token: &str) -> Result<TokenPair, RefreshFailure> {
+            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<AtomicUsize>,