Compare commits

...
12 Commits
Author SHA1 Message Date
ZyyandCodex 0534dd8f26 chore(release): v3.16.0
Co-authored-by: Codex <noreply@openai.com>
2026-09-26 22:02:00 +08:00
ZyyandCodex 525f3d51c1 refactor(store): consolidate unreleased scheduling migration
Fold smart scheduling settings into migration 0019 while preserving all released migrations through 0018. Update the checksum manifest and upgrade test.

Development databases that applied the previous 0019 or 0020 cannot upgrade in place because the migration history differs.

Co-authored-by: Codex <noreply@openai.com>
2026-09-26 21:59:05 +08:00
CharmskyandCodex 33cc427b86 feat(scheduling): add configurable smart scheduling preferences (#295)
Make smart scheduling coefficients configurable and expose optional
higher-weight switchback while preserving the original defaults.

Reuse current quota reset signals, process-local queues, and frozen request
settings while preserving continuation ownership and queue fairness.

Co-authored-by: Codex <noreply@openai.com>
2026-09-26 21:51:28 +08:00
ZyyandCodex d1bf90d5fd fix(frontend): support HTTP plugin pages and consolidate shared helpers
Reuse secure UUID generation for plugin frame channels and sessions on HTTP.
Consolidate utility modules and shorten the plugin views store name.
Declare route roles and default entries in metadata instead of guard mappings.
Keep source integration build output under node_modules/.vite.

Validated frontend lint, production and source builds, browser plugin flows,
and 20 route guard scenarios with simulated sessions and plugin resources.

Fixes #288

Co-authored-by: Codex <noreply@openai.com>
2026-09-26 19:33:17 +08:00
eb23a0fb3b fix(plugins): align management request timeouts (#294)
Use a shared 180-second timeout for management route HTTP calls and the page bridge while keeping the backend RPC limit at 120 seconds. Fix the existing Vue attribute-order lint warning.

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-09-26 19:11:01 +08:00
niceSpriteandCodex ca90f2b586 fix(openai): recover oversized Codex requests before WebSocket send (#287)
Co-authored-by: Codex <noreply@openai.com>
2026-09-26 19:03:40 +08:00
ZyyandCodex cfcf7c486a fix(plugins): detect latest releases and streamline upgrades (#293)
Refresh explicit GitHub release checks while preserving concurrent fetch sharing and fixed-release download caching.

Compare verified package versions and offer a direct upgrade with pinned digests, revision checks, permission disclosure, and configuration fallback.

Co-authored-by: Codex <noreply@openai.com>
2026-09-26 19:00:23 +08:00
b981e42342 feat(plugins): provision owned resources and reconcile account changes (#292)
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-09-26 18:46:12 +08:00
ZyyandCodex 02177af921 fix(accounts): refresh OAuth tokens while scheduling is disabled (#291)
Separate scheduling availability from OAuth candidate selection and credential
state writes across OpenAI, xAI, and the shared account store. Cover disabled
refresh, terminal failures, and in-flight scheduling changes with regression tests.

Co-authored-by: Codex <noreply@openai.com>
2026-09-26 18:32:48 +08:00
ZyyandCodex 2d20f03bcb refactor(admin): simplify diagnostics and refresh shared UI integration
Adopt UI v0.3.0 and its shared Vite adapter, preserve request-volume ordering in diagnostics, stabilize sidebar animation, and track workbench v0.1.15.

Co-authored-by: Codex <noreply@openai.com>
2026-09-26 17:14:54 +08:00
ZyyandCodex c12c600b42 fix(models): scope catalog visibility to discovering accounts (#286)
Preserve OpenAI catalog account sources in runtime snapshots and apply
frozen account scope and model policies when publishing model metadata.
Keep discovery catalogs separate from inference authorization and cover
mixed policies, aliases, group changes, and catalog generation updates.

Co-authored-by: Codex <noreply@openai.com>
2026-09-26 15:51:13 +08:00
ZyyandCodex 1c6b5a8f10 fix(accounts): restore OAuth reauthorization flow
Generate authorization links when reopening account authorization and regenerate them in a single action. Remove the added authorization page shortcut and unused reset event wiring.

Co-authored-by: Codex <noreply@openai.com>
2026-09-26 05:32:04 +08:00
212 changed files with 5382 additions and 697 deletions
-1
View File
@@ -6,7 +6,6 @@
# Frontend
/frontend/node_modules/
/frontend/dist/
/frontend/.vite/
/frontend/.pnpm-store/
/frontend/.pnpm-cache/
*.tsbuildinfo
+1 -1
View File
@@ -148,7 +148,7 @@ AI 进行前端开发时,加载并使用 `$frontend-design` 技能。设计方
| 依赖安全 | Cargo Audit、pnpm Audit | 本次变更的权限、数据与信任边界 |
| 工作流与容器 | Actionlint、相关容器检查 | 改动涉及的运行行为 |
UI 与示例源码分别在独立的 `codex-proxy-ui`、`codex-proxy-plugins` 仓库维护。前端的依赖、锁文件和 ESLint 配置均在 `frontend/` 管理,仓库根目录不维护 Node 包。UI 使用固定 GitHub 标签或提交依赖,锁文件记录实际提交、源码归档地址及完整性摘要,常规开发和 CI 不需要检出同级仓库。安装时 pnpm 按 UI 自身锁文件准备构建依赖,再执行 `prepack` 生成 JS、CSS 和类型声明。开发服务通过 UI 包的 `development` 入口使用 Vue SFC,生产构建使用 `dist`。
UI 与示例源码分别在独立的 `codex-proxy-ui`、`codex-proxy-plugins` 仓库维护。前端的依赖、锁文件和 ESLint 配置均在 `frontend/` 管理,仓库根目录不维护 Node 包。UI 使用固定 GitHub 标签或提交依赖,锁文件记录实际提交、源码归档地址及完整性摘要,常规开发和 CI 不需要检出同级仓库。安装时 pnpm 按 UI 自身锁文件准备构建依赖,再执行 `prepack` 生成 JS、CSS 和类型声明。普通开发与生产构建都使用 UI 编译产物,源码联调通过 `@codex-proxy/ui/vite` 显式开启。
pnpm 版本由 `frontend/package.json` 的 `packageManager` 固定,CI 和 Docker 从该字段读取。`frontend/pnpm-workspace.yaml` 的 `allowBuilds` 按仓库放行自有 UI 包的构建脚本;升级依赖时核对源码并更新锁文件,无需逐次修改构建许可,不开放全部依赖脚本。
+17
View File
@@ -187,6 +187,14 @@ async fn launch(
core.snapshot_control(),
);
or_shutdown!(plugin_runtime, plugin_runtime.bind_client_key_ports(&keys));
let plugin_resources = gateway_admin::initialize_plugin_resources(
store.admin_ports().plugin_resources(),
core.snapshot_control(),
);
or_shutdown!(
plugin_runtime,
plugin_runtime.bind_resource_ports(&plugin_resources)
);
let core = or_shutdown!(plugin_runtime, core.activate().await);
let nested_models = core.nested_model_execution_port();
let affinity_lookup = core.affinity_lookup_port();
@@ -231,6 +239,14 @@ async fn launch(
plugin_runtime,
plugin_runtime.bind_client_key_ports(&plugin_keys)
);
let plugin_resources = gateway_admin::initialize_plugin_resources(
store.admin_ports().plugin_resources(),
core.snapshot_control(),
);
or_shutdown!(
plugin_runtime,
plugin_runtime.bind_resource_ports(&plugin_resources)
);
let mut core = or_shutdown!(plugin_runtime, core.activate().await);
let nested_models = core.nested_model_execution_port();
let affinity_lookup = core.affinity_lookup_port();
@@ -322,6 +338,7 @@ async fn launch(
plan.extend(openai.take_worker_contributions());
plan.extend(xai.take_worker_contributions());
plan.extend(admin.take_worker_contributions());
plan.push(plugin_runtime.maintenance_worker(core.snapshots())?);
or_shutdown!(
plugin_runtime,
host.start_workers(plan, store.worker_leader_lease())
+9
View File
@@ -539,6 +539,15 @@ pub fn initialize_plugin_client_keys(
Arc::new(use_case::plugin_client_keys::DefaultPluginClientKeyAccess::new(service))
}
/// 为 Runtime 组合实例自有资源写入;权限和归属在同一存储事务复核。
#[must_use]
pub fn initialize_plugin_resources(
store: Arc<dyn ports::plugin_resources::PluginResourceStore>,
snapshot: Arc<dyn SnapshotControl>,
) -> Arc<dyn ports::plugin_resources::PluginResourceAccess> {
Arc::new(use_case::plugin_resources::DefaultPluginResourceAccess { store, snapshot })
}
/// Backup Worker 注册:单个可取消 Daemon,owner 固定为 `backup`。
fn backup_worker_contribution(
task: backup::task::BackupTask,
@@ -12,6 +12,7 @@ pub mod import_tasks;
pub mod key_usage;
pub mod observability;
pub mod plugin_client_keys;
pub mod plugin_resources;
pub mod plugins;
pub mod pricing;
pub mod provider_credentials;
@@ -1154,7 +1154,6 @@ pub struct DiagnosticsItem {
pub non_completion_rate: f64,
pub retry_count: u64,
pub retry_rate: f64,
pub impact_score: f64,
pub estimated_cost: Option<DecimalAmount>,
pub attempt_count: u64,
pub total_tokens: u64,
@@ -0,0 +1,42 @@
//! 插件资源归属由宿主调用上下文给出,不能由插件参数选择实例或代次。
use super::Revision;
#[derive(Debug, Clone)]
pub struct PluginResourceOwner {
pub instance_id: String,
pub artifact_sha256: String,
pub revision: Revision,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ManagedResource {
pub id: String,
pub name: String,
pub enabled: bool,
}
pub struct ResourceMutation<T> {
/// 没有事实变化时不递增配置 revision,也不触发发布。
pub revision: Option<Revision>,
pub value: T,
}
pub struct GroupMembersChange {
pub resource_key: String,
pub add: Vec<String>,
pub remove: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GroupMembersChanged {
pub added: u64,
pub removed: u64,
}
/// 插件 Key 的首次创建设置;分组始终从本实例资源键解析。
pub struct ManagedKeyConfig {
pub name: String,
pub limits: gateway_core::policy::RateLimits,
pub budget: gateway_core::engine::budget::ClientBudgetLimits,
}
@@ -38,6 +38,7 @@ pub struct RuntimeSettings {
pub max_waiting_per_account: u32,
pub concurrency_wait_timeout_seconds: u32,
pub responses_max_decompressed_body_bytes: u64,
pub smart_scheduling: gateway_core::account::SmartSchedulingConfig,
pub rotation_strategy: RotationStrategy,
pub min_codex_desktop_version: Option<String>,
pub min_codex_cli_version: Option<String>,
@@ -73,6 +74,7 @@ pub struct ReplaceRuntimeSettings {
pub max_waiting_per_account: u32,
pub concurrency_wait_timeout_seconds: u32,
pub responses_max_decompressed_body_bytes: u64,
pub smart_scheduling: gateway_core::account::SmartSchedulingConfig,
pub rotation_strategy: RotationStrategy,
pub min_codex_desktop_version: Option<String>,
pub min_codex_cli_version: Option<String>,
@@ -6,6 +6,7 @@ pub mod plugin_accounts;
pub mod plugin_client_keys;
pub mod plugin_management;
pub mod plugin_release;
pub mod plugin_resources;
pub mod plugins;
pub mod pricing;
pub mod provider;
@@ -0,0 +1,63 @@
//! 插件自有资源的管理入口与原子持久化端口。
use super::store::AdminStoreResult;
use crate::model::{
AdminError, MutationContext,
account_groups::{CreateAccountGroup, NewAccountGroup},
client_keys::NewClientKey,
plugin_resources::{
GroupMembersChange, GroupMembersChanged, ManagedKeyConfig, ManagedResource,
PluginResourceOwner, ResourceMutation,
},
};
use async_trait::async_trait;
#[async_trait]
pub trait PluginResourceAccess: Send + Sync {
async fn ensure_group(
&self,
owner: &PluginResourceOwner,
resource_key: String,
command: CreateAccountGroup,
context: &MutationContext,
) -> Result<ManagedResource, AdminError>;
async fn ensure_key(
&self,
owner: &PluginResourceOwner,
resource_key: String,
groups: Vec<String>,
command: ManagedKeyConfig,
context: &MutationContext,
) -> Result<ManagedResource, AdminError>;
async fn change_members(
&self,
owner: &PluginResourceOwner,
command: GroupMembersChange,
context: &MutationContext,
) -> Result<GroupMembersChanged, AdminError>;
}
#[async_trait]
pub trait PluginResourceStore: Send + Sync {
async fn ensure_group(
&self,
owner: &PluginResourceOwner,
resource_key: String,
command: NewAccountGroup,
context: &MutationContext,
) -> AdminStoreResult<ResourceMutation<ManagedResource>>;
async fn ensure_key(
&self,
owner: &PluginResourceOwner,
resource_key: String,
groups: Vec<String>,
command: NewClientKey,
context: &MutationContext,
) -> AdminStoreResult<ResourceMutation<ManagedResource>>;
async fn change_members(
&self,
owner: &PluginResourceOwner,
command: GroupMembersChange,
context: &MutationContext,
) -> AdminStoreResult<ResourceMutation<GroupMembersChanged>>;
}
@@ -515,9 +515,15 @@ pub struct AdminStorePorts {
backup: BackupStorePorts,
plugins: Arc<dyn super::plugins::PluginStore>,
plugin_state: Arc<dyn super::plugins::PluginStateStore>,
plugin_resources: Arc<dyn super::plugin_resources::PluginResourceStore>,
}
impl AdminStorePorts {
#[must_use]
pub fn plugin_resources(&self) -> Arc<dyn super::plugin_resources::PluginResourceStore> {
self.plugin_resources.clone()
}
#[must_use]
pub fn plugins(&self) -> Arc<dyn super::plugins::PluginStore> {
Arc::clone(&self.plugins)
@@ -542,6 +548,7 @@ impl AdminStorePorts {
backup: BackupStorePorts,
plugins: Arc<dyn super::plugins::PluginStore>,
plugin_state: Arc<dyn super::plugins::PluginStateStore>,
plugin_resources: Arc<dyn super::plugin_resources::PluginResourceStore>,
) -> Self {
Self {
accounts,
@@ -552,6 +559,7 @@ impl AdminStorePorts {
backup,
plugins,
plugin_state,
plugin_resources,
}
}
@@ -124,9 +124,7 @@ impl ClientKeyService for DefaultClientKeyService {
let plaintext = if let Some(key) = command.custom_key {
key.expose_for_auth().to_owned()
} else {
let mut bytes = [0_u8; 32];
OsRng.fill_bytes(&mut bytes);
format!("sk_{}", URL_SAFE_NO_PAD.encode(bytes))
generate_key()
};
let (config_revision, record) = self
.store
@@ -257,3 +255,10 @@ fn validate_cursor(query: &ClientKeyListQuery) -> Result<(), AdminError> {
Err(AdminError::invalid("Client API Key 游标不合法"))
}
}
// 原生与插件创建共用相同的密钥生成规则。
pub(super) fn generate_key() -> String {
let mut bytes = [0_u8; 32];
OsRng.fill_bytes(&mut bytes);
format!("sk_{}", URL_SAFE_NO_PAD.encode(bytes))
}
@@ -12,6 +12,7 @@ pub mod key_usage;
pub mod observability;
pub mod plugin_accounts;
pub mod plugin_client_keys;
pub mod plugin_resources;
pub(crate) mod plugin_update;
pub mod plugins;
pub mod proxies;
@@ -342,21 +342,13 @@ impl ObservabilityService for DefaultObservabilityService {
let total_requests = items.iter().fold(0_u64, |total, item| {
total.saturating_add(item.request_count)
});
let mut items = items
let items = items
.into_iter()
.map(|item| {
let error_rate = rate_or_zero(item.failure_count, item.request_count);
let non_completion_rate =
rate_or_zero(item.non_completion_count, item.request_count);
let retry_rate = rate_or_zero(item.retry_count, item.request_count);
let impact_score = diagnostic_impact_score(
item.request_count,
total_requests,
error_rate,
non_completion_rate,
retry_rate,
item.first_token_p95_ms,
);
DiagnosticsItem {
key: item.key,
name: item.name,
@@ -372,19 +364,12 @@ impl ObservabilityService for DefaultObservabilityService {
non_completion_rate,
retry_count: item.retry_count,
retry_rate,
impact_score,
estimated_cost: usd_cost(&item.costs),
attempt_count: item.attempt_count,
total_tokens: item.total_tokens,
}
})
.collect::<Vec<_>>();
items.sort_by(|left, right| {
right
.impact_score
.total_cmp(&left.impact_score)
.then_with(|| right.request_count.cmp(&left.request_count))
});
Ok(DiagnosticsResult { dimension, items })
}
@@ -724,25 +709,6 @@ fn decimal(value: &DecimalAmount) -> Option<gateway_core::metering::Decimal> {
value.as_str().parse().ok()
}
fn diagnostic_impact_score(
request_count: u64,
total_requests: u64,
error_rate: f64,
non_completion_rate: f64,
retry_rate: f64,
first_token_p95_ms: Option<u64>,
) -> f64 {
let request_share = rate_or_zero(request_count, total_requests);
let slow_score = first_token_p95_ms
.map_or(0.0, |value| value as f64 / 30_000.0)
.min(1.0);
error_rate * 0.35
+ non_completion_rate * 0.25
+ retry_rate.min(1.0) * 0.20
+ request_share * 0.10
+ slow_score * 0.10
}
fn usd_cost(costs: &[CurrencyCost]) -> Option<DecimalAmount> {
costs
.iter()
@@ -0,0 +1,108 @@
use crate::{
model::{
AdminError, MutationContext,
account_groups::{CreateAccountGroup, NewAccountGroup},
client_keys::NewClientKey,
plugin_resources::{
GroupMembersChange, GroupMembersChanged, ManagedKeyConfig, ManagedResource,
PluginResourceOwner, ResourceMutation,
},
},
ports::plugin_resources::{PluginResourceAccess, PluginResourceStore},
};
use async_trait::async_trait;
use gateway_core::{policy::ClientApiKeyId, routing::AccountGroupId, runtime::SnapshotControl};
use std::sync::Arc;
use uuid::Uuid;
pub(crate) struct DefaultPluginResourceAccess {
pub store: Arc<dyn PluginResourceStore>,
pub snapshot: Arc<dyn SnapshotControl>,
}
impl DefaultPluginResourceAccess {
async fn publish<T>(
&self,
result: crate::ports::store::AdminStoreResult<ResourceMutation<T>>,
) -> Result<T, AdminError> {
let result = result.map_err(|error| super::map_store_error(error, "plugin resource"))?;
if let Some(revision) = result.revision {
super::publish_committed(self.snapshot.as_ref(), revision).await?;
}
Ok(result.value)
}
}
#[async_trait]
impl PluginResourceAccess for DefaultPluginResourceAccess {
async fn ensure_group(
&self,
owner: &PluginResourceOwner,
resource_key: String,
command: CreateAccountGroup,
context: &MutationContext,
) -> Result<ManagedResource, AdminError> {
let id = AccountGroupId::new(format!("grp_{}", Uuid::now_v7().simple()))
.map_err(|_| AdminError::internal("创建账号组 ID 失败"))?;
self.publish(
self.store
.ensure_group(
owner,
resource_key,
NewAccountGroup {
id,
name: command.name,
description: command.description,
color: command.color,
disable_fast: command.disable_fast,
},
context,
)
.await,
)
.await
}
async fn ensure_key(
&self,
owner: &PluginResourceOwner,
resource_key: String,
groups: Vec<String>,
command: ManagedKeyConfig,
context: &MutationContext,
) -> Result<ManagedResource, AdminError> {
let id = ClientApiKeyId::new(format!("key_{}", Uuid::now_v7().simple()))
.map_err(|_| AdminError::internal("创建 Key ID 失败"))?;
self.publish(
self.store
.ensure_key(
owner,
resource_key,
groups,
NewClientKey {
id,
name: command.name,
label: None,
group_ids: Vec::new(),
limits: command.limits,
budget: command.budget,
request_profile_overrides: Default::default(),
plaintext: super::client_keys::generate_key(),
},
context,
)
.await,
)
.await
}
async fn change_members(
&self,
owner: &PluginResourceOwner,
command: GroupMembersChange,
context: &MutationContext,
) -> Result<GroupMembersChanged, AdminError> {
self.publish(self.store.change_members(owner, command, context).await)
.await
}
}
@@ -377,7 +377,11 @@ pub(super) fn default_bindings(metadata: &PluginArtifactMetadata) -> Vec<PluginC
.filter(|(capability, _)| {
!matches!(
capability.as_str(),
"frontend_authentication" | "management" | "command_line" | "model_catalog"
"frontend_authentication"
| "management"
| "command_line"
| "model_catalog"
| "maintenance"
)
})
.flat_map(|(_, contribution)| {
@@ -1105,6 +1105,7 @@ impl SettingsStore for StaticSettingsStore {
max_waiting_per_account: 0,
concurrency_wait_timeout_seconds: 30,
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
smart_scheduling: gateway_core::account::SmartSchedulingConfig::default(),
rotation_strategy: RotationStrategy::Smart,
min_codex_desktop_version: None,
min_codex_cli_version: None,
@@ -38,6 +38,7 @@ fn runtime_settings(enabled: bool, probe_enabled: bool, adaptive: bool) -> Runti
max_waiting_per_account: 0,
concurrency_wait_timeout_seconds: 30,
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
smart_scheduling: gateway_core::account::SmartSchedulingConfig::default(),
rotation_strategy: gateway_admin::model::settings::RotationStrategy::Smart,
min_codex_desktop_version: None,
min_codex_cli_version: None,
@@ -258,6 +258,7 @@ impl AdminHarness {
self.backup,
self.plugin_store,
Arc::new(plugins::TestPluginPorts),
Arc::new(plugins::TestPluginPorts),
),
gateway_admin::AdminRuntimePorts {
plugin_preparation: Arc::new(plugins::TestPluginPorts),
@@ -502,7 +502,31 @@ async fn observability_services_should_calculate_usage_insights_and_diagnostic_s
);
assert_eq!(diagnostics.items[0].request_share, 0.75);
assert_eq!(diagnostics.items[1].request_share, 0.25);
assert!((diagnostics.items[0].impact_score - 0.075).abs() < 1e-9);
}
#[tokio::test]
async fn diagnostics_should_preserve_request_order_when_a_small_sample_fails() {
let now = Utc::now();
let range = observation_range(now);
let store = Arc::new(FixtureObservabilityStore::new(range));
let success = diagnostic("popular", 100);
let mut failure = diagnostic("single_failure", 1);
failure.success_count = 0;
failure.failure_count = 1;
store.replace_diagnostics(vec![success, failure]);
let services = observability_services_with_calculated_billing(store).await;
let result = services
.observability()
.diagnostics(range, UsageFilter::default(), DiagnosticDimension::Model)
.await
.expect("diagnostics");
assert_eq!(result.items[0].key, "popular");
assert_eq!(result.items[1].key, "single_failure");
assert_eq!(result.items[0].error_count, 0);
assert_eq!(result.items[1].error_count, 1);
assert_eq!(result.items[1].error_rate, 1.0);
}
#[tokio::test]
@@ -821,6 +845,7 @@ impl SettingsStore for FixtureSettingsStore {
max_waiting_per_account: 0,
concurrency_wait_timeout_seconds: 30,
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
smart_scheduling: gateway_core::account::SmartSchedulingConfig::default(),
rotation_strategy: RotationStrategy::Smart,
min_codex_desktop_version: None,
min_codex_cli_version: None,
@@ -401,3 +401,46 @@ impl gateway_admin::ports::plugins::PluginDistribution for TestPluginPorts {
Err(AdminError::invalid("unused plugin fixture"))
}
}
#[async_trait::async_trait]
impl gateway_admin::ports::plugin_resources::PluginResourceStore for TestPluginPorts {
async fn ensure_group(
&self,
_: &gateway_admin::model::plugin_resources::PluginResourceOwner,
_: String,
_: gateway_admin::model::account_groups::NewAccountGroup,
_: &gateway_admin::model::MutationContext,
) -> gateway_admin::ports::store::AdminStoreResult<
gateway_admin::model::plugin_resources::ResourceMutation<
gateway_admin::model::plugin_resources::ManagedResource,
>,
> {
unreachable!("resource port is not used by this fixture")
}
async fn ensure_key(
&self,
_: &gateway_admin::model::plugin_resources::PluginResourceOwner,
_: String,
_: Vec<String>,
_: gateway_admin::model::client_keys::NewClientKey,
_: &gateway_admin::model::MutationContext,
) -> gateway_admin::ports::store::AdminStoreResult<
gateway_admin::model::plugin_resources::ResourceMutation<
gateway_admin::model::plugin_resources::ManagedResource,
>,
> {
unreachable!("resource port is not used by this fixture")
}
async fn change_members(
&self,
_: &gateway_admin::model::plugin_resources::PluginResourceOwner,
_: gateway_admin::model::plugin_resources::GroupMembersChange,
_: &gateway_admin::model::MutationContext,
) -> gateway_admin::ports::store::AdminStoreResult<
gateway_admin::model::plugin_resources::ResourceMutation<
gateway_admin::model::plugin_resources::GroupMembersChanged,
>,
> {
unreachable!("resource port is not used by this fixture")
}
}
@@ -225,6 +225,7 @@ async fn settings_should_reject_zero_refresh_margin_before_store_call() {
max_waiting_per_account: 0,
concurrency_wait_timeout_seconds: 30,
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
smart_scheduling: gateway_core::account::SmartSchedulingConfig::default(),
rotation_strategy: RotationStrategy::Smart,
min_codex_desktop_version: None,
min_codex_cli_version: None,
@@ -1294,7 +1294,6 @@ pub(crate) fn diagnostics_view(
non_completion_rate: item.non_completion_rate,
retry_count: item.retry_count,
retry_rate: item.retry_rate,
impact_score: item.impact_score,
estimated_cost: item.estimated_cost.as_ref().map(ToString::to_string),
attempt_count: item.attempt_count,
total_tokens: item.total_tokens,
@@ -758,7 +758,6 @@ pub struct DiagnosticItemView {
pub non_completion_rate: f64,
pub retry_count: u64,
pub retry_rate: f64,
pub impact_score: f64,
pub estimated_cost: Option<String>,
pub attempt_count: u64,
pub total_tokens: u64,
@@ -37,6 +37,7 @@ pub type ProviderRequestProfileUpdates =
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct RuntimeSettingsView {
pub smart_scheduling_defaults: gateway_core::account::SmartSchedulingConfig,
pub provider_request_profiles: ProviderRequestProfiles,
/// 固定兼容字段;值始终从 provider_request_profiles 派生。
pub openai_client_profile: Option<serde_json::Map<String, serde_json::Value>>,
@@ -53,6 +54,7 @@ pub struct RuntimeSettingsView {
pub max_waiting_per_account: u32,
pub concurrency_wait_timeout_seconds: u32,
pub responses_max_decompressed_body_bytes: u64,
pub smart_scheduling: gateway_core::account::SmartSchedulingConfig,
pub rotation_strategy: String,
pub min_codex_desktop_version: Option<String>,
pub min_codex_cli_version: Option<String>,
@@ -95,6 +97,7 @@ pub struct UpdateRuntimeSettingsRequest {
pub max_waiting_per_account: u32,
pub concurrency_wait_timeout_seconds: u32,
pub responses_max_decompressed_body_bytes: u64,
pub smart_scheduling: gateway_core::account::SmartSchedulingConfig,
pub rotation_strategy: String,
pub min_codex_desktop_version: Option<String>,
pub min_codex_cli_version: Option<String>,
@@ -227,6 +230,7 @@ impl UpdateRuntimeSettingsRequest {
max_waiting_per_account: self.max_waiting_per_account,
concurrency_wait_timeout_seconds: self.concurrency_wait_timeout_seconds,
responses_max_decompressed_body_bytes: self.responses_max_decompressed_body_bytes,
smart_scheduling: self.smart_scheduling,
rotation_strategy: RotationStrategy::parse(&self.rotation_strategy)
.ok_or_else(|| WireValidationError::new("rotationStrategy"))?,
min_codex_desktop_version: self.min_codex_desktop_version,
@@ -274,6 +278,8 @@ impl From<RuntimeSettings> for RuntimeSettingsView {
max_waiting_per_account: settings.max_waiting_per_account,
concurrency_wait_timeout_seconds: settings.concurrency_wait_timeout_seconds,
responses_max_decompressed_body_bytes: settings.responses_max_decompressed_body_bytes,
smart_scheduling: settings.smart_scheduling,
smart_scheduling_defaults: gateway_core::account::SmartSchedulingConfig::default(),
rotation_strategy: settings.rotation_strategy.as_str().to_owned(),
min_codex_desktop_version: settings.min_codex_desktop_version,
min_codex_cli_version: settings.min_codex_cli_version,
@@ -169,6 +169,7 @@ impl AdminTestFixture {
gateway_admin::ports::backup::BackupStorePorts::disabled(),
plugin_ports.clone(),
plugin_ports.clone(),
plugin_ports.clone(),
);
let providers: Vec<Arc<dyn ProviderAdmin>> = vec![
Arc::new(UnusedProvider::new("openai", Arc::clone(&provider_error))),
@@ -556,6 +557,7 @@ impl SettingsStore for MemorySettingsStore {
max_waiting_per_account: command.max_waiting_per_account,
concurrency_wait_timeout_seconds: command.concurrency_wait_timeout_seconds,
responses_max_decompressed_body_bytes: command.responses_max_decompressed_body_bytes,
smart_scheduling: command.smart_scheduling,
rotation_strategy: command.rotation_strategy,
min_codex_desktop_version: command.min_codex_desktop_version,
min_codex_cli_version: command.min_codex_cli_version,
@@ -1518,6 +1520,7 @@ fn test_runtime_settings() -> RuntimeSettings {
max_waiting_per_account: 0,
concurrency_wait_timeout_seconds: 30,
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
smart_scheduling: gateway_core::account::SmartSchedulingConfig::default(),
rotation_strategy: RotationStrategy::Smart,
min_codex_desktop_version: None,
min_codex_cli_version: None,
@@ -595,6 +595,7 @@ async fn diagnostics_should_keep_stable_key_and_display_name_contract() {
.await
.expect("diagnostics body");
let value: serde_json::Value = serde_json::from_slice(&body).expect("diagnostics JSON");
assert!(value["data"]["items"][0].get("impactScore").is_none());
assert_eq!(
(
&value["data"]["items"][0]["key"],
@@ -1008,3 +1008,46 @@ async fn version_plan_requires_admin_and_an_explicit_target() {
assert_eq!(response.headers()["cache-control"], "no-store");
}
}
#[async_trait::async_trait]
impl gateway_admin::ports::plugin_resources::PluginResourceStore for TestPluginPorts {
async fn ensure_group(
&self,
_: &gateway_admin::model::plugin_resources::PluginResourceOwner,
_: String,
_: gateway_admin::model::account_groups::NewAccountGroup,
_: &gateway_admin::model::MutationContext,
) -> gateway_admin::ports::store::AdminStoreResult<
gateway_admin::model::plugin_resources::ResourceMutation<
gateway_admin::model::plugin_resources::ManagedResource,
>,
> {
unreachable!("resource port is not used by this fixture")
}
async fn ensure_key(
&self,
_: &gateway_admin::model::plugin_resources::PluginResourceOwner,
_: String,
_: Vec<String>,
_: gateway_admin::model::client_keys::NewClientKey,
_: &gateway_admin::model::MutationContext,
) -> gateway_admin::ports::store::AdminStoreResult<
gateway_admin::model::plugin_resources::ResourceMutation<
gateway_admin::model::plugin_resources::ManagedResource,
>,
> {
unreachable!("resource port is not used by this fixture")
}
async fn change_members(
&self,
_: &gateway_admin::model::plugin_resources::PluginResourceOwner,
_: gateway_admin::model::plugin_resources::GroupMembersChange,
_: &gateway_admin::model::MutationContext,
) -> gateway_admin::ports::store::AdminStoreResult<
gateway_admin::model::plugin_resources::ResourceMutation<
gateway_admin::model::plugin_resources::GroupMembersChanged,
>,
> {
unreachable!("resource port is not used by this fixture")
}
}
@@ -52,6 +52,7 @@ fn update_body() -> Value {
"concurrencyWaitTimeoutSeconds": 30,
"responsesMaxDecompressedBodyBytes": 67108864,
"rotationStrategy": "round_robin",
"smartScheduling": gateway_core::account::SmartSchedulingConfig::default(),
"minCodexDesktopVersion": "26.825.6671",
"minCodexCliVersion": "0.40.0",
"usageRetentionDays": 32,
@@ -70,6 +71,76 @@ fn update_body() -> Value {
})
}
#[tokio::test]
async fn smart_settings_round_trip_and_invalid_updates_leave_the_saved_value_intact() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let mut body = update_body();
let custom = json!({"loadWeight": 2.0, "quotaWeight": 0.0, "healthWeight": 1.0, "latencyWeight": 0.5, "resetWeight": 1.2, "queueWeight": 2.3, "preferHigherWeight": true});
body["smartScheduling"] = custom.clone();
let response = app(fixture.state())
.oneshot(request(
Method::POST,
"/api/admin/settings/update",
Some(body.clone()),
))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let response = response_json(response).await;
assert_eq!(response["data"]["smartScheduling"], custom);
assert_eq!(
response["data"]["smartSchedulingDefaults"],
json!(gateway_core::account::SmartSchedulingConfig::default())
);
let mut invalid_values = vec![
json!(null),
json!({}),
json!({"loadWeight":0,"quotaWeight":0,"healthWeight":0,"latencyWeight":0,"resetWeight":0,"queueWeight":0,"preferHigherWeight":true}),
json!({"loadWeight":0.01,"quotaWeight":1,"healthWeight":1,"latencyWeight":1,"resetWeight":0,"queueWeight":0,"preferHigherWeight":false}),
];
for field in ["resetWeight", "queueWeight"] {
for value in [json!(-0.1), json!(10.1), json!(0.01), json!(null)] {
let mut invalid = custom.clone();
invalid[field] = value;
invalid_values.push(invalid);
}
let mut missing = custom.clone();
missing.as_object_mut().unwrap().remove(field);
invalid_values.push(missing);
}
for invalid in invalid_values {
body["smartScheduling"] = invalid;
let response = app(fixture.state())
.oneshot(request(
Method::POST,
"/api/admin/settings/update",
Some(body.clone()),
))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY);
}
body.as_object_mut().unwrap().remove("smartScheduling");
let response = app(fixture.state())
.oneshot(request(
Method::POST,
"/api/admin/settings/update",
Some(body),
))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY);
let response = app(fixture.state())
.oneshot(request(Method::GET, "/api/admin/settings", None))
.await
.unwrap();
assert_eq!(
response_json(response).await["data"]["smartScheduling"],
custom
);
}
#[test]
fn settings_request_should_reject_unknown_rotation_strategy() {
let mut body = update_body();
@@ -149,6 +220,7 @@ fn settings_response_should_cover_the_full_runtime_settings_contract() {
max_waiting_per_account: 0,
concurrency_wait_timeout_seconds: 30,
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
smart_scheduling: gateway_core::account::SmartSchedulingConfig::default(),
rotation_strategy: RotationStrategy::RoundRobin,
min_codex_desktop_version: Some("26.825.6671".to_owned()),
min_codex_cli_version: Some("0.40.0".to_owned()),
@@ -193,6 +265,8 @@ fn settings_response_should_cover_the_full_runtime_settings_contract() {
"concurrencyWaitTimeoutSeconds": 30,
"responsesMaxDecompressedBodyBytes": 67108864,
"rotationStrategy": "round_robin",
"smartScheduling": gateway_core::account::SmartSchedulingConfig::default(),
"smartSchedulingDefaults": gateway_core::account::SmartSchedulingConfig::default(),
"minCodexDesktopVersion": "26.825.6671",
"minCodexCliVersion": "0.40.0",
"usageRetentionDays": 32,
@@ -257,6 +331,7 @@ fn settings_request_and_response_fields_should_stay_in_lockstep() {
max_waiting_per_account: 0,
concurrency_wait_timeout_seconds: 30,
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
smart_scheduling: request.smart_scheduling,
rotation_strategy: RotationStrategy::parse(&request.rotation_strategy)
.expect("fixture rotation strategy"),
min_codex_desktop_version: request.min_codex_desktop_version,
@@ -290,6 +365,7 @@ fn settings_request_and_response_fields_should_stay_in_lockstep() {
expected_fields.insert("openaiClientProfile".to_owned());
expected_fields.insert("xaiClientProfile".to_owned());
expected_fields.insert("updatedAt".to_owned());
expected_fields.insert("smartSchedulingDefaults".to_owned());
assert_eq!(response_fields, expected_fields);
}
@@ -522,6 +522,24 @@ async fn initial_quota_recovery_delivers_client_projection_instead_of_upstream_s
assert!(error["headers"].get("retry-after").is_none());
}
#[tokio::test]
async fn initial_continuation_recovery_delivers_the_official_client_replay_signal() {
let provider = ProviderError::new(
ProviderErrorKind::ContinuationRecoveryRequired,
UpstreamSendState::NotSent,
)
.with_client_visible_upstream_error(ClientVisibleUpstreamError::new(
"Previous response was not found. Retrying the full request.",
Some("previous_response_not_found".to_owned()),
Some("invalid_request_error".to_owned()),
));
let error = initial_error(EngineError::Provider(provider), Vec::new()).await;
assert_eq!(error["type"], "error");
assert_eq!(error["status"], 400);
assert_eq!(error["error"]["code"], "previous_response_not_found");
assert_eq!(error["error"]["type"], "invalid_request_error");
}
#[tokio::test]
async fn locally_exhausted_account_pool_sends_one_official_usage_limit_error() {
let provider = ProviderError::new(
@@ -13,10 +13,12 @@ mod proxy;
pub use proxy::{InvalidOutboundProxy, OutboundProxy};
pub mod scope;
mod selection;
mod smart_scheduling;
mod store;
pub use error::CredentialError;
pub use model::*;
pub(crate) use selection::smart_score;
pub use selection::*;
pub(crate) use selection::{SMART_SCORE_TOLERANCE, smart_score};
pub use smart_scheduling::{SmartSchedulingConfig, SmartSchedulingConfigError};
pub use store::ProviderAccountStore;
@@ -1053,6 +1053,7 @@ pub struct LoadedCredential {
/// Provider 已计算好时间边界的有界 OAuth refresh 候选查询。
///
/// Store 只负责按持久事实筛选和稳定排序,不拥有提前量或恢复窗口语义。
/// 调度启停不影响凭据续期;停用账号仍按 Provider 的凭据状态与时间边界刷新。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderRefreshQuery {
provider: ProviderKind,
@@ -1117,7 +1118,6 @@ impl ProviderRefreshQuery {
#[must_use]
pub fn contains(&self, account: &ProviderAccount) -> bool {
account.provider() == &self.provider
&& account.enabled()
&& account.has_refresh_token()
&& matches!(
account.credential_state(),
@@ -6,10 +6,12 @@ use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Mutex, RwLock};
use std::time::{Duration, Instant, SystemTime};
use crate::concurrency::ConcurrencyQueuePolicy;
use crate::concurrency::{CapacityWait, ConcurrencyQueuePolicy, QueueRejection};
use crate::identity::ProviderKind;
use super::{AccountConcurrency, AccountStatus, ProviderAccount, ProviderAccountId};
use super::{
AccountConcurrency, AccountStatus, ProviderAccount, ProviderAccountId, SmartSchedulingConfig,
};
/// `runtime_settings.rotation_strategy` 的稳定值。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -47,6 +49,7 @@ impl RotationStrategy {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AccountSelectionPolicy {
strategy: RotationStrategy,
smart_scheduling: SmartSchedulingConfig,
max_concurrent_per_account: AccountConcurrency,
request_interval: Duration,
queue_policy: ConcurrencyQueuePolicy,
@@ -61,6 +64,7 @@ impl AccountSelectionPolicy {
) -> Self {
Self {
strategy,
smart_scheduling: SmartSchedulingConfig::default(),
max_concurrent_per_account: max_concurrent_per_account.into(),
request_interval,
queue_policy: ConcurrencyQueuePolicy {
@@ -86,6 +90,17 @@ impl AccountSelectionPolicy {
self.strategy
}
#[must_use]
pub const fn with_smart_scheduling(mut self, config: SmartSchedulingConfig) -> Self {
self.smart_scheduling = config;
self
}
#[must_use]
pub const fn smart_scheduling(self) -> SmartSchedulingConfig {
self.smart_scheduling
}
#[must_use]
pub const fn max_concurrent_per_account(self) -> AccountConcurrency {
self.max_concurrent_per_account
@@ -510,12 +525,7 @@ impl AccountSelector {
if self.scheduling_blocker(candidate, context).is_some() {
return None;
}
let highest_weight = candidates
.iter()
.filter(|candidate| self.scheduling_blocker(candidate, context).is_none())
.map(|candidate| candidate.account.weight())
.max()?;
let (preferred, _) = self.preferred_decision(candidates, context, highest_weight);
let (preferred, _) = self.preferred_decision(candidates, context);
let preferred = if preferred == PreferredAccountSelection::Hit
&& context.preferred_account.as_ref() != Some(account_id)
{
@@ -585,13 +595,22 @@ impl AccountSelector {
.iter()
.map(|candidate| candidate.account.weight())
.max()?;
let (preferred, preferred_candidate) =
self.preferred_decision(candidates, context, highest_weight);
let (mut preferred, preferred_candidate) = self.preferred_decision(candidates, context);
if let Some(candidate) = preferred_candidate {
return Some(AccountSelection {
candidate,
preferred: PreferredAccountSelection::Hit,
});
// 权重回切只裁决内置策略的软亲和,插件显式选号按实际选择记录结果。
let prefer_higher_weight = context.policy.strategy() == RotationStrategy::Smart
&& context.policy.smart_scheduling().prefer_higher_weight();
if (!context.preferred_account_overrides_weight || prefer_higher_weight)
&& candidate.account.weight() < highest_weight
{
preferred =
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::LowerWeight);
} else {
return Some(AccountSelection {
candidate,
preferred: PreferredAccountSelection::Hit,
});
}
}
eligible.retain(|candidate| candidate.account.weight() == highest_weight);
@@ -629,6 +648,8 @@ impl AccountSelector {
&eligible,
context.policy.max_concurrent_per_account(),
context.round_robin_cursor,
context.policy.smart_scheduling(),
context.now,
)?,
RotationStrategy::Sticky => {
eligible.sort_by_key(|candidate| {
@@ -650,7 +671,6 @@ impl AccountSelector {
&self,
candidates: &'a [AccountCandidate],
context: &AccountSelectionContext,
highest_weight: super::AccountWeight,
) -> (PreferredAccountSelection, Option<&'a AccountCandidate>) {
let Some(preferred) = context.preferred_account.as_ref() else {
return (PreferredAccountSelection::NotRequested, None);
@@ -663,14 +683,6 @@ impl AccountSelector {
};
match self.scheduling_blocker(candidate, context) {
Some(blocker) => (PreferredAccountSelection::Blocked(blocker), None),
None if !context.preferred_account_overrides_weight
&& candidate.account.weight() < highest_weight =>
{
(
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::LowerWeight),
None,
)
}
None => (PreferredAccountSelection::Hit, Some(candidate)),
}
}
@@ -707,6 +719,41 @@ impl AccountSelector {
.collect()
}
/// Provider 先投影续写范围与资格,再将等待候选交给共用评分规则选择队列。
pub async fn wait_for_capacity(
&self,
waiting: &mut CapacityWait<'_, ProviderAccountId>,
keys: &[ProviderAccountId],
candidates: &[AccountCandidate],
context: &AccountSelectionContext,
) -> Result<(), QueueRejection> {
let config = context.policy.smart_scheduling();
let queue_weight = config.weights()[5];
if context.policy.strategy() != RotationStrategy::Smart || queue_weight == 0.0 {
return waiting.wait(keys).await;
}
// 账号信号在队列锁外评分,锁内只读取实时队长并查表,避免扫描候选阻塞其他等待者。
let scores = candidates
.iter()
.map(|candidate| {
(
candidate.account.id(),
smart_score(
candidate,
context.policy.max_concurrent_per_account(),
config,
context.now,
),
)
})
.collect::<HashMap<_, _>>();
waiting
.wait_with_priority(keys, |key, count| {
scores[key] + queue_weight / (1.0 + count as f64)
})
.await
}
pub(crate) fn scheduling_blocker(
&self,
candidate: &AccountCandidate,
@@ -755,14 +802,10 @@ impl AccountSelector {
}
}
const SMART_LOAD_WEIGHT: f64 = 1.0;
const SMART_QUOTA_WEIGHT: f64 = 0.8;
const SMART_FAILURE_WEIGHT: f64 = 1.0;
const SMART_LATENCY_WEIGHT: f64 = 0.5;
// 容忍 5 个百分点的单项负载/失败率差异,避免微小信号波动独占新会话。
pub(crate) const SMART_SCORE_TOLERANCE: f64 = 0.05;
// 首输出 10 秒时延迟得分减半;固定尺度不随其他候选账号变化。
const SMART_LATENCY_HALF_SCORE_MS: f64 = 10_000.0;
// 距重置一小时时得分减半;未知和已过期时间不提供重置奖励。
const SMART_RESET_HALF_SCORE_SECONDS: f64 = 3_600.0;
fn capacity_utilization(
candidate: &AccountCandidate,
@@ -781,16 +824,23 @@ fn select_smart_candidate<'a>(
candidates: &[&'a AccountCandidate],
default_concurrency: AccountConcurrency,
cursor: u64,
config: SmartSchedulingConfig,
now: SystemTime,
) -> Option<&'a AccountCandidate> {
let mut ranked = candidates
.iter()
.map(|candidate| (*candidate, smart_score(candidate, default_concurrency)))
.map(|candidate| {
(
*candidate,
smart_score(candidate, default_concurrency, config, now),
)
})
.collect::<Vec<_>>();
let best_score = ranked
.iter()
.map(|(_, score)| *score)
.max_by(f64::total_cmp)?;
ranked.retain(|(_, score)| best_score - score <= SMART_SCORE_TOLERANCE);
ranked.retain(|(_, score)| best_score - score <= config.score_tolerance());
// 轮换顺序保持稳定,避免分数轻微交错与 cursor 同步后仍反复命中同一账号。
ranked.sort_unstable_by(|(left, _), (right, _)| left.account.id().cmp(right.account.id()));
let index = (cursor % ranked.len() as u64) as usize;
@@ -800,6 +850,8 @@ fn select_smart_candidate<'a>(
pub(crate) fn smart_score(
candidate: &AccountCandidate,
default_concurrency: AccountConcurrency,
config: SmartSchedulingConfig,
now: SystemTime,
) -> f64 {
let load = 1.0 - capacity_utilization(candidate, default_concurrency).clamp(0.0, 1.0);
let quota = candidate
@@ -822,8 +874,26 @@ pub(crate) fn smart_score(
SMART_LATENCY_HALF_SCORE_MS / (SMART_LATENCY_HALF_SCORE_MS + latency as f64)
});
SMART_LOAD_WEIGHT * load
+ SMART_QUOTA_WEIGHT * quota
+ SMART_FAILURE_WEIGHT * failure
+ SMART_LATENCY_WEIGHT * latency
let reset = candidate
.signals
.quota_reset_at
.and_then(|reset| reset.duration_since(now).ok())
.filter(|remaining| !remaining.is_zero())
.map_or(0.0, |remaining| {
SMART_RESET_HALF_SCORE_SECONDS
/ (SMART_RESET_HALF_SCORE_SECONDS + remaining.as_secs_f64())
});
let [
load_weight,
quota_weight,
health_weight,
latency_weight,
reset_weight,
_,
] = config.weights();
load_weight * load
+ quota_weight * quota
+ health_weight * failure
+ latency_weight * latency
+ reset_weight * reset
}
@@ -0,0 +1,138 @@
//! 智能调度参数;固定十分位保证快照可精确比较,浮点数只用于 wire 与评分。
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(try_from = "SmartSchedulingValues", into = "SmartSchedulingValues")]
pub struct SmartSchedulingConfig {
weights: [u8; 6],
prefer_higher_weight: bool,
}
impl Default for SmartSchedulingConfig {
fn default() -> Self {
Self {
weights: [10, 8, 10, 5, 0, 0],
prefer_higher_weight: false,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum SmartSchedulingConfigError {
#[error("{0} must be between 0 and 10 with at most one decimal place")]
InvalidWeight(&'static str),
#[error("at least one smart scheduling weight must be positive")]
EmptyWeights,
}
impl SmartSchedulingConfig {
/// 顺序为负载、剩余额度、健康、首输出延迟、额度重置和排队压力。
pub fn new(
weights: [f64; 6],
prefer_higher_weight: bool,
) -> Result<Self, SmartSchedulingConfigError> {
let mut tenths = [0; 6];
for (index, field) in [
"loadWeight",
"quotaWeight",
"healthWeight",
"latencyWeight",
"resetWeight",
"queueWeight",
]
.into_iter()
.enumerate()
{
let value = weights[index];
if !value.is_finite()
|| !(0.0..=10.0).contains(&value)
|| (value * 10.0).round() / 10.0 != value
{
return Err(SmartSchedulingConfigError::InvalidWeight(field));
}
tenths[index] = (value * 10.0).round() as u8;
}
if tenths == [0; 6] {
return Err(SmartSchedulingConfigError::EmptyWeights);
}
Ok(Self {
weights: tenths,
prefer_higher_weight,
})
}
#[must_use]
pub fn weights(self) -> [f64; 6] {
self.weights.map(|weight| f64::from(weight) / 10.0)
}
#[must_use]
pub const fn prefer_higher_weight(self) -> bool {
self.prefer_higher_weight
}
pub(crate) fn score_tolerance(self) -> f64 {
// 与系数同步缩放,避免仅放大相同比例的系数就改变近似最优候选集合。
// 排队系数只参与入队选择,不能放大立即选号的候选容差。
0.05 * f64::from(
self.weights[..5]
.iter()
.copied()
.map(u16::from)
.sum::<u16>(),
) / 33.0
}
}
#[derive(Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct SmartSchedulingValues {
load_weight: f64,
quota_weight: f64,
health_weight: f64,
latency_weight: f64,
reset_weight: f64,
queue_weight: f64,
prefer_higher_weight: bool,
}
impl TryFrom<SmartSchedulingValues> for SmartSchedulingConfig {
type Error = SmartSchedulingConfigError;
fn try_from(value: SmartSchedulingValues) -> Result<Self, Self::Error> {
Self::new(
[
value.load_weight,
value.quota_weight,
value.health_weight,
value.latency_weight,
value.reset_weight,
value.queue_weight,
],
value.prefer_higher_weight,
)
}
}
impl From<SmartSchedulingConfig> for SmartSchedulingValues {
fn from(value: SmartSchedulingConfig) -> Self {
let [
load_weight,
quota_weight,
health_weight,
latency_weight,
reset_weight,
queue_weight,
] = value.weights();
Self {
load_weight,
quota_weight,
health_weight,
latency_weight,
reset_weight,
queue_weight,
prefer_higher_weight: value.prefer_higher_weight,
}
}
}
+30 -6
View File
@@ -80,6 +80,18 @@ impl<K: Clone + Eq + Hash> ConcurrencyWaitQueue<K> {
keys: &[K],
max_waiting: u32,
deadline: Instant,
) -> Result<WaitTicket<K>, QueueRejection> {
self.enqueue_with_priority(keys, max_waiting, deadline, |_, waiting| -(waiting as f64))
}
/// 在同一锁内读取队长并入队,分数越大越优先,同分沿用候选顺序。
/// 回调只做内存评分,不得阻塞或再次访问本队列。
pub fn enqueue_with_priority(
&self,
keys: &[K],
max_waiting: u32,
deadline: Instant,
mut priority: impl FnMut(&K, usize) -> f64,
) -> Result<WaitTicket<K>, QueueRejection> {
if Instant::now() >= deadline {
return Err(QueueRejection::Timeout);
@@ -95,9 +107,9 @@ impl<K: Clone + Eq + Hash> ConcurrencyWaitQueue<K> {
.iter()
.filter_map(|key| {
let len = state.queues.get(key).map_or(0, VecDeque::len);
(len < max_waiting as usize).then_some((key, len))
(len < max_waiting as usize).then(|| (key, priority(key, len)))
})
.min_by_key(|(_, len)| *len)
.min_by(|(_, left), (_, right)| right.total_cmp(left))
.map(|(key, _)| key.clone())
.ok_or(QueueRejection::Full)?;
let waker = Arc::new(AtomicWaker::new());
@@ -237,6 +249,16 @@ impl<'a, K: Clone + Eq + Hash> CapacityWait<'a, K> {
}
pub async fn wait(&mut self, keys: &[K]) -> Result<(), QueueRejection> {
self.wait_with_priority(keys, |_, waiting| -(waiting as f64))
.await
}
/// 只在首次入队或原账号不再合格时重新选队列,已取得的位置不因评分变化而移动。
pub async fn wait_with_priority(
&mut self,
keys: &[K],
priority: impl FnMut(&K, usize) -> f64,
) -> Result<(), QueueRejection> {
self.started_at.get_or_insert_with(Instant::now);
let deadline = self
.budget
@@ -249,10 +271,12 @@ impl<'a, K: Clone + Eq + Hash> CapacityWait<'a, K> {
self.ticket = None;
}
if self.ticket.is_none() {
self.ticket = Some(
self.queue
.enqueue(keys, self.policy.max_waiting, deadline)?,
);
self.ticket = Some(self.queue.enqueue_with_priority(
keys,
self.policy.max_waiting,
deadline,
priority,
)?);
}
if let Some(ticket) = &self.ticket {
ticket.retry().await?;
@@ -5,7 +5,7 @@ use serde_json::json;
use super::TraceContext;
use crate::account::{
AccountCandidate, AccountSelection, AccountSelectionContext, AccountSelector, RotationStrategy,
SMART_SCORE_TOLERANCE, smart_score,
smart_score,
};
// 单个 trace 事件最多 4 KiB;优先保留实际选中的账号。
@@ -50,7 +50,7 @@ impl TraceContext {
"failureRateBasisPoints": signals.failure_rate_basis_points,
"firstOutputLatencyMs": signals.first_output_latency_ms,
"smartScore": smart.then(|| smart_score(
candidate, context.policy.max_concurrent_per_account()
candidate, context.policy.max_concurrent_per_account(), context.policy.smart_scheduling(), context.now
)),
})
})
@@ -63,7 +63,7 @@ impl TraceContext {
"preferredAccountId": context.preferred_account.as_ref().map(|id| id.as_str()),
"preferredResult": selection.as_ref().map(|s| format!("{:?}", s.preferred())),
"roundRobinCursor": context.round_robin_cursor,
"smartScoreTolerance": smart.then_some(SMART_SCORE_TOLERANCE),
"smartScoreTolerance": smart.then(|| context.policy.smart_scheduling().score_tolerance()),
"candidateCount": candidates.len(),
"omittedCandidates": candidates.len().saturating_sub(observations.len()),
"candidates": observations,
@@ -2118,9 +2118,10 @@ impl ExecutionService for DefaultExecutionService {
.await?
else {
for profile in client.snapshot.public_model_profiles_for_provider(kind) {
if !scope.allows_provider_model(
if !client.snapshot.catalog_model_allowed_for_scope(
kind,
&client.snapshot.mapped_model(profile.model().as_str()),
profile.model(),
scope,
) {
continue;
}
@@ -1,9 +1,10 @@
//! 快照编译与客户端读取需要的 Provider 目录合同,不暴露执行或账号选择能力。
use std::collections::BTreeMap;
use std::collections::{BTreeMap, BTreeSet};
use futures::future::BoxFuture;
use crate::account::ProviderAccountId;
use crate::identity::ProviderKind;
use crate::operation::RawJsonPayload;
@@ -63,6 +64,7 @@ pub struct ProviderModelCapabilities {
upstream_model: UpstreamModelId,
capabilities: ModelCapabilities,
presentation: Option<ModelPresentation>,
catalog_accounts: Option<BTreeSet<ProviderAccountId>>,
}
impl ProviderModelCapabilities {
@@ -72,6 +74,7 @@ impl ProviderModelCapabilities {
upstream_model,
capabilities,
presentation: None,
catalog_accounts: None,
}
}
@@ -81,6 +84,19 @@ impl ProviderModelCapabilities {
self
}
/// 声明发现此模型的账号;仅用于目录可见性,不作为推理能力白名单。
#[must_use]
pub fn with_catalog_accounts(mut self, accounts: BTreeSet<ProviderAccountId>) -> Self {
self.catalog_accounts = Some(accounts);
self
}
/// 未提供账号来源的 Provider 沿用 Provider 级目录;空集合表示无可展示来源。
#[must_use]
pub const fn catalog_accounts(&self) -> Option<&BTreeSet<ProviderAccountId>> {
self.catalog_accounts.as_ref()
}
#[must_use]
pub const fn upstream_model(&self) -> &UpstreamModelId {
&self.upstream_model
@@ -27,6 +27,8 @@ use super::{
const MAXIMUM_CATALOG_STABILITY_ATTEMPTS: usize = 4;
type ModelCatalogAccounts = BTreeMap<ProviderKind, BTreeMap<String, BTreeSet<ProviderAccountId>>>;
/// Store 在一个一致性读取中提供的调度设置事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotSettingsFacts {
@@ -40,6 +42,7 @@ pub struct SnapshotSettingsFacts {
concurrency_wait_timeout_seconds: u32,
responses_max_decompressed_body_bytes: u64,
request_interval_ms: u64,
smart_scheduling: crate::account::SmartSchedulingConfig,
rotation_strategy: String,
model_mappings: BTreeMap<String, String>,
min_codex_desktop_version: Option<String>,
@@ -47,6 +50,15 @@ pub struct SnapshotSettingsFacts {
}
impl SnapshotSettingsFacts {
#[must_use]
pub const fn with_smart_scheduling(
mut self,
config: crate::account::SmartSchedulingConfig,
) -> Self {
self.smart_scheduling = config;
self
}
#[must_use]
pub fn with_pricing(mut self, pricing: crate::metering::PricingOverrides) -> Self {
self.pricing = Arc::new(pricing);
@@ -112,6 +124,7 @@ impl SnapshotSettingsFacts {
concurrency_wait_timeout_seconds: 30,
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
request_interval_ms,
smart_scheduling: crate::account::SmartSchedulingConfig::default(),
rotation_strategy: rotation_strategy.into(),
model_mappings,
min_codex_desktop_version,
@@ -445,9 +458,13 @@ async fn compile_runtime_snapshot(
) -> Result<RuntimeSnapshot, RuntimeSnapshotCompileError> {
// 只有成功取得的完整目录能证明模型缺项;发现型目录和查询失败均交由上游验证。
let mut provider_models = Vec::new();
let mut catalog_accounts = BTreeMap::new();
let mut exhaustive_provider_catalogs = BTreeSet::new();
for provider in &provider_kinds {
if let Some(previous) = previous {
if let Some(accounts) = previous.model_catalog_accounts.get(provider) {
catalog_accounts.insert(provider.clone(), accounts.clone());
}
if previous.exhaustive_provider_catalogs.contains(provider) {
exhaustive_provider_catalogs.insert(provider.clone());
}
@@ -474,6 +491,12 @@ async fn compile_runtime_snapshot(
exhaustive_provider_catalogs.insert(provider.clone());
}
provider_models.extend(models.into_iter().map(|model| {
if let Some(accounts) = model.catalog_accounts() {
catalog_accounts
.entry(provider.clone())
.or_default()
.insert(model.upstream_model().as_str().to_owned(), accounts.clone());
}
let compiled = ProviderModel::new(
provider.clone(),
model.upstream_model().clone(),
@@ -580,6 +603,7 @@ async fn compile_runtime_snapshot(
AccountConcurrency::new(facts.settings.max_concurrent_per_account),
Duration::from_millis(facts.settings.request_interval_ms),
)
.with_smart_scheduling(facts.settings.smart_scheduling)
.with_queue(ConcurrencyQueuePolicy {
max_waiting: facts.settings.max_waiting_per_account,
timeout: queue_timeout,
@@ -665,6 +689,7 @@ async fn compile_runtime_snapshot(
.with_model_mappings(model_mappings)
.with_account_directory(account_directory)
.with_exhaustive_provider_catalogs(exhaustive_provider_catalogs)
.with_model_catalog_accounts(catalog_accounts)
.with_min_codex_client_versions(min_client_versions)
})
}
@@ -686,6 +711,7 @@ pub struct RuntimeSnapshot {
model_mappings: Arc<BTreeMap<String, String>>,
provider_catalog_generations: Arc<BTreeMap<ProviderKind, ProviderCatalogGeneration>>,
exhaustive_provider_catalogs: Arc<BTreeSet<ProviderKind>>,
model_catalog_accounts: Arc<ModelCatalogAccounts>,
account_directory: Arc<RuntimeAccountDirectory>,
client_policies: Arc<BTreeMap<ClientApiKeyId, ClientPolicy>>,
min_codex_client_versions: CodexClientMinVersions,
@@ -868,6 +894,7 @@ impl RuntimeSnapshot {
model_mappings: Arc::new(BTreeMap::new()),
provider_catalog_generations: Arc::new(BTreeMap::new()),
exhaustive_provider_catalogs: Arc::new(exhaustive_provider_catalogs),
model_catalog_accounts: Arc::default(),
account_directory: Arc::new(RuntimeAccountDirectory::default()),
client_policies: Arc::new(client_policy_map),
min_codex_client_versions: CodexClientMinVersions::default(),
@@ -898,6 +925,11 @@ impl RuntimeSnapshot {
self
}
fn with_model_catalog_accounts(mut self, accounts: ModelCatalogAccounts) -> Self {
self.model_catalog_accounts = Arc::new(accounts);
self
}
#[must_use]
pub fn all_account_scope(&self) -> Arc<FrozenAccountScope> {
Arc::new(FrozenAccountScope::new(
@@ -1015,9 +1047,7 @@ impl RuntimeSnapshot {
.flat_map(|provider| {
self.public_models_for_provider(provider)
.into_iter()
.filter(|model| {
scope.allows_provider_model(provider, &self.mapped_model(model.as_str()))
})
.filter(|model| self.catalog_model_allowed_for_scope(provider, model, scope))
})
.collect::<BTreeSet<_>>()
.into_iter()
@@ -1033,9 +1063,7 @@ impl RuntimeSnapshot {
let mut profiles = BTreeMap::new();
for provider in scope.provider_kinds() {
for profile in self.public_model_profiles_for_provider(provider) {
if !scope
.allows_provider_model(provider, &self.mapped_model(profile.model().as_str()))
{
if !self.catalog_model_allowed_for_scope(provider, profile.model(), scope) {
continue;
}
profiles
@@ -1082,10 +1110,30 @@ impl RuntimeSnapshot {
) -> bool {
scope.provider_kinds().iter().any(|provider| {
self.contains_public_model_for_provider(public_model, provider)
&& scope.allows_provider_model(provider, &self.mapped_model(public_model.as_str()))
&& self.catalog_model_allowed_for_scope(provider, public_model, scope)
})
}
pub(crate) fn catalog_model_allowed_for_scope(
&self,
provider: &ProviderKind,
public_model: &PublicModelId,
scope: &FrozenAccountScope,
) -> bool {
let upstream_model = self.mapped_model(public_model.as_str());
match self
.model_catalog_accounts
.get(provider)
.and_then(|models| models.get(&upstream_model))
{
Some(accounts) => accounts.iter().any(|account| {
scope.account_provider(account) == Some(provider)
&& scope.allows_model(account, &upstream_model)
}),
None => scope.allows_provider_model(provider, &upstream_model),
}
}
#[must_use]
pub fn mapped_model(&self, requested: &str) -> String {
if let Some(alias) = self.model_alias(requested) {
@@ -69,6 +69,7 @@ pub struct RuntimeSnapshotUnavailable;
#[derive(Clone, Default)]
pub struct RuntimeSnapshotHandle {
current: Arc<RwLock<Option<Arc<RuntimeSnapshot>>>>,
publications: Arc<SyncMutex<Vec<futures::channel::mpsc::Sender<()>>>>,
}
impl RuntimeSnapshotHandle {
@@ -76,15 +77,38 @@ impl RuntimeSnapshotHandle {
pub fn new(initial: RuntimeSnapshot) -> Self {
Self {
current: Arc::new(RwLock::new(Some(Arc::new(initial)))),
publications: Default::default(),
}
}
pub fn publish(&self, snapshot: RuntimeSnapshot) {
*write_unpoisoned(&self.current) = Some(Arc::new(snapshot));
self.notify_publication();
}
pub fn suspend(&self) {
*write_unpoisoned(&self.current) = None;
self.notify_publication();
}
/// 合并发布与暂停通知;订阅者收到后读取当前快照,不依赖逐条投递。
pub fn subscribe_publications(&self) -> futures::channel::mpsc::Receiver<()> {
let (sender, receiver) = futures::channel::mpsc::channel(0);
self.publications
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.push(sender);
receiver
}
fn notify_publication(&self) {
self.publications
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.retain_mut(|sender| match sender.try_send(()) {
Ok(()) => true,
Err(error) => error.is_full(),
});
}
#[must_use]
@@ -3,6 +3,7 @@ mod model_access;
mod proxy;
mod selection;
mod selection_diagnostics;
mod smart_scheduling;
use std::collections::{BTreeMap, BTreeSet};
use std::num::NonZeroU32;
@@ -0,0 +1,349 @@
use gateway_core::account::{
AccountSelector, PreferredAccountSelection, RotationStrategy, SmartSchedulingConfig,
};
use serde_json::json;
use std::time::{Duration, Instant, SystemTime};
use futures::FutureExt;
use gateway_core::concurrency::{
CapacityWait, ConcurrencyQueuePolicy, ConcurrencyWaitBudget, ConcurrencyWaitQueue,
};
use super::{candidate_with_concurrency, context, weighted_candidate};
#[test]
fn smart_config_rejects_invalid_weights_and_incomplete_wire_values() {
for value in [-0.1, 10.1, 0.01, f64::NAN, f64::INFINITY] {
for dimension in 0..6 {
let mut weights = [1.0; 6];
weights[dimension] = value;
assert!(SmartSchedulingConfig::new(weights, false).is_err());
}
}
assert!(SmartSchedulingConfig::new([0.0; 6], true).is_err());
for value in [json!({}), json!(null), json!({"loadWeight": 1.0})] {
assert!(serde_json::from_value::<SmartSchedulingConfig>(value).is_err());
}
let config = SmartSchedulingConfig::new([0.0, 0.1, 9.9, 10.0, 1.2, 2.3], true).unwrap();
let encoded = serde_json::to_value(config).unwrap();
assert_eq!(
serde_json::from_value::<SmartSchedulingConfig>(encoded).unwrap(),
config
);
}
#[test]
fn reset_weight_rewards_only_future_resets_and_preserves_capacity_checks() {
let mut context = context(RotationStrategy::Smart);
context.policy = context.policy.with_smart_scheduling(
SmartSchedulingConfig::new([0.0, 0.0, 0.0, 0.0, 1.0, 0.0], false).unwrap(),
);
let mut candidates = [
candidate_with_concurrency("acct_unknown", 0, 2),
candidate_with_concurrency("acct_past", 0, 2),
candidate_with_concurrency("acct_now", 0, 2),
candidate_with_concurrency("acct_near", 0, 2),
candidate_with_concurrency("acct_far", 0, 2),
];
candidates[1].signals.quota_reset_at = Some(context.now - Duration::from_secs(1));
candidates[2].signals.quota_reset_at = Some(context.now);
candidates[3].signals.quota_reset_at = Some(context.now + Duration::from_secs(3600));
candidates[4].signals.quota_reset_at = Some(context.now + Duration::from_secs(86_400));
assert_eq!(
AccountSelector
.select(&candidates, &context)
.unwrap()
.candidate()
.account
.id()
.as_str(),
"acct_near"
);
let trace = gateway_core::diagnostics::TraceContext::new("req_reset_scores");
trace.account_selection(&candidates, &context, None);
let snapshot = trace.snapshot().unwrap();
let scores = snapshot["events"][0]["data"]["candidates"]
.as_array()
.unwrap();
for (candidate, expected) in scores.iter().zip([0.0, 0.0, 0.0, 0.5, 0.04]) {
assert_eq!(candidate["smartScore"].as_f64().unwrap(), expected);
}
candidates[3].signals.in_flight = 2;
assert_eq!(
AccountSelector
.select(&candidates, &context)
.unwrap()
.candidate()
.account
.id()
.as_str(),
"acct_far"
);
candidates[3].signals.in_flight = 0;
context.now += Duration::from_secs(3600);
assert_eq!(
AccountSelector
.select(&candidates, &context)
.unwrap()
.candidate()
.account
.id()
.as_str(),
"acct_far"
);
}
#[test]
fn queue_weight_balances_live_pressure_and_health_only_when_enabled_for_smart() {
for (strategy, queue_weight, expects_shortest) in [
(RotationStrategy::Smart, 0.0, true),
(RotationStrategy::Smart, 0.2, false),
(RotationStrategy::Smart, 2.0, true),
(RotationStrategy::RoundRobin, 0.2, true),
(RotationStrategy::QuotaResetPriority, 0.2, true),
(RotationStrategy::Sticky, 0.2, true),
] {
let mut candidates = [
candidate_with_concurrency("acct_healthy", 2, 2),
candidate_with_concurrency("acct_shortest", 2, 2),
];
candidates[1].signals.failure_rate_basis_points = Some(6000);
let mut context = context(strategy);
context.policy = context.policy.with_smart_scheduling(
SmartSchedulingConfig::new([0.0, 0.0, 1.0, 0.0, 0.0, queue_weight], false).unwrap(),
);
let queue = ConcurrencyWaitQueue::default();
let head = queue
.enqueue(
&[candidates[0].account.id().clone()],
2,
Instant::now() + Duration::from_secs(2),
)
.unwrap();
let budget = ConcurrencyWaitBudget::default();
let policy = ConcurrencyQueuePolicy {
max_waiting: 2,
timeout: Duration::from_secs(2),
};
let mut waiting = CapacityWait::new(
&queue,
policy,
SystemTime::now() + Duration::from_secs(2),
&budget,
);
let keys = AccountSelector.wait_candidates(&candidates, &context);
assert!(
AccountSelector
.wait_for_capacity(&mut waiting, &keys, &candidates, &context)
.now_or_never()
.is_none()
);
assert_eq!(
queue.has_waiters(candidates[1].account.id()),
expects_shortest,
"{strategy:?}, {queue_weight}"
);
// 新请求不能绕过已经选定的队列,取消请求同步归还位置。
let newcomer = CapacityWait::new(
&queue,
policy,
SystemTime::now() + Duration::from_secs(2),
&budget,
);
assert!(!newcomer.can_try(candidates[usize::from(expects_shortest)].account.id()));
drop(waiting);
assert!(!queue.has_waiters(candidates[1].account.id()));
drop(head);
assert!(!queue.has_waiters(candidates[0].account.id()));
}
}
#[test]
fn queue_weight_does_not_widen_immediate_selection_tolerance() {
let mut candidates = [
candidate_with_concurrency("acct_a", 0, 2),
candidate_with_concurrency("acct_b", 0, 2),
];
candidates[0].signals.quota_remaining_rank = Some(10000);
candidates[1].signals.quota_remaining_rank = Some(9000);
let mut context = context(RotationStrategy::Smart);
context.round_robin_cursor = 1;
for queue_weight in [0.0, 10.0] {
context.policy = context.policy.with_smart_scheduling(
SmartSchedulingConfig::new([0.0, 1.0, 0.0, 0.0, 0.0, queue_weight], false).unwrap(),
);
assert_eq!(
AccountSelector
.select(&candidates, &context)
.unwrap()
.candidate()
.account
.id()
.as_str(),
"acct_a"
);
}
context.policy = context.policy.with_smart_scheduling(
SmartSchedulingConfig::new([0.0, 0.0, 0.0, 0.0, 0.0, 1.0], false).unwrap(),
);
// 仅开启排队系数时,立即可用的同权重账号按既有游标轮换。
assert_eq!(
AccountSelector
.select(&candidates, &context)
.unwrap()
.candidate()
.account
.id()
.as_str(),
"acct_b"
);
}
#[test]
fn each_smart_weight_controls_its_signal_without_relaxing_capacity() {
let mut candidates = [
candidate_with_concurrency("acct_idle", 0, 10),
candidate_with_concurrency("acct_healthy", 5, 10),
];
candidates[0].signals.quota_remaining_rank = Some(1_000);
candidates[1].signals.quota_remaining_rank = Some(9_000);
candidates[0].signals.failure_rate_basis_points = Some(5_000);
candidates[1].signals.failure_rate_basis_points = Some(0);
candidates[0].signals.first_output_latency_ms = Some(30_000);
candidates[1].signals.first_output_latency_ms = Some(1_000);
for dimension in 0..4 {
let mut weights = [0.0; 6];
weights[dimension] = 1.0;
let mut context = context(RotationStrategy::Smart);
context.policy = context
.policy
.with_smart_scheduling(SmartSchedulingConfig::new(weights, false).unwrap());
let selected = AccountSelector.select(&candidates, &context).unwrap();
assert_eq!(
selected.candidate().account.id().as_str(),
if dimension == 0 {
"acct_idle"
} else {
"acct_healthy"
}
);
let mut saturated = candidates.clone();
saturated[1].signals.in_flight = 10;
assert_eq!(
AccountSelector
.select(&saturated, &context)
.unwrap()
.candidate()
.account
.id()
.as_str(),
"acct_idle"
);
}
}
#[test]
fn proportional_coefficients_preserve_near_best_rotation_and_diagnostics() {
let mut candidates = [
candidate_with_concurrency("acct_a", 0, 100),
candidate_with_concurrency("acct_b", 0, 100),
candidate_with_concurrency("acct_c", 0, 100),
];
// 默认额度权重 0.8 下,9375 正好位于容差边界,9374 刚好超出。
for (candidate, quota) in candidates.iter_mut().zip([10_000, 9_375, 9_374]) {
candidate.signals.quota_remaining_rank = Some(quota);
}
for scale in [1.0, 2.0, 5.0] {
let mut context = context(RotationStrategy::Smart);
context.policy = context.policy.with_smart_scheduling(
SmartSchedulingConfig::new([scale, 0.8 * scale, scale, 0.5 * scale, 0.0, 0.0], false)
.unwrap(),
);
let mut ids = Vec::new();
for cursor in 0..4 {
candidates.reverse();
context.round_robin_cursor = cursor;
ids.push(
AccountSelector
.select(&candidates, &context)
.unwrap()
.candidate()
.account
.id()
.as_str()
.to_owned(),
);
}
assert_eq!(ids, ["acct_a", "acct_b", "acct_a", "acct_b"]);
let trace = gateway_core::diagnostics::TraceContext::new("req_scaled_weights");
trace.account_selection(&candidates, &context, None);
let snapshot = trace.snapshot().unwrap();
let tolerance = snapshot["events"][0]["data"]["smartScoreTolerance"]
.as_f64()
.unwrap();
assert!((tolerance - 0.05 * scale).abs() < 1e-12);
}
}
#[test]
fn switchback_is_opt_in_and_only_affects_smart_soft_affinity() {
let mut candidates = [
weighted_candidate("acct_affinity", 10, 0),
weighted_candidate("acct_primary", 100, 0),
];
for strategy in [
RotationStrategy::Smart,
RotationStrategy::Sticky,
RotationStrategy::RoundRobin,
RotationStrategy::QuotaResetPriority,
] {
let mut context = context(strategy);
context.preferred_account = Some(candidates[0].account.id().clone());
context.preferred_account_overrides_weight = true;
for enabled in [false, true] {
context.policy = context.policy.with_smart_scheduling(
SmartSchedulingConfig::new([1.0, 0.8, 1.0, 0.5, 0.0, 0.0], enabled).unwrap(),
);
let selected = AccountSelector.select(&candidates, &context).unwrap();
let should_switch = enabled && strategy == RotationStrategy::Smart;
assert_eq!(
selected.candidate().account.id(),
candidates[usize::from(should_switch)].account.id()
);
}
}
let mut context = context(RotationStrategy::Smart);
context.policy = context.policy.with_smart_scheduling(
SmartSchedulingConfig::new([1.0, 0.8, 1.0, 0.5, 0.0, 0.0], true).unwrap(),
);
context.preferred_account_overrides_weight = true;
context.preferred_account = Some(candidates[0].account.id().clone());
candidates[1].signals.in_flight = 3;
assert_eq!(
AccountSelector
.select(&candidates, &context)
.unwrap()
.preferred(),
PreferredAccountSelection::Hit
);
candidates[1].signals.in_flight = 0;
context
.excluded_accounts
.insert(candidates[1].account.id().clone());
assert_eq!(
AccountSelector
.select(&candidates, &context)
.unwrap()
.preferred(),
PreferredAccountSelection::Hit
);
context.excluded_accounts.clear();
candidates[0] = weighted_candidate("acct_affinity", 100, 2);
assert_eq!(
AccountSelector
.select(&candidates, &context)
.unwrap()
.preferred(),
PreferredAccountSelection::Hit
);
}
@@ -159,3 +159,75 @@ fn creating_a_waiter_does_not_start_the_budget_before_queueing() {
waiting.wait(&["account"]).await.unwrap();
});
}
#[test]
fn custom_priority_reads_live_counts_and_preserves_limits_and_fifo() {
block_on(async {
let queue = ConcurrencyWaitQueue::default();
let deadline = Instant::now() + Duration::from_secs(2);
let head = queue.enqueue(&["a"], 2, deadline).unwrap();
let mut observed = Vec::new();
let tail = queue
.enqueue_with_priority(&["a", "b"], 2, deadline, |key, count| {
observed.push((*key, count));
if *key == "a" { 10.0 } else { 0.0 }
})
.unwrap();
assert_eq!(observed, [("a", 1), ("b", 0)]);
assert_eq!(*tail.key(), "a");
assert!(tail.turn().now_or_never().is_none());
// 满队列不参与评分,不能因偏好而突破上限。
let fallback = queue
.enqueue_with_priority(&["a", "b"], 2, deadline, |key, count| {
assert_eq!((*key, count), ("b", 0));
0.0
})
.unwrap();
assert_eq!(*fallback.key(), "b");
drop(head);
tail.turn().await.unwrap();
drop((tail, fallback));
let tie = queue
.enqueue_with_priority(&["b", "a"], 2, deadline, |_, _| 1.0)
.unwrap();
assert_eq!(*tie.key(), "b");
});
}
#[test]
fn changed_scores_keep_existing_position_until_account_becomes_ineligible() {
block_on(async {
let queue = ConcurrencyWaitQueue::default();
let budget = ConcurrencyWaitBudget::default();
let policy = ConcurrencyQueuePolicy {
max_waiting: 1,
timeout: Duration::from_secs(2),
};
let mut waiting = CapacityWait::new(
&queue,
policy,
SystemTime::now() + Duration::from_secs(2),
&budget,
);
waiting
.wait_with_priority(&["a", "b"], |key, _| if *key == "a" { 1.0 } else { 0.0 })
.await
.unwrap();
waiting
.wait_with_priority(&["a", "b"], |_, _| {
panic!("existing ticket must not be rescored")
})
.await
.unwrap();
assert!(queue.has_waiters(&"a"));
assert!(!queue.has_waiters(&"b"));
waiting
.wait_with_priority(&["b"], |_, _| 1.0)
.await
.unwrap();
assert!(!queue.has_waiters(&"a"));
assert!(queue.has_waiters(&"b"));
drop(waiting);
assert!(!queue.has_waiters(&"b"));
});
}
@@ -169,6 +169,52 @@ fn delegated_scheduler_keeps_builtin_affinity() {
});
}
#[test]
fn smart_switchback_does_not_override_explicit_plugin_choices_or_their_telemetry() {
block_on(async {
let mut candidates = [candidate("acct_a"), candidate("acct_b")];
candidates[0].account = candidates[0]
.account
.clone()
.with_scheduling(None, AccountWeight::new(10).unwrap());
let mut selection_context = context();
selection_context.policy = AccountSelectionPolicy::new(
RotationStrategy::Smart,
NonZeroU32::new(2).unwrap(),
Duration::ZERO,
)
.with_smart_scheduling(
gateway_core::account::SmartSchedulingConfig::new([1.0, 0.8, 1.0, 0.5, 0.0, 0.0], true)
.unwrap(),
);
for chosen in [Some("acct_a"), Some("acct_b"), None] {
let selected = policy(chosen)
.select_account(
NonZeroU32::new(1).unwrap(),
&ProviderKind::new("openai").unwrap(),
None,
&candidates,
&selection_context,
)
.await
.unwrap()
.unwrap();
assert_eq!(
selected.candidate().account.id().as_str(),
chosen.unwrap_or("acct_b")
);
let expected = match chosen {
Some("acct_a") => PreferredAccountSelection::Hit,
Some(_) => PreferredAccountSelection::OverriddenByPolicy,
None => PreferredAccountSelection::Blocked(
gateway_core::account::AccountSchedulingBlocker::LowerWeight,
),
};
assert_eq!(selected.preferred(), expected);
}
});
}
#[test]
fn scheduler_cannot_select_outside_the_hard_bound_candidate_pool() {
block_on(async {
@@ -362,6 +362,49 @@ fn compiler_accepts_unlimited_default_account_concurrency() {
);
}
#[test]
fn compiled_plans_keep_their_smart_config_after_a_new_snapshot_is_built() {
use gateway_core::account::SmartSchedulingConfig;
let configs = [
SmartSchedulingConfig::default(),
SmartSchedulingConfig::new([0.0, 2.0, 1.0, 0.5, 1.2, 2.3], true).unwrap(),
];
let mut plans = Vec::new();
for (index, config) in configs.into_iter().enumerate() {
let facts = SnapshotFacts::new(
revision(index as u64 + 1),
revision(index as u64 + 1),
SnapshotSettingsFacts::new(3, 0, "smart", BTreeMap::new(), None, None)
.with_smart_scheduling(config),
vec![],
vec![],
vec![SnapshotProviderAccountFacts::new(
gateway_core::account::ProviderAccountId::new("acct_config").unwrap(),
"alpha",
)],
vec![],
);
let compiler = RuntimeSnapshotCompiler::new(
Arc::new(TestSnapshotStore::new(Ok(facts))),
Arc::new(TestCatalog::Unavailable),
);
let snapshot = block_on(compiler.compile()).unwrap();
plans.push(
snapshot
.plan(
&PublicModelId::new("any-model").unwrap(),
&super::operation(),
snapshot.all_account_scope(),
&Default::default(),
)
.unwrap(),
);
}
for (plan, config) in plans.iter().zip(configs) {
assert_eq!(plan.account_selection_policy().smart_scheduling(), config);
}
}
#[test]
fn withdrawn_provider_keeps_accounts_without_becoming_a_route_or_expanding_group_scope() {
use gateway_core::{
@@ -168,19 +168,30 @@ fn committed_account_change_should_publish_without_waiting_for_provider_catalog(
block_on(async {
let store = Arc::new(TestSnapshotStore::new(Ok(scoped_facts(9, true))));
let catalog = Arc::new(TestCatalog::default());
catalog
.models
.lock()
.expect("models lock")
.push(ProviderModelCapabilities::new(
catalog.models.lock().expect("models lock").push(
ProviderModelCapabilities::new(
UpstreamModelId::new("listed-model").expect("model"),
ModelCapabilities::new(Default::default(), None),
));
)
.with_catalog_accounts(std::collections::BTreeSet::from([
ProviderAccountId::new("acct_audit_removed").expect("source account"),
])),
);
let compiler = Arc::new(catalog_compiler(store.clone(), catalog.clone()));
let initial = compiler.compile().await.expect("initial snapshot");
let provider = ProviderKind::new("audit").expect("provider");
let initial_models = initial.public_models_for_provider(&provider);
let handle = RuntimeSnapshotHandle::new(initial);
let initial_scope = initial
.client_policies()
.next()
.unwrap()
.account_scope()
.clone();
assert_eq!(
initial.public_models_for_scope(&initial_scope),
initial_models
);
let handle = RuntimeSnapshotHandle::new(initial.clone());
let publisher = RuntimeSnapshotPublisher::new(
compiler,
handle.clone(),
@@ -200,6 +211,16 @@ fn committed_account_change_should_publish_without_waiting_for_provider_catalog(
initial_models
);
assert_eq!(catalog.queries.load(Ordering::SeqCst), 1);
let committed_scope = committed.client_policies().next().unwrap().account_scope();
assert!(
committed
.public_models_for_scope(committed_scope)
.is_empty()
);
assert_eq!(
initial.public_models_for_scope(&initial_scope),
initial_models
);
assert_ne!(
committed.provider_catalog_generations(),
&catalog.catalog_generations(),
@@ -26,6 +26,7 @@ use super::{
pub(super) struct CachedRelease {
result: Result<ResolvedRelease, AdminError>,
until: Instant,
fetched_at: Instant,
}
#[derive(Clone)]
@@ -57,7 +58,9 @@ impl HttpPluginDistribution {
mut query: GithubReleaseQuery,
credentials: &[SourceCredential],
egress: Option<&PluginDistributionEgress>,
refresh: bool,
) -> Result<ResolvedRelease, AdminError> {
let requested_at = Instant::now();
validate_repository(&query.repository)?;
query.repository.make_ascii_lowercase();
let url = self.release_url(&query)?;
@@ -91,7 +94,11 @@ impl HttpPluginDistribution {
tokio::time::timeout(Duration::from_secs(60), async {
// 同一来源共用锁和结果;取消请求释放锁,下一位查询者可重新发起。
let mut cached = slot.lock().await;
if let Some(entry) = cached.as_ref().filter(|entry| entry.until > Instant::now()) {
// 显式查询刷新成功结果;等待同一次请求的调用者仍共享结果,失败保留短缓存限流。
if let Some(entry) = cached.as_ref().filter(|entry| {
entry.until > Instant::now()
&& (!refresh || entry.fetched_at >= requested_at || entry.result.is_err())
}) {
return entry.result.clone();
}
let _permit = self
@@ -105,9 +112,11 @@ impl HttpPluginDistribution {
} else {
Duration::from_secs(30)
};
let fetched_at = Instant::now();
*cached = Some(CachedRelease {
result: result.clone(),
until: Instant::now() + ttl,
until: fetched_at + ttl,
fetched_at,
});
result
})
@@ -213,7 +222,7 @@ impl HttpPluginDistribution {
credentials: &[SourceCredential],
egress: Option<&PluginDistributionEgress>,
) -> Result<(Url, Option<String>, PluginSource, bool), AdminError> {
let release = self.release(query, credentials, egress).await?;
let release = self.release(query, credentials, egress, false).await?;
let selected = release
.assets
.iter()
@@ -153,7 +153,7 @@ impl PluginDistribution for HttpPluginDistribution {
egress: Option<PluginDistributionEgress>,
) -> Result<PluginRelease, AdminError> {
Ok(self
.release(query, &credentials, egress.as_ref())
.release(query, &credentials, egress.as_ref(), true)
.await?
.view)
}
@@ -109,6 +109,69 @@ async fn concurrent_queries_share_a_single_fetch_and_return_cache_times() {
}
}
#[tokio::test]
async fn explicit_queries_refresh_a_new_release_before_the_cache_expires() {
let server = MockServer::start().await;
let calls = std::sync::atomic::AtomicUsize::new(0);
Mock::given(path("/repos/example/plugins/releases/latest"))
.respond_with(move |_: &wiremock::Request| {
let mut metadata = release();
if calls.fetch_add(1, std::sync::atomic::Ordering::Relaxed) > 0 {
metadata["tag_name"] = json!("model-trace-v0.1.3");
}
ResponseTemplate::new(200).set_body_json(metadata)
})
.expect(2)
.mount(&server)
.await;
let distribution = transport(&server);
let previous = distribution
.query_release(query(None), vec![], None)
.await
.unwrap();
let latest = distribution
.query_release(query(None), vec![], None)
.await
.unwrap();
assert_eq!(previous.tag, "v1.0.0");
assert_eq!(latest.tag, "model-trace-v0.1.3");
assert!(latest.queried_at < previous.expires_at);
}
#[tokio::test]
async fn artifact_downloads_reuse_the_checked_fixed_release() {
let server = MockServer::start().await;
Mock::given(path("/repos/example/plugins/releases/tags/v1.0.0"))
.respond_with(ResponseTemplate::new(200).set_body_json(release()))
.expect(1)
.mount(&server)
.await;
Mock::given(path("/repos/example/plugins/releases/assets/1"))
.respond_with(ResponseTemplate::new(200).set_body_bytes(b"package"))
.mount(&server)
.await;
let distribution = transport(&server);
distribution
.query_release(query(Some("v1.0.0")), vec![], None)
.await
.unwrap();
let downloaded = distribution
.download(
RemotePluginLocation::Github {
repository: "example/plugins".into(),
tag: "v1.0.0".into(),
asset: "example_1.0.0_linux_x86_64.tar.gz".into(),
allow_prerelease: false,
sha256: None,
},
vec![],
None,
)
.await
.unwrap();
assert_eq!(downloaded.sha256, digest(b"package"));
}
#[tokio::test]
async fn rate_limit_applies_to_other_repositories_using_the_same_identity_and_egress() {
let server = MockServer::start().await;
@@ -60,7 +60,7 @@ async fn cancelling_the_fetch_owner_releases_the_shared_slot_for_the_next_query(
response
}
})
.expect(2)
.expect(3)
.mount(&server)
.await;
let distribution = Arc::new(transport(&server));
@@ -84,12 +84,12 @@ async fn cancelling_the_fetch_owner_releases_the_shared_slot_for_the_next_query(
.unwrap();
assert_eq!(recovered.tag, "v1.0.0");
assert_eq!(calls.load(Ordering::SeqCst), 2);
let cached = distribution
let refreshed = distribution
.query_release(query("example/plugins"), vec![], None)
.await
.unwrap();
assert_eq!(cached.queried_at, recovered.queried_at);
assert_eq!(calls.load(Ordering::SeqCst), 2);
assert!(refreshed.queried_at > recovered.queried_at);
assert_eq!(calls.load(Ordering::SeqCst), 3);
}
#[tokio::test]
@@ -101,7 +101,7 @@ async fn cancelling_a_cache_waiter_does_not_cancel_or_poison_the_fetch_owner() {
.set_body_json(release())
.set_delay(Duration::from_millis(500)),
)
.expect(1)
.expect(2)
.mount(&server)
.await;
let distribution = Arc::new(transport(&server));
@@ -124,12 +124,15 @@ async fn cancelling_a_cache_waiter_does_not_cancel_or_poison_the_fetch_owner() {
.unwrap()
.unwrap()
.unwrap();
let cached = distribution
assert_eq!(first.tag, "v1.0.0");
assert_eq!(server.received_requests().await.unwrap().len(), 1);
// 等待者取消不影响持有者完成,也不阻止之后的显式查询刷新结果。
let refreshed = distribution
.query_release(query("example/plugins"), vec![], None)
.await
.unwrap();
assert_eq!(cached.queried_at, first.queried_at);
assert_eq!(server.received_requests().await.unwrap().len(), 1);
assert!(refreshed.queried_at > first.queried_at);
assert_eq!(server.received_requests().await.unwrap().len(), 2);
}
#[tokio::test]
@@ -13,7 +13,8 @@
{ "capability": "web_socket_observer", "versions": [1] },
{ "capability": "usage", "versions": [1] },
{ "capability": "command_line", "versions": [1] },
{ "capability": "management", "versions": [1] }
{ "capability": "management", "versions": [1] },
{ "capability": "maintenance", "versions": [1] }
],
"permissions": [
"network",
@@ -21,6 +22,8 @@
"accounts",
"data",
"requests",
"public_endpoints"
"public_endpoints",
"groups",
"keys"
]
}
@@ -304,7 +304,7 @@ fn credential_facts(
}
}
fn mutation_context(context: &CallContext) -> MutationContext {
pub(super) fn mutation_context(context: &CallContext) -> MutationContext {
MutationContext {
actor: MutationActor::System,
// RPC 回调使用真实 call_id,RPC 结果的后置提交仍使用原始上下文;scope 由宿主签发且全局唯一。
@@ -347,7 +347,7 @@ fn credential_timestamp(value: Option<i64>) -> Result<Option<DateTime<Utc>>, Plu
.transpose()
}
fn encode(value: &impl serde::Serialize) -> Result<RpcReply, PluginFault> {
pub(super) fn encode(value: &impl serde::Serialize) -> Result<RpcReply, PluginFault> {
Ok(RpcReply {
result: serde_json::json!({}),
payload: serde_json::to_vec(value).map_err(|_| invalid())?,
@@ -37,7 +37,12 @@ impl PluginData {
payload: &[u8],
) -> Result<RpcReply, PluginFault> {
// 管理范围的只读授权不继承到客户端请求链,避免借数据查询扩大当前 Key 的范围。
if !self.authorized || !matches!(context.stage, Stage::Management | Stage::CommandLine) {
if !self.authorized
|| !matches!(
context.stage,
Stage::Management | Stage::CommandLine | Stage::Maintenance
)
{
return Err(denied());
}
if params != serde_json::json!({}) {
@@ -7,6 +7,7 @@ mod log;
mod middleware;
mod model;
pub(crate) mod private_state;
mod resources;
mod scope;
use std::{
@@ -16,7 +17,7 @@ use std::{
};
use futures::future::BoxFuture;
use gateway_admin::model::{AdminError, plugins::instances::PluginPermissionGrant};
use gateway_admin::model::AdminError;
use gateway_core::{account::OutboundProxy, lifecycle::CancellationToken};
use gateway_host::outbound::{HttpBody, HttpClient, NetworkPolicy};
use gateway_plugin_sdk::{CallContext, ErrorCode, PluginFault};
@@ -30,6 +31,7 @@ pub(crate) use middleware::{
MiddlewareBinding, MiddlewareBodyAuthority, MiddlewareCompletionBody, MiddlewareInvocation,
};
pub(crate) use model::PluginModelPortSlot;
pub(crate) use resources::PluginResourcePorts;
pub(crate) use scope::NetworkScope;
pub(crate) struct PluginCallbackPorts {
@@ -39,6 +41,7 @@ pub(crate) struct PluginCallbackPorts {
keys: Arc<PluginClientKeyPortSlot>,
models: Arc<PluginModelPortSlot>,
affinity: Arc<PluginAffinityPortSlot>,
resources: Arc<PluginResourcePorts>,
}
impl PluginCallbackPorts {
@@ -49,6 +52,7 @@ impl PluginCallbackPorts {
keys: Arc<PluginClientKeyPortSlot>,
models: Arc<PluginModelPortSlot>,
affinity: Arc<PluginAffinityPortSlot>,
resources: Arc<PluginResourcePorts>,
) -> Self {
Self {
http,
@@ -57,12 +61,14 @@ impl PluginCallbackPorts {
keys,
models,
affinity,
resources,
}
}
}
pub(crate) struct PluginCallbacks {
accounts: Arc<accounts::PluginAccounts>,
resources: Arc<resources::PluginResources>,
data: Arc<data::PluginData>,
keys: Arc<PluginClientKeyPortSlot>,
models_authorized: bool,
@@ -117,14 +123,16 @@ impl HttpStream {
impl PluginCallbacks {
pub(crate) fn new(
grants: &[PluginPermissionGrant],
instance: &gateway_admin::model::plugins::instances::PluginInstance,
maximum_payload: usize,
manifest: &gateway_plugin_sdk::Manifest,
log_slots: Arc<tokio::sync::Semaphore>,
private_state: Arc<private_state::PluginPrivateState>,
ports: PluginCallbackPorts,
) -> Result<Self, AdminError> {
let grants = &instance.grants;
Ok(Self {
resources: Arc::new(resources::PluginResources::new(instance, ports.resources)),
data: Arc::new(data::PluginData::new(ports.accounts.clone(), grants)),
accounts: Arc::new(accounts::PluginAccounts::new(ports.accounts, grants)),
keys: ports.keys,
@@ -374,6 +382,7 @@ impl CallbackHandler for PluginCallbacks {
let log = self.log.clone();
let private_state = self.private_state.clone();
let accounts = self.accounts.clone();
let resources = self.resources.clone();
let data = self.data.clone();
let keys = self.keys.clone();
let models_authorized = self.models_authorized;
@@ -417,6 +426,14 @@ impl CallbackHandler for PluginCallbacks {
if is_private_state {
return private_state.call(&method, params, &payload).await;
}
if matches!(
method.as_str(),
gateway_plugin_sdk::call::resources::GROUP_ENSURE
| gateway_plugin_sdk::call::resources::GROUP_MEMBERS
| gateway_plugin_sdk::call::resources::KEY_ENSURE
) {
return resources.call(&context, &method, params, &payload).await;
}
if method == "host.keys.list" {
if !models_authorized {
return Err(denied());
@@ -0,0 +1,147 @@
use super::{
accounts::{encode, map_admin_error, mutation_context},
denied, invalid,
};
use crate::RpcReply;
use gateway_admin::{
model::{
AdminError,
account_groups::{AccountGroupColor, CreateAccountGroup},
plugin_resources::{GroupMembersChange, ManagedKeyConfig, PluginResourceOwner},
plugins::instances::PluginInstance,
},
ports::plugin_resources::PluginResourceAccess,
};
use gateway_core::{engine::budget::ClientBudgetLimits, policy::RateLimits};
use gateway_plugin_sdk::{CallContext, PluginFault, call::resources};
use std::sync::{Arc, OnceLock, Weak};
pub(crate) struct PluginResourcePorts(OnceLock<Weak<dyn PluginResourceAccess>>);
impl PluginResourcePorts {
pub(crate) const fn new() -> Self {
Self(OnceLock::new())
}
pub(crate) fn bind(&self, port: &Arc<dyn PluginResourceAccess>) -> Result<(), AdminError> {
self.0
.set(Arc::downgrade(port))
.map_err(|_| AdminError::conflict("插件资源端口已经绑定"))
}
}
pub(super) struct PluginResources {
owner: PluginResourceOwner,
ports: Arc<PluginResourcePorts>,
}
impl PluginResources {
pub(super) fn new(instance: &PluginInstance, ports: Arc<PluginResourcePorts>) -> Self {
Self {
owner: PluginResourceOwner {
instance_id: instance.id.clone(),
artifact_sha256: instance.artifact_sha256.clone(),
revision: instance.revision,
},
ports,
}
}
pub(super) async fn call(
&self,
context: &CallContext,
method: &str,
params: serde_json::Value,
payload: &[u8],
) -> Result<RpcReply, PluginFault> {
if params != serde_json::json!({}) {
return Err(invalid());
}
let access = self
.ports
.0
.get()
.and_then(Weak::upgrade)
.ok_or_else(denied)?;
let mutation = mutation_context(context);
match method {
resources::GROUP_ENSURE => {
let request: resources::GroupEnsureRequest =
serde_json::from_slice(payload).map_err(|_| invalid())?;
let value = access
.ensure_group(
&self.owner,
request.resource_key,
CreateAccountGroup {
disable_fast: false,
name: request.name,
description: request.description,
color: AccountGroupColor::parse(&request.color).ok_or_else(invalid)?,
},
&mutation,
)
.await
.map_err(map_admin_error)?;
encode(&resources::ManagedResource {
id: value.id,
name: value.name,
enabled: value.enabled,
})
}
resources::KEY_ENSURE => {
let request: resources::KeyEnsureRequest =
serde_json::from_slice(payload).map_err(|_| invalid())?;
let value = access
.ensure_key(
&self.owner,
request.resource_key,
request.group_resource_keys,
ManagedKeyConfig {
name: request.name,
limits: RateLimits {
max_concurrency: request.max_concurrency,
requests_per_minute: request.requests_per_minute,
},
budget: ClientBudgetLimits {
daily_usd: request
.daily_limit_usd
.parse()
.map_err(|_| invalid())?,
weekly_usd: request
.weekly_limit_usd
.parse()
.map_err(|_| invalid())?,
},
},
&mutation,
)
.await
.map_err(map_admin_error)?;
encode(&resources::ManagedResource {
id: value.id,
name: value.name,
enabled: value.enabled,
})
}
resources::GROUP_MEMBERS => {
let request: resources::GroupMembersChange =
serde_json::from_slice(payload).map_err(|_| invalid())?;
let value = access
.change_members(
&self.owner,
GroupMembersChange {
resource_key: request.resource_key,
add: request.add,
remove: request.remove,
},
&mutation,
)
.await
.map_err(map_admin_error)?;
encode(&resources::GroupMembersChanged {
added: value.added,
removed: value.removed,
})
}
_ => Err(invalid()),
}
}
}
@@ -38,8 +38,11 @@ pub(super) fn validate(
for binding in &instance.bindings {
let resolved = crate::contribution::resolve(manifest, binding)?;
let capability = resolved.capability;
if matches!(capability, Capability::Management | Capability::CommandLine) {
return Err(AdminError::invalid("管理页面与命令行无需功能绑定"));
if matches!(
capability,
Capability::Management | Capability::CommandLine | Capability::Maintenance
) {
return Err(AdminError::invalid("管理页面、命令行与维护无需功能绑定"));
}
let stage: Stage = serde_json::from_value(serde_json::Value::String(binding.stage.clone()))
.map_err(|_| AdminError::invalid("插件调用阶段不合法"))?;
@@ -34,7 +34,6 @@ impl PluginRuntime {
let limits = self.config.package_limits;
let directory = self.config.cache_directory.clone();
let host_version = self.config.host_version.clone();
let grants = instance.grants.clone();
let validation_slot = self
.validators
.clone()
@@ -145,7 +144,7 @@ impl PluginRuntime {
contributes: manifest.contributes.clone(),
};
let callbacks = Arc::new(PluginCallbacks::new(
&grants,
&instance,
self.config.rpc_limits.maximum_buffered_body_bytes,
&manifest,
self.log_slots.clone(),
@@ -157,6 +156,7 @@ impl PluginRuntime {
self.client_key_ports.clone(),
self.model_ports.clone(),
self.affinity_ports.clone(),
self.resource_ports.clone(),
),
)?);
let session = Arc::new(
@@ -241,6 +241,9 @@ impl PluginRuntime {
revision: instance.revision,
session,
private_state,
maintenance: manifest
.contributes
.contains_key(&gateway_plugin_sdk::Capability::Maintenance),
});
Ok(PreparedContributions {
sessions,
@@ -0,0 +1,141 @@
//! 对账只从已发布集合开始;通知合并,实例串行,失败和停用不阻塞其他实例。
use super::{PluginRuntime, PreparedSet, RpcSession};
use futures::{StreamExt as _, future::BoxFuture};
use gateway_core::{
lifecycle::CancellationToken,
runtime::RuntimeSnapshotHandle,
task::{
DaemonRestartPolicy, DaemonTask, WorkerContribution, WorkerDefinitionError, WorkerId,
WorkerKind, WorkerRegistration, WorkerRunnable, WorkerTaskError,
},
};
use gateway_plugin_sdk::Stage;
use std::{sync::Arc, time::Duration};
use tokio::{sync::Notify, task::JoinSet};
const RECONCILE_INTERVAL: Duration = Duration::from_secs(30);
const RETRY_DELAY: Duration = Duration::from_secs(5);
const CALL_TIMEOUT: Duration = Duration::from_secs(30);
struct MaintenanceTask {
runtime: Arc<PluginRuntime>,
snapshots: RuntimeSnapshotHandle,
}
impl PluginRuntime {
/// 由 Host 监督维护任务;CLI 和只读准备不会启动后台业务。
pub fn maintenance_worker(
self: &Arc<Self>,
snapshots: RuntimeSnapshotHandle,
) -> Result<WorkerContribution, WorkerDefinitionError> {
Ok(WorkerContribution::Registration(
WorkerRegistration::try_new(
WorkerId::try_new(
WorkerKind::RuntimeSnapshotReconciliation,
"plugin-maintenance",
)?,
WorkerRunnable::Daemon {
restart: DaemonRestartPolicy::try_new(
Duration::from_secs(1),
Duration::from_secs(30),
)?,
task: Box::new(MaintenanceTask {
runtime: self.clone(),
snapshots,
}),
},
)?,
))
}
}
impl DaemonTask for MaintenanceTask {
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
Box::pin(async move {
let mut publications = self.snapshots.subscribe_publications();
let mut active: Option<Arc<PreparedSet>> = None;
let mut revision = None;
let mut stop = CancellationToken::new();
let mut tasks = JoinSet::new();
let mut notifications = Vec::new();
loop {
let snapshot = self.snapshots.acquire().ok();
let reference = snapshot.as_ref().and_then(|s| s.extensions());
let same_set = active.as_ref().map(|s| &s.id) == reference.map(|s| s.id());
if !same_set {
stop.cancel();
while tasks.join_next().await.is_some() {}
notifications.clear();
stop = CancellationToken::new();
active = match reference {
Some(reference) => {
Some(self.runtime.prepared_set(reference).await.map_err(|_| {
WorkerTaskError::safe("published plugin set is unavailable")
})?)
}
None => None,
};
if let Some(set) = &active {
for instance in set.sessions.iter().filter(|i| i.maintenance) {
let dirty = Arc::new(Notify::new());
dirty.notify_one();
notifications.push(dirty.clone());
tasks.spawn(reconcile(instance.session.clone(), dirty, stop.clone()));
}
}
} else if revision != snapshot.as_ref().map(|s| s.revision()) {
for dirty in &notifications {
dirty.notify_one();
}
}
revision = snapshot.as_ref().map(|s| s.revision());
// 发布变更优先于下一轮维护;不在快照发布锁中等待插件或数据库。
tokio::select! {
biased;
() = cancellation.cancelled() => break,
changed = publications.next() => if changed.is_none() { break; },
result = tasks.join_next(), if !tasks.is_empty() => {
if result.is_some() {
stop.cancel();
while tasks.join_next().await.is_some() {}
return Err(WorkerTaskError::safe("plugin maintenance task stopped unexpectedly"));
}
}
}
}
stop.cancel();
while tasks.join_next().await.is_some() {}
Ok(())
})
}
}
async fn reconcile(session: Arc<RpcSession>, dirty: Arc<Notify>, cancellation: CancellationToken) {
loop {
tokio::select! {
biased;
() = cancellation.cancelled() => return,
() = dirty.notified() => {},
() = tokio::time::sleep(RECONCILE_INTERVAL) => {},
}
let context = session.context(Stage::Maintenance, CALL_TIMEOUT);
let instance_id = context.instance_id.clone();
let result = tokio::select! {
biased;
() = cancellation.cancelled() => return,
result = session.call("plugin.reconcile", context, serde_json::json!({}), Vec::new()) => result,
};
if !result
.is_ok_and(|reply| reply.result == serde_json::json!({}) && reply.payload.is_empty())
{
tracing::warn!(%instance_id, "plugin maintenance failed; retrying after bounded delay");
tokio::select! {
biased;
() = cancellation.cancelled() => return,
() = tokio::time::sleep(RETRY_DELAY) => {},
}
dirty.notify_one();
}
}
}
@@ -1,4 +1,5 @@
mod instance;
mod maintenance;
use std::{
collections::{BTreeMap, BTreeSet},
@@ -95,6 +96,7 @@ pub struct PluginRuntime {
log_slots: Arc<Semaphore>,
pub(super) account_ports: Arc<PluginAccountPortSlot>,
client_key_ports: Arc<PluginClientKeyPortSlot>,
resource_ports: Arc<crate::callback::PluginResourcePorts>,
model_ports: Arc<PluginModelPortSlot>,
affinity_ports: Arc<PluginAffinityPortSlot>,
observers: RequestObserverExtensionIndex,
@@ -159,6 +161,7 @@ struct PreparedInstance {
revision: Revision,
session: Arc<RpcSession>,
private_state: Arc<PluginPrivateState>,
maintenance: bool,
}
impl PreparedSet {
@@ -228,6 +231,7 @@ impl PluginRuntime {
validators: Arc::new(Semaphore::new(2)),
account_ports: Arc::new(PluginAccountPortSlot::new()),
client_key_ports: Arc::new(PluginClientKeyPortSlot::new()),
resource_ports: Arc::new(crate::callback::PluginResourcePorts::new()),
model_ports: Arc::new(PluginModelPortSlot::new()),
affinity_ports: Arc::new(PluginAffinityPortSlot::new()),
observers: RequestObserverExtensionIndex::default(),
@@ -1219,3 +1223,12 @@ fn instance_fingerprint(
let bytes = serde_json::to_vec(&value).map_err(|_| AdminError::invalid("插件配置无法编码"))?;
Ok(hex::encode(Sha256::digest(bytes)))
}
impl PluginRuntime {
pub fn bind_resource_ports(
&self,
access: &Arc<dyn gateway_admin::ports::plugin_resources::PluginResourceAccess>,
) -> Result<(), AdminError> {
self.resource_ports.bind(access)
}
}
@@ -220,6 +220,18 @@ fn callback_allowed(method: &str, stage: Stage, permissions: &[Permission]) -> b
) {
return matches!(stage, Stage::Request | Stage::Attempt);
}
if stage == Stage::Maintenance
&& !matches!(
method,
"host.data.accounts.list"
| "host.data.quota.get"
| "host.groups.ensure"
| "host.groups.change_members"
| "host.keys.ensure"
)
{
return false;
}
let permission = match method {
"host.http.do"
| "host.http.do_stream"
@@ -236,10 +248,29 @@ fn callback_allowed(method: &str, stage: Stage, permissions: &[Permission]) -> b
}
"host.affinity.lookup" => Permission::Requests,
"host.data.accounts.list" | "host.data.quota.get"
if matches!(stage, Stage::Management | Stage::CommandLine) =>
if matches!(
stage,
Stage::Management | Stage::CommandLine | Stage::Maintenance
) =>
{
Permission::Data
}
"host.groups.ensure" | "host.groups.change_members"
if matches!(
stage,
Stage::Management | Stage::CommandLine | Stage::Maintenance
) =>
{
Permission::Groups
}
"host.keys.ensure"
if matches!(
stage,
Stage::Management | Stage::CommandLine | Stage::Maintenance
) =>
{
Permission::Keys
}
_ => return false,
};
permissions.contains(&permission)
@@ -0,0 +1,432 @@
use crate::support::environment::{Environment, account_grant, mutation};
use gateway_admin::{
model::{
account_groups::{AccountGroupColor, CreateAccountGroup},
plugin_resources::{GroupMembersChange, ManagedKeyConfig, PluginResourceOwner},
},
ports::plugin_resources::PluginResourceAccess,
};
use gateway_core::{
lifecycle::CancellationToken,
task::{WorkerContribution, WorkerRunnable},
};
use gateway_plugin_runtime::PluginRuntime;
use serde_json::{Value, json};
use std::{sync::Arc, time::Duration};
fn group(name: &str) -> CreateAccountGroup {
CreateAccountGroup {
name: name.into(),
description: None,
disable_fast: false,
color: AccountGroupColor::parse("#2563EBFF").unwrap(),
}
}
fn key() -> ManagedKeyConfig {
ManagedKeyConfig {
name: " managed key ".into(),
limits: gateway_core::policy::RateLimits::unlimited(),
budget: Default::default(),
}
}
async fn owner(environment: &Environment) -> PluginResourceOwner {
let instance = environment
.store
.admin_ports()
.plugins()
.load_instances()
.await
.unwrap()
.instances
.remove(0);
PluginResourceOwner {
instance_id: instance.id,
artifact_sha256: instance.artifact_sha256,
revision: instance.revision,
}
}
async fn import_account(
environment: &Environment,
core: &gateway_core::CoreBundle,
) -> gateway_core::account::ProviderAccountId {
use gateway_admin::model::provider_credentials::{
PreparedCredentialCreate, PreparedPluginAccountSave, ProviderDocument,
};
let access = gateway_admin::initialize_plugin_accounts(
crate::support::native::admin_registry(),
environment.store.admin_ports().accounts(),
core.snapshot_control(),
);
let result = access
.save(
PreparedPluginAccountSave::Create(PreparedCredentialCreate {
account_id: gateway_core::account::ProviderAccountId::new(format!(
"acct_{}",
uuid::Uuid::new_v4().simple()
))
.unwrap(),
provider_kind: gateway_core::routing::ProviderKind::new("openai").unwrap(),
name: "imported maintenance fixture".into(),
email: None,
upstream_user_id: None,
upstream_account_id: None,
plan_type: None,
authentication_kind: "api_key".into(),
provider_material: ProviderDocument::new(
gateway_core::account::OpaqueProviderData::new(
json!({"key":"test-only"}).as_object().unwrap().clone(),
),
),
model_access: None,
outbound_proxy: None,
has_refresh_token: false,
access_token_expires_at: None,
next_refresh_at: None,
enabled: true,
credential_state: gateway_core::account::CredentialState::Ready,
credential_observed_at: chrono::Utc::now(),
}),
&mutation(),
)
.await
.unwrap();
result.account_id
}
async fn revision(environment: &Environment) -> gateway_admin::model::Revision {
environment
.store
.admin_ports()
.plugins()
.load_instances()
.await
.unwrap()
.config_revision
}
async fn start(
environment: &Environment,
) -> (
Arc<PluginRuntime>,
gateway_core::CoreBundle,
Arc<dyn PluginResourceAccess>,
CancellationToken,
tokio::task::JoinHandle<()>,
) {
let (runtime, core) = environment.runtime().await;
let resources = gateway_admin::initialize_plugin_resources(
environment.store.admin_ports().plugin_resources(),
core.snapshot_control(),
);
runtime.bind_resource_ports(&resources).unwrap();
let WorkerContribution::Registration(registration) =
runtime.maintenance_worker(core.snapshots()).unwrap()
else {
panic!("maintenance registration")
};
let WorkerRunnable::Daemon { task, .. } = registration.runnable else {
panic!("maintenance daemon")
};
let stop = CancellationToken::new();
let cancellation = stop.clone();
let task = tokio::spawn(async move {
task.run(cancellation).await.unwrap();
});
(runtime, core, resources, stop, task)
}
async fn wait_done(path: &std::path::Path, after: usize) -> Value {
tokio::time::timeout(Duration::from_secs(15), async {
loop {
let entries: Vec<Value> = std::fs::read_to_string(path)
.unwrap_or_default()
.lines()
.filter_map(|line| serde_json::from_str(line).ok())
.collect();
let done: Vec<_> = entries
.iter()
.filter(|value| value["phase"] == "done")
.collect();
if done.len() > after {
return (*done.last().unwrap()).clone();
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await
.expect("plugin reconciliation finished")
}
#[tokio::test]
async fn published_reconciliation_provisions_retries_imports_and_restores_without_overwriting_groups()
{
let Some(environment) = Environment::create().await else {
return;
};
let account = environment.account(None).await;
let original = environment.account_group_with_account(&account).await;
let marker = environment.directory.path().join("maintenance.jsonl");
let failed = environment.directory.path().join("failed-once");
environment.install_plugin(json!({"maintenance_fixture":true, "maintenance_marker":marker, "maintenance_fail_once":failed}), vec![account_grant("groups"), account_grant("keys"), account_grant("data")]).await;
let (runtime, core, resources, stop, task) = start(&environment).await;
let first = wait_done(&marker, 0).await;
assert!(failed.exists());
let group_id =
gateway_core::routing::AccountGroupId::new(first["group"].as_str().unwrap()).unwrap();
let members = environment
.store
.admin_ports()
.account_groups()
.load_account_group_members(&[original.clone(), group_id.clone()])
.await
.unwrap();
assert_eq!(
members
.iter()
.filter(|m| m.account_id == account.as_str())
.count(),
2
);
let owner = owner(&environment).await;
let before = revision(&environment).await;
let same = resources
.ensure_group(
&owner,
"pool".into(),
group("ignored replacement name"),
&mutation(),
)
.await
.unwrap();
assert_eq!(same.id, first["group"]);
assert_eq!(same.name, "plugin fixture group");
assert_eq!(before, revision(&environment).await);
let imported = import_account(&environment, &core).await;
tokio::time::timeout(Duration::from_secs(10), async {
loop {
let members = environment
.store
.admin_ports()
.account_groups()
.load_account_group_members(std::slice::from_ref(&group_id))
.await
.unwrap();
if members
.iter()
.any(|member| member.account_id == imported.as_str())
{
break;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await
.unwrap();
stop.cancel();
task.await.unwrap();
runtime.shutdown().await;
environment.release_plugin_accounts(&runtime);
drop(resources);
drop(core);
drop(runtime);
// 离线新增账号后重新启动真实插件子进程;稳定资源键必须复用原分组与 Key。
let offline = environment.account(None).await;
std::fs::write(&marker, "").unwrap();
let (runtime, core, resources, stop, task) = start(&environment).await;
let restored = wait_done(&marker, 0).await;
assert_eq!(restored["group"], first["group"]);
assert_eq!(restored["key"], first["key"]);
let members = environment
.store
.admin_ports()
.account_groups()
.load_account_group_members(std::slice::from_ref(&group_id))
.await
.unwrap();
assert!(members.iter().any(|m| m.account_id == offline.as_str()));
let store = environment.store.admin_ports().plugins();
let snapshot = store.load_instances().await.unwrap();
let mut instance = snapshot.instances[0].clone();
instance.enabled = false;
let disabled = store
.save_instance(instance, snapshot.config_revision, &mutation())
.await
.unwrap();
core.snapshot_control()
.publish_committed(
gateway_core::routing::ConfigRevision::new(disabled.config_revision.get()).unwrap(),
)
.await;
assert!(
resources
.change_members(
&owner,
GroupMembersChange {
resource_key: "pool".into(),
add: vec![],
remove: vec![account.as_str().to_owned()]
},
&mutation()
)
.await
.is_err()
);
stop.cancel();
task.await.unwrap();
runtime.shutdown().await;
environment.release_plugin_accounts(&runtime);
drop(resources);
drop(core);
drop(runtime);
environment.close().await;
}
#[tokio::test]
async fn resource_transactions_enforce_ownership_current_revision_and_idempotent_membership() {
let Some(environment) = Environment::create().await else {
return;
};
environment
.install_plugin(
json!({}),
vec![account_grant("groups"), account_grant("keys")],
)
.await;
let (runtime, core) = environment.runtime().await;
let resources = gateway_admin::initialize_plugin_resources(
environment.store.admin_ports().plugin_resources(),
core.snapshot_control(),
);
let owner = owner(&environment).await;
let account = environment.account(None).await;
let concurrent_context = mutation();
let (left, right) = tokio::join!(
resources.ensure_group(
&owner,
"one".into(),
group("plugin one"),
&concurrent_context
),
resources.ensure_group(
&owner,
"one".into(),
group("concurrent name"),
&concurrent_context
)
);
let one = left.unwrap();
assert_eq!(one, right.unwrap());
let other = environment.account_group_with_account(&account).await;
assert!(
resources
.ensure_group(
&owner,
"collision".into(),
group(&format!("fixture {}", other.as_str())),
&mutation()
)
.await
.is_err()
);
assert!(
resources
.ensure_key(&owner, "key".into(), vec![], key(), &mutation())
.await
.is_err()
);
assert!(
resources
.ensure_key(
&owner,
"key".into(),
vec![other.as_str().to_owned()],
key(),
&mutation()
)
.await
.is_err()
);
let created = resources
.ensure_key(&owner, "key".into(), vec!["one".into()], key(), &mutation())
.await
.unwrap();
assert_eq!(created.name, "managed key");
assert_eq!(
created,
resources
.ensure_key(&owner, "key".into(), vec!["one".into()], key(), &mutation())
.await
.unwrap()
);
let change = || GroupMembersChange {
resource_key: "one".into(),
add: vec![account.as_str().to_owned()],
remove: vec![],
};
assert_eq!(
resources
.change_members(&owner, change(), &mutation())
.await
.unwrap()
.added,
1
);
let before = revision(&environment).await;
assert_eq!(
resources
.change_members(&owner, change(), &mutation())
.await
.unwrap()
.added,
0
);
assert_eq!(revision(&environment).await, before);
let result = resources
.change_members(
&owner,
GroupMembersChange {
resource_key: "one".into(),
add: vec![],
remove: vec![account.as_str().to_owned()],
},
&mutation(),
)
.await
.unwrap();
assert_eq!(result.removed, 1);
let members = environment
.store
.admin_ports()
.account_groups()
.load_account_group_members(&[
other.clone(),
gateway_core::routing::AccountGroupId::new(one.id).unwrap(),
])
.await
.unwrap();
assert_eq!(members.len(), 1);
assert_eq!(members[0].group_id, other);
let store = environment.store.admin_ports().plugins();
let snapshot = store.load_instances().await.unwrap();
let mut instance = snapshot.instances[0].clone();
instance.name = "new configuration revision".into();
store
.save_instance(instance, snapshot.config_revision, &mutation())
.await
.unwrap();
assert!(
resources
.ensure_group(&owner, "stale".into(), group("stale"), &mutation())
.await
.is_err()
);
runtime.shutdown().await;
environment.release_plugin_accounts(&runtime);
drop(resources);
drop(core);
drop(runtime);
environment.close().await;
}
@@ -1,6 +1,7 @@
mod authentication;
mod configuration;
mod data;
mod maintenance;
mod observer;
mod policy;
mod prepare;
@@ -82,7 +82,8 @@ fn assert_exhaustive_capability(capability: Capability) {
| Capability::WebSocketObserver
| Capability::Usage
| Capability::CommandLine
| Capability::Management => {}
| Capability::Management
| Capability::Maintenance => {}
}
}
@@ -93,11 +94,13 @@ fn assert_exhaustive_permission(permission: Permission) {
| Permission::Models
| Permission::Accounts
| Permission::Data
| Permission::PublicEndpoints => {}
| Permission::PublicEndpoints
| Permission::Groups
| Permission::Keys => {}
}
}
fn all_capabilities() -> [Capability; 11] {
fn all_capabilities() -> [Capability; 12] {
[
Capability::FrontendAuthentication,
Capability::Scheduler,
@@ -110,5 +113,6 @@ fn all_capabilities() -> [Capability; 11] {
Capability::Usage,
Capability::CommandLine,
Capability::Management,
Capability::Maintenance,
]
}
@@ -289,3 +289,72 @@ async fn retry_stage_allows_only_logs_and_private_state() {
assert_eq!(callbacks.called.load(Ordering::Relaxed), 2);
session.shutdown(Duration::from_secs(1)).await;
}
#[tokio::test]
async fn managed_resources_require_their_domains_and_control_plane_stages() {
for (method, permission) in [
("host.groups.ensure", Permission::Groups),
("host.groups.change_members", Permission::Groups),
("host.keys.ensure", Permission::Keys),
] {
for granted in [false, true] {
let callbacks = Arc::new(Callbacks::default());
let (_cache, session) = session_with_permissions(
callbacks.clone(),
if granted { vec![permission] } else { vec![] },
)
.await;
for stage in [
Stage::Registration,
Stage::Configuration,
Stage::PublicManagement,
Stage::Request,
Stage::Observation,
Stage::Management,
Stage::CommandLine,
Stage::Maintenance,
] {
let reply = invoke_callback(&session, stage, method).await;
if granted
&& matches!(
stage,
Stage::Management | Stage::CommandLine | Stage::Maintenance
)
{
assert!(reply.is_ok());
} else {
assert_permission_denied(reply);
}
}
session.shutdown(Duration::from_secs(1)).await;
}
}
}
#[tokio::test]
async fn maintenance_has_data_and_state_access_but_cannot_execute_models_or_read_credentials() {
let callbacks = Arc::new(Callbacks::default());
let (_cache, session) = session(callbacks).await;
for method in [
"host.data.accounts.list",
"host.data.quota.get",
"host.state.put",
"host.log",
] {
assert!(
invoke_callback(&session, Stage::Maintenance, method)
.await
.is_ok()
);
}
for method in [
"host.auth.get",
"host.auth.save",
"host.http.do",
"host.model.execute",
"host.keys.list",
] {
assert_permission_denied(invoke_callback(&session, Stage::Maintenance, method).await);
}
session.shutdown(Duration::from_secs(1)).await;
}
@@ -167,6 +167,16 @@ impl Environment {
.map(|grant| serde_json::from_value(json!(grant.permission)).unwrap())
.collect();
let mut contributes = Contributions::new();
if configuration.get("maintenance_fixture").is_some() {
contributes.extend([super::contribution_for_id(
&plugin_id,
Capability::Maintenance,
vec![Stage::Maintenance],
vec![],
vec![],
)]);
}
if configuration.get("command_registration").is_some() {
contributes.extend([super::contribution_for_id(
&plugin_id,
@@ -50,6 +50,7 @@ pub fn contribution_for_id(
Capability::Usage => "usage",
Capability::CommandLine => "commandLine",
Capability::Management => "management",
Capability::Maintenance => "maintenance",
};
(
capability,
@@ -81,6 +81,68 @@ struct Peer {
}
impl Peer {
async fn resource_fixture(
&self,
id: u64,
method: &str,
value: Value,
) -> Result<Value, PluginFault> {
let (result, payload) = self
.callback_payload(id, method, json!({}), serde_json::to_vec(&value).unwrap())
.await?;
assert_eq!(result, json!({}));
Ok(serde_json::from_slice(&payload).unwrap())
}
async fn reconcile_fixture(&self, id: u64) -> Result<(), PluginFault> {
self.append_observation_marker("maintenance_marker", &json!({"phase":"start"}));
let group = self
.resource_fixture(
id,
"host.groups.ensure",
json!({"resource_key":"pool", "name":"plugin fixture group", "color":"#2563EBFF"}),
)
.await?;
if let Some(path) = self.configuration["maintenance_fail_once"].as_str()
&& !std::path::Path::new(path).exists()
{
std::fs::write(path, "failed").unwrap();
return Err(PluginFault::new(ErrorCode::Fault, "fixture retry"));
}
let mut cursor = Value::Null;
loop {
let page = self
.resource_fixture(
id,
"host.data.accounts.list",
json!({"cursor":cursor, "limit":100}),
)
.await?;
let accounts: Vec<_> = page["accounts"]
.as_array()
.unwrap()
.iter()
.map(|account| account["account_id"].clone())
.collect();
self.resource_fixture(
id,
"host.groups.change_members",
json!({"resource_key":"pool", "add":accounts}),
)
.await?;
cursor = page["next_cursor"].clone();
if cursor.is_null() {
break;
}
}
let key = self.resource_fixture(id, "host.keys.ensure", json!({"resource_key":"key", "name":"plugin fixture key", "group_resource_keys":["pool"], "daily_limit_usd":"1", "weekly_limit_usd":"5"})).await?;
self.append_observation_marker(
"maintenance_marker",
&json!({"phase":"done", "group":group["id"], "key":key["id"]}),
);
Ok(())
}
async fn log_fixture(&self, id: u64) {
let Some(entries) = self
.configuration
@@ -297,6 +359,24 @@ impl Peer {
return;
}
match method.as_str() {
"plugin.reconcile" => {
let result = self.reconcile_fixture(id).await;
match result {
Ok(()) => {
self.send(
Message::Result {
id,
result: json!({}),
},
vec![],
)
.await
}
Err(error) => self.send(Message::Error { id, error }, vec![]).await,
}
return;
}
"policy.retry_decision" => {
self.append_observation_marker("retry_marker", &params);
if let Some(delay) = self.configuration["retry_delay_ms"].as_u64() {
@@ -54,6 +54,7 @@ Provider 固定为宿主内置的 OpenAI 与 xAI。插件提供以下扩展能
| `frontend_authentication` | `methods::FRONTEND_IDENTIFIER`、`methods::FRONTEND_AUTHENTICATE` |
| `management` | `PluginBuilder::management`;公开回调另用 `methods::MANAGEMENT_CALLBACK` |
| `command_line` | `PluginBuilder::command_line` |
| `maintenance` | `methods::RECONCILE`;无需功能绑定 |
## 访问域
@@ -65,9 +66,11 @@ Provider 固定为宿主内置的 OpenAI 与 xAI。插件提供以下扩展能
| `network` | 通过 `host.http.*` 使用宿主受管出站网络;仍受统一代理、超时、大小和流控规则约束 |
| `models` | 列出非秘密 Key、按所选 Key 查询模型,以及通过 `host.model.*` 调用模型;调用可能产生消耗 |
| `accounts` | 查询账号、读取原始凭据及创建或替换账号;写入仍经过 revision CAS、审计和发布事务 |
| `data` | 在管理/命令阶段只读全部账号的最小基础信息及已有额度观测,不包含凭据、写入或预测 |
| `data` | 在管理/命令/维护阶段只读全部账号的最小基础信息及已有额度观测,不包含凭据、写入或预测 |
| `requests` | 参与请求/响应处理、路由、调度、观察及亲和查询 |
| `public_endpoints` | 提供无需登录即可访问的已声明静态资源或一次性票据回调 |
| `groups` | 创建本实例分组,允许将所有现有及未来新增账号加入或移出这些分组,保留其他分组关系 |
| `keys` | 创建仅绑定本实例分组的 Key,返回非秘密身份;不授予其他 Key 的修改或明文读取权限 |
`host.log` 和本插件声明的 `host.state.*` 是基础设施,不需要额外 permission。公开调用阶段不开放任何
宿主回调;权限也不能把一个父调用的句柄、流或上下文转移到另一个调用。
@@ -79,7 +82,7 @@ Provider 固定为宿主内置的 OpenAI 与 xAI。插件提供以下扩展能
### 基础事实
纯展示页面声明 `management` 和 `data` 即可,不需要请求处理 binding 或 `accounts` 权限。
`data` 是独立的管理员授权域,只允许 `management`、`command_line` 阶段使用;请求链、注册及公开回调均拒绝。
`data` 是独立的管理员授权域,只允许 `management`、`command_line`、`maintenance` 阶段使用;请求链、注册及公开回调均拒绝。
它可读取全部账号的下列最小事实,不继承或授予客户端 Key 的模型执行权。
| SDK 方法 | 回调 | 查询与结果 |
@@ -118,6 +121,34 @@ for account in page.accounts {
插件可以自行选择其业务需要的账号和 Provider,不需要安装器预先配置允许列表。宿主仍校验账号归属、资格、
凭据版本和并发事务;原始凭据不得进入日志、审计正文或控制元数据。
### 自有资源与维护
声明 `maintenance` 并注册 `methods::RECONCILE`,宿主会在实例启用发布、进程恢复及配置变更后调用
`plugin.reconcile`,并每 30 秒补偿一次。处理器必须幂等;调用可能重复,配置通知会合并,不表示逐条账号事件。
每个实例串行执行,不同实例相互独立;调用限时 30 秒,失败后等待 5 秒重试。维护失败不回滚已经提交的资源,
下一次对账继续补齐。只读校验、准备候选与 CLI 帮助不会启动维护;停用、替换和宿主关闭时取消旧任务。
维护阶段允许日志、私有状态,以及已授权的 `data`、`groups`、`keys` 回调;不开放网络、凭据或模型执行。
其他管理/命令入口也可使用下列资源方法:
| SDK 方法 | 参数与行为 |
| --- | --- |
| `call.host.ensure_group(GroupEnsureRequest)` | `resource_key`、`name`、`color`、可选 `description`;首次创建,之后返回已有 `{id,name,enabled}` |
| `call.host.change_group_members(GroupMembersChange)` | `resource_key`、`add`、`remove`;两列表合计最多 200 个账号 ID,不得重复或交叉,返回实际 `added`/`removed` 数量 |
| `call.host.ensure_key(KeyEnsureRequest)` | `resource_key`、`name`、非空 `group_resource_keys`(最多 64 个)、`max_concurrency`、`requests_per_minute`、`daily_limit_usd`、`weekly_limit_usd`;返回 `{id,name,enabled}`,预算零值沿用原生无限制语义 |
资源键使用 1~64 位小写字母、数字、`_`、`.`、`-`,首位为字母或数字,在实例和资源类型内唯一。
归属由宿主签发,插件不能指定其他实例。`ensure` 只在首次创建时使用属性,不覆盖管理员对名称、启用状态或限制的修改;
同名的管理员资源不会被接管,名称冲突会失败。账号已被删除时增量加入会忽略该 ID。
无实际变化时不会写审计或递增配置版本。创建与归属登记、成员变更、授权复验均在同一事务完成。
插件升级沿用实例资源;停用保留分组和 Key,其原生启用状态保持不变。删除实例解除归属,资源仍由管理员管理;
重新安装为新实例不会接管旧资源。管理员删除自有资源后,仍启用的维护逻辑可在下次对账重新创建。
Key 明文仍通过宿主管理面查看,插件模型调用使用返回的 Key ID。
典型处理器先确保分组存在,再通过 `data` 分页查询账号并增量补齐成员,最后确保 Key 存在。
安装 `groups` 域即授权纳入全部当前及未来账号;插件可按自己的配置筛选账号,但该筛选不构成宿主的权限边界。
### Key、模型与模型调用
`models` 域开放:
@@ -49,6 +49,7 @@ capability 固定并由工具生成:
| `registration` / `retry` | `model_catalog` / `retry_policy` |
| `observation` | `request_lifecycle`、`web_socket_observer`、`usage` |
| `management` / `command_line` | `management` / `command_line` |
| `maintenance` | `maintenance` |
`middleware` 必须显式选择 `request`、`attempt` 或两者;协议格式等真实业务选择也不能省略。
安装清单若携带不同的固定阶段会被拒绝,而不是在加载时静默改写。
@@ -62,8 +63,10 @@ capability 固定并由工具生成:
| `network` | 使用宿主受管网络 |
| `models` | 查询非秘密 Key 与模型并调用模型,可能产生消耗 |
| `accounts` | 查询、读取原始凭据和修改账号 |
| `data` | 仅在管理与命令阶段只读全部账号的基础信息和已有额度观测 |
| `data` | 仅在管理、命令与维护阶段只读全部账号的基础信息和已有额度观测 |
| `requests` | 查看和处理请求、响应、路由、调度与观察事实 |
| `groups` | 创建本实例分组并管理所有当前及未来账号在这些分组中的成员关系 |
| `keys` | 创建绑定本实例分组的 Key,不读取密钥明文 |
| `public_endpoints` | 提供无需登录即可访问的资源或回调 |
安装时统一接受清单声明的域,不再填写逐方法、用途、Key、账号或 Provider 白名单。日志与清单声明的
@@ -10,3 +10,4 @@ pub mod model;
pub mod observation;
pub mod policy;
pub mod registration;
pub mod resources;
@@ -0,0 +1,58 @@
//! 实例自有资源:稳定 resource_key 用于重试,名称只在首次创建时使用。
use serde::{Deserialize, Serialize};
pub const GROUP_ENSURE: &str = "host.groups.ensure";
pub const GROUP_MEMBERS: &str = "host.groups.change_members";
pub const KEY_ENSURE: &str = "host.keys.ensure";
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct GroupEnsureRequest {
pub resource_key: String,
pub name: String,
pub color: String,
#[serde(default)]
pub description: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct GroupMembersChange {
pub resource_key: String,
#[serde(default)]
pub add: Vec<String>,
#[serde(default)]
pub remove: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct KeyEnsureRequest {
pub resource_key: String,
pub name: String,
/// 必须包含至少一个本实例分组,不能创建不受分组约束的 Key。
pub group_resource_keys: Vec<String>,
#[serde(default)]
pub max_concurrency: u64,
#[serde(default)]
pub requests_per_minute: u64,
pub daily_limit_usd: String,
pub weekly_limit_usd: String,
}
/// 返回非秘密身份;调用模型时使用 id,密钥明文仍由管理员管理。
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ManagedResource {
pub id: String,
pub name: String,
pub enabled: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct GroupMembersChanged {
pub added: u64,
pub removed: u64,
}
@@ -17,6 +17,7 @@ pub enum Capability {
Usage,
CommandLine,
Management,
Maintenance,
}
impl Capability {
@@ -35,6 +36,7 @@ impl Capability {
Self::Usage => "usage",
Self::CommandLine => "command_line",
Self::Management => "management",
Self::Maintenance => "maintenance",
}
}
@@ -51,6 +53,7 @@ impl Capability {
Self::RequestLifecycle | Self::WebSocketObserver | Self::Usage => &[Stage::Observation],
Self::CommandLine => &[Stage::CommandLine],
Self::Management => &[Stage::Management],
Self::Maintenance => &[Stage::Maintenance],
}
}
}
@@ -72,6 +75,8 @@ pub enum Stage {
CommandLine,
/// 公共登录回调不继承插件实例的任何宿主回调权限。
PublicManagement,
/// 宿主针对已发布实例签发的幂等维护调用。
Maintenance,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
@@ -144,17 +149,21 @@ pub enum Permission {
Data,
Requests,
PublicEndpoints,
Groups,
Keys,
}
impl Permission {
/// 当前公开访问域,安装摘要与授权校验复用同一集合。
pub const ALL: [Self; 6] = [
pub const ALL: [Self; 8] = [
Self::Network,
Self::Models,
Self::Accounts,
Self::Data,
Self::Requests,
Self::PublicEndpoints,
Self::Groups,
Self::Keys,
];
#[must_use]
@@ -166,6 +175,8 @@ impl Permission {
Self::Data => "data",
Self::Requests => "requests",
Self::PublicEndpoints => "public_endpoints",
Self::Groups => "groups",
Self::Keys => "keys",
}
}
@@ -179,6 +190,8 @@ impl Permission {
Self::Data => "基础数据",
Self::Requests => "请求处理",
Self::PublicEndpoints => "公开入口",
Self::Groups => "专用账号分组",
Self::Keys => "专用 API Key",
}
}
@@ -190,10 +203,14 @@ impl Permission {
Self::Models => "查询模型与 API Key 信息并调用模型,可能产生消耗",
Self::Accounts => "读取和修改账号,包括访问原始凭据",
Self::Data => {
"在管理或命令入口只读查询所有账号的基础信息与已有额度观测,不包含凭据或修改权限"
"在管理、命令或维护入口只读查询所有账号的基础信息与已有额度观测,不包含凭据或修改权限"
}
Self::Requests => "查看和处理请求、响应、路由及账号选择",
Self::PublicEndpoints => "提供无需登录即可访问的资源与回调入口",
Self::Groups => {
"创建本插件的分组,可将所有现有及未来新增账号加入或移出这些分组,不修改其他分组"
}
Self::Keys => "创建绑定本插件分组的 API Key,不读取密钥明文或修改管理员创建的 Key",
}
}
}
@@ -1,11 +1,9 @@
use serde::{Serialize, de::DeserializeOwned};
use crate::{PluginFault, call::data};
use crate::{ErrorCode, PluginFault, call::data};
use super::{HostClient, SessionError};
use super::{HostClient, payload_call};
impl HostClient {
/// 在已授权的 management/command_line 调用中分页读取账号基础事实。
/// 在已授权的 management/command_line/maintenance 调用中分页读取账号基础事实。
///
/// # Errors
///
@@ -14,7 +12,7 @@ impl HostClient {
&self,
query: data::AccountFactsQuery,
) -> Result<data::AccountFactsPage, PluginFault> {
data_call(self, data::ACCOUNTS_LIST, query).await
payload_call(self, data::ACCOUNTS_LIST, query).await
}
/// 读取已有额度观测,不触发上游刷新,也不返回宿主预测结果。
@@ -26,26 +24,6 @@ impl HostClient {
&self,
query: data::QuotaFactsQuery,
) -> Result<data::QuotaFacts, PluginFault> {
data_call(self, data::QUOTA_GET, query).await
payload_call(self, data::QUOTA_GET, query).await
}
}
async fn data_call<T: Serialize, R: DeserializeOwned>(
host: &HostClient,
method: &str,
query: T,
) -> Result<R, PluginFault> {
let invalid = || PluginFault::new(ErrorCode::InvalidInput, "invalid data callback payload");
let reply = host
.call(
method,
serde_json::json!({}),
serde_json::to_vec(&query).map_err(|_| invalid())?,
)
.await
.map_err(SessionError::into_plugin_fault)?;
if reply.result != serde_json::json!({}) {
return Err(invalid());
}
serde_json::from_slice(&reply.payload).map_err(|_| invalid())
}
@@ -4,8 +4,12 @@ mod data;
mod frame;
mod middleware;
mod plugin;
mod resources;
mod session;
use crate::{ErrorCode, PluginFault};
use serde::{Serialize, de::DeserializeOwned};
pub use frame::{read_frame, validate_frame, write_frame};
pub use middleware::{
MiddlewareBody, MiddlewareBodySender, MiddlewareCall, MiddlewareNext, MiddlewarePlugin,
@@ -18,3 +22,23 @@ pub use session::{
CallCancellation, CallFuture, CallReply, HostClient, HostReply, PluginCall, PluginHandler,
PluginSession, ResponseStream, SessionConfig, SessionError, StreamSender,
};
async fn payload_call<T: Serialize, R: DeserializeOwned>(
host: &HostClient,
method: &str,
query: T,
) -> Result<R, PluginFault> {
let invalid = || PluginFault::new(ErrorCode::InvalidInput, "invalid host callback payload");
let reply = host
.call(
method,
serde_json::json!({}),
serde_json::to_vec(&query).map_err(|_| invalid())?,
)
.await
.map_err(SessionError::into_plugin_fault)?;
if reply.result != serde_json::json!({}) {
return Err(invalid());
}
serde_json::from_slice(&reply.payload).map_err(|_| invalid())
}
@@ -126,3 +126,12 @@ pub const STATE_MIGRATE: Method<host::StateMigrationRequest, host::StateMigratio
decode_payload,
encode_payload,
);
/// 启用、恢复、配置变化及周期补偿共用的幂等入口;通知不代表逐条事件。
pub const RECONCILE: Method<Empty, Empty> = Method::new(
"plugin.reconcile",
&[C::Maintenance],
&[S::Maintenance],
decode_metadata_without_payload,
encode_metadata,
);
@@ -16,6 +16,7 @@ pub(super) fn validate(
};
for capability in manifest.contributes.keys() {
match capability {
C::Maintenance => required(methods::RECONCILE.name)?,
C::Middleware => required(crate::call::middleware::HANDLE_METHOD)?,
C::ModelRouter => required(methods::ROUTE_MODEL.name)?,
C::ModelCatalog => required(methods::MODEL_CATALOG_REGISTER.name)?,
@@ -0,0 +1,38 @@
use crate::{PluginFault, call::resources};
use super::{HostClient, payload_call};
impl HostClient {
/// 幂等创建或读取本实例分组,不覆盖管理员修改的现有属性。
///
/// # Errors
/// 权限、实例版本或输入不合法,或者宿主写入失败时返回错误。
pub async fn ensure_group(
&self,
request: resources::GroupEnsureRequest,
) -> Result<resources::ManagedResource, PluginFault> {
payload_call(self, resources::GROUP_ENSURE, request).await
}
/// 增量修改本实例分组的成员;保留账号在其他分组中的成员关系。
///
/// # Errors
/// 权限、实例版本、资源归属或输入不合法,或者宿主写入失败时返回错误。
pub async fn change_group_members(
&self,
request: resources::GroupMembersChange,
) -> Result<resources::GroupMembersChanged, PluginFault> {
payload_call(self, resources::GROUP_MEMBERS, request).await
}
/// 幂等创建或读取仅绑定本实例分组的 Key;不返回密钥明文。
///
/// # Errors
/// 权限、实例版本、资源归属或输入不合法,或者宿主写入失败时返回错误。
pub async fn ensure_key(
&self,
request: resources::KeyEnsureRequest,
) -> Result<resources::ManagedResource, PluginFault> {
payload_call(self, resources::KEY_ENSURE, request).await
}
}
@@ -99,6 +99,7 @@ impl SettingsStore for AdminSettingsStoreAdapter {
concurrency_wait_timeout_seconds: command.concurrency_wait_timeout_seconds,
responses_max_decompressed_body_bytes: command
.responses_max_decompressed_body_bytes,
smart_scheduling: command.smart_scheduling,
rotation_strategy: command.rotation_strategy.as_str().to_owned(),
model_mappings: store_model_mappings(command.model_mappings),
min_codex_desktop_version: command.min_codex_desktop_version,
@@ -137,6 +138,7 @@ impl SettingsStore for AdminSettingsStoreAdapter {
"concurrency_wait_timeout_seconds".to_owned(),
"responses_max_decompressed_body_bytes".to_owned(),
"rotation_strategy".to_owned(),
"smart_scheduling_json".to_owned(),
"min_codex_desktop_version".to_owned(),
"min_codex_cli_version".to_owned(),
"retention".to_owned(),
@@ -247,6 +249,7 @@ pub(crate) fn admin_runtime_settings(
max_waiting_per_account: settings.max_waiting_per_account,
concurrency_wait_timeout_seconds: settings.concurrency_wait_timeout_seconds,
responses_max_decompressed_body_bytes: settings.responses_max_decompressed_body_bytes,
smart_scheduling: settings.smart_scheduling,
rotation_strategy,
min_codex_desktop_version: settings.min_codex_desktop_version,
min_codex_cli_version: settings.min_codex_cli_version,
@@ -179,6 +179,7 @@ async fn connect(
}),
backup_ports(pool.clone(), &config)?,
plugins.clone(),
plugins.clone(),
plugins,
);
@@ -228,7 +228,6 @@ impl AccountGroupStore for PgAccountGroupRepository {
command: NewAccountGroup,
context: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
validate_group_fields(&command.name, command.description.as_deref())?;
let id = command.id.clone();
let audit = mutation_audit(
context,
@@ -244,19 +243,7 @@ impl AccountGroupStore for PgAccountGroupRepository {
let revision = self
.mutate(audit, |transaction| {
Box::pin(async move {
sqlx::query(
"insert into account_groups
(id, name, description, color, disable_fast, enabled, created_at, updated_at)
values ($1, $2, $3, $4, $5, true, now(), now())",
)
.bind(command.id.as_str())
.bind(command.name)
.bind(command.description)
.bind(command.color.as_str())
.bind(command.disable_fast)
.execute(&mut **transaction)
.await
.map_err(|error| map_group_write_error(error, command.id.as_str()))?;
insert_account_group_in_transaction(transaction, &command).await?;
Ok(())
})
})
@@ -669,3 +656,26 @@ fn invalid(message: &str) -> StoreError {
fn unavailable(message: &'static str) -> StoreError {
postgres_unavailable(message)
}
/// 原生管理与插件自有分组共用字段校验和写入规则。
pub(crate) async fn insert_account_group_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
command: &NewAccountGroup,
) -> StoreResult<()> {
validate_group_fields(&command.name, command.description.as_deref())
.map_err(|_| invalid("invalid account group fields"))?;
sqlx::query(
"insert into account_groups
(id, name, description, color, disable_fast, enabled, created_at, updated_at)
values ($1, $2, $3, $4, $5, true, now(), now())",
)
.bind(command.id.as_str())
.bind(&command.name)
.bind(&command.description)
.bind(command.color.as_str())
.bind(command.disable_fast)
.execute(&mut **transaction)
.await
.map_err(|error| map_group_write_error(error, command.id.as_str()))?;
Ok(())
}
@@ -1,6 +1,7 @@
mod artifacts;
mod credentials;
mod instances;
mod resources;
mod sources;
mod state;
@@ -0,0 +1,320 @@
//! 插件资源写入在同一事务中验证当前实例、访问域与资源归属。
use super::super::{
account_groups::insert_account_group_in_transaction,
append_admin_audit_event_in_transaction, bump_config_revision_in_transaction,
client_keys::{NewClientApiKey, insert_client_api_key_in_transaction},
};
use super::PgPluginStore;
use crate::{admin_revision, admin_store_error, mutation_audit};
use async_trait::async_trait;
use gateway_admin::{
model::{
MutationContext,
account_groups::NewAccountGroup,
client_keys::NewClientKey,
plugin_resources::{
GroupMembersChange, GroupMembersChanged, ManagedResource, PluginResourceOwner,
ResourceMutation,
},
},
ports::{
plugin_resources::PluginResourceStore,
store::{AdminStoreError, AdminStoreErrorKind, AdminStoreResult},
},
};
use sqlx::{Postgres, Row as _, Transaction};
use std::collections::BTreeSet;
fn unavailable() -> AdminStoreError {
AdminStoreError::new(
AdminStoreErrorKind::Unavailable,
"plugin resource",
"plugin resource store is unavailable",
)
}
fn denied() -> AdminStoreError {
AdminStoreError::new(
AdminStoreErrorKind::Conflict,
"plugin resource",
"plugin instance or resource authorization changed",
)
}
fn invalid() -> AdminStoreError {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
"plugin resource",
"invalid plugin resource input",
)
}
fn validate_resource_key(key: &str) -> AdminStoreResult<()> {
if key.is_empty()
|| key.len() > 64
|| !key.as_bytes()[0].is_ascii_alphanumeric()
|| !key
.bytes()
.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || b"_.-".contains(&c))
{
return Err(invalid());
}
Ok(())
}
impl PgPluginStore {
async fn resource_transaction(
&self,
owner: &PluginResourceOwner,
permission: &str,
) -> AdminStoreResult<Transaction<'_, Postgres>> {
let mut tx = self.pool.begin().await.map_err(|_| unavailable())?;
// 与全部管理写入保持相同锁顺序;无变化的对账只锁定,不递增 revision。
sqlx::query("select config_revision from runtime_settings where id=1 for update")
.execute(&mut *tx)
.await
.map_err(|_| unavailable())?;
let id = uuid::Uuid::parse_str(&owner.instance_id).map_err(|_| denied())?;
let revision = i64::try_from(owner.revision.get()).map_err(|_| denied())?;
let allowed: bool = sqlx::query_scalar(
"select exists(
select 1 from plugin_instances i
join plugin_artifacts a on a.sha256=i.artifact_sha256
where i.id=$1 and i.enabled and i.revision=$2 and i.artifact_sha256=$3
and a.accepted_at is not null and a.metadata_json->'requestedPermissions' ? $4
)",
)
.bind(id)
.bind(revision)
.bind(&owner.artifact_sha256)
.bind(permission)
.fetch_one(&mut *tx)
.await
.map_err(|_| unavailable())?;
if !allowed {
return Err(denied());
}
Ok(tx)
}
}
async fn finish<T>(
mut tx: Transaction<'_, Postgres>,
context: &MutationContext,
kind: &str,
id: &str,
changed: bool,
value: T,
) -> AdminStoreResult<ResourceMutation<T>> {
let revision = if changed {
let revision = bump_config_revision_in_transaction(&mut tx)
.await
.map_err(|e| admin_store_error("plugin resource", e))?;
let audit = mutation_audit(
context,
"plugin_reconcile",
kind,
id,
vec!["plugin_owned_resource".to_owned()],
);
append_admin_audit_event_in_transaction(&mut tx, audit, revision)
.await
.map_err(|e| admin_store_error("plugin resource", e))?;
Some(admin_revision(revision)?)
} else {
None
};
tx.commit().await.map_err(|_| unavailable())?;
Ok(ResourceMutation { revision, value })
}
fn resource(row: &sqlx::postgres::PgRow) -> AdminStoreResult<ManagedResource> {
Ok(ManagedResource {
id: row.try_get("id").map_err(|_| unavailable())?,
name: row.try_get("name").map_err(|_| unavailable())?,
enabled: row.try_get("enabled").map_err(|_| unavailable())?,
})
}
#[async_trait]
impl PluginResourceStore for PgPluginStore {
async fn ensure_group(
&self,
owner: &PluginResourceOwner,
resource_key: String,
command: NewAccountGroup,
context: &MutationContext,
) -> AdminStoreResult<ResourceMutation<ManagedResource>> {
validate_resource_key(&resource_key)?;
let mut tx = self.resource_transaction(owner, "groups").await?;
let instance = uuid::Uuid::parse_str(&owner.instance_id).map_err(|_| denied())?;
if let Some(row) = sqlx::query(
"select g.id,g.name,g.enabled from plugin_group_resources r
join account_groups g on g.id=r.group_id
where r.instance_id=$1 and r.resource_key=$2",
)
.bind(instance)
.bind(&resource_key)
.fetch_optional(&mut *tx)
.await
.map_err(|_| unavailable())?
{
return finish(tx, context, "account_group", "", false, resource(&row)?).await;
}
insert_account_group_in_transaction(&mut tx, &command)
.await
.map_err(|e| admin_store_error("account group", e))?;
sqlx::query("insert into plugin_group_resources(instance_id,resource_key,group_id) values($1,$2,$3)")
.bind(instance)
.bind(resource_key)
.bind(command.id.as_str())
.execute(&mut *tx)
.await.map_err(|_| unavailable())?;
let value = ManagedResource {
id: command.id.as_str().to_owned(),
name: command.name,
enabled: true,
};
finish(
tx,
context,
"account_group",
command.id.as_str(),
true,
value,
)
.await
}
async fn ensure_key(
&self,
owner: &PluginResourceOwner,
resource_key: String,
groups: Vec<String>,
command: NewClientKey,
context: &MutationContext,
) -> AdminStoreResult<ResourceMutation<ManagedResource>> {
validate_resource_key(&resource_key)?;
if groups.is_empty()
|| groups.len() > 64
|| groups.iter().collect::<BTreeSet<_>>().len() != groups.len()
{
return Err(invalid());
}
for group in &groups {
validate_resource_key(group)?;
}
let mut tx = self.resource_transaction(owner, "keys").await?;
let instance = uuid::Uuid::parse_str(&owner.instance_id).map_err(|_| denied())?;
if let Some(row) = sqlx::query(
"select k.id,k.name,k.enabled from plugin_key_resources r
join client_api_keys k on k.id=r.key_id
where r.instance_id=$1 and r.resource_key=$2",
)
.bind(instance)
.bind(&resource_key)
.fetch_optional(&mut *tx)
.await
.map_err(|_| unavailable())?
{
return finish(tx, context, "client_api_key", "", false, resource(&row)?).await;
}
let group_ids: Vec<String> = sqlx::query_scalar("select group_id from plugin_group_resources where instance_id=$1 and resource_key=any($2::text[])")
.bind(instance)
.bind(&groups)
.fetch_all(&mut *tx)
.await.map_err(|_| unavailable())?;
if group_ids.len() != groups.len() {
return Err(denied());
}
let key = NewClientApiKey {
request_profile_overrides: command.request_profile_overrides,
id: command.id.as_str().to_owned(),
name: command.name,
label: command.label,
key: command.plaintext,
group_ids,
budget: command.budget,
max_concurrency: command.limits.max_concurrency,
requests_per_minute: command.limits.requests_per_minute,
};
insert_client_api_key_in_transaction(&mut tx, &key)
.await
.map_err(|e| admin_store_error("client API key", e))?;
sqlx::query(
"insert into plugin_key_resources(instance_id,resource_key,key_id) values($1,$2,$3)",
)
.bind(instance)
.bind(resource_key)
.bind(&key.id)
.execute(&mut *tx)
.await
.map_err(|_| unavailable())?;
let value = ManagedResource {
id: key.id.clone(),
name: key.name.trim().to_owned(),
enabled: true,
};
finish(tx, context, "client_api_key", &key.id, true, value).await
}
async fn change_members(
&self,
owner: &PluginResourceOwner,
command: GroupMembersChange,
context: &MutationContext,
) -> AdminStoreResult<ResourceMutation<GroupMembersChanged>> {
validate_resource_key(&command.resource_key)?;
let add: BTreeSet<_> = command.add.iter().collect();
let remove: BTreeSet<_> = command.remove.iter().collect();
if command.add.len() + command.remove.len() > 200
|| !add.is_disjoint(&remove)
|| add.len() != command.add.len()
|| remove.len() != command.remove.len()
{
return Err(invalid());
}
let mut tx = self.resource_transaction(owner, "groups").await?;
let instance = uuid::Uuid::parse_str(&owner.instance_id).map_err(|_| denied())?;
let group: String = sqlx::query_scalar(
"select group_id from plugin_group_resources where instance_id=$1 and resource_key=$2",
)
.bind(instance)
.bind(command.resource_key)
.fetch_optional(&mut *tx)
.await
.map_err(|_| unavailable())?
.ok_or_else(denied)?;
// 已删除账号自然不再需要加入;外键负责与并发账号删除保持一致。
let added = sqlx::query(
"insert into account_group_accounts(account_group_id,provider_account_id,created_at)
select $1,id,now() from provider_accounts where id=any($2::text[])
on conflict do nothing",
)
.bind(&group)
.bind(command.add)
.execute(&mut *tx)
.await
.map_err(|_| unavailable())?
.rows_affected();
let removed = sqlx::query(
"delete from account_group_accounts
where account_group_id=$1 and provider_account_id=any($2::text[])",
)
.bind(&group)
.bind(command.remove)
.execute(&mut *tx)
.await
.map_err(|_| unavailable())?
.rows_affected();
finish(
tx,
context,
"account_group",
&group,
added + removed > 0,
GroupMembersChanged { added, removed },
)
.await
}
}
@@ -187,16 +187,16 @@ impl ProviderAccountStore for PgProviderAccountRepository {
has_refresh_token = $7, access_token_expires_at = $8,
next_refresh_at = $9,
credential_state = case
when $10::text is not null and enabled
when $10::text is not null
then $10 else credential_state end,
credential_observed_at = case
when $10::text is not null and enabled
when $10::text is not null
then $11 else credential_observed_at end,
last_error_reason = case
when $10::text is not null and enabled
when $10::text is not null
then $12 else last_error_reason end,
last_error_message = case
when $10::text is not null and enabled
when $10::text is not null
then $13 else last_error_message end,
updated_at = greatest(now(), updated_at, $11)
where id = $1 and credential_revision = $2
@@ -295,23 +295,11 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
update.validate()?;
let result = sqlx::query(
"update provider_accounts
set credential_state = case
when enabled then $3
else credential_state
end,
credential_observed_at = case
when enabled then $4
else credential_observed_at
end,
last_error_reason = case
when enabled then $5
else last_error_reason
end,
last_error_message = case
when enabled then $6
else last_error_message
end,
updated_at = case when enabled then greatest(now(), updated_at, $4) else updated_at end
set credential_state = $3,
credential_observed_at = $4,
last_error_reason = $5,
last_error_message = $6,
updated_at = greatest(now(), updated_at, $4)
where id = $1 and credential_revision = $2
and (credential_observed_at is null or credential_observed_at <= $4)",
)
@@ -834,16 +822,16 @@ pub(crate) async fn rotate_provider_account_in_transaction(
upstream_user_id = case when $11::boolean then $12::text else upstream_user_id end,
upstream_account_id = case when $11::boolean then $13::text else upstream_account_id end,
credential_state = case
when $15::boolean or not enabled then credential_state
when $15::boolean then credential_state
when $11::boolean or credential_state <> 'unknown' then 'ready'
else 'unknown'
end,
credential_observed_at = case
when $15::boolean or not enabled then credential_observed_at
when $15::boolean then credential_observed_at
else now()
end,
last_error_reason = case when enabled and not $15::boolean then null else last_error_reason end,
last_error_message = case when enabled and not $15::boolean then null else last_error_message end,
last_error_reason = case when not $15::boolean then null else last_error_reason end,
last_error_message = case when not $15::boolean then null else last_error_message end,
updated_at = greatest(now(), updated_at)
where id = $1 and provider_kind = $2
and credential_revision = $3
@@ -425,7 +425,6 @@ pub(crate) const REFRESH_CANDIDATES_SELECT: &str = "select auto_location, detect
left join (select id as location_proxy_id, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location
on outbound_proxy_id = location_proxy_id
where provider_kind = $1
and enabled
and has_refresh_token
and credential_state in ('unknown', 'ready')
and access_token_expires_at is not null
@@ -31,6 +31,7 @@ pub struct RuntimeSettings {
pub max_waiting_per_account: u32,
pub concurrency_wait_timeout_seconds: u32,
pub responses_max_decompressed_body_bytes: u64,
pub smart_scheduling: gateway_core::account::SmartSchedulingConfig,
pub rotation_strategy: String,
pub request_location_enabled: bool,
pub request_location: gateway_core::account::RequestLocation,
@@ -133,6 +134,7 @@ pub struct RuntimeSettingsUpdate {
pub max_waiting_per_account: u32,
pub concurrency_wait_timeout_seconds: u32,
pub responses_max_decompressed_body_bytes: u64,
pub smart_scheduling: gateway_core::account::SmartSchedulingConfig,
pub rotation_strategy: String,
pub request_location_enabled: bool,
pub request_location: gateway_core::account::RequestLocation,
@@ -263,7 +265,7 @@ pub(crate) async fn load_runtime_settings_from_pool(pool: &PgPool) -> StoreResul
let row = sqlx::query_as::<_, RuntimeSettingsRow>(
"select provider_request_profiles_json, config_revision, admin_api_key, refresh_margin_seconds, request_location_json, request_location_enabled,
refresh_concurrency, max_concurrent_per_account, request_interval_ms,
rotation_strategy, model_mappings_json, usage_retention_days, ops_event_retention_days,
rotation_strategy, smart_scheduling_json, model_mappings_json, usage_retention_days, ops_event_retention_days,
audit_retention_days, min_codex_desktop_version,
min_codex_cli_version, updated_at, responses_max_decompressed_body_bytes, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds,
account_auto_freeze_enabled, account_auto_freeze_threshold,
@@ -415,7 +417,7 @@ pub(crate) async fn load_runtime_settings_in_transaction(
let row = sqlx::query_as::<_, RuntimeSettingsRow>(
"select provider_request_profiles_json, config_revision, admin_api_key, refresh_margin_seconds, request_location_json, request_location_enabled,
refresh_concurrency, max_concurrent_per_account, request_interval_ms,
rotation_strategy, model_mappings_json, usage_retention_days, ops_event_retention_days,
rotation_strategy, smart_scheduling_json, model_mappings_json, usage_retention_days, ops_event_retention_days,
audit_retention_days, min_codex_desktop_version,
min_codex_cli_version, updated_at, responses_max_decompressed_body_bytes, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds,
account_auto_freeze_enabled, account_auto_freeze_threshold,
@@ -492,6 +494,7 @@ pub(crate) async fn update_runtime_settings_in_transaction(
account_warmup_enabled = $28,
account_warmup_schedule_time = $29,
account_warmup_model = $30,
smart_scheduling_json = $31,
updated_at = now()
where id = 1
returning config_revision",
@@ -538,6 +541,7 @@ pub(crate) async fn update_runtime_settings_in_transaction(
.bind(update.account_warmup_enabled)
.bind(&update.account_warmup_schedule_time)
.bind(update.account_warmup_model.as_deref())
.bind(sqlx::types::Json(update.smart_scheduling))
.fetch_optional(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("update runtime settings in transaction"))?
@@ -596,6 +600,7 @@ struct RuntimeSettingsRow {
refresh_concurrency: i64,
max_concurrent_per_account: i64,
request_interval_ms: i64,
smart_scheduling_json: sqlx::types::Json<gateway_core::account::SmartSchedulingConfig>,
rotation_strategy: String,
request_location_enabled: bool,
request_location_json: sqlx::types::Json<gateway_core::account::RequestLocation>,
@@ -644,6 +649,7 @@ fn runtime_settings_from_row(row: RuntimeSettingsRow) -> StoreResult<RuntimeSett
refresh_concurrency: to_u32(row.refresh_concurrency)?,
max_concurrent_per_account: to_u32(row.max_concurrent_per_account)?,
request_interval_ms: to_u64(row.request_interval_ms)?,
smart_scheduling: row.smart_scheduling_json.0,
rotation_strategy: row.rotation_strategy,
request_location_enabled: row.request_location_enabled,
request_location: row
@@ -33,6 +33,7 @@ pub struct SnapshotRuntimeSettings {
pub max_waiting_per_account: u32,
pub concurrency_wait_timeout_seconds: u32,
pub responses_max_decompressed_body_bytes: u64,
pub smart_scheduling: gateway_core::account::SmartSchedulingConfig,
pub rotation_strategy: String,
pub model_mappings: BTreeMap<String, String>,
pub min_codex_desktop_version: Option<String>,
@@ -162,6 +163,7 @@ impl SnapshotStorePort for PgRuntimeSnapshotRepository {
.with_responses_max_decompressed_body_bytes(
data.settings.responses_max_decompressed_body_bytes,
)
.with_smart_scheduling(data.settings.smart_scheduling)
.with_request_profiles(data.settings.request_profiles)
.with_pricing(data.settings.pricing)
.with_request_location(
@@ -254,6 +256,7 @@ struct SnapshotSettingsRow {
refresh_concurrency: i64,
max_concurrent_per_account: i64,
request_interval_ms: i64,
smart_scheduling_json: sqlx::types::Json<gateway_core::account::SmartSchedulingConfig>,
rotation_strategy: String,
model_mappings_json: sqlx::types::Json<BTreeMap<String, String>>,
min_codex_desktop_version: Option<String>,
@@ -272,7 +275,7 @@ async fn load_settings(
transaction: &mut Transaction<'_, Postgres>,
) -> StoreResult<(Revision, SnapshotRuntimeSettings)> {
let row = sqlx::query_as::<_, SnapshotSettingsRow>(
"select config_revision, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, model_mappings_json, min_codex_desktop_version, min_codex_cli_version, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds, request_location_json, request_location_enabled, responses_max_decompressed_body_bytes, provider_request_profiles_json, pricing_overrides_json, pricing_synced_json from runtime_settings where id = 1",
"select config_revision, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, smart_scheduling_json, model_mappings_json, min_codex_desktop_version, min_codex_cli_version, max_waiting_per_key, max_waiting_per_account, concurrency_wait_timeout_seconds, request_location_json, request_location_enabled, responses_max_decompressed_body_bytes, provider_request_profiles_json, pricing_overrides_json, pricing_synced_json from runtime_settings where id = 1",
)
.fetch_optional(&mut **transaction)
.await
@@ -302,6 +305,7 @@ async fn load_settings(
refresh_concurrency: to_u32(row.refresh_concurrency)?,
max_concurrent_per_account: to_u32(row.max_concurrent_per_account)?,
request_interval_ms: to_u64(row.request_interval_ms)?,
smart_scheduling: row.smart_scheduling_json.0,
rotation_strategy: row.rotation_strategy,
model_mappings: row.model_mappings_json.0,
min_codex_desktop_version: row.min_codex_desktop_version,
@@ -247,8 +247,10 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
"plugin_artifact_credentials",
"plugin_artifact_platforms",
"plugin_artifacts",
"plugin_group_resources",
"plugin_instance_secrets",
"plugin_instances",
"plugin_key_resources",
"plugin_source_credentials",
"plugin_state_generations",
"plugin_state_records",
@@ -114,11 +114,16 @@ async fn refresh_candidates_should_filter_order_and_bound_in_one_query() {
disabled.has_refresh_token = true;
disabled.access_token_expires_at = Some(now);
disabled.enabled = false;
let mut expired = account("acct_refresh_expired", "user-expired");
expired.enabled = false;
expired.has_refresh_token = true;
expired.access_token_expires_at = Some(now - TimeDelta::hours(4));
expired.credential_state = CredentialState::Expired;
let mut other_provider = account("acct_refresh_other", "user-other");
other_provider.provider_kind = "xai".to_owned();
other_provider.has_refresh_token = true;
other_provider.access_token_expires_at = Some(now);
for candidate in [forced, due, retry_later, disabled, other_provider] {
for candidate in [forced, due, retry_later, disabled, expired, other_provider] {
repository
.insert_provider_account(candidate)
.await
@@ -148,7 +153,7 @@ async fn refresh_candidates_should_filter_order_and_bound_in_one_query() {
.iter()
.map(|candidate| candidate.account.id().as_str())
.collect::<Vec<_>>(),
vec!["acct_refresh_forced", "acct_refresh_due"]
vec!["acct_refresh_forced", "acct_refresh_disabled"]
);
assert!(candidates.iter().all(|candidate| {
candidate.credential.expose_to_provider()["access_token"] == "initial-secret"
@@ -2935,7 +2940,83 @@ async fn xai_resettable_usage_limit_state_is_persisted() {
}
#[tokio::test]
async fn disabled_account_preserves_user_state_during_refresh_writes() {
async fn core_refresh_cas_updates_credentials_after_scheduling_is_disabled() {
let Some(database) = TestDatabase::create("disabled_core_refresh").await else {
return;
};
let repository = PgProviderAccountRepository::new(database.pool.clone());
let account_id = ProviderAccountId::new("acct_disabled_core_refresh").unwrap();
repository
.insert_provider_account(account(account_id.as_str(), "user-disabled-core-refresh"))
.await
.unwrap();
let revision = CredentialRevision::new(1).unwrap();
repository
.apply_state_change(AccountStateChange {
account_id: account_id.clone(),
expected_revision: revision,
credential_state: CredentialState::Ready,
observed_at: SystemTime::now(),
error_reason: Some(AccountErrorReason::AccessTokenExpired),
message: Some("previous refresh failure".to_owned()),
})
.await
.unwrap();
let quota = QuotaState::exhausted(QuotaEvidence::UsageLimitReached, SystemTime::now(), None);
repository
.apply_quota_access(QuotaAccessChange {
account_id: account_id.clone(),
expected_revision: revision,
state: quota,
})
.await
.unwrap();
let refreshed = CredentialCasUpdate::new(
account_id.clone(),
revision,
ProviderAccountUpdate {
account_id: account_id.clone(),
name: "unused profile".to_owned(),
email: None,
plan_type: None,
},
plaintext_credential("refreshed-secret"),
true,
Some(SystemTime::now() + Duration::from_secs(3_600)),
None,
)
.unwrap()
.preserving_profile()
.with_account_state(CredentialState::Ready, SystemTime::now(), None, None);
// 模拟刷新在途时停用调度;提交新凭据不能重新启用账号,也不能丢弃刷新结果。
repository.set_enabled(&account_id, false).await.unwrap();
assert!(matches!(
repository
.compare_and_swap_credential(refreshed)
.await
.unwrap(),
CredentialCasOutcome::Updated(_)
));
let current = repository
.load_current_credential(&account_id)
.await
.unwrap();
assert!(!current.account.enabled());
assert_eq!(current.account.name(), account_id.as_str());
assert_eq!(current.account.credential_state(), CredentialState::Ready);
assert_eq!(current.account.last_error_reason(), None);
assert_eq!(current.account.last_error_message(), None);
assert!(current.account.quota().is_exhausted());
assert_eq!(
current.credential.expose_to_provider()["access_token"],
"refreshed-secret"
);
assert_eq!(current.account.revision().get(), 2);
database.close().await;
}
#[tokio::test]
async fn disabled_account_updates_credentials_without_enabling_scheduling() {
let Some(database) = TestDatabase::create("provider_account_disabled_refresh").await else {
return;
};
@@ -2946,7 +3027,7 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
let account_id = ProviderAccountId::new("acct_disabled_refresh").expect("account ID");
let mut disabled = account(account_id.as_str(), "user-disabled-refresh");
disabled.enabled = false;
disabled.credential_state = CredentialState::Expired;
disabled.credential_state = CredentialState::Ready;
repository
.import_provider_accounts(ImportProviderAccounts {
settings: None,
@@ -2964,15 +3045,28 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
repository
.apply_state_change(AccountStateChange {
message: None,
message: Some("refresh token expired".to_owned()),
account_id: account_id.clone(),
expected_revision: CredentialRevision::new(1).expect("credential revision"),
credential_state: CredentialState::Ready,
credential_state: CredentialState::Expired,
observed_at: SystemTime::now(),
error_reason: None,
error_reason: Some(AccountErrorReason::CredentialExpired),
})
.await
.expect("disabled state write is a no-op");
.expect("record rejected credential while scheduling is disabled");
let rejected = repository
.load_current_credential(&account_id)
.await
.unwrap();
assert!(!rejected.account.enabled());
assert_eq!(
rejected.account.credential_state(),
CredentialState::Expired
);
assert_eq!(
rejected.account.last_error_message(),
Some("refresh token expired")
);
repository
.rotate_provider_account(RotateProviderAccount {
settings: None,
@@ -2998,7 +3092,7 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
.await
.expect("load disabled account after refresh writes");
assert!(!current.0);
assert_eq!(current.1, "expired");
assert_eq!(current.1, "ready");
assert_eq!(current.2["access_token"], "disabled-refreshed-secret");
assert_eq!(current.3, 2);
@@ -21,6 +21,7 @@ fn settings_with_margin(refresh_margin_seconds: u64) -> RuntimeSettingsUpdate {
max_waiting_per_account: 0,
concurrency_wait_timeout_seconds: 30,
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
smart_scheduling: gateway_core::account::SmartSchedulingConfig::default(),
rotation_strategy: "smart".to_owned(),
model_mappings: BTreeMap::from([
("gpt-5.4".to_owned(), "gpt-5.5".to_owned()),
@@ -50,6 +51,39 @@ fn runtime_settings_keep_account_rotation_global() {
assert!(settings.validate().is_ok());
}
#[tokio::test]
async fn smart_settings_upgrade_preserves_selection_and_publishes_custom_config() {
use gateway_core::account::SmartSchedulingConfig;
use gateway_store::postgres::{PgRuntimeSnapshotRepository, RuntimeSnapshotRepository};
let Some(database) = TestDatabase::create_through("smart_config", 18).await else {
return;
};
// 升级前不能用包含新列的 Repository,直接写入旧版本已有字段。
sqlx::query("update runtime_settings set rotation_strategy = 'sticky', refresh_margin_seconds = 3600 where id = 1")
.execute(&database.pool).await.unwrap();
super::TEST_MIGRATOR.run(&database.pool).await.unwrap();
let repository = PgRuntimeSettingsRepository::new(database.pool.clone());
let before = repository.load_runtime_settings().await.unwrap();
assert_eq!(before.rotation_strategy, "sticky");
assert_eq!(before.refresh_margin_seconds, 3600);
assert_eq!(before.smart_scheduling, SmartSchedulingConfig::default());
let mut update = settings_with_margin(3600);
update.smart_scheduling =
SmartSchedulingConfig::new([0.0, 2.0, 1.0, 0.5, 1.2, 2.3], true).unwrap();
let expected = update.smart_scheduling;
repository.update_runtime_settings(update).await.unwrap();
let reloaded = repository.load_runtime_settings().await.unwrap();
assert_eq!(reloaded.smart_scheduling, expected);
let snapshot = PgRuntimeSnapshotRepository::new(database.pool.clone())
.load_runtime_snapshot()
.await
.unwrap();
assert_eq!(snapshot.settings.smart_scheduling, expected);
assert!(snapshot.config_revision > before.config_revision);
assert_eq!(before.smart_scheduling, SmartSchedulingConfig::default());
database.close().await;
}
#[test]
fn runtime_settings_require_model_when_warmup_is_enabled() {
let settings = RuntimeSettingsUpdate {
@@ -112,7 +112,13 @@ impl CodexPlanCatalog {
pub struct CodexCredentialCatalogSnapshot {
observed_at: SystemTime,
models: Vec<CodexCatalogModel>,
scope_models: BTreeMap<CodexCatalogScope, Vec<String>>,
scope_catalogs: BTreeMap<CodexCatalogScope, CatalogScopeSnapshot>,
}
#[derive(Clone, PartialEq, Eq)]
struct CatalogScopeSnapshot {
models: Vec<String>,
accounts: BTreeSet<ProviderAccountId>,
}
impl CodexCredentialCatalogSnapshot {
@@ -131,7 +137,25 @@ impl CodexCredentialCatalogSnapshot {
account: &ProviderAccount,
) -> Result<Option<&[String]>, CodexCredentialCatalogError> {
let scope = CodexCatalogScope::for_account(account)?;
Ok(self.scope_models.get(&scope).map(Vec::as_slice))
Ok(self
.scope_catalogs
.get(&scope)
.map(|catalog| catalog.models.as_slice()))
}
/// 将 Provider 私有目录作用域投影为模型来源账号,供 Core 结合冻结政策过滤。
#[must_use]
pub fn model_catalog_accounts(&self) -> BTreeMap<&str, BTreeSet<ProviderAccountId>> {
let mut accounts = BTreeMap::<&str, BTreeSet<ProviderAccountId>>::new();
for catalog in self.scope_catalogs.values() {
for model in &catalog.models {
accounts
.entry(model.as_str())
.or_default()
.extend(catalog.accounts.iter().cloned());
}
}
accounts
}
}
@@ -141,7 +165,7 @@ impl fmt::Debug for CodexCredentialCatalogSnapshot {
.debug_struct("CodexCredentialCatalogSnapshot")
.field("observed_at", &self.observed_at)
.field("model_count", &self.models.len())
.field("scope_count", &self.scope_models.len())
.field("scope_count", &self.scope_catalogs.len())
.finish()
}
}
@@ -585,10 +609,14 @@ impl CodexCredentialCatalogService {
);
let mut union = BTreeMap::<String, CodexCatalogModel>::new();
let mut union_order = Vec::new();
let mut scope_models = BTreeMap::new();
let mut scope_catalogs = BTreeMap::new();
let mut etags = Vec::new();
let mut last_error = None;
for (scope, candidates) in groups {
let accounts = candidates
.iter()
.map(|account| account.id().clone())
.collect();
let fetched = match self.fetch_scope_models(&client, candidates).await {
Ok(fetched) => fetched,
Err(error) => {
@@ -620,7 +648,13 @@ impl CodexCredentialCatalogService {
}
}
}
scope_models.insert(scope, entitlement);
scope_catalogs.insert(
scope,
CatalogScopeSnapshot {
models: entitlement,
accounts,
},
);
etags.extend(fetched.etag);
}
if union.is_empty()
@@ -635,7 +669,7 @@ impl CodexCredentialCatalogService {
.into_iter()
.filter_map(|id| union.remove(&id))
.collect(),
scope_models,
scope_catalogs,
};
Ok(FetchedCatalog { snapshot, etags })
}
@@ -1016,7 +1050,7 @@ fn same_catalog(
left: &CodexCredentialCatalogSnapshot,
right: &CodexCredentialCatalogSnapshot,
) -> bool {
left.models == right.models && left.scope_models == right.scope_models
left.models == right.models && left.scope_catalogs == right.scope_catalogs
}
fn validate_response_etag(etag: &str) -> Result<(), CodexCredentialCatalogError> {
@@ -235,9 +235,6 @@ impl CodexCredentialRepository {
error_reason: Option<AccountErrorReason>,
message: Option<String>,
) -> Result<(), CredentialRepositoryError> {
if !account.enabled() {
return Ok(());
}
let credential_state = oauth_account_state(account, credential_state);
let message = message.filter(|value| !value.trim().is_empty());
let error_reason = if credential_state == CredentialState::Ready {
@@ -639,16 +639,24 @@ impl CodexCredentialSelector {
);
let Some(selection) = selection else {
if !diagnostic && queue_policy.max_waiting > 0 && !wait_candidates.is_empty() {
waiting.wait(&wait_candidates).await.map_err(|error| {
tracing::info!(
request_id = request.attempt.request_id().as_str(),
queue_layer = "account",
queue_wait_ms = waiting.elapsed().as_millis() as u64,
reason = %error,
"OpenAI 账号排队请求被拒绝"
);
CredentialSelectionError::QueueRejected(error)
})?;
AccountSelector
.wait_for_capacity(
&mut waiting,
&wait_candidates,
&candidates,
&wait_context,
)
.await
.map_err(|error| {
tracing::info!(
request_id = request.attempt.request_id().as_str(),
queue_layer = "account",
queue_wait_ms = waiting.elapsed().as_millis() as u64,
reason = %error,
"OpenAI 账号排队请求被拒绝"
);
CredentialSelectionError::QueueRejected(error)
})?;
continue 'capacity;
}
// 只判断本次账号范围;空池、认证失效和租约失败不能伪装成额度耗尽。
@@ -79,7 +79,9 @@ use crate::transport::protocol::responses::{
CodexResponsesRequest, PREVIOUS_RESPONSE_NOT_FOUND_CODE, PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE,
PreviousResponseScope, ResponseEventSignals, TransportRequirement, transport_requirement,
};
use crate::transport::protocol::websocket::WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE;
use crate::transport::protocol::websocket::{
WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE, websocket_response_create_payload_len,
};
use crate::transport::request::{
CodexRequestEncodeError, RequestAccountScope, align_structured_location_fields,
encode_generate_request, scope_request_to_account,
@@ -115,6 +117,8 @@ const PROVIDER_NAME: &str = "openai";
const HTTP_SSE_TRANSPORT: &str = "http_sse";
const HTTP_JSON_TRANSPORT: &str = "http_json";
const WEBSOCKET_TRANSPORT: &str = "websocket";
// 在已观测到的 Codex OAuth 上游 16 MiB 附近消息边界前留出传输 metadata 余量。
const WEBSOCKET_HTTP_FALLBACK_THRESHOLD_BYTES: usize = 15 * 1024 * 1024;
const MAX_COOKIE_HEADER_BYTES: usize = 16 * 1024;
/// 提交边界前最多保留 64 KiB 原始上游 chunk;达到阈值后结束无感换号窗口,
/// 但不会把上游数据改写成协议失败。
@@ -392,10 +396,17 @@ impl Provider for CodexProvider {
let snapshot = self.catalog.synchronize().await.map_err(|_| {
provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
})?;
let mut accounts = snapshot.model_catalog_accounts();
Ok(snapshot
.models()
.iter()
.map(compile_model_capabilities)
.map(|model| {
compile_model_capabilities(model).with_catalog_accounts(
accounts
.remove(model.request_model().as_str())
.unwrap_or_default(),
)
})
.collect())
}
@@ -694,15 +705,29 @@ impl CodexProvider {
{
return Err(continuation_replay_required_error("scope_unavailable"));
}
if upstream_request
.downstream_websocket_connection_id
.is_some()
&& upstream_request.previous_response_id().is_some()
&& upstream_request.generate()
&& previous_session.as_ref().is_some_and(|state| {
state.continuation_scope == OpenAiContinuationScope::ReplayRequired
})
{
// HTTP store=false 没有原生续链,完整历史仍由客户端持有;不能把 delta
// 当作独立新请求发送,也不能用 previous_response_id 猜测上游存储状态。
return Err(continuation_replay_required_error("scope_unavailable"));
}
scope_request_to_account(
&mut upstream_request,
lease.installation_id(),
account_scope,
);
if matches!(
let codex_oauth = matches!(
lease.authentication(),
crate::credential::CodexRuntimeAuthentication::OAuth(_)
) {
);
if codex_oauth {
normalize_selected_codex_downstream_body(
upstream_request.body_mut(),
generate.protocol_payload().context(),
@@ -736,13 +761,49 @@ impl CodexProvider {
&& session_affinity
.as_ref()
.is_some_and(|affinity| self.session_transport_recovery.uses_http(affinity.key()));
let transport = if requirement.requires_websocket() {
let mut transport = if requirement.requires_websocket() {
CodexProviderTransport::PreferWebSocket
} else if context.transport() == AttemptTransport::Fallback || session_http_fallback {
CodexProviderTransport::HttpOnly
} else {
requested_transport
};
if codex_oauth
&& transport == CodexProviderTransport::PreferWebSocket
&& (matches!(
requirement,
TransportRequirement::NewChain | TransportRequirement::WebSocketNewChain
) || (requirement == TransportRequirement::ExactWebSocketContinuation
&& upstream_request
.downstream_websocket_connection_id
.is_some()))
{
let payload_bytes =
websocket_response_create_payload_len(&upstream_request).map_err(|_| {
provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent)
})?;
if payload_bytes >= WEBSOCKET_HTTP_FALLBACK_THRESHOLD_BYTES {
if requirement == TransportRequirement::ExactWebSocketContinuation {
tracing::info!(
request_id = context.request_id().as_str(),
payload_bytes,
threshold_bytes = WEBSOCKET_HTTP_FALLBACK_THRESHOLD_BYTES,
continuation_recovery_action = "client_replay_required",
"OpenAI oversized WebSocket continuation requires full client input before send"
);
return Err(continuation_replay_required_error(
"transport_payload_too_large",
));
}
transport = CodexProviderTransport::HttpOnly;
tracing::info!(
request_id = context.request_id().as_str(),
payload_bytes,
threshold_bytes = WEBSOCKET_HTTP_FALLBACK_THRESHOLD_BYTES,
"OpenAI oversized new-chain request selected HTTP before WebSocket send"
);
}
}
if transport == CodexProviderTransport::PreferWebSocket
&& middleware_headers.iter().any(|header| {
match HeaderValue::from_bytes(header.value()) {
@@ -121,7 +121,8 @@ pub enum TransportRequirement {
HttpRequired,
/// `generate=false + store=false` 预热必须保留在同一条 WebSocket。
ExplicitWebSocketWarmup,
/// 客户端 WebSocket 的非持久化新链,必须在池化连接上建立后续续接状态。
/// 客户端 WebSocket 的非持久化新链,默认在池化连接上建立后续续接状态。
/// Provider 可在发送前对 OAuth 大新链选择 HTTP,后续由客户端完整重放。
WebSocketNewChain,
/// 只能使用持有指定 connection-local response 的精确 WebSocket。
ExactWebSocketContinuation,
@@ -134,7 +135,7 @@ pub enum TransportRequirement {
}
impl TransportRequirement {
/// 是否必须使用 WebSocket,且禁止 HTTP fallback。
/// 默认是否要求 WebSocket;Provider 的 OAuth 大新链预检可在发送前选择 HTTP。
pub fn requires_websocket(self) -> bool {
matches!(
self,
@@ -189,8 +190,8 @@ pub fn transport_requirement(request: &CodexResponsesRequest) -> TransportRequir
TransportRequirement::ExternalUnknown
}
},
// 客户端会在下一轮提交 response ID;HTTP store=false 的成功响应无法
// 在池化 WebSocket 上续接,所以首轮也不能按普通快路径降级到 HTTP。
// HTTP store=false 的成功响应无法在池化 WebSocket 上续接,默认不走 HTTP 快路径;
// Provider 的 OAuth 大新链预检例外通过客户端完整重放恢复后续请求。
None if request.downstream_websocket_connection_id.is_some() && !request.store() => {
TransportRequirement::WebSocketNewChain
}
@@ -225,6 +225,35 @@ pub fn websocket_response_create_payload_text(
})
}
pub(crate) fn websocket_response_create_payload_len(
request: &CodexResponsesRequest,
) -> Result<usize, serde_json::Error> {
struct ByteCount(usize);
impl std::io::Write for ByteCount {
fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
self.0 = self
.0
.checked_add(bytes.len())
.ok_or_else(|| std::io::Error::other("WebSocket payload size overflow"))?;
Ok(bytes.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
let mut count = ByteCount(0);
serde_json::to_writer(
&mut count,
&ResponseCreateFrame {
body: request.body(),
},
)?;
Ok(count.0)
}
fn websocket_payload_keys(request: &CodexResponsesRequest) -> Vec<String> {
let mut keys = Vec::with_capacity(request.body().len() + 1);
keys.push("type".to_string());
@@ -1463,3 +1463,34 @@ async fn disabled_account_can_export_native_catalog_without_a_cached_snapshot()
assert_eq!(models[0].document().protocol(), "codex");
server.verify().await;
}
#[tokio::test]
async fn catalog_account_membership_changes_advance_generation_without_model_changes() {
use std::collections::BTreeSet;
let store = Arc::new(MemoryAccountStore::default());
let first = seed_account(&store, "acct_source_a").await;
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/codex/models"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"models": [{"slug": "gpt-5.4", "display_name": "GPT-5.4"}]
})))
.expect(2)
.mount(&server)
.await;
let service = service_with_catalog_cache(&store, server.uri(), catalog_cache());
let initial = service.synchronize().await.unwrap();
let generation = service.catalog_generation();
let second = seed_account(&store, "acct_source_b").await;
let refreshed = service.refresh_catalogs().await.unwrap();
assert_eq!(initial.models(), refreshed.models());
assert!(service.catalog_generation() > generation);
assert_eq!(
initial.model_catalog_accounts()["gpt-5.4"],
BTreeSet::from([first.id().clone()])
);
assert_eq!(
refreshed.model_catalog_accounts()["gpt-5.4"],
BTreeSet::from([first.id().clone(), second.id().clone()])
);
}
@@ -124,6 +124,80 @@ fn round_robin_attempt() -> AttemptContext {
)
}
#[tokio::test]
async fn smart_reselection_respects_native_and_required_account_boundaries() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_original", "test-original");
create_account(&store, "acct_primary", "test-primary");
store.set_scheduling("acct_original", None, AccountWeight::new(10).unwrap());
store.set_scheduling("acct_primary", None, AccountWeight::new(100).unwrap());
let original_id = ProviderAccountId::new("acct_original").unwrap();
let provider = ProviderKind::new("openai").unwrap();
let affinity = Arc::new(MemorySessionAffinity::default());
let key = ProviderSessionAffinityKey::try_new("smart-config-test").unwrap();
affinity
.bind(&provider, &key, &original_id, Duration::from_secs(60))
.await
.unwrap();
let selector =
selector_with_affinity(&store, Arc::new(TestLeaseCoordinator::default()), affinity);
let request_url = Url::parse(OFFICIAL_CODEX_BASE_URL).unwrap();
for enabled in [false, true] {
for binding in ["soft", "required", "native"] {
let continuation = (binding == "native").then(|| {
ContinuationBinding::Pinned(NativeContinuationPin::new(
PreviousResponseId::new("client_response"),
PreviousResponseId::new("upstream_response"),
ClientApiKeyId::new("key_codex_contract").unwrap(),
provider.clone(),
original_id.clone(),
))
});
let attempt = AttemptContext::new(
RequestAttemptContext::new(
ModelRequestId::new("req_smart_scope").unwrap(),
ClientApiKeyId::new("key_codex_contract").unwrap(),
),
NonZeroU32::new(1).unwrap(),
SystemTime::now() + Duration::from_secs(30),
account_policy().with_smart_scheduling(
gateway_core::account::SmartSchedulingConfig::new(
[1.0, 0.8, 1.0, 0.5, 1.0, 1.0],
enabled,
)
.unwrap(),
),
AccountAttemptContext::new(
BTreeSet::new(),
(binding == "required").then(|| original_id.clone()),
None,
)
.with_account_scope(contract_account_scope()),
continuation,
CancellationToken::new(),
);
let selected = selector
.select(&SelectCodexCredential {
upstream_model: "gpt-5.4",
request_url: &request_url,
attempt: &attempt,
session_affinity_key: Some(&key),
})
.await
.unwrap();
assert_eq!(
selected.account_id().as_str(),
if enabled && binding == "soft" {
"acct_primary"
} else {
"acct_original"
},
"{binding}, switchback={enabled}"
);
}
}
}
fn selector(
store: &Arc<MemoryAccountStore>,
leases: Arc<TestLeaseCoordinator>,
@@ -1776,7 +1850,14 @@ fn queued_attempt(timeout: Duration) -> AttemptContext {
.with_queue(gateway_core::concurrency::ConcurrencyQueuePolicy {
max_waiting: 1,
timeout,
}),
})
.with_smart_scheduling(
gateway_core::account::SmartSchedulingConfig::new(
[1.0, 0.8, 1.0, 0.5, 1.0, 1.0],
false,
)
.unwrap(),
),
AccountAttemptContext::new(BTreeSet::new(), None, None)
.with_account_scope(contract_account_scope()),
None,
@@ -1784,6 +1865,155 @@ fn queued_attempt(timeout: Duration) -> AttemptContext {
)
}
#[test]
fn smart_queue_weight_changes_the_selected_wait_queue() {
use futures::FutureExt;
for queue_weight in [0.0, 0.2] {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "test-healthy");
create_account(&store, "acct_other", "test-unhealthy");
let leases = Arc::new(TestLeaseCoordinator::default());
*leases.busy.lock().unwrap() = true;
let feedback = Arc::new(AccountFeedbackStats::default());
for _ in 0..4 {
feedback.report(
&ProviderKind::new("openai").unwrap(),
&ProviderAccountId::new("acct_other").unwrap(),
AccountAttemptFeedback::Failed {
first_output_ms: None,
},
);
}
let selector = selector_with_runtime(
&store,
leases,
Arc::new(MemorySessionAffinity::default()),
feedback,
Arc::new(MemoryCooldownPort::new()),
);
let make_attempt = |required| {
AttemptContext::new(
RequestAttemptContext::new(
ModelRequestId::new("req_queue_weight").unwrap(),
ClientApiKeyId::new("key_codex_contract").unwrap(),
),
NonZeroU32::new(1).unwrap(),
SystemTime::now() + Duration::from_secs(5),
account_policy()
.with_queue(gateway_core::concurrency::ConcurrencyQueuePolicy {
max_waiting: 2,
timeout: Duration::from_secs(2),
})
.with_smart_scheduling(
gateway_core::account::SmartSchedulingConfig::new(
[0.0, 0.0, 1.0, 0.0, 0.0, queue_weight],
false,
)
.unwrap(),
),
AccountAttemptContext::new(BTreeSet::new(), required, None)
.with_account_scope(contract_account_scope()),
None,
CancellationToken::new(),
)
};
let head_attempt = make_attempt(Some(ProviderAccountId::new("acct_primary").unwrap()));
let next_attempt = make_attempt(None);
let url = Url::parse(OFFICIAL_CODEX_BASE_URL).unwrap();
let head_request = SelectCodexCredential {
upstream_model: "gpt-5.4",
request_url: &url,
attempt: &head_attempt,
session_affinity_key: None,
};
let next_request = SelectCodexCredential {
attempt: &next_attempt,
..head_request
};
let mut head = Box::pin(selector.select(&head_request));
assert!(head.as_mut().now_or_never().is_none());
let mut next = Box::pin(selector.select(&next_request));
assert!(next.as_mut().now_or_never().is_none());
let mut probe = Box::pin(selector.select(&head_request));
match probe.as_mut().now_or_never() {
Some(Err(CredentialSelectionError::QueueRejected(
gateway_core::concurrency::QueueRejection::Full,
))) => assert!(queue_weight > 0.0),
None => assert_eq!(queue_weight, 0.0),
other => panic!("unexpected queue probe: {other:?}"),
}
}
}
#[test]
fn smart_queue_keeps_native_and_required_account_pins_when_another_account_is_free() {
use futures::FutureExt;
for native in [false, true] {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_original", "test-pinned");
create_account(&store, "acct_primary", "test-free");
let original = ProviderAccountId::new("acct_original").unwrap();
let leases = Arc::new(TestLeaseCoordinator::default());
leases
.busy_accounts
.lock()
.unwrap()
.insert(original.clone());
let selector = selector(&store, leases.clone());
let attempt = AttemptContext::new(
RequestAttemptContext::new(
ModelRequestId::new("req_queue_pin").unwrap(),
ClientApiKeyId::new("key_codex_contract").unwrap(),
),
NonZeroU32::new(1).unwrap(),
SystemTime::now() + Duration::from_secs(5),
account_policy()
.with_queue(gateway_core::concurrency::ConcurrencyQueuePolicy {
max_waiting: 1,
timeout: Duration::from_secs(2),
})
.with_smart_scheduling(
gateway_core::account::SmartSchedulingConfig::new(
[1.0, 0.8, 1.0, 0.5, 10.0, 10.0],
true,
)
.unwrap(),
),
AccountAttemptContext::new(BTreeSet::new(), (!native).then(|| original.clone()), None)
.with_account_scope(contract_account_scope()),
native.then(|| {
ContinuationBinding::Pinned(NativeContinuationPin::new(
PreviousResponseId::new("client_response"),
PreviousResponseId::new("upstream_response"),
ClientApiKeyId::new("key_codex_contract").unwrap(),
ProviderKind::new("openai").unwrap(),
original.clone(),
))
}),
CancellationToken::new(),
);
let url = Url::parse(OFFICIAL_CODEX_BASE_URL).unwrap();
let request = SelectCodexCredential {
upstream_model: "gpt-5.4",
request_url: &url,
attempt: &attempt,
session_affinity_key: None,
};
let mut pending = Box::pin(selector.select(&request));
assert!(pending.as_mut().now_or_never().is_none());
leases.busy_accounts.lock().unwrap().clear();
assert_eq!(block_on(pending).unwrap().account_id(), &original);
assert!(
leases
.requests
.lock()
.unwrap()
.iter()
.all(|request| request.account_id() == &original)
);
}
}
#[test]
fn account_queue_is_bounded_and_resumes_after_capacity_is_released() {
use futures::FutureExt;
@@ -241,6 +241,94 @@ async fn seed_refreshable_account(
.await;
}
#[tokio::test]
async fn scheduled_refresh_rotates_tokens_without_enabling_scheduling() {
let store = Arc::new(MemoryAccountStore::default());
let policy = MutableRuntimePolicy::new(Duration::from_secs(5 * 60));
let refresher = SingleUseRefresher::new();
let service = refresh_service(&store, Arc::clone(&refresher), policy);
let account_id = "acct_disabled_refresh";
seed_refreshable_account(&store, account_id, SystemTime::now(), None).await;
let before = store.account(account_id).expect("seeded account");
store
.set_enabled(before.id(), false)
.await
.expect("disable scheduling");
let outcomes = service.refresh_due().await.expect("refresh cycle");
assert!(matches!(
outcomes.as_slice(),
[CodexCredentialRefreshOutcome::Refreshed { account_id: refreshed_id, .. }]
if refreshed_id == account_id
));
assert_eq!(refresher.calls(), 1);
let after = store.account(account_id).expect("refreshed account");
assert!(!after.enabled());
assert_eq!(after.quota(), before.quota());
assert_eq!(
after.status_projection(SystemTime::now(), None).status,
AccountStatus::Disabled
);
let loaded = store
.load_credential(after.id(), after.revision())
.await
.expect("refreshed credential");
let runtime = CodexCredentialCodec::decode(&loaded.credential).expect("runtime credential");
assert_eq!(
runtime
.authentication
.oauth()
.expect("OAuth credential")
.access_token
.expose_secret(),
"refreshed-access-token"
);
}
#[tokio::test]
async fn scheduled_refresh_stops_retrying_rejected_disabled_credentials() {
let store = Arc::new(MemoryAccountStore::default());
let service = CodexCredentialRefreshService::new(
store.repository(),
Arc::new(FailingRefresher {
failure: RefreshFailure::InvalidGrant {
message: None,
upstream: None,
},
}),
Arc::new(RefreshLeases),
Arc::new(RefreshCredentialState),
MutableRuntimePolicy::new(Duration::from_secs(5 * 60)),
);
let account_id = "acct_disabled_rejected";
seed_refreshable_account(&store, account_id, SystemTime::now(), None).await;
let account = store.account(account_id).expect("seeded account");
store
.set_enabled(account.id(), false)
.await
.expect("disable scheduling");
assert!(matches!(
service
.refresh_due()
.await
.expect("refresh cycle")
.as_slice(),
[CodexCredentialRefreshOutcome::Invalidated { .. }]
));
let account = store.account(account_id).expect("rejected account");
assert!(!account.enabled());
assert_eq!(account.credential_state(), CredentialState::Expired);
assert!(
service
.refresh_due()
.await
.expect("next refresh cycle")
.is_empty()
);
}
#[tokio::test]
async fn scheduled_refresh_uses_the_current_margin_without_persisting_a_normal_schedule() {
let store = Arc::new(MemoryAccountStore::default());
@@ -4804,6 +4804,551 @@ async fn websocket_upgrade_required_immediately_enables_session_http_fallback()
server.await.expect("upstream server");
}
#[tokio::test]
async fn oversized_http_new_chain_uses_http_without_disabling_session_websocket() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_provider_contract").await;
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(CAPTURE_COMPLETED_SSE),
)
.expect(1)
.mount(&server)
.await;
let input = json!([
{"type": "message", "role": "user", "content": [
{"type": "input_text", "text": "inspect the screenshot"},
{"type": "input_image", "image_url": format!("data:image/jpeg;base64,{}", "A".repeat(17 * 1024 * 1024))}
]},
{"type": "function_call", "call_id": "call_large", "name": "inspect", "arguments": "{}"},
{"type": "function_call_output", "call_id": "call_large", "output": "done"}
]);
let tools = json!([{"type": "function", "name": "inspect", "parameters": {"type": "object"}}]);
let body = Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), input.clone()),
("tools".to_owned(), tools.clone()),
("session_id".to_owned(), json!("oversized-http-session")),
("stream".to_owned(), json!(true)),
("store".to_owned(), json!(false)),
]);
let provider = provider_with_base_url(&store, server.uri());
let mut stream = Arc::clone(&provider)
.execute(
planned_request(
"openai",
Operation::Generate(GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object("openai", body).unwrap(),
)),
),
context("req_oversized_http", CancellationToken::new()),
)
.await
.expect("prepare oversized HTTP request");
assert_eq!(stream.metadata().transport().as_str(), "http_sse");
let mut completed = false;
while let Some(event) = stream.next().await {
let event = event.expect("oversized HTTP response");
completed |= event
.wire_event()
.is_some_and(|wire| wire.event_type() == Some("response.completed"));
}
assert!(completed);
drop(stream);
let requests = server.received_requests().await.unwrap();
assert_eq!(requests.len(), 1);
assert_eq!(requests[0].method.as_str(), "POST");
let sent = captured_request_body(&requests[0]);
assert_eq!(sent["input"], input);
assert_eq!(sent["tools"], tools);
assert!(sent.get("previous_response_id").is_none());
let next = provider
.execute(
planned_request(
"openai",
Operation::Generate(generate_with_session_context(
"oversized-http-session",
None,
None,
)),
),
context("req_small_after_oversized_http", CancellationToken::new()),
)
.await
.expect("prepare small request in the same session");
assert_eq!(next.metadata().transport().as_str(), "websocket");
}
#[tokio::test]
async fn oversized_http_threshold_counts_utf8_and_json_escaping() {
const THRESHOLD: usize = 15 * 1024 * 1024;
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_provider_contract").await;
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(CAPTURE_COMPLETED_SSE),
)
.expect(1)
.mount(&server)
.await;
let seed_text = "中文\n\"\\";
let mut body = Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
(
"input".to_owned(),
json!([{"type": "message", "role": "user", "content": [{"type": "input_text", "text": seed_text}]}]),
),
("session_id".to_owned(), json!("oversized-boundary-session")),
("stream".to_owned(), json!(true)),
("store".to_owned(), json!(false)),
]);
let provider = provider_with_base_url(&store, server.uri());
let mut seed = Arc::clone(&provider)
.execute(
planned_request(
"openai",
Operation::Generate(GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object("openai", body.clone())
.unwrap()
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(false))])),
)),
),
context("req_size_baseline", CancellationToken::new()),
)
.await
.unwrap();
while let Some(event) = seed.next().await {
event.expect("baseline HTTP response");
}
drop(seed);
let requests = server.received_requests().await.unwrap();
let mut frame = captured_request_body(&requests[0]);
frame["type"] = json!("response.create");
let baseline_len = serde_json::to_vec(&frame).unwrap().len();
for (payload_len, expected) in [(THRESHOLD - 1, "websocket"), (THRESHOLD, "http_sse")] {
let text = format!("{seed_text}{}", "x".repeat(payload_len - baseline_len));
body.get_mut("input").unwrap()[0]["content"][0]["text"] = json!(text);
let stream = Arc::clone(&provider)
.execute(
planned_request(
"openai",
Operation::Generate(GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object("openai", body.clone()).unwrap(),
)),
),
context(
&format!("req_boundary_{payload_len}"),
CancellationToken::new(),
),
)
.await
.expect("prepare boundary request");
assert_eq!(stream.metadata().transport().as_str(), expected);
}
assert_eq!(server.received_requests().await.unwrap().len(), 1);
}
#[tokio::test]
async fn oversized_ws_new_chain_uses_http_but_warmup_and_http_continuation_keep_websocket() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_provider_contract").await;
let server = MockServer::start().await;
let provider = provider_with_base_url(&store, server.uri());
let body = Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!([])),
(
"instructions".to_owned(),
json!("x".repeat(16 * 1024 * 1024)),
),
("store".to_owned(), json!(false)),
]);
for case in ["downstream_websocket", "warmup", "native_continuation"] {
let mut body = body.clone();
let mut request_context = Map::new();
let mut attempt = context(&format!("req_large_{case}"), CancellationToken::new());
match case {
"downstream_websocket" => {
request_context.insert(
"downstream_websocket_connection_id".to_owned(),
json!("ws_large"),
);
}
"warmup" => {
body.insert("generate".to_owned(), json!(false));
}
"native_continuation" => {
body.insert("previous_response_id".to_owned(), json!("client_previous"));
attempt = pinned_continuation_context(
"req_large_native_continuation",
"acct_provider_contract",
"client_previous",
"upstream_previous",
1,
ContinuationAttempt::Native,
);
}
_ => unreachable!(),
}
let mut generate = GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object("openai", body)
.unwrap()
.with_context(request_context),
);
if case == "native_continuation" {
generate = generate.with_provider_session_state(
ProviderSessionState::new(
"openai",
Map::from_iter([
("account_id".to_owned(), json!("acct_provider_contract")),
("continuation_scope".to_owned(), json!("connection_local")),
]),
)
.unwrap(),
);
}
let stream = Arc::clone(&provider)
.execute(
planned_request("openai", Operation::Generate(generate)),
attempt,
)
.await
.expect("prepare required WebSocket request");
let expected = if case == "downstream_websocket" {
"http_sse"
} else {
"websocket"
};
assert_eq!(stream.metadata().transport().as_str(), expected, "{case}");
}
assert!(server.received_requests().await.unwrap().is_empty());
}
#[tokio::test]
async fn oversized_exact_websocket_continuation_requests_client_replay_before_send() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_provider_contract").await;
let server = MockServer::start().await;
let provider = provider_with_base_url(&store, server.uri());
for input_bytes in [128, 16 * 1024 * 1024] {
let payload = ProtocolPayload::json_object("openai", Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!([{"type": "function_call_output", "call_id": "call_large", "output": "x".repeat(input_bytes)}])),
("previous_response_id".to_owned(), json!("client_previous")),
("store".to_owned(), json!(false)),
])).unwrap().with_context(Map::from_iter([
("downstream_websocket_connection_id".to_owned(), json!("ws_native_large")),
]));
let generate = GenerateRequest::from_protocol_payload(payload).with_provider_session_state(
ProviderSessionState::new(
"openai",
Map::from_iter([
("account_id".to_owned(), json!("acct_provider_contract")),
("continuation_scope".to_owned(), json!("connection_local")),
]),
)
.unwrap(),
);
let result = Arc::clone(&provider)
.execute(
planned_request("openai", Operation::Generate(generate)),
pinned_continuation_context(
&format!("req_native_size_{input_bytes}"),
"acct_provider_contract",
"client_previous",
"upstream_previous",
1,
ContinuationAttempt::Native,
),
)
.await;
if input_bytes == 128 {
assert_eq!(
result
.expect("small native continuation")
.metadata()
.transport()
.as_str(),
"websocket"
);
} else {
let error = result
.err()
.expect("large native continuation requires full client input");
assert_eq!(
error.kind(),
ProviderErrorKind::ContinuationRecoveryRequired
);
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
assert_eq!(
error.continuation_recovery_disposition(),
Some(ContinuationRecoveryDisposition::ClientReplayRequired)
);
assert_eq!(
error.client_visible_upstream_error().unwrap().code(),
Some("previous_response_not_found")
);
assert!(!error.allows_pre_delivery_retry());
}
}
assert!(server.received_requests().await.unwrap().is_empty());
}
#[tokio::test]
async fn oversized_api_key_new_chains_preserve_explicit_websocket_transport() {
let upstream = MockServer::start().await;
let store = Arc::new(MemoryAccountStore::default());
store
.seed_api_key(
"acct_provider_contract",
upstream.uri(),
provider_openai::credential::ResponsesTransport::PreferWebsocket,
)
.await;
let provider = provider_with_base_url(&store, upstream.uri());
for downstream_websocket in [false, true] {
let mut payload = ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!([])),
(
"instructions".to_owned(),
json!("x".repeat(16 * 1024 * 1024)),
),
("store".to_owned(), json!(false)),
]),
)
.unwrap();
if downstream_websocket {
payload = payload.with_context(Map::from_iter([(
"downstream_websocket_connection_id".to_owned(),
json!("ws_api_large"),
)]));
}
let stream = Arc::clone(&provider)
.execute(
planned_request(
"openai",
Operation::Generate(GenerateRequest::from_protocol_payload(payload)),
),
context(
&format!("req_large_api_new_{downstream_websocket}"),
CancellationToken::new(),
),
)
.await
.expect("API Key new chain preserves configured transport");
assert_eq!(stream.metadata().transport().as_str(), "websocket");
}
assert!(upstream.received_requests().await.unwrap().is_empty());
}
#[tokio::test]
async fn oversized_api_key_native_continuation_preserves_websocket_transport() {
let upstream = MockServer::start().await;
let store = Arc::new(MemoryAccountStore::default());
store
.seed_api_key(
"acct_provider_contract",
upstream.uri(),
provider_openai::credential::ResponsesTransport::PreferWebsocket,
)
.await;
let payload = ProtocolPayload::json_object("openai", Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!([{"type": "function_call_output", "call_id": "call_api_large", "output": "x".repeat(16 * 1024 * 1024)}])),
("previous_response_id".to_owned(), json!("client_api_previous")),
("store".to_owned(), json!(false)),
])).unwrap().with_context(Map::from_iter([
("downstream_websocket_connection_id".to_owned(), json!("ws_api_large")),
]));
let generate = GenerateRequest::from_protocol_payload(payload).with_provider_session_state(
ProviderSessionState::new(
"openai",
Map::from_iter([
("account_id".to_owned(), json!("acct_provider_contract")),
("continuation_scope".to_owned(), json!("connection_local")),
]),
)
.unwrap(),
);
let stream = provider_with_base_url(&store, upstream.uri())
.execute(
planned_request("openai", Operation::Generate(generate)),
pinned_continuation_context(
"req_large_api_native",
"acct_provider_contract",
"client_api_previous",
"upstream_api_previous",
1,
ContinuationAttempt::Native,
),
)
.await
.expect("API Key native continuation is not rejected by the Codex size policy");
assert_eq!(stream.metadata().transport().as_str(), "websocket");
assert!(upstream.received_requests().await.unwrap().is_empty());
}
#[tokio::test]
async fn oversized_websocket_http_response_continues_with_client_tool_history_replay() {
fn ws_operation(
body: Map<String, Value>,
connection_id: &str,
state: Option<ProviderSessionState>,
) -> Operation {
let payload = ProtocolPayload::json_object("openai", body)
.unwrap()
.with_context(Map::from_iter([(
"downstream_websocket_connection_id".to_owned(),
json!(connection_id),
)]));
let mut generate = GenerateRequest::from_protocol_payload(payload);
if let Some(state) = state {
generate = generate.with_provider_session_state(state);
}
Operation::Generate(generate)
}
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_provider_contract").await;
let server = MockServer::start().await;
Mock::given(method("POST")).and(path("/codex/responses"))
.respond_with(|request: &wiremock::Request| {
let body = captured_request_body(request);
let has_tool_output = body["input"].as_array().unwrap().iter().any(|item| item["type"] == "function_call_output");
let event = if has_tool_output {
CAPTURE_COMPLETED_SSE.to_owned()
} else {
format!("event: response.completed\ndata: {}\n\n", json!({
"type": "response.completed",
"response": {
"id": "resp_large_ws_http", "model": "gpt-5.4", "status": "completed", "store": false,
"output": [{"type": "function_call", "id": "fc_large", "call_id": "call_large", "name": "inspect", "arguments": "{}", "status": "completed"}],
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}
}
}))
};
ResponseTemplate::new(200).insert_header("content-type", "text/event-stream").set_body_string(event)
}).expect(2).mount(&server).await;
let mut body = Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
(
"instructions".to_owned(),
json!("x".repeat(16 * 1024 * 1024)),
),
(
"input".to_owned(),
json!([{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "inspect"}]}]),
),
(
"tools".to_owned(),
json!([{"type": "function", "name": "inspect", "parameters": {"type": "object"}}]),
),
("session_id".to_owned(), json!("large-ws-http-session")),
("store".to_owned(), json!(false)),
]);
let provider = provider_with_base_url(&store, server.uri());
let mut first = Arc::clone(&provider)
.execute(
planned_request("openai", ws_operation(body.clone(), "ws_first", None)),
context("req_large_ws_first", CancellationToken::new()),
)
.await
.expect("prepare large WS request over HTTP");
assert_eq!(first.metadata().transport().as_str(), "http_sse");
let mut state = None;
let mut output = None;
while let Some(event) = first.next().await {
let event = event.expect("HTTP tool response");
if let Some(update) = event.session_update() {
state = Some(update.clone());
}
if let Some(wire) = event.wire_event()
&& wire.event_type() == Some("response.completed")
{
output = wire.data().pointer("/response/output").cloned();
}
}
drop(first);
let state = state.expect("HTTP response publishes replay requirement");
assert_eq!(state.payload()["continuation_scope"], "replay_required");
let tool_output =
json!({"type": "function_call_output", "call_id": "call_large", "output": "done"});
let delta = Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!([tool_output.clone()])),
(
"previous_response_id".to_owned(),
json!("resp_large_ws_http"),
),
("store".to_owned(), json!(false)),
]);
let error = Arc::clone(&provider)
.execute(
planned_request("openai", ws_operation(delta, "ws_first", Some(state))),
pinned_continuation_context(
"req_large_ws_delta",
"acct_provider_contract",
"resp_large_ws_http",
"resp_large_ws_http",
1,
ContinuationAttempt::Native,
),
)
.await
.err()
.expect("HTTP store=false delta requires client replay");
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
assert_eq!(
error.client_visible_upstream_error().unwrap().code(),
Some("previous_response_not_found")
);
assert_eq!(
error.continuation_recovery_disposition(),
Some(ContinuationRecoveryDisposition::ClientReplayRequired)
);
assert_eq!(server.received_requests().await.unwrap().len(), 1);
let input = body.get_mut("input").unwrap().as_array_mut().unwrap();
input.extend(output.unwrap().as_array().unwrap().iter().cloned());
input.push(tool_output);
let mut replay = provider
.execute(
planned_request("openai", ws_operation(body.clone(), "ws_reconnected", None)),
context("req_large_ws_replay", CancellationToken::new()),
)
.await
.expect("client full history replay over HTTP");
assert_eq!(replay.metadata().transport().as_str(), "http_sse");
let mut completed = false;
while let Some(event) = replay.next().await {
completed |= event
.expect("full history HTTP response")
.wire_event()
.is_some_and(|wire| wire.event_type() == Some("response.completed"));
}
assert!(completed);
let requests = server.received_requests().await.unwrap();
assert_eq!(requests.len(), 2);
let sent = captured_request_body(&requests[1]);
assert_eq!(sent["input"], body["input"]);
assert_eq!(sent["tools"], body["tools"]);
assert!(sent.get("previous_response_id").is_none());
}
#[tokio::test]
async fn fallback_attempt_transport_forces_http_sse_for_a_websocket_request() {
let store = Arc::new(MemoryAccountStore::default());
@@ -11442,3 +11987,163 @@ async fn capture_grok_request(
sent.remove("client_metadata");
sent
}
#[tokio::test]
async fn public_catalog_filters_each_api_account_before_union_without_gating_inference() {
use gateway_core::account::{AccountModelAccess, AccountModelAccessMode};
use gateway_core::engine::provider::{Provider, ProviderRegistry};
use gateway_core::routing::snapshot::{
RuntimeSnapshotCompiler, SnapshotFacts, SnapshotProviderAccountFacts,
SnapshotSettingsFacts, SnapshotStoreError, SnapshotStorePort,
};
struct CatalogStore(SnapshotFacts);
impl SnapshotStorePort for CatalogStore {
fn load_snapshot_facts(&self) -> BoxFuture<'_, Result<SnapshotFacts, SnapshotStoreError>> {
Box::pin(async { Ok(self.0.clone()) })
}
fn current_config_revision(
&self,
) -> BoxFuture<'_, Result<ConfigRevision, SnapshotStoreError>> {
Box::pin(async { Ok(self.0.config_revision()) })
}
}
let first = MockServer::start().await;
let second = MockServer::start().await;
let store = Arc::new(MemoryAccountStore::default());
for (id, server, models) in [
(
"acct_aggregator",
&first,
vec!["gpt-5.4", "deepseek-v4.1-flash", "kimi-k2.5"],
),
("acct_gpt", &second, vec!["gpt-5.4"]),
] {
store
.seed_api_key(
id,
server.uri(),
provider_openai::credential::ResponsesTransport::Http,
)
.await;
Mock::given(method("GET"))
.and(path("/models"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"data": models.into_iter().map(|id| json!({"id": id})).collect::<Vec<_>>()
})))
.expect(1)
.mount(server)
.await;
}
let provider = provider_with_base_url(&store, first.uri());
let registry = Arc::new(ProviderRegistry::new([provider as Arc<dyn Provider>]).unwrap());
let policy = |mode, models: &[&str]| {
AccountModelAccess::new(mode, models.iter().map(|id| (*id).to_owned()).collect()).unwrap()
};
for (a_policy, b_policy) in [
(
policy(
AccountModelAccessMode::Allowlist,
&["gpt-5.4", "deepseek-v4.1-flash"],
),
AccountModelAccess::all(),
),
(
policy(AccountModelAccessMode::Denylist, &["kimi-k2.5"]),
AccountModelAccess::all(),
),
(
policy(
AccountModelAccessMode::Allowlist,
&["gpt-5.4", "deepseek-v4.1-flash"],
),
policy(AccountModelAccessMode::Allowlist, &["gpt-5.4"]),
),
] {
let revision = ConfigRevision::new(1).unwrap();
let facts = SnapshotFacts::new(
revision,
revision,
SnapshotSettingsFacts::new(
3,
0,
"smart",
BTreeMap::from([
("gpt-public".into(), "gpt-5.4".into()),
("kimi-public".into(), "kimi-k2.5".into()),
]),
None,
None,
),
vec![],
vec![],
vec![
SnapshotProviderAccountFacts::new(
ProviderAccountId::new("acct_aggregator").unwrap(),
"openai",
)
.with_model_access(a_policy),
SnapshotProviderAccountFacts::new(
ProviderAccountId::new("acct_gpt").unwrap(),
"openai",
)
.with_model_access(b_policy),
],
vec![],
);
let snapshot =
RuntimeSnapshotCompiler::new(Arc::new(CatalogStore(facts)), registry.clone())
.compile()
.await
.unwrap();
let scope = snapshot.all_account_scope();
let listed = snapshot.public_models_for_scope(&scope);
assert_eq!(
listed.iter().map(PublicModelId::as_str).collect::<Vec<_>>(),
["deepseek-v4.1-flash", "gpt-5.4", "gpt-public"]
);
assert_eq!(
snapshot
.public_model_profiles_for_scope(&scope)
.iter()
.map(|profile| profile.model())
.collect::<Vec<_>>(),
listed.iter().collect::<Vec<_>>()
);
for id in ["kimi-k2.5", "kimi-public"] {
let model = PublicModelId::new(id).unwrap();
assert!(!snapshot.contains_public_model_for_scope(&model, &scope));
}
// 目录来源只决定展示,发现型目录仍允许把请求交给政策合规的上游判断。
for id in ["kimi-k2.5", "kimi-public", "not-yet-discovered"] {
let model = PublicModelId::new(id).unwrap();
assert_eq!(
snapshot
.plan(
&model,
&generate_operation(),
scope.clone(),
&RoutingContext::default()
)
.is_ok(),
scope.allows_provider_model(
&ProviderKind::new("openai").unwrap(),
&snapshot.mapped_model(id)
)
);
}
let only_gpt = scope.restricted_to(
&BTreeSet::new(),
&BTreeSet::from([ProviderAccountId::new("acct_gpt").unwrap()]),
);
assert_eq!(
snapshot
.public_models_for_scope(&only_gpt)
.iter()
.map(PublicModelId::as_str)
.collect::<Vec<_>>(),
["gpt-5.4", "gpt-public"]
);
}
}

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