feat(peer): add CPR Link discovery and scoped quota sharing (#566)

* feat(peer): add automatic discovery and scalable quota sharing

Export authorized direct-account pools and budgets through the CPR peer
protocol, with durable upstream discovery and immutable scheduling views.

Show shared-pool node quotas through common quota popovers and lazy account
details. Bound snapshot reads and export concurrency, reduce Redis and SQL
work, and keep only the latest observation for each upstream connection.

Co-authored-by: Codex <noreply@openai.com>

* refactor(peer): centralize upstream discovery and quota projections

Move peer synchronization into Admin, HTTP transport into Host, and expose
neutral Core scheduling signals through explicitly wired provider capabilities.

Preserve provider-visible model catalogs and independent quota disclosure,
with stale windows following the existing OAuth quota layout.

Cover cross-provider discovery, lifecycle fencing, and database leases, and
update the architecture and usage documentation.

Co-authored-by: Codex <noreply@openai.com>

* fix(ui): center usage card loading and empty states

Co-authored-by: Codex <noreply@openai.com>

* fix(accounts): unify peer subscription badges across account views

Share CPR identity rendering through AccountPlanBadge, project upstream metadata into dashboard and proxy account views, and align quota empty states and API badges.

Co-authored-by: Codex <noreply@openai.com>

* fix(peer): bound snapshot work and preserve quota windows

Keep summary reads off full snapshots, reject oversized model expansion early, and yield export work. Preserve daily and weekly unlimited budgets, hide empty token counters, and unify quota empty states.

Co-authored-by: Codex <noreply@openai.com>

---------

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-10-31 15:00:12 +08:00
committed by GitHub
co-authored by Codex Zyy
parent 35e8b73075
commit 58a1f7541e
181 changed files with 10201 additions and 585 deletions
+3
View File
@@ -1705,6 +1705,7 @@ dependencies = [
"cron",
"futures",
"gateway-core",
"gateway-protocol",
"heck",
"hex",
"hkdf",
@@ -1802,6 +1803,7 @@ dependencies = [
"futures",
"gateway-admin",
"gateway-core",
"gateway-protocol",
"hex",
"http",
"http-body-util",
@@ -3323,6 +3325,7 @@ dependencies = [
"futures",
"gateway-admin",
"gateway-core",
"gateway-host",
"gateway-protocol",
"getrandom 0.4.3",
"h2",
@@ -100,7 +100,12 @@ async fn execute(
plugin_runtime,
plugin_runtime.bind_snapshots(&core.snapshots())
);
let management = ManagementPorts::new(&store, providers.admin.clone(), core.snapshot_control());
let management = ManagementPorts::new(
&store,
providers.admin.clone(),
core.snapshot_control(),
providers.peer.service(),
);
or_shutdown!(plugin_runtime, management.bind(&plugin_runtime));
let core = or_shutdown!(plugin_runtime, core.activate().await);
or_shutdown!(
@@ -132,6 +137,7 @@ async fn execute(
providers.admin.clone(),
core.snapshot_control(),
core.account_probe(),
Some(providers.peer.service()),
);
or_shutdown!(plugin_runtime, services.register_accounts(&accounts));
let usage_data = gateway_admin::initialize_usage_data(
@@ -42,6 +42,7 @@ impl ManagementPorts {
store: &gateway_store::StoreBundle,
providers: ProviderAdminRegistry,
snapshot: Arc<dyn SnapshotControl>,
upstream_info: Arc<dyn gateway_admin::ports::peer::PeerUpstreamService>,
) -> Self {
let ports = store.admin_ports();
Self {
@@ -49,6 +50,7 @@ impl ManagementPorts {
providers.clone(),
ports.accounts(),
snapshot.clone(),
Some(upstream_info),
),
keys: gateway_admin::initialize_plugin_client_keys(
providers,
+8 -1
View File
@@ -35,7 +35,12 @@ pub async fn run() -> Result<(), BootstrapError> {
plugin_runtime,
plugin_runtime.bind_snapshots(&core.snapshots())
);
let management = ManagementPorts::new(&store, providers.admin.clone(), core.snapshot_control());
let management = ManagementPorts::new(
&store,
providers.admin.clone(),
core.snapshot_control(),
providers.peer.service(),
);
or_shutdown!(plugin_runtime, management.bind(&plugin_runtime));
let mut core = or_shutdown!(plugin_runtime, core.activate().await);
or_shutdown!(
@@ -77,6 +82,7 @@ pub async fn run() -> Result<(), BootstrapError> {
config.client,
store.admin_ports(),
gateway_admin::AdminRuntimePorts {
upstream_info: Some(providers.peer.service()),
timezone,
service_middleware: service_middleware(core.snapshots(), &plugin_runtime),
providers: providers.admin,
@@ -165,6 +171,7 @@ pub async fn run() -> Result<(), BootstrapError> {
plan.extend(providers.openai.take_worker_contributions());
plan.extend(providers.xai.take_worker_contributions());
plan.extend(admin.take_worker_contributions());
plan.extend(providers.peer.take_worker_contributions());
plan.push(or_shutdown!(
plugin_runtime,
plugin_runtime.maintenance_worker(core.snapshots())
+21 -3
View File
@@ -67,6 +67,7 @@ pub(super) struct Providers {
pub(super) xai: provider_xai::ProviderBundle,
pub(super) core: ProviderRegistry,
pub(super) admin: ProviderAdminRegistry,
pub(super) peer: gateway_admin::PeerUpstreamBundle,
}
impl Providers {
@@ -75,16 +76,33 @@ impl Providers {
environment: &Environment,
) -> Result<Self, BootstrapError> {
let ports = environment.store.provider_ports();
let openai =
provider_openai::initialize(config.with_timezone(environment.timezone), ports.clone())
.await?;
let upstream_reader = Arc::new(gateway_core::account::UpstreamInfoReader::default());
let openai = provider_openai::initialize(
config.with_timezone(environment.timezone),
ports.clone(),
upstream_reader.clone(),
)
.await?;
environment.host.report_startup_ready("OpenAI Provider");
let xai = provider_xai::initialize(ports).await?;
environment.host.report_startup_ready("xAI Provider");
let core = ProviderRegistry::new([openai.core_provider(), xai.core_provider()])?;
let admin = ProviderAdminRegistry::new([openai.admin_provider(), xai.admin_provider()])
.map_err(|_| gateway_admin::model::AdminError::invalid("Provider 管理身份冲突"))?;
let peer = gateway_admin::initialize_peer_upstream(
environment.store.provider_ports().accounts(),
environment.store.upstream_info(),
&admin,
Arc::new(gateway_host::peer::HttpPeerTransport::with_client_builder(
provider_openai::build_reqwest_client_with_custom_ca,
)),
)
.map_err(|_| gateway_admin::model::AdminError::invalid("互联后台任务配置无效"))?;
upstream_reader
.bind(&peer.scheduling_source())
.map_err(|_| gateway_admin::model::AdminError::invalid("互联观测端口重复绑定"))?;
Ok(Self {
peer,
openai,
xai,
core,
@@ -321,6 +321,7 @@ const ADAPTER_PUBLIC_MODULES: &[(&str, &[&str])] = &[
"client_distribution",
"config",
"outbound",
"peer",
"official_plugins",
"pricing",
"plugin_distribution",
@@ -358,11 +359,13 @@ const ALLOWED_INTERNAL_EDGES: &[(&str, &str)] = &[
("codex-proxy-rs", "provider-openai"),
("codex-proxy-rs", "provider-xai"),
("gateway-admin", "gateway-core"),
("gateway-admin", "gateway-protocol"),
("gateway-api", "gateway-admin"),
("gateway-api", "gateway-core"),
("gateway-api", "gateway-protocol"),
("gateway-host", "gateway-admin"),
("gateway-host", "gateway-core"),
("gateway-host", "gateway-protocol"),
("gateway-plugin-runtime", "gateway-host"),
("gateway-plugin-runtime", "gateway-core"),
("gateway-plugin-runtime", "gateway-admin"),
+1
View File
@@ -19,6 +19,7 @@ bytes = "1.11.1"
chrono = { version = "0.4.45", features = ["serde"] }
futures = "0.3.32"
gateway-core = { path = "../gateway-core" }
gateway-protocol = { path = "../gateway-protocol" }
heck = "0.5.0"
hex = "0.4.3"
hkdf = "0.13.0"
+40 -2
View File
@@ -26,6 +26,7 @@ mod use_case;
pub use use_case::plugins::{PluginDistributionPorts, PluginManagementService, PluginsService};
pub use use_case::key_usage::KeyUsageService;
pub use use_case::peer_link::{PeerLinkError, PeerLinkService};
pub use use_case::{
account_groups::AccountGroupService,
@@ -213,6 +214,7 @@ pub struct AdminServices {
proxies: Arc<dyn ProxiesService>,
auth: Arc<dyn AuthService>,
key_usage: Arc<dyn KeyUsageService>,
peer_link: Arc<dyn PeerLinkService>,
accounts: Arc<dyn AccountsService>,
account_groups: Arc<dyn AccountGroupService>,
client_keys: Arc<dyn ClientKeyService>,
@@ -258,6 +260,11 @@ impl AdminServices {
self.key_usage.as_ref()
}
#[must_use]
pub fn peer_link(&self) -> &dyn PeerLinkService {
self.peer_link.as_ref()
}
#[must_use]
pub fn proxies(&self) -> &dyn ProxiesService {
self.proxies.as_ref()
@@ -336,8 +343,13 @@ impl AdminBundle {
}
}
pub use use_case::peer_link::upstream::{
PeerUpstreamBundle, initialize as initialize_peer_upstream,
};
/// 组合根提供给控制面的运行能力;与配置和存储端口分别传入
pub struct AdminRuntimePorts {
pub upstream_info: Option<Arc<dyn ports::peer::PeerUpstreamService>>,
pub timezone: gateway_core::time::DeploymentTimeZone,
pub service_middleware: public_service::PlanSource,
pub plugin_preparation: Arc<dyn ports::plugins::PluginPreparation>,
@@ -388,6 +400,7 @@ async fn initialize_inner(
plugin_accounts: Option<Arc<dyn PluginAccountAccess>>,
) -> Result<AdminBundle, AdminError> {
let AdminRuntimePorts {
upstream_info,
timezone,
service_middleware,
plugin_preparation,
@@ -429,6 +442,7 @@ async fn initialize_inner(
registry.clone(),
snapshot.clone(),
probe.clone(),
upstream_info.clone(),
);
let backup_ports = store.backup();
let backups = Arc::new(DefaultBackupService::new(
@@ -465,9 +479,15 @@ async fn initialize_inner(
store.accounts(),
store.proxies(),
snapshot.clone(),
upstream_info.clone(),
));
let plugin_accounts = plugin_accounts.unwrap_or_else(|| {
initialize_plugin_accounts(registry.clone(), store.accounts(), snapshot.clone())
initialize_plugin_accounts(
registry.clone(),
store.accounts(),
snapshot.clone(),
upstream_info.clone(),
)
});
let import_tasks = use_case::import_tasks::DefaultImportTasksService::new(credentials.clone());
let import_task = use_case::import_tasks::ImportTaskWorker(import_tasks.clone());
@@ -483,6 +503,12 @@ async fn initialize_inner(
public_services.register_accounts(&accounts)?;
public_services.register_usage_data(&usage_data)?;
let services = AdminServices {
peer_link: Arc::new(use_case::peer_link::DefaultPeerLinkService::new(
store.client_keys(),
store.accounts(),
store.account_runtime(),
registry.clone(),
)),
timezone,
public_services: Arc::new(public_services),
plugin_management: Arc::new(PluginManagementService::new(
@@ -505,6 +531,7 @@ async fn initialize_inner(
proxy_probe,
snapshot.clone(),
registry.clone(),
upstream_info.clone(),
)),
auth,
accounts: accounts.clone(),
@@ -526,6 +553,7 @@ async fn initialize_inner(
registry.clone(),
timezone,
system.clone(),
upstream_info,
)),
settings,
usage_data,
@@ -595,9 +623,13 @@ pub fn initialize_plugin_accounts(
providers: ports::provider::ProviderAdminRegistry,
accounts: Arc<dyn ports::store::AccountStore>,
snapshot: Arc<dyn gateway_core::runtime::SnapshotControl>,
upstream_info: Option<Arc<dyn ports::peer::PeerUpstreamService>>,
) -> Arc<dyn PluginAccountAccess> {
Arc::new(use_case::plugin_accounts::DefaultPluginAccountAccess::new(
providers, accounts, snapshot,
providers,
accounts,
snapshot,
upstream_info,
))
}
@@ -674,6 +706,10 @@ fn freeze_recovery_worker_contribution(
/// 服务端与 CLI 复用账号用例,不随命令启动后台调度
#[must_use]
#[expect(
clippy::too_many_arguments,
reason = "账号存储、发布、探测与互联分别通过独立端口注入"
)]
pub fn initialize_accounts(
accounts: Arc<dyn ports::store::AccountStore>,
account_runtime: Arc<dyn ports::store::AccountRuntimeStore>,
@@ -682,6 +718,7 @@ pub fn initialize_accounts(
providers: ProviderAdminRegistry,
snapshot: Arc<dyn SnapshotControl>,
probe: Arc<dyn AccountProbe>,
upstream_info: Option<Arc<dyn ports::peer::PeerUpstreamService>>,
) -> Arc<dyn AccountsService> {
Arc::new(DefaultAccountsService::new(
accounts,
@@ -691,6 +728,7 @@ pub fn initialize_accounts(
providers,
snapshot,
probe,
upstream_info,
))
}
@@ -108,6 +108,7 @@ pub struct ClientKeyListQuery {
/// 不含完整明文 Key 的管理投影
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientKeyRecord {
pub peer_sharing: gateway_core::policy::ClientPeerSharing,
pub request_profile_overrides: ProviderRequestProfileOverrides,
pub id: ClientApiKeyId,
pub name: String,
@@ -167,6 +168,7 @@ impl fmt::Debug for ClientKeySecret {
/// API 提交的 Client Key 创建命令
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CreateClientKey {
pub peer_sharing: gateway_core::policy::ClientPeerSharing,
pub request_profile_overrides: ProviderRequestProfileOverrides,
pub custom_key: Option<PlaintextClientApiKey>,
pub name: String,
@@ -179,6 +181,7 @@ pub struct CreateClientKey {
/// 管理用例生成 ID 与明文后的持久化命令
#[derive(Clone, PartialEq, Eq)]
pub struct NewClientKey {
pub peer_sharing: gateway_core::policy::ClientPeerSharing,
pub request_profile_overrides: ProviderRequestProfileOverrides,
pub id: ClientApiKeyId,
pub name: String,
@@ -204,6 +207,7 @@ impl fmt::Debug for NewClientKey {
/// 修改 Client Key 的公开策略字段
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpdateClientKey {
pub peer_sharing: Option<gateway_core::policy::ClientPeerSharing>,
pub request_profile_override_updates: ProviderRequestProfileOverrideUpdates,
pub id: ClientApiKeyId,
pub name: String,
@@ -14,6 +14,7 @@ pub mod diagnostics;
pub mod import_tasks;
pub mod key_usage;
pub mod observability;
pub mod peer_link;
pub mod plugin_client_keys;
pub mod plugin_resources;
pub mod plugins;
@@ -26,6 +27,7 @@ pub mod reset_credit_automation;
pub mod retention;
pub mod settings;
pub mod system;
pub mod upstream_info;
pub mod usage_data;
/// 管理用例对外返回的稳定错误分类
@@ -648,6 +648,7 @@ pub struct AccountRequestBucket {
/// Dashboard 中一个账号的完整用量事实
#[derive(Debug, Clone, PartialEq)]
pub struct DashboardAccountUsage {
pub upstream: Option<super::upstream_info::ProviderUpstreamInfo>,
pub account_id: String,
pub provider_kind: String,
pub authentication_kind: String,
@@ -0,0 +1,82 @@
//! 当前认证 Key 可公开的本地互联观测,不携带 HTTP 或 Provider 私有字段
use super::provider_credentials::ProviderQuotaWindowRole;
use chrono::{DateTime, Utc};
use gateway_core::engine::budget::ClientBudgetStatus;
/// 调用边界提供的响应资源预算,超限时整体拒绝而不截断授权范围
#[derive(Debug, Clone, Copy)]
pub struct PeerLinkLimits {
pub max_accounts: usize,
pub max_quota_windows: usize,
pub max_models: usize,
pub max_response_bytes: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PeerNodeIdentity {
pub id: String,
pub name: String,
}
#[derive(Debug, Clone, PartialEq)]
pub struct PeerLinkExport {
pub node: PeerNodeIdentity,
pub observed_at: DateTime<Utc>,
pub expires_at: DateTime<Utc>,
pub budget: ClientBudgetStatus,
pub budget_observed_at: DateTime<Utc>,
pub pool: Option<PeerPoolSummary>,
pub quota: Option<Vec<PeerAccountQuota>>,
pub models: Vec<PeerModelSummary>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct PeerPoolSummary {
pub ready: u32,
pub cooldown: u32,
pub exhausted: u32,
pub disabled: u32,
pub error: u32,
pub unknown: u32,
pub plans: Option<std::collections::BTreeMap<String, u32>>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct PeerAccountQuota {
pub id: String,
pub models: Vec<String>,
pub observed_at: Option<DateTime<Utc>>,
pub expires_at: Option<DateTime<Utc>>,
pub windows: Vec<PeerQuotaWindow>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct PeerQuotaWindow {
pub key: String,
pub group: String,
pub label: String,
pub limit_id: Option<String>,
pub limit_name: Option<String>,
pub role: Option<ProviderQuotaWindowRole>,
pub window_seconds: Option<u64>,
pub used_percent: Option<f64>,
pub reset_at: Option<DateTime<Utc>>,
pub limit_reached: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PeerModelState {
Available,
Exhausted,
Unknown,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PeerModelSummary {
pub model: String,
pub complete: bool,
pub state: PeerModelState,
pub observed_at: Option<DateTime<Utc>>,
pub expires_at: Option<DateTime<Utc>>,
}
@@ -790,6 +790,7 @@ impl ProviderQuotaWindowRole {
/// 一个 Provider quota 窗口的公共投影
#[derive(Debug, Clone, PartialEq)]
pub struct ProviderQuotaWindow {
pub upstream: Option<super::upstream_info::UpstreamQuotaWindow>,
pub key: String,
pub group: String,
pub label: String,
@@ -828,8 +829,9 @@ pub struct ProviderQuotaReadRequest<'a> {
}
/// Provider 已解析的 quota 结果及其不透明差异字段
#[derive(Debug, Clone, PartialEq)]
#[derive(Debug, Default, Clone, PartialEq)]
pub struct ProviderQuota {
pub upstream: Option<super::upstream_info::ProviderUpstreamInfo>,
/// 上游额度响应明确提供的套餐,可用于补全账号展示
pub plan_type: Option<String>,
pub observed_at: Option<DateTime<Utc>>,
@@ -1008,6 +1010,21 @@ impl ProviderQuota {
.min_by_key(|(_, period)| *period)
}
/// 互联账号按最长的 Key 消费周期统计本地用量,池内账号窗口不代表此 Key 的周期
#[must_use]
pub fn upstream_usage_window(&self) -> Option<&ProviderQuotaWindow> {
self.windows
.iter()
.filter(|window| {
window
.upstream
.as_ref()
.is_some_and(|upstream| upstream.amount.is_some())
&& window.window_seconds.is_some_and(|seconds| seconds > 0)
})
.max_by_key(|window| window.window_seconds)
}
/// 统计与预测共用窗口归属规则,避免把模型专属桶或短期限流当作账号容量
pub(crate) fn usage_windows(
&self,
@@ -28,6 +28,7 @@ pub struct ProxyListQuery {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProxyAccountRef {
pub upstream: Option<super::upstream_info::ProviderUpstreamInfo>,
pub id: String,
pub name: String,
pub email: Option<String>,
@@ -0,0 +1,93 @@
//! 管理端互联快照的中立展示事实,不携带 Provider 原始协议
use chrono::{DateTime, Utc};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderUpstreamInfo {
pub discovered: bool,
pub node: Option<UpstreamNode>,
pub observed_at: Option<DateTime<Utc>>,
pub stale: bool,
pub connection_revision: String,
pub sync: UpstreamSync,
pub last_error: Option<String>,
pub pool: Option<UpstreamPool>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpstreamNode {
pub id: String,
pub name: String,
pub software: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum UpstreamSyncState {
Idle,
Queued,
Running,
}
impl UpstreamSyncState {
pub const fn as_str(self) -> &'static str {
match self {
Self::Idle => "idle",
Self::Queued => "queued",
Self::Running => "running",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpstreamSync {
pub state: UpstreamSyncState,
pub request_generation: u64,
pub completed_generation: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpstreamRefreshReceipt {
pub account_id: String,
pub connection_revision: String,
pub request_generation: u64,
pub state: UpstreamSyncState,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpstreamPool {
pub account_count: u32,
pub states: UpstreamPoolStates,
pub plans: Vec<UpstreamPlanCount>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpstreamPoolStates {
pub ready: u32,
pub cooldown: u32,
pub exhausted: u32,
pub disabled: u32,
pub error: u32,
pub unknown: u32,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpstreamPlanCount {
pub name: String,
pub count: u32,
}
#[derive(Debug, Clone, PartialEq)]
pub struct UpstreamQuotaWindow {
pub account_id: Option<String>,
pub observed_at: Option<DateTime<Utc>>,
pub stale: bool,
pub amount: Option<UpstreamQuotaAmount>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpstreamQuotaAmount {
pub unit: String,
pub used: String,
pub remaining: Option<String>,
pub total: Option<String>,
}
@@ -2,6 +2,7 @@
pub mod backup;
pub mod client_distribution;
pub mod peer;
pub mod plugin_accounts;
pub mod plugin_client_keys;
pub mod plugin_management;
@@ -0,0 +1,74 @@
//! CPR 连接材料与传输端口;Provider 解释凭据,Host 执行受限网络请求
use std::time::Duration;
use async_trait::async_trait;
use gateway_core::account::{ProviderAccount, UpstreamInfoFailure};
use secrecy::SecretString;
/// 仅供后台请求使用,不进入日志、持久观测或管理响应
pub struct PeerConnection {
pub base_url: String,
pub bearer_token: SecretString,
}
pub struct PeerFetchFailure {
pub kind: UpstreamInfoFailure,
pub retry_after: Option<Duration>,
}
impl From<UpstreamInfoFailure> for PeerFetchFailure {
fn from(kind: UpstreamInfoFailure) -> Self {
Self {
kind,
retry_after: None,
}
}
}
#[async_trait]
pub trait PeerTransport: Send + Sync {
async fn fetch(
&self,
account: &ProviderAccount,
connection: &PeerConnection,
) -> Result<Vec<u8>, PeerFetchFailure>;
}
/// 账号用例读取公共互联观测,刷新动作只提交本地计划
#[async_trait]
pub trait PeerUpstreamService: gateway_core::account::UpstreamInfoSource {
async fn account_changed(
&self,
id: &gateway_core::account::ProviderAccountId,
) -> Result<(), gateway_core::error::StoreError>;
fn forget(&self, id: &gateway_core::account::ProviderAccountId);
async fn request_refresh(
&self,
id: &gateway_core::account::ProviderAccountId,
) -> Result<crate::model::upstream_info::UpstreamRefreshReceipt, gateway_core::error::StoreError>;
async fn read_quota(
&self,
account_id: &gateway_core::account::ProviderAccountId,
) -> Result<
Option<crate::model::provider_credentials::ProviderQuota>,
gateway_core::error::StoreError,
>;
/// 读取版本化摘要,不展开远端节点窗口
async fn read_summary(
&self,
account_id: &gateway_core::account::ProviderAccountId,
) -> Result<
Option<crate::model::provider_credentials::ProviderQuota>,
gateway_core::error::StoreError,
>;
async fn read_quotas(
&self,
requests: &[crate::model::provider_credentials::ProviderQuotaReadRequest<'_>],
) -> Vec<
Result<
crate::model::provider_credentials::ProviderQuota,
super::provider::ProviderAdminError,
>,
>;
}
@@ -179,6 +179,19 @@ pub trait ProviderAdmin: Send + Sync {
}
fn provider_kind(&self) -> &ProviderKind;
/// 与执行 Provider 共用的直连资格,未知实现默认不公开账号
fn is_direct_account(&self, _authentication_kind: &str) -> bool {
false
}
/// 仅依据本地账号事实判断通用模型资格;未知实现不能宣称完整覆盖
fn peer_model_eligible(
&self,
_account: &crate::model::accounts::AccountRecord,
_model: &str,
) -> Option<bool> {
None
}
/// 只读取原生 Provider 能力及账号类型,不执行网络或数据库查询
fn account_capabilities(
@@ -309,6 +322,24 @@ pub trait ProviderAdmin: Send + Sync {
.await
}
/// 仅依据账号类型声明 CPR 连接能力,通用用例不解析 Provider 凭据
fn peer_authentication_kinds(&self) -> &[&str] {
&[]
}
fn supports_peer_connection(&self, authentication_kind: &str) -> bool {
self.peer_authentication_kinds()
.contains(&authentication_kind)
}
/// 从本轮加载的凭据提取连接材料,避免在校验连接版本后另读一份凭据
fn peer_connection(
&self,
_loaded: &gateway_core::account::LoadedCredential,
) -> Result<Option<super::peer::PeerConnection>, ProviderAdminError> {
Ok(None)
}
/// 历史观测的协议字段仅由具体 Provider 解释;不支持时保留累计估算
fn quota_forecast_observation(
&self,
@@ -385,6 +416,10 @@ pub struct ProviderAdminRegistry {
}
impl ProviderAdminRegistry {
pub(crate) fn registered(&self) -> BTreeMap<ProviderKind, Arc<dyn ProviderAdmin>> {
self.providers.as_ref().clone()
}
pub fn pricing_catalog(&self) -> gateway_core::metering::PricingOverrides {
self.providers
.iter()
@@ -131,6 +131,35 @@ pub trait AccountStore:
runtime: AccountRuntimeSnapshot,
) -> AdminStoreResult<Option<AccountPageItem>>;
/// 批量读取指定账号与当前启用分组的交集,分组只用于授权筛选,不加载展示信息
/// 不加载凭据正文或触发 Provider 查询
async fn load_accounts_in_enabled_groups(
&self,
account_ids: &[gateway_core::account::ProviderAccountId],
group_ids: &[gateway_core::routing::AccountGroupId],
) -> AdminStoreResult<Vec<crate::model::accounts::AccountRecord>> {
if group_ids.is_empty() {
return Ok(Vec::new());
}
let mut accounts = Vec::with_capacity(account_ids.len());
for id in account_ids {
if let Some(item) = self
.load_account(id.as_str(), AccountRuntimeSnapshot::default())
.await?
&& item
.account
.groups
.iter()
.any(|group| group.enabled && group_ids.contains(&group.id))
{
let mut account = item.account;
account.groups.clear();
accounts.push(account);
}
}
Ok(accounts)
}
async fn account_exists(&self, account_id: &str) -> AdminStoreResult<bool>;
async fn load_account_usage(
@@ -252,6 +281,14 @@ pub trait AccountRuntimeStore: Send + Sync {
account_ids: &[String],
) -> AdminStoreResult<AccountRuntimeSnapshot>;
/// 只需账号可用状态时不读取在途租约,避免池规模放大容量查询
async fn account_rate_limits(
&self,
account_ids: &[String],
) -> AdminStoreResult<AccountRuntimeSnapshot> {
self.account_runtime(account_ids).await
}
/// 账号错误保护中的账号与其冻结快照;临时限流不包含在内
async fn active_freezes(
&self,
@@ -322,6 +359,16 @@ pub trait AuthStore: Send + Sync {
/// Client API Key 资料读取与管理写入
#[async_trait]
pub trait ClientKeyStore: Send + Sync {
async fn peer_node_identity(
&self,
) -> AdminStoreResult<crate::model::peer_link::PeerNodeIdentity> {
Err(AdminStoreError::new(
AdminStoreErrorKind::Unavailable,
"peer node",
"peer node is unavailable",
))
}
/// 按已验证的 ID 读取资料,不读取完整明文 Key
async fn get_client_key(
&self,
@@ -55,6 +55,13 @@ const CONNECTION_TEST_INPUT: &str = "Reply with exactly OK.";
/// 统一账号页消费的服务
#[async_trait]
pub trait AccountsService: Send + Sync {
async fn refresh_upstream_info(
&self,
_account_ids: Vec<ProviderAccountId>,
) -> Result<Vec<crate::model::upstream_info::UpstreamRefreshReceipt>, AdminError> {
Err(AdminError::invalid("该服务商不支持上游信息查询"))
}
async fn concurrency(
&self,
account_id: ProviderAccountId,
@@ -191,12 +198,17 @@ pub(crate) struct DefaultAccountsService {
providers: ProviderAdminRegistry,
snapshot: Arc<dyn SnapshotControl>,
probe: Arc<dyn AccountProbe>,
upstream_info: Option<Arc<dyn crate::ports::peer::PeerUpstreamService>>,
reset_credit_locks:
Arc<futures::lock::Mutex<BTreeMap<ProviderAccountId, Arc<futures::lock::Mutex<()>>>>>,
}
impl DefaultAccountsService {
#[must_use]
#[expect(
clippy::too_many_arguments,
reason = "账号存储、发布、探测与互联分别通过独立端口注入"
)]
pub(crate) fn new(
accounts: Arc<dyn AccountStore>,
account_runtime: Arc<dyn AccountRuntimeStore>,
@@ -205,6 +217,7 @@ impl DefaultAccountsService {
providers: ProviderAdminRegistry,
snapshot: Arc<dyn SnapshotControl>,
probe: Arc<dyn AccountProbe>,
upstream_info: Option<Arc<dyn crate::ports::peer::PeerUpstreamService>>,
) -> Self {
Self {
accounts,
@@ -214,6 +227,7 @@ impl DefaultAccountsService {
providers,
snapshot,
probe,
upstream_info,
reset_credit_locks: Arc::new(futures::lock::Mutex::new(BTreeMap::new())),
}
}
@@ -231,10 +245,11 @@ impl DefaultAccountsService {
}
async fn read_quotas(
providers: &ProviderAdminRegistry,
&self,
items: &[AccountPageItem],
rolling_usage: &BTreeMap<String, AccountUsage>,
) -> Vec<ProviderQuota> {
let providers = &self.providers;
let mut groups = BTreeMap::<_, Vec<usize>>::new();
for (index, item) in items.iter().enumerate() {
groups
@@ -264,7 +279,12 @@ impl DefaultAccountsService {
}
})
.collect::<Vec<_>>();
let results = provider.read_quotas(&requests).await;
let results = super::peer_link::upstream::read_quotas(
self.upstream_info.as_deref(),
provider.as_ref(),
&requests,
)
.await;
if results.len() != indices.len() {
tracing::warn!(
provider = %kind,
@@ -381,32 +401,51 @@ impl DefaultAccountsService {
async fn load_api_key_usage(
&self,
accounts: &[AccountPageItem],
quotas: &[ProviderQuota],
) -> Result<BTreeMap<String, AccountUsage>, AdminError> {
let now = Utc::now();
// API Key 没有套餐周期;本地累计直接查询账号创建后仍保留的请求记录
let mut usage = BTreeMap::new();
// 互联账号复用已聚合的最长 Key 窗口,普通 API Key 保留累计口径
let windows = accounts
.iter()
.filter(|item| item.account.authentication_kind == "api_key")
.map(|item| AccountUsageWindowQuery {
account_id: item.account.id.clone(),
key: "account-lifetime".to_owned(),
range: TimeRange {
start: item.account.created_at,
end: now,
},
.zip(quotas)
.filter(|(item, _)| item.account.authentication_kind == "api_key")
.filter_map(|(item, quota)| {
if quota
.upstream
.as_ref()
.is_some_and(|upstream| upstream.discovered)
{
if let Some(local_usage) = quota
.upstream_usage_window()
.and_then(|window| window.local_usage.as_ref())
{
usage.insert(item.account.id.clone(), local_usage.clone());
}
return None;
}
Some(AccountUsageWindowQuery {
account_id: item.account.id.clone(),
key: "account-lifetime".to_owned(),
range: TimeRange {
start: item.account.created_at,
end: now,
},
})
})
.collect::<Vec<_>>();
if windows.is_empty() {
return Ok(BTreeMap::new());
return Ok(usage);
}
Ok(self
.accounts
.load_account_usage_by_windows(&windows)
.await
.map_err(|error| map_store_error(error, "API Key account usage"))?
.into_iter()
.map(|result| (result.account_id, result.usage))
.collect())
usage.extend(
self.accounts
.load_account_usage_by_windows(&windows)
.await
.map_err(|error| map_store_error(error, "API Key account usage"))?
.into_iter()
.map(|result| (result.account_id, result.usage)),
);
Ok(usage)
}
async fn load_directory_item(
@@ -428,13 +467,17 @@ impl DefaultAccountsService {
.await
.map_err(|error| map_store_error(error, "rolling account usage"))?;
let rolling_usage = rolling_usage.into_iter().next();
let mut quota = match provider
.quota(ProviderQuotaRequest {
let mut quota = match super::peer_link::upstream::read_quota(
self.upstream_info.as_deref(),
provider.as_ref(),
account,
ProviderQuotaRequest {
account_id: account_id.clone(),
refresh: refresh_quota,
rolling_usage: rolling_usage.clone(),
})
.await
},
)
.await
{
Ok(quota) => quota,
// 凭据更新后的账号投影不要求 Provider 提供额度;显式刷新仍须支持该操作
@@ -458,7 +501,7 @@ impl DefaultAccountsService {
)
.await?;
let usage = self
.load_api_key_usage(std::slice::from_ref(&stored))
.load_api_key_usage(std::slice::from_ref(&stored), std::slice::from_ref(&quota))
.await?
.remove(&stored.account.id)
.or_else(|| {
@@ -486,6 +529,33 @@ impl DefaultAccountsService {
#[async_trait]
impl AccountsService for DefaultAccountsService {
async fn refresh_upstream_info(
&self,
account_ids: Vec<ProviderAccountId>,
) -> Result<Vec<crate::model::upstream_info::UpstreamRefreshReceipt>, AdminError> {
if account_ids.is_empty() || account_ids.len() > 200 {
return Err(AdminError::invalid("请选择要刷新的账号"));
}
let mut results = Vec::with_capacity(account_ids.len());
for id in account_ids {
let (stored, provider) = self.provider_for_account(&id).await?;
if !provider.supports_peer_connection(&stored.account.authentication_kind) {
return Err(AdminError::invalid("该账号不支持上游信息刷新"));
}
let peer = self
.upstream_info
.as_deref()
.ok_or_else(|| AdminError::invalid("上游信息服务未启用"))?;
results.push(peer.request_refresh(&id).await.map_err(|error| {
map_provider_error(
super::peer_link::upstream::store_error(error),
"upstream information refresh",
)
})?);
}
Ok(results)
}
async fn concurrency(
&self,
account_id: ProviderAccountId,
@@ -554,10 +624,10 @@ impl AccountsService for DefaultAccountsService {
.into_iter()
.map(|usage| (usage.account_id.clone(), usage))
.collect::<BTreeMap<_, _>>();
let mut quotas = Self::read_quotas(providers, &page.items, &rolling_usage).await;
let mut quotas = self.read_quotas(&page.items, &rolling_usage).await;
self.attach_quota_local_usage(&page.items, &mut quotas)
.await?;
let mut api_key_usage = self.load_api_key_usage(&page.items).await?;
let mut api_key_usage = self.load_api_key_usage(&page.items, &quotas).await?;
let items = page
.items
.into_iter()
@@ -684,9 +754,12 @@ impl AccountsService for DefaultAccountsService {
"provider credential refresh",
)
.await?;
provider
.account_facts_changed(std::slice::from_ref(&result.account_id))
.await;
super::notify_account_facts_changed(
self.upstream_info.as_deref(),
provider.as_ref(),
std::slice::from_ref(&result.account_id),
)
.await;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
let account = self.load_directory_item(&result.account_id, false).await?;
Ok(AccountRefreshResult {
@@ -726,9 +799,12 @@ impl AccountsService for DefaultAccountsService {
.map_err(|error| map_store_error(error, "enable provider account"))?
.config_revision
};
provider
.account_facts_changed(std::slice::from_ref(&account_id))
.await;
super::notify_account_facts_changed(
self.upstream_info.as_deref(),
provider.as_ref(),
std::slice::from_ref(&account_id),
)
.await;
publish_committed(self.snapshot.as_ref(), config_revision).await?;
let account = self.load_directory_item(&account_id, false).await?;
Ok(AccountRefreshResult {
@@ -752,11 +828,19 @@ impl AccountsService for DefaultAccountsService {
.await
.map_err(|error| map_store_error(error, "provider account"))?;
if !enabled {
provider.account_unavailable(&account_id).await;
}
provider
.account_facts_changed(std::slice::from_ref(&account_id))
super::notify_account_unavailable(
self.upstream_info.as_deref(),
provider.as_ref(),
&account_id,
)
.await;
}
super::notify_account_facts_changed(
self.upstream_info.as_deref(),
provider.as_ref(),
std::slice::from_ref(&account_id),
)
.await;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
@@ -798,10 +882,20 @@ impl AccountsService for DefaultAccountsService {
for (provider, provider_ids) in providers.values() {
if enabled == Some(false) {
for account_id in provider_ids {
provider.account_unavailable(account_id).await;
super::notify_account_unavailable(
self.upstream_info.as_deref(),
provider.as_ref(),
account_id,
)
.await;
}
}
provider.account_facts_changed(provider_ids).await;
super::notify_account_facts_changed(
self.upstream_info.as_deref(),
provider.as_ref(),
provider_ids,
)
.await;
}
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
@@ -831,14 +925,18 @@ impl AccountsService for DefaultAccountsService {
account_id: &ProviderAccountId,
) -> Result<AccountQuotaForecastReport, AdminError> {
let (stored, provider) = self.provider_for_account(account_id).await?;
let quota = provider
.quota(ProviderQuotaRequest {
let quota = super::peer_link::upstream::read_quota(
self.upstream_info.as_deref(),
provider.as_ref(),
&stored.account,
ProviderQuotaRequest {
account_id: account_id.clone(),
refresh: false,
rolling_usage: None,
})
.await
.map_err(|error| map_provider_error(error, "forecast quota snapshot"))?;
},
)
.await
.map_err(|error| map_provider_error(error, "forecast quota snapshot"))?;
let now = Utc::now();
let mut samples = Vec::new();
let mut selected_keys = BTreeSet::new();
@@ -1143,15 +1241,7 @@ fn map_reset_credits_error_after_refresh(
/// 账号目录中单个账号 quota 读取失败时使用的空额度投影
fn empty_quota() -> ProviderQuota {
ProviderQuota {
credits: None,
plan_type: None,
observed_at: None,
refresh_token_expires_at: None,
windows: Vec::new(),
limit_reached: false,
provider_data: None,
}
ProviderQuota::default()
}
fn quota_usage_window(
@@ -165,6 +165,7 @@ impl ClientKeyService for DefaultClientKeyService {
.and_then(|provider| provider.preview_client_profile(profile))
.map_err(|error| super::map_provider_error(error, "client profile"))?;
}
validate_peer_sharing(command.peer_sharing, &command.group_ids)?;
let id = ClientApiKeyId::new(format!("key_{}", Uuid::now_v7().simple()))
.map_err(|_| AdminError::internal("创建 Client API Key ID 失败"))?;
let plaintext = if let Some(key) = command.custom_key {
@@ -176,6 +177,7 @@ impl ClientKeyService for DefaultClientKeyService {
.store
.create_client_key(
NewClientKey {
peer_sharing: command.peer_sharing,
request_profile_overrides: command.request_profile_overrides,
id,
name: command.name,
@@ -211,6 +213,9 @@ impl ClientKeyService for DefaultClientKeyService {
.map_err(|error| super::map_provider_error(error, "client profile"))?;
}
}
if let Some(sharing) = command.peer_sharing {
validate_peer_sharing(sharing, &command.group_ids)?;
}
let id = command.id.clone();
let (config_revision, record) = self
.store
@@ -308,3 +313,16 @@ pub(super) fn generate_key() -> String {
OsRng.fill_bytes(&mut bytes);
format!("sk_{}", URL_SAFE_NO_PAD.encode(bytes))
}
fn validate_peer_sharing(
sharing: gateway_core::policy::ClientPeerSharing,
groups: &[gateway_core::routing::AccountGroupId],
) -> Result<(), AdminError> {
sharing.validate(!groups.is_empty()).map_err(|error| {
use gateway_core::policy::ClientPeerSharingError;
AdminError::invalid(match error {
ClientPeerSharingError::MissingGroupScope => "开启 CPR 互联需要指定账号分组",
ClientPeerSharingError::PlansWithoutPool => "公开套餐分布需要先公开账号池",
})
})
}
@@ -19,13 +19,14 @@ use gateway_core::{account::ProviderAccountId, routing::ProviderKind, runtime::S
use super::{map_store_error, publish_committed};
pub(super) async fn publish_credentials_and_observe_quota(
peer: Option<&dyn crate::ports::peer::PeerUpstreamService>,
provider: &std::sync::Arc<dyn crate::ports::provider::ProviderAdmin>,
snapshot: &dyn SnapshotControl,
revision: crate::model::Revision,
account_ids: &[ProviderAccountId],
request_id: &str,
) -> Result<(), AdminError> {
provider.account_facts_changed(account_ids).await;
notify_account_facts_changed(peer, provider.as_ref(), account_ids).await;
publish_committed(snapshot, revision).await?;
// 凭据已经提交;额度只是可重建的观察,不能拖住管理请求或回滚提交结果
@@ -343,6 +344,7 @@ pub(super) async fn commit_credential_refresh(
}
pub(super) async fn delete_credentials(
peer: Option<&dyn crate::ports::peer::PeerUpstreamService>,
accounts: &dyn AccountStore,
provider: &dyn ProviderAdmin,
command: CredentialDeletion,
@@ -365,11 +367,37 @@ pub(super) async fn delete_credentials(
.await
.map_err(|error| map_store_error(error, resource))?;
for account_id in &account_ids {
provider.account_unavailable(account_id).await;
notify_account_unavailable(peer, provider, account_id).await;
}
provider.account_facts_changed(&account_ids).await;
notify_account_facts_changed(peer, provider, &account_ids).await;
Ok(CredentialDeletionResult {
config_revision: revision,
account_ids,
})
}
pub(super) async fn notify_account_facts_changed(
peer: Option<&dyn crate::ports::peer::PeerUpstreamService>,
provider: &dyn ProviderAdmin,
ids: &[ProviderAccountId],
) {
provider.account_facts_changed(ids).await;
if let Some(peer) = peer {
for id in ids {
if let Err(error) = peer.account_changed(id).await {
tracing::warn!(account_id = %id, error = %error, "upstream information plan reconciliation failed");
}
}
}
}
pub(super) async fn notify_account_unavailable(
peer: Option<&dyn crate::ports::peer::PeerUpstreamService>,
provider: &dyn ProviderAdmin,
id: &ProviderAccountId,
) {
if let Some(peer) = peer {
peer.forget(id);
}
provider.account_unavailable(id).await;
}
@@ -30,6 +30,7 @@ pub struct ProviderCredentials {
accounts: Arc<dyn AccountStore>,
proxies: Arc<dyn crate::ports::proxy::ProxyStore>,
snapshot: Arc<dyn SnapshotControl>,
upstream_info: Option<Arc<dyn crate::ports::peer::PeerUpstreamService>>,
}
impl ProviderCredentials {
@@ -39,12 +40,14 @@ impl ProviderCredentials {
accounts: Arc<dyn AccountStore>,
proxies: Arc<dyn crate::ports::proxy::ProxyStore>,
snapshot: Arc<dyn SnapshotControl>,
upstream_info: Option<Arc<dyn crate::ports::peer::PeerUpstreamService>>,
) -> Self {
Self {
provider,
accounts,
proxies,
snapshot,
upstream_info,
}
}
pub async fn import_document(
@@ -89,6 +92,7 @@ impl ProviderCredentials {
.map_err(|error| map_store_error(error, "provider credential import"))?;
drop(proxy_reservation);
publish_credentials_and_observe_quota(
self.upstream_info.as_deref(),
&self.provider,
self.snapshot.as_ref(),
result.config_revision,
@@ -160,11 +164,19 @@ impl ProviderCredentials {
)
.await?;
if disable_account {
self.provider.account_unavailable(&account_id).await;
}
self.provider
.account_facts_changed(std::slice::from_ref(&result.account_id))
super::super::notify_account_unavailable(
self.upstream_info.as_deref(),
self.provider.as_ref(),
&account_id,
)
.await;
}
super::super::notify_account_facts_changed(
self.upstream_info.as_deref(),
self.provider.as_ref(),
std::slice::from_ref(&result.account_id),
)
.await;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
@@ -174,6 +186,7 @@ impl ProviderCredentials {
command: CredentialDeletion,
) -> Result<CredentialDeletionResult, AdminError> {
let result = delete_credentials(
self.upstream_info.as_deref(),
self.accounts.as_ref(),
self.provider.as_ref(),
command,
@@ -236,6 +249,7 @@ impl super::CredentialsService {
let result = outcome.result;
if outcome.newly_committed {
publish_credentials_and_observe_quota(
self.upstream_info.as_deref(),
provider,
self.snapshot.as_ref(),
result.config_revision,
@@ -19,6 +19,7 @@ pub struct CredentialsService {
accounts: Arc<dyn AccountStore>,
proxies: Arc<dyn ProxyStore>,
snapshot: Arc<dyn SnapshotControl>,
upstream_info: Option<Arc<dyn crate::ports::peer::PeerUpstreamService>>,
}
impl CredentialsService {
@@ -28,12 +29,14 @@ impl CredentialsService {
accounts: Arc<dyn AccountStore>,
proxies: Arc<dyn ProxyStore>,
snapshot: Arc<dyn SnapshotControl>,
upstream_info: Option<Arc<dyn crate::ports::peer::PeerUpstreamService>>,
) -> Self {
Self {
providers,
accounts,
proxies,
snapshot,
upstream_info,
}
}
@@ -47,6 +50,7 @@ impl CredentialsService {
self.accounts.clone(),
self.proxies.clone(),
self.snapshot.clone(),
self.upstream_info.clone(),
))
}
}
@@ -10,6 +10,7 @@ pub mod credentials;
pub mod import_tasks;
pub mod key_usage;
pub mod observability;
pub mod peer_link;
pub mod plugin_accounts;
pub mod plugin_client_keys;
pub mod plugin_resources;
@@ -26,10 +27,10 @@ mod publication;
use credential_mutation::{
commit_authorization, commit_credential_refresh, commit_credential_rotation,
delete_credentials, import_proxy_binding, pending_authorization,
publish_credentials_and_observe_quota, required_credential, required_plugin_credential,
validate_authorization_commit, validate_prepared_import, validate_prepared_rotation,
validate_prepared_rotation_facts,
delete_credentials, import_proxy_binding, notify_account_facts_changed,
notify_account_unavailable, pending_authorization, publish_credentials_and_observe_quota,
required_credential, required_plugin_credential, validate_authorization_commit,
validate_prepared_import, validate_prepared_rotation, validate_prepared_rotation_facts,
};
use error::{map_provider_error, map_store_error};
use publication::publish_committed;
@@ -94,6 +94,7 @@ pub(crate) struct DefaultObservabilityService {
settings: Arc<dyn SettingsStore>,
providers: ProviderAdminRegistry,
system: Arc<dyn super::system::SystemService>,
upstream_info: Option<Arc<dyn crate::ports::peer::PeerUpstreamService>>,
}
impl DefaultObservabilityService {
@@ -105,6 +106,7 @@ impl DefaultObservabilityService {
providers: ProviderAdminRegistry,
timezone: gateway_core::time::DeploymentTimeZone,
system: Arc<dyn super::system::SystemService>,
upstream_info: Option<Arc<dyn crate::ports::peer::PeerUpstreamService>>,
) -> Self {
Self {
timezone,
@@ -113,6 +115,7 @@ impl DefaultObservabilityService {
settings,
providers,
system,
upstream_info,
}
}
@@ -933,6 +936,7 @@ impl DefaultObservabilityService {
.map(|usage| (usage.account_id.clone(), usage))
.collect::<HashMap<_, _>>();
let providers = self.providers.clone();
let peer = self.upstream_info.as_deref();
let quota_reads = accounts.iter().map(|account| {
let rolling_usage = rolling_usage.remove(&account.account_id);
let account_id = ProviderAccountId::new(account.account_id.clone()).ok();
@@ -943,6 +947,11 @@ impl DefaultObservabilityService {
let (Some(account_id), Some(provider)) = (account_id, provider) else {
return None;
};
if let Some(peer) = peer
&& let Ok(Some(quota)) = peer.read_summary(&account_id).await
{
return Some(quota);
}
provider
.quota(ProviderQuotaRequest {
account_id,
@@ -960,6 +969,7 @@ impl DefaultObservabilityService {
quota.apply_limit_reached_display();
account.quota_used_percent = quota.representative_used_percent();
account.quota_window = quota.representative_window().cloned();
account.upstream = quota.upstream;
}
}
}
@@ -0,0 +1,261 @@
//! CPR 互联只读取认证范围内的本地事实,不参与推理准入或上游刷新
mod projection;
mod rate_limit;
pub mod upstream;
use std::{collections::BTreeMap, sync::Arc};
use async_trait::async_trait;
use chrono::{TimeDelta, Utc};
use gateway_core::{engine::execution::AuthenticatedClient, policy::ClientPeerSharing};
use tokio::sync::Semaphore;
use crate::{
model::{
AdminError,
peer_link::{PeerLinkExport, PeerLinkLimits},
provider_credentials::ProviderQuotaReadRequest,
},
ports::{
provider::ProviderAdminRegistry,
store::{AccountRuntimeStore, AccountStore, ClientKeyStore},
},
};
use super::map_store_error;
const MAX_CONCURRENT_SNAPSHOTS: usize = 8;
/// 查询失败保持独立于推理错误,不改变账号或 Key 状态
#[derive(Debug, thiserror::Error)]
pub enum PeerLinkError {
#[error("client API key is unavailable")]
Unauthorized,
#[error("peer sharing is not allowed")]
Forbidden,
#[error("peer query rate limit reached")]
RateLimited { retry_after: u64 },
#[error("peer snapshot exceeds the supported bounds")]
SnapshotTooLarge,
#[error("peer snapshot is unavailable")]
Unavailable(#[from] AdminError),
}
#[async_trait]
pub trait PeerLinkService: Send + Sync {
async fn snapshot(
&self,
client: &AuthenticatedClient,
limits: PeerLinkLimits,
) -> Result<PeerLinkExport, PeerLinkError>;
}
pub(crate) struct DefaultPeerLinkService {
keys: Arc<dyn ClientKeyStore>,
accounts: Arc<dyn AccountStore>,
runtime: Arc<dyn AccountRuntimeStore>,
providers: ProviderAdminRegistry,
rate_limit: rate_limit::QueryRateLimit,
snapshot_slots: Semaphore,
}
impl DefaultPeerLinkService {
pub(crate) fn new(
keys: Arc<dyn ClientKeyStore>,
accounts: Arc<dyn AccountStore>,
runtime: Arc<dyn AccountRuntimeStore>,
providers: ProviderAdminRegistry,
) -> Self {
Self {
keys,
accounts,
runtime,
providers,
rate_limit: Default::default(),
snapshot_slots: Semaphore::new(MAX_CONCURRENT_SNAPSHOTS),
}
}
}
#[async_trait]
impl PeerLinkService for DefaultPeerLinkService {
async fn snapshot(
&self,
client: &AuthenticatedClient,
limits: PeerLinkLimits,
) -> Result<PeerLinkExport, PeerLinkError> {
self.rate_limit.acquire(client.policy().key_id())?;
let frozen_sharing = client.policy().peer_sharing();
// 当前权限只能收窄冻结范围,重查询在读取 Key 分组前拒绝且不排队
let _permit = if frozen_sharing.share_pool || frozen_sharing.share_quota {
Some(
self.snapshot_slots
.try_acquire()
.map_err(|_| PeerLinkError::RateLimited { retry_after: 1 })?,
)
} else {
None
};
let key = self
.keys
.get_client_key(client.policy().key_id())
.await
.map_err(|error| map_store_error(error, "client API key"))?
.filter(|key| key.enabled)
.ok_or(PeerLinkError::Unauthorized)?;
let budget_observed_at = Utc::now();
let sharing = intersect_sharing(frozen_sharing, key.peer_sharing);
if !sharing.enabled || key.groups.is_empty() {
return Err(PeerLinkError::Forbidden);
}
let node = self
.keys
.peer_node_identity()
.await
.map_err(|error| map_store_error(error, "peer node"))?;
let now = Utc::now();
let mut result = PeerLinkExport {
node,
observed_at: now,
expires_at: now + TimeDelta::minutes(5),
budget: key.budget,
budget_observed_at,
pool: None,
quota: None,
models: Vec::new(),
};
if !sharing.share_pool && !sharing.share_quota {
return Ok(result);
}
let ids = client
.policy()
.account_scope()
.allowed_account_ids()
.take(limits.max_accounts.saturating_add(1))
.cloned()
.collect::<Vec<_>>();
if ids.len() > limits.max_accounts {
return Err(PeerLinkError::SnapshotTooLarge);
}
let public_models = if sharing.share_quota {
client
.snapshot()
.public_models_for_scope(client.policy().account_scope())
} else {
Vec::new()
};
if public_models.len() > limits.max_models {
return Err(PeerLinkError::SnapshotTooLarge);
}
let groups = key
.groups
.iter()
.filter(|group| group.enabled)
.map(|group| group.id.clone())
.collect::<Vec<_>>();
let accounts = self
.accounts
.load_accounts_in_enabled_groups(&ids, &groups)
.await
.map_err(|error| map_store_error(error, "peer accounts"))?
.into_iter()
.filter(|account| {
self.providers
.require(&account.provider_kind)
.is_ok_and(|provider| provider.is_direct_account(&account.authentication_kind))
})
.collect::<Vec<_>>();
let account_ids = accounts
.iter()
.map(|account| account.id.clone())
.collect::<Vec<_>>();
let runtime = self
.runtime
.account_rate_limits(&account_ids)
.await
.map_err(|error| map_store_error(error, "peer account runtime"))?;
let mut quotas = BTreeMap::new();
if sharing.share_quota {
let mut window_count = 0usize;
let mut by_provider = BTreeMap::new();
for account in &accounts {
by_provider
.entry(&account.provider_kind)
.or_insert_with(Vec::new)
.push(account);
}
for (kind, accounts) in by_provider {
let Ok(provider) = self.providers.require(kind) else {
continue;
};
let requests = accounts
.iter()
.map(|account| ProviderQuotaReadRequest {
account,
rolling_usage: None,
})
.collect::<Vec<_>>();
for (account, quota) in accounts
.into_iter()
.zip(provider.read_quotas(&requests).await)
{
if let Ok(quota) = quota {
window_count = window_count.saturating_add(quota.windows.len());
if window_count > limits.max_quota_windows {
return Err(PeerLinkError::SnapshotTooLarge);
}
quotas.insert(account.id.clone(), quota);
}
}
}
}
let now = Utc::now();
result.observed_at = now;
result.expires_at = now + TimeDelta::minutes(5);
if sharing.share_pool {
result.pool = Some(projection::pool_summary(
&accounts,
&runtime,
&self.providers,
sharing.share_plans,
now,
));
}
if sharing.share_quota {
let (quota, mut models) = projection::quota_and_models(
client,
&accounts,
&quotas,
&self.providers,
&runtime,
&public_models,
now,
limits.max_response_bytes,
)
.await?;
// 发布窗口内缺失或已变更的账号不能被静默剔除后推导整个模型池耗尽
if accounts.len() != ids.len() {
for model in &mut models {
model.complete = false;
if model.state == crate::model::peer_link::PeerModelState::Exhausted {
model.state = crate::model::peer_link::PeerModelState::Unknown;
}
}
}
result.quota = Some(quota);
result.models = models;
}
Ok(result)
}
}
fn intersect_sharing(frozen: ClientPeerSharing, current: ClientPeerSharing) -> ClientPeerSharing {
ClientPeerSharing {
enabled: frozen.enabled && current.enabled,
share_pool: frozen.share_pool && current.share_pool,
share_quota: frozen.share_quota && current.share_quota,
share_plans: frozen.share_plans && current.share_plans,
}
}
@@ -0,0 +1,260 @@
//! 从已授权的本地账号事实派生公开汇总,保留原始观测时间
use std::collections::BTreeMap;
use chrono::{DateTime, TimeDelta, Utc};
use gateway_core::{
account::{
AccountStatus, AccountStatusFacts, ProviderAccountId, QuotaAccessState,
resolve_account_status,
},
engine::execution::AuthenticatedClient,
};
use sha2::{Digest, Sha256};
use crate::{
model::{
accounts::{AccountRecord, AccountRuntimeSnapshot},
peer_link::{
PeerAccountQuota, PeerModelState, PeerModelSummary, PeerPoolSummary, PeerQuotaWindow,
},
provider_credentials::ProviderQuota,
},
ports::provider::ProviderAdminRegistry,
};
pub(super) fn pool_summary(
accounts: &[AccountRecord],
runtime: &AccountRuntimeSnapshot,
providers: &ProviderAdminRegistry,
share_plans: bool,
now: DateTime<Utc>,
) -> PeerPoolSummary {
let mut summary = PeerPoolSummary {
plans: share_plans.then(BTreeMap::new),
..Default::default()
};
for account in accounts {
match status(account, runtime, now) {
AccountStatus::Normal => summary.ready += 1,
AccountStatus::RateLimited => summary.cooldown += 1,
AccountStatus::QuotaExhausted => summary.exhausted += 1,
AccountStatus::Disabled => summary.disabled += 1,
AccountStatus::Error => summary.error += 1,
}
if let Some(plans) = &mut summary.plans
&& let (Some(plan), Ok(provider)) = (
&account.plan_type,
providers.require(&account.provider_kind),
)
{
*plans.entry(provider.plan_type_display(plan)).or_default() += 1;
}
}
summary
}
#[expect(
clippy::too_many_arguments,
reason = "投影显式接收已授权事实和调用边界的响应预算"
)]
pub(super) async fn quota_and_models(
client: &AuthenticatedClient,
accounts: &[AccountRecord],
quotas: &BTreeMap<String, ProviderQuota>,
providers: &ProviderAdminRegistry,
runtime: &AccountRuntimeSnapshot,
models: &[gateway_core::routing::PublicModelId],
now: DateTime<Utc>,
max_response_bytes: usize,
) -> Result<(Vec<PeerAccountQuota>, Vec<PeerModelSummary>), super::PeerLinkError> {
let snapshot = client.snapshot();
let scope = client.policy().account_scope();
let mut observations = BTreeMap::<&str, ModelObservations>::new();
let mut model_bytes = 0usize;
let mut result = Vec::with_capacity(accounts.len());
for (index, account) in accounts.iter().enumerate() {
// 大池投影不能连续占住异步执行线程,取消请求也要能及时释放并发名额
if index % 16 == 0 {
tokio::task::yield_now().await;
}
let observation = model_observation(account, runtime, now);
let Ok(id) = ProviderAccountId::new(account.id.clone()) else {
continue;
};
let provider = providers.require(&account.provider_kind).ok();
let mut account_models = Vec::new();
for model in models {
let upstream = snapshot.mapped_model(model.as_str());
if !snapshot.allows_public_model_for_account(model, scope, &id)
|| !account.model_access.allows(&upstream)
{
continue;
}
let eligible = provider
.as_ref()
.and_then(|provider| provider.peer_model_eligible(account, &upstream));
if eligible == Some(false) {
continue;
}
observations
.entry(model.as_str())
.or_default()
.observe(eligible.and(observation));
if eligible == Some(true) {
// 模型名是响应字节数的下界;提前拒绝必然超限的交叉展开,不误拒稀疏授权
model_bytes = model_bytes.saturating_add(model.as_str().len());
if model_bytes > max_response_bytes {
return Err(super::PeerLinkError::SnapshotTooLarge);
}
account_models.push(model.as_str().to_owned());
}
}
let quota = quotas.get(&account.id);
let observed_at = quota.and_then(|quota| quota.observed_at);
let expires_at = observed_at.map(|observed| {
quota
.into_iter()
.flat_map(|quota| &quota.windows)
.filter_map(|window| window.reset_at)
.fold(observed + TimeDelta::minutes(30), DateTime::min)
});
// 标识只在当前 Key 内稳定,不能用原始账号 ID 跨节点或跨授权关联身份
let mut digest = Sha256::new();
digest.update(b"cpr.peer.account\0");
digest.update(client.policy().plaintext_key().expose_for_auth().as_bytes());
digest.update([0]);
digest.update(account.id.as_bytes());
let id = hex::encode(&digest.finalize()[..16]);
result.push(PeerAccountQuota {
id,
models: account_models,
observed_at,
expires_at,
windows: quota
.into_iter()
.flat_map(|quota| &quota.windows)
.map(|window| PeerQuotaWindow {
key: window.key.clone(),
group: window.group.clone(),
label: window.label.clone(),
limit_id: window.limit_id.clone(),
limit_name: window.limit_name.clone(),
role: window.role,
window_seconds: window.window_seconds,
used_percent: window.used_percent,
reset_at: window.reset_at,
limit_reached: window.limit_reached,
})
.collect(),
});
}
let summaries = models
.iter()
.map(|model| {
let values = observations.remove(model.as_str()).unwrap_or_default();
let complete = values.observed_at.is_some() && !values.incomplete;
let state = if values.available {
PeerModelState::Available
} else if complete {
PeerModelState::Exhausted
} else {
PeerModelState::Unknown
};
PeerModelSummary {
model: model.as_str().to_owned(),
complete,
state,
observed_at: values.observed_at,
expires_at: values.expires_at,
}
})
.collect();
Ok((result, summaries))
}
#[derive(Clone, Copy)]
struct ModelObservation {
state: PeerModelState,
observed_at: DateTime<Utc>,
expires_at: DateTime<Utc>,
}
#[derive(Default)]
struct ModelObservations {
incomplete: bool,
available: bool,
observed_at: Option<DateTime<Utc>>,
expires_at: Option<DateTime<Utc>>,
}
impl ModelObservations {
fn observe(&mut self, observation: Option<ModelObservation>) {
let Some(observation) = observation else {
self.incomplete = true;
return;
};
self.available |= observation.state == PeerModelState::Available;
self.observed_at = Some(self.observed_at.map_or(observation.observed_at, |current| {
current.min(observation.observed_at)
}));
self.expires_at = Some(self.expires_at.map_or(observation.expires_at, |current| {
current.min(observation.expires_at)
}));
}
}
fn model_observation(
account: &AccountRecord,
runtime: &AccountRuntimeSnapshot,
now: DateTime<Utc>,
) -> Option<ModelObservation> {
if !matches!(
status(account, runtime, now),
AccountStatus::Normal | AccountStatus::QuotaExhausted
) {
return None;
}
let state = match account.quota.access() {
QuotaAccessState::Allowed => PeerModelState::Available,
QuotaAccessState::Exhausted => PeerModelState::Exhausted,
QuotaAccessState::Unknown => return None,
};
let observed_at = DateTime::<Utc>::from(account.quota.observed_at()?);
let expires_at = account
.quota
.reset_at()
.map(DateTime::<Utc>::from)
.map_or(observed_at + TimeDelta::minutes(30), |reset| {
reset.min(observed_at + TimeDelta::minutes(30))
})
.min(now + TimeDelta::minutes(5));
if observed_at > now || expires_at <= now {
return None;
}
Some(ModelObservation {
state,
observed_at,
expires_at,
})
}
fn status(
account: &AccountRecord,
runtime: &AccountRuntimeSnapshot,
now: DateTime<Utc>,
) -> AccountStatus {
resolve_account_status(
&AccountStatusFacts {
enabled: account.enabled,
credential_state: account.credential_state,
access_token_expires_at: account.access_token_expires_at.map(Into::into),
quota: account.quota,
cooldown: runtime.cooldown.get(&account.id).copied(),
last_error_reason: account.last_error_reason,
last_error_message: None,
},
now.into(),
)
.status
}
@@ -0,0 +1,46 @@
//! 互联只读查询的短期限流,不占用推理请求配额
use super::PeerLinkError;
use gateway_core::policy::ClientApiKeyId;
use std::{
collections::{BTreeMap, VecDeque},
sync::Mutex,
time::{Duration, Instant},
};
const WINDOW: Duration = Duration::from_secs(60);
const REQUESTS_PER_WINDOW: usize = 30;
#[derive(Default)]
pub(super) struct QueryRateLimit(Mutex<BTreeMap<ClientApiKeyId, VecDeque<Instant>>>);
impl QueryRateLimit {
pub(super) fn acquire(&self, id: &ClientApiKeyId) -> Result<(), PeerLinkError> {
let now = Instant::now();
let mut entries = self
.0
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
entries.retain(|_, requests| {
while requests
.front()
.is_some_and(|time| now.duration_since(*time) >= WINDOW)
{
requests.pop_front();
}
!requests.is_empty()
});
let requests = entries.entry(id.clone()).or_default();
if requests.len() >= REQUESTS_PER_WINDOW {
let retry_after = requests.front().map_or(1, |time| {
WINDOW
.saturating_sub(now.duration_since(*time))
.as_secs()
.saturating_add(1)
});
return Err(PeerLinkError::RateLimited { retry_after });
}
requests.push_back(now);
Ok(())
}
}
@@ -0,0 +1,134 @@
//! 账号查询合并原生额度与公共互联投影,保持列表摘要和详情窗口边界
use crate::{
model::{
accounts::AccountRecord,
provider_credentials::{ProviderQuota, ProviderQuotaReadRequest, ProviderQuotaRequest},
},
ports::provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind},
};
use gateway_core::account::{CredentialRevision, ProviderAccountId, UpstreamConnection};
fn connection(account: &AccountRecord) -> Result<UpstreamConnection, ProviderAdminError> {
let invalid = || ProviderAdminError::new(ProviderAdminErrorKind::Invalid);
Ok(UpstreamConnection::new(
ProviderAccountId::new(account.id.clone()).map_err(|_| invalid())?,
CredentialRevision::new(account.credential_revision.get()).map_err(|_| invalid())?,
account.effective_outbound_proxy(),
))
}
pub(crate) fn store_error(error: gateway_core::error::StoreError) -> ProviderAdminError {
use gateway_core::error::StoreErrorKind;
let kind = match error.kind() {
StoreErrorKind::Conflict => ProviderAdminErrorKind::Conflict,
StoreErrorKind::InvalidData | StoreErrorKind::InvalidState => {
ProviderAdminErrorKind::NotFound
}
StoreErrorKind::Unavailable => ProviderAdminErrorKind::Unavailable,
_ => ProviderAdminErrorKind::Internal,
};
ProviderAdminError::new(kind).with_source(error)
}
pub(crate) async fn read_quota(
peer: Option<&dyn crate::ports::peer::PeerUpstreamService>,
provider: &dyn ProviderAdmin,
account: &AccountRecord,
request: ProviderQuotaRequest,
) -> Result<ProviderQuota, ProviderAdminError> {
if !request.refresh
&& provider.supports_peer_connection(&account.authentication_kind)
&& let Some(peer) = peer
{
return peer
.read_quota(&request.account_id)
.await
.map(|quota| quota.unwrap_or_else(ProviderQuota::default))
.map_err(store_error);
}
provider.quota(request).await
}
pub(crate) async fn read_quotas(
peer: Option<&dyn crate::ports::peer::PeerUpstreamService>,
provider: &dyn ProviderAdmin,
requests: &[ProviderQuotaReadRequest<'_>],
) -> Vec<Result<ProviderQuota, ProviderAdminError>> {
let mut results = provider.read_quotas(requests).await;
let Some(peer) = peer else {
return results;
};
if results.len() != requests.len() {
return results;
}
let indices = requests
.iter()
.enumerate()
.filter_map(|(index, request)| {
provider
.supports_peer_connection(&request.account.authentication_kind)
.then_some(index)
})
.collect::<Vec<_>>();
let queries = indices
.iter()
.map(|&index| ProviderQuotaReadRequest {
account: requests[index].account,
rolling_usage: requests[index].rolling_usage,
})
.collect::<Vec<_>>();
for (index, quota) in indices.into_iter().zip(peer.read_quotas(&queries).await) {
results[index] = quota;
}
results
}
/// 列表只加载轻量摘要,连接身份过期时丢弃旧观测
pub(super) async fn read_peer_quotas(
service: &super::PeerInfoService,
requests: &[ProviderQuotaReadRequest<'_>],
) -> Vec<Result<ProviderQuota, ProviderAdminError>> {
let connections = requests
.iter()
.map(|request| connection(request.account))
.collect::<Vec<_>>();
let ids = connections
.iter()
.filter_map(|connection| {
connection
.as_ref()
.ok()
.map(|value| value.account_id.clone())
})
.collect::<Vec<_>>();
if ids.is_empty() {
return connections
.into_iter()
.map(|connection| connection.map(|_| ProviderQuota::default()))
.collect();
}
let summaries = service
.read_summaries(&ids)
.await
.map(|records| {
records
.into_iter()
.map(|(connection, quota)| (connection.account_id.clone(), (connection, quota)))
.collect::<std::collections::BTreeMap<_, _>>()
})
.map_err(store_error);
connections
.into_iter()
.map(|connection| {
connection.and_then(|connection| {
Ok(summaries
.as_ref()
.map_err(Clone::clone)?
.get(&connection.account_id)
.filter(|(observed, _)| *observed == connection)
.map_or_else(ProviderQuota::default, |(_, quota)| quota.clone()))
})
})
.collect()
}
@@ -0,0 +1,100 @@
//! 公共 CPR 上游同步、管理投影与调度观测,连接凭据解释由 Provider 提供
mod accounts;
mod projection;
mod scheduling;
mod service;
mod snapshot;
mod worker;
use std::sync::Arc;
use crate::{
model::{provider_credentials::ProviderQuota, upstream_info::UpstreamRefreshReceipt},
ports::{
peer::{PeerTransport, PeerUpstreamService},
provider::ProviderAdminRegistry,
},
};
pub(crate) use accounts::{read_quota, read_quotas, store_error};
use gateway_core::{
account::{ProviderAccountId, ProviderAccountStore, UpstreamInfoSource, UpstreamInfoStore},
error::StoreError,
task::{WorkerContribution, WorkerDefinitionError},
};
use service::PeerInfoService;
/// 组合根独立注册公共互联 Worker,不属于任何具体 Provider 的任务集合
pub struct PeerUpstreamBundle {
service: Arc<PeerInfoService>,
workers: Vec<WorkerContribution>,
}
impl PeerUpstreamBundle {
pub fn service(&self) -> Arc<dyn PeerUpstreamService> {
self.service.clone()
}
pub fn scheduling_source(&self) -> Arc<dyn UpstreamInfoSource> {
self.service.clone()
}
pub fn take_worker_contributions(&mut self) -> Vec<WorkerContribution> {
std::mem::take(&mut self.workers)
}
}
pub fn initialize(
accounts: Arc<dyn ProviderAccountStore>,
store: Arc<dyn UpstreamInfoStore>,
providers: &ProviderAdminRegistry,
transport: Arc<dyn PeerTransport>,
) -> Result<PeerUpstreamBundle, WorkerDefinitionError> {
let service = Arc::new(PeerInfoService::new(accounts, store, providers, transport));
let worker = service.worker()?;
Ok(PeerUpstreamBundle {
service,
workers: vec![worker],
})
}
#[async_trait::async_trait]
impl PeerUpstreamService for PeerInfoService {
async fn account_changed(&self, id: &ProviderAccountId) -> Result<(), StoreError> {
PeerInfoService::account_changed(self, id).await
}
fn forget(&self, id: &ProviderAccountId) {
PeerInfoService::forget(self, id);
}
async fn request_refresh(
&self,
id: &ProviderAccountId,
) -> Result<UpstreamRefreshReceipt, StoreError> {
PeerInfoService::request_refresh(self, id)
.await
.map(|plan| Self::receipt(&plan))
}
async fn read_quota(
&self,
account_id: &ProviderAccountId,
) -> Result<Option<ProviderQuota>, StoreError> {
PeerInfoService::read_quota(self, account_id).await
}
async fn read_summary(
&self,
account_id: &ProviderAccountId,
) -> Result<Option<ProviderQuota>, StoreError> {
PeerInfoService::read_summary(self, account_id).await
}
async fn read_quotas(
&self,
requests: &[crate::model::provider_credentials::ProviderQuotaReadRequest<'_>],
) -> Vec<Result<ProviderQuota, crate::ports::provider::ProviderAdminError>> {
accounts::read_peer_quotas(self, requests).await
}
}
@@ -0,0 +1,238 @@
//! 将已保存的互联观测投影为管理端额度与同步事实
use std::time::SystemTime;
use crate::model::{
provider_credentials::{
ProviderQuota, ProviderQuotaWindow, ProviderQuotaWindowRole, QuotaLocalUsageAttribution,
},
upstream_info::{
ProviderUpstreamInfo, UpstreamNode, UpstreamPlanCount, UpstreamPool, UpstreamPoolStates,
UpstreamQuotaAmount, UpstreamQuotaWindow, UpstreamRefreshReceipt, UpstreamSync,
UpstreamSyncState,
},
};
use gateway_core::account::{UpstreamInfoPlan, UpstreamInfoRecord};
use gateway_protocol::peer::{
PeerComponentStatus, PeerLink, PeerPoolDisclosure, PeerQuotaDisclosure,
};
use super::{PeerInfoService, snapshot};
impl PeerInfoService {
pub(crate) fn project(record: &UpstreamInfoRecord, now: SystemTime) -> ProviderQuota {
let peer = record.snapshot.as_ref().and_then(snapshot::decode);
Self::project_peer(&record.plan, peer.as_ref(), now)
}
pub(super) fn project_peer(
plan: &UpstreamInfoPlan,
peer: Option<&PeerLink>,
now: SystemTime,
) -> ProviderQuota {
let observed_at = peer
.as_ref()
.and_then(|peer| snapshot::datetime(&peer.observed_at));
let stale = peer.as_ref().is_none_or(|peer| {
!snapshot::component_fresh(Some(&peer.observed_at), Some(&peer.expires_at), now)
});
let info = ProviderUpstreamInfo {
discovered: plan.last_success_at.is_some(),
node: peer.as_ref().map(|peer| UpstreamNode {
id: peer.node.id.clone(),
name: peer.node.name.clone(),
software: peer.node.software.clone(),
}),
observed_at,
stale,
connection_revision: plan.connection.revision.clone(),
sync: sync_state(plan, now),
last_error: plan.last_error.map(|error| error.as_str().to_owned()),
pool: peer
.as_ref()
.filter(|peer| peer.pool.disclosure == PeerPoolDisclosure::Summary)
.and_then(|peer| {
let states = peer.pool.states.as_ref()?;
Some(UpstreamPool {
account_count: peer.pool.account_count?,
states: UpstreamPoolStates {
ready: states.ready,
cooldown: states.cooldown,
exhausted: states.exhausted,
disabled: states.disabled,
error: states.error,
unknown: states.unknown,
},
plans: peer
.pool
.plans
.as_deref()
.unwrap_or_default()
.iter()
.map(|plan| UpstreamPlanCount {
name: plan.name.clone(),
count: plan.count,
})
.collect(),
})
}),
};
let mut windows = Vec::new();
if let Some(peer) = &peer {
if peer.usage.status == PeerComponentStatus::Available
&& let Some(usage) = &peer.usage.data
{
for (key, label, seconds, window) in [
("daily", "日额度", 86400, &usage.daily),
("weekly", "周额度", 604800, &usage.weekly),
] {
// 不限额仍有日、周账本,只是不计算占用百分比
let total = window
.total
.as_deref()
.and_then(snapshot::decimal)
.filter(|value| *value > 0.0);
let used_percent = snapshot::decimal(&window.used)
.zip(total)
.map(|(used, total)| (used / total * 100.0).clamp(0.0, 100.0));
let window_observed = peer
.usage
.observed_at
.as_deref()
.and_then(snapshot::datetime);
let window_stale = stale
|| !peer
.usage
.observed_at
.as_deref()
.is_some_and(|value| snapshot::fresh(value, now))
|| window
.resets_at
.as_deref()
.and_then(snapshot::timestamp)
.is_some_and(|reset| reset <= now);
windows.push(ProviderQuotaWindow {
upstream: Some(UpstreamQuotaWindow {
account_id: None,
observed_at: window_observed,
stale: window_stale,
amount: Some(UpstreamQuotaAmount {
unit: usage.unit.clone(),
used: window.used.clone(),
total: window.total.clone(),
remaining: window.remaining.clone(),
}),
}),
key: format!("peer:budget:{key}"),
group: "general".to_owned(),
label: label.to_owned(),
limit_id: None,
limit_name: None,
role: Some(if key == "weekly" {
ProviderQuotaWindowRole::Secondary
} else {
ProviderQuotaWindowRole::Primary
}),
// Key 消费窗口覆盖此连接的全部请求,本机用量按同一时间边界聚合
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(seconds),
used_percent,
reset_at: window.resets_at.as_deref().and_then(snapshot::datetime),
limit_reached: !window_stale && snapshot::budget_exhausted(window, now),
local_usage: None,
provider_data: None,
});
}
}
// 池内窗口各自保留账号归属,不能将不同套餐相加为统一消费额度
if peer.quota.disclosure == PeerQuotaDisclosure::Accounts {
for (index, account) in peer.quota.accounts.iter().enumerate() {
let window_observed =
account.observed_at.as_deref().and_then(snapshot::datetime);
let window_stale = stale
|| !snapshot::component_fresh(
account.observed_at.as_deref(),
account.expires_at.as_deref(),
now,
);
for window in &account.windows {
let window_stale = window_stale
|| window
.reset_at
.as_deref()
.and_then(snapshot::timestamp)
.is_some_and(|reset| reset <= now);
let scope = format!(
"peer:{}:{}:{}",
account.id.len(),
account.id,
window.limit_id.as_deref().unwrap_or(&window.group)
);
windows.push(ProviderQuotaWindow {
upstream: Some(UpstreamQuotaWindow {
account_id: Some(account.id.clone()),
observed_at: window_observed,
stale: window_stale,
amount: None,
}),
key: format!("{scope}:{}", window.key),
group: window.group.clone(),
label: window.label.clone(),
limit_id: Some(scope),
limit_name: Some(match window.limit_name.as_deref() {
Some(name) => format!("账号 {} · {name}", index + 1),
None => format!("账号 {}", index + 1),
}),
role: match window.role.as_deref() {
Some("primary") => Some(ProviderQuotaWindowRole::Primary),
Some("secondary") => Some(ProviderQuotaWindowRole::Secondary),
Some("monthly") => Some(ProviderQuotaWindowRole::Monthly),
_ => None,
},
local_usage_attribution: QuotaLocalUsageAttribution::Unavailable,
window_seconds: window.window_seconds,
used_percent: window.used_percent,
reset_at: window.reset_at.as_deref().and_then(snapshot::datetime),
limit_reached: !window_stale && window.limit_reached,
local_usage: None,
provider_data: None,
});
}
}
}
}
ProviderQuota {
upstream: Some(info),
plan_type: None,
observed_at,
refresh_token_expires_at: None,
windows,
credits: None,
limit_reached: false,
provider_data: None,
}
}
pub(crate) fn receipt(record: &UpstreamInfoPlan) -> UpstreamRefreshReceipt {
UpstreamRefreshReceipt {
account_id: record.connection.account_id.as_str().to_owned(),
connection_revision: record.connection.revision.clone(),
request_generation: record.request_generation,
state: sync_state(record, SystemTime::now()).state,
}
}
}
fn sync_state(record: &UpstreamInfoPlan, now: SystemTime) -> UpstreamSync {
UpstreamSync {
state: if record.claimed_until.is_some_and(|expiry| expiry > now) {
UpstreamSyncState::Running
} else if record.request_generation > record.completed_generation {
UpstreamSyncState::Queued
} else {
UpstreamSyncState::Idle
},
request_generation: record.request_generation,
completed_generation: record.completed_generation,
}
}
@@ -0,0 +1,195 @@
//! 将互联观测压缩为按模型索引的调度事实,保留各层独立的有效期
use std::{collections::HashMap, sync::Arc, time::SystemTime};
use gateway_protocol::peer::{PeerComponentStatus, PeerLink, PeerModelCoverage, PeerModelState};
use super::snapshot;
use gateway_core::account::UpstreamSchedulingSignal;
#[derive(Clone, Copy)]
struct Validity {
observed_at: SystemTime,
expires_at: SystemTime,
}
impl Validity {
fn component(observed: Option<&str>, expires: Option<&str>) -> Option<Self> {
let observed_at = observed.and_then(snapshot::timestamp)?;
let expires_at = match expires {
Some(value) => snapshot::timestamp(value)?,
None => observed_at + snapshot::FRESH_FOR,
};
(expires_at <= observed_at + snapshot::MAX_AGE).then_some(Self {
observed_at,
expires_at,
})
}
fn fresh(self, now: SystemTime) -> bool {
self.observed_at <= now + std::time::Duration::from_secs(60) && self.expires_at > now
}
}
struct QuotaWindow {
remaining_percent: f64,
reset_at: Option<SystemTime>,
}
struct AccountQuota {
validity: Validity,
windows: Vec<QuotaWindow>,
}
impl AccountQuota {
fn remaining_percent(&self, now: SystemTime) -> Option<f64> {
if !self.validity.fresh(now) {
return None;
}
self.windows
.iter()
.filter(|window| window.reset_at.is_none_or(|reset| reset > now))
.map(|window| window.remaining_percent)
.reduce(f64::min)
}
}
struct ModelObservation {
validity: Validity,
exhausted: bool,
account: Option<Arc<AccountQuota>>,
}
pub(super) struct SchedulingObservation {
validity: Validity,
exhausted_budgets: Vec<Validity>,
models: HashMap<String, Option<ModelObservation>>,
}
impl SchedulingObservation {
pub(super) fn new(peer: &PeerLink, now: SystemTime) -> Option<Self> {
let validity = Validity::component(Some(&peer.observed_at), Some(&peer.expires_at))?;
let mut exhausted_budgets = Vec::with_capacity(2);
if peer.usage.status == PeerComponentStatus::Available
&& let Some(usage) = &peer.usage.data
&& let Some(usage_validity) =
Validity::component(peer.usage.observed_at.as_deref(), None)
{
for window in [&usage.daily, &usage.weekly] {
if snapshot::budget_exhausted(window, now) {
let mut validity = usage_validity;
if let Some(reset) = window.resets_at.as_deref().and_then(snapshot::timestamp) {
validity.expires_at = validity.expires_at.min(reset);
}
exhausted_budgets.push(validity);
}
}
}
// 同一账号重复列出模型不增加候选数,多账号覆盖不计算池平均百分比
let mut model_accounts = HashMap::<&str, Option<usize>>::new();
for (index, account) in peer.quota.accounts.iter().enumerate() {
for model in &account.models {
model_accounts
.entry(model)
.and_modify(|candidate| {
if *candidate != Some(index) {
*candidate = None;
}
})
.or_insert(Some(index));
}
}
let mut accounts = HashMap::<usize, Option<Arc<AccountQuota>>>::new();
let mut models = HashMap::with_capacity(peer.models.len());
for model in &peer.models {
models.entry(model.model.clone()).or_insert_with(|| {
if model.coverage != PeerModelCoverage::Complete
|| model.state == PeerModelState::Unknown
{
return None;
}
let validity =
Validity::component(model.observed_at.as_deref(), model.expires_at.as_deref())?;
let account = model_accounts
.get(model.model.as_str())
.copied()
.flatten()
.and_then(|index| {
accounts
.entry(index)
.or_insert_with(|| {
let account = &peer.quota.accounts[index];
let validity = Validity::component(
account.observed_at.as_deref(),
account.expires_at.as_deref(),
)?;
Some(Arc::new(AccountQuota {
validity,
windows: account
.windows
.iter()
.filter_map(|window| {
Some(QuotaWindow {
remaining_percent: 100.0 - window.used_percent?,
reset_at: window
.reset_at
.as_deref()
.and_then(snapshot::timestamp),
})
})
.collect(),
}))
})
.clone()
});
Some(ModelObservation {
validity,
exhausted: model.state == PeerModelState::Exhausted,
account,
})
});
}
Some(Self {
validity,
exhausted_budgets,
models,
})
}
pub(super) fn fresh(&self, now: SystemTime) -> bool {
self.validity.fresh(now)
&& now
.duration_since(self.validity.observed_at)
.unwrap_or_default()
<= snapshot::FRESH_FOR
}
pub(super) fn signal(&self, model: &str, now: SystemTime) -> Option<UpstreamSchedulingSignal> {
if !self.fresh(now) {
return None;
}
if self
.exhausted_budgets
.iter()
.any(|validity| validity.fresh(now))
{
return Some(UpstreamSchedulingSignal {
exhausted: true,
remaining_percent: None,
});
}
let model = self.models.get(model)?.as_ref()?;
if !model.validity.fresh(now) {
return None;
}
Some(UpstreamSchedulingSignal {
exhausted: model.exhausted,
remaining_percent: model
.account
.as_ref()
.and_then(|account| account.remaining_percent(now)),
})
}
}
@@ -0,0 +1,518 @@
//! 后台计划协调与仅供 Smart 排序的内存观测
use std::{
collections::{BTreeMap, HashMap},
sync::{Arc, RwLock},
time::{Duration, SystemTime},
};
use futures::{StreamExt as _, TryStreamExt as _, stream};
use gateway_core::{
account::{
LoadedCredential, ProviderAccount, ProviderAccountId, ProviderAccountStore,
UpstreamConnection, UpstreamInfoClaim, UpstreamInfoCompletion, UpstreamInfoFailure,
UpstreamInfoPlan, UpstreamInfoRecord, UpstreamInfoSource, UpstreamInfoStore,
UpstreamSchedulingSignal,
},
error::{StoreError, StoreErrorKind},
routing::ProviderKind,
task::WorkerCycleContext,
};
use crate::model::provider_credentials::ProviderQuota;
use gateway_protocol::peer::PeerLink;
use crate::ports::{
peer::{PeerConnection, PeerTransport},
provider::{ProviderAdmin, ProviderAdminRegistry},
};
use super::{scheduling::SchedulingObservation, snapshot};
const BATCH_SIZE: u32 = 100;
const FETCH_BATCH: u32 = 20;
const SNAPSHOT_READ_BATCH: usize = 2;
const CONCURRENCY: usize = 4;
const MINIMUM_INTERVAL: Duration = Duration::from_secs(15);
type CachedSummary = (UpstreamInfoPlan, Option<Arc<PeerLink>>);
struct CachedObservation {
connection: UpstreamConnection,
completed_generation: u64,
last_success_at: Option<SystemTime>,
checked_at: SystemTime,
scheduling: Option<Arc<SchedulingObservation>>,
summary: Option<Arc<PeerLink>>,
}
impl CachedObservation {
fn new(
plan: &UpstreamInfoPlan,
scheduling: Option<Arc<SchedulingObservation>>,
summary: Option<Arc<PeerLink>>,
now: SystemTime,
) -> Self {
Self {
connection: plan.connection.clone(),
completed_generation: plan.completed_generation,
last_success_at: plan.last_success_at,
checked_at: now,
scheduling,
summary,
}
}
fn matches(&self, plan: &UpstreamInfoPlan) -> bool {
self.connection == plan.connection
&& self.completed_generation == plan.completed_generation
&& self.last_success_at == plan.last_success_at
// 周期刷新不增加请求代次,撤销授权仍须立即丢弃此前的观测
&& (plan.last_error != Some(UpstreamInfoFailure::Unauthorized) || self.summary.is_none())
}
}
pub(crate) struct PeerInfoService {
accounts: Arc<dyn ProviderAccountStore>,
store: Arc<dyn UpstreamInfoStore>,
providers: BTreeMap<ProviderKind, Arc<dyn ProviderAdmin>>,
transport: Arc<dyn PeerTransport>,
cursors: tokio::sync::Mutex<BTreeMap<ProviderKind, Option<ProviderAccountId>>>,
cache_refresh: tokio::sync::Mutex<()>,
observations: RwLock<HashMap<ProviderAccountId, CachedObservation>>,
}
impl PeerInfoService {
pub(crate) fn new(
accounts: Arc<dyn ProviderAccountStore>,
store: Arc<dyn UpstreamInfoStore>,
providers: &ProviderAdminRegistry,
transport: Arc<dyn PeerTransport>,
) -> Self {
Self {
accounts,
store,
providers: providers.registered(),
transport,
cursors: tokio::sync::Mutex::new(BTreeMap::new()),
cache_refresh: tokio::sync::Mutex::new(()),
observations: RwLock::new(HashMap::new()),
}
}
pub(crate) async fn account_changed(&self, id: &ProviderAccountId) -> Result<(), StoreError> {
let loaded = self.accounts.load_current_credential(id).await?;
self.reconcile(&loaded).await
}
pub(crate) fn forget(&self, id: &ProviderAccountId) {
let previous = self
.observations
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.remove(id);
drop(previous);
}
async fn reconcile(&self, loaded: &LoadedCredential) -> Result<(), StoreError> {
let connection = UpstreamConnection::for_account(&loaded.account);
let enabled = self.connection(loaded).is_some();
self.store
.reconcile(
&connection,
enabled && loaded.account.enabled(),
SystemTime::now(),
)
.await?;
let _refresh = self.cache_refresh.lock().await;
let mut observations = self
.observations
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let previous = if !enabled
|| !loaded.account.enabled()
|| observations
.get(loaded.account.id())
.is_some_and(|cached| cached.connection != connection)
{
observations.remove(loaded.account.id())
} else {
None
};
drop(observations);
drop(previous);
Ok(())
}
pub(crate) async fn request_refresh(
&self,
id: &ProviderAccountId,
) -> Result<UpstreamInfoPlan, StoreError> {
let loaded = self.accounts.load_current_credential(id).await?;
if self.connection(&loaded).is_none() {
return Err(StoreError::new(StoreErrorKind::InvalidData));
}
self.reconcile(&loaded).await?;
self.store
.request_refresh(
&UpstreamConnection::for_account(&loaded.account),
SystemTime::now(),
MINIMUM_INTERVAL,
)
.await
}
pub(crate) async fn read_quota(
&self,
id: &ProviderAccountId,
) -> Result<Option<ProviderQuota>, StoreError> {
let Some(connection) = self.current_connection(id).await? else {
return Ok(None);
};
Ok(self
.read(std::slice::from_ref(id))
.await?
.iter()
.find(|record| record.plan.connection == connection)
.map(|record| Self::project(record, SystemTime::now())))
}
pub(crate) async fn read_summary(
&self,
id: &ProviderAccountId,
) -> Result<Option<ProviderQuota>, StoreError> {
let Some(connection) = self.current_connection(id).await? else {
return Ok(None);
};
Ok(self
.read_summaries(std::slice::from_ref(id))
.await?
.into_iter()
.find(|(current, _)| *current == connection)
.map(|(_, quota)| quota))
}
async fn current_connection(
&self,
id: &ProviderAccountId,
) -> Result<Option<UpstreamConnection>, StoreError> {
let account = self
.accounts
.get_account(id)
.await?
.ok_or_else(|| StoreError::new(StoreErrorKind::InvalidData))?;
if !self
.providers
.get(account.provider())
.is_some_and(|provider| {
provider.supports_peer_connection(account.authentication_kind())
})
{
return Ok(None);
}
Ok(Some(UpstreamConnection::for_account(&account)))
}
pub(crate) async fn read(
&self,
ids: &[ProviderAccountId],
) -> Result<Vec<UpstreamInfoRecord>, StoreError> {
self.store.read(ids).await
}
pub(super) async fn synchronize(&self, context: &WorkerCycleContext) -> Result<(), StoreError> {
let now = SystemTime::now();
// 过期观测保留轻量版本,避免每轮重读旧 JSON;长期未扫描到的账号释放缓存
let expired = self
.observations
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.extract_if(|_, cached| {
now.duration_since(cached.checked_at).unwrap_or_default() > snapshot::MAX_AGE
})
.collect::<Vec<_>>();
drop(expired);
let mut failure = None;
for provider in self.providers.keys() {
if context.cancellation().is_cancelled() {
break;
}
// 一个 Provider 的账号读取失败不阻止其他 Provider 本轮同步
if let Err(error) = self.synchronize_provider(provider, context).await {
failure.get_or_insert(error);
}
}
failure.map_or(Ok(()), Err)
}
fn connection(&self, loaded: &LoadedCredential) -> Option<PeerConnection> {
self.providers
.get(loaded.account.provider())?
.peer_connection(loaded)
.ok()
.flatten()
}
async fn synchronize_provider(
&self,
provider: &ProviderKind,
context: &WorkerCycleContext,
) -> Result<(), StoreError> {
let mut cursors = self.cursors.lock().await;
let cursor = cursors.entry(provider.clone()).or_default();
let candidates = self
.store
.list_candidates(
provider,
self.providers[provider].peer_authentication_kinds(),
cursor.as_ref(),
BATCH_SIZE,
)
.await?;
*cursor = if candidates.len() < BATCH_SIZE as usize {
None
} else {
candidates.last().map(|loaded| loaded.account.id().clone())
};
drop(cursors);
let mut ids = Vec::with_capacity(candidates.len());
for loaded in candidates {
if context.cancellation().is_cancelled() {
return Ok(());
}
match self.reconcile(&loaded).await {
Ok(()) => ids.push(loaded.account.id().clone()),
Err(error)
if matches!(
error.kind(),
StoreErrorKind::Conflict | StoreErrorKind::InvalidData
) => {}
Err(error) => return Err(error),
}
}
self.refresh_cache(&ids).await?;
let pending = self
.store
.claim_due(
provider,
SystemTime::now(),
FETCH_BATCH,
Duration::from_secs(60),
)
.await?;
stream::iter(pending)
.map(|record| self.synchronize_one(record))
.buffer_unordered(CONCURRENCY)
.try_collect::<Vec<_>>()
.await?;
Ok(())
}
async fn synchronize_one(&self, claim: UpstreamInfoClaim) -> Result<(), StoreError> {
let record = claim.plan;
let loaded = match self
.accounts
.load_current_credential(&record.connection.account_id)
.await
{
Ok(loaded) => loaded,
Err(error)
if matches!(
error.kind(),
StoreErrorKind::Conflict | StoreErrorKind::InvalidData
) =>
{
return Ok(());
}
Err(error) => return Err(error),
};
if UpstreamConnection::for_account(&loaded.account) != record.connection
|| !loaded.account.enabled()
{
return self.reconcile(&loaded).await;
}
let Some(authentication) = self.connection(&loaded) else {
return self.reconcile(&loaded).await;
};
let (result, retry_after) =
match self.transport.fetch(&loaded.account, &authentication).await {
Ok(bytes) => (snapshot::parse(&bytes, SystemTime::now()), None),
Err(failure) => (Err(failure.kind), failure.retry_after),
};
let now = SystemTime::now();
let delay = match result {
Ok(_) => Duration::from_secs(240 + stable_jitter(loaded.account.id().as_str())),
Err(UpstreamInfoFailure::Unauthorized | UpstreamInfoFailure::Unsupported) => {
Duration::from_secs(30 * 60)
}
Err(_) => Duration::from_secs(
30_u64
.saturating_mul(1_u64 << record.consecutive_failures.min(6))
.min(30 * 60),
),
}
.max(retry_after.unwrap_or_default());
let completion = UpstreamInfoCompletion {
connection: record.connection.clone(),
request_generation: record.request_generation,
claim_token: claim.claim_token,
completed_at: now,
next_attempt_at: now + delay,
result,
};
if self.store.complete(completion).await? {
self.refresh_cache(std::slice::from_ref(loaded.account.id()))
.await?;
}
Ok(())
}
pub(crate) async fn read_summaries(
&self,
ids: &[ProviderAccountId],
) -> Result<Vec<(UpstreamConnection, ProviderQuota)>, StoreError> {
let plans = self.store.read_plans(ids).await?;
let cached = {
let observations = self
.observations
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
plans
.into_iter()
.map(|plan| {
let cached = observations.get(&plan.connection.account_id)?;
cached
.matches(&plan)
.then(|| (plan, cached.summary.clone()))
})
.collect::<Option<Vec<_>>>()
};
let summaries = match cached {
Some(summaries) => summaries,
None => self.refresh_cache(ids).await?,
};
let now = SystemTime::now();
Ok(summaries
.into_iter()
.map(|(plan, summary)| {
let quota = Self::project_peer(&plan, summary.as_deref(), now);
(plan.connection, quota)
})
.collect())
}
async fn refresh_cache(
&self,
ids: &[ProviderAccountId],
) -> Result<Vec<CachedSummary>, StoreError> {
let _refresh = self.cache_refresh.lock().await;
let mut plans = self.store.read_plans(ids).await?;
let now = SystemTime::now();
let mut changed = Vec::new();
let mut replaced = Vec::new();
{
let mut observations = self
.observations
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
for plan in &plans {
if let Some(cached) = observations.get_mut(&plan.connection.account_id)
&& cached.matches(plan)
{
cached.checked_at = now;
continue;
}
if plan.last_success_at.is_none()
|| plan.last_error == Some(UpstreamInfoFailure::Unauthorized)
{
replaced.push(observations.insert(
plan.connection.account_id.clone(),
CachedObservation::new(plan, None, None, now),
));
} else {
changed.push(plan.connection.account_id.clone());
}
}
}
// 调度投影可能被其他请求持有,替换后的最后一次释放也留在锁外
drop(replaced);
// 完整观测可接近协议大小上限,投影后及时释放,避免同时物化整页 JSON
for batch in changed.chunks(SNAPSHOT_READ_BATCH) {
for record in self.store.read(batch).await? {
if let Some(plan) = plans
.iter_mut()
.find(|plan| plan.connection.account_id == record.plan.connection.account_id)
{
*plan = record.plan.clone();
}
self.cache(&record);
}
}
let observations = self
.observations
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
Ok(plans
.into_iter()
.map(|plan| {
let summary = observations
.get(&plan.connection.account_id)
.filter(|cached| cached.matches(&plan))
.and_then(|cached| cached.summary.clone());
(plan, summary)
})
.collect())
}
fn cache(&self, record: &UpstreamInfoRecord) {
let now = SystemTime::now();
let mut peer = record.snapshot.as_ref().and_then(snapshot::decode);
let scheduling = peer
.as_ref()
.and_then(|peer| SchedulingObservation::new(peer, now))
.filter(|observation| observation.fresh(now))
.map(Arc::new);
let summary = peer.take().map(|mut peer| {
peer.quota = Default::default();
peer.models = Vec::new();
Arc::new(peer)
});
let cached = CachedObservation::new(&record.plan, scheduling, summary, now);
let previous = self
.observations
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(record.plan.connection.account_id.clone(), cached);
drop(previous);
}
}
impl UpstreamInfoSource for PeerInfoService {
/// 数据面只取不可变投影;模型查询与有效期判断不持有共享锁
fn signal(
&self,
account: &ProviderAccount,
model: &str,
now: SystemTime,
) -> Option<UpstreamSchedulingSignal> {
let connection = UpstreamConnection::for_account(account);
let scheduling = {
let observations = self
.observations
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let cached = observations.get(account.id())?;
if cached.connection != connection {
return None;
}
Arc::clone(cached.scheduling.as_ref()?)
};
scheduling.signal(model, now)
}
}
fn stable_jitter(id: &str) -> u64 {
id.bytes().fold(0_u64, |hash, byte| {
hash.wrapping_mul(31).wrapping_add(u64::from(byte))
}) % 60
}
@@ -0,0 +1,292 @@
//! 互联响应的有界校验、能力授权与时间有效性解释
use std::time::{Duration, SystemTime};
use chrono::{DateTime, Utc};
use gateway_core::account::{OpaqueProviderData, UpstreamInfoFailure};
use gateway_protocol::peer::{
MAX_ACCOUNTS, MAX_MODELS, MAX_QUOTA_WINDOWS, OBJECT, PeerLink, PeerPoolDisclosure,
SCHEMA_VERSION,
};
use serde::Deserialize as _;
pub(super) const MAX_AGE: Duration = Duration::from_secs(30 * 60);
pub(super) const FRESH_FOR: Duration = Duration::from_secs(5 * 60);
pub(super) fn timestamp(value: &str) -> Option<SystemTime> {
DateTime::parse_from_rfc3339(value)
.ok()
.map(SystemTime::from)
}
pub(super) fn datetime(value: &str) -> Option<DateTime<Utc>> {
DateTime::parse_from_rfc3339(value)
.ok()
.map(|value| value.with_timezone(&Utc))
}
pub(super) fn fresh(observed: &str, now: SystemTime) -> bool {
timestamp(observed).is_some_and(|observed| {
observed <= now + Duration::from_secs(60)
&& now.duration_since(observed).unwrap_or_default() <= FRESH_FOR
})
}
pub(super) fn component_fresh(
observed: Option<&str>,
expires: Option<&str>,
now: SystemTime,
) -> bool {
let Some(observed) = observed.and_then(timestamp) else {
return false;
};
let expiry = match expires {
Some(value) => match timestamp(value) {
Some(expiry) => expiry,
None => return false,
},
None => observed + FRESH_FOR,
};
observed <= now + Duration::from_secs(60) && expiry > now && expiry <= observed + MAX_AGE
}
pub(super) fn decode(data: &OpaqueProviderData) -> Option<PeerLink> {
PeerLink::deserialize(
serde::de::value::MapDeserializer::<_, serde_json::Error>::new(
data.expose_to_provider()
.iter()
.map(|(key, value)| (key.as_str(), value)),
),
)
.ok()
}
pub(super) fn parse(
bytes: &[u8],
now: SystemTime,
) -> Result<OpaqueProviderData, UpstreamInfoFailure> {
let mut peer: PeerLink =
serde_json::from_slice(bytes).map_err(|_| UpstreamInfoFailure::InvalidResponse)?;
let valid_text = |value: &str, max: usize| {
!value.is_empty() && value.len() <= max && !value.chars().any(char::is_control)
};
let observed = timestamp(&peer.observed_at).ok_or(UpstreamInfoFailure::InvalidResponse)?;
let expires = timestamp(&peer.expires_at).ok_or(UpstreamInfoFailure::InvalidResponse)?;
if peer.object != OBJECT
|| peer.schema_version != SCHEMA_VERSION
|| !valid_text(&peer.node.id, 256)
|| !valid_text(&peer.node.name, 256)
|| !valid_text(&peer.node.software, 128)
|| observed > now + Duration::from_secs(60)
|| now.duration_since(observed).unwrap_or_default() > MAX_AGE
|| expires <= observed
|| expires > observed + MAX_AGE
|| peer.capabilities.len() > 32
|| peer
.quota
.accounts
.iter()
.map(|account| account.windows.len())
.sum::<usize>()
> MAX_QUOTA_WINDOWS
|| peer.quota.accounts.len() > MAX_ACCOUNTS
|| peer.models.len() > MAX_MODELS
{
return Err(UpstreamInfoFailure::InvalidResponse);
}
// 未授权能力不能借同名数据块越过逐层公开边界
if !peer.capabilities.iter().any(|value| value == "key_usage") {
peer.usage.data = None;
}
if !peer
.capabilities
.iter()
.any(|value| value == "pool_summary")
{
peer.pool = Default::default();
}
if !peer
.capabilities
.iter()
.any(|value| value == "quota_windows")
{
peer.quota = Default::default();
}
if !peer
.capabilities
.iter()
.any(|value| value == "model_availability")
{
peer.models.clear();
}
if peer.pool.disclosure == PeerPoolDisclosure::Summary {
let count = peer
.pool
.account_count
.ok_or(UpstreamInfoFailure::InvalidResponse)?;
let states = peer
.pool
.states
.as_ref()
.ok_or(UpstreamInfoFailure::InvalidResponse)?;
let state_total = [
states.ready,
states.cooldown,
states.exhausted,
states.disabled,
states.error,
states.unknown,
]
.into_iter()
.map(u64::from)
.sum::<u64>();
if u64::from(count) > MAX_ACCOUNTS as u64 || state_total != u64::from(count) {
return Err(UpstreamInfoFailure::InvalidResponse);
}
if let Some(plans) = &peer.pool.plans
&& (plans.len() > MAX_ACCOUNTS
|| plans
.iter()
.any(|plan| !valid_text(&plan.name, 256) || plan.count > count)
|| plans.iter().map(|plan| u64::from(plan.count)).sum::<u64>() > u64::from(count))
{
return Err(UpstreamInfoFailure::InvalidResponse);
}
}
if let Some(usage) = &peer.usage.data {
if !valid_text(&usage.unit, 16) {
return Err(UpstreamInfoFailure::InvalidResponse);
}
for window in [&usage.daily, &usage.weekly] {
if [
Some(&window.used),
window.total.as_ref(),
window.remaining.as_ref(),
]
.into_iter()
.flatten()
.any(|value| decimal(value).is_none())
|| window
.resets_at
.as_ref()
.is_some_and(|value| timestamp(value).is_none())
{
return Err(UpstreamInfoFailure::InvalidResponse);
}
}
}
for account in &peer.quota.accounts {
if !valid_text(&account.id, 256)
|| account.models.len() > MAX_MODELS
|| account.windows.iter().any(|window| {
!valid_text(&window.key, 256)
|| !valid_text(&window.label, 256)
|| !valid_text(&window.group, 128)
|| window
.limit_id
.as_deref()
.is_some_and(|value| !valid_text(value, 256))
|| window
.limit_name
.as_deref()
.is_some_and(|value| !valid_text(value, 256))
|| window
.used_percent
.is_some_and(|value| !value.is_finite() || !(0.0..=100.0).contains(&value))
|| window
.reset_at
.as_ref()
.is_some_and(|value| timestamp(value).is_none())
})
{
return Err(UpstreamInfoFailure::InvalidResponse);
}
}
for model in &peer.models {
if !valid_text(&model.model, 256)
|| model
.observed_at
.as_deref()
.is_some_and(|value| timestamp(value).is_none())
|| model
.expires_at
.as_deref()
.is_some_and(|value| timestamp(value).is_none())
{
return Err(UpstreamInfoFailure::InvalidResponse);
}
}
match serde_json::to_value(peer) {
Ok(serde_json::Value::Object(object)) => Ok(OpaqueProviderData::new(object)),
_ => Err(UpstreamInfoFailure::InvalidResponse),
}
}
pub(super) fn decimal(value: &str) -> Option<f64> {
if value.is_empty()
|| value.len() > 64
|| value
.bytes()
.any(|byte| !byte.is_ascii_digit() && byte != b'.')
|| value.bytes().filter(|byte| *byte == b'.').count() > 1
{
return None;
}
value
.parse::<f64>()
.ok()
.filter(|value| value.is_finite() && *value >= 0.0)
}
pub(super) fn budget_exhausted(
window: &gateway_protocol::peer::PeerBudgetWindow,
now: SystemTime,
) -> bool {
let Some(total) = window.total.as_deref() else {
return false;
};
if compare_decimal(total, "0") != std::cmp::Ordering::Greater
|| window
.resets_at
.as_deref()
.and_then(timestamp)
.is_some_and(|reset| reset <= now)
{
return false;
}
window
.remaining
.as_deref()
.is_some_and(|remaining| compare_decimal(remaining, "0") == std::cmp::Ordering::Equal)
|| compare_decimal(&window.used, total) != std::cmp::Ordering::Less
}
// 金额触顶比较保留十进制精度,百分比展示的浮点近似不能成为耗尽依据
fn compare_decimal(left: &str, right: &str) -> std::cmp::Ordering {
let parts = |value: &str| {
let (integer, fractional) = value.split_once('.').unwrap_or((value, ""));
(
integer.trim_start_matches('0').to_owned(),
fractional.to_owned(),
)
};
let (left_integer, left_fractional) = parts(left);
let (right_integer, right_fractional) = parts(right);
left_integer
.len()
.cmp(&right_integer.len())
.then_with(|| left_integer.cmp(&right_integer))
.then_with(|| {
let width = left_fractional.len().max(right_fractional.len());
left_fractional
.bytes()
.chain(std::iter::repeat(b'0'))
.take(width)
.cmp(
right_fractional
.bytes()
.chain(std::iter::repeat(b'0'))
.take(width),
)
})
}
@@ -0,0 +1,50 @@
//! 将互联同步交给 Host 周期监督并响应取消
use std::{sync::Arc, time::Duration};
use futures::future::BoxFuture;
use gateway_core::task::{
ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerDefinitionError, WorkerId,
WorkerKind, WorkerLeaseRequest, WorkerRegistration, WorkerRunnable, WorkerSchedule,
WorkerTaskError,
};
use super::PeerInfoService;
impl PeerInfoService {
pub(crate) fn worker(self: &Arc<Self>) -> Result<WorkerContribution, WorkerDefinitionError> {
let id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, "peer-info")?;
let ttl = Duration::from_secs(120);
let schedule = WorkerSchedule::try_new(
Duration::from_secs(15),
Duration::from_secs(1),
Duration::from_secs(60),
ttl,
Duration::from_secs(30),
)?;
let lease = WorkerLeaseRequest::try_new(id.clone(), ttl)?;
Ok(WorkerContribution::Registration(
WorkerRegistration::try_new(
id,
WorkerRunnable::Scheduled {
schedule,
lease: Some(lease),
task: Box::new(PeerInfoTask(Arc::clone(self))),
},
)?,
))
}
}
struct PeerInfoTask(Arc<PeerInfoService>);
impl ScheduledTask for PeerInfoTask {
fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
Box::pin(async move {
tokio::select! {
() = context.cancellation().cancelled() => Ok(()),
result = self.0.synchronize(&context) => result.map_err(|error| WorkerTaskError::safe("upstream information synchronization failed").with_source(error)),
}
})
}
}
@@ -29,6 +29,7 @@ pub(crate) struct DefaultPluginAccountAccess {
providers: ProviderAdminRegistry,
accounts: Arc<dyn AccountStore>,
snapshot: Arc<dyn SnapshotControl>,
upstream_info: Option<Arc<dyn crate::ports::peer::PeerUpstreamService>>,
}
impl DefaultPluginAccountAccess {
@@ -37,11 +38,13 @@ impl DefaultPluginAccountAccess {
providers: ProviderAdminRegistry,
accounts: Arc<dyn AccountStore>,
snapshot: Arc<dyn SnapshotControl>,
upstream_info: Option<Arc<dyn crate::ports::peer::PeerUpstreamService>>,
) -> Self {
Self {
providers,
accounts,
snapshot,
upstream_info,
}
}
@@ -51,14 +54,19 @@ impl DefaultPluginAccountAccess {
refresh: bool,
) -> Result<crate::model::provider_credentials::ProviderQuota, AdminError> {
let account = self.get_runtime(account_id).await?;
self.provider(&account.provider_kind)?
.quota(crate::model::provider_credentials::ProviderQuotaRequest {
let provider = self.provider(&account.provider_kind)?;
super::peer_link::upstream::read_quota(
self.upstream_info.as_deref(),
provider.as_ref(),
&account,
crate::model::provider_credentials::ProviderQuotaRequest {
account_id: account_id.clone(),
refresh,
rolling_usage: None,
})
.await
.map_err(|error| map_provider_error(error, "plugin quota facts"))
},
)
.await
.map_err(|error| map_provider_error(error, "plugin quota facts"))
}
fn provider(
@@ -165,6 +173,7 @@ impl PluginAccountAccess for DefaultPluginAccountAccess {
.credential
.credential_revision;
publish_credentials_and_observe_quota(
self.upstream_info.as_deref(),
&provider,
self.snapshot.as_ref(),
result.config_revision,
@@ -208,6 +217,7 @@ impl PluginAccountAccess for DefaultPluginAccountAccess {
.credential_revision
.ok_or_else(|| AdminError::internal("账号保存未返回凭据版本"))?;
publish_credentials_and_observe_quota(
self.upstream_info.as_deref(),
&provider,
self.snapshot.as_ref(),
result.config_revision,
@@ -83,6 +83,7 @@ impl PluginResourceAccess for DefaultPluginResourceAccess {
resource_key,
groups,
NewClientKey {
peer_sharing: Default::default(),
id,
name: command.name,
label: None,
@@ -87,6 +87,7 @@ pub(crate) struct DefaultProxiesService {
probe: Arc<dyn ProxyProbe>,
snapshot: Arc<dyn SnapshotControl>,
providers: ProviderAdminRegistry,
upstream_info: Option<Arc<dyn crate::ports::peer::PeerUpstreamService>>,
test_slots: Semaphore,
}
@@ -96,12 +97,14 @@ impl DefaultProxiesService {
probe: Arc<dyn ProxyProbe>,
snapshot: Arc<dyn SnapshotControl>,
providers: ProviderAdminRegistry,
upstream_info: Option<Arc<dyn crate::ports::peer::PeerUpstreamService>>,
) -> Self {
Self {
store,
probe,
snapshot,
providers,
upstream_info,
test_slots: Semaphore::new(4),
}
}
@@ -226,19 +229,30 @@ impl ProxiesService for DefaultProxiesService {
.map(|mut account| async {
// 与账号目录一致,缺失套餐时读取已有额度快照;限制并发且不触发上游刷新
let mut cached_quota = None;
if explicit_plan_type(account.plan_type.as_deref()).is_none()
&& let Ok(kind) = ProviderKind::new(account.provider_kind.clone())
if let Ok(kind) = ProviderKind::new(account.provider_kind.clone())
&& let Ok(provider) = self.providers.require(&kind)
&& let Ok(account_id) = ProviderAccountId::new(account.id.clone())
&& let Ok(quota) = provider
.quota(ProviderQuotaRequest {
account_id,
refresh: false,
rolling_usage: None,
})
.await
{
cached_quota = Some(quota);
if provider.supports_peer_connection(&account.authentication_kind)
&& let Some(peer) = self.upstream_info.as_deref()
{
account.upstream = peer
.read_summary(&account_id)
.await
.ok()
.flatten()
.and_then(|quota| quota.upstream);
}
if explicit_plan_type(account.plan_type.as_deref()).is_none() {
cached_quota = provider
.quota(ProviderQuotaRequest {
account_id,
refresh: false,
rolling_usage: None,
})
.await
.ok();
}
}
account.plan_type_display = self.providers.resolve_account_plan(
&account.provider_kind,
@@ -107,6 +107,7 @@ fn provider_document_debug_should_not_expose_opaque_material() {
#[test]
fn representative_quota_should_prefer_short_window_and_highest_usage() {
let quota = ProviderQuota {
upstream: None,
credits: None,
plan_type: None,
observed_at: None,
@@ -129,6 +130,7 @@ fn representative_quota_should_prefer_account_wide_window_over_model_specific_wi
let mut model_specific = quota_window("shortTerm", Some(18_000), Some(64.0));
model_specific.local_usage_attribution = QuotaLocalUsageAttribution::Unavailable;
let quota = ProviderQuota {
upstream: None,
credits: None,
plan_type: None,
observed_at: None,
@@ -167,11 +169,13 @@ fn usage_window_should_not_use_daily_rolling_usage_for_weekly_statistics() {
models: Vec::new(),
};
let quota = ProviderQuota {
upstream: None,
credits: None,
plan_type: None,
observed_at: None,
refresh_token_expires_at: None,
windows: vec![ProviderQuotaWindow {
upstream: None,
key: "free-rolling-24h".to_owned(),
group: "shortTerm".to_owned(),
label: "日限额".to_owned(),
@@ -196,6 +200,7 @@ fn usage_window_should_not_use_daily_rolling_usage_for_weekly_statistics() {
#[test]
fn exhausted_quota_should_project_full_usage_to_only_the_representative_window() {
let mut quota = ProviderQuota {
upstream: None,
credits: None,
plan_type: None,
observed_at: None,
@@ -225,6 +230,7 @@ fn exhausted_quota_should_preserve_the_provider_identified_reached_window() {
let mut reached = quota_window("monthly", Some(2_592_000), Some(98.0));
reached.limit_reached = true;
let mut quota = ProviderQuota {
upstream: None,
credits: None,
plan_type: None,
observed_at: None,
@@ -247,6 +253,7 @@ fn quota_window(
used_percent: Option<f64>,
) -> ProviderQuotaWindow {
ProviderQuotaWindow {
upstream: None,
key: group.to_owned(),
group: group.to_owned(),
label: group.to_owned(),
@@ -86,6 +86,7 @@ fn selected_usage_window(quota: &ProviderQuota) -> Option<(&str, AccountUsagePer
fn usage_quota(windows: Vec<ProviderQuotaWindow>) -> ProviderQuota {
ProviderQuota {
upstream: None,
credits: None,
plan_type: None,
observed_at: None,
@@ -98,6 +99,7 @@ fn usage_quota(windows: Vec<ProviderQuotaWindow>) -> ProviderQuota {
fn usage_window(key: &str, group: &str, seconds: u64) -> ProviderQuotaWindow {
ProviderQuotaWindow {
upstream: None,
key: key.to_owned(),
group: group.to_owned(),
label: key.to_owned(),
@@ -136,6 +138,7 @@ fn usage_window(key: &str, group: &str, seconds: u64) -> ProviderQuotaWindow {
#[test]
fn dashboard_quota_should_preserve_unknown_and_actual_window_facts() {
let mut quota = ProviderQuota {
upstream: None,
credits: None,
plan_type: None,
observed_at: None,
@@ -148,6 +151,7 @@ fn dashboard_quota_should_preserve_unknown_and_actual_window_facts() {
assert!(quota.representative_used_percent().is_none());
let rolling = ProviderQuotaWindow {
upstream: None,
key: "provider_rolling".to_owned(),
group: "requests".to_owned(),
label: "Provider rolling window".to_owned(),
@@ -180,6 +184,7 @@ fn dashboard_quota_should_preserve_unknown_and_actual_window_facts() {
quota.windows.insert(
0,
ProviderQuotaWindow {
upstream: None,
key: "model_specific".to_owned(),
used_percent: Some(100.0),
local_usage_attribution: QuotaLocalUsageAttribution::Unavailable,
@@ -15,6 +15,7 @@ fn now() -> DateTime<Utc> {
fn window(key: &str, days: u64) -> ProviderQuotaWindow {
ProviderQuotaWindow {
upstream: None,
key: key.to_owned(),
group: if days > 7 { "monthly" } else { "shortTerm" }.to_owned(),
label: key.to_owned(),
@@ -58,6 +59,7 @@ fn window(key: &str, days: u64) -> ProviderQuotaWindow {
fn quota(windows: Vec<ProviderQuotaWindow>) -> ProviderQuota {
ProviderQuota {
upstream: None,
credits: None,
plan_type: None,
observed_at: Some(now()),
@@ -1473,6 +1473,7 @@ async fn plugin_account_save_publishes_native_credential_change() {
ProviderAdminRegistry::new([provider.clone() as Arc<dyn ProviderAdmin>]).unwrap(),
store.clone(),
snapshot,
None,
);
let result = access
.save(
@@ -1516,6 +1517,7 @@ async fn plugin_account_reads_resolve_the_authoritative_provider_from_the_accoun
ProviderAdminRegistry::new([registered_provider as Arc<dyn ProviderAdmin>]).unwrap(),
store,
Arc::new(RecordingPluginAccountPublication(events.clone())),
None,
);
let account_id = ProviderAccountId::new("acct_test").expect("account ID");
@@ -1563,6 +1565,7 @@ async fn plugin_account_save_rejects_stale_authentication_kind_before_commit() {
ProviderAdminRegistry::new([provider as Arc<dyn ProviderAdmin>]).unwrap(),
store.clone(),
Arc::new(RecordingPluginAccountPublication(events.clone())),
None,
);
let error = access
@@ -1692,9 +1695,11 @@ async fn accounts_recover_disabled_should_only_enable_scheduling_and_preserve_qu
);
}
let quota = ProviderQuota {
upstream: None,
credits: None,
observed_at: Some(Utc::now()),
windows: vec![ProviderQuotaWindow {
upstream: None,
key: "primary".to_owned(),
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
@@ -1996,6 +2001,7 @@ async fn accounts_should_fill_missing_plan_from_quota_without_overriding_known_s
] {
let provider = FakeProviderAdmin::new("openai", events());
provider.set_quota(ProviderQuota {
upstream: None,
plan_type: quota_plan.map(str::to_owned),
..empty_quota()
});
@@ -2039,11 +2045,13 @@ async fn accounts_list_should_degrade_quota_failure_to_empty_window_without_drop
let events = events();
let openai = FakeProviderAdmin::new("openai", events.clone());
openai.set_quota(ProviderQuota {
upstream: None,
credits: None,
plan_type: None,
observed_at: Some(Utc::now()),
refresh_token_expires_at: None,
windows: vec![ProviderQuotaWindow {
upstream: None,
key: "primary".to_owned(),
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
@@ -2197,11 +2205,13 @@ async fn accounts_list_should_map_unknown_credential_to_error_not_normal() {
async fn accounts_list_should_not_derive_rate_limited_from_provider_quota_view() {
let provider = FakeProviderAdmin::new("openai", events());
provider.set_quota(ProviderQuota {
upstream: None,
credits: None,
plan_type: None,
observed_at: Some(Utc::now()),
refresh_token_expires_at: None,
windows: vec![ProviderQuotaWindow {
upstream: None,
key: "primary".to_owned(),
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
@@ -2246,11 +2256,13 @@ async fn accounts_list_should_not_derive_rate_limited_from_provider_quota_view()
async fn accounts_list_should_not_derive_exhaustion_from_provider_quota_view() {
let provider = FakeProviderAdmin::new("openai", events());
provider.set_quota(ProviderQuota {
upstream: None,
credits: None,
plan_type: None,
observed_at: Some(Utc::now()),
refresh_token_expires_at: None,
windows: vec![ProviderQuotaWindow {
upstream: None,
key: "primary".to_owned(),
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
@@ -2298,10 +2310,12 @@ async fn quota_forecast_reads_raw_snapshot_and_limits_usage_to_observation_time(
let observed = now - TimeDelta::hours(1);
let reset = now + TimeDelta::days(1);
provider.set_quota(ProviderQuota {
upstream: None,
credits: None,
observed_at: Some(observed),
limit_reached: true,
windows: vec![ProviderQuotaWindow {
upstream: None,
key: "week".to_owned(),
group: "shortTerm".to_owned(),
label: "周额度".to_owned(),
@@ -2390,6 +2404,7 @@ async fn quota_forecast_only_queries_the_two_selected_source_windows() {
.map(|index| {
let monthly = index % 2 == 1;
ProviderQuotaWindow {
upstream: None,
key: if monthly {
format!("month-{index}")
} else {
@@ -2411,6 +2426,7 @@ async fn quota_forecast_only_queries_the_two_selected_source_windows() {
})
.collect();
provider.set_quota(ProviderQuota {
upstream: None,
plan_type: Some("pro".to_owned()),
observed_at: Some(observed),
windows,
@@ -2462,9 +2478,11 @@ async fn quota_forecast_same_window_reversal_keeps_recorded_usage_in_the_report(
let reset = now + TimeDelta::days(1);
let provider = FakeProviderAdmin::new("openai", events());
provider.set_quota(ProviderQuota {
upstream: None,
plan_type: Some("pro".to_owned()),
observed_at: Some(observed),
windows: vec![ProviderQuotaWindow {
upstream: None,
key: "week".to_owned(),
group: "shortTerm".to_owned(),
label: "周额度".to_owned(),
@@ -2540,10 +2558,12 @@ async fn quota_forecast_mid_cycle_sampling_accepts_small_reset_jitter_but_not_a_
let added = now - TimeDelta::hours(5);
let provider = FakeProviderAdmin::new("openai", events());
provider.set_quota(ProviderQuota {
upstream: None,
credits: None,
plan_type: Some("pro".to_owned()),
observed_at: Some(observed),
windows: vec![ProviderQuotaWindow {
upstream: None,
key: "week".to_owned(),
group: "shortTerm".to_owned(),
label: "周额度".to_owned(),
@@ -2636,6 +2656,163 @@ async fn quota_forecast_mid_cycle_sampling_accepts_small_reset_jitter_but_not_a_
assert_eq!(cycle.estimated_tokens, Some(679));
}
#[tokio::test]
async fn peer_account_usage_should_follow_budget_windows_without_borrowing_pool_windows() {
use gateway_admin::model::upstream_info::{
ProviderUpstreamInfo, UpstreamQuotaAmount, UpstreamQuotaWindow, UpstreamSync,
UpstreamSyncState,
};
for longest in [86_400, 604_800] {
let provider = FakeProviderAdmin::new("openai", events());
let mut account = account_record("openai");
account.authentication_kind = "api_key".to_owned();
account.created_at = Utc::now() - TimeDelta::days(60);
let store = FakeAccountStore::with_account(account, events());
let reset_at = Utc::now() + TimeDelta::hours(1);
let mut quota = empty_quota();
quota.upstream = Some(ProviderUpstreamInfo {
discovered: true,
node: None,
observed_at: Some(Utc::now()),
stale: false,
connection_revision: "1".to_owned(),
sync: UpstreamSync {
state: UpstreamSyncState::Idle,
request_generation: 1,
completed_generation: 1,
},
last_error: None,
pool: None,
});
quota.windows = [86_400, 604_800, 30 * 86_400]
.into_iter()
.filter(|seconds| *seconds <= longest || *seconds == 30 * 86_400)
.map(|seconds| ProviderQuotaWindow {
upstream: Some(UpstreamQuotaWindow {
account_id: None,
observed_at: Some(Utc::now()),
stale: false,
amount: (seconds <= longest).then(|| UpstreamQuotaAmount {
unit: "USD".to_owned(),
used: "1".to_owned(),
remaining: Some("9".to_owned()),
total: Some("10".to_owned()),
}),
}),
key: seconds.to_string(),
group: "general".to_owned(),
label: "额度".to_owned(),
limit_id: None,
limit_name: None,
role: None,
local_usage_attribution: if seconds <= longest {
QuotaLocalUsageAttribution::AccountWide
} else {
QuotaLocalUsageAttribution::Unavailable
},
window_seconds: Some(seconds),
used_percent: Some(10.0),
reset_at: Some(reset_at),
limit_reached: false,
local_usage: None,
provider_data: None,
})
.collect();
provider.set_quota(quota.clone());
let window_usage = [86_400, 604_800]
.into_iter()
.filter(|seconds| *seconds <= longest)
.map(|seconds| AccountUsageWindowResult {
account_id: "acct_test".to_owned(),
key: seconds.to_string(),
usage: quota_local_usage("acct_test", if seconds == longest { 1_234 } else { 123 }),
})
.collect::<Vec<_>>();
store.set_quota_window_usage(window_usage.clone());
let services = accounts_service(provider.clone(), store.clone()).await;
let page = services
.accounts()
.list(AccountListQuery {
page: 1,
page_size: gateway_admin::model::PageSize::new(20).unwrap(),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
})
.await
.unwrap();
assert_eq!(
page.items[0].usage.as_ref().unwrap().total_tokens,
Some(1_234)
);
let queries = store.quota_window_queries();
assert_eq!(queries.len(), window_usage.len());
for (query, expected) in queries.iter().zip(&window_usage) {
assert_eq!(query.key, expected.key);
assert_eq!(
query.range.start,
reset_at - TimeDelta::seconds(query.key.parse().unwrap())
);
assert_eq!(query.range.end, reset_at);
let window = page.items[0]
.quota
.windows
.iter()
.find(|window| window.key == query.key)
.unwrap();
assert_eq!(window.local_usage.as_ref(), Some(&expected.usage));
assert_eq!(window.used_percent, Some(10.0));
assert_eq!(
window
.upstream
.as_ref()
.unwrap()
.amount
.as_ref()
.unwrap()
.used,
"1"
);
}
assert!(
page.items[0]
.quota
.windows
.last()
.unwrap()
.local_usage
.is_none()
);
let id = ProviderAccountId::new("acct_test").unwrap();
let detail = services.accounts().quota(&id, false).await.unwrap();
assert_eq!(detail.usage, page.items[0].usage);
assert_eq!(detail.quota.windows, page.items[0].quota.windows);
assert_eq!(store.quota_window_queries()[0].range, queries[0].range);
quota
.windows
.iter_mut()
.find(|window| window.window_seconds == Some(longest))
.unwrap()
.reset_at = None;
provider.set_quota(quota);
store.set_quota_window_usage(
window_usage
.into_iter()
.filter(|usage| usage.key != longest.to_string())
.collect(),
);
let unopened = services.accounts().quota(&id, false).await.unwrap();
assert!(
unopened.usage.is_none(),
"未开启的最长窗口不能借用日窗口或历史累计"
);
}
}
#[tokio::test]
async fn api_key_list_and_detail_should_accumulate_local_usage_without_subscription_windows() {
let provider = FakeProviderAdmin::new("openai", events());
@@ -2687,11 +2864,13 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
let provider = FakeProviderAdmin::new("openai", events());
let reset_at = Utc::now() + TimeDelta::hours(1);
provider.set_quota(ProviderQuota {
upstream: None,
credits: None,
plan_type: None,
observed_at: Some(Utc::now()),
refresh_token_expires_at: None,
windows: vec![ProviderQuotaWindow {
upstream: None,
key: "primary".to_owned(),
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
@@ -2756,6 +2935,7 @@ async fn accounts_list_and_quota_refresh_should_select_the_same_weekly_or_monthl
let account_id = ProviderAccountId::new("acct_test").expect("account id");
let reset_at = Utc::now() + TimeDelta::days(1);
let short = ProviderQuotaWindow {
upstream: None,
key: "short".to_owned(),
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
@@ -2771,12 +2951,14 @@ async fn accounts_list_and_quota_refresh_should_select_the_same_weekly_or_monthl
provider_data: None,
};
let week = ProviderQuotaWindow {
upstream: None,
key: "week".to_owned(),
label: "周额度".to_owned(),
window_seconds: Some(7 * 86_400),
..short.clone()
};
let month = ProviderQuotaWindow {
upstream: None,
key: "month".to_owned(),
group: "monthly".to_owned(),
label: "月额度".to_owned(),
@@ -2788,6 +2970,7 @@ async fn accounts_list_and_quota_refresh_should_select_the_same_weekly_or_monthl
(vec![short, month], "month", 3000, 30),
] {
provider.set_quota(ProviderQuota {
upstream: None,
windows,
..empty_quota()
});
@@ -2845,12 +3028,14 @@ async fn accounts_list_should_not_attach_account_usage_to_model_specific_quota_w
let provider = FakeProviderAdmin::new("openai", events());
let reset_at = Utc::now() + TimeDelta::days(7);
provider.set_quota(ProviderQuota {
upstream: None,
credits: None,
plan_type: None,
observed_at: Some(Utc::now()),
refresh_token_expires_at: None,
windows: vec![
ProviderQuotaWindow {
upstream: None,
key: "core-primary".to_owned(),
group: "weekly".to_owned(),
label: "周额度".to_owned(),
@@ -2866,6 +3051,7 @@ async fn accounts_list_should_not_attach_account_usage_to_model_specific_quota_w
provider_data: None,
},
ProviderQuotaWindow {
upstream: None,
key: "codex-bengalfox-primary".to_owned(),
group: "weekly".to_owned(),
label: "周额度".to_owned(),
@@ -3147,6 +3333,7 @@ pub(super) fn document() -> ProviderDocument {
fn empty_quota() -> ProviderQuota {
ProviderQuota {
upstream: None,
credits: None,
plan_type: None,
observed_at: None,
@@ -3591,6 +3778,7 @@ async fn plugin_quota_refresh_uses_native_provider_without_changing_configuratio
ProviderAdminRegistry::new([provider.clone() as Arc<dyn ProviderAdmin>]).unwrap(),
store,
Arc::new(RecordingPluginAccountPublication(events.clone())),
None,
);
let id = ProviderAccountId::new("acct_test").unwrap();
access.get_quota(&id).await.unwrap();
@@ -726,7 +726,9 @@ async fn threshold_services(
fn weekly_quota(used_percent: f64) -> ProviderQuota {
ProviderQuota {
upstream: None,
windows: vec![ProviderQuotaWindow {
upstream: None,
key: "week".into(),
group: "shortTerm".into(),
label: "周额度".into(),
@@ -870,6 +872,7 @@ async fn automation_defers_missing_invalid_or_failed_weekly_quota() {
(empty_quota(), None),
(
ProviderQuota {
upstream: None,
windows: vec![short, model, month],
..empty_quota()
},
@@ -877,6 +880,7 @@ async fn automation_defers_missing_invalid_or_failed_weekly_quota() {
),
(
ProviderQuota {
upstream: None,
windows: vec![missing_percent],
..empty_quota()
},
@@ -96,6 +96,7 @@ impl ClientKeyStore for TestClientKeyStore {
}
self.plaintexts.lock().unwrap().push(key.plaintext);
let record = ClientKeyRecord {
peer_sharing: key.peer_sharing,
request_profile_overrides: Default::default(),
id: key.id,
name: key.name,
@@ -234,6 +235,7 @@ async fn plugin_client_key_list_projects_only_public_identity_and_uses_stable_na
list_response: Mutex::new(Some(ClientKeyPage {
config_revision: Revision::new(7).expect("revision"),
items: vec![ClientKeyRecord {
peer_sharing: Default::default(),
request_profile_overrides: Default::default(),
id: key_id.clone(),
name: "Public identity".to_owned(),
@@ -370,6 +372,7 @@ async fn duplicate_names_report_the_same_actionable_conflict_on_create_and_updat
.update(
&mutation_context(),
UpdateClientKey {
peer_sharing: None,
request_profile_override_updates: Default::default(),
id: ClientApiKeyId::new("key_existing").unwrap(),
name: "Migration".to_owned(),
@@ -390,6 +393,7 @@ async fn duplicate_names_report_the_same_actionable_conflict_on_create_and_updat
fn create_command(key: Option<&str>) -> CreateClientKey {
CreateClientKey {
peer_sharing: Default::default(),
request_profile_overrides: Default::default(),
custom_key: key.map(|value| PlaintextClientApiKey::new(value).unwrap()),
name: "Migration".to_owned(),
@@ -400,6 +404,53 @@ fn create_command(key: Option<&str>) -> CreateClientKey {
}
}
#[tokio::test]
async fn peer_sharing_requires_explicit_groups_and_independent_plan_authorization_before_writing() {
use gateway_core::policy::ClientPeerSharing;
let store = Arc::new(TestClientKeyStore::default());
let services = super::AdminHarness::new()
.client_keys(store.clone())
.build()
.await;
for sharing in [
ClientPeerSharing {
enabled: true,
..Default::default()
},
ClientPeerSharing {
share_plans: true,
..Default::default()
},
] {
let mut command = create_command(None);
command.peer_sharing = sharing;
let error = services
.client_keys()
.create(&mutation_context(), command)
.await
.unwrap_err();
assert_eq!(error.kind(), AdminErrorKind::Invalid);
}
assert!(store.plaintexts.lock().unwrap().is_empty());
let mut command = create_command(None);
command.peer_sharing = ClientPeerSharing {
enabled: true,
share_pool: true,
share_quota: true,
share_plans: true,
};
command.group_ids.push(
gateway_core::routing::AccountGroupId::new(format!("grp_{}", "1".repeat(32))).unwrap(),
);
let expected = command.peer_sharing;
let result = services
.client_keys()
.create(&mutation_context(), command)
.await
.unwrap();
assert_eq!(result.secret.record.peer_sharing, expected);
}
fn mutation_context() -> MutationContext {
MutationContext {
actor: MutationActor::System,
@@ -436,6 +487,7 @@ async fn duplicate_group_ids_are_rejected_before_key_creation_or_update() {
.update(
&mutation_context(),
UpdateClientKey {
peer_sharing: None,
request_profile_override_updates: Default::default(),
id: ClientApiKeyId::new("key_existing").unwrap(),
name: "Migration".to_owned(),
@@ -10,6 +10,7 @@ mod credentials;
mod freeze_recovery;
mod import_tasks;
mod observability;
mod peer_link;
mod plugin_update;
mod plugins;
mod proxies;
@@ -286,6 +287,7 @@ impl AdminHarness {
Arc::new(plugins::TestPluginPorts),
),
gateway_admin::AdminRuntimePorts {
upstream_info: None,
timezone: Default::default(),
service_middleware: self.service_middleware,
plugin_preparation: Arc::new(plugins::TestPluginPorts),
@@ -909,6 +911,7 @@ impl SettingsStore for UnavailableStore {
}
struct UnavailableProvider {
peer_enabled: bool,
kind: ProviderKind,
dashboard_profile: Option<DashboardWireProfile>,
calculated_billing: Option<gateway_admin::model::observability::CalculatedBillingBreakdown>,
@@ -931,6 +934,7 @@ impl UnavailableProvider {
fn new(kind: &str) -> Self {
Self {
kind: ProviderKind::new(kind).expect("provider kind"),
peer_enabled: false,
dashboard_profile: None,
calculated_billing: None,
}
@@ -943,6 +947,28 @@ pub(super) fn test_provider(kind: &str) -> Arc<dyn ProviderAdmin> {
#[async_trait]
impl ProviderAdmin for UnavailableProvider {
fn peer_authentication_kinds(&self) -> &[&str] {
if self.peer_enabled {
&["peer_test"]
} else {
&[]
}
}
fn peer_connection(
&self,
loaded: &gateway_core::account::LoadedCredential,
) -> Result<Option<gateway_admin::ports::peer::PeerConnection>, ProviderAdminError> {
Ok(
(self.peer_enabled && loaded.account.authentication_kind() == "peer_test").then(|| {
gateway_admin::ports::peer::PeerConnection {
base_url: format!("https://{}.invalid/v1", self.kind),
bearer_token: secrecy::SecretString::from("test-peer-only"),
}
}),
)
}
fn provider_kind(&self) -> &ProviderKind {
&self.kind
}
@@ -1031,6 +1057,7 @@ impl ProviderAdmin for UnavailableProvider {
pub(super) fn dashboard_profile_provider() -> Arc<dyn ProviderAdmin> {
Arc::new(UnavailableProvider {
peer_enabled: false,
kind: ProviderKind::new("openai").expect("provider kind"),
dashboard_profile: Some(DashboardWireProfile {
provider: "openai".to_owned(),
@@ -1074,6 +1101,7 @@ pub(super) fn calculated_billing_provider() -> Arc<dyn ProviderAdmin> {
.expect("test billing amount"),
};
Arc::new(UnavailableProvider {
peer_enabled: false,
kind: ProviderKind::new("openai").expect("provider kind"),
dashboard_profile: None,
calculated_billing: Some(
@@ -995,6 +995,7 @@ async fn dashboard_should_preserve_provider_plan_names() {
*store.account_usage.lock().unwrap() = plans
.iter()
.map(|(plan, _)| DashboardAccountUsage {
upstream: None,
account_id: format!("account-{plan}"),
provider_kind: "openai".to_owned(),
authentication_kind: "oauth".to_owned(),
@@ -0,0 +1,3 @@
//! 互联共享与上游同步用例测试
mod upstream;
@@ -0,0 +1,398 @@
//! 公共互联按 Provider 声明发现连接,不将认证种类或调度观测绑在 OpenAI 上
use super::super::UnavailableProvider;
use async_trait::async_trait;
use chrono::Utc;
use gateway_admin::{
initialize_peer_upstream,
ports::{
peer::{PeerConnection, PeerFetchFailure, PeerTransport},
provider::{ProviderAdmin, ProviderAdminRegistry},
},
};
use gateway_core::{
account::*,
error::{StoreError, StoreErrorKind},
lifecycle::CancellationToken,
routing::ProviderKind,
task::{WorkerContribution, WorkerCycleContext, WorkerRunnable},
};
use std::{
collections::BTreeMap,
sync::{Arc, Mutex},
time::{Duration, SystemTime},
};
#[tokio::test]
async fn discovery_uses_provider_capabilities_and_publishes_neutral_signals() {
for failed_provider in [None, Some("openai")] {
let providers = ["openai", "xai"].map(|kind| {
let mut provider = UnavailableProvider::new(kind);
provider.peer_enabled = true;
Arc::new(provider) as Arc<dyn ProviderAdmin>
});
let registry = ProviderAdminRegistry::new(providers).unwrap();
let accounts = Arc::new(TestAccounts(vec![
account("acct_first", "openai", "peer_test"),
account("acct_second", "xai", "peer_test"),
account("acct_native", "xai", "oauth"),
]));
let store = Arc::new(TestPlans {
accounts: accounts.clone(),
records: Mutex::default(),
failed_provider,
});
let transport = Arc::new(TestTransport(Mutex::new(Vec::new())));
let mut bundle =
initialize_peer_upstream(accounts.clone(), store, &registry, transport.clone())
.unwrap();
let reader = UpstreamInfoReader::default();
assert!(
reader
.signal(&accounts.0[0], "model", SystemTime::now())
.is_none()
);
let source = bundle.scheduling_source();
reader.bind(&source).unwrap();
assert!(reader.bind(&source).is_err());
drop(source);
let WorkerContribution::Registration(worker) = bundle.take_worker_contributions().remove(0)
else {
panic!("scheduled worker")
};
let WorkerRunnable::Scheduled { ref task, .. } = worker.runnable else {
panic!("scheduled worker")
};
let result = task
.run_cycle(WorkerCycleContext::new(
worker.id.clone(),
None,
CancellationToken::new(),
))
.await;
assert_eq!(result.is_err(), failed_provider.is_some());
assert_eq!(
*transport.0.lock().unwrap(),
if failed_provider.is_some() {
vec!["acct_second"]
} else {
vec!["acct_first", "acct_second"]
}
);
for account in &accounts.0[..2] {
if Some(account.provider().as_str()) == failed_provider {
assert!(reader.signal(account, "model", SystemTime::now()).is_none());
continue;
}
let observed = bundle
.service()
.read_quota(account.id())
.await
.unwrap()
.unwrap();
assert_eq!(
observed.upstream.unwrap().node.unwrap().name,
account.provider().as_str()
);
assert!(
reader
.signal(account, "model", SystemTime::now())
.unwrap()
.exhausted
);
}
assert!(
reader
.signal(&accounts.0[2], "model", SystemTime::now())
.is_none()
);
drop(worker);
drop(bundle);
assert!(
reader
.signal(&accounts.0[0], "model", SystemTime::now())
.is_none()
);
}
}
fn account(id: &str, provider: &str, authentication: &str) -> ProviderAccount {
ProviderAccount::new(
ProviderAccountId::new(id).unwrap(),
ProviderKind::new(provider).unwrap(),
id.to_owned(),
None,
authentication.to_owned(),
CredentialRevision::new(1).unwrap(),
None,
)
.with_account_facts(
true,
CredentialState::Ready,
QuotaState::unknown(),
None,
None,
)
}
struct TestTransport(Mutex<Vec<String>>);
#[async_trait]
impl PeerTransport for TestTransport {
async fn fetch(
&self,
account: &ProviderAccount,
connection: &PeerConnection,
) -> Result<Vec<u8>, PeerFetchFailure> {
assert_eq!(
connection.base_url,
format!("https://{}.invalid/v1", account.provider())
);
self.0
.lock()
.unwrap()
.push(account.id().as_str().to_owned());
let now = Utc::now();
Ok(serde_json::to_vec(&serde_json::json!({
"object":"cpr.peer", "schema_version":1,
"node":{"id":account.id().as_str(),"name":account.provider().as_str(),"software":"cpr"},
"observed_at":now.to_rfc3339(),
"expires_at":(now + chrono::Duration::minutes(5)).to_rfc3339(),
"capabilities":["key_usage"],
"usage":{"status":"available","observed_at":now.to_rfc3339(),"data":{"unit":"USD",
"daily":{"total":"1","used":"1","remaining":"0","resetsAt":(now + chrono::Duration::days(1)).to_rfc3339()},
"weekly":{"total":"0","used":"0","remaining":"0"}}},
"pool":{"disclosure":"none"}, "quota":{"disclosure":"none"}, "models":[]
})).unwrap())
}
}
struct TestPlans {
failed_provider: Option<&'static str>,
accounts: Arc<TestAccounts>,
records: Mutex<BTreeMap<ProviderAccountId, UpstreamInfoRecord>>,
}
#[async_trait]
impl UpstreamInfoStore for TestPlans {
async fn list_candidates(
&self,
provider: &ProviderKind,
kinds: &[&str],
after: Option<&ProviderAccountId>,
limit: u32,
) -> Result<Vec<LoadedCredential>, StoreError> {
if Some(provider.as_str()) == self.failed_provider {
return Err(StoreError::new(StoreErrorKind::Unavailable));
}
Ok(self
.accounts
.0
.iter()
.filter(|account| {
account.provider() == provider
&& kinds.contains(&account.authentication_kind())
&& after.is_none_or(|id| account.id() > id)
})
.take(limit as usize)
.map(loaded)
.collect())
}
async fn reconcile(
&self,
connection: &UpstreamConnection,
enabled: bool,
now: SystemTime,
) -> Result<(), StoreError> {
assert!(enabled);
self.records
.lock()
.unwrap()
.entry(connection.account_id.clone())
.or_insert_with(|| UpstreamInfoRecord {
plan: UpstreamInfoPlan {
connection: connection.clone(),
request_generation: 1,
completed_generation: 0,
next_attempt_at: now,
claimed_until: None,
last_attempt_at: None,
last_success_at: None,
last_error: None,
consecutive_failures: 0,
},
snapshot: None,
});
Ok(())
}
async fn request_refresh(
&self,
_: &UpstreamConnection,
_: SystemTime,
_: Duration,
) -> Result<UpstreamInfoPlan, StoreError> {
panic!("automatic discovery does not enqueue a manual refresh")
}
async fn read_plans(
&self,
ids: &[ProviderAccountId],
) -> Result<Vec<UpstreamInfoPlan>, StoreError> {
Ok(self
.read(ids)
.await?
.into_iter()
.map(|record| record.plan)
.collect())
}
async fn read(&self, ids: &[ProviderAccountId]) -> Result<Vec<UpstreamInfoRecord>, StoreError> {
let records = self.records.lock().unwrap();
Ok(ids
.iter()
.filter_map(|id| records.get(id).cloned())
.collect())
}
async fn claim_due(
&self,
provider: &ProviderKind,
now: SystemTime,
limit: u32,
ttl: Duration,
) -> Result<Vec<UpstreamInfoClaim>, StoreError> {
Ok(self
.records
.lock()
.unwrap()
.values_mut()
.filter(|record| {
self.accounts.0.iter().any(|account| {
account.id() == &record.plan.connection.account_id
&& account.provider() == provider
}) && record.plan.next_attempt_at <= now
})
.take(limit as usize)
.map(|record| {
record.plan.claimed_until = Some(now + ttl);
record.plan.last_attempt_at = Some(now);
UpstreamInfoClaim {
plan: record.plan.clone(),
claim_token: uuid::Uuid::new_v4(),
}
})
.collect())
}
async fn complete(&self, completion: UpstreamInfoCompletion) -> Result<bool, StoreError> {
let mut records = self.records.lock().unwrap();
let record = records.get_mut(&completion.connection.account_id).unwrap();
assert_eq!(record.plan.connection, completion.connection);
record.snapshot = Some(completion.result.unwrap());
record.plan.completed_generation = completion.request_generation;
record.plan.last_success_at = Some(completion.completed_at);
record.plan.claimed_until = None;
record.plan.next_attempt_at = completion.next_attempt_at;
Ok(true)
}
}
struct TestAccounts(Vec<ProviderAccount>);
fn loaded(account: &ProviderAccount) -> LoadedCredential {
LoadedCredential {
account: account.clone(),
credential: PlaintextCredential::new(Default::default()),
}
}
#[async_trait]
impl ProviderAccountStore for TestAccounts {
async fn create_account(&self, _account: NewProviderAccount) -> Result<(), StoreError> {
panic!("unused account operation")
}
async fn get_account(
&self,
account_id: &ProviderAccountId,
) -> Result<Option<ProviderAccount>, StoreError> {
Ok(self
.0
.iter()
.find(|account| account.id() == account_id)
.cloned())
}
async fn list_accounts(&self) -> Result<Vec<ProviderAccount>, StoreError> {
panic!("unused account operation")
}
async fn list_for_provider(
&self,
_provider: &ProviderKind,
) -> Result<Vec<ProviderAccount>, StoreError> {
panic!("unused account operation")
}
async fn list_refresh_candidates(
&self,
_query: ProviderRefreshQuery,
) -> Result<Vec<LoadedCredential>, StoreError> {
panic!("unused account operation")
}
async fn load_credential(
&self,
_account: &ProviderAccountId,
_expected_revision: CredentialRevision,
) -> Result<LoadedCredential, StoreError> {
panic!("unused account operation")
}
async fn load_current_credential(
&self,
id: &ProviderAccountId,
) -> Result<LoadedCredential, StoreError> {
self.0
.iter()
.find(|account| account.id() == id)
.map(loaded)
.ok_or_else(|| StoreError::new(StoreErrorKind::InvalidData))
}
async fn compare_and_swap_credential(
&self,
_update: CredentialCasUpdate,
) -> Result<CredentialCasOutcome, StoreError> {
panic!("unused account operation")
}
async fn get_quotas(
&self,
_accounts: &[ProviderAccountId],
) -> Result<Vec<QuotaObservation>, StoreError> {
panic!("unused account operation")
}
async fn compare_and_swap_quota(
&self,
_observation: QuotaObservation,
) -> Result<QuotaWriteOutcome, StoreError> {
panic!("unused account operation")
}
async fn touch_quota_observation(
&self,
_touch: QuotaObservationTouch,
) -> Result<QuotaWriteOutcome, StoreError> {
panic!("unused account operation")
}
async fn apply_quota_access(
&self,
_change: QuotaAccessChange,
) -> Result<QuotaWriteOutcome, StoreError> {
panic!("unused account operation")
}
async fn apply_state_change(&self, _change: AccountStateChange) -> Result<(), StoreError> {
panic!("unused account operation")
}
async fn update_account(&self, _update: ProviderAccountUpdate) -> Result<(), StoreError> {
panic!("unused account operation")
}
async fn set_enabled(
&self,
_account: &ProviderAccountId,
_enabled: bool,
) -> Result<(), StoreError> {
panic!("unused account operation")
}
async fn delete_account(&self, _account: &ProviderAccountId) -> Result<(), StoreError> {
panic!("unused account operation")
}
}
@@ -203,6 +203,7 @@ async fn linked_accounts_share_plan_resolution_and_only_read_cached_quota() {
] {
let provider = FakeProviderAdmin::new("openai", events());
provider.set_quota(ProviderQuota {
upstream: None,
credits: None,
plan_type: cached.map(str::to_owned),
observed_at: None,
@@ -215,6 +216,7 @@ async fn linked_accounts_share_plan_resolution_and_only_read_cached_quota() {
.provider(provider.clone())
.proxies(Arc::new(TestProxies {
accounts: Some(vec![ProxyAccountRef {
upstream: None,
id: "acct_plan".to_owned(),
name: "套餐测试".to_owned(),
email: None,
@@ -17,6 +17,10 @@ where
get(super::concurrency::list::<S>),
)
.route("/api/admin/accounts/detail", get(account_detail::<S>))
.route(
"/api/admin/accounts/upstream-info/refresh",
post(super::upstream_info::refresh::<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>))
@@ -44,7 +44,9 @@ use serde::{Deserialize, Deserializer, Serialize};
use serde_json::{Map, Value};
use uuid::{Uuid, Version};
use super::presenter::{format_compact_number, format_decimal_currency, format_number};
use super::presenter::{
format_compact_number, format_decimal_currency, format_number, format_percentage,
};
use super::{
AdminAuth, AdminEnvelope, AdminError, AdminJson, AdminQuery, AdminResponse, PageMeta,
WireValidationError,
@@ -66,11 +68,14 @@ mod credentials;
mod handlers;
mod import_tasks;
mod presenter;
mod upstream_info;
mod wire;
pub use credentials::*;
pub use handlers::{profile_avatar_response, router};
pub(super) use presenter::quota_window_view;
pub(crate) use upstream_info::view as upstream_info_view;
pub use upstream_info::{AccountQuotaAmountView, AccountUpstreamInfoView};
pub use wire::*;
use credentials::{
@@ -92,11 +92,31 @@ pub(super) fn account_view(
let usage_period = quota.usage_window().map(|(_, period)| period);
let mut usage = account_usage_view(usage, usage_period, now, time);
if account.authentication_kind == "api_key" {
usage.window_kind = AccountUsageWindowKind::General;
usage.window_kind = if quota
.upstream
.as_ref()
.is_some_and(|upstream| upstream.discovered)
{
match quota
.upstream_usage_window()
.and_then(|window| window.window_seconds)
{
Some(86_400) => AccountUsageWindowKind::Daily,
Some(604_800) => AccountUsageWindowKind::Weekly,
_ => AccountUsageWindowKind::Quota,
}
} else {
AccountUsageWindowKind::General
};
}
let upstream = quota
.upstream
.clone()
.map(|value| super::upstream_info::view(value, now, time));
let (mut quota, refresh_token_expires_at) = account_quota_view(quota, cooldown, now, time);
quota.protection_rules = protection_rules;
AccountView {
upstream,
capabilities: capabilities.into(),
id: account.id.clone(),
name: account.name,
@@ -277,6 +297,7 @@ pub(crate) fn quota_window_view(
time: crate::time::TimePresenter,
) -> AccountQuotaWindowView {
let ProviderQuotaWindow {
upstream,
key,
group,
label,
@@ -295,6 +316,16 @@ pub(crate) fn quota_window_view(
.as_ref()
.map_or_else(|| label.clone(), |name| format!("{name} · {label}"));
AccountQuotaWindowView {
source_account_id: upstream.as_ref().and_then(|value| value.account_id.clone()),
amount: upstream
.as_ref()
.and_then(|value| value.amount.as_ref())
.map(super::upstream_info::amount_view),
observed_at_display: upstream
.as_ref()
.and_then(|value| value.observed_at)
.map(|value| time.datetime(&value)),
stale: upstream.as_ref().is_some_and(|value| value.stale),
label_display,
window_label_display: label,
key,
@@ -304,8 +335,7 @@ pub(crate) fn quota_window_view(
role: role.map(|value| value.as_str().to_owned()),
window_seconds,
used_percent,
used_percent_display: used_percent
.map_or_else(|| "—".to_owned(), |value| format!("{value:.1}%")),
used_percent_display: used_percent.map_or_else(|| "—".to_owned(), format_percentage),
limit_reached,
local_usage: local_usage
.as_ref()
@@ -0,0 +1,191 @@
//! 互联快照展示与异步刷新 HTTP 合同
use gateway_admin::model::upstream_info::{ProviderUpstreamInfo, UpstreamQuotaAmount};
use super::*;
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountUpstreamInfoView {
pub discovered: bool,
pub node: Option<UpstreamNodeView>,
pub observed_at: Option<String>,
pub observed_at_display: String,
pub observed_at_relative: Option<crate::time::RelativeTimeView>,
pub stale: bool,
pub connection_revision: String,
pub sync: UpstreamSyncView,
pub last_error: Option<String>,
pub pool: Option<UpstreamPoolView>,
}
#[derive(Debug, Clone, Serialize)]
pub struct UpstreamNodeView {
pub id: String,
pub name: String,
pub software: String,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct UpstreamSyncView {
pub state: String,
pub request_generation: u64,
pub completed_generation: u64,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct UpstreamPoolView {
pub account_count: u32,
pub states: UpstreamPoolStatesView,
pub plans: Vec<UpstreamPlanView>,
}
#[derive(Debug, Clone, Serialize)]
pub struct UpstreamPoolStatesView {
pub ready: u32,
pub cooldown: u32,
pub exhausted: u32,
pub disabled: u32,
pub error: u32,
pub unknown: u32,
}
#[derive(Debug, Clone, Serialize)]
pub struct UpstreamPlanView {
pub name: String,
pub count: u32,
}
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct RefreshReceiptView {
account_id: String,
connection_revision: String,
request_generation: u64,
state: String,
}
#[derive(Serialize)]
struct RefreshData {
items: Vec<RefreshReceiptView>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountQuotaAmountView {
pub used_display: String,
pub remaining_display: String,
pub total_display: String,
}
pub(super) fn amount_view(amount: &UpstreamQuotaAmount) -> AccountQuotaAmountView {
let display = |value: Option<&str>| {
value.map_or_else(
|| "∞".to_owned(),
|value| crate::admin::presenter::format_decimal_currency(value, &amount.unit),
)
};
AccountQuotaAmountView {
used_display: display(Some(&amount.used)),
remaining_display: display(amount.remaining.as_deref()),
total_display: display(amount.total.as_deref()),
}
}
pub(crate) fn view(
info: ProviderUpstreamInfo,
now: DateTime<Utc>,
time: crate::time::TimePresenter,
) -> AccountUpstreamInfoView {
AccountUpstreamInfoView {
discovered: info.discovered,
node: info.node.map(|node| UpstreamNodeView {
id: node.id,
name: node.name,
software: node.software,
}),
observed_at: info.observed_at.map(|value| time.rfc3339(&value)),
observed_at_display: info
.observed_at
.map_or_else(|| "—".to_owned(), |value| time.datetime(&value)),
observed_at_relative: info.observed_at.map(|value| time.relative_view(value, now)),
stale: info.stale,
connection_revision: info.connection_revision,
sync: UpstreamSyncView {
state: info.sync.state.as_str().to_owned(),
request_generation: info.sync.request_generation,
completed_generation: info.sync.completed_generation,
},
last_error: info.last_error,
pool: info.pool.map(|pool| UpstreamPoolView {
account_count: pool.account_count,
states: UpstreamPoolStatesView {
ready: pool.states.ready,
cooldown: pool.states.cooldown,
exhausted: pool.states.exhausted,
disabled: pool.states.disabled,
error: pool.states.error,
unknown: pool.states.unknown,
},
plans: pool
.plans
.into_iter()
.map(|plan| UpstreamPlanView {
name: plan.name,
count: plan.count,
})
.collect(),
}),
}
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub(super) struct RefreshRequest {
account_ids: Vec<String>,
}
pub(super) async fn refresh<S>(
_auth: AdminAuth,
State(state): State<S>,
AdminJson(request): AdminJson<RefreshRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: crate::auth::SessionState + Send + Sync,
{
if request.account_ids.is_empty() || request.account_ids.len() > 200 {
return Err(map_wire_error(WireValidationError::new("accountIds")));
}
let ids = request
.account_ids
.into_iter()
.collect::<BTreeSet<_>>()
.into_iter()
.map(|id| {
require_account_id(&id, "accountIds")?;
ProviderAccountId::new(id).map_err(|_| WireValidationError::new("accountIds"))
})
.collect::<Result<Vec<_>, _>>()
.map_err(map_wire_error)?;
let items = state
.admin_services()
.accounts()
.refresh_upstream_info(ids)
.await
.map_err(map_service_error)?;
let items = items
.into_iter()
.map(|item| RefreshReceiptView {
account_id: item.account_id,
connection_revision: item.connection_revision,
request_generation: item.request_generation,
state: item.state.as_str().to_owned(),
})
.collect();
Ok(AdminResponse::new(
StatusCode::ACCEPTED,
AdminEnvelope::ok(RefreshData { items }),
))
}
@@ -266,6 +266,7 @@ pub struct AccountSummaryView {
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountView {
pub upstream: Option<super::AccountUpstreamInfoView>,
pub capabilities: AccountCapabilitiesView,
pub outbound_proxy_id: Option<String>,
pub outbound_proxy_endpoint: Option<String>,
@@ -412,6 +413,11 @@ pub struct AccountQuotaCreditsView {
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountQuotaWindowView {
#[serde(skip_serializing_if = "Option::is_none")]
pub source_account_id: Option<String>,
pub amount: Option<super::AccountQuotaAmountView>,
pub observed_at_display: Option<String>,
pub stale: bool,
pub key: String,
pub group: String,
pub limit_id: Option<String>,
@@ -433,6 +439,7 @@ pub struct AccountQuotaWindowView {
#[serde(rename_all = "snake_case")]
pub enum AccountUsageWindowKind {
General,
Daily,
Weekly,
Monthly,
Quota,
@@ -152,6 +152,8 @@ impl ClientKeySort {
#[derive(Clone, PartialEq, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct CreateClientKeyRequest {
#[serde(default)]
peer_sharing: gateway_core::policy::ClientPeerSharing,
#[serde(default)]
provider_request_profile_overrides: ProviderRequestProfileOverrides,
openai_client_profile_override: Option<serde_json::Map<String, serde_json::Value>>,
@@ -190,6 +192,7 @@ impl CreateClientKeyRequest {
self.xai_client_profile_override,
)?;
Ok(CreateClientKey {
peer_sharing: self.peer_sharing,
request_profile_overrides,
custom_key: self
.custom_key
@@ -217,6 +220,7 @@ impl CreateClientKeyRequest {
#[derive(Debug, Clone, PartialEq, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct UpdateClientKeyRequest {
peer_sharing: Option<gateway_core::policy::ClientPeerSharing>,
#[serde(default)]
provider_request_profile_overrides: ProviderRequestProfileOverrideUpdates,
#[serde(default, deserialize_with = "deserialize_profile_override")]
@@ -248,6 +252,7 @@ impl UpdateClientKeyRequest {
self.xai_client_profile_override,
)?;
Ok(UpdateClientKey {
peer_sharing: self.peer_sharing,
request_profile_override_updates,
id: client_key_id(self.id, "clientKeyMutationNotFound")?,
name: self.name,
@@ -334,6 +339,7 @@ impl ClientKeyMutationRequest {
#[derive(Debug, Clone, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ClientKeyView {
peer_sharing: gateway_core::policy::ClientPeerSharing,
provider_request_profile_overrides: ProviderRequestProfileOverrides,
/// 固定兼容字段;值始终从 provider_request_profile_overrides 派生
openai_client_profile_override: Option<serde_json::Map<String, serde_json::Value>>,
@@ -384,6 +390,7 @@ impl From<(ClientKeyRecord, crate::time::TimePresenter)> for ClientKeyView {
.cloned(),
xai_client_profile_override: provider_request_profile_overrides.get("xai").cloned(),
provider_request_profile_overrides,
peer_sharing: record.peer_sharing,
id: record.id.to_string(),
name: record.name,
label: record.label,
@@ -919,6 +919,9 @@ pub(crate) fn dashboard_view(
),
};
account_usage_views.push(DashboardAccountUsageView {
upstream: credential.upstream.map(|info| {
crate::admin::accounts::upstream_info_view(info, chrono::Utc::now(), time)
}),
id: credential.account_id.clone(),
provider: credential.provider_kind.clone(),
authentication_kind: credential.authentication_kind.clone(),
@@ -402,6 +402,7 @@ pub enum DashboardAccountMetric {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct DashboardAccountUsageView {
pub upstream: Option<crate::admin::accounts::AccountUpstreamInfoView>,
pub id: String,
pub provider: String,
pub authentication_kind: String,
@@ -41,11 +41,11 @@ pub fn format_compact_number(value: u64) -> String {
}
/// 货币展示格式化
/// USD 常规金额保留两位,小于 1 美元时最多保留四位
/// USD 常规金额最多保留两位,小于 1 美元时最多保留四位,小数末尾不补零
#[must_use]
pub fn format_decimal_currency(amount: &str, currency: &str) -> String {
if currency != "USD" {
return format!("{currency} {amount}");
return format!("{currency} {}", trim_decimal_zeros(amount));
}
const DECIMAL_SCALE: u128 = 10_000_000_000;
@@ -60,13 +60,21 @@ pub fn format_decimal_currency(amount: &str, currency: &str) -> String {
let display_scale = 10_u128.pow(precision);
let whole = rounded / display_scale;
let fraction = rounded % display_scale;
let mut display = format!("{whole}.{fraction:0width$}", width = precision as usize);
while display.ends_with('0')
&& display
.split_once('.')
.is_some_and(|(_, fraction)| fraction.len() > 2)
{
display.pop();
let display = format!("{whole}.{fraction:0width$}", width = precision as usize);
format!("${}", trim_decimal_zeros(&display))
}
/// 百分比最多保留一位小数,整数百分比不补零
#[must_use]
pub(super) fn format_percentage(value: f64) -> String {
let display = format!("{value:.1}");
format!("{}%", trim_decimal_zeros(&display))
}
fn trim_decimal_zeros(value: &str) -> &str {
if value.contains('.') {
value.trim_end_matches('0').trim_end_matches('.')
} else {
value
}
format!("${display}")
}
@@ -135,6 +135,7 @@ impl From<ProxyTestResult> for ProxyTestView {
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct ProxyAccountView {
upstream: Option<super::accounts::AccountUpstreamInfoView>,
id: String,
name: String,
email: Option<String>,
@@ -322,6 +323,13 @@ where
.items
.into_iter()
.map(|account| ProxyAccountView {
upstream: account.upstream.map(|info| {
super::accounts::upstream_info_view(
info,
chrono::Utc::now(),
crate::time::TimePresenter::new(state.admin_services().timezone()),
)
}),
id: account.id,
name: account.name,
email: account.email,
@@ -8,6 +8,7 @@ pub mod images;
pub(crate) mod live;
pub(crate) mod middleware;
pub mod models;
mod peer_link;
pub mod responses;
pub mod router;
pub mod search;
@@ -0,0 +1,123 @@
//! Bearer Key 的 CPR 互联导出路由,只消费 Admin 的本地观测
mod wire;
use axum::{
Router,
extract::{Query, State},
http::{HeaderMap, HeaderValue, StatusCode, Uri, header},
middleware,
response::{IntoResponse, Response},
routing::get,
};
use gateway_admin::PeerLinkError;
use gateway_admin::model::peer_link::PeerLinkLimits;
use gateway_core::engine::execution::AuthenticatedClient;
use gateway_protocol::peer::{MAX_ACCOUNTS, MAX_MODELS, MAX_QUOTA_WINDOWS, MAX_RESPONSE_BYTES};
use serde::Deserialize;
use super::{
auth::{client_access_error_response, verify_client},
error::{missing_client_api_key_response, openai_error_response},
};
use crate::{
ApiState,
auth::SessionState as _,
openai::middleware::{self as plugin_middleware, RequestInput},
};
pub(super) fn router() -> Router<ApiState> {
Router::new()
.route(super::router::PEER_LINK_PATH, get(peer_link))
.layer(middleware::map_response(no_store))
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct PeerQuery {}
async fn peer_link(State(state): State<ApiState>, headers: HeaderMap, uri: Uri) -> Response {
let client = match verify_client(state.openai(), &headers).await {
Ok(client) => client,
Err(error) => return client_access_error_response(error.into()),
};
plugin_middleware::query_response(
state.openai().execution(),
client,
RequestInput::query(super::router::PEER_LINK_PATH.to_owned(), None, &headers),
move |client| peer_response(state, client, uri),
)
.await
}
async fn peer_response(state: ApiState, client: AuthenticatedClient, uri: Uri) -> Response {
if Query::<PeerQuery>::try_from_uri(&uri).is_err() {
return openai_error_response(
StatusCode::BAD_REQUEST,
"Peer query does not accept query parameters.",
"invalid_request_error",
"invalid_peer_query",
)
.into_response();
}
match state
.admin_services()
.peer_link()
.snapshot(
&client,
PeerLinkLimits {
max_accounts: MAX_ACCOUNTS,
max_quota_windows: MAX_QUOTA_WINDOWS,
max_models: MAX_MODELS,
max_response_bytes: MAX_RESPONSE_BYTES,
},
)
.await
{
Ok(snapshot) => match serde_json::to_vec(&wire::response(snapshot)) {
Ok(body) if body.len() <= MAX_RESPONSE_BYTES => {
([(header::CONTENT_TYPE, "application/json")], body).into_response()
}
_ => unavailable(),
},
Err(PeerLinkError::Unauthorized) => missing_client_api_key_response().into_response(),
Err(PeerLinkError::Forbidden) => openai_error_response(
StatusCode::FORBIDDEN,
"CPR Link sharing is not enabled for this API key.",
"permission_error",
"peer_sharing_disabled",
)
.into_response(),
Err(PeerLinkError::RateLimited { retry_after }) => {
let mut response = openai_error_response(
StatusCode::TOO_MANY_REQUESTS,
"Peer query rate limit reached.",
"rate_limit_error",
"peer_query_rate_limited",
)
.into_response();
response
.headers_mut()
.insert(header::RETRY_AFTER, HeaderValue::from(retry_after));
response
}
Err(PeerLinkError::Unavailable(_) | PeerLinkError::SnapshotTooLarge) => unavailable(),
}
}
fn unavailable() -> Response {
openai_error_response(
StatusCode::SERVICE_UNAVAILABLE,
"Peer snapshot is currently unavailable.",
"server_error",
"peer_snapshot_unavailable",
)
.into_response()
}
async fn no_store(mut response: Response) -> Response {
response
.headers_mut()
.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store"));
response
}
@@ -0,0 +1,136 @@
//! 将已脱敏的互联领域观测投影为版本化协议
use chrono::{DateTime, Utc};
use gateway_admin::model::peer_link::{PeerLinkExport, PeerModelState as ModelState};
use gateway_core::metering::Decimal;
use gateway_protocol::peer::{
OBJECT, PeerBudgetWindow, PeerComponentStatus, PeerKeyUsage, PeerLink, PeerModelAvailability,
PeerModelCoverage, PeerModelState, PeerNode, PeerPlanCount, PeerPool, PeerPoolDisclosure,
PeerPoolStates, PeerQuota, PeerQuotaAccount, PeerQuotaDisclosure, PeerQuotaWindow, PeerUsage,
SCHEMA_VERSION,
};
use std::time::SystemTime;
pub(super) fn response(export: PeerLinkExport) -> PeerLink {
let mut capabilities = vec!["key_usage".to_owned()];
let pool = export.pool.map_or_else(PeerPool::default, |pool| {
capabilities.push("pool_summary".to_owned());
PeerPool {
disclosure: PeerPoolDisclosure::Summary,
observed_at: Some(export.observed_at.to_rfc3339()),
account_count: Some(
pool.ready
+ pool.cooldown
+ pool.exhausted
+ pool.disabled
+ pool.error
+ pool.unknown,
),
states: Some(PeerPoolStates {
ready: pool.ready,
cooldown: pool.cooldown,
exhausted: pool.exhausted,
disabled: pool.disabled,
error: pool.error,
unknown: pool.unknown,
}),
plans: pool.plans.map(|plans| {
plans
.into_iter()
.map(|(name, count)| PeerPlanCount { name, count })
.collect()
}),
}
});
let quota = export.quota.map_or_else(PeerQuota::default, |accounts| {
capabilities.extend(["quota_windows".to_owned(), "model_availability".to_owned()]);
PeerQuota {
disclosure: PeerQuotaDisclosure::Accounts,
accounts: accounts
.into_iter()
.map(|account| PeerQuotaAccount {
id: account.id,
models: account.models,
observed_at: account.observed_at.map(|time| time.to_rfc3339()),
expires_at: account.expires_at.map(|time| time.to_rfc3339()),
windows: account
.windows
.into_iter()
.map(|window| PeerQuotaWindow {
key: window.key,
group: window.group,
label: window.label,
limit_id: window.limit_id,
limit_name: window.limit_name,
role: window.role.map(|role| role.as_str().to_owned()),
window_seconds: window.window_seconds,
used_percent: window.used_percent,
reset_at: window.reset_at.map(|time| time.to_rfc3339()),
limit_reached: window.limit_reached,
})
.collect(),
})
.collect(),
}
});
PeerLink {
object: OBJECT.to_owned(),
schema_version: SCHEMA_VERSION,
node: PeerNode {
id: export.node.id,
name: export.node.name,
software: "codex-proxy-rs".to_owned(),
},
capabilities,
observed_at: export.observed_at.to_rfc3339(),
expires_at: export.expires_at.to_rfc3339(),
usage: PeerUsage {
status: PeerComponentStatus::Available,
observed_at: Some(export.budget_observed_at.to_rfc3339()),
data: Some(PeerKeyUsage {
unit: "USD".to_owned(),
daily: budget_window(
export.budget.limits.daily_usd,
export.budget.daily_used_usd,
export.budget.daily_resets_at,
),
weekly: budget_window(
export.budget.limits.weekly_usd,
export.budget.weekly_used_usd,
export.budget.weekly_resets_at,
),
}),
},
pool,
quota,
models: export
.models
.into_iter()
.map(|model| PeerModelAvailability {
model: model.model,
coverage: if model.complete {
PeerModelCoverage::Complete
} else {
PeerModelCoverage::Partial
},
state: match model.state {
ModelState::Available => PeerModelState::Available,
ModelState::Exhausted => PeerModelState::Exhausted,
ModelState::Unknown => PeerModelState::Unknown,
},
observed_at: model.observed_at.map(|time| time.to_rfc3339()),
expires_at: model.expires_at.map(|time| time.to_rfc3339()),
})
.collect(),
}
}
fn budget_window(total: Decimal, used: Decimal, resets_at: Option<SystemTime>) -> PeerBudgetWindow {
let limited = total != Decimal::ZERO;
PeerBudgetWindow {
total: limited.then(|| total.canonical()),
used: used.canonical(),
remaining: limited.then(|| total.saturating_sub(used).canonical()),
resets_at: resets_at.map(|time| DateTime::<Utc>::from(time).to_rfc3339()),
}
}
@@ -27,6 +27,7 @@ pub(super) const SEARCH_PATH: &str = "/v1/alpha/search";
pub(super) const MODELS_PATH: &str = "/v1/models";
pub(super) const MODEL_DETAIL_PATH: &str = "/v1/models/{model_id}";
pub(super) const USAGE_PATH: &str = "/v1/usage";
pub(super) const PEER_LINK_PATH: &str = "/v1/cpr/link";
/// 构造 OpenAI 客户端协议路由
pub(crate) fn router() -> Router<ApiState> {
@@ -40,6 +41,7 @@ pub(crate) fn router() -> Router<ApiState> {
// 官方 OpenAI 模型详情合同使用 path ID;它不属于 Admin API 约束
.route(MODEL_DETAIL_PATH, get(model_detail))
.merge(usage::router())
.merge(super::peer_link::router())
.merge(live::router())
// OpenAI 数据面正文属于客户端/上游协议;代理不能用私有大小上限提前拒绝
// 上游本可接受的未来 payload
@@ -9,6 +9,73 @@ use tower::ServiceExt as _;
use super::super::{AdminTestFixture, AdminTestState};
#[tokio::test]
async fn upstream_refresh_requires_admin_and_returns_an_accepted_generation() {
use gateway_core::account::{AccountStatusFacts, CredentialState, QuotaState};
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
*fixture.account.lock().unwrap() = Some(account_page_item(
AccountStatusFacts {
enabled: true,
credential_state: CredentialState::Ready,
access_token_expires_at: None,
quota: QuotaState::unknown(),
cooldown: None,
last_error_reason: None,
last_error_message: None,
},
chrono::Utc::now(),
"api_key",
));
for (authenticated, ids, expected) in [
(
false,
serde_json::json!(["acct_cooldown"]),
StatusCode::UNAUTHORIZED,
),
(true, serde_json::json!([]), StatusCode::BAD_REQUEST),
(
true,
serde_json::json!(["invalid"]),
StatusCode::BAD_REQUEST,
),
(
true,
serde_json::json!(["acct_cooldown", "acct_cooldown"]),
StatusCode::ACCEPTED,
),
] {
let mut request = Request::builder()
.method("POST")
.uri("/api/admin/accounts/upstream-info/refresh")
.header("x-request-id", "req_upstream_refresh")
.header(header::CONTENT_TYPE, "application/json");
if authenticated {
request = request.header(header::COOKIE, "cpr_session=valid-session");
}
let response = admin::router::<AdminTestState>()
.with_state(fixture.state())
.oneshot(
request
.body(Body::from(
serde_json::json!({"accountIds":ids}).to_string(),
))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), expected);
if expected == StatusCode::ACCEPTED {
let body = to_bytes(response.into_body(), 8192).await.unwrap();
let value: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(
value["data"]["items"],
serde_json::json!([{"accountId":"acct_cooldown", "connectionRevision":"connection-7", "requestGeneration":7, "state":"queued"}])
);
}
}
}
#[tokio::test]
async fn account_model_reads_and_refreshes_preserve_provider_observation_time() {
use chrono::{DateTime, Utc};
@@ -307,6 +374,7 @@ async fn account_quota_refresh_time_uses_server_clock_and_deployment_timezone()
),
] {
*fixture.provider_quota.lock().unwrap() = Some(ProviderQuota {
upstream: None,
plan_type: None,
observed_at,
refresh_token_expires_at: None,
@@ -527,10 +527,15 @@ mod response {
#[test]
fn account_usage_currency_should_limit_usd_to_four_fraction_digits() {
assert_eq!(format_decimal_currency("0.1204956", "USD"), "$0.1205");
assert_eq!(format_decimal_currency("0.99996", "USD"), "$1.00");
assert_eq!(format_decimal_currency("0.99996", "USD"), "$1");
assert_eq!(format_decimal_currency("0.0300", "USD"), "$0.03");
assert_eq!(format_decimal_currency("0.00", "USD"), "$0");
assert_eq!(format_decimal_currency("100.0000", "USD"), "$100");
assert_eq!(format_decimal_currency("12.3000", "USD"), "$12.3");
assert_eq!(format_decimal_currency("12.3456", "USD"), "$12.35");
assert_eq!(format_decimal_currency("0.1204956", "CNY"), "CNY 0.1204956");
assert_eq!(format_decimal_currency("12.3000", "CNY"), "CNY 12.3");
assert_eq!(format_decimal_currency("100", "CNY"), "CNY 100");
}
}
@@ -70,7 +70,7 @@ fn quota_forecast_projection_formats_capacity_and_preserves_null_zero() {
assert_eq!(value["forecasts"][1]["estimatedTokens"], 0);
assert_eq!(value["forecasts"][1]["estimatedTokensDisplay"], "0");
assert_eq!(value["forecasts"][1]["estimatedUsd"], 0.0);
assert_eq!(value["forecasts"][1]["estimatedUsdDisplay"], "$0.00");
assert_eq!(value["forecasts"][1]["estimatedUsdDisplay"], "$0");
assert_eq!(value["forecasts"][1]["period"], "monthly");
assert_eq!(value["forecasts"][1]["targetDays"], 30.0);
assert_eq!(value["forecasts"][1]["extrapolated"], true);
@@ -434,6 +434,7 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() {
.expect("valid time");
let view = ClientKeyView::from((
gateway_admin::model::client_keys::ClientKeyRecord {
peer_sharing: Default::default(),
request_profile_overrides: Default::default(),
budget: Default::default(),
id: gateway_core::policy::ClientApiKeyId::new("key_visible").expect("Client Key ID"),
@@ -1,6 +1,7 @@
//! 管理控制面 HTTP 接口的测试入口
mod plugins;
mod upstream_info;
use std::{
collections::BTreeMap,
sync::{
@@ -97,6 +98,11 @@ mod settings;
mod system;
mod wire;
pub(super) struct ProviderQuotaGate {
pub entered: tokio::sync::Semaphore,
pub release: tokio::sync::Semaphore,
}
pub(super) struct AdminTestFixture {
pub services: AdminServices,
plugin_ports: Arc<plugins::TestPluginPorts>,
@@ -113,6 +119,7 @@ pub(super) struct AdminTestFixture {
pub dashboard_summary_range: Arc<Mutex<Option<TimeRange>>>,
pub provider_error: Arc<Mutex<Option<ProviderAdminError>>>,
pub provider_quota: Arc<Mutex<Option<ProviderQuota>>>,
pub provider_quota_gate: Arc<Mutex<Option<Arc<ProviderQuotaGate>>>>,
pub provider_models: Arc<Mutex<Option<ProviderModels>>>,
pub account: Arc<Mutex<Option<AccountPageItem>>>,
pub account_usage: Arc<Mutex<Vec<AccountUsage>>>,
@@ -190,18 +197,21 @@ impl AdminTestFixture {
plugin_ports.clone(),
);
let provider_quota = Arc::new(Mutex::new(None));
let provider_quota_gate = Arc::new(Mutex::new(None));
let provider_models = Arc::new(Mutex::new(None));
let providers: Vec<Arc<dyn ProviderAdmin>> = vec![
Arc::new(UnusedProvider::new(
"openai",
Arc::clone(&provider_error),
Arc::clone(&provider_quota),
Arc::clone(&provider_quota_gate),
Arc::clone(&provider_models),
)),
Arc::new(UnusedProvider::new(
"xai",
Arc::clone(&provider_error),
Arc::clone(&provider_quota),
Arc::clone(&provider_quota_gate),
Arc::clone(&provider_models),
)),
];
@@ -215,6 +225,9 @@ impl AdminTestFixture {
ClientConfig::default(),
stores,
gateway_admin::AdminRuntimePorts {
upstream_info: Some(Arc::new(upstream_info::TestPeerInfo(Arc::clone(
&provider_quota,
)))),
timezone,
service_middleware: std::sync::Arc::new(|| None),
plugin_preparation: plugin_ports.clone(),
@@ -251,6 +264,7 @@ impl AdminTestFixture {
dashboard_summary_range,
provider_error,
provider_quota,
provider_quota_gate,
provider_models,
account,
account_usage,
@@ -997,6 +1011,15 @@ fn mutation(
#[async_trait]
impl ClientKeyStore for MemoryClientKeyStore {
async fn peer_node_identity(
&self,
) -> AdminStoreResult<gateway_admin::model::peer_link::PeerNodeIdentity> {
Ok(gateway_admin::model::peer_link::PeerNodeIdentity {
id: "test-peer-node".to_owned(),
name: "测试节点".to_owned(),
})
}
async fn update_client_key_budget_limits(
&self,
_: gateway_admin::model::client_keys::UpdateClientKeyBudgetLimits,
@@ -1071,6 +1094,7 @@ impl ClientKeyStore for MemoryClientKeyStore {
let now = Utc::now();
Ok(Some(ClientKeySecret::new(
ClientKeyRecord {
peer_sharing: Default::default(),
request_profile_overrides: Default::default(),
budget: Default::default(),
id: id.clone(),
@@ -1612,6 +1636,7 @@ struct UnusedProvider {
kind: ProviderKind,
error: Arc<Mutex<Option<ProviderAdminError>>>,
quota: Arc<Mutex<Option<ProviderQuota>>>,
quota_gate: Arc<Mutex<Option<Arc<ProviderQuotaGate>>>>,
models: Arc<Mutex<Option<ProviderModels>>>,
}
@@ -1620,12 +1645,14 @@ impl UnusedProvider {
kind: &str,
error: Arc<Mutex<Option<ProviderAdminError>>>,
quota: Arc<Mutex<Option<ProviderQuota>>>,
quota_gate: Arc<Mutex<Option<Arc<ProviderQuotaGate>>>>,
models: Arc<Mutex<Option<ProviderModels>>>,
) -> Self {
Self {
kind: ProviderKind::new(kind).expect("provider kind"),
error,
quota,
quota_gate,
models,
}
}
@@ -1633,6 +1660,22 @@ impl UnusedProvider {
#[async_trait]
impl ProviderAdmin for UnusedProvider {
fn is_direct_account(&self, authentication_kind: &str) -> bool {
authentication_kind == "oauth"
}
fn peer_model_eligible(
&self,
_: &gateway_admin::model::accounts::AccountRecord,
_: &str,
) -> Option<bool> {
Some(true)
}
fn peer_authentication_kinds(&self) -> &[&str] {
&["api_key"]
}
fn pricing_catalog(&self) -> gateway_admin::model::pricing::ProviderPricingCatalog {
serde_json::from_value(serde_json::json!({
"builtin-model": {
@@ -1721,6 +1764,11 @@ impl ProviderAdmin for UnusedProvider {
&self,
_: gateway_admin::model::provider_credentials::ProviderQuotaRequest,
) -> Result<ProviderQuota, ProviderAdminError> {
let gate = self.quota_gate.lock().unwrap().clone();
if let Some(gate) = gate {
gate.entered.add_permits(1);
gate.release.acquire().await.unwrap().forget();
}
if let Some(quota) = self.quota.lock().expect("provider quota").clone() {
return Ok(quota);
}
@@ -1770,6 +1770,7 @@ async fn dashboard_summary_should_identify_metrics_without_localized_labels() {
.map(|(index, (requests, tokens, last_used_at))| {
let account_id = format!("acct_metric_{index}");
DashboardAccountUsage {
upstream: None,
account_id: account_id.clone(),
provider_kind: "openai".to_owned(),
authentication_kind: "oauth".to_owned(),
@@ -1796,6 +1797,7 @@ async fn dashboard_summary_should_identify_metrics_without_localized_labels() {
request_buckets: Vec::new(),
quota_used_percent: None,
quota_window: requests.map(|request_count| ProviderQuotaWindow {
upstream: None,
key: "rolling".to_owned(),
group: "shortTerm".to_owned(),
label: "滚动窗口".to_owned(),
@@ -1969,6 +1971,7 @@ async fn dashboard_relative_time_preserves_server_anchor_boundaries_and_deployme
provider_accounts: Default::default(),
trend: Vec::new(),
account_usage: vec![DashboardAccountUsage {
upstream: None,
account_id: "acct_relative_time".to_owned(),
provider_kind: "openai".to_owned(),
authentication_kind: "oauth".to_owned(),
@@ -192,6 +192,7 @@ impl ProxyStore for MemoryProxies {
}
Ok(ProxyAccountPage {
items: vec![ProxyAccountRef {
upstream: None,
id: "acct_linked".to_owned(),
name: "工作账号".to_owned(),
email: Some("work@example.invalid".to_owned()),
@@ -456,7 +457,7 @@ async fn proxy_routes_save_reload_test_rename_and_delete_without_exposing_creden
json!([{
"id":"acct_linked", "name":"工作账号", "email":"work@example.invalid",
"provider":"openai", "enabled":true, "authenticationKind":"oauth",
"planType":"plus", "planTypeDisplay":"Plus", "groups":[]
"planType":"plus", "planTypeDisplay":"Plus", "groups":[], "upstream":null
}])
);
@@ -0,0 +1,58 @@
//! 账号 API 通过公共互联端口读观测与提交刷新,不由具体 Provider 执行
use super::*;
use gateway_admin::model::provider_credentials::{ProviderQuota, ProviderQuotaReadRequest};
use gateway_admin::ports::peer::PeerUpstreamService;
use gateway_core::account::{ProviderAccount, UpstreamInfoSource, UpstreamSchedulingSignal};
use gateway_core::error::StoreError;
use std::time::SystemTime;
pub(super) struct TestPeerInfo(pub Arc<Mutex<Option<ProviderQuota>>>);
impl UpstreamInfoSource for TestPeerInfo {
fn signal(
&self,
_: &ProviderAccount,
_: &str,
_: SystemTime,
) -> Option<UpstreamSchedulingSignal> {
None
}
}
#[async_trait]
impl PeerUpstreamService for TestPeerInfo {
async fn account_changed(&self, _: &ProviderAccountId) -> Result<(), StoreError> {
Ok(())
}
fn forget(&self, _: &ProviderAccountId) {}
async fn request_refresh(
&self,
account_id: &ProviderAccountId,
) -> Result<gateway_admin::model::upstream_info::UpstreamRefreshReceipt, StoreError> {
Ok(
gateway_admin::model::upstream_info::UpstreamRefreshReceipt {
account_id: account_id.as_str().to_owned(),
connection_revision: "connection-7".to_owned(),
request_generation: 7,
state: gateway_admin::model::upstream_info::UpstreamSyncState::Queued,
},
)
}
async fn read_quota(&self, _: &ProviderAccountId) -> Result<Option<ProviderQuota>, StoreError> {
Ok(self.0.lock().unwrap().clone())
}
async fn read_summary(
&self,
_: &ProviderAccountId,
) -> Result<Option<ProviderQuota>, StoreError> {
Ok(self.0.lock().unwrap().clone())
}
async fn read_quotas(
&self,
requests: &[ProviderQuotaReadRequest<'_>],
) -> Vec<Result<ProviderQuota, ProviderAdminError>> {
let quota = self.0.lock().unwrap().clone().unwrap_or_default();
requests.iter().map(|_| Ok(quota.clone())).collect()
}
}
@@ -18,6 +18,7 @@ pub(super) async fn fixture() -> AdminTestFixture {
let fixture = key_fixture().await;
let now = Utc::now();
*fixture.client_key.lock().unwrap() = Some(ClientKeyRecord {
peer_sharing: Default::default(),
request_profile_overrides: Default::default(),
id: ClientApiKeyId::new("key-42").unwrap(),
name: "Development".to_owned(),
@@ -9,6 +9,7 @@ mod images;
mod live;
mod middleware;
mod models;
mod peer_link;
mod responses;
mod router;
mod search;
@@ -0,0 +1,234 @@
//! 原生目录来源仅用于展示,互联额度覆盖仍按账号的模型候选资格计算
use super::*;
use futures::future::BoxFuture;
use gateway_core::{
account::{AccountModelAccess, AccountModelAccessMode},
operation::OperationKind,
routing::{
ProviderCatalogGeneration, ProviderCatalogPort, ProviderCatalogUnavailable,
ProviderModelCapabilities,
snapshot::{
RuntimeSnapshotCompiler, SnapshotAccountGroupFacts, SnapshotAccountGroupMemberFacts,
SnapshotClientPolicyFacts, SnapshotFacts, SnapshotProviderAccountFacts,
SnapshotStoreError, SnapshotStorePort,
},
},
};
struct CatalogFixture(SnapshotFacts, usize);
impl SnapshotStorePort for CatalogFixture {
fn load_snapshot_facts(&self) -> BoxFuture<'_, Result<SnapshotFacts, SnapshotStoreError>> {
Box::pin(async { Ok(self.0.clone()) })
}
fn current_config_revision(&self) -> BoxFuture<'_, Result<ConfigRevision, SnapshotStoreError>> {
Box::pin(async { Ok(ConfigRevision::new(1).unwrap()) })
}
}
impl ProviderCatalogPort for CatalogFixture {
fn is_direct_account(&self, _: &ProviderKind, authentication_kind: &str) -> bool {
authentication_kind == "oauth"
}
fn catalog_generations(&self) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
[(
ProviderKind::new("openai").unwrap(),
ProviderCatalogGeneration::new(1),
)]
.into()
}
fn query_model_capabilities(
&self,
_: &ProviderKind,
) -> BoxFuture<'_, Result<Vec<ProviderModelCapabilities>, ProviderCatalogUnavailable>> {
Box::pin(async {
Ok((0..self.1)
.map(|index| {
ProviderModelCapabilities::new(
UpstreamModelId::new(if index == 0 {
"model-a".to_owned()
} else {
format!("model-{index}")
})
.unwrap(),
ModelCapabilities::new([OperationKind::Generate].into(), None),
)
.with_catalog_accounts(
[ProviderAccountId::new("acct_catalog_source").unwrap()].into(),
)
})
.collect())
})
}
}
#[tokio::test]
async fn native_quota_coverage_uses_routing_eligibility_and_current_model_permissions() {
let (fixture, _) = fixture(sharing(), RAW_KEY).await;
let snapshot = catalog_snapshot(&fixture, 1).await;
let app = router_with_snapshot(&fixture, snapshot);
// 目录来源账号当前不可读,另一个未贡献目录的账号仍有可用观测,不能被漏掉
let body = response_json(query(&app, PATH, Some(RAW_KEY)).await).await;
assert_eq!(
body["quota"]["accounts"][0]["models"],
json!(["alias", "model-a"])
);
for model in body["models"].as_array().unwrap() {
assert_eq!(model["state"], "available");
assert_eq!(model["coverage"], "partial");
}
fixture
.account
.lock()
.unwrap()
.as_mut()
.unwrap()
.account
.model_access =
AccountModelAccess::new(AccountModelAccessMode::Denylist, vec!["model-a".into()]).unwrap();
let body = response_json(query(&app, PATH, Some(RAW_KEY)).await).await;
assert_eq!(body["quota"]["accounts"][0]["models"], json!([]));
for model in body["models"].as_array().unwrap() {
assert_eq!(model["state"], "unknown");
assert_eq!(model["coverage"], "partial");
}
}
async fn catalog_snapshot(fixture: &AdminTestFixture, models: usize) -> RuntimeSnapshot {
let group = fixture.client_key.lock().unwrap().as_ref().unwrap().groups[0].clone();
let account_ids = ["acct_catalog_source", "acct_private_sentinel"]
.map(|id| ProviderAccountId::new(id).unwrap());
let facts = Arc::new(CatalogFixture(
SnapshotFacts::new(
ConfigRevision::new(1).unwrap(),
ConfigRevision::new(1).unwrap(),
gateway_core::settings::SettingsValues::new(
2,
"smart",
[("alias".into(), "model-a".into())].into(),
None,
None,
),
vec![
SnapshotClientPolicyFacts::new(
ClientApiKeyId::new("key-42").unwrap(),
PlaintextClientApiKey::new(RAW_KEY).unwrap(),
vec![group.id.clone()],
RateLimits::unlimited(),
)
.with_peer_sharing(sharing()),
],
vec![SnapshotAccountGroupFacts::new(
group.id.clone(),
group.name,
true,
)],
account_ids
.iter()
.map(|id| {
SnapshotProviderAccountFacts::new(id.clone(), "openai")
.with_authentication_kind("oauth".into())
})
.collect(),
account_ids
.iter()
.map(|id| SnapshotAccountGroupMemberFacts::new(group.id.clone(), id.clone()))
.collect(),
),
models,
));
RuntimeSnapshotCompiler::new(facts.clone(), facts)
.compile()
.await
.unwrap()
}
#[tokio::test]
async fn native_catalog_accepts_the_model_limit_and_rejects_overflow() {
let (fixture, _) = fixture(sharing(), RAW_KEY).await;
// 目录还公开一个 alias,边界按公开名称总数计算
for (count, expected) in [
(gateway_protocol::peer::MAX_MODELS - 1, StatusCode::OK),
(
gateway_protocol::peer::MAX_MODELS,
StatusCode::SERVICE_UNAVAILABLE,
),
] {
let app = router_with_snapshot(&fixture, catalog_snapshot(&fixture, count).await);
let response = query(&app, PATH, Some(RAW_KEY)).await;
assert_eq!(response.status(), expected);
if expected == StatusCode::OK {
let bytes = axum::body::to_bytes(
response.into_body(),
gateway_protocol::peer::MAX_RESPONSE_BYTES,
)
.await
.unwrap();
let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
assert_eq!(
body["models"].as_array().unwrap().len(),
gateway_protocol::peer::MAX_MODELS
);
} else {
assert!(response_json(response).await.get("pool").is_none());
}
}
}
#[tokio::test]
async fn model_expansion_honors_the_byte_budget_and_current_permissions() {
use gateway_admin::{PeerLinkError, model::peer_link::PeerLinkLimits};
use gateway_core::engine::execution::ExecutionService as _;
let (fixture, _) = fixture(sharing(), RAW_KEY).await;
let execution = execution_with_snapshot(catalog_snapshot(&fixture, 1).await);
let client = execution.authenticate(RAW_KEY).unwrap();
let limits = PeerLinkLimits {
max_accounts: 3000,
max_models: 2000,
max_quota_windows: 6000,
max_response_bytes: 12,
};
let snapshot = fixture
.services
.peer_link()
.snapshot(&client, limits)
.await
.unwrap();
assert_eq!(snapshot.quota.unwrap()[0].models, ["alias", "model-a"]);
let limits = PeerLinkLimits {
max_response_bytes: 11,
..limits
};
assert!(matches!(
fixture.services.peer_link().snapshot(&client, limits).await,
Err(PeerLinkError::SnapshotTooLarge)
));
fixture
.account
.lock()
.unwrap()
.as_mut()
.unwrap()
.account
.model_access =
AccountModelAccess::new(AccountModelAccessMode::Denylist, vec!["model-a".into()]).unwrap();
let snapshot = fixture
.services
.peer_link()
.snapshot(&client, limits)
.await
.unwrap();
assert!(snapshot.quota.unwrap()[0].models.is_empty());
assert!(
snapshot
.models
.iter()
.all(|model| model.state == gateway_admin::model::peer_link::PeerModelState::Unknown)
);
}
@@ -0,0 +1,534 @@
//! 验证互联导出的授权范围、过期观测和独立于推理准入的只读合同
mod catalog;
use crate::{
admin::AdminTestFixture,
support::{RAW_KEY, empty_request, key_fixture, response_json},
};
use axum::{
Router,
http::{Method, StatusCode, header},
response::Response,
};
use chrono::{TimeDelta, Utc};
use gateway_admin::model::{
Revision,
account_groups::{AccountGroupColor, AccountGroupRef},
accounts::{AccountCapacity, AccountPageItem, AccountRecord},
provider_credentials::{
ProviderQuota, ProviderQuotaWindow, ProviderQuotaWindowRole, QuotaLocalUsageAttribution,
},
};
use gateway_core::{
account::{
AccountStatusFacts, CredentialState, ProviderAccountId, QuotaEvidence, QuotaState,
resolve_account_status,
},
engine::{execution::DefaultExecutionService, provider::ProviderRegistry},
policy::{ClientApiKeyId, ClientPeerSharing, ClientPolicy, PlaintextClientApiKey, RateLimits},
routing::{
AccountGroupId, ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities,
ProviderKind, ProviderModel, RoutingGroupSnapshot, RuntimeAccount, RuntimeAccountDirectory,
RuntimeSnapshot, UpstreamModelId,
},
runtime::RuntimeSnapshotHandle,
};
use serde_json::json;
use std::{
collections::{BTreeMap, BTreeSet},
sync::Arc,
};
use tower::ServiceExt as _;
const PATH: &str = "/v1/cpr/link";
fn sharing() -> ClientPeerSharing {
ClientPeerSharing {
enabled: true,
share_pool: true,
share_quota: true,
share_plans: false,
}
}
async fn fixture(sharing: ClientPeerSharing, plaintext: &str) -> (AdminTestFixture, Router) {
let fixture = key_fixture().await;
let now = Utc::now();
let group = AccountGroupRef {
id: AccountGroupId::new(format!("grp_{}", "1".repeat(32))).unwrap(),
name: "private-group-sentinel".into(),
color: AccountGroupColor::parse("#2563EBFF").unwrap(),
enabled: true,
};
let id = ClientApiKeyId::new("key-42").unwrap();
let mut key = fixture
.services
.client_keys()
.reveal(&id)
.await
.unwrap()
.record;
key.peer_sharing = sharing;
key.groups = vec![group.clone()];
key.name = "private-key-sentinel".into();
key.budget.limits.daily_usd = "1".parse().unwrap();
key.budget.daily_used_usd = "2".parse().unwrap();
*fixture.client_key.lock().unwrap() = Some(key);
let provider = ProviderKind::new("openai").unwrap();
let facts = AccountStatusFacts {
enabled: true,
credential_state: CredentialState::Ready,
access_token_expires_at: None,
quota: QuotaState::allowed(now.into()),
cooldown: None,
last_error_reason: None,
last_error_message: None,
};
*fixture.account.lock().unwrap() = Some(AccountPageItem {
protection_rules: Vec::new(),
capacity: AccountCapacity {
used_slots: None,
total_slots: None,
},
projection: resolve_account_status(&facts, now.into()),
account: AccountRecord {
id: "acct_private_sentinel".into(),
provider_kind: provider,
groups: vec![group.clone()],
name: "private-account-sentinel".into(),
notes: None,
email: Some("private-email-sentinel@example.invalid".into()),
upstream_user_id: None,
upstream_account_id: None,
plan_type: Some("pro".into()),
authentication_kind: "oauth".into(),
credential_revision: Revision::new(1).unwrap(),
has_refresh_token: true,
access_token_expires_at: None,
next_refresh_at: None,
enabled: true,
concurrency_limit: None,
weight: Default::default(),
model_access: Default::default(),
outbound_proxy_id: None,
outbound_proxy: None,
outbound_proxy_enabled: true,
credential_state: facts.credential_state,
credential_observed_at: now,
quota: facts.quota,
last_error_reason: None,
last_error_message: None,
created_at: now,
updated_at: now,
},
});
*fixture.provider_quota.lock().unwrap() = Some(ProviderQuota {
upstream: None,
plan_type: Some("pro".into()),
observed_at: Some(now),
refresh_token_expires_at: None,
windows: vec![ProviderQuotaWindow {
upstream: None,
key: "weekly".into(),
group: "shortTerm".into(),
label: "周额度".into(),
limit_id: None,
limit_name: None,
role: Some(ProviderQuotaWindowRole::Secondary),
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(604800),
used_percent: Some(47.0),
reset_at: Some(now + TimeDelta::days(2)),
limit_reached: false,
local_usage: None,
provider_data: None,
}],
credits: None,
limit_reached: false,
provider_data: None,
});
let app = router_with_sharing(&fixture, sharing, plaintext);
(fixture, app)
}
fn router_with_sharing(
fixture: &AdminTestFixture,
sharing: ClientPeerSharing,
plaintext: &str,
) -> Router {
let group = fixture.client_key.lock().unwrap().as_ref().unwrap().groups[0].clone();
let provider = ProviderKind::new("openai").unwrap();
let id = ClientApiKeyId::new("key-42").unwrap();
let directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
ProviderAccountId::new("acct_private_sentinel").unwrap(),
RuntimeAccount::new(provider.clone(), BTreeSet::from([group.id.clone()]))
.with_direct_authentication_kind(Some("oauth".into())),
)])));
let scope = FrozenAccountScope::new(
directory.clone(),
ClientRoutingScope::restricted(
vec![RoutingGroupSnapshot::new(
group.id.clone(),
"test".to_owned(),
)],
BTreeSet::from([group.id]),
BTreeSet::from([provider.clone()]),
)
.unwrap(),
);
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).unwrap(),
gateway_core::settings::SettingsValues::new(2, "smart", Default::default(), None, None),
vec![provider.clone()],
vec![ProviderModel::new(
provider,
UpstreamModelId::new("model-a").unwrap(),
ModelCapabilities::new(
BTreeSet::from([gateway_core::operation::OperationKind::Generate]),
Some(16000),
),
)],
vec![
ClientPolicy::new(
id,
PlaintextClientApiKey::new(plaintext).unwrap(),
Arc::new(scope),
true,
RateLimits::unlimited(),
)
.with_peer_sharing(sharing),
],
)
.unwrap()
.with_account_directory(directory);
router_with_snapshot(fixture, snapshot)
}
fn router_with_snapshot(fixture: &AdminTestFixture, snapshot: RuntimeSnapshot) -> Router {
let execution = execution_with_snapshot(snapshot);
super::api_router_with_admin_and_execution(fixture.services.clone(), execution)
}
fn execution_with_snapshot(snapshot: RuntimeSnapshot) -> Arc<DefaultExecutionService> {
Arc::new(DefaultExecutionService::new(
RuntimeSnapshotHandle::new(snapshot),
Arc::new(super::UnusedExecutionStore),
ProviderRegistry::default(),
Arc::new(super::UnusedAdmissions),
Arc::new(super::UnusedContinuation),
Arc::new(super::IgnoredClientApiKeyUsage),
Arc::new(crate::support::RecordingDiagnostics::default()),
))
}
async fn query(app: &Router, path: &str, plaintext: Option<&str>) -> Response {
let mut request = empty_request(Method::GET, path);
if let Some(key) = plaintext {
request.headers_mut().insert(
header::AUTHORIZATION,
format!("Bearer {key}").parse().unwrap(),
);
}
let response = app.clone().oneshot(request).await.unwrap();
assert_eq!(response.headers()[header::CACHE_CONTROL], "no-store");
response
}
#[tokio::test]
async fn exports_scoped_observations_even_when_the_key_budget_is_exhausted_without_writing_usage() {
let (fixture, app) = fixture(sharing(), RAW_KEY).await;
let before = fixture.client_key.lock().unwrap().clone();
let response = query(&app, PATH, Some(RAW_KEY)).await;
assert_eq!(response.status(), StatusCode::OK);
let body = response_json(response).await;
assert_eq!(body["object"], "cpr.peer");
assert_eq!(body["usage"]["data"]["daily"]["remaining"], "0");
assert_eq!(body["pool"]["account_count"], 1);
assert!(body["pool"].get("plans").is_none());
assert_eq!(
body["quota"]["accounts"][0]["windows"][0]["used_percent"],
47.0
);
assert_eq!(body["models"][0]["state"], "available");
assert_eq!(body["models"][0]["coverage"], "complete");
assert!(!body.to_string().contains("sentinel"));
assert!(!body.to_string().contains(RAW_KEY));
assert_eq!(*fixture.client_key.lock().unwrap(), before);
assert!(fixture.observations.lock().unwrap().summaries.is_empty());
}
#[tokio::test]
async fn requires_authentication_explicit_sharing_and_bound_groups_and_rejects_query_overrides() {
let (fixture, app) = fixture(sharing(), RAW_KEY).await;
assert_eq!(
query(&app, PATH, None).await.status(),
StatusCode::UNAUTHORIZED
);
assert_eq!(
query(&app, &format!("{PATH}?account=private"), Some(RAW_KEY))
.await
.status(),
StatusCode::BAD_REQUEST
);
fixture
.client_key
.lock()
.unwrap()
.as_mut()
.unwrap()
.peer_sharing
.enabled = false;
assert_eq!(
query(&app, PATH, Some(RAW_KEY)).await.status(),
StatusCode::FORBIDDEN
);
{
let mut key = fixture.client_key.lock().unwrap();
let key = key.as_mut().unwrap();
key.peer_sharing.enabled = true;
key.groups.clear();
}
assert_eq!(
query(&app, PATH, Some(RAW_KEY)).await.status(),
StatusCode::FORBIDDEN
);
fixture.client_key.lock().unwrap().as_mut().unwrap().enabled = false;
assert_eq!(
query(&app, PATH, Some(RAW_KEY)).await.status(),
StatusCode::UNAUTHORIZED
);
}
#[tokio::test]
async fn disabled_or_changed_group_bindings_remove_accounts_from_the_frozen_scope() {
let (fixture, app) = fixture(sharing(), RAW_KEY).await;
let original_group = fixture.client_key.lock().unwrap().as_ref().unwrap().groups[0].clone();
fixture.client_key.lock().unwrap().as_mut().unwrap().groups[0].enabled = false;
let response = query(&app, PATH, Some(RAW_KEY)).await;
assert_eq!(response.status(), StatusCode::OK);
let body = response_json(response).await;
assert_eq!(body["pool"]["account_count"], 0);
{
let mut key = fixture.client_key.lock().unwrap();
let group = &mut key.as_mut().unwrap().groups[0];
*group = original_group.clone();
group.id = AccountGroupId::new(format!("grp_{}", "2".repeat(32))).unwrap();
}
let body = response_json(query(&app, PATH, Some(RAW_KEY)).await).await;
assert_eq!(body["pool"]["account_count"], 0);
fixture.client_key.lock().unwrap().as_mut().unwrap().groups[0] = original_group;
let body = response_json(query(&app, PATH, Some(RAW_KEY)).await).await;
assert_eq!(body["pool"]["account_count"], 1);
}
#[tokio::test]
async fn current_field_permissions_and_direct_account_facts_narrow_the_frozen_scope() {
let (fixture, app) = fixture(sharing(), RAW_KEY).await;
fixture
.client_key
.lock()
.unwrap()
.as_mut()
.unwrap()
.peer_sharing
.share_quota = false;
let body = response_json(query(&app, PATH, Some(RAW_KEY)).await).await;
assert_eq!(body["quota"], json!({ "disclosure": "none" }));
assert!(body.get("models").is_none());
fixture
.account
.lock()
.unwrap()
.as_mut()
.unwrap()
.account
.authentication_kind = "api_key".into();
let body = response_json(query(&app, PATH, Some(RAW_KEY)).await).await;
assert_eq!(body["pool"]["account_count"], 0);
}
#[tokio::test]
async fn stale_observations_remain_visible_but_cannot_claim_exhausted_model_coverage() {
let (fixture, app) = fixture(sharing(), RAW_KEY).await;
let old = Utc::now() - TimeDelta::hours(1);
fixture
.provider_quota
.lock()
.unwrap()
.as_mut()
.unwrap()
.observed_at = Some(old);
fixture
.account
.lock()
.unwrap()
.as_mut()
.unwrap()
.account
.quota = QuotaState::exhausted(QuotaEvidence::ProviderDenied, old.into(), None);
let body = response_json(query(&app, PATH, Some(RAW_KEY)).await).await;
assert_eq!(
body["quota"]["accounts"][0]["observed_at"],
old.to_rfc3339()
);
assert_eq!(body["models"][0]["state"], "unknown");
assert_eq!(body["models"][0]["coverage"], "partial");
fixture
.account
.lock()
.unwrap()
.as_mut()
.unwrap()
.account
.quota = QuotaState::exhausted(QuotaEvidence::ProviderDenied, Utc::now().into(), None);
let body = response_json(query(&app, PATH, Some(RAW_KEY)).await).await;
assert_eq!(body["models"][0]["state"], "exhausted");
assert_eq!(body["models"][0]["coverage"], "complete");
}
#[tokio::test]
async fn account_identifiers_are_isolated_per_key_and_peer_queries_have_separate_rate_limits() {
let (_, app) = fixture(sharing(), RAW_KEY).await;
let body = response_json(query(&app, PATH, Some(RAW_KEY)).await).await;
let (_, other) = fixture(sharing(), "different-peer-key").await;
let other_body = response_json(query(&other, PATH, Some("different-peer-key")).await).await;
assert_ne!(
body["quota"]["accounts"][0]["id"],
other_body["quota"]["accounts"][0]["id"]
);
for _ in 1..30 {
assert_eq!(
query(&app, PATH, Some(RAW_KEY)).await.status(),
StatusCode::OK
);
}
let response = query(&app, PATH, Some(RAW_KEY)).await;
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
assert!(response.headers().contains_key(header::RETRY_AFTER));
assert_eq!(
query(&app, "/v1/usage", Some(RAW_KEY)).await.status(),
StatusCode::OK
);
}
#[tokio::test]
async fn oversized_snapshots_are_rejected_without_returning_a_truncated_pool() {
let (fixture, app) = fixture(sharing(), RAW_KEY).await;
let window = fixture
.provider_quota
.lock()
.unwrap()
.as_ref()
.unwrap()
.windows[0]
.clone();
fixture
.provider_quota
.lock()
.unwrap()
.as_mut()
.unwrap()
.windows = vec![window.clone(); gateway_protocol::peer::MAX_QUOTA_WINDOWS + 1];
let response = query(&app, PATH, Some(RAW_KEY)).await;
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
assert!(response_json(response).await.get("pool").is_none());
let mut oversized = window;
oversized.label = "x".repeat(gateway_protocol::peer::MAX_RESPONSE_BYTES);
fixture
.provider_quota
.lock()
.unwrap()
.as_mut()
.unwrap()
.windows = vec![oversized];
assert_eq!(
query(&app, PATH, Some(RAW_KEY)).await.status(),
StatusCode::SERVICE_UNAVAILABLE
);
}
#[tokio::test]
async fn heavy_snapshots_reject_excess_concurrency_and_release_slots_on_completion_or_cancellation()
{
let (fixture, app) = fixture(sharing(), RAW_KEY).await;
let gate = Arc::new(crate::admin::ProviderQuotaGate {
entered: tokio::sync::Semaphore::new(0),
release: tokio::sync::Semaphore::new(0),
});
*fixture.provider_quota_gate.lock().unwrap() = Some(gate.clone());
let mut requests = Vec::new();
for _ in 0..8 {
let app = app.clone();
requests.push(tokio::spawn(async move {
query(&app, PATH, Some(RAW_KEY)).await
}));
}
tokio::time::timeout(
std::time::Duration::from_secs(2),
gate.entered.acquire_many(8),
)
.await
.expect("eight snapshots entered quota aggregation")
.unwrap()
.forget();
let current_key = fixture.client_key.lock().unwrap().take();
let response = tokio::time::timeout(
std::time::Duration::from_secs(2),
query(&app, PATH, Some(RAW_KEY)),
)
.await
.expect("excess snapshot must not wait for a slot");
*fixture.client_key.lock().unwrap() = current_key;
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
assert_eq!(response.headers()[header::RETRY_AFTER], "1");
assert_eq!(
response_json(response).await["error"]["code"],
"peer_query_rate_limited"
);
assert_eq!(
query(&app, "/v1/usage", Some(RAW_KEY)).await.status(),
StatusCode::OK
);
let budget_app = router_with_sharing(
&fixture,
ClientPeerSharing {
enabled: true,
..Default::default()
},
RAW_KEY,
);
let budget = tokio::time::timeout(
std::time::Duration::from_secs(2),
query(&budget_app, PATH, Some(RAW_KEY)),
)
.await
.expect("budget-only queries do not consume heavy-query slots");
assert_eq!(budget.status(), StatusCode::OK);
assert_eq!(response_json(budget).await["pool"]["disclosure"], "none");
let cancelled = requests.pop().unwrap();
cancelled.abort();
assert!(cancelled.await.unwrap_err().is_cancelled());
let app_for_replacement = app.clone();
requests.push(tokio::spawn(async move {
query(&app_for_replacement, PATH, Some(RAW_KEY)).await
}));
tokio::time::timeout(std::time::Duration::from_secs(2), gate.entered.acquire())
.await
.expect("cancelled snapshot released its slot")
.unwrap()
.forget();
gate.release.add_permits(8);
for request in requests {
assert_eq!(request.await.unwrap().status(), StatusCode::OK);
}
gate.release.add_permits(1);
assert_eq!(
query(&app, PATH, Some(RAW_KEY)).await.status(),
StatusCode::OK
);
}
@@ -19,10 +19,12 @@ pub mod scope;
mod selection;
mod smart_scheduling;
mod store;
mod upstream;
pub use error::CredentialError;
pub use model::*;
pub(crate) use selection::smart_score;
pub use selection::*;
pub use smart_scheduling::{SmartSchedulingConfig, SmartSchedulingConfigError};
pub use store::ProviderAccountStore;
pub use store::{ProviderAccountStore, UpstreamInfoStore};
pub use upstream::*;
@@ -48,6 +48,7 @@ pub struct RuntimeAccount {
provider_kind: ProviderKind,
group_ids: Arc<BTreeSet<AccountGroupId>>,
model_access: super::AccountModelAccess,
direct_authentication_kind: Option<String>,
}
impl RuntimeAccount {
@@ -57,6 +58,7 @@ impl RuntimeAccount {
provider_kind,
group_ids: Arc::new(group_ids),
model_access: super::AccountModelAccess::all(),
direct_authentication_kind: None,
}
}
@@ -66,6 +68,23 @@ impl RuntimeAccount {
self
}
/// Provider 确认该认证方式直接使用自身资源,而非继续转供另一个节点
#[must_use]
pub fn with_direct_authentication_kind(mut self, kind: Option<String>) -> Self {
self.direct_authentication_kind = kind;
self
}
#[must_use]
pub const fn is_direct_account(&self) -> bool {
self.direct_authentication_kind.is_some()
}
#[must_use]
pub fn direct_authentication_kind(&self) -> Option<&str> {
self.direct_authentication_kind.as_deref()
}
#[must_use]
pub const fn model_access(&self) -> &super::AccountModelAccess {
&self.model_access
@@ -248,6 +267,7 @@ impl ClientRoutingScope {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FrozenAccountScope {
allow_image_generation: bool,
direct_accounts_only: bool,
fast_mode: FastMode,
request_profiles: BTreeMap<ProviderKind, super::OpaqueProviderData>,
directory: Arc<RuntimeAccountDirectory>,
@@ -258,6 +278,31 @@ pub struct FrozenAccountScope {
}
impl FrozenAccountScope {
/// 枚举冻结权限实际允许的账号,保留所有既有直连和分组约束
pub fn allowed_account_ids(&self) -> impl Iterator<Item = &ProviderAccountId> {
self.directory.accounts.keys().filter(|id| self.allows(id))
}
/// 互联限制与原账号范围取交集,分组变更和请求局部重选均不能扩大它
#[must_use]
pub fn with_direct_accounts_only(mut self, enabled: bool) -> Self {
self.direct_accounts_only = enabled;
if enabled {
self.provider_kinds = Arc::new(
self.provider_kinds
.iter()
.filter(|provider| {
self.directory.accounts.iter().any(|(id, account)| {
account.provider_kind() == *provider && self.allows(id)
})
})
.cloned()
.collect(),
);
}
self
}
/// Key 绑定分组合并后的生图权限,与账号成员资格一同冻结
#[must_use]
pub const fn with_image_generation(mut self, allowed: bool) -> Self {
@@ -310,6 +355,7 @@ impl FrozenAccountScope {
};
Self {
allow_image_generation: true,
direct_accounts_only: false,
fast_mode: FastMode::Default,
request_profiles: BTreeMap::new(),
directory,
@@ -369,7 +415,10 @@ impl FrozenAccountScope {
let Some(account) = self.directory.account(account_id) else {
return false;
};
if self.excluded_account_ids.contains(account_id)
if (self.direct_accounts_only
&& (!account.is_direct_account()
|| matches!(self.client_scope, ClientRoutingScope::AllAccounts)))
|| self.excluded_account_ids.contains(account_id)
|| !self.provider_kinds.contains(account.provider_kind())
|| self
.allowed_account_ids
@@ -398,6 +447,16 @@ impl FrozenAccountScope {
.is_some_and(|account| account.model_access.allows(upstream_model))
}
/// 当前凭据类型必须与已授权的直连事实相同,快照发布间隙不能转成代理账号
#[must_use]
pub fn allows_current_account(&self, account: &super::ProviderAccount) -> bool {
self.allows(account.id())
&& (!self.direct_accounts_only
|| self.directory.account(account.id()).is_some_and(|frozen| {
frozen.direct_authentication_kind() == Some(account.authentication_kind())
}))
}
/// 目录按整个授权账号池过滤,不能只看用于获取元数据的账号政策
#[must_use]
pub fn allows_provider_model(&self, provider: &ProviderKind, upstream_model: &str) -> bool {
@@ -848,7 +848,7 @@ impl AccountSelector {
if context
.account_scope
.as_ref()
.is_some_and(|scope| !scope.allows(candidate.account.id()))
.is_some_and(|scope| !scope.allows_current_account(&candidate.account))
{
return Some(AccountSchedulingBlocker::OutsideClientScope);
}
@@ -1,6 +1,7 @@
//! Provider 账号持久化端口
use async_trait::async_trait;
use std::time::{Duration, SystemTime};
use crate::error::StoreError;
use crate::identity::ProviderKind;
@@ -9,7 +10,8 @@ use super::{
AccountStateChange, CredentialCasOutcome, CredentialCasUpdate, CredentialRevision,
LoadedCredential, NewProviderAccount, ProviderAccount, ProviderAccountId,
ProviderAccountUpdate, ProviderRefreshQuery, QuotaAccessChange, QuotaObservation,
QuotaObservationTouch, QuotaWriteOutcome,
QuotaObservationTouch, QuotaWriteOutcome, UpstreamConnection, UpstreamInfoClaim,
UpstreamInfoCompletion, UpstreamInfoPlan, UpstreamInfoRecord,
};
/// `provider_accounts` 的数据库中立端口
@@ -87,3 +89,56 @@ pub trait ProviderAccountStore: Send + Sync {
async fn delete_account(&self, account: &ProviderAccountId) -> Result<(), StoreError>;
}
/// 与 Provider 无关的计划持久化端口,所有方法均只访问本地存储
#[async_trait]
pub trait UpstreamInfoStore: Send + Sync {
/// 按账号 ID 游标扫描 Provider 声明支持互联的认证种类;通用存储不判断认证协议
async fn list_candidates(
&self,
provider: &ProviderKind,
authentication_kinds: &[&str],
after: Option<&ProviderAccountId>,
limit: u32,
) -> Result<Vec<LoadedCredential>, StoreError>;
/// 新连接创建立即到期的计划;身份变化清空旧观测;停用则删除计划与观测
/// 写入前须按当前账号事实核对连接身份;匹配失败返回 Conflict
async fn reconcile(
&self,
connection: &UpstreamConnection,
enabled: bool,
now: SystemTime,
) -> Result<(), StoreError>;
/// 请求代次合并;正在领取或排队的计划不重复递增,失败退避不能被手动绕过
async fn request_refresh(
&self,
connection: &UpstreamConnection,
now: SystemTime,
minimum_interval: Duration,
) -> Result<UpstreamInfoPlan, StoreError>;
/// 批量读取同步元数据,用于核对缓存观测版本,不读取快照正文
async fn read_plans(
&self,
accounts: &[ProviderAccountId],
) -> Result<Vec<UpstreamInfoPlan>, StoreError>;
async fn read(
&self,
accounts: &[ProviderAccountId],
) -> Result<Vec<UpstreamInfoRecord>, StoreError>;
/// 原子领取到期计划并设置租期;每次领取生成不能复用的领取标识
async fn claim_due(
&self,
provider: &ProviderKind,
now: SystemTime,
limit: u32,
lease_duration: Duration,
) -> Result<Vec<UpstreamInfoClaim>, StoreError>;
/// 返回 false 代表连接、领取标识或代次已经变化;Unauthorized 必须同时清空观测
async fn complete(&self, completion: UpstreamInfoCompletion) -> Result<bool, StoreError>;
}
@@ -0,0 +1,161 @@
//! 上游互联观测与后台计划;原始协议由控制面解释,不改变账号自身额度状态
use std::{
sync::{Arc, OnceLock, Weak},
time::SystemTime,
};
use sha2::{Digest, Sha256};
use super::{CredentialRevision, OpaqueProviderData, ProviderAccount, ProviderAccountId};
/// 请求发出时的连接身份;出口摘要包含托管代理解析后的认证信息,但不持久化明文
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpstreamConnection {
pub account_id: ProviderAccountId,
pub credential_revision: CredentialRevision,
pub revision: String,
}
impl UpstreamConnection {
#[must_use]
pub fn for_account(account: &ProviderAccount) -> Self {
Self::new(
account.id().clone(),
account.revision(),
account.outbound_proxy(),
)
}
#[must_use]
pub fn new(
account_id: ProviderAccountId,
credential_revision: CredentialRevision,
proxy: Option<&super::OutboundProxy>,
) -> Self {
let mut hash = Sha256::new();
hash.update(credential_revision.get().to_be_bytes());
if let Some(proxy) = proxy {
hash.update(proxy.expose_url().as_bytes());
}
Self {
account_id,
credential_revision,
revision: hash
.finalize()
.iter()
.map(|byte| format!("{byte:02x}"))
.collect(),
}
}
}
/// 可展示的受控同步原因,不保存上游正文、URL 或凭据
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum UpstreamInfoFailure {
Unauthorized,
Unsupported,
RateLimited,
InvalidResponse,
Unavailable,
}
impl UpstreamInfoFailure {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Unauthorized => "unauthorized",
Self::Unsupported => "unsupported",
Self::RateLimited => "rate_limited",
Self::InvalidResponse => "invalid_response",
Self::Unavailable => "unavailable",
}
}
}
/// 同步元数据独立于观测正文,计划查询与领取不加载完整快照
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpstreamInfoPlan {
pub connection: UpstreamConnection,
pub request_generation: u64,
pub completed_generation: u64,
pub next_attempt_at: SystemTime,
pub claimed_until: Option<SystemTime>,
pub last_attempt_at: Option<SystemTime>,
pub last_success_at: Option<SystemTime>,
pub last_error: Option<UpstreamInfoFailure>,
pub consecutive_failures: u32,
}
/// 同一连接最后成功观测与独立同步计划,失败不覆盖旧观测
#[derive(Debug, Clone, PartialEq)]
pub struct UpstreamInfoRecord {
pub plan: UpstreamInfoPlan,
pub snapshot: Option<OpaqueProviderData>,
}
/// 每次领取使用独立标识,计划删除重建也不能接受旧执行者的结果
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpstreamInfoClaim {
pub plan: UpstreamInfoPlan,
pub claim_token: uuid::Uuid,
}
/// 完整观测只在连接身份、领取标识与代次仍匹配时写入
#[derive(Debug, Clone, PartialEq)]
pub struct UpstreamInfoCompletion {
pub connection: UpstreamConnection,
pub request_generation: u64,
pub claim_token: uuid::Uuid,
pub completed_at: SystemTime,
pub next_attempt_at: SystemTime,
pub result: Result<OpaqueProviderData, UpstreamInfoFailure>,
}
/// 调度侧只消费已发布的中立观测,不读取协议、凭据或持久存储
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct UpstreamSchedulingSignal {
pub exhausted: bool,
pub remaining_percent: Option<f64>,
}
pub trait UpstreamInfoSource: Send + Sync {
fn signal(
&self,
account: &ProviderAccount,
model: &str,
now: SystemTime,
) -> Option<UpstreamSchedulingSignal>;
}
/// 组合根在管理服务就绪后绑定一次;弱引用避免 Provider 与管理服务互相保活
#[derive(Default)]
pub struct UpstreamInfoReader {
source: OnceLock<Weak<dyn UpstreamInfoSource>>,
}
impl UpstreamInfoReader {
pub fn bind(
&self,
source: &Arc<dyn UpstreamInfoSource>,
) -> Result<(), UpstreamInfoAlreadyBound> {
self.source
.set(Arc::downgrade(source))
.map_err(|_| UpstreamInfoAlreadyBound)
}
}
impl UpstreamInfoSource for UpstreamInfoReader {
fn signal(
&self,
account: &ProviderAccount,
model: &str,
now: SystemTime,
) -> Option<UpstreamSchedulingSignal> {
self.source.get()?.upgrade()?.signal(account, model, now)
}
}
#[derive(Debug, thiserror::Error)]
#[error("upstream information source is already bound")]
pub struct UpstreamInfoAlreadyBound;
@@ -1152,6 +1152,11 @@ impl fmt::Debug for ProviderRequest {
/// 新的 attempt 再次调用
#[async_trait]
pub trait Provider: Send + Sync {
/// 声明可用于互联共享的直连认证方式,不执行网络查询
fn is_direct_account(&self, _authentication_kind: &str) -> bool {
false
}
/// Provider 解释自己的隐私规则,Core 仅冻结编译结果
fn compile_privacy_policy(
&self,
@@ -1371,6 +1376,12 @@ impl ProviderRegistry {
}
impl ProviderCatalogPort for ProviderRegistry {
fn is_direct_account(&self, provider: &ProviderKind, authentication_kind: &str) -> bool {
self.providers
.get(provider)
.is_some_and(|provider| provider.is_direct_account(authentication_kind))
}
fn model_catalog_is_exhaustive(&self, provider: &ProviderKind) -> bool {
self.providers
.get(provider)
@@ -3,6 +3,9 @@
//! Client API Key 冻结账号分组权限;模型名称不参与权限判断
mod client_version;
mod peer_sharing;
pub use peer_sharing::{ClientPeerSharing, ClientPeerSharingError};
pub use client_version::{
ClientVersionRejection, CodexClientKind, CodexClientMinVersions, CodexClientVersion,
@@ -119,6 +122,7 @@ pub struct ClientSettings {
/// 从 `client_api_keys` 冻结的公开准入事实
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientPolicy {
peer_sharing: ClientPeerSharing,
name_snapshot: Option<String>,
defaults: Arc<ClientSettings>,
key_id: ClientApiKeyId,
@@ -164,6 +168,7 @@ impl ClientPolicy {
limits: RateLimits,
) -> Self {
Self {
peer_sharing: ClientPeerSharing::default(),
name_snapshot: None,
defaults: Arc::new(ClientSettings {
request_profiles: account_scope.request_profiles().clone(),
@@ -184,6 +189,23 @@ impl ClientPolicy {
self
}
#[must_use]
pub fn with_peer_sharing(mut self, sharing: ClientPeerSharing) -> Self {
self.account_scope = Arc::new(
self.account_scope
.as_ref()
.clone()
.with_direct_accounts_only(sharing.enabled),
);
self.peer_sharing = sharing;
self
}
#[must_use]
pub const fn peer_sharing(&self) -> ClientPeerSharing {
self.peer_sharing
}
#[must_use]
pub fn name_snapshot(&self) -> Option<&str> {
self.name_snapshot.as_deref()
@@ -0,0 +1,35 @@
//! Client Key 对互联节点公开本地观测的逐项授权
use serde::{Deserialize, Serialize};
/// 公开权限默认关闭,不因协议识别或预算查询自动开启
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ClientPeerSharing {
#[serde(default)]
pub enabled: bool,
#[serde(default)]
pub share_pool: bool,
#[serde(default)]
pub share_quota: bool,
#[serde(default)]
pub share_plans: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ClientPeerSharingError {
MissingGroupScope,
PlansWithoutPool,
}
impl ClientPeerSharing {
pub const fn validate(self, has_bound_groups: bool) -> Result<(), ClientPeerSharingError> {
if self.enabled && !has_bound_groups {
return Err(ClientPeerSharingError::MissingGroupScope);
}
if self.share_plans && !self.share_pool {
return Err(ClientPeerSharingError::PlansWithoutPool);
}
Ok(())
}
}
@@ -149,6 +149,11 @@ pub struct ProviderCatalogUnavailable;
/// 快照编译与对账使用的对象安全目录端口
pub trait ProviderCatalogPort: Send + Sync {
/// Provider 解释认证类型,快照只冻结直连事实;未知扩展不默认获得转供资格
fn is_direct_account(&self, _provider: &ProviderKind, _authentication_kind: &str) -> bool {
false
}
/// 目录是否完整到足以根据缺项拒绝请求;发现型目录交由上游验证模型名
fn model_catalog_is_exhaustive(&self, _provider: &ProviderKind) -> bool {
true
@@ -10,7 +10,8 @@ use crate::account::{FastMode, ProviderAccountId};
use crate::concurrency::ConcurrencyQueuePolicy;
use crate::operation::{Operation, OperationKind};
use crate::policy::{
ClientApiKeyId, ClientPolicy, CodexClientMinVersions, PlaintextClientApiKey, RateLimits,
ClientApiKeyId, ClientPeerSharing, ClientPolicy, CodexClientMinVersions, PlaintextClientApiKey,
RateLimits,
};
use crate::settings::{SettingsValues, compiled::CompiledSettings};
use crate::validation::RoutingError;
@@ -29,6 +30,7 @@ type ModelCatalogAccounts = BTreeMap<ProviderKind, BTreeMap<String, BTreeSet<Pro
/// Store 读取到的一个启用 Client API Key 策略事实
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotClientPolicyFacts {
peer_sharing: ClientPeerSharing,
name_snapshot: Option<String>,
request_profiles: BTreeMap<ProviderKind, crate::account::OpaqueProviderData>,
key_id: ClientApiKeyId,
@@ -38,6 +40,12 @@ pub struct SnapshotClientPolicyFacts {
}
impl SnapshotClientPolicyFacts {
#[must_use]
pub const fn with_peer_sharing(mut self, sharing: ClientPeerSharing) -> Self {
self.peer_sharing = sharing;
self
}
#[must_use]
pub fn with_name_snapshot(mut self, name: Option<String>) -> Self {
self.name_snapshot = name;
@@ -61,6 +69,7 @@ impl SnapshotClientPolicyFacts {
limits: RateLimits,
) -> Self {
Self {
peer_sharing: ClientPeerSharing::default(),
key_id,
name_snapshot: None,
request_profiles: BTreeMap::new(),
@@ -112,6 +121,7 @@ pub struct SnapshotProviderAccountFacts {
account_id: ProviderAccountId,
provider_kind: String,
model_access: crate::account::AccountModelAccess,
authentication_kind: String,
}
impl SnapshotProviderAccountFacts {
@@ -121,11 +131,18 @@ impl SnapshotProviderAccountFacts {
account_id,
provider_kind: provider_kind.into(),
model_access: crate::account::AccountModelAccess::all(),
authentication_kind: String::new(),
}
}
}
impl SnapshotProviderAccountFacts {
#[must_use]
pub fn with_authentication_kind(mut self, authentication_kind: String) -> Self {
self.authentication_kind = authentication_kind;
self
}
#[must_use]
pub fn with_model_access(mut self, model_access: crate::account::AccountModelAccess) -> Self {
self.model_access = model_access;
@@ -472,12 +489,17 @@ async fn compile_runtime_snapshot(
for account in facts.provider_accounts {
let provider_kind = ProviderKind::new(account.provider_kind)
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)?;
let direct_account =
catalogs.is_direct_account(&provider_kind, &account.authentication_kind);
// 账号归属独立于当前执行器集合;Provider 撤下后仍保留账号和分组,路由只选已注册执行器
if accounts
.insert(
account.account_id.clone(),
RuntimeAccount::new(provider_kind, BTreeSet::new())
.with_model_access(account.model_access),
.with_model_access(account.model_access)
.with_direct_authentication_kind(
direct_account.then_some(account.authentication_kind),
),
)
.is_some()
{
@@ -503,7 +525,10 @@ async fn compile_runtime_snapshot(
.get_mut(&account_id)
.ok_or(RuntimeSnapshotCompileError::InvalidData)?;
*account = RuntimeAccount::new(account.provider_kind().clone(), group_ids)
.with_model_access(account.model_access().clone());
.with_model_access(account.model_access().clone())
.with_direct_authentication_kind(
account.direct_authentication_kind().map(str::to_owned),
);
}
let account_directory = Arc::new(RuntimeAccountDirectory::new(accounts));
@@ -559,7 +584,8 @@ async fn compile_runtime_snapshot(
true,
policy.limits,
)
.with_name_snapshot(policy.name_snapshot),
.with_name_snapshot(policy.name_snapshot)
.with_peer_sharing(policy.peer_sharing),
);
}
@@ -1014,6 +1040,43 @@ impl RuntimeSnapshot {
&& !self.is_known_upstream_model(&upstream_model))
}
/// 通用模型候选资格,不以原生目录来源限制推理,也不替代操作能力与实时选号检查
#[must_use]
pub fn allows_public_model_for_account(
&self,
public_model: &PublicModelId,
scope: &FrozenAccountScope,
account: &ProviderAccountId,
) -> bool {
let Some(provider) = scope.account_provider(account) else {
return false;
};
let upstream = self.mapped_model(public_model.as_str());
if !scope.allows_model(account, &upstream)
|| !self.contains_public_model_for_provider(public_model, provider)
{
return false;
}
let Ok(upstream) = UpstreamModelId::from_client_wire(upstream) else {
return false;
};
let native = self
.provider_models
.get(provider)
.is_some_and(|models| models.contains_key(&upstream));
// 仅由适配器贡献的模型仍受声明账号范围限制,不能借发现型目录扩大资格
native
|| !self
.contributed_models()
.iter()
.any(|model| &model.provider == provider && &model.id == public_model)
|| self.contributed_models().iter().any(|model| {
&model.provider == provider
&& &model.id == public_model
&& model.accounts.contains(account)
})
}
fn is_known_upstream_model(&self, upstream_model: &UpstreamModelId) -> bool {
self.provider_models
.values()
@@ -3,6 +3,7 @@
mod location;
mod model_access;
mod proxy;
mod scope;
mod selection;
mod selection_diagnostics;
mod smart_scheduling;
@@ -0,0 +1,95 @@
//! 互联直连限制与分组、模型和请求局部范围的交集
use super::*;
use gateway_core::account::{AccountModelAccess, AccountModelAccessMode};
use gateway_core::routing::{AccountGroupId, RoutingGroupSnapshot};
#[test]
fn direct_scope_intersects_groups_models_and_request_restrictions() {
let group = AccountGroupId::new("grp_11111111111111111111111111111111").unwrap();
let provider = ProviderKind::new("openai").unwrap();
let direct = ProviderAccountId::new("acct_direct").unwrap();
let forwarded = ProviderAccountId::new("acct_forwarded").unwrap();
let unshared = ProviderAccountId::new("acct_unshared").unwrap();
let directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([
(
direct.clone(),
RuntimeAccount::new(provider.clone(), BTreeSet::from([group.clone()]))
.with_direct_authentication_kind(Some("oauth".into()))
.with_model_access(
AccountModelAccess::new(
AccountModelAccessMode::Allowlist,
vec!["gpt-5".into()],
)
.unwrap(),
),
),
(
forwarded.clone(),
RuntimeAccount::new(provider.clone(), BTreeSet::from([group.clone()])),
),
(
unshared.clone(),
RuntimeAccount::new(provider.clone(), BTreeSet::new())
.with_direct_authentication_kind(Some("oauth".into())),
),
])));
let scope = FrozenAccountScope::new(
directory.clone(),
ClientRoutingScope::restricted(
vec![RoutingGroupSnapshot::new(group.clone(), "共享组".into())],
BTreeSet::from([group]),
BTreeSet::from([provider]),
)
.unwrap(),
)
.with_direct_accounts_only(true);
assert!(scope.allows_model(&direct, "gpt-5"));
assert!(!scope.allows_model(&direct, "unshared-model"));
assert!(!scope.allows(&forwarded));
assert!(!scope.allows(&unshared));
let restricted = scope.restricted_to(&BTreeSet::new(), &BTreeSet::from([forwarded]));
assert!(restricted.provider_kinds().is_empty());
assert!(!restricted.allows(&direct));
assert!(!scope.excluding_account(direct.clone()).allows(&direct));
assert!(
!FrozenAccountScope::new(directory, ClientRoutingScope::AllAccounts)
.with_direct_accounts_only(true)
.allows(&direct)
);
}
#[test]
fn direct_scope_rejects_current_proxy_credentials_before_snapshot_publication() {
let group = AccountGroupId::new("grp_11111111111111111111111111111111").unwrap();
let provider = ProviderKind::new("openai").unwrap();
let id = ProviderAccountId::new("acct_direct").unwrap();
let directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
id.clone(),
RuntimeAccount::new(provider.clone(), BTreeSet::from([group.clone()]))
.with_direct_authentication_kind(Some("oauth".into())),
)])));
let scope = FrozenAccountScope::new(
directory,
ClientRoutingScope::restricted(
vec![RoutingGroupSnapshot::new(group.clone(), "共享组".into())],
BTreeSet::from([group]),
BTreeSet::from([provider.clone()]),
)
.unwrap(),
)
.with_direct_accounts_only(true);
let current = |kind: &str| {
ProviderAccount::new(
id.clone(),
provider.clone(),
"direct".into(),
None,
kind.into(),
CredentialRevision::new(2).unwrap(),
None,
)
};
assert!(scope.allows_current_account(&current("oauth")));
assert!(!scope.allows_current_account(&current("api_key")));
}
@@ -183,6 +183,23 @@ fn same_name_native_and_adapted_sources_keep_account_capabilities_separate() {
let scope = native.all_account_scope();
let b_only = Arc::new(scope.restricted_to(&Default::default(), &[b.clone()].into()));
assert!(native.public_models_for_scope(&b_only).is_empty());
assert!(native.allows_public_model_for_account(&id, &b_only, &b));
let c = ProviderAccountId::new("acct_c").unwrap();
assert!(!native.allows_public_model_for_account(&id, &scope, &c));
assert!(!native.allows_public_model_for_account(&id, &b_only, &c));
assert!(
native
.plan(
&id,
&super::operation(),
b_only.clone(),
&Default::default()
)
.unwrap()
.candidates()[0]
.account_scope()
.allows(&b)
);
let snapshot = native.with_extensions(Some(ExtensionSetReference::new(
ExtensionSetId::new("adapted".into()).unwrap(),
Arc::new(ModelLease(vec![contribution.clone()])),
@@ -277,7 +294,10 @@ fn adapted_models_limit_visibility_and_routing_to_capable_authorized_accounts()
use gateway_core::{
account::ProviderAccountId,
operation::{Feature, GenerateRequest, Operation, ProtocolPayload},
routing::{ContributedModel, SupportLevel},
routing::{
ClientRoutingScope, ContributedModel, FrozenAccountScope, RuntimeAccount,
RuntimeAccountDirectory, SupportLevel,
},
};
let account = ProviderAccountId::new("acct_openai").unwrap();
let provider = ProviderKind::new("openai").unwrap();
@@ -296,6 +316,16 @@ fn adapted_models_limit_visibility_and_routing_to_capable_authorized_accounts()
)));
let scope = super::account_scope();
assert!(snapshot.public_models_for_scope(&scope).contains(&id));
assert!(snapshot.allows_public_model_for_account(&id, &scope, &account));
let undeclared = ProviderAccountId::new("acct_undeclared").unwrap();
let all = FrozenAccountScope::new(
Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
undeclared.clone(),
RuntimeAccount::new(ProviderKind::new("openai").unwrap(), Default::default()),
)]))),
ClientRoutingScope::all_accounts(),
);
assert!(!snapshot.allows_public_model_for_account(&id, &all, &undeclared));
assert!(
snapshot
.public_model_profiles_for_scope(&scope)
@@ -1349,6 +1379,89 @@ fn facts(config_revision: u64, observed_current_revision: u64) -> SnapshotFacts
facts_with_min_versions(config_revision, observed_current_revision, None, None)
}
#[test]
fn peer_policy_compilation_keeps_only_explicitly_shared_direct_accounts() {
use gateway_core::{
account::ProviderAccountId, policy::ClientPeerSharing, routing::AccountGroupId,
};
struct DirectCatalog;
impl ProviderCatalogPort for DirectCatalog {
fn is_direct_account(&self, provider: &ProviderKind, kind: &str) -> bool {
provider.as_str() == "alpha" && kind == "oauth"
}
fn catalog_generations(&self) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
BTreeMap::from([(
ProviderKind::new("alpha").unwrap(),
ProviderCatalogGeneration::new(1),
)])
}
fn query_model_capabilities(
&self,
_: &ProviderKind,
) -> BoxFuture<'_, Result<Vec<ProviderModelCapabilities>, ProviderCatalogUnavailable>>
{
Box::pin(async { Ok(Vec::new()) })
}
}
let group = AccountGroupId::new("grp_11111111111111111111111111111111").unwrap();
let direct = ProviderAccountId::new("acct_direct").unwrap();
let proxy = ProviderAccountId::new("acct_proxy").unwrap();
for (bound, enabled, expected) in [
(true, true, true),
(true, false, false),
(false, true, false),
] {
let facts = SnapshotFacts::new(
revision(1),
revision(1),
SettingsValues::new(3, "smart", BTreeMap::new(), None, None),
vec![
SnapshotClientPolicyFacts::new(
ClientApiKeyId::new("key_peer").unwrap(),
PlaintextClientApiKey::new("sk_peer_test").unwrap(),
if bound {
vec![group.clone()]
} else {
Vec::new()
},
RateLimits::unlimited(),
)
.with_peer_sharing(ClientPeerSharing {
enabled: true,
..Default::default()
}),
],
vec![SnapshotAccountGroupFacts::new(
group.clone(),
"共享组".into(),
enabled,
)],
vec![
SnapshotProviderAccountFacts::new(direct.clone(), "alpha")
.with_authentication_kind("oauth".into()),
SnapshotProviderAccountFacts::new(proxy.clone(), "alpha")
.with_authentication_kind("api_key".into()),
],
vec![
SnapshotAccountGroupMemberFacts::new(group.clone(), direct.clone()),
SnapshotAccountGroupMemberFacts::new(group.clone(), proxy.clone()),
],
);
let snapshot = block_on(
RuntimeSnapshotCompiler::new(
Arc::new(TestSnapshotStore::new(Ok(facts))),
Arc::new(DirectCatalog),
)
.compile(),
)
.unwrap();
let scope = snapshot.client_policies().next().unwrap().account_scope();
assert_eq!(scope.allows(&direct), expected);
assert!(!scope.allows(&proxy));
assert_eq!(scope.provider_kinds().is_empty(), !expected);
}
}
fn facts_with_min_versions(
config_revision: u64,
observed_current_revision: u64,
+1
View File
@@ -18,6 +18,7 @@ flate2 = "1.1.9"
futures = "0.3.32"
gateway-admin = { path = "../gateway-admin" }
gateway-core = { path = "../gateway-core" }
gateway-protocol = { path = "../gateway-protocol" }
hex = "0.4.3"
http = "1.4.0"
http-body-util = "0.1.3"
+1
View File
@@ -6,6 +6,7 @@ pub mod config;
mod logging;
pub mod official_plugins;
pub mod outbound;
pub mod peer;
pub mod plugin_distribution;
pub mod pricing;
pub mod process;
+109
View File
@@ -0,0 +1,109 @@
//! 互联请求复用账号出口和证书策略,重定向不会携带密钥跨站点转发
use std::time::{Duration, SystemTime};
use futures::StreamExt as _;
use gateway_core::account::{ProviderAccount, UpstreamInfoFailure};
use secrecy::ExposeSecret as _;
use gateway_admin::ports::peer::{PeerConnection, PeerFetchFailure, PeerTransport};
use std::sync::Arc;
type ClientBuilder = dyn Fn(reqwest::ClientBuilder) -> Result<reqwest::Client, ()> + Send + Sync;
/// 组合根注入证书策略;出口、超时、重定向和响应上限由 Host 统一执行
pub struct HttpPeerTransport {
build_client: Arc<ClientBuilder>,
}
impl HttpPeerTransport {
pub fn with_client_builder<F, E>(build: F) -> Self
where
F: Fn(reqwest::ClientBuilder) -> Result<reqwest::Client, E> + Send + Sync + 'static,
{
Self {
build_client: Arc::new(move |builder| build(builder).map_err(|_| ())),
}
}
}
use gateway_protocol::peer::MAX_RESPONSE_BYTES as BODY_LIMIT;
const REQUEST_TIMEOUT: Duration = Duration::from_secs(5);
#[async_trait::async_trait]
impl PeerTransport for HttpPeerTransport {
async fn fetch(
&self,
account: &ProviderAccount,
authentication: &PeerConnection,
) -> Result<Vec<u8>, PeerFetchFailure> {
let mut builder = reqwest::Client::builder()
.no_proxy()
.redirect(reqwest::redirect::Policy::none())
.connect_timeout(REQUEST_TIMEOUT)
.timeout(REQUEST_TIMEOUT);
if let Some(proxy) = account.outbound_proxy() {
builder = builder.proxy(
reqwest::Proxy::all(proxy.expose_url())
.map_err(|_| UpstreamInfoFailure::Unavailable)?,
);
}
let client = (self.build_client)(builder).map_err(|_| UpstreamInfoFailure::Unavailable)?;
let response = client
.get(format!(
"{}/cpr/link",
authentication.base_url.trim_end_matches('/')
))
.bearer_auth(authentication.bearer_token.expose_secret())
.header(reqwest::header::ACCEPT, "application/json")
.timeout(REQUEST_TIMEOUT)
.send()
.await
.map_err(|_| UpstreamInfoFailure::Unavailable)?;
if !response.status().is_success() {
let kind = match response.status().as_u16() {
401 | 403 => UpstreamInfoFailure::Unauthorized,
404 | 405 | 501 => UpstreamInfoFailure::Unsupported,
429 => UpstreamInfoFailure::RateLimited,
_ => UpstreamInfoFailure::Unavailable,
};
let retry_after = response
.headers()
.get(reqwest::header::RETRY_AFTER)
.and_then(|value| value.to_str().ok())
.and_then(|value| {
value
.parse::<u64>()
.ok()
.map(Duration::from_secs)
.or_else(|| {
chrono::DateTime::parse_from_rfc2822(value)
.ok()
.and_then(|time| {
SystemTime::from(time)
.duration_since(SystemTime::now())
.ok()
})
})
})
.map(|value| value.min(Duration::from_secs(24 * 60 * 60)));
return Err(PeerFetchFailure { kind, retry_after });
}
if response
.content_length()
.is_some_and(|length| length > BODY_LIMIT as u64)
{
return Err(UpstreamInfoFailure::InvalidResponse.into());
}
let mut bytes = Vec::new();
let mut stream = response.bytes_stream();
while let Some(chunk) = stream.next().await {
let chunk = chunk.map_err(|_| UpstreamInfoFailure::Unavailable)?;
if bytes.len().saturating_add(chunk.len()) > BODY_LIMIT {
return Err(UpstreamInfoFailure::InvalidResponse.into());
}
bytes.extend_from_slice(&chunk);
}
Ok(bytes)
}
}
@@ -52,6 +52,7 @@ async fn key_facts_read_current_database_groups_without_secrets() {
.client_keys()
.update_client_key(
UpdateClientKey {
peer_sharing: None,
request_profile_override_updates: Default::default(),
id: ClientApiKeyId::new(id).unwrap(),
name: "facts".into(),
@@ -131,6 +131,7 @@ impl PluginAccountAccess for Facts {
self.reads.fetch_add(1, Ordering::SeqCst);
assert_eq!(account.as_str(), "acct_facts");
Ok(ProviderQuota {
upstream: None,
credits: Some(ProviderQuotaCredits {
has_credits: true,
unlimited: false,
@@ -140,6 +141,7 @@ impl PluginAccountAccess for Facts {
observed_at: None,
refresh_token_expires_at: None,
windows: vec![ProviderQuotaWindow {
upstream: None,
key: "weekly".into(),
group: "private group".into(),
label: "private label".into(),
@@ -65,6 +65,7 @@ async fn import_account(
crate::support::native::admin_registry(),
environment.store.admin_ports().accounts(),
core.snapshot_control(),
None,
);
let result = access
.save(
@@ -99,6 +99,7 @@ impl Environment {
.client_keys()
.create_client_key(
NewClientKey {
peer_sharing: Default::default(),
request_profile_overrides: BTreeMap::new(),
id: ClientApiKeyId::new(id).unwrap(),
name: format!("fixture {id}"),
@@ -494,6 +495,7 @@ impl Environment {
super::native::admin_registry(),
self.store.admin_ports().accounts(),
core.snapshot_control(),
None,
);
runtime.bind_account_ports(&access).unwrap();
self.plugin_accounts
@@ -582,6 +584,7 @@ impl Environment {
ClientConfig::default(),
self.store.admin_ports(),
gateway_admin::AdminRuntimePorts {
upstream_info: None,
timezone: Default::default(),
service_middleware: {
let snapshots = core.snapshots();
@@ -4,3 +4,6 @@
/// OpenAI wire contract
pub mod openai;
/// CPR 互联只读观测协议
pub mod peer;
+184
View File
@@ -0,0 +1,184 @@
//! CPR 互联的版本化只读线协议,不包含授权、采集或调度策略
use serde::{Deserialize, Serialize};
pub const OBJECT: &str = "cpr.peer";
pub const SCHEMA_VERSION: u32 = 1;
pub const MAX_ACCOUNTS: usize = 3_000;
pub const MAX_QUOTA_WINDOWS: usize = 6_000;
pub const MAX_MODELS: usize = 2_000;
pub const MAX_RESPONSE_BYTES: usize = 4 * 1024 * 1024;
/// 可选能力按当前认证主体公开,未知能力可以忽略
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct PeerLink {
pub object: String,
pub schema_version: u32,
pub node: PeerNode,
pub capabilities: Vec<String>,
pub observed_at: String,
pub expires_at: String,
pub usage: PeerUsage,
pub pool: PeerPool,
#[serde(default)]
pub quota: PeerQuota,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub models: Vec<PeerModelAvailability>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PeerNode {
pub id: String,
pub name: String,
pub software: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PeerComponentStatus {
Available,
Unavailable,
#[serde(other)]
Unknown,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PeerUsage {
pub status: PeerComponentStatus,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub observed_at: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub data: Option<PeerKeyUsage>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PeerKeyUsage {
pub unit: String,
pub daily: PeerBudgetWindow,
pub weekly: PeerBudgetWindow,
}
/// 金额保留十进制字符串,不限额时 total 与 remaining 为 null
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PeerBudgetWindow {
pub total: Option<String>,
pub used: String,
pub remaining: Option<String>,
#[serde(rename = "resetsAt")]
pub resets_at: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PeerPoolDisclosure {
Summary,
#[default]
#[serde(other)]
None,
}
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct PeerPool {
pub disclosure: PeerPoolDisclosure,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub observed_at: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub account_count: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub states: Option<PeerPoolStates>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub plans: Option<Vec<PeerPlanCount>>,
}
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct PeerPoolStates {
pub ready: u32,
pub cooldown: u32,
pub exhausted: u32,
pub disabled: u32,
pub error: u32,
pub unknown: u32,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PeerPlanCount {
pub name: String,
pub count: u32,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PeerQuotaDisclosure {
Accounts,
#[default]
#[serde(other)]
None,
}
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
pub struct PeerQuota {
pub disclosure: PeerQuotaDisclosure,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub accounts: Vec<PeerQuotaAccount>,
}
/// 匿名标识在授权 Key 内稳定,窗口仍保留账号与模型归属
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct PeerQuotaAccount {
pub id: String,
pub models: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub observed_at: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub expires_at: Option<String>,
pub windows: Vec<PeerQuotaWindow>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct PeerQuotaWindow {
pub key: String,
pub group: String,
pub label: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub limit_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub limit_name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub role: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub window_seconds: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub used_percent: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reset_at: Option<String>,
#[serde(default)]
pub limit_reached: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PeerModelCoverage {
Complete,
#[serde(other)]
Partial,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PeerModelState {
Available,
Exhausted,
#[serde(other)]
Unknown,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PeerModelAvailability {
pub model: String,
pub coverage: PeerModelCoverage,
pub state: PeerModelState,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub observed_at: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub expires_at: Option<String>,
}
@@ -11,6 +11,7 @@ pub struct StoreBundle {
admin_ports: AdminStorePorts,
core_ports: CoreStorePorts,
provider_ports: ProviderStorePorts,
upstream_info: Arc<dyn gateway_core::account::UpstreamInfoStore>,
retention: Arc<dyn gateway_admin::ports::retention::RetentionStore>,
worker_leader_lease: Arc<dyn WorkerLeaderLeasePort>,
health_probes: Vec<Arc<dyn HealthProbe>>,
@@ -36,6 +37,11 @@ impl StoreBundle {
self.core_ports.clone()
}
#[must_use]
pub fn upstream_info(&self) -> Arc<dyn gateway_core::account::UpstreamInfoStore> {
Arc::clone(&self.upstream_info)
}
#[must_use]
pub fn provider_ports(&self) -> ProviderStorePorts {
self.provider_ports.clone()
@@ -143,6 +149,8 @@ async fn connect(
postgres::PgProviderAccountRepository::new(pool.clone())
.with_account_error_protection(cooldowns.clone()),
);
let upstream_info: Arc<dyn gateway_core::account::UpstreamInfoStore> =
provider_accounts.clone();
let account_store: Arc<dyn ProviderAccountStore> = provider_accounts;
let credential_leases =
@@ -331,6 +339,7 @@ async fn connect(
admin_ports,
core_ports,
provider_ports,
upstream_info,
retention,
worker_leader_lease,
health_probes,
@@ -32,7 +32,7 @@ use gateway_core::{
engine::budget::{ClientBudgetLimits, ClientBudgetStatus},
engine::execution::ClientApiKeyUsageSink,
metering::Decimal,
policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits},
policy::{ClientApiKeyId, ClientPeerSharing, PlaintextClientApiKey, RateLimits},
routing::{AccountGroupId, ProviderKind},
task::{DaemonTask, WorkerDaemonContext, WorkerTaskError},
};
@@ -53,6 +53,7 @@ const CLIENT_API_KEY_LAST_USED_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientApiKeySnapshot {
pub peer_sharing: ClientPeerSharing,
pub name_snapshot: Option<String>,
pub request_profiles: std::collections::BTreeMap<
gateway_core::routing::ProviderKind,
@@ -73,6 +74,7 @@ impl ClientApiKeySnapshot {
requests_per_minute: i64,
) -> StoreResult<Self> {
Ok(Self {
peer_sharing: ClientPeerSharing::default(),
name_snapshot: None,
request_profiles: std::collections::BTreeMap::new(),
id: ClientApiKeyId::new(id)
@@ -118,6 +120,7 @@ impl fmt::Debug for ClientApiKeySecret {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientApiKeyRecord {
pub peer_sharing: ClientPeerSharing,
pub request_profile_overrides: BTreeMap<ProviderKind, OpaqueProviderData>,
pub id: String,
pub name: String,
@@ -279,6 +282,7 @@ pub struct ClientApiKeyPage {
#[derive(Clone)]
pub struct NewClientApiKey {
pub peer_sharing: ClientPeerSharing,
pub request_profile_overrides: BTreeMap<ProviderKind, OpaqueProviderData>,
pub id: String,
pub name: String,
@@ -307,12 +311,16 @@ impl NewClientApiKey {
require_nonempty(ENTITY, "name", &self.name)?;
validate_group_ids(&self.group_ids)?;
validate_request_profiles(self.request_profile_overrides.values())?;
self.peer_sharing
.validate(!self.group_ids.is_empty())
.map_err(|_| invalid("invalid peer sharing scope"))?;
validate_key(&self.key)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpdateClientApiKeyDetails {
pub peer_sharing: Option<ClientPeerSharing>,
pub request_profile_override_updates: BTreeMap<ProviderKind, Option<OpaqueProviderData>>,
pub id: String,
pub name: String,
@@ -378,7 +386,7 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
query.validate()?;
let total = count_client_api_keys(&self.pool, query.search.as_deref()).await?;
let mut statement = QueryBuilder::<Postgres>::new(
"select k.id, k.name, k.label, k.provider_request_profiles_json as request_profile_overrides,
"select k.id, k.name, k.label, k.peer_sharing_json, k.provider_request_profiles_json as request_profile_overrides,
left(k.key, least(10, length(k.key) / 2)) as prefix, k.enabled,
k.max_concurrency, k.requests_per_minute, k.last_used_at, k.created_at,
k.updated_at, '[]'::jsonb as groups, '{}'::text[] as provider_kinds
@@ -438,7 +446,7 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
async fn get_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeyRecord>> {
require_nonempty(ENTITY, "id", id)?;
let record = sqlx::query(
"select k.id, k.name, k.label, k.provider_request_profiles_json as request_profile_overrides,
"select k.id, k.name, k.label, k.peer_sharing_json, k.provider_request_profiles_json as request_profile_overrides,
left(k.key, least(10, length(k.key) / 2)) as prefix, k.enabled,
k.max_concurrency, k.requests_per_minute, k.last_used_at, k.created_at,
k.updated_at, coalesce(groups.groups, '[]'::jsonb) as groups,
@@ -722,6 +730,20 @@ impl PgAdminClientKeyStore {
#[async_trait]
impl ClientKeyStore for PgAdminClientKeyStore {
async fn peer_node_identity(
&self,
) -> AdminStoreResult<gateway_admin::model::peer_link::PeerNodeIdentity> {
let (id, name) = sqlx::query_as::<_, (String, String)>(
"select id::text, name from peer_node where singleton",
)
.fetch_one(&self.keys.pool)
.await
.map_err(|source| {
admin_store_error("peer node", postgres_unavailable("load peer node", source))
})?;
Ok(gateway_admin::model::peer_link::PeerNodeIdentity { id, name })
}
async fn update_client_key_budget_limits(
&self,
command: gateway_admin::model::client_keys::UpdateClientKeyBudgetLimits,
@@ -889,6 +911,7 @@ impl ClientKeyStore for PgAdminClientKeyStore {
.control_plane
.create_client_api_key(
NewClientApiKey {
peer_sharing: command.peer_sharing,
request_profile_overrides: command.request_profile_overrides,
id: id.as_str().to_owned(),
name: command.name,
@@ -918,6 +941,7 @@ impl ClientKeyStore for PgAdminClientKeyStore {
"daily_limit_usd",
"weekly_limit_usd",
"provider_request_profiles_json",
"peer_sharing_json",
]
.into_iter()
.map(str::to_owned)
@@ -939,6 +963,7 @@ impl ClientKeyStore for PgAdminClientKeyStore {
.control_plane
.update_client_api_key(
UpdateClientApiKeyDetails {
peer_sharing: command.peer_sharing,
request_profile_override_updates: command.request_profile_override_updates,
id: id.as_str().to_owned(),
name: command.name,
@@ -966,6 +991,7 @@ impl ClientKeyStore for PgAdminClientKeyStore {
"daily_limit_usd",
"weekly_limit_usd",
"provider_request_profiles_json",
"peer_sharing_json",
]
.into_iter()
.map(str::to_owned)
@@ -1096,6 +1122,7 @@ fn admin_client_key_cursor(cursor: ClientApiKeyCursor) -> AdminStoreResult<Admin
fn admin_client_key_record(record: ClientApiKeyRecord) -> AdminStoreResult<AdminClientKeyRecord> {
Ok(AdminClientKeyRecord {
peer_sharing: record.peer_sharing,
request_profile_overrides: record.request_profile_overrides,
id: ClientApiKeyId::new(record.id)
.map_err(|_| admin_store_error(ENTITY, invalid("invalid client key id")))?,
@@ -1148,8 +1175,8 @@ pub(crate) async fn insert_client_api_key_in_transaction(
sqlx::query(
"insert into client_api_keys (
id, name, label, key, enabled, max_concurrency, requests_per_minute,
last_used_at, created_at, updated_at, daily_limit_usd, weekly_limit_usd, provider_request_profiles_json
) values ($1, $2, $3, $4, true, $5, $6, null, now(), now(), $7::text::numeric, $8::text::numeric, $9::jsonb)",
last_used_at, created_at, updated_at, daily_limit_usd, weekly_limit_usd, provider_request_profiles_json, peer_sharing_json
) values ($1, $2, $3, $4, true, $5, $6, null, now(), now(), $7::text::numeric, $8::text::numeric, $9::jsonb, $10::jsonb)",
)
.bind(&key.id)
.bind(key.name.trim())
@@ -1160,6 +1187,7 @@ pub(crate) async fn insert_client_api_key_in_transaction(
.bind(key.budget.daily_usd.canonical())
.bind(key.budget.weekly_usd.canonical())
.bind(sqlx::types::Json(request_profile_overrides))
.bind(sqlx::types::Json(key.peer_sharing))
.execute(&mut **transaction)
.await
.map_err(|error| {
@@ -1187,6 +1215,28 @@ pub(crate) async fn update_client_api_key_in_transaction(
) -> StoreResult<()> {
key.validate()?;
ensure_client_key_name_available(transaction, &key.id, key.name.trim()).await?;
// 控制面行锁已串行化配置写入,按同一事务里的保留配置验证分组修改
let sharing = match key.peer_sharing {
Some(sharing) => sharing,
None => {
sqlx::query_scalar::<_, sqlx::types::Json<ClientPeerSharing>>(
"select peer_sharing_json from client_api_keys where id = $1",
)
.bind(&key.id)
.fetch_optional(&mut **transaction)
.await
.map_err(|source| postgres_unavailable("load client peer sharing", source))?
.ok_or_else(|| StoreError::NotFound {
source: None,
entity: ENTITY,
id: key.id.clone(),
})?
.0
}
};
sharing
.validate(!key.group_ids.is_empty())
.map_err(|_| invalid("invalid peer sharing scope"))?;
let removed_profiles = key
.request_profile_override_updates
.keys()
@@ -1211,7 +1261,8 @@ pub(crate) async fn update_client_api_key_in_transaction(
daily_limit_usd = coalesce($6::text::numeric, daily_limit_usd),
weekly_limit_usd = coalesce($7::text::numeric, weekly_limit_usd),
provider_request_profiles_json = (provider_request_profiles_json - $8::text[])
|| $9::jsonb
|| $9::jsonb,
peer_sharing_json = coalesce($10::jsonb, peer_sharing_json)
where id = $1",
)
.bind(&key.id)
@@ -1223,6 +1274,7 @@ pub(crate) async fn update_client_api_key_in_transaction(
.bind(key.weekly_limit_usd.map(|amount| amount.canonical()))
.bind(removed_profiles)
.bind(sqlx::types::Json(replacement_profiles))
.bind(key.peer_sharing.map(sqlx::types::Json))
.execute(&mut **transaction)
.await
.map_err(|source| postgres_unavailable("update client API key in transaction", source))?;
@@ -1420,6 +1472,10 @@ fn client_record_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<ClientApiK
.map_err(|source| invalid("invalid groups").with_source(source))?;
Ok(
ClientApiKeyRecord {
peer_sharing: row
.try_get::<sqlx::types::Json<ClientPeerSharing>, _>("peer_sharing_json")
.map_err(|source| invalid("invalid peer sharing").with_source(source))?
.0,
request_profile_overrides:
decode_request_profiles(
row.try_get::<sqlx::types::Json<

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