diff --git a/backend/crates/gateway-core/src/provider_ports/mod.rs b/backend/crates/gateway-core/src/provider_ports/mod.rs index 1ec8f2e85..d5c78bc07 100644 --- a/backend/crates/gateway-core/src/provider_ports/mod.rs +++ b/backend/crates/gateway-core/src/provider_ports/mod.rs @@ -12,6 +12,7 @@ use crate::engine::credential::{ AccountAvailability, AccountFeedbackStats, AccountRuntimeSignals, CredentialRevision, OpaqueProviderData, ProviderAccountId, ProviderAccountStore, }; +use crate::error::{IdentifierError, validate_text}; use crate::routing::ProviderKind; const MAX_PENDING_FLOW_TTL: Duration = Duration::from_secs(30 * 60); @@ -338,25 +339,44 @@ impl ProviderRefreshLeaseRequest { } } -/// Opaque catalog cache 的平台、账号与 credential revision 作用域。 +/// Provider 定义的 catalog cache 作用域。 +/// +/// Core 不解释其值;例如 Provider 可以使用套餐、区域或产品线作为共享目录边界。 +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ProviderCatalogScope(String); + +impl ProviderCatalogScope { + /// 创建一个稳定且可用于 Redis 隔离键的 Provider-owned 作用域。 + /// + /// # Errors + /// + /// 空值、过长文本或控制字符会被拒绝。 + pub fn new(value: impl Into) -> Result { + let value = value.into(); + validate_text(&value, 128, false, None)?; + Ok(Self(value)) + } + + /// 返回 Provider-owned 作用域文本。 + #[must_use] + pub fn as_str(&self) -> &str { + &self.0 + } +} + +/// Opaque catalog cache 的 Provider 与 Provider-owned 作用域。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct ProviderCatalogCacheKey { provider_kind: ProviderKind, - account_id: ProviderAccountId, - credential_revision: CredentialRevision, + scope: ProviderCatalogScope, } impl ProviderCatalogCacheKey { #[must_use] - pub const fn new( - provider_kind: ProviderKind, - account_id: ProviderAccountId, - credential_revision: CredentialRevision, - ) -> Self { + pub const fn new(provider_kind: ProviderKind, scope: ProviderCatalogScope) -> Self { Self { provider_kind, - account_id, - credential_revision, + scope, } } @@ -366,13 +386,8 @@ impl ProviderCatalogCacheKey { } #[must_use] - pub const fn account_id(&self) -> &ProviderAccountId { - &self.account_id - } - - #[must_use] - pub const fn credential_revision(&self) -> CredentialRevision { - self.credential_revision + pub const fn scope(&self) -> &ProviderCatalogScope { + &self.scope } } diff --git a/backend/crates/gateway-store/src/redis/credential_state.rs b/backend/crates/gateway-store/src/redis/credential_state.rs index 7129416a8..def5ff188 100644 --- a/backend/crates/gateway-store/src/redis/credential_state.rs +++ b/backend/crates/gateway-store/src/redis/credential_state.rs @@ -71,18 +71,17 @@ pub trait CredentialStateRepository: Send + Sync { async fn clear_credential_state(&self, provider_account_id: &str) -> StoreResult; } -/// Provider 账号目录 cache 的隔离键。Revision 变化后不会读取旧目录。 +/// Provider 目录 cache 的隔离键。具体作用域由对应 Provider 决定。 #[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderAccountCatalogCacheKey { +pub struct RedisProviderCatalogCacheKey { pub provider_kind: String, - pub provider_account_id: String, - pub credential_revision: Revision, + pub catalog_scope: String, } -impl ProviderAccountCatalogCacheKey { +impl RedisProviderCatalogCacheKey { pub fn validate(&self) -> StoreResult<()> { require_nonempty( - "provider account catalog cache", + "provider catalog cache", "provider_kind", &self.provider_kind, )?; @@ -92,25 +91,25 @@ impl ProviderAccountCatalogCacheKey { return Err(catalog_invalid("provider_kind must be a stable slug")); } require_nonempty( - "provider account catalog cache", - "provider_account_id", - &self.provider_account_id, + "provider catalog cache", + "catalog_scope", + &self.catalog_scope, ) } } #[async_trait] -pub trait ProviderAccountCatalogCacheRepository: Send + Sync { - async fn replace_provider_account_catalog( +pub trait ProviderCatalogCacheRepository: Send + Sync { + async fn replace_provider_catalog( &self, - key: &ProviderAccountCatalogCacheKey, + key: &RedisProviderCatalogCacheKey, catalog: &OpaqueProviderData, ttl_seconds: u64, ) -> StoreResult<()>; - async fn get_provider_account_catalog( + async fn get_provider_catalog( &self, - key: &ProviderAccountCatalogCacheKey, + key: &RedisProviderCatalogCacheKey, ) -> StoreResult>; } @@ -133,15 +132,12 @@ impl RedisCredentialStateRepository { Ok(format!("{}:account:{fingerprint}:state", self.namespace)) } - fn catalog_key(&self, key: &ProviderAccountCatalogCacheKey) -> StoreResult { + fn catalog_key(&self, key: &RedisProviderCatalogCacheKey) -> StoreResult { key.validate()?; - let fingerprint = - resource_fingerprint("provider account catalog cache", &key.provider_account_id)?; + let fingerprint = resource_fingerprint("provider catalog cache", &key.catalog_scope)?; Ok(format!( - "{}:provider:{}:account:{{{fingerprint}}}:catalog:{}", - self.namespace, - key.provider_kind, - key.credential_revision.get() + "{}:provider:{}:catalog:{{{fingerprint}}}", + self.namespace, key.provider_kind, )) } } @@ -224,10 +220,10 @@ impl CredentialStateRepository for RedisCredentialStateRepository { } #[async_trait] -impl ProviderAccountCatalogCacheRepository for RedisCredentialStateRepository { - async fn replace_provider_account_catalog( +impl ProviderCatalogCacheRepository for RedisCredentialStateRepository { + async fn replace_provider_catalog( &self, - key: &ProviderAccountCatalogCacheKey, + key: &RedisProviderCatalogCacheKey, catalog: &OpaqueProviderData, ttl_seconds: u64, ) -> StoreResult<()> { @@ -252,20 +248,20 @@ impl ProviderAccountCatalogCacheRepository for RedisCredentialStateRepository { .arg(ttl_ms) .query_async::<()>(&mut connection) .await - .map_err(|_| redis_unavailable("replace provider account catalog cache"))?; + .map_err(|_| redis_unavailable("replace provider catalog cache"))?; Ok(()) } - async fn get_provider_account_catalog( + async fn get_provider_catalog( &self, - key: &ProviderAccountCatalogCacheKey, + key: &RedisProviderCatalogCacheKey, ) -> StoreResult> { let mut connection = self.connection.clone(); let payload = redis::cmd("GET") .arg(self.catalog_key(key)?) .query_async::>>(&mut connection) .await - .map_err(|_| redis_unavailable("read provider account catalog cache"))?; + .map_err(|_| redis_unavailable("read provider catalog cache"))?; let Some(payload) = payload else { return Ok(None); }; @@ -358,13 +354,11 @@ impl ProviderCatalogCachePort for RedisCredentialStateRepository { if ttl_seconds == 0 || ttl.subsec_nanos() != 0 { return Err(provider_invalid("validate catalog cache TTL")); } - let key = ProviderAccountCatalogCacheKey { + let key = RedisProviderCatalogCacheKey { provider_kind: key.provider_kind().as_str().to_owned(), - provider_account_id: key.account_id().as_str().to_owned(), - credential_revision: Revision::new(key.credential_revision().get()) - .map_err(|_| provider_invalid("encode catalog cache key"))?, + catalog_scope: key.scope().as_str().to_owned(), }; - ProviderAccountCatalogCacheRepository::replace_provider_account_catalog( + ProviderCatalogCacheRepository::replace_provider_catalog( self, &key, catalog, @@ -381,13 +375,11 @@ impl ProviderCatalogCachePort for RedisCredentialStateRepository { ) -> futures::future::BoxFuture<'a, Result, ProviderStoreError>> { Box::pin(async move { - let key = ProviderAccountCatalogCacheKey { + let key = RedisProviderCatalogCacheKey { provider_kind: key.provider_kind().as_str().to_owned(), - provider_account_id: key.account_id().as_str().to_owned(), - credential_revision: Revision::new(key.credential_revision().get()) - .map_err(|_| provider_invalid("encode catalog cache key"))?, + catalog_scope: key.scope().as_str().to_owned(), }; - ProviderAccountCatalogCacheRepository::get_provider_account_catalog(self, &key) + ProviderCatalogCacheRepository::get_provider_catalog(self, &key) .await .map_err(|_| provider_unavailable("read catalog cache")) }) @@ -403,7 +395,7 @@ fn invalid(message: &str) -> StoreError { fn catalog_invalid(message: &str) -> StoreError { StoreError::InvalidData { - entity: "provider account catalog cache", + entity: "provider catalog cache", message: message.to_owned(), } } diff --git a/backend/crates/gateway-store/tests/redis/credential_state.rs b/backend/crates/gateway-store/tests/redis/credential_state.rs index cd6ce8c67..39d97d4dc 100644 --- a/backend/crates/gateway-store/tests/redis/credential_state.rs +++ b/backend/crates/gateway-store/tests/redis/credential_state.rs @@ -2,14 +2,14 @@ use chrono::Utc; use gateway_store::{ Revision, redis::{ - CredentialStateCache, ProviderAccountCatalogCacheKey, - ProviderAccountCatalogCacheRepository, RedisCredentialStateRepository, + CredentialStateCache, ProviderCatalogCacheRepository, RedisCredentialStateRepository, + RedisProviderCatalogCacheKey, }, }; #[test] fn credential_state_adapter_implements_opaque_catalog_cache_port() { - fn assert_port() {} + fn assert_port() {} assert_port::(); } @@ -26,11 +26,10 @@ fn credential_state_rejects_provider_specific_status() { } #[test] -fn provider_catalog_cache_key_requires_provider_and_account() { - let key = ProviderAccountCatalogCacheKey { +fn provider_catalog_cache_key_requires_provider_and_scope() { + let key = RedisProviderCatalogCacheKey { provider_kind: "xai".to_owned(), - provider_account_id: String::new(), - credential_revision: Revision::new(1).expect("positive revision"), + catalog_scope: String::new(), }; assert!(key.validate().is_err()); } diff --git a/backend/crates/providers/openai/src/admin.rs b/backend/crates/providers/openai/src/admin.rs index 543101849..8568b27d8 100644 --- a/backend/crates/providers/openai/src/admin.rs +++ b/backend/crates/providers/openai/src/admin.rs @@ -440,24 +440,21 @@ impl ProviderAdmin for OpenAiAdminProvider { refresh: bool, ) -> Result { let account = self.account(account_id).await?; - let models = if refresh { + let catalog = if refresh { self.catalog - .synchronize_account(account_id) + .refresh_account_catalog(account_id) .await .map_err(map_catalog_error)? } else { self.catalog - .cached_account_models(account_id, account.revision()) + .cached_or_refresh_account_catalog(&account) + .await .map_err(map_catalog_error)? - .unwrap_or_default() }; - let observed_at = self - .catalog - .cached() - .map_err(map_catalog_error)? - .map(|snapshot| DateTime::::from(snapshot.observed_at())); - let models = models - .into_iter() + let models = catalog + .models() + .iter() + .cloned() .map(|model| { let id = UpstreamModelId::new(model.clone()) .map_err(|_| provider_admin_error(ProviderAdminErrorKind::Internal))?; @@ -466,7 +463,7 @@ impl ProviderAdmin for OpenAiAdminProvider { .collect::, ProviderAdminError>>()?; Ok(ProviderModels { models, - observed_at, + observed_at: Some(catalog.observed_at()), }) } diff --git a/backend/crates/providers/openai/src/credential/catalog.rs b/backend/crates/providers/openai/src/credential/catalog.rs index 8e2f5f319..ad9b7dd6e 100644 --- a/backend/crates/providers/openai/src/credential/catalog.rs +++ b/backend/crates/providers/openai/src/credential/catalog.rs @@ -1,14 +1,19 @@ -//! Codex model entitlement 的实时查询与可重建进程内 cache;不落 PostgreSQL。 +//! Codex 账号 entitlement 的进程内快照,以及套餐级 Redis 模型目录 cache。 -use std::collections::{BTreeMap, btree_map::Entry}; +use std::collections::{BTreeMap, BTreeSet, btree_map::Entry}; use std::fmt; use std::sync::{Arc, Mutex, RwLock}; -use std::time::SystemTime; +use std::time::{Duration, SystemTime}; +use chrono::{DateTime, Utc}; use gateway_core::engine::credential::{ - AccountAvailability, CredentialRevision, ProviderAccount, ProviderAccountId, + AccountAvailability, CredentialRevision, OpaqueProviderData, ProviderAccount, ProviderAccountId, }; use gateway_core::engine::provider::ProviderCatalogGeneration; +use gateway_core::provider_ports::{ + ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCatalogScope, +}; +use gateway_core::routing::{ProviderKind, UpstreamModelId}; use secrecy::ExposeSecret; use thiserror::Error; use tokio::sync::Notify; @@ -20,6 +25,72 @@ use crate::transport::profile::CodexWireProfileState; use crate::transport::{CodexBackendClient, CodexCatalogModel, CodexRequestContext}; const MAX_RESPONSE_ETAG_BYTES: usize = 256; +const PLAN_CATALOG_CACHE_TTL: Duration = Duration::from_secs(5 * 60); +const MAX_CATALOG_FETCH_ATTEMPTS: usize = 3; +const MAX_PLAN_CATALOG_MODELS: usize = 2_048; + +/// OpenAI 以套餐划分的模型目录作用域。 +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] +pub struct CodexCatalogScope(ProviderCatalogScope); + +impl CodexCatalogScope { + /// 从账号已验证的套餐事实构造目录作用域。 + pub fn for_account(account: &ProviderAccount) -> Result { + let plan = account + .plan_type() + .map(str::trim) + .filter(|plan| !plan.is_empty()) + .map(str::to_ascii_lowercase) + .unwrap_or_else(|| "unknown".to_owned()); + ProviderCatalogScope::new(format!("plan:{plan}")) + .map(Self) + .map_err(|_| CodexCredentialCatalogError::InvalidCredentialData) + } + + /// 返回稳定的 Provider-owned 作用域。 + #[must_use] + pub fn as_str(&self) -> &str { + self.0.as_str() + } +} + +/// Redis TTL cache 中一条可重建的套餐模型目录。 +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CodexPlanCatalog { + scope: CodexCatalogScope, + observed_at: DateTime, + models: Vec, +} + +impl CodexPlanCatalog { + #[must_use] + pub const fn new( + scope: CodexCatalogScope, + observed_at: DateTime, + models: Vec, + ) -> Self { + Self { + scope, + observed_at, + models, + } + } + + #[must_use] + pub const fn scope(&self) -> &CodexCatalogScope { + &self.scope + } + + #[must_use] + pub const fn observed_at(&self) -> DateTime { + self.observed_at + } + + #[must_use] + pub fn models(&self) -> &[String] { + &self.models + } +} #[derive(Clone)] pub struct CodexCredentialCatalogSnapshot { @@ -102,6 +173,7 @@ struct CatalogCacheState { } struct FetchedAccountModels { + account_id: ProviderAccountId, revision: CredentialRevision, models: Vec, etag: Option, @@ -125,6 +197,7 @@ pub struct CodexCredentialCatalogService { http: reqwest::Client, base_url: String, agent_identity: Arc, + plan_catalog_cache: Arc, cache: Arc>, etags: Arc>, etag_notification: Arc, @@ -137,6 +210,7 @@ impl CodexCredentialCatalogService { http: reqwest::Client, base_url: String, agent_identity: Arc, + plan_catalog_cache: Arc, ) -> Self { Self { repository, @@ -144,6 +218,7 @@ impl CodexCredentialCatalogService { http, base_url, agent_identity, + plan_catalog_cache, cache: Arc::new(RwLock::new(CatalogCacheState::default())), etags: Arc::new(Mutex::new(CatalogEtagState::default())), etag_notification: Arc::new(Notify::new()), @@ -197,11 +272,23 @@ impl CodexCredentialCatalogService { })) } - /// 只刷新指定账号的 realtime catalog,并以本地 revision 原子合并账号事实。 - pub async fn synchronize_account( + /// 优先读取套餐目录 cache;缺失时才用当前账号所属套餐的有限候选集实时填充。 + pub async fn cached_or_refresh_account_catalog( + &self, + account: &ProviderAccount, + ) -> Result { + let scope = CodexCatalogScope::for_account(account)?; + if let Some(catalog) = self.read_plan_catalog(&scope).await? { + return Ok(catalog); + } + self.refresh_account_catalog(account.id()).await + } + + /// 实时刷新指定账号所属套餐的模型集合,并覆盖可重建 Redis cache。 + pub async fn refresh_account_catalog( &self, account_id: &ProviderAccountId, - ) -> Result, CodexCredentialCatalogError> { + ) -> Result { let account = self .repository .store() @@ -210,27 +297,39 @@ impl CodexCredentialCatalogService { .map_err(|_| CodexCredentialCatalogError::InvalidCredentialData)? .filter(|account| account.provider().as_str() == "openai") .ok_or(CodexCredentialCatalogError::NoEligibleCredential)?; + let scope = CodexCatalogScope::for_account(&account)?; + let mut candidates = + catalog_candidates_by_scope(self.repository.list_for_provider().await?)? + .remove(&scope) + .unwrap_or_default(); + if candidates.is_empty() { + return Err(CodexCredentialCatalogError::NoEligibleCredential); + } + candidates.sort_by(|left, right| { + let left_preferred = left.id() == account_id; + let right_preferred = right.id() == account_id; + right_preferred + .cmp(&left_preferred) + .then_with(|| left.id().cmp(right.id())) + }); let client = CodexBackendClient::new( self.http.clone(), self.base_url.clone(), self.profile.clone(), ); - let cache_revision = self.cache_revision()?; - let fetched = self.fetch_account_models(&client, &account).await?; - let entitlement = fetched - .models - .iter() - .map(|model| model.request_model().as_str().to_owned()) - .collect::>(); - self.replace_account_cache( - cache_revision, - account_id, - fetched.revision, - fetched.models, - entitlement.clone(), - )?; - self.record_applied_catalog_etags(fetched.etag)?; - Ok(entitlement) + let fetched = self.fetch_scope_models(&client, candidates).await?; + let catalog = CodexPlanCatalog::new(scope, Utc::now(), model_ids(&fetched.models)); + self.replace_plan_catalog(&catalog).await?; + Ok(catalog) + } + + /// 读取当前账号所属套餐的目录 cache,不触发上游请求。 + pub async fn read_account_catalog( + &self, + account: &ProviderAccount, + ) -> Result, CodexCredentialCatalogError> { + self.read_plan_catalog(&CodexCatalogScope::for_account(account)?) + .await } pub async fn synchronize( @@ -246,14 +345,7 @@ impl CodexCredentialCatalogService { async fn fetch_catalog(&self) -> Result { let accounts = self.repository.list_for_provider().await?; - let now = SystemTime::now(); - let accounts = accounts - .into_iter() - .filter(|account| eligible_catalog_account(account, now)) - .collect::>(); - if accounts.is_empty() { - return Err(CodexCredentialCatalogError::NoEligibleCredential); - } + let groups = catalog_candidates_by_scope(accounts)?; let client = CodexBackendClient::new( self.http.clone(), self.base_url.clone(), @@ -263,12 +355,17 @@ impl CodexCredentialCatalogService { let mut union_order = Vec::new(); let mut account_models = BTreeMap::new(); let mut etags = Vec::new(); - for account in accounts { - let fetched = self.fetch_account_models(&client, &account).await?; - let mut entitlement = Vec::with_capacity(fetched.models.len()); + for (scope, candidates) in groups { + let fetched = self.fetch_scope_models(&client, candidates).await?; + let entitlement = model_ids(&fetched.models); + self.replace_plan_catalog(&CodexPlanCatalog::new( + scope, + Utc::now(), + entitlement.clone(), + )) + .await?; for model in &fetched.models { let id = model.request_model().as_str().to_owned(); - entitlement.push(id.clone()); match union.entry(id) { Entry::Vacant(entry) => { union_order.push(entry.key().clone()); @@ -281,7 +378,7 @@ impl CodexCredentialCatalogService { } } account_models.insert( - account.id().clone(), + fetched.account_id, CodexAccountEntitlement { revision: fetched.revision, models: entitlement, @@ -301,6 +398,55 @@ impl CodexCredentialCatalogService { Ok(FetchedCatalog { snapshot, etags }) } + async fn fetch_scope_models( + &self, + client: &CodexBackendClient, + candidates: Vec, + ) -> Result { + let mut last_error = None; + for account in candidates.into_iter().take(MAX_CATALOG_FETCH_ATTEMPTS) { + match self.fetch_account_models(client, &account).await { + Ok(fetched) => return Ok(fetched), + Err(error) => last_error = Some(error), + } + } + Err(last_error.unwrap_or(CodexCredentialCatalogError::NoEligibleCredential)) + } + + async fn replace_plan_catalog( + &self, + catalog: &CodexPlanCatalog, + ) -> Result<(), CodexCredentialCatalogError> { + let key = self.plan_catalog_key(catalog.scope())?; + let document = encode_plan_catalog(catalog); + self.plan_catalog_cache + .replace(&key, &document, PLAN_CATALOG_CACHE_TTL) + .await + .map_err(|_| CodexCredentialCatalogError::Cache) + } + + async fn read_plan_catalog( + &self, + scope: &CodexCatalogScope, + ) -> Result, CodexCredentialCatalogError> { + let key = self.plan_catalog_key(scope)?; + self.plan_catalog_cache + .read(&key) + .await + .map_err(|_| CodexCredentialCatalogError::Cache)? + .map(|document| decode_plan_catalog(scope, document)) + .transpose() + } + + fn plan_catalog_key( + &self, + scope: &CodexCatalogScope, + ) -> Result { + let provider_kind = + ProviderKind::new("openai").map_err(|_| CodexCredentialCatalogError::Cache)?; + Ok(ProviderCatalogCacheKey::new(provider_kind, scope.0.clone())) + } + async fn fetch_account_models( &self, client: &CodexBackendClient, @@ -354,84 +500,13 @@ impl CodexCredentialCatalogService { } let snapshot = result.map_err(|_| CodexCredentialCatalogError::Upstream)?; Ok(FetchedAccountModels { + account_id: prepared.account.id().clone(), revision: prepared.account.revision(), models: snapshot.models().to_vec(), etag: snapshot.etag().map(str::to_owned), }) } - fn replace_account_cache( - &self, - expected_cache_revision: u64, - account_id: &ProviderAccountId, - revision: CredentialRevision, - models: Vec, - entitlement: Vec, - ) -> Result<(), CodexCredentialCatalogError> { - let mut cache = self - .cache - .write() - .map_err(|_| CodexCredentialCatalogError::Cache)?; - if cache.revision != expected_cache_revision { - return Err(CodexCredentialCatalogError::ConcurrentUpdate); - } - let mut union = BTreeMap::::new(); - let mut union_order = Vec::new(); - let mut account_models = BTreeMap::new(); - if let Some(existing) = cache.snapshot.as_ref() { - for model in &existing.models { - let id = model.request_model().as_str().to_owned(); - union_order.push(id.clone()); - union.insert(id, model.clone()); - } - account_models = existing.account_models.clone(); - } - for model in models { - let id = model.request_model().as_str().to_owned(); - match union.entry(id) { - Entry::Vacant(entry) => { - union_order.push(entry.key().clone()); - entry.insert(model); - } - Entry::Occupied(entry) if entry.get() == &model => {} - Entry::Occupied(_) => { - return Err(CodexCredentialCatalogError::ConflictingModelFacts); - } - } - } - account_models.insert( - account_id.clone(), - CodexAccountEntitlement { - revision, - models: entitlement, - }, - ); - union.retain(|model, _| { - account_models - .values() - .any(|entitlement| entitlement.models.iter().any(|item| item == model)) - }); - union_order.retain(|id| union.contains_key(id)); - let snapshot = CodexCredentialCatalogSnapshot { - observed_at: SystemTime::now(), - models: union_order - .into_iter() - .filter_map(|id| union.remove(&id)) - .collect(), - account_models, - }; - let changed = cache - .snapshot - .as_ref() - .is_none_or(|existing| !same_catalog(existing, &snapshot)); - cache.snapshot = Some(snapshot); - cache.revision = cache.revision.saturating_add(1); - if changed { - cache.generation = cache.generation.saturating_add(1); - } - Ok(()) - } - pub fn invalidate(&self) -> Result<(), CodexCredentialCatalogError> { let mut cache = self .cache @@ -478,7 +553,17 @@ impl CodexCredentialCatalogService { } } - /// 忽略当前 cache,按 ETag 变化强制生成一份完整新快照。 + /// 立即刷新所有套餐目录,但不接管 ETag daemon 已认领的刷新状态。 + /// + /// 周期 worker 与 ETag daemon 可以同时请求目录;只有后者可以完成 + /// `inflight` ETag 的状态转换,避免周期刷新错误地确认另一个请求的版本。 + pub async fn refresh_catalogs( + &self, + ) -> Result { + self.refresh_inner().await + } + + /// 忽略当前 cache,按已认领的 ETag 变化强制生成一份完整新快照。 pub async fn refresh( &self, ) -> Result { @@ -585,6 +670,105 @@ impl CodexCredentialCatalogService { } } +fn model_ids(models: &[CodexCatalogModel]) -> Vec { + models + .iter() + .map(|model| model.request_model().as_str().to_owned()) + .collect() +} + +fn catalog_candidates_by_scope( + accounts: Vec, +) -> Result>, CodexCredentialCatalogError> { + let now = SystemTime::now(); + let mut groups = BTreeMap::>::new(); + for account in accounts + .into_iter() + .filter(|account| eligible_catalog_account(account, now)) + { + let scope = CodexCatalogScope::for_account(&account)?; + groups.entry(scope).or_default().push(account); + } + for candidates in groups.values_mut() { + candidates.sort_by(|left, right| left.id().cmp(right.id())); + } + if groups.is_empty() { + return Err(CodexCredentialCatalogError::NoEligibleCredential); + } + Ok(groups) +} + +fn encode_plan_catalog(catalog: &CodexPlanCatalog) -> OpaqueProviderData { + let mut document = serde_json::Map::new(); + document.insert("version".to_owned(), serde_json::Value::from(1)); + document.insert( + "scope".to_owned(), + serde_json::Value::String(catalog.scope().as_str().to_owned()), + ); + document.insert( + "observedAt".to_owned(), + serde_json::Value::String(catalog.observed_at().to_rfc3339()), + ); + document.insert( + "models".to_owned(), + serde_json::Value::Array( + catalog + .models() + .iter() + .cloned() + .map(serde_json::Value::String) + .collect(), + ), + ); + OpaqueProviderData::new(document) +} + +fn decode_plan_catalog( + scope: &CodexCatalogScope, + document: OpaqueProviderData, +) -> Result { + let mut fields = document.into_inner(); + if fields.remove("version").and_then(|value| value.as_u64()) != Some(1) { + return Err(CodexCredentialCatalogError::Cache); + } + if fields + .remove("scope") + .and_then(|value| value.as_str().map(ToOwned::to_owned)) + .as_deref() + != Some(scope.as_str()) + { + return Err(CodexCredentialCatalogError::Cache); + } + let observed_at = fields + .remove("observedAt") + .and_then(|value| value.as_str().map(ToOwned::to_owned)) + .and_then(|value| DateTime::parse_from_rfc3339(&value).ok()) + .map(|value| value.with_timezone(&Utc)) + .ok_or(CodexCredentialCatalogError::Cache)?; + let values = fields + .remove("models") + .and_then(|value| value.as_array().cloned()) + .filter(|models| !models.is_empty() && models.len() <= MAX_PLAN_CATALOG_MODELS) + .ok_or(CodexCredentialCatalogError::Cache)?; + if !fields.is_empty() { + return Err(CodexCredentialCatalogError::Cache); + } + let mut seen = BTreeSet::new(); + let mut models = Vec::with_capacity(values.len()); + for value in values { + let value = value + .as_str() + .map(ToOwned::to_owned) + .ok_or(CodexCredentialCatalogError::Cache)?; + let model = UpstreamModelId::new(value).map_err(|_| CodexCredentialCatalogError::Cache)?; + if !seen.insert(model.as_str().to_owned()) { + return Err(CodexCredentialCatalogError::Cache); + } + models.push(model.as_str().to_owned()); + } + Ok(CodexPlanCatalog::new(scope.clone(), observed_at, models)) +} + fn same_catalog( left: &CodexCredentialCatalogSnapshot, right: &CodexCredentialCatalogSnapshot, diff --git a/backend/crates/providers/openai/src/credential/mod.rs b/backend/crates/providers/openai/src/credential/mod.rs index c495e03da..5318ef031 100644 --- a/backend/crates/providers/openai/src/credential/mod.rs +++ b/backend/crates/providers/openai/src/credential/mod.rs @@ -34,7 +34,8 @@ pub use admin::{ PreparedCodexCredentialRotationGuard, RotateManagedCodexCredential, }; pub use catalog::{ - CodexCredentialCatalogError, CodexCredentialCatalogService, CodexCredentialCatalogSnapshot, + CodexCatalogScope, CodexCredentialCatalogError, CodexCredentialCatalogService, + CodexCredentialCatalogSnapshot, CodexPlanCatalog, }; pub use cookie::{CodexCookiePolicy, CookiePolicyError}; pub use identity::{ diff --git a/backend/crates/providers/openai/src/lib.rs b/backend/crates/providers/openai/src/lib.rs index ec84c8bf1..e5447724b 100644 --- a/backend/crates/providers/openai/src/lib.rs +++ b/backend/crates/providers/openai/src/lib.rs @@ -102,6 +102,7 @@ pub async fn initialize( http.clone(), config.base_url().to_owned(), Arc::clone(&agent_identity), + ports.catalog_cache(), )); let quota = Arc::new(CodexCredentialQuotaService::new( repository.clone(), diff --git a/backend/crates/providers/openai/src/provider.rs b/backend/crates/providers/openai/src/provider.rs index 72f90556a..bd3d80c27 100644 --- a/backend/crates/providers/openai/src/provider.rs +++ b/backend/crates/providers/openai/src/provider.rs @@ -40,11 +40,12 @@ use serde_json::{Map, Value, json}; use url::Url; use crate::credential::{ - CodexAccountFailure, CodexAgentIdentityTaskService, CodexCredentialCatalogService, - CodexCredentialLease, CodexCredentialQuotaService, CodexCredentialRefreshOutcome, - CodexCredentialRefreshService, CodexCredentialSelector, CodexCyberPolicyScope, - CodexQuotaRefreshPolicy, CredentialSelectionError, RuntimeCodexCookie, SelectCodexCredential, - derive_codex_cyber_policy_session_key, derive_codex_session_affinity_key, + CodexAccountFailure, CodexAgentIdentityTaskService, CodexCredentialCatalogError, + CodexCredentialCatalogService, CodexCredentialLease, CodexCredentialQuotaService, + CodexCredentialRefreshOutcome, CodexCredentialRefreshService, CodexCredentialSelector, + CodexCyberPolicyScope, CodexQuotaRefreshPolicy, CredentialSelectionError, RuntimeCodexCookie, + SelectCodexCredential, derive_codex_cyber_policy_session_key, + derive_codex_session_affinity_key, }; use crate::transport::canonical::{ CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalOutcome, @@ -2116,7 +2117,10 @@ pub(crate) fn worker_contributions( WorkerContribution::Registration(scheduled_registration( quota_id, quota_refresh_policy.interval(), - Box::new(OpenAiQuotaTask { quota }), + Box::new(OpenAiQuotaTask { + quota, + catalog: Arc::clone(&catalog), + }), )?), WorkerContribution::Registration(WorkerRegistration::try_new( etag_id, @@ -2208,6 +2212,7 @@ impl ScheduledTask for OpenAiOAuthRefreshTask { struct OpenAiQuotaTask { quota: Arc, + catalog: Arc, } struct OpenAiCatalogEtagTask { @@ -2264,13 +2269,33 @@ impl ScheduledTask for OpenAiQuotaTask { if context.cancellation().is_cancelled() { return Ok(()); } + let mut failures = false; match self.quota.synchronize().await { Ok(summary) if summary.has_operational_failures() => { tracing::warn!("OpenAI quota cycle contained operational failures"); - Ok(()) } - Ok(_) => Ok(()), - Err(_) => Err(WorkerTaskError::safe("OpenAI quota synchronization failed")), + Ok(_) => {} + Err(_) => { + failures = true; + tracing::warn!("OpenAI quota synchronization failed"); + } + } + if context.cancellation().is_cancelled() { + return Ok(()); + } + match self.catalog.refresh_catalogs().await { + Ok(_) | Err(CodexCredentialCatalogError::NoEligibleCredential) => {} + Err(error) => { + failures = true; + tracing::warn!(error = %error, "OpenAI model catalog refresh failed"); + } + } + if failures { + Err(WorkerTaskError::safe( + "OpenAI quota or catalog synchronization failed", + )) + } else { + Ok(()) } }) } diff --git a/backend/crates/providers/openai/tests/admin.rs b/backend/crates/providers/openai/tests/admin.rs index 7a2acf8da..2b783fb9e 100644 --- a/backend/crates/providers/openai/tests/admin.rs +++ b/backend/crates/providers/openai/tests/admin.rs @@ -393,9 +393,15 @@ async fn openai_admin_provider_projects_cached_quota_models_and_canonical_export .expect("stored account"); let record = account_record(&account); let config = valid_config(); + let catalog_cache = Arc::new(TestCatalogCache::default()); + catalog_cache.seed("plan:pro", ["gpt-5.4"]); let bundle = provider_openai::initialize( config.config.clone(), - provider_ports_with(Arc::clone(&store), Arc::new(TestOAuthPending::default())), + provider_ports_with_catalog( + Arc::clone(&store), + Arc::new(TestOAuthPending::default()), + catalog_cache, + ), ) .await .expect("OpenAI bundle"); @@ -435,7 +441,7 @@ async fn openai_admin_provider_projects_cached_quota_models_and_canonical_export .models(&account_id, false) .await .expect("cached models"); - assert!(models.models.is_empty()); + assert_eq!(models.models[0].id.as_str(), "gpt-5.4"); let loaded = store .load_credential(account.id(), account.revision()) .await @@ -595,13 +601,21 @@ fn provider_ports() -> ProviderStorePorts { fn provider_ports_with( accounts: Arc, pending: Arc, +) -> ProviderStorePorts { + provider_ports_with_catalog(accounts, pending, Arc::new(TestCatalogCache::default())) +} + +fn provider_ports_with_catalog( + accounts: Arc, + pending: Arc, + catalog_cache: Arc, ) -> ProviderStorePorts { ProviderStorePorts::new( accounts, Arc::new(TestLeaseCoordinator::default()), Arc::new(MemorySessionAffinity::default()), Arc::new(MemorySessionExclusions::default()), - Arc::new(TestCatalogCache::default()), + catalog_cache, Arc::new(TestCredentialState), Arc::new(TestCooldown), Arc::new(TestRuntimePolicy), @@ -682,7 +696,7 @@ impl ProviderCatalogCachePort for TestCatalogCache { self.values .lock() .expect("catalog cache") - .insert(key.account_id().to_string(), catalog.clone()); + .insert(key.scope().as_str().to_owned(), catalog.clone()); Ok(()) }) } @@ -696,12 +710,37 @@ impl ProviderCatalogCachePort for TestCatalogCache { .values .lock() .expect("catalog cache") - .get(key.account_id().as_str()) + .get(key.scope().as_str()) .cloned()) }) } } +impl TestCatalogCache { + fn seed(&self, scope: &str, models: impl IntoIterator) { + let mut document = Map::new(); + document.insert("version".to_owned(), Value::from(1)); + document.insert("scope".to_owned(), Value::String(scope.to_owned())); + document.insert( + "observedAt".to_owned(), + Value::String(Utc::now().to_rfc3339()), + ); + document.insert( + "models".to_owned(), + Value::Array( + models + .into_iter() + .map(|model| Value::String(model.to_owned())) + .collect(), + ), + ); + self.values + .lock() + .expect("catalog cache") + .insert(scope.to_owned(), OpaqueProviderData::new(document)); + } +} + struct TestCredentialState; impl ProviderCredentialStatePort for TestCredentialState { diff --git a/backend/crates/providers/openai/tests/credential/catalog.rs b/backend/crates/providers/openai/tests/credential/catalog.rs index 15aec4d03..d02f99bb4 100644 --- a/backend/crates/providers/openai/tests/credential/catalog.rs +++ b/backend/crates/providers/openai/tests/credential/catalog.rs @@ -1,13 +1,39 @@ use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use chrono::{TimeZone as _, Utc}; use futures::executor::block_on; -use gateway_core::engine::credential::ProviderAccountId; +use gateway_core::engine::credential::{ProviderAccount, ProviderAccountId}; +use gateway_core::provider_ports::ProviderCatalogCachePort; use provider_openai::OFFICIAL_CODEX_BASE_URL; -use provider_openai::credential::{CodexCredentialCatalogError, CodexCredentialCatalogService}; +use provider_openai::credential::{ + CodexCredentialCatalogError, CodexCredentialCatalogService, ImportCodexOAuthCredential, +}; use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState}; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate}; -use crate::support::{MemoryAccountStore, agent_identity_service}; +use crate::support::{MemoryAccountStore, agent_identity_service, catalog_cache, profile, secret}; + +const OFFICIAL_FIXTURE: &[u8] = + include_bytes!("../transport/fixtures/official_models_snapshot.json"); + +struct ReplacingCatalogResponder { + calls: Arc, +} + +impl Respond for ReplacingCatalogResponder { + fn respond(&self, _: &Request) -> ResponseTemplate { + let body = if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { + OFFICIAL_FIXTURE.to_vec() + } else { + br#"{"models":[{"slug":"gpt-5.5","display_name":"GPT-5.5"}]}"#.to_vec() + }; + ResponseTemplate::new(200) + .insert_header("content-type", "application/json") + .set_body_raw(body, "application/json") + } +} fn wire_profile() -> CodexWireProfileState { CodexWireProfileState::new(CodexWireProfile { @@ -27,6 +53,14 @@ fn wire_profile() -> CodexWireProfileState { } fn service(store: &Arc) -> CodexCredentialCatalogService { + service_with_catalog_cache(store, OFFICIAL_CODEX_BASE_URL.to_owned(), catalog_cache()) +} + +fn service_with_catalog_cache( + store: &Arc, + base_url: String, + catalog_cache: Arc, +) -> CodexCredentialCatalogService { CodexCredentialCatalogService::new( store.repository(), wire_profile(), @@ -34,11 +68,26 @@ fn service(store: &Arc) -> CodexCredentialCatalogService { .no_proxy() .build() .expect("client"), - OFFICIAL_CODEX_BASE_URL.to_owned(), + base_url, agent_identity_service(store), + catalog_cache, ) } +async fn seed_account(store: &Arc, account_id: &str) -> ProviderAccount { + store + .seed_oauth_credential(ImportCodexOAuthCredential { + account_id: account_id.to_owned(), + name: account_id.to_owned(), + secret: secret(&format!("access-{account_id}")), + verified_account: profile(&format!("chatgpt-{account_id}")), + next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)), + enabled: true, + }) + .await; + store.account(account_id).expect("seeded account") +} + #[test] fn catalog_starts_without_a_snapshot_or_generation() { let store = Arc::new(MemoryAccountStore::default()); @@ -66,7 +115,7 @@ async fn missing_account_refresh_fails_before_network_io() { let account = ProviderAccountId::new("acct_missing").expect("account id"); let error = service(&store) - .synchronize_account(&account) + .refresh_account_catalog(&account) .await .expect_err("missing account"); @@ -76,6 +125,75 @@ async fn missing_account_refresh_fails_before_network_io() { )); } +#[tokio::test] +async fn plan_catalog_cache_is_shared_and_manual_refresh_replaces_it() { + let store = Arc::new(MemoryAccountStore::default()); + let first = seed_account(&store, "acct_catalog_a").await; + let second = seed_account(&store, "acct_catalog_b").await; + let server = MockServer::start().await; + let calls = Arc::new(AtomicUsize::new(0)); + Mock::given(method("GET")) + .and(path("/codex/models")) + .respond_with(ReplacingCatalogResponder { + calls: Arc::clone(&calls), + }) + .expect(2) + .mount(&server) + .await; + let service = service_with_catalog_cache(&store, server.uri(), catalog_cache()); + + let initial = service + .cached_or_refresh_account_catalog(&first) + .await + .expect("cold cache fetch"); + let cached = service + .cached_or_refresh_account_catalog(&second) + .await + .expect("shared cache hit"); + let refreshed = service + .refresh_account_catalog(second.id()) + .await + .expect("manual refresh"); + let replaced = service + .cached_or_refresh_account_catalog(&first) + .await + .expect("replaced cache hit"); + + assert_eq!(initial.models(), ["gpt-5.4"]); + assert_eq!(cached.models(), ["gpt-5.4"]); + assert_eq!(refreshed.models(), ["gpt-5.5"]); + assert_eq!(replaced.models(), ["gpt-5.5"]); + assert_eq!(calls.load(Ordering::SeqCst), 2); + server.verify().await; +} + +#[tokio::test] +async fn plan_catalog_refresh_stops_after_three_failed_accounts() { + let store = Arc::new(MemoryAccountStore::default()); + let accounts = [ + seed_account(&store, "acct_catalog_limit_a").await, + seed_account(&store, "acct_catalog_limit_b").await, + seed_account(&store, "acct_catalog_limit_c").await, + seed_account(&store, "acct_catalog_limit_d").await, + ]; + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/codex/models")) + .respond_with(ResponseTemplate::new(503)) + .expect(3) + .mount(&server) + .await; + let service = service_with_catalog_cache(&store, server.uri(), catalog_cache()); + + let error = service + .refresh_account_catalog(accounts[0].id()) + .await + .expect_err("three failed accounts stop the refresh"); + + assert!(matches!(error, CodexCredentialCatalogError::Upstream)); + server.verify().await; +} + #[tokio::test] async fn response_etag_change_is_deduplicated_and_queued_once() { let store = Arc::new(MemoryAccountStore::default()); @@ -99,6 +217,38 @@ async fn response_etag_change_is_deduplicated_and_queued_once() { .expect("queued ETag refresh"); } +#[tokio::test] +async fn periodic_catalog_refresh_does_not_complete_an_etag_refresh_it_did_not_claim() { + let store = Arc::new(MemoryAccountStore::default()); + let service = service(&store); + + assert!( + service + .observe_response_etag("\"models-v2\"") + .expect("new ETag") + ); + tokio::time::timeout( + std::time::Duration::from_secs(1), + service.wait_for_etag_refresh(), + ) + .await + .expect("ETag daemon claims the refresh"); + + assert!(matches!( + service.refresh_catalogs().await, + Err(CodexCredentialCatalogError::NoEligibleCredential) + )); + assert!( + tokio::time::timeout( + std::time::Duration::from_millis(50), + service.wait_for_etag_refresh(), + ) + .await + .is_err(), + "periodic refresh must not requeue or finish the daemon-owned ETag" + ); +} + #[test] fn invalid_response_etag_is_rejected_without_touching_the_catalog() { let store = Arc::new(MemoryAccountStore::default()); diff --git a/backend/crates/providers/openai/tests/credential/contract.rs b/backend/crates/providers/openai/tests/credential/contract.rs index 2d3127562..e57084455 100644 --- a/backend/crates/providers/openai/tests/credential/contract.rs +++ b/backend/crates/providers/openai/tests/credential/contract.rs @@ -27,7 +27,7 @@ use url::Url; use crate::support::{ MemoryAccountStore, MemorySessionAffinity, MemorySessionExclusions, TestLeaseCoordinator, - account_policy, agent_identity_service, profile, secret, + account_policy, agent_identity_service, catalog_cache, profile, secret, }; fn create_account(store: &Arc, id: &str, token: &str) { @@ -129,6 +129,7 @@ fn selector_with_runtime( http.clone(), OFFICIAL_CODEX_BASE_URL.to_owned(), Arc::clone(&agent_identity), + catalog_cache(), )); let quota = Arc::new(CodexCredentialQuotaService::new( store.repository(), diff --git a/backend/crates/providers/openai/tests/provider/contract.rs b/backend/crates/providers/openai/tests/provider/contract.rs index ab0aaf34b..49e918405 100644 --- a/backend/crates/providers/openai/tests/provider/contract.rs +++ b/backend/crates/providers/openai/tests/provider/contract.rs @@ -28,7 +28,7 @@ use serde_json::{Map, json}; use crate::support::{ MemoryAccountStore, MemorySessionAffinity, MemorySessionExclusions, TestLeaseCoordinator, - account_policy, agent_identity_service_with_pool, profile, secret, + account_policy, agent_identity_service_with_pool, catalog_cache, profile, secret, }; fn wire_profile() -> CodexWireProfileState { @@ -66,6 +66,7 @@ fn provider_with_affinity( http.clone(), OFFICIAL_CODEX_BASE_URL.to_owned(), Arc::clone(&agent_identity), + catalog_cache(), )); let quota = Arc::new(CodexCredentialQuotaService::new( store.repository(), diff --git a/backend/crates/providers/openai/tests/support.rs b/backend/crates/providers/openai/tests/support.rs index 4ad52cdc3..4aa36a825 100644 --- a/backend/crates/providers/openai/tests/support.rs +++ b/backend/crates/providers/openai/tests/support.rs @@ -10,16 +10,17 @@ use async_trait::async_trait; use futures::future::BoxFuture; use gateway_core::engine::credential::{ AccountAvailability, AccountRuntimeSignals, AccountStateChange, CredentialCasOutcome, - CredentialCasUpdate, CredentialRevision, LoadedCredential, NewProviderAccount, ProviderAccount, - ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, QuotaObservation, - QuotaWriteOutcome, + CredentialCasUpdate, CredentialRevision, LoadedCredential, NewProviderAccount, + OpaqueProviderData, ProviderAccount, ProviderAccountId, ProviderAccountStore, + ProviderAccountUpdate, QuotaObservation, QuotaWriteOutcome, }; use gateway_core::error::{StoreError, StoreErrorKind}; use gateway_core::provider_ports::{ - ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshPolicy, - ProviderRuntimePolicyPort, ProviderSchedulingLeaseRequest, ProviderSchedulingState, - ProviderSessionAffinityKey, ProviderSessionAffinityPort, ProviderSessionExclusionPort, - ProviderSessionExclusions, ProviderStoreError, + ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderLeaseAcquisition, ProviderLeasePort, + ProviderLeaseRequest, ProviderRefreshPolicy, ProviderRuntimePolicyPort, + ProviderSchedulingLeaseRequest, ProviderSchedulingState, ProviderSessionAffinityKey, + ProviderSessionAffinityPort, ProviderSessionExclusionPort, ProviderSessionExclusions, + ProviderStoreError, }; use gateway_core::routing::ProviderKind; use provider_openai::credential::{ @@ -579,6 +580,46 @@ impl ProviderSessionExclusionPort for MemorySessionExclusions { } } +#[derive(Default)] +pub(crate) struct MemoryCatalogCache { + values: Mutex>, +} + +impl ProviderCatalogCachePort for MemoryCatalogCache { + fn replace<'a>( + &'a self, + key: &'a ProviderCatalogCacheKey, + catalog: &'a OpaqueProviderData, + _ttl: Duration, + ) -> BoxFuture<'a, Result<(), ProviderStoreError>> { + Box::pin(async move { + self.values + .lock() + .expect("catalog cache") + .insert(key.scope().as_str().to_owned(), catalog.clone()); + Ok(()) + }) + } + + fn read<'a>( + &'a self, + key: &'a ProviderCatalogCacheKey, + ) -> BoxFuture<'a, Result, ProviderStoreError>> { + Box::pin(async move { + Ok(self + .values + .lock() + .expect("catalog cache") + .get(key.scope().as_str()) + .cloned()) + }) + } +} + +pub(crate) fn catalog_cache() -> Arc { + Arc::new(MemoryCatalogCache::default()) +} + pub(crate) fn profile(account_id: &str) -> CodexAccountProfile { let now = chrono::Utc::now(); CodexAccountProfile { diff --git a/backend/crates/providers/xai/src/admin.rs b/backend/crates/providers/xai/src/admin.rs index 6e86e8640..f47ae86f1 100644 --- a/backend/crates/providers/xai/src/admin.rs +++ b/backend/crates/providers/xai/src/admin.rs @@ -574,23 +574,15 @@ impl ProviderAdmin for XaiAdminProvider { ) -> Result { let account = self.account(account_id).await?; let catalog = if refresh { - Some( - self.catalog - .refresh_account_catalog(account_id) - .await - .map_err(map_catalog_error)?, - ) - } else { self.catalog - .read_account_catalog(account_id, account.revision()) + .refresh_account_catalog(account_id) + .await + .map_err(map_catalog_error)? + } else { + self.catalog + .cached_or_refresh_account_catalog(&account) .await .map_err(map_catalog_error)? - }; - let Some(catalog) = catalog else { - return Ok(ProviderModels { - models: Vec::new(), - observed_at: None, - }); }; let models = catalog .seed() diff --git a/backend/crates/providers/xai/src/credential/catalog.rs b/backend/crates/providers/xai/src/credential/catalog.rs index 939cef37d..c81ddc9f8 100644 --- a/backend/crates/providers/xai/src/credential/catalog.rs +++ b/backend/crates/providers/xai/src/credential/catalog.rs @@ -7,13 +7,14 @@ use std::time::{Duration, Instant, SystemTime}; use async_trait::async_trait; use chrono::{DateTime, Utc}; -use futures::{StreamExt, TryStreamExt, stream}; use gateway_core::engine::credential::{ AccountAvailability, AccountQuotaSignals, CredentialRevision, OpaqueProviderData, ProviderAccount, ProviderAccountId, QuotaObservation, }; use gateway_core::engine::provider::ProviderCatalogGeneration; -use gateway_core::provider_ports::{ProviderCatalogCacheKey, ProviderCatalogCachePort}; +use gateway_core::provider_ports::{ + ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCatalogScope, +}; use gateway_core::routing::{ConfigRevision, ProviderKind}; use tokio::sync::Mutex; @@ -27,8 +28,8 @@ use crate::{ parse_grok_billing, }; -const MAX_CONCURRENT_CATALOG_REQUESTS: usize = 8; const CATALOG_CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(5 * 60); +const MAX_CATALOG_FETCH_ATTEMPTS: usize = 3; const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60); const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5); @@ -599,39 +600,56 @@ impl fmt::Debug for GrokCredentialCatalogSeed { } } -/// Redis/内存 TTL cache 中的一条可重建 account catalog。 +/// xAI 以套餐划分的模型目录作用域。 +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] +pub struct GrokCatalogScope(ProviderCatalogScope); + +impl GrokCatalogScope { + /// 从账号已验证的套餐事实构造目录作用域。 + pub fn for_account(account: &ProviderAccount) -> Result { + let plan = account + .plan_type() + .map(str::trim) + .filter(|plan| !plan.is_empty()) + .map(str::to_ascii_lowercase) + .unwrap_or_else(|| "unknown".to_owned()); + ProviderCatalogScope::new(format!("plan:{plan}")) + .map(Self) + .map_err(|_| GrokCatalogCacheError::InvalidData) + } + + /// 返回稳定的 Provider-owned 作用域。 + #[must_use] + pub fn as_str(&self) -> &str { + self.0.as_str() + } +} + +/// Redis/内存 TTL cache 中的一条可重建套餐 catalog。 #[derive(Clone, Debug, PartialEq, Eq)] -pub struct GrokAccountCatalog { - account_id: ProviderAccountId, - revision: CredentialRevision, +pub struct GrokPlanCatalog { + scope: GrokCatalogScope, observed_at: DateTime, seed: GrokCredentialCatalogSeed, } -impl GrokAccountCatalog { +impl GrokPlanCatalog { #[must_use] pub const fn new( - account_id: ProviderAccountId, - revision: CredentialRevision, + scope: GrokCatalogScope, observed_at: DateTime, seed: GrokCredentialCatalogSeed, ) -> Self { Self { - account_id, - revision, + scope, observed_at, seed, } } #[must_use] - pub const fn account_id(&self) -> &ProviderAccountId { - &self.account_id - } - - #[must_use] - pub const fn revision(&self) -> CredentialRevision { - self.revision + pub const fn scope(&self) -> &GrokCatalogScope { + &self.scope } #[must_use] @@ -648,18 +666,16 @@ impl GrokAccountCatalog { /// Provider-owned catalog cache;实现只能保存可重建 TTL 数据。 #[async_trait] pub trait GrokCredentialCatalogCache: Send + Sync { - async fn replace(&self, catalog: GrokAccountCatalog) -> Result<(), GrokCatalogCacheError>; + async fn replace(&self, catalog: GrokPlanCatalog) -> Result<(), GrokCatalogCacheError>; async fn read( &self, - account_id: &ProviderAccountId, - revision: CredentialRevision, - ) -> Result, GrokCatalogCacheError>; + scope: &GrokCatalogScope, + ) -> Result, GrokCatalogCacheError>; async fn observed_model_support( &self, - account_id: &ProviderAccountId, - revision: CredentialRevision, + scope: &GrokCatalogScope, model: &str, ) -> Result, GrokCatalogCacheError>; } @@ -679,17 +695,17 @@ impl GrokCatalogCache { }) } - fn key( - &self, - account_id: &ProviderAccountId, - revision: CredentialRevision, - ) -> ProviderCatalogCacheKey { - ProviderCatalogCacheKey::new(self.provider_kind.clone(), account_id.clone(), revision) + fn key(&self, scope: &GrokCatalogScope) -> ProviderCatalogCacheKey { + ProviderCatalogCacheKey::new(self.provider_kind.clone(), scope.0.clone()) } - fn encode(catalog: &GrokAccountCatalog) -> OpaqueProviderData { + fn encode(catalog: &GrokPlanCatalog) -> OpaqueProviderData { let mut document = serde_json::Map::new(); document.insert("version".to_owned(), serde_json::Value::from(1)); + document.insert( + "scope".to_owned(), + serde_json::Value::String(catalog.scope().as_str().to_owned()), + ); document.insert( "observedAt".to_owned(), serde_json::Value::String(catalog.observed_at().to_rfc3339()), @@ -716,14 +732,21 @@ impl GrokCatalogCache { } fn decode( - account_id: &ProviderAccountId, - revision: CredentialRevision, + scope: &GrokCatalogScope, document: OpaqueProviderData, - ) -> Result { + ) -> Result { let mut fields = document.into_inner(); if fields.remove("version").and_then(|value| value.as_u64()) != Some(1) { return Err(GrokCatalogCacheError::InvalidData); } + if fields + .remove("scope") + .and_then(|value| value.as_str().map(ToOwned::to_owned)) + .as_deref() + != Some(scope.as_str()) + { + return Err(GrokCatalogCacheError::InvalidData); + } let observed_at = fields .remove("observedAt") .and_then(|value| value.as_str().map(ToOwned::to_owned)) @@ -752,21 +775,16 @@ impl GrokCatalogCache { } let seed = GrokCredentialCatalogSeed::new(models, etag) .map_err(|_| GrokCatalogCacheError::InvalidData)?; - Ok(GrokAccountCatalog::new( - account_id.clone(), - revision, - observed_at, - seed, - )) + Ok(GrokPlanCatalog::new(scope.clone(), observed_at, seed)) } } #[async_trait] impl GrokCredentialCatalogCache for GrokCatalogCache { - async fn replace(&self, catalog: GrokAccountCatalog) -> Result<(), GrokCatalogCacheError> { + async fn replace(&self, catalog: GrokPlanCatalog) -> Result<(), GrokCatalogCacheError> { self.port .replace( - &self.key(catalog.account_id(), catalog.revision()), + &self.key(catalog.scope()), &Self::encode(&catalog), CATALOG_CACHE_TTL, ) @@ -776,25 +794,23 @@ impl GrokCredentialCatalogCache for GrokCatalogCache { async fn read( &self, - account_id: &ProviderAccountId, - revision: CredentialRevision, - ) -> Result, GrokCatalogCacheError> { + scope: &GrokCatalogScope, + ) -> Result, GrokCatalogCacheError> { self.port - .read(&self.key(account_id, revision)) + .read(&self.key(scope)) .await .map_err(|_| GrokCatalogCacheError::Unavailable)? - .map(|document| Self::decode(account_id, revision, document)) + .map(|document| Self::decode(scope, document)) .transpose() } async fn observed_model_support( &self, - account_id: &ProviderAccountId, - revision: CredentialRevision, + scope: &GrokCatalogScope, model: &str, ) -> Result, GrokCatalogCacheError> { Ok(self - .read(account_id, revision) + .read(scope) .await? .map(|catalog| catalog.seed().permits(model))) } @@ -808,12 +824,12 @@ pub enum GrokCatalogCacheError { InvalidData, } -/// 一次 Provider 同步得到的账号目录和严格模型并集。 +/// 一次 Provider 同步得到的套餐目录和严格模型并集。 #[derive(Clone, Debug)] pub struct GrokCredentialCatalogSnapshot { config_revision: ConfigRevision, observed_at: DateTime, - accounts: Vec, + plans: Vec, models: Vec, } @@ -829,8 +845,8 @@ impl GrokCredentialCatalogSnapshot { } #[must_use] - pub fn accounts(&self) -> &[GrokAccountCatalog] { - &self.accounts + pub fn plans(&self) -> &[GrokPlanCatalog] { + &self.plans } #[must_use] @@ -917,61 +933,80 @@ impl GrokCredentialCatalogService { pub async fn cache_seed( &self, - account_id: ProviderAccountId, - revision: CredentialRevision, + account_id: &ProviderAccountId, seed: GrokCredentialCatalogSeed, ) -> Result<(), GrokCredentialCatalogError> { + let account = self + .repository + .load_current(account_id) + .await + .map_err(|_| GrokCredentialCatalogError::Store)? + .account; + let scope = GrokCatalogScope::for_account(&account) + .map_err(|_| GrokCredentialCatalogError::InvalidCredentialData)?; self.cache - .replace(GrokAccountCatalog::new( - account_id, - revision, - Utc::now(), - seed, - )) + .replace(GrokPlanCatalog::new(scope, Utc::now(), seed)) .await .map_err(|_| GrokCredentialCatalogError::Cache) } - /// 实时刷新一个指定账号的模型集合并写入可重建 cache。 + /// 优先读取套餐目录 cache;缺失时才用当前账号所属套餐的有限候选集实时填充。 + pub async fn cached_or_refresh_account_catalog( + &self, + account: &ProviderAccount, + ) -> Result { + let scope = GrokCatalogScope::for_account(account) + .map_err(|_| GrokCredentialCatalogError::InvalidCredentialData)?; + if let Some(catalog) = self + .cache + .read(&scope) + .await + .map_err(|_| GrokCredentialCatalogError::Cache)? + { + return Ok(catalog); + } + self.refresh_account_catalog(account.id()).await + } + + /// 实时刷新指定账号所属套餐的模型集合,并覆盖可重建 cache。 pub async fn refresh_account_catalog( &self, account_id: &ProviderAccountId, - ) -> Result { - let candidate = self + ) -> Result { + let candidates = self .repository - .load_current(account_id) + .list_loaded_for_provider() .await .map_err(|_| GrokCredentialCatalogError::Store)?; - if !eligible_catalog_candidate(&candidate) { + let mut groups = catalog_candidates_by_scope(candidates)?; + let Some(scope) = groups.iter().find_map(|(scope, candidates)| { + candidates + .iter() + .any(|candidate| candidate.account.id() == account_id) + .then(|| scope.clone()) + }) else { return Err(GrokCredentialCatalogError::NoEligibleCredential); - } - let revision = candidate.account.revision(); - let seed = self - .fetch_seed( - candidate.access_token, - SecretValue::new(candidate.account.upstream_user_id().to_owned()), - candidate - .account - .email() - .map(|email| SecretValue::new(email.to_owned())), - ) - .await?; - let catalog = GrokAccountCatalog::new(account_id.clone(), revision, Utc::now(), seed); - self.cache - .replace(catalog.clone()) - .await - .map_err(|_| GrokCredentialCatalogError::Cache)?; - Ok(catalog) + }; + let mut candidates = groups.remove(&scope).unwrap_or_default(); + candidates.sort_by(|left, right| { + let left_preferred = left.account.id() == account_id; + let right_preferred = right.account.id() == account_id; + right_preferred + .cmp(&left_preferred) + .then_with(|| left.account.id().cmp(right.account.id())) + }); + self.refresh_scope_catalog(scope, candidates).await } - /// 读取一个指定 revision 的已验证模型 cache,不触发上游请求。 + /// 读取当前账号所属套餐的目录 cache,不触发上游请求。 pub async fn read_account_catalog( &self, - account_id: &ProviderAccountId, - revision: CredentialRevision, - ) -> Result, GrokCredentialCatalogError> { + account: &ProviderAccount, + ) -> Result, GrokCredentialCatalogError> { + let scope = GrokCatalogScope::for_account(account) + .map_err(|_| GrokCredentialCatalogError::InvalidCredentialData)?; self.cache - .read(account_id, revision) + .read(&scope) .await .map_err(|_| GrokCredentialCatalogError::Cache) } @@ -985,7 +1020,7 @@ impl GrokCredentialCatalogService { Ok(GrokCredentialCatalogSnapshot { config_revision, observed_at: catalog.observed_at, - accounts: catalog.accounts, + plans: catalog.plans, models: catalog.models, }) } @@ -1004,34 +1039,25 @@ impl GrokCredentialCatalogService { .into_iter() .filter(eligible_catalog_candidate) .collect::>(); - if candidates.is_empty() { - return Err(GrokCredentialCatalogError::NoEligibleCredential); + let groups = catalog_candidates_by_scope(candidates)?; + let mut fetched = Vec::with_capacity(groups.len()); + for (scope, candidates) in groups { + fetched.push(self.fetch_scope_catalog(scope, candidates).await?); } - - let wire_profile = self.wire_profile.clone(); - let mut fetched = stream::iter(candidates.into_iter().map(|candidate| { - let client = Arc::clone(&self.client); - let wire_profile = wire_profile.clone(); - async move { fetch_candidate_catalog(client, candidate, wire_profile).await } - })) - .buffer_unordered(MAX_CONCURRENT_CATALOG_REQUESTS) - .try_collect::>() - .await?; - fetched.sort_by(|left, right| left.account_id.cmp(&right.account_id)); + fetched.sort_by(|left, right| left.scope.cmp(&right.scope)); let models = strict_model_union(&fetched)?; let observed_at = Utc::now(); - let accounts = fetched + let plans = fetched .into_iter() - .map(|fetched| GrokAccountCatalog { - account_id: fetched.account_id, - revision: fetched.revision, + .map(|fetched| GrokPlanCatalog { + scope: fetched.scope, observed_at, seed: fetched.seed, }) .collect::>(); - for account in &accounts { + for plan in &plans { self.cache - .replace(account.clone()) + .replace(plan.clone()) .await .map_err(|_| GrokCredentialCatalogError::Cache)?; } @@ -1039,11 +1065,47 @@ impl GrokCredentialCatalogService { Ok(FetchedProviderCatalog { observed_at, - accounts, + plans, models, }) } + async fn refresh_scope_catalog( + &self, + scope: GrokCatalogScope, + candidates: Vec, + ) -> Result { + let fetched = self.fetch_scope_catalog(scope.clone(), candidates).await?; + let catalog = GrokPlanCatalog::new(scope, Utc::now(), fetched.seed); + self.cache + .replace(catalog.clone()) + .await + .map_err(|_| GrokCredentialCatalogError::Cache)?; + Ok(catalog) + } + + async fn fetch_scope_catalog( + &self, + scope: GrokCatalogScope, + candidates: Vec, + ) -> Result { + let mut last_error = None; + for candidate in candidates.into_iter().take(MAX_CATALOG_FETCH_ATTEMPTS) { + match fetch_candidate_catalog( + Arc::clone(&self.client), + scope.clone(), + candidate, + self.wire_profile.clone(), + ) + .await + { + Ok(catalog) => return Ok(catalog), + Err(error) => last_error = Some(error), + } + } + Err(last_error.unwrap_or(GrokCredentialCatalogError::NoEligibleCredential)) + } + fn publish_models(&self, models: &[GrokCatalogModel]) { let mut published = self .published @@ -1069,18 +1131,35 @@ impl fmt::Debug for GrokCredentialCatalogService { } struct FetchedCredentialCatalog { - account_id: ProviderAccountId, - revision: CredentialRevision, + scope: GrokCatalogScope, snapshot: GrokModelCatalogSnapshot, seed: GrokCredentialCatalogSeed, } struct FetchedProviderCatalog { observed_at: DateTime, - accounts: Vec, + plans: Vec, models: Vec, } +fn catalog_candidates_by_scope( + candidates: Vec, +) -> Result>, GrokCredentialCatalogError> { + let mut groups = BTreeMap::>::new(); + for candidate in candidates.into_iter().filter(eligible_catalog_candidate) { + let scope = GrokCatalogScope::for_account(&candidate.account) + .map_err(|_| GrokCredentialCatalogError::InvalidCredentialData)?; + groups.entry(scope).or_default().push(candidate); + } + for candidates in groups.values_mut() { + candidates.sort_by(|left, right| left.account.id().cmp(right.account.id())); + } + if groups.is_empty() { + return Err(GrokCredentialCatalogError::NoEligibleCredential); + } + Ok(groups) +} + fn eligible_catalog_candidate(candidate: &LoadedGrokCredential) -> bool { let account = &candidate.account; let now = SystemTime::now(); @@ -1104,6 +1183,7 @@ fn eligible_catalog_candidate(candidate: &LoadedGrokCredential) -> bool { async fn fetch_candidate_catalog( client: Arc, + scope: GrokCatalogScope, candidate: LoadedGrokCredential, wire_profile: XaiWireProfileState, ) -> Result { @@ -1123,8 +1203,7 @@ async fn fetch_candidate_catalog( .map_err(|_| GrokCredentialCatalogError::Upstream)?; let seed = GrokCredentialCatalogSeed::from_snapshot(&snapshot)?; Ok(FetchedCredentialCatalog { - account_id: candidate.account.id().clone(), - revision: candidate.account.revision(), + scope, snapshot, seed, }) diff --git a/backend/crates/providers/xai/src/credential/mod.rs b/backend/crates/providers/xai/src/credential/mod.rs index 92ba782a9..07bc9e33d 100644 --- a/backend/crates/providers/xai/src/credential/mod.rs +++ b/backend/crates/providers/xai/src/credential/mod.rs @@ -26,11 +26,11 @@ pub use config::{ }; pub use catalog::{ - GROK_FREE_ROLLING_WINDOW_SECONDS, GrokAccountCatalog, GrokBillingPresentation, - GrokCatalogCache, GrokCatalogCacheError, GrokCredentialCatalogCache, + GROK_FREE_ROLLING_WINDOW_SECONDS, GrokBillingPresentation, GrokCatalogCache, + GrokCatalogCacheError, GrokCatalogScope, GrokCredentialCatalogCache, GrokCredentialCatalogError, GrokCredentialCatalogSeed, GrokCredentialCatalogService, - GrokCredentialCatalogSnapshot, GrokCredentialQuotaService, GrokQuotaError, GrokQuotaPeriodKind, - GrokQuotaSnapshot, + GrokCredentialCatalogSnapshot, GrokCredentialQuotaService, GrokPlanCatalog, GrokQuotaError, + GrokQuotaPeriodKind, GrokQuotaSnapshot, }; pub use discovery::DiscoveryDocument; pub use error::{ diff --git a/backend/crates/providers/xai/src/credential/refresh.rs b/backend/crates/providers/xai/src/credential/refresh.rs index ca44c26de..3cad66f9c 100644 --- a/backend/crates/providers/xai/src/credential/refresh.rs +++ b/backend/crates/providers/xai/src/credential/refresh.rs @@ -636,10 +636,7 @@ impl GrokCredentialRefreshService { ) .await { - let _ = self - .catalog - .cache_seed(account_id.clone(), record.credential_revision, seed) - .await; + let _ = self.catalog.cache_seed(&account_id, seed).await; } Ok(GrokCredentialRefreshOutcome::Refreshed { account_id, diff --git a/backend/crates/providers/xai/src/credential/selector.rs b/backend/crates/providers/xai/src/credential/selector.rs index 69c38fcc8..6616597da 100644 --- a/backend/crates/providers/xai/src/credential/selector.rs +++ b/backend/crates/providers/xai/src/credential/selector.rs @@ -14,7 +14,7 @@ use gateway_core::provider_ports::{ }; use gateway_core::routing::ProviderKind; -use super::catalog::{GrokCredentialCatalogCache, GrokCredentialQuotaService}; +use super::catalog::{GrokCatalogScope, GrokCredentialCatalogCache, GrokCredentialQuotaService}; use super::repository::GrokCredentialRepository; use super::types::{GrokCredentialAvailability, UpdateGrokCredentialState}; use crate::{ @@ -69,14 +69,14 @@ impl GrokAccountSessionSelector { self.quota.prepare_scheduling(&accounts).await; let mut catalog_eligible = Vec::new(); for account in accounts { - let support = self - .catalog_cache - .observed_model_support( - account.id(), - account.revision(), - request.upstream_model().as_str(), - ) - .await; + let support = match GrokCatalogScope::for_account(&account) { + Ok(scope) => { + self.catalog_cache + .observed_model_support(&scope, request.upstream_model().as_str()) + .await + } + Err(error) => Err(error), + }; if !matches!(support, Ok(Some(false))) { catalog_eligible.push(account); } diff --git a/backend/crates/providers/xai/src/lib.rs b/backend/crates/providers/xai/src/lib.rs index 3ff625a9a..c0556d045 100644 --- a/backend/crates/providers/xai/src/lib.rs +++ b/backend/crates/providers/xai/src/lib.rs @@ -27,8 +27,8 @@ pub use credential::{ AllowedRedirectUri, AuthorizationCallback, AuthorizationCodeGrant, CallbackRejection, ConfigError, CreateGrokCredential, DiscoveryDocument, DueGrokCredential, FailClosedTokenVerifier, FailureClass, FormField, FormValue, GROK_FREE_ROLLING_WINDOW_SECONDS, - GrokAccountCatalog, GrokAccountExport, GrokAccountProfile, GrokAccountSessionSelector, - GrokBillingPresentation, GrokCatalogCache, GrokCatalogCacheError, GrokCredentialAdmin, + GrokAccountExport, GrokAccountProfile, GrokAccountSessionSelector, GrokBillingPresentation, + GrokCatalogCache, GrokCatalogCacheError, GrokCatalogScope, GrokCredentialAdmin, GrokCredentialAvailability, GrokCredentialCatalogCache, GrokCredentialCatalogError, GrokCredentialCatalogSeed, GrokCredentialCatalogService, GrokCredentialCatalogSnapshot, GrokCredentialLifecycle, GrokCredentialQuotaService, GrokCredentialRecord, @@ -37,10 +37,10 @@ pub use credential::{ GrokCredentialRepository, GrokCredentialRepositoryError, GrokOAuthClient, GrokOAuthConfig, GrokOAuthImportCandidate, GrokOAuthImportDocument, GrokOAuthImportEntry, GrokOAuthImportError, GrokOAuthImportMetadata, GrokOAuthImportTokens, GrokOAuthRefreshClient, GrokOAuthSecret, - GrokQuotaError, GrokQuotaPeriodKind, GrokQuotaSnapshot, GrokRefreshFailure, GrokRefreshTokens, - HttpHeader, HttpMethod, OAuthError, OAuthErrorCode, OAuthHttpRequest, OAuthHttpResponse, - OAuthHttpTransport, OAuthOperation, OAuthPrincipal, OFFICIAL_CLIENT_ID, OFFICIAL_ISSUER, - OFFICIAL_REDIRECT_URI, OFFICIAL_SCOPES, PendingAuthorization, Pkce, + GrokPlanCatalog, GrokQuotaError, GrokQuotaPeriodKind, GrokQuotaSnapshot, GrokRefreshFailure, + GrokRefreshTokens, HttpHeader, HttpMethod, OAuthError, OAuthErrorCode, OAuthHttpRequest, + OAuthHttpResponse, OAuthHttpTransport, OAuthOperation, OAuthPrincipal, OFFICIAL_CLIENT_ID, + OFFICIAL_ISSUER, OFFICIAL_REDIRECT_URI, OFFICIAL_SCOPES, PendingAuthorization, Pkce, PreparedGrokCredentialRotation, PreparedGrokCredentialRotationGuard, ProtocolViolation, RedirectUriAllowlist, RefreshTokenGrant, RefreshedTokenSet, ReqwestOidcTokenVerifier, RotateGrokCredential, RotateManagedGrokCredential, SecretValue, TokenCandidate, diff --git a/backend/crates/providers/xai/tests/admin.rs b/backend/crates/providers/xai/tests/admin.rs index fef00fd98..d70d7f904 100644 --- a/backend/crates/providers/xai/tests/admin.rs +++ b/backend/crates/providers/xai/tests/admin.rs @@ -412,9 +412,15 @@ async fn xai_admin_provider_projects_cached_quota_models_and_canonical_export() seed_input(&store, &input).await.expect("create account"); let account = store.account(&input.account_id).expect("stored account"); let record = account_record(&account); + let catalog_cache = Arc::new(TestCatalogCache::default()); + catalog_cache.seed("plan:standard", ["grok-4.5"]); let bundle = provider_xai::initialize( xai_config(), - provider_ports_with(Arc::clone(&store), Arc::new(TestOAuthPending::default())), + provider_ports_with_catalog( + Arc::clone(&store), + Arc::new(TestOAuthPending::default()), + catalog_cache, + ), ) .await .expect("xAI bundle"); @@ -457,7 +463,7 @@ async fn xai_admin_provider_projects_cached_quota_models_and_canonical_export() .models(account.id(), false) .await .expect("cached models"); - assert!(models.models.is_empty()); + assert_eq!(models.models[0].id.as_str(), "grok-4.5"); let loaded = store .load_credential(account.id(), account.revision()) .await @@ -578,13 +584,21 @@ fn provider_ports() -> ProviderStorePorts { fn provider_ports_with( accounts: Arc, pending: Arc, +) -> ProviderStorePorts { + provider_ports_with_catalog(accounts, pending, Arc::new(TestCatalogCache::default())) +} + +fn provider_ports_with_catalog( + accounts: Arc, + pending: Arc, + catalog_cache: Arc, ) -> ProviderStorePorts { ProviderStorePorts::new( accounts, Arc::new(TestLeases), Arc::new(TestSessionAffinity), Arc::new(TestSessionExclusions), - Arc::new(TestCatalogCache::default()), + catalog_cache, Arc::new(TestCredentialState), Arc::new(TestCooldown), Arc::new(TestRuntimePolicy), @@ -689,7 +703,7 @@ impl ProviderCatalogCachePort for TestCatalogCache { self.values .lock() .expect("catalog cache") - .insert(key.account_id().to_string(), catalog.clone()); + .insert(key.scope().as_str().to_owned(), catalog.clone()); Ok(()) }) } @@ -703,12 +717,37 @@ impl ProviderCatalogCachePort for TestCatalogCache { .values .lock() .expect("catalog cache") - .get(key.account_id().as_str()) + .get(key.scope().as_str()) .cloned()) }) } } +impl TestCatalogCache { + fn seed(&self, scope: &str, models: impl IntoIterator) { + let mut document = Map::new(); + document.insert("version".to_owned(), Value::from(1)); + document.insert("scope".to_owned(), Value::String(scope.to_owned())); + document.insert( + "observedAt".to_owned(), + Value::String(Utc::now().to_rfc3339()), + ); + document.insert( + "models".to_owned(), + Value::Array( + models + .into_iter() + .map(|model| Value::String(model.to_owned())) + .collect(), + ), + ); + self.values + .lock() + .expect("catalog cache") + .insert(scope.to_owned(), OpaqueProviderData::new(document)); + } +} + struct TestCredentialState; impl ProviderCredentialStatePort for TestCredentialState { diff --git a/backend/crates/providers/xai/tests/credential/catalog.rs b/backend/crates/providers/xai/tests/credential/catalog.rs index 485382c6a..a6ecc4012 100644 --- a/backend/crates/providers/xai/tests/credential/catalog.rs +++ b/backend/crates/providers/xai/tests/credential/catalog.rs @@ -13,11 +13,11 @@ use gateway_core::routing::ConfigRevision; use provider_xai::{ GrokBillingRequest, GrokBillingTransport, GrokBillingTransportError, GrokBillingTransportErrorKind, GrokBillingTransportFuture, GrokBillingTransportResponse, - GrokCredentialCatalogCache, GrokCredentialCatalogError, GrokCredentialCatalogSeed, - GrokCredentialRepository, GrokModelCatalogRequest, GrokModelCatalogTransport, - GrokModelCatalogTransportError, GrokModelCatalogTransportErrorKind, - GrokModelCatalogTransportFuture, GrokModelCatalogTransportResponse, GrokQuotaError, - SecretValue, + GrokCatalogScope, GrokCredentialCatalogCache, GrokCredentialCatalogError, + GrokCredentialCatalogSeed, GrokCredentialRepository, GrokModelCatalogRequest, + GrokModelCatalogTransport, GrokModelCatalogTransportError, GrokModelCatalogTransportErrorKind, + GrokModelCatalogTransportFuture, GrokModelCatalogTransportResponse, GrokPlanCatalog, + GrokQuotaError, SecretValue, }; use crate::support::{ @@ -28,6 +28,7 @@ const OFFICIAL_FIXTURE: &[u8] = include_bytes!("../transport/catalog/fixtures/official_grok_models_snapshot.json"); struct QueueCatalogTransport { + calls: AtomicUsize, responses: Mutex>>, } @@ -35,6 +36,7 @@ struct QueueCatalogTransport { impl QueueCatalogTransport { fn from_bodies(bodies: impl IntoIterator>) -> Arc { Arc::new(Self { + calls: AtomicUsize::new(0), responses: Mutex::new( bodies .into_iter() @@ -46,15 +48,32 @@ impl QueueCatalogTransport { fn failure() -> Arc { Arc::new(Self { + calls: AtomicUsize::new(0), responses: Mutex::new(VecDeque::from([Err(GrokModelCatalogTransportError::new( GrokModelCatalogTransportErrorKind::Unavailable, ))])), }) } + + fn from_results( + responses: impl IntoIterator< + Item = Result, + >, + ) -> Arc { + Arc::new(Self { + calls: AtomicUsize::new(0), + responses: Mutex::new(responses.into_iter().collect()), + }) + } + + fn calls(&self) -> usize { + self.calls.load(Ordering::SeqCst) + } } impl GrokModelCatalogTransport for QueueCatalogTransport { fn execute(&self, _: GrokModelCatalogRequest) -> GrokModelCatalogTransportFuture<'_> { + self.calls.fetch_add(1, Ordering::SeqCst); let response = self .responses .lock() @@ -185,8 +204,8 @@ async fn concurrent_cold_scheduling_hydration_reads_quota_once() { } #[tokio::test] -async fn synchronization_caches_each_account_and_returns_strict_union() { - let (_, repository) = +async fn synchronization_caches_each_plan_and_returns_strict_union() { + let (store, repository) = repository_with_accounts(&[("catalog-a", "subject-a"), ("catalog-b", "subject-b")]).await; let cache = MemoryGrokCatalogCache::shared(); let cache_port: Arc = cache.clone(); @@ -206,27 +225,26 @@ async fn synchronization_caches_each_account_and_returns_strict_union() { service.query_models().await.expect("same catalog sync"); assert_eq!(service.catalog_generation().get(), 1); - assert_eq!(snapshot.accounts().len(), 2); + assert_eq!(snapshot.plans().len(), 1); assert_eq!(snapshot.models().len(), 1); assert_eq!(snapshot.models()[0].request_model().as_str(), "grok-4.5"); - for id in [account_id("catalog-a"), account_id("catalog-b")] { - assert_eq!( - cache - .observed_model_support( - &id, - gateway_core::engine::credential::CredentialRevision::new(1).expect("revision"), - "grok-4.5", - ) - .await - .expect("cache lookup"), - Some(true) - ); - } + let account = store + .account(&account_id("catalog-a")) + .expect("created account"); + let scope = GrokCatalogScope::for_account(&account).expect("catalog scope"); + assert_eq!( + cache + .observed_model_support(&scope, "grok-4.5") + .await + .expect("cache lookup"), + Some(true) + ); } #[tokio::test] async fn single_account_catalog_refresh_and_read_use_provider_cache_boundary() { - let (_, repository) = repository_with_accounts(&[("account-models", "subject-models")]).await; + let (store, repository) = + repository_with_accounts(&[("account-models", "subject-models")]).await; let cache = MemoryGrokCatalogCache::shared(); let cache_port: Arc = cache; let service = crate::support::grok_catalog_service( @@ -242,8 +260,9 @@ async fn single_account_catalog_refresh_and_read_use_provider_cache_boundary() { let cached = service .read_account_catalog( - &account_id("account-models"), - CredentialRevision::new(1).expect("revision"), + &store + .account(&account_id("account-models")) + .expect("created account"), ) .await .expect("read cache") @@ -253,7 +272,7 @@ async fn single_account_catalog_refresh_and_read_use_provider_cache_boundary() { #[tokio::test] async fn single_account_catalog_read_miss_does_not_call_upstream() { - let (_, repository) = + let (store, repository) = repository_with_accounts(&[("account-models-miss", "subject-models")]).await; let service = crate::support::grok_catalog_service( repository, @@ -263,8 +282,9 @@ async fn single_account_catalog_read_miss_does_not_call_upstream() { assert!( service .read_account_catalog( - &account_id("account-models-miss"), - CredentialRevision::new(1).expect("revision"), + &store + .account(&account_id("account-models-miss")) + .expect("created account"), ) .await .expect("read cache") @@ -272,6 +292,124 @@ async fn single_account_catalog_read_miss_does_not_call_upstream() { ); } +#[tokio::test] +async fn cached_account_catalog_uses_shared_plan_cache_without_upstream_call() { + let (store, repository) = + repository_with_accounts(&[("catalog-cache-hit", "subject-cache-hit")]).await; + let account = store + .account(&account_id("catalog-cache-hit")) + .expect("created account"); + let scope = GrokCatalogScope::for_account(&account).expect("catalog scope"); + let cache = MemoryGrokCatalogCache::shared(); + cache + .replace(GrokPlanCatalog::new( + scope, + chrono::Utc::now(), + GrokCredentialCatalogSeed::new(["grok-4.5"], None).expect("seed"), + )) + .await + .expect("cache catalog"); + let transport = QueueCatalogTransport::from_bodies([]); + let service = crate::support::grok_catalog_service(repository, transport.clone(), cache); + + let catalog = service + .cached_or_refresh_account_catalog(&account) + .await + .expect("cached catalog"); + + assert_eq!(catalog.seed().models(), ["grok-4.5"]); + assert_eq!(transport.calls(), 0); +} + +#[tokio::test] +async fn manual_catalog_refresh_replaces_the_shared_plan_cache() { + let (store, repository) = + repository_with_accounts(&[("catalog-manual-refresh", "subject-manual-refresh")]).await; + let account = store + .account(&account_id("catalog-manual-refresh")) + .expect("created account"); + let scope = GrokCatalogScope::for_account(&account).expect("catalog scope"); + let cache = MemoryGrokCatalogCache::shared(); + cache + .replace(GrokPlanCatalog::new( + scope, + chrono::Utc::now(), + GrokCredentialCatalogSeed::new(["grok-before-refresh"], None).expect("seed"), + )) + .await + .expect("cache catalog"); + let transport = QueueCatalogTransport::from_bodies([OFFICIAL_FIXTURE.to_vec()]); + let service = crate::support::grok_catalog_service(repository, transport.clone(), cache); + + let refreshed = service + .refresh_account_catalog(account.id()) + .await + .expect("manual refresh"); + + assert_eq!(refreshed.seed().models(), ["grok-4.5"]); + assert_eq!(transport.calls(), 1); +} + +#[tokio::test] +async fn catalog_refresh_falls_back_within_the_same_plan() { + let (_, repository) = repository_with_accounts(&[ + ("catalog-fallback-a", "subject-fallback-a"), + ("catalog-fallback-b", "subject-fallback-b"), + ]) + .await; + let transport = QueueCatalogTransport::from_results([ + Err(GrokModelCatalogTransportError::new( + GrokModelCatalogTransportErrorKind::Unavailable, + )), + Ok(GrokModelCatalogTransportResponse::new( + OFFICIAL_FIXTURE.to_vec(), + None, + )), + ]); + let service = crate::support::grok_catalog_service( + repository, + transport.clone(), + MemoryGrokCatalogCache::shared(), + ); + + let catalog = service + .refresh_account_catalog(&account_id("catalog-fallback-a")) + .await + .expect("second account fills the plan catalog"); + + assert_eq!(catalog.seed().models(), ["grok-4.5"]); + assert_eq!(transport.calls(), 2); +} + +#[tokio::test] +async fn catalog_refresh_stops_after_three_failed_accounts_in_one_plan() { + let (_, repository) = repository_with_accounts(&[ + ("catalog-limit-a", "subject-limit-a"), + ("catalog-limit-b", "subject-limit-b"), + ("catalog-limit-c", "subject-limit-c"), + ("catalog-limit-d", "subject-limit-d"), + ]) + .await; + let transport = QueueCatalogTransport::from_results((0..3).map(|_| { + Err(GrokModelCatalogTransportError::new( + GrokModelCatalogTransportErrorKind::Unavailable, + )) + })); + let service = crate::support::grok_catalog_service( + repository, + transport.clone(), + MemoryGrokCatalogCache::shared(), + ); + + let error = service + .refresh_account_catalog(&account_id("catalog-limit-a")) + .await + .expect_err("three failures stop the refresh"); + + assert!(matches!(error, GrokCredentialCatalogError::Upstream)); + assert_eq!(transport.calls(), 3); +} + #[tokio::test] async fn disabled_accounts_are_not_sent_to_catalog_transport() { let (store, repository) = repository_with_accounts(&[("disabled", "subject-disabled")]).await; @@ -342,8 +480,19 @@ async fn one_upstream_failure_rejects_the_whole_catalog_cycle() { #[tokio::test] async fn conflicting_facts_for_same_slug_fail_closed() { - let (_, repository) = + let (store, repository) = repository_with_accounts(&[("conflict-a", "subject-a"), ("conflict-b", "subject-b")]).await; + let account_id = account_id("conflict-b"); + let account = store.account(&account_id).expect("created account"); + store + .update_account(ProviderAccountUpdate { + account_id, + name: account.name().to_owned(), + email: account.email().map(str::to_owned), + plan_type: Some("premium".to_owned()), + }) + .await + .expect("separate plan catalog"); let mut conflicting: serde_json::Value = serde_json::from_slice(OFFICIAL_FIXTURE).expect("fixture JSON"); conflicting["data"][0]["name"] = serde_json::json!("Different name"); diff --git a/backend/crates/providers/xai/tests/credential/selector.rs b/backend/crates/providers/xai/tests/credential/selector.rs index 94f1ff277..a1033f93b 100644 --- a/backend/crates/providers/xai/tests/credential/selector.rs +++ b/backend/crates/providers/xai/tests/credential/selector.rs @@ -6,17 +6,18 @@ use chrono::Utc; use gateway_core::engine::credential::{ AccountAttemptFeedback, AccountAvailability, AccountFeedbackStats, AccountRuntimeSignals, AccountSelectionPolicy, CredentialCasOutcome, CredentialRevision, OpaqueProviderData, - ProviderAccountId, ProviderAccountStore, QuotaObservation, QuotaWriteOutcome, RotationStrategy, + ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, QuotaObservation, + QuotaWriteOutcome, RotationStrategy, }; use gateway_core::provider_ports::{ ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderSchedulingState, ProviderStoreError, }; use provider_xai::{ - GrokAccountCatalog, GrokAccountSessionSelector, GrokBillingRequest, GrokBillingTransport, + GrokAccountSessionSelector, GrokBillingRequest, GrokBillingTransport, GrokBillingTransportError, GrokBillingTransportErrorKind, GrokBillingTransportFuture, - GrokCredentialAdmin, GrokCredentialAvailability, GrokCredentialCatalogCache, - GrokCredentialCatalogSeed, GrokCredentialFailure, GrokCredentialRepository, + GrokCatalogScope, GrokCredentialAdmin, GrokCredentialAvailability, GrokCredentialCatalogCache, + GrokCredentialCatalogSeed, GrokCredentialFailure, GrokCredentialRepository, GrokPlanCatalog, GrokSessionSelection, GrokSessionSelector, GrokSessionSelectorError, RotateManagedGrokCredential, UpdateGrokCredentialState, }; @@ -111,10 +112,10 @@ impl SelectorFixture { }) .await .expect("ready account"); + let account = store.account(&input.account_id).expect("created account"); cache - .replace(GrokAccountCatalog::new( - input.account_id.clone(), - CredentialRevision::new(1).expect("revision"), + .replace(GrokPlanCatalog::new( + GrokCatalogScope::for_account(&account).expect("catalog scope"), Utc::now(), GrokCredentialCatalogSeed::new(["grok-4.5"], None).expect("catalog"), )) @@ -452,16 +453,47 @@ async fn stale_catalog_revision_does_not_block_transparent_request() { #[tokio::test] async fn explicit_catalog_non_membership_excludes_only_unsupported_account() { let fixture = SelectorFixture::new(&["aaa-unsupported", "zzz-supported"]).await; + let supported_id = account_id("zzz-supported"); + let supported = fixture + .store + .account(&supported_id) + .expect("supported account"); + fixture + .store + .update_account(ProviderAccountUpdate { + account_id: supported_id.clone(), + name: supported.name().to_owned(), + email: supported.email().map(str::to_owned), + plan_type: Some("premium".to_owned()), + }) + .await + .expect("move supported account to another plan"); + let unsupported = fixture + .store + .account(&account_id("aaa-unsupported")) + .expect("unsupported account"); + let supported = fixture + .store + .account(&supported_id) + .expect("supported account"); fixture .cache - .replace(GrokAccountCatalog::new( - account_id("aaa-unsupported"), - CredentialRevision::new(1).expect("revision"), + .replace(GrokPlanCatalog::new( + GrokCatalogScope::for_account(&unsupported).expect("catalog scope"), Utc::now(), GrokCredentialCatalogSeed::new(["grok-other"], None).expect("catalog"), )) .await .expect("replace catalog"); + fixture + .cache + .replace(GrokPlanCatalog::new( + GrokCatalogScope::for_account(&supported).expect("catalog scope"), + Utc::now(), + GrokCredentialCatalogSeed::new(["grok-4.5"], None).expect("catalog"), + )) + .await + .expect("cache supported plan"); let session = fixture .selector diff --git a/backend/crates/providers/xai/tests/support.rs b/backend/crates/providers/xai/tests/support.rs index 9afd9cd01..155dc05ce 100644 --- a/backend/crates/providers/xai/tests/support.rs +++ b/backend/crates/providers/xai/tests/support.rs @@ -21,12 +21,12 @@ use gateway_core::provider_ports::{ }; use gateway_core::routing::ProviderKind; use provider_xai::{ - CreateGrokCredential, GROK_BILLING_URL, GrokAccountCatalog, GrokAccountProfile, - GrokBillingTransport, GrokCatalogCacheError, GrokCredentialAdmin, GrokCredentialAvailability, + CreateGrokCredential, GROK_BILLING_URL, GrokAccountProfile, GrokBillingTransport, + GrokCatalogCacheError, GrokCatalogScope, GrokCredentialAdmin, GrokCredentialAvailability, GrokCredentialCatalogCache, GrokCredentialCatalogService, GrokCredentialQuotaService, GrokCredentialRepository, GrokCredentialRepositoryError, GrokEndpointPolicy, - GrokModelCatalogTransport, GrokOAuthSecret, GrokReqwestTransportBuildError, SecretValue, - XaiConfig, XaiWireProfileState, + GrokModelCatalogTransport, GrokOAuthSecret, GrokPlanCatalog, GrokReqwestTransportBuildError, + SecretValue, XaiConfig, XaiWireProfileState, }; use reqwest::Client; use reqwest::redirect::Policy; @@ -474,7 +474,7 @@ fn rebuild_account(previous: &ProviderAccount, replacement: AccountReplacement) #[derive(Default)] pub struct MemoryGrokCatalogCache { - entries: Mutex>, + entries: Mutex>, } impl MemoryGrokCatalogCache { @@ -485,29 +485,25 @@ impl MemoryGrokCatalogCache { #[async_trait] impl GrokCredentialCatalogCache for MemoryGrokCatalogCache { - async fn replace(&self, catalog: GrokAccountCatalog) -> Result<(), GrokCatalogCacheError> { - lock(&self.entries).insert((catalog.account_id().clone(), catalog.revision()), catalog); + async fn replace(&self, catalog: GrokPlanCatalog) -> Result<(), GrokCatalogCacheError> { + lock(&self.entries).insert(catalog.scope().clone(), catalog); Ok(()) } async fn read( &self, - account_id: &ProviderAccountId, - revision: CredentialRevision, - ) -> Result, GrokCatalogCacheError> { - Ok(lock(&self.entries) - .get(&(account_id.clone(), revision)) - .cloned()) + scope: &GrokCatalogScope, + ) -> Result, GrokCatalogCacheError> { + Ok(lock(&self.entries).get(scope).cloned()) } async fn observed_model_support( &self, - account_id: &ProviderAccountId, - revision: CredentialRevision, + scope: &GrokCatalogScope, model: &str, ) -> Result, GrokCatalogCacheError> { Ok(lock(&self.entries) - .get(&(account_id.clone(), revision)) + .get(scope) .map(|catalog| catalog.seed().permits(model))) } }