mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 00:05:49 +00:00
feat: add native account refresh scheduler
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
@@ -13,7 +13,7 @@ Update this file before each feature commit.
|
||||
| Logging and pagination | Completed | `eb0341e` + `fc16c48` + `b0c10f2` | `cargo test event_logs_are_cursor_paginated`; `cargo test admin_logs_are_cursor_paginated_and_include_request_id`; `cargo test rotating_writer_splits_files_when_size_limit_is_reached`; full `fmt/test/clippy` before commit | Adds cursor pagination, `Page<T>` camelCase serialization, admin `/admin/logs` query API, event log repository, runtime tracing initialization, HTTP trace layer, and date/size-based file rotation. |
|
||||
| TLS headers and fingerprint | Completed | `1bdd08c` | `cargo test codex_headers_include_desktop_identity_and_turn_state`; `cargo tree \| rg 'reqwest\|rustls'`; full `fmt/test/clippy` before commit | Adds Codex Desktop fingerprint model, exact identity/request headers, and pinned reqwest/rustls client builder. |
|
||||
| Cookie persistence | Completed | `4710926` + `5d3cc6e` | `cargo test cookie_jar_captures_and_replays_account_scoped_cookies`; `cargo test cookie_repository_encrypts_and_replays_account_scoped_cookies`; full `fmt/test/clippy` before commit | Adds account-scoped in-memory Cookie jar and encrypted SQLite Cookie repository; Cookies are replayed only for the same account/domain. |
|
||||
| Account pool and refresh | Completed | `0f97c5c` | `cargo test account_pool_skips_expired_disabled_banned_and_quota_exhausted_accounts`; full `fmt/test/clippy` before commit | Adds account statuses, acquisition filter, lifecycle events, token pair, OAuth config, and native refresh policy shell with no legacy compatibility layer. |
|
||||
| Account pool and refresh | Completed | `0f97c5c` + refresh scheduler commit | `cargo test account_pool_skips_expired_disabled_banned_and_quota_exhausted_accounts`; `cargo test --test account_refresh_test`; full `fmt/test/clippy` before commit | Adds account statuses, acquisition filter, lifecycle events, token pair, OAuth config, and native refresh scheduler for before-expiry refresh, 401-triggered refresh, refresh-token preservation, failure status transitions, and bounded concurrency. |
|
||||
| Translation | Completed | `eff182e` | `cargo test chat_completion_translates_to_codex_response_request`; full `fmt/test/clippy` before commit | Adds OpenAI Chat request types, Codex Responses request type, translation, and OpenAI-compatible error body helper. |
|
||||
| HTTP routes | Completed | `ab0ad8f` + `3830b21` | `cargo test v1_requires_client_api_key_not_admin_cookie`; `cargo test models_route_returns_openai_compatible_codex_model_list`; full `fmt/test/clippy` before commit | Adds Axum router, health route, `/v1/models`, `/v1` unauthorized OpenAI error body, global request-id response header, and binary listener. |
|
||||
| Admin/frontend response contract | Completed | `1c34ea7` + `b0c10f2` | `cargo test admin_envelope_serializes_lower_camel_case_request_id`; `cargo test admin_logs_reject_missing_admin_session_cookie`; `cargo test v1_response_has_request_id_header_without_admin_body_field`; full `fmt/test/clippy` before commit | Adds reusable admin envelopes with body codes, lower camelCase JSON, paginated list shape, requestId bodies, admin session cookie enforcement, and keeps `/v1/*` OpenAI-compatible. |
|
||||
|
||||
@@ -2953,6 +2953,7 @@ git commit -m "docs: document codex-proxy-rs scope"
|
||||
## Audit Follow-ups
|
||||
|
||||
- [x] Add `POST /admin/login` so admin login is a password-backed session flow, not a client API key flow. The route issues an HttpOnly `cpr_admin_session` cookie, returns the admin envelope with lower camelCase `expiresAt`, and rejects `Bearer cpr_...` as admin credentials.
|
||||
- [x] Replace the refresh policy shell with a native refresh scheduler. It refreshes before expiry, refreshes immediately after an upstream 401 trigger, preserves the existing refresh token when the refresh response omits one, maps refresh failures to account statuses, and enforces configured concurrency with a semaphore.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -33,6 +33,7 @@ create table if not exists accounts (
|
||||
plan_type text,
|
||||
access_token_cipher text not null,
|
||||
refresh_token_cipher text,
|
||||
access_token_expires_at text,
|
||||
status text not null,
|
||||
quota_json text,
|
||||
quota_fetched_at text,
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use chrono::Utc;
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
@@ -21,6 +21,7 @@ pub struct Account {
|
||||
pub plan_type: Option<String>,
|
||||
pub access_token: String,
|
||||
pub refresh_token: Option<String>,
|
||||
pub access_token_expires_at: Option<DateTime<Utc>>,
|
||||
pub status: AccountStatus,
|
||||
pub added_at: String,
|
||||
pub last_used_at: Option<String>,
|
||||
@@ -37,6 +38,7 @@ impl Account {
|
||||
plan_type: None,
|
||||
access_token: format!("token-{id}"),
|
||||
refresh_token: Some(format!("refresh-{id}")),
|
||||
access_token_expires_at: Some(Utc::now() + Duration::hours(1)),
|
||||
status,
|
||||
added_at: Utc::now().to_rfc3339(),
|
||||
last_used_at: None,
|
||||
|
||||
@@ -1,5 +1,143 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use thiserror::Error;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use crate::{
|
||||
accounts::model::{Account, AccountStatus},
|
||||
auth::token::TokenPair,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct RefreshPolicy {
|
||||
pub refresh_margin_seconds: u64,
|
||||
pub refresh_concurrency: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum RefreshTrigger {
|
||||
BeforeExpiry,
|
||||
Unauthorized,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
|
||||
pub enum RefreshFailure {
|
||||
#[error("refresh token is invalid or expired")]
|
||||
InvalidGrant,
|
||||
#[error("account quota is exhausted")]
|
||||
QuotaExhausted,
|
||||
#[error("account is banned")]
|
||||
Banned,
|
||||
#[error("account is disabled")]
|
||||
Disabled,
|
||||
#[error("refresh transport failed")]
|
||||
Transport,
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum RefreshError {
|
||||
#[error("refresh task semaphore closed")]
|
||||
ConcurrencyClosed,
|
||||
#[error("refresh transport failed")]
|
||||
Transport,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait TokenRefresher: Send + Sync + 'static {
|
||||
async fn refresh(&self, refresh_token: &str) -> Result<TokenPair, RefreshFailure>;
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RefreshScheduler<C> {
|
||||
policy: RefreshPolicy,
|
||||
client: Arc<C>,
|
||||
semaphore: Arc<Semaphore>,
|
||||
}
|
||||
|
||||
impl<C> RefreshScheduler<C>
|
||||
where
|
||||
C: TokenRefresher,
|
||||
{
|
||||
pub fn new(policy: RefreshPolicy, client: C) -> Self {
|
||||
let concurrency = policy.refresh_concurrency.max(1) as usize;
|
||||
Self {
|
||||
policy,
|
||||
client: Arc::new(client),
|
||||
semaphore: Arc::new(Semaphore::new(concurrency)),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn refresh_account_at(
|
||||
&self,
|
||||
account: &Account,
|
||||
trigger: RefreshTrigger,
|
||||
now: DateTime<Utc>,
|
||||
) -> Result<Account, RefreshError> {
|
||||
if !self.should_refresh(account, trigger, now) {
|
||||
return Ok(account.clone());
|
||||
}
|
||||
let Some(refresh_token) = account.refresh_token.as_deref() else {
|
||||
let mut expired = account.clone();
|
||||
expired.status = AccountStatus::Expired;
|
||||
return Ok(expired);
|
||||
};
|
||||
|
||||
let _permit = self
|
||||
.semaphore
|
||||
.acquire()
|
||||
.await
|
||||
.map_err(|_| RefreshError::ConcurrencyClosed)?;
|
||||
// 中文注释:刷新响应未返回 refresh_token 时必须保留旧值,避免把账号永久刷新能力清空。
|
||||
match self.client.refresh(refresh_token).await {
|
||||
Ok(token_pair) => Ok(apply_token_pair(account, token_pair)),
|
||||
Err(RefreshFailure::Transport) => Err(RefreshError::Transport),
|
||||
Err(error) => Ok(apply_refresh_failure(account, error)),
|
||||
}
|
||||
}
|
||||
|
||||
fn should_refresh(
|
||||
&self,
|
||||
account: &Account,
|
||||
trigger: RefreshTrigger,
|
||||
now: DateTime<Utc>,
|
||||
) -> bool {
|
||||
if account.status != AccountStatus::Active {
|
||||
return false;
|
||||
}
|
||||
match trigger {
|
||||
RefreshTrigger::Unauthorized => true,
|
||||
RefreshTrigger::BeforeExpiry => account
|
||||
.access_token_expires_at
|
||||
.is_some_and(|expires_at| expires_at <= now + self.refresh_margin()),
|
||||
}
|
||||
}
|
||||
|
||||
fn refresh_margin(&self) -> Duration {
|
||||
let seconds = self.policy.refresh_margin_seconds.min(86_400 * 7) as i64;
|
||||
Duration::seconds(seconds)
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_token_pair(account: &Account, token_pair: TokenPair) -> Account {
|
||||
let mut refreshed = account.clone();
|
||||
refreshed.access_token = token_pair.access_token;
|
||||
if let Some(refresh_token) = token_pair.refresh_token {
|
||||
refreshed.refresh_token = Some(refresh_token);
|
||||
}
|
||||
refreshed.status = AccountStatus::Active;
|
||||
refreshed
|
||||
}
|
||||
|
||||
fn apply_refresh_failure(account: &Account, failure: RefreshFailure) -> Account {
|
||||
let mut updated = account.clone();
|
||||
updated.status = match failure {
|
||||
RefreshFailure::InvalidGrant => AccountStatus::Expired,
|
||||
RefreshFailure::QuotaExhausted => AccountStatus::QuotaExhausted,
|
||||
RefreshFailure::Banned => AccountStatus::Banned,
|
||||
RefreshFailure::Disabled => AccountStatus::Disabled,
|
||||
RefreshFailure::Transport => AccountStatus::Active,
|
||||
};
|
||||
updated
|
||||
}
|
||||
|
||||
@@ -0,0 +1,163 @@
|
||||
use std::{
|
||||
sync::{
|
||||
atomic::{AtomicUsize, Ordering},
|
||||
Arc,
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{Duration as ChronoDuration, Utc};
|
||||
use tokio::time::sleep;
|
||||
|
||||
use codex_proxy_rs::{
|
||||
accounts::model::{Account, AccountStatus},
|
||||
auth::{
|
||||
refresh::{
|
||||
RefreshFailure, RefreshPolicy, RefreshScheduler, RefreshTrigger, TokenRefresher,
|
||||
},
|
||||
token::TokenPair,
|
||||
},
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_scheduler_refreshes_before_expiry_and_preserves_refresh_token() {
|
||||
let now = Utc::now();
|
||||
let mut account = Account::test("acct_1", AccountStatus::Active);
|
||||
account.access_token_expires_at = Some(now + ChronoDuration::seconds(60));
|
||||
account.refresh_token = Some("rt_keep".to_string());
|
||||
let scheduler = RefreshScheduler::new(
|
||||
RefreshPolicy {
|
||||
refresh_margin_seconds: 300,
|
||||
refresh_concurrency: 2,
|
||||
},
|
||||
StaticRefreshClient {
|
||||
result: Ok(TokenPair {
|
||||
access_token: "new-access".to_string(),
|
||||
refresh_token: None,
|
||||
}),
|
||||
},
|
||||
);
|
||||
|
||||
let refreshed = scheduler
|
||||
.refresh_account_at(&account, RefreshTrigger::BeforeExpiry, now)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(refreshed.access_token, "new-access");
|
||||
assert_eq!(refreshed.refresh_token.as_deref(), Some("rt_keep"));
|
||||
assert_eq!(refreshed.status, AccountStatus::Active);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_scheduler_refreshes_immediately_after_unauthorized() {
|
||||
let now = Utc::now();
|
||||
let mut account = Account::test("acct_1", AccountStatus::Active);
|
||||
account.access_token_expires_at = Some(now + ChronoDuration::hours(6));
|
||||
let scheduler = RefreshScheduler::new(
|
||||
RefreshPolicy {
|
||||
refresh_margin_seconds: 300,
|
||||
refresh_concurrency: 2,
|
||||
},
|
||||
StaticRefreshClient {
|
||||
result: Ok(TokenPair {
|
||||
access_token: "new-access-after-401".to_string(),
|
||||
refresh_token: Some("rt_rotated".to_string()),
|
||||
}),
|
||||
},
|
||||
);
|
||||
|
||||
let refreshed = scheduler
|
||||
.refresh_account_at(&account, RefreshTrigger::Unauthorized, now)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(refreshed.access_token, "new-access-after-401");
|
||||
assert_eq!(refreshed.refresh_token.as_deref(), Some("rt_rotated"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_scheduler_maps_refresh_failures_to_account_status() {
|
||||
let now = Utc::now();
|
||||
let account = Account::test("acct_1", AccountStatus::Active);
|
||||
let scheduler = RefreshScheduler::new(
|
||||
RefreshPolicy {
|
||||
refresh_margin_seconds: 300,
|
||||
refresh_concurrency: 2,
|
||||
},
|
||||
StaticRefreshClient {
|
||||
result: Err(RefreshFailure::Banned),
|
||||
},
|
||||
);
|
||||
|
||||
let refreshed = scheduler
|
||||
.refresh_account_at(&account, RefreshTrigger::Unauthorized, now)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(refreshed.status, AccountStatus::Banned);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_scheduler_limits_refresh_concurrency() {
|
||||
let now = Utc::now();
|
||||
let client = CountingRefreshClient::default();
|
||||
let scheduler = RefreshScheduler::new(
|
||||
RefreshPolicy {
|
||||
refresh_margin_seconds: 300,
|
||||
refresh_concurrency: 2,
|
||||
},
|
||||
client.clone(),
|
||||
);
|
||||
|
||||
let mut handles = Vec::new();
|
||||
for index in 0..5 {
|
||||
let mut account = Account::test(&format!("acct_{index}"), AccountStatus::Active);
|
||||
account.access_token_expires_at = Some(now + ChronoDuration::seconds(1));
|
||||
let scheduler = scheduler.clone();
|
||||
handles.push(tokio::spawn(async move {
|
||||
scheduler
|
||||
.refresh_account_at(&account, RefreshTrigger::BeforeExpiry, now)
|
||||
.await
|
||||
.unwrap()
|
||||
}));
|
||||
}
|
||||
|
||||
for handle in handles {
|
||||
let refreshed = handle.await.unwrap();
|
||||
assert_eq!(refreshed.status, AccountStatus::Active);
|
||||
}
|
||||
assert_eq!(client.max_seen.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct StaticRefreshClient {
|
||||
result: Result<TokenPair, RefreshFailure>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl TokenRefresher for StaticRefreshClient {
|
||||
async fn refresh(&self, _refresh_token: &str) -> Result<TokenPair, RefreshFailure> {
|
||||
self.result.clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
struct CountingRefreshClient {
|
||||
active: Arc<AtomicUsize>,
|
||||
max_seen: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl TokenRefresher for CountingRefreshClient {
|
||||
async fn refresh(&self, refresh_token: &str) -> Result<TokenPair, RefreshFailure> {
|
||||
let current = self.active.fetch_add(1, Ordering::SeqCst) + 1;
|
||||
self.max_seen.fetch_max(current, Ordering::SeqCst);
|
||||
sleep(Duration::from_millis(25)).await;
|
||||
self.active.fetch_sub(1, Ordering::SeqCst);
|
||||
Ok(TokenPair {
|
||||
access_token: format!("access-{refresh_token}"),
|
||||
refresh_token: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user