mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-11 00:01:15 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
431992a7a4 | ||
|
|
d3143406dc | ||
|
|
478f2d1f21 |
@@ -422,6 +422,21 @@ fn loaded(account: &ProviderAccount) -> LoadedCredential {
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderAccountStore for TestAccounts {
|
||||
async fn load_account_catalogs(
|
||||
&self,
|
||||
_: &gateway_core::routing::ProviderKind,
|
||||
) -> Result<Vec<gateway_core::account::StoredAccountCatalog>, gateway_core::error::StoreError>
|
||||
{
|
||||
Ok(Vec::new())
|
||||
}
|
||||
async fn save_account_catalog(
|
||||
&self,
|
||||
_: &gateway_core::account::AccountCatalogObservation,
|
||||
_: &[u8],
|
||||
) -> Result<bool, gateway_core::error::StoreError> {
|
||||
unreachable!("fixture does not discover account catalogs")
|
||||
}
|
||||
|
||||
async fn create_account(&self, _account: NewProviderAccount) -> Result<(), StoreError> {
|
||||
panic!("unused account operation")
|
||||
}
|
||||
|
||||
@@ -28,5 +28,7 @@ pub use model::*;
|
||||
pub(crate) use selection::smart_score;
|
||||
pub use selection::*;
|
||||
pub use smart_scheduling::{SmartSchedulingConfig, SmartSchedulingConfigError};
|
||||
pub use store::ProviderAccountStore;
|
||||
pub use store::{
|
||||
AccountCatalogObservation, AccountCatalogOmissions, ProviderAccountStore, StoredAccountCatalog,
|
||||
};
|
||||
pub use upstream::*;
|
||||
|
||||
@@ -12,9 +12,45 @@ use super::{
|
||||
QuotaObservationTouch, QuotaObservationVersion, QuotaWriteOutcome,
|
||||
};
|
||||
|
||||
/// 一个账号最后成功的目录观测;正文由 Provider 解释,归属不能按相同正文合并
|
||||
#[derive(Clone)]
|
||||
pub struct AccountCatalogObservation {
|
||||
pub account_id: ProviderAccountId,
|
||||
pub revision: CredentialRevision,
|
||||
pub plan: Option<String>,
|
||||
pub etag: Option<String>,
|
||||
pub observed_at: std::time::SystemTime,
|
||||
}
|
||||
|
||||
/// 批量恢复时每份不同的正文只传输、解析一次
|
||||
pub struct StoredAccountCatalog {
|
||||
pub accounts: Vec<AccountCatalogObservation>,
|
||||
pub document: bytes::Bytes,
|
||||
}
|
||||
|
||||
/// 批量选号的存储优化:只有账号自己的当前版本、有效时间和目录摘要都匹配才可跳过
|
||||
pub struct AccountCatalogOmissions {
|
||||
pub documents: Vec<[u8; 32]>,
|
||||
pub exceptions: Vec<ProviderAccountId>,
|
||||
pub observed_after: std::time::SystemTime,
|
||||
pub as_of: std::time::SystemTime,
|
||||
}
|
||||
|
||||
/// `provider_accounts` 的数据库中立端口
|
||||
#[async_trait]
|
||||
pub trait ProviderAccountStore: Send + Sync {
|
||||
async fn load_account_catalogs(
|
||||
&self,
|
||||
provider: &ProviderKind,
|
||||
) -> Result<Vec<StoredAccountCatalog>, StoreError>;
|
||||
|
||||
/// 账号删除或凭据已变化时返回 false,不发布迟到的观测
|
||||
async fn save_account_catalog(
|
||||
&self,
|
||||
observation: &AccountCatalogObservation,
|
||||
document: &[u8],
|
||||
) -> Result<bool, StoreError>;
|
||||
|
||||
async fn create_account(&self, account: NewProviderAccount) -> Result<(), StoreError>;
|
||||
|
||||
async fn get_account(
|
||||
@@ -36,6 +72,16 @@ pub trait ProviderAccountStore: Send + Sync {
|
||||
account_ids: &[ProviderAccountId],
|
||||
) -> Result<Vec<ProviderAccount>, StoreError>;
|
||||
|
||||
/// 存储可以下推目录筛选,也可以返回完整候选;调用方始终复核最终账号
|
||||
async fn list_catalog_candidates(
|
||||
&self,
|
||||
provider: &ProviderKind,
|
||||
accounts: &[ProviderAccountId],
|
||||
_omissions: &AccountCatalogOmissions,
|
||||
) -> Result<Vec<ProviderAccount>, StoreError> {
|
||||
self.list_for_provider_accounts(provider, accounts).await
|
||||
}
|
||||
|
||||
/// 一次有界查询返回账号事实和 revision-fenced 明文 credential
|
||||
async fn list_refresh_candidates(
|
||||
&self,
|
||||
|
||||
@@ -14,6 +14,21 @@ fn loaded_credential_from_record(
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
async fn load_account_catalogs(
|
||||
&self,
|
||||
provider: &ProviderKind,
|
||||
) -> Result<Vec<gateway_core::account::StoredAccountCatalog>, CoreStoreError> {
|
||||
self.load_model_catalogs(provider).await
|
||||
}
|
||||
|
||||
async fn save_account_catalog(
|
||||
&self,
|
||||
observation: &gateway_core::account::AccountCatalogObservation,
|
||||
document: &[u8],
|
||||
) -> Result<bool, CoreStoreError> {
|
||||
self.save_model_catalog(observation, document).await
|
||||
}
|
||||
|
||||
async fn create_account(&self, account: CoreNewProviderAccount) -> Result<(), CoreStoreError> {
|
||||
if account.account.revision().get() != 1 {
|
||||
return Err(CoreStoreError::new(CoreStoreErrorKind::InvalidData));
|
||||
@@ -88,6 +103,29 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn list_catalog_candidates(
|
||||
&self,
|
||||
provider: &ProviderKind,
|
||||
accounts: &[CoreProviderAccountId],
|
||||
omissions: &gateway_core::account::AccountCatalogOmissions,
|
||||
) -> Result<Vec<CoreProviderAccount>, CoreStoreError> {
|
||||
let ids = accounts
|
||||
.iter()
|
||||
.map(CoreProviderAccountId::as_str)
|
||||
.collect::<Vec<_>>();
|
||||
self.list_provider_account_summaries(
|
||||
Some(provider.as_str()),
|
||||
false,
|
||||
Some(&ids),
|
||||
Some(omissions),
|
||||
)
|
||||
.await
|
||||
.map_err(core_store_error)?
|
||||
.into_iter()
|
||||
.map(core_account_from_summary)
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn list_for_provider_accounts(
|
||||
&self,
|
||||
provider: &ProviderKind,
|
||||
@@ -100,12 +138,17 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
.iter()
|
||||
.map(CoreProviderAccountId::as_str)
|
||||
.collect::<Vec<_>>();
|
||||
self.list_provider_account_summaries(Some(provider.as_str()), false, Some(&account_ids))
|
||||
.await
|
||||
.map_err(core_store_error)?
|
||||
.into_iter()
|
||||
.map(core_account_from_summary)
|
||||
.collect()
|
||||
self.list_provider_account_summaries(
|
||||
Some(provider.as_str()),
|
||||
false,
|
||||
Some(&account_ids),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(core_store_error)?
|
||||
.into_iter()
|
||||
.map(core_account_from_summary)
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn list_refresh_candidates(
|
||||
|
||||
@@ -58,6 +58,7 @@ mod admin_queries;
|
||||
mod authorization;
|
||||
mod core_adapter;
|
||||
mod mapping;
|
||||
mod model_catalogs;
|
||||
mod quota;
|
||||
mod quota_forecast;
|
||||
mod repository;
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
//! 账号目录的版本约束写入和按相同正文合并的批量恢复
|
||||
|
||||
use super::{
|
||||
CoreCredentialRevision, CoreProviderAccountId, CoreStoreError, CoreStoreErrorKind, DateTime,
|
||||
PgProviderAccountRepository, ProviderKind, Utc,
|
||||
};
|
||||
use gateway_core::account::{AccountCatalogObservation, StoredAccountCatalog};
|
||||
|
||||
impl PgProviderAccountRepository {
|
||||
pub(super) async fn load_model_catalogs(
|
||||
&self,
|
||||
provider: &ProviderKind,
|
||||
) -> Result<Vec<StoredAccountCatalog>, CoreStoreError> {
|
||||
// 先聚合轻量归属,再按主键读取代表正文,避免一万个相同账号展开一万份 TOAST
|
||||
let rows = sqlx::query_as::<_, (Vec<String>, Vec<i64>, Vec<Option<String>>, Vec<Option<String>>, Vec<DateTime<Utc>>, Vec<u8>)>(
|
||||
"select g.ids, g.revisions, g.plans, g.etags, g.times, d.document from (
|
||||
select min(c.account_id) representative, array_agg(c.account_id order by c.account_id) ids,
|
||||
array_agg(c.credential_revision order by c.account_id) revisions,
|
||||
array_agg(c.plan_type order by c.account_id) plans, array_agg(c.etag order by c.account_id) etags, array_agg(c.observed_at order by c.account_id) times
|
||||
from provider_account_model_catalogs c join provider_accounts a on a.id=c.account_id
|
||||
where a.provider_kind=$1 and a.enabled and a.credential_state in ('unknown','ready')
|
||||
and (a.access_token_expires_at is null or a.access_token_expires_at > now())
|
||||
and a.credential_revision=c.credential_revision and a.plan_type is not distinct from c.plan_type
|
||||
group by c.document_hash
|
||||
) g join provider_account_model_catalogs d on d.account_id=g.representative")
|
||||
.bind(provider.as_str()).fetch_all(&self.pool).await.map_err(unavailable)?;
|
||||
rows.into_iter()
|
||||
.map(|(ids, revisions, plans, etags, times, document)| {
|
||||
let accounts = ids
|
||||
.into_iter()
|
||||
.zip(revisions)
|
||||
.zip(plans)
|
||||
.zip(etags)
|
||||
.zip(times)
|
||||
.map(|((((id, revision), plan), etag), time)| {
|
||||
Ok(AccountCatalogObservation {
|
||||
account_id: CoreProviderAccountId::new(id).map_err(|e| {
|
||||
CoreStoreError::caused_by(CoreStoreErrorKind::InvalidData, e)
|
||||
})?,
|
||||
revision: CoreCredentialRevision::new(
|
||||
u64::try_from(revision).map_err(|_| {
|
||||
CoreStoreError::new(CoreStoreErrorKind::InvalidData)
|
||||
})?,
|
||||
)
|
||||
.map_err(|e| {
|
||||
CoreStoreError::caused_by(CoreStoreErrorKind::InvalidData, e)
|
||||
})?,
|
||||
plan,
|
||||
etag,
|
||||
observed_at: time.into(),
|
||||
})
|
||||
})
|
||||
.collect::<Result<_, CoreStoreError>>()?;
|
||||
Ok(StoredAccountCatalog {
|
||||
accounts,
|
||||
document: document.into(),
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(super) async fn save_model_catalog(
|
||||
&self,
|
||||
observation: &AccountCatalogObservation,
|
||||
document: &[u8],
|
||||
) -> Result<bool, CoreStoreError> {
|
||||
let result = sqlx::query("insert into provider_account_model_catalogs (account_id, credential_revision, plan_type, observed_at, document, etag)
|
||||
select a.id, a.credential_revision, a.plan_type, $4, $5, $6 from provider_accounts a
|
||||
where a.id=$1 and a.credential_revision=$2 and a.plan_type is not distinct from $3 for share
|
||||
on conflict (account_id) do update set credential_revision=excluded.credential_revision, plan_type=excluded.plan_type,
|
||||
etag=excluded.etag, observed_at=excluded.observed_at, document=excluded.document
|
||||
where provider_account_model_catalogs.credential_revision < excluded.credential_revision or
|
||||
(provider_account_model_catalogs.credential_revision=excluded.credential_revision and provider_account_model_catalogs.observed_at <= excluded.observed_at)")
|
||||
.bind(observation.account_id.as_str()).bind(i64::try_from(observation.revision.get()).map_err(|_|CoreStoreError::new(CoreStoreErrorKind::InvalidData))?)
|
||||
.bind(observation.plan.as_deref()).bind(DateTime::<Utc>::from(observation.observed_at)).bind(document).bind(observation.etag.as_deref())
|
||||
.execute(&self.pool).await.map_err(unavailable)?;
|
||||
Ok(result.rows_affected() == 1)
|
||||
}
|
||||
}
|
||||
|
||||
fn unavailable(source: sqlx::Error) -> CoreStoreError {
|
||||
CoreStoreError::caused_by(CoreStoreErrorKind::Unavailable, source)
|
||||
}
|
||||
@@ -104,6 +104,7 @@ impl PgProviderAccountRepository {
|
||||
provider_kind: Option<&str>,
|
||||
include_disabled: bool,
|
||||
account_ids: Option<&[&str]>,
|
||||
omissions: Option<&gateway_core::account::AccountCatalogOmissions>,
|
||||
) -> StoreResult<Vec<ProviderAccountSummary>> {
|
||||
let mut query = sqlx::QueryBuilder::<Postgres>::new(
|
||||
"select coalesce(proxy_enabled, true) as outbound_proxy_enabled, coalesce(max_websocket_connections, 0) as outbound_proxy_max_websocket_connections, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_id, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id,
|
||||
@@ -129,6 +130,24 @@ impl PgProviderAccountRepository {
|
||||
.push_bind(account_ids)
|
||||
.push("::text[])");
|
||||
}
|
||||
if let Some(omit) = omissions {
|
||||
let hashes = omit
|
||||
.documents
|
||||
.iter()
|
||||
.map(<[u8; 32]>::as_slice)
|
||||
.collect::<Vec<_>>();
|
||||
let exceptions = omit
|
||||
.exceptions
|
||||
.iter()
|
||||
.map(|id| id.as_str())
|
||||
.collect::<Vec<_>>();
|
||||
query.push(" and not exists (select 1 from provider_account_model_catalogs c
|
||||
where c.account_id=provider_accounts.id and c.credential_revision=provider_accounts.credential_revision
|
||||
and c.plan_type is not distinct from provider_accounts.plan_type and c.document_hash=any(")
|
||||
.push_bind(hashes).push("::bytea[]) and c.observed_at > ").push_bind(DateTime::<Utc>::from(omit.observed_after))
|
||||
.push(" and c.observed_at <= ").push_bind(DateTime::<Utc>::from(omit.as_of))
|
||||
.push(" and not (c.account_id=any(").push_bind(exceptions).push("::text[])))");
|
||||
}
|
||||
query.push(" order by provider_kind, name, id");
|
||||
let rows = query
|
||||
.build()
|
||||
@@ -215,7 +234,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
provider_kind: Option<&str>,
|
||||
include_disabled: bool,
|
||||
) -> StoreResult<Vec<ProviderAccountSummary>> {
|
||||
self.list_provider_account_summaries(provider_kind, include_disabled, None)
|
||||
self.list_provider_account_summaries(provider_kind, include_disabled, None, None)
|
||||
.await
|
||||
}
|
||||
|
||||
|
||||
@@ -103,3 +103,101 @@ async fn scoped_account_list_filters_ids_provider_and_live_enabled_state() {
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bulk_catalog_restore_keeps_ten_thousand_owners_and_fences_late_writes() {
|
||||
use gateway_core::account::{AccountCatalogObservation, CredentialRevision};
|
||||
let Some(database) = TestDatabase::create("account_catalog_bulk").await else {
|
||||
return;
|
||||
};
|
||||
let repo = PgProviderAccountRepository::new(database.pool.clone());
|
||||
sqlx::query("insert into provider_accounts (id,provider_kind,name,authentication_kind,provider_credentials_json,has_refresh_token,credential_state,credential_observed_at,created_at,updated_at,plan_type)
|
||||
select 'acct_'||i,'openai','source '||i,'oauth','{}',false,'ready',now(),now(),now(),'pro' from generate_series(1,10000) i").execute(&database.pool).await.unwrap();
|
||||
let observation = AccountCatalogObservation {
|
||||
account_id: ProviderAccountId::new("acct_1").unwrap(),
|
||||
revision: CredentialRevision::new(1).unwrap(),
|
||||
observed_at: std::time::SystemTime::now(),
|
||||
etag: Some("v1".into()),
|
||||
plan: Some("pro".into()),
|
||||
};
|
||||
let document = vec![b'x'; 64 * 1024];
|
||||
assert!(
|
||||
repo.save_account_catalog(&observation, &document)
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
sqlx::query("insert into provider_account_model_catalogs (account_id,credential_revision,plan_type,observed_at,document)
|
||||
select a.id,1,'pro',now(),c.document from provider_accounts a cross join provider_account_model_catalogs c where a.id<>'acct_1'").execute(&database.pool).await.unwrap();
|
||||
let provider = ProviderKind::new("openai").unwrap();
|
||||
let start = std::time::Instant::now();
|
||||
let group = repo.load_account_catalogs(&provider).await.unwrap();
|
||||
eprintln!(
|
||||
"catalog bulk restore: 10000 accounts, {} documents, {} ms",
|
||||
group.len(),
|
||||
start.elapsed().as_millis()
|
||||
);
|
||||
assert_eq!(group.len(), 1);
|
||||
assert_eq!(group[0].accounts.len(), 10000);
|
||||
assert_eq!(&group[0].document[..], &document);
|
||||
let older = AccountCatalogObservation {
|
||||
observed_at: observation.observed_at - std::time::Duration::from_secs(1),
|
||||
..observation.clone()
|
||||
};
|
||||
assert!(!repo.save_account_catalog(&older, b"late").await.unwrap());
|
||||
sqlx::query("update provider_accounts set credential_revision=2 where id='acct_1'")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
!repo
|
||||
.save_account_catalog(&observation, b"old revision")
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
sqlx::query("update provider_accounts set enabled=false where id='acct_2'")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query("delete from provider_accounts where id='acct_3'")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let group = repo.load_account_catalogs(&provider).await.unwrap();
|
||||
assert_eq!(group[0].accounts.len(), 9997);
|
||||
assert!(
|
||||
group[0]
|
||||
.accounts
|
||||
.iter()
|
||||
.all(|a| !matches!(a.account_id.as_str(), "acct_1" | "acct_2" | "acct_3"))
|
||||
);
|
||||
// 数据库只能排除有当前版本、未过期的账号证据;显式例外仍进入选号
|
||||
let hash: Vec<u8> = sqlx::query_scalar("select sha256($1::bytea)")
|
||||
.bind(&document)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let now = std::time::SystemTime::now();
|
||||
sqlx::query("update provider_account_model_catalogs set observed_at=now()-interval '1 hour' where account_id='acct_4'").execute(&database.pool).await.unwrap();
|
||||
let filter = gateway_core::account::AccountCatalogOmissions {
|
||||
documents: vec![hash.try_into().unwrap()],
|
||||
exceptions: vec![ProviderAccountId::new("acct_5").unwrap()],
|
||||
observed_after: now - std::time::Duration::from_secs(1800),
|
||||
as_of: now,
|
||||
};
|
||||
let ids = (1..=10000)
|
||||
.map(|n| ProviderAccountId::new(format!("acct_{n}")).unwrap())
|
||||
.collect::<Vec<_>>();
|
||||
let candidates = repo
|
||||
.list_catalog_candidates(&provider, &ids, &filter)
|
||||
.await
|
||||
.unwrap();
|
||||
let ids = candidates
|
||||
.iter()
|
||||
.map(|a| a.id().as_str())
|
||||
.collect::<std::collections::BTreeSet<_>>();
|
||||
assert_eq!(
|
||||
ids,
|
||||
std::collections::BTreeSet::from(["acct_1", "acct_4", "acct_5"])
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
@@ -297,7 +297,7 @@ impl ProviderAdmin for OpenAiAdminProvider {
|
||||
return;
|
||||
}
|
||||
self.quota.invalidate_scheduling(account_ids);
|
||||
if let Err(error) = self.catalog.invalidate() {
|
||||
if let Err(error) = self.catalog.invalidate_accounts(account_ids).await {
|
||||
tracing::warn!(
|
||||
account_count = account_ids.len(),
|
||||
error = %error,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,861 @@
|
||||
//! OpenAI 账号模型目录的聚合、缓存与调度资格
|
||||
|
||||
mod refresh;
|
||||
|
||||
use std::collections::{BTreeMap, BTreeSet, btree_map::Entry};
|
||||
use std::fmt;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use futures::{StreamExt, stream};
|
||||
use gateway_core::account::AccountCatalogObservation;
|
||||
use gateway_core::account::{
|
||||
CredentialRevision, CredentialState, OpaqueProviderData, ProviderAccount, ProviderAccountId,
|
||||
};
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCatalogScope,
|
||||
};
|
||||
use gateway_core::routing::{
|
||||
FrozenAccountScope, ModelPresentation, ProviderCatalogGeneration, ProviderKind,
|
||||
ProviderModelContent, ProviderModelDescriptor, UpstreamModelId,
|
||||
};
|
||||
use secrecy::ExposeSecret;
|
||||
use thiserror::Error;
|
||||
use tokio::sync::Notify;
|
||||
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::repository::{CodexCredentialRepository, CredentialRepositoryError};
|
||||
use crate::transport::profile::CodexWireProfileState;
|
||||
use crate::transport::{CodexBackendClient, CodexCatalogModel, CodexRequestContext};
|
||||
|
||||
const MAX_RESPONSE_ETAG_BYTES: usize = 256;
|
||||
const ACCOUNT_CATALOG_CACHE_TTL: Duration = Duration::from_secs(5 * 60);
|
||||
// 调度证据覆盖默认 15 分钟刷新周期及抖动,长期未刷新时恢复上游验证
|
||||
const SCHEDULING_CATALOG_MAX_AGE: Duration = Duration::from_secs(30 * 60);
|
||||
const MAX_ACCOUNT_CATALOG_MODELS: usize = 2_048;
|
||||
const MODEL_CATALOG_TIMEOUT: Duration = Duration::from_secs(15);
|
||||
|
||||
const REFRESH_CONCURRENCY: usize = 32;
|
||||
// 同时到期或导入大量账号时也限制起始速率,不能把连接并发上限当作请求速率上限
|
||||
const REFRESH_SPACING: Duration = Duration::from_millis(10);
|
||||
// 小账号池直接在内存复核,避免为少数账号增加数据库联表
|
||||
const MIN_CATALOG_PUSHDOWN: usize = 128;
|
||||
const PUBLICATION_BATCH: usize = 128;
|
||||
|
||||
/// 同一套餐也可能按账号开放模型,换凭据或套餐后必须重新取证
|
||||
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
|
||||
pub struct CodexCatalogScope {
|
||||
account_id: ProviderAccountId,
|
||||
revision: CredentialRevision,
|
||||
plan: Option<String>,
|
||||
}
|
||||
|
||||
impl CodexCatalogScope {
|
||||
#[must_use]
|
||||
pub fn for_account(account: &ProviderAccount) -> Self {
|
||||
Self {
|
||||
account_id: account.id().clone(),
|
||||
revision: account.revision(),
|
||||
plan: account
|
||||
.plan_type()
|
||||
.map(str::trim)
|
||||
.filter(|plan| !plan.is_empty())
|
||||
.map(str::to_ascii_lowercase),
|
||||
}
|
||||
}
|
||||
|
||||
/// Redis 作用域长度有上限,不直接拼接可变长度的账号与套餐文本
|
||||
pub fn cache_scope(&self) -> Result<ProviderCatalogScope, CodexCredentialCatalogError> {
|
||||
use sha2::{Digest, Sha256};
|
||||
let mut digest = Sha256::new();
|
||||
digest.update(self.account_id.as_str());
|
||||
digest.update([0]);
|
||||
digest.update(self.plan.as_deref().unwrap_or_default());
|
||||
ProviderCatalogScope::new(format!(
|
||||
"account:{}:{}",
|
||||
hex::encode(digest.finalize()),
|
||||
self.revision.get()
|
||||
))
|
||||
.map_err(|_| CodexCredentialCatalogError::Cache)
|
||||
}
|
||||
}
|
||||
|
||||
/// Redis TTL cache 中一条可重建的账号模型目录
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct CodexAccountCatalog {
|
||||
scope: ProviderCatalogScope,
|
||||
observed_at: DateTime<Utc>,
|
||||
models: Vec<String>,
|
||||
}
|
||||
|
||||
impl CodexAccountCatalog {
|
||||
#[must_use]
|
||||
pub const fn new(
|
||||
scope: ProviderCatalogScope,
|
||||
observed_at: DateTime<Utc>,
|
||||
models: Vec<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
scope,
|
||||
observed_at,
|
||||
models,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn scope(&self) -> &ProviderCatalogScope {
|
||||
&self.scope
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn observed_at(&self) -> DateTime<Utc> {
|
||||
self.observed_at
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn models(&self) -> &[String] {
|
||||
&self.models
|
||||
}
|
||||
}
|
||||
|
||||
struct CatalogContent {
|
||||
models: Vec<CodexCatalogModel>,
|
||||
ids: Vec<String>,
|
||||
digest: [u8; 32],
|
||||
}
|
||||
|
||||
impl CatalogContent {
|
||||
fn new(models: Vec<CodexCatalogModel>, document: &[u8]) -> Self {
|
||||
use sha2::Digest;
|
||||
Self {
|
||||
digest: sha2::Sha256::digest(document).into(),
|
||||
ids: model_ids(&models),
|
||||
models,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct AccountEntry {
|
||||
observed_at: SystemTime,
|
||||
content: Arc<CatalogContent>,
|
||||
invalidated: Arc<std::sync::atomic::AtomicBool>,
|
||||
etag: Option<String>,
|
||||
}
|
||||
|
||||
type Entries = BTreeMap<CodexCatalogScope, AccountEntry>;
|
||||
|
||||
pub struct CodexCredentialCatalogSnapshot {
|
||||
models: Vec<CodexCatalogModel>,
|
||||
entries: Entries,
|
||||
sources: Vec<(Arc<CatalogContent>, Vec<ProviderAccountId>)>,
|
||||
}
|
||||
|
||||
impl CodexCredentialCatalogSnapshot {
|
||||
fn new(entries: Entries) -> Self {
|
||||
let mut models: Vec<CodexCatalogModel> = Vec::new();
|
||||
let mut indices = BTreeMap::<&str, usize>::new();
|
||||
let mut sources = BTreeMap::<_, (Arc<CatalogContent>, Vec<ProviderAccountId>)>::new();
|
||||
for (scope, entry) in &entries {
|
||||
let (_, members) = sources
|
||||
.entry(entry.content.digest)
|
||||
.or_insert_with(|| (entry.content.clone(), Vec::new()));
|
||||
members.push(scope.account_id.clone());
|
||||
if members.len() > 1 {
|
||||
continue;
|
||||
}
|
||||
for model in &entry.content.models {
|
||||
match indices.entry(model.request_model().as_str()) {
|
||||
Entry::Vacant(slot) => {
|
||||
slot.insert(models.len());
|
||||
models.push(model.clone());
|
||||
}
|
||||
Entry::Occupied(slot)
|
||||
if model.document().protocol() == "codex"
|
||||
&& models[*slot.get()].document().protocol() != "codex" =>
|
||||
{
|
||||
models[*slot.get()] = model.clone()
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
Self {
|
||||
models,
|
||||
entries,
|
||||
sources: sources.into_values().collect(),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn models(&self) -> &[CodexCatalogModel] {
|
||||
&self.models
|
||||
}
|
||||
|
||||
pub(super) fn excludes(&self, account: &ProviderAccount, model: &str, now: SystemTime) -> bool {
|
||||
self.entries
|
||||
.get(&CodexCatalogScope::for_account(account))
|
||||
.is_some_and(|e| {
|
||||
!e.invalidated.load(std::sync::atomic::Ordering::Relaxed)
|
||||
&& now
|
||||
.duration_since(e.observed_at)
|
||||
.is_ok_and(|age| age < SCHEDULING_CATALOG_MAX_AGE)
|
||||
&& !e.content.ids.iter().any(|id| id == model)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn account_models(&self, account: &ProviderAccount) -> Option<&[String]> {
|
||||
self.entries
|
||||
.get(&CodexCatalogScope::for_account(account))
|
||||
.map(|e| e.content.ids.as_slice())
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn model_catalog_accounts(&self) -> BTreeMap<&str, BTreeSet<ProviderAccountId>> {
|
||||
let mut accounts = BTreeMap::<&str, BTreeSet<ProviderAccountId>>::new();
|
||||
for (scope, entry) in &self.entries {
|
||||
for model in &entry.content.ids {
|
||||
accounts
|
||||
.entry(model.as_str())
|
||||
.or_default()
|
||||
.insert(scope.account_id.clone());
|
||||
}
|
||||
}
|
||||
accounts
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for CodexCredentialCatalogSnapshot {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CodexCredentialCatalogSnapshot")
|
||||
.field("model_count", &self.models.len())
|
||||
.field("scope_count", &self.entries.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Error)]
|
||||
pub enum CodexCredentialCatalogError {
|
||||
#[error("Codex model catalog has no eligible account")]
|
||||
NoEligibleCredential,
|
||||
#[error("Codex model catalog account data is invalid")]
|
||||
InvalidCredentialData,
|
||||
#[error("Codex model catalog upstream query failed: {detail}")]
|
||||
Upstream { detail: String },
|
||||
#[error("Codex model catalog cache is unavailable")]
|
||||
Cache,
|
||||
#[error("Codex model catalog ETag is invalid")]
|
||||
InvalidEtag,
|
||||
#[error("Codex model catalog changed during refresh")]
|
||||
ConcurrentUpdate,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct CatalogCacheState {
|
||||
restored: bool,
|
||||
revision: u64,
|
||||
generation: u64,
|
||||
snapshot: Option<Arc<CodexCredentialCatalogSnapshot>>,
|
||||
pending: BTreeMap<CodexCatalogScope, String>,
|
||||
}
|
||||
|
||||
struct FetchedAccountModels {
|
||||
models: Vec<CodexCatalogModel>,
|
||||
document: bytes::Bytes,
|
||||
etag: Option<String>,
|
||||
}
|
||||
|
||||
impl From<CredentialRepositoryError> for CodexCredentialCatalogError {
|
||||
fn from(_: CredentialRepositoryError) -> Self {
|
||||
Self::InvalidCredentialData
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct CodexCredentialCatalogService {
|
||||
repository: CodexCredentialRepository,
|
||||
profile: CodexWireProfileState,
|
||||
http: reqwest::Client,
|
||||
base_url: String,
|
||||
account_catalog_cache: Arc<dyn ProviderCatalogCachePort>,
|
||||
cache: Arc<RwLock<CatalogCacheState>>,
|
||||
refresh_notification: Arc<Notify>,
|
||||
refresh_lock: Arc<tokio::sync::Mutex<()>>,
|
||||
publication: Arc<std::sync::Mutex<()>>,
|
||||
}
|
||||
|
||||
impl CodexCredentialCatalogService {
|
||||
pub fn new(
|
||||
repository: CodexCredentialRepository,
|
||||
profile: CodexWireProfileState,
|
||||
http: reqwest::Client,
|
||||
base_url: String,
|
||||
account_catalog_cache: Arc<dyn ProviderCatalogCachePort>,
|
||||
) -> Self {
|
||||
Self {
|
||||
repository,
|
||||
profile,
|
||||
http,
|
||||
base_url,
|
||||
account_catalog_cache,
|
||||
cache: Arc::new(RwLock::new(CatalogCacheState::default())),
|
||||
refresh_notification: Arc::new(Notify::new()),
|
||||
refresh_lock: Arc::new(tokio::sync::Mutex::new(())),
|
||||
publication: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 原生条目取自后台发现的账号目录;用当前账号版本和冻结的 Key 政策双重过滤
|
||||
pub async fn client_model_catalog(
|
||||
&self,
|
||||
scope: &FrozenAccountScope,
|
||||
) -> Result<Vec<ProviderModelDescriptor>, CodexCredentialCatalogError> {
|
||||
let snapshot = self
|
||||
.cached()?
|
||||
.ok_or(CodexCredentialCatalogError::NoEligibleCredential)?;
|
||||
// 正文相同的账号只需核对每个模型在当前政策下的首选来源
|
||||
let mut candidates = BTreeSet::new();
|
||||
for (content, members) in &snapshot.sources {
|
||||
if let Some(id) = members.iter().find(|id| scope.allows(id)) {
|
||||
candidates.insert(id.clone());
|
||||
} else {
|
||||
continue;
|
||||
}
|
||||
for id in &content.ids {
|
||||
if let Some(account) = members
|
||||
.iter()
|
||||
.find(|account| scope.allows_model(account, id))
|
||||
{
|
||||
candidates.insert(account.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
let ids = candidates.into_iter().collect::<Vec<_>>();
|
||||
let mut accounts = self.repository.list_for_accounts(&ids).await?;
|
||||
let now = SystemTime::now();
|
||||
// 代表账号刚被停用或轮换时,仅这一次扩大本地资格查询,仍不查询上游
|
||||
if accounts.len() != ids.len()
|
||||
|| accounts.iter().any(|a| {
|
||||
!eligible_catalog_account(a, now)
|
||||
|| !snapshot
|
||||
.entries
|
||||
.contains_key(&CodexCatalogScope::for_account(a))
|
||||
})
|
||||
{
|
||||
let ids = snapshot
|
||||
.entries
|
||||
.keys()
|
||||
.filter(|s| scope.allows(&s.account_id))
|
||||
.map(|s| s.account_id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
accounts = self.repository.list_for_accounts(&ids).await?;
|
||||
}
|
||||
let mut selected = BTreeMap::new();
|
||||
let mut available = false;
|
||||
for account in accounts.iter().filter(|a| eligible_catalog_account(a, now)) {
|
||||
let Some(entry) = snapshot
|
||||
.entries
|
||||
.get(&CodexCatalogScope::for_account(account))
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
available = true;
|
||||
let api = account.authentication_kind() == super::CODEX_AUTHENTICATION_KIND_API_KEY;
|
||||
for (position, model) in entry.content.models.iter().enumerate() {
|
||||
if !scope.allows_model(account.id(), model.request_model().as_str()) {
|
||||
continue;
|
||||
}
|
||||
let order = (
|
||||
api,
|
||||
model.document().protocol() != "codex",
|
||||
account.id(),
|
||||
position,
|
||||
);
|
||||
let slot = selected
|
||||
.entry(model.request_model())
|
||||
.or_insert((order, model));
|
||||
if order < slot.0 {
|
||||
*slot = (order, model);
|
||||
}
|
||||
}
|
||||
}
|
||||
if !available {
|
||||
return Err(CodexCredentialCatalogError::NoEligibleCredential);
|
||||
}
|
||||
let mut models: Vec<_> = selected.into_values().collect();
|
||||
models.sort_unstable_by_key(|(order, _)| *order);
|
||||
Ok(models
|
||||
.into_iter()
|
||||
.map(|(_, model)| {
|
||||
let id = model.request_model().as_str();
|
||||
ProviderModelDescriptor {
|
||||
model: model.request_model().clone(),
|
||||
content: if model.document().protocol() == "codex" {
|
||||
ProviderModelContent::Native(model.document().clone())
|
||||
} else {
|
||||
ProviderModelContent::Adapted(
|
||||
ModelPresentation::new(Some(id.to_owned()), None)
|
||||
.with_agent_tools(true, false)
|
||||
.with_hidden(is_specialized_api_model(id)),
|
||||
)
|
||||
},
|
||||
}
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
pub(crate) fn messages_model_capabilities(
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
model: &UpstreamModelId,
|
||||
) -> Option<gateway_protocol::anthropic::MessagesModelCapabilities> {
|
||||
let snapshot = self.cached().ok()??;
|
||||
snapshot
|
||||
.entries
|
||||
.get(&CodexCatalogScope::for_account(account))?
|
||||
.content
|
||||
.models
|
||||
.iter()
|
||||
.find(|entry| entry.request_model() == model)?
|
||||
.messages_capabilities()
|
||||
.cloned()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn catalog_generation(&self) -> ProviderCatalogGeneration {
|
||||
let cache = self
|
||||
.cache
|
||||
.read()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
ProviderCatalogGeneration::new(cache.generation)
|
||||
}
|
||||
|
||||
pub fn cached(
|
||||
&self,
|
||||
) -> Result<Option<Arc<CodexCredentialCatalogSnapshot>>, CodexCredentialCatalogError> {
|
||||
Ok(self
|
||||
.cache
|
||||
.read()
|
||||
.map_err(|_| CodexCredentialCatalogError::Cache)?
|
||||
.snapshot
|
||||
.as_ref()
|
||||
.cloned())
|
||||
}
|
||||
|
||||
/// 读取账号当前凭据对应的模型目录;不触发网络
|
||||
pub fn cached_account_models(
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
) -> Result<Option<Vec<String>>, CodexCredentialCatalogError> {
|
||||
let Some(snapshot) = self.cached()? else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(snapshot.account_models(account).map(<[String]>::to_vec))
|
||||
}
|
||||
|
||||
/// 大池中相同正文只传递一个摘要;存储只跳过有当前账号证据的来源
|
||||
pub(super) fn omissions(
|
||||
&self,
|
||||
model: &str,
|
||||
ids: &[ProviderAccountId],
|
||||
exempt: impl Fn(&ProviderAccountId) -> bool,
|
||||
) -> Option<gateway_core::account::AccountCatalogOmissions> {
|
||||
if ids.len() < MIN_CATALOG_PUSHDOWN {
|
||||
return None;
|
||||
}
|
||||
let cache = self.cache.read().ok()?;
|
||||
let snapshot = cache.snapshot.as_ref()?.clone();
|
||||
let pending = cache
|
||||
.pending
|
||||
.keys()
|
||||
.map(|s| s.account_id.clone())
|
||||
.collect::<BTreeSet<_>>();
|
||||
drop(cache);
|
||||
let documents = snapshot
|
||||
.sources
|
||||
.iter()
|
||||
.filter(|(c, _)| !c.ids.iter().any(|id| id == model))
|
||||
.map(|(c, _)| c.digest)
|
||||
.collect::<Vec<_>>();
|
||||
if documents.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let as_of = SystemTime::now();
|
||||
Some(gateway_core::account::AccountCatalogOmissions {
|
||||
documents,
|
||||
exceptions: ids
|
||||
.iter()
|
||||
.filter(|id| pending.contains(*id) || exempt(id))
|
||||
.cloned()
|
||||
.collect(),
|
||||
observed_after: as_of.checked_sub(SCHEDULING_CATALOG_MAX_AGE)?,
|
||||
as_of,
|
||||
})
|
||||
}
|
||||
|
||||
/// 优先读取账号目录缓存;缺失时只查询该账号
|
||||
pub async fn cached_or_refresh_account_catalog(
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
) -> Result<CodexAccountCatalog, CodexCredentialCatalogError> {
|
||||
if let Some(catalog) = self.read_account_catalog(account).await? {
|
||||
return Ok(catalog);
|
||||
}
|
||||
self.refresh_account_catalog(account.id()).await
|
||||
}
|
||||
|
||||
/// 实时刷新指定账号的模型集合,包括管理员显式查询的停用账号
|
||||
pub async fn refresh_account_catalog(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Result<CodexAccountCatalog, CodexCredentialCatalogError> {
|
||||
let account = self
|
||||
.repository
|
||||
.store()
|
||||
.get_account(account_id)
|
||||
.await
|
||||
.map_err(|_| CodexCredentialCatalogError::InvalidCredentialData)?
|
||||
.filter(|account| account.provider().as_str() == "openai")
|
||||
.ok_or(CodexCredentialCatalogError::NoEligibleCredential)?;
|
||||
let scope = CodexCatalogScope::for_account(&account);
|
||||
let client = CodexBackendClient::new(
|
||||
self.http.clone(),
|
||||
self.base_url.clone(),
|
||||
self.profile.clone(),
|
||||
);
|
||||
let revision = self.state_revision()?;
|
||||
let (catalog, entry) = self
|
||||
.fetch_observation(&client, &account, &scope, None)
|
||||
.await?;
|
||||
if eligible_catalog_account(&account, SystemTime::now())
|
||||
&& let Err(error) = self.merge_observations(revision, vec![(scope, entry)])
|
||||
{
|
||||
if !matches!(error, CodexCredentialCatalogError::ConcurrentUpdate) {
|
||||
return Err(error);
|
||||
}
|
||||
self.schedule_refresh();
|
||||
}
|
||||
Ok(catalog)
|
||||
}
|
||||
|
||||
async fn fetch_observation(
|
||||
&self,
|
||||
client: &CodexBackendClient,
|
||||
account: &ProviderAccount,
|
||||
scope: &CodexCatalogScope,
|
||||
etag: Option<String>,
|
||||
) -> Result<(CodexAccountCatalog, AccountEntry), CodexCredentialCatalogError> {
|
||||
let observed_at = SystemTime::now();
|
||||
let fetched = self.fetch_account_models(client, account).await?;
|
||||
let observed = AccountCatalogObservation {
|
||||
account_id: account.id().clone(),
|
||||
revision: account.revision(),
|
||||
plan: account.plan_type().map(str::to_owned),
|
||||
observed_at,
|
||||
etag: etag.clone().or(fetched.etag),
|
||||
};
|
||||
if !self
|
||||
.repository
|
||||
.store()
|
||||
.save_account_catalog(&observed, &fetched.document)
|
||||
.await
|
||||
.map_err(|_| CodexCredentialCatalogError::Cache)?
|
||||
{
|
||||
return Err(CodexCredentialCatalogError::ConcurrentUpdate);
|
||||
}
|
||||
let content = Arc::new(CatalogContent::new(fetched.models, &fetched.document));
|
||||
// Redis 仅为管理端短期缓存,失败不能撤销已持久化的目录
|
||||
let catalog = CodexAccountCatalog::new(
|
||||
scope.cache_scope()?,
|
||||
observed_at.into(),
|
||||
content.ids.clone(),
|
||||
);
|
||||
if let Err(error) = self.replace_account_catalog(&catalog).await {
|
||||
tracing::warn!(%error,"OpenAI account catalog cache write failed");
|
||||
}
|
||||
Ok((
|
||||
catalog,
|
||||
AccountEntry {
|
||||
observed_at,
|
||||
content,
|
||||
etag: observed.etag,
|
||||
invalidated: Default::default(),
|
||||
},
|
||||
))
|
||||
}
|
||||
/// 读取目标账号的原生目录条目;供管理端按账号导出目录文件使用
|
||||
///
|
||||
/// 按账号直接请求上游,避免其他账号的同名模型替换目标账号的原生对象;
|
||||
/// 停用账号也不在常规快照候选中,仍允许管理员显式导出
|
||||
/// 返回顺序沿用上游目录顺序,条目内容保持上游原样,不做别名改写或字段裁剪
|
||||
pub async fn account_catalog_documents(
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
) -> Result<(Vec<CodexCatalogModel>, SystemTime), CodexCredentialCatalogError> {
|
||||
let client = CodexBackendClient::new(
|
||||
self.http.clone(),
|
||||
self.base_url.clone(),
|
||||
self.profile.clone(),
|
||||
);
|
||||
let fetched = self.fetch_account_models(&client, account).await?;
|
||||
if fetched.models.is_empty() {
|
||||
return Err(CodexCredentialCatalogError::NoEligibleCredential);
|
||||
}
|
||||
Ok((fetched.models, SystemTime::now()))
|
||||
}
|
||||
|
||||
/// 读取当前账号的目录缓存,不触发上游请求
|
||||
pub async fn read_account_catalog(
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
) -> Result<Option<CodexAccountCatalog>, CodexCredentialCatalogError> {
|
||||
let scope = CodexCatalogScope::for_account(account);
|
||||
let key = scope.cache_scope()?;
|
||||
if let Some(snapshot) = self.cached()?
|
||||
&& let Some(entry) = snapshot.entries.get(&scope)
|
||||
{
|
||||
return Ok(Some(CodexAccountCatalog::new(
|
||||
key,
|
||||
entry.observed_at.into(),
|
||||
entry.content.ids.clone(),
|
||||
)));
|
||||
}
|
||||
self.read_cached_catalog(&key).await
|
||||
}
|
||||
|
||||
async fn replace_account_catalog(
|
||||
&self,
|
||||
catalog: &CodexAccountCatalog,
|
||||
) -> Result<(), CodexCredentialCatalogError> {
|
||||
let key = self.account_catalog_key(catalog.scope())?;
|
||||
let document = encode_account_catalog(catalog);
|
||||
self.account_catalog_cache
|
||||
.replace(&key, &document, ACCOUNT_CATALOG_CACHE_TTL)
|
||||
.await
|
||||
.map_err(|_| CodexCredentialCatalogError::Cache)
|
||||
}
|
||||
|
||||
async fn read_cached_catalog(
|
||||
&self,
|
||||
scope: &ProviderCatalogScope,
|
||||
) -> Result<Option<CodexAccountCatalog>, CodexCredentialCatalogError> {
|
||||
let key = self.account_catalog_key(scope)?;
|
||||
self.account_catalog_cache
|
||||
.read(&key)
|
||||
.await
|
||||
.map_err(|_| CodexCredentialCatalogError::Cache)?
|
||||
.map(|document| decode_account_catalog(scope, document))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn account_catalog_key(
|
||||
&self,
|
||||
scope: &ProviderCatalogScope,
|
||||
) -> Result<ProviderCatalogCacheKey, CodexCredentialCatalogError> {
|
||||
let provider_kind =
|
||||
ProviderKind::new("openai").map_err(|_| CodexCredentialCatalogError::Cache)?;
|
||||
Ok(ProviderCatalogCacheKey::new(provider_kind, scope.clone()))
|
||||
}
|
||||
|
||||
async fn fetch_account_models(
|
||||
&self,
|
||||
client: &CodexBackendClient,
|
||||
account: &ProviderAccount,
|
||||
) -> Result<FetchedAccountModels, CodexCredentialCatalogError> {
|
||||
tokio::time::timeout(MODEL_CATALOG_TIMEOUT, async {
|
||||
let credential = self
|
||||
.repository
|
||||
.load_runtime_credential(account)
|
||||
.await
|
||||
.map_err(|_| CodexCredentialCatalogError::InvalidCredentialData)?;
|
||||
let request_id = format!("catalog_{}", Uuid::now_v7().simple());
|
||||
let authorization = credential
|
||||
.authentication
|
||||
.authorization_header()
|
||||
.map_err(|_| CodexCredentialCatalogError::InvalidCredentialData)?;
|
||||
let mut client = client
|
||||
.for_account(account)
|
||||
.map_err(|error| CodexCredentialCatalogError::Upstream {
|
||||
detail: error.to_string(),
|
||||
})?
|
||||
.with_responses_api_base_url(credential.authentication.responses_api_base_url());
|
||||
if let super::CodexRuntimeAuthentication::ApiKey(auth) = &credential.authentication
|
||||
&& auth.configuration.protocol == super::ApiKeyProtocol::Messages
|
||||
{
|
||||
client = client.with_messages_api_base_url(&auth.configuration.base_url);
|
||||
}
|
||||
let result = client
|
||||
.fetch_models_with_context(
|
||||
CodexRequestContext::auxiliary(
|
||||
authorization.expose_secret(),
|
||||
account.upstream_account_id(),
|
||||
&request_id,
|
||||
None,
|
||||
),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
let snapshot = result.map_err(|error| CodexCredentialCatalogError::Upstream {
|
||||
detail: error.to_string(),
|
||||
})?;
|
||||
let (models, document, etag) = snapshot.into_parts();
|
||||
Ok(FetchedAccountModels {
|
||||
models,
|
||||
document,
|
||||
etag,
|
||||
})
|
||||
})
|
||||
.await
|
||||
.map_err(|_| CodexCredentialCatalogError::Upstream {
|
||||
detail: "model catalog timed out".to_owned(),
|
||||
})?
|
||||
}
|
||||
}
|
||||
|
||||
fn is_specialized_api_model(model: &str) -> bool {
|
||||
// 普通目录不提供 Codex 可见性;仅按已知的专用模型 ID 隐藏,保留未知 ID 和直接调用
|
||||
// 依据 https://developers.openai.com/api/docs/models 与 Codex 的语音通话模型
|
||||
matches!(
|
||||
model,
|
||||
"dall-e-2"
|
||||
| "dall-e-3"
|
||||
| "chatgpt-image-latest"
|
||||
| "gpt-image-1"
|
||||
| "gpt-image-1-mini"
|
||||
| "gpt-image-1.5"
|
||||
| "gpt-image-1.5-2025-12-16"
|
||||
| "gpt-image-2"
|
||||
| "gpt-image-2-2026-04-21"
|
||||
| "gpt-image-2.5-flare"
|
||||
| "gpt-image-2.5-flare-2026-09-08"
|
||||
| "gpt-image-2.5-sunburst"
|
||||
| "gpt-image-2.5-sunburst-2026-09-08"
|
||||
| "gpt-image-2.6-flare"
|
||||
| "gpt-image-2.6-sunburst"
|
||||
| "gpt-live-1"
|
||||
| "gpt-live-1-codex"
|
||||
| "gpt-live-transcribe"
|
||||
| "gpt-realtime"
|
||||
| "gpt-realtime-2025-08-28"
|
||||
| "gpt-realtime-1.5"
|
||||
| "gpt-realtime-2"
|
||||
| "gpt-realtime-2.1"
|
||||
| "gpt-realtime-2.1-mini"
|
||||
| "gpt-realtime-mini"
|
||||
| "gpt-realtime-mini-2025-10-06"
|
||||
| "gpt-realtime-mini-2025-12-15"
|
||||
| "gpt-realtime-translate"
|
||||
| "gpt-realtime-whisper"
|
||||
| "gpt-4o-realtime-preview"
|
||||
| "gpt-4o-realtime-preview-2024-10-01"
|
||||
| "gpt-4o-realtime-preview-2024-12-17"
|
||||
| "gpt-4o-realtime-preview-2025-06-03"
|
||||
| "gpt-4o-mini-realtime-preview"
|
||||
| "gpt-4o-mini-realtime-preview-2024-12-17"
|
||||
| "tts-1"
|
||||
| "tts-1-hd"
|
||||
| "gpt-4o-mini-tts"
|
||||
| "gpt-4o-mini-tts-2025-03-20"
|
||||
| "gpt-4o-mini-tts-2025-12-15"
|
||||
)
|
||||
}
|
||||
|
||||
fn model_ids(models: &[CodexCatalogModel]) -> Vec<String> {
|
||||
models
|
||||
.iter()
|
||||
.map(|model| model.request_model().as_str().to_owned())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn encode_account_catalog(catalog: &CodexAccountCatalog) -> 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_account_catalog(
|
||||
scope: &ProviderCatalogScope,
|
||||
document: OpaqueProviderData,
|
||||
) -> Result<CodexAccountCatalog, 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_ACCOUNT_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(CodexAccountCatalog::new(scope.clone(), observed_at, models))
|
||||
}
|
||||
|
||||
fn validate_response_etag(etag: &str) -> Result<(), CodexCredentialCatalogError> {
|
||||
if etag.is_empty() || etag.len() > MAX_RESPONSE_ETAG_BYTES || etag.chars().any(char::is_control)
|
||||
{
|
||||
return Err(CodexCredentialCatalogError::InvalidEtag);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn eligible_catalog_account(account: &ProviderAccount, now: SystemTime) -> bool {
|
||||
account.enabled()
|
||||
&& account
|
||||
.access_token_expires_at()
|
||||
.is_none_or(|expires_at| expires_at > now)
|
||||
&& matches!(
|
||||
account.credential_state(),
|
||||
CredentialState::Unknown | CredentialState::Ready
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,366 @@
|
||||
//! 账号目录的批量恢复、变更通知和限速刷新
|
||||
|
||||
use super::*;
|
||||
|
||||
impl CodexCredentialCatalogService {
|
||||
/// 启动只恢复已成功落库、仍匹配账号身份的目录
|
||||
pub async fn restore(&self) -> Result<(), CodexCredentialCatalogError> {
|
||||
let revision = self.state_revision()?;
|
||||
let provider =
|
||||
ProviderKind::new("openai").map_err(|_| CodexCredentialCatalogError::Cache)?;
|
||||
let groups = self
|
||||
.repository
|
||||
.store()
|
||||
.load_account_catalogs(&provider)
|
||||
.await
|
||||
.map_err(|_| CodexCredentialCatalogError::Cache)?;
|
||||
let mut entries = Entries::new();
|
||||
for group in groups {
|
||||
let decoded =
|
||||
match crate::transport::CodexModelCatalogSnapshot::restore(&group.document) {
|
||||
Ok(decoded) => decoded,
|
||||
Err(error) => {
|
||||
tracing::warn!(%error,"OpenAI stored model catalog is invalid");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let (models, _, _) = decoded.into_parts();
|
||||
let content = Arc::new(CatalogContent::new(models, &group.document));
|
||||
for a in group.accounts {
|
||||
let scope = CodexCatalogScope {
|
||||
account_id: a.account_id,
|
||||
revision: a.revision,
|
||||
plan: a
|
||||
.plan
|
||||
.map(|p| p.trim().to_ascii_lowercase())
|
||||
.filter(|p| !p.is_empty()),
|
||||
};
|
||||
entries.insert(
|
||||
scope,
|
||||
AccountEntry {
|
||||
observed_at: a.observed_at,
|
||||
content: content.clone(),
|
||||
etag: a.etag,
|
||||
invalidated: Default::default(),
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
let _writer = self
|
||||
.publication
|
||||
.lock()
|
||||
.map_err(|_| CodexCredentialCatalogError::Cache)?;
|
||||
self.publish(revision, entries)?;
|
||||
self.cache
|
||||
.write()
|
||||
.map_err(|_| CodexCredentialCatalogError::Cache)?
|
||||
.restored = true;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) fn state_revision(&self) -> Result<u64, CodexCredentialCatalogError> {
|
||||
self.cache
|
||||
.read()
|
||||
.map(|c| c.revision)
|
||||
.map_err(|_| CodexCredentialCatalogError::Cache)
|
||||
}
|
||||
|
||||
pub async fn invalidate_accounts(
|
||||
&self,
|
||||
accounts: &[ProviderAccountId],
|
||||
) -> Result<(), CodexCredentialCatalogError> {
|
||||
let ids = accounts.iter().collect::<BTreeSet<_>>();
|
||||
let current = self
|
||||
.repository
|
||||
.list_for_accounts(accounts)
|
||||
.await
|
||||
.inspect_err(|_| self.schedule_refresh())?;
|
||||
let current = current
|
||||
.iter()
|
||||
.filter(|a| eligible_catalog_account(a, SystemTime::now()))
|
||||
.map(CodexCatalogScope::for_account)
|
||||
.collect::<BTreeSet<_>>();
|
||||
let _writer = self
|
||||
.publication
|
||||
.lock()
|
||||
.map_err(|_| CodexCredentialCatalogError::Cache)?;
|
||||
let revision = {
|
||||
let mut cache = self
|
||||
.cache
|
||||
.write()
|
||||
.map_err(|_| CodexCredentialCatalogError::Cache)?;
|
||||
cache.revision = cache.revision.wrapping_add(1);
|
||||
cache
|
||||
.pending
|
||||
.retain(|scope, _| !ids.contains(&scope.account_id) || current.contains(scope));
|
||||
cache.revision
|
||||
};
|
||||
if let Some(snapshot) = self.cached()? {
|
||||
let keep = |scope: &CodexCatalogScope| {
|
||||
!ids.contains(&scope.account_id) || current.contains(scope)
|
||||
};
|
||||
if snapshot.entries.keys().any(|s| !keep(s)) {
|
||||
let mut entries = snapshot.entries.clone();
|
||||
entries.retain(|s, _| keep(s));
|
||||
self.publish(revision, entries)?;
|
||||
}
|
||||
}
|
||||
self.schedule_refresh();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn observe_response_etag(
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
etag: &str,
|
||||
) -> Result<bool, CodexCredentialCatalogError> {
|
||||
validate_response_etag(etag)?;
|
||||
let scope = CodexCatalogScope::for_account(account);
|
||||
let mut state = self
|
||||
.cache
|
||||
.write()
|
||||
.map_err(|_| CodexCredentialCatalogError::Cache)?;
|
||||
if state.pending.get(&scope).map(String::as_str) == Some(etag)
|
||||
|| !state.pending.contains_key(&scope)
|
||||
&& state
|
||||
.snapshot
|
||||
.as_ref()
|
||||
.and_then(|s| s.entries.get(&scope))
|
||||
.and_then(|e| e.etag.as_deref())
|
||||
== Some(etag)
|
||||
{
|
||||
return Ok(false);
|
||||
}
|
||||
if let Some(entry) = state.snapshot.as_ref().and_then(|s| s.entries.get(&scope)) {
|
||||
entry
|
||||
.invalidated
|
||||
.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||
}
|
||||
state.pending.insert(scope, etag.to_owned());
|
||||
drop(state);
|
||||
self.refresh_notification.notify_one();
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
pub(crate) fn schedule_refresh(&self) {
|
||||
self.refresh_notification.notify_one();
|
||||
}
|
||||
|
||||
pub async fn wait_for_refresh(&self) {
|
||||
self.refresh_notification.notified().await;
|
||||
}
|
||||
|
||||
/// 只合并成功的账号;手动刷新和后台任务同时完成时,保留较新的观测
|
||||
pub(super) fn merge_observations(
|
||||
&self,
|
||||
revision: u64,
|
||||
batch: Vec<(CodexCatalogScope, AccountEntry)>,
|
||||
) -> Result<(), CodexCredentialCatalogError> {
|
||||
let _writer = self
|
||||
.publication
|
||||
.lock()
|
||||
.map_err(|_| CodexCredentialCatalogError::Cache)?;
|
||||
if self.state_revision()? != revision {
|
||||
return Err(CodexCredentialCatalogError::ConcurrentUpdate);
|
||||
}
|
||||
let mut entries = self
|
||||
.cached()?
|
||||
.map(|snapshot| snapshot.entries.clone())
|
||||
.unwrap_or_default();
|
||||
for (scope, entry) in batch {
|
||||
if entries
|
||||
.get(&scope)
|
||||
.is_none_or(|old| old.observed_at < entry.observed_at)
|
||||
{
|
||||
entries.insert(scope, entry);
|
||||
}
|
||||
}
|
||||
self.publish(revision, entries)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn publish(
|
||||
&self,
|
||||
revision: u64,
|
||||
entries: Entries,
|
||||
) -> Result<Option<Arc<CodexCredentialCatalogSnapshot>>, CodexCredentialCatalogError> {
|
||||
let next =
|
||||
(!entries.is_empty()).then(|| Arc::new(CodexCredentialCatalogSnapshot::new(entries)));
|
||||
let old = self.cached()?;
|
||||
let changed = match (&old, &next) {
|
||||
(Some(a), Some(b)) => {
|
||||
a.entries.len() != b.entries.len()
|
||||
|| a.entries.iter().zip(&b.entries).any(|((sa, a), (sb, b))| {
|
||||
sa != sb
|
||||
|| (!Arc::ptr_eq(&a.content, &b.content)
|
||||
&& a.content.digest != b.content.digest)
|
||||
})
|
||||
}
|
||||
_ => old.is_some() != next.is_some(),
|
||||
};
|
||||
let mut cache = self
|
||||
.cache
|
||||
.write()
|
||||
.map_err(|_| CodexCredentialCatalogError::Cache)?;
|
||||
if cache.revision != revision {
|
||||
return Err(CodexCredentialCatalogError::ConcurrentUpdate);
|
||||
}
|
||||
if let Some(snapshot) = &next {
|
||||
cache.pending.retain(|scope, tag| {
|
||||
let Some(entry) = snapshot.entries.get(scope) else {
|
||||
return true;
|
||||
};
|
||||
if entry.etag.as_deref() == Some(tag)
|
||||
&& !entry.invalidated.load(std::sync::atomic::Ordering::Relaxed)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
entry
|
||||
.invalidated
|
||||
.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||
true
|
||||
});
|
||||
}
|
||||
cache.snapshot = next.clone();
|
||||
if changed {
|
||||
cache.generation = cache.generation.wrapping_add(1);
|
||||
}
|
||||
Ok(next)
|
||||
}
|
||||
|
||||
/// 一次对账只探测缺失、过期或有新 ETag 的账号;成功逐批发布,失败不更新观测时间
|
||||
pub async fn refresh_catalogs(
|
||||
&self,
|
||||
) -> Result<Arc<CodexCredentialCatalogSnapshot>, CodexCredentialCatalogError> {
|
||||
let _writer = self.refresh_lock.lock().await;
|
||||
// 启动时数据库暂不可用,应先重试恢复,不能把整池误判为没有历史
|
||||
if !self
|
||||
.cache
|
||||
.read()
|
||||
.map_err(|_| CodexCredentialCatalogError::Cache)?
|
||||
.restored
|
||||
{
|
||||
self.restore().await?;
|
||||
}
|
||||
let revision = self.state_revision()?;
|
||||
let now = SystemTime::now();
|
||||
let accounts = self.repository.list_for_provider().await?;
|
||||
let current: BTreeMap<_, _> = accounts
|
||||
.into_iter()
|
||||
.filter(|a| eligible_catalog_account(a, now))
|
||||
.map(|a| (CodexCatalogScope::for_account(&a), a))
|
||||
.collect();
|
||||
let pending = {
|
||||
let mut cache = self
|
||||
.cache
|
||||
.write()
|
||||
.map_err(|_| CodexCredentialCatalogError::Cache)?;
|
||||
cache.pending.retain(|scope, _| current.contains_key(scope));
|
||||
cache.pending.clone()
|
||||
};
|
||||
let previous = self.cached()?;
|
||||
let mut entries = previous
|
||||
.as_ref()
|
||||
.map(|s| s.entries.clone())
|
||||
.unwrap_or_default();
|
||||
entries.retain(|scope, _| current.contains_key(scope));
|
||||
if previous
|
||||
.as_ref()
|
||||
.is_some_and(|s| s.entries.len() != entries.len())
|
||||
{
|
||||
let _writer = self
|
||||
.publication
|
||||
.lock()
|
||||
.map_err(|_| CodexCredentialCatalogError::Cache)?;
|
||||
let mut latest = self
|
||||
.cached()?
|
||||
.map(|s| s.entries.clone())
|
||||
.unwrap_or_default();
|
||||
latest.retain(|scope, _| current.contains_key(scope));
|
||||
self.publish(revision, latest)?;
|
||||
}
|
||||
let mut contents: BTreeMap<_, _> = entries
|
||||
.values()
|
||||
.map(|e| (e.content.digest, e.content.clone()))
|
||||
.collect();
|
||||
let due = current
|
||||
.into_iter()
|
||||
.filter(|(scope, _)| {
|
||||
pending.contains_key(scope)
|
||||
|| !entries.get(scope).is_some_and(|e| {
|
||||
now.duration_since(e.observed_at)
|
||||
.is_ok_and(|age| age < SCHEDULING_CATALOG_MAX_AGE)
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
if due.is_empty() {
|
||||
return self
|
||||
.cached()?
|
||||
.ok_or(CodexCredentialCatalogError::NoEligibleCredential);
|
||||
}
|
||||
let client = CodexBackendClient::new(
|
||||
self.http.clone(),
|
||||
self.base_url.clone(),
|
||||
self.profile.clone(),
|
||||
);
|
||||
let mut pace = tokio::time::interval(REFRESH_SPACING);
|
||||
pace.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
|
||||
let starts = stream::unfold(pace, |mut clock| async move {
|
||||
clock.tick().await;
|
||||
Some(((), clock))
|
||||
});
|
||||
let fetches = stream::iter(due)
|
||||
.zip(starts)
|
||||
.map(|((scope, account), ())| {
|
||||
let client = &client;
|
||||
let etag = pending.get(&scope).cloned();
|
||||
async move {
|
||||
let (_, entry) = self
|
||||
.fetch_observation(client, &account, &scope, etag.clone())
|
||||
.await?;
|
||||
Ok((scope, entry))
|
||||
}
|
||||
})
|
||||
.buffer_unordered(REFRESH_CONCURRENCY);
|
||||
tokio::pin!(fetches);
|
||||
let mut batch = Vec::with_capacity(PUBLICATION_BATCH);
|
||||
let mut completed = 0_usize;
|
||||
let mut failures = 0_usize;
|
||||
let mut last_error = None;
|
||||
while let Some(result) = fetches.next().await {
|
||||
match result {
|
||||
Ok((scope, mut entry)) => {
|
||||
entry.content = contents
|
||||
.entry(entry.content.digest)
|
||||
.or_insert_with(|| entry.content.clone())
|
||||
.clone();
|
||||
batch.push((scope, entry));
|
||||
completed += 1;
|
||||
if completed == 1 || completed.is_multiple_of(PUBLICATION_BATCH) {
|
||||
self.merge_observations(revision, std::mem::take(&mut batch))?;
|
||||
}
|
||||
}
|
||||
Err(error) => {
|
||||
failures += 1;
|
||||
last_error = Some(error);
|
||||
}
|
||||
}
|
||||
}
|
||||
if !batch.is_empty() {
|
||||
self.merge_observations(revision, batch)?;
|
||||
}
|
||||
if let Some(error) = &last_error {
|
||||
tracing::warn!(successes=completed,failures,%error,"OpenAI account catalog refresh incomplete");
|
||||
}
|
||||
if pending
|
||||
.keys()
|
||||
.any(|k| self.cache.read().is_ok_and(|c| c.pending.contains_key(k)))
|
||||
&& let Some(error) = last_error.clone()
|
||||
{
|
||||
self.refresh_notification.notify_one();
|
||||
return Err(error);
|
||||
}
|
||||
self.cached()?
|
||||
.ok_or(last_error.unwrap_or(CodexCredentialCatalogError::NoEligibleCredential))
|
||||
}
|
||||
}
|
||||
@@ -752,7 +752,39 @@ impl CodexCredentialSelector {
|
||||
.into_iter()
|
||||
.flat_map(|scope| scope.allowed_account_ids().cloned())
|
||||
.collect::<Vec<_>>();
|
||||
self.repository.list_for_accounts(&account_ids).await?
|
||||
let omissions = match (model, request.attempt.account_scope()) {
|
||||
(CodexSelectionModel::Requested(model), Some(scope)) => {
|
||||
self.catalog.omissions(model, &account_ids, |id| {
|
||||
scope.directory().account(id).is_some_and(|a| {
|
||||
a.model_access().mode() == AccountModelAccessMode::Allowlist
|
||||
}) || adapter_model.as_ref().is_some_and(|m| {
|
||||
!matches!(
|
||||
request.attempt.upstream_adapter_for_account(
|
||||
&self.provider_kind,
|
||||
m,
|
||||
id
|
||||
),
|
||||
Ok(None)
|
||||
)
|
||||
})
|
||||
})
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
if let Some(filter) = omissions {
|
||||
let candidates = self
|
||||
.repository
|
||||
.store()
|
||||
.list_catalog_candidates(&self.provider_kind, &account_ids, &filter)
|
||||
.await
|
||||
.map_err(CredentialRepositoryError::from)?;
|
||||
if request.attempt.trace().is_enabled() {
|
||||
request.attempt.trace().record("account.catalog_candidates",serde_json::json!({"scopedCount":account_ids.len(),"candidateCount":candidates.len()}));
|
||||
}
|
||||
candidates
|
||||
} else {
|
||||
self.repository.list_for_accounts(&account_ids).await?
|
||||
}
|
||||
};
|
||||
let mut model_access_rejected = 0_usize;
|
||||
let accounts = accounts
|
||||
@@ -778,20 +810,14 @@ impl CodexCredentialSelector {
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let mut eligible = Vec::with_capacity(accounts.len());
|
||||
let catalog_excluded = match (model, request.attempt.account_scope()) {
|
||||
(CodexSelectionModel::Requested(model), Some(scope)) if !diagnostic => {
|
||||
// 候选已通过冻结的模型政策,白名单精确命中的 ID 交给上游验证
|
||||
let mut excluded =
|
||||
self.catalog
|
||||
.scheduling_exclusions(&accounts, model, SystemTime::now());
|
||||
excluded.retain(|id| {
|
||||
scope.directory().account(id).is_none_or(|account| {
|
||||
account.model_access().mode() != AccountModelAccessMode::Allowlist
|
||||
})
|
||||
});
|
||||
excluded
|
||||
}
|
||||
_ => Default::default(),
|
||||
let catalog = match (model, request.attempt.account_scope()) {
|
||||
(CodexSelectionModel::Requested(model), Some(scope)) if !diagnostic => self
|
||||
.catalog
|
||||
.cached()
|
||||
.ok()
|
||||
.flatten()
|
||||
.map(|snapshot| (snapshot, model, scope, SystemTime::now())),
|
||||
_ => None,
|
||||
};
|
||||
let mut catalog_rejected = 0_usize;
|
||||
let mut native_accounts = std::collections::BTreeSet::new();
|
||||
@@ -816,7 +842,15 @@ impl CodexCredentialSelector {
|
||||
}
|
||||
// 受管适配器拥有其上游模型合同,不套用原生目录与拒绝反馈
|
||||
if adapter.is_none() {
|
||||
if catalog_excluded.contains(account.id()) {
|
||||
if catalog
|
||||
.as_ref()
|
||||
.is_some_and(|(snapshot, model, scope, now)| {
|
||||
snapshot.excludes(&account, model, *now)
|
||||
&& scope.directory().account(account.id()).is_none_or(|a| {
|
||||
a.model_access().mode() != AccountModelAccessMode::Allowlist
|
||||
})
|
||||
})
|
||||
{
|
||||
catalog_rejected += 1;
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -153,6 +153,9 @@ pub async fn initialize(
|
||||
config.base_url().to_owned(),
|
||||
ports.catalog_cache(),
|
||||
));
|
||||
if let Err(error) = catalog.restore().await {
|
||||
tracing::warn!(%error,"OpenAI model catalog could not be restored");
|
||||
}
|
||||
let quota = Arc::new(CodexCredentialQuotaService::new(
|
||||
repository.clone(),
|
||||
profile.clone(),
|
||||
|
||||
@@ -943,7 +943,7 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
yield ProviderEvent::observation(observation);
|
||||
}
|
||||
if let Some(etag) = response.response_metadata.models_etag.as_deref()
|
||||
&& let Err(error) = catalog.observe_response_etag(etag)
|
||||
&& let Err(error) = catalog.observe_response_etag(&active_account,etag)
|
||||
{
|
||||
tracing::warn!(
|
||||
error = %error,
|
||||
|
||||
@@ -49,12 +49,10 @@ impl CodexProvider {
|
||||
gateway_protocol::anthropic::strip_foreign_thinking_for_messages(&mut body);
|
||||
body
|
||||
} else {
|
||||
let capabilities = tokio::select! {
|
||||
biased;
|
||||
_ = context.cancellation().cancelled() => return Err(provider_error(ProviderErrorKind::Cancelled, UpstreamSendState::NotSent)),
|
||||
_ = context.deadline().wait() => return Err(provider_error(ProviderErrorKind::Timeout, UpstreamSendState::NotSent)),
|
||||
result = self.catalog.messages_model_capabilities(selected.lease.account(), &model) => result.ok().flatten().unwrap_or_default(),
|
||||
};
|
||||
let capabilities = self
|
||||
.catalog
|
||||
.messages_model_capabilities(selected.lease.account(), &model)
|
||||
.unwrap_or_default();
|
||||
let mut body = responses_to_messages_request(original, &capabilities)
|
||||
.map_err(|_| invalid_request())?;
|
||||
// Responses 的 HTTP JSON 和 SSE 共用事件输出,跨协议上游统一请求事件流
|
||||
|
||||
@@ -536,9 +536,16 @@ impl Provider for CodexProvider {
|
||||
async fn query_model_capabilities(
|
||||
&self,
|
||||
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
|
||||
let snapshot = self.catalog.synchronize().await.map_err(|_| {
|
||||
provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
|
||||
})?;
|
||||
// 能力发现读取已发布事实;冷启动和配置失效由 Provider 后台任务重建目录
|
||||
let snapshot = self
|
||||
.catalog
|
||||
.cached()
|
||||
.map_err(|_| {
|
||||
provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
|
||||
})?
|
||||
.ok_or_else(|| {
|
||||
provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
|
||||
})?;
|
||||
let mut accounts = snapshot.model_catalog_accounts();
|
||||
Ok(snapshot
|
||||
.models()
|
||||
@@ -555,11 +562,11 @@ impl Provider for CodexProvider {
|
||||
|
||||
async fn query_client_model_catalog(
|
||||
&self,
|
||||
client_api_key_id: &gateway_core::policy::ClientApiKeyId,
|
||||
_client_api_key_id: &gateway_core::policy::ClientApiKeyId,
|
||||
scope: &gateway_core::account::scope::FrozenAccountScope,
|
||||
protocol: &str,
|
||||
client_version: &str,
|
||||
peer_route: Option<&reqwest::header::HeaderValue>,
|
||||
_client_version: &str,
|
||||
_peer_route: Option<&reqwest::header::HeaderValue>,
|
||||
) -> Result<
|
||||
Option<Vec<gateway_core::routing::ProviderModelDescriptor>>,
|
||||
gateway_core::routing::ProviderCatalogUnavailable,
|
||||
@@ -568,7 +575,7 @@ impl Provider for CodexProvider {
|
||||
return Ok(None);
|
||||
}
|
||||
self.catalog
|
||||
.client_model_catalog(client_api_key_id, scope, client_version, peer_route)
|
||||
.client_model_catalog(scope)
|
||||
.await
|
||||
.map(Some)
|
||||
.map_err(|_| gateway_core::routing::ProviderCatalogUnavailable)
|
||||
|
||||
@@ -46,7 +46,10 @@ pub(crate) fn worker_contributions(
|
||||
contributions.push(WorkerContribution::Registration(scheduled_registration(
|
||||
refresh_id,
|
||||
OAUTH_REFRESH_INTERVAL,
|
||||
Box::new(OpenAiOAuthRefreshTask { service: refresh }),
|
||||
Box::new(OpenAiOAuthRefreshTask {
|
||||
service: refresh,
|
||||
catalog: Arc::clone(&catalog),
|
||||
}),
|
||||
)?));
|
||||
}
|
||||
contributions.extend([
|
||||
@@ -132,6 +135,7 @@ pub(super) fn scheduled_registration(
|
||||
}
|
||||
|
||||
pub(super) struct OpenAiOAuthRefreshTask {
|
||||
catalog: Arc<CodexCredentialCatalogService>,
|
||||
service: Arc<CodexCredentialRefreshService>,
|
||||
}
|
||||
|
||||
@@ -144,6 +148,7 @@ impl ScheduledTask for OpenAiOAuthRefreshTask {
|
||||
let outcomes = self.service.refresh_due().await.map_err(|error| {
|
||||
WorkerTaskError::safe("OpenAI OAuth refresh failed").with_source(error)
|
||||
})?;
|
||||
let mut changed = Vec::new();
|
||||
let mut refreshed = 0_u64;
|
||||
let mut invalidated = 0_u64;
|
||||
let mut banned = 0_u64;
|
||||
@@ -155,9 +160,18 @@ impl ScheduledTask for OpenAiOAuthRefreshTask {
|
||||
let mut failed_accounts = Vec::new();
|
||||
for outcome in &outcomes {
|
||||
match outcome {
|
||||
CodexCredentialRefreshOutcome::Refreshed { .. } => refreshed += 1,
|
||||
CodexCredentialRefreshOutcome::Invalidated { .. } => invalidated += 1,
|
||||
CodexCredentialRefreshOutcome::Banned { .. } => banned += 1,
|
||||
CodexCredentialRefreshOutcome::Refreshed { account_id, .. } => {
|
||||
refreshed += 1;
|
||||
changed.push(account_id);
|
||||
}
|
||||
CodexCredentialRefreshOutcome::Invalidated { account_id } => {
|
||||
invalidated += 1;
|
||||
changed.push(account_id);
|
||||
}
|
||||
CodexCredentialRefreshOutcome::Banned { account_id } => {
|
||||
banned += 1;
|
||||
changed.push(account_id);
|
||||
}
|
||||
CodexCredentialRefreshOutcome::Transient { account_id } => {
|
||||
transient += 1;
|
||||
transient_accounts.push(account_id);
|
||||
@@ -172,6 +186,21 @@ impl ScheduledTask for OpenAiOAuthRefreshTask {
|
||||
}
|
||||
}
|
||||
}
|
||||
if !changed.is_empty() {
|
||||
let ids = changed
|
||||
.into_iter()
|
||||
.map(|id| gateway_core::account::ProviderAccountId::new(id.as_str()))
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map_err(|error| {
|
||||
WorkerTaskError::safe("OpenAI 账号 ID 无效").with_source(error)
|
||||
})?;
|
||||
self.catalog
|
||||
.invalidate_accounts(&ids)
|
||||
.await
|
||||
.map_err(|error| {
|
||||
WorkerTaskError::safe("OpenAI 凭据变更后目录失效失败").with_source(error)
|
||||
})?;
|
||||
}
|
||||
if !outcomes.is_empty() {
|
||||
tracing::info!(
|
||||
refreshed,
|
||||
@@ -306,12 +335,21 @@ impl DaemonTask for OpenAiCatalogEtagTask {
|
||||
loop {
|
||||
tokio::select! {
|
||||
() = cancellation.cancelled() => return Ok(()),
|
||||
() = self.catalog.wait_for_etag_refresh() => {},
|
||||
() = self.catalog.wait_for_refresh() => {},
|
||||
};
|
||||
// 失败的 ETag 已重新入队,交回 Host 退避后再领取,避免立即重试
|
||||
self.catalog.refresh().await.map_err(|error| {
|
||||
WorkerTaskError::safe("OpenAI 模型目录 ETag 刷新失败").with_source(error)
|
||||
})?;
|
||||
let result = tokio::select! {
|
||||
()=cancellation.cancelled()=>return Ok(()),
|
||||
result=self.catalog.refresh_catalogs()=>result,
|
||||
};
|
||||
match result {
|
||||
Ok(_) | Err(CodexCredentialCatalogError::NoEligibleCredential) => {}
|
||||
Err(error) => {
|
||||
self.catalog.schedule_refresh();
|
||||
return Err(
|
||||
WorkerTaskError::safe("OpenAI 模型目录刷新失败").with_source(error)
|
||||
);
|
||||
}
|
||||
}
|
||||
context.succeeded();
|
||||
}
|
||||
})
|
||||
|
||||
@@ -237,13 +237,40 @@ impl CodexCatalogModel {
|
||||
}
|
||||
|
||||
/// 一次完整成功的 Codex 远端模型快照
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub struct CodexModelCatalogSnapshot {
|
||||
document: Bytes,
|
||||
models: Vec<CodexCatalogModel>,
|
||||
etag: Option<String>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for CodexModelCatalogSnapshot {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("CodexModelCatalogSnapshot")
|
||||
.field("models", &self.models)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl CodexModelCatalogSnapshot {
|
||||
pub(crate) fn into_parts(self) -> (Vec<CodexCatalogModel>, Bytes, Option<String>) {
|
||||
(self.models, self.document, self.etag)
|
||||
}
|
||||
|
||||
/// 落库的目录沿用相同的上游校验器;协议标记区分普通 API 与 Messages 能力
|
||||
pub(crate) fn restore(document: &[u8]) -> Result<Self, CodexModelCatalogError> {
|
||||
match document.split_first() {
|
||||
Some((0, body)) => parse_codex_model_catalog(body, None),
|
||||
Some((1, body)) => {
|
||||
parse_api_model_catalog(body, None, OpenAiUpstreamProtocol::ResponsesApi)
|
||||
}
|
||||
Some((2, body)) => {
|
||||
parse_api_model_catalog(body, None, OpenAiUpstreamProtocol::MessagesApi)
|
||||
}
|
||||
_ => Err(CodexModelCatalogError::InvalidWire),
|
||||
}
|
||||
}
|
||||
|
||||
/// 返回本轮全部模型
|
||||
#[must_use]
|
||||
pub fn models(&self) -> &[CodexCatalogModel] {
|
||||
@@ -289,7 +316,7 @@ pub enum CodexModelCatalogError {
|
||||
InvalidEtag,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[derive(Debug, Deserialize, serde::Serialize)]
|
||||
struct CodexModelsWire {
|
||||
models: Vec<serde_json::Value>,
|
||||
}
|
||||
@@ -361,6 +388,7 @@ fn parse_native_model_catalog(
|
||||
return Err(CodexModelCatalogError::TooManyModels);
|
||||
}
|
||||
|
||||
let document = archive(protocol, &wire)?;
|
||||
let etag = etag.map(validate_etag).transpose()?;
|
||||
let mut seen = BTreeSet::new();
|
||||
let mut models = Vec::with_capacity(wire.models.len());
|
||||
@@ -380,7 +408,11 @@ fn parse_native_model_catalog(
|
||||
}
|
||||
models.push(model);
|
||||
}
|
||||
Ok(CodexModelCatalogSnapshot { models, etag })
|
||||
Ok(CodexModelCatalogSnapshot {
|
||||
document,
|
||||
models,
|
||||
etag,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn catalog_etag(headers: &HeaderMap) -> Result<Option<String>, CodexModelCatalogError> {
|
||||
@@ -617,12 +649,30 @@ pub(super) fn parse_api_model_catalog(
|
||||
#[serde(flatten)]
|
||||
fields: serde_json::Map<String, serde_json::Value>,
|
||||
}
|
||||
let value: serde_json::Value =
|
||||
let mut value: serde_json::Value =
|
||||
serde_json::from_slice(body).map_err(|_| CodexModelCatalogError::InvalidWire)?;
|
||||
// 部分 Responses 上游按客户端版本返回完整 Codex 目录;不能降格为只有 ID 的画像
|
||||
if value.get("models").is_some() {
|
||||
return parse_native_model_catalog(body, etag, OpenAiUpstreamProtocol::ResponsesApi);
|
||||
}
|
||||
if let Some(top) = value.as_object_mut() {
|
||||
top.retain(|key, _| key == "data");
|
||||
}
|
||||
if let Some(items) = value
|
||||
.get_mut("data")
|
||||
.and_then(serde_json::Value::as_array_mut)
|
||||
{
|
||||
for item in items {
|
||||
if let Some(fields) = item.as_object_mut() {
|
||||
fields.retain(|key, _| {
|
||||
matches!(key.as_str(), "id" | "display_name")
|
||||
|| protocol == OpenAiUpstreamProtocol::MessagesApi
|
||||
&& matches!(key.as_str(), "capabilities" | "max_tokens")
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
let document = archive(protocol, &value)?;
|
||||
let wire: ApiModels =
|
||||
serde_json::from_value(value).map_err(|_| CodexModelCatalogError::InvalidWire)?;
|
||||
if wire.data.is_empty() {
|
||||
@@ -668,7 +718,21 @@ pub(super) fn parse_api_model_catalog(
|
||||
});
|
||||
}
|
||||
Ok(CodexModelCatalogSnapshot {
|
||||
document,
|
||||
models,
|
||||
etag: etag.map(validate_etag).transpose()?,
|
||||
})
|
||||
}
|
||||
|
||||
fn archive(
|
||||
protocol: OpenAiUpstreamProtocol,
|
||||
value: &impl serde::Serialize,
|
||||
) -> Result<Bytes, CodexModelCatalogError> {
|
||||
let mut bytes = vec![match protocol {
|
||||
OpenAiUpstreamProtocol::Codex => 0,
|
||||
OpenAiUpstreamProtocol::ResponsesApi => 1,
|
||||
OpenAiUpstreamProtocol::MessagesApi => 2,
|
||||
}];
|
||||
serde_json::to_writer(&mut bytes, value).map_err(|_| CodexModelCatalogError::InvalidWire)?;
|
||||
Ok(bytes.into())
|
||||
}
|
||||
|
||||
+202
-626
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,239 @@
|
||||
//! 验证大账号池与慢目录上游不会阻塞启动时的能力读取
|
||||
|
||||
use super::*;
|
||||
use gateway_core::engine::provider::{Provider, ProviderRegistry};
|
||||
use gateway_core::routing::snapshot::{
|
||||
RuntimeSnapshotCompiler, SnapshotFacts, SnapshotProviderAccountFacts, SnapshotStoreError,
|
||||
SnapshotStorePort,
|
||||
};
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use tokio::sync::Notify;
|
||||
use wiremock::{Request, Respond};
|
||||
|
||||
struct StartupStore(SnapshotFacts);
|
||||
|
||||
impl SnapshotStorePort for StartupStore {
|
||||
fn load_snapshot_facts(&self) -> BoxFuture<'_, Result<SnapshotFacts, SnapshotStoreError>> {
|
||||
Box::pin(async { Ok(self.0.clone()) })
|
||||
}
|
||||
|
||||
fn current_config_revision(&self) -> BoxFuture<'_, Result<ConfigRevision, SnapshotStoreError>> {
|
||||
Box::pin(async { Ok(self.0.config_revision()) })
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cold_catalog_with_ten_thousand_accounts_never_fetches_during_capability_reads() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let mut accounts = Vec::new();
|
||||
for index in 0..10_000 {
|
||||
let id = format!("acct_startup_{index:05}");
|
||||
create_account(&store, &id).await;
|
||||
accounts.push(SnapshotProviderAccountFacts::new(
|
||||
ProviderAccountId::new(id).unwrap(),
|
||||
"openai",
|
||||
));
|
||||
}
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.respond_with(ResponseTemplate::new(503).set_delay(Duration::from_secs(30)))
|
||||
.expect(0)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let revision = ConfigRevision::new(1).unwrap();
|
||||
let facts = Arc::new(StartupStore(SnapshotFacts::new(
|
||||
revision,
|
||||
revision,
|
||||
SettingsValues::new(3, "smart", BTreeMap::new(), None, None),
|
||||
vec![],
|
||||
vec![],
|
||||
accounts,
|
||||
vec![],
|
||||
)));
|
||||
|
||||
// 重建 Provider 模拟重启;Core 的读取只报告目录未知,不触发全量账号扫描
|
||||
for _ in 0..3 {
|
||||
let provider = provider_with_base_url(&store, server.uri());
|
||||
let error =
|
||||
tokio::time::timeout(Duration::from_secs(1), provider.query_model_capabilities())
|
||||
.await
|
||||
.expect("startup capability read must not wait for upstream")
|
||||
.expect_err("a cold catalog is unknown");
|
||||
assert_eq!(error.kind(), ProviderErrorKind::Unavailable);
|
||||
let registry = Arc::new(ProviderRegistry::new([provider as Arc<dyn Provider>]).unwrap());
|
||||
let compiler = RuntimeSnapshotCompiler::new(facts.clone(), registry);
|
||||
let snapshot = tokio::time::timeout(Duration::from_secs(5), compiler.compile())
|
||||
.await
|
||||
.expect("startup snapshot must not wait for the account scan")
|
||||
.expect("cold catalog must not prevent snapshot publication");
|
||||
assert!(snapshot.public_models().is_empty());
|
||||
}
|
||||
server.verify().await;
|
||||
}
|
||||
|
||||
struct SlowCatalog {
|
||||
calls: Arc<AtomicUsize>,
|
||||
first_batch: Arc<Notify>,
|
||||
}
|
||||
|
||||
impl Respond for SlowCatalog {
|
||||
fn respond(&self, _: &Request) -> ResponseTemplate {
|
||||
if self.calls.fetch_add(1, Ordering::SeqCst) + 1 == 32 {
|
||||
self.first_batch.notify_one();
|
||||
}
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_raw(OFFICIAL_FIXTURE.to_vec(), "application/json")
|
||||
.set_delay(Duration::from_secs(30))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn slow_background_catalog_is_bounded_and_does_not_block_capability_reads() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
for index in 0..64 {
|
||||
create_account(&store, &format!("acct_background_{index:02}")).await;
|
||||
}
|
||||
let server = MockServer::start().await;
|
||||
let calls = Arc::new(AtomicUsize::new(0));
|
||||
let first_batch = Arc::new(Notify::new());
|
||||
Mock::given(method("GET"))
|
||||
.respond_with(SlowCatalog {
|
||||
calls: Arc::clone(&calls),
|
||||
first_batch: Arc::clone(&first_batch),
|
||||
})
|
||||
.mount(&server)
|
||||
.await;
|
||||
let (provider, _, _, catalog) = provider_and_quota_with_runtime_ports(
|
||||
&store,
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
server.uri(),
|
||||
Arc::new(TestLeaseCoordinator::default()),
|
||||
u32::try_from(DEFAULT_STREAM_MAX_RETRIES).unwrap(),
|
||||
Arc::new(MemoryCooldownPort::new()),
|
||||
crate::support::runtime_policy(),
|
||||
);
|
||||
let refresh = {
|
||||
let catalog = Arc::clone(&catalog);
|
||||
tokio::spawn(async move { catalog.refresh_catalogs().await })
|
||||
};
|
||||
tokio::time::timeout(Duration::from_secs(5), first_batch.notified())
|
||||
.await
|
||||
.expect("background scan must make progress beyond four accounts");
|
||||
assert_eq!(calls.load(Ordering::SeqCst), 32);
|
||||
assert!(
|
||||
tokio::time::timeout(Duration::from_secs(1), provider.query_model_capabilities(),)
|
||||
.await
|
||||
.expect("capability read must not join a slow refresh")
|
||||
.is_err()
|
||||
);
|
||||
|
||||
// 取消整轮扫描后刷新锁必须释放,下一轮后台任务仍能完成
|
||||
refresh.abort();
|
||||
assert!(refresh.await.unwrap_err().is_cancelled());
|
||||
server.reset().await;
|
||||
Mock::given(method("GET"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_raw(OFFICIAL_FIXTURE.to_vec(), "application/json"),
|
||||
)
|
||||
.expect(64)
|
||||
.mount(&server)
|
||||
.await;
|
||||
tokio::time::timeout(Duration::from_secs(5), catalog.refresh_catalogs())
|
||||
.await
|
||||
.expect("refresh cancellation must release the lock")
|
||||
.expect("catalog refresh");
|
||||
let generation = provider.catalog_generation();
|
||||
for _ in 0..3 {
|
||||
assert!(provider.query_model_capabilities().await.is_ok());
|
||||
}
|
||||
assert_eq!(provider.catalog_generation(), generation);
|
||||
server.verify().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn issue_666_same_plan_special_model_remains_visible_and_routes_only_to_its_source_after_restart()
|
||||
{
|
||||
let server = MockServer::start().await;
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let source = "acct_provider_contract";
|
||||
let other = "acct_scope_new";
|
||||
for (id, models) in [
|
||||
(source, vec!["common", "special-666"]),
|
||||
(other, vec!["common"]),
|
||||
] {
|
||||
create_account(&store, id).await;
|
||||
Mock::given(method("GET")).and(path("/codex/models")).and(header("chatgpt-account-id",format!("chatgpt-{id}")))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"models":models.iter().map(|id|json!({"slug":id,"display_name":id})).collect::<Vec<_>>()})))
|
||||
.expect(1).mount(&server).await;
|
||||
}
|
||||
// 调度本来优先选 B,只有 A 自己的目录能授权特殊模型
|
||||
store.set_scheduling(
|
||||
other,
|
||||
None,
|
||||
gateway_core::account::AccountWeight::new(100).unwrap(),
|
||||
);
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.and(header("chatgpt-account-id", format!("chatgpt-{source}")))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_raw(CAPTURE_COMPLETED_SSE, "text/event-stream"),
|
||||
)
|
||||
.expect(2)
|
||||
.mount(&server)
|
||||
.await;
|
||||
for restarted in [false, true] {
|
||||
let (provider, _, _, catalog) = provider_and_quota_with_runtime_ports(
|
||||
&store,
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
server.uri(),
|
||||
Arc::new(TestLeaseCoordinator::default()),
|
||||
u32::try_from(DEFAULT_STREAM_MAX_RETRIES).unwrap(),
|
||||
Arc::new(MemoryCooldownPort::new()),
|
||||
crate::support::runtime_policy(),
|
||||
);
|
||||
catalog.restore().await.unwrap();
|
||||
if !restarted {
|
||||
assert!(provider.query_model_capabilities().await.is_err());
|
||||
}
|
||||
catalog.refresh_catalogs().await.unwrap();
|
||||
assert!(
|
||||
catalog
|
||||
.client_model_catalog(&contract_account_scope())
|
||||
.await
|
||||
.unwrap()
|
||||
.iter()
|
||||
.any(|m| m.model.as_str() == "special-666")
|
||||
);
|
||||
let sources = catalog.cached().unwrap().unwrap();
|
||||
assert_eq!(
|
||||
sources.model_catalog_accounts()["special-666"],
|
||||
std::collections::BTreeSet::from([gateway_core::account::ProviderAccountId::new(
|
||||
source
|
||||
)
|
||||
.unwrap()])
|
||||
);
|
||||
let op = Operation::Generate(GenerateRequest::from_protocol_payload(
|
||||
ProtocolPayload::json_object(
|
||||
"openai",
|
||||
json!({"model":"special-666","input":"hello"})
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone(),
|
||||
)
|
||||
.unwrap()
|
||||
.with_context(Map::from_iter([("use_websocket".into(), json!(false))])),
|
||||
));
|
||||
let mut stream = provider
|
||||
.execute(
|
||||
planned_request_for_model("openai", op, "special-666"),
|
||||
context("req_issue_666", CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(stream.metadata().provider_account_id().as_str(), source);
|
||||
while let Some(event) = stream.next().await {
|
||||
event.unwrap();
|
||||
}
|
||||
}
|
||||
server.verify().await;
|
||||
}
|
||||
@@ -500,7 +500,7 @@ async fn translated_thinking_uses_selected_account_catalog_and_preserves_fast_be
|
||||
.mount(server)
|
||||
.await;
|
||||
}
|
||||
let provider = provider(&store);
|
||||
let provider = provider_with_catalog_and_base_url(&store, adaptive.uri()).await;
|
||||
for _ in 0..2 {
|
||||
for account in ["acct_provider_contract", "acct_manual"] {
|
||||
let body = json!({"model":"unrelated-client-alias","input":"hello","max_output_tokens":4096,"reasoning":{"effort":"xhigh"},"tools":[{"type":"function","name":"echo"}],"parallel_tool_calls":false,"service_tier":"priority"});
|
||||
@@ -558,50 +558,20 @@ async fn translated_thinking_uses_selected_account_catalog_and_preserves_fast_be
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancelled_capability_lookup_never_sends_generation() {
|
||||
async fn cancelled_messages_request_never_sends_generation_or_catalog_lookup() {
|
||||
let server = MockServer::start().await;
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
seed(&store, &server, ApiKeyProtocol::Messages).await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/models"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_delay(Duration::from_secs(10))
|
||||
.set_body_json(json!({"data":[]})),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let cancellation = CancellationToken::new();
|
||||
let cancelled = cancellation.clone();
|
||||
let task = tokio::spawn(async move {
|
||||
provider(&store)
|
||||
.execute(
|
||||
planned_request("openai", generate_operation()),
|
||||
context("req_cancel_catalog", cancelled),
|
||||
)
|
||||
.await
|
||||
});
|
||||
tokio::time::timeout(Duration::from_secs(2), async {
|
||||
while server.received_requests().await.unwrap().is_empty() {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
cancellation.cancel();
|
||||
let result = tokio::time::timeout(Duration::from_secs(1), task)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let result = provider(&store)
|
||||
.execute(
|
||||
planned_request("openai", generate_operation()),
|
||||
context("req_cancel_catalog", cancellation),
|
||||
)
|
||||
.await;
|
||||
assert!(matches!(result,Err(error) if error.kind()==ProviderErrorKind::Cancelled));
|
||||
assert!(
|
||||
server
|
||||
.received_requests()
|
||||
.await
|
||||
.unwrap()
|
||||
.iter()
|
||||
.all(|request| request.method == "GET")
|
||||
);
|
||||
assert!(server.received_requests().await.unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -759,7 +729,7 @@ async fn oversized_messages_json_is_rejected_before_wire_delivery() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn slow_catalog_does_not_block_generation_and_populates_next_turn() {
|
||||
async fn cold_directory_never_queries_models_during_generation() {
|
||||
let server = MockServer::start().await;
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
seed(&store, &server, ApiKeyProtocol::Messages).await;
|
||||
@@ -772,7 +742,7 @@ async fn slow_catalog_does_not_block_generation_and_populates_next_turn() {
|
||||
json!({"data":[{"id":"gpt-5.4","max_tokens":2048}],"has_more":false}),
|
||||
),
|
||||
)
|
||||
.expect(1)
|
||||
.expect(0)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("POST"))
|
||||
@@ -782,7 +752,7 @@ async fn slow_catalog_does_not_block_generation_and_populates_next_turn() {
|
||||
.mount(&server)
|
||||
.await;
|
||||
let provider = provider(&store);
|
||||
for i in 0..2 {
|
||||
for _ in 0..2 {
|
||||
let mut stream = provider
|
||||
.clone()
|
||||
.execute(
|
||||
@@ -794,9 +764,6 @@ async fn slow_catalog_does_not_block_generation_and_populates_next_turn() {
|
||||
while let Some(event) = stream.next().await {
|
||||
event.unwrap();
|
||||
}
|
||||
if i == 0 {
|
||||
tokio::time::sleep(Duration::from_secs(2)).await;
|
||||
}
|
||||
}
|
||||
let limits = server
|
||||
.received_requests()
|
||||
@@ -810,6 +777,6 @@ async fn slow_catalog_does_not_block_generation_and_populates_next_turn() {
|
||||
.unwrap()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(limits, [32000, 2048]);
|
||||
assert_eq!(limits, [32000, 32000]);
|
||||
server.verify().await;
|
||||
}
|
||||
|
||||
@@ -4,6 +4,7 @@ mod account_eligibility;
|
||||
mod account_isolation;
|
||||
mod cache_preference;
|
||||
mod capacity;
|
||||
mod catalog_startup;
|
||||
mod cookies;
|
||||
mod correlation;
|
||||
mod error_details;
|
||||
@@ -1207,7 +1208,7 @@ fn provider_and_quota_with_affinity_and_base_url_and_leases(
|
||||
leases: Arc<TestLeaseCoordinator>,
|
||||
stream_max_retries: u32,
|
||||
) -> (Arc<CodexProvider>, Arc<CodexCredentialQuotaService>) {
|
||||
let (provider, quota, _) = provider_and_quota_with_runtime_ports(
|
||||
let (provider, quota, _, _) = provider_and_quota_with_runtime_ports(
|
||||
store,
|
||||
session_affinity,
|
||||
base_url,
|
||||
@@ -1219,6 +1220,44 @@ fn provider_and_quota_with_affinity_and_base_url_and_leases(
|
||||
(provider, quota)
|
||||
}
|
||||
|
||||
async fn provider_with_catalog_and_base_url(
|
||||
store: &Arc<MemoryAccountStore>,
|
||||
base_url: String,
|
||||
) -> Arc<CodexProvider> {
|
||||
provider_and_quota_with_catalog(
|
||||
store,
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
base_url,
|
||||
Arc::new(TestLeaseCoordinator::default()),
|
||||
u32::try_from(DEFAULT_STREAM_MAX_RETRIES).unwrap(),
|
||||
)
|
||||
.await
|
||||
.0
|
||||
}
|
||||
|
||||
async fn provider_and_quota_with_catalog(
|
||||
store: &Arc<MemoryAccountStore>,
|
||||
session_affinity: Arc<MemorySessionAffinity>,
|
||||
base_url: String,
|
||||
leases: Arc<TestLeaseCoordinator>,
|
||||
stream_max_retries: u32,
|
||||
) -> (Arc<CodexProvider>, Arc<CodexCredentialQuotaService>) {
|
||||
let (provider, quota, _, catalog) = provider_and_quota_with_runtime_ports(
|
||||
store,
|
||||
session_affinity,
|
||||
base_url,
|
||||
leases,
|
||||
stream_max_retries,
|
||||
Arc::new(MemoryCooldownPort::new()),
|
||||
crate::support::runtime_policy(),
|
||||
);
|
||||
catalog
|
||||
.refresh_catalogs()
|
||||
.await
|
||||
.expect("background catalog");
|
||||
(provider, quota)
|
||||
}
|
||||
|
||||
fn provider_and_quota_with_runtime_ports(
|
||||
store: &Arc<MemoryAccountStore>,
|
||||
session_affinity: Arc<MemorySessionAffinity>,
|
||||
@@ -1231,6 +1270,7 @@ fn provider_and_quota_with_runtime_ports(
|
||||
Arc<CodexProvider>,
|
||||
Arc<CodexCredentialQuotaService>,
|
||||
Arc<CodexWebSocketPool>,
|
||||
Arc<CodexCredentialCatalogService>,
|
||||
) {
|
||||
let profile = wire_profile();
|
||||
let http = reqwest::Client::builder()
|
||||
@@ -1269,7 +1309,7 @@ fn provider_and_quota_with_runtime_ports(
|
||||
|
||||
let provider = CodexProvider::new(
|
||||
selector,
|
||||
catalog,
|
||||
Arc::clone(&catalog),
|
||||
Arc::clone("a),
|
||||
account_feedback,
|
||||
http,
|
||||
@@ -1279,7 +1319,7 @@ fn provider_and_quota_with_runtime_ports(
|
||||
stream_max_retries,
|
||||
)
|
||||
.expect("official OpenAI provider");
|
||||
(Arc::new(provider), quota, websocket_pool)
|
||||
(Arc::new(provider), quota, websocket_pool, catalog)
|
||||
}
|
||||
|
||||
pub(super) async fn create_account(store: &Arc<MemoryAccountStore>, id: &str) {
|
||||
@@ -2832,7 +2872,7 @@ async fn image_endpoints_bypass_only_the_text_catalog_and_preserve_the_current_c
|
||||
.await;
|
||||
}
|
||||
|
||||
let provider = provider_with_base_url(&store, server.uri());
|
||||
let provider = provider_with_catalog_and_base_url(&store, server.uri()).await;
|
||||
let catalog = provider
|
||||
.query_model_capabilities()
|
||||
.await
|
||||
@@ -9138,7 +9178,7 @@ fn provider_with_cooldowns(
|
||||
base_url: String,
|
||||
cooldowns: Arc<MemoryCooldownPort>,
|
||||
) -> (Arc<CodexProvider>, Arc<CodexWebSocketPool>) {
|
||||
let (provider, _, pool) = provider_and_quota_with_runtime_ports(
|
||||
let (provider, _, pool, _) = provider_and_quota_with_runtime_ports(
|
||||
store,
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
base_url,
|
||||
@@ -10782,7 +10822,7 @@ async fn provider_compiles_catalog_presentation_for_codex_models() {
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let provider = provider_with_base_url(&store, server.uri());
|
||||
let provider = provider_with_catalog_and_base_url(&store, server.uri()).await;
|
||||
|
||||
let capabilities = provider
|
||||
.query_model_capabilities()
|
||||
@@ -10847,14 +10887,7 @@ async fn provider_compiles_catalog_presentation_for_codex_models() {
|
||||
})
|
||||
.expect("document");
|
||||
assert_eq!(document, original["models"][0]);
|
||||
assert!(
|
||||
server
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("requests")
|
||||
.iter()
|
||||
.any(|request| request.url.query() == Some("client_version=0.154.0"))
|
||||
);
|
||||
assert_eq!(server.received_requests().await.unwrap().len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -10881,7 +10914,7 @@ async fn provider_preserves_independent_catalog_context_windows() {
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let provider = provider_with_base_url(&store, server.uri());
|
||||
let provider = provider_with_catalog_and_base_url(&store, server.uri()).await;
|
||||
|
||||
let capabilities = provider
|
||||
.query_model_capabilities()
|
||||
@@ -10916,7 +10949,7 @@ async fn provider_routes_catalog_model_when_supported_in_api_is_false() {
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let provider = provider_with_base_url(&store, server.uri());
|
||||
let provider = provider_with_catalog_and_base_url(&store, server.uri()).await;
|
||||
|
||||
let capabilities = provider
|
||||
.query_model_capabilities()
|
||||
@@ -11616,7 +11649,8 @@ async fn api_key_responses_forward_lite_memgen_and_native_compaction_with_catalo
|
||||
.expect(1)
|
||||
.mount(&upstream)
|
||||
.await;
|
||||
let provider = provider(&store);
|
||||
let provider =
|
||||
provider_with_catalog_and_base_url(&store, OFFICIAL_CODEX_BASE_URL.to_owned()).await;
|
||||
provider.query_model_capabilities().await.unwrap();
|
||||
let body = json!({"model":"gpt-5.4","input":[{"type":"compaction_trigger"}],"reasoning":{"effort":"future-effort"},"tools":[{"type":"future-tool"}],"future_option":{"nested":true}});
|
||||
let operation = Operation::Generate(GenerateRequest::from_protocol_payload(
|
||||
@@ -11712,19 +11746,16 @@ async fn api_key_default_http_uses_own_prefix_plain_json_and_only_own_authentica
|
||||
provider_openai::credential::ResponsesTransport::Http,
|
||||
)
|
||||
.await;
|
||||
// 后台发现不协商客户端版本;客户端目录独立请求并按实际版本缓存
|
||||
for query in [None, Some("client_version=1.0.0")] {
|
||||
Mock::given(method("GET"))
|
||||
.and(path(format!("{prefix}/models")))
|
||||
.and(move |request: &wiremock::Request| request.url.query() == query)
|
||||
.and(header("authorization", "Bearer sk-api-test-only"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_json(json!({"data":[{"id":"gpt-5.4"}]})),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&upstream)
|
||||
.await;
|
||||
}
|
||||
Mock::given(method("GET"))
|
||||
.and(path(format!("{prefix}/models")))
|
||||
.and(move |request: &wiremock::Request| request.url.query().is_none())
|
||||
.and(header("authorization", "Bearer sk-api-test-only"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_json(json!({"data":[{"id":"gpt-5.4"}]})),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&upstream)
|
||||
.await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path(format!("{prefix}/responses")))
|
||||
.and(header("authorization", "Bearer sk-api-test-only"))
|
||||
@@ -11734,13 +11765,14 @@ async fn api_key_default_http_uses_own_prefix_plain_json_and_only_own_authentica
|
||||
.expect(1)
|
||||
.mount(&upstream)
|
||||
.await;
|
||||
let (provider, quota) = provider_and_quota_with_affinity_and_base_url_and_leases(
|
||||
let (provider, quota) = provider_and_quota_with_catalog(
|
||||
&store,
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
oauth.uri(),
|
||||
Arc::new(TestLeaseCoordinator::default()),
|
||||
DEFAULT_STREAM_MAX_RETRIES as u32,
|
||||
);
|
||||
)
|
||||
.await;
|
||||
provider
|
||||
.query_model_capabilities()
|
||||
.await
|
||||
@@ -11806,12 +11838,10 @@ async fn api_key_default_http_uses_own_prefix_plain_json_and_only_own_authentica
|
||||
.iter()
|
||||
.filter(|request| request.method == "GET")
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(model_requests.len(), 2);
|
||||
for query in [None, Some("client_version=1.0.0")] {
|
||||
let model_request = model_requests
|
||||
.iter()
|
||||
.find(|request| request.url.query() == query)
|
||||
.expect("separate background and client catalog requests");
|
||||
assert_eq!(model_requests.len(), 1);
|
||||
{
|
||||
let model_request = model_requests[0];
|
||||
assert!(model_request.url.query().is_none());
|
||||
assert_eq!(
|
||||
model_request.headers["authorization"],
|
||||
"Bearer sk-api-test-only"
|
||||
@@ -13313,7 +13343,7 @@ async fn public_catalog_filters_each_api_account_before_union_and_preserves_rout
|
||||
.mount(server)
|
||||
.await;
|
||||
}
|
||||
let provider = provider_with_base_url(&store, first.uri());
|
||||
let provider = provider_with_catalog_and_base_url(&store, first.uri()).await;
|
||||
let registry = Arc::new(ProviderRegistry::new([provider as Arc<dyn Provider>]).unwrap());
|
||||
let policy = |mode, models: &[&str]| {
|
||||
AccountModelAccess::new(mode, models.iter().map(|id| (*id).to_owned()).collect()).unwrap()
|
||||
|
||||
+42
-5
@@ -222,12 +222,14 @@ async fn account_allowlists_override_catalog_omissions_but_weight_affinity_and_p
|
||||
store.set_scheduling("acct_scope_new", None, AccountWeight::new(100).unwrap());
|
||||
let affinity = Arc::new(MemorySessionAffinity::default());
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
let provider = provider_with_affinity_and_base_url_and_leases(
|
||||
let (provider, _) = provider_and_quota_with_catalog(
|
||||
&store,
|
||||
affinity.clone(),
|
||||
native.uri(),
|
||||
leases.clone(),
|
||||
);
|
||||
u32::try_from(DEFAULT_STREAM_MAX_RETRIES).unwrap(),
|
||||
)
|
||||
.await;
|
||||
provider.query_model_capabilities().await.unwrap();
|
||||
for (model, expected) in [
|
||||
("gpt-6-luna", "acct_scope_new"),
|
||||
@@ -366,7 +368,7 @@ async fn only_a_current_account_catalog_can_exclude_a_model() {
|
||||
use gateway_core::account::{CredentialCasOutcome, CredentialCasUpdate, ProviderAccountUpdate};
|
||||
use provider_openai::credential::ResponsesTransport;
|
||||
|
||||
for scenario in ["missing", "unknown_plan", "rotated"] {
|
||||
for scenario in ["missing", "unknown_plan", "rotated", "expired"] {
|
||||
let native = MockServer::start().await;
|
||||
let api = MockServer::start().await;
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
@@ -396,7 +398,7 @@ async fn only_a_current_account_catalog_can_exclude_a_model() {
|
||||
store
|
||||
.seed_api_key("acct_scope_new", api.uri(), ResponsesTransport::Http)
|
||||
.await;
|
||||
let response = if scenario == "missing" {
|
||||
let response = if matches!(scenario, "missing" | "expired") {
|
||||
ResponseTemplate::new(503)
|
||||
} else {
|
||||
ResponseTemplate::new(200).set_body_json(json!({"data": [{"id": "gpt-6-luna"}]}))
|
||||
@@ -421,7 +423,42 @@ async fn only_a_current_account_catalog_can_exclude_a_model() {
|
||||
.expect(if scenario == "unknown_plan" { 0 } else { 1 })
|
||||
.mount(upstream)
|
||||
.await;
|
||||
let provider = provider_with_base_url(&store, native.uri());
|
||||
let provider = if scenario == "expired" {
|
||||
use gateway_core::account::AccountCatalogObservation;
|
||||
let a = store.account("acct_scope_new").unwrap();
|
||||
store
|
||||
.save_account_catalog(
|
||||
&AccountCatalogObservation {
|
||||
account_id: a.id().clone(),
|
||||
revision: a.revision(),
|
||||
plan: a.plan_type().map(str::to_owned),
|
||||
etag: None,
|
||||
observed_at: SystemTime::now() - Duration::from_secs(3600),
|
||||
},
|
||||
b"\x01{\"data\":[{\"id\":\"gpt-6-luna\"}]}",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let (provider, _, _, catalog) = provider_and_quota_with_runtime_ports(
|
||||
&store,
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
native.uri(),
|
||||
Arc::new(TestLeaseCoordinator::default()),
|
||||
u32::try_from(DEFAULT_STREAM_MAX_RETRIES).unwrap(),
|
||||
Arc::new(MemoryCooldownPort::new()),
|
||||
crate::support::runtime_policy(),
|
||||
);
|
||||
catalog.restore().await.unwrap();
|
||||
// B 的失败不能清空上次成功结果,也不能继续以旧的缺失条目拒绝请求
|
||||
catalog.refresh_catalogs().await.unwrap();
|
||||
assert_eq!(
|
||||
catalog.cached_account_models(&a).unwrap().unwrap(),
|
||||
["gpt-6-luna"]
|
||||
);
|
||||
provider
|
||||
} else {
|
||||
provider_with_catalog_and_base_url(&store, native.uri()).await
|
||||
};
|
||||
provider.query_model_capabilities().await.unwrap();
|
||||
if scenario == "rotated" {
|
||||
let account = store.account("acct_scope_new").unwrap();
|
||||
|
||||
@@ -53,6 +53,15 @@ type CredentialLoadHook = dyn for<'a> Fn(
|
||||
#[derive(Default)]
|
||||
pub(crate) struct MemoryAccountStore {
|
||||
accounts: Mutex<BTreeMap<ProviderAccountId, StoredAccount>>,
|
||||
catalogs: Mutex<
|
||||
BTreeMap<
|
||||
ProviderAccountId,
|
||||
(
|
||||
gateway_core::account::AccountCatalogObservation,
|
||||
bytes::Bytes,
|
||||
),
|
||||
>,
|
||||
>,
|
||||
account_reads: AtomicUsize,
|
||||
provider_list_reads: AtomicUsize,
|
||||
quota_reads: AtomicUsize,
|
||||
@@ -296,6 +305,62 @@ impl MemoryAccountStore {
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderAccountStore for MemoryAccountStore {
|
||||
async fn load_account_catalogs(
|
||||
&self,
|
||||
provider: &ProviderKind,
|
||||
) -> Result<Vec<gateway_core::account::StoredAccountCatalog>, StoreError> {
|
||||
let accounts = self.accounts.lock().unwrap();
|
||||
let mut groups = BTreeMap::<bytes::Bytes, Vec<_>>::new();
|
||||
for (id, (observation, document)) in self.catalogs.lock().unwrap().iter() {
|
||||
if accounts.get(id).is_some_and(|row| {
|
||||
row.account.provider() == provider
|
||||
&& row.account.enabled()
|
||||
&& matches!(
|
||||
row.account.credential_state(),
|
||||
CredentialState::Ready | CredentialState::Unknown
|
||||
)
|
||||
&& row
|
||||
.account
|
||||
.access_token_expires_at()
|
||||
.is_none_or(|ts| ts > SystemTime::now())
|
||||
&& row.account.revision() == observation.revision
|
||||
&& row.account.plan_type() == observation.plan.as_deref()
|
||||
}) {
|
||||
groups
|
||||
.entry(document.clone())
|
||||
.or_default()
|
||||
.push(observation.clone());
|
||||
}
|
||||
}
|
||||
Ok(groups
|
||||
.into_iter()
|
||||
.map(
|
||||
|(document, accounts)| gateway_core::account::StoredAccountCatalog {
|
||||
accounts,
|
||||
document,
|
||||
},
|
||||
)
|
||||
.collect())
|
||||
}
|
||||
async fn save_account_catalog(
|
||||
&self,
|
||||
observation: &gateway_core::account::AccountCatalogObservation,
|
||||
document: &[u8],
|
||||
) -> Result<bool, StoreError> {
|
||||
let rows = self.accounts.lock().unwrap();
|
||||
if !rows.get(&observation.account_id).is_some_and(|a| {
|
||||
a.account.revision() == observation.revision
|
||||
&& a.account.plan_type() == observation.plan.as_deref()
|
||||
}) {
|
||||
return Ok(false);
|
||||
}
|
||||
self.catalogs.lock().unwrap().insert(
|
||||
observation.account_id.clone(),
|
||||
(observation.clone(), bytes::Bytes::copy_from_slice(document)),
|
||||
);
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn create_account(&self, input: NewProviderAccount) -> Result<(), StoreError> {
|
||||
let mut accounts = self.accounts.lock().expect("account store lock");
|
||||
if accounts.contains_key(input.account.id()) {
|
||||
@@ -383,13 +448,13 @@ impl ProviderAccountStore for MemoryAccountStore {
|
||||
));
|
||||
}
|
||||
// 与 Postgres 实现的调度列表语义一致:停用账号不进入常规候选
|
||||
Ok(self
|
||||
.accounts
|
||||
.lock()
|
||||
.expect("account store lock")
|
||||
.values()
|
||||
let accounts = self.accounts.lock().unwrap();
|
||||
Ok(account_ids
|
||||
.iter()
|
||||
.collect::<std::collections::BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.filter_map(|id| accounts.get(id))
|
||||
.filter(|stored| stored.account.provider() == provider && stored.account.enabled())
|
||||
.filter(|stored| account_ids.contains(stored.account.id()))
|
||||
.map(|stored| stored.account.clone())
|
||||
.collect())
|
||||
}
|
||||
|
||||
@@ -308,6 +308,21 @@ impl MemoryProviderAccountStore {
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderAccountStore for MemoryProviderAccountStore {
|
||||
async fn load_account_catalogs(
|
||||
&self,
|
||||
_: &gateway_core::routing::ProviderKind,
|
||||
) -> Result<Vec<gateway_core::account::StoredAccountCatalog>, gateway_core::error::StoreError>
|
||||
{
|
||||
Ok(Vec::new())
|
||||
}
|
||||
async fn save_account_catalog(
|
||||
&self,
|
||||
_: &gateway_core::account::AccountCatalogObservation,
|
||||
_: &[u8],
|
||||
) -> Result<bool, gateway_core::error::StoreError> {
|
||||
unreachable!("fixture does not discover account catalogs")
|
||||
}
|
||||
|
||||
async fn create_account(&self, account: NewProviderAccount) -> Result<(), StoreError> {
|
||||
let mut accounts = lock(&self.accounts);
|
||||
if accounts.contains_key(account.account.id()) {
|
||||
|
||||
@@ -38,3 +38,4 @@ b215ea90d5c8753aeb8192b17ffe65f7a95c1d620aaeaf9fc4337d8cf6f2634e 0037_retry_aft
|
||||
321f4a09367717eb13c3cf5dbfd1f4b014e31c615a74bf6e3b09b8e810d87198 0038_preserve_client_key_ledger.sql
|
||||
a6b4e81d38e5d70776a75eba22749354573ebcb186160ff9eb45fd91157b3e0d 0039_queue_timeouts_quota_credits_and_installation_seed.sql
|
||||
cd9e31400fc6b0d49a6515048c104341fca2cf44f290d89c401b41f94ead84a3 0040_account_scheduling_key_budgets_and_peer_routing.sql
|
||||
aeff99fee7cef5a553ae49ffe27a48f49e31ccfb3fbc61a2482c9b109478202d 0041_account_model_catalogs.sql
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
-- 每个账号只保留最后一次成功观测,凭据轮换后旧目录不能成为新凭据的权限证据
|
||||
create table provider_account_model_catalogs (
|
||||
account_id text primary key references provider_accounts(id) on delete cascade,
|
||||
credential_revision bigint not null check (credential_revision > 0),
|
||||
plan_type text,
|
||||
etag text,
|
||||
observed_at timestamptz not null,
|
||||
document bytea not null,
|
||||
document_hash bytea generated always as (sha256(document)) stored
|
||||
);
|
||||
+5
-5
@@ -386,7 +386,7 @@ Responses 的 `allowed_tools` 只向 Messages 上游声明本轮允许的工具
|
||||
跨协议推理设置读取顶层 `thinking`、`output_config.effort` 或 `reasoning.effort`。
|
||||
转向 Messages 时按所选账号模型目录声明的能力选择 adaptive 或手动预算,并约束输出上限和 effort 等级。
|
||||
不支持关闭的 adaptive 模型将 `none` 降到最低可用 effort;手动预算至少为 1024、严格小于输出上限,不足时关闭推理。
|
||||
生成请求最多等待目录 500 ms;未及时取得目录或未声明能力时沿用 CPA 手动预算,`auto` 使用 8192,不根据模型别名猜测能力。
|
||||
生成请求只读取账号已发现的能力;目录缺失或未声明能力时沿用 CPA 手动预算,`auto` 使用 8192,不根据模型别名猜测能力。
|
||||
强制工具与手动推理冲突时显式关闭推理;紧接着只回传工具结果且未指定 effort 时,沿用无思考模式。
|
||||
adaptive 路径将普通强制选择降为 `auto`,保留串行约束。
|
||||
仅设置 `parallel_tool_calls:false` 也会生成串行工具选择,`none` 不携带并行选项。
|
||||
@@ -679,7 +679,7 @@ HTML 或截断正文当作 message。
|
||||
未携带上述协议头时,`GET /v1/models` 默认返回 OpenAI 兼容列表 `{"object": "list", "data": [...]}`;请求携带非空
|
||||
`client_version` query 参数(Codex 客户端)时改为返回 Codex 专用目录合同 `{"models": [...]}`
|
||||
|
||||
OpenAI Provider 按客户端传入的 `client_version` 请求上游目录,完整保留每个模型 JSON 对象,包括
|
||||
OpenAI Provider 使用后台按账号发现的目录,客户端版本只选择输出格式。目录完整保留每个原生模型 JSON 对象,包括
|
||||
`base_instructions`、`model_messages`、`service_tiers`、工具与能力字段,以及未知嵌套字段、显式 `null`
|
||||
和字段缺失的区别。
|
||||
API Key 上游返回完整 Codex `models` 目录时沿用该合同;仅返回普通 `data` 模型列表时使用通用画像,
|
||||
@@ -699,12 +699,12 @@ OpenAI 聚合当前 Key 范围内合格 OAuth 和 API Key 账号的目录。每
|
||||
同名模型优先采用 OAuth 原生对象,其次采用 API Key 的完整 Codex 对象,再使用普通模型 ID。
|
||||
同类来源按账号 ID 排序,保留各来源的上游条目顺序,不跨账号拼接字段
|
||||
|
||||
部分账号失败或超时时,返回已成功读取的目录并集;没有成功读取任何目录时返回 `503 model_catalog_unavailable`。
|
||||
后台刷新失败或超时保留该账号最后成功的结果。当前范围没有任何可用的已知目录时返回 `503 model_catalog_unavailable`。
|
||||
成功读取但全部条目被账号政策过滤时返回空列表。
|
||||
目录保留完整的原生 Codex 模型对象。对于 OpenAI API Key 普通模型列表和 xAI 目录生成的通用画像,
|
||||
当前 Key 可访问的适配器若声明了同名模型,其画像在原位置替换通用画像。仅由插件提供的可见模型按声明顺序追加到所属 Provider 的条目之后
|
||||
|
||||
原生目录可能返回缓存结果,成功缓存有效期为 5 分钟。每次查询都检查账号资格和 Client Key 范围。
|
||||
服务启动时批量恢复账号目录;目录读取只检查本地账号资格和 Client Key 范围,不等待上游查询。
|
||||
成功目录响应带 `Cache-Control: private, no-store`,不借用上游 ETag 标识经过选择/别名/聚合后的正文
|
||||
|
||||
`service_tiers` 的 `name` 用于生成 `/fast` 等命令,`id` 是请求使用的 `service_tier` 值,二者不能互换。
|
||||
@@ -2414,7 +2414,7 @@ models.dev 同步只导入可表示为当前文本 Token 计价的 OpenAI/xAI
|
||||
|
||||
`openaiClientProfile` 保存通用选择,首次默认 `MacOS · Desktop · 自动最新`。
|
||||
兼容字段的更新语义见上节。初始化不读取 YAML 身份字段。
|
||||
该配置作用于 Client Key 的 OpenAI 模型请求与原生模型目录,适用于 HTTP/SSE、WebSocket、Images 和 Search。
|
||||
该配置作用于 Client Key 的 OpenAI 模型请求,适用于 HTTP/SSE、WebSocket、Images 和 Search。
|
||||
不改变 xAI、入站客户端版本门禁、账号认证或后台 Desktop 专属操作
|
||||
|
||||
自定义配置使用 `{ "mode": "custom", "userAgent": "完整 UA" }`。
|
||||
|
||||
@@ -329,13 +329,14 @@ Provider 模型能力、目录代次与 `ProviderCatalogPort` 由 `routing::cata
|
||||
Provider 可为模型提供来源账号集合;Core 在公开目录查询时结合冻结账号范围与政策过滤来源,
|
||||
配置发布复用已缓存的来源事实。来源集合只约束目录展示,不直接作为推理账号白名单;Provider 独占上游资格判断
|
||||
|
||||
OpenAI 在 PostgreSQL 保存账号最近一次成功的模型目录、凭据版本、观测时间及 ETag。启动时批量恢复,按不同正文解析一次,保留每个账号的独立归属。没有历史目录也能发布路由快照;能力读取与推理不等待上游查询。
|
||||
后台只刷新目录缺失、已满 30 分钟或收到新 ETag 的账号;凭据轮换后按新版本重新发现。每轮最多并发 32 个,新请求至少间隔 10 ms。失败保留最后成功的结果,过期资料只用于展示,不能据此排除候选账号
|
||||
|
||||
客户端原生模型目录同样由 Provider 拥有,通过 `ExecutionService` 和现有 Registry 的只读调用传递。
|
||||
`routing::catalog` 的原生条目只包含模型 ID 与不透明协议正文,Core 负责冻结账号范围和整对象模型映射,
|
||||
不解析上游字段;API 负责协议输出,不直接依赖具体 Provider。原生目录失败不会退回通用画像,只有不提供
|
||||
该协议原生目录的 Provider(如 xAI 的 Codex 适配)才走画像转换。目录读取不创建计量请求或推理 attempt。
|
||||
OpenAI 的账号/凭据 revision/客户端版本隔离、并发合并、TTL、容量与失效均归属 credential catalog
|
||||
service;客户端原生对象保存在有界进程缓存中,与账号级 Redis 模型 ID 缓存分开管理。
|
||||
该缓存最多保留 32 项,原生正文总预算为 32 MiB;单次结果超过预算时仍正常交付,但不驻留缓存
|
||||
OpenAI 从账号目录生成原生列表;先按冻结政策选择来源,再读取本地账号状态核对凭据版本。同一份正文可以复用,但不会把一个账号的模型授予同套餐的其他账号。客户端版本和 Key 身份不触发额外发现
|
||||
|
||||
全局请求位置及开关由运行设置持久化;快照仅在开关开启时生成全局覆盖,并随
|
||||
`RuntimeSnapshot → RoutingPlan → AttemptContext` 冻结传递,关闭时保留客户端原有字段;
|
||||
@@ -517,7 +518,7 @@ OpenAI Provider 统一处理 Responses 的账号归属与 turn state 隔离,HT
|
||||
turn state 不参与选号,也不授权其他账号状态或 response ID 的复用
|
||||
|
||||
OpenAI 的 OAuth 与 API Key 共用现有账号和事务。API Key 的 Base URL、密钥、上游协议和传输策略属于 Provider 凭据 JSON,
|
||||
随 credential revision 更新;普通详情只投影非敏感连接设置。API Key 客户端目录按账号、凭据版本和客户端版本隔离,
|
||||
随 credential revision 更新;普通详情只投影非敏感连接设置。API Key 目录按账号和凭据版本隔离,
|
||||
标准 `data` 模型列表通过通用画像输出,完整 Codex `models` 目录保留原生对象。通用账号层按 Provider 提交的
|
||||
credential state 调度,不以是否存在上游用户 ID 推断可用性;OAuth 未完成身份投影时由 Provider 保持 `unknown`。
|
||||
状态恢复和未补齐身份的凭据轮换保留 `unknown`。
|
||||
@@ -532,7 +533,7 @@ OpenAI 模型目录用于发现与原生 Responses 账号资格筛选。
|
||||
同协议业务响应使用原始 wire,跨协议才调用转换器;客户端错误兼容与原生续写额度恢复按明确例外处理,见 [OpenAI 透传](api.md#openai-透传)。
|
||||
`transport/correlation` 在出站副本上按本次安装身份统一映射关联标识。
|
||||
Responses 的账号归属判断及换号头部清理由 Provider 统一执行,各传输消费同一出站请求。
|
||||
原生安装身份由 OpenAI Provider 按 Key 种子、Key ID 与选中账号 ID 确定性派生,无需额外映射表;未配置种子时取自账号,与 UA 继承方式无关。模型目录缓存按生效身份隔离,Live 后续阶段使用引导时冻结的种子;插件的身份生成、持久化及关联字段改写通过既有中间件链路自行完成
|
||||
原生安装身份由 OpenAI Provider 按 Key 种子、Key ID 与选中账号 ID 确定性派生,无需额外映射表;未配置种子时取自账号,与 UA 继承方式无关。Live 后续阶段使用引导时冻结的种子;插件的身份生成、持久化及关联字段改写通过既有中间件链路自行完成
|
||||
内部选号、亲和、连接池与本地诊断继续读取原始身份;具体字段边界见 [API 协议](api.md#请求头与身份)
|
||||
canonical facts 从同一数据旁路提取,用于路由、恢复判断、观测和计费。
|
||||
WebSocket transport 按完整文本消息交付,wire 同时保留原文与可选解析视图;同协议出口使用原文,
|
||||
@@ -600,8 +601,7 @@ OpenAI / xAI 导入用例。并发槽位由单个 Worker 统一管理,条目
|
||||
连接池按出口隔离。修改或启停代理推进配置 revision,使后续请求使用新出口;已执行请求可沿原连接完成。
|
||||
代理变更不推进 `credential_revision`,不会使进行中的令牌刷新因凭据版本冲突而丢失结果。
|
||||
xAI 的推理连接池继续按账号绑定隔离;OAuth、目录/计费辅助请求在各自 transport 内复用有界出口
|
||||
client,OIDC 的 JWKS 缓存与单飞归属对应出口状态。自动刷新提交凭据后的目录预热重新读取当前账号
|
||||
出口,不从 token-only 路径绕过代理;JWKS 过期或获取失败仍不使用 stale fallback
|
||||
client,OIDC 的 JWKS 缓存与单飞归属对应出口状态。凭据变更后由目录任务按账号当前出口重新发现;JWKS 过期或获取失败仍不使用 stale fallback
|
||||
|
||||
代理独立保存和测试,通过 `outboundProxyId` 绑定账号;账号保留解析后的 URL 供 Provider 使用。
|
||||
出口测试结果属于诊断信息,不作为选择、授权或导入绑定的准入条件。
|
||||
@@ -763,7 +763,7 @@ OpenAI Provider 的账号资格模块统一解释请求类型、认证与套餐
|
||||
|
||||
原生 Responses 在评分、亲和和租约前检查冻结的账号模型政策。OpenAI 账号白名单精确命中目标 ID 时允许上游验证,
|
||||
其余 OpenAI 账号及 xAI 账号按有效目录排除未列出的模型。OpenAI 目录按来源账号、凭据版本和套餐隔离;同套餐账号可以有不同的模型。
|
||||
OpenAI 使用未满 30 分钟的进程目录快照,xAI 按套餐去重读取 TTL 目录缓存,选号不请求上游目录。
|
||||
OpenAI 只使用未满 30 分钟且尚未收到变化通知的账号目录排除候选。大池可在读取账号时按目录摘要筛选,数据库逐账号核对版本、套餐和观测时间;白名单、待刷新通知及受管适配器不参与该筛选。xAI 按套餐去重读取 TTL 目录缓存,选号不请求上游目录。
|
||||
没有有效目录时保留上游验证,不从其他账号借用 OpenAI 模型资格,也不从其他套餐借用 xAI 模型资格
|
||||
|
||||
原生 OpenAI Responses 按[官方 Codex 0.166.0](https://github.com/openai/codex/blob/6ff1e1ca3b30fc8b914720681fcf5fe62e95f856/codex-rs/codex-api/src/error.rs#L73)的结构化 `model_not_found` 错误码与状态合同识别模型拒绝,
|
||||
@@ -1006,7 +1006,7 @@ PostgreSQL 周期对账才是正确性基础
|
||||
| 日志、OAuth 恢复记录、在线更新状态、备份暂存 | `.runtime/` | 部署节点本地运行文件 |
|
||||
| 重置卡库存与消费结果 | 对应 Provider 的上游 | 本地到期计划不能证明卡片仍可用;手动查询直连上游,自动消费前重新核对 |
|
||||
| 重置卡自动使用配置、发现时间与未决幂等键 | PostgreSQL 账号记录 | Admin 拥有计划与重试编排;账号删除时一并清理,重启沿用已保存时间及幂等键 |
|
||||
| Provider 公开模型与官方发布资料 | Provider/runtime cache | 由官方目录或发布源刷新,与 PostgreSQL 中的用户身份选择分别管理 |
|
||||
| Provider 公开模型与官方发布资料 | Provider 的账号目录表或运行时缓存 | 由官方目录或发布源刷新,与 PostgreSQL 中的用户身份选择分别管理 |
|
||||
| Windows 安装包临时直链 | Host 进程内短缓存 | 按需解析、严格校验、到期前丢弃;不写 PostgreSQL/Redis,也不代理包字节 |
|
||||
|
||||
账号对外状态不是独立列,而是 PostgreSQL credential/quota 事实与 Redis cooldown 的统一投影:
|
||||
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
# v3.45.0
|
||||
# v3.45.1
|
||||
|
||||
## 新增功能
|
||||
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
version: 3.45.0
|
||||
version: 3.45.1
|
||||
published:
|
||||
'3':
|
||||
stable: 3.44.0
|
||||
rc: 3.44.0
|
||||
beta: 3.44.0
|
||||
alpha: 3.44.0
|
||||
stable: 3.45.0
|
||||
rc: 3.45.0
|
||||
beta: 3.45.0
|
||||
alpha: 3.45.0
|
||||
|
||||
Reference in New Issue
Block a user