From fc388bb799e635a2e6c5ccb6c02c2bf31bba76c6 Mon Sep 17 00:00:00 2001 From: Zyy Date: Mon, 31 Aug 2026 02:21:44 +0800 Subject: [PATCH] refactor(providers): bound credential refresh queries Move refresh candidate filtering, ordering, exclusions, and limits into one PostgreSQL query for OpenAI and xAI. Keep refresh policy in each provider and preserve the final credential revision fence without adding a scheduling cache or catalog table. Co-authored-by: Codex --- .../gateway-core/src/engine/credential.rs | 90 +++++++++++++++++++ .../provider_accounts/core_adapter.rs | 43 ++++++++- .../src/postgres/provider_accounts/mod.rs | 3 +- .../src/postgres/provider_accounts/rows.rs | 23 +++++ .../tests/postgres/provider_accounts.rs | 74 ++++++++++++++- .../openai/src/credential/refresh.rs | 70 +++++---------- .../openai/src/credential/repository.rs | 26 +++++- .../crates/providers/openai/tests/support.rs | 28 +++++- .../providers/xai/src/credential/refresh.rs | 70 ++++++--------- .../xai/src/credential/repository.rs | 20 ++--- backend/crates/providers/xai/tests/support.rs | 26 +++++- 11 files changed, 356 insertions(+), 117 deletions(-) diff --git a/backend/crates/gateway-core/src/engine/credential.rs b/backend/crates/gateway-core/src/engine/credential.rs index 4a3893187..c15ee6fb0 100644 --- a/backend/crates/gateway-core/src/engine/credential.rs +++ b/backend/crates/gateway-core/src/engine/credential.rs @@ -910,6 +910,90 @@ pub struct LoadedCredential { pub credential: PlaintextCredential, } +/// Provider 已计算好时间边界的有界 OAuth refresh 候选查询。 +/// +/// Store 只负责按持久事实筛选和稳定排序，不拥有提前量或恢复窗口语义。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderRefreshQuery { + provider: ProviderKind, + refresh_due_before: SystemTime, + force_due_before: SystemTime, + observed_at: SystemTime, + excluded_account_ids: Vec, + limit: NonZeroU32, +} + +impl ProviderRefreshQuery { + #[must_use] + pub fn new( + provider: ProviderKind, + refresh_due_before: SystemTime, + force_due_before: SystemTime, + observed_at: SystemTime, + excluded_account_ids: Vec, + limit: NonZeroU32, + ) -> Self { + Self { + provider, + refresh_due_before, + force_due_before, + observed_at, + excluded_account_ids, + limit, + } + } + + #[must_use] + pub const fn provider(&self) -> &ProviderKind { + &self.provider + } + + #[must_use] + pub const fn refresh_due_before(&self) -> SystemTime { + self.refresh_due_before + } + + #[must_use] + pub const fn force_due_before(&self) -> SystemTime { + self.force_due_before + } + + #[must_use] + pub const fn observed_at(&self) -> SystemTime { + self.observed_at + } + + #[must_use] + pub fn excluded_account_ids(&self) -> &[ProviderAccountId] { + &self.excluded_account_ids + } + + #[must_use] + pub const fn limit(&self) -> NonZeroU32 { + self.limit + } + + /// 判断内存 Store 中的账号是否满足与 PostgreSQL 相同的候选谓词。 + #[must_use] + pub fn contains(&self, account: &ProviderAccount) -> bool { + account.provider() == &self.provider + && account.enabled() + && account.has_refresh_token() + && matches!( + account.credential_state(), + CredentialState::Unknown | CredentialState::Ready + ) + && !self.excluded_account_ids.contains(account.id()) + && account.access_token_expires_at().is_some_and(|expires_at| { + expires_at <= self.force_due_before + || (expires_at <= self.refresh_due_before + && account + .next_refresh_at() + .is_none_or(|retry_at| retry_at <= self.observed_at)) + }) + } +} + /// Admin/Provider import 创建账号时的一次性明文输入。 #[derive(Clone, PartialEq)] pub struct NewProviderAccount { @@ -1191,6 +1275,12 @@ pub trait ProviderAccountStore: Send + Sync { provider: &ProviderKind, ) -> Result, StoreError>; + /// 一次有界查询返回账号事实和 revision-fenced 明文 credential。 + async fn list_refresh_candidates( + &self, + query: ProviderRefreshQuery, + ) -> Result, StoreError>; + async fn load_credential( &self, account: &ProviderAccountId, diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/core_adapter.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/core_adapter.rs index 10fc6faa1..5ed314108 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/core_adapter.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/core_adapter.rs @@ -2,6 +2,15 @@ use super::*; +fn loaded_credential_from_record( + record: ProviderAccountRecord, +) -> Result { + Ok(LoadedCredential { + account: core_account_from_summary(record.summary)?, + credential: PlaintextCredential::new(record.provider_credentials_json.fields().clone()), + }) +} + #[async_trait] impl ProviderAccountStore for PgProviderAccountRepository { async fn create_account(&self, account: CoreNewProviderAccount) -> Result<(), CoreStoreError> { @@ -72,6 +81,35 @@ impl ProviderAccountStore for PgProviderAccountRepository { .collect() } + async fn list_refresh_candidates( + &self, + query: CoreProviderRefreshQuery, + ) -> Result, CoreStoreError> { + let excluded_account_ids = query + .excluded_account_ids() + .iter() + .map(|account_id| account_id.as_str()) + .collect::>(); + let rows = sqlx::query(REFRESH_CANDIDATES_SELECT) + .bind(query.provider().as_str()) + .bind(DateTime::::from(query.refresh_due_before())) + .bind(DateTime::::from(query.force_due_before())) + .bind(DateTime::::from(query.observed_at())) + .bind(excluded_account_ids) + .bind(i64::from(query.limit().get())) + .fetch_all(&self.pool) + .await + .map_err(|_| CoreStoreError::new(CoreStoreErrorKind::Unavailable))?; + rows.into_iter() + .map(account_record_from_row) + .map(|record| { + record + .map_err(core_store_error) + .and_then(loaded_credential_from_record) + }) + .collect() + } + async fn load_credential( &self, account: &CoreProviderAccountId, @@ -93,10 +131,7 @@ impl ProviderAccountStore for PgProviderAccountRepository { .await .map_err(core_store_error)? .ok_or_else(|| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?; - Ok(LoadedCredential { - account: core_account_from_summary(record.summary)?, - credential: PlaintextCredential::new(record.provider_credentials_json.fields().clone()), - }) + loaded_credential_from_record(record) } async fn compare_and_swap_credential( diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/mod.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/mod.rs index 27b0c31d0..f975b4260 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/mod.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/mod.rs @@ -43,7 +43,8 @@ use gateway_core::engine::credential::{ NewProviderAccount as CoreNewProviderAccount, OpaqueProviderData, PlaintextCredential, ProviderAccount as CoreProviderAccount, ProviderAccountId as CoreProviderAccountId, ProviderAccountIdentity, ProviderAccountStore, - ProviderAccountUpdate as CoreProviderAccountUpdate, QuotaAccessChange, QuotaAccessState, + ProviderAccountUpdate as CoreProviderAccountUpdate, + ProviderRefreshQuery as CoreProviderRefreshQuery, QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaObservation, QuotaObservationTouch, QuotaState, QuotaWriteOutcome, }; use gateway_core::error::{StoreError as CoreStoreError, StoreErrorKind as CoreStoreErrorKind}; diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs index d256dd927..bcee1a3e0 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs @@ -455,6 +455,29 @@ pub(crate) const ACCOUNT_SELECT_BY_IDS: &str = "select id, provider_kind, name, where id = any($1::text[]) and provider_kind = $2 order by id"; +pub(crate) const REFRESH_CANDIDATES_SELECT: &str = "select id, provider_kind, name, email, upstream_user_id, + upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision, + has_refresh_token, access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, credential_state, + provider_quota_json, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at, + last_error_reason, last_error_message, + credential_observed_at, quota_observed_at, created_at, updated_at + from provider_accounts + where provider_kind = $1 + and enabled + and has_refresh_token + and credential_state in ('unknown', 'ready') + and access_token_expires_at is not null + and ( + access_token_expires_at <= $3 + or ( + access_token_expires_at <= $2 + and (next_refresh_at is null or next_refresh_at <= $4) + ) + ) + and not (id = any($5::text[])) + order by access_token_expires_at, id + limit $6"; + pub(crate) fn account_record_from_row( row: sqlx::postgres::PgRow, ) -> StoreResult { diff --git a/backend/crates/gateway-store/tests/postgres/provider_accounts.rs b/backend/crates/gateway-store/tests/postgres/provider_accounts.rs index 1158a3f7b..5a71ed7c1 100644 --- a/backend/crates/gateway-store/tests/postgres/provider_accounts.rs +++ b/backend/crates/gateway-store/tests/postgres/provider_accounts.rs @@ -1,4 +1,7 @@ -use std::time::{Duration, SystemTime}; +use std::{ + num::NonZeroU32, + time::{Duration, SystemTime}, +}; use chrono::{TimeDelta, Utc}; use gateway_admin::{ @@ -23,8 +26,8 @@ use gateway_core::engine::credential::{ AccountErrorReason, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate, CredentialRevision, CredentialState, OpaqueProviderData, PlaintextCredential, ProviderAccountId, ProviderAccountIdentity, ProviderAccountStore, ProviderAccountUpdate, - QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaObservation, QuotaObservationTouch, - QuotaState, QuotaWriteOutcome, + ProviderRefreshQuery, QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaObservation, + QuotaObservationTouch, QuotaState, QuotaWriteOutcome, }; use gateway_core::routing::{AccountGroupId, ProviderKind}; use gateway_store::{ @@ -84,6 +87,71 @@ fn postgres_provider_account_adapter_implements_core_port() { assert_terminal_admin_port::(); } +#[tokio::test] +async fn refresh_candidates_should_filter_order_and_bound_in_one_query() { + let Some(database) = TestDatabase::create("provider_refresh_candidates").await else { + return; + }; + let repository = PgProviderAccountRepository::new(database.pool.clone()); + let now = Utc::now(); + let mut forced = account("acct_refresh_forced", "user-forced"); + forced.has_refresh_token = true; + forced.access_token_expires_at = Some(now - TimeDelta::hours(3)); + forced.next_refresh_at = Some(now + TimeDelta::hours(1)); + let mut due = account("acct_refresh_due", "user-due"); + due.has_refresh_token = true; + due.access_token_expires_at = Some(now); + let mut retry_later = account("acct_refresh_retry_later", "user-retry-later"); + retry_later.has_refresh_token = true; + retry_later.access_token_expires_at = Some(now); + retry_later.next_refresh_at = Some(now + TimeDelta::hours(1)); + let mut disabled = account("acct_refresh_disabled", "user-disabled"); + disabled.has_refresh_token = true; + disabled.access_token_expires_at = Some(now); + disabled.enabled = false; + 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] { + repository + .insert_provider_account(candidate) + .await + .expect("insert refresh candidate fixture"); + } + + let observed_at = SystemTime::from(now); + let query = ProviderRefreshQuery::new( + ProviderKind::new("openai").expect("Provider kind"), + observed_at + .checked_add(Duration::from_secs(5 * 60)) + .expect("refresh boundary"), + observed_at + .checked_sub(Duration::from_secs(2 * 60 * 60)) + .expect("force boundary"), + observed_at, + Vec::new(), + NonZeroU32::new(2).expect("positive limit"), + ); + let candidates = repository + .list_refresh_candidates(query) + .await + .expect("load bounded refresh candidates"); + + assert_eq!( + candidates + .iter() + .map(|candidate| candidate.account.id().as_str()) + .collect::>(), + vec!["acct_refresh_forced", "acct_refresh_due"] + ); + assert!(candidates.iter().all(|candidate| { + candidate.credential.expose_to_provider()["access_token"] == "initial-secret" + })); + + database.close().await; +} + #[tokio::test] async fn provider_account_should_allow_missing_upstream_user_id() { let Some(database) = TestDatabase::create("provider_account_missing_upstream_identity").await diff --git a/backend/crates/providers/openai/src/credential/refresh.rs b/backend/crates/providers/openai/src/credential/refresh.rs index dc285770d..008a4fde7 100644 --- a/backend/crates/providers/openai/src/credential/refresh.rs +++ b/backend/crates/providers/openai/src/credential/refresh.rs @@ -1,13 +1,14 @@ //! Codex AT/RT 刷新状态机；Redis lease + ProviderAccountStore CAS，无 SQL。 use std::collections::BTreeSet; +use std::num::NonZeroU32; use std::sync::Arc; use std::time::{Duration, SystemTime}; use chrono::{DateTime, Utc}; use futures::{StreamExt as _, stream}; use gateway_core::engine::credential::{ - AccountErrorReason, CredentialState, ProviderAccount, ProviderAccountId, + AccountErrorReason, CredentialState, ProviderAccount, ProviderAccountId, ProviderRefreshQuery, }; use gateway_core::provider_ports::{ ProviderCredentialStatePort, ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, @@ -15,6 +16,7 @@ use gateway_core::provider_ports::{ ProviderRuntimePolicyPort, ProviderStoreError, provider_refresh_backoff_at, provider_refresh_retry_at, }; +use gateway_core::routing::ProviderKind; use secrecy::{ExposeSecret, SecretString}; use thiserror::Error; @@ -132,28 +134,6 @@ impl<'a> RefreshFailureContext<'a> { } } -fn refresh_due_at( - account: &ProviderAccount, - policy: ProviderRefreshPolicy, - now: SystemTime, -) -> Option { - if refresh_recovery_window_exhausted(account.access_token_expires_at(), now) { - // 即使 retry-not-before 异常地落在 deadline 之后，也必须进入终态化路径。 - return refresh_recovery_deadline(account.access_token_expires_at()); - } - let access_token_expires_at = account.access_token_expires_at()?; - if !policy.is_refresh_due(access_token_expires_at, now) { - return None; - } - if account - .next_refresh_at() - .is_some_and(|retry_at| retry_at > now) - { - return None; - } - Some(access_token_expires_at) -} - pub struct DueCodexCredential { pub account: ProviderAccount, pub secret: CodexOAuthSecret, @@ -410,31 +390,25 @@ impl CodexCredentialRefreshService { CodexCredentialRefreshError, > { let now = SystemTime::now(); - let mut accounts = self - .repository - .store() - .list_accounts() - .await - .map_err(CredentialRepositoryError::from)?; - accounts.retain(|account| { - account.provider().as_str() == PROVIDER_NAME - && account.enabled() - && account.has_refresh_token() - && !excluded.contains(account.id()) - && matches!( - account.credential_state(), - CredentialState::Unknown | CredentialState::Ready - ) - && refresh_due_at(account, policy, now).is_some() - }); - accounts - .sort_by_key(|account| (refresh_due_at(account, policy, now), account.id().clone())); - accounts.truncate(MAX_REFRESH_BATCH as usize); - let mut due = Vec::with_capacity(accounts.len()); + let provider = ProviderKind::new(PROVIDER_NAME) + .map_err(|_| CredentialRepositoryError::InvalidCredentialData)?; + let limit = NonZeroU32::new(MAX_REFRESH_BATCH) + .ok_or(CredentialRepositoryError::InvalidCredentialData)?; + let query = ProviderRefreshQuery::new( + provider, + now.checked_add(policy.margin()).unwrap_or(now), + now.checked_sub(REFRESH_RECOVERY_WINDOW) + .unwrap_or(SystemTime::UNIX_EPOCH), + now, + excluded.iter().cloned().collect(), + limit, + ); + let candidates = self.repository.list_refresh_candidates(query).await?; + let mut due = Vec::with_capacity(candidates.len()); let mut failures = Vec::new(); - for account in accounts { - let account_id = account.id().to_string(); - match self.repository.load_runtime_credential(&account).await { + for loaded in candidates { + let account_id = loaded.account.id().to_string(); + match self.repository.decode_runtime_credential(&loaded) { Ok(runtime) if runtime .authentication @@ -445,7 +419,7 @@ impl CodexCredentialRefreshService { unreachable!("OAuth runtime authentication was checked above") }; due.push(DueCodexCredential { - account, + account: loaded.account, secret: secret.clone(), }); } diff --git a/backend/crates/providers/openai/src/credential/repository.rs b/backend/crates/providers/openai/src/credential/repository.rs index e3ea1277b..37c07b78b 100644 --- a/backend/crates/providers/openai/src/credential/repository.rs +++ b/backend/crates/providers/openai/src/credential/repository.rs @@ -5,8 +5,8 @@ use std::time::SystemTime; use gateway_core::engine::credential::{ AccountErrorReason, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate, - CredentialRevision, CredentialState, ProviderAccount, ProviderAccountStore, - ProviderAccountUpdate, + CredentialRevision, CredentialState, LoadedCredential, ProviderAccount, ProviderAccountStore, + ProviderAccountUpdate, ProviderRefreshQuery, }; use gateway_core::routing::ProviderKind; use secrecy::ExposeSecret; @@ -129,6 +129,26 @@ impl CodexCredentialRepository { .map_err(Into::into) } + pub async fn list_refresh_candidates( + &self, + query: ProviderRefreshQuery, + ) -> Result, CredentialRepositoryError> { + self.store + .list_refresh_candidates(query) + .await + .map_err(Into::into) + } + + pub fn decode_runtime_credential( + &self, + loaded: &LoadedCredential, + ) -> Result { + if loaded.account.provider().as_str() != PROVIDER_NAME { + return Err(CredentialRepositoryError::InvalidCredentialData); + } + CodexCredentialCodec::decode(&loaded.credential).map_err(Into::into) + } + pub async fn load_runtime_credential( &self, account: &ProviderAccount, @@ -143,7 +163,7 @@ impl CodexCredentialRepository { if loaded.account != *account { return Err(CredentialRepositoryError::RevisionConflict); } - CodexCredentialCodec::decode(&loaded.credential).map_err(Into::into) + self.decode_runtime_credential(&loaded) } pub async fn load_complete_data( diff --git a/backend/crates/providers/openai/tests/support.rs b/backend/crates/providers/openai/tests/support.rs index 2dd1d36dd..4e6626fdf 100644 --- a/backend/crates/providers/openai/tests/support.rs +++ b/backend/crates/providers/openai/tests/support.rs @@ -13,7 +13,8 @@ use gateway_core::engine::credential::{ AccountWeight, CredentialCasOutcome, CredentialCasUpdate, CredentialCasUpdateParts, CredentialRevision, CredentialState, LoadedCredential, NewProviderAccount, OpaqueProviderData, ProviderAccount, ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, - QuotaAccessChange, QuotaObservation, QuotaObservationTouch, QuotaState, QuotaWriteOutcome, + ProviderRefreshQuery, QuotaAccessChange, QuotaObservation, QuotaObservationTouch, QuotaState, + QuotaWriteOutcome, }; use gateway_core::error::{StoreError, StoreErrorKind}; use gateway_core::policy::ClientApiKeyId; @@ -175,6 +176,31 @@ impl ProviderAccountStore for MemoryAccountStore { .collect()) } + async fn list_refresh_candidates( + &self, + query: ProviderRefreshQuery, + ) -> Result, StoreError> { + let mut candidates = self + .accounts + .lock() + .expect("account store lock") + .values() + .filter(|stored| query.contains(&stored.account)) + .map(|stored| LoadedCredential { + account: stored.account.clone(), + credential: stored.credential.clone(), + }) + .collect::>(); + candidates.sort_by_key(|loaded| { + ( + loaded.account.access_token_expires_at(), + loaded.account.id().clone(), + ) + }); + candidates.truncate(query.limit().get() as usize); + Ok(candidates) + } + async fn load_credential( &self, account: &ProviderAccountId, diff --git a/backend/crates/providers/xai/src/credential/refresh.rs b/backend/crates/providers/xai/src/credential/refresh.rs index 1704ff098..f459758c0 100644 --- a/backend/crates/providers/xai/src/credential/refresh.rs +++ b/backend/crates/providers/xai/src/credential/refresh.rs @@ -2,6 +2,7 @@ use std::collections::BTreeSet; use std::fmt; +use std::num::NonZeroU32; use std::sync::Arc; use std::time::{Duration, SystemTime}; @@ -10,6 +11,7 @@ use chrono::{DateTime, Utc}; use futures::{StreamExt as _, stream}; use gateway_core::engine::credential::{ AccountErrorReason, CredentialRevision, CredentialState, LoadedCredential, ProviderAccountId, + ProviderRefreshQuery, }; use gateway_core::provider_ports::{ ProviderCredentialStatePort, ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, @@ -17,10 +19,11 @@ use gateway_core::provider_ports::{ ProviderRuntimePolicyPort, ProviderStoreError, provider_refresh_backoff_at, provider_refresh_retry_at, }; +use gateway_core::routing::ProviderKind; use super::catalog::GrokCredentialCatalogService; use super::repository::{ - GrokCredentialAdmin, GrokCredentialRepository, GrokCredentialRepositoryError, + GrokCredentialAdmin, GrokCredentialRepository, GrokCredentialRepositoryError, loaded_from_core, }; use super::types::{ GrokAccountProfile, GrokOAuthSecret, PreparedGrokCredentialRotation, RotateGrokCredential, @@ -33,6 +36,7 @@ use crate::{ const MAX_REFRESH_BATCH: u32 = 100; const MAX_REFRESH_EXCLUSIONS: usize = 400; +const PROVIDER_NAME: &str = "xai"; const MAX_SECRET_BYTES: usize = 64 * 1_024; const DISCOVERY_NEGATIVE_CACHE_TTL: Duration = Duration::from_secs(5); /// 连续失败计数窗口；静默满窗后计数过期归零。 @@ -865,31 +869,33 @@ impl GrokCredentialRepository { .cloned() .collect::>(); let now = SystemTime::now(); - let mut accounts = self - .list_all_accounts() - .await? - .into_iter() - .filter(|account| !excluded.contains(account.id()) && account_due(account, policy, now)) - .collect::>(); - accounts.sort_by_key(|account| { - ( - account.access_token_expires_at(), - account.next_refresh_at(), - account.id().clone(), - ) - }); - accounts.truncate(MAX_REFRESH_BATCH as usize); + let provider = ProviderKind::new(PROVIDER_NAME) + .map_err(|_| GrokCredentialRepositoryError::InvalidCredentialData)?; + let limit = NonZeroU32::new(MAX_REFRESH_BATCH) + .ok_or(GrokCredentialRepositoryError::InvalidCredentialData)?; + let query = ProviderRefreshQuery::new( + provider, + now.checked_add(policy.margin()).unwrap_or(now), + now.checked_sub(REFRESH_RECOVERY_WINDOW) + .unwrap_or(SystemTime::UNIX_EPOCH), + now, + excluded.into_iter().collect(), + limit, + ); + let candidates = self.list_refresh_candidates(query).await?; - let mut due = Vec::with_capacity(accounts.len()); + let mut due = Vec::with_capacity(candidates.len()); let mut failed_account_ids = Vec::new(); - for account in accounts { - let loaded = match self.load(account.id(), account.revision()).await { + for candidate in candidates { + let account_id = candidate.account.id().clone(); + let loaded = match loaded_from_core(candidate) { Ok(loaded) => loaded, Err(_) => { - failed_account_ids.push(account.id().clone()); + failed_account_ids.push(account_id); continue; } }; + let account = loaded.account; let Some(subject) = account.upstream_user_id() else { failed_account_ids.push(account.id().clone()); continue; @@ -915,32 +921,6 @@ impl GrokCredentialRepository { } } -fn account_due( - account: &gateway_core::engine::credential::ProviderAccount, - policy: ProviderRefreshPolicy, - now: SystemTime, -) -> bool { - account.enabled() - && account.has_refresh_token() - && !matches!( - account.credential_state(), - CredentialState::Expired | CredentialState::Banned | CredentialState::Invalid - ) - && (refresh_recovery_window_exhausted(account.access_token_expires_at(), now) - || (access_token_refresh_due(account.access_token_expires_at(), policy, now) - && account - .next_refresh_at() - .is_none_or(|retry_at| retry_at <= now))) -} - -fn access_token_refresh_due( - access_token_expires_at: Option, - policy: ProviderRefreshPolicy, - now: SystemTime, -) -> bool { - access_token_expires_at.is_some_and(|expires_at| policy.is_refresh_due(expires_at, now)) -} - fn refreshed_tokens(tokens: RefreshedTokenSet) -> Result { let expires_in = tokens.expires_in().ok_or(GrokRefreshFailure::Rejected)?; Ok(GrokRefreshTokens { diff --git a/backend/crates/providers/xai/src/credential/repository.rs b/backend/crates/providers/xai/src/credential/repository.rs index 799385ea4..08ca8cdcf 100644 --- a/backend/crates/providers/xai/src/credential/repository.rs +++ b/backend/crates/providers/xai/src/credential/repository.rs @@ -334,6 +334,16 @@ impl GrokCredentialRepository { Self { store } } + pub(crate) async fn list_refresh_candidates( + &self, + query: gateway_core::engine::credential::ProviderRefreshQuery, + ) -> Result, GrokCredentialRepositoryError> { + self.store + .list_refresh_candidates(query) + .await + .map_err(map_store_error) + } + /// 读取不含 secret 的生命周期投影；Provider 仍是明文 credential JSON 的唯一解析者。 pub async fn read_lifecycle( &self, @@ -506,16 +516,6 @@ impl GrokCredentialRepository { Ok(accounts) } - pub(crate) async fn list_all_accounts( - &self, - ) -> Result, GrokCredentialRepositoryError> { - let accounts = self.store.list_accounts().await.map_err(map_store_error)?; - Ok(accounts - .into_iter() - .filter(|account| account.provider().as_str() == XAI_PROVIDER_KIND) - .collect()) - } - pub(crate) async fn load_current( &self, account_id: &ProviderAccountId, diff --git a/backend/crates/providers/xai/tests/support.rs b/backend/crates/providers/xai/tests/support.rs index d1e297298..ec1250d5f 100644 --- a/backend/crates/providers/xai/tests/support.rs +++ b/backend/crates/providers/xai/tests/support.rs @@ -11,8 +11,8 @@ use gateway_core::engine::credential::{ AccountConcurrencyLimit, AccountErrorReason, AccountStateChange, AccountWeight, CredentialCasOutcome, CredentialCasUpdate, CredentialCasUpdateParts, CredentialRevision, CredentialState, LoadedCredential, NewProviderAccount, PlaintextCredential, ProviderAccount, - ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, QuotaAccessChange, - QuotaObservation, QuotaObservationTouch, QuotaState, QuotaWriteOutcome, + ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, ProviderRefreshQuery, + QuotaAccessChange, QuotaObservation, QuotaObservationTouch, QuotaState, QuotaWriteOutcome, }; use gateway_core::error::{StoreError, StoreErrorKind}; use gateway_core::provider_ports::{ @@ -299,6 +299,28 @@ impl ProviderAccountStore for MemoryProviderAccountStore { .collect()) } + async fn list_refresh_candidates( + &self, + query: ProviderRefreshQuery, + ) -> Result, StoreError> { + let mut candidates = lock(&self.accounts) + .values() + .filter(|stored| query.contains(&stored.account)) + .map(|stored| LoadedCredential { + account: stored.account.clone(), + credential: stored.credential.clone(), + }) + .collect::>(); + candidates.sort_by_key(|loaded| { + ( + loaded.account.access_token_expires_at(), + loaded.account.id().clone(), + ) + }); + candidates.truncate(query.limit().get() as usize); + Ok(candidates) + } + async fn load_credential( &self, account: &ProviderAccountId,