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,