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

refactor(gateway): remove provider instance layer

Browse Source 
Collapse routing, health, circuit state, account pools, and admin APIs onto the fixed OpenAI and xAI Provider boundary while keeping credential runtime state account-scoped.

Remove configurable provider-instance and routing-catalog persistence, use fixed official endpoints in provider adapters, align account routes and frontend contracts, and mirror provider tests under their source modules.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-25 00:33:12 +08:00
co-authored by Codex
parent acaf98512c
commit 70f8afb122
134 changed files with 2828 additions and 10770 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-admin/src/lib.rs
+4 -15
View File
@@ -16,9 +16,9 @@ pub mod ports;
mod use_case;
pub use use_case::{
accounts::AccountsService, auth::AuthService, catalog::CatalogService,
client_keys::ClientKeyService, observability::ObservabilityService, openai::OpenAiService,
settings::SettingsService, system::SystemService, xai::XaiService,
accounts::AccountsService, auth::AuthService, client_keys::ClientKeyService,
observability::ObservabilityService, openai::OpenAiService, settings::SettingsService,
system::SystemService, xai::XaiService,
};
use model::{AdminError, AdminErrorKind};
@@ -28,7 +28,7 @@ use ports::{
system::SystemOperations,
};
use use_case::{
accounts::DefaultAccountsService, auth::DefaultAuthService, catalog::DefaultCatalogService,
accounts::DefaultAccountsService, auth::DefaultAuthService,
client_keys::DefaultClientKeyService, observability::DefaultObservabilityService,
openai::DefaultOpenAiService, settings::DefaultSettingsService, system::DefaultSystemService,
xai::DefaultXaiService,
@@ -140,7 +140,6 @@ pub enum AdminConfigError {
pub struct AdminServices {
auth: Arc<dyn AuthService>,
accounts: Arc<dyn AccountsService>,
catalog: Arc<dyn CatalogService>,
client_keys: Arc<dyn ClientKeyService>,
observability: Arc<dyn ObservabilityService>,
settings: Arc<dyn SettingsService>,
@@ -160,11 +159,6 @@ impl AdminServices {
self.accounts.as_ref()
}
#[must_use]
pub fn catalog(&self) -> &dyn CatalogService {
self.catalog.as_ref()
}
#[must_use]
pub fn client_keys(&self) -> &dyn ClientKeyService {
self.client_keys.as_ref()
@@ -242,7 +236,6 @@ pub async fn initialize(
let accounts = Arc::new(DefaultAccountsService::new(
store.accounts(),
store.catalog(),
store.settings(),
registry.clone(),
snapshot.clone(),
@@ -251,10 +244,6 @@ pub async fn initialize(
let services = AdminServices {
auth,
accounts,
catalog: Arc::new(DefaultCatalogService::new(
store.catalog(),
snapshot.clone(),
)),
client_keys: Arc::new(DefaultClientKeyService::new(
store.client_keys(),
snapshot.clone(),
backend/crates/gateway-admin/src/model/accounts.rs
+4 -5
View File
@@ -5,7 +5,7 @@ use std::pin::Pin;
use chrono::{DateTime, Utc};
use futures::Stream;
use gateway_core::routing::{ProviderInstanceId, ProviderKind};
use gateway_core::routing::ProviderKind;
use super::{PageSize, Revision};
@@ -72,7 +72,6 @@ pub struct AccountListQuery {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountRecord {
pub id: String,
pub provider_instance_id: ProviderInstanceId,
pub provider_kind: ProviderKind,
pub name: String,
pub email: Option<String>,
@@ -169,11 +168,11 @@ pub struct SetAccountEnabled {
pub enabled: bool,
}
/// 账号删除命令。
/// 账号批量删除命令。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DeleteAccount {
pub struct DeleteAccounts {
pub expected_config_revision: Revision,
pub account_id: String,
pub account_ids: Vec<String>,
}
/// 账号连接测试的语义事件。
backend/crates/gateway-admin/src/model/catalog.rs
-89
View File
@@ -1,89 +0,0 @@
//! Provider instance 目录的 Command 与 Result。
use chrono::{DateTime, Utc};
use gateway_core::routing::{ProviderInstanceId, ProviderKind};
use super::{PageSize, Revision};
/// 一个可管理的 Provider instance。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderInstance {
pub id: ProviderInstanceId,
pub provider_kind: ProviderKind,
pub name: String,
pub base_url: String,
pub enabled: bool,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
/// Provider instance 列表结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderInstanceCatalog {
pub config_revision: Revision,
pub items: Vec<ProviderInstance>,
}
/// Provider instance 详情及其所属控制面 revision。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderInstanceDetail {
pub config_revision: Revision,
pub item: ProviderInstance,
}
/// Provider instance 的稳定游标查询。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CatalogListQuery {
pub cursor: Option<ProviderInstanceId>,
pub page_size: PageSize,
}
/// Provider instance 分页结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderInstancePage {
pub config_revision: Revision,
pub items: Vec<ProviderInstance>,
pub next_cursor: Option<ProviderInstanceId>,
}
/// 创建 Provider instance。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CreateProviderInstance {
pub expected_config_revision: Revision,
pub id: ProviderInstanceId,
pub provider_kind: ProviderKind,
pub name: String,
pub base_url: String,
}
/// 修改 Provider instance 的公共字段。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpdateProviderInstance {
pub expected_config_revision: Revision,
pub id: ProviderInstanceId,
pub name: String,
pub base_url: String,
}
/// 修改 Provider instance 的启用状态。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SetProviderInstanceEnabled {
pub expected_config_revision: Revision,
pub id: ProviderInstanceId,
pub enabled: bool,
}
/// 删除 Provider instance。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DeleteProviderInstance {
pub expected_config_revision: Revision,
pub id: ProviderInstanceId,
}
/// Provider instance 写入后的状态。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderInstanceMutation {
pub config_revision: Revision,
pub instance: Option<ProviderInstance>,
}
backend/crates/gateway-admin/src/model/mod.rs
-1
View File
@@ -4,7 +4,6 @@ use std::num::{NonZeroU16, NonZeroU64};
pub mod accounts;
pub mod auth;
pub mod catalog;
pub mod client_keys;
pub mod observability;
pub mod provider_credentials;
backend/crates/gateway-admin/src/model/observability.rs
-4
View File
@@ -457,7 +457,6 @@ pub struct DashboardAccountModelUsage {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DashboardAccountUsage {
pub account_id: String,
pub provider_instance_id: String,
pub provider_kind: String,
pub name: String,
pub email: Option<String>,
@@ -506,7 +505,6 @@ pub struct UsageRecord {
pub client_transport: String,
pub requested_model_id: String,
pub input_token_estimate: u64,
pub provider_instance_id: Option<String>,
pub provider_kind: Option<String>,
pub provider_account_ref: Option<String>,
pub provider_account_name: Option<String>,
@@ -579,7 +577,6 @@ pub struct UsageAttempt {
pub attempt_index: u32,
pub component: String,
pub operation: String,
pub provider_instance_id: Option<String>,
pub provider_kind: Option<String>,
pub provider_account_ref: Option<String>,
pub upstream_model_id: Option<String>,
@@ -663,7 +660,6 @@ pub struct OpsError {
pub client_api_key_ref: Option<String>,
pub component: String,
pub operation: String,
pub provider_instance_id: Option<String>,
pub provider_kind: Option<String>,
pub provider_account_ref: Option<String>,
pub upstream_model_id: Option<String>,
backend/crates/gateway-admin/src/model/provider_credentials.rs
+17 -14
View File
@@ -5,7 +5,7 @@ use std::fmt;
use chrono::{DateTime, Utc};
use gateway_core::{
engine::credential::{OpaqueProviderData, ProviderAccountId},
routing::{ProviderInstanceId, ProviderKind, UpstreamModelId},
routing::{ProviderKind, UpstreamModelId},
};
use super::{
@@ -51,7 +51,6 @@ pub struct CredentialCursor {
/// Provider credential 列表查询。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CredentialListQuery {
pub provider_instance_id: Option<ProviderInstanceId>,
pub availability: Option<CredentialAvailabilityFilter>,
pub enabled: Option<bool>,
pub window: CredentialListWindow,
@@ -103,7 +102,6 @@ pub struct CredentialDetails {
pub struct ImportCredentials {
pub context: MutationContext,
pub expected_config_revision: Revision,
pub provider_instance_id: ProviderInstanceId,
pub document: ProviderDocument,
}
@@ -113,7 +111,6 @@ impl fmt::Debug for ImportCredentials {
.debug_struct("ImportCredentials")
.field("context", &self.context)
.field("expected_config_revision", &self.expected_config_revision)
.field("provider_instance_id", &self.provider_instance_id)
.field("document", &self.document)
.finish()
}
@@ -126,9 +123,8 @@ pub struct CredentialImportResult {
pub credential_ids: Vec<ProviderAccountId>,
}
/// Provider 解析导入文档时只接收实例与不透明文档，不接触 revision 或审计上下文。
/// Provider 解析导入文档时只接收不透明文档，不接触 revision 或审计上下文。
pub struct PrepareCredentialImport {
pub provider_instance_id: ProviderInstanceId,
pub document: ProviderDocument,
}
@@ -136,7 +132,6 @@ impl fmt::Debug for PrepareCredentialImport {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("PrepareCredentialImport")
.field("provider_instance_id", &self.provider_instance_id)
.field("document", &self.document)
.finish()
}
@@ -146,7 +141,6 @@ impl fmt::Debug for PrepareCredentialImport {
#[derive(Debug, Clone, PartialEq)]
pub struct PreparedCredentialCreate {
pub account_id: ProviderAccountId,
pub provider_instance_id: ProviderInstanceId,
pub provider_kind: ProviderKind,
pub name: String,
pub email: Option<String>,
@@ -168,7 +162,6 @@ pub struct PreparedCredentialCreate {
#[derive(Debug, Clone, PartialEq)]
pub struct PreparedCredentialImport {
pub provider_kind: ProviderKind,
pub provider_instance_id: ProviderInstanceId,
pub credentials: Vec<PreparedCredentialCreate>,
}
@@ -249,11 +242,9 @@ impl fmt::Debug for AuthorizationOwnerBinding {
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AuthorizationMutationTarget {
Create {
provider_instance_id: ProviderInstanceId,
name: String,
},
Reauthorize {
provider_instance_id: ProviderInstanceId,
account_id: ProviderAccountId,
expected_credential_revision: Revision,
},
@@ -310,7 +301,6 @@ impl PendingAuthorizationMutation {
pub struct StartAuthorization {
pub context: MutationContext,
pub expected_config_revision: Revision,
pub provider_instance_id: ProviderInstanceId,
pub name: String,
pub reauthorization: Option<ReauthorizationTarget>,
}
@@ -427,6 +417,21 @@ pub struct CredentialMutationResult {
pub credential_revision: Option<Revision>,
}
/// 同一 Provider 管理范围内的 credential 批量删除。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CredentialDeletion {
pub context: MutationContext,
pub expected_config_revision: Revision,
pub account_ids: Vec<ProviderAccountId>,
}
/// Credential 批量删除提交结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CredentialDeletionResult {
pub config_revision: Revision,
pub account_ids: Vec<ProviderAccountId>,
}
/// Provider-owned token 轮换命令。
pub struct RotateCredential {
pub mutation: CredentialMutation,
@@ -459,7 +464,6 @@ impl fmt::Debug for PrepareCredentialRotation {
#[derive(Debug, Clone, PartialEq)]
pub struct PreparedCredentialRotationFacts {
pub account_id: ProviderAccountId,
pub provider_instance_id: ProviderInstanceId,
pub provider_kind: ProviderKind,
pub expected_credential_revision: Revision,
pub name: String,
@@ -613,7 +617,6 @@ impl fmt::Debug for ProviderExport {
#[derive(Debug, Clone, PartialEq)]
pub struct AccountDirectoryItem {
pub account: AccountRecord,
pub provider_instance_name: String,
pub status: AccountStatus,
pub usage: Option<super::accounts::AccountUsage>,
pub quota: ProviderQuota,
backend/crates/gateway-admin/src/ports/store.rs
+3 -54
View File
@@ -9,15 +9,10 @@ use async_trait::async_trait;
use crate::model::{
MutationContext, Revision,
accounts::{
AccountListQuery, AccountPage, AccountRecord, AccountUsage, DeleteAccount,
AccountListQuery, AccountPage, AccountRecord, AccountUsage, DeleteAccounts,
SetAccountEnabled,
},
auth::{AdminAuditEvent, AdminSession},
catalog::{
CreateProviderInstance, DeleteProviderInstance, ProviderInstanceCatalog,
ProviderInstanceDetail, ProviderInstanceMutation, SetProviderInstanceEnabled,
UpdateProviderInstance,
},
client_keys::{
ClientKeyListQuery, ClientKeyPage, ClientKeyRecord, ClientKeySecret, DeleteClientKey,
NewClientKey, SetClientKeyEnabled, UpdateClientKey,
@@ -142,9 +137,9 @@ pub trait AccountStore: Send + Sync {
context: &MutationContext,
) -> AdminStoreResult<Revision>;
async fn delete_account(
async fn delete_accounts(
&self,
command: DeleteAccount,
command: DeleteAccounts,
context: &MutationContext,
) -> AdminStoreResult<Revision>;
@@ -194,44 +189,6 @@ pub trait AuthStore: Send + Sync {
async fn append_audit_event(&self, event: AdminAuditEvent) -> AdminStoreResult<()>;
}
/// Provider instance 目录与 revision CAS 写入。
#[async_trait]
pub trait CatalogStore: Send + Sync {
async fn list_provider_instances(
&self,
include_disabled: bool,
) -> AdminStoreResult<ProviderInstanceCatalog>;
async fn load_provider_instance(
&self,
id: &gateway_core::routing::ProviderInstanceId,
) -> AdminStoreResult<Option<ProviderInstanceDetail>>;
async fn create_provider_instance(
&self,
command: CreateProviderInstance,
context: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation>;
async fn update_provider_instance(
&self,
command: UpdateProviderInstance,
context: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation>;
async fn set_provider_instance_enabled(
&self,
command: SetProviderInstanceEnabled,
context: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation>;
async fn delete_provider_instance(
&self,
command: DeleteProviderInstance,
context: &MutationContext,
) -> AdminStoreResult<Revision>;
}
/// Client API Key 管理及其 revision CAS 写入。
#[async_trait]
pub trait ClientKeyStore: Send + Sync {
@@ -334,7 +291,6 @@ pub trait SettingsStore: Send + Sync {
pub struct AdminStorePorts {
accounts: Arc<dyn AccountStore>,
auth: Arc<dyn AuthStore>,
catalog: Arc<dyn CatalogStore>,
client_keys: Arc<dyn ClientKeyStore>,
observability: Arc<dyn ObservabilityStore>,
settings: Arc<dyn SettingsStore>,
@@ -345,7 +301,6 @@ impl AdminStorePorts {
pub fn new(
accounts: Arc<dyn AccountStore>,
auth: Arc<dyn AuthStore>,
catalog: Arc<dyn CatalogStore>,
client_keys: Arc<dyn ClientKeyStore>,
observability: Arc<dyn ObservabilityStore>,
settings: Arc<dyn SettingsStore>,
@@ -353,7 +308,6 @@ impl AdminStorePorts {
Self {
accounts,
auth,
catalog,
client_keys,
observability,
settings,
@@ -370,11 +324,6 @@ impl AdminStorePorts {
self.auth.clone()
}
#[must_use]
pub fn catalog(&self) -> Arc<dyn CatalogStore> {
self.catalog.clone()
}
#[must_use]
pub fn client_keys(&self) -> Arc<dyn ClientKeyStore> {
self.client_keys.clone()
backend/crates/gateway-admin/src/use_case/accounts.rs
+9 -41
View File
@@ -29,7 +29,7 @@ use crate::{
},
ports::{
provider::ProviderAdminRegistry,
store::{AccountStore, CatalogStore, SettingsStore},
store::{AccountStore, SettingsStore},
},
};
@@ -79,7 +79,6 @@ pub trait AccountsService: Send + Sync {
pub(crate) struct DefaultAccountsService {
accounts: Arc<dyn AccountStore>,
catalog: Arc<dyn CatalogStore>,
settings: Arc<dyn SettingsStore>,
providers: ProviderAdminRegistry,
snapshot: Arc<dyn SnapshotControl>,
@@ -90,7 +89,6 @@ impl DefaultAccountsService {
#[must_use]
pub(crate) const fn new(
accounts: Arc<dyn AccountStore>,
catalog: Arc<dyn CatalogStore>,
settings: Arc<dyn SettingsStore>,
providers: ProviderAdminRegistry,
snapshot: Arc<dyn SnapshotControl>,
@@ -98,7 +96,6 @@ impl DefaultAccountsService {
) -> Self {
Self {
accounts,
catalog,
settings,
providers,
snapshot,
@@ -146,17 +143,6 @@ impl DefaultAccountsService {
.load_runtime_settings()
.await
.map_err(|error| map_store_error(error, "runtime settings"))?;
let instance = self
.catalog
.load_provider_instance(&account.provider_instance_id)
.await
.map_err(|error| map_store_error(error, "Provider instance"))?
.ok_or_else(|| AdminError::not_found("Provider instance was not found"))?;
if instance.item.provider_kind != account.provider_kind {
return Err(AdminError::invalid(
"Provider account and instance kinds do not match",
));
}
let now = Utc::now();
let retained_range = retained_usage_range(now, settings.usage_retention_days);
let rolling_range = TimeRange {
@@ -188,7 +174,6 @@ impl DefaultAccountsService {
.await
.map_err(|error| map_provider_error(error, "provider quota"))?;
Ok(AccountDirectoryItem {
provider_instance_name: instance.item.name,
status: account_status(&account, now),
usage: usage.into_iter().next(),
account,
@@ -200,10 +185,9 @@ impl DefaultAccountsService {
#[async_trait]
impl AccountsService for DefaultAccountsService {
async fn list(&self, query: AccountListQuery) -> Result<AccountDirectoryPage, AdminError> {
let (page, settings, instances) = futures::try_join!(
let (page, settings) = futures::try_join!(
self.accounts.list_accounts(query),
self.settings.load_runtime_settings(),
self.catalog.list_provider_instances(true),
)
.map_err(|error| map_store_error(error, "account directory"))?;
let now = Utc::now();
@@ -230,11 +214,6 @@ impl AccountsService for DefaultAccountsService {
.into_iter()
.map(|usage| (usage.account_id.clone(), usage))
.collect::<BTreeMap<_, _>>();
let instances = instances
.items
.into_iter()
.map(|instance| (instance.id.clone(), instance))
.collect::<BTreeMap<_, _>>();
let quotas = futures::future::join_all(page.items.iter().map(|account| async {
let account_id = ProviderAccountId::new(account.id.clone())
.map_err(|_| AdminError::invalid("Invalid provider account ID"))?;
@@ -258,24 +237,13 @@ impl AccountsService for DefaultAccountsService {
.items
.into_iter()
.zip(quotas)
.map(|(account, quota)| {
let instance = instances
.get(&account.provider_instance_id)
.ok_or_else(|| AdminError::not_found("Provider instance was not found"))?;
if instance.provider_kind != account.provider_kind {
return Err(AdminError::invalid(
"Provider account and instance kinds do not match",
));
}
Ok(AccountDirectoryItem {
provider_instance_name: instance.name.clone(),
status: account_status(&account, now),
usage: usage.get(&account.id).cloned(),
account,
quota,
})
.map(|(account, quota)| AccountDirectoryItem {
status: account_status(&account, now),
usage: usage.get(&account.id).cloned(),
account,
quota,
})
.collect::<Result<Vec<_>, AdminError>>()?;
.collect();
Ok(AccountDirectoryPage {
config_revision: page.config_revision,
items,
@@ -416,7 +384,7 @@ impl AccountsService for DefaultAccountsService {
let result = probe
.probe(AccountProbeRequest {
account_id,
provider_instance_id: account.provider_instance_id,
provider_kind: account.provider_kind,
upstream_model,
operation,
})
backend/crates/gateway-admin/src/use_case/catalog.rs
-159
View File
@@ -1,159 +0,0 @@
//! Provider instance 目录用例。
use std::sync::Arc;
use async_trait::async_trait;
use gateway_core::routing::{ProviderInstanceId, snapshot::SnapshotControl};
use crate::{
model::{
AdminError, MutationContext, Revision,
catalog::{
CatalogListQuery, CreateProviderInstance, DeleteProviderInstance,
ProviderInstanceDetail, ProviderInstanceMutation, ProviderInstancePage,
SetProviderInstanceEnabled, UpdateProviderInstance,
},
},
ports::store::CatalogStore,
};
use super::{map_store_error, publish_committed};
/// API 消费的 Provider instance 管理服务。
#[async_trait]
pub trait CatalogService: Send + Sync {
async fn list(&self, query: CatalogListQuery) -> Result<ProviderInstancePage, AdminError>;
async fn get(&self, id: &ProviderInstanceId) -> Result<ProviderInstanceDetail, AdminError>;
async fn create(
&self,
context: &MutationContext,
command: CreateProviderInstance,
) -> Result<ProviderInstanceMutation, AdminError>;
async fn update(
&self,
context: &MutationContext,
command: UpdateProviderInstance,
) -> Result<ProviderInstanceMutation, AdminError>;
async fn set_enabled(
&self,
context: &MutationContext,
command: SetProviderInstanceEnabled,
) -> Result<ProviderInstanceMutation, AdminError>;
async fn delete(
&self,
context: &MutationContext,
command: DeleteProviderInstance,
) -> Result<Revision, AdminError>;
}
pub(crate) struct DefaultCatalogService {
store: Arc<dyn CatalogStore>,
snapshot: Arc<dyn SnapshotControl>,
}
impl DefaultCatalogService {
#[must_use]
pub(crate) fn new(store: Arc<dyn CatalogStore>, snapshot: Arc<dyn SnapshotControl>) -> Self {
Self { store, snapshot }
}
}
#[async_trait]
impl CatalogService for DefaultCatalogService {
async fn list(&self, query: CatalogListQuery) -> Result<ProviderInstancePage, AdminError> {
let catalog = self
.store
.list_provider_instances(true)
.await
.map_err(|error| map_store_error(error, "provider instance catalog"))?;
let start = match query.cursor.as_ref() {
Some(cursor) => catalog
.items
.iter()
.position(|item| &item.id == cursor)
.map(|index| index + 1)
.ok_or_else(|| AdminError::invalid("Invalid provider instance cursor"))?,
None => 0,
};
let limit = usize::from(query.page_size.get());
let has_more = catalog.items.len().saturating_sub(start) > limit;
let items = catalog
.items
.into_iter()
.skip(start)
.take(limit)
.collect::<Vec<_>>();
let next_cursor = has_more
.then(|| items.last().map(|item| item.id.clone()))
.flatten();
Ok(ProviderInstancePage {
config_revision: catalog.config_revision,
items,
next_cursor,
})
}
async fn get(&self, id: &ProviderInstanceId) -> Result<ProviderInstanceDetail, AdminError> {
self.store
.load_provider_instance(id)
.await
.map_err(|error| map_store_error(error, "provider instance"))?
.ok_or_else(|| AdminError::not_found("Provider instance was not found"))
}
async fn create(
&self,
context: &MutationContext,
command: CreateProviderInstance,
) -> Result<ProviderInstanceMutation, AdminError> {
let result = self
.store
.create_provider_instance(command, context)
.await
.map_err(|error| map_store_error(error, "provider instance"))?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
async fn update(
&self,
context: &MutationContext,
command: UpdateProviderInstance,
) -> Result<ProviderInstanceMutation, AdminError> {
let result = self
.store
.update_provider_instance(command, context)
.await
.map_err(|error| map_store_error(error, "provider instance"))?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
async fn set_enabled(
&self,
context: &MutationContext,
command: SetProviderInstanceEnabled,
) -> Result<ProviderInstanceMutation, AdminError> {
let result = self
.store
.set_provider_instance_enabled(command, context)
.await
.map_err(|error| map_store_error(error, "provider instance"))?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
async fn delete(
&self,
context: &MutationContext,
command: DeleteProviderInstance,
) -> Result<Revision, AdminError> {
let revision = self
.store
.delete_provider_instance(command, context)
.await
.map_err(|error| map_store_error(error, "provider instance"))?;
publish_committed(self.snapshot.as_ref(), revision).await?;
Ok(revision)
}
}
backend/crates/gateway-admin/src/use_case/mod.rs
+31 -46
View File
@@ -2,7 +2,6 @@
pub mod accounts;
pub mod auth;
pub mod catalog;
pub mod client_keys;
pub mod observability;
pub mod openai;
@@ -13,12 +12,13 @@ pub mod xai;
use crate::{
model::{
AdminError, AdminErrorKind, MutationContext,
accounts::{DeleteAccount, SetAccountEnabled},
accounts::{DeleteAccounts, SetAccountEnabled},
provider_credentials::{
AuthorizationMutationTarget, CredentialDetails, CredentialMutation,
CredentialMutationResult, CredentialRotationCommit, PendingAuthorizationMutation,
PreparedAuthorizationCommit, PreparedAuthorizationCredential, PreparedCredentialImport,
PreparedCredentialRotation, StartAuthorization,
AuthorizationMutationTarget, CredentialDeletion, CredentialDeletionResult,
CredentialDetails, CredentialMutation, CredentialMutationResult,
CredentialRotationCommit, PendingAuthorizationMutation, PreparedAuthorizationCommit,
PreparedAuthorizationCredential, PreparedCredentialImport, PreparedCredentialRotation,
StartAuthorization,
},
},
ports::{
@@ -93,21 +93,17 @@ async fn pending_authorization(
Some(target) => {
let details =
required_credential(accounts, provider_kind, &target.account_id, resource).await?;
if details.credential.provider_instance_id != command.provider_instance_id
|| details.credential.credential_revision != target.credential_revision
{
if details.credential.credential_revision != target.credential_revision {
return Err(AdminError::conflict(format!(
"{resource} reauthorization target is stale"
)));
}
AuthorizationMutationTarget::Reauthorize {
provider_instance_id: command.provider_instance_id.clone(),
account_id: target.account_id.clone(),
expected_credential_revision: target.credential_revision,
}
}
None => AuthorizationMutationTarget::Create {
provider_instance_id: command.provider_instance_id.clone(),
name: command.name.clone(),
},
};
@@ -123,16 +119,14 @@ async fn pending_authorization(
fn validate_prepared_import(
provider_kind: &ProviderKind,
provider_instance_id: &gateway_core::routing::ProviderInstanceId,
prepared: &PreparedCredentialImport,
resource: &'static str,
) -> Result<(), AdminError> {
if prepared.provider_kind != *provider_kind
|| prepared.provider_instance_id != *provider_instance_id
|| prepared.credentials.iter().any(|credential| {
credential.provider_kind != *provider_kind
|| credential.provider_instance_id != *provider_instance_id
})
|| prepared
.credentials
.iter()
.any(|credential| credential.provider_kind != *provider_kind)
{
return Err(AdminError::conflict(format!(
"{resource} prepared facts do not match the requested Provider scope"
@@ -149,7 +143,6 @@ fn validate_prepared_rotation(
let facts = prepared.facts();
if facts.account_id.as_str() != account.id.as_str()
|| facts.provider_kind != account.provider_kind
|| facts.provider_instance_id != account.provider_instance_id
|| facts.expected_credential_revision != account.credential_revision
{
return Err(AdminError::conflict(format!(
@@ -174,18 +167,11 @@ fn validate_authorization_commit(
}
let matches_target = match (prepared.pending.target(), &prepared.credential) {
(
AuthorizationMutationTarget::Create {
provider_instance_id,
..
},
AuthorizationMutationTarget::Create { .. },
PreparedAuthorizationCredential::Create(credential),
) => {
credential.provider_kind == *provider_kind
&& credential.provider_instance_id == *provider_instance_id
}
) => credential.provider_kind == *provider_kind,
(
AuthorizationMutationTarget::Reauthorize {
provider_instance_id,
account_id,
expected_credential_revision,
},
@@ -193,7 +179,6 @@ fn validate_authorization_commit(
) => {
let facts = credential.facts();
facts.provider_kind == *provider_kind
&& facts.provider_instance_id == *provider_instance_id
&& facts.account_id == *account_id
&& facts.expected_credential_revision == *expected_credential_revision
}
@@ -322,34 +307,34 @@ async fn set_credential_enabled(
})
}
async fn delete_credential(
async fn delete_credentials(
accounts: &dyn AccountStore,
provider: &dyn ProviderAdmin,
command: CredentialMutation,
command: CredentialDeletion,
resource: &'static str,
) -> Result<CredentialMutationResult, AdminError> {
required_credential(
accounts,
provider.provider_kind(),
&command.account_id,
resource,
)
.await?;
let account_id = command.account_id;
) -> Result<CredentialDeletionResult, AdminError> {
for account_id in &command.account_ids {
required_credential(accounts, provider.provider_kind(), account_id, resource).await?;
}
let account_ids = command.account_ids;
let revision = accounts
.delete_account(
DeleteAccount {
.delete_accounts(
DeleteAccounts {
expected_config_revision: command.expected_config_revision,
account_id: account_id.as_str().to_owned(),
account_ids: account_ids
.iter()
.map(|account_id| account_id.as_str().to_owned())
.collect(),
},
&command.context,
)
.await
.map_err(|error| map_store_error(error, resource))?;
provider.account_unavailable(&account_id).await;
Ok(CredentialMutationResult {
for account_id in &account_ids {
provider.account_unavailable(account_id).await;
}
Ok(CredentialDeletionResult {
config_revision: revision,
account_id,
credential_revision: None,
account_ids,
})
}
backend/crates/gateway-admin/src/use_case/openai.rs
+8 -10
View File
@@ -9,7 +9,8 @@ use crate::{
model::{
AdminError,
provider_credentials::{
AuthorizationStarted, CompleteAuthorization, CredentialDetails, CredentialImportCommit,
AuthorizationStarted, CompleteAuthorization, CredentialDeletion,
CredentialDeletionResult, CredentialDetails, CredentialImportCommit,
CredentialImportResult, CredentialListQuery, CredentialMutation,
CredentialMutationResult, CredentialPage, ImportCredentials, PrepareCredentialImport,
PrepareCredentialRotation, RotateCredential, StartAuthorization,
@@ -19,7 +20,7 @@ use crate::{
};
use super::{
commit_authorization, commit_credential_rotation, delete_credential, map_provider_error,
commit_authorization, commit_credential_rotation, delete_credentials, map_provider_error,
map_store_error, pending_authorization, publish_committed, required_credential,
set_credential_enabled, validate_authorization_commit, validate_prepared_import,
validate_prepared_rotation,
@@ -59,8 +60,8 @@ pub trait OpenAiService: Send + Sync {
) -> Result<CredentialMutationResult, AdminError>;
async fn delete(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError>;
command: CredentialDeletion,
) -> Result<CredentialDeletionResult, AdminError>;
}
pub(crate) struct DefaultOpenAiService {
@@ -112,18 +113,15 @@ impl OpenAiService for DefaultOpenAiService {
) -> Result<CredentialImportResult, AdminError> {
let context = command.context;
let expected_config_revision = command.expected_config_revision;
let provider_instance_id = command.provider_instance_id;
let prepared = self
.provider
.prepare_import(PrepareCredentialImport {
provider_instance_id: provider_instance_id.clone(),
document: command.document,
})
.await
.map_err(|error| map_provider_error(error, "OpenAI credential import"))?;
validate_prepared_import(
self.provider.provider_kind(),
&provider_instance_id,
&prepared,
"OpenAI credential import",
)?;
@@ -262,9 +260,9 @@ impl OpenAiService for DefaultOpenAiService {
async fn delete(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError> {
let result = delete_credential(
command: CredentialDeletion,
) -> Result<CredentialDeletionResult, AdminError> {
let result = delete_credentials(
self.accounts.as_ref(),
self.provider.as_ref(),
command,
backend/crates/gateway-admin/src/use_case/xai.rs
+10 -13
View File
@@ -9,17 +9,17 @@ use crate::{
model::{
AdminError,
provider_credentials::{
AuthorizationStarted, CompleteAuthorization, CredentialImportCommit,
CredentialImportResult, CredentialListQuery, CredentialMutation,
CredentialMutationResult, CredentialPage, ImportCredentials, PrepareCredentialImport,
StartAuthorization,
AuthorizationStarted, CompleteAuthorization, CredentialDeletion,
CredentialDeletionResult, CredentialImportCommit, CredentialImportResult,
CredentialListQuery, CredentialMutation, CredentialMutationResult, CredentialPage,
ImportCredentials, PrepareCredentialImport, StartAuthorization,
},
},
ports::{provider::ProviderAdmin, store::AccountStore},
};
use super::{
commit_authorization, delete_credential, map_provider_error, map_store_error,
commit_authorization, delete_credentials, map_provider_error, map_store_error,
pending_authorization, publish_committed, set_credential_enabled,
validate_authorization_commit, validate_prepared_import,
};
@@ -50,8 +50,8 @@ pub trait XaiService: Send + Sync {
) -> Result<CredentialMutationResult, AdminError>;
async fn delete(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError>;
command: CredentialDeletion,
) -> Result<CredentialDeletionResult, AdminError>;
}
pub(crate) struct DefaultXaiService {
@@ -90,18 +90,15 @@ impl XaiService for DefaultXaiService {
) -> Result<CredentialImportResult, AdminError> {
let context = command.context;
let expected_config_revision = command.expected_config_revision;
let provider_instance_id = command.provider_instance_id;
let prepared = self
.provider
.prepare_import(PrepareCredentialImport {
provider_instance_id: provider_instance_id.clone(),
document: command.document,
})
.await
.map_err(|error| map_provider_error(error, "xAI credential import"))?;
validate_prepared_import(
self.provider.provider_kind(),
&provider_instance_id,
&prepared,
"xAI credential import",
)?;
@@ -198,9 +195,9 @@ impl XaiService for DefaultXaiService {
async fn delete(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError> {
let result = delete_credential(
command: CredentialDeletion,
) -> Result<CredentialDeletionResult, AdminError> {
let result = delete_credentials(
self.accounts.as_ref(),
self.provider.as_ref(),
command,
backend/crates/gateway-admin/tests/main.rs
+1 -4
View File
@@ -12,9 +12,7 @@ use gateway_admin::{
provider_credentials::{CredentialCommitGuard, ProviderDocument},
settings::AdminApiKey,
},
ports::store::{
AccountStore, AuthStore, CatalogStore, ClientKeyStore, ObservabilityStore, SettingsStore,
},
ports::store::{AccountStore, AuthStore, ClientKeyStore, ObservabilityStore, SettingsStore},
};
mod use_case;
@@ -98,7 +96,6 @@ fn all_store_capability_traits_should_be_object_safe() {
assert_object_safe::<dyn AccountStore>();
assert_object_safe::<dyn AuthStore>();
assert_object_safe::<dyn CatalogStore>();
assert_object_safe::<dyn ClientKeyStore>();
assert_object_safe::<dyn ObservabilityStore>();
assert_object_safe::<dyn SettingsStore>();
backend/crates/gateway-admin/tests/use_case/accounts.rs
+13 -105
View File
@@ -9,7 +9,7 @@ use gateway_core::{
probe::{AccountProbe, AccountProbeRequest, AccountProbeResult},
},
error::GatewayError,
routing::{ProviderInstanceId, ProviderKind},
routing::ProviderKind,
};
use gateway_admin::{
@@ -18,12 +18,7 @@ use gateway_admin::{
MutationContext, Revision,
accounts::{
AccountAvailability, AccountListQuery, AccountPage, AccountRecord, AccountSummary,
AccountUsage, DeleteAccount, SetAccountEnabled,
},
catalog::{
CreateProviderInstance, DeleteProviderInstance, ProviderInstance,
ProviderInstanceCatalog, ProviderInstanceDetail, ProviderInstanceMutation,
SetProviderInstanceEnabled, UpdateProviderInstance,
AccountUsage, DeleteAccounts, SetAccountEnabled,
},
observability::TimeRange,
provider_credentials::{
@@ -47,8 +42,7 @@ use gateway_admin::{
ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind, ProviderAdminRegistry,
},
store::{
AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStoreResult, CatalogStore,
SettingsStore,
AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStoreResult, SettingsStore,
},
},
};
@@ -101,7 +95,6 @@ impl FakeProviderAdmin {
PreparedCredentialRotation::new(
PreparedCredentialRotationFacts {
account_id: ProviderAccountId::new(account.id.clone()).expect("account ID"),
provider_instance_id: account.provider_instance_id.clone(),
provider_kind: account.provider_kind.clone(),
expected_credential_revision: account.credential_revision,
name: account.name.clone(),
@@ -153,18 +146,13 @@ impl ProviderAdmin for FakeProviderAdmin {
async fn prepare_import(
&self,
command: PrepareCredentialImport,
_command: PrepareCredentialImport,
) -> Result<PreparedCredentialImport, ProviderAdminError> {
self.record("provider.prepare_import");
self.require_available()?;
Ok(PreparedCredentialImport {
provider_kind: self.kind.clone(),
provider_instance_id: command.provider_instance_id.clone(),
credentials: vec![prepared_create(
self.kind.clone(),
command.provider_instance_id,
"prepared-import",
)],
credentials: vec![prepared_create(self.kind.clone(), "prepared-import")],
})
}
@@ -198,22 +186,15 @@ impl ProviderAdmin for FakeProviderAdmin {
return Err(ProviderAdminError::new(ProviderAdminErrorKind::NotFound));
}
let credential = match pending.target() {
AuthorizationMutationTarget::Create {
provider_instance_id,
name,
} => PreparedAuthorizationCredential::Create(prepared_create(
self.kind.clone(),
provider_instance_id.clone(),
name,
)),
AuthorizationMutationTarget::Create { name } => {
PreparedAuthorizationCredential::Create(prepared_create(self.kind.clone(), name))
}
AuthorizationMutationTarget::Reauthorize {
provider_instance_id,
account_id,
expected_credential_revision,
} => {
let mut account = account_record(self.kind.as_str());
account.id = account_id.as_str().to_owned();
account.provider_instance_id = provider_instance_id.clone();
account.credential_revision = *expected_credential_revision;
PreparedAuthorizationCredential::Reauthorize(self.prepared_rotation(&account))
}
@@ -502,9 +483,9 @@ impl AccountStore for FakeAccountStore {
Ok(revision(2))
}
async fn delete_account(
async fn delete_accounts(
&self,
_: DeleteAccount,
_: DeleteAccounts,
context: &MutationContext,
) -> AdminStoreResult<Revision> {
self.record("store.delete");
@@ -574,62 +555,6 @@ impl SettingsStore for StaticSettingsStore {
}
}
struct StaticCatalogStore {
instance: ProviderInstance,
}
#[async_trait]
impl CatalogStore for StaticCatalogStore {
async fn list_provider_instances(&self, _: bool) -> AdminStoreResult<ProviderInstanceCatalog> {
Ok(ProviderInstanceCatalog {
config_revision: revision(1),
items: vec![self.instance.clone()],
})
}
async fn load_provider_instance(
&self,
id: &ProviderInstanceId,
) -> AdminStoreResult<Option<ProviderInstanceDetail>> {
Ok((id == &self.instance.id).then(|| ProviderInstanceDetail {
config_revision: revision(1),
item: self.instance.clone(),
}))
}
async fn create_provider_instance(
&self,
_: CreateProviderInstance,
_: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation> {
Err(store_unavailable())
}
async fn update_provider_instance(
&self,
_: UpdateProviderInstance,
_: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation> {
Err(store_unavailable())
}
async fn set_provider_instance_enabled(
&self,
_: SetProviderInstanceEnabled,
_: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation> {
Err(store_unavailable())
}
async fn delete_provider_instance(
&self,
_: DeleteProviderInstance,
_: &MutationContext,
) -> AdminStoreResult<Revision> {
Err(store_unavailable())
}
}
struct RecordingGuard {
events: EventLog,
finished: bool,
@@ -731,7 +656,7 @@ async fn accounts_refresh_should_keep_guard_through_store_commit() {
.await
.expect("refresh credential");
assert_eq!(result.config_revision, revision(2));
assert_eq!(result.account.provider_instance_name, "Test Provider");
assert_eq!(result.account.account.provider_kind.as_str(), "openai");
assert_eq!(
result.account.status,
gateway_admin::model::accounts::AccountStatus::Active
@@ -771,7 +696,7 @@ async fn accounts_list_should_return_complete_directory_semantics() {
assert_eq!(page.summary.total, 1);
assert_eq!(page.summary.active, 1);
let account = page.items.first().expect("account item");
assert_eq!(account.provider_instance_name, "Test Provider");
assert_eq!(account.account.provider_kind.as_str(), "openai");
assert_eq!(
account.status,
gateway_admin::model::accounts::AccountStatus::Active
@@ -856,7 +781,6 @@ pub(super) fn account_record(kind: &str) -> AccountRecord {
let now = Utc::now();
AccountRecord {
id: "acct_test".to_owned(),
provider_instance_id: ProviderInstanceId::new(format!("inst_{kind}")).expect("instance ID"),
provider_kind: ProviderKind::new(kind).expect("provider kind"),
name: "test account".to_owned(),
email: Some("test@example.invalid".to_owned()),
@@ -882,15 +806,10 @@ pub(super) fn revision(value: u64) -> Revision {
Revision::new(value).expect("positive revision")
}
fn prepared_create(
provider_kind: ProviderKind,
provider_instance_id: ProviderInstanceId,
name: &str,
) -> PreparedCredentialCreate {
fn prepared_create(provider_kind: ProviderKind, name: &str) -> PreparedCredentialCreate {
let now = Utc::now();
PreparedCredentialCreate {
account_id: ProviderAccountId::new("acct_prepared").expect("prepared account ID"),
provider_instance_id,
provider_kind,
name: name.to_owned(),
email: Some("prepared@example.invalid".to_owned()),
@@ -923,19 +842,8 @@ async fn accounts_service(
provider: Arc<FakeProviderAdmin>,
store: Arc<FakeAccountStore>,
) -> AdminServices {
let now = Utc::now();
let instance = ProviderInstance {
id: store.account.provider_instance_id.clone(),
provider_kind: store.account.provider_kind.clone(),
name: "Test Provider".to_owned(),
base_url: "https://example.invalid".to_owned(),
enabled: true,
created_at: now,
updated_at: now,
};
super::AdminHarness::new()
.accounts(store)
.catalog(Arc::new(StaticCatalogStore { instance }))
.settings(Arc::new(StaticSettingsStore))
.provider(provider)
.probe(Arc::new(SuccessfulAccountProbe))
backend/crates/gateway-admin/tests/use_case/catalog.rs
-144
View File
@@ -1,144 +0,0 @@
use std::sync::Arc;
use async_trait::async_trait;
use chrono::Utc;
use gateway_core::routing::{ProviderInstanceId, ProviderKind};
use gateway_admin::{
model::{
MutationContext, PageSize, Revision,
catalog::{
CatalogListQuery, CreateProviderInstance, DeleteProviderInstance, ProviderInstance,
ProviderInstanceCatalog, ProviderInstanceDetail, ProviderInstanceMutation,
SetProviderInstanceEnabled, UpdateProviderInstance,
},
},
ports::store::{AdminStoreError, AdminStoreErrorKind, AdminStoreResult, CatalogStore},
};
struct CatalogFixtureStore {
items: Vec<ProviderInstance>,
}
#[async_trait]
impl CatalogStore for CatalogFixtureStore {
async fn list_provider_instances(&self, _: bool) -> AdminStoreResult<ProviderInstanceCatalog> {
Ok(ProviderInstanceCatalog {
config_revision: Revision::new(7).expect("revision"),
items: self.items.clone(),
})
}
async fn load_provider_instance(
&self,
id: &ProviderInstanceId,
) -> AdminStoreResult<Option<ProviderInstanceDetail>> {
Ok(self
.items
.iter()
.find(|item| &item.id == id)
.cloned()
.map(|item| ProviderInstanceDetail {
config_revision: Revision::new(7).expect("revision"),
item,
}))
}
async fn create_provider_instance(
&self,
_: CreateProviderInstance,
_: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation> {
Err(unused())
}
async fn update_provider_instance(
&self,
_: UpdateProviderInstance,
_: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation> {
Err(unused())
}
async fn set_provider_instance_enabled(
&self,
_: SetProviderInstanceEnabled,
_: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation> {
Err(unused())
}
async fn delete_provider_instance(
&self,
_: DeleteProviderInstance,
_: &MutationContext,
) -> AdminStoreResult<Revision> {
Err(unused())
}
}
#[tokio::test]
async fn catalog_cursor_should_resume_after_exact_instance() {
let items = ["inst_a", "inst_b", "inst_c"]
.into_iter()
.map(instance)
.collect();
let services = super::AdminHarness::new()
.catalog(Arc::new(CatalogFixtureStore { items }))
.build()
.await;
let page = services
.catalog()
.list(CatalogListQuery {
cursor: Some(ProviderInstanceId::new("inst_a").expect("cursor")),
page_size: PageSize::new(1).expect("page size"),
})
.await
.expect("catalog page");
assert_eq!(page.items[0].id.as_str(), "inst_b");
assert_eq!(
page.next_cursor.as_ref().map(ProviderInstanceId::as_str),
Some("inst_b")
);
}
#[tokio::test]
async fn catalog_detail_should_keep_its_control_plane_revision() {
let services = super::AdminHarness::new()
.catalog(Arc::new(CatalogFixtureStore {
items: vec![instance("inst_detail")],
}))
.build()
.await;
let detail = services
.catalog()
.get(&ProviderInstanceId::new("inst_detail").expect("instance ID"))
.await
.expect("catalog detail");
assert_eq!(detail.config_revision.get(), 7);
assert_eq!(detail.item.id.as_str(), "inst_detail");
}
fn instance(id: &str) -> ProviderInstance {
ProviderInstance {
id: ProviderInstanceId::new(id).expect("instance ID"),
provider_kind: ProviderKind::new("openai").expect("provider"),
name: id.to_owned(),
base_url: "https://example.com".to_owned(),
enabled: true,
created_at: Utc::now(),
updated_at: Utc::now(),
}
}
fn unused() -> AdminStoreError {
AdminStoreError::new(
AdminStoreErrorKind::Unavailable,
"catalog",
"unused in this test",
)
}
backend/crates/gateway-admin/tests/use_case/mod.rs
+5 -65
View File
@@ -1,6 +1,5 @@
mod accounts;
mod auth;
mod catalog;
mod client_keys;
mod observability;
mod openai;
@@ -17,15 +16,10 @@ use gateway_admin::{
model::{
MutationContext, Revision,
accounts::{
AccountListQuery, AccountPage, AccountRecord, AccountUsage, DeleteAccount,
AccountListQuery, AccountPage, AccountRecord, AccountUsage, DeleteAccounts,
SetAccountEnabled,
},
auth::{AdminAuditEvent, AdminSession},
catalog::{
CreateProviderInstance, DeleteProviderInstance, ProviderInstanceCatalog,
ProviderInstanceDetail, ProviderInstanceMutation, SetProviderInstanceEnabled,
UpdateProviderInstance,
},
client_keys::{
ClientKeyListQuery, ClientKeyPage, ClientKeyRecord, ClientKeySecret, DeleteClientKey,
NewClientKey, SetClientKeyEnabled, UpdateClientKey,
@@ -52,7 +46,7 @@ use gateway_admin::{
provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind},
store::{
AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts, AdminStoreResult,
AuthStore, CatalogStore, ClientKeyStore, ObservabilityStore, SettingsStore,
AuthStore, ClientKeyStore, ObservabilityStore, SettingsStore,
},
system::{
SystemOperationError, SystemOperationErrorKind, SystemOperations,
@@ -67,14 +61,13 @@ use gateway_core::{
},
error::{GatewayError, GatewayErrorKind},
policy::ClientApiKeyId,
routing::{ConfigRevision, ProviderInstanceId, ProviderKind, snapshot::SnapshotControl},
routing::{ConfigRevision, ProviderKind, snapshot::SnapshotControl},
};
pub(super) struct AdminHarness {
default_password: String,
accounts: Arc<dyn AccountStore>,
auth: Arc<dyn AuthStore>,
catalog: Arc<dyn CatalogStore>,
client_keys: Arc<dyn ClientKeyStore>,
observability: Arc<dyn ObservabilityStore>,
settings: Arc<dyn SettingsStore>,
@@ -90,7 +83,6 @@ impl AdminHarness {
default_password: "strong-test-password".to_owned(),
accounts: unavailable.clone(),
auth: Arc::new(BootstrapAuthStore::default()),
catalog: unavailable.clone(),
client_keys: unavailable.clone(),
observability: unavailable.clone(),
settings: unavailable,
@@ -118,11 +110,6 @@ impl AdminHarness {
self
}
pub(super) fn catalog(mut self, store: Arc<dyn CatalogStore>) -> Self {
self.catalog = store;
self
}
pub(super) fn client_keys(mut self, store: Arc<dyn ClientKeyStore>) -> Self {
self.client_keys = store;
self
@@ -165,7 +152,6 @@ impl AdminHarness {
AdminStorePorts::new(
self.accounts,
self.auth,
self.catalog,
self.client_keys,
self.observability,
self.settings,
@@ -322,9 +308,9 @@ impl AccountStore for UnavailableStore {
Err(unavailable("account enabled"))
}
async fn delete_account(
async fn delete_accounts(
&self,
_: DeleteAccount,
_: DeleteAccounts,
_: &MutationContext,
) -> AdminStoreResult<Revision> {
Err(unavailable("account delete"))
@@ -339,52 +325,6 @@ impl AccountStore for UnavailableStore {
}
}
#[async_trait]
impl CatalogStore for UnavailableStore {
async fn list_provider_instances(&self, _: bool) -> AdminStoreResult<ProviderInstanceCatalog> {
Err(unavailable("provider catalog"))
}
async fn load_provider_instance(
&self,
_: &ProviderInstanceId,
) -> AdminStoreResult<Option<ProviderInstanceDetail>> {
Err(unavailable("provider instance"))
}
async fn create_provider_instance(
&self,
_: CreateProviderInstance,
_: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation> {
Err(unavailable("provider create"))
}
async fn update_provider_instance(
&self,
_: UpdateProviderInstance,
_: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation> {
Err(unavailable("provider update"))
}
async fn set_provider_instance_enabled(
&self,
_: SetProviderInstanceEnabled,
_: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation> {
Err(unavailable("provider enabled"))
}
async fn delete_provider_instance(
&self,
_: DeleteProviderInstance,
_: &MutationContext,
) -> AdminStoreResult<Revision> {
Err(unavailable("provider delete"))
}
}
#[async_trait]
impl ClientKeyStore for UnavailableStore {
async fn list_client_keys(&self, _: ClientKeyListQuery) -> AdminStoreResult<ClientKeyPage> {
backend/crates/gateway-admin/tests/use_case/openai.rs
+10 -12
View File
@@ -1,13 +1,12 @@
use std::sync::Arc;
use gateway_core::engine::credential::ProviderAccountId;
use gateway_core::routing::ProviderInstanceId;
use gateway_admin::{
AdminServices,
model::provider_credentials::{
AuthorizationMutationTarget, CompleteAuthorization, CredentialMutation, ImportCredentials,
ReauthorizationTarget, StartAuthorization,
AuthorizationMutationTarget, CompleteAuthorization, CredentialDeletion, CredentialMutation,
ImportCredentials, ReauthorizationTarget, StartAuthorization,
},
ports::provider::ProviderAdminErrorKind,
};
@@ -69,7 +68,7 @@ async fn openai_delete_should_commit_then_release_provider_resources() {
services
.openai()
.delete(mutation("acct_test"))
.delete(deletion("acct_test"))
.await
.expect("delete credential");
@@ -116,7 +115,6 @@ async fn openai_import_should_prepare_before_atomic_store_commit() {
.import_document(ImportCredentials {
context: context("import-openai"),
expected_config_revision: revision(1),
provider_instance_id: instance(),
document: document(),
})
.await
@@ -142,7 +140,6 @@ async fn openai_import_provider_error_should_not_touch_store_transaction() {
.import_document(ImportCredentials {
context: context("import-openai-error"),
expected_config_revision: revision(1),
provider_instance_id: instance(),
document: document(),
})
.await
@@ -162,7 +159,6 @@ async fn openai_reauthorization_should_restore_pending_envelope_and_hold_guard_t
.start_authorization(StartAuthorization {
context: context("oauth-start-openai"),
expected_config_revision: revision(7),
provider_instance_id: instance(),
name: "reauthorize".to_owned(),
reauthorization: Some(ReauthorizationTarget {
account_id: ProviderAccountId::new("acct_test").expect("account ID"),
@@ -182,11 +178,9 @@ async fn openai_reauthorization_should_restore_pending_envelope_and_hold_guard_t
assert!(matches!(
pending.target(),
AuthorizationMutationTarget::Reauthorize {
provider_instance_id,
account_id,
expected_credential_revision,
} if provider_instance_id == &instance()
&& account_id.as_str() == "acct_test"
} if account_id.as_str() == "acct_test"
&& *expected_credential_revision == revision(1)
));
assert!(!format!("{pending:?}").contains("admin-test"));
@@ -231,6 +225,10 @@ fn mutation(account_id: &str) -> CredentialMutation {
}
}
fn instance() -> ProviderInstanceId {
ProviderInstanceId::new("inst_openai").expect("instance ID")
fn deletion(account_id: &str) -> CredentialDeletion {
CredentialDeletion {
context: context("request-openai"),
expected_config_revision: revision(1),
account_ids: vec![ProviderAccountId::new(account_id).expect("account ID")],
}
}
backend/crates/gateway-admin/tests/use_case/xai.rs
+1 -2
View File
@@ -1,6 +1,6 @@
use std::sync::Arc;
use gateway_core::{engine::credential::ProviderAccountId, routing::ProviderInstanceId};
use gateway_core::engine::credential::ProviderAccountId;
use gateway_admin::{
AdminServices,
@@ -47,7 +47,6 @@ async fn xai_import_should_prepare_before_atomic_store_commit() {
.import_document(ImportCredentials {
context: context("import-xai"),
expected_config_revision: revision(3),
provider_instance_id: ProviderInstanceId::new("inst_xai").expect("instance ID"),
document: document(),
})
.await
backend/crates/gateway-api/src/admin/accounts.rs
+741 -32
View File
@@ -12,6 +12,7 @@ use axum::{
},
routing::{get, post},
};
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use chrono::{DateTime, FixedOffset, Utc};
use futures::{Stream, StreamExt as _};
use gateway_admin::model::{
@@ -23,15 +24,18 @@ use gateway_admin::model::{
},
provider_credentials::{
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
AuthorizationStarted, CompleteAuthorization, CredentialDeletion, CredentialDeletionResult,
CredentialImportResult, CredentialMutation, CredentialMutationResult, ImportCredentials,
ProviderDocument, ProviderModels, ProviderQuota, ProviderQuotaWindow,
ReauthorizationTarget, RotateCredential, StartAuthorization,
},
};
use gateway_core::{
engine::credential::ProviderAccountId,
engine::credential::{OpaqueProviderData, ProviderAccountId},
routing::{ProviderKind, UpstreamModelId},
};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use serde_json::{Map, Value};
use super::{
AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, PageMeta,
@@ -41,6 +45,13 @@ use super::{
const DEFAULT_PAGE_SIZE: u32 = 50;
const MAX_PAGE_SIZE: u32 = 200;
const MAX_SEARCH_BYTES: usize = 256;
const MAX_ID_BYTES: usize = 256;
const MAX_NAME_BYTES: usize = 512;
const MAX_IMPORT_DATA_BYTES: usize = 64 * 1024 * 1024;
const MAX_ACCESS_TOKEN_BYTES: usize = 16 * 1024;
const MAX_REFRESH_TOKEN_BYTES: usize = 64 * 1024;
const MAX_CALLBACK_URL_BYTES: usize = 64 * 1024;
const MAX_ACCOUNT_DELETE_BATCH: usize = 200;
/// 账号列表查询参数。
#[derive(Debug, Clone, Default, Deserialize)]
@@ -173,8 +184,6 @@ pub struct AccountView {
pub id: String,
pub name: String,
pub provider: String,
pub provider_instance_id: String,
pub provider_instance_name: String,
pub resource_ref: String,
pub email: Option<String>,
pub account_id: Option<String>,
@@ -307,27 +316,27 @@ pub struct CurrencyCostView {
/// 账号详情类 GET 的固定 ID query。
#[derive(Debug, Clone, Deserialize)]
#[serde(deny_unknown_fields)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AccountIdQuery {
pub id: String,
pub account_id: String,
}
impl AccountIdQuery {
pub fn validate(&self) -> Result<(), WireValidationError> {
require_account_id(&self.id)
require_account_id(&self.account_id, "accountId")
}
fn into_id(self) -> Result<ProviderAccountId, WireValidationError> {
self.validate()?;
ProviderAccountId::new(self.id).map_err(|_| WireValidationError::new("id"))
ProviderAccountId::new(self.account_id).map_err(|_| WireValidationError::new("accountId"))
}
}
/// 敏感导出的固定 query；IDs 使用逗号分隔，禁止隐式导出全部账号。
#[derive(Debug, Clone, Deserialize)]
#[serde(deny_unknown_fields)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AccountExportQuery {
pub ids: String,
pub account_ids: String,
pub confirm: String,
}
@@ -337,41 +346,47 @@ impl AccountExportQuery {
return Err(WireValidationError::new("confirm"));
}
let ids = self
.ids
.account_ids
.split(',')
.map(str::trim)
.filter(|id| !id.is_empty())
.map(ToOwned::to_owned)
.collect::<Vec<_>>();
if ids.is_empty() || ids.len() > 200 || ids.iter().any(|id| require_account_id(id).is_err())
if ids.is_empty()
|| ids.len() > 200
|| ids
.iter()
.any(|id| require_account_id(id, "accountIds").is_err())
{
return Err(WireValidationError::new("ids"));
return Err(WireValidationError::new("accountIds"));
}
let unique = ids.iter().collect::<BTreeSet<_>>();
if unique.len() != ids.len() {
return Err(WireValidationError::new("ids"));
return Err(WireValidationError::new("accountIds"));
}
ids.into_iter()
.map(|id| ProviderAccountId::new(id).map_err(|_| WireValidationError::new("ids")))
.map(|id| {
ProviderAccountId::new(id).map_err(|_| WireValidationError::new("accountIds"))
})
.collect()
}
}
/// 账号运行期动作。
#[derive(Debug, Clone, Deserialize)]
#[serde(deny_unknown_fields)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AccountActionRequest {
pub id: String,
pub account_id: String,
}
impl AccountActionRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
require_account_id(&self.id)
require_account_id(&self.account_id, "accountId")
}
fn into_id(self) -> Result<ProviderAccountId, WireValidationError> {
self.validate()?;
ProviderAccountId::new(self.id).map_err(|_| WireValidationError::new("id"))
ProviderAccountId::new(self.account_id).map_err(|_| WireValidationError::new("accountId"))
}
}
@@ -379,13 +394,13 @@ impl AccountActionRequest {
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AccountRefreshRequest {
pub id: String,
pub account_id: String,
pub expected_config_revision: u64,
}
impl AccountRefreshRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
require_account_id(&self.id)?;
require_account_id(&self.account_id, "accountId")?;
if self.expected_config_revision == 0 {
return Err(WireValidationError::new("expectedConfigRevision"));
}
@@ -397,7 +412,8 @@ impl AccountRefreshRequest {
Ok((
Revision::new(self.expected_config_revision)
.map_err(|_| WireValidationError::new("expectedConfigRevision"))?,
ProviderAccountId::new(self.id).map_err(|_| WireValidationError::new("id"))?,
ProviderAccountId::new(self.account_id)
.map_err(|_| WireValidationError::new("accountId"))?,
))
}
}
@@ -406,13 +422,13 @@ impl AccountRefreshRequest {
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AccountTestQuery {
pub id: String,
pub account_id: String,
pub model_id: String,
}
impl AccountTestQuery {
pub fn validate(&self) -> Result<(), WireValidationError> {
require_account_id(&self.id)?;
require_account_id(&self.account_id, "accountId")?;
if self.model_id.trim().is_empty() || self.model_id.chars().any(char::is_control) {
return Err(WireValidationError::new("modelId"));
}
@@ -422,7 +438,8 @@ impl AccountTestQuery {
fn into_command(self) -> Result<(ProviderAccountId, UpstreamModelId), WireValidationError> {
self.validate()?;
Ok((
ProviderAccountId::new(self.id).map_err(|_| WireValidationError::new("id"))?,
ProviderAccountId::new(self.account_id)
.map_err(|_| WireValidationError::new("accountId"))?,
UpstreamModelId::new(self.model_id).map_err(|_| WireValidationError::new("modelId"))?,
))
}
@@ -536,26 +553,413 @@ impl From<DomainConnectionTestEvent> for AccountConnectionTestEvent {
}
}
fn require_account_id(value: &str) -> Result<(), WireValidationError> {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum AccountProvider {
OpenAi,
Xai,
}
impl AccountProvider {
fn parse(value: &str) -> Result<Self, WireValidationError> {
match value.trim() {
"openai" => Ok(Self::OpenAi),
"xai" => Ok(Self::Xai),
_ => Err(WireValidationError::new("provider")),
}
}
}
/// Provider-owned 账号导入请求；公共 API 不解释 `data` 内部字段。
#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AccountImportRequest {
pub provider: String,
pub expected_config_revision: u64,
pub data: Value,
}
impl AccountImportRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
AccountProvider::parse(&self.provider)?;
require_positive(self.expected_config_revision, "expectedConfigRevision")?;
if !self.data.is_object()
|| serde_json::to_vec(&self.data)
.map_or(true, |encoded| encoded.len() > MAX_IMPORT_DATA_BYTES)
{
return Err(WireValidationError::new("data"));
}
Ok(())
}
fn into_command(
self,
context: gateway_admin::model::MutationContext,
) -> Result<(AccountProvider, ImportCredentials), WireValidationError> {
self.validate()?;
let provider = AccountProvider::parse(&self.provider)?;
Ok((
provider,
ImportCredentials {
context,
expected_config_revision: revision(
self.expected_config_revision,
"expectedConfigRevision",
)?,
document: provider_document(self.data, "data")?,
},
))
}
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct StartAccountAuthorizationRequest {
pub provider: String,
pub expected_config_revision: u64,
pub name: String,
pub account_id: Option<String>,
pub expected_credential_revision: Option<u64>,
}
impl StartAccountAuthorizationRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
AccountProvider::parse(&self.provider)?;
require_positive(self.expected_config_revision, "expectedConfigRevision")?;
require_text(&self.name, MAX_NAME_BYTES, "name")?;
match (
self.account_id.as_deref(),
self.expected_credential_revision,
) {
(None, None) => Ok(()),
(Some(account_id), Some(credential_revision)) => {
require_account_id(account_id, "accountId")?;
require_positive(credential_revision, "expectedCredentialRevision")
}
_ => Err(WireValidationError::new("reauthorization")),
}
}
fn into_command(
self,
context: gateway_admin::model::MutationContext,
) -> Result<(AccountProvider, StartAuthorization), WireValidationError> {
self.validate()?;
let provider = AccountProvider::parse(&self.provider)?;
let reauthorization = match (self.account_id, self.expected_credential_revision) {
(None, None) => None,
(Some(account_id), Some(credential_revision)) => Some(ReauthorizationTarget {
account_id: ProviderAccountId::new(account_id)
.map_err(|_| WireValidationError::new("accountId"))?,
credential_revision: revision(credential_revision, "expectedCredentialRevision")?,
}),
_ => return Err(WireValidationError::new("reauthorization")),
};
Ok((
provider,
StartAuthorization {
context,
expected_config_revision: revision(
self.expected_config_revision,
"expectedConfigRevision",
)?,
name: self.name,
reauthorization,
},
))
}
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct CompleteAccountAuthorizationRequest {
pub provider: String,
pub flow_id: String,
pub callback_url: String,
}
impl CompleteAccountAuthorizationRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
let provider = AccountProvider::parse(&self.provider)?;
match provider {
AccountProvider::OpenAi => {
if !URL_SAFE_NO_PAD
.decode(&self.flow_id)
.is_ok_and(|decoded| decoded.len() == 32)
{
return Err(WireValidationError::new("flowId"));
}
}
AccountProvider::Xai => require_wire_id(&self.flow_id, "flowId")?,
}
require_text(&self.callback_url, MAX_CALLBACK_URL_BYTES, "callbackUrl")
}
fn into_command(
self,
context: gateway_admin::model::MutationContext,
) -> Result<(AccountProvider, CompleteAuthorization), WireValidationError> {
self.validate()?;
let provider = AccountProvider::parse(&self.provider)?;
Ok((
provider,
CompleteAuthorization {
context,
flow_id: self.flow_id,
callback_url: self.callback_url,
},
))
}
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AccountMutationRequest {
pub provider: String,
pub account_id: String,
pub expected_config_revision: u64,
}
impl AccountMutationRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
AccountProvider::parse(&self.provider)?;
require_account_id(&self.account_id, "accountId")?;
require_positive(self.expected_config_revision, "expectedConfigRevision")
}
fn into_command(
self,
context: gateway_admin::model::MutationContext,
) -> Result<(AccountProvider, CredentialMutation), WireValidationError> {
self.validate()?;
let provider = AccountProvider::parse(&self.provider)?;
Ok((
provider,
CredentialMutation {
context,
expected_config_revision: revision(
self.expected_config_revision,
"expectedConfigRevision",
)?,
account_id: ProviderAccountId::new(self.account_id)
.map_err(|_| WireValidationError::new("accountId"))?,
},
))
}
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AccountDeletionRequest {
pub provider: String,
pub account_ids: Vec<String>,
pub expected_config_revision: u64,
}
impl AccountDeletionRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
AccountProvider::parse(&self.provider)?;
if self.account_ids.is_empty() || self.account_ids.len() > MAX_ACCOUNT_DELETE_BATCH {
return Err(WireValidationError::new("accountIds"));
}
let mut unique = BTreeSet::new();
for account_id in &self.account_ids {
require_account_id(account_id, "accountIds")?;
if !unique.insert(account_id.as_str()) {
return Err(WireValidationError::new("accountIds"));
}
}
require_positive(self.expected_config_revision, "expectedConfigRevision")
}
fn into_command(
self,
context: gateway_admin::model::MutationContext,
) -> Result<(AccountProvider, CredentialDeletion), WireValidationError> {
self.validate()?;
let provider = AccountProvider::parse(&self.provider)?;
Ok((
provider,
CredentialDeletion {
context,
expected_config_revision: revision(
self.expected_config_revision,
"expectedConfigRevision",
)?,
account_ids: self
.account_ids
.into_iter()
.map(ProviderAccountId::new)
.collect::<Result<Vec<_>, _>>()
.map_err(|_| WireValidationError::new("accountIds"))?,
},
))
}
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct RotateAccountRequest {
pub provider: String,
pub account_id: String,
pub expected_config_revision: u64,
pub expected_credential_revision: u64,
pub access_token: String,
pub refresh_token: Option<String>,
}
impl RotateAccountRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
if AccountProvider::parse(&self.provider)? != AccountProvider::OpenAi {
return Err(WireValidationError::new("provider"));
}
require_account_id(&self.account_id, "accountId")?;
require_positive(self.expected_config_revision, "expectedConfigRevision")?;
require_positive(
self.expected_credential_revision,
"expectedCredentialRevision",
)?;
validate_oauth_material(&self.access_token, self.refresh_token.as_deref())
}
fn into_command(
self,
context: gateway_admin::model::MutationContext,
) -> Result<RotateCredential, WireValidationError> {
self.validate()?;
let mut material = Map::new();
material.insert("access_token".to_owned(), Value::String(self.access_token));
material.insert(
"refresh_token".to_owned(),
self.refresh_token.map_or(Value::Null, Value::String),
);
Ok(RotateCredential {
mutation: CredentialMutation {
context,
expected_config_revision: revision(
self.expected_config_revision,
"expectedConfigRevision",
)?,
account_id: ProviderAccountId::new(self.account_id)
.map_err(|_| WireValidationError::new("accountId"))?,
},
expected_credential_revision: revision(
self.expected_credential_revision,
"expectedCredentialRevision",
)?,
provider_material: ProviderDocument::new(OpaqueProviderData::new(material)),
})
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountImportData {
pub config_revision: u64,
pub imported_count: usize,
pub account_ids: Vec<String>,
}
impl From<CredentialImportResult> for AccountImportData {
fn from(result: CredentialImportResult) -> Self {
let account_ids = result
.credential_ids
.into_iter()
.map(|account_id| account_id.to_string())
.collect::<Vec<_>>();
Self {
config_revision: result.config_revision.get(),
imported_count: account_ids.len(),
account_ids,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountAuthorizationData {
pub flow_id: String,
pub authorization_url: String,
pub expires_at: DateTime<Utc>,
}
impl From<AuthorizationStarted> for AccountAuthorizationData {
fn from(started: AuthorizationStarted) -> Self {
Self {
flow_id: started.flow_id,
authorization_url: started.authorization_url,
expires_at: started.expires_at,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountMutationData {
pub config_revision: u64,
pub account_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub credential_revision: Option<u64>,
}
impl From<CredentialMutationResult> for AccountMutationData {
fn from(result: CredentialMutationResult) -> Self {
Self {
config_revision: result.config_revision.get(),
account_id: result.account_id.to_string(),
credential_revision: result.credential_revision.map(Revision::get),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountDeletionData {
pub config_revision: u64,
pub deleted_count: usize,
pub account_ids: Vec<String>,
}
impl From<CredentialDeletionResult> for AccountDeletionData {
fn from(result: CredentialDeletionResult) -> Self {
let account_ids = result
.account_ids
.into_iter()
.map(|account_id| account_id.to_string())
.collect::<Vec<_>>();
Self {
config_revision: result.config_revision.get(),
deleted_count: account_ids.len(),
account_ids,
}
}
}
fn require_account_id(value: &str, field: &'static str) -> Result<(), WireValidationError> {
if value.trim().is_empty()
|| value.len() > 128
|| value.chars().any(char::is_control)
|| !value.starts_with("acct_")
{
return Err(WireValidationError::new("id"));
return Err(WireValidationError::new(field));
}
Ok(())
}
/// 构造固定 GET 账号目录路由。
/// 构造统一账号管理路由。
pub fn router<S>() -> Router<S>
where
S: AdminSessionState + Clone + Send + Sync + 'static,
{
Router::new()
.route("/api/admin/accounts", get(list_accounts::<S>))
.route("/api/admin/accounts/detail", get(account_detail::<S>))
.route("/api/admin/accounts/export", get(export_accounts::<S>))
.route("/api/admin/accounts/import", post(import_accounts::<S>))
.route("/api/admin/accounts/refresh", post(refresh_account::<S>))
.route("/api/admin/accounts/rotate", post(rotate_account::<S>))
.route("/api/admin/accounts/enable", post(enable_account::<S>))
.route("/api/admin/accounts/disable", post(disable_account::<S>))
.route("/api/admin/accounts/delete", post(delete_accounts::<S>))
.route("/api/admin/accounts/quota", get(account_quota::<S>))
.route(
"/api/admin/accounts/quota/refresh",
@@ -567,9 +971,17 @@ where
post(refresh_account_models::<S>),
)
.route(
"/api/admin/accounts/test",
"/api/admin/accounts/connection-test",
get(test_account_connection::<S>),
)
.route(
"/api/admin/accounts/oauth/start",
post(start_account_authorization::<S>),
)
.route(
"/api/admin/accounts/oauth/complete",
post(complete_account_authorization::<S>),
)
}
async fn list_accounts<S>(
@@ -593,6 +1005,27 @@ where
Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data)))
}
async fn account_detail<S>(
_auth: AdminAuth,
State(state): State<S>,
Query(query): Query<AccountIdQuery>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let account_id = query.into_id().map_err(map_wire_error)?;
let result = state
.admin_services()
.accounts()
.quota(&account_id, false)
.await
.map_err(map_service_error)?;
let data = AccountQuotaData {
account: account_view(result, Utc::now()),
};
Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data)))
}
async fn export_accounts<S>(
auth: AdminAuth,
State(state): State<S>,
@@ -616,6 +1049,208 @@ where
Ok(response)
}
async fn import_accounts<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<AccountImportRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let (provider, command) = request
.into_command(auth.context().mutation_context())
.map_err(map_wire_error)?;
let result = match provider {
AccountProvider::OpenAi => {
state
.admin_services()
.openai()
.import_document(command)
.await
}
AccountProvider::Xai => state.admin_services().xai().import_document(command).await,
}
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::CREATED,
AdminEnvelope::ok(AccountImportData::from(result)),
))
}
async fn start_account_authorization<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<StartAccountAuthorizationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let (provider, command) = request
.into_command(auth.context().mutation_context())
.map_err(map_wire_error)?;
let result = match provider {
AccountProvider::OpenAi => {
state
.admin_services()
.openai()
.start_authorization(command)
.await
}
AccountProvider::Xai => {
state
.admin_services()
.xai()
.start_authorization(command)
.await
}
}
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::CREATED,
AdminEnvelope::ok(AccountAuthorizationData::from(result)),
))
}
async fn complete_account_authorization<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<CompleteAccountAuthorizationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let (provider, command) = request
.into_command(auth.context().mutation_context())
.map_err(map_wire_error)?;
let result = match provider {
AccountProvider::OpenAi => {
state
.admin_services()
.openai()
.complete_authorization(command)
.await
}
AccountProvider::Xai => {
state
.admin_services()
.xai()
.complete_authorization(command)
.await
}
}
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::CREATED,
AdminEnvelope::ok(AccountMutationData::from(result)),
))
}
async fn rotate_account<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<RotateAccountRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let command = request
.into_command(auth.context().mutation_context())
.map_err(map_wire_error)?;
let result = state
.admin_services()
.openai()
.rotate(command)
.await
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(AccountMutationData::from(result)),
))
}
async fn enable_account<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<AccountMutationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
mutate_account(auth, state, request, AccountMutationAction::Enable).await
}
async fn disable_account<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<AccountMutationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
mutate_account(auth, state, request, AccountMutationAction::Disable).await
}
#[derive(Clone, Copy)]
enum AccountMutationAction {
Enable,
Disable,
}
async fn mutate_account<S>(
auth: AdminAuth,
state: S,
request: AccountMutationRequest,
action: AccountMutationAction,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let (provider, command) = request
.into_command(auth.context().mutation_context())
.map_err(map_wire_error)?;
let result = match (provider, action) {
(AccountProvider::OpenAi, AccountMutationAction::Enable) => {
state.admin_services().openai().enable(command).await
}
(AccountProvider::OpenAi, AccountMutationAction::Disable) => {
state.admin_services().openai().disable(command).await
}
(AccountProvider::Xai, AccountMutationAction::Enable) => {
state.admin_services().xai().enable(command).await
}
(AccountProvider::Xai, AccountMutationAction::Disable) => {
state.admin_services().xai().disable(command).await
}
}
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(AccountMutationData::from(result)),
))
}
async fn delete_accounts<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<AccountDeletionRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let (provider, command) = request
.into_command(auth.context().mutation_context())
.map_err(map_wire_error)?;
let result = match provider {
AccountProvider::OpenAi => state.admin_services().openai().delete(command).await,
AccountProvider::Xai => state.admin_services().xai().delete(command).await,
}
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(AccountDeletionData::from(result)),
))
}
async fn refresh_account<S>(
auth: AdminAuth,
State(state): State<S>,
@@ -796,7 +1431,6 @@ fn account_models_data(result: ProviderModels) -> AccountModelsData {
fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
let AccountDirectoryItem {
account,
provider_instance_name,
status,
usage,
quota,
@@ -810,8 +1444,6 @@ fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
id: account.id.clone(),
name: account.name,
provider: account.provider_kind.to_string(),
provider_instance_id: account.provider_instance_id.to_string(),
provider_instance_name,
resource_ref: account.id,
email: account.email,
account_id: account.upstream_account_id,
@@ -1069,6 +1701,83 @@ fn empty_account_usage() -> AccountUsageView {
}
}
fn revision(value: u64, field: &'static str) -> Result<Revision, WireValidationError> {
Revision::new(value).map_err(|_| WireValidationError::new(field))
}
fn provider_document(
value: Value,
field: &'static str,
) -> Result<ProviderDocument, WireValidationError> {
match value {
Value::Object(document) => Ok(ProviderDocument::new(OpaqueProviderData::new(document))),
_ => Err(WireValidationError::new(field)),
}
}
fn require_wire_id(value: &str, field: &'static str) -> Result<(), WireValidationError> {
require_text(value, MAX_ID_BYTES, field)?;
if value.starts_with("__")
|| !value
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'.' | b':'))
{
return Err(WireValidationError::new(field));
}
Ok(())
}
fn require_text(
value: &str,
max_bytes: usize,
field: &'static str,
) -> Result<(), WireValidationError> {
if value.trim().is_empty() || value.len() > max_bytes || value.chars().any(char::is_control) {
return Err(WireValidationError::new(field));
}
Ok(())
}
fn require_positive(value: u64, field: &'static str) -> Result<(), WireValidationError> {
if value == 0 || i64::try_from(value).is_err() {
return Err(WireValidationError::new(field));
}
Ok(())
}
fn validate_oauth_material(
access_token: &str,
refresh_token: Option<&str>,
) -> Result<(), WireValidationError> {
if access_token.len() > MAX_ACCESS_TOKEN_BYTES
|| !valid_visible_ascii(access_token)
|| !valid_compact_jwt_shape(access_token)
{
return Err(WireValidationError::new("accessToken"));
}
if refresh_token.is_some_and(|token| {
token.len() > MAX_REFRESH_TOKEN_BYTES
|| !valid_visible_ascii(token)
|| token == access_token
}) {
return Err(WireValidationError::new("refreshToken"));
}
Ok(())
}
fn valid_compact_jwt_shape(value: &str) -> bool {
let mut segments = value.split('.');
matches!(
(segments.next(), segments.next(), segments.next(), segments.next()),
(Some(header), Some(payload), Some(signature), None)
if !header.is_empty() && !payload.is_empty() && !signature.is_empty()
)
}
fn valid_visible_ascii(value: &str) -> bool {
!value.is_empty() && value.bytes().all(|byte| (0x21..=0x7e).contains(&byte))
}
fn provider_document_value(document: ProviderDocument) -> Value {
Value::Object(document.into_provider_data().into_inner())
}
backend/crates/gateway-api/src/admin/auth.rs
+2 -2
View File
@@ -207,9 +207,9 @@ where
S: AdminSessionState + Clone + Send + Sync + 'static,
{
Router::new()
.route("/api/admin/login", post(login::<S>))
.route("/api/admin/auth/login", post(login::<S>))
.route("/api/admin/auth/status", get(session_status::<S>))
.route("/api/admin/logout", post(logout::<S>))
.route("/api/admin/auth/logout", post(logout::<S>))
}
async fn login<S>(
backend/crates/gateway-api/src/admin/catalog.rs
-523
View File
@@ -1,523 +0,0 @@
//! Provider instance 管理边界。
use axum::{
Json, Router,
extract::{Query, State},
http::StatusCode,
response::IntoResponse,
routing::{get, post},
};
use chrono::{DateTime, Utc};
use gateway_admin::model::{
PageSize, Revision,
catalog::{
CatalogListQuery as CatalogQuery, CreateProviderInstance, DeleteProviderInstance,
ProviderInstance, ProviderInstanceDetail, ProviderInstanceMutation, ProviderInstancePage,
SetProviderInstanceEnabled, UpdateProviderInstance,
},
};
use gateway_core::routing::{ProviderInstanceId, ProviderKind};
use serde::{Deserialize, Serialize};
use super::{
AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, WireValidationError,
wire::map_admin_service_error,
};
const DEFAULT_PAGE_SIZE: u16 = 50;
const MAX_PAGE_SIZE: u16 = 200;
const MAX_CURSOR_BYTES: usize = 1024;
const MAX_TEXT_BYTES: usize = 2048;
#[derive(Debug, Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct CatalogListQuery {
pub cursor: Option<String>,
pub limit: Option<u16>,
}
impl CatalogListQuery {
pub fn validate(&self) -> Result<(), WireValidationError> {
validate_cursor(self.cursor.as_deref())?;
validate_limit(self.limit)
}
#[must_use]
pub fn page_size(&self) -> u16 {
self.limit.unwrap_or(DEFAULT_PAGE_SIZE)
}
fn into_command(self) -> Result<CatalogQuery, WireValidationError> {
self.validate()?;
let page_size =
PageSize::new(self.page_size()).map_err(|_| WireValidationError::new("limit"))?;
let cursor = self
.cursor
.map(ProviderInstanceId::new)
.transpose()
.map_err(|_| WireValidationError::new("catalogCursor"))?;
Ok(CatalogQuery { cursor, page_size })
}
}
#[derive(Debug, Clone, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct CatalogIdQuery {
pub id: String,
}
impl CatalogIdQuery {
pub fn validate(&self) -> Result<(), WireValidationError> {
require_id(&self.id, "id")
}
fn into_id(self) -> Result<ProviderInstanceId, WireValidationError> {
self.validate()?;
ProviderInstanceId::new(self.id)
.map_err(|_| WireValidationError::new("catalogDetailNotFound"))
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct CatalogMutationData {
pub config_revision: u64,
pub id: String,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct CreateProviderInstanceRequest {
pub id: String,
pub expected_config_revision: u64,
pub provider_kind: String,
pub name: String,
pub base_url: String,
}
impl CreateProviderInstanceRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
require_provider_instance_id(&self.id, "id")?;
require_revision(self.expected_config_revision)?;
require_provider_kind(&self.provider_kind)?;
require_text(&self.name, "name")?;
require_text(&self.base_url, "baseUrl")
}
fn into_command(self) -> Result<CreateProviderInstance, WireValidationError> {
self.validate()?;
Ok(CreateProviderInstance {
expected_config_revision: revision(self.expected_config_revision)?,
id: ProviderInstanceId::new(self.id).map_err(|_| WireValidationError::new("id"))?,
provider_kind: ProviderKind::new(self.provider_kind)
.map_err(|_| WireValidationError::new("providerKind"))?,
name: self.name,
base_url: self.base_url,
})
}
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct UpdateProviderInstanceRequest {
pub id: String,
pub expected_config_revision: u64,
pub name: String,
pub base_url: String,
}
impl UpdateProviderInstanceRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
require_provider_instance_id(&self.id, "id")?;
require_revision(self.expected_config_revision)?;
require_text(&self.name, "name")?;
require_text(&self.base_url, "baseUrl")
}
fn into_command(self) -> Result<UpdateProviderInstance, WireValidationError> {
self.validate()?;
Ok(UpdateProviderInstance {
expected_config_revision: revision(self.expected_config_revision)?,
id: ProviderInstanceId::new(self.id).map_err(|_| WireValidationError::new("id"))?,
name: self.name,
base_url: self.base_url,
})
}
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct CatalogRevisionRequest {
pub id: String,
pub expected_config_revision: u64,
}
impl CatalogRevisionRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
require_id(&self.id, "id")?;
require_revision(self.expected_config_revision)
}
fn into_parts(self) -> Result<(Revision, ProviderInstanceId), WireValidationError> {
self.validate()?;
Ok((
revision(self.expected_config_revision)?,
ProviderInstanceId::new(self.id)
.map_err(|_| WireValidationError::new("catalogMutationNotFound"))?,
))
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ProviderInstanceView {
pub id: String,
pub provider_kind: String,
pub name: String,
pub base_url: String,
pub enabled: bool,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
impl From<ProviderInstance> for ProviderInstanceView {
fn from(instance: ProviderInstance) -> Self {
Self {
id: instance.id.to_string(),
provider_kind: instance.provider_kind.to_string(),
name: instance.name,
base_url: instance.base_url,
enabled: instance.enabled,
created_at: instance.created_at,
updated_at: instance.updated_at,
}
}
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ProviderInstanceListData {
pub config_revision: u64,
pub items: Vec<ProviderInstanceView>,
pub next_cursor: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ProviderInstanceDetailData {
pub config_revision: u64,
pub item: ProviderInstanceView,
}
impl From<ProviderInstancePage> for ProviderInstanceListData {
fn from(page: ProviderInstancePage) -> Self {
Self {
config_revision: page.config_revision.get(),
items: page.items.into_iter().map(Into::into).collect(),
next_cursor: page.next_cursor.map(|cursor| cursor.to_string()),
}
}
}
impl From<ProviderInstanceDetail> for ProviderInstanceDetailData {
fn from(detail: ProviderInstanceDetail) -> Self {
Self {
config_revision: detail.config_revision.get(),
item: detail.item.into(),
}
}
}
pub fn router<S>() -> Router<S>
where
S: AdminSessionState + Clone + Send + Sync + 'static,
{
Router::new()
.route(
"/api/admin/provider-instances",
get(list_provider_instances::<S>),
)
.route(
"/api/admin/provider-instances/detail",
get(provider_instance::<S>),
)
.route(
"/api/admin/provider-instances/create",
post(create_provider_instance::<S>),
)
.route(
"/api/admin/provider-instances/update",
post(update_provider_instance::<S>),
)
.route(
"/api/admin/provider-instances/enable",
post(enable_provider_instance::<S>),
)
.route(
"/api/admin/provider-instances/disable",
post(disable_provider_instance::<S>),
)
.route(
"/api/admin/provider-instances/delete",
post(delete_provider_instance::<S>),
)
}
async fn list_provider_instances<S>(
_auth: AdminAuth,
State(state): State<S>,
Query(query): Query<CatalogListQuery>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let result = state
.admin_services()
.catalog()
.list(query.into_command().map_err(map_wire_error)?)
.await
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(ProviderInstanceListData::from(result)),
))
}
async fn provider_instance<S>(
_auth: AdminAuth,
State(state): State<S>,
Query(query): Query<CatalogIdQuery>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let id = query.into_id().map_err(map_wire_error)?;
let result = state
.admin_services()
.catalog()
.get(&id)
.await
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(ProviderInstanceDetailData::from(result)),
))
}
async fn create_provider_instance<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<CreateProviderInstanceRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let command = request.into_command().map_err(map_wire_error)?;
let id = command.id.to_string();
let result = state
.admin_services()
.catalog()
.create(&auth.context().mutation_context(), command)
.await
.map_err(map_service_error)?;
mutation_response(StatusCode::CREATED, result, id)
}
async fn update_provider_instance<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<UpdateProviderInstanceRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let command = request.into_command().map_err(map_wire_error)?;
let id = command.id.to_string();
let result = state
.admin_services()
.catalog()
.update(&auth.context().mutation_context(), command)
.await
.map_err(map_service_error)?;
mutation_response(StatusCode::OK, result, id)
}
async fn enable_provider_instance<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<CatalogRevisionRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
set_provider_instance_enabled(auth, state, request, true).await
}
async fn disable_provider_instance<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<CatalogRevisionRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
set_provider_instance_enabled(auth, state, request, false).await
}
async fn set_provider_instance_enabled<S>(
auth: AdminAuth,
state: S,
request: CatalogRevisionRequest,
enabled: bool,
) -> Result<AdminResponse<AdminEnvelope<CatalogMutationData>>, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let (expected_config_revision, id) = request.into_parts().map_err(map_wire_error)?;
let response_id = id.to_string();
let result = state
.admin_services()
.catalog()
.set_enabled(
&auth.context().mutation_context(),
SetProviderInstanceEnabled {
expected_config_revision,
id,
enabled,
},
)
.await
.map_err(map_service_error)?;
mutation_response(StatusCode::OK, result, response_id)
}
async fn delete_provider_instance<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<CatalogRevisionRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let (expected_config_revision, id) = request.into_parts().map_err(map_wire_error)?;
let response_id = id.to_string();
let config_revision = state
.admin_services()
.catalog()
.delete(
&auth.context().mutation_context(),
DeleteProviderInstance {
expected_config_revision,
id,
},
)
.await
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(CatalogMutationData {
config_revision: config_revision.get(),
id: response_id,
}),
))
}
fn mutation_response(
status: StatusCode,
result: ProviderInstanceMutation,
fallback_id: String,
) -> Result<AdminResponse<AdminEnvelope<CatalogMutationData>>, AdminError> {
let id = result
.instance
.map_or(fallback_id, |instance| instance.id.to_string());
Ok(AdminResponse::new(
status,
AdminEnvelope::ok(CatalogMutationData {
config_revision: result.config_revision.get(),
id,
}),
))
}
fn require_revision(value: u64) -> Result<(), WireValidationError> {
if value == 0 || i64::try_from(value).is_err() {
return Err(WireValidationError::new("expectedConfigRevision"));
}
Ok(())
}
fn revision(value: u64) -> Result<Revision, WireValidationError> {
require_revision(value)?;
Revision::new(value).map_err(|_| WireValidationError::new("expectedConfigRevision"))
}
fn require_id(value: &str, field: &'static str) -> Result<(), WireValidationError> {
require_text(value, field)?;
if value.starts_with("__")
|| !value
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'.' | b':'))
{
return Err(WireValidationError::new(field));
}
Ok(())
}
fn require_provider_instance_id(
value: &str,
field: &'static str,
) -> Result<(), WireValidationError> {
ProviderInstanceId::new(value).map_or_else(|_| Err(WireValidationError::new(field)), |_| Ok(()))
}
fn require_text(value: &str, field: &'static str) -> Result<(), WireValidationError> {
if value.trim().is_empty()
|| value.len() > MAX_TEXT_BYTES
|| value.chars().any(char::is_control)
{
return Err(WireValidationError::new(field));
}
Ok(())
}
fn require_provider_kind(value: &str) -> Result<(), WireValidationError> {
if matches!(value, "openai" | "xai") {
Ok(())
} else {
Err(WireValidationError::new("providerKind"))
}
}
fn validate_cursor(cursor: Option<&str>) -> Result<(), WireValidationError> {
if cursor.is_some_and(|value| value.is_empty() || value.len() > MAX_CURSOR_BYTES) {
Err(WireValidationError::new("cursor"))
} else {
Ok(())
}
}
fn validate_limit(limit: Option<u16>) -> Result<(), WireValidationError> {
if limit.is_some_and(|value| value == 0 || value > MAX_PAGE_SIZE) {
Err(WireValidationError::new("limit"))
} else {
Ok(())
}
}
fn map_wire_error(error: WireValidationError) -> AdminError {
match error.field() {
"catalogCursor" => AdminError::bad_request("Invalid catalog cursor"),
"catalogDetailNotFound" => AdminError::not_found("Provider instance was not found"),
"catalogMutationNotFound" => AdminError::not_found("provider instance was not found"),
field => AdminError::bad_request(format!("Invalid field: {field}")),
}
}
fn map_service_error(error: gateway_admin::model::AdminError) -> AdminError {
if error.kind() == gateway_admin::model::AdminErrorKind::Invalid
&& error.to_string() == "Invalid provider instance cursor"
{
return AdminError::bad_request("Invalid catalog cursor");
}
map_admin_service_error(error, "Catalog unavailable")
}
backend/crates/gateway-api/src/admin/client_keys.rs
+3 -2
View File
@@ -665,9 +665,10 @@ where
S: AdminSessionState + Clone + Send + Sync + 'static,
{
Router::new()
.route("/api/admin/client-keys", get(list_client_keys::<S>))
.route(
"/api/admin/client-keys",
get(list_client_keys::<S>).post(create_client_key::<S>),
"/api/admin/client-keys/create",
post(create_client_key::<S>),
)
.route("/api/admin/client-keys/reveal", get(reveal_client_key::<S>))
.route(
backend/crates/gateway-api/src/admin/mod.rs
-6
View File
@@ -9,14 +9,11 @@ use axum::{
pub mod accounts;
pub mod auth;
pub mod catalog;
pub mod client_keys;
pub mod observability;
pub mod openai;
pub mod settings;
pub mod system;
pub mod wire;
pub mod xai;
pub use auth::{AdminAuth, AdminSessionState};
pub use wire::{
@@ -32,13 +29,10 @@ where
Router::new()
.merge(accounts::router::<S>())
.merge(auth::router::<S>())
.merge(catalog::router::<S>())
.merge(client_keys::router::<S>())
.merge(openai::router::<S>())
.merge(observability::router::<S>())
.merge(settings::router::<S>())
.merge(system::router::<S>())
.merge(xai::router::<S>())
.layer(middleware::map_response(no_store))
}
backend/crates/gateway-api/src/admin/observability.rs
+3 -9
View File
@@ -642,7 +642,6 @@ pub struct UsageAttemptView {
pub attempt_index: u32,
pub trigger: String,
pub provider: String,
pub provider_instance_id: String,
pub model: String,
pub transport: String,
pub send_state: String,
@@ -1183,7 +1182,6 @@ pub struct OpsErrorMetadataView {
pub source: String,
pub component: String,
pub attempt_id: Option<String>,
pub provider_instance_id: Option<String>,
pub account_label: Option<String>,
}
@@ -1204,14 +1202,14 @@ where
get(usage_records_summary::<S>),
)
.route(
"/api/admin/usage/records/insights/overview",
"/api/admin/usage/insights/overview",
get(usage_insights_overview::<S>),
)
.route(
"/api/admin/usage/records/insights/diagnostics",
"/api/admin/usage/insights/diagnostics",
get(usage_insights_diagnostics::<S>),
)
.route("/api/admin/ops/errors", get(ops_errors::<S>))
.route("/api/admin/operations/errors", get(ops_errors::<S>))
}
async fn dashboard_summary<S>(
@@ -1744,9 +1742,6 @@ fn usage_attempt_view(attempt: domain::UsageAttempt) -> UsageAttemptView {
provider: attempt
.provider_kind
.unwrap_or_else(|| "unknown".to_owned()),
provider_instance_id: attempt
.provider_instance_id
.unwrap_or_else(|| "unknown".to_owned()),
model: attempt
.upstream_model_id
.unwrap_or_else(|| "unknown".to_owned()),
@@ -1843,7 +1838,6 @@ fn ops_error_view(error: domain::OpsError) -> OpsErrorView {
source: error.source,
component: error.component,
attempt_id: None,
provider_instance_id: error.provider_instance_id,
account_label: None,
},
created_at: error.occurred_at,
backend/crates/gateway-api/src/admin/openai.rs
-895
View File
@@ -1,895 +0,0 @@
//! Codex OAuth credential 管理端 wire contract。
use axum::{
Json, Router,
extract::{Query, State},
http::StatusCode,
response::IntoResponse,
routing::{get, post},
};
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use chrono::{DateTime, Utc};
use gateway_admin::model::{
AdminError as AdminServiceError, AdminErrorKind, PageSize, Revision,
accounts::{AccountAvailability, AccountRecord},
provider_credentials::{
AuthorizationStarted, CompleteAuthorization, CredentialAvailabilityFilter,
CredentialCursor, CredentialDetails, CredentialImportResult, CredentialListQuery,
CredentialListWindow, CredentialMutation, CredentialMutationResult, CredentialPage,
ImportCredentials, ProviderDocument, ReauthorizationTarget, RotateCredential,
StartAuthorization,
},
};
use gateway_core::{
engine::credential::{OpaqueProviderData, ProviderAccountId},
routing::ProviderInstanceId,
};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::admin::{
AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, WireValidationError,
};
pub const DEFAULT_PAGE_SIZE: u16 = 50;
pub const MAX_PAGE_SIZE: u16 = 200;
const MAX_ID_BYTES: usize = 256;
const MAX_NAME_BYTES: usize = 512;
const MAX_CURSOR_BYTES: usize = 1_024;
const MAX_IMPORT_DOCUMENT_BYTES: usize = 64 * 1_024 * 1_024;
const MAX_ACCESS_TOKEN_BYTES: usize = 16 * 1_024;
const MAX_REFRESH_TOKEN_BYTES: usize = 64 * 1_024;
const MAX_CALLBACK_URL_BYTES: usize = 16 * 1_024;
/// Codex credential 列表查询。
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ListCredentialsQuery {
pub provider_instance_id: Option<String>,
pub availability: Option<String>,
pub enabled: Option<bool>,
pub cursor: Option<String>,
pub limit: Option<u16>,
}
impl ListCredentialsQuery {
/// 校验筛选与游标边界。
pub fn validate(&self) -> Result<(), WireValidationError> {
if let Some(provider_instance_id) = self.provider_instance_id.as_deref() {
require_id(provider_instance_id, "providerInstanceId")?;
}
if self.availability.as_deref().is_some_and(|value| {
!matches!(
value,
"unknown" | "ready" | "cooldown" | "exhausted" | "invalid"
)
}) {
return Err(WireValidationError::new("availability"));
}
if self
.cursor
.as_deref()
.is_some_and(|value| value.is_empty() || value.len() > MAX_CURSOR_BYTES)
{
return Err(WireValidationError::new("cursor"));
}
if self
.limit
.is_some_and(|limit| limit == 0 || limit > MAX_PAGE_SIZE)
{
return Err(WireValidationError::new("limit"));
}
Ok(())
}
fn into_command(self) -> Result<CredentialListQuery, WireValidationError> {
self.validate()?;
Ok(CredentialListQuery {
provider_instance_id: self
.provider_instance_id
.map(ProviderInstanceId::new)
.transpose()
.map_err(|_| WireValidationError::new("providerInstanceId"))?,
availability: self
.availability
.as_deref()
.map(parse_availability)
.transpose()?,
enabled: self.enabled,
window: CredentialListWindow::Page {
cursor: self.cursor.as_deref().map(decode_cursor).transpose()?,
page_size: PageSize::new(self.limit.unwrap_or(DEFAULT_PAGE_SIZE))
.map_err(|_| WireValidationError::new("limit"))?,
},
})
}
}
/// Codex credential 详情查询。
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct CredentialDetailsQuery {
pub credential_id: String,
}
impl CredentialDetailsQuery {
/// 校验 credential ID。
pub fn validate(&self) -> Result<(), WireValidationError> {
require_id(&self.credential_id, "credentialId")
}
fn into_account_id(self) -> Result<ProviderAccountId, WireValidationError> {
self.validate()?;
ProviderAccountId::new(self.credential_id)
.map_err(|_| WireValidationError::new("credentialId"))
}
}
/// OpenAI Provider-owned 账号导入文档。
#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ImportCredentialsDocumentRequest {
pub expected_config_revision: u64,
pub provider_instance_id: String,
pub document: Value,
}
impl ImportCredentialsDocumentRequest {
/// Wire 只约束文档大小；结构识别与 token 验证完全归 Provider。
pub fn validate(&self) -> Result<(), WireValidationError> {
require_positive(self.expected_config_revision, "expectedConfigRevision")?;
require_id(&self.provider_instance_id, "providerInstanceId")?;
if !self.document.is_object()
|| serde_json::to_vec(&self.document)
.map_or(true, |encoded| encoded.len() > MAX_IMPORT_DOCUMENT_BYTES)
{
return Err(WireValidationError::new("document"));
}
Ok(())
}
fn into_command(
self,
context: gateway_admin::model::MutationContext,
) -> Result<ImportCredentials, WireValidationError> {
self.validate()?;
Ok(ImportCredentials {
context,
expected_config_revision: revision(
self.expected_config_revision,
"expectedConfigRevision",
)?,
provider_instance_id: ProviderInstanceId::new(self.provider_instance_id)
.map_err(|_| WireValidationError::new("providerInstanceId"))?,
document: provider_document(self.document, "document")?,
})
}
}
/// 启动 Codex Authorization Code + PKCE；可创建新账号或绑定既有账号重新授权。
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct StartOAuthAuthorizationRequest {
pub expected_config_revision: u64,
pub provider_instance_id: String,
pub name: String,
pub credential_id: Option<String>,
pub expected_credential_revision: Option<u64>,
}
impl StartOAuthAuthorizationRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
require_positive(self.expected_config_revision, "expectedConfigRevision")?;
require_id(&self.provider_instance_id, "providerInstanceId")?;
require_text(&self.name, MAX_NAME_BYTES, "name")?;
match (
self.credential_id.as_deref(),
self.expected_credential_revision,
) {
(None, None) => {}
(Some(credential_id), Some(revision)) => {
require_id(credential_id, "credentialId")?;
require_positive(revision, "expectedCredentialRevision")?;
}
_ => return Err(WireValidationError::new("reauthorization")),
}
Ok(())
}
fn into_command(
self,
context: gateway_admin::model::MutationContext,
) -> Result<StartAuthorization, WireValidationError> {
self.validate()?;
let reauthorization = match (self.credential_id, self.expected_credential_revision) {
(None, None) => None,
(Some(account_id), Some(credential_revision)) => Some(ReauthorizationTarget {
account_id: ProviderAccountId::new(account_id)
.map_err(|_| WireValidationError::new("credentialId"))?,
credential_revision: revision(credential_revision, "expectedCredentialRevision")?,
}),
_ => return Err(WireValidationError::new("reauthorization")),
};
Ok(StartAuthorization {
context,
expected_config_revision: revision(
self.expected_config_revision,
"expectedConfigRevision",
)?,
provider_instance_id: ProviderInstanceId::new(self.provider_instance_id)
.map_err(|_| WireValidationError::new("providerInstanceId"))?,
name: self.name,
reauthorization,
})
}
}
/// 完成 Codex Authorization Code flow；flow ID 与 callback 都放在 POST body。
#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct CompleteOAuthAuthorizationRequest {
pub flow_id: String,
pub callback_url: String,
}
impl CompleteOAuthAuthorizationRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
if !URL_SAFE_NO_PAD
.decode(&self.flow_id)
.is_ok_and(|decoded| decoded.len() == 32)
{
return Err(WireValidationError::new("flowId"));
}
if self.callback_url.is_empty()
|| self.callback_url.len() > MAX_CALLBACK_URL_BYTES
|| self.callback_url.chars().any(char::is_control)
{
return Err(WireValidationError::new("callbackUrl"));
}
Ok(())
}
fn into_command(
self,
context: gateway_admin::model::MutationContext,
) -> Result<CompleteAuthorization, WireValidationError> {
self.validate()?;
Ok(CompleteAuthorization {
context,
flow_id: self.flow_id,
callback_url: self.callback_url,
})
}
}
/// Codex OAuth token 轮换请求。
#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct RotateCredentialRequest {
pub credential_id: String,
pub expected_config_revision: u64,
pub expected_credential_revision: u64,
pub access_token: String,
pub refresh_token: Option<String>,
}
impl RotateCredentialRequest {
/// 校验 ID、revision 与 OAuth material 的 wire 形状。
pub fn validate(&self) -> Result<(), WireValidationError> {
require_id(&self.credential_id, "credentialId")?;
require_positive(self.expected_config_revision, "expectedConfigRevision")?;
require_positive(
self.expected_credential_revision,
"expectedCredentialRevision",
)?;
validate_oauth_material(&self.access_token, self.refresh_token.as_deref())
}
fn into_command(
self,
context: gateway_admin::model::MutationContext,
) -> Result<RotateCredential, WireValidationError> {
self.validate()?;
let mut material = Map::new();
material.insert("access_token".to_owned(), Value::String(self.access_token));
material.insert(
"refresh_token".to_owned(),
self.refresh_token.map_or(Value::Null, Value::String),
);
Ok(RotateCredential {
mutation: CredentialMutation {
context,
expected_config_revision: revision(
self.expected_config_revision,
"expectedConfigRevision",
)?,
account_id: ProviderAccountId::new(self.credential_id)
.map_err(|_| WireValidationError::new("credentialId"))?,
},
expected_credential_revision: revision(
self.expected_credential_revision,
"expectedCredentialRevision",
)?,
provider_material: ProviderDocument::new(OpaqueProviderData::new(material)),
})
}
}
/// Codex credential 生命周期 mutation 请求。
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct CredentialMutationRequest {
pub credential_id: String,
pub expected_config_revision: u64,
}
impl CredentialMutationRequest {
/// 校验 ID 与配置 revision。
pub fn validate(&self) -> Result<(), WireValidationError> {
require_id(&self.credential_id, "credentialId")?;
require_positive(self.expected_config_revision, "expectedConfigRevision")
}
fn into_command(
self,
context: gateway_admin::model::MutationContext,
) -> Result<CredentialMutation, WireValidationError> {
self.validate()?;
Ok(CredentialMutation {
context,
expected_config_revision: revision(
self.expected_config_revision,
"expectedConfigRevision",
)?,
account_id: ProviderAccountId::new(self.credential_id)
.map_err(|_| WireValidationError::new("credentialId"))?,
})
}
}
/// 列表游标的 JSON 形状；HTTP 组合层负责 base64url 编解码。
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct CredentialCursorWire {
pub created_at: DateTime<Utc>,
pub credential_id: String,
}
/// Admin API 可安全返回的 Codex credential 视图。
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct CodexCredentialView {
pub id: String,
pub provider_instance_id: String,
pub name: String,
pub email: Option<String>,
pub upstream_user_id: String,
pub upstream_account_id: Option<String>,
pub enabled: bool,
pub credential_revision: i64,
pub has_refresh_token: bool,
pub availability: String,
pub availability_reason: Option<String>,
pub plan_type: Option<String>,
pub access_token_expires_at: DateTime<Utc>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub cooldown_until: Option<DateTime<Utc>>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
/// Codex credential 列表响应数据。
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct CodexCredentialListData {
pub config_revision: u64,
pub items: Vec<CodexCredentialView>,
pub next_cursor: Option<String>,
}
/// Codex credential 详情响应数据。
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct CodexCredentialDetailsData {
pub config_revision: u64,
pub credential: CodexCredentialView,
}
/// 已提交的结构配置 mutation 响应数据。
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct CodexCredentialMutationData {
pub config_revision: u64,
pub credential_id: String,
}
/// Provider 正式文档一次原子导入的结果。
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct CodexCredentialsDocumentImportData {
pub config_revision: u64,
pub credential_ids: Vec<String>,
}
/// Authorization Code + PKCE 启动响应。
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct CodexOAuthAuthorizationStartedData {
pub flow_id: String,
pub authorization_url: String,
pub expires_at: DateTime<Utc>,
}
/// 已提交的 token 轮换响应数据。
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct CodexCredentialRotationData {
pub credential_id: String,
pub credential_revision: i64,
}
impl From<AccountRecord> for CodexCredentialView {
fn from(account: AccountRecord) -> Self {
Self {
id: account.id,
provider_instance_id: account.provider_instance_id.to_string(),
name: account.name,
email: account.email,
upstream_user_id: account.upstream_user_id,
upstream_account_id: account.upstream_account_id,
enabled: account.enabled,
credential_revision: i64::try_from(account.credential_revision.get())
.unwrap_or(i64::MAX),
has_refresh_token: account.has_refresh_token,
availability: codex_availability(account.availability).to_owned(),
availability_reason: account.availability_reason,
plan_type: account.plan_type,
access_token_expires_at: account.access_token_expires_at,
next_refresh_at: account.next_refresh_at,
cooldown_until: account.cooldown_until,
created_at: account.created_at,
updated_at: account.updated_at,
}
}
}
impl From<CredentialDetails> for CodexCredentialDetailsData {
fn from(details: CredentialDetails) -> Self {
Self {
config_revision: details.config_revision.get(),
credential: details.credential.into(),
}
}
}
impl From<CredentialImportResult> for CodexCredentialsDocumentImportData {
fn from(result: CredentialImportResult) -> Self {
Self {
config_revision: result.config_revision.get(),
credential_ids: result
.credential_ids
.into_iter()
.map(|id| id.to_string())
.collect(),
}
}
}
impl From<AuthorizationStarted> for CodexOAuthAuthorizationStartedData {
fn from(started: AuthorizationStarted) -> Self {
Self {
flow_id: started.flow_id,
authorization_url: started.authorization_url,
expires_at: started.expires_at,
}
}
}
impl From<CredentialMutationResult> for CodexCredentialMutationData {
fn from(result: CredentialMutationResult) -> Self {
Self {
config_revision: result.config_revision.get(),
credential_id: result.account_id.to_string(),
}
}
}
fn validate_oauth_material(
access_token: &str,
refresh_token: Option<&str>,
) -> Result<(), WireValidationError> {
if access_token.len() > MAX_ACCESS_TOKEN_BYTES
|| !valid_visible_ascii(access_token)
|| !valid_compact_jwt_shape(access_token)
{
return Err(WireValidationError::new("accessToken"));
}
if refresh_token.is_some_and(|token| {
token.len() > MAX_REFRESH_TOKEN_BYTES
|| !valid_visible_ascii(token)
|| token == access_token
}) {
return Err(WireValidationError::new("refreshToken"));
}
Ok(())
}
fn valid_compact_jwt_shape(value: &str) -> bool {
let mut segments = value.split('.');
matches!(
(segments.next(), segments.next(), segments.next(), segments.next()),
(Some(header), Some(payload), Some(signature), None)
if !header.is_empty() && !payload.is_empty() && !signature.is_empty()
)
}
fn valid_visible_ascii(value: &str) -> bool {
!value.is_empty() && value.bytes().all(|byte| (0x21..=0x7e).contains(&byte))
}
fn require_id(value: &str, field: &'static str) -> Result<(), WireValidationError> {
require_text(value, MAX_ID_BYTES, field)?;
if value.starts_with("__") {
return Err(WireValidationError::new(field));
}
Ok(())
}
fn require_text(
value: &str,
max_bytes: usize,
field: &'static str,
) -> Result<(), WireValidationError> {
if value.trim().is_empty() || value.len() > max_bytes || value.chars().any(char::is_control) {
return Err(WireValidationError::new(field));
}
Ok(())
}
fn require_positive(value: u64, field: &'static str) -> Result<(), WireValidationError> {
if value == 0 || i64::try_from(value).is_err() {
return Err(WireValidationError::new(field));
}
Ok(())
}
/// 构造固定 GET/POST Codex credential 管理路由。
pub fn router<S>() -> Router<S>
where
S: AdminSessionState + Clone + Send + Sync + 'static,
{
Router::new()
.route("/api/admin/openai/credentials", get(list_credentials::<S>))
.route("/api/admin/openai/credential", get(credential_details::<S>))
.route(
"/api/admin/openai/credentials/import-document",
post(import_credentials_document::<S>),
)
.route(
"/api/admin/openai/oauth/authorization/start",
post(start_oauth_authorization::<S>),
)
.route(
"/api/admin/openai/oauth/authorization/complete",
post(complete_oauth_authorization::<S>),
)
.route(
"/api/admin/openai/credentials/rotate",
post(rotate_credential::<S>),
)
.route(
"/api/admin/openai/credentials/enable",
post(enable_credential::<S>),
)
.route(
"/api/admin/openai/credentials/disable",
post(disable_credential::<S>),
)
.route(
"/api/admin/openai/credentials/delete",
post(delete_credential::<S>),
)
}
async fn list_credentials<S>(
_auth: AdminAuth,
State(state): State<S>,
Query(query): Query<ListCredentialsQuery>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let command = query.into_command().map_err(map_wire_error)?;
let page = state
.admin_services()
.openai()
.list(command)
.await
.map_err(map_service_error)?;
ok(StatusCode::OK, credential_list_data(page)?)
}
async fn credential_details<S>(
_auth: AdminAuth,
State(state): State<S>,
Query(query): Query<CredentialDetailsQuery>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let account_id = query.into_account_id().map_err(map_wire_error)?;
let details = state
.admin_services()
.openai()
.details(&account_id)
.await
.map_err(map_service_error)?;
ok(StatusCode::OK, CodexCredentialDetailsData::from(details))
}
async fn import_credentials_document<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<ImportCredentialsDocumentRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let command = request
.into_command(auth.context().mutation_context())
.map_err(map_wire_error)?;
let result = state
.admin_services()
.openai()
.import_document(command)
.await
.map_err(map_service_error)?;
ok(
StatusCode::CREATED,
CodexCredentialsDocumentImportData::from(result),
)
}
async fn start_oauth_authorization<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<StartOAuthAuthorizationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let command = request
.into_command(auth.context().mutation_context())
.map_err(map_wire_error)?;
let result = state
.admin_services()
.openai()
.start_authorization(command)
.await
.map_err(map_service_error)?;
ok(
StatusCode::CREATED,
CodexOAuthAuthorizationStartedData::from(result),
)
}
async fn complete_oauth_authorization<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<CompleteOAuthAuthorizationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let command = request
.into_command(auth.context().mutation_context())
.map_err(map_wire_error)?;
let result = state
.admin_services()
.openai()
.complete_authorization(command)
.await
.map_err(map_service_error)?;
ok(
StatusCode::CREATED,
CodexCredentialMutationData::from(result),
)
}
async fn rotate_credential<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<RotateCredentialRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let command = request
.into_command(auth.context().mutation_context())
.map_err(map_wire_error)?;
let result = state
.admin_services()
.openai()
.rotate(command)
.await
.map_err(map_service_error)?;
ok(StatusCode::OK, credential_rotation_data(result)?)
}
async fn enable_credential<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<CredentialMutationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
mutate_credential(auth, state, request, CredentialAction::Enable).await
}
async fn disable_credential<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<CredentialMutationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
mutate_credential(auth, state, request, CredentialAction::Disable).await
}
async fn delete_credential<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<CredentialMutationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
mutate_credential(auth, state, request, CredentialAction::Delete).await
}
#[derive(Clone, Copy)]
enum CredentialAction {
Enable,
Disable,
Delete,
}
async fn mutate_credential<S>(
auth: AdminAuth,
state: S,
request: CredentialMutationRequest,
action: CredentialAction,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let command = request
.into_command(auth.context().mutation_context())
.map_err(map_wire_error)?;
let service = state.admin_services().openai();
let result = match action {
CredentialAction::Enable => service.enable(command).await,
CredentialAction::Disable => service.disable(command).await,
CredentialAction::Delete => service.delete(command).await,
}
.map_err(map_service_error)?;
ok(StatusCode::OK, CodexCredentialMutationData::from(result))
}
fn ok<T: Serialize>(
status: StatusCode,
data: T,
) -> Result<AdminResponse<AdminEnvelope<T>>, AdminError> {
Ok(AdminResponse::new(status, AdminEnvelope::ok(data)))
}
fn map_wire_error(error: WireValidationError) -> AdminError {
AdminError::bad_request(format!("Invalid Codex credential field: {}", error.field()))
}
fn map_service_error(error: AdminServiceError) -> AdminError {
match error.kind() {
AdminErrorKind::Invalid => AdminError::bad_request(error.to_string()),
AdminErrorKind::Unauthorized => AdminError::admin_session_required(),
AdminErrorKind::NotFound => AdminError::not_found(error.to_string()),
AdminErrorKind::Conflict => AdminError::conflict(error.to_string()),
AdminErrorKind::RateLimited => AdminError::too_many_login_attempts(),
AdminErrorKind::BadGateway => AdminError::bad_gateway(error.to_string()),
AdminErrorKind::Unavailable => AdminError::service_unavailable(error.to_string()),
AdminErrorKind::Internal => AdminError::internal(error.to_string()),
}
}
fn revision(value: u64, field: &'static str) -> Result<Revision, WireValidationError> {
Revision::new(value).map_err(|_| WireValidationError::new(field))
}
fn provider_document(
value: Value,
field: &'static str,
) -> Result<ProviderDocument, WireValidationError> {
match value {
Value::Object(document) => Ok(ProviderDocument::new(OpaqueProviderData::new(document))),
_ => Err(WireValidationError::new(field)),
}
}
fn parse_availability(value: &str) -> Result<CredentialAvailabilityFilter, WireValidationError> {
match value {
"unknown" => Ok(CredentialAvailabilityFilter::Exact(
AccountAvailability::Unknown,
)),
"ready" => Ok(CredentialAvailabilityFilter::Exact(
AccountAvailability::Ready,
)),
"cooldown" => Ok(CredentialAvailabilityFilter::Exact(
AccountAvailability::Cooldown,
)),
"exhausted" => Ok(CredentialAvailabilityFilter::Exact(
AccountAvailability::QuotaExhausted,
)),
"invalid" => Ok(CredentialAvailabilityFilter::AnyOf(vec![
AccountAvailability::Expired,
AccountAvailability::Banned,
AccountAvailability::Invalid,
])),
_ => Err(WireValidationError::new("availability")),
}
}
fn decode_cursor(value: &str) -> Result<CredentialCursor, WireValidationError> {
let encoded = URL_SAFE_NO_PAD
.decode(value)
.map_err(|_| WireValidationError::new("cursor"))?;
let wire: CredentialCursorWire =
serde_json::from_slice(&encoded).map_err(|_| WireValidationError::new("cursor"))?;
require_id(&wire.credential_id, "cursor")?;
Ok(CredentialCursor {
created_at: wire.created_at,
account_id: ProviderAccountId::new(wire.credential_id)
.map_err(|_| WireValidationError::new("cursor"))?,
})
}
fn encode_cursor(cursor: &CredentialCursor) -> Result<String, AdminError> {
let encoded = serde_json::to_vec(&CredentialCursorWire {
created_at: cursor.created_at,
credential_id: cursor.account_id.to_string(),
})
.map_err(|_| AdminError::internal("Failed to encode credential cursor"))?;
Ok(URL_SAFE_NO_PAD.encode(encoded))
}
fn credential_list_data(page: CredentialPage) -> Result<CodexCredentialListData, AdminError> {
Ok(CodexCredentialListData {
config_revision: page.config_revision.get(),
items: page.items.into_iter().map(Into::into).collect(),
next_cursor: page.next_cursor.as_ref().map(encode_cursor).transpose()?,
})
}
fn credential_rotation_data(
result: CredentialMutationResult,
) -> Result<CodexCredentialRotationData, AdminError> {
let credential_revision = result
.credential_revision
.ok_or_else(|| AdminError::internal("Credential revision is missing"))?;
Ok(CodexCredentialRotationData {
credential_id: result.account_id.to_string(),
credential_revision: i64::try_from(credential_revision.get()).unwrap_or(i64::MAX),
})
}
const fn codex_availability(value: AccountAvailability) -> &'static str {
match value {
AccountAvailability::Unknown => "unknown",
AccountAvailability::Ready => "ready",
AccountAvailability::Cooldown => "cooldown",
AccountAvailability::QuotaExhausted => "exhausted",
AccountAvailability::Expired
| AccountAvailability::Banned
| AccountAvailability::Invalid => "invalid",
}
}
backend/crates/gateway-api/src/admin/settings.rs
+7 -5
View File
@@ -174,13 +174,15 @@ where
S: AdminSessionState + Clone + Send + Sync + 'static,
{
Router::new()
.route(
"/api/admin/settings",
get(settings::<S>).post(update_settings::<S>),
)
.route("/api/admin/settings", get(settings::<S>))
.route("/api/admin/settings/update", post(update_settings::<S>))
.route(
"/api/admin/settings/admin-api-key",
get(admin_api_key_status::<S>).post(delete_admin_api_key::<S>),
get(admin_api_key_status::<S>),
)
.route(
"/api/admin/settings/admin-api-key/delete",
post(delete_admin_api_key::<S>),
)
.route(
"/api/admin/settings/admin-api-key/regenerate",
backend/crates/gateway-api/src/admin/system.rs
+3 -3
View File
@@ -216,13 +216,13 @@ where
{
Router::new()
.route("/api/admin/system/version", get(version::<S>))
.route("/api/admin/system/update-detail", get(update_detail::<S>))
.route("/api/admin/system/update/detail", get(update_detail::<S>))
.route(
"/api/admin/system/update-events",
"/api/admin/system/update/events",
get(update_event_stream::<S>),
)
.route("/api/admin/system/update", post(perform_update::<S>))
.route("/api/admin/system/update-status", get(update_status::<S>))
.route("/api/admin/system/update/status", get(update_status::<S>))
.route("/api/admin/system/rollback", post(rollback::<S>))
.route("/api/admin/system/restart", post(restart::<S>))
}
backend/crates/gateway-api/src/admin/xai.rs
-597
View File
@@ -1,597 +0,0 @@
//! xAI 管理端 wire contract。
use axum::{
Json, Router,
extract::{Query, State},
http::StatusCode,
response::IntoResponse,
routing::{get, post},
};
use chrono::{DateTime, Utc};
use gateway_admin::model::{
AdminError as AdminServiceError, AdminErrorKind, Revision,
accounts::{AccountAvailability, AccountRecord},
provider_credentials::{
AuthorizationStarted, CompleteAuthorization, CredentialImportResult, CredentialListQuery,
CredentialListWindow, CredentialMutation, CredentialMutationResult, CredentialPage,
ImportCredentials, ProviderDocument, StartAuthorization,
},
};
use gateway_core::{
engine::credential::{OpaqueProviderData, ProviderAccountId},
routing::ProviderInstanceId,
};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::admin::{
AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, WireValidationError,
};
const MAX_IMPORT_DOCUMENT_BYTES: usize = 64 * 1024 * 1024;
/// xAI Provider-owned 正式导入文档；API 不解释文档内部字段。
#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct XaiCredentialImportDocumentRequest {
pub expected_config_revision: u64,
pub provider_instance_id: String,
pub document: Value,
}
impl XaiCredentialImportDocumentRequest {
/// 只校验公共 revision、instance 和请求体边界；Provider 负责全部格式语义。
pub fn validate(&self) -> Result<(), WireValidationError> {
require_positive(self.expected_config_revision, "expectedConfigRevision")?;
require_id(&self.provider_instance_id, "providerInstanceId")?;
if !self.document.is_object()
|| serde_json::to_vec(&self.document)
.map_or(true, |encoded| encoded.len() > MAX_IMPORT_DOCUMENT_BYTES)
{
return Err(WireValidationError::new("document"));
}
Ok(())
}
fn into_command(
self,
context: gateway_admin::model::MutationContext,
) -> Result<ImportCredentials, WireValidationError> {
self.validate()?;
Ok(ImportCredentials {
context,
expected_config_revision: revision(
self.expected_config_revision,
"expectedConfigRevision",
)?,
provider_instance_id: ProviderInstanceId::new(self.provider_instance_id)
.map_err(|_| WireValidationError::new("providerInstanceId"))?,
document: provider_document(self.document, "document")?,
})
}
}
/// 批量导入成功响应；不包含任何 token、subject、邮箱或 client ID。
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct XaiCredentialImportData {
pub config_revision: u64,
pub imported_count: usize,
pub credential_ids: Vec<String>,
}
impl XaiCredentialImportData {
#[must_use]
pub fn new(config_revision: u64, credential_ids: Vec<String>) -> Self {
Self {
config_revision,
imported_count: credential_ids.len(),
credential_ids,
}
}
}
const MAX_ID_BYTES: usize = 256;
const MAX_NAME_BYTES: usize = 512;
const MAX_CALLBACK_FIELD_BYTES: usize = 64 * 1024;
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ListCredentialsQuery {
pub provider_instance_id: Option<String>,
}
impl ListCredentialsQuery {
pub fn validate(&self) -> Result<(), WireValidationError> {
if let Some(id) = self.provider_instance_id.as_deref() {
require_id(id, "providerInstanceId")?;
}
Ok(())
}
fn into_command(self) -> Result<CredentialListQuery, WireValidationError> {
self.validate()?;
Ok(CredentialListQuery {
provider_instance_id: self
.provider_instance_id
.map(ProviderInstanceId::new)
.transpose()
.map_err(|_| WireValidationError::new("providerInstanceId"))?,
availability: None,
enabled: None,
window: CredentialListWindow::All,
})
}
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct NewCredentialRequest {
pub expected_config_revision: u64,
pub provider_instance_id: String,
pub name: String,
}
impl NewCredentialRequest {
fn validate(&self) -> Result<(), WireValidationError> {
require_positive(self.expected_config_revision, "expectedConfigRevision")?;
require_id(&self.provider_instance_id, "providerInstanceId")?;
require_text(&self.name, MAX_NAME_BYTES, "name")?;
Ok(())
}
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct StartAuthorizationRequest {
#[serde(flatten)]
pub credential: NewCredentialRequest,
}
impl StartAuthorizationRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
self.credential.validate()
}
fn into_command(
self,
context: gateway_admin::model::MutationContext,
) -> Result<StartAuthorization, WireValidationError> {
self.validate()?;
Ok(StartAuthorization {
context,
expected_config_revision: revision(
self.credential.expected_config_revision,
"expectedConfigRevision",
)?,
provider_instance_id: ProviderInstanceId::new(self.credential.provider_instance_id)
.map_err(|_| WireValidationError::new("providerInstanceId"))?,
name: self.credential.name,
reauthorization: None,
})
}
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct CompleteAuthorizationRequest {
pub flow_id: String,
pub callback_url: String,
}
impl CompleteAuthorizationRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
require_id(&self.flow_id, "flowId")?;
require_text(&self.callback_url, MAX_CALLBACK_FIELD_BYTES, "callbackUrl")
}
fn into_command(
self,
context: gateway_admin::model::MutationContext,
) -> Result<CompleteAuthorization, WireValidationError> {
self.validate()?;
Ok(CompleteAuthorization {
context,
flow_id: self.flow_id,
callback_url: self.callback_url,
})
}
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct CredentialMutationRequest {
pub credential_id: String,
pub expected_config_revision: u64,
}
impl CredentialMutationRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
require_id(&self.credential_id, "credentialId")?;
require_positive(self.expected_config_revision, "expectedConfigRevision")
}
fn into_command(
self,
context: gateway_admin::model::MutationContext,
) -> Result<CredentialMutation, WireValidationError> {
self.validate()?;
Ok(CredentialMutation {
context,
expected_config_revision: revision(
self.expected_config_revision,
"expectedConfigRevision",
)?,
account_id: ProviderAccountId::new(self.credential_id)
.map_err(|_| WireValidationError::new("credentialId"))?,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AuthorizationStartData {
pub flow_id: String,
pub authorization_url: String,
pub expires_at: DateTime<Utc>,
}
/// xAI OAuth credential 的安全管理视图。
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct XaiCredentialViewData {
pub id: String,
pub provider_instance_id: String,
pub name: String,
pub email: Option<String>,
pub upstream_user_id: String,
pub upstream_account_id: Option<String>,
pub plan_type: Option<String>,
pub enabled: bool,
pub credential_revision: i64,
pub has_refresh_token: bool,
pub availability: String,
pub availability_reason: Option<String>,
pub access_token_expires_at: DateTime<Utc>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub cooldown_until: Option<DateTime<Utc>>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
/// xAI OAuth credential 列表响应。
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct XaiCredentialListData {
pub config_revision: u64,
pub items: Vec<XaiCredentialViewData>,
}
/// Credential 生命周期 mutation 响应。
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct XaiCredentialMutationData {
pub config_revision: u64,
pub credential_id: String,
}
impl From<AccountRecord> for XaiCredentialViewData {
fn from(account: AccountRecord) -> Self {
Self {
id: account.id,
provider_instance_id: account.provider_instance_id.to_string(),
name: account.name,
email: account.email,
upstream_user_id: account.upstream_user_id,
upstream_account_id: account.upstream_account_id,
plan_type: account.plan_type,
enabled: account.enabled,
credential_revision: i64::try_from(account.credential_revision.get())
.unwrap_or(i64::MAX),
has_refresh_token: account.has_refresh_token,
availability: availability_name(account.availability).to_owned(),
availability_reason: account.availability_reason,
access_token_expires_at: account.access_token_expires_at,
next_refresh_at: account.next_refresh_at,
cooldown_until: account.cooldown_until,
created_at: account.created_at,
updated_at: account.updated_at,
}
}
}
impl From<CredentialPage> for XaiCredentialListData {
fn from(page: CredentialPage) -> Self {
Self {
config_revision: page.config_revision.get(),
items: page.items.into_iter().map(Into::into).collect(),
}
}
}
impl From<CredentialImportResult> for XaiCredentialImportData {
fn from(result: CredentialImportResult) -> Self {
Self::new(
result.config_revision.get(),
result
.credential_ids
.into_iter()
.map(|id| id.to_string())
.collect(),
)
}
}
impl From<AuthorizationStarted> for AuthorizationStartData {
fn from(started: AuthorizationStarted) -> Self {
Self {
flow_id: started.flow_id,
authorization_url: started.authorization_url,
expires_at: started.expires_at,
}
}
}
impl From<CredentialMutationResult> for XaiCredentialMutationData {
fn from(result: CredentialMutationResult) -> Self {
Self {
config_revision: result.config_revision.get(),
credential_id: result.account_id.to_string(),
}
}
}
/// 构造固定 GET/POST xAI OAuth 管理路由。
pub fn router<S>() -> Router<S>
where
S: AdminSessionState + Clone + Send + Sync + 'static,
{
Router::new()
.route("/api/admin/xai/credentials", get(list_credentials::<S>))
.route(
"/api/admin/xai/credentials/import-document",
post(import_credentials_document::<S>),
)
.route(
"/api/admin/xai/oauth/authorization/start",
post(start_authorization_code_flow::<S>),
)
.route(
"/api/admin/xai/oauth/authorization/complete",
post(complete_authorization_code_flow::<S>),
)
.route(
"/api/admin/xai/credentials/disable",
post(disable_credential::<S>),
)
.route(
"/api/admin/xai/credentials/enable",
post(enable_credential::<S>),
)
.route(
"/api/admin/xai/credentials/delete",
post(delete_credential::<S>),
)
}
async fn list_credentials<S>(
_auth: AdminAuth,
State(state): State<S>,
Query(query): Query<ListCredentialsQuery>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let command = query.into_command().map_err(map_wire_error)?;
let page = state
.admin_services()
.xai()
.list(command)
.await
.map_err(map_service_error)?;
ok(StatusCode::OK, XaiCredentialListData::from(page))
}
async fn import_credentials_document<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<XaiCredentialImportDocumentRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let command = request
.into_command(auth.context().mutation_context())
.map_err(map_wire_error)?;
let result = state
.admin_services()
.xai()
.import_document(command)
.await
.map_err(map_service_error)?;
ok(StatusCode::CREATED, XaiCredentialImportData::from(result))
}
async fn start_authorization_code_flow<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<StartAuthorizationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let command = request
.into_command(auth.context().mutation_context())
.map_err(map_wire_error)?;
let started = state
.admin_services()
.xai()
.start_authorization(command)
.await
.map_err(map_service_error)?;
require_id(&started.flow_id, "flowId").map_err(map_wire_error)?;
ok(StatusCode::CREATED, AuthorizationStartData::from(started))
}
async fn complete_authorization_code_flow<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<CompleteAuthorizationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let command = request
.into_command(auth.context().mutation_context())
.map_err(map_wire_error)?;
let result = state
.admin_services()
.xai()
.complete_authorization(command)
.await
.map_err(map_service_error)?;
ok(StatusCode::CREATED, XaiCredentialMutationData::from(result))
}
async fn disable_credential<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<CredentialMutationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
mutate_credential(auth, state, request, CredentialAction::Disable).await
}
async fn enable_credential<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<CredentialMutationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
mutate_credential(auth, state, request, CredentialAction::Enable).await
}
async fn delete_credential<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<CredentialMutationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
mutate_credential(auth, state, request, CredentialAction::Delete).await
}
#[derive(Clone, Copy)]
enum CredentialAction {
Disable,
Enable,
Delete,
}
async fn mutate_credential<S>(
auth: AdminAuth,
state: S,
request: CredentialMutationRequest,
action: CredentialAction,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let command = request
.into_command(auth.context().mutation_context())
.map_err(map_wire_error)?;
let service = state.admin_services().xai();
let result = match action {
CredentialAction::Disable => service.disable(command).await,
CredentialAction::Enable => service.enable(command).await,
CredentialAction::Delete => service.delete(command).await,
}
.map_err(map_service_error)?;
ok(StatusCode::OK, XaiCredentialMutationData::from(result))
}
fn ok<T: Serialize>(
status: StatusCode,
data: T,
) -> Result<AdminResponse<AdminEnvelope<T>>, AdminError> {
Ok(AdminResponse::new(status, AdminEnvelope::ok(data)))
}
fn map_wire_error(error: WireValidationError) -> AdminError {
AdminError::bad_request(format!("Invalid xAI OAuth field: {}", error.field()))
}
fn map_service_error(error: AdminServiceError) -> AdminError {
match error.kind() {
AdminErrorKind::Invalid => AdminError::bad_request(error.to_string()),
AdminErrorKind::Unauthorized => AdminError::admin_session_required(),
AdminErrorKind::NotFound => AdminError::not_found(error.to_string()),
AdminErrorKind::Conflict => AdminError::conflict(error.to_string()),
AdminErrorKind::RateLimited => AdminError::too_many_login_attempts(),
AdminErrorKind::BadGateway | AdminErrorKind::Unavailable => {
AdminError::bad_gateway(error.to_string())
}
AdminErrorKind::Internal => AdminError::internal(error.to_string()),
}
}
fn revision(value: u64, field: &'static str) -> Result<Revision, WireValidationError> {
Revision::new(value).map_err(|_| WireValidationError::new(field))
}
fn provider_document(
value: Value,
field: &'static str,
) -> Result<ProviderDocument, WireValidationError> {
match value {
Value::Object(document) => Ok(ProviderDocument::new(OpaqueProviderData::new(document))),
_ => Err(WireValidationError::new(field)),
}
}
const fn availability_name(value: AccountAvailability) -> &'static str {
match value {
AccountAvailability::Unknown => "unknown",
AccountAvailability::Ready => "ready",
AccountAvailability::Cooldown => "cooldown",
AccountAvailability::QuotaExhausted => "quota_exhausted",
AccountAvailability::Expired => "expired",
AccountAvailability::Banned => "banned",
AccountAvailability::Invalid => "invalid",
}
}
fn require_id(value: &str, field: &'static str) -> Result<(), WireValidationError> {
require_text(value, MAX_ID_BYTES, field)?;
if value.starts_with("__")
|| !value
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'.' | b':'))
{
return Err(WireValidationError::new(field));
}
Ok(())
}
fn require_text(
value: &str,
max_bytes: usize,
field: &'static str,
) -> Result<(), WireValidationError> {
if value.trim().is_empty() || value.len() > max_bytes || value.chars().any(char::is_control) {
return Err(WireValidationError::new(field));
}
Ok(())
}
fn require_positive(value: u64, field: &'static str) -> Result<(), WireValidationError> {
if value == 0 || i64::try_from(value).is_err() {
return Err(WireValidationError::new(field));
}
Ok(())
}
backend/crates/gateway-api/src/openai/error.rs
+2 -2
View File
@@ -77,7 +77,7 @@ pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError {
{
GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no upstream Provider instance or account is currently available",
"no upstream provider is currently available for this request",
)
}
EngineError::Provider(error) => GatewayError::from_provider(error),
@@ -90,7 +90,7 @@ pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError {
EngineError::ProviderNotRegistered { .. } | EngineError::EmptyRoutingPlan => {
GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no upstream Provider instance or account is currently available",
"no upstream provider is currently available for this request",
)
}
EngineError::Store(_)
backend/crates/gateway-api/tests/admin/accounts.rs
+98 -13
View File
@@ -102,19 +102,23 @@ mod response {
}
mod actions {
use gateway_admin::model::accounts::{
AccountConnectionTestEvent as DomainConnectionTestEvent, AccountStatus,
use gateway_admin::model::{
Revision,
accounts::{AccountConnectionTestEvent as DomainConnectionTestEvent, AccountStatus},
provider_credentials::{CredentialDeletionResult, CredentialImportResult},
};
use gateway_api::admin::accounts::{
AccountActionRequest, AccountConnectionTestEvent, AccountExportData, AccountExportQuery,
AccountIdQuery, AccountRefreshRequest, AccountTestQuery,
AccountActionRequest, AccountConnectionTestEvent, AccountDeletionData,
AccountDeletionRequest, AccountExportData, AccountExportQuery, AccountIdQuery,
AccountImportData, AccountImportRequest, AccountRefreshRequest, AccountTestQuery,
};
use gateway_core::engine::credential::ProviderAccountId;
use serde_json::json;
#[test]
fn export_should_require_explicit_unique_ids_and_confirmation() {
let valid: AccountExportQuery = serde_json::from_value(json!({
"ids": "acct_1,acct_2",
"accountIds": "acct_1,acct_2",
"confirm": "export_sensitive_accounts"
}))
.expect("decode export query");
@@ -129,9 +133,9 @@ mod actions {
);
for query in [
json!({ "ids": "", "confirm": "export_sensitive_accounts" }),
json!({ "ids": "acct_1,acct_1", "confirm": "export_sensitive_accounts" }),
json!({ "ids": "acct_1", "confirm": "yes" }),
json!({ "accountIds": "", "confirm": "export_sensitive_accounts" }),
json!({ "accountIds": "acct_1,acct_1", "confirm": "export_sensitive_accounts" }),
json!({ "accountIds": "acct_1", "confirm": "yes" }),
] {
assert!(
serde_json::from_value::<AccountExportQuery>(query)
@@ -145,13 +149,13 @@ mod actions {
#[test]
fn account_actions_should_require_frozen_account_ids_and_revision() {
let id: AccountIdQuery =
serde_json::from_value(json!({ "id": "acct_1" })).expect("decode ID query");
serde_json::from_value(json!({ "accountId": "acct_1" })).expect("decode ID query");
assert!(id.validate().is_ok());
let action: AccountActionRequest =
serde_json::from_value(json!({ "id": "legacy-id" })).expect("decode action");
assert_eq!(action.validate().unwrap_err().field(), "id");
serde_json::from_value(json!({ "accountId": "legacy-id" })).expect("decode action");
assert_eq!(action.validate().unwrap_err().field(), "accountId");
let refresh: AccountRefreshRequest = serde_json::from_value(json!({
"id": "acct_1",
"accountId": "acct_1",
"expectedConfigRevision": 0
}))
.expect("decode refresh");
@@ -161,10 +165,91 @@ mod actions {
);
}
#[test]
fn account_import_should_use_provider_and_opaque_data_fields() {
let valid: AccountImportRequest = serde_json::from_value(json!({
"provider": "openai",
"expectedConfigRevision": 7,
"data": {
"providerOwnedUnknownField": {"nested": [1, 2, 3]},
"accounts": [{"credentials": {"access_token": "provider-validates-this"}}]
}
}))
.expect("decode account import");
assert!(valid.validate().is_ok());
let invalid: AccountImportRequest = serde_json::from_value(json!({
"provider": "xai",
"expectedConfigRevision": 7,
"data": []
}))
.expect("decode invalid account import");
assert_eq!(invalid.validate().unwrap_err().field(), "data");
assert!(
serde_json::from_value::<AccountImportRequest>(json!({
"provider": "openai",
"expectedConfigRevision": 7,
"document": {}
}))
.is_err()
);
}
#[test]
fn account_deletion_should_validate_one_provider_batch_and_emit_account_ids() {
let request: AccountDeletionRequest = serde_json::from_value(json!({
"provider": "xai",
"accountIds": ["acct_1", "acct_2"],
"expectedConfigRevision": 8
}))
.expect("decode account deletion");
assert!(request.validate().is_ok());
let duplicate: AccountDeletionRequest = serde_json::from_value(json!({
"provider": "xai",
"accountIds": ["acct_1", "acct_1"],
"expectedConfigRevision": 8
}))
.expect("decode duplicate account deletion");
assert_eq!(duplicate.validate().unwrap_err().field(), "accountIds");
let response = AccountDeletionData::from(CredentialDeletionResult {
config_revision: Revision::new(9).expect("revision"),
account_ids: vec![
ProviderAccountId::new("acct_1").expect("account ID"),
ProviderAccountId::new("acct_2").expect("account ID"),
],
});
assert_eq!(
serde_json::to_value(response).expect("serialize account deletion"),
json!({
"configRevision": 9,
"deletedCount": 2,
"accountIds": ["acct_1", "acct_2"]
})
);
}
#[test]
fn account_import_response_should_emit_account_ids() {
let response = AccountImportData::from(CredentialImportResult {
config_revision: Revision::new(8).expect("revision"),
credential_ids: vec![ProviderAccountId::new("acct_imported").expect("account ID")],
});
assert_eq!(
serde_json::to_value(response).expect("serialize account import"),
json!({
"configRevision": 8,
"importedCount": 1,
"accountIds": ["acct_imported"]
})
);
}
#[test]
fn connection_test_should_require_model_in_query() {
let query: AccountTestQuery = serde_json::from_value(json!({
"id": "acct_1",
"accountId": "acct_1",
"modelId": " "
}))
.expect("decode connection test query");
backend/crates/gateway-api/tests/admin/catalog.rs
-128
View File
@@ -1,128 +0,0 @@
use chrono::Utc;
use gateway_api::admin::catalog::{
CatalogListQuery, CreateProviderInstanceRequest, ProviderInstanceView,
};
use serde_json::json;
fn provider_request(provider_kind: &str) -> CreateProviderInstanceRequest {
serde_json::from_value(json!({
"id": format!("inst_{provider_kind}"),
"expectedConfigRevision": 1,
"providerKind": provider_kind,
"name": provider_kind,
"baseUrl": "https://example.invalid"
}))
.expect("decode provider instance")
}
#[test]
fn provider_instance_should_accept_openai_platform() {
assert!(provider_request("openai").validate().is_ok());
}
#[test]
fn provider_instance_should_accept_xai_platform() {
assert!(provider_request("xai").validate().is_ok());
}
#[test]
fn provider_instance_should_reject_removed_scheduling_fields() {
let result = serde_json::from_value::<CreateProviderInstanceRequest>(json!({
"id": "inst_openai",
"expectedConfigRevision": 1,
"providerKind": "openai",
"name": "OpenAI",
"baseUrl": "https://chatgpt.com",
"maxConcurrency": 3
}));
assert!(result.is_err());
}
#[test]
fn provider_instance_should_reject_removed_model_routing_fields() {
let result = serde_json::from_value::<CreateProviderInstanceRequest>(json!({
"id": "inst_openai",
"expectedConfigRevision": 1,
"providerKind": "openai",
"name": "OpenAI",
"baseUrl": "https://chatgpt.com",
"upstreamModelId": "gpt-5.4"
}));
assert!(result.is_err());
}
#[test]
fn provider_instance_should_reject_unknown_provider_kind() {
let request = provider_request("unknown");
assert_eq!(request.validate().unwrap_err().field(), "providerKind");
}
#[test]
fn provider_instance_should_reject_id_without_core_prefix() {
let request: CreateProviderInstanceRequest = serde_json::from_value(json!({
"id": "codex-official",
"expectedConfigRevision": 1,
"providerKind": "openai",
"name": "Codex",
"baseUrl": "https://chatgpt.com/backend-api"
}))
.expect("decode provider instance");
assert_eq!(request.validate().unwrap_err().field(), "id");
}
#[test]
fn provider_instance_should_require_positive_revision() {
let mut request = provider_request("openai");
request.expected_config_revision = 0;
assert_eq!(
request.validate().unwrap_err().field(),
"expectedConfigRevision"
);
}
#[test]
fn catalog_page_size_should_remain_bounded() {
let query = CatalogListQuery {
cursor: None,
limit: Some(201),
};
assert_eq!(query.validate().unwrap_err().field(), "limit");
}
#[test]
fn provider_instance_view_should_not_emit_provider_config_json() {
let value = serde_json::to_value(ProviderInstanceView {
id: "inst_openai".into(),
provider_kind: "openai".into(),
name: "Codex".into(),
base_url: "https://chatgpt.com".into(),
enabled: true,
created_at: Utc::now(),
updated_at: Utc::now(),
})
.expect("serialize provider instance");
assert!(value.get("configJson").is_none());
}
#[test]
fn provider_instance_view_should_emit_platform_explicitly() {
let value = serde_json::to_value(ProviderInstanceView {
id: "inst_xai".into(),
provider_kind: "xai".into(),
name: "xAI".into(),
base_url: "https://grok.com".into(),
enabled: true,
created_at: Utc::now(),
updated_at: Utc::now(),
})
.expect("serialize provider instance");
assert_eq!(value["providerKind"], "xai");
}
backend/crates/gateway-api/tests/admin/mod.rs
+5 -61
View File
@@ -14,15 +14,10 @@ use gateway_admin::{
model::{
MutationContext, Revision,
accounts::{
AccountListQuery, AccountPage, AccountRecord, AccountUsage, DeleteAccount,
AccountListQuery, AccountPage, AccountRecord, AccountUsage, DeleteAccounts,
SetAccountEnabled,
},
auth::{AdminAuditEvent, AdminSession},
catalog::{
CreateProviderInstance, DeleteProviderInstance, ProviderInstanceCatalog,
ProviderInstanceDetail, ProviderInstanceMutation, SetProviderInstanceEnabled,
UpdateProviderInstance,
},
client_keys::{
ClientKeyListQuery, ClientKeyPage, ClientKeyRecord, ClientKeySecret, DeleteClientKey,
NewClientKey, SetClientKeyEnabled, UpdateClientKey,
@@ -51,7 +46,7 @@ use gateway_admin::{
provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind},
store::{
AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts, AdminStoreResult,
AuthStore, CatalogStore, ClientKeyStore, ObservabilityStore, SettingsStore,
AuthStore, ClientKeyStore, ObservabilityStore, SettingsStore,
},
system::{
SystemOperationError, SystemOperationErrorKind, SystemOperations,
@@ -68,21 +63,17 @@ use gateway_core::{
error::GatewayError,
policy::{ClientApiKeyId, RateLimits},
routing::{
ConfigRevision, ProviderInstanceId, ProviderKind, PublicModelId, UpstreamModelId,
snapshot::SnapshotControl,
ConfigRevision, ProviderKind, PublicModelId, UpstreamModelId, snapshot::SnapshotControl,
},
};
mod accounts;
mod auth;
mod catalog;
mod client_keys;
mod observability;
mod openai;
mod settings;
mod system;
mod wire;
mod xai;
pub(super) struct AdminTestFixture {
pub services: AdminServices,
@@ -104,7 +95,6 @@ impl AdminTestFixture {
let stores = AdminStorePorts::new(
unused.clone(),
auth.clone(),
unused.clone(),
client_keys.clone(),
unused,
settings.clone(),
@@ -515,9 +505,9 @@ impl AccountStore for UnusedStore {
Err(unavailable("account enabled"))
}
async fn delete_account(
async fn delete_accounts(
&self,
_: DeleteAccount,
_: DeleteAccounts,
_: &MutationContext,
) -> AdminStoreResult<Revision> {
Err(unavailable("account delete"))
@@ -532,52 +522,6 @@ impl AccountStore for UnusedStore {
}
}
#[async_trait]
impl CatalogStore for UnusedStore {
async fn list_provider_instances(&self, _: bool) -> AdminStoreResult<ProviderInstanceCatalog> {
Err(unavailable("provider catalog"))
}
async fn load_provider_instance(
&self,
_: &ProviderInstanceId,
) -> AdminStoreResult<Option<ProviderInstanceDetail>> {
Err(unavailable("provider instance"))
}
async fn create_provider_instance(
&self,
_: CreateProviderInstance,
_: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation> {
Err(unavailable("provider create"))
}
async fn update_provider_instance(
&self,
_: UpdateProviderInstance,
_: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation> {
Err(unavailable("provider update"))
}
async fn set_provider_instance_enabled(
&self,
_: SetProviderInstanceEnabled,
_: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation> {
Err(unavailable("provider enabled"))
}
async fn delete_provider_instance(
&self,
_: DeleteProviderInstance,
_: &MutationContext,
) -> AdminStoreResult<Revision> {
Err(unavailable("provider delete"))
}
}
#[async_trait]
impl ObservabilityStore for UnusedStore {
async fn dashboard_summary(&self, _: TimeRange) -> AdminStoreResult<DashboardObservation> {
backend/crates/gateway-api/tests/admin/observability.rs
-1
View File
@@ -289,7 +289,6 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
client_transport: "http_sse".to_owned(),
requested_model_id: "grok-4.5".to_owned(),
input_token_estimate: 1,
provider_instance_id: None,
provider_kind: Some("xai".to_owned()),
provider_account_ref: None,
provider_account_name: None,
backend/crates/gateway-api/tests/admin/openai.rs
-202
View File
@@ -1,202 +0,0 @@
use chrono::{TimeZone, Utc};
use gateway_admin::model::{
Revision,
accounts::{AccountAvailability, AccountRecord},
};
use gateway_api::admin::openai::{
CodexCredentialView, CompleteOAuthAuthorizationRequest, CredentialDetailsQuery,
CredentialMutationRequest, ImportCredentialsDocumentRequest, ListCredentialsQuery,
RotateCredentialRequest, StartOAuthAuthorizationRequest,
};
use gateway_core::routing::{ProviderInstanceId, ProviderKind};
fn compact_jwt() -> String {
"header.payload.signature".to_owned()
}
#[test]
fn document_import_accepts_provider_owned_object() {
let request: ImportCredentialsDocumentRequest = serde_json::from_value(serde_json::json!({
"expectedConfigRevision": 7,
"providerInstanceId": "inst_openai",
"document": {"accounts": [{"token": compact_jwt()}]}
}))
.expect("deserialize Provider import document");
assert_eq!(request.validate(), Ok(()));
}
#[test]
fn document_import_rejects_unknown_outer_fields() {
let result = serde_json::from_value::<ImportCredentialsDocumentRequest>(serde_json::json!({
"expectedConfigRevision": 7,
"providerInstanceId": "inst_openai",
"document": {"accounts": []},
"maxConcurrency": 2
}));
assert!(result.is_err());
}
#[test]
fn document_import_rejects_non_object_document() {
let array: ImportCredentialsDocumentRequest = serde_json::from_value(serde_json::json!({
"expectedConfigRevision": 7,
"providerInstanceId": "inst_openai",
"document": []
}))
.expect("deserialize array document");
assert_eq!(array.validate().unwrap_err().field(), "document");
}
#[test]
fn document_import_keeps_provider_document_opaque_to_api() {
let request: ImportCredentialsDocumentRequest = serde_json::from_value(serde_json::json!({
"expectedConfigRevision": 7,
"providerInstanceId": "inst_openai",
"document": {
"providerOwnedUnknownField": {"nested": [1, 2, 3]},
"accounts": [{"credentials": {"access_token": "provider-validates-this"}}]
}
}))
.expect("deserialize opaque Provider document");
assert_eq!(request.validate(), Ok(()));
}
#[test]
fn list_query_should_validate_status_pagination_and_reserved_ids() {
let valid: ListCredentialsQuery = serde_json::from_value(serde_json::json!({
"providerInstanceId": "inst_openai",
"availability": "cooldown",
"enabled": false,
"cursor": "cursor",
"limit": 200
}))
.expect("deserialize list query");
let invalid_status: ListCredentialsQuery = serde_json::from_value(serde_json::json!({
"availability": "refreshing"
}))
.expect("deserialize invalid status query");
let invalid_limit: ListCredentialsQuery = serde_json::from_value(serde_json::json!({
"limit": 201
}))
.expect("deserialize invalid page size query");
assert_eq!(valid.validate(), Ok(()));
assert_eq!(
invalid_status.validate().unwrap_err().field(),
"availability"
);
assert_eq!(invalid_limit.validate().unwrap_err().field(), "limit");
}
#[test]
fn detail_and_mutation_should_reject_missing_invalid_or_overflowing_identity() {
let detail = CredentialDetailsQuery {
credential_id: "__reserved__".to_owned(),
};
let mutation = CredentialMutationRequest {
credential_id: "cred_codex".to_owned(),
expected_config_revision: 0,
};
let rotation: RotateCredentialRequest = serde_json::from_value(serde_json::json!({
"credentialId": "cred_codex",
"expectedConfigRevision": 1,
"expectedCredentialRevision": u64::MAX,
"accessToken": compact_jwt()
}))
.expect("deserialize overflowing revision");
assert_eq!(detail.validate().unwrap_err().field(), "credentialId");
assert_eq!(
mutation.validate().unwrap_err().field(),
"expectedConfigRevision"
);
assert_eq!(
rotation.validate().unwrap_err().field(),
"expectedCredentialRevision"
);
}
#[test]
fn oauth_wire_should_keep_flow_id_and_callback_in_post_body_only() {
let start: StartOAuthAuthorizationRequest = serde_json::from_value(serde_json::json!({
"expectedConfigRevision": 7,
"providerInstanceId": "inst_openai",
"name": "OAuth browser flow",
"credentialId": "acct_existing",
"expectedCredentialRevision": 9
}))
.expect("start wire");
let complete: CompleteOAuthAuthorizationRequest = serde_json::from_value(serde_json::json!({
"flowId": "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
"callbackUrl": "http://localhost:1455/auth/callback?code=code&state=state"
}))
.expect("complete wire");
let client_state =
serde_json::from_value::<StartOAuthAuthorizationRequest>(serde_json::json!({
"expectedConfigRevision": 7,
"providerInstanceId": "inst_openai",
"name": "OAuth browser flow",
"state": "client-controlled"
}))
.expect_err("state must remain server-side");
assert_eq!(start.validate(), Ok(()));
assert_eq!(complete.validate(), Ok(()));
assert!(client_state.to_string().contains("unknown field `state`"));
}
#[test]
fn oauth_reauthorization_requires_account_and_revision_together() {
let missing_revision: StartOAuthAuthorizationRequest =
serde_json::from_value(serde_json::json!({
"expectedConfigRevision": 7,
"providerInstanceId": "inst_openai",
"name": "OAuth browser flow",
"credentialId": "acct_existing"
}))
.expect("deserialize incomplete reauthorization");
assert_eq!(
missing_revision.validate().unwrap_err().field(),
"reauthorization"
);
}
#[test]
fn credential_result_should_map_to_the_existing_safe_wire_view() {
let timestamp = Utc
.with_ymd_and_hms(2026, 7, 18, 3, 0, 0)
.single()
.expect("valid fixture timestamp");
let view = CodexCredentialView::from(AccountRecord {
id: "acct_openai_safe".to_owned(),
provider_instance_id: ProviderInstanceId::new("inst_openai").expect("provider instance ID"),
provider_kind: ProviderKind::new("openai").expect("provider kind"),
name: "Codex OAuth".to_owned(),
email: Some("verified@example.com".to_owned()),
upstream_user_id: "subject_openai".to_owned(),
upstream_account_id: None,
plan_type: Some("pro".to_owned()),
credential_revision: Revision::new(2).expect("credential revision"),
has_refresh_token: true,
access_token_expires_at: timestamp,
next_refresh_at: Some(timestamp),
enabled: true,
availability: AccountAvailability::Banned,
availability_reason: Some("account disabled".to_owned()),
cooldown_until: None,
availability_observed_at: timestamp,
quota_observed_at: None,
created_at: timestamp,
updated_at: timestamp,
});
let rendered = serde_json::to_value(view).expect("serialize safe credential view");
assert_eq!(rendered["availability"], "invalid");
assert_eq!(rendered["credentialRevision"], 2);
assert!(rendered.get("accessToken").is_none());
}
backend/crates/gateway-api/tests/admin/settings.rs
+2 -2
View File
@@ -117,7 +117,7 @@ async fn settings_post_should_replace_provider_scoped_model_mappings() {
let response = app(fixture.state())
.oneshot(request(
Method::POST,
"/api/admin/settings",
"/api/admin/settings/update",
Some(update_body()),
))
.await
@@ -196,7 +196,7 @@ async fn admin_key_delete_should_use_fixed_post_path() {
let response = app(fixture.state())
.oneshot(request(
Method::POST,
"/api/admin/settings/admin-api-key",
"/api/admin/settings/admin-api-key/delete",
None,
))
.await
backend/crates/gateway-api/tests/admin/xai.rs
-208
View File
@@ -1,208 +0,0 @@
use chrono::{TimeZone, Utc};
use gateway_admin::model::{
Revision,
accounts::{AccountAvailability, AccountRecord},
};
use gateway_api::admin::xai::{
AuthorizationStartData, CompleteAuthorizationRequest, NewCredentialRequest,
StartAuthorizationRequest, XaiCredentialImportData, XaiCredentialImportDocumentRequest,
XaiCredentialListData, XaiCredentialViewData,
};
use gateway_core::routing::{ProviderInstanceId, ProviderKind};
use serde_json::{Value, json};
fn valid_request() -> Value {
json!({
"expectedConfigRevision": 7,
"providerInstanceId": "inst_xai",
"document": {
"type": "oauth-account-bundle",
"version": 1,
"exported_at": "2026-07-18T02:47:01Z",
"proxies": [],
"accounts": [{
"name": "Grok OAuth",
"platform": "grok",
"type": "oauth",
"credentials": {
"access_token": "access-token-secret",
"refresh_token": "refresh-token-secret",
"token_type": "Bearer",
"expires_at": "2026-07-18T06:47:01Z",
"email": "ignored@example.com",
"base_url": "https://cli-chat-proxy.grok.com/v1",
"id_token": "header.payload.signature",
"client_id": "official-client",
"scope": "openid offline_access grok-cli:access api:access"
},
"extra": {"email": "also-ignored@example.com"},
"concurrency": 1,
"priority": 0
}]
}
})
}
#[test]
fn import_document_accepts_opaque_provider_owned_shape() {
let request: XaiCredentialImportDocumentRequest =
serde_json::from_value(valid_request()).expect("deserialize Provider document");
request
.validate()
.expect("validate Provider document boundary");
}
#[test]
fn import_document_keeps_inner_schema_owned_by_provider() {
let mut request = valid_request();
request["document"]["accounts"][0]["credentials"]["provider_future_field"] =
json!({"nested": true});
let request = serde_json::from_value::<XaiCredentialImportDocumentRequest>(request)
.expect("API must not decode Provider fields");
assert_eq!(request.validate(), Ok(()));
}
#[test]
fn import_document_rejects_unknown_outer_field() {
let mut request = valid_request();
request["schemaVersion"] = json!(1);
let error = serde_json::from_value::<XaiCredentialImportDocumentRequest>(request)
.err()
.expect("unknown API field must fail");
assert!(error.to_string().contains("unknown field `schemaVersion`"));
}
#[test]
fn import_document_rejects_non_object_document() {
let mut request = valid_request();
request["document"] = json!([]);
let request: XaiCredentialImportDocumentRequest =
serde_json::from_value(request).expect("deserialize generic JSON");
assert_eq!(request.validate().unwrap_err().field(), "document");
}
#[test]
fn import_response_should_never_contain_identity_or_secret_fields() {
let response = XaiCredentialImportData::new(8, vec!["cred_xai_safe".to_owned()]);
let rendered = serde_json::to_string(&response).expect("serialize response");
assert_eq!(
rendered,
r#"{"configRevision":8,"importedCount":1,"credentialIds":["cred_xai_safe"]}"#
);
}
#[test]
fn credential_list_wire_should_expose_only_safe_view_fields() {
let timestamp = Utc
.with_ymd_and_hms(2026, 7, 18, 3, 0, 0)
.single()
.expect("valid fixture timestamp");
let response = XaiCredentialListData {
config_revision: 8,
items: vec![XaiCredentialViewData {
id: "cred_xai_safe".to_owned(),
provider_instance_id: "inst_xai".to_owned(),
name: "Grok OAuth".to_owned(),
email: Some("verified@example.com".to_owned()),
upstream_user_id: "subject_xai".to_owned(),
upstream_account_id: None,
plan_type: Some("pro".to_owned()),
enabled: true,
credential_revision: 2,
has_refresh_token: true,
availability: "ready".to_owned(),
availability_reason: None,
access_token_expires_at: timestamp,
next_refresh_at: Some(timestamp),
cooldown_until: None,
created_at: timestamp,
updated_at: timestamp,
}],
};
let rendered = serde_json::to_string(&response).expect("serialize typed list response");
assert!(rendered.contains("credentialRevision"));
assert!(rendered.contains("accessTokenExpiresAt"));
for forbidden in [
"\"accessToken\":",
"\"refreshToken\":",
"\"idToken\":",
"\"subject\":",
] {
assert!(!rendered.contains(forbidden), "wire leaked {forbidden}");
}
}
#[test]
fn authorization_start_wire_contains_only_browser_flow_fields() {
let response = AuthorizationStartData {
flow_id: "flow_xai_safe".to_owned(),
authorization_url: "https://auth.x.ai/oauth2/auth?state=redacted".to_owned(),
expires_at: Utc::now(),
};
let rendered = serde_json::to_string(&response).expect("serialize authorization response");
assert!(rendered.contains("authorizationUrl"));
assert!(!rendered.contains("codeVerifier"));
assert!(!rendered.contains("nonce"));
}
#[test]
fn authorization_requests_require_server_owned_flow_and_full_callback_url() {
let start = StartAuthorizationRequest {
credential: NewCredentialRequest {
expected_config_revision: 7,
provider_instance_id: "inst_xai".to_owned(),
name: "xAI OAuth".to_owned(),
},
};
start.validate().expect("authorization start is valid");
let complete: CompleteAuthorizationRequest = serde_json::from_value(json!({
"flowId": "flow_xai_safe",
"callbackUrl": "http://127.0.0.1:56121/callback?code=redacted&state=redacted"
}))
.expect("deserialize callback completion");
complete.validate().expect("callback completion is valid");
}
#[test]
fn credential_result_should_map_to_the_existing_safe_wire_view() {
let timestamp = Utc
.with_ymd_and_hms(2026, 7, 18, 3, 0, 0)
.single()
.expect("valid fixture timestamp");
let view = XaiCredentialViewData::from(AccountRecord {
id: "acct_xai_safe".to_owned(),
provider_instance_id: ProviderInstanceId::new("inst_xai").expect("provider instance ID"),
provider_kind: ProviderKind::new("xai").expect("provider kind"),
name: "Grok OAuth".to_owned(),
email: Some("verified@example.com".to_owned()),
upstream_user_id: "subject_xai".to_owned(),
upstream_account_id: None,
plan_type: Some("pro".to_owned()),
credential_revision: Revision::new(2).expect("credential revision"),
has_refresh_token: true,
access_token_expires_at: timestamp,
next_refresh_at: Some(timestamp),
enabled: true,
availability: AccountAvailability::QuotaExhausted,
availability_reason: None,
cooldown_until: None,
availability_observed_at: timestamp,
quota_observed_at: None,
created_at: timestamp,
updated_at: timestamp,
});
let rendered = serde_json::to_value(view).expect("serialize safe credential view");
assert_eq!(rendered["availability"], "quota_exhausted");
assert_eq!(rendered["credentialRevision"], 2);
assert!(rendered.get("refreshToken").is_none());
}
backend/crates/gateway-api/tests/architecture.rs
-6
View File
@@ -35,15 +35,12 @@ fn source_tree_should_match_frozen_machine_manifest() {
let mut expected = vec![
"src/admin/accounts.rs",
"src/admin/auth.rs",
"src/admin/catalog.rs",
"src/admin/client_keys.rs",
"src/admin/openai.rs",
"src/admin/mod.rs",
"src/admin/observability.rs",
"src/admin/settings.rs",
"src/admin/system.rs",
"src/admin/wire.rs",
"src/admin/xai.rs",
"src/health.rs",
"src/lib.rs",
"src/openai/auth.rs",
@@ -72,15 +69,12 @@ fn test_tree_should_match_frozen_rust_mirror() {
let mut expected = vec![
"tests/admin/accounts.rs",
"tests/admin/auth.rs",
"tests/admin/catalog.rs",
"tests/admin/client_keys.rs",
"tests/admin/openai.rs",
"tests/admin/mod.rs",
"tests/admin/observability.rs",
"tests/admin/settings.rs",
"tests/admin/system.rs",
"tests/admin/wire.rs",
"tests/admin/xai.rs",
"tests/architecture.rs",
"tests/health.rs",
"tests/main.rs",
backend/crates/gateway-api/tests/openai/error.rs
+1 -1
View File
@@ -128,7 +128,7 @@ fn local_provider_capacity_exhaustion_should_map_to_service_unavailable() {
gateway_error_from_engine(&error),
GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no upstream Provider instance or account is currently available"
"no upstream provider is currently available for this request"
)
);
}
backend/crates/gateway-api/tests/openai/mod.rs
+7 -15
View File
@@ -34,8 +34,8 @@ use gateway_core::lifecycle::{ConnectionDraining, ConnectionGuard, ConnectionLif
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::snapshot::RuntimeSnapshotHandle;
use gateway_core::routing::{
ConfigRevision, InstanceHealth, ModelCapabilities, ProviderInstance, ProviderInstanceId,
ProviderKind, ProviderModel, RuntimeSnapshot, UpstreamModelId,
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, RuntimeSnapshot,
UpstreamModelId,
};
pub(super) async fn api_router(execution: Arc<dyn ExecutionService>) -> axum::Router {
@@ -80,8 +80,6 @@ pub(super) fn authenticated_client_for_provider(
fn snapshot(plaintext: &str, provider_name: &str) -> RuntimeSnapshot {
let provider = ProviderKind::new(provider_name).expect("provider");
let instance_id =
ProviderInstanceId::new(format!("inst_{provider_name}_api_test")).expect("instance");
let capabilities = ModelCapabilities::new(
BTreeSet::from([gateway_core::operation::OperationKind::Generate]),
128_000,
@@ -94,18 +92,12 @@ fn snapshot(plaintext: &str, provider_name: &str) -> RuntimeSnapshot {
NonZeroU32::new(2).expect("concurrency"),
Duration::from_millis(1),
),
vec![ProviderInstance::new(
instance_id.clone(),
provider.clone(),
"https://api.example.invalid".to_owned(),
true,
InstanceHealth::Healthy,
)],
vec![provider.clone()],
["model-a", "model-b"]
.into_iter()
.map(|model| {
ProviderModel::new(
instance_id.clone(),
provider.clone(),
UpstreamModelId::new(model).expect("model"),
capabilities.clone(),
)
@@ -230,21 +222,21 @@ struct UnusedCircuits;
impl ProviderCircuitPort for UnusedCircuits {
fn decision<'a>(
&'a self,
_: &'a ProviderInstanceId,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
Box::pin(async { unreachable!("authentication fixture does not execute") })
}
fn observe_failure<'a>(
&'a self,
_: &'a ProviderInstanceId,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { unreachable!("authentication fixture does not execute") })
}
fn observe_success<'a>(
&'a self,
_: &'a ProviderInstanceId,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { unreachable!("authentication fixture does not execute") })
}
backend/crates/gateway-core/src/engine/continuation.rs
+4 -18
View File
@@ -10,7 +10,7 @@ use futures::future::BoxFuture;
use crate::engine::credential::ProviderAccountId;
use crate::error::{IdentifierError, SafeUpstreamValue, validate_text};
use crate::policy::ClientApiKeyId;
use crate::routing::{ProviderInstanceId, ProviderKind};
use crate::routing::ProviderKind;
/// 客户端传入的 previous response ID。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -46,7 +46,7 @@ pub enum NativeContinuationScope {
/// 外层已解析并认证的 native previous-response pin。
///
/// 该值不代表数据库 transcript；它只阻止 native handle 被发送到错误的
/// Provider、instance 或账号。
/// Provider 或账号。
#[derive(Clone, PartialEq, Eq)]
pub struct NativeContinuationPin {
/// 客户端提交、仅供 Store 在调用方隔离下查找的网关 response ID。
@@ -54,7 +54,6 @@ pub struct NativeContinuationPin {
/// Store 从已提交成功请求解析出的 Provider 原生 response handle。
upstream_response_id: SafeUpstreamValue,
provider: ProviderKind,
instance: ProviderInstanceId,
account: ProviderAccountId,
scope: NativeContinuationScope,
}
@@ -65,14 +64,12 @@ impl NativeContinuationPin {
previous_response_id: PreviousResponseId,
upstream_response_id: SafeUpstreamValue,
provider: ProviderKind,
instance: ProviderInstanceId,
account: ProviderAccountId,
) -> Self {
Self {
previous_response_id,
upstream_response_id,
provider,
instance,
account,
scope: NativeContinuationScope::ConnectionLocal,
}
@@ -101,11 +98,6 @@ impl NativeContinuationPin {
&self.provider
}
#[must_use]
pub const fn instance(&self) -> &ProviderInstanceId {
&self.instance
}
#[must_use]
pub const fn account(&self) -> &ProviderAccountId {
&self.account
@@ -118,13 +110,8 @@ impl NativeContinuationPin {
/// 校验本次 route/account 选择没有破坏 native pin。
#[must_use]
pub fn matches(
&self,
provider: &ProviderKind,
instance: &ProviderInstanceId,
account: &ProviderAccountId,
) -> bool {
self.provider == *provider && self.instance == *instance && self.account == *account
pub fn matches(&self, provider: &ProviderKind, account: &ProviderAccountId) -> bool {
self.provider == *provider && self.account == *account
}
}
@@ -135,7 +122,6 @@ impl fmt::Debug for NativeContinuationPin {
.field("previous_response_id", &"<redacted>")
.field("upstream_response_id", &"<redacted>")
.field("provider", &self.provider)
.field("instance", &self.instance)
.field("account", &self.account)
.field("scope", &self.scope)
.finish()
backend/crates/gateway-core/src/engine/coordinator.rs
+205 -261
View File
@@ -1,4 +1,4 @@
//! 唯一的账号重试、Provider instance 切换、发送与下游 commit barrier owner。
//! 唯一的账号重试、发送与下游 commit barrier owner。
use std::collections::{BTreeSet, VecDeque};
use std::num::NonZeroU32;
@@ -89,8 +89,6 @@ where
commit_policy,
account_state_owner,
cancellation,
candidate_index: 0,
last_attempt_candidate_index: None,
attempts: 0,
excluded_accounts: BTreeSet::new(),
credential_recovery_attempted_accounts: BTreeSet::new(),
@@ -151,7 +149,7 @@ enum StreamCommitPolicy {
/// API 可逐事件消费的 Core 执行会话。
///
/// API 只能提交下游 delivery 边界；账号重试、Provider instance 切换、断流终结与
/// API 只能提交下游 delivery 边界；账号重试、断流终结与
/// `model_requests` 写回均留在本类型内。
pub struct ResponseExecutionSession<S: ?Sized> {
engine: Arc<GatewayEngine<S>>,
@@ -167,8 +165,6 @@ pub struct ResponseExecutionSession<S: ?Sized> {
commit_policy: StreamCommitPolicy,
account_state_owner: Option<ProviderAccountStateOwner>,
cancellation: CancellationToken,
candidate_index: usize,
last_attempt_candidate_index: Option<usize>,
attempts: u32,
excluded_accounts: BTreeSet<crate::engine::credential::ProviderAccountId>,
credential_recovery_attempted_accounts: BTreeSet<crate::engine::credential::ProviderAccountId>,
@@ -443,220 +439,199 @@ where
}
async fn prepare_attempt(&mut self) -> Result<(), EngineError> {
loop {
if self.attempts >= self.plan.max_attempts().get() {
return Err(EngineError::EmptyRoutingPlan);
}
let Some(candidate) = self.plan.candidates().get(self.candidate_index).cloned() else {
let error = GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no upstream account or Provider instance is available",
);
self.finish_failure(FailureFinalization {
outcome: ExecutionOutcome::Failed,
send_state: self.current_send_state(),
error,
upstream_status_code: None,
provider_error_code: None,
retry_after_ms: None,
provider_response_id: None,
})
.await?;
return Err(EngineError::EmptyRoutingPlan);
};
let next_attempt = self
.attempts
.checked_add(1)
.and_then(NonZeroU32::new)
.ok_or(EngineError::EmptyRoutingPlan)?;
let context = AttemptContext::new(
RequestAttemptContext::new(
self.request_id.clone(),
self.client_api_key_ref.clone(),
),
next_attempt,
self.deadline,
self.plan.account_selection_policy(),
AccountAttemptContext::new(
self.excluded_accounts.clone(),
self.required_account.clone(),
self.account_state_owner.clone(),
)
.with_credential_recovery_attempted(
self.required_account.as_ref().is_some_and(|account| {
self.credential_recovery_attempted_accounts
.contains(account)
}),
),
self.continuation.clone(),
self.cancellation.clone(),
)
.with_continuation_attempt(self.continuation_attempt);
let trigger = if self.attempts == 0 {
AttemptTrigger::Initial
} else if self.last_attempt_candidate_index == Some(self.candidate_index) {
AttemptTrigger::AccountRetry
} else {
AttemptTrigger::InstanceFallback
};
let provider = self
.engine
.providers()
.get(candidate.provider())
.cloned()
.ok_or_else(|| EngineError::ProviderNotRegistered {
provider: candidate.provider().as_str().to_owned(),
})?;
let provider_request = ProviderRequest::new(self.operation.clone(), candidate.clone());
let stream = match poll_provider(
provider,
provider_request,
context,
self.cancellation.clone(),
self.deadline,
)
.await
{
ProviderBoundary::Cancelled => {
self.finish_interruption(EngineError::Cancelled).await?;
return Err(EngineError::Cancelled);
}
ProviderBoundary::Deadline => {
self.record_provider_failure(
candidate.instance().clone(),
ProviderErrorKind::Timeout,
);
self.finish_interruption(EngineError::Deadline).await?;
return Err(EngineError::Deadline);
}
ProviderBoundary::Result(result) => match *result {
Ok(stream) => stream,
Err(error) => {
if error.kind() != ProviderErrorKind::Unavailable
|| error.send_state() != UpstreamSendState::NotSent
{
self.record_provider_failure(
candidate.instance().clone(),
error.kind(),
);
}
let provider_proved_replay_safe = provider_proved_replay_safe(&error);
if self.can_fallback_instance(provider_proved_replay_safe) {
self.candidate_index += 1;
continue;
}
let engine_error = provider_engine_error(&error);
self.finish_provider_error(&error).await?;
return Err(engine_error);
}
},
};
if !stream.metadata().confirms(&candidate) {
self.record_provider_failure(
candidate.instance().clone(),
ProviderErrorKind::Protocol,
);
let error = GatewayError::new(
GatewayErrorKind::Internal,
"provider metadata did not match the frozen candidate",
);
self.finish_failure(FailureFinalization {
outcome: ExecutionOutcome::Failed,
send_state: self.current_send_state(),
error,
upstream_status_code: None,
provider_error_code: None,
retry_after_ms: None,
provider_response_id: None,
})
.await?;
return Err(EngineError::ProviderMetadataMismatch);
}
let metadata = stream.metadata().clone();
if self
.required_account
.as_ref()
.is_some_and(|required| metadata.provider_account_id() != Some(required))
{
let error = GatewayError::new(
GatewayErrorKind::Internal,
"provider did not use the required account",
);
self.finish_failure(FailureFinalization {
outcome: ExecutionOutcome::Failed,
send_state: self.current_send_state(),
error,
upstream_status_code: None,
provider_error_code: None,
retry_after_ms: None,
provider_response_id: None,
})
.await?;
return Err(EngineError::RequiredAccountMismatch);
}
if let Some(pin) = self
.continuation
.as_ref()
.and_then(ContinuationBinding::pinned)
&& self.continuation_attempt == ContinuationAttempt::Native
&& !metadata.provider_account_id().is_some_and(|account| {
pin.matches(metadata.provider(), metadata.instance(), account)
})
{
let error = GatewayError::new(
GatewayErrorKind::Internal,
"native continuation binding did not match selected account",
);
self.finish_failure(FailureFinalization {
outcome: ExecutionOutcome::Failed,
send_state: self.current_send_state(),
error,
upstream_status_code: None,
provider_error_code: None,
retry_after_ms: None,
provider_response_id: None,
})
.await?;
return Err(EngineError::ContinuationPinMismatch);
}
if self.account_state_owner.is_none()
&& let Some(account) = metadata.provider_account_id()
{
self.account_state_owner = Some(ProviderAccountStateOwner::new(
metadata.provider().clone(),
metadata.instance().clone(),
account.clone(),
));
}
self.engine
.store()
.record_attempt(AttemptRecord {
request_id: self.request_id.clone(),
attempt_count: next_attempt,
trigger,
provider_instance_id: metadata.instance().clone(),
provider_kind: metadata.provider().clone(),
provider_account_id: metadata.provider_account_id().cloned(),
provider_account_ref: metadata.provider_account_id().cloned(),
upstream_model_id: metadata.upstream_model().clone(),
upstream_transport: metadata.transport().as_str().to_owned(),
http_version: None,
})
.await?;
self.attempts = next_attempt.get();
self.last_attempt_candidate_index = Some(self.candidate_index);
self.current = Some(CurrentAttempt {
stream,
metadata,
trigger,
index: next_attempt,
started_at: SystemTime::now(),
send_observed: false,
response_observation: None,
});
return Ok(());
if self.attempts >= self.plan.max_attempts().get() {
return Err(EngineError::EmptyRoutingPlan);
}
let Some(candidate) = self.plan.candidates().first().cloned() else {
let error = GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no upstream Provider is available",
);
self.finish_failure(FailureFinalization {
outcome: ExecutionOutcome::Failed,
send_state: self.current_send_state(),
error,
upstream_status_code: None,
provider_error_code: None,
retry_after_ms: None,
provider_response_id: None,
})
.await?;
return Err(EngineError::EmptyRoutingPlan);
};
let next_attempt = self
.attempts
.checked_add(1)
.and_then(NonZeroU32::new)
.ok_or(EngineError::EmptyRoutingPlan)?;
let context = AttemptContext::new(
RequestAttemptContext::new(self.request_id.clone(), self.client_api_key_ref.clone()),
next_attempt,
self.deadline,
self.plan.account_selection_policy(),
AccountAttemptContext::new(
self.excluded_accounts.clone(),
self.required_account.clone(),
self.account_state_owner.clone(),
)
.with_credential_recovery_attempted(
self.required_account.as_ref().is_some_and(|account| {
self.credential_recovery_attempted_accounts
.contains(account)
}),
),
self.continuation.clone(),
self.cancellation.clone(),
)
.with_continuation_attempt(self.continuation_attempt);
let trigger = if self.attempts == 0 {
AttemptTrigger::Initial
} else {
AttemptTrigger::AccountRetry
};
let provider = self
.engine
.providers()
.get(candidate.provider())
.cloned()
.ok_or_else(|| EngineError::ProviderNotRegistered {
provider: candidate.provider().as_str().to_owned(),
})?;
let provider_request = ProviderRequest::new(self.operation.clone(), candidate.clone());
let stream = match poll_provider(
provider,
provider_request,
context,
self.cancellation.clone(),
self.deadline,
)
.await
{
ProviderBoundary::Cancelled => {
self.finish_interruption(EngineError::Cancelled).await?;
return Err(EngineError::Cancelled);
}
ProviderBoundary::Deadline => {
self.record_provider_failure(
candidate.provider().clone(),
ProviderErrorKind::Timeout,
);
self.finish_interruption(EngineError::Deadline).await?;
return Err(EngineError::Deadline);
}
ProviderBoundary::Result(result) => match *result {
Ok(stream) => stream,
Err(error) => {
if error.kind() != ProviderErrorKind::Unavailable
|| error.send_state() != UpstreamSendState::NotSent
{
self.record_provider_failure(candidate.provider().clone(), error.kind());
}
let engine_error = provider_engine_error(&error);
self.finish_provider_error(&error).await?;
return Err(engine_error);
}
},
};
if !stream.metadata().confirms(&candidate) {
self.record_provider_failure(candidate.provider().clone(), ProviderErrorKind::Protocol);
let error = GatewayError::new(
GatewayErrorKind::Internal,
"provider metadata did not match the frozen candidate",
);
self.finish_failure(FailureFinalization {
outcome: ExecutionOutcome::Failed,
send_state: self.current_send_state(),
error,
upstream_status_code: None,
provider_error_code: None,
retry_after_ms: None,
provider_response_id: None,
})
.await?;
return Err(EngineError::ProviderMetadataMismatch);
}
let metadata = stream.metadata().clone();
if self
.required_account
.as_ref()
.is_some_and(|required| metadata.provider_account_id() != Some(required))
{
let error = GatewayError::new(
GatewayErrorKind::Internal,
"provider did not use the required account",
);
self.finish_failure(FailureFinalization {
outcome: ExecutionOutcome::Failed,
send_state: self.current_send_state(),
error,
upstream_status_code: None,
provider_error_code: None,
retry_after_ms: None,
provider_response_id: None,
})
.await?;
return Err(EngineError::RequiredAccountMismatch);
}
if let Some(pin) = self
.continuation
.as_ref()
.and_then(ContinuationBinding::pinned)
&& self.continuation_attempt == ContinuationAttempt::Native
&& !metadata
.provider_account_id()
.is_some_and(|account| pin.matches(metadata.provider(), account))
{
let error = GatewayError::new(
GatewayErrorKind::Internal,
"native continuation binding did not match selected account",
);
self.finish_failure(FailureFinalization {
outcome: ExecutionOutcome::Failed,
send_state: self.current_send_state(),
error,
upstream_status_code: None,
provider_error_code: None,
retry_after_ms: None,
provider_response_id: None,
})
.await?;
return Err(EngineError::ContinuationPinMismatch);
}
if self.account_state_owner.is_none()
&& let Some(account) = metadata.provider_account_id()
{
self.account_state_owner = Some(ProviderAccountStateOwner::new(
metadata.provider().clone(),
account.clone(),
));
}
self.engine
.store()
.record_attempt(AttemptRecord {
request_id: self.request_id.clone(),
attempt_count: next_attempt,
trigger,
provider_kind: metadata.provider().clone(),
provider_account_id: metadata.provider_account_id().cloned(),
provider_account_ref: metadata.provider_account_id().cloned(),
upstream_model_id: metadata.upstream_model().clone(),
upstream_transport: metadata.transport().as_str().to_owned(),
http_version: None,
})
.await?;
self.attempts = next_attempt.get();
self.current = Some(CurrentAttempt {
stream,
metadata,
trigger,
index: next_attempt,
started_at: SystemTime::now(),
send_observed: false,
response_observation: None,
});
Ok(())
}
async fn observe_event(&mut self, event: &GatewayEvent) -> Result<(), EngineError> {
@@ -757,7 +732,7 @@ where
/// 返回 `true` 表示调用方必须丢弃本 attempt 已收集的未提交事件。
async fn handle_stream_error(&mut self, error: ProviderError) -> Result<bool, EngineError> {
let current = self.current.take().ok_or(EngineError::EmptyRoutingPlan)?;
self.record_provider_failure(current.metadata.instance().clone(), error.kind());
self.record_provider_failure(current.metadata.provider().clone(), error.kind());
let send_state = if current.send_observed {
UpstreamSendState::Sent
} else {
@@ -810,18 +785,12 @@ where
{
self.excluded_accounts.insert(account.clone());
}
if error.kind() == ProviderErrorKind::Unsupported
&& self.can_fallback_instance(provider_proved_replay_safe)
{
self.candidate_index += 1;
}
self.engine
.store()
.record_intermediate_failure(IntermediateFailure {
request_id: self.request_id.clone(),
attempt_index: current.index,
trigger: current.trigger,
instance_id: current.metadata.instance().clone(),
provider_kind: current.metadata.provider().clone(),
account_id: current.metadata.provider_account_id().cloned(),
upstream_model_id: current.metadata.upstream_model().clone(),
@@ -907,29 +876,6 @@ where
self.before_commit.clear();
}
fn can_fallback_instance(&self, provider_proved_replay_safe: bool) -> bool {
if self.required_account.is_some()
|| self.candidate_index + 1 >= self.plan.candidates().len()
{
return false;
}
let next = &self.plan.candidates()[self.candidate_index + 1];
match self.continuation_attempt {
ContinuationAttempt::None => {
self.operation.retry_safety() == RetrySafety::Idempotent
|| provider_proved_replay_safe
}
ContinuationAttempt::ReplayAny => {
provider_proved_replay_safe
&& self
.operation
.provider_session_state(next.provider().as_str())
.is_some()
}
ContinuationAttempt::Native | ContinuationAttempt::ReplayOwner => false,
}
}
async fn finish_success(&mut self) -> Result<(), EngineError> {
if self.finalized {
return Ok(());
@@ -1172,36 +1118,34 @@ where
}
fn record_current_provider_success(&mut self) {
let instance_id = self
let provider_kind = self
.current
.as_ref()
.map(|current| current.metadata.instance().clone());
if let Some(instance_id) = instance_id {
.map(|current| current.metadata.provider().clone());
if let Some(provider_kind) = provider_kind {
self.provider_attempt_outcomes
.push(ProviderAttemptOutcome::Succeeded {
provider_instance_id: instance_id,
});
.push(ProviderAttemptOutcome::Succeeded { provider_kind });
}
}
fn record_current_provider_failure(&mut self, error_kind: ProviderErrorKind) {
let instance_id = self
let provider_kind = self
.current
.as_ref()
.map(|current| current.metadata.instance().clone());
if let Some(instance_id) = instance_id {
self.record_provider_failure(instance_id, error_kind);
.map(|current| current.metadata.provider().clone());
if let Some(provider_kind) = provider_kind {
self.record_provider_failure(provider_kind, error_kind);
}
}
fn record_provider_failure(
&mut self,
provider_instance_id: crate::routing::ProviderInstanceId,
provider_kind: crate::routing::ProviderKind,
error_kind: ProviderErrorKind,
) {
self.provider_attempt_outcomes
.push(ProviderAttemptOutcome::Failed {
provider_instance_id,
provider_kind,
error_kind,
});
}
backend/crates/gateway-core/src/engine/credential.rs
+4 -12
View File
@@ -11,7 +11,7 @@ use serde_json::{Map, Value};
use thiserror::Error;
use crate::error::{IdentifierError, StoreError, validate_text};
use crate::routing::{ProviderInstanceId, ProviderKind};
use crate::routing::ProviderKind;
/// `provider_accounts.id` 的核心值对象。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -192,7 +192,6 @@ impl AccountAvailability {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccount {
id: ProviderAccountId,
instance: ProviderInstanceId,
provider: ProviderKind,
name: String,
email: Option<String>,
@@ -209,11 +208,10 @@ pub struct ProviderAccount {
}
impl ProviderAccount {
/// 创建账号快照；Provider 身份一致性由数据库复合 FK 保证。
/// 创建账号快照。
#[must_use]
pub const fn new(
id: ProviderAccountId,
instance: ProviderInstanceId,
provider: ProviderKind,
name: String,
upstream_user_id: String,
@@ -222,7 +220,6 @@ impl ProviderAccount {
) -> Self {
Self {
id,
instance,
provider,
name,
email: None,
@@ -281,11 +278,6 @@ impl ProviderAccount {
&self.id
}
#[must_use]
pub const fn instance(&self) -> &ProviderInstanceId {
&self.instance
}
#[must_use]
pub const fn provider(&self) -> &ProviderKind {
&self.provider
@@ -585,9 +577,9 @@ pub trait ProviderAccountStore: Send + Sync {
async fn list_accounts(&self) -> Result<Vec<ProviderAccount>, StoreError>;
async fn list_for_instance(
async fn list_for_provider(
&self,
instance: &ProviderInstanceId,
provider: &ProviderKind,
) -> Result<Vec<ProviderAccount>, StoreError>;
async fn load_credential(
backend/crates/gateway-core/src/engine/execution.rs
+76 -76
View File
@@ -19,15 +19,16 @@ use crate::engine::coordinator::ResponseExecutionSession;
use crate::engine::probe::{AccountProbe, AccountProbeRequest, AccountProbeResult};
use crate::engine::provider::ProviderRegistry;
use crate::engine::{
AttemptCoordinator, CancellationToken, CommitRequirement, CoordinatedEvent, EngineError,
ExecutionStore, GatewayEngine, ModelRequestId, NewModelRequest, ProviderAttemptOutcome,
AttemptCoordinator, AttemptRecord, CancellationToken, CommitRequirement, CoordinatedEvent,
EngineError, ExecutionStore, GatewayEngine, IntermediateFailure, ModelRequestFinalization,
ModelRequestId, NewModelRequest, ProviderAttemptOutcome, RecoveryReport, UpstreamSendState,
};
use crate::error::{GatewayError, GatewayErrorKind, ProviderErrorKind};
use crate::error::{GatewayError, GatewayErrorKind, ProviderErrorKind, StoreError};
use crate::event::{GatewayEvent, ProviderEvent, ProviderResponseHeader};
use crate::operation::{Operation, ReasoningEffort};
use crate::policy::{ClientApiKeyId, ClientPolicy};
use crate::routing::snapshot::RuntimeSnapshotHandle;
use crate::routing::{ProviderInstanceId, PublicModelId, RoutingContext, RuntimeSnapshot};
use crate::routing::{ProviderKind, PublicModelId, RoutingContext, RuntimeSnapshot};
const MODEL_REQUEST_DEADLINE: Duration = Duration::from_secs(10 * 60);
@@ -156,21 +157,22 @@ pub struct ProviderCircuitError;
pub trait ProviderCircuitPort: Send + Sync {
fn decision<'a>(
&'a self,
provider_instance_id: &'a ProviderInstanceId,
provider_kind: &'a ProviderKind,
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>>;
fn observe_failure<'a>(
&'a self,
provider_instance_id: &'a ProviderInstanceId,
provider_kind: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>>;
fn observe_success<'a>(
&'a self,
provider_instance_id: &'a ProviderInstanceId,
provider_kind: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>>;
}
pub struct DefaultExecutionService {
snapshots: RuntimeSnapshotHandle,
coordinator: Arc<AttemptCoordinator<dyn ExecutionStore>>,
probe_coordinator: Arc<AttemptCoordinator<dyn ExecutionStore>>,
providers: ProviderRegistry,
admissions: Arc<dyn ClientAdmissionPort>,
circuits: Arc<dyn ProviderCircuitPort>,
@@ -188,9 +190,12 @@ impl DefaultExecutionService {
continuation: Arc<dyn NativeContinuationPort>,
) -> Self {
let engine = GatewayEngine::<dyn ExecutionStore>::new(execution, providers.clone());
let transient: Arc<dyn ExecutionStore> = Arc::new(TransientExecutionStore);
let probe_engine = GatewayEngine::<dyn ExecutionStore>::new(transient, providers.clone());
Self {
snapshots,
coordinator: Arc::new(AttemptCoordinator::new(engine)),
probe_coordinator: Arc::new(AttemptCoordinator::new(probe_engine)),
providers,
admissions,
circuits,
@@ -210,12 +215,7 @@ impl DefaultExecutionService {
})?;
let request_id = new_request_id()?;
let routing_context = self
.route_context(
&request.client.snapshot,
request.client.policy.provider_kind(),
&request.public_model,
&request.operation,
)
.route_context(request.client.policy.provider_kind())
.await?;
let plan = request
.client
@@ -377,50 +377,20 @@ impl DefaultExecutionService {
async fn route_context(
&self,
snapshot: &RuntimeSnapshot,
provider_kind: &crate::routing::ProviderKind,
model: &PublicModelId,
operation: &Operation,
provider_kind: &ProviderKind,
) -> Result<RoutingContext, GatewayError> {
let allowed_instances = snapshot.instance_ids_for_provider(provider_kind);
let preliminary = snapshot
.plan(
model,
operation,
&RoutingContext {
provider_kind: Some(provider_kind.clone()),
allowed_instances: Some(allowed_instances.clone()),
..RoutingContext::default()
},
let blocked_providers = match self.circuits.decision(provider_kind).await.map_err(|_| {
GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"provider health state is temporarily unavailable",
)
.map_err(map_routing_error)?;
let mut blocked_instances = BTreeSet::new();
let mut checked = BTreeSet::new();
for candidate in preliminary.candidates() {
if !checked.insert(candidate.instance().clone()) {
continue;
}
match self
.circuits
.decision(candidate.instance())
.await
.map_err(|_| {
GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"provider health state is temporarily unavailable",
)
})? {
ProviderCircuitDecision::Allow => {}
ProviderCircuitDecision::BlockedUntil(_) => {
blocked_instances.insert(candidate.instance().clone());
}
}
}
})? {
ProviderCircuitDecision::Allow => BTreeSet::new(),
ProviderCircuitDecision::BlockedUntil(_) => BTreeSet::from([provider_kind.clone()]),
};
Ok(RoutingContext {
provider_kind: Some(provider_kind.clone()),
allowed_instances: Some(allowed_instances),
blocked_instances,
..RoutingContext::default()
blocked_providers,
})
}
@@ -430,7 +400,7 @@ impl DefaultExecutionService {
) -> Result<AccountProbeResult, GatewayError> {
let AccountProbeRequest {
account_id,
provider_instance_id,
provider_kind,
upstream_model,
operation,
} = request;
@@ -440,22 +410,12 @@ impl DefaultExecutionService {
"runtime snapshot is unavailable",
)
})?;
let provider_kind = snapshot
.provider_for_instance(&provider_instance_id)
.cloned()
.ok_or_else(|| {
GatewayError::new(
GatewayErrorKind::Unsupported,
"Provider instance is unavailable",
)
})?;
let public_model =
PublicModelId::new(upstream_model.as_str().to_owned()).map_err(|_| {
GatewayError::new(GatewayErrorKind::Unsupported, "requested model is invalid")
})?;
let routing_context = RoutingContext {
provider_kind: Some(provider_kind),
allowed_instances: Some(BTreeSet::from([provider_instance_id.clone()])),
..RoutingContext::default()
};
let plan = snapshot
@@ -478,7 +438,7 @@ impl DefaultExecutionService {
config_revision: plan.config_revision(),
protocol: "admin_connection_test".to_owned(),
operation: operation.kind(),
endpoint: "/api/admin/accounts/test".to_owned(),
endpoint: "/api/admin/accounts/connection-test".to_owned(),
client_transport: ClientTransport::InternalProbe.as_str().to_owned(),
requested_model: public_model,
input_token_estimate,
@@ -494,7 +454,7 @@ impl DefaultExecutionService {
deadline_at,
};
let mut session = self
.coordinator
.probe_coordinator
.start(
new_request,
operation,
@@ -529,6 +489,52 @@ impl DefaultExecutionService {
}
}
struct TransientExecutionStore;
#[async_trait::async_trait]
impl ExecutionStore for TransientExecutionStore {
async fn create_model_request(&self, _: NewModelRequest) -> Result<(), StoreError> {
Ok(())
}
async fn record_attempt(&self, _: AttemptRecord) -> Result<(), StoreError> {
Ok(())
}
async fn mark_send_state(
&self,
_: &ModelRequestId,
_: UpstreamSendState,
) -> Result<(), StoreError> {
Ok(())
}
async fn mark_downstream_committed(
&self,
_: &ModelRequestId,
_: SystemTime,
_: Option<u16>,
) -> Result<(), StoreError> {
Ok(())
}
async fn record_client_status(&self, _: &ModelRequestId, _: u16) -> Result<(), StoreError> {
Ok(())
}
async fn record_intermediate_failure(&self, _: IntermediateFailure) -> Result<(), StoreError> {
Ok(())
}
async fn finalize_model_request(&self, _: ModelRequestFinalization) -> Result<(), StoreError> {
Ok(())
}
async fn recover_expired(&self, _: SystemTime) -> Result<RecoveryReport, StoreError> {
Ok(RecoveryReport::default())
}
}
impl ExecutionService for DefaultExecutionService {
fn authenticate(
&self,
@@ -758,15 +764,9 @@ async fn publish_provider_attempt_outcomes(
) {
for outcome in outcomes {
let result = match outcome.error_kind() {
None => {
circuits
.observe_success(outcome.provider_instance_id())
.await
}
None => circuits.observe_success(outcome.provider_kind()).await,
Some(kind) if provider_failure_affects_circuit(kind) => {
circuits
.observe_failure(outcome.provider_instance_id())
.await
circuits.observe_failure(outcome.provider_kind()).await
}
Some(_) => continue,
};
@@ -806,7 +806,7 @@ fn map_routing_error(error: crate::error::RoutingError) -> GatewayError {
match error {
crate::error::RoutingError::NoCapableProvider { .. } => GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no Provider instance can execute this request",
"no provider can execute this request",
),
_ => GatewayError::new(
GatewayErrorKind::Internal,
@@ -827,7 +827,7 @@ pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError {
EngineError::EmptyRoutingPlan | EngineError::ProviderNotRegistered { .. } => {
GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no Provider instance is available",
"no provider is available",
)
}
_ => GatewayError::new(GatewayErrorKind::Internal, "request execution failed"),
backend/crates/gateway-core/src/engine/mod.rs
+15 -52
View File
@@ -31,9 +31,7 @@ use crate::error::{
use crate::event::ProviderEvent;
use crate::operation::OperationKind;
use crate::policy::ClientApiKeyId;
use crate::routing::{
ConfigRevision, ProviderInstanceId, ProviderKind, PublicModelId, UpstreamModelId,
};
use crate::routing::{ConfigRevision, ProviderKind, PublicModelId, UpstreamModelId};
/// `model_requests.id`。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -105,7 +103,6 @@ impl ExecutionOutcome {
pub enum AttemptTrigger {
Initial,
AccountRetry,
InstanceFallback,
}
impl AttemptTrigger {
@@ -114,39 +111,30 @@ impl AttemptTrigger {
match self {
Self::Initial => "initial",
Self::AccountRetry => "account_retry",
Self::InstanceFallback => "instance_fallback",
}
}
}
/// 一次实际上游调用对 Provider instance 健康度产生的事实。
/// 一次实际上游调用对 Provider 健康度产生的事实。
///
/// 该事实只描述调用结果，不在 Core 中定义 circuit 策略或持久化方式。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ProviderAttemptOutcome {
/// 上游流自然完成且通过 canonical event 序列校验。
Succeeded {
provider_instance_id: ProviderInstanceId,
},
Succeeded { provider_kind: ProviderKind },
/// 上游打开或流式阶段返回了稳定 Provider 错误。
Failed {
provider_instance_id: ProviderInstanceId,
provider_kind: ProviderKind,
error_kind: ProviderErrorKind,
},
}
impl ProviderAttemptOutcome {
/// 返回本次调用实际归属的 Provider instance。
/// 返回本次调用实际归属的 Provider。
#[must_use]
pub const fn provider_instance_id(&self) -> &ProviderInstanceId {
pub const fn provider_kind(&self) -> &ProviderKind {
match self {
Self::Succeeded {
provider_instance_id,
}
| Self::Failed {
provider_instance_id,
..
} => provider_instance_id,
Self::Succeeded { provider_kind } | Self::Failed { provider_kind, .. } => provider_kind,
}
}
@@ -290,46 +278,28 @@ impl AccountAttemptContext {
/// 请求中 Provider 账号绑定状态的唯一归属。
///
/// `turn_state` 等 opaque 状态只能发送给创建它的 Provider、instance 与账号；
/// `turn_state` 等 opaque 状态只能发送给创建它的 Provider 与账号；
/// Core 在首次真实选号后冻结该事实，Provider 据此决定是否清理跨账号状态。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccountStateOwner {
provider: ProviderKind,
instance: ProviderInstanceId,
account: ProviderAccountId,
}
impl ProviderAccountStateOwner {
#[must_use]
pub const fn new(
provider: ProviderKind,
instance: ProviderInstanceId,
account: ProviderAccountId,
) -> Self {
Self {
provider,
instance,
account,
}
pub const fn new(provider: ProviderKind, account: ProviderAccountId) -> Self {
Self { provider, account }
}
#[must_use]
pub fn from_continuation(pin: &NativeContinuationPin) -> Self {
Self::new(
pin.provider().clone(),
pin.instance().clone(),
pin.account().clone(),
)
Self::new(pin.provider().clone(), pin.account().clone())
}
#[must_use]
pub fn matches(
&self,
provider: &ProviderKind,
instance: &ProviderInstanceId,
account: &ProviderAccountId,
) -> bool {
self.provider == *provider && self.instance == *instance && self.account == *account
pub fn matches(&self, provider: &ProviderKind, account: &ProviderAccountId) -> bool {
self.provider == *provider && self.account == *account
}
#[must_use]
@@ -337,11 +307,6 @@ impl ProviderAccountStateOwner {
&self.provider
}
#[must_use]
pub const fn instance(&self) -> &ProviderInstanceId {
&self.instance
}
#[must_use]
pub const fn account(&self) -> &ProviderAccountId {
&self.account
@@ -534,7 +499,6 @@ pub struct AttemptRecord {
pub request_id: ModelRequestId,
pub attempt_count: NonZeroU32,
pub trigger: AttemptTrigger,
pub provider_instance_id: ProviderInstanceId,
pub provider_kind: ProviderKind,
pub provider_account_id: Option<ProviderAccountId>,
pub provider_account_ref: Option<ProviderAccountId>,
@@ -543,13 +507,12 @@ pub struct AttemptRecord {
pub http_version: Option<String>,
}
/// 需要解释换号或 Provider instance 切换的中间失败。
/// 需要解释换号的中间失败。
#[derive(Debug)]
pub struct IntermediateFailure {
pub request_id: ModelRequestId,
pub attempt_index: NonZeroU32,
pub trigger: AttemptTrigger,
pub instance_id: ProviderInstanceId,
pub provider_kind: ProviderKind,
pub account_id: Option<ProviderAccountId>,
pub upstream_model_id: UpstreamModelId,
@@ -695,7 +658,7 @@ pub enum EngineError {
Store(#[from] StoreError),
#[error("provider `{provider}` is not registered")]
ProviderNotRegistered { provider: String },
#[error("provider metadata did not match the selected Provider instance and account")]
#[error("provider metadata did not match the selected Provider and account")]
ProviderMetadataMismatch,
#[error("native continuation pin did not match the selected account")]
ContinuationPinMismatch,
backend/crates/gateway-core/src/engine/probe.rs
+2 -2
View File
@@ -3,13 +3,13 @@
use futures::future::BoxFuture;
use crate::error::GatewayError;
use crate::routing::{ProviderInstanceId, UpstreamModelId};
use crate::routing::{ProviderKind, UpstreamModelId};
use crate::{engine::credential::ProviderAccountId, operation::Operation};
#[derive(Debug, Clone, PartialEq)]
pub struct AccountProbeRequest {
pub account_id: ProviderAccountId,
pub provider_instance_id: ProviderInstanceId,
pub provider_kind: ProviderKind,
pub upstream_model: UpstreamModelId,
pub operation: Operation,
}
backend/crates/gateway-core/src/engine/provider.rs
+13 -24
View File
@@ -17,10 +17,7 @@ use crate::error::{
};
use crate::event::{EventSequenceValidator, ProviderEvent};
use crate::operation::Operation;
use crate::routing::{
ModelCapabilities, ProviderCandidate, ProviderInstance, ProviderInstanceId, ProviderKind,
UpstreamModelId,
};
use crate::routing::{ModelCapabilities, ProviderCandidate, ProviderKind, UpstreamModelId};
/// Box 只出现在 Provider Registry 的统一 event envelope 边界。
pub type EventStream =
@@ -121,7 +118,6 @@ impl UpstreamTransport {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderCallMetadata {
provider: ProviderKind,
instance: ProviderInstanceId,
upstream_model: UpstreamModelId,
resource: ProviderResource,
upstream_request_id: Option<SafeUpstreamValue>,
@@ -133,14 +129,12 @@ impl ProviderCallMetadata {
#[must_use]
pub const fn new(
provider: ProviderKind,
instance: ProviderInstanceId,
upstream_model: UpstreamModelId,
resource: ProviderResource,
transport: UpstreamTransport,
) -> Self {
Self {
provider,
instance,
upstream_model,
resource,
upstream_request_id: None,
@@ -161,12 +155,6 @@ impl ProviderCallMetadata {
&self.provider
}
/// 返回 instance。
#[must_use]
pub const fn instance(&self) -> &ProviderInstanceId {
&self.instance
}
/// 返回实际模型；必须与冻结 target 一致。
#[must_use]
pub const fn upstream_model(&self) -> &UpstreamModelId {
@@ -206,9 +194,7 @@ impl ProviderCallMetadata {
/// 确认 metadata 没有替换请求计划中冻结的 Provider 候选。
#[must_use]
pub fn confirms(&self, candidate: &ProviderCandidate) -> bool {
candidate.provider() == &self.provider
&& candidate.instance() == &self.instance
&& candidate.upstream_model() == &self.upstream_model
candidate.provider() == &self.provider && candidate.upstream_model() == &self.upstream_model
}
}
@@ -386,7 +372,6 @@ impl fmt::Debug for ProviderRequest {
.debug_struct("ProviderRequest")
.field("operation", &self.operation)
.field("provider", self.candidate.provider())
.field("instance", self.candidate.instance())
.field("upstream_model", self.candidate.upstream_model())
.finish()
}
@@ -410,14 +395,13 @@ pub trait Provider: Send + Sync {
ProviderRequestObservation::default()
}
/// 查询 instance 的实时模型目录，并由 Provider 自己编译能力事实。
/// 查询当前 Provider 的实时模型目录，并由 Provider 自己编译能力事实。
///
/// # Errors
///
/// 目录 transport、认证或 Provider 协议失败时返回稳定错误。
async fn query_model_capabilities(
&self,
instance: &ProviderInstance,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError>;
/// 选择一个未被排除的资源并返回 cold [`ProviderStream`]。
@@ -525,6 +509,11 @@ impl ProviderRegistry {
self.providers.get(provider)
}
/// 返回全部已注册 Provider。
pub fn provider_kinds(&self) -> impl Iterator<Item = &ProviderKind> {
self.providers.keys()
}
#[must_use]
pub fn request_observation(
&self,
@@ -547,22 +536,22 @@ impl ProviderRegistry {
.collect()
}
/// 通过编译期 Provider adapter 查询并编译 instance 的实时能力目录。
/// 通过编译期 Provider adapter 查询并编译实时能力目录。
///
/// # Errors
///
/// Provider 未注册或目录查询失败时返回错误。
pub async fn query_model_capabilities(
&self,
instance: &ProviderInstance,
provider_kind: &ProviderKind,
) -> Result<Vec<ProviderModelCapabilities>, ProviderCatalogError> {
let provider = self.providers.get(instance.provider()).ok_or_else(|| {
let provider = self.providers.get(provider_kind).ok_or_else(|| {
ProviderCatalogError::NotRegistered {
provider: instance.provider().as_str().to_owned(),
provider: provider_kind.as_str().to_owned(),
}
})?;
provider
.query_model_capabilities(instance)
.query_model_capabilities()
.await
.map_err(ProviderCatalogError::Query)
}
backend/crates/gateway-core/src/error.rs
+2 -2
View File
@@ -87,8 +87,8 @@ pub enum RoutingError {
/// 实体 ID。
id: String,
},
/// 固定平台内没有可执行本次请求的 Provider instance。
#[error("no provider instance can execute model `{model}`")]
/// 固定平台内没有可执行本次请求的 Provider。
#[error("no provider can execute model `{model}`")]
NoCapableProvider {
/// 客户端提交的模型名称。
model: String,
backend/crates/gateway-core/src/provider_ports/mod.rs
+13 -108
View File
@@ -12,7 +12,7 @@ use crate::engine::credential::{
AccountAvailability, AccountRuntimeSignals, CredentialRevision, OpaqueProviderData,
ProviderAccountId, ProviderAccountStore,
};
use crate::routing::{ProviderInstanceId, ProviderKind};
use crate::routing::ProviderKind;
const MAX_PENDING_FLOW_TTL: Duration = Duration::from_secs(30 * 60);
@@ -44,7 +44,7 @@ impl ProviderStoreError {
}
}
/// 一次 Provider instance 的完整可重建调度状态。
/// 一个 Provider 的完整可重建调度状态。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderSchedulingState {
signals: BTreeMap<ProviderAccountId, AccountRuntimeSignals>,
@@ -85,7 +85,7 @@ impl ProviderSchedulingState {
/// 请求级账号 lease 的全部中立事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderSchedulingLeaseRequest {
provider_instance_id: ProviderInstanceId,
provider_kind: ProviderKind,
account_id: ProviderAccountId,
credential_revision: CredentialRevision,
max_concurrent: NonZeroU32,
@@ -96,7 +96,7 @@ pub struct ProviderSchedulingLeaseRequest {
impl ProviderSchedulingLeaseRequest {
#[must_use]
pub const fn new(
provider_instance_id: ProviderInstanceId,
provider_kind: ProviderKind,
account_id: ProviderAccountId,
credential_revision: CredentialRevision,
max_concurrent: NonZeroU32,
@@ -104,7 +104,7 @@ impl ProviderSchedulingLeaseRequest {
deadline: SystemTime,
) -> Self {
Self {
provider_instance_id,
provider_kind,
account_id,
credential_revision,
max_concurrent,
@@ -114,8 +114,8 @@ impl ProviderSchedulingLeaseRequest {
}
#[must_use]
pub const fn provider_instance_id(&self) -> &ProviderInstanceId {
&self.provider_instance_id
pub const fn provider_kind(&self) -> &ProviderKind {
&self.provider_kind
}
#[must_use]
@@ -177,7 +177,7 @@ pub enum ProviderLeaseRequest {
pub trait ProviderLeasePort: Send + Sync {
fn load_state<'a>(
&'a self,
provider_instance_id: &'a ProviderInstanceId,
provider_kind: &'a ProviderKind,
accounts: &'a [ProviderAccountId],
) -> BoxFuture<'a, Result<ProviderSchedulingState, ProviderStoreError>>;
@@ -187,7 +187,7 @@ pub trait ProviderLeasePort: Send + Sync {
) -> BoxFuture<'_, Result<ProviderLeaseAcquisition, ProviderStoreError>>;
}
/// 所有 Provider、所有实例共享的 OAuth refresh 并发容量。
/// 所有 Provider 共享的 OAuth refresh 并发容量。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ProviderRefreshCapacityRequest {
max_concurrent: NonZeroU32,
@@ -286,96 +286,6 @@ pub trait ProviderCatalogCachePort: Send + Sync {
) -> BoxFuture<'a, Result<Option<OpaqueProviderData>, ProviderStoreError>>;
}
/// Provider worker 读取的中立 instance 事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderInstanceConfig {
id: ProviderInstanceId,
provider_kind: ProviderKind,
base_url: String,
enabled: bool,
}
impl ProviderInstanceConfig {
#[must_use]
pub const fn new(
id: ProviderInstanceId,
provider_kind: ProviderKind,
base_url: String,
enabled: bool,
) -> Self {
Self {
id,
provider_kind,
base_url,
enabled,
}
}
#[must_use]
pub const fn id(&self) -> &ProviderInstanceId {
&self.id
}
#[must_use]
pub const fn provider_kind(&self) -> &ProviderKind {
&self.provider_kind
}
#[must_use]
pub fn base_url(&self) -> &str {
&self.base_url
}
#[must_use]
pub const fn enabled(&self) -> bool {
self.enabled
}
}
pub trait ProviderInstanceCatalogPort: Send + Sync {
fn list_instances<'a>(
&'a self,
provider_kind: &'a ProviderKind,
include_disabled: bool,
) -> BoxFuture<'a, Result<Vec<ProviderInstanceConfig>, ProviderStoreError>>;
}
/// Provider 的两种目录能力：instance 事实与 provider-owned model cache。
#[derive(Clone)]
pub struct ProviderCatalogPorts {
instances: Arc<dyn ProviderInstanceCatalogPort>,
model_cache: Arc<dyn ProviderCatalogCachePort>,
}
impl ProviderCatalogPorts {
#[must_use]
pub fn new(
instances: Arc<dyn ProviderInstanceCatalogPort>,
model_cache: Arc<dyn ProviderCatalogCachePort>,
) -> Self {
Self {
instances,
model_cache,
}
}
#[must_use]
pub fn instances(&self) -> Arc<dyn ProviderInstanceCatalogPort> {
Arc::clone(&self.instances)
}
#[must_use]
pub fn model_cache(&self) -> Arc<dyn ProviderCatalogCachePort> {
Arc::clone(&self.model_cache)
}
}
impl fmt::Debug for ProviderCatalogPorts {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("ProviderCatalogPorts([CAPABILITIES])")
}
}
/// Redis 中可重建的账号状态投影。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderCredentialState {
@@ -745,7 +655,7 @@ pub trait OAuthPendingFlowPort: Send + Sync {
pub struct ProviderStorePorts {
accounts: Arc<dyn ProviderAccountStore>,
leases: Arc<dyn ProviderLeasePort>,
catalog: ProviderCatalogPorts,
catalog_cache: Arc<dyn ProviderCatalogCachePort>,
credential_state: Arc<dyn ProviderCredentialStatePort>,
cooldowns: Arc<dyn ProviderCooldownPort>,
runtime_policy: Arc<dyn ProviderRuntimePolicyPort>,
@@ -757,7 +667,7 @@ impl ProviderStorePorts {
pub fn new(
accounts: Arc<dyn ProviderAccountStore>,
leases: Arc<dyn ProviderLeasePort>,
catalog: ProviderCatalogPorts,
catalog_cache: Arc<dyn ProviderCatalogCachePort>,
credential_state: Arc<dyn ProviderCredentialStatePort>,
cooldowns: Arc<dyn ProviderCooldownPort>,
runtime_policy: Arc<dyn ProviderRuntimePolicyPort>,
@@ -766,7 +676,7 @@ impl ProviderStorePorts {
Self {
accounts,
leases,
catalog,
catalog_cache,
credential_state,
cooldowns,
runtime_policy,
@@ -784,14 +694,9 @@ impl ProviderStorePorts {
Arc::clone(&self.leases)
}
#[must_use]
pub fn instances(&self) -> Arc<dyn ProviderInstanceCatalogPort> {
self.catalog.instances()
}
#[must_use]
pub fn catalog_cache(&self) -> Arc<dyn ProviderCatalogCachePort> {
self.catalog.model_cache()
Arc::clone(&self.catalog_cache)
}
#[must_use]
backend/crates/gateway-core/src/routing/mod.rs
+16 -137
View File
@@ -1,4 +1,4 @@
//! Provider 平台、实例、模型目录、精确模型映射与请求级候选计划。
//! Provider、模型目录、精确模型映射与请求级候选计划。
pub mod snapshot;
@@ -15,30 +15,6 @@ use crate::operation::{CapabilityRequirements, Feature, OperationKind};
const MAX_REQUEST_ATTEMPTS: u32 = 32;
/// 一个具体 Provider endpoint 的 ID。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ProviderInstanceId(String);
impl ProviderInstanceId {
/// 校验并创建 Provider instance ID。
pub fn new(value: impl Into<String>) -> Result<Self, IdentifierError> {
let value = value.into();
validate_text(&value, 128, false, Some("inst_"))?;
Ok(Self(value))
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Display for ProviderInstanceId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.0)
}
}
/// 编译进二进制的 Provider adapter slug。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ProviderKind(String);
@@ -214,151 +190,59 @@ impl ModelCapabilities {
}
}
/// Provider instance 的可重建运行健康状态。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum InstanceHealth {
Healthy,
Degraded,
Unavailable,
CircuitOpen,
Saturated,
}
impl InstanceHealth {
const fn is_routable(self, allow_degraded: bool) -> bool {
matches!(self, Self::Healthy) || (allow_degraded && matches!(self, Self::Degraded))
}
}
/// `provider_instances` 与实时健康状态组成的只读快照。
/// 一个 Provider 实时发现的上游模型能力。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderInstance {
id: ProviderInstanceId,
pub struct ProviderModel {
provider: ProviderKind,
base_url: String,
enabled: bool,
health: InstanceHealth,
upstream_model: UpstreamModelId,
capabilities: ModelCapabilities,
}
impl ProviderInstance {
impl ProviderModel {
#[must_use]
pub const fn new(
id: ProviderInstanceId,
provider: ProviderKind,
base_url: String,
enabled: bool,
health: InstanceHealth,
upstream_model: UpstreamModelId,
capabilities: ModelCapabilities,
) -> Self {
Self {
id,
provider,
base_url,
enabled,
health,
upstream_model,
capabilities,
}
}
#[must_use]
pub const fn id(&self) -> &ProviderInstanceId {
&self.id
}
#[must_use]
pub const fn provider(&self) -> &ProviderKind {
&self.provider
}
#[must_use]
pub fn base_url(&self) -> &str {
&self.base_url
}
#[must_use]
pub const fn enabled(&self) -> bool {
self.enabled
}
#[must_use]
pub const fn health(&self) -> InstanceHealth {
self.health
}
}
/// 一个 instance 实时发现的上游模型能力。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderModel {
instance: ProviderInstanceId,
upstream_model: UpstreamModelId,
capabilities: ModelCapabilities,
}
impl ProviderModel {
#[must_use]
pub const fn new(
instance: ProviderInstanceId,
upstream_model: UpstreamModelId,
capabilities: ModelCapabilities,
) -> Self {
Self {
instance,
upstream_model,
capabilities,
}
}
#[must_use]
pub const fn instance(&self) -> &ProviderInstanceId {
&self.instance
}
#[must_use]
pub const fn upstream_model(&self) -> &UpstreamModelId {
&self.upstream_model
}
}
/// Provider circuit 的 request-scoped observation fence。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ProviderInstanceObservationToken {
pub epoch: u64,
pub fencing_token: u64,
pub probe_fencing_token: Option<u64>,
}
/// 本次请求选择 Provider instance 时使用的动态过滤事实。
/// 本次请求选择 Provider 时使用的动态过滤事实。
#[derive(Debug, Clone, Default)]
pub struct RoutingContext {
/// 已认证 Client API Key 绑定的平台；模型名称不参与平台猜测。
pub provider_kind: Option<ProviderKind>,
pub allowed_instances: Option<BTreeSet<ProviderInstanceId>>,
pub allow_degraded: bool,
pub blocked_instances: BTreeSet<ProviderInstanceId>,
pub provider_observation_tokens: BTreeMap<ProviderInstanceId, ProviderInstanceObservationToken>,
pub blocked_providers: BTreeSet<ProviderKind>,
}
/// 已绑定 Provider、instance 与真实上游模型的请求候选。
/// 已绑定 Provider 与真实上游模型的请求候选。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderCandidate {
instance: ProviderInstance,
provider: ProviderKind,
upstream_model: UpstreamModelId,
emulated_features: BTreeSet<Feature>,
observation_token: Option<ProviderInstanceObservationToken>,
}
impl ProviderCandidate {
#[must_use]
pub const fn instance_snapshot(&self) -> &ProviderInstance {
&self.instance
}
#[must_use]
pub const fn instance(&self) -> &ProviderInstanceId {
self.instance.id()
}
#[must_use]
pub const fn provider(&self) -> &ProviderKind {
self.instance.provider()
&self.provider
}
#[must_use]
@@ -370,14 +254,9 @@ impl ProviderCandidate {
pub const fn emulated_features(&self) -> &BTreeSet<Feature> {
&self.emulated_features
}
#[must_use]
pub const fn observation_token(&self) -> Option<ProviderInstanceObservationToken> {
self.observation_token
}
}
/// 一次请求冻结的平台内 Provider instance 尝试顺序。
/// 一次请求冻结的 Provider 尝试顺序。
#[derive(Debug, Clone)]
pub struct RoutingPlan {
config_revision: ConfigRevision,
backend/crates/gateway-core/src/routing/snapshot.rs
+41 -137
View File
@@ -24,9 +24,8 @@ use crate::task::{
};
use super::{
ConfigRevision, InstanceHealth, ModelCapabilities, ProviderCandidate, ProviderInstance,
ProviderInstanceId, ProviderKind, ProviderModel, PublicModelId, RoutingContext, RoutingPlan,
UpstreamModelId,
ConfigRevision, ModelCapabilities, ProviderCandidate, ProviderKind, ProviderModel,
PublicModelId, RoutingContext, RoutingPlan, UpstreamModelId,
};
const RECONCILIATION_INTERVAL: Duration = Duration::from_secs(5);
@@ -62,32 +61,6 @@ impl SnapshotSettingsFacts {
}
}
/// Store 读取到的一个 Provider instance 配置事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotProviderInstanceFacts {
id: String,
provider_kind: String,
base_url: String,
enabled: bool,
}
impl SnapshotProviderInstanceFacts {
#[must_use]
pub fn new(
id: impl Into<String>,
provider_kind: impl Into<String>,
base_url: impl Into<String>,
enabled: bool,
) -> Self {
Self {
id: id.into(),
provider_kind: provider_kind.into(),
base_url: base_url.into(),
enabled,
}
}
}
/// Store 读取到的一个启用 Client API Key 策略事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotClientPolicyFacts {
@@ -120,7 +93,6 @@ pub struct SnapshotFacts {
config_revision: ConfigRevision,
observed_current_revision: ConfigRevision,
settings: SnapshotSettingsFacts,
provider_instances: Vec<SnapshotProviderInstanceFacts>,
client_policies: Vec<SnapshotClientPolicyFacts>,
}
@@ -130,14 +102,12 @@ impl SnapshotFacts {
config_revision: ConfigRevision,
observed_current_revision: ConfigRevision,
settings: SnapshotSettingsFacts,
provider_instances: Vec<SnapshotProviderInstanceFacts>,
client_policies: Vec<SnapshotClientPolicyFacts>,
) -> Self {
Self {
config_revision,
observed_current_revision,
settings,
provider_instances,
client_policies,
}
}
@@ -226,7 +196,7 @@ impl RuntimeSnapshotCompiler {
Self { store, providers }
}
/// 读取一个 revision，并为 instance 查询实时模型目录。
/// 读取一个 revision，并为已注册 Provider 查询实时模型目录。
pub async fn compile(&self) -> Result<RuntimeSnapshot, RuntimeSnapshotCompileError> {
for _ in 0..MAXIMUM_CATALOG_STABILITY_ATTEMPTS {
let catalog_generations = self.providers.catalog_generations();
@@ -252,32 +222,17 @@ async fn compile_runtime_snapshot(
facts: SnapshotFacts,
providers: &ProviderRegistry,
) -> Result<RuntimeSnapshot, RuntimeSnapshotCompileError> {
let mut instances_by_id = BTreeMap::new();
for record in facts.provider_instances {
let id = ProviderInstanceId::new(record.id)
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)?;
let instance = ProviderInstance::new(
id.clone(),
ProviderKind::new(record.provider_kind)
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)?,
record.base_url,
record.enabled,
InstanceHealth::Healthy,
);
if instances_by_id.insert(id, instance).is_some() {
return Err(RuntimeSnapshotCompileError::InvalidData);
}
}
let provider_kinds = providers.provider_kinds().cloned().collect::<Vec<_>>();
// 实时目录只提供公开模型与能力提示；查询失败时保留 instance 透传语义。
// 实时目录只提供公开模型与能力提示；查询失败时保留 Provider 透传语义。
let mut provider_models = Vec::new();
for instance in instances_by_id.values() {
let Ok(models) = providers.query_model_capabilities(instance).await else {
for provider in &provider_kinds {
let Ok(models) = providers.query_model_capabilities(provider).await else {
continue;
};
provider_models.extend(models.into_iter().map(|model| {
ProviderModel::new(
instance.id().clone(),
provider.clone(),
model.upstream_model().clone(),
model.capabilities().clone(),
)
@@ -326,7 +281,7 @@ async fn compile_runtime_snapshot(
RuntimeSnapshot::new(
facts.config_revision,
selection_policy,
instances_by_id.into_values().collect(),
provider_kinds,
provider_models,
client_policies,
)
@@ -339,44 +294,42 @@ async fn compile_runtime_snapshot(
pub struct RuntimeSnapshot {
revision: ConfigRevision,
account_selection_policy: AccountSelectionPolicy,
instances: Arc<BTreeMap<ProviderInstanceId, ProviderInstance>>,
provider_models:
Arc<BTreeMap<ProviderInstanceId, BTreeMap<UpstreamModelId, ModelCapabilities>>>,
providers: Arc<BTreeSet<ProviderKind>>,
provider_models: Arc<BTreeMap<ProviderKind, BTreeMap<UpstreamModelId, ModelCapabilities>>>,
provider_model_mappings: Arc<BTreeMap<ProviderKind, BTreeMap<String, String>>>,
provider_catalog_generations: Arc<BTreeMap<ProviderKind, ProviderCatalogGeneration>>,
client_policies: Arc<BTreeMap<ClientApiKeyId, ClientPolicy>>,
}
impl RuntimeSnapshot {
/// 校验 instance、实时模型目录和 Client API Key，并构建快照。
/// 校验 Provider、实时模型目录和 Client API Key，并构建快照。
pub fn new(
revision: ConfigRevision,
account_selection_policy: AccountSelectionPolicy,
instances: Vec<ProviderInstance>,
providers: Vec<ProviderKind>,
provider_models: Vec<ProviderModel>,
client_policies: Vec<ClientPolicy>,
) -> Result<Self, RoutingError> {
let mut instance_map = BTreeMap::new();
for instance in instances {
let id = instance.id().clone();
if instance_map.insert(id.clone(), instance).is_some() {
let mut provider_set = BTreeSet::new();
for provider in providers {
if !provider_set.insert(provider.clone()) {
return Err(RoutingError::DuplicateEntity {
entity: "provider instance",
id: id.to_string(),
entity: "provider",
id: provider.to_string(),
});
}
}
let mut model_map =
BTreeMap::<ProviderInstanceId, BTreeMap<UpstreamModelId, ModelCapabilities>>::new();
BTreeMap::<ProviderKind, BTreeMap<UpstreamModelId, ModelCapabilities>>::new();
for model in provider_models {
if !instance_map.contains_key(model.instance()) {
if !provider_set.contains(model.provider()) {
return Err(RoutingError::NotFound {
entity: "provider instance",
id: model.instance().to_string(),
entity: "provider",
id: model.provider().to_string(),
});
}
let models = model_map.entry(model.instance).or_default();
let models = model_map.entry(model.provider).or_default();
let upstream_model = model.upstream_model;
if models
.insert(upstream_model.clone(), model.capabilities)
@@ -404,7 +357,7 @@ impl RuntimeSnapshot {
Ok(Self {
revision,
account_selection_policy,
instances: Arc::new(instance_map),
providers: Arc::new(provider_set),
provider_models: Arc::new(model_map),
provider_model_mappings: Arc::new(BTreeMap::new()),
provider_catalog_generations: Arc::new(BTreeMap::new()),
@@ -446,13 +399,7 @@ impl RuntimeSnapshot {
#[must_use]
pub fn public_models_for_provider(&self, provider: &ProviderKind) -> Vec<PublicModelId> {
let mut models = BTreeSet::new();
for (instance_id, discovered) in self.provider_models.iter() {
let Some(instance) = self.instances.get(instance_id) else {
continue;
};
if !instance.enabled || instance.provider() != provider {
continue;
}
if let Some(discovered) = self.provider_models.get(provider) {
models.extend(
discovered
.keys()
@@ -471,12 +418,7 @@ impl RuntimeSnapshot {
#[must_use]
pub fn public_models(&self) -> Vec<PublicModelId> {
let providers = self
.instances
.values()
.map(|instance| instance.provider().clone())
.collect::<BTreeSet<_>>();
providers
self.providers
.iter()
.flat_map(|provider| self.public_models_for_provider(provider))
.collect::<BTreeSet<_>>()
@@ -484,36 +426,14 @@ impl RuntimeSnapshot {
.collect()
}
/// 透明代理不以目录白名单拒绝模型；只要求该平台至少有一个启用 instance。
/// 透明代理不以目录白名单拒绝模型；只要求该 Provider 已注册。
#[must_use]
pub fn contains_public_model_for_provider(
&self,
_public_model: &PublicModelId,
provider: &ProviderKind,
) -> bool {
self.instances
.values()
.any(|instance| instance.enabled && instance.provider() == provider)
}
#[must_use]
pub fn instance_ids_for_provider(
&self,
provider: &ProviderKind,
) -> BTreeSet<ProviderInstanceId> {
self.instances
.iter()
.filter_map(|(id, instance)| {
(instance.enabled && instance.provider() == provider).then_some(id.clone())
})
.collect()
}
#[must_use]
pub fn provider_for_instance(&self, instance_id: &ProviderInstanceId) -> Option<&ProviderKind> {
self.instances
.get(instance_id)
.map(ProviderInstance::provider)
self.providers.contains(provider)
}
#[must_use]
@@ -535,22 +455,24 @@ impl RuntimeSnapshot {
operation: &Operation,
context: &RoutingContext,
) -> Result<RoutingPlan, RoutingError> {
let upstream_model = context.provider_kind.as_ref().map_or_else(
|| public_model.as_str().to_owned(),
|provider| self.mapped_model(provider, public_model.as_str()),
);
let upstream_model =
UpstreamModelId::new(upstream_model).map_err(|_| RoutingError::InvalidIdentifier)?;
let requirements = operation.capability_requirements();
let mut candidates = Vec::new();
for instance in self.instances.values() {
if !self.instance_is_eligible(instance, context) {
for provider in self.providers.iter() {
if context
.provider_kind
.as_ref()
.is_some_and(|expected| expected != provider)
|| context.blocked_providers.contains(provider)
{
continue;
}
let upstream_model =
UpstreamModelId::new(self.mapped_model(provider, public_model.as_str()))
.map_err(|_| RoutingError::InvalidIdentifier)?;
let emulated_features = match self
.provider_models
.get(instance.id())
.get(provider)
.and_then(|models| models.get(&upstream_model))
{
Some(capabilities) => {
@@ -562,13 +484,9 @@ impl RuntimeSnapshot {
None => BTreeSet::new(),
};
candidates.push(ProviderCandidate {
instance: instance.clone(),
upstream_model: upstream_model.clone(),
provider: provider.clone(),
upstream_model,
emulated_features,
observation_token: context
.provider_observation_tokens
.get(instance.id())
.copied(),
});
}
@@ -588,20 +506,6 @@ impl RuntimeSnapshot {
candidates: Arc::from(candidates),
})
}
fn instance_is_eligible(&self, instance: &ProviderInstance, context: &RoutingContext) -> bool {
instance.enabled
&& instance.health.is_routable(context.allow_degraded)
&& context
.provider_kind
.as_ref()
.is_none_or(|provider| instance.provider() == provider)
&& !context.blocked_instances.contains(instance.id())
&& context
.allowed_instances
.as_ref()
.is_none_or(|allowed| allowed.contains(instance.id()))
}
}
/// 请求级冻结失败；此状态必须 fail closed。
backend/crates/gateway-core/tests/engine/continuation.rs
+1 -3
View File
@@ -3,14 +3,13 @@ use gateway_core::engine::continuation::{
};
use gateway_core::engine::credential::ProviderAccountId;
use gateway_core::error::SafeUpstreamValue;
use gateway_core::routing::{ProviderInstanceId, ProviderKind};
use gateway_core::routing::ProviderKind;
fn pin() -> NativeContinuationPin {
NativeContinuationPin::new(
PreviousResponseId::new("response-private").expect("valid response"),
SafeUpstreamValue::new("upstream-private").expect("valid upstream response"),
ProviderKind::new("openai").expect("valid provider"),
ProviderInstanceId::new("inst_openai").expect("valid instance"),
ProviderAccountId::new("acct_codex").expect("valid account"),
)
}
@@ -35,7 +34,6 @@ fn external_binding_debug_should_redact_previous_response_id() {
fn native_pin_should_reject_different_account() {
assert!(!pin().matches(
&ProviderKind::new("openai").expect("valid provider"),
&ProviderInstanceId::new("inst_openai").expect("valid instance"),
&ProviderAccountId::new("acct_other").expect("valid account"),
));
}
backend/crates/gateway-core/tests/engine/coordinator.rs
+39 -184
View File
@@ -38,9 +38,8 @@ use gateway_core::operation::{
};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::{
ConfigRevision, InstanceHealth, ModelCapabilities, ProviderInstance, ProviderInstanceId,
ProviderKind, ProviderModel, PublicModelId, RoutingContext, RoutingPlan, RuntimeSnapshot,
UpstreamModelId,
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
RoutingPlan, RuntimeSnapshot, UpstreamModelId,
};
use serde_json::{Map, Value};
@@ -245,7 +244,6 @@ impl Provider for ScriptedProvider {
async fn query_model_capabilities(
&self,
_instance: &ProviderInstance,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
Ok(Vec::new())
}
@@ -268,7 +266,6 @@ impl Provider for ScriptedProvider {
let candidate = request.candidate();
let metadata = ProviderCallMetadata::new(
candidate.provider().clone(),
candidate.instance().clone(),
candidate.upstream_model().clone(),
ProviderResource::Account {
id: ProviderAccountId::new(account_id).expect("account id"),
@@ -293,7 +290,6 @@ impl Provider for ScriptedProvider {
let candidate = request.candidate();
let metadata = ProviderCallMetadata::new(
candidate.provider().clone(),
candidate.instance().clone(),
candidate.upstream_model().clone(),
ProviderResource::Account {
id: ProviderAccountId::new(account_id).expect("account id"),
@@ -375,29 +371,11 @@ fn image_stream(image_output_tokens: Option<u64>) -> Vec<Result<GatewayEvent, Pr
]
}
fn plan(operation: &Operation, instance_count: u32) -> RoutingPlan {
fn plan(operation: &Operation) -> RoutingPlan {
let provider = ProviderKind::new("openai").expect("provider");
let public_model = PublicModelId::new("gpt-5").expect("public model");
let capabilities =
ModelCapabilities::new(BTreeSet::from([operation.kind()]), 128_000, Some(32_000));
let mut instances = Vec::new();
let mut provider_models = Vec::new();
for index in 1..=instance_count {
let instance_id =
ProviderInstanceId::new(format!("inst_openai_{index}")).expect("instance id");
instances.push(ProviderInstance::new(
instance_id.clone(),
provider.clone(),
format!("https://openai-{index}.example"),
true,
InstanceHealth::Healthy,
));
provider_models.push(ProviderModel::new(
instance_id,
UpstreamModelId::new("gpt-5").expect("upstream model"),
capabilities.clone(),
));
}
RuntimeSnapshot::new(
ConfigRevision::new(1).expect("config revision"),
AccountSelectionPolicy::new(
@@ -405,8 +383,12 @@ fn plan(operation: &Operation, instance_count: u32) -> RoutingPlan {
NonZeroU32::new(2).expect("account concurrency"),
Duration::from_millis(50),
),
instances,
provider_models,
vec![provider.clone()],
vec![ProviderModel::new(
provider.clone(),
UpstreamModelId::new("gpt-5").expect("upstream model"),
capabilities,
)],
Vec::new(),
)
.expect("snapshot")
@@ -469,7 +451,7 @@ fn terminal_non_idempotent_failure(
continuation: Option<ContinuationBinding>,
) -> (Arc<FakeStore>, Arc<ScriptedProvider>) {
let operation = operation(RetrySafety::NonIdempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::Stream {
account_id: "acct_first",
@@ -500,7 +482,7 @@ fn terminal_non_idempotent_failure(
#[test]
fn success_updates_one_model_request_and_persists_usage() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, store, _) = coordinator(vec![Script::Stream {
account_id: "acct_one",
items: complete_stream(Some(12)),
@@ -535,7 +517,7 @@ fn success_updates_one_model_request_and_persists_usage() {
assert_eq!(
session.provider_attempt_outcomes(),
&[ProviderAttemptOutcome::Succeeded {
provider_instance_id: ProviderInstanceId::new("inst_openai_1").expect("instance"),
provider_kind: ProviderKind::new("openai").expect("provider"),
}]
);
let state = store.state.lock().expect("store lock");
@@ -561,7 +543,7 @@ fn success_updates_one_model_request_and_persists_usage() {
fn finalized_image_request(image_output_tokens: Option<u64>) -> FinalState {
let operation = image_generate_operation();
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, store, _) = coordinator(vec![Script::Stream {
account_id: "acct_image",
items: image_stream(image_output_tokens),
@@ -605,7 +587,7 @@ fn image_request_without_output_should_persist_failure() {
#[test]
fn failed_image_request_should_persist_failure() {
let operation = image_generate_operation();
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, store, _) = coordinator(vec![Script::Stream {
account_id: "acct_image",
items: vec![Err(ProviderError::new(
@@ -634,7 +616,7 @@ fn failed_image_request_should_persist_failure() {
#[test]
fn response_observation_is_persisted_but_never_delivered() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let observation = ProviderResponseObservation::new(
UpstreamTransport::new("http_sse").expect("actual transport"),
)
@@ -691,7 +673,7 @@ fn response_observation_is_persisted_but_never_delivered() {
#[test]
fn unknown_wire_event_before_response_identity_is_discarded_with_retried_attempt() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let unknown = ProtocolWireEvent::json(
"openai",
Some("response.future_event".to_owned()),
@@ -737,7 +719,7 @@ fn unknown_wire_event_before_response_identity_is_discarded_with_retried_attempt
#[test]
fn discarded_attempt_observation_does_not_leak_into_retry_result() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let first_observation = ProviderResponseObservation::new(
UpstreamTransport::new("websocket").expect("first transport"),
)
@@ -818,7 +800,7 @@ fn discarded_attempt_observation_does_not_leak_into_retry_result() {
#[test]
fn websocket_success_keeps_client_http_status_absent() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, store, _) = coordinator(vec![Script::Stream {
account_id: "acct_one",
items: complete_stream(None),
@@ -844,7 +826,7 @@ fn websocket_success_keeps_client_http_status_absent() {
#[test]
fn authenticated_native_continuation_reaches_every_attempt_context() {
let operation = generate_operation();
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, _, provider) = coordinator(vec![Script::Stream {
account_id: "acct_one",
items: complete_stream(None),
@@ -853,7 +835,6 @@ fn authenticated_native_continuation_reaches_every_attempt_context() {
PreviousResponseId::new("previous-secret-id").expect("previous response"),
SafeUpstreamValue::new("provider-native-id").expect("upstream response"),
ProviderKind::new("openai").expect("provider"),
ProviderInstanceId::new("inst_openai_1").expect("instance"),
ProviderAccountId::new("acct_one").expect("account"),
);
@@ -895,7 +876,6 @@ fn authenticated_native_continuation_reaches_every_attempt_context() {
.expect("continuation account-state owner");
assert!(owner.matches(
&ProviderKind::new("openai").expect("provider"),
&ProviderInstanceId::new("inst_openai_1").expect("instance"),
&ProviderAccountId::new("acct_one").expect("account"),
));
}
@@ -910,7 +890,7 @@ fn native_continuation_replays_owner_before_safely_switching_account() {
let operation = Operation::Generate(generate.with_provider_session_state(
ProviderSessionState::new("openai", payload).expect("provider session state"),
));
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, _, provider) = coordinator(vec![
Script::Stream {
account_id: "acct_one",
@@ -943,7 +923,6 @@ fn native_continuation_replays_owner_before_safely_switching_account() {
PreviousResponseId::new("previous-secret-id").expect("previous response"),
SafeUpstreamValue::new("provider-native-id").expect("upstream response"),
ProviderKind::new("openai").expect("provider"),
ProviderInstanceId::new("inst_openai_1").expect("instance"),
ProviderAccountId::new("acct_one").expect("account"),
);
@@ -983,7 +962,7 @@ fn native_continuation_replays_owner_before_safely_switching_account() {
#[test]
fn required_account_reaches_provider_and_matching_metadata_succeeds() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 2);
let route_plan = plan(&operation);
let (coordinator, _, provider) = coordinator(vec![Script::Stream {
account_id: "acct_required",
items: complete_stream(None),
@@ -1009,7 +988,7 @@ fn required_account_reaches_provider_and_matching_metadata_succeeds() {
#[test]
fn provider_metadata_for_another_account_fails_closed() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 2);
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::Stream {
account_id: "acct_wrong",
@@ -1043,93 +1022,10 @@ fn provider_metadata_for_another_account_fails_closed() {
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Failed);
}
#[test]
fn required_account_disables_instance_fallback_before_stream() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 2);
let (coordinator, store, provider) = coordinator(vec![
Script::Error(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::NotSent,
)),
Script::Stream {
account_id: "acct_required",
items: complete_stream(None),
},
]);
let required = ProviderAccountId::new("acct_required").expect("account id");
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
Some(required.clone()),
None,
CancellationToken::new(),
))
.expect("start execution");
let error = block_on(session.collect_uncommitted())
.expect_err("required account must not switch targets");
assert!(matches!(
error,
gateway_core::engine::EngineError::Provider(_)
));
let contexts = provider.contexts.lock().expect("contexts lock");
assert_eq!(contexts.len(), 1);
assert_eq!(contexts[0].required_account(), Some(&required));
assert_eq!(provider.scripts.lock().expect("scripts lock").len(), 1);
let state = store.state.lock().expect("store lock");
assert!(state.attempts.is_empty());
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Failed);
}
#[test]
fn local_selection_failure_falls_back_without_instance_failure_observation() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 2);
let (coordinator, store, provider) = coordinator(vec![
Script::Error(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::NotSent,
)),
Script::Stream {
account_id: "acct_second",
items: complete_stream(None),
},
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
block_on(session.collect_uncommitted()).expect("fallback response");
block_on(session.commit_downstream(Some(200))).expect("commit response");
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 2);
assert_eq!(
session.provider_attempt_outcomes(),
&[ProviderAttemptOutcome::Succeeded {
provider_instance_id: ProviderInstanceId::new("inst_openai_2")
.expect("second instance"),
}]
);
let state = store.state.lock().expect("store lock");
assert_eq!(state.attempts.len(), 1);
assert_eq!(
state.attempts[0].provider_instance_id,
ProviderInstanceId::new("inst_openai_2").expect("second instance")
);
}
#[test]
fn required_account_disables_account_retry_after_stream_creation() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::Stream {
account_id: "acct_required",
@@ -1162,7 +1058,7 @@ fn required_account_disables_account_retry_after_stream_creation() {
assert_eq!(
session.provider_attempt_outcomes(),
&[ProviderAttemptOutcome::Failed {
provider_instance_id: ProviderInstanceId::new("inst_openai_1").expect("instance"),
provider_kind: ProviderKind::new("openai").expect("provider"),
error_kind: ProviderErrorKind::Unavailable,
}]
);
@@ -1177,7 +1073,7 @@ fn required_account_disables_account_retry_after_stream_creation() {
#[test]
fn latest_provider_reported_cost_is_persisted_as_known_usd_total() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let events = vec![
Ok(GatewayEvent::Started(ResponseMeta::new(
"native-cost",
@@ -1220,7 +1116,7 @@ fn latest_provider_reported_cost_is_persisted_as_known_usd_total() {
#[test]
fn calculated_cost_is_persisted_when_provider_does_not_report_cost() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let events = vec![
Ok(GatewayEvent::Started(ResponseMeta::new(
"calculated-cost",
@@ -1258,7 +1154,7 @@ fn calculated_cost_is_persisted_when_provider_does_not_report_cost() {
#[test]
fn provider_reported_cost_should_not_be_replaced_by_calculated_cost() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let events = vec![
Ok(GatewayEvent::Started(ResponseMeta::new(
"reported-cost",
@@ -1307,7 +1203,7 @@ fn provider_reported_cost_should_not_be_replaced_by_calculated_cost() {
#[test]
fn discarded_attempt_cost_never_leaks_into_retry_result() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, store, _) = coordinator(vec![
Script::Stream {
account_id: "acct_first",
@@ -1352,7 +1248,7 @@ fn discarded_attempt_cost_never_leaks_into_retry_result() {
#[test]
fn pre_commit_failure_excludes_account_and_retries_same_target() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::Stream {
account_id: "acct_first",
@@ -1394,7 +1290,6 @@ fn pre_commit_failure_excludes_account_and_retries_same_target() {
.account_state_owner()
.is_some_and(|owner| owner.matches(
&ProviderKind::new("openai").expect("provider"),
&ProviderInstanceId::new("inst_openai_1").expect("instance"),
&ProviderAccountId::new("acct_first").expect("account"),
))
);
@@ -1417,7 +1312,7 @@ fn pre_commit_failure_excludes_account_and_retries_same_target() {
#[test]
fn recovered_credential_retries_the_same_account_exactly_once() {
let operation = operation(RetrySafety::NonIdempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::Stream {
account_id: "acct_first",
@@ -1462,7 +1357,7 @@ fn recovered_credential_retries_the_same_account_exactly_once() {
#[test]
fn non_idempotent_explicit_429_rejection_rotates_account_before_output() {
let operation = operation(RetrySafety::NonIdempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::Stream {
account_id: "acct_first",
@@ -1542,7 +1437,7 @@ fn explicit_429_with_ambiguous_send_state_is_not_retried() {
#[test]
fn explicit_429_after_structural_event_should_retry_before_commit() {
let operation = operation(RetrySafety::NonIdempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::Stream {
account_id: "acct_first",
@@ -1589,7 +1484,6 @@ fn native_continuation_explicit_429_is_not_retried() {
PreviousResponseId::new("previous-secret-id").expect("previous response"),
SafeUpstreamValue::new("provider-native-id").expect("upstream response"),
ProviderKind::new("openai").expect("provider"),
ProviderInstanceId::new("inst_openai_1").expect("instance"),
ProviderAccountId::new("acct_first").expect("account"),
);
let (store, provider) = terminal_non_idempotent_failure(
@@ -1632,49 +1526,10 @@ fn external_continuation_explicit_429_is_not_retried() {
assert_eq!(state.intermediate_failures, 0);
}
#[test]
fn non_idempotent_not_sent_failure_can_fallback_target() {
let operation = operation(RetrySafety::NonIdempotent);
let route_plan = plan(&operation, 2);
let (coordinator, store, provider) = coordinator(vec![
Script::Error(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::NotSent,
)),
Script::Stream {
account_id: "acct_second",
items: complete_stream(None),
},
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
block_on(session.collect_uncommitted()).expect("safe instance fallback succeeds");
block_on(session.commit_downstream(Some(200))).expect("commit response");
let contexts = provider.contexts.lock().expect("contexts lock");
assert_eq!(contexts.len(), 2);
assert!(contexts[0].account_state_owner().is_none());
assert!(contexts[1].account_state_owner().is_none());
let state = store.state.lock().expect("store lock");
assert_eq!(state.attempts.len(), 1);
assert_eq!(
state.attempts[0].provider_instance_id,
ProviderInstanceId::new("inst_openai_2").expect("second instance")
);
}
#[test]
fn ambiguous_send_state_stops_retry() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::Stream {
account_id: "acct_first",
@@ -1720,7 +1575,7 @@ fn ambiguous_send_state_stops_retry() {
#[test]
fn structural_event_before_replay_safe_failure_should_switch_account_before_commit() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::Stream {
account_id: "acct_first",
@@ -1763,7 +1618,7 @@ fn structural_event_before_replay_safe_failure_should_switch_account_before_comm
assert_eq!(
session.provider_attempt_outcomes(),
&[ProviderAttemptOutcome::Failed {
provider_instance_id: ProviderInstanceId::new("inst_openai_1").expect("instance"),
provider_kind: ProviderKind::new("openai").expect("provider"),
error_kind: ProviderErrorKind::Transport,
}]
);
@@ -1779,7 +1634,7 @@ fn structural_event_before_replay_safe_failure_should_switch_account_before_comm
#[test]
fn cancellation_before_pending_delivery_commit_reaches_terminal_state() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, store, _) = coordinator(vec![Script::Stream {
account_id: "acct_first",
items: complete_stream(None),
@@ -1811,9 +1666,9 @@ fn cancellation_before_pending_delivery_commit_reaches_terminal_state() {
}
#[test]
fn local_unavailability_before_stream_does_not_create_attempt_or_instance_failure() {
fn local_unavailability_before_stream_does_not_create_attempt_or_provider_failure() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, store, _) = coordinator(vec![Script::Error(
ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
.with_upstream_code(SafeUpstreamValue::new("overloaded").expect("provider code"))
@@ -1850,7 +1705,7 @@ fn local_unavailability_before_stream_does_not_create_attempt_or_instance_failur
#[test]
fn expired_deadline_finalizes_without_calling_provider() {
let operation = operation(RetrySafety::Idempotent);
let route_plan = plan(&operation, 1);
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![]);
let error = match block_on(coordinator.start(
backend/crates/gateway-core/tests/engine/credential.rs
+1 -2
View File
@@ -10,12 +10,11 @@ use gateway_core::engine::credential::{
CredentialRevision, OpaqueProviderData, PlaintextCredential, ProviderAccount,
ProviderAccountId, ProviderAccountUpdate, RotationStrategy,
};
use gateway_core::routing::{ProviderInstanceId, ProviderKind};
use gateway_core::routing::ProviderKind;
fn account(id: &str) -> ProviderAccount {
ProviderAccount::new(
ProviderAccountId::new(id).expect("valid account"),
ProviderInstanceId::new("inst_test").expect("valid instance"),
ProviderKind::new("openai").expect("valid provider"),
id.to_owned(),
format!("user-{id}"),
backend/crates/gateway-core/tests/engine/execution.rs
+226 -3
View File
@@ -1,8 +1,47 @@
use gateway_core::engine::execution::provider_failure_affects_circuit;
use gateway_core::error::ProviderErrorKind;
use std::{
collections::BTreeSet,
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
time::{Duration, SystemTime},
};
use async_trait::async_trait;
use futures::{executor::block_on, future::BoxFuture};
use gateway_core::engine::admission::{
ClientAdmissionDecision, ClientAdmissionError, ClientAdmissionPort, ClientAdmissionRecovery,
ClientAdmissionRequest, ClientAdmissionRestoreResult,
};
use gateway_core::engine::continuation::{
NativeContinuationPin, NativeContinuationPort, NativeContinuationStoreError, PreviousResponseId,
};
use gateway_core::engine::credential::{
AccountSelectionPolicy, ProviderAccountId, RotationStrategy,
};
use gateway_core::engine::execution::{
DefaultExecutionService, ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
provider_failure_affects_circuit,
};
use gateway_core::engine::probe::{AccountProbe, AccountProbeRequest};
use gateway_core::engine::provider::ProviderRegistry;
use gateway_core::engine::{
AttemptRecord, ExecutionStore, IntermediateFailure, ModelRequestFinalization, ModelRequestId,
NewModelRequest, RecoveryReport, UpstreamSendState,
};
use gateway_core::error::{GatewayErrorKind, ProviderErrorKind, StoreError};
use gateway_core::operation::{
ContentPart, GenerateRequest, Message, MessageRole, Operation, OperationKind,
};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::snapshot::RuntimeSnapshotHandle;
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, RuntimeSnapshot,
UpstreamModelId,
};
#[test]
fn only_instance_attributable_failures_should_affect_circuit() {
fn only_provider_attributable_failures_should_affect_circuit() {
assert!(provider_failure_affects_circuit(ProviderErrorKind::Timeout));
assert!(provider_failure_affects_circuit(
ProviderErrorKind::Transport
@@ -14,3 +53,187 @@ fn only_instance_attributable_failures_should_affect_circuit() {
ProviderErrorKind::InvalidRequest
));
}
#[test]
fn account_probe_should_not_write_to_the_persistent_execution_store() {
let store = Arc::new(TrackingExecutionStore::default());
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(probe_snapshot()),
store.clone(),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(UnusedContinuation),
);
let error = block_on(service.probe(AccountProbeRequest {
account_id: ProviderAccountId::new("acct_probe").expect("account ID"),
provider_kind: ProviderKind::new("openai").expect("provider kind"),
upstream_model: UpstreamModelId::new("gpt-probe").expect("model ID"),
operation: probe_operation(),
}))
.expect_err("empty Provider registry should stop the probe after it starts");
assert_eq!(error.kind(), GatewayErrorKind::NoAvailableProvider);
assert!(!store.touched.load(Ordering::SeqCst));
}
#[derive(Default)]
struct TrackingExecutionStore {
touched: AtomicBool,
}
impl TrackingExecutionStore {
fn touch(&self) {
self.touched.store(true, Ordering::SeqCst);
}
}
#[async_trait]
impl ExecutionStore for TrackingExecutionStore {
async fn create_model_request(&self, _: NewModelRequest) -> Result<(), StoreError> {
self.touch();
Ok(())
}
async fn record_attempt(&self, _: AttemptRecord) -> Result<(), StoreError> {
self.touch();
Ok(())
}
async fn mark_send_state(
&self,
_: &ModelRequestId,
_: UpstreamSendState,
) -> Result<(), StoreError> {
self.touch();
Ok(())
}
async fn mark_downstream_committed(
&self,
_: &ModelRequestId,
_: SystemTime,
_: Option<u16>,
) -> Result<(), StoreError> {
self.touch();
Ok(())
}
async fn record_client_status(&self, _: &ModelRequestId, _: u16) -> Result<(), StoreError> {
self.touch();
Ok(())
}
async fn record_intermediate_failure(&self, _: IntermediateFailure) -> Result<(), StoreError> {
self.touch();
Ok(())
}
async fn finalize_model_request(&self, _: ModelRequestFinalization) -> Result<(), StoreError> {
self.touch();
Ok(())
}
async fn recover_expired(&self, _: SystemTime) -> Result<RecoveryReport, StoreError> {
self.touch();
Ok(RecoveryReport::default())
}
}
struct UnusedAdmissions;
impl ClientAdmissionPort for UnusedAdmissions {
fn admit(
&self,
_: ClientAdmissionRequest,
) -> BoxFuture<'_, Result<ClientAdmissionDecision, ClientAdmissionError>> {
Box::pin(async { Ok(ClientAdmissionDecision::Granted) })
}
fn release<'a>(
&'a self,
_: &'a ClientApiKeyId,
_: &'a ModelRequestId,
) -> BoxFuture<'a, Result<bool, ClientAdmissionError>> {
Box::pin(async { Ok(true) })
}
fn restore(
&self,
_: ClientAdmissionRecovery,
) -> BoxFuture<'_, Result<ClientAdmissionRestoreResult, ClientAdmissionError>> {
Box::pin(async { Ok(ClientAdmissionRestoreResult::default()) })
}
}
struct UnusedCircuits;
impl ProviderCircuitPort for UnusedCircuits {
fn decision<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<ProviderCircuitDecision, ProviderCircuitError>> {
Box::pin(async { Ok(ProviderCircuitDecision::Allow) })
}
fn observe_failure<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { Ok(()) })
}
fn observe_success<'a>(
&'a self,
_: &'a ProviderKind,
) -> BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async { Ok(()) })
}
}
struct UnusedContinuation;
impl NativeContinuationPort for UnusedContinuation {
fn resolve<'a>(
&'a self,
_: &'a ClientApiKeyId,
_: &'a PreviousResponseId,
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>> {
Box::pin(async { Ok(None) })
}
}
fn probe_snapshot() -> RuntimeSnapshot {
let provider = ProviderKind::new("openai").expect("provider kind");
let capabilities = ModelCapabilities::new(
BTreeSet::from([OperationKind::Generate]),
128_000,
Some(16_000),
);
RuntimeSnapshot::new(
ConfigRevision::new(1).expect("config revision"),
AccountSelectionPolicy::new(
RotationStrategy::Smart,
std::num::NonZeroU32::new(1).expect("concurrency"),
Duration::from_millis(1),
),
vec![provider.clone()],
vec![ProviderModel::new(
provider,
UpstreamModelId::new("gpt-probe").expect("model ID"),
capabilities,
)],
Vec::new(),
)
.expect("probe snapshot")
}
fn probe_operation() -> Operation {
let message = Message::new(
MessageRole::User,
vec![ContentPart::Text("ping".to_owned())],
)
.expect("probe message");
Operation::Generate(GenerateRequest::new(vec![message]).expect("probe request"))
}
backend/crates/gateway-core/tests/engine/provider.rs
+3 -15
View File
@@ -13,10 +13,7 @@ use gateway_core::engine::provider::{
use gateway_core::engine::{AttemptContext, UpstreamSendState};
use gateway_core::error::{IdentifierError, ProviderError, ProviderErrorKind};
use gateway_core::operation::OperationKind;
use gateway_core::routing::{
InstanceHealth, ModelCapabilities, ProviderInstance, ProviderInstanceId, ProviderKind,
UpstreamModelId,
};
use gateway_core::routing::{ModelCapabilities, ProviderKind, UpstreamModelId};
struct NamedProvider(&'static str);
@@ -32,7 +29,6 @@ impl Provider for NamedProvider {
async fn query_model_capabilities(
&self,
_instance: &ProviderInstance,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
Ok(vec![ProviderModelCapabilities::new(
UpstreamModelId::new("live-model").expect("model"),
@@ -78,15 +74,8 @@ fn registry_should_query_provider_compiled_model_capabilities() {
.register(Arc::new(NamedProvider("openai")))
.expect("provider");
let registry = builder.build();
let instance = ProviderInstance::new(
ProviderInstanceId::new("inst_openai").expect("instance"),
ProviderKind::new("openai").expect("provider kind"),
"https://api.example".to_owned(),
true,
InstanceHealth::Healthy,
);
let models = futures::executor::block_on(registry.query_model_capabilities(&instance))
let provider = ProviderKind::new("openai").expect("provider kind");
let models = futures::executor::block_on(registry.query_model_capabilities(&provider))
.expect("live catalog");
assert_eq!(models[0].upstream_model().as_str(), "live-model");
@@ -113,7 +102,6 @@ fn provider_stream_should_release_owned_lease_on_drop() {
let released = Arc::new(AtomicBool::new(false));
let metadata = ProviderCallMetadata::new(
ProviderKind::new("openai").expect("valid provider"),
ProviderInstanceId::new("inst_openai").expect("valid instance"),
UpstreamModelId::new("gpt-5").expect("valid model"),
ProviderResource::Anonymous(ResourceId::none()),
UpstreamTransport::new("http_sse").expect("valid transport"),
backend/crates/gateway-core/tests/routing/mod.rs
+20 -38
View File
@@ -9,8 +9,8 @@ use gateway_core::operation::{
};
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::{
ConfigRevision, InstanceHealth, ModelCapabilities, ProviderInstance, ProviderInstanceId,
ProviderKind, ProviderModel, PublicModelId, RoutingContext, RuntimeSnapshot, UpstreamModelId,
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
RuntimeSnapshot, UpstreamModelId,
};
mod snapshot;
@@ -31,19 +31,9 @@ fn capabilities() -> ModelCapabilities {
)
}
fn instance(id: &str, provider: &str) -> ProviderInstance {
ProviderInstance::new(
ProviderInstanceId::new(id).expect("valid instance"),
ProviderKind::new(provider).expect("valid provider"),
format!("https://{provider}.example.invalid"),
true,
InstanceHealth::Healthy,
)
}
fn model(instance: &str, name: &str, capabilities: ModelCapabilities) -> ProviderModel {
fn model(provider: &str, name: &str, capabilities: ModelCapabilities) -> ProviderModel {
ProviderModel::new(
ProviderInstanceId::new(instance).expect("valid instance"),
ProviderKind::new(provider).expect("valid provider"),
UpstreamModelId::new(name).expect("valid model"),
capabilities,
)
@@ -73,12 +63,12 @@ fn snapshot() -> RuntimeSnapshot {
ConfigRevision::new(1).expect("revision"),
scheduling(),
vec![
instance("inst_openai", "openai"),
instance("inst_xai", "xai"),
ProviderKind::new("openai").expect("provider"),
ProviderKind::new("xai").expect("provider"),
],
vec![
model("inst_openai", "gpt-5.5", capabilities()),
model("inst_xai", "grok-4.5", capabilities()),
model("openai", "gpt-5.5", capabilities()),
model("xai", "grok-4.5", capabilities()),
],
Vec::new(),
)
@@ -100,17 +90,12 @@ fn config_revision_should_reject_zero() {
assert!(ConfigRevision::new(0).is_err());
}
#[test]
fn provider_instance_id_should_require_prefix() {
assert!(ProviderInstanceId::new("missing").is_err());
}
#[test]
fn snapshot_should_publish_only_enabled_plaintext_client_policies() {
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("revision"),
scheduling(),
vec![instance("inst_openai", "openai")],
vec![ProviderKind::new("openai").expect("provider")],
Vec::new(),
vec![
client_policy("key_enabled", "sk_enabled", true),
@@ -126,12 +111,12 @@ fn snapshot_should_publish_only_enabled_plaintext_client_policies() {
}
#[test]
fn snapshot_should_reject_model_for_missing_instance() {
fn snapshot_should_reject_model_for_missing_provider() {
let result = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("revision"),
scheduling(),
Vec::new(),
vec![model("inst_missing", "gpt-5.5", capabilities())],
vec![model("missing", "gpt-5.5", capabilities())],
Vec::new(),
);
@@ -139,14 +124,14 @@ fn snapshot_should_reject_model_for_missing_instance() {
}
#[test]
fn snapshot_should_reject_duplicate_instance_model() {
fn snapshot_should_reject_duplicate_provider_model() {
let result = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("revision"),
scheduling(),
vec![instance("inst_openai", "openai")],
vec![ProviderKind::new("openai").expect("provider")],
vec![
model("inst_openai", "gpt-5.5", capabilities()),
model("inst_openai", "gpt-5.5", capabilities()),
model("openai", "gpt-5.5", capabilities()),
model("openai", "gpt-5.5", capabilities()),
],
Vec::new(),
);
@@ -212,7 +197,7 @@ fn unmapped_model_should_pass_through_unchanged() {
}
#[test]
fn blocked_instance_should_be_filtered() {
fn blocked_provider_should_be_filtered() {
let snapshot = snapshot();
let error = snapshot
.plan(
@@ -220,10 +205,7 @@ fn blocked_instance_should_be_filtered() {
&operation(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
blocked_instances: BTreeSet::from([
ProviderInstanceId::new("inst_openai").expect("instance")
]),
..RoutingContext::default()
blocked_providers: BTreeSet::from([ProviderKind::new("openai").expect("provider")]),
},
)
.expect_err("blocked platform has no candidate");
@@ -239,9 +221,9 @@ fn known_unsupported_operation_should_not_be_bypassed() {
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("revision"),
scheduling(),
vec![instance("inst_openai", "openai")],
vec![ProviderKind::new("openai").expect("provider")],
vec![model(
"inst_openai",
"openai",
"gpt-known-unsupported",
ModelCapabilities::new(BTreeSet::new(), 0, None),
)],
@@ -305,7 +287,7 @@ fn public_catalog_should_include_discovered_models_and_aliases() {
}
#[test]
fn unknown_model_should_be_accepted_when_platform_has_instance() {
fn unknown_model_should_be_accepted_when_provider_is_registered() {
assert!(snapshot().contains_public_model_for_provider(
&PublicModelId::new("future-model").expect("model"),
&ProviderKind::new("openai").expect("provider"),
backend/crates/gateway-core/tests/routing/snapshot.rs
+42 -13
View File
@@ -20,13 +20,12 @@ use gateway_core::policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::snapshot::{
RuntimeSnapshotCompileError, RuntimeSnapshotCompiler, RuntimeSnapshotHandle,
RuntimeSnapshotPublisher, SnapshotClientPolicyFacts, SnapshotControl, SnapshotFacts,
SnapshotProviderInstanceFacts, SnapshotRevisionStream, SnapshotSettingsFacts,
SnapshotStoreError, SnapshotStorePort, SnapshotSubscriptionError, SnapshotSubscriptionPort,
runtime_revision_needs_refresh,
SnapshotRevisionStream, SnapshotSettingsFacts, SnapshotStoreError, SnapshotStorePort,
SnapshotSubscriptionError, SnapshotSubscriptionPort, runtime_revision_needs_refresh,
};
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderInstance, ProviderKind, PublicModelId,
RuntimeSnapshot, UpstreamModelId,
ConfigRevision, ModelCapabilities, ProviderKind, PublicModelId, RuntimeSnapshot,
UpstreamModelId,
};
use gateway_core::task::WorkerKind;
@@ -87,6 +86,39 @@ struct PublishingCatalogProvider {
queries: AtomicUsize,
}
struct UnavailableCatalogProvider;
#[async_trait]
impl Provider for UnavailableCatalogProvider {
fn name(&self) -> &'static str {
"alpha"
}
fn catalog_generation(&self) -> ProviderCatalogGeneration {
ProviderCatalogGeneration::new(0)
}
async fn query_model_capabilities(
&self,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
Err(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::NotSent,
))
}
async fn execute(
&self,
_: ProviderRequest,
_: AttemptContext,
) -> Result<ProviderStream, ProviderError> {
Err(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::NotSent,
))
}
}
#[async_trait]
impl Provider for PublishingCatalogProvider {
fn name(&self) -> &'static str {
@@ -99,7 +131,6 @@ impl Provider for PublishingCatalogProvider {
async fn query_model_capabilities(
&self,
_: &ProviderInstance,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
if self.queries.fetch_add(1, Ordering::SeqCst) == 0 {
self.generation.store(1, Ordering::SeqCst);
@@ -146,7 +177,11 @@ fn compiler_should_reject_revision_changed_during_consistent_read() {
#[test]
fn compiler_should_preserve_passthrough_when_provider_catalog_is_unavailable() {
let compiler = compiler(Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))));
let providers =
ProviderRegistry::new([Arc::new(UnavailableCatalogProvider) as Arc<dyn Provider>])
.expect("provider registry");
let compiler =
RuntimeSnapshotCompiler::new(Arc::new(TestSnapshotStore::new(Ok(facts(3, 3)))), providers);
let snapshot = block_on(compiler.compile()).expect("compile snapshot");
let provider = ProviderKind::new("alpha").expect("provider");
@@ -302,12 +337,6 @@ fn facts(config_revision: u64, observed_current_revision: u64) -> SnapshotFacts
BTreeMap::from([("public-model".to_owned(), "upstream-model".to_owned())]),
)]),
),
vec![SnapshotProviderInstanceFacts::new(
"inst_alpha",
"alpha",
"https://alpha.example.invalid",
true,
)],
vec![SnapshotClientPolicyFacts::new(
ClientApiKeyId::new("key_one").expect("key ID"),
PlaintextClientApiKey::new("sk_test").expect("plaintext key"),
backend/crates/gateway-store/src/lib.rs
+3 -8
View File
@@ -1,6 +1,6 @@
//! 多 Provider 网关的 PostgreSQL 持久化与 Redis 协调 adapter。
//!
//! 业务规则与 port 由 `gateway-core` 拥有。本 crate 只负责把终态的八张业务表
//! 业务规则与 port 由 `gateway-core` 拥有。本 crate 只负责把终态的七张业务表
//! 和可丢失 Redis 状态映射为明确的基础设施操作。
use std::sync::Arc;
@@ -19,7 +19,7 @@ use gateway_admin::ports::store::{
};
use gateway_core::CoreStorePorts;
use gateway_core::health::{HealthProbe, HealthState};
use gateway_core::provider_ports::{ProviderCatalogPorts, ProviderStorePorts};
use gateway_core::provider_ports::ProviderStorePorts;
use gateway_core::task::{
ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerDisabledReason, WorkerId,
WorkerKind, WorkerLeaderLeasePort, WorkerLeaseRequest, WorkerRegistration, WorkerRunnable,
@@ -256,10 +256,6 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
redis_connection.clone(),
REDIS_NAMESPACE,
)?);
let provider_catalog = ProviderCatalogPorts::new(
Arc::new(postgres::PgConfigCatalogRepository::new(pool.clone())),
credential_state.clone(),
);
let runtime_policy = Arc::new(postgres::PgRuntimeSettingsRepository::new(pool.clone()));
let oauth_pending = Arc::new(redis::RedisOAuthPendingFlowRepository::new(
redis_connection.clone(),
@@ -276,7 +272,6 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
REDIS_NAMESPACE,
)?,
}),
Arc::new(postgres::PgAdminCatalogStore::new(pool.clone())),
Arc::new(postgres::PgAdminClientKeyStore::new(pool.clone())),
Arc::new(postgres::PgAdminObservabilityStore::new(pool.clone())),
Arc::new(AdminSettingsStoreAdapter {
@@ -311,7 +306,7 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
let provider_ports = ProviderStorePorts::new(
account_store,
provider_leases,
provider_catalog,
credential_state.clone(),
credential_state,
cooldowns,
runtime_policy,
backend/crates/gateway-store/src/postgres/config_catalog.rs
-593
View File
@@ -1,593 +0,0 @@
//! `provider_instances` 的 PostgreSQL owner。
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use gateway_admin::{
model::{
MutationContext,
catalog::{
CreateProviderInstance as AdminCreateProviderInstance,
DeleteProviderInstance as AdminDeleteProviderInstance,
ProviderInstance as AdminProviderInstance, ProviderInstanceCatalog,
ProviderInstanceDetail, ProviderInstanceMutation, SetProviderInstanceEnabled,
UpdateProviderInstance as AdminUpdateProviderInstance,
},
},
ports::store::{AdminStoreResult, CatalogStore},
};
use gateway_core::provider_ports::{
ProviderInstanceCatalogPort, ProviderInstanceConfig, ProviderStoreError, ProviderStoreErrorKind,
};
use gateway_core::routing::{ProviderInstanceId, ProviderKind};
use sqlx::{PgPool, Postgres, Transaction};
use crate::{
StoreError, StoreResult, admin_revision, admin_store_error, mutation_audit,
postgres_unavailable, require_nonempty, store_revision,
};
use super::{ControlPlaneRepository, PgControlPlaneRepository};
const ENTITY: &str = "provider instance";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderInstanceRecord {
pub id: String,
pub provider_kind: String,
pub name: String,
pub base_url: String,
pub enabled: bool,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NewProviderInstance {
pub id: String,
pub provider_kind: String,
pub name: String,
pub base_url: String,
}
impl NewProviderInstance {
pub fn validate(&self) -> StoreResult<()> {
validate_instance_fields(&self.id, &self.provider_kind, &self.name, &self.base_url)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpdateProviderInstance {
pub id: String,
pub name: String,
pub base_url: String,
pub enabled: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpdateProviderInstanceDetails {
pub id: String,
pub name: String,
pub base_url: String,
}
impl UpdateProviderInstanceDetails {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "id", &self.id)?;
require_nonempty(ENTITY, "name", &self.name)?;
require_nonempty(ENTITY, "base_url", &self.base_url)
}
}
#[async_trait]
pub trait ConfigCatalogRepository: Send + Sync {
async fn get_provider_instance(&self, id: &str) -> StoreResult<Option<ProviderInstanceRecord>>;
async fn list_provider_instances(
&self,
include_disabled: bool,
) -> StoreResult<Vec<ProviderInstanceRecord>>;
async fn insert_provider_instance(&self, instance: NewProviderInstance) -> StoreResult<()>;
async fn update_provider_instance(&self, instance: UpdateProviderInstance)
-> StoreResult<bool>;
async fn delete_provider_instance(&self, id: &str) -> StoreResult<bool>;
}
#[derive(Clone)]
pub struct PgConfigCatalogRepository {
pool: PgPool,
}
impl PgConfigCatalogRepository {
#[must_use]
pub const fn new(pool: PgPool) -> Self {
Self { pool }
}
}
#[async_trait]
impl ConfigCatalogRepository for PgConfigCatalogRepository {
async fn get_provider_instance(&self, id: &str) -> StoreResult<Option<ProviderInstanceRecord>> {
require_nonempty(ENTITY, "id", id)?;
let row = sqlx::query_as::<
_,
(
String,
String,
String,
String,
bool,
DateTime<Utc>,
DateTime<Utc>,
),
>(
"select id, provider_kind, name, base_url, enabled, created_at, updated_at
from provider_instances where id = $1",
)
.bind(id)
.fetch_optional(&self.pool)
.await
.map_err(|_| postgres_unavailable("get provider instance"))?;
Ok(row.map(provider_instance_from_row))
}
async fn list_provider_instances(
&self,
include_disabled: bool,
) -> StoreResult<Vec<ProviderInstanceRecord>> {
let rows = sqlx::query_as::<
_,
(
String,
String,
String,
String,
bool,
DateTime<Utc>,
DateTime<Utc>,
),
>(
"select id, provider_kind, name, base_url, enabled, created_at, updated_at
from provider_instances
where $1 or enabled
order by provider_kind, name, id",
)
.bind(include_disabled)
.fetch_all(&self.pool)
.await
.map_err(|_| postgres_unavailable("list provider instances"))?;
Ok(rows.into_iter().map(provider_instance_from_row).collect())
}
async fn insert_provider_instance(&self, instance: NewProviderInstance) -> StoreResult<()> {
instance.validate()?;
sqlx::query(
"insert into provider_instances
(id, provider_kind, name, base_url, enabled, created_at, updated_at)
values ($1, $2, $3, $4, true, now(), now())",
)
.bind(instance.id)
.bind(instance.provider_kind)
.bind(instance.name)
.bind(instance.base_url)
.execute(&self.pool)
.await
.map_err(|_| postgres_unavailable("insert provider instance"))?;
Ok(())
}
async fn update_provider_instance(
&self,
instance: UpdateProviderInstance,
) -> StoreResult<bool> {
require_nonempty(ENTITY, "id", &instance.id)?;
require_nonempty(ENTITY, "name", &instance.name)?;
require_nonempty(ENTITY, "base_url", &instance.base_url)?;
let result = sqlx::query(
"update provider_instances
set name = $2, base_url = $3, enabled = $4, updated_at = now()
where id = $1",
)
.bind(instance.id)
.bind(instance.name)
.bind(instance.base_url)
.bind(instance.enabled)
.execute(&self.pool)
.await
.map_err(|_| postgres_unavailable("update provider instance"))?;
Ok(result.rows_affected() == 1)
}
async fn delete_provider_instance(&self, id: &str) -> StoreResult<bool> {
require_nonempty(ENTITY, "id", id)?;
let result = sqlx::query("delete from provider_instances where id = $1")
.bind(id)
.execute(&self.pool)
.await
.map_err(|_| postgres_unavailable("delete provider instance"))?;
Ok(result.rows_affected() == 1)
}
}
impl ProviderInstanceCatalogPort for PgConfigCatalogRepository {
fn list_instances<'a>(
&'a self,
provider_kind: &'a ProviderKind,
include_disabled: bool,
) -> futures::future::BoxFuture<'a, Result<Vec<ProviderInstanceConfig>, ProviderStoreError>>
{
Box::pin(async move {
self.list_provider_instances(include_disabled)
.await
.map_err(|_| {
ProviderStoreError::new(
ProviderStoreErrorKind::Unavailable,
"list Provider instances",
)
})?
.into_iter()
.filter(|instance| instance.provider_kind == provider_kind.as_str())
.map(|instance| {
let id = ProviderInstanceId::new(instance.id).map_err(|_| {
ProviderStoreError::new(
ProviderStoreErrorKind::InvalidData,
"decode Provider instance ID",
)
})?;
let kind = ProviderKind::new(instance.provider_kind).map_err(|_| {
ProviderStoreError::new(
ProviderStoreErrorKind::InvalidData,
"decode Provider kind",
)
})?;
Ok(ProviderInstanceConfig::new(
id,
kind,
instance.base_url,
instance.enabled,
))
})
.collect()
})
}
}
/// Admin 用例所需的 Provider instance 事务能力。
#[derive(Clone)]
pub struct PgAdminCatalogStore {
catalog: PgConfigCatalogRepository,
control_plane: PgControlPlaneRepository,
}
impl PgAdminCatalogStore {
#[must_use]
pub fn new(pool: PgPool) -> Self {
Self {
catalog: PgConfigCatalogRepository::new(pool.clone()),
control_plane: PgControlPlaneRepository::new(pool),
}
}
async fn revision(&self) -> AdminStoreResult<gateway_admin::model::Revision> {
self.control_plane
.load_control_plane()
.await
.map_err(|error| admin_store_error("provider instance", error))
.and_then(|snapshot| admin_revision(snapshot.settings.config_revision))
}
async fn load_mutation(
&self,
revision: crate::Revision,
id: &ProviderInstanceId,
) -> AdminStoreResult<ProviderInstanceMutation> {
let instance = self
.catalog
.get_provider_instance(id.as_str())
.await
.map_err(|error| admin_store_error("provider instance", error))?
.map(admin_provider_instance)
.transpose()?;
Ok(ProviderInstanceMutation {
config_revision: admin_revision(revision)?,
instance,
})
}
}
#[async_trait]
impl CatalogStore for PgAdminCatalogStore {
async fn list_provider_instances(
&self,
include_disabled: bool,
) -> AdminStoreResult<ProviderInstanceCatalog> {
let config_revision = self.revision().await?;
let items = self
.catalog
.list_provider_instances(include_disabled)
.await
.map_err(|error| admin_store_error("provider instance", error))?
.into_iter()
.map(admin_provider_instance)
.collect::<AdminStoreResult<Vec<_>>>()?;
Ok(ProviderInstanceCatalog {
config_revision,
items,
})
}
async fn load_provider_instance(
&self,
id: &ProviderInstanceId,
) -> AdminStoreResult<Option<ProviderInstanceDetail>> {
let (revision, instance) = futures::try_join!(self.revision(), async {
self.catalog
.get_provider_instance(id.as_str())
.await
.map_err(|error| admin_store_error("provider instance", error))?
.map(admin_provider_instance)
.transpose()
},)?;
Ok(instance.map(|item| ProviderInstanceDetail {
config_revision: revision,
item,
}))
}
async fn create_provider_instance(
&self,
command: AdminCreateProviderInstance,
context: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation> {
let id = command.id;
let revision = self
.control_plane
.create_provider_instance(
store_revision(command.expected_config_revision)?,
NewProviderInstance {
id: id.as_str().to_owned(),
provider_kind: command.provider_kind.as_str().to_owned(),
name: command.name,
base_url: command.base_url,
},
mutation_audit(
context,
"create",
"provider_instance",
id.as_str(),
vec![
"provider_kind".to_owned(),
"name".to_owned(),
"base_url".to_owned(),
"enabled".to_owned(),
],
),
)
.await
.map_err(|error| admin_store_error("provider instance", error))?;
self.load_mutation(revision, &id).await
}
async fn update_provider_instance(
&self,
command: AdminUpdateProviderInstance,
context: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation> {
let id = command.id;
let revision = self
.control_plane
.update_provider_instance(
store_revision(command.expected_config_revision)?,
UpdateProviderInstanceDetails {
id: id.as_str().to_owned(),
name: command.name,
base_url: command.base_url,
},
mutation_audit(
context,
"update",
"provider_instance",
id.as_str(),
vec!["name".to_owned(), "base_url".to_owned()],
),
)
.await
.map_err(|error| admin_store_error("provider instance", error))?;
self.load_mutation(revision, &id).await
}
async fn set_provider_instance_enabled(
&self,
command: SetProviderInstanceEnabled,
context: &MutationContext,
) -> AdminStoreResult<ProviderInstanceMutation> {
let id = command.id;
let revision = self
.control_plane
.set_provider_instance_enabled(
store_revision(command.expected_config_revision)?,
id.as_str(),
command.enabled,
mutation_audit(
context,
if command.enabled { "enable" } else { "disable" },
"provider_instance",
id.as_str(),
vec!["enabled".to_owned()],
),
)
.await
.map_err(|error| admin_store_error("provider instance", error))?;
self.load_mutation(revision, &id).await
}
async fn delete_provider_instance(
&self,
command: AdminDeleteProviderInstance,
context: &MutationContext,
) -> AdminStoreResult<gateway_admin::model::Revision> {
self.control_plane
.delete_provider_instance(
store_revision(command.expected_config_revision)?,
command.id.as_str(),
mutation_audit(
context,
"delete",
"provider_instance",
command.id.as_str(),
Vec::new(),
),
)
.await
.map_err(|error| admin_store_error("provider instance", error))
.and_then(admin_revision)
}
}
fn admin_provider_instance(
record: ProviderInstanceRecord,
) -> AdminStoreResult<AdminProviderInstance> {
Ok(AdminProviderInstance {
id: ProviderInstanceId::new(record.id).map_err(|_| {
admin_store_error(
"provider instance",
StoreError::InvalidData {
entity: ENTITY,
message: "persisted instance ID is invalid".to_owned(),
},
)
})?,
provider_kind: ProviderKind::new(record.provider_kind).map_err(|_| {
admin_store_error(
"provider instance",
StoreError::InvalidData {
entity: ENTITY,
message: "persisted provider kind is invalid".to_owned(),
},
)
})?,
name: record.name,
base_url: record.base_url,
enabled: record.enabled,
created_at: record.created_at,
updated_at: record.updated_at,
})
}
pub(crate) async fn insert_provider_instance_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
instance: &NewProviderInstance,
) -> StoreResult<()> {
instance.validate()?;
sqlx::query(
"insert into provider_instances
(id, provider_kind, name, base_url, enabled, created_at, updated_at)
values ($1, $2, $3, $4, true, now(), now())",
)
.bind(&instance.id)
.bind(&instance.provider_kind)
.bind(&instance.name)
.bind(&instance.base_url)
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("insert provider instance in transaction"))?;
Ok(())
}
pub(crate) async fn update_provider_instance_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
instance: &UpdateProviderInstanceDetails,
) -> StoreResult<()> {
instance.validate()?;
let result = sqlx::query(
"update provider_instances
set name = $2, base_url = $3, updated_at = now()
where id = $1",
)
.bind(&instance.id)
.bind(&instance.name)
.bind(&instance.base_url)
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("update provider instance in transaction"))?;
require_changed(result.rows_affected(), &instance.id)
}
pub(crate) async fn set_provider_instance_enabled_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
id: &str,
enabled: bool,
) -> StoreResult<()> {
require_nonempty(ENTITY, "id", id)?;
let result =
sqlx::query("update provider_instances set enabled = $2, updated_at = now() where id = $1")
.bind(id)
.bind(enabled)
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("set provider instance state in transaction"))?;
require_changed(result.rows_affected(), id)
}
pub(crate) async fn delete_provider_instance_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
id: &str,
) -> StoreResult<()> {
require_nonempty(ENTITY, "id", id)?;
let result = sqlx::query("delete from provider_instances where id = $1")
.bind(id)
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("delete provider instance in transaction"))?;
require_changed(result.rows_affected(), id)
}
fn require_changed(rows_affected: u64, id: &str) -> StoreResult<()> {
if rows_affected == 1 {
Ok(())
} else {
Err(StoreError::NotFound {
entity: ENTITY,
id: id.to_owned(),
})
}
}
type ProviderInstanceRow = (
String,
String,
String,
String,
bool,
DateTime<Utc>,
DateTime<Utc>,
);
fn provider_instance_from_row(row: ProviderInstanceRow) -> ProviderInstanceRecord {
ProviderInstanceRecord {
id: row.0,
provider_kind: row.1,
name: row.2,
base_url: row.3,
enabled: row.4,
created_at: row.5,
updated_at: row.6,
}
}
fn validate_instance_fields(id: &str, kind: &str, name: &str, base_url: &str) -> StoreResult<()> {
require_nonempty(ENTITY, "id", id)?;
require_nonempty(ENTITY, "provider_kind", kind)?;
require_nonempty(ENTITY, "name", name)?;
require_nonempty(ENTITY, "base_url", base_url)?;
let valid_kind = kind.bytes().enumerate().all(|(index, byte)| {
byte.is_ascii_lowercase()
|| byte.is_ascii_digit()
|| (index > 0 && matches!(byte, b'-' | b'_'))
});
if !valid_kind {
return Err(StoreError::InvalidData {
entity: ENTITY,
message: "provider_kind must be a stable lowercase slug".to_owned(),
});
}
Ok(())
}
backend/crates/gateway-store/src/postgres/execution.rs
+8 -14
View File
@@ -129,7 +129,6 @@ impl NewModelRequest {
pub struct ModelRequestAttemptStart {
pub model_request_id: String,
pub attempt_count: u32,
pub provider_instance_id: String,
pub provider_kind: String,
pub provider_account_id: Option<String>,
pub provider_account_ref: Option<String>,
@@ -142,7 +141,6 @@ impl ModelRequestAttemptStart {
pub fn validate(&self) -> StoreResult<()> {
for (field, value) in [
("model_request_id", self.model_request_id.as_str()),
("provider_instance_id", self.provider_instance_id.as_str()),
("provider_kind", self.provider_kind.as_str()),
("upstream_model_id", self.upstream_model_id.as_str()),
("upstream_transport", self.upstream_transport.as_str()),
@@ -388,21 +386,19 @@ impl ModelRequestRepository for PgExecutionStore {
attempt.validate()?;
let count = sqlx::query_scalar::<_, i32>(
"update model_requests
set provider_instance_id = $2,
provider_kind = $3,
provider_account_id = $4,
provider_account_ref = $5,
upstream_model_id = $6,
upstream_transport = $7,
http_version = $8,
attempt_count = $9,
set provider_kind = $2,
provider_account_id = $3,
provider_account_ref = $4,
upstream_model_id = $5,
upstream_transport = $6,
http_version = $7,
attempt_count = $8,
upstream_send_state = 'not_sent'
where id = $1 and outcome = 'running' and downstream_committed_at is null
and $9 = attempt_count + 1
and $8 = attempt_count + 1
returning attempt_count",
)
.bind(&attempt.model_request_id)
.bind(attempt.provider_instance_id)
.bind(attempt.provider_kind)
.bind(attempt.provider_account_id)
.bind(attempt.provider_account_ref)
@@ -665,7 +661,6 @@ impl ExecutionStore for PgExecutionStore {
.begin_model_request_attempt(ModelRequestAttemptStart {
model_request_id: attempt.request_id.as_str().to_owned(),
attempt_count: expected_count,
provider_instance_id: attempt.provider_instance_id.as_str().to_owned(),
provider_kind: attempt.provider_kind.as_str().to_owned(),
provider_account_id: attempt
.provider_account_id
@@ -751,7 +746,6 @@ impl ExecutionStore for PgExecutionStore {
level: super::OpsEventLevel::Warning,
component: "routing".to_owned(),
operation: failure.trigger.as_str().to_owned(),
provider_instance_id: Some(failure.instance_id.as_str().to_owned()),
provider_kind: Some(failure.provider_kind.as_str().to_owned()),
provider_account_id: failure.account_id.as_ref().map(|id| id.as_str().to_owned()),
provider_account_ref: failure.account_id.as_ref().map(|id| id.as_str().to_owned()),
backend/crates/gateway-store/src/postgres/history.rs
+17 -28
View File
@@ -10,7 +10,7 @@ use gateway_core::{
credential::ProviderAccountId,
},
error::SafeUpstreamValue,
routing::{ProviderInstanceId, ProviderKind},
routing::ProviderKind,
};
use sqlx::PgPool;
@@ -22,7 +22,6 @@ pub struct ModelRequestHistoryRecord {
pub client_api_key_ref: String,
pub requested_model_id: String,
pub provider_kind: Option<String>,
pub provider_instance_id: Option<String>,
pub provider_account_ref: Option<String>,
pub upstream_model_id: Option<String>,
pub outcome: String,
@@ -112,16 +111,11 @@ impl ModelRequestHistoryRepository for PgHistoryRepository {
"caller_client_api_key_ref",
caller_client_api_key_ref,
)?;
let row = sqlx::query_as::<_, (String, String, String, String)>(
"select mr.provider_kind, mr.provider_instance_id,
mr.provider_account_id, mr.upstream_response_id
let row = sqlx::query_as::<_, (String, String, String)>(
"select mr.provider_kind, mr.provider_account_id, mr.upstream_response_id
from model_requests mr
join provider_instances instance
on instance.id = mr.provider_instance_id
and instance.provider_kind = mr.provider_kind
join provider_accounts account
on account.id = mr.provider_account_id
and account.provider_instance_id = mr.provider_instance_id
and account.provider_kind = mr.provider_kind
where mr.client_response_id = $1
and mr.client_api_key_ref = $2
@@ -168,8 +162,8 @@ async fn fetch_history(
"id" => {
sqlx::query_as::<_, HistoryRow>(
"select id, client_api_key_ref, requested_model_id, provider_kind,
provider_instance_id, provider_account_ref,
upstream_model_id, outcome, client_response_id, upstream_response_id,
provider_account_ref, upstream_model_id, outcome,
client_response_id, upstream_response_id,
started_at, completed_at
from model_requests where id = $1",
)
@@ -190,8 +184,8 @@ async fn fetch_history_by_client_response(
) -> StoreResult<Option<ModelRequestHistoryRecord>> {
let row = sqlx::query_as::<_, HistoryRow>(
"select id, client_api_key_ref, requested_model_id, provider_kind,
provider_instance_id, provider_account_ref,
upstream_model_id, outcome, client_response_id, upstream_response_id,
provider_account_ref, upstream_model_id, outcome,
client_response_id, upstream_response_id,
started_at, completed_at
from model_requests
where client_response_id = $1 and client_api_key_ref = $2",
@@ -211,7 +205,6 @@ type HistoryRow = (
Option<String>,
Option<String>,
Option<String>,
Option<String>,
String,
Option<String>,
Option<String>,
@@ -225,36 +218,32 @@ fn history_from_row(row: HistoryRow) -> ModelRequestHistoryRecord {
client_api_key_ref: row.1,
requested_model_id: row.2,
provider_kind: row.3,
provider_instance_id: row.4,
provider_account_ref: row.5,
upstream_model_id: row.6,
outcome: row.7,
client_response_id: row.8,
upstream_response_id: row.9,
started_at: row.10,
completed_at: row.11,
provider_account_ref: row.4,
upstream_model_id: row.5,
outcome: row.6,
client_response_id: row.7,
upstream_response_id: row.8,
started_at: row.9,
completed_at: row.10,
}
}
fn native_pin_from_row(
client_response_id: &str,
row: (String, String, String, String),
row: (String, String, String),
) -> StoreResult<NativeContinuationPin> {
let previous_response_id = PreviousResponseId::new(client_response_id.to_owned())
.map_err(|_| invalid_native_pin("invalid client response ID"))?;
let provider =
ProviderKind::new(row.0).map_err(|_| invalid_native_pin("invalid provider kind"))?;
let upstream_response_id = SafeUpstreamValue::new(row.3)
let upstream_response_id = SafeUpstreamValue::new(row.2)
.map_err(|_| invalid_native_pin("invalid upstream response ID"))?;
let instance = ProviderInstanceId::new(row.1)
.map_err(|_| invalid_native_pin("invalid provider instance ID"))?;
let account = ProviderAccountId::new(row.2)
let account = ProviderAccountId::new(row.1)
.map_err(|_| invalid_native_pin("invalid provider account ID"))?;
Ok(NativeContinuationPin::new(
previous_response_id,
upstream_response_id,
provider,
instance,
account,
)
.with_scope(NativeContinuationScope::Persisted))
backend/crates/gateway-store/src/postgres/mod.rs
-108
View File
@@ -8,7 +8,6 @@ use crate::{Revision, StoreBackend, StoreError, StoreResult, postgres_unavailabl
mod admin_security_audit;
mod admission_recovery;
mod client_keys;
mod config_catalog;
mod execution;
mod history;
mod observability;
@@ -21,7 +20,6 @@ mod snapshot;
pub use admin_security_audit::*;
pub use admission_recovery::*;
pub use client_keys::*;
pub use config_catalog::*;
pub use execution::*;
pub use history::*;
pub use observability::*;
@@ -73,35 +71,6 @@ pub trait ControlPlaneRepository: Send + Sync {
replacement: ControlPlaneReplacement,
) -> StoreResult<ControlPlaneSnapshot>;
async fn create_provider_instance(
&self,
expected_revision: Revision,
instance: NewProviderInstance,
audit: AdminAuditEvent,
) -> StoreResult<Revision>;
async fn update_provider_instance(
&self,
expected_revision: Revision,
instance: UpdateProviderInstanceDetails,
audit: AdminAuditEvent,
) -> StoreResult<Revision>;
async fn set_provider_instance_enabled(
&self,
expected_revision: Revision,
id: &str,
enabled: bool,
audit: AdminAuditEvent,
) -> StoreResult<Revision>;
async fn delete_provider_instance(
&self,
expected_revision: Revision,
id: &str,
audit: AdminAuditEvent,
) -> StoreResult<Revision>;
async fn create_client_api_key(
&self,
expected_revision: Revision,
@@ -205,66 +174,6 @@ impl ControlPlaneRepository for PgControlPlaneRepository {
}
}
async fn create_provider_instance(
&self,
expected_revision: Revision,
instance: NewProviderInstance,
audit: AdminAuditEvent,
) -> StoreResult<Revision> {
self.apply_targeted_mutation(
expected_revision,
ControlPlaneMutation::CreateProviderInstance(instance),
audit,
)
.await
}
async fn update_provider_instance(
&self,
expected_revision: Revision,
instance: UpdateProviderInstanceDetails,
audit: AdminAuditEvent,
) -> StoreResult<Revision> {
self.apply_targeted_mutation(
expected_revision,
ControlPlaneMutation::UpdateProviderInstance(instance),
audit,
)
.await
}
async fn set_provider_instance_enabled(
&self,
expected_revision: Revision,
id: &str,
enabled: bool,
audit: AdminAuditEvent,
) -> StoreResult<Revision> {
self.apply_targeted_mutation(
expected_revision,
ControlPlaneMutation::SetProviderInstanceEnabled {
id: id.to_owned(),
enabled,
},
audit,
)
.await
}
async fn delete_provider_instance(
&self,
expected_revision: Revision,
id: &str,
audit: AdminAuditEvent,
) -> StoreResult<Revision> {
self.apply_targeted_mutation(
expected_revision,
ControlPlaneMutation::DeleteProviderInstance(id.to_owned()),
audit,
)
.await
}
async fn create_client_api_key(
&self,
expected_revision: Revision,
@@ -327,10 +236,6 @@ impl ControlPlaneRepository for PgControlPlaneRepository {
}
enum ControlPlaneMutation {
CreateProviderInstance(NewProviderInstance),
UpdateProviderInstance(UpdateProviderInstanceDetails),
SetProviderInstanceEnabled { id: String, enabled: bool },
DeleteProviderInstance(String),
CreateClientApiKey(NewClientApiKey),
UpdateClientApiKey(UpdateClientApiKeyDetails),
SetClientApiKeyEnabled { id: String, enabled: bool },
@@ -353,19 +258,6 @@ impl PgControlPlaneRepository {
let revision =
bump_config_revision_in_transaction(&mut transaction, expected_revision).await?;
match mutation {
ControlPlaneMutation::CreateProviderInstance(instance) => {
insert_provider_instance_in_transaction(&mut transaction, &instance).await?;
}
ControlPlaneMutation::UpdateProviderInstance(instance) => {
update_provider_instance_in_transaction(&mut transaction, &instance).await?;
}
ControlPlaneMutation::SetProviderInstanceEnabled { id, enabled } => {
set_provider_instance_enabled_in_transaction(&mut transaction, &id, enabled)
.await?;
}
ControlPlaneMutation::DeleteProviderInstance(id) => {
delete_provider_instance_in_transaction(&mut transaction, &id).await?;
}
ControlPlaneMutation::CreateClientApiKey(key) => {
insert_client_api_key_in_transaction(&mut transaction, &key).await?;
}
backend/crates/gateway-store/src/postgres/observability.rs
+8 -21
View File
@@ -377,7 +377,6 @@ pub struct ProviderAccountMetrics {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccountUsageObservation {
pub account_id: String,
pub provider_instance_id: String,
pub provider_kind: String,
pub name: String,
pub email: Option<String>,
@@ -479,7 +478,6 @@ pub struct UsageRecord {
pub client_transport: String,
pub requested_model_id: String,
pub input_token_estimate: u64,
pub provider_instance_id: Option<String>,
pub provider_kind: Option<String>,
pub provider_account_ref: Option<String>,
pub provider_account_name: Option<String>,
@@ -549,7 +547,6 @@ pub struct UsageAttemptObservation {
pub attempt_index: u32,
pub component: String,
pub operation: String,
pub provider_instance_id: Option<String>,
pub provider_kind: Option<String>,
pub provider_account_ref: Option<String>,
pub upstream_model_id: Option<String>,
@@ -621,7 +618,6 @@ pub struct OpsErrorRecord {
pub client_api_key_ref: Option<String>,
pub component: String,
pub operation: String,
pub provider_instance_id: Option<String>,
pub provider_kind: Option<String>,
pub provider_account_ref: Option<String>,
pub upstream_model_id: Option<String>,
@@ -1189,7 +1185,6 @@ fn admin_dashboard_account_usage(
) -> AdminStoreResult<admin_observability::DashboardAccountUsage> {
Ok(admin_observability::DashboardAccountUsage {
account_id: usage.account_id,
provider_instance_id: usage.provider_instance_id,
provider_kind: usage.provider_kind,
name: usage.name,
email: usage.email,
@@ -1310,7 +1305,6 @@ fn admin_usage_record(record: UsageRecord) -> AdminStoreResult<admin_observabili
client_transport: record.client_transport,
requested_model_id: record.requested_model_id,
input_token_estimate: record.input_token_estimate,
provider_instance_id: record.provider_instance_id,
provider_kind: record.provider_kind,
provider_account_ref: record.provider_account_ref,
provider_account_name: record.provider_account_name,
@@ -1399,7 +1393,6 @@ fn admin_usage_attempt(
attempt_index: attempt.attempt_index,
component: attempt.component,
operation: attempt.operation,
provider_instance_id: attempt.provider_instance_id,
provider_kind: attempt.provider_kind,
provider_account_ref: attempt.provider_account_ref,
upstream_model_id: attempt.upstream_model_id,
@@ -1487,7 +1480,6 @@ fn admin_ops_error(error: OpsErrorRecord) -> admin_observability::OpsError {
client_api_key_ref: error.client_api_key_ref,
component: error.component,
operation: error.operation,
provider_instance_id: error.provider_instance_id,
provider_kind: error.provider_kind,
provider_account_ref: error.provider_account_ref,
upstream_model_id: error.upstream_model_id,
@@ -2046,7 +2038,7 @@ async fn provider_account_usage(
}
let mut statement = QueryBuilder::<Postgres>::new(
"select pa.id, pa.provider_instance_id, pa.provider_kind, pa.name, pa.email,
"select pa.id, pa.provider_kind, pa.name, pa.email,
pa.plan_type, pa.enabled, pa.availability,
count(mr.id)::bigint as request_count,
count(mr.id) filter (where mr.outcome = 'succeeded')::bigint as success_count,
@@ -2082,7 +2074,7 @@ async fn provider_account_usage(
statement.push("::text[])");
}
statement.push(
" group by pa.id, pa.provider_instance_id, pa.provider_kind, pa.name, pa.email,
" group by pa.id, pa.provider_kind, pa.name, pa.email,
pa.plan_type, pa.enabled, pa.availability
order by max(mr.started_at) desc nulls last, pa.name, pa.id limit ",
);
@@ -2262,7 +2254,6 @@ fn provider_account_from_row(
) -> StoreResult<ProviderAccountUsageObservation> {
Ok(ProviderAccountUsageObservation {
account_id: get(row, "id")?,
provider_instance_id: get(row, "provider_instance_id")?,
provider_kind: get(row, "provider_kind")?,
name: get(row, "name")?,
email: get(row, "email")?,
@@ -2291,7 +2282,7 @@ fn provider_account_from_row(
const USAGE_RECORD_SELECT: &str =
"select mr.id, mr.client_api_key_ref, mr.config_revision, mr.protocol, mr.operation,
mr.endpoint, mr.client_transport, mr.requested_model_id,
mr.input_token_estimate, mr.provider_instance_id, mr.provider_kind, mr.provider_account_ref,
mr.input_token_estimate, mr.provider_kind, mr.provider_account_ref,
pa.name as provider_account_name, pa.email as provider_account_email,
mr.upstream_model_id, mr.upstream_transport, mr.http_version, mr.websocket_pool,
mr.attempt_count, mr.upstream_send_state, mr.downstream_committed_at,
@@ -2405,7 +2396,7 @@ async fn usage_record_detail(pool: &PgPool, request_id: &str) -> StoreResult<Usa
let request = usage_record_from_row(&row)?;
let rows = sqlx::query(
"select id, attempt_index, component, operation,
provider_instance_id, provider_kind, provider_account_ref, upstream_model_id,
provider_kind, provider_account_ref, upstream_model_id,
failure_kind, status_code, provider_error_code, retry_after_ms,
upstream_request_id, latency_ms, message, created_at
from ops_events where model_request_id = $1
@@ -2434,7 +2425,6 @@ fn intermediate_attempt_from_row(
attempt_index: to_u32(get(row, "attempt_index")?)?,
component: get(row, "component")?,
operation: get(row, "operation")?,
provider_instance_id: get(row, "provider_instance_id")?,
provider_kind: get(row, "provider_kind")?,
provider_account_ref: get(row, "provider_account_ref")?,
upstream_model_id: get(row, "upstream_model_id")?,
@@ -2469,7 +2459,6 @@ fn final_attempt_from_request(request: &UsageRecord) -> UsageAttemptObservation
attempt_index: request.attempt_count,
component: "model_request".to_owned(),
operation: request.operation.clone(),
provider_instance_id: request.provider_instance_id.clone(),
provider_kind: request.provider_kind.clone(),
provider_account_ref: request.provider_account_ref.clone(),
upstream_model_id: request.upstream_model_id.clone(),
@@ -2622,7 +2611,7 @@ const OPS_ERRORS_CTE: &str = "with errors as (
mr.id as event_id, mr.id as request_id,
nullif(mr.attempt_count, 0) as attempt_index,
mr.client_api_key_ref, 'model_request'::text as component, mr.operation,
mr.provider_instance_id, mr.provider_kind,
mr.provider_kind,
mr.provider_account_ref, mr.upstream_model_id, mr.upstream_transport,
coalesce(mr.error_kind, 'failed') as failure_kind,
coalesce(mr.upstream_status_code, mr.client_status_code) as status_code,
@@ -2635,7 +2624,7 @@ const OPS_ERRORS_CTE: &str = "with errors as (
union all
select 'ops_event'::text, oe.id, oe.model_request_id, oe.attempt_index,
mr.client_api_key_ref, oe.component, oe.operation,
oe.provider_instance_id, oe.provider_kind, oe.provider_account_ref,
oe.provider_kind, oe.provider_account_ref,
oe.upstream_model_id, null::text, oe.failure_kind, oe.status_code,
oe.provider_error_code, mr.client_response_id, oe.upstream_request_id,
oe.latency_ms, oe.message, oe.occurrence_count, oe.created_at,
@@ -2649,8 +2638,8 @@ async fn list_ops_errors(pool: &PgPool, query: OpsErrorQuery) -> StoreResult<Ops
let mut statement = QueryBuilder::<Postgres>::new(OPS_ERRORS_CTE);
statement.push(
" select source, event_id, request_id, attempt_index, client_api_key_ref,
component, operation, provider_instance_id,
provider_kind, provider_account_ref, upstream_model_id, upstream_transport,
component, operation, provider_kind, provider_account_ref,
upstream_model_id, upstream_transport,
failure_kind, status_code, provider_error_code, client_response_id,
upstream_request_id, latency_ms, message, occurrence_count, occurred_at,
stable_sort_id
@@ -2775,7 +2764,6 @@ fn usage_record_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<UsageRecord
client_transport: get(row, "client_transport")?,
requested_model_id: get(row, "requested_model_id")?,
input_token_estimate: unsigned(row, "input_token_estimate")?,
provider_instance_id: get(row, "provider_instance_id")?,
provider_kind: get(row, "provider_kind")?,
provider_account_ref: get(row, "provider_account_ref")?,
provider_account_name: get(row, "provider_account_name")?,
@@ -2843,7 +2831,6 @@ fn ops_error_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<OpsErrorRecord
client_api_key_ref: get(row, "client_api_key_ref")?,
component: get(row, "component")?,
operation: get(row, "operation")?,
provider_instance_id: get(row, "provider_instance_id")?,
provider_kind: get(row, "provider_kind")?,
provider_account_ref: get(row, "provider_account_ref")?,
upstream_model_id: get(row, "upstream_model_id")?,
backend/crates/gateway-store/src/postgres/ops_events.rs
+2 -5
View File
@@ -32,7 +32,6 @@ pub struct OpsEvent {
pub level: OpsEventLevel,
pub component: String,
pub operation: String,
pub provider_instance_id: Option<String>,
pub provider_kind: Option<String>,
pub provider_account_id: Option<String>,
pub provider_account_ref: Option<String>,
@@ -110,13 +109,12 @@ impl OpsEventRepository for PgOpsEventRepository {
sqlx::query(
"insert into ops_events (
id, model_request_id, attempt_index, level, component, operation,
provider_instance_id, provider_kind,
provider_account_id, provider_account_ref, upstream_model_id,
provider_kind, provider_account_id, provider_account_ref, upstream_model_id,
failure_kind, status_code, provider_error_code, retry_after_ms,
upstream_request_id, latency_ms, message, occurrence_count, created_at
) values (
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13,
$14, $15, $16, $17, $18, $19, $20
$14, $15, $16, $17, $18, $19
)",
)
.bind(event.id)
@@ -131,7 +129,6 @@ impl OpsEventRepository for PgOpsEventRepository {
.bind(event.level.as_str())
.bind(event.component)
.bind(event.operation)
.bind(event.provider_instance_id)
.bind(event.provider_kind)
.bind(event.provider_account_id)
.bind(event.provider_account_ref)
backend/crates/gateway-store/src/postgres/provider_accounts.rs
+53 -82
View File
@@ -16,7 +16,7 @@ use gateway_admin::{
AccountListQuery as AdminAccountListQuery, AccountModelUsage, AccountPage,
AccountRecord, AccountSort as AdminAccountSort,
AccountSortField as AdminAccountSortField, AccountStatus as AdminAccountStatus,
AccountSummary, AccountUsage, DeleteAccount, SetAccountEnabled,
AccountSummary, AccountUsage, DeleteAccounts, SetAccountEnabled,
SortDirection as AdminSortDirection,
},
observability::{
@@ -43,7 +43,7 @@ use gateway_core::engine::credential::{
QuotaWriteOutcome,
};
use gateway_core::error::{StoreError as CoreStoreError, StoreErrorKind as CoreStoreErrorKind};
use gateway_core::routing::{ProviderInstanceId, ProviderKind};
use gateway_core::routing::ProviderKind;
use crate::{
ConflictKind, JsonObject, Revision, StoreError, StoreResult, admin_revision, admin_store_error,
@@ -67,18 +67,15 @@ const ADMIN_USAGE_CHUNK_SIZE: usize = 200;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccountAdminScope {
pub provider_kind: String,
pub provider_instance_id: String,
}
impl ProviderAccountAdminScope {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "provider_kind", &self.provider_kind)?;
require_nonempty(ENTITY, "provider_instance_id", &self.provider_instance_id)
require_nonempty(ENTITY, "provider_kind", &self.provider_kind)
}
fn contains(&self, account: &NewProviderAccount) -> bool {
account.provider_kind == self.provider_kind
&& account.provider_instance_id == self.provider_instance_id
}
}
@@ -124,7 +121,6 @@ impl ProviderAccountAvailability {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccountSummary {
pub id: String,
pub provider_instance_id: String,
pub provider_kind: String,
pub name: String,
pub email: Option<String>,
@@ -169,7 +165,6 @@ impl fmt::Debug for ProviderAccountRecord {
#[derive(Clone)]
pub struct NewProviderAccount {
pub id: String,
pub provider_instance_id: String,
pub provider_kind: String,
pub name: String,
pub email: Option<String>,
@@ -200,7 +195,6 @@ impl fmt::Debug for NewProviderAccount {
formatter
.debug_struct("NewProviderAccount")
.field("id", &self.id)
.field("provider_instance_id", &self.provider_instance_id)
.field("provider_kind", &self.provider_kind)
.field("name", &self.name)
.field("enabled", &self.enabled)
@@ -214,7 +208,6 @@ impl fmt::Debug for NewProviderAccount {
impl NewProviderAccount {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "id", &self.id)?;
require_nonempty(ENTITY, "provider_instance_id", &self.provider_instance_id)?;
require_nonempty(ENTITY, "provider_kind", &self.provider_kind)?;
require_nonempty(ENTITY, "name", &self.name)?;
require_nonempty(ENTITY, "upstream_user_id", &self.upstream_user_id)?;
@@ -413,7 +406,7 @@ pub trait ProviderAccountRepository: Send + Sync {
async fn load_provider_account(&self, id: &str) -> StoreResult<Option<ProviderAccountRecord>>;
async fn list_provider_accounts(
&self,
provider_instance_id: Option<&str>,
provider_kind: Option<&str>,
include_disabled: bool,
) -> StoreResult<Vec<ProviderAccountSummary>>;
async fn insert_provider_account(&self, account: NewProviderAccount) -> StoreResult<()>;
@@ -500,20 +493,20 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
async fn list_provider_accounts(
&self,
provider_instance_id: Option<&str>,
provider_kind: Option<&str>,
include_disabled: bool,
) -> StoreResult<Vec<ProviderAccountSummary>> {
let rows = sqlx::query(
"select id, provider_instance_id, provider_kind, name, email, upstream_user_id,
"select id, provider_kind, name, email, upstream_user_id,
upstream_account_id, plan_type, credential_revision, has_refresh_token,
access_token_expires_at, next_refresh_at, enabled, availability,
availability_reason, cooldown_until, availability_observed_at,
quota_observed_at, created_at, updated_at
from provider_accounts
where ($1::text is null or provider_instance_id = $1) and ($2 or enabled)
order by provider_kind, provider_instance_id, name, id",
where ($1::text is null or provider_kind = $1) and ($2 or enabled)
order by provider_kind, name, id",
)
.bind(provider_instance_id)
.bind(provider_kind)
.bind(include_disabled)
.fetch_all(&self.pool)
.await
@@ -525,18 +518,17 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
account.validate()?;
sqlx::query(
"insert into provider_accounts (
id, provider_instance_id, provider_kind, name, email, upstream_user_id,
id, provider_kind, name, email, upstream_user_id,
upstream_account_id, plan_type, provider_credentials_json, credential_revision,
has_refresh_token, access_token_expires_at, next_refresh_at, enabled,
availability, availability_reason, cooldown_until, provider_quota_json,
availability_observed_at, quota_observed_at, created_at, updated_at
) values (
$1, $2, $3, $4, $5, $6, $7, $8, $9, 1, $10, $11, $12, $13,
$14, null, $15, null, $16, null, now(), greatest(now(), $16)
$1, $2, $3, $4, $5, $6, $7, $8, 1, $9, $10, $11, $12,
$13, null, $14, null, $15, null, now(), greatest(now(), $15)
)",
)
.bind(account.id)
.bind(account.provider_instance_id)
.bind(account.provider_kind)
.bind(account.name)
.bind(account.email)
@@ -734,7 +726,6 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
validate_admin_account_ids(&account_ids)?;
let rows = sqlx::query(ACCOUNT_SELECT_BY_IDS)
.bind(&account_ids)
.bind(&scope.provider_instance_id)
.bind(&scope.provider_kind)
.fetch_all(&self.pool)
.await
@@ -960,7 +951,6 @@ impl PgAdminAccountStore {
})?;
Ok(ProviderAccountAdminScope {
provider_kind: record.summary.provider_kind,
provider_instance_id: record.summary.provider_instance_id,
})
}
@@ -972,7 +962,6 @@ impl PgAdminAccountStore {
action: &str,
) -> AdminStoreResult<CredentialImportResult> {
let provider_kind = prepared.provider_kind.as_str().to_owned();
let provider_instance_id = prepared.provider_instance_id.as_str().to_owned();
let accounts = prepared
.credentials
.into_iter()
@@ -985,15 +974,14 @@ impl PgAdminAccountStore {
store_revision(expected_config_revision)?,
ImportProviderAccounts {
scope: ProviderAccountAdminScope {
provider_kind,
provider_instance_id: provider_instance_id.clone(),
provider_kind: provider_kind.clone(),
},
accounts,
audit: mutation_audit(
context,
action,
"provider_account",
&provider_instance_id,
&provider_kind,
vec!["credentials".to_owned()],
),
},
@@ -1027,7 +1015,6 @@ impl PgAdminAccountStore {
let account_id = prepared.account_id.clone();
let scope = ProviderAccountAdminScope {
provider_kind: prepared.provider_kind.as_str().to_owned(),
provider_instance_id: prepared.provider_instance_id.as_str().to_owned(),
};
let rotation = self
.accounts
@@ -1171,14 +1158,10 @@ impl AccountStore for PgAdminAccountStore {
provider_kind: &ProviderKind,
query: CredentialListQuery,
) -> AdminStoreResult<CredentialPage> {
let provider_instance_id = query
.provider_instance_id
.as_ref()
.map(ProviderInstanceId::as_str);
let (control_plane, accounts) = futures::try_join!(
self.control_plane.load_control_plane(),
self.accounts
.list_provider_accounts(provider_instance_id, true),
.list_provider_accounts(Some(provider_kind.as_str()), true),
)
.map_err(|error| admin_store_error(ENTITY, error))?;
let mut accounts = accounts
@@ -1340,7 +1323,6 @@ impl AccountStore for PgAdminAccountStore {
expected_config_revision,
PreparedCredentialImport {
provider_kind: credential.provider_kind.clone(),
provider_instance_id: credential.provider_instance_id.clone(),
credentials: vec![credential],
},
context,
@@ -1433,23 +1415,35 @@ impl AccountStore for PgAdminAccountStore {
.and_then(admin_revision)
}
async fn delete_account(
async fn delete_accounts(
&self,
command: DeleteAccount,
command: DeleteAccounts,
context: &MutationContext,
) -> AdminStoreResult<AdminRevision> {
let scope = self.required_scope(&command.account_id).await?;
let first_account_id = command.account_ids.first().ok_or_else(|| {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
ENTITY,
"account deletion requires at least one account ID",
)
})?;
let scope = self.required_scope(first_account_id).await?;
let audit_target = if command.account_ids.len() == 1 {
first_account_id.clone()
} else {
"provider_accounts".to_owned()
};
self.accounts
.delete_provider_accounts_admin(
store_revision(command.expected_config_revision)?,
DeleteProviderAccounts {
scope,
account_ids: vec![command.account_id.clone()],
account_ids: command.account_ids,
audit: mutation_audit(
context,
"delete",
"provider_account",
&command.account_id,
&audit_target,
Vec::new(),
),
},
@@ -1677,15 +1671,6 @@ const fn admin_account_status_name(status: AdminAccountStatus) -> &'static str {
fn admin_account_record(summary: ProviderAccountSummary) -> AdminStoreResult<AccountRecord> {
Ok(AccountRecord {
id: summary.id,
provider_instance_id: ProviderInstanceId::new(summary.provider_instance_id).map_err(
|_| {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
ENTITY,
"persisted Provider instance ID is invalid",
)
},
)?,
provider_kind: ProviderKind::new(summary.provider_kind).map_err(|_| {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
@@ -1716,7 +1701,6 @@ fn admin_account_record(summary: ProviderAccountSummary) -> AdminStoreResult<Acc
fn prepared_account(credential: PreparedCredentialCreate) -> StoreResult<NewProviderAccount> {
Ok(NewProviderAccount {
id: credential.account_id.as_str().to_owned(),
provider_instance_id: credential.provider_instance_id.as_str().to_owned(),
provider_kind: credential.provider_kind.as_str().to_owned(),
name: credential.name,
email: credential.email,
@@ -1856,14 +1840,14 @@ async fn upsert_provider_account_in_transaction(
account.validate()?;
let imported_id = sqlx::query_scalar::<_, String>(
"insert into provider_accounts (
id, provider_instance_id, provider_kind, name, email, upstream_user_id,
id, provider_kind, name, email, upstream_user_id,
upstream_account_id, plan_type, provider_credentials_json, credential_revision,
has_refresh_token, access_token_expires_at, next_refresh_at, enabled,
availability, availability_reason, cooldown_until, provider_quota_json,
availability_observed_at, quota_observed_at, created_at, updated_at
) values (
$1, $2, $3, $4, $5, $6, $7, $8, $9, 1, $10, $11, $12, $13,
$14, $15, $16, null, $17, null, now(), greatest(now(), $17)
$1, $2, $3, $4, $5, $6, $7, $8, 1, $9, $10, $11, $12,
$13, $14, $15, null, $16, null, now(), greatest(now(), $16)
)
on conflict (
provider_kind,
@@ -1886,11 +1870,9 @@ async fn upsert_provider_account_in_transaction(
availability_observed_at = excluded.availability_observed_at,
quota_observed_at = null,
updated_at = greatest(now(), excluded.availability_observed_at)
where provider_accounts.provider_instance_id = excluded.provider_instance_id
returning id",
)
.bind(&account.id)
.bind(&account.provider_instance_id)
.bind(&account.provider_kind)
.bind(&account.name)
.bind(&account.email)
@@ -1938,21 +1920,20 @@ async fn rotate_provider_account_in_transaction(
) -> StoreResult<Revision> {
let next = sqlx::query_scalar::<_, i64>(
"update provider_accounts
set name = $5,
email = $6,
plan_type = $7,
provider_credentials_json = $8,
set name = $4,
email = $5,
plan_type = $6,
provider_credentials_json = $7,
credential_revision = credential_revision + 1,
has_refresh_token = $9,
access_token_expires_at = $10,
next_refresh_at = $11,
has_refresh_token = $8,
access_token_expires_at = $9,
next_refresh_at = $10,
updated_at = now()
where id = $1 and provider_instance_id = $2 and provider_kind = $3
and credential_revision = $4
where id = $1 and provider_kind = $2
and credential_revision = $3
returning credential_revision",
)
.bind(&update.account_id)
.bind(&scope.provider_instance_id)
.bind(&scope.provider_kind)
.bind(to_i64(update.expected_revision.get())?)
.bind(&profile.name)
@@ -1980,11 +1961,10 @@ async fn set_provider_account_enabled_in_transaction(
enabled: bool,
) -> StoreResult<()> {
let result = sqlx::query(
"update provider_accounts set enabled = $4, updated_at = now()
where id = $1 and provider_instance_id = $2 and provider_kind = $3",
"update provider_accounts set enabled = $3, updated_at = now()
where id = $1 and provider_kind = $2",
)
.bind(account_id)
.bind(&scope.provider_instance_id)
.bind(&scope.provider_kind)
.bind(enabled)
.execute(&mut **transaction)
@@ -2000,12 +1980,10 @@ async fn delete_provider_accounts_in_transaction(
) -> StoreResult<()> {
let deleted = sqlx::query_scalar::<_, String>(
"delete from provider_accounts
where id = any($1::text[]) and provider_instance_id = $2
and provider_kind = $3
where id = any($1::text[]) and provider_kind = $2
returning id",
)
.bind(account_ids)
.bind(&scope.provider_instance_id)
.bind(&scope.provider_kind)
.fetch_all(&mut **transaction)
.await
@@ -2098,7 +2076,6 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.map_err(core_store_error)?;
self.insert_provider_account(NewProviderAccount {
id: account.account.id().as_str().to_owned(),
provider_instance_id: account.account.instance().as_str().to_owned(),
provider_kind: account.account.provider().as_str().to_owned(),
name: account.account.name().to_owned(),
email: account.account.email().map(str::to_owned),
@@ -2141,11 +2118,11 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.collect()
}
async fn list_for_instance(
async fn list_for_provider(
&self,
instance: &ProviderInstanceId,
provider: &ProviderKind,
) -> Result<Vec<CoreProviderAccount>, CoreStoreError> {
self.list_provider_accounts(Some(instance.as_str()), false)
self.list_provider_accounts(Some(provider.as_str()), false)
.await
.map_err(core_store_error)?
.into_iter()
@@ -2365,22 +2342,20 @@ impl ProviderAccountStore for PgProviderAccountRepository {
}
}
const ACCOUNT_SELECT: &str =
"select id, provider_instance_id, provider_kind, name, email, upstream_user_id,
const ACCOUNT_SELECT: &str = "select id, provider_kind, name, email, upstream_user_id,
upstream_account_id, plan_type, provider_credentials_json, credential_revision,
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, availability,
availability_reason, cooldown_until, provider_quota_json,
availability_observed_at, quota_observed_at, created_at, updated_at
from provider_accounts where id = $1";
const ACCOUNT_SELECT_BY_IDS: &str =
"select id, provider_instance_id, provider_kind, name, email, upstream_user_id,
const ACCOUNT_SELECT_BY_IDS: &str = "select id, provider_kind, name, email, upstream_user_id,
upstream_account_id, plan_type, provider_credentials_json, credential_revision,
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, availability,
availability_reason, cooldown_until, provider_quota_json,
availability_observed_at, quota_observed_at, created_at, updated_at
from provider_accounts
where id = any($1::text[]) and provider_instance_id = $2 and provider_kind = $3
where id = any($1::text[]) and provider_kind = $2
order by id";
fn account_record_from_row(row: sqlx::postgres::PgRow) -> StoreResult<ProviderAccountRecord> {
@@ -2407,15 +2382,12 @@ fn core_account_from_summary(
) -> Result<CoreProviderAccount, CoreStoreError> {
let id = CoreProviderAccountId::new(summary.id)
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
let instance = ProviderInstanceId::new(summary.provider_instance_id)
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
let provider = ProviderKind::new(summary.provider_kind)
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
let revision = CoreCredentialRevision::new(summary.credential_revision.get())
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
Ok(CoreProviderAccount::new(
id,
instance,
provider,
summary.name,
summary.upstream_user_id,
@@ -2490,7 +2462,6 @@ fn account_summary_from_row(row: sqlx::postgres::PgRow) -> StoreResult<ProviderA
.map_err(|_| invalid("invalid availability"))?;
Ok(ProviderAccountSummary {
id: get(&row, "id")?,
provider_instance_id: get(&row, "provider_instance_id")?,
provider_kind: get(&row, "provider_kind")?,
name: get(&row, "name")?,
email: get(&row, "email")?,
backend/crates/gateway-store/src/postgres/snapshot.rs
+4 -124
View File
@@ -7,18 +7,15 @@ use chrono::{DateTime, Utc};
use gateway_core::routing::{
ConfigRevision,
snapshot::{
SnapshotClientPolicyFacts, SnapshotFacts, SnapshotProviderInstanceFacts,
SnapshotSettingsFacts, SnapshotStoreError, SnapshotStorePort,
SnapshotClientPolicyFacts, SnapshotFacts, SnapshotSettingsFacts, SnapshotStoreError,
SnapshotStorePort,
},
};
use sqlx::{PgPool, Postgres, Row, Transaction};
use sqlx::{PgPool, Postgres, Transaction};
use crate::{Revision, StoreError, StoreResult, postgres_unavailable};
use super::{
ClientApiKeySnapshot, ProviderAccountAvailability, ProviderAccountSummary,
ProviderInstanceRecord,
};
use super::ClientApiKeySnapshot;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotRuntimeSettings {
@@ -43,8 +40,6 @@ pub struct RuntimeSnapshotData {
pub config_revision: Revision,
pub observed_current_revision: Revision,
pub settings: SnapshotRuntimeSettings,
pub provider_instances: Vec<ProviderInstanceRecord>,
pub provider_accounts: Vec<ProviderAccountSummary>,
pub client_api_keys: Vec<ClientApiKeySnapshot>,
}
@@ -81,8 +76,6 @@ impl RuntimeSnapshotRepository for PgRuntimeSnapshotRepository {
.map_err(|_| postgres_unavailable("configure runtime snapshot transaction"))?;
let (config_revision, settings) = load_settings(&mut transaction).await?;
let provider_instances = load_instances(&mut transaction).await?;
let provider_accounts = load_accounts(&mut transaction).await?;
let client_api_keys = load_client_keys(&mut transaction).await?;
transaction
.commit()
@@ -95,8 +88,6 @@ impl RuntimeSnapshotRepository for PgRuntimeSnapshotRepository {
config_revision,
observed_current_revision,
settings,
provider_instances,
provider_accounts,
client_api_keys,
})
}
@@ -153,18 +144,6 @@ impl SnapshotStorePort for PgRuntimeSnapshotRepository {
data.settings.rotation_strategy,
data.settings.provider_model_mappings,
);
let provider_instances = data
.provider_instances
.into_iter()
.map(|instance| {
SnapshotProviderInstanceFacts::new(
instance.id,
instance.provider_kind,
instance.base_url,
instance.enabled,
)
})
.collect();
let client_policies = data
.client_api_keys
.into_iter()
@@ -181,7 +160,6 @@ impl SnapshotStorePort for PgRuntimeSnapshotRepository {
config_revision,
observed_current_revision,
settings,
provider_instances,
client_policies,
))
})
@@ -243,84 +221,6 @@ async fn load_settings(
))
}
async fn load_instances(
transaction: &mut Transaction<'_, Postgres>,
) -> StoreResult<Vec<ProviderInstanceRecord>> {
let rows = sqlx::query_as::<
_,
(
String,
String,
String,
String,
bool,
DateTime<Utc>,
DateTime<Utc>,
),
>(
"select id, provider_kind, name, base_url, enabled, created_at, updated_at
from provider_instances where enabled order by provider_kind, name, id",
)
.fetch_all(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("load snapshot provider instances"))?;
Ok(rows
.into_iter()
.map(|row| ProviderInstanceRecord {
id: row.0,
provider_kind: row.1,
name: row.2,
base_url: row.3,
enabled: row.4,
created_at: row.5,
updated_at: row.6,
})
.collect())
}
async fn load_accounts(
transaction: &mut Transaction<'_, Postgres>,
) -> StoreResult<Vec<ProviderAccountSummary>> {
let rows = sqlx::query(
"select id, provider_instance_id, provider_kind, name, email, upstream_user_id,
upstream_account_id, plan_type, credential_revision, has_refresh_token,
access_token_expires_at, next_refresh_at, enabled, availability,
availability_reason, cooldown_until, availability_observed_at,
quota_observed_at, created_at, updated_at
from provider_accounts where enabled order by provider_kind, provider_instance_id, id",
)
.fetch_all(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("load snapshot provider accounts"))?;
rows.into_iter()
.map(|row| {
let availability: String = snapshot_get(&row, "availability")?;
Ok(ProviderAccountSummary {
id: snapshot_get(&row, "id")?,
provider_instance_id: snapshot_get(&row, "provider_instance_id")?,
provider_kind: snapshot_get(&row, "provider_kind")?,
name: snapshot_get(&row, "name")?,
email: snapshot_get(&row, "email")?,
upstream_user_id: snapshot_get(&row, "upstream_user_id")?,
upstream_account_id: snapshot_get(&row, "upstream_account_id")?,
plan_type: snapshot_get(&row, "plan_type")?,
credential_revision: revision_from_i64(snapshot_get(&row, "credential_revision")?)?,
has_refresh_token: snapshot_get(&row, "has_refresh_token")?,
access_token_expires_at: snapshot_get(&row, "access_token_expires_at")?,
next_refresh_at: snapshot_get(&row, "next_refresh_at")?,
enabled: snapshot_get(&row, "enabled")?,
availability: parse_availability(&availability)?,
availability_reason: snapshot_get(&row, "availability_reason")?,
cooldown_until: snapshot_get(&row, "cooldown_until")?,
availability_observed_at: snapshot_get(&row, "availability_observed_at")?,
quota_observed_at: snapshot_get(&row, "quota_observed_at")?,
created_at: snapshot_get(&row, "created_at")?,
updated_at: snapshot_get(&row, "updated_at")?,
})
})
.collect()
}
async fn load_client_keys(
transaction: &mut Transaction<'_, Postgres>,
) -> StoreResult<Vec<ClientApiKeySnapshot>> {
@@ -336,26 +236,6 @@ async fn load_client_keys(
.collect()
}
fn parse_availability(value: &str) -> StoreResult<ProviderAccountAvailability> {
match value {
"unknown" => Ok(ProviderAccountAvailability::Unknown),
"ready" => Ok(ProviderAccountAvailability::Ready),
"cooldown" => Ok(ProviderAccountAvailability::Cooldown),
"quota_exhausted" => Ok(ProviderAccountAvailability::QuotaExhausted),
"expired" => Ok(ProviderAccountAvailability::Expired),
"banned" => Ok(ProviderAccountAvailability::Banned),
"invalid" => Ok(ProviderAccountAvailability::Invalid),
_ => Err(invalid("unknown provider account availability")),
}
}
fn snapshot_get<'r, T>(row: &'r sqlx::postgres::PgRow, column: &'static str) -> StoreResult<T>
where
T: sqlx::Decode<'r, sqlx::Postgres> + sqlx::Type<sqlx::Postgres>,
{
row.try_get(column).map_err(|_| invalid(column))
}
fn revision_from_i64(value: i64) -> StoreResult<Revision> {
Revision::new(to_u64(value)?)
}
backend/crates/gateway-store/src/redis/credential_cooldown.rs
+4 -4
View File
@@ -15,7 +15,7 @@ use gateway_core::{
provider_ports::{
ProviderCooldown, ProviderCooldownPort, ProviderStoreError, ProviderStoreErrorKind,
},
routing::ProviderInstanceId,
routing::ProviderKind,
};
use redis::{Script, aio::ConnectionManager};
@@ -331,11 +331,11 @@ impl ProviderAccountStore for CooldownCachingProviderAccountStore {
self.authoritative.list_accounts().await
}
async fn list_for_instance(
async fn list_for_provider(
&self,
instance: &ProviderInstanceId,
provider: &ProviderKind,
) -> Result<Vec<ProviderAccount>, CoreStoreError> {
self.authoritative.list_for_instance(instance).await
self.authoritative.list_for_provider(provider).await
}
async fn load_credential(
backend/crates/gateway-store/src/redis/credential_leases.rs
+15 -19
View File
@@ -1,4 +1,4 @@
//! Provider instance、账号与 OAuth refresh 的 Redis lease/fencing。
//! Provider、账号与 OAuth refresh 的 Redis lease/fencing。
use std::collections::BTreeMap;
use std::fmt;
@@ -14,7 +14,7 @@ use gateway_core::provider_ports::{
ProviderRefreshCapacityRequest, ProviderSchedulingLeaseRequest, ProviderSchedulingState,
ProviderStoreError, ProviderStoreErrorKind,
};
use gateway_core::routing::ProviderInstanceId;
use gateway_core::routing::ProviderKind;
use redis::{Script, aio::ConnectionManager};
use uuid::Uuid;
@@ -125,7 +125,7 @@ return tostring(cursor - 1)
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CredentialLeaseScope {
ProviderInstance,
Provider,
ProviderAccount,
OAuthRefreshCapacity,
OAuthRefresh,
@@ -135,7 +135,7 @@ pub enum CredentialLeaseScope {
impl CredentialLeaseScope {
const fn as_str(self) -> &'static str {
match self {
Self::ProviderInstance => "instance",
Self::Provider => "provider",
Self::ProviderAccount => "account",
Self::OAuthRefreshCapacity => "refresh-capacity",
Self::OAuthRefresh => "refresh",
@@ -329,14 +329,10 @@ impl RedisCredentialLeaseRepository {
}))
}
/// 原子推进跨进程共享的 Provider instance 调度游标。
pub async fn advance_scheduling_cursor(&self, provider_instance_id: &str) -> StoreResult<u64> {
require_nonempty(
"provider scheduling cursor",
"provider_instance_id",
provider_instance_id,
)?;
let key = self.scheduling_cursor_key(provider_instance_id)?;
/// 原子推进跨进程共享的 Provider 调度游标。
pub async fn advance_scheduling_cursor(&self, provider_kind: &str) -> StoreResult<u64> {
require_nonempty("provider scheduling cursor", "provider_kind", provider_kind)?;
let key = self.scheduling_cursor_key(provider_kind)?;
let mut connection = self.connection.clone();
let cursor = Script::new(ADVANCE_SCHEDULING_CURSOR_SCRIPT)
.key(key)
@@ -373,10 +369,10 @@ impl RedisCredentialLeaseRepository {
))
}
fn scheduling_cursor_key(&self, provider_instance_id: &str) -> StoreResult<String> {
let fingerprint = resource_fingerprint("provider scheduling cursor", provider_instance_id)?;
fn scheduling_cursor_key(&self, provider_kind: &str) -> StoreResult<String> {
let fingerprint = resource_fingerprint("provider scheduling cursor", provider_kind)?;
Ok(format!(
"{}:scheduler:instance:{{{fingerprint}}}:cursor",
"{}:scheduler:provider:{{{fingerprint}}}:cursor",
self.namespace
))
}
@@ -479,10 +475,10 @@ impl RedisProviderLeaseCoordinator {
async fn next_scheduling_cursor(
&self,
provider_instance_id: &ProviderInstanceId,
provider_kind: &ProviderKind,
) -> Result<u64, ProviderStoreError> {
self.repository
.advance_scheduling_cursor(provider_instance_id.as_str())
.advance_scheduling_cursor(provider_kind.as_str())
.await
.map_err(|_| provider_unavailable("advance scheduling cursor"))
}
@@ -590,12 +586,12 @@ impl RedisProviderLeaseCoordinator {
impl ProviderLeasePort for RedisProviderLeaseCoordinator {
fn load_state<'a>(
&'a self,
provider_instance_id: &'a ProviderInstanceId,
provider_kind: &'a ProviderKind,
accounts: &'a [ProviderAccountId],
) -> futures::future::BoxFuture<'a, Result<ProviderSchedulingState, ProviderStoreError>> {
Box::pin(async move {
let signals = self.load_signals(accounts).await?;
let round_robin_cursor = self.next_scheduling_cursor(provider_instance_id).await?;
let round_robin_cursor = self.next_scheduling_cursor(provider_kind).await?;
Ok(ProviderSchedulingState::new(
signals,
None,
backend/crates/gateway-store/src/redis/provider_circuit.rs
+23 -35
View File
@@ -1,4 +1,4 @@
//! Provider instance 的可重建 Redis circuit。
//! Provider 的可重建 Redis circuit。
use std::num::NonZeroU32;
use std::time::Duration;
@@ -8,7 +8,7 @@ use chrono::{DateTime, Utc};
use gateway_core::engine::execution::{
ProviderCircuitDecision as CoreCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
};
use gateway_core::routing::ProviderInstanceId;
use gateway_core::routing::ProviderKind;
use redis::{Script, aio::ConnectionManager};
use crate::{StoreError, StoreResult, redis_unavailable, require_nonempty};
@@ -68,13 +68,13 @@ pub struct ProviderCircuitObservation {
pub trait ProviderCircuitRepository: Send + Sync {
async fn provider_circuit_decision(
&self,
provider_instance_id: &str,
provider_kind: &str,
) -> StoreResult<ProviderCircuitDecision>;
async fn observe_provider_failure(
&self,
provider_instance_id: &str,
provider_kind: &str,
) -> StoreResult<ProviderCircuitObservation>;
async fn observe_provider_success(&self, provider_instance_id: &str) -> StoreResult<()>;
async fn observe_provider_success(&self, provider_kind: &str) -> StoreResult<()>;
}
#[derive(Clone)]
@@ -100,9 +100,9 @@ impl RedisProviderCircuitRepository {
})
}
fn key(&self, provider_instance_id: &str) -> StoreResult<String> {
let fingerprint = resource_fingerprint("provider circuit", provider_instance_id)?;
Ok(format!("{}:instance:{fingerprint}:circuit", self.namespace))
fn key(&self, provider_kind: &str) -> StoreResult<String> {
let fingerprint = resource_fingerprint("provider circuit", provider_kind)?;
Ok(format!("{}:provider:{fingerprint}:circuit", self.namespace))
}
}
@@ -110,16 +110,12 @@ impl RedisProviderCircuitRepository {
impl ProviderCircuitRepository for RedisProviderCircuitRepository {
async fn provider_circuit_decision(
&self,
provider_instance_id: &str,
provider_kind: &str,
) -> StoreResult<ProviderCircuitDecision> {
require_nonempty(
"provider circuit",
"provider_instance_id",
provider_instance_id,
)?;
require_nonempty("provider circuit", "provider_kind", provider_kind)?;
let mut connection = self.connection.clone();
let (allow, until): (i64, String) = Script::new(DECISION_SCRIPT)
.key(self.key(provider_instance_id)?)
.key(self.key(provider_kind)?)
.invoke_async(&mut connection)
.await
.map_err(|_| redis_unavailable("read provider circuit"))?;
@@ -132,18 +128,14 @@ impl ProviderCircuitRepository for RedisProviderCircuitRepository {
async fn observe_provider_failure(
&self,
provider_instance_id: &str,
provider_kind: &str,
) -> StoreResult<ProviderCircuitObservation> {
require_nonempty(
"provider circuit",
"provider_instance_id",
provider_instance_id,
)?;
require_nonempty("provider circuit", "provider_kind", provider_kind)?;
let duration_ms = u64::try_from(self.policy.open_duration.as_millis())
.map_err(|_| invalid("open duration is too large"))?;
let mut connection = self.connection.clone();
let (failures, until): (String, String) = Script::new(FAILURE_SCRIPT)
.key(self.key(provider_instance_id)?)
.key(self.key(provider_kind)?)
.arg(self.policy.failure_threshold.get())
.arg(duration_ms)
.invoke_async(&mut connection)
@@ -163,15 +155,11 @@ impl ProviderCircuitRepository for RedisProviderCircuitRepository {
})
}
async fn observe_provider_success(&self, provider_instance_id: &str) -> StoreResult<()> {
require_nonempty(
"provider circuit",
"provider_instance_id",
provider_instance_id,
)?;
async fn observe_provider_success(&self, provider_kind: &str) -> StoreResult<()> {
require_nonempty("provider circuit", "provider_kind", provider_kind)?;
let mut connection = self.connection.clone();
redis::cmd("DEL")
.arg(self.key(provider_instance_id)?)
.arg(self.key(provider_kind)?)
.query_async::<i64>(&mut connection)
.await
.map_err(|_| redis_unavailable("reset provider circuit"))?;
@@ -182,10 +170,10 @@ impl ProviderCircuitRepository for RedisProviderCircuitRepository {
impl ProviderCircuitPort for RedisProviderCircuitRepository {
fn decision<'a>(
&'a self,
provider_instance_id: &'a ProviderInstanceId,
provider_kind: &'a ProviderKind,
) -> futures::future::BoxFuture<'a, Result<CoreCircuitDecision, ProviderCircuitError>> {
Box::pin(async move {
self.provider_circuit_decision(provider_instance_id.as_str())
self.provider_circuit_decision(provider_kind.as_str())
.await
.map(|decision| match decision {
ProviderCircuitDecision::Allow => CoreCircuitDecision::Allow,
@@ -199,10 +187,10 @@ impl ProviderCircuitPort for RedisProviderCircuitRepository {
fn observe_failure<'a>(
&'a self,
provider_instance_id: &'a ProviderInstanceId,
provider_kind: &'a ProviderKind,
) -> futures::future::BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async move {
self.observe_provider_failure(provider_instance_id.as_str())
self.observe_provider_failure(provider_kind.as_str())
.await
.map(|_| ())
.map_err(|_| ProviderCircuitError)
@@ -211,10 +199,10 @@ impl ProviderCircuitPort for RedisProviderCircuitRepository {
fn observe_success<'a>(
&'a self,
provider_instance_id: &'a ProviderInstanceId,
provider_kind: &'a ProviderKind,
) -> futures::future::BoxFuture<'a, Result<(), ProviderCircuitError>> {
Box::pin(async move {
self.observe_provider_success(provider_instance_id.as_str())
self.observe_provider_success(provider_kind.as_str())
.await
.map_err(|_| ProviderCircuitError)
})
backend/crates/gateway-store/tests/postgres/config_catalog.rs
-12
View File
@@ -1,12 +0,0 @@
use gateway_store::postgres::NewProviderInstance;
#[test]
fn provider_instance_rejects_unstable_kind() {
let instance = NewProviderInstance {
id: "instance-1".to_owned(),
provider_kind: "xAI!".to_owned(),
name: "xAI".to_owned(),
base_url: "https://example.invalid".to_owned(),
};
assert!(instance.validate().is_err());
}
backend/crates/gateway-store/tests/postgres/history.rs
+4 -15
View File
@@ -12,7 +12,6 @@ fn history_contains_no_payload_field() {
client_api_key_ref: "key-1".to_owned(),
requested_model_id: "coding".to_owned(),
provider_kind: None,
provider_instance_id: None,
provider_account_ref: None,
upstream_model_id: None,
outcome: "failed".to_owned(),
@@ -142,24 +141,14 @@ async fn native_pin_requires_committed_success_but_not_a_specific_transport() {
}
async fn seed_native_history(pool: &sqlx::PgPool) -> Result<(), sqlx::Error> {
sqlx::query(
"insert into provider_instances (
id, provider_kind, name, base_url, enabled, created_at, updated_at
) values (
'inst_history', 'openai', 'history', 'https://example.invalid',
true, now(), now()
)",
)
.execute(pool)
.await?;
sqlx::query(
"insert into provider_accounts (
id, provider_instance_id, provider_kind, name, upstream_user_id,
id, provider_kind, name, upstream_user_id,
provider_credentials_json, credential_revision, has_refresh_token,
access_token_expires_at, enabled, availability, availability_observed_at,
created_at, updated_at
) values (
'acct_history', 'inst_history', 'openai', 'history', 'user-history',
'acct_history', 'openai', 'history', 'user-history',
'{}'::jsonb, 1, false, now() + interval '1 day', true, 'ready',
now(), now(), now()
)",
@@ -170,14 +159,14 @@ async fn seed_native_history(pool: &sqlx::PgPool) -> Result<(), sqlx::Error> {
"insert into model_requests (
id, client_api_key_ref, config_revision, protocol, operation, endpoint,
client_transport, requested_model_id, input_token_estimate,
provider_instance_id, provider_kind, provider_account_id,
provider_kind, provider_account_id,
provider_account_ref, upstream_model_id, upstream_transport, attempt_count,
upstream_send_state, downstream_committed_at, outcome, client_status_code,
upstream_status_code, client_response_id, upstream_response_id,
cost_source, started_at, deadline_at, completed_at
) values (
'req_history', 'key_owner', 1, 'openai', 'responses', '/v1/responses',
'http_sse', 'history-model', 1, 'inst_history', 'openai', 'acct_history',
'http_sse', 'history-model', 1, 'openai', 'acct_history',
'acct_history', 'history-upstream',
'websocket', 1, 'sent', now(), 'succeeded', 200, 200,
'resp_gateway_visible', 'resp_upstream_native', 'unavailable',
backend/crates/gateway-store/tests/postgres/mod.rs
-1
View File
@@ -10,7 +10,6 @@ use uuid::Uuid;
mod admin_security_audit;
mod admission_recovery;
mod client_keys;
mod config_catalog;
mod execution;
mod history;
mod observability;
backend/crates/gateway-store/tests/postgres/observability.rs
+8 -19
View File
@@ -517,25 +517,14 @@ async fn seed_observability_facts(
pool: &PgPool,
now: chrono::DateTime<Utc>,
) -> Result<(), sqlx::Error> {
sqlx::query(
"insert into provider_instances (
id, provider_kind, name, base_url, enabled, created_at, updated_at
) values (
'inst_observe', 'openai', 'observability', 'https://example.invalid',
true, $1, $1
)",
)
.bind(now)
.execute(pool)
.await?;
sqlx::query(
"insert into provider_accounts (
id, provider_instance_id, provider_kind, name, email, upstream_user_id,
id, provider_kind, name, email, upstream_user_id,
upstream_account_id, plan_type, provider_credentials_json, credential_revision,
has_refresh_token, access_token_expires_at, enabled, availability,
availability_observed_at, created_at, updated_at
) values (
'acct_observe', 'inst_observe', 'openai', 'primary', 'account@example.invalid',
'acct_observe', 'openai', 'primary', 'account@example.invalid',
'user-observe', null, 'pro', '{}'::jsonb, 1, false, $1 + interval '1 day',
true, 'ready', $1, $1, $1
)",
@@ -547,7 +536,7 @@ async fn seed_observability_facts(
"insert into model_requests (
id, client_api_key_ref, config_revision, protocol, operation, endpoint,
client_transport, requested_model_id, input_token_estimate,
provider_instance_id, provider_kind, provider_account_id,
provider_kind, provider_account_id,
provider_account_ref, upstream_model_id, upstream_transport, websocket_pool,
attempt_count,
upstream_send_state, downstream_committed_at, outcome, client_status_code,
@@ -559,7 +548,7 @@ async fn seed_observability_facts(
started_at, deadline_at, completed_at
) values (
'req_observe_success', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
'http_sse', 'public-model', 100, 'inst_observe', 'openai', 'acct_observe',
'http_sse', 'public-model', 100, 'openai', 'acct_observe',
'acct_observe', 'upstream-model',
'http_sse', 'reuse', 1, 'sent', $1 - interval '19 minutes', 'succeeded', 200, 200,
'resp_observe_success', 'upstream_req_success', 'upstream_resp_success',
@@ -575,7 +564,7 @@ async fn seed_observability_facts(
"insert into model_requests (
id, client_api_key_ref, config_revision, protocol, operation, endpoint,
client_transport, requested_model_id, input_token_estimate,
provider_instance_id, provider_kind, provider_account_id,
provider_kind, provider_account_id,
provider_account_ref, upstream_model_id, upstream_transport, attempt_count,
upstream_send_state, outcome, client_status_code, upstream_status_code,
error_kind, provider_error_code, error_message, retry_after_ms,
@@ -584,7 +573,7 @@ async fn seed_observability_facts(
started_at, deadline_at, completed_at
) values (
'req_observe_failed', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
'http_sse', 'public-model', 80, 'inst_observe', 'openai', 'acct_observe',
'http_sse', 'public-model', 80, 'openai', 'acct_observe',
'acct_observe', 'upstream-model',
'http_sse', 2, 'sent', 'failed', 502, 429, 'rate_limited', 'rate_limit',
'upstream limited', 1000, 0, 0, true, false, 'unavailable', 700,
@@ -597,12 +586,12 @@ async fn seed_observability_facts(
sqlx::query(
"insert into ops_events (
id, model_request_id, attempt_index, level, component, operation,
provider_instance_id, provider_kind, provider_account_id,
provider_kind, provider_account_id,
provider_account_ref, upstream_model_id, failure_kind, status_code,
provider_error_code, retry_after_ms, latency_ms, message, occurrence_count, created_at
) values (
'ops_observe_retry', 'req_observe_failed', 1, 'warning', 'routing', 'responses',
'inst_observe', 'openai', 'acct_observe', 'acct_observe',
'openai', 'acct_observe', 'acct_observe',
'upstream-model', 'rate_limited', 429, 'rate_limit', 1000, 300,
'first account was limited', 1, $1 - interval '9 minutes 30 seconds'
)",
backend/crates/gateway-store/tests/postgres/ops_events.rs
-1
View File
@@ -10,7 +10,6 @@ fn request_scoped_ops_event_requires_attempt_index() {
level: OpsEventLevel::Warning,
component: "routing".to_owned(),
operation: "fallback".to_owned(),
provider_instance_id: None,
provider_kind: None,
provider_account_id: None,
provider_account_ref: None,
backend/crates/gateway-store/tests/postgres/provider_accounts.rs
+31 -88
View File
@@ -5,7 +5,7 @@ use gateway_admin::{
model::{
MutationActor, MutationContext, PageSize, Revision as AdminRevision,
accounts::{
AccountListQuery, AccountSort, AccountSortField, AccountStatus, DeleteAccount,
AccountListQuery, AccountSort, AccountSortField, AccountStatus, DeleteAccounts,
SetAccountEnabled as AdminSetAccountEnabled, SortDirection,
},
observability::TimeRange,
@@ -52,9 +52,6 @@ async fn core_quota_batch_reads_only_observed_accounts_in_one_contract_call() {
let Some(database) = TestDatabase::create("provider_account_quota_batch").await else {
return;
};
seed_instance(&database.pool, "inst_openai_admin", "openai")
.await
.expect("seed OpenAI instance");
let repository = PgProviderAccountRepository::new(database.pool.clone());
for id in ["acct_quota_a", "acct_quota_b", "acct_quota_empty"] {
repository
@@ -117,19 +114,12 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
let Some(database) = TestDatabase::create("provider_account_terminal_list").await else {
return;
};
seed_instance(&database.pool, "inst_openai_admin", "openai")
.await
.expect("seed OpenAI instance");
seed_instance(&database.pool, "inst_xai_admin", "xai")
.await
.expect("seed xAI instance");
let repository = PgProviderAccountRepository::new(database.pool.clone());
let now = Utc::now();
let mut alpha = account("acct_alpha", "user-alpha");
alpha.email = Some("alpha@example.invalid".to_owned());
let mut beta = account("acct_beta", "user-beta");
beta.provider_instance_id = "inst_xai_admin".to_owned();
beta.provider_kind = "xai".to_owned();
beta.email = Some("beta@example.invalid".to_owned());
beta.availability = ProviderAccountAvailability::Banned;
@@ -150,7 +140,6 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
ModelRequestSeed {
request_id: "req_alpha_recent",
account_id: "acct_alpha",
provider_instance_id: "inst_openai_admin",
provider_kind: "openai",
model: "gpt-list",
total_tokens: 10,
@@ -165,7 +154,6 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
ModelRequestSeed {
request_id: "req_beta_recent",
account_id: "acct_beta",
provider_instance_id: "inst_xai_admin",
provider_kind: "xai",
model: "grok-list",
total_tokens: 50,
@@ -180,7 +168,6 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
ModelRequestSeed {
request_id: "req_beta_expired_retention",
account_id: "acct_beta",
provider_instance_id: "inst_xai_admin",
provider_kind: "xai",
model: "grok-list",
total_tokens: 500,
@@ -195,7 +182,6 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
ModelRequestSeed {
request_id: "req_charlie_recent",
account_id: "acct_charlie",
provider_instance_id: "inst_openai_admin",
provider_kind: "openai",
model: "gpt-list",
total_tokens: 60,
@@ -308,16 +294,12 @@ async fn terminal_credential_list_preserves_grouped_filters_and_unpaged_collecti
let Some(database) = TestDatabase::create("provider_credential_terminal_windows").await else {
return;
};
seed_instance(&database.pool, "inst_xai_admin", "xai")
.await
.expect("seed xAI instance");
let repository = PgProviderAccountRepository::new(database.pool.clone());
for index in 0..205_u16 {
let mut credential = account(
&format!("acct_xai_{index:03}"),
&format!("user-xai-{index:03}"),
);
credential.provider_instance_id = "inst_xai_admin".to_owned();
credential.provider_kind = "xai".to_owned();
credential.availability = match index {
0 => ProviderAccountAvailability::Expired,
@@ -337,7 +319,6 @@ async fn terminal_credential_list_preserves_grouped_filters_and_unpaged_collecti
.list_credentials(
&provider,
CredentialListQuery {
provider_instance_id: None,
availability: None,
enabled: None,
window: CredentialListWindow::All,
@@ -352,7 +333,6 @@ async fn terminal_credential_list_preserves_grouped_filters_and_unpaged_collecti
.list_credentials(
&provider,
CredentialListQuery {
provider_instance_id: None,
availability: None,
enabled: None,
window: CredentialListWindow::Page {
@@ -370,7 +350,6 @@ async fn terminal_credential_list_preserves_grouped_filters_and_unpaged_collecti
.list_credentials(
&provider,
CredentialListQuery {
provider_instance_id: None,
availability: Some(CredentialAvailabilityFilter::AnyOf(vec![
gateway_admin::model::accounts::AccountAvailability::Expired,
gateway_admin::model::accounts::AccountAvailability::Banned,
@@ -392,9 +371,6 @@ async fn terminal_admin_usage_chunks_large_selections_and_preserves_exact_costs(
let Some(database) = TestDatabase::create("provider_account_terminal_usage").await else {
return;
};
seed_instance(&database.pool, "inst_openai_admin", "openai")
.await
.expect("seed OpenAI instance");
let repository = PgProviderAccountRepository::new(database.pool.clone());
repository
.insert_provider_account(account("acct_usage_exact", "user-usage-exact"))
@@ -406,7 +382,6 @@ async fn terminal_admin_usage_chunks_large_selections_and_preserves_exact_costs(
ModelRequestSeed {
request_id: "req_usage_exact",
account_id: "acct_usage_exact",
provider_instance_id: "inst_openai_admin",
provider_kind: "openai",
model: "gpt-exact",
total_tokens: 18,
@@ -448,9 +423,6 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
let Some(database) = TestDatabase::create("provider_account_terminal_mutation").await else {
return;
};
seed_instance(&database.pool, "inst_openai_admin", "openai")
.await
.expect("seed OpenAI instance");
PgProviderAccountRepository::new(database.pool.clone())
.insert_provider_account(account("acct_terminal_mutation", "user-terminal-mutation"))
.await
@@ -482,10 +454,10 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
assert!(!enabled);
let error = store
.delete_account(
DeleteAccount {
.delete_accounts(
DeleteAccounts {
expected_config_revision: AdminRevision::new(1).expect("stale revision"),
account_id: "acct_terminal_mutation".to_owned(),
account_ids: vec!["acct_terminal_mutation".to_owned()],
},
&context,
)
@@ -504,10 +476,10 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
assert_eq!(audit_count, 1);
let revision = store
.delete_account(
DeleteAccount {
.delete_accounts(
DeleteAccounts {
expected_config_revision: revision,
account_id: "acct_terminal_mutation".to_owned(),
account_ids: vec!["acct_terminal_mutation".to_owned()],
},
&context,
)
@@ -537,23 +509,29 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
}
#[tokio::test]
async fn terminal_admin_delete_removes_an_enabled_account_atomically() {
async fn terminal_admin_delete_removes_enabled_accounts_in_one_transaction() {
let Some(database) = TestDatabase::create("provider_account_enabled_delete").await else {
return;
};
seed_instance(&database.pool, "inst_openai_admin", "openai")
.await
.expect("seed OpenAI instance");
PgProviderAccountRepository::new(database.pool.clone())
.insert_provider_account(account("acct_enabled_delete", "user-enabled-delete"))
.await
.expect("insert enabled account");
let repository = PgProviderAccountRepository::new(database.pool.clone());
for (account_id, upstream_user_id) in [
("acct_enabled_delete_a", "user-enabled-delete-a"),
("acct_enabled_delete_b", "user-enabled-delete-b"),
] {
repository
.insert_provider_account(account(account_id, upstream_user_id))
.await
.expect("insert enabled account");
}
let revision = PgAdminAccountStore::new(database.pool.clone())
.delete_account(
DeleteAccount {
.delete_accounts(
DeleteAccounts {
expected_config_revision: AdminRevision::new(1).expect("initial revision"),
account_id: "acct_enabled_delete".to_owned(),
account_ids: vec![
"acct_enabled_delete_a".to_owned(),
"acct_enabled_delete_b".to_owned(),
],
},
&MutationContext {
actor: MutationActor::System,
@@ -565,7 +543,8 @@ async fn terminal_admin_delete_removes_an_enabled_account_atomically() {
assert_eq!(revision.get(), 2);
assert_eq!(
account_count(&database.pool, "acct_enabled_delete").await,
account_count(&database.pool, "acct_enabled_delete_a").await
+ account_count(&database.pool, "acct_enabled_delete_b").await,
0
);
database.close().await;
@@ -576,13 +555,9 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
let Some(database) = TestDatabase::create("provider_account_admin_upsert").await else {
return;
};
seed_instance(&database.pool, "inst_openai_admin", "openai")
.await
.expect("seed Codex instance");
let repository = PgProviderAccountRepository::new(database.pool.clone());
let scope = ProviderAccountAdminScope {
provider_kind: "openai".to_owned(),
provider_instance_id: "inst_openai_admin".to_owned(),
};
let imported = repository
.import_provider_accounts(
@@ -671,14 +646,10 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
let Some(database) = TestDatabase::create("provider_account_core_refresh").await else {
return;
};
seed_instance(&database.pool, "inst_xai_refresh", "xai")
.await
.expect("seed xAI instance");
let repository = PgProviderAccountRepository::new(database.pool.clone());
repository
.insert_provider_account(NewProviderAccount {
id: "acct_core_refresh".to_owned(),
provider_instance_id: "inst_xai_refresh".to_owned(),
provider_kind: "xai".to_owned(),
name: "before refresh".to_owned(),
email: Some("before@example.invalid".to_owned()),
@@ -789,16 +760,9 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
let Some(database) = TestDatabase::create("provider_account_admin").await else {
return;
};
seed_instance(&database.pool, "inst_openai_admin", "openai")
.await
.expect("seed Codex instance");
seed_instance(&database.pool, "inst_xai_admin", "xai")
.await
.expect("seed xAI instance");
let repository = PgProviderAccountRepository::new(database.pool.clone());
let scope = ProviderAccountAdminScope {
provider_kind: "openai".to_owned(),
provider_instance_id: "inst_openai_admin".to_owned(),
};
let mut cooldown_account = account("acct_admin_b", "user-admin-b");
@@ -862,7 +826,6 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
let wrong_scope = ProviderAccountAdminScope {
provider_kind: "xai".to_owned(),
provider_instance_id: "inst_xai_admin".to_owned(),
};
let wrong_scope_error = repository
.rotate_provider_account(
@@ -949,7 +912,6 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
fn account(id: &str, upstream_user_id: &str) -> NewProviderAccount {
NewProviderAccount {
id: id.to_owned(),
provider_instance_id: "inst_openai_admin".to_owned(),
provider_kind: "openai".to_owned(),
name: id.to_owned(),
email: Some(format!("{id}@example.invalid")),
@@ -1024,7 +986,6 @@ fn audit(id: &str, action: &str, entity_ref: &str) -> AdminAuditEvent {
struct ModelRequestSeed<'a> {
request_id: &'a str,
account_id: &'a str,
provider_instance_id: &'a str,
provider_kind: &'a str,
model: &'a str,
total_tokens: i64,
@@ -1040,7 +1001,7 @@ async fn seed_model_request(
"insert into model_requests (
id, client_api_key_ref, config_revision, protocol, operation, endpoint,
client_transport, requested_model_id, input_token_estimate,
provider_instance_id, provider_kind, provider_account_id,
provider_kind, provider_account_id,
provider_account_ref, upstream_model_id, upstream_transport, attempt_count,
upstream_send_state, outcome, client_status_code, upstream_status_code,
input_tokens, output_tokens, cached_tokens, cache_write_tokens, reasoning_tokens,
@@ -1048,15 +1009,14 @@ async fn seed_model_request(
started_at, deadline_at, completed_at
) values (
$1, 'key-provider-account-test', 1, 'openai', 'responses', '/v1/responses',
'http_sse', $5, 1, $3, $4, $2, $2, $5, 'http_sse', 1,
'sent', 'succeeded', 200, 200, $6, 0, 0, 0, 0,
$6, 'provider_reported', $7::numeric, 'USD', $8,
$8 + interval '5 minutes', $8 + interval '1 second'
'http_sse', $4, 1, $3, $2, $2, $4, 'http_sse', 1,
'sent', 'succeeded', 200, 200, $5, 0, 0, 0, 0,
$5, 'provider_reported', $6::numeric, 'USD', $7,
$7 + interval '5 minutes', $7 + interval '1 second'
)",
)
.bind(seed.request_id)
.bind(seed.account_id)
.bind(seed.provider_instance_id)
.bind(seed.provider_kind)
.bind(seed.model)
.bind(seed.total_tokens)
@@ -1067,23 +1027,6 @@ async fn seed_model_request(
Ok(())
}
async fn seed_instance(
pool: &sqlx::PgPool,
id: &str,
provider_kind: &str,
) -> Result<(), sqlx::Error> {
sqlx::query(
"insert into provider_instances (
id, provider_kind, name, base_url, enabled, created_at, updated_at
) values ($1, $2, $1, 'https://example.invalid', true, now(), now())",
)
.bind(id)
.bind(provider_kind)
.execute(pool)
.await?;
Ok(())
}
async fn current_revision(pool: &sqlx::PgPool) -> i64 {
sqlx::query_scalar("select config_revision from runtime_settings where id = 1")
.fetch_one(pool)
backend/crates/gateway-store/tests/redis/credential_cooldown.rs
+4 -5
View File
@@ -13,7 +13,7 @@ use gateway_core::{
QuotaObservation, QuotaWriteOutcome,
},
error::{StoreError as CoreStoreError, StoreErrorKind},
routing::{ProviderInstanceId, ProviderKind},
routing::ProviderKind,
};
use gateway_store::{
Revision, StoreBackend, StoreError, StoreResult,
@@ -58,11 +58,11 @@ impl ProviderAccountStore for CooldownTestAccountStore {
Ok(self.accounts.clone())
}
async fn list_for_instance(
async fn list_for_provider(
&self,
_instance: &ProviderInstanceId,
_provider: &ProviderKind,
) -> Result<Vec<ProviderAccount>, CoreStoreError> {
unreachable!("cooldown tests do not list one instance")
unreachable!("cooldown tests do not list one provider")
}
async fn load_credential(
@@ -377,7 +377,6 @@ fn ready_test_account(account_id: &str) -> ProviderAccount {
fn test_provider_account(account_id: &str) -> ProviderAccount {
ProviderAccount::new(
ProviderAccountId::new(account_id).expect("valid account ID"),
ProviderInstanceId::new("inst_cooldown_test").expect("valid instance ID"),
ProviderKind::new("openai").expect("valid provider kind"),
"Cooldown test".to_owned(),
"upstream-user".to_owned(),
backend/crates/providers/openai/src/admin.rs
+23 -107
View File
@@ -38,12 +38,9 @@ use gateway_core::operation::{
};
use gateway_core::provider_ports::{
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingFlowPort, OAuthPendingPutOutcome,
OAuthPendingTakeOutcome, ProviderInstanceCatalogPort, ProviderRuntimePolicyPort,
ProviderStoreError, ProviderStoreErrorKind,
};
use gateway_core::routing::{
InstanceHealth, ProviderInstance, ProviderInstanceId, ProviderKind, UpstreamModelId,
OAuthPendingTakeOutcome, ProviderRuntimePolicyPort, ProviderStoreError, ProviderStoreErrorKind,
};
use gateway_core::routing::{ProviderKind, UpstreamModelId};
use secrecy::{ExposeSecret as _, SecretString};
use serde::Deserialize;
use serde_json::{Map, Number, Value};
@@ -68,7 +65,7 @@ use crate::transport::profile::{
};
const PROVIDER_NAME: &str = "openai";
const PENDING_DOCUMENT_SCHEMA_VERSION: u64 = 1;
const PENDING_DOCUMENT_SCHEMA_VERSION: u64 = 2;
const MAX_ACCESS_TOKEN_BYTES: usize = 16 * 1024;
const MAX_REFRESH_TOKEN_BYTES: usize = 64 * 1024;
@@ -77,7 +74,6 @@ pub(crate) struct OpenAiAdminProvider {
provider_kind: ProviderKind,
profile: CodexWireProfileState,
accounts: Arc<dyn ProviderAccountStore>,
instances: Arc<dyn ProviderInstanceCatalogPort>,
credentials: Arc<CodexCredentialAdminService>,
verifier: Arc<dyn CodexAccountIdentityVerifier>,
oauth: Arc<dyn CodexOAuthAdmin>,
@@ -103,7 +99,6 @@ impl OpenAiAdminProvider {
provider_kind: ProviderKind,
profile: CodexWireProfileState,
accounts: Arc<dyn ProviderAccountStore>,
instances: Arc<dyn ProviderInstanceCatalogPort>,
services: OpenAiAdminServices,
websocket_pool: Arc<CodexWebSocketPool>,
desktop_release: CodexDesktopReleaseStatus,
@@ -112,7 +107,6 @@ impl OpenAiAdminProvider {
provider_kind,
profile,
accounts,
instances,
credentials: services.credentials,
verifier: services.verifier,
oauth: services.oauth,
@@ -136,30 +130,6 @@ impl OpenAiAdminProvider {
.ok_or_else(|| provider_admin_error(ProviderAdminErrorKind::NotFound))
}
async fn instance(
&self,
account: &ProviderAccount,
) -> Result<ProviderInstance, ProviderAdminError> {
let instance = self
.instances
.list_instances(&self.provider_kind, true)
.await
.map_err(map_provider_store_error)?
.into_iter()
.find(|instance| instance.id() == account.instance())
.ok_or_else(|| provider_admin_error(ProviderAdminErrorKind::NotFound))?;
if instance.provider_kind() != &self.provider_kind {
return Err(provider_admin_error(ProviderAdminErrorKind::Internal));
}
Ok(ProviderInstance::new(
instance.id().clone(),
instance.provider_kind().clone(),
instance.base_url().to_owned(),
instance.enabled(),
InstanceHealth::Healthy,
))
}
async fn preserve_existing_installation_id(
&self,
mut incoming: NewProviderAccount,
@@ -173,7 +143,6 @@ impl OpenAiAdminProvider {
return Ok(incoming);
};
if existing.provider() != incoming.account.provider()
|| existing.instance() != incoming.account.instance()
|| existing.upstream_user_id() != incoming.account.upstream_user_id()
|| existing.upstream_account_id() != incoming.account.upstream_account_id()
{
@@ -283,10 +252,9 @@ impl ProviderAdmin for OpenAiAdminProvider {
) -> Result<PreparedCredentialImport, ProviderAdminError> {
let prepared = self
.credentials
.prepare_import_document(
command.provider_instance_id.clone(),
Value::Object(command.document.into_provider_data().into_inner()),
)
.prepare_import_document(Value::Object(
command.document.into_provider_data().into_inner(),
))
.await
.map_err(map_credential_admin_error)?;
let observed_at = Utc::now();
@@ -297,7 +265,6 @@ impl ProviderAdmin for OpenAiAdminProvider {
}
Ok(PreparedCredentialImport {
provider_kind: self.provider_kind.clone(),
provider_instance_id: command.provider_instance_id,
credentials,
})
}
@@ -340,21 +307,16 @@ impl ProviderAdmin for OpenAiAdminProvider {
PreparedAuthorizationCredential::Create(prepared_create(credential, Utc::now())?)
}
CompletedCodexOAuthCredential::Reauthorize(credential) => {
let (provider_instance_id, provider_kind) = match completed.mutation.target() {
AuthorizationMutationTarget::Reauthorize {
provider_instance_id,
..
} => (
provider_instance_id.clone(),
completed.mutation.provider_kind().clone(),
),
let provider_kind = match completed.mutation.target() {
AuthorizationMutationTarget::Reauthorize { .. } => {
completed.mutation.provider_kind().clone()
}
AuthorizationMutationTarget::Create { .. } => {
return Err(provider_admin_error(ProviderAdminErrorKind::Internal));
}
};
PreparedAuthorizationCredential::Reauthorize(prepared_rotation(
credential,
provider_instance_id,
provider_kind,
)?)
}
@@ -426,11 +388,7 @@ impl ProviderAdmin for OpenAiAdminProvider {
next_refresh_at,
})
.map_err(map_credential_admin_error)?;
prepared_rotation(
prepared,
command.account.provider_instance_id,
command.account.provider_kind,
)
prepared_rotation(prepared, command.account.provider_kind)
}
async fn prepare_refresh(
@@ -447,11 +405,7 @@ impl ProviderAdmin for OpenAiAdminProvider {
.manual_refresh(account_id, revision)
.await
.map_err(map_credential_admin_error)?;
prepared_rotation(
prepared,
command.account.provider_instance_id,
command.account.provider_kind,
)
prepared_rotation(prepared, command.account.provider_kind)
}
async fn quota(
@@ -463,12 +417,11 @@ impl ProviderAdmin for OpenAiAdminProvider {
refresh,
rolling_usage: _,
} = request;
let account = self.account(&account_id).await?;
self.account(&account_id).await?;
let snapshot = if refresh {
let instance = self.instance(&account).await?;
Some(
self.quota
.refresh_account(&instance, &account_id)
.refresh_account(&account_id)
.await
.map_err(map_quota_error)?,
)
@@ -487,21 +440,20 @@ impl ProviderAdmin for OpenAiAdminProvider {
refresh: bool,
) -> Result<ProviderModels, ProviderAdminError> {
let account = self.account(account_id).await?;
let instance = self.instance(&account).await?;
let models = if refresh {
self.catalog
.synchronize_account(&instance, account_id)
.synchronize_account(account_id)
.await
.map_err(map_catalog_error)?
} else {
self.catalog
.cached_account_models(instance.id(), account_id, account.revision())
.cached_account_models(account_id, account.revision())
.map_err(map_catalog_error)?
.unwrap_or_default()
};
let observed_at = self
.catalog
.cached(instance.id())
.cached()
.map_err(map_catalog_error)?
.map(|snapshot| DateTime::<Utc>::from(snapshot.observed_at()));
let models = models
@@ -573,7 +525,6 @@ fn prepared_create(
let availability = admin_availability(account.availability());
Ok(PreparedCredentialCreate {
account_id: account.id().clone(),
provider_instance_id: account.instance().clone(),
provider_kind: account.provider().clone(),
name: account.name().to_owned(),
email: account.email().map(str::to_owned),
@@ -594,7 +545,6 @@ fn prepared_create(
fn prepared_rotation(
prepared: crate::credential::PreparedCodexCredentialRotation,
provider_instance_id: ProviderInstanceId,
provider_kind: ProviderKind,
) -> Result<PreparedCredentialRotation, ProviderAdminError> {
let (profile, credential, guard) = prepared.into_parts();
@@ -615,7 +565,6 @@ fn prepared_rotation(
Ok(PreparedCredentialRotation::new(
PreparedCredentialRotationFacts {
account_id,
provider_instance_id,
provider_kind,
expected_credential_revision,
name: profile.name,
@@ -678,7 +627,6 @@ fn account_from_record(account: &AccountRecord) -> Result<ProviderAccount, Provi
.map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?;
Ok(ProviderAccount::new(
account_id,
account.provider_instance_id.clone(),
account.provider_kind.clone(),
account.name.clone(),
account.upstream_user_id.clone(),
@@ -742,7 +690,6 @@ fn dashboard_desktop_release(
fn account_matches_record(account: &ProviderAccount, record: &AccountRecord) -> bool {
account.id().as_str() == record.id
&& account.instance() == &record.provider_instance_id
&& account.provider() == &record.provider_kind
&& account.revision().get() == record.credential_revision.get()
&& account.name() == record.name
@@ -981,7 +928,6 @@ struct PendingDocument {
flow_id: String,
owner_ref: String,
started_request_ref: String,
provider_instance_id: String,
name: String,
expires_at: DateTime<Utc>,
state: String,
@@ -1007,11 +953,9 @@ struct PendingMutationDocument {
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
enum PendingTargetDocument {
Create {
provider_instance_id: String,
name: String,
},
Reauthorize {
provider_instance_id: String,
account_id: String,
expected_credential_revision: u64,
},
@@ -1039,10 +983,6 @@ fn encode_pending(pending: &CodexPendingAuthorization) -> Map<String, Value> {
"started_request_ref".to_owned(),
Value::String(pending.started_request_ref().to_owned()),
);
document.insert(
"provider_instance_id".to_owned(),
Value::String(pending.provider_instance_id().to_owned()),
);
document.insert("name".to_owned(), Value::String(pending.name().to_owned()));
document.insert(
"expires_at".to_owned(),
@@ -1111,27 +1051,15 @@ fn encode_mutation(mutation: &PendingAuthorizationMutation) -> Map<String, Value
fn encode_target(target: &AuthorizationMutationTarget) -> Map<String, Value> {
let mut document = Map::new();
match target {
AuthorizationMutationTarget::Create {
provider_instance_id,
name,
} => {
AuthorizationMutationTarget::Create { name } => {
document.insert("kind".to_owned(), Value::String("create".to_owned()));
document.insert(
"provider_instance_id".to_owned(),
Value::String(provider_instance_id.to_string()),
);
document.insert("name".to_owned(), Value::String(name.clone()));
}
AuthorizationMutationTarget::Reauthorize {
provider_instance_id,
account_id,
expected_credential_revision,
} => {
document.insert("kind".to_owned(), Value::String("reauthorize".to_owned()));
document.insert(
"provider_instance_id".to_owned(),
Value::String(provider_instance_id.to_string()),
);
document.insert(
"account_id".to_owned(),
Value::String(account_id.to_string()),
@@ -1174,7 +1102,6 @@ fn decode_pending(
flow_id: document.flow_id,
owner_ref: document.owner_ref,
started_request_ref: document.started_request_ref,
provider_instance_id: document.provider_instance_id,
name: document.name,
expires_at: document.expires_at,
state: SecretString::from(document.state),
@@ -1197,21 +1124,11 @@ fn decode_mutation(
let provider_kind = ProviderKind::new(document.provider_kind)
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?;
let target = match document.target {
PendingTargetDocument::Create {
provider_instance_id,
name,
} => AuthorizationMutationTarget::Create {
provider_instance_id: ProviderInstanceId::new(provider_instance_id)
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?,
name,
},
PendingTargetDocument::Create { name } => AuthorizationMutationTarget::Create { name },
PendingTargetDocument::Reauthorize {
provider_instance_id,
account_id,
expected_credential_revision,
} => AuthorizationMutationTarget::Reauthorize {
provider_instance_id: ProviderInstanceId::new(provider_instance_id)
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?,
account_id: ProviderAccountId::new(account_id)
.map_err(|_| CodexOAuthPendingStoreError::InvalidValue)?,
expected_credential_revision: Revision::new(expected_credential_revision)
@@ -1352,7 +1269,7 @@ fn map_oauth_error(error: CodexOAuthAdminError) -> ProviderAdminError {
fn map_quota_error(error: CodexCredentialQuotaError) -> ProviderAdminError {
use CodexCredentialQuotaError as Error;
provider_admin_error(match error {
Error::InvalidInstance | Error::InvalidCredentialData => ProviderAdminErrorKind::Invalid,
Error::InvalidCredentialData => ProviderAdminErrorKind::Invalid,
Error::NotFound => ProviderAdminErrorKind::NotFound,
Error::RevisionConflict => ProviderAdminErrorKind::Conflict,
Error::Repository(_) | Error::Store | Error::Upstream => {
@@ -1364,10 +1281,9 @@ fn map_quota_error(error: CodexCredentialQuotaError) -> ProviderAdminError {
fn map_catalog_error(error: CodexCredentialCatalogError) -> ProviderAdminError {
use CodexCredentialCatalogError as Error;
provider_admin_error(match error {
Error::InvalidInstance
| Error::InvalidCredentialData
| Error::ConflictingModelFacts
| Error::InvalidEtag => ProviderAdminErrorKind::Invalid,
Error::InvalidCredentialData | Error::ConflictingModelFacts | Error::InvalidEtag => {
ProviderAdminErrorKind::Invalid
}
Error::NoEligibleCredential => ProviderAdminErrorKind::NotFound,
Error::ConcurrentUpdate => ProviderAdminErrorKind::Conflict,
Error::Upstream | Error::Cache => ProviderAdminErrorKind::Unavailable,
backend/crates/providers/openai/src/credential/admin.rs
+1 -8
View File
@@ -19,7 +19,7 @@ use gateway_core::provider_ports::{
ProviderRefreshCapacityRequest, ProviderRefreshLeaseRequest, ProviderRefreshPolicy,
ProviderRuntimePolicyPort,
};
use gateway_core::routing::{ProviderInstanceId, ProviderKind};
use gateway_core::routing::ProviderKind;
use secrecy::{ExposeSecret, SecretString};
use serde::Serialize;
use serde_json::Value;
@@ -65,7 +65,6 @@ const CPR_ACCOUNT_KEYS: &[&str] = &[
pub struct ImportCodexOAuthCredential {
pub account_id: String,
pub provider_instance_id: String,
pub name: String,
pub secret: CodexOAuthSecret,
pub verified_account: CodexAccountProfile,
@@ -78,7 +77,6 @@ impl std::fmt::Debug for ImportCodexOAuthCredential {
formatter
.debug_struct("ImportCodexOAuthCredential")
.field("account_id", &self.account_id)
.field("provider_instance_id", &self.provider_instance_id)
.field("name", &self.name)
.field("secret", &"<redacted>")
.field("verified_account", &self.verified_account)
@@ -346,8 +344,6 @@ impl CodexCredentialAdmin {
) -> Result<NewProviderAccount, CodexCredentialAdminError> {
let account_id = ProviderAccountId::new(input.account_id)
.map_err(|_| CodexCredentialAdminError::InvalidInput)?;
let instance = ProviderInstanceId::new(input.provider_instance_id)
.map_err(|_| CodexCredentialAdminError::InvalidInput)?;
let provider = ProviderKind::new(PROVIDER_NAME)
.map_err(|_| CodexCredentialAdminError::InvalidInput)?;
if input.name.trim().is_empty() {
@@ -363,7 +359,6 @@ impl CodexCredentialAdmin {
.map_err(|_| CodexCredentialAdminError::InvalidCredential)?;
let account = ProviderAccount::new(
account_id,
instance,
provider,
input.name,
upstream_user_id,
@@ -626,7 +621,6 @@ impl CodexCredentialAdminService {
/// 按文档结构严格判别账号布局，并归一到唯一 `NewProviderAccount` 写入路径。
pub async fn prepare_import_document(
&self,
provider_instance_id: ProviderInstanceId,
payload: Value,
) -> Result<PreparedCodexAccountImport, CodexCredentialAdminError> {
if serde_json::to_vec(&payload)
@@ -673,7 +667,6 @@ impl CodexCredentialAdminService {
)?;
let mut prepared = CodexCredentialAdmin.prepare_import(ImportCodexOAuthCredential {
account_id,
provider_instance_id: provider_instance_id.to_string(),
name: candidate
.name
.filter(|name| !name.trim().is_empty())
backend/crates/providers/openai/src/credential/catalog.rs
+63 -121
View File
@@ -9,14 +9,13 @@ use gateway_core::engine::credential::{
AccountAvailability, CredentialRevision, ProviderAccount, ProviderAccountId,
};
use gateway_core::engine::provider::ProviderCatalogGeneration;
use gateway_core::routing::{ProviderInstance, ProviderInstanceId};
use secrecy::ExposeSecret;
use thiserror::Error;
use tokio::sync::Notify;
use uuid::Uuid;
use super::repository::{CodexCredentialRepository, CredentialRepositoryError};
use crate::provider::{CodexOriginPolicy, CodexProviderInstanceConfig};
use crate::provider::OFFICIAL_CODEX_BASE_URL;
use crate::transport::profile::CodexWireProfileState;
use crate::transport::{CodexBackendClient, CodexCatalogModel, CodexRequestContext};
@@ -24,7 +23,6 @@ const MAX_RESPONSE_ETAG_BYTES: usize = 256;
#[derive(Clone)]
pub struct CodexCredentialCatalogSnapshot {
provider_instance_id: ProviderInstanceId,
observed_at: SystemTime,
models: Vec<CodexCatalogModel>,
account_models: BTreeMap<ProviderAccountId, CodexAccountEntitlement>,
@@ -37,11 +35,6 @@ struct CodexAccountEntitlement {
}
impl CodexCredentialCatalogSnapshot {
#[must_use]
pub const fn provider_instance_id(&self) -> &ProviderInstanceId {
&self.provider_instance_id
}
#[must_use]
pub const fn observed_at(&self) -> SystemTime {
self.observed_at
@@ -69,7 +62,6 @@ impl fmt::Debug for CodexCredentialCatalogSnapshot {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("CodexCredentialCatalogSnapshot")
.field("provider_instance_id", &self.provider_instance_id)
.field("observed_at", &self.observed_at)
.field("model_count", &self.models.len())
.field("account_count", &self.account_models.len())
@@ -79,8 +71,6 @@ impl fmt::Debug for CodexCredentialCatalogSnapshot {
#[derive(Debug, Error)]
pub enum CodexCredentialCatalogError {
#[error("Codex model catalog instance is invalid")]
InvalidInstance,
#[error("Codex model catalog has no eligible account")]
NoEligibleCredential,
#[error("Codex model catalog account data is invalid")]
@@ -99,16 +89,16 @@ pub enum CodexCredentialCatalogError {
#[derive(Default)]
struct CatalogEtagState {
applied: BTreeMap<ProviderInstanceId, String>,
pending: BTreeMap<ProviderInstanceId, String>,
inflight: BTreeMap<ProviderInstanceId, String>,
applied: Option<String>,
pending: Option<String>,
inflight: Option<String>,
}
#[derive(Default)]
struct CatalogCacheState {
revision: u64,
generation: u64,
snapshots: BTreeMap<ProviderInstanceId, CodexCredentialCatalogSnapshot>,
snapshot: Option<CodexCredentialCatalogSnapshot>,
}
struct FetchedAccountModels {
@@ -116,7 +106,7 @@ struct FetchedAccountModels {
etag: Option<String>,
}
struct FetchedInstanceCatalog {
struct FetchedCatalog {
snapshot: CodexCredentialCatalogSnapshot,
etags: Vec<String>,
}
@@ -132,7 +122,6 @@ pub struct CodexCredentialCatalogService {
repository: CodexCredentialRepository,
profile: CodexWireProfileState,
http: reqwest::Client,
origin_policy: Arc<dyn CodexOriginPolicy>,
cache: Arc<RwLock<CatalogCacheState>>,
etags: Arc<Mutex<CatalogEtagState>>,
etag_notification: Arc<Notify>,
@@ -143,13 +132,11 @@ impl CodexCredentialCatalogService {
repository: CodexCredentialRepository,
profile: CodexWireProfileState,
http: reqwest::Client,
origin_policy: Arc<dyn CodexOriginPolicy>,
) -> Self {
Self {
repository,
profile,
http,
origin_policy,
cache: Arc::new(RwLock::new(CatalogCacheState::default())),
etags: Arc::new(Mutex::new(CatalogEtagState::default())),
etag_notification: Arc::new(Notify::new()),
@@ -167,25 +154,23 @@ impl CodexCredentialCatalogService {
pub fn cached(
&self,
instance: &ProviderInstanceId,
) -> Result<Option<CodexCredentialCatalogSnapshot>, CodexCredentialCatalogError> {
Ok(self
.cache
.read()
.map_err(|_| CodexCredentialCatalogError::Cache)?
.snapshots
.get(instance)
.snapshot
.as_ref()
.cloned())
}
/// 读取单账号当前仍新鲜的模型 entitlement；不触发网络。
pub fn cached_account_models(
&self,
instance: &ProviderInstanceId,
account_id: &ProviderAccountId,
revision: CredentialRevision,
) -> Result<Option<Vec<String>>, CodexCredentialCatalogError> {
Ok(self.cached(instance)?.and_then(|snapshot| {
Ok(self.cached()?.and_then(|snapshot| {
snapshot
.account_models(account_id, revision)
.map(<[String]>::to_vec)
@@ -194,12 +179,11 @@ impl CodexCredentialCatalogService {
pub fn observed_model_support(
&self,
instance: &ProviderInstanceId,
account_id: &ProviderAccountId,
revision: CredentialRevision,
model: &str,
) -> Result<Option<bool>, CodexCredentialCatalogError> {
Ok(self.cached(instance)?.and_then(|snapshot| {
Ok(self.cached()?.and_then(|snapshot| {
snapshot
.account_models(account_id, revision)
.map(|models| models.iter().any(|candidate| candidate == model))
@@ -209,25 +193,19 @@ impl CodexCredentialCatalogService {
/// 只刷新指定账号的 realtime catalog，并以本地 revision 原子合并账号事实。
pub async fn synchronize_account(
&self,
instance: &ProviderInstance,
account_id: &ProviderAccountId,
) -> Result<Vec<String>, CodexCredentialCatalogError> {
let config =
CodexProviderInstanceConfig::from_snapshot(instance, self.origin_policy.as_ref())
.map_err(|_| CodexCredentialCatalogError::InvalidInstance)?;
let account = self
.repository
.store()
.get_account(account_id)
.await
.map_err(|_| CodexCredentialCatalogError::InvalidCredentialData)?
.filter(|account| {
account.provider().as_str() == "openai" && account.instance() == config.id()
})
.filter(|account| account.provider().as_str() == "openai")
.ok_or(CodexCredentialCatalogError::NoEligibleCredential)?;
let client = CodexBackendClient::new(
self.http.clone(),
config.base_url().as_str(),
OFFICIAL_CODEX_BASE_URL,
self.profile.clone(),
);
let cache_revision = self.cache_revision()?;
@@ -239,36 +217,28 @@ impl CodexCredentialCatalogService {
.collect::<Vec<_>>();
self.replace_account_cache(
cache_revision,
config.id(),
account_id,
account.revision(),
fetched.models,
entitlement.clone(),
)?;
self.record_applied_catalog_etags(config.id(), fetched.etag)?;
self.record_applied_catalog_etags(fetched.etag)?;
Ok(entitlement)
}
pub async fn synchronize_instance(
pub async fn synchronize(
&self,
instance: &ProviderInstance,
) -> Result<CodexCredentialCatalogSnapshot, CodexCredentialCatalogError> {
if let Some(cached) = self.cached(instance.id())? {
if let Some(cached) = self.cached()? {
return Ok(cached);
}
let cache_revision = self.cache_revision()?;
let fetched = self.fetch_instance_catalog(instance).await?;
self.commit_instance_catalog(cache_revision, fetched)
let fetched = self.fetch_catalog().await?;
self.commit_catalog(cache_revision, fetched)
}
async fn fetch_instance_catalog(
&self,
instance: &ProviderInstance,
) -> Result<FetchedInstanceCatalog, CodexCredentialCatalogError> {
let config =
CodexProviderInstanceConfig::from_snapshot(instance, self.origin_policy.as_ref())
.map_err(|_| CodexCredentialCatalogError::InvalidInstance)?;
let accounts = self.repository.list_for_instance(config.id()).await?;
async fn fetch_catalog(&self) -> Result<FetchedCatalog, CodexCredentialCatalogError> {
let accounts = self.repository.list_for_provider().await?;
let now = SystemTime::now();
let accounts = accounts
.into_iter()
@@ -279,7 +249,7 @@ impl CodexCredentialCatalogService {
}
let client = CodexBackendClient::new(
self.http.clone(),
config.base_url().as_str(),
OFFICIAL_CODEX_BASE_URL,
self.profile.clone(),
);
let mut union = BTreeMap::<String, CodexCatalogModel>::new();
@@ -314,7 +284,6 @@ impl CodexCredentialCatalogService {
}
let observed_at = SystemTime::now();
let snapshot = CodexCredentialCatalogSnapshot {
provider_instance_id: config.id().clone(),
observed_at,
models: union_order
.into_iter()
@@ -322,7 +291,7 @@ impl CodexCredentialCatalogService {
.collect(),
account_models,
};
Ok(FetchedInstanceCatalog { snapshot, etags })
Ok(FetchedCatalog { snapshot, etags })
}
async fn fetch_account_models(
@@ -363,7 +332,6 @@ impl CodexCredentialCatalogService {
fn replace_account_cache(
&self,
expected_cache_revision: u64,
instance: &ProviderInstanceId,
account_id: &ProviderAccountId,
revision: CredentialRevision,
models: Vec<CodexCatalogModel>,
@@ -379,7 +347,7 @@ impl CodexCredentialCatalogService {
let mut union = BTreeMap::<String, CodexCatalogModel>::new();
let mut union_order = Vec::new();
let mut account_models = BTreeMap::new();
if let Some(existing) = cache.snapshots.get(instance) {
if let Some(existing) = cache.snapshot.as_ref() {
for model in &existing.models {
let id = model.request_model().as_str().to_owned();
union_order.push(id.clone());
@@ -414,7 +382,6 @@ impl CodexCredentialCatalogService {
});
union_order.retain(|id| union.contains_key(id));
let snapshot = CodexCredentialCatalogSnapshot {
provider_instance_id: instance.clone(),
observed_at: SystemTime::now(),
models: union_order
.into_iter()
@@ -423,10 +390,10 @@ impl CodexCredentialCatalogService {
account_models,
};
let changed = cache
.snapshots
.get(instance)
.snapshot
.as_ref()
.is_none_or(|existing| !same_catalog(existing, &snapshot));
cache.snapshots.insert(instance.clone(), snapshot);
cache.snapshot = Some(snapshot);
cache.revision = cache.revision.saturating_add(1);
if changed {
cache.generation = cache.generation.saturating_add(1);
@@ -434,15 +401,12 @@ impl CodexCredentialCatalogService {
Ok(())
}
pub fn invalidate(
&self,
instance: &ProviderInstanceId,
) -> Result<(), CodexCredentialCatalogError> {
pub fn invalidate(&self) -> Result<(), CodexCredentialCatalogError> {
let mut cache = self
.cache
.write()
.map_err(|_| CodexCredentialCatalogError::Cache)?;
if cache.snapshots.remove(instance).is_some() {
if cache.snapshot.take().is_some() {
cache.revision = cache.revision.saturating_add(1);
cache.generation = cache.generation.saturating_add(1);
}
@@ -450,33 +414,20 @@ impl CodexCredentialCatalogService {
}
/// 记录普通 Responses 响应声明的目录版本；相同版本只触发一次。
pub fn observe_response_etag(
&self,
instance: &ProviderInstanceId,
etag: &str,
) -> Result<bool, CodexCredentialCatalogError> {
pub fn observe_response_etag(&self, etag: &str) -> Result<bool, CodexCredentialCatalogError> {
validate_response_etag(etag)?;
let changed = {
let mut state = self
.etags
.lock()
.map_err(|_| CodexCredentialCatalogError::Cache)?;
let already_observed = state
.applied
.get(instance)
.is_some_and(|value| value == etag)
|| state
.pending
.get(instance)
.is_some_and(|value| value == etag)
|| state
.inflight
.get(instance)
.is_some_and(|value| value == etag);
let already_observed = state.applied.as_deref() == Some(etag)
|| state.pending.as_deref() == Some(etag)
|| state.inflight.as_deref() == Some(etag);
if already_observed {
false
} else {
state.pending.insert(instance.clone(), etag.to_owned());
state.pending = Some(etag.to_owned());
true
}
};
@@ -486,39 +437,35 @@ impl CodexCredentialCatalogService {
Ok(changed)
}
/// 等待并合并一批需要强制刷新的 Provider instance。
pub async fn wait_for_etag_refresh(&self) -> Vec<ProviderInstanceId> {
/// 等待并认领一次需要强制刷新的 Provider 目录。
pub async fn wait_for_etag_refresh(&self) {
loop {
let pending = self.begin_pending_etag_refreshes();
if !pending.is_empty() {
return pending;
if self.begin_pending_etag_refresh() {
return;
}
self.etag_notification.notified().await;
}
}
/// 忽略当前 cache，按 ETag 变化强制生成一份完整新快照。
pub async fn refresh_instance(
pub async fn refresh(
&self,
instance: &ProviderInstance,
) -> Result<CodexCredentialCatalogSnapshot, CodexCredentialCatalogError> {
let result = self.refresh_instance_inner(instance).await;
self.finish_etag_refresh(instance.id(), result.is_ok())?;
let result = self.refresh_inner().await;
self.finish_etag_refresh(result.is_ok())?;
result
}
async fn refresh_instance_inner(
async fn refresh_inner(
&self,
instance: &ProviderInstance,
) -> Result<CodexCredentialCatalogSnapshot, CodexCredentialCatalogError> {
let cache_revision = self.cache_revision()?;
let fetched = self.fetch_instance_catalog(instance).await?;
self.commit_instance_catalog(cache_revision, fetched)
let fetched = self.fetch_catalog().await?;
self.commit_catalog(cache_revision, fetched)
}
fn record_applied_catalog_etags(
&self,
instance: &ProviderInstanceId,
etags: impl IntoIterator<Item = String>,
) -> Result<(), CodexCredentialCatalogError> {
let mut distinct = etags.into_iter().collect::<std::collections::BTreeSet<_>>();
@@ -532,39 +479,37 @@ impl CodexCredentialCatalogService {
self.etags
.lock()
.map_err(|_| CodexCredentialCatalogError::Cache)?
.applied
.insert(instance.clone(), etag);
.applied = Some(etag);
Ok(())
}
fn begin_pending_etag_refreshes(&self) -> Vec<ProviderInstanceId> {
fn begin_pending_etag_refresh(&self) -> bool {
let Ok(mut state) = self.etags.lock() else {
return Vec::new();
return false;
};
let pending = std::mem::take(&mut state.pending);
let instances = pending.keys().cloned().collect();
state.inflight.extend(pending);
instances
let Some(pending) = state.pending.take() else {
return false;
};
state.inflight = Some(pending);
true
}
fn finish_etag_refresh(
&self,
instance: &ProviderInstanceId,
succeeded: bool,
) -> Result<(), CodexCredentialCatalogError> {
fn finish_etag_refresh(&self, succeeded: bool) -> Result<(), CodexCredentialCatalogError> {
let should_retry = {
let mut state = self
.etags
.lock()
.map_err(|_| CodexCredentialCatalogError::Cache)?;
let Some(etag) = state.inflight.remove(instance) else {
let Some(etag) = state.inflight.take() else {
return Ok(());
};
if succeeded {
state.applied.insert(instance.clone(), etag);
state.applied = Some(etag);
false
} else {
state.pending.entry(instance.clone()).or_insert(etag);
if state.pending.is_none() {
state.pending = Some(etag);
}
true
}
};
@@ -581,12 +526,11 @@ impl CodexCredentialCatalogService {
.map_err(|_| CodexCredentialCatalogError::Cache)
}
fn commit_instance_catalog(
fn commit_catalog(
&self,
expected_cache_revision: u64,
fetched: FetchedInstanceCatalog,
fetched: FetchedCatalog,
) -> Result<CodexCredentialCatalogSnapshot, CodexCredentialCatalogError> {
let instance = fetched.snapshot.provider_instance_id.clone();
let snapshot = fetched.snapshot;
let mut cache = self
.cache
@@ -596,16 +540,16 @@ impl CodexCredentialCatalogService {
return Err(CodexCredentialCatalogError::ConcurrentUpdate);
}
let changed = cache
.snapshots
.get(&instance)
.snapshot
.as_ref()
.is_none_or(|existing| !same_catalog(existing, &snapshot));
cache.snapshots.insert(instance.clone(), snapshot.clone());
cache.snapshot = Some(snapshot.clone());
cache.revision = cache.revision.saturating_add(1);
if changed {
cache.generation = cache.generation.saturating_add(1);
}
drop(cache);
self.record_applied_catalog_etags(&instance, fetched.etags)?;
self.record_applied_catalog_etags(fetched.etags)?;
Ok(snapshot)
}
}
@@ -614,9 +558,7 @@ fn same_catalog(
left: &CodexCredentialCatalogSnapshot,
right: &CodexCredentialCatalogSnapshot,
) -> bool {
left.provider_instance_id == right.provider_instance_id
&& left.models == right.models
&& left.account_models == right.account_models
left.models == right.models && left.account_models == right.account_models
}
fn validate_response_etag(etag: &str) -> Result<(), CodexCredentialCatalogError> {
backend/crates/providers/openai/src/credential/oauth.rs
+7 -39
View File
@@ -131,7 +131,6 @@ pub struct CodexPendingAuthorization {
flow_id: String,
owner_ref: String,
started_request_ref: String,
provider_instance_id: String,
name: String,
expires_at: DateTime<Utc>,
state: SecretString,
@@ -145,7 +144,6 @@ pub struct StoredCodexPendingAuthorization {
pub flow_id: String,
pub owner_ref: String,
pub started_request_ref: String,
pub provider_instance_id: String,
pub name: String,
pub expires_at: DateTime<Utc>,
pub state: SecretString,
@@ -177,7 +175,6 @@ impl CodexPendingAuthorization {
flow_id: input.flow_id,
owner_ref: input.owner_ref,
started_request_ref: input.started_request_ref,
provider_instance_id: input.provider_instance_id,
name: input.name,
expires_at: input.expires_at,
state: input.state,
@@ -189,7 +186,6 @@ impl CodexPendingAuthorization {
if !valid_text(&pending.flow_id)
|| !valid_text(&pending.owner_ref)
|| !valid_text(&pending.started_request_ref)
|| !valid_text(&pending.provider_instance_id)
|| !valid_text(&pending.name)
|| pending.expires_at <= Utc::now()
|| !valid_secret(pending.state.expose_secret())
@@ -217,11 +213,6 @@ impl CodexPendingAuthorization {
&self.started_request_ref
}
#[must_use]
pub fn provider_instance_id(&self) -> &str {
&self.provider_instance_id
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
@@ -265,7 +256,6 @@ impl fmt::Debug for CodexPendingAuthorization {
.field("flow_id", &"<redacted>")
.field("owner_ref", &"<redacted>")
.field("started_request_ref", &"<redacted>")
.field("provider_instance_id", &self.provider_instance_id)
.field("name", &self.name)
.field("expires_at", &self.expires_at)
.field("state", &"<redacted>")
@@ -372,13 +362,9 @@ impl CodexOAuthAdmin for CodexOAuthAdminService {
&self,
command: StartCodexOAuthAuthorization,
) -> Result<CodexOAuthAuthorizationStarted, CodexOAuthAdminError> {
let (provider_instance_id, name, reauthorization) = match command.mutation.target() {
AuthorizationMutationTarget::Create {
provider_instance_id,
name,
} => (provider_instance_id.to_string(), name.clone(), None),
let (name, reauthorization) = match command.mutation.target() {
AuthorizationMutationTarget::Create { name } => (name.clone(), None),
AuthorizationMutationTarget::Reauthorize {
provider_instance_id,
account_id,
expected_credential_revision,
} => {
@@ -391,12 +377,10 @@ impl CodexOAuthAdmin for CodexOAuthAdminService {
.map_err(map_store_error)?;
if current.account.provider().as_str() != "openai"
|| current.account.id() != account_id
|| current.account.instance() != provider_instance_id
{
return Err(CodexOAuthAdminError::NotFound);
}
(
current.account.instance().to_string(),
current.account.name().to_owned(),
Some(CodexOAuthReauthorizationTarget {
account_id: account_id.clone(),
@@ -405,13 +389,8 @@ impl CodexOAuthAdmin for CodexOAuthAdminService {
)
}
};
self.start_pending(
command.mutation,
provider_instance_id,
name,
reauthorization,
)
.await
self.start_pending(command.mutation, name, reauthorization)
.await
}
async fn complete_authorization(
@@ -487,7 +466,6 @@ impl CodexOAuthAdmin for CodexOAuthAdminService {
self.credentials
.prepare_import(ImportCodexOAuthCredential {
account_id,
provider_instance_id: pending.provider_instance_id,
name: pending.name,
secret,
verified_account: profile,
@@ -508,11 +486,10 @@ impl CodexOAuthAdminService {
async fn start_pending(
&self,
mutation: PendingAuthorizationMutation,
provider_instance_id: String,
name: String,
reauthorization: Option<CodexOAuthReauthorizationTarget>,
) -> Result<CodexOAuthAuthorizationStarted, CodexOAuthAdminError> {
if !valid_text(&provider_instance_id) || !valid_text(&name) {
if !valid_text(&name) {
return Err(CodexOAuthAdminError::InvalidInput);
}
let expires_at = Utc::now()
@@ -524,7 +501,6 @@ impl CodexOAuthAdminService {
flow_id: random_secret()?,
owner_ref,
started_request_ref,
provider_instance_id,
name,
expires_at,
state: SecretString::from(random_secret()?),
@@ -699,23 +675,15 @@ fn pending_mutation_matches(pending: &CodexPendingAuthorization) -> bool {
return false;
}
match (pending.mutation.target(), pending.reauthorization.as_ref()) {
(
AuthorizationMutationTarget::Create {
provider_instance_id,
name,
},
None,
) => provider_instance_id.as_str() == pending.provider_instance_id && name == &pending.name,
(AuthorizationMutationTarget::Create { name }, None) => name == &pending.name,
(
AuthorizationMutationTarget::Reauthorize {
provider_instance_id,
account_id,
expected_credential_revision,
},
Some(target),
) => {
provider_instance_id.as_str() == pending.provider_instance_id
&& account_id == target.account_id()
account_id == target.account_id()
&& expected_credential_revision.get() == target.credential_revision().get()
}
_ => false,
backend/crates/providers/openai/src/credential/quota.rs
+6 -24
View File
@@ -9,7 +9,6 @@ use gateway_core::engine::credential::{
AccountAvailability, AccountQuotaSignals, CredentialRevision, OpaqueProviderData,
ProviderAccount, ProviderAccountId, ProviderAccountStore, QuotaObservation, QuotaWriteOutcome,
};
use gateway_core::routing::ProviderInstance;
use gateway_protocol::openai::events::{
ParsedRateLimits, RateLimitDetails, RateLimitWindow, parse_rate_limit_headers,
};
@@ -20,7 +19,7 @@ use thiserror::Error;
use tokio::sync::Mutex;
use uuid::Uuid;
use crate::provider::{CodexOriginPolicy, CodexProviderInstanceConfig};
use crate::provider::OFFICIAL_CODEX_BASE_URL;
use crate::transport::profile::CodexWireProfileState;
use crate::transport::{CodexBackendClient, CodexClientError, CodexRequestContext};
@@ -173,8 +172,6 @@ impl CodexQuotaSyncSummary {
#[derive(Debug, Error)]
pub enum CodexCredentialQuotaError {
#[error("Codex quota instance is invalid")]
InvalidInstance,
#[error("Codex quota response is invalid")]
InvalidCredentialData,
#[error(transparent)]
@@ -200,7 +197,6 @@ pub struct CodexCredentialQuotaService {
store: Arc<dyn ProviderAccountStore>,
profile: CodexWireProfileState,
http: Client,
origin_policy: Arc<dyn CodexOriginPolicy>,
scheduling: CodexQuotaSchedulingProjection,
}
@@ -377,14 +373,12 @@ impl CodexCredentialQuotaService {
repository: CodexCredentialRepository,
profile: CodexWireProfileState,
http: Client,
origin_policy: Arc<dyn CodexOriginPolicy>,
) -> Self {
Self {
store: Arc::clone(repository.store()),
repository,
profile,
http,
origin_policy,
scheduling: CodexQuotaSchedulingProjection::default(),
}
}
@@ -433,19 +427,13 @@ impl CodexCredentialQuotaService {
self.scheduling.signals(account)
}
pub async fn synchronize_instance(
&self,
instance: &gateway_core::routing::ProviderInstance,
) -> Result<CodexQuotaSyncSummary, CodexCredentialQuotaError> {
let config =
CodexProviderInstanceConfig::from_snapshot(instance, self.origin_policy.as_ref())
.map_err(|_| CodexCredentialQuotaError::InvalidInstance)?;
pub async fn synchronize(&self) -> Result<CodexQuotaSyncSummary, CodexCredentialQuotaError> {
let client = CodexBackendClient::new(
self.http.clone(),
config.base_url().as_str(),
OFFICIAL_CODEX_BASE_URL,
self.profile.clone(),
);
let accounts = self.repository.list_for_instance(config.id()).await?;
let accounts = self.repository.list_for_provider().await?;
let mut summary = CodexQuotaSyncSummary::default();
let now = SystemTime::now();
for account in accounts
@@ -677,24 +665,18 @@ impl CodexCredentialQuotaService {
/// 只刷新指定账号，revision-fenced 写入动态 Provider JSON 后返回解析快照。
pub async fn refresh_account(
&self,
instance: &ProviderInstance,
account_id: &ProviderAccountId,
) -> Result<CodexAccountQuotaSnapshot, CodexCredentialQuotaError> {
let config =
CodexProviderInstanceConfig::from_snapshot(instance, self.origin_policy.as_ref())
.map_err(|_| CodexCredentialQuotaError::InvalidInstance)?;
let account = self
.store
.get_account(account_id)
.await?
.filter(|account| {
account.provider().as_str() == "openai" && account.instance() == config.id()
})
.filter(|account| account.provider().as_str() == "openai")
.ok_or(CodexCredentialQuotaError::NotFound)?;
let runtime = self.repository.load_runtime_credential(&account).await?;
let client = CodexBackendClient::new(
self.http.clone(),
config.base_url().as_str(),
OFFICIAL_CODEX_BASE_URL,
self.profile.clone(),
);
let request_id = format!("quota_{}", Uuid::now_v7().simple());
backend/crates/providers/openai/src/credential/repository.rs
+5 -4
View File
@@ -9,7 +9,7 @@ use gateway_core::engine::credential::{
CredentialRevision, LoadedCredential, ProviderAccount, ProviderAccountId, ProviderAccountStore,
ProviderAccountUpdate,
};
use gateway_core::routing::ProviderInstanceId;
use gateway_core::routing::ProviderKind;
use secrecy::ExposeSecret;
use thiserror::Error;
@@ -157,12 +157,13 @@ impl CodexCredentialRepository {
cas_revision(self.store.compare_and_swap_credential(update).await?)
}
pub async fn list_for_instance(
pub async fn list_for_provider(
&self,
instance: &ProviderInstanceId,
) -> Result<Vec<ProviderAccount>, CredentialRepositoryError> {
let provider = ProviderKind::new(PROVIDER_NAME)
.map_err(|_| CredentialRepositoryError::InvalidCredentialData)?;
self.store
.list_for_instance(instance)
.list_for_provider(&provider)
.await
.map_err(Into::into)
}
backend/crates/providers/openai/src/credential/selector.rs
+9 -16
View File
@@ -14,7 +14,7 @@ use gateway_core::provider_ports::{
ProviderLeaseAcquisition, ProviderLeaseGuard, ProviderLeasePort, ProviderLeaseRequest,
ProviderSchedulingLeaseRequest, ProviderStoreError,
};
use gateway_core::routing::ProviderInstanceId;
use gateway_core::routing::ProviderKind;
use secrecy::ExposeSecret;
use thiserror::Error;
use url::Url;
@@ -25,7 +25,6 @@ use super::quota::CodexCredentialQuotaService;
use super::repository::{CodexCredentialRepository, CredentialRepositoryError};
use super::types::{CodexCookie, CodexCookieCaptureOutcome, CodexOAuthSecret, RuntimeCodexCookie};
const PROVIDER_NAME: &str = "openai";
const DEFAULT_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60);
const CLOUDFLARE_PATH_BLOCK_COOLDOWN: Duration = Duration::from_secs(30);
const CLOUDFLARE_RECOVERY_STALE_AFTER: Duration = Duration::from_secs(60 * 60);
@@ -76,13 +75,13 @@ enum CookieRecovery {
}
pub struct SelectCodexCredential<'a> {
pub provider_instance_id: &'a ProviderInstanceId,
pub upstream_model: &'a str,
pub request_url: &'a Url,
pub attempt: &'a AttemptContext,
}
pub struct CodexCredentialSelector {
provider_kind: ProviderKind,
repository: CodexCredentialRepository,
leases: Arc<dyn ProviderLeasePort>,
catalog: Arc<CodexCredentialCatalogService>,
@@ -94,6 +93,7 @@ pub struct CodexCredentialSelector {
impl CodexCredentialSelector {
#[must_use]
pub fn new(
provider_kind: ProviderKind,
repository: CodexCredentialRepository,
leases: Arc<dyn ProviderLeasePort>,
catalog: Arc<CodexCredentialCatalogService>,
@@ -101,6 +101,7 @@ impl CodexCredentialSelector {
cookie_policy: CodexCookiePolicy,
) -> Self {
Self {
provider_kind,
repository,
leases,
catalog,
@@ -114,18 +115,13 @@ impl CodexCredentialSelector {
&self,
request: &SelectCodexCredential<'_>,
) -> Result<CodexCredentialLease, CredentialSelectionError> {
let accounts = self
.repository
.list_for_instance(request.provider_instance_id)
.await?;
let accounts = self.repository.list_for_provider().await?;
let accounts = accounts
.into_iter()
.filter(|account| {
account.provider().as_str() == PROVIDER_NAME
&& account.instance() == request.provider_instance_id
account.provider() == &self.provider_kind
&& !matches!(
self.catalog.observed_model_support(
request.provider_instance_id,
account.id(),
account.revision(),
request.upstream_model,
@@ -141,7 +137,7 @@ impl CodexCredentialSelector {
.collect::<Vec<_>>();
let scheduling = self
.leases
.load_state(request.provider_instance_id, &account_ids)
.load_state(&self.provider_kind, &account_ids)
.await?;
let round_robin_cursor = scheduling.round_robin_cursor();
let candidates = accounts
@@ -171,10 +167,7 @@ impl CodexCredentialSelector {
ContinuationAttempt::ReplayOwner => request
.attempt
.account_state_owner()
.filter(|owner| {
owner.provider().as_str() == PROVIDER_NAME
&& owner.instance() == request.provider_instance_id
})
.filter(|owner| owner.provider() == &self.provider_kind)
.map(|owner| owner.account().clone()),
ContinuationAttempt::None | ContinuationAttempt::ReplayAny => None,
};
@@ -219,7 +212,7 @@ impl CodexCredentialSelector {
.leases
.try_acquire(ProviderLeaseRequest::Scheduling(
ProviderSchedulingLeaseRequest::new(
request.provider_instance_id.clone(),
self.provider_kind.clone(),
account.id().clone(),
account.revision(),
policy.max_concurrent_per_account(),
backend/crates/providers/openai/src/lib.rs
+20 -27
View File
@@ -32,8 +32,8 @@ use crate::transport::{CodexWebSocketPool, build_reqwest_client};
pub use config::{CodexWireProfileConfig, OpenAiConfig, OpenAiConfigError};
pub use provider::{
CodexOriginPolicy, CodexProvider, CodexProviderConfigError, CodexProviderInstanceConfig,
CodexProviderTransport, OFFICIAL_CODEX_BASE_PATH, OfficialCodexOriginPolicy,
CodexProvider, CodexProviderConfigError, CodexProviderTransport, OFFICIAL_CODEX_BASE_PATH,
OFFICIAL_CODEX_BASE_URL,
};
pub mod credential;
@@ -64,10 +64,8 @@ pub async fn initialize(
let accounts: Arc<dyn ProviderAccountStore> = ports.accounts();
let leases = ports.leases();
let runtime_policy = ports.runtime_policy();
let instances = ports.instances();
let profile = config.wire_profile_state();
let http = build_reqwest_client().map_err(|_| OpenAiInitializeError::Transport)?;
let origin_policy: Arc<dyn CodexOriginPolicy> = Arc::new(OfficialCodexOriginPolicy);
let desktop_release = Arc::new(CodexDesktopReleaseService::new(
profile.clone(),
Arc::new(
@@ -88,15 +86,14 @@ pub async fn initialize(
repository.clone(),
profile.clone(),
http.clone(),
Arc::clone(&origin_policy),
));
let quota = Arc::new(CodexCredentialQuotaService::new(
repository.clone(),
profile.clone(),
http.clone(),
Arc::clone(&origin_policy),
));
let selector = Arc::new(CodexCredentialSelector::new(
provider_kind.clone(),
repository.clone(),
Arc::clone(&leases),
Arc::clone(&catalog),
@@ -104,15 +101,17 @@ pub async fn initialize(
CodexCookiePolicy::official().map_err(|_| OpenAiInitializeError::CookiePolicy)?,
));
let websocket_pool = Arc::new(CodexWebSocketPool::default());
let core_provider: Arc<dyn Provider> = Arc::new(CodexProvider::new(
selector,
Arc::clone(&catalog),
Arc::clone(&quota),
http,
profile.clone(),
Arc::clone(&websocket_pool),
origin_policy,
));
let core_provider: Arc<dyn Provider> = Arc::new(
CodexProvider::new(
selector,
Arc::clone(&catalog),
Arc::clone(&quota),
http,
profile.clone(),
Arc::clone(&websocket_pool),
)
.map_err(OpenAiInitializeError::Provider)?,
);
let token_client = Arc::new(
credential::token_client::official_openai_token_client()
@@ -156,10 +155,9 @@ pub async fn initialize(
CodexCredentialAdmin,
));
let admin_provider: Arc<dyn ProviderAdmin> = Arc::new(OpenAiAdminProvider::new(
provider_kind.clone(),
provider_kind,
profile,
accounts,
Arc::clone(&instances),
OpenAiAdminServices {
credentials: credential_admin,
verifier: identity,
@@ -171,16 +169,9 @@ pub async fn initialize(
websocket_pool,
desktop_release_status,
));
let worker_contributions = provider::worker_contributions(
refresh,
quota,
catalog,
cli_release,
desktop_release,
instances,
provider_kind,
)
.map_err(|_| OpenAiInitializeError::Worker)?;
let worker_contributions =
provider::worker_contributions(refresh, quota, catalog, cli_release, desktop_release)
.map_err(|_| OpenAiInitializeError::Worker)?;
Ok(ProviderBundle {
core_provider,
@@ -217,6 +208,8 @@ pub enum OpenAiInitializeError {
InvalidProviderKind,
#[error("OpenAI transport could not initialize")]
Transport,
#[error(transparent)]
Provider(CodexProviderConfigError),
#[error("OpenAI cookie policy could not initialize")]
CookiePolicy,
#[error("OpenAI token client could not initialize")]
backend/crates/providers/openai/src/provider.rs
+50 -298
View File
@@ -1,8 +1,8 @@
//! Codex 的 `gateway-core` Provider adapter。
use std::collections::{BTreeSet, HashMap};
use std::collections::BTreeSet;
use std::fmt;
use std::sync::{Arc, Mutex};
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use async_trait::async_trait;
@@ -22,11 +22,7 @@ use gateway_core::event::{
ProviderResponseTimings, UpstreamHttpVersion, WebSocketPoolKind,
};
use gateway_core::operation::{GenerateRequest, Operation, OperationKind, ProviderSessionState};
use gateway_core::provider_ports::ProviderInstanceCatalogPort;
use gateway_core::routing::{
InstanceHealth, ModelCapabilities, ProviderInstance, ProviderInstanceId, ProviderKind,
UpstreamModelId,
};
use gateway_core::routing::{ModelCapabilities, ProviderKind, UpstreamModelId};
use gateway_core::task::{
DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerContribution, WorkerCycleContext,
WorkerDefinitionError, WorkerId, WorkerKind, WorkerLeaseRequest, WorkerRegistration,
@@ -36,7 +32,7 @@ use reqwest::Client;
use secrecy::{ExposeSecret, SecretString};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use url::{Host, Url};
use url::Url;
use crate::credential::{
CodexAccountFailure, CodexCredentialCatalogService, CodexCredentialLease,
@@ -70,8 +66,8 @@ const PROVIDER_NAME: &str = "openai";
const HTTP_SSE_TRANSPORT: &str = "http_sse";
const WEBSOCKET_TRANSPORT: &str = "websocket";
const MAX_COOKIE_HEADER_BYTES: usize = 16 * 1024;
const OFFICIAL_CODEX_HOST: &str = "chatgpt.com";
pub const OFFICIAL_CODEX_BASE_PATH: &str = "/backend-api";
pub const OFFICIAL_CODEX_BASE_URL: &str = "https://chatgpt.com/backend-api";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CodexProviderTransport {
@@ -89,92 +85,18 @@ impl CodexProviderTransport {
}
}
/// Provider base URL 的显式信任边界。
///
/// 生产组装注入 [`OfficialCodexOriginPolicy`]；Provider 只依赖这个端口，
/// 不包含部署环境或地址类型分支。
pub trait CodexOriginPolicy: Send + Sync {
fn allows(&self, url: &Url) -> bool;
}
#[derive(Debug, Default)]
pub struct OfficialCodexOriginPolicy;
impl CodexOriginPolicy for OfficialCodexOriginPolicy {
fn allows(&self, url: &Url) -> bool {
url.scheme() == "https"
&& matches!(url.host(), Some(Host::Domain(host)) if host == OFFICIAL_CODEX_HOST)
&& url.port().is_none()
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CodexProviderInstanceConfig {
id: ProviderInstanceId,
base_url: Url,
transport: CodexProviderTransport,
}
impl CodexProviderInstanceConfig {
pub fn from_snapshot(
instance: &ProviderInstance,
origin_policy: &dyn CodexOriginPolicy,
) -> Result<Self, CodexProviderConfigError> {
if instance.provider().as_str() != PROVIDER_NAME {
return Err(CodexProviderConfigError::ProviderMismatch);
}
let mut base_url = Url::parse(instance.base_url())
.map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?;
normalize_and_validate_base_url(&mut base_url, origin_policy)?;
Ok(Self {
id: instance.id().clone(),
base_url,
transport: CodexProviderTransport::PreferWebSocket,
})
}
#[must_use]
pub const fn id(&self) -> &ProviderInstanceId {
&self.id
}
#[must_use]
pub const fn base_url(&self) -> &Url {
&self.base_url
}
#[must_use]
pub const fn transport(&self) -> CodexProviderTransport {
self.transport
}
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum CodexProviderConfigError {
#[error("provider instance does not belong to Codex")]
ProviderMismatch,
#[error("Codex provider base URL is invalid")]
#[error("official Codex provider URL is invalid")]
InvalidBaseUrl,
#[error("Codex provider base URL is not allowed")]
UnsafeBaseUrl,
}
#[derive(Clone)]
struct CompiledInstance {
config: CodexProviderInstanceConfig,
responses_url: Url,
client: CodexBackendClient,
}
pub struct CodexProvider {
selector: Arc<CodexCredentialSelector>,
catalog: Arc<CodexCredentialCatalogService>,
quota: Arc<CodexCredentialQuotaService>,
http: Client,
profile: CodexWireProfileState,
origin_policy: Arc<dyn CodexOriginPolicy>,
websocket_pool: Arc<CodexWebSocketPool>,
compiled_instances: Mutex<HashMap<String, CompiledInstance>>,
client: CodexBackendClient,
responses_url: Url,
}
impl CodexProvider {
@@ -185,59 +107,19 @@ impl CodexProvider {
http: Client,
profile: CodexWireProfileState,
websocket_pool: Arc<CodexWebSocketPool>,
origin_policy: Arc<dyn CodexOriginPolicy>,
) -> Self {
Self {
) -> Result<Self, CodexProviderConfigError> {
let responses_url =
Url::parse(&endpoint_url(OFFICIAL_CODEX_BASE_URL, CODEX_RESPONSES_PATH))
.map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?;
let client = CodexBackendClient::new(http, OFFICIAL_CODEX_BASE_URL, profile)
.with_websocket_pool(websocket_pool);
Ok(Self {
selector,
catalog,
quota,
http,
profile,
origin_policy,
websocket_pool,
compiled_instances: Mutex::new(HashMap::new()),
}
}
pub fn validate_instance(
instance: &ProviderInstance,
) -> Result<CodexProviderInstanceConfig, CodexProviderConfigError> {
CodexProviderInstanceConfig::from_snapshot(instance, &OfficialCodexOriginPolicy)
}
fn compile_instance(
&self,
snapshot: &ProviderInstance,
) -> Result<CompiledInstance, ProviderError> {
let config =
CodexProviderInstanceConfig::from_snapshot(snapshot, self.origin_policy.as_ref())
.map_err(map_instance_config_error)?;
let responses_url = Url::parse(&endpoint_url(
config.base_url.as_str(),
CODEX_RESPONSES_PATH,
))
.map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?;
let mut instances = self
.compiled_instances
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(compiled) = instances.get(config.id().as_str())
&& compiled.config == config
{
return Ok(compiled.clone());
}
let compiled = CompiledInstance {
client: CodexBackendClient::new(
self.http.clone(),
config.base_url.as_str(),
self.profile.clone(),
)
.with_websocket_pool(Arc::clone(&self.websocket_pool)),
config,
client,
responses_url,
};
instances.insert(compiled.config.id().as_str().to_owned(), compiled.clone());
Ok(compiled)
})
}
}
@@ -247,7 +129,6 @@ impl fmt::Debug for CodexProvider {
.debug_struct("CodexProvider")
.field("selector", &"<account-selector>")
.field("catalog", &"<ttl-catalog>")
.field("origin_policy", &"[ORIGIN_POLICY]")
.finish()
}
}
@@ -291,15 +172,10 @@ impl Provider for CodexProvider {
async fn query_model_capabilities(
&self,
instance: &ProviderInstance,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
let snapshot = self
.catalog
.synchronize_instance(instance)
.await
.map_err(|_| {
provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
})?;
let snapshot = self.catalog.synchronize().await.map_err(|_| {
provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
})?;
Ok(snapshot
.models()
.iter()
@@ -319,14 +195,6 @@ impl Provider for CodexProvider {
));
}
let candidate = request.candidate();
let instance_snapshot = candidate.instance_snapshot();
if instance_snapshot.id() != candidate.instance() {
return Err(provider_error(
ProviderErrorKind::Protocol,
UpstreamSendState::NotSent,
));
}
let instance = self.compile_instance(instance_snapshot)?;
if context.cancellation().is_cancelled() {
return Err(provider_error(
ProviderErrorKind::Cancelled,
@@ -351,15 +219,14 @@ impl Provider for CodexProvider {
encode_generate_request(generate, candidate.upstream_model().as_str())
.map_err(map_request_error)?;
let request_input = upstream_request.input().to_vec();
let transport = selected_transport(&request, &instance.config)?;
let transport = selected_transport(&request)?;
apply_transport(&mut upstream_request, transport);
let lease = self
.selector
.select(&SelectCodexCredential {
provider_instance_id: candidate.instance(),
upstream_model: candidate.upstream_model().as_str(),
request_url: &instance.responses_url,
request_url: &self.responses_url,
attempt: &context,
})
.await
@@ -367,11 +234,12 @@ impl Provider for CodexProvider {
let lease = Arc::new(lease);
let provider_kind = ProviderKind::new(PROVIDER_NAME)
.map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?;
let cross_account = context.account_state_owner().is_some_and(|owner| {
!owner.matches(&provider_kind, candidate.instance(), lease.account_id())
}) || previous_session
.as_ref()
.is_some_and(|state| state.account_id != lease.account_id().as_str());
let cross_account = context
.account_state_owner()
.is_some_and(|owner| !owner.matches(&provider_kind, lease.account_id()))
|| previous_session
.as_ref()
.is_some_and(|state| state.account_id != lease.account_id().as_str());
let replay_previous_response = matches!(
context.continuation_attempt(),
ContinuationAttempt::ReplayOwner | ContinuationAttempt::ReplayAny
@@ -448,7 +316,6 @@ impl Provider for CodexProvider {
}
let metadata = ProviderCallMetadata::new(
provider_kind,
candidate.instance().clone(),
candidate.upstream_model().clone(),
ProviderResource::Account {
id: lease.account_id().clone(),
@@ -469,8 +336,8 @@ impl Provider for CodexProvider {
continuation_scope: None,
});
let events = cold_response_stream(ColdResponse {
client: instance.client,
response_origin: instance.responses_url,
client: self.client.clone(),
response_origin: self.responses_url.clone(),
request: upstream_request,
upstream_model: candidate.upstream_model().clone(),
transport_policy: transport,
@@ -478,7 +345,6 @@ impl Provider for CodexProvider {
selector: Arc::clone(&self.selector),
quota: Arc::clone(&self.quota),
catalog: Arc::clone(&self.catalog),
provider_instance_id: candidate.instance().clone(),
lease: Arc::clone(&lease),
session_capture,
});
@@ -496,7 +362,6 @@ struct ColdResponse {
selector: Arc<CodexCredentialSelector>,
quota: Arc<CodexCredentialQuotaService>,
catalog: Arc<CodexCredentialCatalogService>,
provider_instance_id: ProviderInstanceId,
lease: Arc<CodexCredentialLease>,
session_capture: Option<OpenAiSessionCapture>,
}
@@ -688,7 +553,6 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
selector,
quota,
catalog,
provider_instance_id,
lease,
mut session_capture,
} = response;
@@ -764,10 +628,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
)?);
synchronize_passive_quota(&quota, &lease, &response.rate_limit_headers).await;
if let Some(etag) = response.response_metadata.models_etag.as_deref()
&& let Err(error) = catalog.observe_response_etag(&provider_instance_id, etag)
&& let Err(error) = catalog.observe_response_etag(etag)
{
tracing::warn!(
provider_instance_id = %provider_instance_id,
error = %error,
"OpenAI model ETag observation was rejected"
);
@@ -1045,11 +908,8 @@ fn compile_model_capabilities(model: &CodexCatalogModel) -> ProviderModelCapabil
ProviderModelCapabilities::new(model.request_model().clone(), capabilities)
}
fn selected_transport(
request: &ProviderRequest,
instance: &CodexProviderInstanceConfig,
) -> Result<CodexProviderTransport, ProviderError> {
let mut transport = instance.transport();
fn selected_transport(request: &ProviderRequest) -> Result<CodexProviderTransport, ProviderError> {
let mut transport = CodexProviderTransport::PreferWebSocket;
if let Some(value) = request
.operation()
.provider_options(PROVIDER_NAME)
@@ -1193,16 +1053,6 @@ fn map_request_error(error: CodexRequestEncodeError) -> ProviderError {
provider_error(kind, UpstreamSendState::NotSent)
}
fn map_instance_config_error(error: CodexProviderConfigError) -> ProviderError {
let kind = match error {
CodexProviderConfigError::ProviderMismatch => ProviderErrorKind::InvalidRequest,
CodexProviderConfigError::InvalidBaseUrl | CodexProviderConfigError::UnsafeBaseUrl => {
ProviderErrorKind::Protocol
}
};
provider_error(kind, UpstreamSendState::NotSent)
}
fn map_selection_error(error: CredentialSelectionError) -> ProviderError {
match error {
CredentialSelectionError::CapacityUnavailable { retry_after } => {
@@ -1584,30 +1434,6 @@ fn remaining(deadline: SystemTime) -> Option<Duration> {
.filter(|remaining| !remaining.is_zero())
}
fn normalize_and_validate_base_url(
url: &mut Url,
origin_policy: &dyn CodexOriginPolicy,
) -> Result<(), CodexProviderConfigError> {
if url.cannot_be_a_base()
|| url.host().is_none()
|| !url.username().is_empty()
|| url.password().is_some()
|| url.query().is_some()
|| url.fragment().is_some()
{
return Err(CodexProviderConfigError::UnsafeBaseUrl);
}
if !origin_policy.allows(url) {
return Err(CodexProviderConfigError::UnsafeBaseUrl);
}
let normalized_path = url.path().trim_end_matches('/');
if normalized_path != OFFICIAL_CODEX_BASE_PATH {
return Err(CodexProviderConfigError::UnsafeBaseUrl);
}
url.set_path(OFFICIAL_CODEX_BASE_PATH);
Ok(())
}
const WORKER_INITIAL_BACKOFF: Duration = Duration::from_secs(1);
const WORKER_MAXIMUM_BACKOFF: Duration = Duration::from_secs(60);
const WORKER_LEASE_TTL: Duration = Duration::from_secs(15 * 60);
@@ -1624,8 +1450,6 @@ pub(crate) fn worker_contributions(
catalog: Arc<CodexCredentialCatalogService>,
cli_release: Arc<CodexCliReleaseService>,
desktop_release: Arc<CodexDesktopReleaseService>,
instances: Arc<dyn ProviderInstanceCatalogPort>,
provider_kind: ProviderKind,
) -> Result<Vec<WorkerContribution>, WorkerDefinitionError> {
let refresh_id = WorkerId::try_new(WorkerKind::OAuthRefresh, PROVIDER_NAME)?;
let quota_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, PROVIDER_NAME)?;
@@ -1643,11 +1467,7 @@ pub(crate) fn worker_contributions(
WorkerContribution::Registration(scheduled_registration(
quota_id,
QUOTA_REFRESH_INTERVAL,
Box::new(OpenAiQuotaTask {
instances: Arc::clone(&instances),
quota,
provider_kind: provider_kind.clone(),
}),
Box::new(OpenAiQuotaTask { quota }),
)?),
WorkerContribution::Registration(WorkerRegistration::try_new(
etag_id,
@@ -1656,11 +1476,7 @@ pub(crate) fn worker_contributions(
WORKER_INITIAL_BACKOFF,
WORKER_MAXIMUM_BACKOFF,
)?,
task: Box::new(OpenAiCatalogEtagTask {
instances,
catalog,
provider_kind,
}),
task: Box::new(OpenAiCatalogEtagTask { catalog }),
},
)?),
WorkerContribution::Registration(scheduled_registration(
@@ -1741,15 +1557,11 @@ impl ScheduledTask for OpenAiOAuthRefreshTask {
}
struct OpenAiQuotaTask {
instances: Arc<dyn ProviderInstanceCatalogPort>,
quota: Arc<CodexCredentialQuotaService>,
provider_kind: ProviderKind,
}
struct OpenAiCatalogEtagTask {
instances: Arc<dyn ProviderInstanceCatalogPort>,
catalog: Arc<CodexCredentialCatalogService>,
provider_kind: ProviderKind,
}
struct OpenAiDesktopReleaseTask {
@@ -1799,48 +1611,16 @@ impl ScheduledTask for OpenAiDesktopReleaseTask {
impl ScheduledTask for OpenAiQuotaTask {
fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
Box::pin(async move {
let instances = self
.instances
.list_instances(&self.provider_kind, false)
.await
.map_err(|_| WorkerTaskError::safe("OpenAI Provider instances unavailable"))?;
let mut failures = 0_u64;
for config in instances {
if context.cancellation().is_cancelled() {
return Ok(());
}
if !config.enabled() || config.provider_kind() != &self.provider_kind {
failures = failures.saturating_add(1);
continue;
}
let instance = ProviderInstance::new(
config.id().clone(),
config.provider_kind().clone(),
config.base_url().to_owned(),
config.enabled(),
InstanceHealth::Healthy,
);
match self.quota.synchronize_instance(&instance).await {
Ok(summary) if summary.has_operational_failures() => {
tracing::warn!(
provider_instance_id = %instance.id(),
"OpenAI quota cycle contained operational failures"
);
}
Ok(_) => {}
Err(_) => {
failures = failures.saturating_add(1);
tracing::warn!(
provider_instance_id = %instance.id(),
"OpenAI quota synchronization failed"
);
}
}
if context.cancellation().is_cancelled() {
return Ok(());
}
if failures == 0 {
Ok(())
} else {
Err(WorkerTaskError::safe("OpenAI quota synchronization failed"))
match self.quota.synchronize().await {
Ok(summary) if summary.has_operational_failures() => {
tracing::warn!("OpenAI quota cycle contained operational failures");
Ok(())
}
Ok(_) => Ok(()),
Err(_) => Err(WorkerTaskError::safe("OpenAI quota synchronization failed")),
}
})
}
@@ -1853,43 +1633,15 @@ impl DaemonTask for OpenAiCatalogEtagTask {
) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
Box::pin(async move {
loop {
let changed = tokio::select! {
tokio::select! {
() = cancellation.cancelled() => return Ok(()),
changed = self.catalog.wait_for_etag_refresh() => changed,
() = self.catalog.wait_for_etag_refresh() => {},
};
let instances = self
.instances
.list_instances(&self.provider_kind, false)
.await
.map_err(|_| {
WorkerTaskError::safe(
"OpenAI Provider instances unavailable for model ETag refresh",
)
})?;
for config in instances {
if cancellation.is_cancelled() {
return Ok(());
}
if !changed.contains(config.id())
|| !config.enabled()
|| config.provider_kind() != &self.provider_kind
{
continue;
}
let instance = ProviderInstance::new(
config.id().clone(),
config.provider_kind().clone(),
config.base_url().to_owned(),
true,
InstanceHealth::Healthy,
if let Err(error) = self.catalog.refresh().await {
tracing::warn!(
error = %error,
"OpenAI model catalog ETag refresh failed"
);
if let Err(error) = self.catalog.refresh_instance(&instance).await {
tracing::warn!(
provider_instance_id = %instance.id(),
error = %error,
"OpenAI model catalog ETag refresh failed"
);
}
}
}
})
backend/crates/providers/openai/tests/admin.rs
+12 -52
View File
@@ -37,14 +37,13 @@ use gateway_core::operation::{
use gateway_core::policy::ClientApiKeyId;
use gateway_core::provider_ports::{
NewOAuthPendingFlow, OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingTakeOutcome,
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCatalogPorts, ProviderCooldown,
ProviderCooldownPort, ProviderCredentialState, ProviderCredentialStatePort,
ProviderInstanceCatalogPort, ProviderInstanceConfig, ProviderRefreshPolicy,
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCooldown, ProviderCooldownPort,
ProviderCredentialState, ProviderCredentialStatePort, ProviderRefreshPolicy,
ProviderRuntimePolicyPort, ProviderStoreError, ProviderStorePorts,
};
use gateway_core::routing::{
ConfigRevision, InstanceHealth, ModelCapabilities, ProviderInstance, ProviderInstanceId,
ProviderKind, ProviderModel, PublicModelId, RoutingContext, RuntimeSnapshot, UpstreamModelId,
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
RuntimeSnapshot, UpstreamModelId,
};
use gateway_core::task::{WorkerContribution, WorkerKind, WorkerRunnable};
use provider_openai::config::{CodexWireProfileConfig, OpenAiConfig};
@@ -53,7 +52,7 @@ use provider_openai::transport::profile::APPCAST_POLL_INTERVAL;
use serde_json::{Map, Value, json};
use std::collections::BTreeSet;
use crate::support::{MemoryAccountStore, TestLeaseCoordinator, instance_id, profile, secret};
use crate::support::{MemoryAccountStore, TestLeaseCoordinator, profile, secret};
#[tokio::test]
#[ignore = "requires CODEX_REAL_ACCOUNT_FIXTURE and consumes live OpenAI quota"]
@@ -75,7 +74,6 @@ async fn real_openai_conversation_crosses_production_provider_boundaries() {
let prepared = bundle
.admin_provider()
.prepare_import(PrepareCredentialImport {
provider_instance_id: instance_id(),
document: ProviderDocument::new(OpaqueProviderData::new(document)),
})
.await
@@ -87,16 +85,9 @@ async fn real_openai_conversation_crosses_production_provider_boundaries() {
.expect("seed verified OpenAI account");
}
let instance = ProviderInstance::new(
instance_id(),
ProviderKind::new("openai").expect("provider"),
"https://chatgpt.com/backend-api".to_owned(),
true,
InstanceHealth::Healthy,
);
let models = bundle
.core_provider()
.query_model_capabilities(&instance)
.query_model_capabilities()
.await
.expect("query live OpenAI catalog");
let model = models.first().expect("one live generation model");
@@ -104,7 +95,7 @@ async fn real_openai_conversation_crosses_production_provider_boundaries() {
.admin_provider()
.connection_test_operation(model.upstream_model(), "Reply with exactly CPR_REAL_OK.")
.expect("real connection operation");
let request = planned_real_request(instance, model.upstream_model().clone(), operation);
let request = planned_real_request(model.upstream_model().clone(), operation);
let mut stream = bundle
.core_provider()
.execute(request, real_attempt_context())
@@ -133,7 +124,6 @@ fn prepared_account(
) -> NewProviderAccount {
let account = ProviderAccount::new(
value.account_id,
value.provider_instance_id,
value.provider_kind,
value.name,
value.upstream_user_id,
@@ -154,11 +144,8 @@ fn prepared_account(
}
}
fn planned_real_request(
instance: ProviderInstance,
upstream_model: UpstreamModelId,
operation: Operation,
) -> ProviderRequest {
fn planned_real_request(upstream_model: UpstreamModelId, operation: Operation) -> ProviderRequest {
let provider = ProviderKind::new("openai").expect("provider");
let public_model = PublicModelId::new(upstream_model.as_str()).expect("public model");
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("revision"),
@@ -167,9 +154,9 @@ fn planned_real_request(
NonZeroU32::new(1).expect("concurrency"),
Duration::ZERO,
),
vec![instance.clone()],
vec![provider.clone()],
vec![ProviderModel::new(
instance.id().clone(),
provider,
upstream_model,
ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), 1_000_000, None),
)],
@@ -351,7 +338,6 @@ async fn openai_admin_provider_persists_the_full_pending_envelope_and_binds_owne
Revision::new(7).expect("revision"),
ProviderKind::new("openai").expect("provider"),
AuthorizationMutationTarget::Create {
provider_instance_id: instance_id(),
name: "OAuth account".to_owned(),
},
AuthorizationOwnerBinding::from_context(&start_context),
@@ -401,7 +387,6 @@ async fn openai_admin_provider_projects_cached_quota_models_and_canonical_export
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: "acct_admin_projection".to_owned(),
provider_instance_id: instance_id().to_string(),
name: "admin projection".to_owned(),
secret: secret("admin-projection-access"),
verified_account: profile("chatgpt-admin-projection"),
@@ -487,7 +472,6 @@ async fn openai_admin_provider_rejects_unprepared_mutations_before_store_commit(
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: "acct_admin_invalid".to_owned(),
provider_instance_id: instance_id().to_string(),
name: "admin invalid".to_owned(),
secret: secret("admin-invalid-access"),
verified_account: profile("chatgpt-admin-invalid"),
@@ -506,7 +490,6 @@ async fn openai_admin_provider_rejects_unprepared_mutations_before_store_commit(
let admin = bundle.admin_provider();
let import_error = admin
.prepare_import(PrepareCredentialImport {
provider_instance_id: instance_id(),
document: ProviderDocument::new(OpaqueProviderData::new(Map::new())),
})
.await
@@ -544,10 +527,7 @@ fn provider_ports_with(
ProviderStorePorts::new(
accounts,
Arc::new(TestLeaseCoordinator::default()),
ProviderCatalogPorts::new(
Arc::new(TestInstanceCatalog),
Arc::new(TestCatalogCache::default()),
),
Arc::new(TestCatalogCache::default()),
Arc::new(TestCredentialState),
Arc::new(TestCooldown),
Arc::new(TestRuntimePolicy),
@@ -559,7 +539,6 @@ fn account_record(account: &ProviderAccount) -> AccountRecord {
let now = Utc::now();
AccountRecord {
id: account.id().to_string(),
provider_instance_id: account.instance().clone(),
provider_kind: account.provider().clone(),
name: account.name().to_owned(),
email: account.email().map(str::to_owned),
@@ -600,25 +579,6 @@ fn valid_config() -> OpenAiConfig {
}
}
struct TestInstanceCatalog;
impl ProviderInstanceCatalogPort for TestInstanceCatalog {
fn list_instances<'a>(
&'a self,
provider_kind: &'a ProviderKind,
_include_disabled: bool,
) -> BoxFuture<'a, Result<Vec<ProviderInstanceConfig>, ProviderStoreError>> {
Box::pin(async move {
Ok(vec![ProviderInstanceConfig::new(
ProviderInstanceId::new("inst_openai_primary").expect("instance"),
provider_kind.clone(),
"https://chatgpt.com/backend-api".to_owned(),
true,
)])
})
}
}
#[derive(Default)]
struct TestCatalogCache {
values: Mutex<BTreeMap<String, OpaqueProviderData>>,
backend/crates/providers/openai/tests/credential/admin.rs
+89 -114
View File
@@ -23,14 +23,11 @@ use provider_openai::credential::{
use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState};
use secrecy::{ExposeSecret, SecretString};
use crate::support::{
MemoryAccountStore, codex_account, instance_id, profile, runtime_policy, secret,
};
use crate::support::{MemoryAccountStore, codex_account, profile, runtime_policy, secret};
fn import(id: &str, token: &str) -> ImportCodexOAuthCredential {
ImportCodexOAuthCredential {
account_id: id.to_owned(),
provider_instance_id: instance_id().to_string(),
name: format!("name-{id}"),
secret: secret(token),
verified_account: profile(&format!("chatgpt-{id}")),
@@ -270,7 +267,7 @@ struct ManualLeases {
impl ProviderLeasePort for ManualLeases {
fn load_state<'a>(
&'a self,
_: &'a gateway_core::routing::ProviderInstanceId,
_: &'a gateway_core::routing::ProviderKind,
_: &'a [ProviderAccountId],
) -> BoxFuture<
'a,
@@ -321,7 +318,6 @@ async fn manual_refresh_fixture(
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: "acct_manual_refresh".to_owned(),
provider_instance_id: instance_id().to_string(),
name: "manual refresh".to_owned(),
secret: secret("old-access"),
verified_account: profile("chatgpt-acct_manual_refresh"),
@@ -507,26 +503,23 @@ fn unused_import_refresher() -> Arc<ManualRefresher> {
async fn formal_cpr_import_is_strict_and_uses_the_single_core_write_shape() {
let refresher = unused_import_refresher();
let prepared = import_service(refresher.clone())
.prepare_import_document(
instance_id(),
serde_json::json!({
"sourceFormat": "cpr",
"accounts": [{
"id": "acct_cpr_import",
"email": "cpr@example.com",
"accountId": "chatgpt-cpr",
"userId": "user-chatgpt-cpr",
"label": "CPR account",
"planType": "pro",
"token": "Bearer token-cpr",
"refreshToken": "refresh-cpr",
"accessTokenExpiresAt": "2100-01-01T00:00:00+00:00",
"status": "disabled",
"addedAt": "2026-07-18T10:47:01+08:00",
"updatedAt": "2026-07-19T11:00:00+08:00"
}]
}),
)
.prepare_import_document(serde_json::json!({
"sourceFormat": "cpr",
"accounts": [{
"id": "acct_cpr_import",
"email": "cpr@example.com",
"accountId": "chatgpt-cpr",
"userId": "user-chatgpt-cpr",
"label": "CPR account",
"planType": "pro",
"token": "Bearer token-cpr",
"refreshToken": "refresh-cpr",
"accessTokenExpiresAt": "2100-01-01T00:00:00+00:00",
"status": "disabled",
"addedAt": "2026-07-18T10:47:01+08:00",
"updatedAt": "2026-07-19T11:00:00+08:00"
}]
}))
.await
.expect("CPR import");
assert_eq!(prepared.accounts().len(), 1);
@@ -539,13 +532,10 @@ async fn formal_cpr_import_is_strict_and_uses_the_single_core_write_shape() {
assert!(refresher.seen.lock().expect("seen tokens").is_empty());
let error = import_service(unused_import_refresher())
.prepare_import_document(
instance_id(),
serde_json::json!({
"sourceFormat": "cpr",
"accounts": [{"token": "token-cpr", "unexpected": true}]
}),
)
.prepare_import_document(serde_json::json!({
"sourceFormat": "cpr",
"accounts": [{"token": "token-cpr", "unexpected": true}]
}))
.await
.expect_err("unknown CPR account key");
assert_eq!(error, CodexCredentialAdminError::InvalidInput);
@@ -554,18 +544,15 @@ async fn formal_cpr_import_is_strict_and_uses_the_single_core_write_shape() {
#[tokio::test]
async fn cpr_import_uses_verified_token_identity_instead_of_stale_export_projections() {
let error = import_service(unused_import_refresher())
.prepare_import_document(
instance_id(),
serde_json::json!({
"sourceFormat": "cpr",
"accounts": [{
"id": "acct_verified_identity",
"accountId": "user-chatgpt-cpr",
"userId": "user-chatgpt-cpr",
"token": "token-cpr"
}]
}),
)
.prepare_import_document(serde_json::json!({
"sourceFormat": "cpr",
"accounts": [{
"id": "acct_verified_identity",
"accountId": "user-chatgpt-cpr",
"userId": "user-chatgpt-cpr",
"token": "token-cpr"
}]
}))
.await
.expect_err("stale document identity must not override authenticated identity");
@@ -575,23 +562,20 @@ async fn cpr_import_uses_verified_token_identity_instead_of_stale_export_project
#[tokio::test]
async fn cpr_batch_allows_distinct_users_in_the_same_workspace() {
let prepared = import_service(unused_import_refresher())
.prepare_import_document(
instance_id(),
serde_json::json!({
"sourceFormat": "cpr",
"accounts": [{
"id": "acct_shared_a",
"accountId": "chatgpt-shared",
"userId": "user-a",
"token": "token-shared-a"
}, {
"id": "acct_shared_b",
"accountId": "chatgpt-shared",
"userId": "user-b",
"token": "token-shared-b"
}]
}),
)
.prepare_import_document(serde_json::json!({
"sourceFormat": "cpr",
"accounts": [{
"id": "acct_shared_a",
"accountId": "chatgpt-shared",
"userId": "user-a",
"token": "token-shared-a"
}, {
"id": "acct_shared_b",
"accountId": "chatgpt-shared",
"userId": "user-b",
"token": "token-shared-b"
}]
}))
.await
.expect("distinct users sharing one workspace are separate credentials");
@@ -604,23 +588,20 @@ async fn cpr_batch_allows_distinct_users_in_the_same_workspace() {
);
let duplicate = import_service(unused_import_refresher())
.prepare_import_document(
instance_id(),
serde_json::json!({
"sourceFormat": "cpr",
"accounts": [{
"id": "acct_duplicate_a",
"accountId": "chatgpt-shared",
"userId": "user-a",
"token": "token-shared-a"
}, {
"id": "acct_duplicate_b",
"accountId": "chatgpt-shared",
"userId": "user-a",
"token": "token-shared-a"
}]
}),
)
.prepare_import_document(serde_json::json!({
"sourceFormat": "cpr",
"accounts": [{
"id": "acct_duplicate_a",
"accountId": "chatgpt-shared",
"userId": "user-a",
"token": "token-shared-a"
}, {
"id": "acct_duplicate_b",
"accountId": "chatgpt-shared",
"userId": "user-a",
"token": "token-shared-a"
}]
}))
.await
.expect_err("the same upstream user and workspace must stay unique");
assert_eq!(duplicate, CodexCredentialAdminError::InvalidInput);
@@ -629,43 +610,37 @@ async fn cpr_batch_allows_distinct_users_in_the_same_workspace() {
#[tokio::test]
async fn credential_bundle_and_auth_document_normalize_to_the_same_core_accounts() {
let bundle = import_service(unused_import_refresher())
.prepare_import_document(
instance_id(),
serde_json::json!({
"exported_at": "2026-07-03T15:46:38.717Z",
"proxies": [],
"accounts": [{
"name": "bundle@example.com",
"platform": "openai",
"type": "oauth",
"credentials": {
"at": "token-bundle",
"refresh_token": "refresh-bundle",
"chatgpt_account_id": "chatgpt-bundle",
"chatgpt_user_id": "user-chatgpt-bundle"
},
"concurrency": 3,
"priority": 50
}]
}),
)
.prepare_import_document(serde_json::json!({
"exported_at": "2026-07-03T15:46:38.717Z",
"proxies": [],
"accounts": [{
"name": "bundle@example.com",
"platform": "openai",
"type": "oauth",
"credentials": {
"at": "token-bundle",
"refresh_token": "refresh-bundle",
"chatgpt_account_id": "chatgpt-bundle",
"chatgpt_user_id": "user-chatgpt-bundle"
},
"concurrency": 3,
"priority": 50
}]
}))
.await
.expect("credential bundle import");
let auth_document = import_service(unused_import_refresher())
.prepare_import_document(
instance_id(),
serde_json::json!({
"accounts": [{
"type": "openai",
"access_token": "token-auth-document",
"refresh_token": "refresh-auth-document",
"chatgpt_account_id": "chatgpt-auth-document",
"chatgpt_user_id": "user-chatgpt-auth-document",
"email": "auth-document@example.com",
"label": "Auth document"
}]
}),
)
.prepare_import_document(serde_json::json!({
"accounts": [{
"type": "openai",
"access_token": "token-auth-document",
"refresh_token": "refresh-auth-document",
"chatgpt_account_id": "chatgpt-auth-document",
"chatgpt_user_id": "user-chatgpt-auth-document",
"email": "auth-document@example.com",
"label": "Auth document"
}]
}))
.await
.expect("auth document import");
@@ -738,7 +713,7 @@ async fn real_cpr_fixture_import_contract() {
runtime_policy(),
);
let prepared = service
.prepare_import_document(instance_id(), payload)
.prepare_import_document(payload)
.await
.expect("real CPR import");
backend/crates/providers/openai/tests/credential/catalog.rs
+65 -294
View File
@@ -1,21 +1,12 @@
use std::sync::Arc;
use std::time::SystemTime;
use chrono::{TimeZone, Utc};
use chrono::{TimeZone as _, Utc};
use futures::executor::block_on;
use gateway_core::engine::credential::{
AccountAvailability, AccountStateChange, CredentialRevision, ProviderAccountId,
ProviderAccountStore as _,
};
use gateway_core::routing::{InstanceHealth, ProviderInstance, ProviderKind};
use provider_openai::credential::{
CodexCredentialCatalogError, CodexCredentialCatalogService, ImportCodexOAuthCredential,
};
use gateway_core::engine::credential::ProviderAccountId;
use provider_openai::credential::{CodexCredentialCatalogError, CodexCredentialCatalogService};
use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState};
use wiremock::matchers::{header, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
use crate::support::{MemoryAccountStore, instance_id, loopback_origin_policy, profile, secret};
use crate::support::MemoryAccountStore;
fn wire_profile() -> CodexWireProfileState {
CodexWireProfileState::new(CodexWireProfile {
@@ -26,7 +17,7 @@ fn wire_profile() -> CodexWireProfileState {
os_type: "linux".to_owned(),
os_version: "6.8".to_owned(),
arch: "x86_64".to_owned(),
terminal: "xterm".to_owned(),
terminal: "catalog-contract".to_owned(),
verified_at: Utc
.with_ymd_and_hms(2026, 7, 18, 0, 0, 0)
.single()
@@ -34,30 +25,6 @@ fn wire_profile() -> CodexWireProfileState {
})
}
fn instance(base_url: &str) -> ProviderInstance {
ProviderInstance::new(
instance_id(),
ProviderKind::new("openai").expect("provider"),
base_url.to_owned(),
true,
InstanceHealth::Healthy,
)
}
async fn create(store: &Arc<MemoryAccountStore>, id: &str, token: &str) {
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: id.to_owned(),
provider_instance_id: instance_id().to_string(),
name: id.to_owned(),
secret: secret(token),
verified_account: profile(&format!("chatgpt-{id}")),
next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
})
.await;
}
fn service(store: &Arc<MemoryAccountStore>) -> CodexCredentialCatalogService {
CodexCredentialCatalogService::new(
store.repository(),
@@ -66,275 +33,79 @@ fn service(store: &Arc<MemoryAccountStore>) -> CodexCredentialCatalogService {
.no_proxy()
.build()
.expect("client"),
loopback_origin_policy(),
)
}
fn catalog(model: &str, display: &str) -> serde_json::Value {
serde_json::json!({
"models": [{
"slug": model,
"display_name": display,
"supported_in_api": true,
"supported_reasoning_levels": [{"effort": "medium"}],
"context_window": 128000
}]
})
}
fn account_id(value: &str) -> ProviderAccountId {
ProviderAccountId::new(value).expect("account id")
}
fn revision() -> CredentialRevision {
CredentialRevision::new(1).expect("revision")
}
#[tokio::test]
async fn realtime_catalog_builds_union_and_per_account_entitlements() {
let server = MockServer::start().await;
for (token, model) in [("token-one", "gpt-5.4"), ("token-two", "gpt-5.5")] {
Mock::given(method("GET"))
.and(path("/backend-api/codex/models"))
.and(header("authorization", format!("Bearer {token}")))
.respond_with(ResponseTemplate::new(200).set_body_json(catalog(model, model)))
.expect(1)
.mount(&server)
.await;
}
#[test]
fn catalog_starts_without_a_snapshot_or_generation() {
let store = Arc::new(MemoryAccountStore::default());
create(&store, "acct_one", "token-one").await;
create(&store, "acct_two", "token-two").await;
let snapshot = service(&store)
.synchronize_instance(&instance(&format!("{}/backend-api", server.uri())))
.await
.expect("synchronize catalog");
assert_eq!(snapshot.models().len(), 2);
assert_eq!(
snapshot.account_models(&account_id("acct_one"), revision()),
Some(&["gpt-5.4".to_owned()][..])
);
assert_eq!(
snapshot.account_models(&account_id("acct_two"), revision()),
Some(&["gpt-5.5".to_owned()][..])
);
}
#[tokio::test]
async fn realtime_catalog_should_preserve_the_official_model_order() {
let server = MockServer::start().await;
let official_order = ["gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.4"];
Mock::given(method("GET"))
.and(path("/backend-api/codex/models"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"models": official_order
.iter()
.enumerate()
.map(|(index, model)| serde_json::json!({
"slug": model,
"display_name": model,
"priority": index + 1,
"supported_in_api": true,
"supported_reasoning_levels": [{"effort": "medium"}],
"context_window": 128000
}))
.collect::<Vec<_>>()
})))
.expect(1)
.mount(&server)
.await;
let store = Arc::new(MemoryAccountStore::default());
create(&store, "acct_order", "token-order").await;
let snapshot = service(&store)
.synchronize_instance(&instance(&format!("{}/backend-api", server.uri())))
.await
.expect("ordered catalog");
assert_eq!(
snapshot
.models()
.iter()
.map(|model| model.request_model().as_str())
.collect::<Vec<_>>(),
official_order
);
assert_eq!(
snapshot.account_models(&account_id("acct_order"), revision()),
Some(official_order.map(str::to_owned).as_slice())
);
}
#[tokio::test]
async fn catalog_is_reused_until_explicit_invalidation_or_etag_change() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/backend-api/codex/models"))
.respond_with(ResponseTemplate::new(200).set_body_json(catalog("gpt-5.4", "GPT-5.4")))
.expect(1)
.mount(&server)
.await;
let store = Arc::new(MemoryAccountStore::default());
create(&store, "acct_one", "token-one").await;
let service = service(&store);
let instance = instance(&format!("{}/backend-api", server.uri()));
let first = service
.synchronize_instance(&instance)
.await
.expect("first query");
let second = service
.synchronize_instance(&instance)
.await
.expect("cached query");
assert_eq!(first.observed_at(), second.observed_at());
assert!(service.cached(instance.id()).expect("cache read").is_some());
}
#[tokio::test]
async fn single_account_refresh_queries_only_that_account_and_updates_its_cache() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/backend-api/codex/models"))
.and(header("authorization", "Bearer token-two"))
.respond_with(ResponseTemplate::new(200).set_body_json(catalog("gpt-5.5", "GPT-5.5")))
.expect(1)
.mount(&server)
.await;
let store = Arc::new(MemoryAccountStore::default());
create(&store, "acct_one", "token-one").await;
create(&store, "acct_two", "token-two").await;
let service = service(&store);
let instance = instance(&format!("{}/backend-api", server.uri()));
let selected = ProviderAccountId::new("acct_two").expect("account id");
let models = service
.synchronize_account(&instance, &selected)
.await
.expect("single account catalog");
assert_eq!(models, ["gpt-5.5"]);
assert_eq!(
service
.cached_account_models(instance.id(), &selected, revision())
.expect("cache read"),
Some(vec!["gpt-5.5".to_owned()])
);
assert_eq!(
service
.cached_account_models(instance.id(), &account_id("acct_one"), revision(),)
.expect("other cache read"),
None
);
}
#[tokio::test]
async fn invalidation_forces_a_new_realtime_query() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/backend-api/codex/models"))
.respond_with(ResponseTemplate::new(200).set_body_json(catalog("gpt-5.4", "GPT-5.4")))
.expect(2)
.mount(&server)
.await;
let store = Arc::new(MemoryAccountStore::default());
create(&store, "acct_one", "token-one").await;
let service = service(&store);
let instance = instance(&format!("{}/backend-api", server.uri()));
assert!(service.cached().expect("cache read").is_none());
assert_eq!(service.catalog_generation().get(), 0);
service
.synchronize_instance(&instance)
.await
.expect("first query");
assert_eq!(service.catalog_generation().get(), 1);
service.invalidate(instance.id()).expect("invalidate");
assert_eq!(service.catalog_generation().get(), 2);
service
.synchronize_instance(&instance)
.await
.expect("second query");
assert_eq!(service.catalog_generation().get(), 3);
}
#[tokio::test]
async fn response_etag_change_is_deduplicated_and_queued_for_refresh() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/backend-api/codex/models"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("etag", "\"models-v1\"")
.set_body_json(catalog("gpt-5.4", "GPT-5.4")),
)
.expect(1)
.mount(&server)
.await;
let store = Arc::new(MemoryAccountStore::default());
create(&store, "acct_one", "token-one").await;
let service = service(&store);
let instance = instance(&format!("{}/backend-api", server.uri()));
service
.synchronize_instance(&instance)
.await
.expect("initial catalog");
assert!(
!service
.observe_response_etag(instance.id(), "\"models-v1\"")
.expect("same ETag")
);
assert!(
service
.observe_response_etag(instance.id(), "\"models-v2\"")
.expect("changed ETag")
);
assert!(
!service
.observe_response_etag(instance.id(), "\"models-v2\"")
.expect("duplicate ETag")
);
assert_eq!(
service.wait_for_etag_refresh().await,
vec![instance.id().clone()]
);
}
#[tokio::test]
async fn quota_exhausted_account_remains_eligible_for_catalog_discovery() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/backend-api/codex/models"))
.respond_with(ResponseTemplate::new(200).set_body_json(catalog("gpt-5.4", "GPT-5.4")))
.expect(1)
.mount(&server)
.await;
let store = Arc::new(MemoryAccountStore::default());
create(&store, "acct_quota", "token-one").await;
let account = store.account("acct_quota").expect("created account");
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
.expect("mark quota exhausted");
let snapshot = service(&store)
.synchronize_instance(&instance(&format!("{}/backend-api", server.uri())))
.await
.expect("discover catalog through quota exhausted account");
assert_eq!(snapshot.models().len(), 1);
}
#[test]
fn catalog_without_schedulable_accounts_fails_before_network_io() {
fn catalog_without_openai_accounts_fails_before_network_io() {
let store = Arc::new(MemoryAccountStore::default());
let error =
block_on(service(&store).synchronize_instance(&instance("http://127.0.0.1:9/backend-api")))
.expect_err("empty account set");
let error = block_on(service(&store).synchronize()).expect_err("empty account set");
assert!(matches!(
error,
CodexCredentialCatalogError::NoEligibleCredential
));
}
#[tokio::test]
async fn missing_account_refresh_fails_before_network_io() {
let store = Arc::new(MemoryAccountStore::default());
let account = ProviderAccountId::new("acct_missing").expect("account id");
let error = service(&store)
.synchronize_account(&account)
.await
.expect_err("missing account");
assert!(matches!(
error,
CodexCredentialCatalogError::NoEligibleCredential
));
}
#[tokio::test]
async fn response_etag_change_is_deduplicated_and_queued_once() {
let store = Arc::new(MemoryAccountStore::default());
let service = service(&store);
assert!(
service
.observe_response_etag("\"models-v2\"")
.expect("new ETag")
);
assert!(
!service
.observe_response_etag("\"models-v2\"")
.expect("duplicate ETag")
);
tokio::time::timeout(
std::time::Duration::from_secs(1),
service.wait_for_etag_refresh(),
)
.await
.expect("queued ETag refresh");
}
#[test]
fn invalid_response_etag_is_rejected_without_touching_the_catalog() {
let store = Arc::new(MemoryAccountStore::default());
let service = service(&store);
let error = service
.observe_response_etag("invalid\netag")
.expect_err("invalid ETag");
assert!(matches!(error, CodexCredentialCatalogError::InvalidEtag));
assert!(service.cached().expect("cache read").is_none());
assert_eq!(service.catalog_generation().get(), 0);
}
backend/crates/providers/openai/tests/credential/contract.rs
+7 -23
View File
@@ -12,24 +12,21 @@ use gateway_core::engine::{
AccountAttemptContext, AttemptContext, CancellationToken, ModelRequestId, RequestAttemptContext,
};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::ProviderKind;
use provider_openai::credential::{
CodexAccountFailure, CodexCookiePolicy, CodexCredentialCatalogService, CodexCredentialCodec,
CodexCredentialQuotaService, CodexCredentialSelector, CredentialSelectionError,
ImportCodexOAuthCredential, SelectCodexCredential,
};
use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState};
use provider_openai::{CodexOriginPolicy, OfficialCodexOriginPolicy};
use secrecy::ExposeSecret;
use url::Url;
use crate::support::{
MemoryAccountStore, TestLeaseCoordinator, account_policy, instance_id, profile, secret,
};
use crate::support::{MemoryAccountStore, TestLeaseCoordinator, account_policy, profile, secret};
fn create_account(store: &Arc<MemoryAccountStore>, id: &str, token: &str) {
block_on(store.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: id.to_owned(),
provider_instance_id: instance_id().to_string(),
name: id.to_owned(),
secret: secret(token),
verified_account: profile(&format!("chatgpt-{id}")),
@@ -95,20 +92,18 @@ fn selector(
verified_at: Utc::now(),
});
let http = reqwest::Client::builder().build().expect("HTTP client");
let origin: Arc<dyn CodexOriginPolicy> = Arc::new(OfficialCodexOriginPolicy);
let catalog = Arc::new(CodexCredentialCatalogService::new(
store.repository(),
profile.clone(),
http.clone(),
Arc::clone(&origin),
));
let quota = Arc::new(CodexCredentialQuotaService::new(
store.repository(),
profile,
http,
origin,
));
CodexCredentialSelector::new(
ProviderKind::new("openai").expect("provider"),
store.repository(),
leases,
catalog,
@@ -231,7 +226,6 @@ fn selector_uses_frozen_global_account_policy_for_lease() {
let lease =
block_on(
selector.select(&SelectCodexCredential {
provider_instance_id: &instance_id(),
upstream_model: "gpt-5.4",
request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses")
.expect("request URL"),
@@ -252,7 +246,10 @@ fn selector_uses_frozen_global_account_policy_for_lease() {
.expect("runtime credential");
assert_eq!(runtime.installation_id, installation_id);
let requests = leases.requests.lock().expect("lease requests lock");
assert_eq!(requests[0].provider_instance_id(), &instance_id());
assert_eq!(
requests[0].provider_kind(),
&ProviderKind::new("openai").expect("provider")
);
assert_eq!(requests[0].account_id(), lease.account_id());
assert_eq!(
requests[0].credential_revision(),
@@ -276,7 +273,6 @@ fn selector_round_robin_cursor_advances_across_requests() {
for _ in 0..4 {
let attempt = round_robin_attempt();
let lease = block_on(selector.select(&SelectCodexCredential {
provider_instance_id: &instance_id(),
upstream_model: "gpt-5.4",
request_url: &request_url,
attempt: &attempt,
@@ -303,7 +299,6 @@ fn selector_honors_attempt_local_account_exclusion() {
let lease =
block_on(
selector.select(&SelectCodexCredential {
provider_instance_id: &instance_id(),
upstream_model: "gpt-5.4",
request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses")
.expect("request URL"),
@@ -325,7 +320,6 @@ fn selector_uses_only_the_required_account() {
let lease =
block_on(
selector.select(&SelectCodexCredential {
provider_instance_id: &instance_id(),
upstream_model: "gpt-5.4",
request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses")
.expect("request URL"),
@@ -348,7 +342,6 @@ fn unavailable_required_account_never_falls_back() {
let error =
block_on(
selector.select(&SelectCodexCredential {
provider_instance_id: &instance_id(),
upstream_model: "gpt-5.4",
request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses")
.expect("request URL"),
@@ -373,7 +366,6 @@ fn selector_returns_capacity_error_when_every_redis_lease_is_busy() {
let error =
block_on(
selector.select(&SelectCodexCredential {
provider_instance_id: &instance_id(),
upstream_model: "gpt-5.4",
request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses")
.expect("request URL"),
@@ -398,7 +390,6 @@ fn credential_expired_failure_marks_unified_account_expired() {
let lease =
block_on(
selector.select(&SelectCodexCredential {
provider_instance_id: &instance_id(),
upstream_model: "gpt-5.4",
request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses")
.expect("request URL"),
@@ -427,7 +418,6 @@ fn identity_verification_failure_isolates_only_selected_account() {
let lease =
block_on(
selector.select(&SelectCodexCredential {
provider_instance_id: &instance_id(),
upstream_model: "gpt-5.4",
request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses")
.expect("request URL"),
@@ -466,7 +456,6 @@ fn cloudflare_challenge_backoff_escalates_and_success_resets_it() {
let lease =
block_on(
selector.select(&SelectCodexCredential {
provider_instance_id: &instance_id(),
upstream_model: "gpt-5.4",
request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses")
.expect("request URL"),
@@ -532,7 +521,6 @@ fn repeated_cloudflare_path_block_marks_only_the_affected_account_invalid() {
let lease =
block_on(
selector.select(&SelectCodexCredential {
provider_instance_id: &instance_id(),
upstream_model: "gpt-5.4",
request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses")
.expect("request URL"),
@@ -572,7 +560,6 @@ fn cloudflare_challenge_expires_provider_owned_cookies_at_cooldown_boundary() {
Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("request URL");
let first_attempt = attempt_with_required(BTreeSet::new(), Some(required.clone()));
let first = block_on(selector.select(&SelectCodexCredential {
provider_instance_id: &instance_id(),
upstream_model: "gpt-5.4",
request_url: &request_url,
attempt: &first_attempt,
@@ -587,7 +574,6 @@ fn cloudflare_challenge_expires_provider_owned_cookies_at_cooldown_boundary() {
let second_attempt = attempt_with_required(BTreeSet::new(), Some(required));
let second = block_on(selector.select(&SelectCodexCredential {
provider_instance_id: &instance_id(),
upstream_model: "gpt-5.4",
request_url: &request_url,
attempt: &second_attempt,
@@ -620,7 +606,6 @@ fn cloudflare_path_block_deletes_provider_owned_cookies() {
Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("request URL");
let first_attempt = attempt_with_required(BTreeSet::new(), Some(required.clone()));
let first = block_on(selector.select(&SelectCodexCredential {
provider_instance_id: &instance_id(),
upstream_model: "gpt-5.4",
request_url: &request_url,
attempt: &first_attempt,
@@ -635,7 +620,6 @@ fn cloudflare_path_block_deletes_provider_owned_cookies() {
let second_attempt = attempt_with_required(BTreeSet::new(), Some(required));
let second = block_on(selector.select(&SelectCodexCredential {
provider_instance_id: &instance_id(),
upstream_model: "gpt-5.4",
request_url: &request_url,
attempt: &second_attempt,
backend/crates/providers/openai/tests/credential/oauth.rs
+1 -6
View File
@@ -24,7 +24,7 @@ use secrecy::{ExposeSecret, SecretString};
use sha2::{Digest as _, Sha256};
use url::Url;
use crate::support::{MemoryAccountStore, instance_id, profile, runtime_policy, secret};
use crate::support::{MemoryAccountStore, profile, runtime_policy, secret};
#[derive(Default)]
struct PendingStore {
@@ -45,7 +45,6 @@ impl CodexOAuthPendingStore for PendingStore {
flow_id: pending.flow_id().to_owned(),
owner_ref: pending.owner_ref().to_owned(),
started_request_ref: pending.started_request_ref().to_owned(),
provider_instance_id: pending.provider_instance_id().to_owned(),
name: pending.name().to_owned(),
expires_at: pending.expires_at(),
state: pending.state().clone(),
@@ -176,7 +175,6 @@ fn create_mutation(request_id: &str) -> PendingAuthorizationMutation {
Revision::new(1).expect("revision"),
ProviderKind::new("openai").expect("provider"),
AuthorizationMutationTarget::Create {
provider_instance_id: instance_id(),
name: "OAuth Account".to_owned(),
},
AuthorizationOwnerBinding::from_context(&context),
@@ -189,7 +187,6 @@ fn reauthorization_mutation(request_id: &str, account_id: &str) -> PendingAuthor
Revision::new(1).expect("revision"),
ProviderKind::new("openai").expect("provider"),
AuthorizationMutationTarget::Reauthorize {
provider_instance_id: instance_id(),
account_id: ProviderAccountId::new(account_id).expect("account id"),
expected_credential_revision: Revision::new(1).expect("revision"),
},
@@ -317,7 +314,6 @@ async fn reauthorization_binds_account_revision_and_returns_only_prepared_rotati
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: "acct_oauth_reauth".to_owned(),
provider_instance_id: instance_id().to_string(),
name: "reauthorize".to_owned(),
secret: secret("old-oauth-access"),
verified_account: profile("chatgpt-oauth"),
@@ -414,7 +410,6 @@ async fn reauthorization_rejects_identity_rebinding() {
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: "acct_oauth_rebind".to_owned(),
provider_instance_id: instance_id().to_string(),
name: "reauthorize".to_owned(),
secret: secret("old-oauth-access"),
verified_account: profile("chatgpt-oauth"),
backend/crates/providers/openai/tests/credential/quota.rs
+71 -313
View File
@@ -1,22 +1,20 @@
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use std::time::SystemTime;
use chrono::{TimeZone, Utc};
use chrono::{TimeZone as _, Utc};
use futures::future::join_all;
use gateway_core::engine::credential::{
AccountAvailability, AccountStateChange, ProviderAccountStore,
AccountAvailability, OpaqueProviderData, ProviderAccountStore as _, QuotaObservation,
QuotaWriteOutcome,
};
use gateway_core::routing::{InstanceHealth, ProviderInstance, ProviderKind};
use provider_openai::credential::{
CodexCredentialQuotaService, CodexQuotaSyncSummary, CodexQuotaWindowKind,
ImportCodexOAuthCredential, parse_codex_quota_usage,
CodexCredentialQuotaService, CodexQuotaWindowKind, ImportCodexOAuthCredential,
parse_codex_quota_usage,
};
use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState};
use serde_json::json;
use wiremock::matchers::{header, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
use crate::support::{MemoryAccountStore, instance_id, loopback_origin_policy, profile, secret};
use crate::support::{MemoryAccountStore, profile, secret};
#[test]
fn parser_extracts_dynamic_windows_without_a_fixed_database_shape() {
@@ -32,6 +30,7 @@ fn parser_extracts_dynamic_windows_without_a_fixed_database_shape() {
}]
}))
.expect("valid dynamic quota");
assert_eq!(fact.remaining_percent(), Some(20));
assert_eq!(
fact.resets_at().map(|value| value.timestamp()),
@@ -49,6 +48,7 @@ fn parser_treats_any_confirmed_provider_limit_as_exhausted() {
}]
}))
.expect("valid quota");
assert!(fact.exhausted());
}
@@ -63,6 +63,7 @@ fn parser_does_not_infer_exhaustion_from_unknown_credit_fields() {
}
}))
.expect("recognized credits object");
assert!(!fact.exhausted());
}
@@ -108,6 +109,7 @@ fn parser_rejects_wrong_known_field_type_without_echoing_body() {
"rate_limit": {"allowed": marker}
}))
.expect_err("known field type must be strict");
assert!(!format!("{error:?} {error}").contains(marker));
}
@@ -125,7 +127,7 @@ fn wire_profile() -> CodexWireProfileState {
os_type: "linux".to_owned(),
os_version: "6.8".to_owned(),
arch: "x86_64".to_owned(),
terminal: "xterm".to_owned(),
terminal: "quota-contract".to_owned(),
verified_at: Utc
.with_ymd_and_hms(2026, 7, 18, 0, 0, 0)
.single()
@@ -141,24 +143,26 @@ fn quota_service(store: &Arc<MemoryAccountStore>) -> CodexCredentialQuotaService
.no_proxy()
.build()
.expect("client"),
loopback_origin_policy(),
)
}
async fn create_account(store: &Arc<MemoryAccountStore>, account_id: &str) {
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: account_id.to_owned(),
name: account_id.to_owned(),
secret: secret(&format!("token-{account_id}")),
verified_account: profile(&format!("chatgpt-{account_id}")),
next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
})
.await;
}
#[tokio::test]
async fn concurrent_cold_scheduling_hydration_reads_quota_once() {
let store = Arc::new(MemoryAccountStore::default());
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: "acct_hydration".to_owned(),
provider_instance_id: instance_id().to_string(),
name: "hydration".to_owned(),
secret: secret("hydration-token"),
verified_account: profile("chatgpt-acct_hydration"),
next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
})
.await;
create_account(&store, "acct_hydration").await;
let account = store.account("acct_hydration").expect("created account");
let service = quota_service(&store);
@@ -168,8 +172,10 @@ async fn concurrent_cold_scheduling_hydration_reads_quota_once() {
}
#[tokio::test]
async fn quota_service_stores_raw_provider_json_and_projects_common_state() {
let server = MockServer::start().await;
async fn persisted_provider_quota_projects_dynamic_windows_without_network_io() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_quota").await;
let account = store.account("acct_quota").expect("created account");
let raw = json!({
"rate_limit": {
"allowed": true,
@@ -182,8 +188,7 @@ async fn quota_service_stores_raw_provider_json_and_projects_common_state() {
"used_percent": 42,
"reset_at": 1_900_604_800,
"limit_window_seconds": 604_800
},
"provider_specific_future_window": {"used": "opaque"}
}
},
"additional_rate_limits": null,
"spend_control": {
@@ -195,177 +200,39 @@ async fn quota_service_stores_raw_provider_json_and_projects_common_state() {
},
"provider_specific_root": {"opaque": [1, 2, 3]}
});
Mock::given(method("GET"))
.and(path("/backend-api/wham/usage"))
.and(header("authorization", "Bearer quota-token"))
.respond_with(ResponseTemplate::new(200).set_body_json(&raw))
.expect(2)
.mount(&server)
.await;
let store = Arc::new(MemoryAccountStore::default());
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: "acct_quota".to_owned(),
provider_instance_id: instance_id().to_string(),
name: "quota".to_owned(),
secret: secret("quota-token"),
verified_account: profile("chatgpt-acct_quota"),
next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
let outcome = store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
raw.as_object().expect("quota object").clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await;
let service = quota_service(&store);
let instance = ProviderInstance::new(
instance_id(),
ProviderKind::new("openai").expect("provider"),
format!("{}/backend-api", server.uri()),
true,
InstanceHealth::Healthy,
);
let summary = service
.synchronize_instance(&instance)
.await
.expect("synchronize quota");
assert_eq!(summary.updated, 1, "unexpected quota summary: {summary:?}");
let account = store.account("acct_quota").expect("account");
assert_eq!(account.availability(), AccountAvailability::Ready);
let parsed = service
.expect("persist quota");
assert_eq!(outcome, QuotaWriteOutcome::Updated);
let snapshot = quota_service(&store)
.read_account(account.id())
.await
.expect("read quota")
.expect("quota snapshot");
assert_eq!(parsed.account_id(), account.id());
assert_eq!(parsed.credential_revision(), account.revision());
assert_eq!(parsed.windows().len(), 3);
assert_eq!(parsed.windows()[0].kind(), CodexQuotaWindowKind::ShortTerm);
assert_eq!(parsed.windows()[0].window_seconds(), Some(18_000));
assert_eq!(parsed.windows()[0].source(), "core");
assert_eq!(parsed.windows()[1].kind(), CodexQuotaWindowKind::Weekly);
assert_eq!(parsed.windows()[1].window_seconds(), Some(604_800));
assert_eq!(parsed.windows()[2].kind(), CodexQuotaWindowKind::Monthly);
assert_eq!(parsed.windows()[2].source(), "spend_control");
let refreshed = service
.refresh_account(&instance, account.id())
.await
.expect("refresh one account");
assert_eq!(refreshed.account_id(), account.id());
assert_eq!(refreshed.windows()[0].used_percent(), Some(37.0));
let observation = store
.get_quotas(std::slice::from_ref(account.id()))
.await
.expect("quota read")
.pop()
.expect("quota observation");
assert_eq!(snapshot.windows().len(), 3);
assert_eq!(
observation
.quota
.expect("raw quota")
.expose_to_provider()
.get("provider_specific_root"),
raw.get("provider_specific_root")
);
}
#[tokio::test]
async fn quota_refresh_success_must_not_clear_newer_future_cooldown() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/backend-api/wham/usage"))
.respond_with(
ResponseTemplate::new(200)
.set_delay(Duration::from_millis(100))
.set_body_json(json!({
"rate_limit": {
"allowed": true,
"primary_window": {"used_percent": 10}
}
})),
)
.expect(1)
.mount(&server)
.await;
let store = Arc::new(MemoryAccountStore::default());
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: "acct_quota_cooldown".to_owned(),
provider_instance_id: instance_id().to_string(),
name: "quota cooldown".to_owned(),
secret: secret("quota-cooldown-token"),
verified_account: profile("chatgpt-acct_quota_cooldown"),
next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
})
.await;
let account = store.account("acct_quota_cooldown").expect("quota account");
let service = quota_service(&store);
let instance = ProviderInstance::new(
instance_id(),
ProviderKind::new("openai").expect("provider"),
format!("{}/backend-api", server.uri()),
true,
InstanceHealth::Healthy,
);
let cooldown_until = SystemTime::now() + Duration::from_secs(3_600);
let synchronize = service.synchronize_instance(&instance);
let apply_newer_cooldown = async {
tokio::time::sleep(Duration::from_millis(20)).await;
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::Cooldown,
reason: Some("model_request_rate_limited".to_owned()),
cooldown_until: Some(cooldown_until),
observed_at: SystemTime::now(),
})
.await
.expect("apply newer cooldown");
};
let (summary, ()) = tokio::join!(synchronize, apply_newer_cooldown);
assert_eq!(summary.expect("quota synchronization").updated, 1);
let refreshed = store
.account("acct_quota_cooldown")
.expect("refreshed account");
assert_eq!(refreshed.availability(), AccountAvailability::Cooldown);
assert_eq!(refreshed.cooldown_until(), Some(cooldown_until));
}
#[tokio::test]
async fn non_exhausted_quota_refresh_should_restore_quota_exhausted_account() {
let (summary, availability, cooldown_until) = synchronize_account_from_state(
"acct_quota_recovered",
AccountAvailability::QuotaExhausted,
None,
ResponseTemplate::new(200).set_body_json(json!({
"rate_limit": {
"allowed": true,
"primary_window": {"used_percent": 10}
}
})),
)
.await;
assert_eq!(
(summary.updated, availability, cooldown_until),
(1, AccountAvailability::Ready, None)
snapshot.windows()[0].kind(),
CodexQuotaWindowKind::ShortTerm
);
assert_eq!(snapshot.windows()[0].window_seconds(), Some(18_000));
assert_eq!(snapshot.windows()[1].kind(), CodexQuotaWindowKind::Weekly);
assert_eq!(snapshot.windows()[2].kind(), CodexQuotaWindowKind::Monthly);
}
#[tokio::test]
async fn passive_rate_limit_headers_update_quota_and_account_state_with_revision_fence() {
let store = Arc::new(MemoryAccountStore::default());
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: "acct_passive_quota".to_owned(),
provider_instance_id: instance_id().to_string(),
name: "passive quota".to_owned(),
secret: secret("passive-quota-token"),
verified_account: profile("chatgpt-acct_passive_quota"),
next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
})
.await;
create_account(&store, "acct_passive_quota").await;
let account = store.account("acct_passive_quota").expect("account");
let service = quota_service(&store);
let reset_at = 1_900_000_000_i64;
@@ -390,142 +257,33 @@ async fn passive_rate_limit_headers_update_quota_and_account_state_with_revision
.account("acct_passive_quota")
.expect("current account");
assert_eq!(current.availability(), AccountAvailability::QuotaExhausted);
let observation = store
.get_quotas(std::slice::from_ref(account.id()))
let snapshot = service
.read_account(account.id())
.await
.expect("quota")
.pop()
.expect("quota observation");
let quota = serde_json::Value::Object(
observation
.quota
.expect("quota JSON")
.expose_to_provider()
.clone(),
);
.expect("read quota")
.expect("quota snapshot");
assert_eq!(snapshot.windows()[0].window_seconds(), Some(18_000));
assert_eq!(
quota
.pointer("/rate_limit/primary_window/limit_window_seconds")
.and_then(serde_json::Value::as_u64),
Some(18_000)
);
assert_eq!(
quota
.pointer("/rate_limit/primary_window/reset_at")
.and_then(serde_json::Value::as_i64),
snapshot.windows()[0]
.reset_at()
.map(|value| value.timestamp()),
Some(reset_at)
);
}
#[tokio::test]
async fn exhausted_quota_refresh_should_upgrade_future_cooldown() {
let cooldown_until = SystemTime::now() + Duration::from_secs(3_600);
let (summary, availability, projected_cooldown) = synchronize_account_from_state(
"acct_quota_exhausted",
AccountAvailability::Cooldown,
Some(cooldown_until),
ResponseTemplate::new(200).set_body_json(json!({
"rate_limit": {
"allowed": false,
"limit_reached": true,
"primary_window": {"used_percent": 100}
}
})),
)
.await;
assert_eq!(
(summary.exhausted, availability, projected_cooldown),
(1, AccountAvailability::QuotaExhausted, None)
);
}
#[tokio::test]
async fn structured_deactivated_workspace_rejection_should_ban_account() {
let (summary, availability, cooldown_until) = synchronize_account_from_state(
"acct_quota_deactivated",
AccountAvailability::Ready,
None,
ResponseTemplate::new(402)
.set_body_json(json!({"detail": {"code": "deactivated_workspace"}})),
)
.await;
assert_eq!(
(summary.invalid, availability, cooldown_until),
(1, AccountAvailability::Banned, None)
);
}
#[tokio::test]
async fn generic_payment_required_rejection_should_exhaust_quota() {
let (summary, availability, cooldown_until) = synchronize_account_from_state(
"acct_quota_payment_required",
AccountAvailability::Ready,
None,
ResponseTemplate::new(402).set_body_json(json!({"detail": {"code": "payment_required"}})),
)
.await;
assert_eq!(
(summary.exhausted, availability, cooldown_until),
(1, AccountAvailability::QuotaExhausted, None)
);
}
async fn synchronize_account_from_state(
account_id: &str,
availability: AccountAvailability,
cooldown_until: Option<SystemTime>,
response: ResponseTemplate,
) -> (
CodexQuotaSyncSummary,
AccountAvailability,
Option<SystemTime>,
) {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/backend-api/wham/usage"))
.respond_with(response)
.expect(1)
.mount(&server)
.await;
async fn synchronize_without_accounts_is_a_noop_before_network_io() {
let store = Arc::new(MemoryAccountStore::default());
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: account_id.to_owned(),
provider_instance_id: instance_id().to_string(),
name: account_id.to_owned(),
secret: secret(&format!("token-{account_id}")),
verified_account: profile(&format!("chatgpt-{account_id}")),
next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
})
.await;
let account = store.account(account_id).expect("quota account");
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability,
reason: Some("previous_state".to_owned()),
cooldown_until,
observed_at: SystemTime::now(),
})
let summary = quota_service(&store)
.synchronize()
.await
.expect("apply initial account state");
let service = quota_service(&store);
let instance = ProviderInstance::new(
instance_id(),
ProviderKind::new("openai").expect("provider"),
format!("{}/backend-api", server.uri()),
true,
InstanceHealth::Healthy,
);
let summary = service
.synchronize_instance(&instance)
.await
.expect("synchronize quota");
let account = store.account(account_id).expect("refreshed quota account");
(summary, account.availability(), account.cooldown_until())
.expect("empty synchronization");
assert_eq!(summary.updated, 0);
assert_eq!(summary.exhausted, 0);
assert_eq!(summary.invalid, 0);
assert_eq!(summary.cooldown, 0);
assert_eq!(summary.transient, 0);
assert_eq!(summary.stale, 0);
}
backend/crates/providers/openai/tests/credential/refresh.rs
+2 -5
View File
@@ -17,7 +17,7 @@ use provider_openai::credential::{
};
use secrecy::{ExposeSecret, SecretString};
use crate::support::{MemoryAccountStore, instance_id, profile, runtime_policy, secret};
use crate::support::{MemoryAccountStore, profile, runtime_policy, secret};
struct Refresher {
outcomes: Mutex<VecDeque<Result<TokenPair, RefreshFailure>>>,
@@ -177,7 +177,7 @@ impl CodexAccountIdentityVerifier for SignedOnlyIdentity {
impl ProviderLeasePort for RefreshLeases {
fn load_state<'a>(
&'a self,
_: &'a gateway_core::routing::ProviderInstanceId,
_: &'a gateway_core::routing::ProviderKind,
_: &'a [gateway_core::engine::credential::ProviderAccountId],
) -> BoxFuture<
'a,
@@ -227,7 +227,6 @@ async fn setup(
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: "acct_refresh".to_owned(),
provider_instance_id: instance_id().to_string(),
name: "refresh".to_owned(),
secret: secret("old-access"),
verified_account: profile("chatgpt-acct_refresh"),
@@ -361,7 +360,6 @@ async fn unavailable_usage_preserves_rotated_tokens_and_persists_backoff() {
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: "acct_signed_only".to_owned(),
provider_instance_id: instance_id().to_string(),
name: "signed only".to_owned(),
secret: secret("old-access"),
verified_account: account_profile,
@@ -535,7 +533,6 @@ async fn malformed_account_refresh_does_not_stop_later_accounts() {
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: account_id.to_owned(),
provider_instance_id: instance_id().to_string(),
name: account_id.to_owned(),
secret: secret(account_id),
verified_account: profile(&format!("chatgpt-{account_id}")),
backend/crates/providers/openai/tests/provider/config.rs
-83
View File
@@ -1,83 +0,0 @@
use gateway_core::routing::{InstanceHealth, ProviderInstance, ProviderInstanceId, ProviderKind};
use provider_openai::{
CodexProviderConfigError, CodexProviderInstanceConfig, CodexProviderTransport,
OFFICIAL_CODEX_BASE_PATH, OfficialCodexOriginPolicy,
};
use crate::support::LoopbackCodexOriginPolicy;
fn instance(provider: &str, base_url: &str) -> ProviderInstance {
ProviderInstance::new(
ProviderInstanceId::new("inst_openai").expect("instance id"),
ProviderKind::new(provider).expect("provider"),
base_url.to_owned(),
true,
InstanceHealth::Healthy,
)
}
#[test]
fn official_instance_requires_exact_https_origin_and_base_path() {
let compiled = CodexProviderInstanceConfig::from_snapshot(
&instance("openai", "https://chatgpt.com/backend-api/"),
&OfficialCodexOriginPolicy,
)
.expect("official endpoint");
assert_eq!(compiled.base_url().path(), OFFICIAL_CODEX_BASE_PATH);
assert_eq!(
compiled.transport(),
CodexProviderTransport::PreferWebSocket
);
}
#[test]
fn official_policy_rejects_lookalike_private_and_metadata_origins() {
for endpoint in [
"https://api.chatgpt.com/backend-api",
"https://chatgpt.com.evil.invalid/backend-api",
"https://chatgpt.com:8443/backend-api",
"https://chatgpt.com/other",
"https://169.254.169.254/backend-api",
"https://10.0.0.1/backend-api",
"http://chatgpt.com/backend-api",
] {
assert_eq!(
CodexProviderInstanceConfig::from_snapshot(
&instance("openai", endpoint),
&OfficialCodexOriginPolicy,
),
Err(CodexProviderConfigError::UnsafeBaseUrl),
"endpoint must fail closed: {endpoint}"
);
}
}
#[test]
fn loopback_policy_requires_explicit_numeric_loopback_host() {
let numeric = instance("openai", "http://127.0.0.1:43123/backend-api");
assert_eq!(
CodexProviderInstanceConfig::from_snapshot(&numeric, &OfficialCodexOriginPolicy),
Err(CodexProviderConfigError::UnsafeBaseUrl)
);
assert!(
CodexProviderInstanceConfig::from_snapshot(&numeric, &LoopbackCodexOriginPolicy).is_ok()
);
assert_eq!(
CodexProviderInstanceConfig::from_snapshot(
&instance("openai", "http://localhost:43123/backend-api"),
&LoopbackCodexOriginPolicy,
),
Err(CodexProviderConfigError::UnsafeBaseUrl)
);
}
#[test]
fn non_codex_instance_is_rejected_before_transport_initialization() {
assert_eq!(
CodexProviderInstanceConfig::from_snapshot(
&instance("xai", "https://chatgpt.com/backend-api"),
&OfficialCodexOriginPolicy,
),
Err(CodexProviderConfigError::ProviderMismatch)
);
}
backend/crates/providers/openai/tests/provider/contract.rs
+192 -2520
View File
File diff suppressed because it is too large Load Diff
backend/crates/providers/openai/tests/provider/mod.rs
-1
View File
@@ -1,2 +1 @@
mod config;
mod contract;
backend/crates/providers/openai/tests/support.rs
+5 -28
View File
@@ -20,14 +20,12 @@ use gateway_core::provider_ports::{
ProviderRuntimePolicyPort, ProviderSchedulingLeaseRequest, ProviderSchedulingState,
ProviderStoreError,
};
use gateway_core::routing::{ProviderInstanceId, ProviderKind};
use provider_openai::CodexOriginPolicy;
use gateway_core::routing::ProviderKind;
use provider_openai::credential::{
CodexAccountProfile, CodexCredentialAdmin, CodexCredentialRepository, CodexOAuthSecret,
ImportCodexOAuthCredential,
};
use secrecy::SecretString;
use url::{Host, Url};
#[derive(Clone)]
struct StoredAccount {
@@ -110,16 +108,16 @@ impl ProviderAccountStore for MemoryAccountStore {
.collect())
}
async fn list_for_instance(
async fn list_for_provider(
&self,
instance: &ProviderInstanceId,
provider: &ProviderKind,
) -> Result<Vec<ProviderAccount>, StoreError> {
Ok(self
.accounts
.lock()
.expect("account store lock")
.values()
.filter(|stored| stored.account.instance() == instance)
.filter(|stored| stored.account.provider() == provider)
.map(|stored| stored.account.clone())
.collect())
}
@@ -294,21 +292,6 @@ const fn store_error(kind: StoreErrorKind) -> StoreError {
StoreError::new(kind)
}
#[derive(Debug)]
pub(crate) struct LoopbackCodexOriginPolicy;
impl CodexOriginPolicy for LoopbackCodexOriginPolicy {
fn allows(&self, url: &Url) -> bool {
url.scheme() == "http"
&& matches!(url.host(), Some(Host::Ipv4(host)) if host.is_loopback())
&& url.port().is_some()
}
}
pub(crate) fn loopback_origin_policy() -> Arc<dyn CodexOriginPolicy> {
Arc::new(LoopbackCodexOriginPolicy)
}
struct AccountRebuild {
revision: CredentialRevision,
enabled: bool,
@@ -330,7 +313,6 @@ fn rebuild_account(current: &ProviderAccount, rebuild: AccountRebuild) -> Provid
});
ProviderAccount::new(
current.id().clone(),
current.instance().clone(),
current.provider().clone(),
name,
current.upstream_user_id().to_owned(),
@@ -360,7 +342,7 @@ pub(crate) struct TestLeaseCoordinator {
impl ProviderLeasePort for TestLeaseCoordinator {
fn load_state<'a>(
&'a self,
_provider_instance_id: &'a ProviderInstanceId,
_provider_kind: &'a ProviderKind,
accounts: &'a [ProviderAccountId],
) -> BoxFuture<'a, Result<ProviderSchedulingState, ProviderStoreError>> {
Box::pin(async move {
@@ -460,14 +442,9 @@ pub(crate) fn account_policy() -> gateway_core::engine::credential::AccountSelec
)
}
pub(crate) fn instance_id() -> ProviderInstanceId {
ProviderInstanceId::new("inst_openai_primary").expect("instance id")
}
pub(crate) fn codex_account(id: &str) -> ProviderAccount {
ProviderAccount::new(
ProviderAccountId::new(id).expect("account id"),
instance_id(),
ProviderKind::new("openai").expect("provider"),
id.to_owned(),
format!("user-chatgpt-{id}"),
backend/crates/providers/xai/src/admin.rs
+10 -60
View File
@@ -42,7 +42,7 @@ use gateway_core::provider_ports::{
OAuthPendingTakeOutcome, ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderStoreError,
ProviderStoreErrorKind,
};
use gateway_core::routing::{ProviderInstanceId, ProviderKind, UpstreamModelId};
use gateway_core::routing::{ProviderKind, UpstreamModelId};
use serde::Deserialize;
use serde_json::{Map, Number, Value};
use sha2::{Digest as _, Sha256};
@@ -62,7 +62,7 @@ use crate::credential::{
};
use crate::transport::{GROK_CLI_BASE_URL, XAI_PROVIDER_NAME, grok_billing_breakdown};
const PENDING_SCHEMA_VERSION: u64 = 1;
const PENDING_SCHEMA_VERSION: u64 = 2;
const PENDING_TTL: TimeDelta = TimeDelta::minutes(30);
const MAX_PENDING_TEXT_BYTES: usize = 512;
@@ -325,7 +325,6 @@ impl ProviderAdmin for XaiAdminProvider {
Uuid::now_v7().simple()
))
.map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?,
provider_instance_id: command.provider_instance_id.clone(),
name,
email,
upstream_account_id: None,
@@ -340,7 +339,6 @@ impl ProviderAdmin for XaiAdminProvider {
}
Ok(PreparedCredentialImport {
provider_kind: self.provider_kind.clone(),
provider_instance_id: command.provider_instance_id,
credentials,
})
}
@@ -353,7 +351,6 @@ impl ProviderAdmin for XaiAdminProvider {
return Err(provider_error(ProviderAdminErrorKind::Invalid));
}
if let AuthorizationMutationTarget::Reauthorize {
provider_instance_id,
account_id,
expected_credential_revision,
} = pending.target()
@@ -365,9 +362,7 @@ impl ProviderAdmin for XaiAdminProvider {
.load_credential(account_id, revision)
.await
.map_err(map_store_error)?;
if current.account.provider() != &self.provider_kind
|| current.account.instance() != provider_instance_id
{
if current.account.provider() != &self.provider_kind {
return Err(provider_error(ProviderAdminErrorKind::NotFound));
}
}
@@ -410,10 +405,7 @@ impl ProviderAdmin for XaiAdminProvider {
.await
.map_err(map_provider_store_error)?;
let credential = match stored.mutation.target() {
AuthorizationMutationTarget::Create {
provider_instance_id,
name,
} => {
AuthorizationMutationTarget::Create { name } => {
let prepared = GrokCredentialAdmin
.prepare_verified_account(
&VerifiedGrokAccount {
@@ -422,7 +414,6 @@ impl ProviderAdmin for XaiAdminProvider {
Uuid::now_v7().simple()
))
.map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?,
provider_instance_id: provider_instance_id.clone(),
name: name.clone(),
email: None,
upstream_account_id: None,
@@ -436,7 +427,6 @@ impl ProviderAdmin for XaiAdminProvider {
PreparedAuthorizationCredential::Create(prepared_create(prepared, Utc::now())?)
}
AuthorizationMutationTarget::Reauthorize {
provider_instance_id,
account_id,
expected_credential_revision,
} => {
@@ -447,15 +437,12 @@ impl ProviderAdmin for XaiAdminProvider {
.load_credential(account_id, revision)
.await
.map_err(map_store_error)?;
if current.account.provider() != &self.provider_kind
|| current.account.instance() != provider_instance_id
{
if current.account.provider() != &self.provider_kind {
return Err(provider_error(ProviderAdminErrorKind::NotFound));
}
let prepared = verified_rotation(current, tokens, policy)?;
PreparedAuthorizationCredential::Reauthorize(prepared_rotation(
prepared,
provider_instance_id.clone(),
self.provider_kind.clone(),
)?)
}
@@ -524,11 +511,7 @@ impl ProviderAdmin for XaiAdminProvider {
.await
.map_err(map_provider_store_error)?;
let prepared = verified_rotation(current, tokens, policy)?;
prepared_rotation(
prepared,
command.account.provider_instance_id,
command.account.provider_kind,
)
prepared_rotation(prepared, command.account.provider_kind)
}
async fn prepare_refresh(
@@ -549,11 +532,7 @@ impl ProviderAdmin for XaiAdminProvider {
.prepare_manual_refresh(&account_id, revision)
.await
.map_err(map_refresh_error)?;
prepared_rotation(
prepared,
command.account.provider_instance_id,
command.account.provider_kind,
)
prepared_rotation(prepared, command.account.provider_kind)
}
async fn quota(
@@ -735,7 +714,6 @@ fn prepared_create(
} = prepared;
Ok(PreparedCredentialCreate {
account_id: account.id().clone(),
provider_instance_id: account.instance().clone(),
provider_kind: account.provider().clone(),
name: account.name().to_owned(),
email: account.email().map(str::to_owned),
@@ -756,7 +734,6 @@ fn prepared_create(
fn prepared_rotation(
prepared: PreparedGrokCredentialRotation,
provider_instance_id: ProviderInstanceId,
provider_kind: ProviderKind,
) -> Result<PreparedCredentialRotation, ProviderAdminError> {
let (profile, credential, guard) = prepared.into_parts();
@@ -775,7 +752,6 @@ fn prepared_rotation(
Ok(PreparedCredentialRotation::new(
PreparedCredentialRotationFacts {
account_id,
provider_instance_id,
provider_kind,
expected_credential_revision: Revision::new(expected_revision.get())
.map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?,
@@ -812,7 +788,6 @@ fn account_from_record(account: &AccountRecord) -> Result<ProviderAccount, Provi
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?;
Ok(ProviderAccount::new(
account_id,
account.provider_instance_id.clone(),
account.provider_kind.clone(),
account.name.clone(),
account.upstream_user_id.clone(),
@@ -837,7 +812,6 @@ fn account_from_record(account: &AccountRecord) -> Result<ProviderAccount, Provi
fn account_matches_record(account: &ProviderAccount, record: &AccountRecord) -> bool {
account.id().as_str() == record.id
&& account.instance() == &record.provider_instance_id
&& account.provider() == &record.provider_kind
&& account.revision().get() == record.credential_revision.get()
&& account.name() == record.name
@@ -1029,11 +1003,9 @@ struct PendingMutationDocument {
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
enum PendingTargetDocument {
Create {
provider_instance_id: String,
name: String,
},
Reauthorize {
provider_instance_id: String,
account_id: String,
expected_credential_revision: u64,
},
@@ -1101,27 +1073,15 @@ fn encode_mutation(mutation: &PendingAuthorizationMutation) -> Map<String, Value
fn encode_target(target: &AuthorizationMutationTarget) -> Map<String, Value> {
let mut document = Map::new();
match target {
AuthorizationMutationTarget::Create {
provider_instance_id,
name,
} => {
AuthorizationMutationTarget::Create { name } => {
document.insert("kind".to_owned(), Value::String("create".to_owned()));
document.insert(
"provider_instance_id".to_owned(),
Value::String(provider_instance_id.to_string()),
);
document.insert("name".to_owned(), Value::String(name.clone()));
}
AuthorizationMutationTarget::Reauthorize {
provider_instance_id,
account_id,
expected_credential_revision,
} => {
document.insert("kind".to_owned(), Value::String("reauthorize".to_owned()));
document.insert(
"provider_instance_id".to_owned(),
Value::String(provider_instance_id.to_string()),
);
document.insert(
"account_id".to_owned(),
Value::String(account_id.to_string()),
@@ -1194,21 +1154,11 @@ fn decode_mutation(
let provider_kind = ProviderKind::new(document.provider_kind)
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?;
let target = match document.target {
PendingTargetDocument::Create {
provider_instance_id,
name,
} => AuthorizationMutationTarget::Create {
provider_instance_id: ProviderInstanceId::new(provider_instance_id)
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?,
name,
},
PendingTargetDocument::Create { name } => AuthorizationMutationTarget::Create { name },
PendingTargetDocument::Reauthorize {
provider_instance_id,
account_id,
expected_credential_revision,
} => AuthorizationMutationTarget::Reauthorize {
provider_instance_id: ProviderInstanceId::new(provider_instance_id)
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?,
account_id: ProviderAccountId::new(account_id)
.map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?,
expected_credential_revision: Revision::new(expected_credential_revision)
@@ -1410,7 +1360,7 @@ fn map_quota_error(error: GrokQuotaError) -> ProviderAdminError {
fn map_catalog_error(error: GrokCredentialCatalogError) -> ProviderAdminError {
use GrokCredentialCatalogError as Error;
provider_error(match error {
Error::InvalidInstance | Error::InvalidCredentialData | Error::ConflictingModelFacts => {
Error::InvalidCredentialData | Error::ConflictingModelFacts => {
ProviderAdminErrorKind::Invalid
}
Error::NoEligibleCredential => ProviderAdminErrorKind::NotFound,
backend/crates/providers/xai/src/credential/catalog.rs
+18 -44
View File
@@ -14,7 +14,7 @@ use gateway_core::engine::credential::{
};
use gateway_core::engine::provider::ProviderCatalogGeneration;
use gateway_core::provider_ports::{ProviderCatalogCacheKey, ProviderCatalogCachePort};
use gateway_core::routing::{ConfigRevision, ProviderInstance, ProviderKind};
use gateway_core::routing::{ConfigRevision, ProviderKind};
use tokio::sync::Mutex;
use super::repository::{GrokCredentialRepository, LoadedGrokCredential};
@@ -22,9 +22,9 @@ use super::types::{GrokCredentialAvailability, UpdateGrokCredentialState};
use crate::XaiWireProfileState;
use crate::transport::catalog::{MAX_CATALOG_MODELS, valid_model_slug, validate_etag};
use crate::{
GrokBillingClient, GrokBillingTransport, GrokBuildProvider, GrokCatalogModel,
GrokModelCatalogClient, GrokModelCatalogSession, GrokModelCatalogSnapshot,
GrokModelCatalogTransport, SecretValue, parse_grok_billing,
GrokBillingClient, GrokBillingTransport, GrokCatalogModel, GrokModelCatalogClient,
GrokModelCatalogSession, GrokModelCatalogSnapshot, GrokModelCatalogTransport, SecretValue,
parse_grok_billing,
};
const MAX_CONCURRENT_CATALOG_REQUESTS: usize = 8;
@@ -805,11 +805,10 @@ pub enum GrokCatalogCacheError {
InvalidData,
}
/// 一次 instance 同步得到的账号目录和严格模型并集。
/// 一次 Provider 同步得到的账号目录和严格模型并集。
#[derive(Clone, Debug)]
pub struct GrokCredentialCatalogSnapshot {
config_revision: ConfigRevision,
provider_instance_id: String,
observed_at: DateTime<Utc>,
accounts: Vec<GrokAccountCatalog>,
models: Vec<GrokCatalogModel>,
@@ -821,11 +820,6 @@ impl GrokCredentialCatalogSnapshot {
self.config_revision
}
#[must_use]
pub fn provider_instance_id(&self) -> &str {
&self.provider_instance_id
}
#[must_use]
pub const fn observed_at(&self) -> DateTime<Utc> {
self.observed_at
@@ -844,8 +838,6 @@ impl GrokCredentialCatalogSnapshot {
#[derive(Debug, thiserror::Error)]
pub enum GrokCredentialCatalogError {
#[error("Grok model catalog instance configuration is invalid")]
InvalidInstance,
#[error("Grok model catalog has no eligible OAuth account")]
NoEligibleCredential,
#[error("Grok model catalog credential snapshot is stale")]
@@ -874,7 +866,7 @@ pub struct GrokCredentialCatalogService {
#[derive(Default)]
struct PublishedCatalogState {
generation: u64,
models: BTreeMap<String, Vec<GrokCatalogModel>>,
models: Vec<GrokCatalogModel>,
}
impl GrokCredentialCatalogService {
@@ -982,15 +974,13 @@ impl GrokCredentialCatalogService {
}
/// 实时拉取全部 eligible OAuth account；结果只写可重建 TTL cache，不写 PostgreSQL。
pub async fn synchronize_instance(
pub async fn synchronize(
&self,
instance: &ProviderInstance,
config_revision: ConfigRevision,
) -> Result<GrokCredentialCatalogSnapshot, GrokCredentialCatalogError> {
let catalog = self.fetch_and_cache(instance).await?;
let catalog = self.fetch_and_cache().await?;
Ok(GrokCredentialCatalogSnapshot {
config_revision,
provider_instance_id: catalog.provider_instance_id,
observed_at: catalog.observed_at,
accounts: catalog.accounts,
models: catalog.models,
@@ -998,22 +988,14 @@ impl GrokCredentialCatalogService {
}
/// Provider Registry 构建 RuntimeSnapshot 时使用的实时能力目录。
pub async fn query_instance_models(
&self,
instance: &ProviderInstance,
) -> Result<Vec<GrokCatalogModel>, GrokCredentialCatalogError> {
Ok(self.fetch_and_cache(instance).await?.models)
pub async fn query_models(&self) -> Result<Vec<GrokCatalogModel>, GrokCredentialCatalogError> {
Ok(self.fetch_and_cache().await?.models)
}
async fn fetch_and_cache(
&self,
instance: &ProviderInstance,
) -> Result<FetchedInstanceCatalog, GrokCredentialCatalogError> {
let instance_config = GrokBuildProvider::validate_instance(instance)
.map_err(|_| GrokCredentialCatalogError::InvalidInstance)?;
async fn fetch_and_cache(&self) -> Result<FetchedProviderCatalog, GrokCredentialCatalogError> {
let candidates = self
.repository
.list_loaded_for_instance(instance_config.id())
.list_loaded_for_provider()
.await
.map_err(|_| GrokCredentialCatalogError::Store)?
.into_iter()
@@ -1050,31 +1032,24 @@ impl GrokCredentialCatalogService {
.await
.map_err(|_| GrokCredentialCatalogError::Cache)?;
}
self.publish_models(instance_config.id().as_str(), &models);
self.publish_models(&models);
Ok(FetchedInstanceCatalog {
provider_instance_id: instance_config.id().as_str().to_owned(),
Ok(FetchedProviderCatalog {
observed_at,
accounts,
models,
})
}
fn publish_models(&self, instance: &str, models: &[GrokCatalogModel]) {
fn publish_models(&self, models: &[GrokCatalogModel]) {
let mut published = self
.published
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if published
.models
.get(instance)
.is_some_and(|current| current == models)
{
if published.models == models {
return;
}
published
.models
.insert(instance.to_owned(), models.to_vec());
published.models = models.to_vec();
published.generation = published.generation.saturating_add(1);
}
}
@@ -1097,8 +1072,7 @@ struct FetchedCredentialCatalog {
seed: GrokCredentialCatalogSeed,
}
struct FetchedInstanceCatalog {
provider_instance_id: String,
struct FetchedProviderCatalog {
observed_at: DateTime<Utc>,
accounts: Vec<GrokAccountCatalog>,
models: Vec<GrokCatalogModel>,
backend/crates/providers/xai/src/credential/refresh.rs
-5
View File
@@ -16,7 +16,6 @@ use gateway_core::provider_ports::{
ProviderRefreshCapacityRequest, ProviderRefreshLeaseRequest, ProviderRefreshPolicy,
ProviderRuntimePolicyPort, ProviderStoreError, provider_refresh_retry_at,
};
use gateway_core::routing::ProviderInstanceId;
use super::catalog::GrokCredentialCatalogService;
use super::repository::{
@@ -41,7 +40,6 @@ const REFRESH_RETRY_DELAY: Duration = Duration::from_secs(10 * 60);
/// 一个到期且已按 revision 读取明文 RT 的 xAI account。
pub struct DueGrokCredential {
account_id: ProviderAccountId,
provider_instance_id: ProviderInstanceId,
credential_revision: CredentialRevision,
refresh_token: SecretValue,
id_token: Option<SecretValue>,
@@ -75,7 +73,6 @@ impl fmt::Debug for DueGrokCredential {
formatter
.debug_struct("DueGrokCredential")
.field("account_id", &self.account_id)
.field("provider_instance_id", &self.provider_instance_id)
.field("credential_revision", &self.credential_revision)
.field("refresh_token", &"[REDACTED]")
.field("id_token", &self.id_token.as_ref().map(|_| "[REDACTED]"))
@@ -327,7 +324,6 @@ impl GrokCredentialRefreshService {
let policy = self.runtime_policy.load_refresh_policy().await?;
let credential = DueGrokCredential {
account_id: account_id.clone(),
provider_instance_id: loaded.account.instance().clone(),
credential_revision,
refresh_token: loaded.refresh_token,
id_token: loaded.id_token,
@@ -801,7 +797,6 @@ impl GrokCredentialRepository {
};
due.push(DueGrokCredential {
account_id: account.id().clone(),
provider_instance_id: account.instance().clone(),
credential_revision: account.revision(),
refresh_token: loaded.refresh_token,
id_token: loaded.id_token,
backend/crates/providers/xai/src/credential/repository.rs
+10 -13
View File
@@ -13,7 +13,7 @@ use gateway_core::engine::credential::{
};
use gateway_core::error::StoreErrorKind;
use gateway_core::provider_ports::ProviderRefreshPolicy;
use gateway_core::routing::{ProviderInstanceId, ProviderKind};
use gateway_core::routing::ProviderKind;
use serde_json::Value;
use thiserror::Error;
@@ -36,7 +36,6 @@ const MAX_IMPORT_BATCH: usize = 200;
/// 已跨过官方 OAuth/OIDC 验证边界、等待转为 Core 创建命令的账号。
pub struct VerifiedGrokAccount {
pub account_id: ProviderAccountId,
pub provider_instance_id: ProviderInstanceId,
pub name: String,
pub email: Option<String>,
pub upstream_account_id: Option<String>,
@@ -50,7 +49,6 @@ impl std::fmt::Debug for VerifiedGrokAccount {
formatter
.debug_struct("VerifiedGrokAccount")
.field("account_id", &self.account_id)
.field("provider_instance_id", &self.provider_instance_id)
.field("name", &self.name)
.field("email", &self.email.as_ref().map(|_| "[REDACTED]"))
.field(
@@ -157,7 +155,6 @@ impl GrokCredentialAdmin {
.map_err(|_| GrokCredentialRepositoryError::InvalidInput("provider_kind"))?;
let account = ProviderAccount::new(
input.account_id.clone(),
input.provider_instance_id.clone(),
provider,
input.name.clone(),
input.account.subject.clone(),
@@ -204,7 +201,6 @@ impl GrokCredentialAdmin {
.map_err(|_| GrokCredentialRepositoryError::InvalidInput("next_refresh_at"))?;
self.prepare_import(&CreateGrokCredential {
account_id: input.account_id.clone(),
provider_instance_id: input.provider_instance_id.clone(),
name: input.name.clone(),
secret: GrokOAuthSecret {
access_token: input.tokens.access_token().clone(),
@@ -379,7 +375,6 @@ impl GrokCredentialRepository {
Ok(GrokCredentialRecord {
account_id: input.account_id.clone(),
provider_instance_id: current.account.instance().clone(),
credential_revision: revision,
})
}
@@ -476,14 +471,15 @@ impl GrokCredentialRepository {
Ok(loaded)
}
/// 读取 instance 下的全部 xAI account；credential 逐行按 revision fence 加载。
pub(crate) async fn list_loaded_for_instance(
/// 读取 xAI Provider 的全部账号；credential 逐行按 revision fence 加载。
pub(crate) async fn list_loaded_for_provider(
&self,
instance: &ProviderInstanceId,
) -> Result<Vec<LoadedGrokCredential>, GrokCredentialRepositoryError> {
let provider = ProviderKind::new(XAI_PROVIDER_KIND)
.map_err(|_| GrokCredentialRepositoryError::InvalidCredentialData)?;
let accounts = self
.store
.list_for_instance(instance)
.list_for_provider(&provider)
.await
.map_err(map_store_error)?;
let mut loaded = Vec::with_capacity(accounts.len());
@@ -495,13 +491,14 @@ impl GrokCredentialRepository {
Ok(loaded)
}
pub(crate) async fn list_accounts_for_instance(
pub(crate) async fn list_accounts_for_provider(
&self,
instance: &ProviderInstanceId,
) -> Result<Vec<ProviderAccount>, GrokCredentialRepositoryError> {
let provider = ProviderKind::new(XAI_PROVIDER_KIND)
.map_err(|_| GrokCredentialRepositoryError::InvalidCredentialData)?;
let accounts = self
.store
.list_for_instance(instance)
.list_for_provider(&provider)
.await
.map_err(map_store_error)?;
for account in &accounts {

Some files were not shown because too many files have changed in this diff Show More