fix: constrain account token refresh lifecycle

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-07-10 12:58:04 +08:00
co-authored by Codex
parent 89ff4be281
commit 68e76d8138
30 changed files with 531 additions and 524 deletions
+5 -3
View File
@@ -47,7 +47,7 @@ use crate::{
},
},
runtime::state::AppState,
upstream::accounts::model::AccountStatus,
upstream::accounts::{model::AccountStatus, token_refresh::token_refresh_status_eligible},
};
const ACCOUNT_STATS_PAGE_LIMIT: u32 = 200;
@@ -319,6 +319,8 @@ struct AccountQuotaUsageWindow {
impl AccountListStats {
fn data_for(&self, account: AdminAccountMetadata) -> AdminAccountData {
let account_id = account.id.clone();
let token_refreshing = token_refresh_status_eligible(account.status)
&& self.refreshing_account_ids.contains(&account_id);
AdminAccountData::from_parts(
account,
self.usage_by_account.get(&account_id),
@@ -327,7 +329,7 @@ impl AccountListStats {
.get(&account_id)
.cloned()
.unwrap_or_default(),
self.refreshing_account_ids.contains(&account_id),
token_refreshing,
)
}
}
@@ -958,7 +960,7 @@ async fn account_list_stats(
}
fn account_display_status(status: AccountStatus, token_refreshing: bool) -> &'static str {
if token_refreshing && matches!(status, AccountStatus::Active | AccountStatus::Expired) {
if token_refreshing {
"refreshing"
} else {
status.as_str()
@@ -14,7 +14,7 @@ use crate::{
json::{NumberedPage, Page},
time::elapsed_millis_i64,
},
upstream::accounts::store::AccountStore,
upstream::accounts::{model::AccountStatus, store::AccountStore},
};
use super::{
@@ -284,6 +284,13 @@ impl AdminAccountService {
duration_ms: elapsed_millis_i64(started_at),
};
match account.status {
AccountStatus::Active | AccountStatus::QuotaExhausted => {}
AccountStatus::Expired => return Ok(skipped("account expired")),
AccountStatus::Disabled => return Ok(skipped("manually disabled")),
AccountStatus::Banned => return Ok(skipped("account banned")),
}
let Some(refresh_token) = account
.refresh_token
.as_ref()
@@ -291,9 +298,6 @@ impl AdminAccountService {
else {
return Ok(skipped("no refresh token"));
};
if account.status == crate::upstream::accounts::model::AccountStatus::Disabled {
return Ok(skipped("manually disabled"));
}
let lease_now = Utc::now();
let lease_owner = format!(
+137 -95
View File
@@ -6,14 +6,25 @@ use secrecy::ExposeSecret;
use serde_json::{json, Value};
use tokio::sync::mpsc;
use crate::upstream::{
accounts::{model::Account, store::StoredAccount},
models::{info::CodexModelInfo, service::ModelRefreshPlanAccount},
protocol::{
responses::CodexResponsesRequest,
sse::{encode_sse_event, parse_sse_events, SseEvent},
use crate::{
proxy::dispatch::{
errors::{
auth_failure_account_status, is_auth_upstream_error, is_quota_exhausted_upstream_error,
},
responses::sse_failure::sse_failure_account_status,
},
upstream::{
accounts::{
model::{Account, AccountStatus},
store::StoredAccount,
},
models::{info::CodexModelInfo, service::ModelRefreshPlanAccount},
protocol::{
responses::{CodexResponsesRequest, ResponsesSseFailure},
sse::{encode_sse_event, parse_sse_events, SseEvent},
},
transport::{is_banned_upstream_error, CodexClientError, CodexRequestContext},
},
transport::CodexRequestContext,
};
use super::{contracts::AdminAccountError, AdminAccountService};
@@ -27,6 +38,27 @@ pub struct AccountModelOption {
pub label: String,
}
struct AccountTestOutcome {
event: Value,
status: Option<AccountStatus>,
}
impl AccountTestOutcome {
fn success() -> Self {
Self {
event: json!({ "type": "test_complete", "success": true }),
status: Some(AccountStatus::Active),
}
}
fn error(error: impl Into<String>, status: Option<AccountStatus>) -> Self {
Self {
event: json!({ "type": "error", "error": error.into() }),
status,
}
}
}
impl AdminAccountService {
pub async fn account_models(
&self,
@@ -97,6 +129,8 @@ impl AdminAccountService {
.flatten();
let installation_id = self.installation_id.clone();
let codex = self.codex.clone();
let service = self.clone();
let stored_account_id = account.id.clone();
let request_id = uuid::Uuid::new_v4().to_string();
let (tx, rx) = mpsc::channel::<Bytes>(16);
@@ -137,19 +171,20 @@ impl AdminAccountService {
session_id: None,
};
match codex.create_response_stream(&request, context).await {
Ok(response) => process_upstream_test_stream(response.body, tx).await,
Err(error) => {
send_test_event(
&tx,
json!({
"type": "error",
"error": error.to_string()
}),
)
let outcome = match codex.create_response_stream(&request, context).await {
Ok(response) => process_upstream_test_stream(response.body, &tx).await,
Err(error) => AccountTestOutcome::error(
error.to_string(),
account_status_from_client_error(&error),
),
};
if let Some(status) = outcome.status {
service
.apply_connection_test_status(&stored_account_id, status)
.await;
}
}
send_test_event(&tx, outcome.event).await;
});
let stream = futures::stream::unfold(rx, |mut rx| async {
@@ -159,6 +194,50 @@ impl AdminAccountService {
});
Ok(Box::pin(stream))
}
async fn apply_connection_test_status(&self, account_id: &str, status: AccountStatus) {
let current = match self.store.get(account_id).await {
Ok(Some(account)) => account,
Ok(None) => return,
Err(error) => {
tracing::warn!(
account_id,
error = %error,
"failed to inspect account after connection test"
);
return;
}
};
if current.status == AccountStatus::Disabled {
return;
}
match self.store.set_status(account_id, status).await {
Ok(true) => {}
Ok(false) => return,
Err(error) => {
tracing::warn!(
account_id,
status = %status,
error = %error,
"failed to persist account status after connection test"
);
return;
}
}
if matches!(status, AccountStatus::Expired | AccountStatus::Banned) {
if let Err(error) = self.store.set_next_refresh_at(account_id, None).await {
tracing::warn!(
account_id,
error = %error,
"failed to clear token refresh schedule after connection test"
);
}
}
self.sync_account_pool_best_effort(account_id, "connection test")
.await;
}
}
fn account_model_option(model: &CodexModelInfo) -> AccountModelOption {
@@ -239,69 +318,61 @@ fn test_responses_request(model: String) -> CodexResponsesRequest {
async fn process_upstream_test_stream(
mut body: crate::upstream::transport::CodexBackendSseStream,
tx: mpsc::Sender<Bytes>,
) {
tx: &mpsc::Sender<Bytes>,
) -> AccountTestOutcome {
let mut buffer = String::new();
while let Some(chunk) = body.next().await {
let chunk = match chunk {
Ok(bytes) => bytes,
Err(error) => {
send_test_event(&tx, json!({ "type": "error", "error": error.to_string() })).await;
return;
return AccountTestOutcome::error(
error.to_string(),
account_status_from_client_error(&error),
);
}
};
buffer.push_str(&String::from_utf8_lossy(&chunk));
while let Some(frame) = take_sse_frame(&mut buffer) {
if process_sse_frame(&tx, &frame).await {
return;
if let Some(outcome) = process_sse_frame(tx, &frame).await {
return outcome;
}
}
}
if !buffer.trim().is_empty() && process_sse_frame(&tx, &buffer).await {
return;
if !buffer.trim().is_empty() {
if let Some(outcome) = process_sse_frame(tx, &buffer).await {
return outcome;
}
}
send_test_event(
&tx,
json!({
"type": "error",
"error": "Stream ended before response.completed"
}),
)
.await;
AccountTestOutcome::error("Stream ended before response.completed", None)
}
async fn process_sse_frame(tx: &mpsc::Sender<Bytes>, frame: &str) -> bool {
async fn process_sse_frame(tx: &mpsc::Sender<Bytes>, frame: &str) -> Option<AccountTestOutcome> {
let events = match parse_sse_events(frame) {
Ok(events) => events,
Err(error) => {
send_test_event(
tx,
json!({
"type": "error",
"error": error.to_string()
}),
)
.await;
return true;
return Some(AccountTestOutcome::error(error.to_string(), None));
}
};
for event in events {
if process_sse_event(tx, &event).await {
return true;
if let Some(outcome) = process_sse_event(tx, &event).await {
return Some(outcome);
}
}
false
None
}
async fn process_sse_event(tx: &mpsc::Sender<Bytes>, event: &SseEvent) -> bool {
async fn process_sse_event(
tx: &mpsc::Sender<Bytes>,
event: &SseEvent,
) -> Option<AccountTestOutcome> {
let value: Value = match serde_json::from_str(&event.data) {
Ok(value) => value,
Err(_) => return false,
Err(_) => return None,
};
match value.get("type").and_then(Value::as_str) {
Some("response.output_text.delta") => {
@@ -310,57 +381,28 @@ async fn process_sse_event(tx: &mpsc::Sender<Bytes>, event: &SseEvent) -> bool {
send_test_event(tx, json!({ "type": "content", "text": delta })).await;
}
}
false
None
}
Some("response.completed" | "response.done") => {
send_test_event(tx, json!({ "type": "test_complete", "success": true })).await;
true
Some("response.completed" | "response.done") => Some(AccountTestOutcome::success()),
Some(event_name @ ("response.failed" | "error")) => {
let failure = ResponsesSseFailure::from_event(event_name, &value);
let status = sse_failure_account_status(&failure);
Some(AccountTestOutcome::error(failure.message, status))
}
Some("response.failed") => {
send_test_event(
tx,
json!({
"type": "error",
"error": response_failure_message(&value)
}),
)
.await;
true
}
Some("error") => {
send_test_event(
tx,
json!({
"type": "error",
"error": error_event_message(&value)
}),
)
.await;
true
}
_ => false,
_ => None,
}
}
fn response_failure_message(value: &Value) -> String {
value
.pointer("/response/error/message")
.or_else(|| value.pointer("/error/message"))
.or_else(|| value.get("message"))
.and_then(Value::as_str)
.filter(|message| !message.trim().is_empty())
.unwrap_or("OpenAI response failed")
.to_string()
}
fn error_event_message(value: &Value) -> String {
value
.pointer("/error/message")
.or_else(|| value.get("message"))
.and_then(Value::as_str)
.filter(|message| !message.trim().is_empty())
.unwrap_or("Unknown upstream error")
.to_string()
fn account_status_from_client_error(error: &CodexClientError) -> Option<AccountStatus> {
if is_banned_upstream_error(error) {
Some(AccountStatus::Banned)
} else if is_quota_exhausted_upstream_error(error) {
Some(AccountStatus::QuotaExhausted)
} else if is_auth_upstream_error(error) {
Some(auth_failure_account_status(error))
} else {
None
}
}
fn take_sse_frame(buffer: &mut String) -> Option<String> {
+7 -3
View File
@@ -42,7 +42,10 @@ use crate::{
},
},
runtime::state::AppState,
upstream::accounts::model::{Account, AccountStatus},
upstream::accounts::{
model::{Account, AccountStatus},
token_refresh::token_refresh_status_eligible,
},
};
const DASHBOARD_ACCOUNT_USAGE_LIMIT: u32 = 4;
@@ -825,8 +828,9 @@ fn account_pool_summary(
};
for account in accounts {
match account.status {
AccountStatus::Active | AccountStatus::Expired
if refreshing_account_ids.contains(&account.id) =>
status
if token_refresh_status_eligible(status)
&& refreshing_account_ids.contains(&account.id) =>
{
summary.refreshing += 1;
}
+1 -1
View File
@@ -92,7 +92,7 @@ pub struct AuthConfig {
impl Default for AuthConfig {
fn default() -> Self {
Self {
refresh_margin_seconds: 300,
refresh_margin_seconds: 3600,
refresh_enabled: true,
refresh_concurrency: 2,
max_concurrent_per_account: 3,
+1 -1
View File
@@ -42,7 +42,7 @@ impl Default for AdminSettings {
fn default() -> Self {
Self {
model_aliases: BTreeMap::new(),
refresh_margin_seconds: 300,
refresh_margin_seconds: 3600,
refresh_concurrency: 2,
max_concurrent_per_account: 3,
request_interval_ms: 50,
@@ -1,28 +0,0 @@
use std::sync::Arc;
use crate::upstream::{
accounts::{model::AccountStatus, token_refresh::RuntimeTokenRefreshService},
token_client::OpenAiTokenClient,
};
pub(super) fn trigger_refresh_after_auth_failure(
token_refresh: &Arc<RuntimeTokenRefreshService<OpenAiTokenClient>>,
account_id: &str,
status: AccountStatus,
) {
if status != AccountStatus::Expired {
return;
}
let token_refresh = Arc::clone(token_refresh);
let account_id = account_id.to_string();
tokio::spawn(async move {
if let Err(error) = token_refresh.trigger_account_refresh_now(&account_id).await {
tracing::warn!(
account_id = %account_id,
error = %error,
"reactive token refresh after upstream auth failure failed"
);
}
});
}
-1
View File
@@ -1,6 +1,5 @@
//! 请求调度模块(账号选择、回退、恢复、会话亲和性)。
mod auth_recovery;
pub mod cloudflare;
pub mod errors;
mod exhaustion;
+1 -1
View File
@@ -6,6 +6,6 @@ mod event_recording;
mod live_stream;
mod prefetch;
pub(crate) mod service;
mod sse_failure;
pub(crate) mod sse_failure;
mod stream_lifecycle;
mod trace;
@@ -22,7 +22,6 @@ use crate::{
usage_record_service::AdminUsageRecordService,
},
proxy::dispatch::{
auth_recovery::trigger_refresh_after_auth_failure,
cloudflare::{
cloudflare_challenge_error_message, cloudflare_path_block_error_message,
is_cloudflare_challenge_upstream_error, is_cloudflare_path_block_upstream_error,
@@ -56,7 +55,6 @@ use crate::{
upstream::accounts::{
model::{Account, AccountStatus},
pool::{AccountAcquireRequest, RuntimeAccountPoolService},
token_refresh::RuntimeTokenRefreshService,
},
upstream::{
models::service::ModelService,
@@ -67,7 +65,6 @@ use crate::{
CollectedResponse,
},
},
token_client::OpenAiTokenClient,
transport::{
backend_transport_for_response_request, is_banned_upstream_error, CodexBackendClient,
CodexBackendResponse,
@@ -112,7 +109,6 @@ pub struct ResponseDispatchService {
reasoning_replay: Arc<Mutex<ReasoningReplayCache>>,
usage_records: Arc<AdminUsageRecordService>,
ops_errors: Arc<AdminOpsErrorLogService>,
token_refresh: Arc<RuntimeTokenRefreshService<OpenAiTokenClient>>,
installation_id: Option<String>,
cloudflare: CloudflareRecovery,
}
@@ -124,7 +120,6 @@ pub(crate) struct ResponseDispatchServiceParts {
pub session_affinity: Arc<RuntimeSessionAffinityService>,
pub usage_records: Arc<AdminUsageRecordService>,
pub ops_errors: Arc<AdminOpsErrorLogService>,
pub token_refresh: Arc<RuntimeTokenRefreshService<OpenAiTokenClient>>,
pub installation_id: Option<String>,
pub cloudflare: CloudflareRecovery,
}
@@ -182,7 +177,6 @@ impl ResponseDispatchService {
)))),
usage_records: parts.usage_records,
ops_errors: parts.ops_errors,
token_refresh: parts.token_refresh,
installation_id: parts.installation_id,
cloudflare: parts.cloudflare,
}
@@ -713,11 +707,6 @@ impl ResponseDispatchService {
self.account_pool
.set_status(&release_account_id, account_status)
.await;
trigger_refresh_after_auth_failure(
&self.token_refresh,
&release_account_id,
account_status,
);
excluded_account_ids.push(release_account_id);
continue;
}
@@ -769,11 +758,6 @@ impl ResponseDispatchService {
self.account_pool
.set_status(&release_account_id, account_status)
.await;
trigger_refresh_after_auth_failure(
&self.token_refresh,
&release_account_id,
account_status,
);
excluded_account_ids.push(release_account_id);
}
Err(error) if is_cloudflare_challenge_upstream_error(&error) => {
@@ -1244,11 +1228,6 @@ impl ResponseDispatchService {
self.account_pool
.set_status(&release_account_id, account_status)
.await;
trigger_refresh_after_auth_failure(
&self.token_refresh,
&release_account_id,
account_status,
);
excluded_account_ids.push(release_account_id);
continue;
}
@@ -1415,11 +1394,6 @@ impl ResponseDispatchService {
self.account_pool
.set_status(&release_account_id, account_status)
.await;
trigger_refresh_after_auth_failure(
&self.token_refresh,
&release_account_id,
account_status,
);
excluded_account_ids.push(release_account_id);
self.account_pool.release(&account.id).await;
continue;
@@ -1527,11 +1501,6 @@ impl ResponseDispatchService {
self.account_pool
.set_status(&release_account_id, account_status)
.await;
trigger_refresh_after_auth_failure(
&self.token_refresh,
&release_account_id,
account_status,
);
excluded_account_ids.push(release_account_id);
}
Err(error) if is_cloudflare_challenge_upstream_error(&error) => {
@@ -1798,11 +1767,6 @@ impl ResponseDispatchService {
self.account_pool
.set_status(&release_account_id, account_status)
.await;
trigger_refresh_after_auth_failure(
&self.token_refresh,
&release_account_id,
account_status,
);
excluded_account_ids.push(release_account_id);
}
Err(error) if is_cloudflare_challenge_upstream_error(&error) => {
@@ -52,7 +52,7 @@ where
failure.upstream_code.as_deref().is_some_and(code_signal) || message_signal(&failure.message)
}
pub(super) fn is_quota_exhausted_sse_failure(failure: &ResponsesSseFailure) -> bool {
pub(crate) fn is_quota_exhausted_sse_failure(failure: &ResponsesSseFailure) -> bool {
sse_failure_matches_parts(
failure,
|code| matches!(code, "quota_exceeded" | "insufficient_quota"),
@@ -80,7 +80,7 @@ fn is_auth_sse_failure_message(message: &str) -> bool {
|| message.contains("token expired")
}
pub(super) fn is_auth_sse_failure(failure: &ResponsesSseFailure) -> bool {
pub(crate) fn is_auth_sse_failure(failure: &ResponsesSseFailure) -> bool {
sse_failure_matches_parts(
failure,
is_auth_sse_failure_code,
@@ -107,7 +107,7 @@ pub(super) fn client_error_invalid_reasoning_replay(error: &CodexClientError) ->
)
}
pub(super) fn auth_sse_failure_account_status(failure: &ResponsesSseFailure) -> AccountStatus {
pub(crate) fn auth_sse_failure_account_status(failure: &ResponsesSseFailure) -> AccountStatus {
if sse_failure_matches(failure, is_banned_auth_signal) {
AccountStatus::Banned
} else {
@@ -115,6 +115,16 @@ pub(super) fn auth_sse_failure_account_status(failure: &ResponsesSseFailure) ->
}
}
pub(crate) fn sse_failure_account_status(failure: &ResponsesSseFailure) -> Option<AccountStatus> {
if is_quota_exhausted_sse_failure(failure) {
Some(AccountStatus::QuotaExhausted)
} else if is_auth_sse_failure(failure) {
Some(auth_sse_failure_account_status(failure))
} else {
None
}
}
pub(super) fn first_sse_failure(
prefetched: &[u8],
) -> Result<Option<ResponsesSseFailure>, SseError> {
-1
View File
@@ -278,7 +278,6 @@ impl Services {
session_affinity: session_affinity.clone(),
usage_records: usage_records.clone(),
ops_errors: ops_errors.clone(),
token_refresh: token_refresh.clone(),
installation_id: installation_id.clone(),
cloudflare: cloudflare_recovery,
}));
@@ -73,7 +73,6 @@ where
info!(
scheduled = summary.scheduled,
immediate = summary.immediate,
recovery_scheduled = summary.recovery_scheduled,
replaced = summary.replaced,
"token 刷新定时器已调度"
);
@@ -244,15 +244,6 @@ fn jittered_refresh_margin_seconds(account_id: &str, margin_seconds: u64) -> i64
.clamp(0.0, i64::MAX as f64) as i64
}
/// 触发刷新动作的原因。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RefreshTrigger {
/// 在访问令牌即将过期前触发刷新。
BeforeExpiry,
/// 在上游返回未授权后立即刷新。
Unauthorized,
}
/// 上游刷新失败后的领域结果。
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum RefreshFailure {
@@ -315,7 +306,7 @@ where
self.policy.update(policy);
}
/// 在给定时间点按触发原因刷新账号。
/// 在给定时间点按到期策略刷新账号。
///
/// 当账号不需要刷新时返回原账号快照;当刷新失败但属于可映射的领域错误时,
/// 返回更新过状态的账号。
@@ -326,10 +317,9 @@ where
pub async fn refresh_account_at(
&self,
account: &Account,
trigger: RefreshTrigger,
now: DateTime<Utc>,
) -> Result<Account, RefreshError> {
if !self.should_refresh(account, trigger, now) {
if !self.should_refresh_account_at(account, now) {
return Ok(account.clone());
}
@@ -353,29 +343,10 @@ where
}
}
/// 判断账号在给定触发原因下是否需要刷新。
pub fn should_refresh_account_at(
&self,
account: &Account,
trigger: RefreshTrigger,
now: DateTime<Utc>,
) -> bool {
self.should_refresh(account, trigger, now)
}
fn should_refresh(
&self,
account: &Account,
trigger: RefreshTrigger,
now: DateTime<Utc>,
) -> bool {
match trigger {
RefreshTrigger::Unauthorized => can_refresh_after_auth_failure(account.status),
RefreshTrigger::BeforeExpiry => {
can_refresh_before_expiry(account.status)
&& self.should_refresh_before_expiry(account, now)
}
}
/// 判断账号在给定时间点是否需要刷新。
pub fn should_refresh_account_at(&self, account: &Account, now: DateTime<Utc>) -> bool {
token_refresh_status_eligible(account.status)
&& self.should_refresh_before_expiry(account, now)
}
fn should_refresh_before_expiry(&self, account: &Account, now: DateTime<Utc>) -> bool {
@@ -398,27 +369,13 @@ where
}
}
fn can_refresh_after_auth_failure(status: AccountStatus) -> bool {
matches!(
status,
AccountStatus::Active | AccountStatus::Expired | AccountStatus::QuotaExhausted
)
}
fn can_refresh_before_expiry(status: AccountStatus) -> bool {
pub(crate) fn token_refresh_status_eligible(status: AccountStatus) -> bool {
matches!(
status,
AccountStatus::Active | AccountStatus::QuotaExhausted
)
}
fn status_after_successful_token_refresh(status: AccountStatus) -> AccountStatus {
match status {
AccountStatus::Expired => AccountStatus::Active,
status => status,
}
}
/// 将新的 token 对应用到账号快照上。
pub fn apply_token_pair(account: &Account, token_pair: TokenPair) -> Account {
let mut refreshed = account.clone();
@@ -429,7 +386,6 @@ pub fn apply_token_pair(account: &Account, token_pair: TokenPair) -> Account {
refreshed.refresh_token = Some(refresh_token);
}
refreshed.status = status_after_successful_token_refresh(account.status);
refreshed
}
@@ -23,8 +23,8 @@ use crate::upstream::accounts::{
StoredAccount,
},
token_refresh::{
jittered_refresh_at, jwt_expiration, RefreshError, RefreshLeaseStore,
RefreshLeaseStoreError, RefreshPolicy, RefreshScheduler, RefreshTrigger,
jittered_refresh_at, jwt_expiration, token_refresh_status_eligible, RefreshError,
RefreshLeaseStore, RefreshLeaseStoreError, RefreshPolicy, RefreshScheduler,
RuntimeRefreshPolicy, TokenRefresher,
},
};
@@ -55,7 +55,6 @@ where
struct ScheduledRefreshTimer {
scheduled_at: DateTime<Utc>,
trigger: RefreshTrigger,
handle: JoinHandle<()>,
}
@@ -183,50 +182,21 @@ where
scanned: accounts.len(),
..TokenTimerSummary::default()
};
let mut expired_recovery_index = 0;
for account in accounts {
if self.skip_account_without_refresh_token(&account).await? {
if self.clear_refresh_schedule_if_ineligible(&account).await? {
summary.skipped += 1;
continue;
}
if matches!(
account.status,
AccountStatus::Disabled | AccountStatus::Banned
) {
summary.skipped += 1;
continue;
}
let timer = match account.status {
AccountStatus::Expired => {
let delay = match account.next_refresh_at {
Some(next_refresh_at) if next_refresh_at > now => {
duration_until(now, next_refresh_at)
}
_ => Duration::from_secs(30 + expired_recovery_index * 2),
};
expired_recovery_index += 1;
Some((delay, RefreshTrigger::Unauthorized))
}
_ => self
.timer_delay_for_account(&account, now)
.map(|delay| (delay, RefreshTrigger::BeforeExpiry)),
};
let Some((delay, trigger)) = timer else {
let Some(delay) = self.timer_delay_for_account(&account, now) else {
summary.skipped += 1;
continue;
};
if delay.is_zero() && account.status != AccountStatus::Expired {
if delay.is_zero() {
if self.clear_scheduled_timer(&account.id).await {
summary.replaced += 1;
}
let outcome = self
.refresh_scheduled_account(&account.id, trigger, now)
.await?;
let outcome = self.refresh_scheduled_account(&account.id, now).await?;
if matches!(outcome, TokenRefreshOutcome::Skipped) {
summary.skipped += 1;
} else {
@@ -239,7 +209,7 @@ where
self.persist_next_refresh_at_if_changed(&account, Some(scheduled_at))
.await?;
let schedule_result = self
.schedule_account_timer(account.id.clone(), delay, trigger, scheduled_at)
.schedule_account_timer(account.id.clone(), delay, scheduled_at)
.await;
match schedule_result {
ScheduleTimerResult::Unchanged => {
@@ -250,9 +220,7 @@ where
ScheduleTimerResult::Scheduled => {}
}
if account.status == AccountStatus::Expired {
summary.recovery_scheduled += 1;
} else if delay.is_zero() {
if delay.is_zero() {
summary.immediate += 1;
} else {
summary.scheduled += 1;
@@ -267,38 +235,6 @@ where
self.timers.lock().await.len()
}
/// 立即刷新被上游认证失败命中的账号。
pub async fn trigger_account_refresh_now(
&self,
account_id: &str,
) -> TokenRefreshServiceResult<()> {
let Some(account) = self.store.get(account_id).await? else {
return Err(TokenRefreshServiceError::AccountNotFound(
account_id.to_string(),
));
};
if account.refresh_token.is_none() {
self.clear_scheduled_timer(account_id).await;
if account.next_refresh_at.is_some() {
persist_next_refresh_at(&self.store, account_id, None).await?;
}
return Ok(());
}
if matches!(
account.status,
AccountStatus::Disabled | AccountStatus::Banned
) {
return Ok(());
}
self.clear_scheduled_timer(account_id).await;
persist_next_refresh_at(&self.store, account_id, None).await?;
self.refresh_scheduled_account(account_id, RefreshTrigger::Unauthorized, Utc::now())
.await
.map(|_| ())
}
/// 在指定时间点执行一次到期刷新扫描。
pub async fn refresh_due_accounts_once_at(
&self,
@@ -315,7 +251,7 @@ where
};
for account in accounts {
if self.skip_account_without_refresh_token(&account).await? {
if self.clear_refresh_schedule_if_ineligible(&account).await? {
summary.skipped += 1;
continue;
}
@@ -325,10 +261,10 @@ where
continue;
}
let Some(trigger) = self.refresh_trigger_for_account(&account, now) else {
if !self.scheduler.should_refresh_account_at(&account, now) {
summary.skipped += 1;
continue;
};
}
if !self.try_mark_in_flight(&account.id).await {
summary.skipped += 1;
@@ -349,7 +285,7 @@ where
}
let outcome = self
.refresh_account_with_status_transitions(&account, trigger, now)
.refresh_account_with_status_transitions(&account, now)
.await;
self.release_in_flight(&account.id).await;
@@ -384,11 +320,11 @@ where
(refresh_at - now).to_std().ok()
}
async fn skip_account_without_refresh_token(
async fn clear_refresh_schedule_if_ineligible(
&self,
account: &Account,
) -> TokenRefreshServiceResult<bool> {
if account.refresh_token.is_some() {
if account.refresh_token.is_some() && token_refresh_status_eligible(account.status) {
return Ok(false);
}
@@ -426,15 +362,15 @@ where
&self,
account_id: String,
delay: Duration,
trigger: RefreshTrigger,
scheduled_at: DateTime<Utc>,
) -> Pin<Box<dyn Future<Output = ScheduleTimerResult> + Send + '_>> {
Box::pin(async move {
{
let timers = self.timers.lock().await;
if timers.get(&account_id).is_some_and(|timer| {
timer.scheduled_at == scheduled_at && timer.trigger == trigger
}) {
if timers
.get(&account_id)
.is_some_and(|timer| timer.scheduled_at == scheduled_at)
{
return ScheduleTimerResult::Unchanged;
}
}
@@ -448,7 +384,7 @@ where
}
task.timers.lock().await.remove(&timer_account_id);
if let Err(error) = task
.refresh_scheduled_account(&timer_account_id, trigger, scheduled_at)
.refresh_scheduled_account(&timer_account_id, scheduled_at)
.await
{
warn!(
@@ -463,7 +399,6 @@ where
account_id,
ScheduledRefreshTimer {
scheduled_at,
trigger,
handle,
},
);
@@ -495,7 +430,6 @@ where
async fn refresh_scheduled_account(
&self,
account_id: &str,
trigger: RefreshTrigger,
now: DateTime<Utc>,
) -> TokenRefreshServiceResult<TokenRefreshOutcome> {
let Some(account) = self.store.get(account_id).await? else {
@@ -504,7 +438,7 @@ where
));
};
let account = stored_account_to_refresh_account(account);
if self.skip_account_without_refresh_token(&account).await? {
if self.clear_refresh_schedule_if_ineligible(&account).await? {
return Ok(TokenRefreshOutcome::Skipped);
}
@@ -529,7 +463,7 @@ where
}
let outcome = self
.refresh_account_with_status_transitions(&account, trigger, now)
.refresh_account_with_status_transitions(&account, now)
.await;
self.release_in_flight(&account.id).await;
@@ -563,40 +497,20 @@ where
}
let scheduled_at = scheduled_at_from_delay(now, delay);
self.schedule_account_timer(
account.id.clone(),
delay,
RefreshTrigger::BeforeExpiry,
scheduled_at,
)
.await;
}
fn refresh_trigger_for_account(
&self,
account: &Account,
now: DateTime<Utc>,
) -> Option<RefreshTrigger> {
if account.status == AccountStatus::Expired {
return Some(RefreshTrigger::Unauthorized);
}
self.scheduler
.should_refresh_account_at(account, RefreshTrigger::BeforeExpiry, now)
.then_some(RefreshTrigger::BeforeExpiry)
self.schedule_account_timer(account.id.clone(), delay, scheduled_at)
.await;
}
async fn refresh_account_with_status_transitions(
&self,
account: &Account,
trigger: RefreshTrigger,
now: DateTime<Utc>,
) -> TokenRefreshServiceResult<TokenRefreshOutcome> {
let max_attempts = self.retry_delays.len() + 1;
let mut permanent_failures = 0;
for attempt_index in 0..max_attempts {
let Some(attempt) = self.prepare_refresh_attempt(account, trigger, now).await? else {
let Some(attempt) = self.prepare_refresh_attempt(account, now).await? else {
return Ok(TokenRefreshOutcome::Skipped);
};
@@ -606,11 +520,7 @@ where
return Ok(TokenRefreshOutcome::StatusUpdated);
}
match self
.scheduler
.refresh_account_at(&attempt, trigger, now)
.await
{
match self.scheduler.refresh_account_at(&attempt, now).await {
Ok(updated) if is_permanent_refresh_failure_status(updated.status) => {
permanent_failures += 1;
if permanent_failures < PERMANENT_FAILURE_CONFIRMATION_THRESHOLD
@@ -631,7 +541,6 @@ where
persist_token_update(&self.store, &updated).await?;
info!(
account_id = %updated.id,
trigger = ?trigger,
attempt = attempt_index + 1,
status = %updated.status,
access_token_expires_at = ?updated.access_token_expires_at,
@@ -645,12 +554,7 @@ where
Err(RefreshError::RetryableTransport | RefreshError::Transport)
if self.sleep_before_retry(&attempt.id, attempt_index).await => {}
Err(RefreshError::RetryableTransport | RefreshError::Transport) => {
persist_status(
&self.store,
&attempt.id,
status_after_temporary_refresh_failure(attempt.status),
)
.await?;
persist_status(&self.store, &attempt.id, attempt.status).await?;
let recovery_delay = stable_jittered_duration(
&attempt.id,
Duration::from_secs(RECOVERY_DELAY_SECONDS as u64),
@@ -672,7 +576,6 @@ where
async fn prepare_refresh_attempt(
&self,
account: &Account,
trigger: RefreshTrigger,
now: DateTime<Utc>,
) -> TokenRefreshServiceResult<Option<Account>> {
let Some(stored) = self.store.get(&account.id).await? else {
@@ -681,19 +584,14 @@ where
));
};
let attempt = stored_account_to_refresh_account(stored);
if matches!(
attempt.status,
AccountStatus::Disabled | AccountStatus::Banned
) || attempt.next_refresh_at.is_some_and(|value| value > now)
if !token_refresh_status_eligible(attempt.status)
|| attempt.next_refresh_at.is_some_and(|value| value > now)
|| attempt.refresh_token != account.refresh_token
{
return Ok(None);
}
if !self
.scheduler
.should_refresh_account_at(&attempt, trigger, now)
{
if !self.scheduler.should_refresh_account_at(&attempt, now) {
return Ok(None);
}
@@ -787,8 +685,6 @@ pub struct TokenTimerSummary {
pub scheduled: usize,
/// 已调度的立即刷新定时器数。
pub immediate: usize,
/// 已调度的过期账号恢复定时器数。
pub recovery_scheduled: usize,
/// 无需调度的账号数。
pub skipped: usize,
/// 被替换的既有定时器数。
@@ -798,7 +694,7 @@ pub struct TokenTimerSummary {
impl TokenTimerSummary {
/// 返回本轮调度发生变化的定时器数量。
pub fn changed(&self) -> usize {
self.scheduled + self.immediate + self.recovery_scheduled
self.scheduled + self.immediate
}
}
@@ -894,14 +790,6 @@ fn is_permanent_refresh_failure_status(status: AccountStatus) -> bool {
matches!(status, AccountStatus::Expired | AccountStatus::Banned)
}
fn status_after_temporary_refresh_failure(status: AccountStatus) -> AccountStatus {
match status {
AccountStatus::Expired | AccountStatus::QuotaExhausted => status,
AccountStatus::Active => AccountStatus::Active,
AccountStatus::Disabled | AccountStatus::Banned => status,
}
}
fn refresh_lease_owner() -> String {
format!("runtime-token-refresh:{}", Uuid::new_v4().simple())
}
+1 -1
View File
@@ -200,7 +200,7 @@ pub struct ResponsesSseFailure {
}
impl ResponsesSseFailure {
fn from_event(event: &str, value: &Value) -> Self {
pub(crate) fn from_event(event: &str, value: &Value) -> Self {
Self {
event: event.to_string(),
message: failure_message(value).unwrap_or_else(|| "Codex upstream SSE failed".into()),
+81
View File
@@ -473,6 +473,87 @@ async fn admin_account_refresh_should_skip_disabled_account_without_consuming_re
);
}
#[tokio::test]
async fn admin_account_refresh_should_skip_expired_account_without_consuming_refresh_token() {
let server = wiremock::MockServer::start().await;
wiremock::Mock::given(wiremock::matchers::method("POST"))
.and(wiremock::matchers::path("/oauth/token"))
.respond_with(wiremock::ResponseTemplate::new(200).set_body_json(json!({
"access_token": test_jwt(
"chatgpt-expired-refresh",
Some("user-expired-refresh"),
Some("expired-refresh@example.com"),
Some("plus"),
),
"refresh_token": "refresh-expired-rotated"
})))
.expect(0)
.mount(&server)
.await;
let (app, _state, pool, _dir) = admin_accounts_test_app_with_oauth_token_endpoint(
"admin-account-refresh-expired-skip.sqlite",
124,
format!("{}/oauth/token", server.uri()),
)
.await;
seed_account(
&pool,
NewAccount {
id: "acct_refresh_expired".to_string(),
email: Some("expired-refresh@example.com".to_string()),
account_id: Some("chatgpt-expired-refresh".to_string()),
user_id: Some("user-expired-refresh".to_string()),
label: None,
plan_type: None,
access_token: SecretString::new(
test_jwt(
"chatgpt-expired-refresh",
Some("user-expired-refresh"),
Some("expired-refresh@example.com"),
Some("plus"),
)
.into(),
),
refresh_token: Some(SecretString::new("refresh-expired".to_string().into())),
access_token_expires_at: None,
status: AccountStatus::Expired,
added_at: None,
},
)
.await;
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/api/admin/accounts/refresh")
.header("content-type", "application/json")
.header("cookie", "cpr_admin_session=session_1")
.body(Body::from(
json!({"id": "acct_refresh_expired"}).to_string(),
))
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let body = response_json(response).await;
let stored = SqliteAccountStore::new(pool)
.get("acct_refresh_expired")
.await
.unwrap()
.unwrap();
assert_eq!(status, StatusCode::OK);
assert_eq!(body["data"]["result"], "skipped");
assert_eq!(body["data"]["error"], "account expired");
assert_eq!(stored.status, AccountStatus::Expired);
assert_eq!(
stored.refresh_token.unwrap().expose_secret(),
"refresh-expired"
);
}
#[tokio::test]
async fn admin_account_refresh_should_skip_account_without_refresh_token() {
let (app, _state, pool, _dir) =
+45
View File
@@ -118,6 +118,51 @@ async fn admin_accounts_list_should_derive_refreshing_display_status_from_active
assert_eq!(item["tokenRefreshing"], true);
}
#[tokio::test]
async fn admin_accounts_list_should_ignore_refresh_lease_for_expired_account() {
let (app, _state, pool, _dir) =
admin_accounts_test_app("admin-accounts-expired-lease.sqlite", 72).await;
sqlx::query("insert into accounts (id, email, access_token, refresh_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?)")
.bind("acct_expired")
.bind("expired@example.com")
.bind("access-token")
.bind("refresh-token")
.bind("expired")
.bind("2026-07-09T00:00:00Z")
.bind("2026-07-09T00:00:00Z")
.execute(&pool)
.await
.unwrap();
sqlx::query("insert into account_refresh_leases (account_id, owner, expires_at, updated_at) values (?, ?, ?, ?)")
.bind("acct_expired")
.bind("stale-owner")
.bind("2099-01-01T00:00:00Z")
.bind("2026-07-09T00:00:00Z")
.execute(&pool)
.await
.unwrap();
let response = app
.oneshot(
Request::builder()
.method("GET")
.uri("/api/admin/accounts?page=1&pageSize=10")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_accounts_expired_lease")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let body = response_json(response).await;
let item = &body["data"]["items"][0];
assert_eq!(item["status"], "expired");
assert_eq!(item["displayStatus"], "expired");
assert_eq!(item["tokenRefreshing"], false);
}
#[tokio::test]
async fn admin_accounts_list_should_include_usage_quota_and_model_stats() {
let (app, _state, pool, _dir) =
+121
View File
@@ -320,6 +320,11 @@ async fn account_test_stream_should_translate_upstream_responses_sse() {
)
.await;
seed_test_account(&pool).await;
let store = codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool.clone());
store
.set_status("acct_test", AccountStatus::QuotaExhausted)
.await
.expect("pre-test status should persist");
let response = app
.oneshot(
@@ -370,4 +375,120 @@ async fn account_test_stream_should_translate_upstream_responses_sse() {
assert_eq!(upstream_body["model"], "gpt-5.5");
assert_eq!(upstream_body["stream"], true);
assert_eq!(upstream_body["store"], false);
assert_eq!(
store
.get("acct_test")
.await
.unwrap()
.expect("account should exist")
.status,
AccountStatus::Active
);
}
#[tokio::test]
async fn account_test_stream_should_mark_expired_after_auth_failure() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/backend-api/codex/responses"))
.respond_with(ResponseTemplate::new(401).set_body_json(json!({
"error": {
"code": "token_expired",
"message": "token expired"
}
})))
.mount(&server)
.await;
let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url(
"admin-account-test-expired.sqlite",
95,
format!("{}/backend-api", server.uri()),
)
.await;
seed_test_account(&pool).await;
let response = app
.oneshot(
Request::builder()
.method("GET")
.uri("/api/admin/accounts/test?id=acct_test&modelId=gpt-5.5")
.header("cookie", "cpr_admin_session=session_1")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let body = String::from_utf8(
to_bytes(response.into_body(), usize::MAX)
.await
.unwrap()
.to_vec(),
)
.unwrap();
let events = test_events(&body);
let store = codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool);
let stored = store.get("acct_test").await.unwrap().unwrap();
assert_eq!(events.last().unwrap()["type"], "error");
assert_eq!(stored.status, AccountStatus::Expired);
assert!(stored.next_refresh_at.is_none());
}
#[tokio::test]
async fn account_test_stream_should_mark_quota_exhausted_after_failed_sse() {
let server = MockServer::start().await;
let upstream_sse = encode_sse_event(
"",
&json!({
"type": "response.failed",
"response": {
"error": {
"code": "quota_exceeded",
"message": "quota exhausted"
}
}
})
.to_string(),
);
Mock::given(method("POST"))
.and(path("/backend-api/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(upstream_sse),
)
.mount(&server)
.await;
let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url(
"admin-account-test-quota.sqlite",
96,
format!("{}/backend-api", server.uri()),
)
.await;
seed_test_account(&pool).await;
let response = app
.oneshot(
Request::builder()
.method("GET")
.uri("/api/admin/accounts/test?id=acct_test&modelId=gpt-5.5")
.header("cookie", "cpr_admin_session=session_1")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let body = String::from_utf8(
to_bytes(response.into_body(), usize::MAX)
.await
.unwrap()
.to_vec(),
)
.unwrap();
let events = test_events(&body);
let store = codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool);
let stored = store.get("acct_test").await.unwrap().unwrap();
assert_eq!(events.last().unwrap()["type"], "error");
assert_eq!(stored.status, AccountStatus::QuotaExhausted);
}
+1 -1
View File
@@ -81,7 +81,7 @@ fn default_config_keeps_only_codex_backend() {
assert_eq!(cfg.server.host, "0.0.0.0");
assert_eq!(cfg.api.base_url, "https://chatgpt.com/backend-api");
assert!(cfg.model_aliases.is_empty());
assert_eq!(cfg.auth.refresh_margin_seconds, 300);
assert_eq!(cfg.auth.refresh_margin_seconds, 3600);
assert_eq!(cfg.auth.rotation_strategy, "smart");
assert_eq!(cfg.auth.oauth_client_id, "app_EMoamEEZ73f0CkXaXp7hrann");
assert_eq!(
@@ -69,7 +69,7 @@ async fn token_refresh_task_should_keep_business_status_before_refresher_call()
}
#[tokio::test]
async fn token_refresh_task_should_recover_expired_account_after_restart() {
async fn token_refresh_task_should_skip_expired_account_after_restart() {
let dir = tempfile::tempdir().expect("temp dir");
let db = dir.path().join("token-refresh-recovery.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
@@ -78,7 +78,6 @@ async fn token_refresh_task_should_recover_expired_account_after_restart() {
let store = SqliteAccountStore::new(pool);
let now = Utc.with_ymd_and_hms(2026, 6, 19, 10, 0, 0).unwrap();
let old_access_token = test_jwt((now + Duration::hours(1)).timestamp());
let new_access_token = test_jwt((now + Duration::hours(2)).timestamp());
store
.insert(NewAccount {
id: "acct-expired-recovery".to_string(),
@@ -95,37 +94,30 @@ async fn token_refresh_task_should_recover_expired_account_after_restart() {
})
.await
.expect("account should be inserted");
let refresher = StaticTokenRefresher {
response: Arc::new(Mutex::new(Ok(TokenPair {
access_token: new_access_token.clone(),
refresh_token: None,
}))),
};
let refresher = CountingTokenRefresher::default();
let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new(
store.clone(),
RefreshPolicy {
refresh_margin_seconds: 300,
refresh_concurrency: 1,
},
refresher,
refresher.clone(),
);
let summary = task
.refresh_due_accounts_once_at(now)
.await
.expect("expired account should recover");
.expect("expired account should be skipped");
let stored = store
.get("acct-expired-recovery")
.await
.expect("account should load")
.expect("account should exist");
assert_eq!(summary.refreshed, 1);
assert_eq!(stored.status, AccountStatus::Active);
assert_eq!(
stored.access_token.expose_secret(),
new_access_token.as_str()
);
assert_eq!(summary.skipped, 1);
assert_eq!(refresher.calls.load(Ordering::SeqCst), 0);
assert_eq!(stored.status, AccountStatus::Expired);
assert!(stored.next_refresh_at.is_none());
}
#[tokio::test]
@@ -215,7 +215,7 @@ async fn token_refresh_task_should_schedule_future_per_account_timer_without_ref
}
#[tokio::test]
async fn token_refresh_task_should_refresh_immediately_after_unauthorized_even_with_future_timer() {
async fn token_refresh_task_should_cancel_future_timer_without_refreshing_expired_account() {
let dir = tempfile::tempdir().expect("temp dir");
let db = dir.path().join("token-refresh-unauthorized-now.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
@@ -233,7 +233,7 @@ async fn token_refresh_task_should_refresh_immediately_after_unauthorized_even_w
user_id: Some("user-unauthorized-now".to_string()),
label: None,
plan_type: Some("plus".to_string()),
access_token: SecretString::new(old_access_token.into()),
access_token: SecretString::new(old_access_token.clone().into()),
refresh_token: Some(SecretString::new(
"refresh-unauthorized-now".to_string().into(),
)),
@@ -269,76 +269,23 @@ async fn token_refresh_task_should_refresh_immediately_after_unauthorized_even_w
.await
.expect("expired status should persist"));
task.trigger_account_refresh_now("acct-refresh-unauthorized-now")
task.schedule_account_timers_once_at(now)
.await
.expect("unauthorized refresh should run immediately");
.expect("expired account schedule should be cleared");
let stored = store
.get("acct-refresh-unauthorized-now")
.await
.expect("account should load")
.expect("account should exist");
assert_eq!(refresher.calls.load(Ordering::SeqCst), 1);
assert_eq!(stored.status, AccountStatus::Active);
assert_eq!(
stored.access_token.expose_secret(),
new_access_token.as_str()
);
assert!(stored.next_refresh_at.is_some());
assert_eq!(task.scheduled_timer_count().await, 1);
}
#[tokio::test]
async fn token_refresh_task_should_skip_unauthorized_refresh_without_refresh_token() {
let dir = tempfile::tempdir().expect("temp dir");
let db = dir.path().join("token-refresh-unauthorized-no-rt.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
.await
.expect("sqlite pool");
let store = SqliteAccountStore::new(pool);
let now = Utc::now();
let old_access_token = test_jwt((now + Duration::minutes(10)).timestamp());
store
.insert(NewAccount {
id: "acct-refresh-unauthorized-no-rt".to_string(),
email: Some("unauthorized-no-rt@example.com".to_string()),
account_id: Some("chatgpt-unauthorized-no-rt".to_string()),
user_id: Some("user-unauthorized-no-rt".to_string()),
label: None,
plan_type: Some("plus".to_string()),
access_token: SecretString::new(old_access_token.clone().into()),
refresh_token: None,
added_at: None,
access_token_expires_at: Some(now + Duration::minutes(10)),
status: AccountStatus::Expired,
})
.await
.expect("account should be inserted");
let refresher = CountingTokenRefresher::default();
let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new(
store.clone(),
RefreshPolicy {
refresh_margin_seconds: 300,
refresh_concurrency: 1,
},
refresher.clone(),
);
task.trigger_account_refresh_now("acct-refresh-unauthorized-no-rt")
.await
.expect("unauthorized refresh should no-op without refresh token");
let stored = store
.get("acct-refresh-unauthorized-no-rt")
.await
.expect("account should load")
.expect("account should exist");
assert_eq!(refresher.calls.load(Ordering::SeqCst), 0);
assert_eq!(stored.status, AccountStatus::Expired);
assert_eq!(
stored.access_token.expose_secret(),
old_access_token.as_str()
);
assert!(stored.next_refresh_at.is_none());
assert_eq!(task.scheduled_timer_count().await, 0);
}
#[tokio::test]
@@ -213,9 +213,9 @@ async fn token_refresh_task_should_refresh_quota_exhausted_account_without_clear
}
#[tokio::test]
async fn token_refresh_task_should_not_clear_quota_limit_when_expired_token_refreshes() {
async fn token_refresh_task_should_not_clear_quota_limit_when_exhausted_token_refreshes() {
let dir = tempfile::tempdir().expect("temp dir");
let db = dir.path().join("token-refresh-expired-quota.sqlite");
let db = dir.path().join("token-refresh-exhausted-quota.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
.await
.expect("sqlite pool");
@@ -226,21 +226,21 @@ async fn token_refresh_task_should_not_clear_quota_limit_when_expired_token_refr
let new_access_token = test_jwt(new_expires_at.timestamp());
store
.insert(NewAccount {
id: "acct-expired-quota".to_string(),
email: Some("expired-quota@example.com".to_string()),
account_id: Some("chatgpt-expired-quota".to_string()),
user_id: Some("user-expired-quota".to_string()),
id: "acct-exhausted-quota".to_string(),
email: Some("exhausted-quota@example.com".to_string()),
account_id: Some("chatgpt-exhausted-quota".to_string()),
user_id: Some("user-exhausted-quota".to_string()),
label: None,
plan_type: Some("plus".to_string()),
access_token: SecretString::new(
test_jwt((now - Duration::seconds(30)).timestamp()).into(),
),
refresh_token: Some(SecretString::new(
"refresh-expired-quota".to_string().into(),
"refresh-exhausted-quota".to_string().into(),
)),
added_at: None,
access_token_expires_at: Some(now - Duration::seconds(30)),
status: AccountStatus::Expired,
status: AccountStatus::QuotaExhausted,
})
.await
.expect("account should be inserted");
@@ -248,7 +248,7 @@ async fn token_refresh_task_should_not_clear_quota_limit_when_expired_token_refr
"update accounts set quota_limit_reached = 1, quota_cooldown_until = ? where id = ?",
)
.bind(cooldown_until.to_rfc3339())
.bind("acct-expired-quota")
.bind("acct-exhausted-quota")
.execute(&pool)
.await
.unwrap();
@@ -270,16 +270,16 @@ async fn token_refresh_task_should_not_clear_quota_limit_when_expired_token_refr
let summary = task
.refresh_due_accounts_once_at(now)
.await
.expect("expired quota-limited account should refresh token");
.expect("quota-exhausted account should refresh token");
let stored = store
.get("acct-expired-quota")
.get("acct-exhausted-quota")
.await
.expect("account should load")
.expect("account should exist");
let quota_state: (i64, Option<String>) = sqlx::query_as(
"select quota_limit_reached, quota_cooldown_until from accounts where id = ?",
)
.bind("acct-expired-quota")
.bind("acct-exhausted-quota")
.fetch_one(&pool)
.await
.unwrap();
@@ -370,7 +370,7 @@ async fn refresh_scheduler_should_refresh_before_expiry_and_preserve_refresh_tok
use chrono::{Duration as ChronoDuration, Utc};
use codex_proxy_rs::upstream::accounts::model::AccountStatus;
use codex_proxy_rs::upstream::accounts::token_refresh::{
RefreshFailure, RefreshPolicy, RefreshScheduler, RefreshTrigger, TokenPair, TokenRefresher,
RefreshFailure, RefreshPolicy, RefreshScheduler, TokenPair, TokenRefresher,
};
#[derive(Clone)]
@@ -403,7 +403,7 @@ async fn refresh_scheduler_should_refresh_before_expiry_and_preserve_refresh_tok
);
let refreshed = scheduler
.refresh_account_at(&account, RefreshTrigger::BeforeExpiry, now)
.refresh_account_at(&account, now)
.await
.expect("refresh should succeed");
@@ -1,10 +1,11 @@
<script setup lang="ts">
import { KeyRound, LoaderCircle, MoreHorizontal, Power, RefreshCw, Trash2, Wifi } from '@lucide/vue'
import { computed } from 'vue'
import BaseButton from '@/components/base/BaseButton.vue'
import BasePopover from '@/components/base/BasePopover.vue'
defineProps<{
const props = defineProps<{
account: any
deleting: boolean
refreshing: boolean
@@ -13,6 +14,12 @@ defineProps<{
scheduleLabel: string
}>()
const canRefreshToken = computed(
() =>
props.account.hasRefreshToken &&
(props.account.status === 'active' || props.account.status === 'quota_exhausted'),
)
const emit = defineEmits<{
delete: [account: any]
test: [account: any]
@@ -54,7 +61,7 @@ const emit = defineEmits<{
测试连接
</button>
<button
v-if="account.hasRefreshToken"
v-if="canRefreshToken"
type="button"
class="flex h-8.5 w-full items-center gap-2 rounded-(--cp-input-radius-small) border-0 bg-transparent px-3 text-left text-[13px] leading-none font-[650] text-(--cp-text-primary) transition-colors hover:bg-(--cp-default-bg-hover) disabled:cursor-not-allowed disabled:text-(--cp-disabled-text)"
:disabled="refreshing"
@@ -57,7 +57,11 @@ interface AccountModelsResponse {
}>
}
export function useAccountConnectionTest() {
interface AccountConnectionTestOptions {
onResult?: () => void
}
export function useAccountConnectionTest(options: AccountConnectionTestOptions = {}) {
const showConnectionTestModal = shallowRef(false)
const testingAccount = shallowRef<ConnectionTestAccount | null>(null)
const connectionTestStatus = shallowRef<ConnectionTestStatus>('idle')
@@ -270,6 +274,7 @@ export function useAccountConnectionTest() {
appendConnectionTestLog(connectionTestError.value, 'danger')
finishConnectionTest('error')
}
options.onResult?.()
clearConnectionTestRun()
return
}
@@ -277,6 +282,7 @@ export function useAccountConnectionTest() {
connectionTestError.value = event.error || '测试连接失败'
appendConnectionTestLog(connectionTestError.value, 'danger')
finishConnectionTest('error')
options.onResult?.()
clearConnectionTestRun()
}
}
@@ -1,5 +1,5 @@
import { clamp } from 'es-toolkit'
import { computed, onMounted, onUnmounted, ref, watch, type Ref } from 'vue'
import { computed, onMounted, ref, type Ref } from 'vue'
import dayjs from 'dayjs'
import {
@@ -77,8 +77,6 @@ export function useAccountMutations(options: {
const authorizingOAuth = authorizingOAuthAction.loading
const batchDeleting = batchDeletingAction.loading
const exportingAccounts = exportingAccountsAction.loading
let tokenRefreshPollTimer: ReturnType<typeof window.setInterval> | undefined
const createForm = ref({
mode: 'oauth',
tokenText: '',
@@ -404,44 +402,10 @@ export function useAccountMutations(options: {
return `cpr-accounts-selected-${selectedCount}-${dayjs().format('YYYY-MM-DD')}.json`
}
function hasTokenRefreshingAccount() {
return accounts.value.some((account) => account.tokenRefreshing)
}
function startTokenRefreshPolling() {
if (tokenRefreshPollTimer !== undefined) return
tokenRefreshPollTimer = window.setInterval(() => {
if (!loading.value && hasTokenRefreshingAccount()) {
void loadAccounts()
}
}, 2000)
}
function stopTokenRefreshPolling() {
if (tokenRefreshPollTimer === undefined) return
window.clearInterval(tokenRefreshPollTimer)
tokenRefreshPollTimer = undefined
}
watch(
() => hasTokenRefreshingAccount(),
(hasRefreshing) => {
if (hasRefreshing) {
startTokenRefreshPolling()
} else {
stopTokenRefreshPolling()
}
},
)
onMounted(() => {
loadAccounts()
})
onUnmounted(() => {
stopTokenRefreshPolling()
})
return {
loading,
accounts,
+21 -19
View File
@@ -24,25 +24,6 @@ import AccountStatusBadge from './components/AccountStatusBadge.vue'
import AccountTableActions from './components/AccountTableActions.vue'
import AccountUsagePanel from './components/AccountUsagePanel.vue'
const {
showConnectionTestModal,
testingAccount,
connectionTestStatus,
connectionTestModel,
connectionTestLogs,
connectionTestError,
connectionTestStartedAt,
connectionTestFinishedAt,
connectionTestDurationMs,
testingConnectionIds,
loadingConnectionTestModels,
connectionTestSelectedModel,
connectionTestModelOptions,
connectionTestStatusView,
openConnectionTest,
handleTestConnection,
} = useAccountConnectionTest()
const selectedIds = ref<Set<string>>(new Set())
const totalAccounts = ref(0)
@@ -95,6 +76,27 @@ const {
totalAccounts,
})
const {
showConnectionTestModal,
testingAccount,
connectionTestStatus,
connectionTestModel,
connectionTestLogs,
connectionTestError,
connectionTestStartedAt,
connectionTestFinishedAt,
connectionTestDurationMs,
testingConnectionIds,
loadingConnectionTestModels,
connectionTestSelectedModel,
connectionTestModelOptions,
connectionTestStatusView,
openConnectionTest,
handleTestConnection,
} = useAccountConnectionTest({
onResult: () => void loadAccounts(),
})
const {
expandedAccountIds,
allSelected,
@@ -29,7 +29,10 @@ const requestIntervalMs = defineModel<string>('requestIntervalMs', { required: t
</BaseInput>
</BaseFormItem>
<BaseFormItem label="提前刷新秒数" description="Token 过期前多少秒触发刷新。">
<BaseFormItem
label="提前刷新秒数"
description="Token 过期前多少秒触发刷新,默认 3600 秒(1 小时)。"
>
<BaseInput v-model="refreshMarginSeconds" type="number">
<template #prefix>
<Timer class="size-4" />
+2 -2
View File
@@ -59,7 +59,7 @@ const { copy } = useClipboard()
const form = reactive<SettingsForm>({
modelAliases: {},
refreshMarginSeconds: 300,
refreshMarginSeconds: 3600,
refreshConcurrency: 2,
maxConcurrentPerAccount: 3,
requestIntervalMs: 50,
@@ -145,7 +145,7 @@ function rowsToAliases(rows: AliasRow[]) {
function applySettings(data: any) {
form.modelAliases = data.modelAliases || {}
form.refreshMarginSeconds = Number(data.refreshMarginSeconds ?? 300)
form.refreshMarginSeconds = Number(data.refreshMarginSeconds ?? 3600)
form.refreshConcurrency = Number(data.refreshConcurrency ?? 2)
form.maxConcurrentPerAccount = Number(data.maxConcurrentPerAccount ?? 3)
form.requestIntervalMs = Number(data.requestIntervalMs ?? 50)