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

fix(accounts): refresh OAuth tokens while scheduling is disabled (#291)

Browse Source 
Separate scheduling availability from OAuth candidate selection and credential
state writes across OpenAI, xAI, and the shared account store. Cover disabled
refresh, terminal failures, and in-flight scheduling changes with regression tests.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-09-26 18:32:48 +08:00
committed by GitHub
co-authored by Codex
parent 2d20f03bcb
commit 02177af921
9 changed files with 272 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-core/src/account/model.rs
+1 -1
View File
@@ -1053,6 +1053,7 @@ pub struct LoadedCredential {
/// Provider 已计算好时间边界的有界 OAuth refresh 候选查询。
///
/// Store 只负责按持久事实筛选和稳定排序，不拥有提前量或恢复窗口语义。
/// 调度启停不影响凭据续期；停用账号仍按 Provider 的凭据状态与时间边界刷新。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderRefreshQuery {
provider: ProviderKind,
@@ -1117,7 +1118,6 @@ impl ProviderRefreshQuery {
#[must_use]
pub fn contains(&self, account: &ProviderAccount) -> bool {
account.provider() == &self.provider
&& account.enabled()
&& account.has_refresh_token()
&& matches!(
account.credential_state(),
backend/crates/gateway-store/src/postgres/provider_accounts/core_adapter.rs
+4 -4
View File
@@ -187,16 +187,16 @@ impl ProviderAccountStore for PgProviderAccountRepository {
has_refresh_token = $7, access_token_expires_at = $8,
next_refresh_at = $9,
credential_state = case
when $10::text is not null and enabled
when $10::text is not null
then $10 else credential_state end,
credential_observed_at = case
when $10::text is not null and enabled
when $10::text is not null
then $11 else credential_observed_at end,
last_error_reason = case
when $10::text is not null and enabled
when $10::text is not null
then $12 else last_error_reason end,
last_error_message = case
when $10::text is not null and enabled
when $10::text is not null
then $13 else last_error_message end,
updated_at = greatest(now(), updated_at, $11)
where id = $1 and credential_revision = $2
backend/crates/gateway-store/src/postgres/provider_accounts/repository.rs
+9 -21
View File
@@ -295,23 +295,11 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
update.validate()?;
let result = sqlx::query(
"update provider_accounts
set credential_state = case
when enabled then $3
else credential_state
end,
credential_observed_at = case
when enabled then $4
else credential_observed_at
end,
last_error_reason = case
when enabled then $5
else last_error_reason
end,
last_error_message = case
when enabled then $6
else last_error_message
end,
updated_at = case when enabled then greatest(now(), updated_at, $4) else updated_at end
set credential_state = $3,
credential_observed_at = $4,
last_error_reason = $5,
last_error_message = $6,
updated_at = greatest(now(), updated_at, $4)
where id = $1 and credential_revision = $2
and (credential_observed_at is null or credential_observed_at <= $4)",
)
@@ -834,16 +822,16 @@ pub(crate) async fn rotate_provider_account_in_transaction(
upstream_user_id = case when $11::boolean then $12::text else upstream_user_id end,
upstream_account_id = case when $11::boolean then $13::text else upstream_account_id end,
credential_state = case
when $15::boolean or not enabled then credential_state
when $15::boolean then credential_state
when $11::boolean or credential_state <> 'unknown' then 'ready'
else 'unknown'
end,
credential_observed_at = case
when $15::boolean or not enabled then credential_observed_at
when $15::boolean then credential_observed_at
else now()
end,
last_error_reason = case when enabled and not $15::boolean then null else last_error_reason end,
last_error_message = case when enabled and not $15::boolean then null else last_error_message end,
last_error_reason = case when not $15::boolean then null else last_error_reason end,
last_error_message = case when not $15::boolean then null else last_error_message end,
updated_at = greatest(now(), updated_at)
where id = $1 and provider_kind = $2
and credential_revision = $3
backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs
-1
View File
@@ -425,7 +425,6 @@ pub(crate) const REFRESH_CANDIDATES_SELECT: &str = "select auto_location, detect
left join (select id as location_proxy_id, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location
on outbound_proxy_id = location_proxy_id
where provider_kind = $1
and enabled
and has_refresh_token
and credential_state in ('unknown', 'ready')
and access_token_expires_at is not null
backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs
+103 -9
View File
@@ -114,11 +114,16 @@ async fn refresh_candidates_should_filter_order_and_bound_in_one_query() {
disabled.has_refresh_token = true;
disabled.access_token_expires_at = Some(now);
disabled.enabled = false;
let mut expired = account("acct_refresh_expired", "user-expired");
expired.enabled = false;
expired.has_refresh_token = true;
expired.access_token_expires_at = Some(now - TimeDelta::hours(4));
expired.credential_state = CredentialState::Expired;
let mut other_provider = account("acct_refresh_other", "user-other");
other_provider.provider_kind = "xai".to_owned();
other_provider.has_refresh_token = true;
other_provider.access_token_expires_at = Some(now);
for candidate in [forced, due, retry_later, disabled, other_provider] {
for candidate in [forced, due, retry_later, disabled, expired, other_provider] {
repository
.insert_provider_account(candidate)
.await
@@ -148,7 +153,7 @@ async fn refresh_candidates_should_filter_order_and_bound_in_one_query() {
.iter()
.map(|candidate| candidate.account.id().as_str())
.collect::<Vec<_>>(),
vec!["acct_refresh_forced", "acct_refresh_due"]
vec!["acct_refresh_forced", "acct_refresh_disabled"]
);
assert!(candidates.iter().all(|candidate| {
candidate.credential.expose_to_provider()["access_token"] == "initial-secret"
@@ -2935,7 +2940,83 @@ async fn xai_resettable_usage_limit_state_is_persisted() {
}
#[tokio::test]
async fn disabled_account_preserves_user_state_during_refresh_writes() {
async fn core_refresh_cas_updates_credentials_after_scheduling_is_disabled() {
let Some(database) = TestDatabase::create("disabled_core_refresh").await else {
return;
};
let repository = PgProviderAccountRepository::new(database.pool.clone());
let account_id = ProviderAccountId::new("acct_disabled_core_refresh").unwrap();
repository
.insert_provider_account(account(account_id.as_str(), "user-disabled-core-refresh"))
.await
.unwrap();
let revision = CredentialRevision::new(1).unwrap();
repository
.apply_state_change(AccountStateChange {
account_id: account_id.clone(),
expected_revision: revision,
credential_state: CredentialState::Ready,
observed_at: SystemTime::now(),
error_reason: Some(AccountErrorReason::AccessTokenExpired),
message: Some("previous refresh failure".to_owned()),
})
.await
.unwrap();
let quota = QuotaState::exhausted(QuotaEvidence::UsageLimitReached, SystemTime::now(), None);
repository
.apply_quota_access(QuotaAccessChange {
account_id: account_id.clone(),
expected_revision: revision,
state: quota,
})
.await
.unwrap();
let refreshed = CredentialCasUpdate::new(
account_id.clone(),
revision,
ProviderAccountUpdate {
account_id: account_id.clone(),
name: "unused profile".to_owned(),
email: None,
plan_type: None,
},
plaintext_credential("refreshed-secret"),
true,
Some(SystemTime::now() + Duration::from_secs(3_600)),
None,
)
.unwrap()
.preserving_profile()
.with_account_state(CredentialState::Ready, SystemTime::now(), None, None);
// 模拟刷新在途时停用调度；提交新凭据不能重新启用账号，也不能丢弃刷新结果。
repository.set_enabled(&account_id, false).await.unwrap();
assert!(matches!(
repository
.compare_and_swap_credential(refreshed)
.await
.unwrap(),
CredentialCasOutcome::Updated(_)
));
let current = repository
.load_current_credential(&account_id)
.await
.unwrap();
assert!(!current.account.enabled());
assert_eq!(current.account.name(), account_id.as_str());
assert_eq!(current.account.credential_state(), CredentialState::Ready);
assert_eq!(current.account.last_error_reason(), None);
assert_eq!(current.account.last_error_message(), None);
assert!(current.account.quota().is_exhausted());
assert_eq!(
current.credential.expose_to_provider()["access_token"],
"refreshed-secret"
);
assert_eq!(current.account.revision().get(), 2);
database.close().await;
}
#[tokio::test]
async fn disabled_account_updates_credentials_without_enabling_scheduling() {
let Some(database) = TestDatabase::create("provider_account_disabled_refresh").await else {
return;
};
@@ -2946,7 +3027,7 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
let account_id = ProviderAccountId::new("acct_disabled_refresh").expect("account ID");
let mut disabled = account(account_id.as_str(), "user-disabled-refresh");
disabled.enabled = false;
disabled.credential_state = CredentialState::Expired;
disabled.credential_state = CredentialState::Ready;
repository
.import_provider_accounts(ImportProviderAccounts {
settings: None,
@@ -2964,15 +3045,28 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
repository
.apply_state_change(AccountStateChange {
message: None,
message: Some("refresh token expired".to_owned()),
account_id: account_id.clone(),
expected_revision: CredentialRevision::new(1).expect("credential revision"),
credential_state: CredentialState::Ready,
credential_state: CredentialState::Expired,
observed_at: SystemTime::now(),
error_reason: None,
error_reason: Some(AccountErrorReason::CredentialExpired),
})
.await
.expect("disabled state write is a no-op");
.expect("record rejected credential while scheduling is disabled");
let rejected = repository
.load_current_credential(&account_id)
.await
.unwrap();
assert!(!rejected.account.enabled());
assert_eq!(
rejected.account.credential_state(),
CredentialState::Expired
);
assert_eq!(
rejected.account.last_error_message(),
Some("refresh token expired")
);
repository
.rotate_provider_account(RotateProviderAccount {
settings: None,
@@ -2998,7 +3092,7 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
.await
.expect("load disabled account after refresh writes");
assert!(!current.0);
assert_eq!(current.1, "expired");
assert_eq!(current.1, "ready");
assert_eq!(current.2["access_token"], "disabled-refreshed-secret");
assert_eq!(current.3, 2);
backend/crates/providers/openai/src/credential/repository.rs
-3
View File
@@ -235,9 +235,6 @@ impl CodexCredentialRepository {
error_reason: Option<AccountErrorReason>,
message: Option<String>,
) -> Result<(), CredentialRepositoryError> {
if !account.enabled() {
return Ok(());
}
let credential_state = oauth_account_state(account, credential_state);
let message = message.filter(|value| !value.trim().is_empty());
let error_reason = if credential_state == CredentialState::Ready {
backend/crates/providers/openai/tests/credential/refresh.rs
+88
View File
@@ -241,6 +241,94 @@ async fn seed_refreshable_account(
.await;
}
#[tokio::test]
async fn scheduled_refresh_rotates_tokens_without_enabling_scheduling() {
let store = Arc::new(MemoryAccountStore::default());
let policy = MutableRuntimePolicy::new(Duration::from_secs(5 * 60));
let refresher = SingleUseRefresher::new();
let service = refresh_service(&store, Arc::clone(&refresher), policy);
let account_id = "acct_disabled_refresh";
seed_refreshable_account(&store, account_id, SystemTime::now(), None).await;
let before = store.account(account_id).expect("seeded account");
store
.set_enabled(before.id(), false)
.await
.expect("disable scheduling");
let outcomes = service.refresh_due().await.expect("refresh cycle");
assert!(matches!(
outcomes.as_slice(),
[CodexCredentialRefreshOutcome::Refreshed { account_id: refreshed_id, .. }]
if refreshed_id == account_id
));
assert_eq!(refresher.calls(), 1);
let after = store.account(account_id).expect("refreshed account");
assert!(!after.enabled());
assert_eq!(after.quota(), before.quota());
assert_eq!(
after.status_projection(SystemTime::now(), None).status,
AccountStatus::Disabled
);
let loaded = store
.load_credential(after.id(), after.revision())
.await
.expect("refreshed credential");
let runtime = CodexCredentialCodec::decode(&loaded.credential).expect("runtime credential");
assert_eq!(
runtime
.authentication
.oauth()
.expect("OAuth credential")
.access_token
.expose_secret(),
"refreshed-access-token"
);
}
#[tokio::test]
async fn scheduled_refresh_stops_retrying_rejected_disabled_credentials() {
let store = Arc::new(MemoryAccountStore::default());
let service = CodexCredentialRefreshService::new(
store.repository(),
Arc::new(FailingRefresher {
failure: RefreshFailure::InvalidGrant {
message: None,
upstream: None,
},
}),
Arc::new(RefreshLeases),
Arc::new(RefreshCredentialState),
MutableRuntimePolicy::new(Duration::from_secs(5 * 60)),
);
let account_id = "acct_disabled_rejected";
seed_refreshable_account(&store, account_id, SystemTime::now(), None).await;
let account = store.account(account_id).expect("seeded account");
store
.set_enabled(account.id(), false)
.await
.expect("disable scheduling");
assert!(matches!(
service
.refresh_due()
.await
.expect("refresh cycle")
.as_slice(),
[CodexCredentialRefreshOutcome::Invalidated { .. }]
));
let account = store.account(account_id).expect("rejected account");
assert!(!account.enabled());
assert_eq!(account.credential_state(), CredentialState::Expired);
assert!(
service
.refresh_due()
.await
.expect("next refresh cycle")
.is_empty()
);
}
#[tokio::test]
async fn scheduled_refresh_uses_the_current_margin_without_persisting_a_normal_schedule() {
let store = Arc::new(MemoryAccountStore::default());
backend/crates/providers/xai/tests/credential/refresh.rs
+64
View File
@@ -547,6 +547,70 @@ async fn set_token_deadlines(
));
}
#[tokio::test]
async fn scheduled_refresh_rotates_tokens_without_enabling_scheduling() {
let input = due_input("disabled-refresh");
let id = input.account_id.clone();
let (store, _, _, service) =
fixture(input, [Ok(success_tokens(Some("new-refresh")))], true).await;
let before = store.account(&id).expect("seeded account");
store
.set_enabled(&id, false)
.await
.expect("disable scheduling");
let outcomes = service.refresh_due().await.expect("refresh cycle");
assert!(matches!(
outcomes.as_slice(),
[GrokCredentialRefreshOutcome::Refreshed { account_id, .. }] if account_id == &id
));
let after = store.account(&id).expect("refreshed account");
assert!(!after.enabled());
assert_eq!(after.quota(), before.quota());
assert_eq!(
after.status_projection(SystemTime::now(), None).status,
gateway_core::account::AccountStatus::Disabled
);
let credential = store.credential(&id).expect("refreshed credential");
assert_eq!(credential_object(&credential)["access_token"], "new-access");
assert_eq!(
credential_object(&credential)["refresh_token"],
"new-refresh"
);
}
#[tokio::test]
async fn scheduled_refresh_stops_retrying_rejected_disabled_credentials() {
let input = due_input("disabled-rejected");
let id = input.account_id.clone();
let (store, _, _, service) =
fixture(input, [Err(GrokRefreshFailure::InvalidGrant)], true).await;
store
.set_enabled(&id, false)
.await
.expect("disable scheduling");
assert!(matches!(
service
.refresh_due()
.await
.expect("refresh cycle")
.as_slice(),
[GrokCredentialRefreshOutcome::Invalidated { .. }]
));
let account = store.account(&id).expect("rejected account");
assert!(!account.enabled());
assert_eq!(account.credential_state(), CredentialState::Expired);
assert!(
service
.refresh_due()
.await
.expect("next refresh cycle")
.is_empty()
);
}
#[tokio::test]
async fn successful_refresh_rotates_plaintext_tokens_once() {
let input = due_input("success");
docs/api.md
+3
View File
@@ -835,6 +835,9 @@ OAuth start 使用：
在该窗口的 `resetAt + 2 分钟` 后主动复核。后台每 30 秒检查触发条件；同一重置边界复核后仍未更新时
回到 30 分钟重试，避免旧 reset 持续触发请求。各窗口独立确认恢复，时间到期本身不会直接解除账号
耗尽或将展示用量归零。
- 账号 `enabled` 控制是否参与请求调度，不控制 OAuth 凭据续期。OpenAI 与 xAI 的停用账号仍按各自刷新
策略维护 Token；刷新结果更新凭据状态，但不会启用调度。无 refresh token 或凭据已进入失效、无效、
封禁状态的账号不参加后台自动续期。
- `POST /accounts/recover` 对停用账号只将 `enabled` 改为 `true`，保留已有额度快照、凭据、错误和 Redis
cooldown；启用后仍按这些事实投影状态，不会把已有错误或耗尽改成正常。对已启用账号则执行强制恢复：
清除 Redis cooldown 和已保存的额度/错误，恢复为可调度 credential。两条路径均不访问上游；强制恢复