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

feat(catalog): share provider model catalogs by plan

Browse Source 
Serve connection-test model lists from a provider-plan Redis cache, refresh manually through bounded same-plan candidates, and keep periodic catalog refreshes isolated from OpenAI ETag handling.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-25 00:33:12 +08:00
co-authored by Codex
parent a9f9125caf
commit 8eea10d509
23 changed files with 1156 additions and 426 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-core/src/provider_ports/mod.rs
+32 -17
View File
@@ -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<String>) -> Result<Self, IdentifierError> {
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
}
}
backend/crates/gateway-store/src/redis/credential_state.rs
+31 -39
View File
@@ -71,18 +71,17 @@ pub trait CredentialStateRepository: Send + Sync {
async fn clear_credential_state(&self, provider_account_id: &str) -> StoreResult<bool>;
}
/// 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<Option<OpaqueProviderData>>;
}
@@ -133,15 +132,12 @@ impl RedisCredentialStateRepository {
Ok(format!("{}:account:{fingerprint}:state", self.namespace))
}
fn catalog_key(&self, key: &ProviderAccountCatalogCacheKey) -> StoreResult<String> {
fn catalog_key(&self, key: &RedisProviderCatalogCacheKey) -> StoreResult<String> {
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<Option<OpaqueProviderData>> {
let mut connection = self.connection.clone();
let payload = redis::cmd("GET")
.arg(self.catalog_key(key)?)
.query_async::<Option<Vec<u8>>>(&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<Option<OpaqueProviderData>, 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(),
}
}
backend/crates/gateway-store/tests/redis/credential_state.rs
+6 -7
View File
@@ -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<T: ProviderAccountCatalogCacheRepository>() {}
fn assert_port<T: ProviderCatalogCacheRepository>() {}
assert_port::<RedisCredentialStateRepository>();
}
@@ -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());
}
backend/crates/providers/openai/src/admin.rs
+9 -12
View File
@@ -440,24 +440,21 @@ impl ProviderAdmin for OpenAiAdminProvider {
refresh: bool,
) -> Result<ProviderModels, ProviderAdminError> {
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::<Utc>::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::<Result<Vec<_>, ProviderAdminError>>()?;
Ok(ProviderModels {
models,
observed_at,
observed_at: Some(catalog.observed_at()),
})
}
backend/crates/providers/openai/src/credential/catalog.rs
+293 -109
View File
@@ -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<Self, CodexCredentialCatalogError> {
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<Utc>,
models: Vec<String>,
}
impl CodexPlanCatalog {
#[must_use]
pub const fn new(
scope: CodexCatalogScope,
observed_at: DateTime<Utc>,
models: Vec<String>,
) -> 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<Utc> {
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<CodexCatalogModel>,
etag: Option<String>,
@@ -125,6 +197,7 @@ pub struct CodexCredentialCatalogService {
http: reqwest::Client,
base_url: String,
agent_identity: Arc<CodexAgentIdentityTaskService>,
plan_catalog_cache: Arc<dyn ProviderCatalogCachePort>,
cache: Arc<RwLock<CatalogCacheState>>,
etags: Arc<Mutex<CatalogEtagState>>,
etag_notification: Arc<Notify>,
@@ -137,6 +210,7 @@ impl CodexCredentialCatalogService {
http: reqwest::Client,
base_url: String,
agent_identity: Arc<CodexAgentIdentityTaskService>,
plan_catalog_cache: Arc<dyn ProviderCatalogCachePort>,
) -> 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<CodexPlanCatalog, CodexCredentialCatalogError> {
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<Vec<String>, CodexCredentialCatalogError> {
) -> Result<CodexPlanCatalog, CodexCredentialCatalogError> {
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::<Vec<_>>();
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<Option<CodexPlanCatalog>, 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<FetchedCatalog, CodexCredentialCatalogError> {
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::<Vec<_>>();
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<ProviderAccount>,
) -> Result<FetchedAccountModels, CodexCredentialCatalogError> {
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<Option<CodexPlanCatalog>, 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<ProviderCatalogCacheKey, CodexCredentialCatalogError> {
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<CodexCatalogModel>,
entitlement: Vec<String>,
) -> 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::<String, CodexCatalogModel>::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<CodexCredentialCatalogSnapshot, CodexCredentialCatalogError> {
self.refresh_inner().await
}
/// 忽略当前 cache，按已认领的 ETag 变化强制生成一份完整新快照。
pub async fn refresh(
&self,
) -> Result<CodexCredentialCatalogSnapshot, CodexCredentialCatalogError> {
@@ -585,6 +670,105 @@ impl CodexCredentialCatalogService {
}
}
fn model_ids(models: &[CodexCatalogModel]) -> Vec<String> {
models
.iter()
.map(|model| model.request_model().as_str().to_owned())
.collect()
}
fn catalog_candidates_by_scope(
accounts: Vec<ProviderAccount>,
) -> Result<BTreeMap<CodexCatalogScope, Vec<ProviderAccount>>, CodexCredentialCatalogError> {
let now = SystemTime::now();
let mut groups = BTreeMap::<CodexCatalogScope, Vec<ProviderAccount>>::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<CodexPlanCatalog, CodexCredentialCatalogError> {
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,
backend/crates/providers/openai/src/credential/mod.rs
+2 -1
View File
@@ -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::{
backend/crates/providers/openai/src/lib.rs
+1
View File
@@ -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(),
backend/crates/providers/openai/src/provider.rs
+34 -9
View File
@@ -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<CodexCredentialQuotaService>,
catalog: Arc<CodexCredentialCatalogService>,
}
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(())
}
})
}
backend/crates/providers/openai/tests/admin.rs
+44 -5
View File
@@ -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<MemoryAccountStore>,
pending: Arc<TestOAuthPending>,
) -> ProviderStorePorts {
provider_ports_with_catalog(accounts, pending, Arc::new(TestCatalogCache::default()))
}
fn provider_ports_with_catalog(
accounts: Arc<MemoryAccountStore>,
pending: Arc<TestOAuthPending>,
catalog_cache: Arc<TestCatalogCache>,
) -> 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<Item = &'static str>) {
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 {
backend/crates/providers/openai/tests/credential/catalog.rs
+155 -5
View File
@@ -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<AtomicUsize>,
}
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<MemoryAccountStore>) -> CodexCredentialCatalogService {
service_with_catalog_cache(store, OFFICIAL_CODEX_BASE_URL.to_owned(), catalog_cache())
}
fn service_with_catalog_cache(
store: &Arc<MemoryAccountStore>,
base_url: String,
catalog_cache: Arc<dyn ProviderCatalogCachePort>,
) -> CodexCredentialCatalogService {
CodexCredentialCatalogService::new(
store.repository(),
wire_profile(),
@@ -34,11 +68,26 @@ fn service(store: &Arc<MemoryAccountStore>) -> 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<MemoryAccountStore>, 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());
backend/crates/providers/openai/tests/credential/contract.rs
+2 -1
View File
@@ -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<MemoryAccountStore>, 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(),
backend/crates/providers/openai/tests/provider/contract.rs
+2 -1
View File
@@ -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(),
backend/crates/providers/openai/tests/support.rs
+48 -7
View File
@@ -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<BTreeMap<String, OpaqueProviderData>>,
}
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<Option<OpaqueProviderData>, ProviderStoreError>> {
Box::pin(async move {
Ok(self
.values
.lock()
.expect("catalog cache")
.get(key.scope().as_str())
.cloned())
})
}
}
pub(crate) fn catalog_cache() -> Arc<dyn ProviderCatalogCachePort> {
Arc::new(MemoryCatalogCache::default())
}
pub(crate) fn profile(account_id: &str) -> CodexAccountProfile {
let now = chrono::Utc::now();
CodexAccountProfile {
backend/crates/providers/xai/src/admin.rs
+6 -14
View File
@@ -574,23 +574,15 @@ impl ProviderAdmin for XaiAdminProvider {
) -> Result<ProviderModels, ProviderAdminError> {
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()
backend/crates/providers/xai/src/credential/catalog.rs
+196 -117
View File
@@ -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<Self, GrokCatalogCacheError> {
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<Utc>,
seed: GrokCredentialCatalogSeed,
}
impl GrokAccountCatalog {
impl GrokPlanCatalog {
#[must_use]
pub const fn new(
account_id: ProviderAccountId,
revision: CredentialRevision,
scope: GrokCatalogScope,
observed_at: DateTime<Utc>,
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<Option<GrokAccountCatalog>, GrokCatalogCacheError>;
scope: &GrokCatalogScope,
) -> Result<Option<GrokPlanCatalog>, GrokCatalogCacheError>;
async fn observed_model_support(
&self,
account_id: &ProviderAccountId,
revision: CredentialRevision,
scope: &GrokCatalogScope,
model: &str,
) -> Result<Option<bool>, 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<GrokAccountCatalog, GrokCatalogCacheError> {
) -> Result<GrokPlanCatalog, GrokCatalogCacheError> {
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<Option<GrokAccountCatalog>, GrokCatalogCacheError> {
scope: &GrokCatalogScope,
) -> Result<Option<GrokPlanCatalog>, 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<Option<bool>, 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<Utc>,
accounts: Vec<GrokAccountCatalog>,
plans: Vec<GrokPlanCatalog>,
models: Vec<GrokCatalogModel>,
}
@@ -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<GrokPlanCatalog, GrokCredentialCatalogError> {
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<GrokAccountCatalog, GrokCredentialCatalogError> {
let candidate = self
) -> Result<GrokPlanCatalog, GrokCredentialCatalogError> {
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<Option<GrokAccountCatalog>, GrokCredentialCatalogError> {
account: &ProviderAccount,
) -> Result<Option<GrokPlanCatalog>, 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::<Vec<_>>();
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::<Vec<_>>()
.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::<Vec<_>>();
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<LoadedGrokCredential>,
) -> Result<GrokPlanCatalog, GrokCredentialCatalogError> {
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<LoadedGrokCredential>,
) -> Result<FetchedCredentialCatalog, GrokCredentialCatalogError> {
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<Utc>,
accounts: Vec<GrokAccountCatalog>,
plans: Vec<GrokPlanCatalog>,
models: Vec<GrokCatalogModel>,
}
fn catalog_candidates_by_scope(
candidates: Vec<LoadedGrokCredential>,
) -> Result<BTreeMap<GrokCatalogScope, Vec<LoadedGrokCredential>>, GrokCredentialCatalogError> {
let mut groups = BTreeMap::<GrokCatalogScope, Vec<LoadedGrokCredential>>::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<GrokModelCatalogClient>,
scope: GrokCatalogScope,
candidate: LoadedGrokCredential,
wire_profile: XaiWireProfileState,
) -> Result<FetchedCredentialCatalog, GrokCredentialCatalogError> {
@@ -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,
})
backend/crates/providers/xai/src/credential/mod.rs
+4 -4
View File
@@ -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::{
backend/crates/providers/xai/src/credential/refresh.rs
+1 -4
View File
@@ -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,
backend/crates/providers/xai/src/credential/selector.rs
+9 -9
View File
@@ -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);
}
backend/crates/providers/xai/src/lib.rs
+6 -6
View File
@@ -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,
backend/crates/providers/xai/tests/admin.rs
+44 -5
View File
@@ -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<MemoryProviderAccountStore>,
pending: Arc<TestOAuthPending>,
) -> ProviderStorePorts {
provider_ports_with_catalog(accounts, pending, Arc::new(TestCatalogCache::default()))
}
fn provider_ports_with_catalog(
accounts: Arc<MemoryProviderAccountStore>,
pending: Arc<TestOAuthPending>,
catalog_cache: Arc<TestCatalogCache>,
) -> 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<Item = &'static str>) {
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 {
backend/crates/providers/xai/tests/credential/catalog.rs
+177 -28
View File
@@ -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<VecDeque<Result<GrokModelCatalogTransportResponse, GrokModelCatalogTransportError>>>,
}
@@ -35,6 +36,7 @@ struct QueueCatalogTransport {
impl QueueCatalogTransport {
fn from_bodies(bodies: impl IntoIterator<Item = Vec<u8>>) -> Arc<Self> {
Arc::new(Self {
calls: AtomicUsize::new(0),
responses: Mutex::new(
bodies
.into_iter()
@@ -46,15 +48,32 @@ impl QueueCatalogTransport {
fn failure() -> Arc<Self> {
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<GrokModelCatalogTransportResponse, GrokModelCatalogTransportError>,
>,
) -> Arc<Self> {
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<dyn GrokCredentialCatalogCache> = 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<dyn GrokCredentialCatalogCache> = 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");
backend/crates/providers/xai/tests/credential/selector.rs
+42 -10
View File
@@ -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
backend/crates/providers/xai/tests/support.rs
+12 -16
View File
@@ -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<BTreeMap<(ProviderAccountId, CredentialRevision), GrokAccountCatalog>>,
entries: Mutex<BTreeMap<GrokCatalogScope, GrokPlanCatalog>>,
}
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<Option<GrokAccountCatalog>, GrokCatalogCacheError> {
Ok(lock(&self.entries)
.get(&(account_id.clone(), revision))
.cloned())
scope: &GrokCatalogScope,
) -> Result<Option<GrokPlanCatalog>, GrokCatalogCacheError> {
Ok(lock(&self.entries).get(scope).cloned())
}
async fn observed_model_support(
&self,
account_id: &ProviderAccountId,
revision: CredentialRevision,
scope: &GrokCatalogScope,
model: &str,
) -> Result<Option<bool>, GrokCatalogCacheError> {
Ok(lock(&self.entries)
.get(&(account_id.clone(), revision))
.get(scope)
.map(|catalog| catalog.seed().permits(model)))
}
}