Compare commits

...
25 Commits
Author SHA1 Message Date
ZyyandCodex 3915aa259b test(accounts): align capacity tests with module layout
Co-authored-by: Codex <noreply@openai.com>
2026-09-27 22:53:03 +08:00
ZyyandCodex d310286e33 chore(release): v3.17.0
Co-authored-by: Codex <noreply@openai.com>
2026-09-27 22:41:22 +08:00
Zyy ea4e3059db feat(accounts): show current concurrency capacity 2026-09-27 22:36:51 +08:00
3bf7030fba fix(openai): widen and trace stream replay grace (#304)
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-09-27 19:51:41 +08:00
5d5965c6ca feat(plugin): expose non-secret client key group facts (#302)
Co-authored-by: root <root@DESKTOP-PRVCERP.localdomain>
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-09-27 19:47:35 +08:00
ZyyandCodex 0710ed1f7c feat(plugins): expose native budgets and quota observations (#300)
Expose independent key budget management and quota observation capabilities
through the public plugin SDK, runtime authorization, and native admin services.
Reuse existing budget transactions and Provider quota refresh paths without
adding workflow-specific state or automatic account-to-key orchestration.

Cover authorization, partial budget updates, resets, quota refresh projection,
transaction rollback, and SDK request behavior. Remove ignored external probes.

Co-authored-by: Codex <noreply@openai.com>
2026-09-27 18:44:32 +08:00
4634d8dc6d fix(openai): clarify quota failures without changing credential state (#301)
Preserve quota HTTP status and bounded error codes, and expose actionable
static messages through the existing admin error response. Record background
quota failures in structured logs without mutating credential snapshots.

Cover OAuth terminal results during concurrent quota rejection, successful
recovery, account fact preservation, and safe HTTP error propagation.

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-09-27 15:31:37 +08:00
ZyyandCodex 8a0759905e fix(frontend): stabilize loading states and unify key usage actions
Keep client identity fields visible while presets load and preserve dashboard
content during silent refreshes. Align key usage toolbar backgrounds, spin the
refresh icon, and disable logout without replacing its icon.

Co-authored-by: Codex <noreply@openai.com>
2026-09-27 13:27:13 +08:00
ZyyandCodex e830b2ef48 fix(openai): isolate account context and scope websocket interrupts (#298)
Filter account-owned routing and identity context at API and Provider boundaries. Bind WebSocket interrupts to the active response on its original account and socket, preserving valid continuations.

Co-authored-by: Codex <noreply@openai.com>
2026-09-27 05:27:07 +08:00
ae8da00459 fix(openai): preserve affinity while waiting for account capacity (#297)
Classify saturated account snapshots as capacity unavailable and prefer
the bound account's existing queue when account waiting is enabled.
Preserve explicit policy picks, child affinity delegation and queue limits.

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-09-27 03:30:09 +08:00
ZyyandCodex 7bbea42fba docs(redis): clarify ACL setup and add regression coverage
Document ACL usernames in connection URLs and the Compose URL override.
Run Store tests with a dedicated ACL user and the default user disabled
after the existing password-authenticated suite.

Co-authored-by: Codex <noreply@openai.com>
2026-09-27 02:49:41 +08:00
ZyyandCodex 2e0669d462 fix(frontend): distinguish historical system update failures
Keep previous operation failures available as history without marking the current installation as failed. Preserve current task errors, response-loss recovery, and restart status.

Fixes #236

Co-authored-by: Codex <noreply@openai.com>
2026-09-27 02:08:48 +08:00
ZyyandCodex 2df9494e99 fix(frontend): stabilize system update progress layout
Co-authored-by: Codex <noreply@openai.com>
2026-09-26 22:26:10 +08:00
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
295 changed files with 10955 additions and 1079 deletions
+13 -2
View File
@@ -19,7 +19,8 @@ description: Codex Proxy RS 仓库开发与问题审查入口。用户要求“
| 模块职责、调用链、状态归属 | [系统架构](../../../docs/architecture.md) |
| Rust 开发 | 加载并使用 `$rust-best-practices`,结合 [系统架构](../../../docs/architecture.md) 与 [项目约定](../../../CONTRIBUTING.md#项目约定) |
| HTTP 路由、请求与响应、导入导出 | [API 文档](../../../docs/api.md) |
| 开发独立网关插件 | [cpr-plugin-dev](../cpr-plugin-dev/SKILL.md):公开 SDK、能力选择、示例、打包与验证 |
| 插件 SDK、Runtime 或宿主能力扩展 | 先执行下方「插件架构边界」,结合 [系统架构](../../../docs/architecture.md) 定位原生能力归属 |
| 开发独立网关插件 | 先执行下方「插件架构边界」,再使用 [cpr-plugin-dev](../cpr-plugin-dev/SKILL.md):公开 SDK、能力选择、示例、打包与验证 |
| 前端开发、管理端组件、主题 Token、视觉验证 | 按下方「前端任务」加载设计约束,结合 [管理端主题](../../../docs/theme.md) 与 [界面验证](../../../CONTRIBUTING.md#界面验证) |
| 常规文档更新 | [文档职责与更新条件](../../../CONTRIBUTING.md#文档职责与更新条件):维护当前说明,在原章节就地修订,不记录变更历史 |
| 部署、客户端配置、备份与恢复 | [部署说明](../../../deploy/README.md) |
@@ -41,6 +42,16 @@ description: Codex Proxy RS 仓库开发与问题审查入口。用户要求“
加载并使用 `$rust-best-practices`,先按架构文档定位事实归属与既有调用链,实现后逐项执行 [后端自审](../../../docs/architecture.md#后端自审) 中与本次改动相关的检查。
交付说明实际复用位置、重要边界的处理与验证缺口,不用“符合最佳实践”或“CI 已通过”代替自审结论。
## 插件架构边界
**底线:宿主可以暴露底层能力,但不能替插件完成业务。** 开发或审查插件 SDK、Runtime、宿主接口及独立网关插件时,必须先明确宿主与插件各自的职责;具体架构依据见 [系统架构](../../../docs/architecture.md)。
- 宿主提供独立、可组合的受控原生能力,负责授权、资源隔离、事务和审计,复用已有服务与状态归属。
插件自身的触发条件、跨资源业务关联、策略和流程编排由插件实现;不能将某个 Issue 的使用场景直接做成宿主的组合流程,或为其增加专属业务状态与补偿逻辑。
- 开发独立网关插件时,在插件内通过公开 SDK 组合用户要求的业务。发现能力缺口时,先界定可独立使用的基础接口,不能把整个插件业务搬进宿主。
- 基础接口必须说明输入输出、权限、事务、失败及重复调用语义;这些合同不应夹带对插件业务工作流或私有数据模型的要求。
- 自审完整调用链、协议、测试与文档,检查是否混入插件业务判断、自动串联或专属状态;仅更换通用名称不算消除业务耦合。
## 执行顺序
1. 从仓库根目录确认工作区、当前任务范围和相关运行实例;保留已有改动。
@@ -49,7 +60,7 @@ description: Codex Proxy RS 仓库开发与问题审查入口。用户要求“
优先依据适用的官方源码确定或评估方案。结合上表读取所需文档,先找可复用的同类实现,再按授权修改或审查所属模块。
4. 文档和运行结果不一致时,追查当前源码、配置与版本,区分实现缺陷和文档过期;
按 [文档职责与更新条件](../../../CONTRIBUTING.md#文档职责与更新条件) 判断是否修订,不为每次改动追加文档或发布说明。
5. 复核相对目标基线的完整差异及上下游:开发时清理本次引入的职责越界、重复逻辑、冗余状态和无依据的分支;仅审查时报告问题与依据。页面与后端分别按上述任务入口执行,不等待 PR 审查代做开发自审。
5. 复核相对目标基线的完整差异及上下游:开发时清理本次引入的职责越界、重复逻辑、冗余状态和无依据的分支;仅审查时报告问题与依据。页面与后端分别按上述任务入口执行,插件相关变更还须逐项核对「插件架构边界」,不等待 PR 审查代做开发自审。
6. 交付前按文档职责逐段复核最终文档 diff:新增内容是否回答当前读者的使用或维护问题,是否属于该文档,
是否与其他位置重复,脱离本次变更背景能否独立理解。仅解释实现经过、前后对比或验证结果的内容移到
任务报告或 PR;仍影响升级、恢复的说明保留适用条件与操作。没有文档差异时,只核对行为变化是否使现有说明失真。
+15
View File
@@ -146,6 +146,21 @@ jobs:
working-directory: backend
run: cargo test --test main --locked
- name: Store tests with Redis ACL authentication
working-directory: backend
shell: bash
run: |
set -euo pipefail
redis_acl_password="$(openssl rand -hex 24)"
echo "::add-mask::$redis_acl_password"
docker exec "${{ job.services.redis.id }}" redis-cli -e --no-auth-warning \
-u "$CPR_TEST_REDIS_URL" --user default ACL SETUSER cpr_acl reset on \
">$redis_acl_password" '~*' '&*' '+@all' >/dev/null
docker exec --env REDISCLI_AUTH="$redis_acl_password" "${{ job.services.redis.id }}" \
redis-cli -e --user cpr_acl ACL SETUSER default off >/dev/null
export CPR_TEST_REDIS_URL="redis://cpr_acl:${redis_acl_password}@127.0.0.1:6379"
cargo test -p gateway-store --test main --locked
frontend:
if: inputs.run_frontend
runs-on: ubuntu-latest
-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())
+10 -1
View File
@@ -527,7 +527,7 @@ pub fn initialize_plugin_accounts(
))
}
/// 为 Runtime 创建只暴露非秘密 Client Key 目录的窄端口。
/// 为 Runtime 创建非秘密 Client Key 目录与预算重置的窄端口。
#[must_use]
pub fn initialize_plugin_client_keys(
providers: ports::provider::ProviderAdminRegistry,
@@ -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,
@@ -223,6 +223,16 @@ pub struct AccountPage {
pub struct AccountPageItem {
pub account: AccountRecord,
pub projection: AccountStatusProjection,
pub capacity: AccountCapacity,
}
/// 网关配置的账号并发上限与查询时的占用,不包含排队请求或上游隐藏限制。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AccountCapacity {
/// 实时租约读取失败时为 `None`,不能当作空闲。
pub used_slots: Option<u64>,
/// 应用账号覆盖或全局默认值后的上限;`None` 表示不限。
pub total_slots: Option<u64>,
}
/// 统一账号目录的全局状态计数,不受当前筛选和分页影响。
@@ -10,7 +10,10 @@ use gateway_core::{
routing::{AccountGroupId, ProviderKind},
};
use super::{AdminModelError, Revision, account_groups::AccountGroupRef};
use super::{
AdminModelError, Revision, account_groups::AccountGroupRef,
plugin_resources::PluginResourceOwner,
};
pub type ProviderRequestProfileOverrides =
BTreeMap<ProviderKind, gateway_core::account::OpaqueProviderData>;
@@ -219,6 +222,21 @@ pub struct ResetClientKeyBudget {
pub period: ClientKeyBudgetPeriod,
}
/// 仅更新指定金额上限;省略的周期保持不变,不修改已用金额或窗口。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpdateClientKeyBudgetLimits {
pub id: ClientApiKeyId,
pub daily_limit_usd: Option<gateway_core::metering::Decimal>,
pub weekly_limit_usd: Option<gateway_core::metering::Decimal>,
}
/// 预算变更的调用来源;插件身份由 Runtime 给出,不能从插件请求反序列化。
#[derive(Debug, Clone)]
pub enum ClientKeyBudgetMutationOrigin {
Admin,
Plugin(PluginResourceOwner),
}
/// 删除 Client Key。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DeleteClientKey {
@@ -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,
@@ -1,9 +1,18 @@
//! 插件可见的 Client Key 非秘密资源投影。
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::AccountGroupId;
use super::PageSize;
/// 插件可读取的当前 Key 范围,不含密钥或配置秘密。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PluginClientKeyFacts {
pub id: ClientApiKeyId,
pub enabled: bool,
pub group_ids: Vec<AccountGroupId>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PluginClientKeyCursor {
pub name: String,
@@ -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,
}
@@ -1227,6 +1227,7 @@ impl fmt::Debug for ProviderExport {
#[derive(Debug, Clone, PartialEq)]
pub struct AccountDirectoryItem {
pub account: AccountRecord,
pub capacity: super::accounts::AccountCapacity,
pub capabilities: super::accounts::ProviderAccountCapabilities,
/// Provider 提供的套餐展示名称;未识别到套餐时为空。
pub plan_type_display: Option<String>,
@@ -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;
@@ -34,6 +34,14 @@ pub trait PluginAccountAccess: Send + Sync {
Err(AdminError::unavailable("插件额度事实查询暂不可用"))
}
/// 通过 Provider 原生管理路径刷新额度;不向插件暴露凭据或执行额度重置。
async fn refresh_quota(
&self,
_account_id: &ProviderAccountId,
) -> Result<crate::model::provider_credentials::ProviderQuota, AdminError> {
Err(AdminError::unavailable("插件额度观测刷新暂不可用"))
}
async fn save(
&self,
command: PreparedPluginAccountSave,
@@ -1,14 +1,35 @@
//! 插件发现 Client Key 公开身份的窄端口。
//! 插件 Client Key 非秘密目录与受控预算管理端口。
use async_trait::async_trait;
use gateway_core::{engine::budget::ClientBudgetStatus, policy::ClientApiKeyId};
use crate::model::{
AdminError,
plugin_client_keys::{PluginClientKeyListQuery, PluginClientKeyPage},
AdminError, MutationContext,
client_keys::{ResetClientKeyBudget, UpdateClientKeyBudgetLimits},
plugin_client_keys::{PluginClientKeyFacts, PluginClientKeyListQuery, PluginClientKeyPage},
plugin_resources::PluginResourceOwner,
};
#[async_trait]
pub trait PluginClientKeyAccess: Send + Sync {
async fn facts(&self, id: &ClientApiKeyId) -> Result<PluginClientKeyFacts, AdminError>;
async fn budget(&self, id: &ClientApiKeyId) -> Result<ClientBudgetStatus, AdminError>;
async fn update_budget_limits(
&self,
owner: &PluginResourceOwner,
command: UpdateClientKeyBudgetLimits,
context: &MutationContext,
) -> Result<ClientApiKeyId, AdminError>;
async fn reset_budget(
&self,
owner: &PluginResourceOwner,
command: ResetClientKeyBudget,
context: &MutationContext,
) -> Result<ClientApiKeyId, AdminError>;
async fn list(
&self,
query: PluginClientKeyListQuery,
@@ -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>>;
}
@@ -22,8 +22,9 @@ use crate::model::{
},
auth::{AdminAuditEvent, AuthSession},
client_keys::{
ClientKeyListQuery, ClientKeyPage, ClientKeyRecord, ClientKeySecret, DeleteClientKey,
NewClientKey, ResetClientKeyBudget, SetClientKeyEnabled, UpdateClientKey,
ClientKeyBudgetMutationOrigin, ClientKeyListQuery, ClientKeyPage, ClientKeyRecord,
ClientKeySecret, DeleteClientKey, NewClientKey, ResetClientKeyBudget, SetClientKeyEnabled,
UpdateClientKey, UpdateClientKeyBudgetLimits,
},
observability::{
DashboardObservation, DashboardRuntimeSlots, DiagnosticDimension, DiagnosticObservation,
@@ -332,10 +333,19 @@ pub trait ClientKeyStore: Send + Sync {
context: &MutationContext,
) -> AdminStoreResult<Revision>;
/// 局部更新预算上限,保留其他策略和账本;无变化时不产生配置版本或审计。
async fn update_client_key_budget_limits(
&self,
command: UpdateClientKeyBudgetLimits,
origin: ClientKeyBudgetMutationOrigin,
context: &MutationContext,
) -> AdminStoreResult<Option<Revision>>;
/// 仅修改运行时账本并原子记录审计,不推进配置版本。
async fn reset_client_key_budget(
&self,
command: ResetClientKeyBudget,
origin: ClientKeyBudgetMutationOrigin,
context: &MutationContext,
) -> AdminStoreResult<()>;
}
@@ -515,9 +525,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 +558,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 +569,7 @@ impl AdminStorePorts {
backup,
plugins,
plugin_state,
plugin_resources,
}
}
@@ -363,6 +363,7 @@ impl DefaultAccountsService {
.and_then(|(window, _)| window.local_usage.clone())
});
Ok(AccountDirectoryItem {
capacity: stored.capacity,
capabilities: provider
.account_capabilities(account_id, &stored.account.authentication_kind),
plan_type_display: self.providers.resolve_account_plan(
@@ -402,6 +403,18 @@ impl AccountsService for DefaultAccountsService {
.iter()
.map(|item| item.account.id.clone())
.collect::<Vec<_>>();
// 只读取当前页的租约占用;观测失败不能把账号误报为空闲或拖垮目录。
let in_flight = if ids.is_empty() {
None
} else {
match self.account_runtime.account_runtime(&ids).await {
Ok(runtime) => runtime.in_flight,
Err(error) => {
tracing::warn!(error = %error, "account capacity projection is unavailable");
None
}
}
};
let rolling_usage = self
.accounts
.load_account_usage(rolling_range, &ids)
@@ -472,6 +485,12 @@ impl AccountsService for DefaultAccountsService {
.and_then(|(window, _)| window.local_usage.clone())
});
Ok(AccountDirectoryItem {
capacity: crate::model::accounts::AccountCapacity {
used_slots: in_flight
.as_ref()
.map(|counts| counts.get(&item.account.id).copied().unwrap_or(0)),
..item.capacity
},
capabilities,
plan_type_display: providers.resolve_account_plan(
item.account.provider_kind.as_str(),
@@ -4,8 +4,9 @@ use std::sync::Arc;
use async_trait::async_trait;
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::runtime::SnapshotControl;
use gateway_core::{
engine::budget::ClientBudgetStatus, policy::ClientApiKeyId, runtime::SnapshotControl,
};
use rand_core::{OsRng, RngCore as _};
use uuid::Uuid;
@@ -13,10 +14,10 @@ use crate::{
model::{
AdminError, MutationContext,
client_keys::{
ClientKeyCursorValue, ClientKeyListQuery, ClientKeyMutation, ClientKeyPage,
ClientKeySecret, ClientKeySortField, CreateClientKey, CreatedClientKey,
DeleteClientKey, NewClientKey, ResetClientKeyBudget, SetClientKeyEnabled,
UpdateClientKey,
ClientKeyBudgetMutationOrigin, ClientKeyCursorValue, ClientKeyListQuery,
ClientKeyMutation, ClientKeyPage, ClientKeyRecord, ClientKeySecret, ClientKeySortField,
CreateClientKey, CreatedClientKey, DeleteClientKey, NewClientKey, ResetClientKeyBudget,
SetClientKeyEnabled, UpdateClientKey, UpdateClientKeyBudgetLimits,
},
},
ports::store::{AdminStoreError, AdminStoreErrorKind, ClientKeyStore},
@@ -27,6 +28,7 @@ use super::{map_store_error, publish_committed};
/// API 消费的 Client Key 管理服务。
#[async_trait]
pub trait ClientKeyService: Send + Sync {
async fn get(&self, id: &ClientApiKeyId) -> Result<ClientKeyRecord, AdminError>;
async fn list(&self, query: ClientKeyListQuery) -> Result<ClientKeyPage, AdminError>;
async fn reveal(&self, id: &ClientApiKeyId) -> Result<ClientKeySecret, AdminError>;
async fn create(
@@ -49,10 +51,18 @@ pub trait ClientKeyService: Send + Sync {
context: &MutationContext,
command: DeleteClientKey,
) -> Result<ClientKeyMutation, AdminError>;
async fn budget(&self, id: &ClientApiKeyId) -> Result<ClientBudgetStatus, AdminError>;
async fn update_budget_limits(
&self,
context: &MutationContext,
command: UpdateClientKeyBudgetLimits,
origin: ClientKeyBudgetMutationOrigin,
) -> Result<ClientApiKeyId, AdminError>;
async fn reset_budget(
&self,
context: &MutationContext,
command: ResetClientKeyBudget,
origin: ClientKeyBudgetMutationOrigin,
) -> Result<ClientApiKeyId, AdminError>;
}
@@ -79,14 +89,48 @@ impl DefaultClientKeyService {
#[async_trait]
impl ClientKeyService for DefaultClientKeyService {
async fn get(&self, id: &ClientApiKeyId) -> Result<ClientKeyRecord, AdminError> {
self.store
.get_client_key(id)
.await
.map_err(|error| map_store_error(error, "client API key"))?
.ok_or_else(|| AdminError::not_found("Client API Key 不存在"))
}
async fn budget(&self, id: &ClientApiKeyId) -> Result<ClientBudgetStatus, AdminError> {
self.get(id).await.map(|key| key.budget)
}
async fn update_budget_limits(
&self,
context: &MutationContext,
command: UpdateClientKeyBudgetLimits,
origin: ClientKeyBudgetMutationOrigin,
) -> Result<ClientApiKeyId, AdminError> {
if command.daily_limit_usd.is_none() && command.weekly_limit_usd.is_none() {
return Err(AdminError::invalid("至少指定一个预算上限"));
}
let id = command.id.clone();
if let Some(revision) = self
.store
.update_client_key_budget_limits(command, origin, context)
.await
.map_err(|error| map_store_error(error, "client API key"))?
{
publish_committed(self.snapshot.as_ref(), revision).await?;
}
Ok(id)
}
async fn reset_budget(
&self,
context: &MutationContext,
command: ResetClientKeyBudget,
origin: ClientKeyBudgetMutationOrigin,
) -> Result<ClientApiKeyId, AdminError> {
let id = command.id.clone();
self.store
.reset_client_key_budget(command, context)
.reset_client_key_budget(command, origin, context)
.await
.map_err(|error| map_store_error(error, "client API key"))?;
Ok(id)
@@ -124,9 +168,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 +299,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()
@@ -45,6 +45,22 @@ impl DefaultPluginAccountAccess {
}
}
async fn quota(
&self,
account_id: &ProviderAccountId,
refresh: bool,
) -> Result<crate::model::provider_credentials::ProviderQuota, AdminError> {
let account = self.get_runtime(account_id).await?;
self.provider(&account.provider_kind)?
.quota(crate::model::provider_credentials::ProviderQuotaRequest {
account_id: account_id.clone(),
refresh,
rolling_usage: None,
})
.await
.map_err(|error| map_provider_error(error, "plugin quota facts"))
}
fn provider(
&self,
provider_kind: &ProviderKind,
@@ -99,16 +115,16 @@ impl PluginAccountAccess for DefaultPluginAccountAccess {
&self,
account_id: &ProviderAccountId,
) -> Result<crate::model::provider_credentials::ProviderQuota, AdminError> {
let account = self.get_runtime(account_id).await?;
self.provider(&account.provider_kind)?
.quota(crate::model::provider_credentials::ProviderQuotaRequest {
account_id: account_id.clone(),
refresh: false,
rolling_usage: None,
})
.await
.map_err(|error| map_provider_error(error, "plugin quota facts"))
self.quota(account_id, false).await
}
async fn refresh_quota(
&self,
account_id: &ProviderAccountId,
) -> Result<crate::model::provider_credentials::ProviderQuota, AdminError> {
self.quota(account_id, true).await
}
async fn save(
&self,
command: PreparedPluginAccountSave,
@@ -1,19 +1,23 @@
//! 插件 Client Key 目录;复用管理服务并收窄为非秘密投影。
//! 插件 Client Key 管理;复用管理服务,只开放非秘密目录与预算操作。
use std::sync::Arc;
use async_trait::async_trait;
use gateway_core::{engine::budget::ClientBudgetStatus, policy::ClientApiKeyId};
use crate::{
model::{
AdminError,
AdminError, MutationContext,
client_keys::{
ClientKeyCursor, ClientKeyCursorValue, ClientKeyListQuery, ClientKeyPageSize,
ClientKeySort, ClientKeySortField, SortDirection,
ClientKeyBudgetMutationOrigin, ClientKeyCursor, ClientKeyCursorValue,
ClientKeyListQuery, ClientKeyPageSize, ClientKeySort, ClientKeySortField,
ResetClientKeyBudget, SortDirection, UpdateClientKeyBudgetLimits,
},
plugin_client_keys::{
PluginClientKey, PluginClientKeyCursor, PluginClientKeyListQuery, PluginClientKeyPage,
PluginClientKey, PluginClientKeyCursor, PluginClientKeyFacts, PluginClientKeyListQuery,
PluginClientKeyPage,
},
plugin_resources::PluginResourceOwner,
},
ports::plugin_client_keys::PluginClientKeyAccess,
use_case::client_keys::ClientKeyService,
@@ -31,6 +35,49 @@ impl DefaultPluginClientKeyAccess {
#[async_trait]
impl PluginClientKeyAccess for DefaultPluginClientKeyAccess {
async fn facts(&self, id: &ClientApiKeyId) -> Result<PluginClientKeyFacts, AdminError> {
let key = self.service.get(id).await?;
Ok(PluginClientKeyFacts {
id: key.id,
enabled: key.enabled,
group_ids: key.groups.into_iter().map(|group| group.id).collect(),
})
}
async fn budget(&self, id: &ClientApiKeyId) -> Result<ClientBudgetStatus, AdminError> {
self.service.budget(id).await
}
async fn update_budget_limits(
&self,
owner: &PluginResourceOwner,
command: UpdateClientKeyBudgetLimits,
context: &MutationContext,
) -> Result<ClientApiKeyId, AdminError> {
self.service
.update_budget_limits(
context,
command,
ClientKeyBudgetMutationOrigin::Plugin(owner.clone()),
)
.await
}
async fn reset_budget(
&self,
owner: &PluginResourceOwner,
command: ResetClientKeyBudget,
context: &MutationContext,
) -> Result<ClientApiKeyId, AdminError> {
self.service
.reset_budget(
context,
command,
ClientKeyBudgetMutationOrigin::Plugin(owner.clone()),
)
.await
}
async fn list(
&self,
query: PluginClientKeyListQuery,
@@ -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)| {
@@ -1,4 +1,7 @@
use std::sync::{Arc, Mutex};
use std::{
collections::BTreeMap,
sync::{Arc, Mutex},
};
use async_trait::async_trait;
use chrono::{TimeDelta, Utc};
@@ -55,7 +58,8 @@ use gateway_admin::{
ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind, ProviderAdminRegistry,
},
store::{
AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStoreResult, SettingsStore,
AccountRuntimeStore, AccountStore, AdminStoreError, AdminStoreErrorKind,
AdminStoreResult, SettingsStore,
},
},
};
@@ -668,6 +672,12 @@ impl FakeAccountStore {
last_error_message: account.last_error_message.clone(),
};
AccountPageItem {
capacity: gateway_admin::model::accounts::AccountCapacity {
used_slots: None,
total_slots: account
.concurrency_limit
.map(|limit| u64::from(limit.get())),
},
account,
projection: resolve_account_status(&facts, std::time::SystemTime::now()),
}
@@ -1105,6 +1115,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,
@@ -3404,3 +3415,136 @@ async fn unregistered_provider_accounts_have_no_live_management_capabilities() {
);
assert_eq!(page.items[0].capabilities, Default::default());
}
#[tokio::test]
async fn plugin_quota_refresh_uses_native_provider_without_changing_configuration() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
let access = gateway_admin::initialize_plugin_accounts(
ProviderAdminRegistry::new([provider.clone() as Arc<dyn ProviderAdmin>]).unwrap(),
store,
Arc::new(RecordingPluginAccountPublication(events.clone())),
);
let id = ProviderAccountId::new("acct_test").unwrap();
access.get_quota(&id).await.unwrap();
access.refresh_quota(&id).await.unwrap();
assert_eq!(
provider.quota_requests(),
vec![
ProviderQuotaRequest {
account_id: id.clone(),
refresh: false,
rolling_usage: None
},
ProviderQuotaRequest {
account_id: id,
refresh: true,
rolling_usage: None
},
]
);
assert!(!recorded(&events).contains(&"snapshot.publish_committed"));
}
struct CapacityRuntime {
counts: Option<BTreeMap<String, u64>>,
fail: bool,
reads: Mutex<Vec<Vec<String>>>,
}
#[async_trait]
impl AccountRuntimeStore for CapacityRuntime {
async fn active_rate_limits(&self) -> AdminStoreResult<AccountRuntimeSnapshot> {
Ok(AccountRuntimeSnapshot::default())
}
async fn account_runtime(&self, ids: &[String]) -> AdminStoreResult<AccountRuntimeSnapshot> {
self.reads.lock().unwrap().push(ids.to_vec());
if self.fail {
return Err(store_unavailable());
}
Ok(AccountRuntimeSnapshot {
in_flight: self.counts.clone(),
..Default::default()
})
}
async fn active_freezes(
&self,
) -> AdminStoreResult<BTreeMap<String, gateway_admin::model::accounts::AccountFreeze>> {
Ok(BTreeMap::new())
}
async fn capacity_peaks(&self, _: &[String]) -> AdminStoreResult<BTreeMap<String, u32>> {
Ok(BTreeMap::new())
}
async fn finish_freeze(
&self,
_: &str,
_: &gateway_admin::model::accounts::AccountFreeze,
_: Option<chrono::DateTime<Utc>>,
) -> AdminStoreResult<bool> {
Ok(false)
}
}
#[tokio::test]
async fn account_capacity_should_batch_page_ids_and_distinguish_idle_from_unavailable() {
for (counts, fail, expected) in [
(
Some(BTreeMap::from([("acct_test".to_owned(), 3)])),
false,
Some(3),
),
(Some(BTreeMap::new()), false, Some(0)),
(None, false, None),
(None, true, None),
] {
for empty in [false, true] {
let runtime = Arc::new(CapacityRuntime {
counts: counts.clone(),
fail,
reads: Mutex::default(),
});
let mut account = account_record("openai");
account.concurrency_limit = gateway_core::account::AccountConcurrencyLimit::new(5);
let store = FakeAccountStore::with_account(account, events());
if empty {
store.set_accounts(Vec::new());
}
let services = super::AdminHarness::new()
.accounts(store)
.account_runtime(runtime.clone())
.settings(Arc::new(StaticSettingsStore))
.provider(FakeProviderAdmin::new("openai", events()))
.build()
.await;
let page = services
.accounts()
.list(AccountListQuery {
page: 1,
page_size: gateway_admin::model::PageSize::new(20).unwrap(),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
})
.await
.unwrap();
if empty {
assert!(page.items.is_empty());
assert!(runtime.reads.lock().unwrap().is_empty());
} else {
assert_eq!(
*runtime.reads.lock().unwrap(),
[vec!["acct_test".to_owned()]]
);
assert_eq!(page.items[0].capacity.used_slots, expected);
assert_eq!(page.items[0].capacity.total_slots, Some(5));
}
}
}
}
@@ -33,9 +33,19 @@ struct TestClientKeyStore {
#[async_trait]
impl ClientKeyStore for TestClientKeyStore {
async fn update_client_key_budget_limits(
&self,
_: gateway_admin::model::client_keys::UpdateClientKeyBudgetLimits,
_: gateway_admin::model::client_keys::ClientKeyBudgetMutationOrigin,
_: &MutationContext,
) -> AdminStoreResult<Option<Revision>> {
Err(unused())
}
async fn reset_client_key_budget(
&self,
command: gateway_admin::model::client_keys::ResetClientKeyBudget,
_: gateway_admin::model::client_keys::ClientKeyBudgetMutationOrigin,
_: &MutationContext,
) -> AdminStoreResult<()> {
self.resets.lock().unwrap().push(command);
@@ -147,7 +157,11 @@ async fn reset_budget_forwards_scope_and_returns_only_key_identity() {
};
let result = services
.client_keys()
.reset_budget(&mutation_context(), command.clone())
.reset_budget(
&mutation_context(),
command.clone(),
gateway_admin::model::client_keys::ClientKeyBudgetMutationOrigin::Admin,
)
.await
.unwrap();
assert_eq!(result, command.id);
@@ -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),
@@ -643,9 +644,19 @@ impl AccountRuntimeStore for UnavailableStore {
#[async_trait]
impl ClientKeyStore for UnavailableStore {
async fn update_client_key_budget_limits(
&self,
_: gateway_admin::model::client_keys::UpdateClientKeyBudgetLimits,
_: gateway_admin::model::client_keys::ClientKeyBudgetMutationOrigin,
_: &MutationContext,
) -> AdminStoreResult<Option<Revision>> {
Err(unavailable("client key budget limits"))
}
async fn reset_client_key_budget(
&self,
_: gateway_admin::model::client_keys::ResetClientKeyBudget,
_: gateway_admin::model::client_keys::ClientKeyBudgetMutationOrigin,
_: &MutationContext,
) -> AdminStoreResult<()> {
Err(unavailable("client key budget reset"))
@@ -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,
@@ -60,6 +60,7 @@ pub(super) fn account_models_data(result: ProviderModels) -> AccountModelsData {
pub(super) fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
let AccountDirectoryItem {
account,
capacity,
capabilities,
plan_type_display,
projection,
@@ -109,6 +110,10 @@ pub(super) fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> Ac
error_message: projection.error_message,
enabled: account.enabled,
concurrency_limit: account.concurrency_limit.map(|limit| limit.get()),
capacity: AccountCapacityView {
used_slots: capacity.used_slots,
total_slots: capacity.total_slots,
},
weight: account.weight.get(),
model_access: account.model_access,
outbound_proxy_endpoint: account
@@ -290,6 +290,7 @@ pub struct AccountView {
pub error_message: Option<String>,
pub enabled: bool,
pub concurrency_limit: Option<u32>,
pub capacity: AccountCapacityView,
pub weight: u16,
pub model_access: gateway_core::account::AccountModelAccess,
pub access_token_expires_at: Option<String>,
@@ -304,6 +305,13 @@ pub struct AccountView {
pub usage: AccountUsageView,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountCapacityView {
pub used_slots: Option<u64>,
pub total_slots: Option<u64>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountCapabilitiesView {
@@ -977,6 +977,7 @@ where
.reset_budget(
&auth.context().mutation_context(),
payload.into_command().map_err(map_wire_error)?,
gateway_admin::model::client_keys::ClientKeyBudgetMutationOrigin::Admin,
)
.await
.map_err(map_service_error)?;
@@ -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,
@@ -584,6 +584,8 @@ fn passthrough_header_name(name: &str, connection_headers: &[String]) -> bool {
| "chatgpt-project-id"
| "openai-organization"
| "openai-project"
| "x-openai-account-routing-override"
| "x-openai-fedramp"
// 上游指纹必须由运行时画像统一生成,客户端 originator/User-Agent/version
// 不能作为不透明头透传覆盖,避免不同下游客户端暴露不一致的设备指纹。
| "originator"
@@ -1,6 +1,7 @@
//! 下游客户端 WebSocket 的单 owner pump 与有界收发边界。
use std::{
collections::VecDeque,
fmt,
sync::{
Arc,
@@ -14,7 +15,7 @@ use futures::{Sink, SinkExt, Stream, StreamExt};
use gateway_core::lifecycle::CancellationToken;
use thiserror::Error;
use tokio::{
sync::{mpsc, oneshot},
sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot},
task::JoinHandle,
time::{Instant, timeout},
};
@@ -48,6 +49,21 @@ pub enum ConnectionEvent {
Exited(PumpExitReason),
}
/// 接收队列与活动期间暂存的请求共用容量;移动帧不释放它占用的名额。
pub(super) struct PendingConnectionEvent {
pub(super) event: ConnectionEvent,
_permit: Option<OwnedSemaphorePermit>,
}
impl PendingConnectionEvent {
fn exited(reason: PumpExitReason) -> Self {
Self {
event: ConnectionEvent::Exited(reason),
_permit: None,
}
}
}
/// 下游写入阶段;名称刻意使用 write,而不是暗示客户端已消费的 delivery。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FramePhase {
@@ -213,7 +229,8 @@ pub struct ResponsesWebSocketConnection {
opened_at: Instant,
expired: Arc<AtomicBool>,
commands: Option<mpsc::Sender<ConnectionCommand>>,
incoming: mpsc::Receiver<ConnectionEvent>,
incoming: mpsc::Receiver<PendingConnectionEvent>,
deferred: VecDeque<PendingConnectionEvent>,
exited: oneshot::Receiver<PumpExitReason>,
pump_task: Option<JoinHandle<()>>,
stats: Arc<ConnectionStats>,
@@ -251,17 +268,43 @@ impl ResponsesWebSocketConnection {
/// 等待下一个需要业务层处理的客户端事件。
pub async fn next_event(&mut self) -> Option<ConnectionEvent> {
if self.exit_reason.is_none() {
match self.exited.try_recv() {
Ok(reason) => {
self.exit_reason = Some(reason);
}
Err(tokio::sync::oneshot::error::TryRecvError::Closed) => {
self.exit_reason = Some(PumpExitReason::PumpStopped);
}
Err(tokio::sync::oneshot::error::TryRecvError::Empty) => {}
}
}
if let Some(reason) = self.exit_reason {
return Some(ConnectionEvent::Exited(reason));
}
if let Some(event) = self.deferred.pop_front() {
return Some(event.event);
}
self.next_active_event().await.map(|event| event.event)
}
/// 活动响应期间只消费新到的帧,已排队的下一轮请求继续保持串行。
pub(super) async fn next_active_event(&mut self) -> Option<PendingConnectionEvent> {
if let Some(reason) = self.exit_reason {
return Some(PendingConnectionEvent::exited(reason));
}
let event = tokio::select! {
biased;
reason = &mut self.exited => {
Some(ConnectionEvent::Exited(reason.unwrap_or(PumpExitReason::PumpStopped)))
Some(PendingConnectionEvent::exited(reason.unwrap_or(PumpExitReason::PumpStopped)))
}
event = self.incoming.recv() => event,
};
if let Some(ConnectionEvent::Exited(reason)) = event.as_ref() {
if let Some(PendingConnectionEvent {
event: ConnectionEvent::Exited(reason),
..
}) = event.as_ref()
{
self.exit_reason.get_or_insert(*reason);
} else if event.is_none() {
self.exit_reason.get_or_insert(PumpExitReason::PumpStopped);
@@ -269,6 +312,10 @@ impl ResponsesWebSocketConnection {
event
}
pub(super) fn defer(&mut self, event: PendingConnectionEvent) {
self.deferred.push_back(event);
}
/// 等待连接退出,不消费留给后续串行请求的业务帧。
pub async fn wait_for_exit(&mut self) -> PumpExitReason {
if let Some(reason) = self.exit_reason {
@@ -494,6 +541,7 @@ where
expired,
commands: Some(command_tx),
incoming: incoming_rx,
deferred: VecDeque::new(),
exited: exit_rx,
pump_task: Some(pump_task),
stats,
@@ -506,7 +554,7 @@ where
async fn run_pump<S, E>(
mut socket: S,
mut commands: mpsc::Receiver<ConnectionCommand>,
incoming: mpsc::Sender<ConnectionEvent>,
incoming: mpsc::Sender<PendingConnectionEvent>,
exited: oneshot::Sender<PumpExitReason>,
connection_id: Arc<str>,
cancellation: CancellationToken,
@@ -518,6 +566,7 @@ async fn run_pump<S, E>(
S: Stream<Item = Result<Message, E>> + Sink<Message, Error = E> + Unpin,
E: fmt::Display,
{
let event_slots = Arc::new(Semaphore::new(INBOUND_EVENT_BUFFER));
let deadline = tokio::time::sleep_until(opened_at + config.max_age);
tokio::pin!(deadline);
let mut deadline_elapsed = false;
@@ -536,11 +585,9 @@ async fn run_pump<S, E>(
expired.store(true, Ordering::Release);
// 满队列本身已能唤醒空闲协调层;到期标志仍阻止启动下一轮。
// 不能因无法追加 Expired 通知而中断正在收尾的响应。
match incoming.try_send(ConnectionEvent::Expired) {
Ok(()) | Err(mpsc::error::TrySendError::Full(_)) => {}
Err(mpsc::error::TrySendError::Closed(_)) => {
break PumpExitReason::CoordinatorDropped;
}
match emit_incoming(&incoming, &event_slots, ConnectionEvent::Expired) {
Ok(()) | Err(PumpExitReason::InboundOverload) => {}
Err(reason) => break reason,
}
}
command = commands.recv() => {
@@ -608,12 +655,12 @@ async fn run_pump<S, E>(
}
match message {
Some(Ok(Message::Text(payload))) => {
if let Err(reason) = emit_incoming(&incoming, ConnectionEvent::Text(payload.to_string())) {
if let Err(reason) = emit_incoming(&incoming, &event_slots, ConnectionEvent::Text(payload.to_string())) {
break reason;
}
}
Some(Ok(Message::Binary(_))) => {
if let Err(reason) = emit_incoming(&incoming, ConnectionEvent::Binary) {
if let Err(reason) = emit_incoming(&incoming, &event_slots, ConnectionEvent::Binary) {
break reason;
}
}
@@ -662,11 +709,20 @@ async fn run_pump<S, E>(
}
fn emit_incoming(
incoming: &mpsc::Sender<ConnectionEvent>,
incoming: &mpsc::Sender<PendingConnectionEvent>,
slots: &Arc<Semaphore>,
event: ConnectionEvent,
) -> Result<(), PumpExitReason> {
incoming.try_send(event).map_err(|error| match error {
mpsc::error::TrySendError::Full(_) => PumpExitReason::InboundOverload,
mpsc::error::TrySendError::Closed(_) => PumpExitReason::CoordinatorDropped,
})
let permit = Arc::clone(slots)
.try_acquire_owned()
.map_err(|_| PumpExitReason::InboundOverload)?;
incoming
.try_send(PendingConnectionEvent {
event,
_permit: Some(permit),
})
.map_err(|error| match error {
mpsc::error::TrySendError::Full(_) => PumpExitReason::InboundOverload,
mpsc::error::TrySendError::Closed(_) => PumpExitReason::CoordinatorDropped,
})
}
@@ -12,6 +12,7 @@ use gateway_core::engine::middleware::{
MiddlewareBody, MiddlewareError, MiddlewareFrame, MiddlewareFraming, MiddlewareHeader,
MiddlewareResponse,
};
use gateway_core::engine::response_control::ResponseControl;
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::event::ProviderResponseHeader;
use gateway_core::operation::ProviderSessionState;
@@ -22,8 +23,10 @@ use crate::openai::responses::validation::{ResponseValidationFacts, ResponsesDel
use super::{
super::{DecodedResponsesRequest, OpenAiResponsesEncoder, ProtocolErrorBody},
connection::{FramePhase, ResponsesWebSocketConnection, WriteContext},
protocol::{error_event, initial_engine_error_event, response_metadata_event},
connection::{ConnectionEvent, FramePhase, ResponsesWebSocketConnection, WriteContext},
protocol::{
decode_response_interrupt, error_event, initial_engine_error_event, response_metadata_event,
},
};
#[derive(Clone, Copy, PartialEq, Eq)]
@@ -144,6 +147,7 @@ pub(super) async fn forward_response(
replay: &mut ConnectionReplaySnapshot,
capture: ReplayCaptureHandle,
validation: ResponseValidationFacts,
response_control: ResponseControl,
) -> ForwardOutcome {
let (protocol, status, headers, mut body, _) = response.into_parts();
if protocol != "openai" || StatusCode::from_u16(status).is_err() {
@@ -156,7 +160,9 @@ pub(super) async fn forward_response(
ProviderResponseHeader::new(name, value)
})
.collect::<Vec<_>>();
let first = match next_body_input(connection, body.as_mut()).await {
let first = match next_body_input(connection, body.as_mut(), &response_control, &request_id)
.await
{
BodyInput::Frame(Ok(Some(frame))) => frame,
BodyInput::Frame(Ok(None)) => {
return fail_body(connection, body, MiddlewareError::InvalidState, &request_id).await;
@@ -200,10 +206,33 @@ pub(super) async fn forward_response(
loop {
let frame = match current.take() {
Some(frame) => frame,
None => match next_body_input(connection, body.as_mut()).await {
BodyInput::Frame(Ok(Some(frame))) => frame,
BodyInput::Frame(Ok(None)) if terminal_seen && body.is_finalized() => {
if delivery_validator.finish_websocket_delivery().is_err() {
None => {
match next_body_input(connection, body.as_mut(), &response_control, &request_id)
.await
{
BodyInput::Frame(Ok(Some(frame))) => frame,
BodyInput::Frame(Ok(None)) if terminal_seen && body.is_finalized() => {
if delivery_validator.finish_websocket_delivery().is_err() {
return fail_body(
connection,
body,
MiddlewareError::InvalidState,
&request_id,
)
.await;
}
{
let captured = capture
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(response_id) = captured.response_id.clone() {
replay.commit(response_id, captured.provider_state.clone());
}
}
body.close().await;
return ForwardOutcome::Continue;
}
BodyInput::Frame(Ok(None)) => {
return fail_body(
connection,
body,
@@ -212,29 +241,15 @@ pub(super) async fn forward_response(
)
.await;
}
{
let captured = capture
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(response_id) = captured.response_id.clone() {
replay.commit(response_id, captured.provider_state.clone());
}
BodyInput::Frame(Err(error)) => {
return fail_body(connection, body, error, &request_id).await;
}
BodyInput::Disconnect => {
detach_body(body);
return ForwardOutcome::Disconnect;
}
body.close().await;
return ForwardOutcome::Continue;
}
BodyInput::Frame(Ok(None)) => {
return fail_body(connection, body, MiddlewareError::InvalidState, &request_id)
.await;
}
BodyInput::Frame(Err(error)) => {
return fail_body(connection, body, error, &request_id).await;
}
BodyInput::Disconnect => {
detach_body(body);
return ForwardOutcome::Disconnect;
}
},
}
};
if terminal_seen || validate_frame(&frame).is_err() {
return fail_body(connection, body, MiddlewareError::InvalidState, &request_id).await;
@@ -293,11 +308,45 @@ enum BodyInput {
async fn next_body_input(
connection: &mut ResponsesWebSocketConnection,
body: &mut dyn MiddlewareBody,
response_control: &ResponseControl,
request_id: &Arc<str>,
) -> BodyInput {
tokio::select! {
biased;
_ = connection.wait_for_exit() => BodyInput::Disconnect,
frame = body.next_frame() => BodyInput::Frame(frame),
// 处理控制帧时继续持有同一个读取 future,不能取消正在加工正文的中间件。
let frame = body.next_frame();
tokio::pin!(frame);
loop {
tokio::select! {
event = connection.next_active_event() => {
let Some(event) = event else { return BodyInput::Disconnect; };
match &event.event {
ConnectionEvent::Text(payload) => {
let error = match decode_response_interrupt(payload) {
Ok(Some(response_id)) => response_control.interrupt(&response_id).err().map(|_| {
super::super::RequestDecodeError::InvalidValue { field: "response_id".to_owned() }.protocol_body()
}),
Ok(None) => {
connection.defer(event);
continue;
}
Err(error) => Some(error.protocol_body()),
};
if let Some(error) = error
&& send_protocol_error(connection, StatusCode::BAD_REQUEST, error, request_id).await == ForwardOutcome::Disconnect
{
return BodyInput::Disconnect;
}
}
ConnectionEvent::Expired => {
// 允许当前响应收尾,外层仍会在下一轮准入前关闭过期连接。
}
ConnectionEvent::Binary => {
connection.defer(event);
}
ConnectionEvent::Exited(_) => return BodyInput::Disconnect,
}
}
frame = &mut frame => return BodyInput::Frame(frame),
}
}
}
@@ -254,6 +254,7 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
}
};
let request_id = Arc::<str>::from(prepared.request_id().to_string());
let response_control = prepared.response_control();
let capture = new_replay_capture();
let validation = ResponseValidationFacts::default();
let input = HttpMiddlewareInput {
@@ -343,6 +344,7 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
&mut replay,
capture,
validation,
response_control,
)
.await
== ForwardOutcome::Disconnect
@@ -31,6 +31,44 @@ pub fn decode_response_create_with_context(
decode_response_create_inner(payload, request_headers)
}
/// 中断只引用活动响应;普通创建帧仍交给原有解码与串行准入路径。
pub(super) fn decode_response_interrupt(
payload: &str,
) -> Result<Option<String>, ResponseCreateFrameError> {
#[derive(serde::Deserialize)]
struct FrameType {
#[serde(rename = "type")]
message_type: Option<String>,
}
// 非控制帧仍在原来的串行请求边界校验,不能抢先拒绝排队的普通请求。
// 只投影类型,避免为排队的大型 response.create 再构造完整 JSON 树。
let Ok(frame) = serde_json::from_str::<FrameType>(payload) else {
return Ok(None);
};
if frame.message_type.as_deref() != Some("response.interrupt") {
return Ok(None);
}
let value: Value =
serde_json::from_str(payload).map_err(|_| ResponseCreateFrameError::InvalidJson)?;
if value.get("mode").and_then(Value::as_str) != Some("discard_partial_items") {
return Err(ResponseCreateFrameError::Request(
RequestDecodeError::InvalidValue {
field: "mode".to_owned(),
},
));
}
let id = value
.get("response_id")
.and_then(Value::as_str)
.filter(|id| !id.is_empty())
.ok_or_else(|| {
ResponseCreateFrameError::Request(RequestDecodeError::InvalidValue {
field: "response_id".to_owned(),
})
})?;
Ok(Some(id.to_owned()))
}
fn decode_response_create_inner(
payload: &str,
request_headers: &OpenAiRequestHeaders,
@@ -337,7 +337,39 @@ mod provider {
use super::*;
#[tokio::test]
async fn provider_public_errors_survive_import_and_refresh_http_handlers() {
async fn unsupported_quota_should_preserve_account_capacity_wire() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
for (used_slots, total_slots) in [(Some(3), Some(5)), (Some(0), None), (None, Some(10))] {
let mut stored = account("openai");
stored.capacity = gateway_admin::model::accounts::AccountCapacity {
used_slots,
total_slots,
};
*fixture.account.lock().unwrap() = Some(stored);
let mut request = request(
Method::GET,
"/api/admin/accounts/quota?accountId=acct_error_test",
Body::empty(),
);
request
.headers_mut()
.insert(header::COOKIE, SESSION_COOKIE.parse().unwrap());
let response = app(fixture.state()).oneshot(request).await.unwrap();
let (status, _, body) = response_json(response).await;
assert_eq!(status, StatusCode::OK);
assert_eq!(
body["data"]["account"]["capacity"],
json!({
"usedSlots": used_slots,
"totalSlots": total_slots,
})
);
}
}
#[tokio::test]
async fn provider_public_errors_survive_import_credential_and_quota_refresh_handlers() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
for provider in ["openai", "xai"] {
@@ -383,6 +415,10 @@ mod provider {
"/api/admin/accounts/import",
json!({"provider": provider, "data": {"accessToken": "at-synthetic-test"}}),
),
(
"/api/admin/accounts/quota/refresh",
json!({"accountId": "acct_error_test"}),
),
] {
let response = send(&fixture, path, body).await;
assert_eq!(response.headers()["x-request-id"], "req_admin_errors");
@@ -468,6 +504,10 @@ mod provider {
last_error_message: None,
};
AccountPageItem {
capacity: gateway_admin::model::accounts::AccountCapacity {
used_slots: None,
total_slots: None,
},
account: AccountRecord {
notes: None,
model_access: Default::default(),
+26 -1
View File
@@ -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,
@@ -861,9 +863,23 @@ fn mutation(
#[async_trait]
impl ClientKeyStore for MemoryClientKeyStore {
async fn update_client_key_budget_limits(
&self,
_: gateway_admin::model::client_keys::UpdateClientKeyBudgetLimits,
_: gateway_admin::model::client_keys::ClientKeyBudgetMutationOrigin,
_: &MutationContext,
) -> AdminStoreResult<Option<Revision>> {
Err(AdminStoreError::new(
AdminStoreErrorKind::Unavailable,
"client key",
"unused budget update",
))
}
async fn reset_client_key_budget(
&self,
command: gateway_admin::model::client_keys::ResetClientKeyBudget,
_: gateway_admin::model::client_keys::ClientKeyBudgetMutationOrigin,
_: &MutationContext,
) -> AdminStoreResult<()> {
use gateway_admin::model::client_keys::ClientKeyBudgetPeriod;
@@ -1024,6 +1040,9 @@ impl AccountStore for UnusedStore {
_: TimeRange,
_: &[String],
) -> AdminStoreResult<Vec<AccountUsage>> {
if self.account.lock().expect("account").is_some() {
return Ok(Vec::new());
}
Err(unavailable("account usage"))
}
@@ -1409,7 +1428,12 @@ impl ProviderAdmin for UnusedProvider {
&self,
_: gateway_admin::model::provider_credentials::ProviderQuotaRequest,
) -> Result<ProviderQuota, ProviderAdminError> {
Err(unsupported_provider())
Err(self
.error
.lock()
.expect("provider error")
.clone()
.unwrap_or_else(unsupported_provider))
}
async fn models(
@@ -1518,6 +1542,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);
}
@@ -143,6 +143,7 @@ fn test_tree_should_match_frozen_rust_mirror() {
"tests/openai/responses/request.rs",
"tests/openai/responses/websocket/connection.rs",
"tests/openai/responses/websocket/forward.rs",
"tests/openai/responses/websocket/interrupt.rs",
"tests/openai/responses/websocket/mod.rs",
"tests/openai/responses/websocket/protocol.rs",
"tests/openai/router.rs",
@@ -313,6 +313,10 @@ fn decoder_should_preserve_ordinary_request_headers_as_opaque_multivalues() {
HeaderValue::from_static("client-attestation"),
);
headers.insert("x-oai-is", HeaderValue::from_static("client-is"));
for name in ["x-openai-account-routing-override", "x-openai-fedramp"] {
headers.append(name, HeaderValue::from_static("first"));
headers.append(name, HeaderValue::from_static("second"));
}
headers.insert(
"x-oai-is-update",
HeaderValue::from_static("client-is-update"),
@@ -381,6 +385,8 @@ fn decoder_should_preserve_ordinary_request_headers_as_opaque_multivalues() {
"cookie",
"chatgpt-account-id",
"chatgpt-project-id",
"x-openai-account-routing-override",
"x-openai-fedramp",
// 上游指纹由运行时画像统一生成,客户端不得覆盖。
"user-agent",
"originator",
@@ -916,6 +916,33 @@ async fn queued_binary_frame_should_be_rejected_after_the_active_response_finish
server.abort();
}
#[tokio::test]
async fn requests_deferred_during_active_response_share_the_inbound_queue_limit() {
let (trace, mut socket, server) = start_active_response().await;
for _ in 0..33 {
if socket
.send(ClientMessage::Text(
json!({"type":"response.create","model":"model-a","input":"queued"})
.to_string()
.into(),
))
.await
.is_err()
{
break;
}
}
tokio::time::timeout(
Duration::from_secs(5),
trace.cancellation_observed.notified(),
)
.await
.expect("queue overload cancels the active execution");
assert_eq!(trace.starts.load(Ordering::Acquire), 1);
assert!(trace.cancelled.load(Ordering::Acquire));
server.abort();
}
#[tokio::test]
async fn client_close_should_cancel_active_execution_without_starting_queued_requests() {
let (trace, mut socket, server) = start_active_response().await;
@@ -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(
@@ -0,0 +1,240 @@
use super::*;
use gateway_core::engine::response_control::{ResponseControl, ResponseInterruptError};
#[derive(Default)]
struct InterruptProvider {
calls: AtomicUsize,
controls: Mutex<Vec<ResponseControl>>,
}
#[async_trait]
impl Provider for InterruptProvider {
fn name(&self) -> &'static str {
"openai"
}
fn catalog_generation(&self) -> ProviderCatalogGeneration {
ProviderCatalogGeneration::default()
}
async fn query_model_capabilities(
&self,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
Ok(Vec::new())
}
async fn execute(
self: Arc<Self>,
request: ProviderRequest,
context: AttemptContext,
) -> Result<ProviderStream, ProviderError> {
let number = self.calls.fetch_add(1, Ordering::SeqCst);
let id = format!("resp_interrupt_{number}");
let control = context
.response_control()
.expect("root control propagated through Core")
.clone();
let active = control.activate(id.clone()).unwrap();
self.controls.lock().unwrap().push(control);
let candidate = request.candidate();
let metadata = ProviderCallMetadata::new(
candidate.provider().clone(),
candidate.upstream_model().unwrap().clone(),
ProviderAccountId::new("acct_api_test").unwrap(),
UpstreamTransport::new("websocket").unwrap(),
);
let body = futures::stream::unfold(
(0, Some(active), id),
|(phase, mut active, id)| async move {
let (event_type, status, event) = match phase {
0 => (
"response.created",
"in_progress",
GatewayEvent::Started(ResponseMeta::new(&id, "model-a")),
),
1 => {
active.as_ref().unwrap().requested().await;
drop(active.take());
(
"response.incomplete",
"incomplete",
GatewayEvent::Completed(ResponseMeta::new(&id, "model-a")),
)
}
_ => return None,
};
let mut response = json!({"id":id,"model":"model-a","status":status,"output":[]});
if phase == 1 {
response["incomplete_details"] = json!({"reason":"interrupted"});
}
let event = ProviderEvent::canonical_with_wire(
vec![event],
ProtocolWireEvent::json(
"openai",
Some(event_type.to_owned()),
json!({
"type":event_type, "response":response
}),
)
.unwrap(),
);
Some((Ok(event), (phase + 1, active, id)))
},
);
Ok(ProviderStream::new(metadata, body, ()))
}
}
#[derive(Default)]
struct InterruptAdmissions {
active: AtomicUsize,
}
impl ClientAdmissionPort for InterruptAdmissions {
fn abandon(&self, _: &ClientApiKeyId, _: &ModelRequestId) {
self.active.fetch_sub(1, Ordering::SeqCst);
}
fn admit(
&self,
_: ClientAdmissionRequest,
) -> BoxFuture<'_, Result<ClientAdmissionDecision, ClientAdmissionError>> {
Box::pin(async {
self.active.fetch_add(1, Ordering::SeqCst);
Ok(ClientAdmissionDecision::Granted)
})
}
fn release<'a>(
&'a self,
_: &'a ClientApiKeyId,
_: &'a ModelRequestId,
) -> BoxFuture<'a, Result<bool, ClientAdmissionError>> {
Box::pin(async {
self.active.fetch_sub(1, Ordering::SeqCst);
Ok(true)
})
}
fn restore(
&self,
_: ClientAdmissionRecovery,
) -> BoxFuture<'_, Result<ClientAdmissionRestoreResult, ClientAdmissionError>> {
Box::pin(async { unreachable!("no recovery in live connection test") })
}
}
type TestSocket =
tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
struct Server(tokio::task::JoinHandle<()>);
impl Drop for Server {
fn drop(&mut self) {
self.0.abort();
}
}
async fn send(socket: &mut TestSocket, value: Value) {
socket
.send(ClientMessage::Text(value.to_string().into()))
.await
.unwrap();
}
async fn next(socket: &mut TestSocket) -> Value {
loop {
let message = tokio::time::timeout(Duration::from_secs(5), socket.next())
.await
.expect("response deadline")
.unwrap()
.unwrap();
let value: Value = serde_json::from_str(message.to_text().unwrap()).unwrap();
if value["type"] != "response.metadata" {
return value;
}
}
}
#[tokio::test]
async fn interrupt_is_live_scoped_and_keeps_queued_creates_serial() {
let provider = Arc::new(InterruptProvider::default());
let admissions = Arc::new(InterruptAdmissions::default());
let execution = Arc::new(DefaultExecutionService::new(
RuntimeSnapshotHandle::new(snapshot("sk_ws_interrupt", "openai")),
Arc::new(SettlementPorts::default()),
ProviderRegistry::new([provider.clone() as Arc<dyn Provider>]).unwrap(),
admissions.clone(),
Arc::new(UnusedContinuation),
Arc::new(IgnoredClientApiKeyUsage),
));
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let _server = Server(tokio::spawn(async move {
axum::serve(listener, app).await.unwrap()
}));
let mut request = format!("ws://{address}/v1/responses")
.into_client_request()
.unwrap();
request.headers_mut().insert(
AUTHORIZATION,
HeaderValue::from_static("Bearer sk_ws_interrupt"),
);
let (mut socket, _) = connect_async(request).await.unwrap();
let create = json!({"type":"response.create","model":"model-a","input":"hello"});
let interrupt =
|id: &str, mode: &str| json!({"type":"response.interrupt","response_id":id,"mode":mode});
send(
&mut socket,
interrupt("resp_interrupt_0", "discard_partial_items"),
)
.await;
assert_eq!(next(&mut socket).await["type"], "error");
assert_eq!(provider.calls.load(Ordering::SeqCst), 0);
send(&mut socket, create.clone()).await;
assert_eq!(next(&mut socket).await["type"], "response.created");
send(&mut socket, create).await;
for invalid in [
interrupt("resp_other", "discard_partial_items"),
interrupt("resp_interrupt_0", "unknown"),
interrupt("", "discard_partial_items"),
] {
send(&mut socket, invalid).await;
let error = next(&mut socket).await;
assert_eq!(error["type"], "error");
assert_eq!(error["status"], 400);
assert_eq!(provider.calls.load(Ordering::SeqCst), 1);
}
send(
&mut socket,
interrupt("resp_interrupt_0", "discard_partial_items"),
)
.await;
let terminal = next(&mut socket).await;
assert_eq!(terminal["type"], "response.incomplete");
assert_eq!(
terminal["response"]["incomplete_details"]["reason"],
"interrupted"
);
assert_eq!(
next(&mut socket).await["response"]["id"],
"resp_interrupt_1"
);
assert_eq!(provider.calls.load(Ordering::SeqCst), 2);
send(
&mut socket,
interrupt("resp_interrupt_0", "discard_partial_items"),
)
.await;
assert_eq!(next(&mut socket).await["type"], "error");
assert_eq!(
provider.controls.lock().unwrap()[0].interrupt("resp_interrupt_0"),
Err(ResponseInterruptError::Unavailable)
);
socket.close(None).await.unwrap();
tokio::time::timeout(Duration::from_secs(5), async {
while admissions.active.load(Ordering::SeqCst) != 0 {
tokio::task::yield_now().await;
}
})
.await
.expect("disconnect releases active execution");
assert_eq!(
provider.controls.lock().unwrap()[1].interrupt("resp_interrupt_1"),
Err(ResponseInterruptError::Unavailable)
);
}
@@ -1,5 +1,6 @@
mod connection;
mod forward;
mod interrupt;
mod protocol;
use std::sync::{
@@ -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
@@ -464,6 +479,7 @@ pub enum PreferredAccountSelection {
pub struct AccountSelection<'a> {
candidate: &'a AccountCandidate,
preferred: PreferredAccountSelection,
policy_choice: bool,
}
impl<'a> AccountSelection<'a> {
@@ -476,6 +492,12 @@ impl<'a> AccountSelection<'a> {
pub const fn preferred(self) -> PreferredAccountSelection {
self.preferred
}
/// 显式策略选号不受 Provider 的软亲和等待覆盖;委托内置选号不属于显式选择。
#[must_use]
pub const fn is_policy_choice(self) -> bool {
self.policy_choice
}
}
/// 同一 target 内唯一的账号排序器。
@@ -510,12 +532,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)
{
@@ -526,6 +543,7 @@ impl AccountSelector {
Some(AccountSelection {
candidate,
preferred,
policy_choice: true,
})
}
@@ -585,13 +603,23 @@ 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,
policy_choice: false,
});
}
}
eligible.retain(|candidate| candidate.account.weight() == highest_weight);
@@ -629,6 +657,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| {
@@ -643,6 +673,7 @@ impl AccountSelector {
Some(AccountSelection {
candidate,
preferred,
policy_choice: false,
})
}
@@ -650,7 +681,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 +693,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 +729,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 +812,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 +834,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 +860,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 +884,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,
@@ -52,6 +52,7 @@ enum AccountSelection {
}
pub(super) struct CoordinationExtensions {
response_control: Option<super::response_control::ResponseControl>,
continuation: Option<ContinuationBinding>,
observation: Option<RequestObservationDispatch>,
request_policy: Option<super::policy::RequestPolicyContext>,
@@ -65,12 +66,21 @@ pub(super) struct CoordinationExtensions {
}
impl CoordinationExtensions {
pub(super) fn with_response_control(
mut self,
control: Option<super::response_control::ResponseControl>,
) -> Self {
self.response_control = control;
self
}
pub(super) fn new(
continuation: Option<ContinuationBinding>,
observation: Option<RequestObservationDispatch>,
) -> Self {
Self {
continuation,
response_control: None,
observation,
request_policy: None,
execution_effects: None,
@@ -230,6 +240,7 @@ where
cancellation: CancellationToken,
) -> Result<ResponseExecutionSession<S>, EngineError> {
let CoordinationExtensions {
response_control,
continuation,
observation: request_observation,
request_policy,
@@ -295,6 +306,7 @@ where
.fuse(),
pending_request: Some(request),
request_persisted: false,
response_control,
operation,
plan,
request_policy,
@@ -409,6 +421,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
deadline_timer: Fuse<Delay>,
pending_request: Option<NewModelRequest>,
request_persisted: bool,
response_control: Option<super::response_control::ResponseControl>,
operation: Operation,
plan: RoutingPlan,
request_policy: Option<super::policy::RequestPolicyContext>,
@@ -996,6 +1009,7 @@ where
}
let context = AttemptContext::new(
RequestAttemptContext::new(self.request_id.clone(), self.client_api_key_ref.clone())
.with_response_control(self.response_control.clone())
.with_request_profile(self.request_profiles.get(candidate.provider()).cloned())
.with_disable_fast(self.plan.disable_fast())
.with_pricing(self.plan.pricing())
@@ -157,6 +157,7 @@ pub struct StartExecution {
/// 中间件短路时本值直接释放;调用 `next` 后必须原样交回
/// [`ExecutionService::start_prepared`],不能重新认证为另一个 Key。
pub struct PreparedRootExecution {
response_control: super::response_control::ResponseControl,
client: AuthenticatedClient,
request_id: ModelRequestId,
started_at: SystemTime,
@@ -167,6 +168,11 @@ pub struct PreparedRootExecution {
}
impl PreparedRootExecution {
#[must_use]
pub fn response_control(&self) -> super::response_control::ResponseControl {
self.response_control.clone()
}
fn new(client: AuthenticatedClient) -> Result<Self, GatewayError> {
let started_at = SystemTime::now();
let deadline_at = started_at
@@ -178,6 +184,7 @@ impl PreparedRootExecution {
let execution_effects_baseline = execution_effects.epoch();
Ok(Self {
client,
response_control: super::response_control::ResponseControl::default(),
request_id: new_request_id()?,
started_at,
deadline_at,
@@ -292,6 +299,7 @@ struct PendingStartExecution {
}
struct AuthorizedExecution {
response_control: Option<super::response_control::ResponseControl>,
account_scope: Arc<FrozenAccountScope>,
deadline_at: SystemTime,
cancellation: CancellationToken,
@@ -894,6 +902,7 @@ impl DefaultExecutionService {
) -> Result<StartedExecution, GatewayError> {
let PreparedRootExecution {
client,
response_control,
request_id,
started_at,
deadline_at,
@@ -902,6 +911,7 @@ impl DefaultExecutionService {
execution_effects_baseline,
} = prepared;
let authorization = AuthorizedExecution {
response_control: Some(response_control),
account_scope: Arc::clone(client.policy.account_scope()),
deadline_at,
cancellation,
@@ -1155,6 +1165,7 @@ impl DefaultExecutionService {
.map(|graph| Arc::clone(&graph.effects));
let extensions =
CoordinationExtensions::new(continuation, request_observation.clone())
.with_response_control(authorization.response_control.clone())
.with_request_policy(request_policy)
.with_execution_effects(
execution_effects,
@@ -1511,6 +1522,7 @@ impl DefaultExecutionService {
metadata.user_agent = None;
let authorization = AuthorizedExecution {
account_scope,
response_control: None,
deadline_at: parent.deadline_at,
cancellation: parent.cancellation.child_token(),
extension_scope,
@@ -1676,6 +1688,7 @@ impl DefaultExecutionService {
},
AuthorizedExecution {
account_scope,
response_control: None,
deadline_at: authority.deadline_at,
cancellation: authority.cancellation.child_token(),
extension_scope: authority.extension_scope.clone(),
@@ -2118,9 +2131,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;
}
@@ -14,6 +14,7 @@ pub mod observation;
pub mod policy;
pub mod probe;
pub mod provider;
pub mod response_control;
pub use coordinator::{AttemptCoordinator, ResponseExecutionSession};
@@ -323,6 +324,7 @@ impl ContinuationAttempt {
/// Provider 每次执行可见的 request-local context。
#[derive(Debug, Clone)]
pub struct RequestAttemptContext {
response_control: Option<response_control::ResponseControl>,
pricing: Arc<crate::metering::PricingOverrides>,
request_profile: Option<crate::account::OpaqueProviderData>,
disable_fast: bool,
@@ -343,6 +345,15 @@ pub struct RequestAttemptContext {
}
impl RequestAttemptContext {
#[must_use]
pub fn with_response_control(
mut self,
control: Option<response_control::ResponseControl>,
) -> Self {
self.response_control = control;
self
}
#[must_use]
pub fn with_pricing(mut self, pricing: Arc<crate::metering::PricingOverrides>) -> Self {
self.pricing = pricing;
@@ -376,6 +387,7 @@ impl RequestAttemptContext {
#[must_use]
pub fn new(request_id: ModelRequestId, client_api_key_ref: ClientApiKeyId) -> Self {
Self {
response_control: None,
request_id,
client_api_key_ref,
request_profile: None,
@@ -527,6 +539,11 @@ pub struct AttemptContext {
}
impl AttemptContext {
#[must_use]
pub fn response_control(&self) -> Option<&response_control::ResponseControl> {
self.request.response_control.as_ref()
}
#[must_use]
pub fn pricing(&self) -> &crate::metering::PricingOverrides {
&self.request.pricing
@@ -0,0 +1,80 @@
//! 当前执行的响应中断信号;只由持有活动响应的 Provider 注册,不查找其他执行或连接。
use std::fmt;
use std::sync::{Arc, Mutex, Weak};
use crate::lifecycle::CancellationToken;
/// 单次根执行的控制入口;重试不继承已结束响应的控制目标。
#[derive(Clone, Default)]
pub struct ResponseControl {
active: Arc<Mutex<Weak<InterruptTarget>>>,
}
impl fmt::Debug for ResponseControl {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("ResponseControl([REDACTED])")
}
}
struct InterruptTarget {
response_id: String,
requested: CancellationToken,
}
/// Provider 持有的活动响应控制权;释放后客户端无法再向该响应发出中断。
pub struct ActiveResponseInterrupt(Arc<InterruptTarget>);
/// 控制帧不能作用于当前执行时的稳定原因,不携带其他响应身份。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ResponseInterruptError {
Unavailable,
ResponseMismatch,
}
impl ResponseControl {
/// 在上游确认响应 ID 后注册;同一执行只允许一个存活的响应 owner。
pub fn activate(&self, response_id: String) -> Option<ActiveResponseInterrupt> {
let mut active = self
.active
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if active.upgrade().is_some() || response_id.is_empty() {
return None;
}
let target = Arc::new(InterruptTarget {
response_id,
requested: CancellationToken::new(),
});
*active = Arc::downgrade(&target);
Some(ActiveResponseInterrupt(target))
}
/// 重复中断合并为一次信号;ID 必须匹配本执行当前的上游响应。
pub fn interrupt(&self, response_id: &str) -> Result<(), ResponseInterruptError> {
let active = self
.active
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let target = active
.upgrade()
.ok_or(ResponseInterruptError::Unavailable)?;
if target.response_id != response_id {
return Err(ResponseInterruptError::ResponseMismatch);
}
target.requested.cancel();
Ok(())
}
}
impl ActiveResponseInterrupt {
pub fn response_id(&self) -> &str {
&self.0.response_id
}
pub fn requested(&self) -> impl Future<Output = ()> + Send + 'static {
// 等待者只持有信号,不能延长活动响应 owner 的生命周期。
let requested = self.0.requested.clone();
async move { requested.cancelled().await }
}
}
@@ -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"));
});
}
@@ -8,6 +8,7 @@ mod middleware;
mod policy;
mod probe;
mod provider;
mod response_control;
use gateway_core::engine::AttemptTrigger;
@@ -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 {
@@ -0,0 +1,43 @@
use futures::FutureExt;
use gateway_core::engine::response_control::{ResponseControl, ResponseInterruptError};
#[test]
fn interrupt_requires_the_current_response_owner_and_does_not_extend_its_lifetime() {
let control = ResponseControl::default();
assert_eq!(
control.interrupt("resp_a"),
Err(ResponseInterruptError::Unavailable)
);
assert!(control.activate(String::new()).is_none());
let active = control.activate("resp_a".to_owned()).unwrap();
let requested = active.requested();
assert!(control.activate("resp_b".to_owned()).is_none());
assert_eq!(
control.interrupt("resp_b"),
Err(ResponseInterruptError::ResponseMismatch)
);
assert!(active.requested().now_or_never().is_none());
control.interrupt("resp_a").unwrap();
control.interrupt("resp_a").unwrap();
assert_eq!(requested.now_or_never(), Some(()));
let retained_waiter = active.requested();
drop(active);
assert_eq!(
control.interrupt("resp_a"),
Err(ResponseInterruptError::Unavailable)
);
let next = control.activate("resp_b".to_owned()).unwrap();
assert!(next.requested().now_or_never().is_none());
drop(retained_waiter);
}
#[test]
fn matching_response_ids_do_not_share_control_between_root_executions() {
let first = ResponseControl::default();
let second = ResponseControl::default();
let active_first = first.activate("resp_same".to_owned()).unwrap();
let active_second = second.activate("resp_same".to_owned()).unwrap();
first.interrupt("resp_same").unwrap();
assert_eq!(active_first.requested().now_or_never(), Some(()));
assert!(active_second.requested().now_or_never().is_none());
}
@@ -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;
@@ -237,54 +300,3 @@ async fn oversized_metadata_is_rejected_before_json_parsing() {
.is_err()
);
}
#[tokio::test]
#[ignore = "实际访问 GitHub 公共仓库;通过 --ignored 显式运行"]
async fn live_github_resolves_latest_then_pins_the_same_release() {
let distribution = gateway_host::plugin_distribution::HttpPluginDistribution::new(
std::sync::Arc::new(gateway_host::outbound::HttpClient::new().unwrap()),
)
.unwrap();
let mut query = GithubReleaseQuery {
repository: "cli/cli".into(),
tag: None,
allow_prerelease: false,
};
let latest = distribution
.query_release(query.clone(), vec![], None)
.await
.unwrap();
assert!(!latest.prerelease);
assert!(!latest.assets.is_empty());
query.tag = Some(latest.tag.clone());
let pinned = distribution
.query_release(query, vec![], None)
.await
.unwrap();
assert_eq!(pinned.tag, latest.tag);
let checksums = pinned
.assets
.iter()
.find(|asset| asset.name.ends_with("checksums.txt"))
.unwrap();
let expected = checksums
.sha256
.clone()
.expect("GitHub Release checksum asset must have a published digest");
let downloaded = distribution
.download(
RemotePluginLocation::Github {
repository: pinned.repository,
tag: pinned.tag,
asset: checksums.name.clone(),
allow_prerelease: false,
sha256: Some(expected.clone()),
},
vec![],
None,
)
.await
.unwrap();
assert_eq!(downloaded.sha256, expected);
assert!(!downloaded.archive.is_empty());
}
@@ -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,10 @@
"accounts",
"data",
"requests",
"public_endpoints"
"public_endpoints",
"groups",
"keys",
"key_budgets",
"quota_observations"
]
}
@@ -5,7 +5,7 @@ use std::sync::{Arc, OnceLock, Weak};
use chrono::{DateTime, Utc};
use gateway_admin::{
model::{
AdminError, AdminErrorKind, MutationActor, MutationContext, PageSize, Revision,
AdminError, PageSize, Revision,
plugins::instances::PluginPermissionGrant,
provider_credentials::{
PluginAccountListQuery, PreparedCredentialCreate, PreparedCredentialRotationFacts,
@@ -26,7 +26,11 @@ use gateway_plugin_sdk::{
},
};
use super::{NetworkScope, invalid};
use super::{
NetworkScope,
admin::{encode, map_admin_error, mutation_context},
invalid,
};
use crate::RpcReply;
pub(crate) struct PluginAccountPortSlot {
@@ -304,17 +308,6 @@ fn credential_facts(
}
}
fn mutation_context(context: &CallContext) -> MutationContext {
MutationContext {
actor: MutationActor::System,
// RPC 回调使用真实 call_id,RPC 结果的后置提交仍使用原始上下文;scope 由宿主签发且全局唯一。
request_id: format!(
"plugin:{}:scope:{}:call:{}",
context.instance_id, context.resource_scope_id, context.call_id
),
}
}
fn decode<T: serde::de::DeserializeOwned>(payload: &[u8]) -> Result<T, PluginFault> {
serde_json::from_slice(payload).map_err(|_| invalid())
}
@@ -347,30 +340,9 @@ fn credential_timestamp(value: Option<i64>) -> Result<Option<DateTime<Utc>>, Plu
.transpose()
}
fn encode(value: &impl serde::Serialize) -> Result<RpcReply, PluginFault> {
Ok(RpcReply {
result: serde_json::json!({}),
payload: serde_json::to_vec(value).map_err(|_| invalid())?,
})
}
fn denied() -> PluginFault {
PluginFault::new(
ErrorCode::PermissionDenied,
"account callback is not authorized",
)
}
pub(super) fn map_admin_error(error: AdminError) -> PluginFault {
let code = match error.kind() {
AdminErrorKind::Invalid => ErrorCode::InvalidInput,
AdminErrorKind::Unauthorized | AdminErrorKind::Forbidden => ErrorCode::PermissionDenied,
AdminErrorKind::NotFound => ErrorCode::Rejected,
AdminErrorKind::Conflict => ErrorCode::Conflict,
AdminErrorKind::RateLimited => ErrorCode::Capacity,
AdminErrorKind::UpstreamResultUnknown => ErrorCode::Uncertain,
AdminErrorKind::BadGateway => ErrorCode::Upstream,
AdminErrorKind::Unavailable | AdminErrorKind::Internal => ErrorCode::Fault,
};
PluginFault::new(code, "account callback failed")
}
@@ -0,0 +1,38 @@
//! Admin 端口回调共享的审计上下文、结果编码与错误映射。
use gateway_admin::model::{AdminError, AdminErrorKind, MutationActor, MutationContext};
use gateway_plugin_sdk::{CallContext, ErrorCode, PluginFault};
use crate::RpcReply;
pub(super) fn mutation_context(context: &CallContext) -> MutationContext {
MutationContext {
actor: MutationActor::System,
// 回调沿用宿主签发的调用上下文,插件不能自报审计身份。
request_id: format!(
"plugin:{}:scope:{}:call:{}",
context.instance_id, context.resource_scope_id, context.call_id
),
}
}
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(|_| super::invalid())?,
})
}
pub(super) fn map_admin_error(error: AdminError) -> PluginFault {
let code = match error.kind() {
AdminErrorKind::Invalid => ErrorCode::InvalidInput,
AdminErrorKind::Unauthorized | AdminErrorKind::Forbidden => ErrorCode::PermissionDenied,
AdminErrorKind::NotFound => ErrorCode::Rejected,
AdminErrorKind::Conflict => ErrorCode::Conflict,
AdminErrorKind::RateLimited => ErrorCode::Capacity,
AdminErrorKind::UpstreamResultUnknown => ErrorCode::Uncertain,
AdminErrorKind::BadGateway => ErrorCode::Upstream,
AdminErrorKind::Unavailable | AdminErrorKind::Internal => ErrorCode::Fault,
};
PluginFault::new(code, "admin callback failed")
}
@@ -4,28 +4,35 @@ use gateway_admin::model::{
PageSize, plugins::instances::PluginPermissionGrant,
provider_credentials::PluginAccountListQuery,
};
use gateway_core::{account::ProviderAccountId, routing::ProviderKind};
use gateway_core::{account::ProviderAccountId, policy::ClientApiKeyId, routing::ProviderKind};
use gateway_plugin_sdk::{CallContext, PluginFault, Stage, call::data};
use super::{
accounts::{PluginAccountPortSlot, map_admin_error},
denied, invalid,
accounts::PluginAccountPortSlot, admin::map_admin_error, denied, invalid,
keys::PluginClientKeyPortSlot,
};
use crate::RpcReply;
pub(super) struct PluginData {
accounts: Arc<PluginAccountPortSlot>,
authorized: bool,
keys: Arc<PluginClientKeyPortSlot>,
data_authorized: bool,
quota_authorized: bool,
}
impl PluginData {
pub(super) fn new(
accounts: Arc<PluginAccountPortSlot>,
keys: Arc<PluginClientKeyPortSlot>,
grants: &[PluginPermissionGrant],
) -> Self {
Self {
accounts,
authorized: grants.iter().any(|grant| grant.permission == "data"),
keys,
data_authorized: grants.iter().any(|grant| grant.permission == "data"),
quota_authorized: grants
.iter()
.any(|grant| grant.permission == "quota_observations"),
}
}
@@ -36,16 +43,48 @@ impl PluginData {
params: serde_json::Value,
payload: &[u8],
) -> Result<RpcReply, PluginFault> {
// 管理范围的只读授权不继承到客户端请求链,避免借数据查询扩大当前 Key 的范围。
if !self.authorized || !matches!(context.stage, Stage::Management | Stage::CommandLine) {
// 基础事实读取不隐含上游访问权;管理范围也不继承到客户端请求链。
let authorized = match method {
data::ACCOUNTS_LIST | data::KEYS_GET => self.data_authorized,
data::QUOTA_GET => self.data_authorized || self.quota_authorized,
data::QUOTA_REFRESH => self.quota_authorized,
_ => false,
};
if !authorized
|| !matches!(
context.stage,
Stage::Management | Stage::CommandLine | Stage::Maintenance
)
{
return Err(denied());
}
if params != serde_json::json!({}) {
return Err(invalid());
}
let accounts = self.accounts.upgrade().map_err(map_admin_error)?;
let payload = match method {
data::KEYS_GET => {
let query: data::ClientKeyFactsQuery =
serde_json::from_slice(payload).map_err(|_| invalid())?;
let id = ClientApiKeyId::new(query.client_key_id).map_err(|_| invalid())?;
let key = self
.keys
.upgrade()?
.facts(&id)
.await
.map_err(map_admin_error)?;
serde_json::to_vec(&data::ClientKeyFacts {
schema_version: 1,
client_key_id: key.id.as_str().to_owned(),
enabled: key.enabled,
group_ids: key
.group_ids
.into_iter()
.map(|id| id.as_str().to_owned())
.collect(),
})
}
data::ACCOUNTS_LIST => {
let accounts = self.accounts.upgrade().map_err(map_admin_error)?;
let query: data::AccountFactsQuery =
serde_json::from_slice(payload).map_err(|_| invalid())?;
let page = accounts
@@ -84,14 +123,17 @@ impl PluginData {
next_cursor: page.next_cursor.map(|id| id.as_str().to_owned()),
})
}
data::QUOTA_GET => {
data::QUOTA_GET | data::QUOTA_REFRESH => {
let accounts = self.accounts.upgrade().map_err(map_admin_error)?;
let query: data::QuotaFactsQuery =
serde_json::from_slice(payload).map_err(|_| invalid())?;
let account_id = ProviderAccountId::new(query.account_id).map_err(|_| invalid())?;
let quota = accounts
.get_quota(&account_id)
.await
.map_err(map_admin_error)?;
let quota = if method == data::QUOTA_REFRESH {
accounts.refresh_quota(&account_id).await
} else {
accounts.get_quota(&account_id).await
}
.map_err(map_admin_error)?;
serde_json::to_vec(&data::QuotaFacts {
schema_version: 1,
account_id: account_id.as_str().to_owned(),
@@ -1,4 +1,4 @@
//! Key 目录只传递非秘密身份,模型规则仍由 Core 在执行时复核。
//! Key 目录与预算回调;实例身份由宿主冻结,写入授权在存储事务复验。
use std::sync::{Arc, OnceLock, Weak};
@@ -6,19 +6,127 @@ use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use gateway_admin::{
model::{
AdminError, PageSize,
client_keys::{ClientKeyBudgetPeriod, ResetClientKeyBudget, UpdateClientKeyBudgetLimits},
plugin_client_keys::{PluginClientKeyCursor, PluginClientKeyListQuery},
plugin_resources::PluginResourceOwner,
plugins::instances::PluginInstance,
},
ports::plugin_client_keys::PluginClientKeyAccess,
};
use gateway_core::policy::ClientApiKeyId;
use gateway_plugin_sdk::{
PluginFault,
call::host::{ClientKey, KeyListRequest, KeyListResult},
CallContext, PluginFault,
call::{
host::{ClientKey, KeyListRequest, KeyListResult},
key_budgets,
},
};
use super::{denied, invalid};
use super::{
admin::{encode, map_admin_error, mutation_context},
denied, invalid,
};
use crate::RpcReply;
pub(super) struct PluginClientKeys {
owner: PluginResourceOwner,
ports: Arc<PluginClientKeyPortSlot>,
}
impl PluginClientKeys {
pub(super) fn new(instance: &PluginInstance, ports: Arc<PluginClientKeyPortSlot>) -> 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 method == "host.keys.list" {
return self.ports.list(params, payload).await;
}
if params != serde_json::json!({}) {
return Err(invalid());
}
let access = self.ports.upgrade()?;
match method {
key_budgets::GET => {
let request: key_budgets::GetKeyBudgetRequest =
serde_json::from_slice(payload).map_err(|_| invalid())?;
let id = ClientApiKeyId::new(request.client_key_id).map_err(|_| invalid())?;
let budget = access.budget(&id).await.map_err(map_admin_error)?;
encode(&key_budgets::KeyBudget {
client_key_id: id.as_str().to_owned(),
daily_limit_usd: budget.limits.daily_usd.canonical(),
weekly_limit_usd: budget.limits.weekly_usd.canonical(),
daily_used_usd: budget.daily_used_usd.canonical(),
weekly_used_usd: budget.weekly_used_usd.canonical(),
daily_resets_at_ms: budget
.daily_resets_at
.map(|time| chrono::DateTime::<chrono::Utc>::from(time).timestamp_millis()),
weekly_resets_at_ms: budget
.weekly_resets_at
.map(|time| chrono::DateTime::<chrono::Utc>::from(time).timestamp_millis()),
})
}
key_budgets::UPDATE_LIMITS => {
let request: key_budgets::UpdateKeyBudgetLimitsRequest =
serde_json::from_slice(payload).map_err(|_| invalid())?;
let command = UpdateClientKeyBudgetLimits {
id: ClientApiKeyId::new(request.client_key_id).map_err(|_| invalid())?,
daily_limit_usd: request
.daily_limit_usd
.map(|value| value.parse())
.transpose()
.map_err(|_| invalid())?,
weekly_limit_usd: request
.weekly_limit_usd
.map(|value| value.parse())
.transpose()
.map_err(|_| invalid())?,
};
let id = access
.update_budget_limits(&self.owner, command, &mutation_context(context))
.await
.map_err(map_admin_error)?;
encode(&key_budgets::UpdateKeyBudgetLimitsResult {
client_key_id: id.as_str().to_owned(),
})
}
key_budgets::RESET => {
let request: key_budgets::ResetKeyBudgetRequest =
serde_json::from_slice(payload).map_err(|_| invalid())?;
let command = ResetClientKeyBudget {
id: ClientApiKeyId::new(request.client_key_id).map_err(|_| invalid())?,
period: match request.period {
key_budgets::BudgetPeriod::Daily => ClientKeyBudgetPeriod::Daily,
key_budgets::BudgetPeriod::Weekly => ClientKeyBudgetPeriod::Weekly,
key_budgets::BudgetPeriod::All => ClientKeyBudgetPeriod::All,
},
};
let id = access
.reset_budget(&self.owner, command, &mutation_context(context))
.await
.map_err(map_admin_error)?;
encode(&key_budgets::ResetKeyBudgetResult {
client_key_id: id.as_str().to_owned(),
})
}
_ => Err(denied()),
}
}
}
pub(crate) struct PluginClientKeyPortSlot {
access: OnceLock<Weak<dyn PluginClientKeyAccess>>,
}
@@ -33,7 +141,11 @@ impl PluginClientKeyPortSlot {
pub(crate) fn bind(&self, access: &Arc<dyn PluginClientKeyAccess>) -> Result<(), AdminError> {
self.access
.set(Arc::downgrade(access))
.map_err(|_| AdminError::conflict("插件 Key 目录端口已经绑定"))
.map_err(|_| AdminError::conflict("插件 Client Key 端口已经绑定"))
}
pub(super) fn upgrade(&self) -> Result<Arc<dyn PluginClientKeyAccess>, PluginFault> {
self.access.get().and_then(Weak::upgrade).ok_or_else(denied)
}
pub(super) async fn list(
@@ -61,15 +173,11 @@ impl PluginClientKeyPortSlot {
})
})
.transpose()?;
let access = self
.access
.get()
.and_then(Weak::upgrade)
.ok_or_else(denied)?;
let access = self.upgrade()?;
let page = access
.list(PluginClientKeyListQuery { cursor, limit })
.await
.map_err(super::accounts::map_admin_error)?;
.map_err(map_admin_error)?;
let next_cursor = page
.next_cursor
.map(|cursor| {
@@ -137,7 +137,8 @@ fn connection_options(headers: &[MiddlewareHeader]) -> Vec<String> {
}
fn sensitive_or_identity_header(name: &str) -> bool {
name.contains("auth")
name == "x-openai-fedramp"
|| name.contains("auth")
|| name.contains("credential")
|| name.contains("secret")
|| name.contains("token")
@@ -1,4 +1,5 @@
mod accounts;
mod admin;
mod affinity;
mod data;
mod http;
@@ -7,6 +8,7 @@ mod log;
mod middleware;
mod model;
pub(crate) mod private_state;
mod resources;
mod scope;
use std::{
@@ -16,7 +18,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 +32,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 +42,7 @@ pub(crate) struct PluginCallbackPorts {
keys: Arc<PluginClientKeyPortSlot>,
models: Arc<PluginModelPortSlot>,
affinity: Arc<PluginAffinityPortSlot>,
resources: Arc<PluginResourcePorts>,
}
impl PluginCallbackPorts {
@@ -49,6 +53,7 @@ impl PluginCallbackPorts {
keys: Arc<PluginClientKeyPortSlot>,
models: Arc<PluginModelPortSlot>,
affinity: Arc<PluginAffinityPortSlot>,
resources: Arc<PluginResourcePorts>,
) -> Self {
Self {
http,
@@ -57,15 +62,16 @@ 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,
keys: Arc<keys::PluginClientKeys>,
affinity: Arc<affinity::PluginAffinity>,
log: Arc<log::PluginLog>,
models: Arc<model::PluginModels>,
@@ -117,18 +123,23 @@ 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 {
data: Arc::new(data::PluginData::new(ports.accounts.clone(), grants)),
resources: Arc::new(resources::PluginResources::new(instance, ports.resources)),
data: Arc::new(data::PluginData::new(
ports.accounts.clone(),
ports.keys.clone(),
grants,
)),
accounts: Arc::new(accounts::PluginAccounts::new(ports.accounts, grants)),
keys: ports.keys,
models_authorized: grants.iter().any(|grant| grant.permission == "models"),
keys: Arc::new(keys::PluginClientKeys::new(instance, ports.keys)),
affinity: Arc::new(affinity::PluginAffinity::new(ports.affinity, grants)),
log: Arc::new(
log::PluginLog::new(manifest, log_slots)
@@ -374,9 +385,9 @@ 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;
let models = self.models.clone();
let affinity = self.affinity.clone();
Box::pin(async move {
@@ -417,11 +428,22 @@ impl CallbackHandler for PluginCallbacks {
if is_private_state {
return private_state.call(&method, params, &payload).await;
}
if method == "host.keys.list" {
if !models_authorized {
return Err(denied());
}
return keys.list(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 matches!(
method.as_str(),
"host.keys.list"
| gateway_plugin_sdk::call::key_budgets::RESET
| gateway_plugin_sdk::call::key_budgets::GET
| gateway_plugin_sdk::call::key_budgets::UPDATE_LIMITS
) {
return keys.call(&context, &method, params, &payload).await;
}
if method == "host.models.list" {
return models.call(&context, &call, &method, params, payload).await;
@@ -444,7 +466,9 @@ impl CallbackHandler for PluginCallbacks {
if matches!(
method.as_str(),
gateway_plugin_sdk::call::data::ACCOUNTS_LIST
| gateway_plugin_sdk::call::data::KEYS_GET
| gateway_plugin_sdk::call::data::QUOTA_GET
| gateway_plugin_sdk::call::data::QUOTA_REFRESH
) {
return data.call(&context, &method, params, &payload).await;
}
@@ -0,0 +1,147 @@
use super::{
admin::{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(),
@@ -261,7 +265,7 @@ impl PluginRuntime {
self.account_ports.bind(access)
}
/// Key 目录由 Admin 组合并保活,Runtime 不取得明文访问端口。
/// Key 目录与预算端口由 Admin 组合并保活,Runtime 不取得明文访问端口。
pub fn bind_client_key_ports(
&self,
access: &Arc<dyn PluginClientKeyAccess>,
@@ -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,38 @@ 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.keys.get"
| "host.data.quota.get"
| "host.quota_observations.refresh"
| "host.groups.ensure"
| "host.groups.change_members"
| "host.keys.ensure"
| "host.keys.list"
| "host.keys.reset_budget"
| "host.keys.get_budget"
| "host.keys.update_budget_limits"
)
{
return false;
}
if method == "host.data.quota.get" {
return matches!(
stage,
Stage::Management | Stage::CommandLine | Stage::Maintenance
) && (permissions.contains(&Permission::Data)
|| permissions.contains(&Permission::QuotaObservations));
}
if method == "host.keys.list" {
return (stage != Stage::Maintenance && permissions.contains(&Permission::Models))
|| (matches!(
stage,
Stage::Management | Stage::CommandLine | Stage::Maintenance
) && permissions.contains(&Permission::KeyBudgets));
}
let permission = match method {
"host.http.do"
| "host.http.do_stream"
@@ -229,17 +261,51 @@ fn callback_allowed(method: &str, stage: Stage, permissions: &[Permission]) -> b
| "host.model.execute_stream"
| "host.model.stream_read"
| "host.model.stream_close"
| "host.models.list"
| "host.keys.list" => Permission::Models,
| "host.models.list" => Permission::Models,
"host.auth.list" | "host.auth.get_runtime" | "host.auth.get" | "host.auth.save" => {
Permission::Accounts
}
"host.affinity.lookup" => Permission::Requests,
"host.data.accounts.list" | "host.data.quota.get"
if matches!(stage, Stage::Management | Stage::CommandLine) =>
"host.data.accounts.list" | "host.data.keys.get"
if matches!(
stage,
Stage::Management | Stage::CommandLine | Stage::Maintenance
) =>
{
Permission::Data
}
"host.quota_observations.refresh"
if matches!(
stage,
Stage::Management | Stage::CommandLine | Stage::Maintenance
) =>
{
Permission::QuotaObservations
}
"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
}
"host.keys.reset_budget" | "host.keys.get_budget" | "host.keys.update_budget_limits"
if matches!(
stage,
Stage::Management | Stage::CommandLine | Stage::Maintenance
) =>
{
Permission::KeyBudgets
}
_ => return false,
};
permissions.contains(&permission)
@@ -0,0 +1,110 @@
use gateway_admin::{
model::{
client_keys::{SetClientKeyEnabled, UpdateClientKey},
plugins::{instances::PluginPermissionGrant, management::PluginManagementRequest},
},
ports::plugin_management::PluginManagement,
};
use gateway_core::policy::{ClientApiKeyId, RateLimits};
use serde_json::{Value, json};
#[tokio::test]
async fn key_facts_read_current_database_groups_without_secrets() {
let Some(environment) = crate::support::environment::Environment::create().await else {
eprintln!("SKIP: plugin integration environment absent");
return;
};
let account = environment.account(None).await;
let group = environment.account_group_with_account(&account).await;
let id = "key_facts_fixture";
environment.client_key(id, "sk-test-key-facts-secret").await;
let (runtime, core) = environment.plugin(json!({
"management_registration":{"routes":[{"method":"GET","path":"facts","request_content_types":[],"response_content_types":["application/json"]}]},
"data_queries":[
{"method":"host.data.keys.get","query":{"client_key_id":id}},
{"method":"host.data.keys.get","query":{"client_key_id":"key_missing"}},
{"method":"host.data.keys.get","query":{"client_key_id":id,"secret":true}}
]
}), vec![PluginPermissionGrant { permission: "data".into() }]).await;
let access = gateway_admin::initialize_plugin_client_keys(
crate::support::native::admin_registry(),
environment.store.admin_ports().client_keys(),
core.snapshot_control(),
);
runtime.bind_client_key_ports(&access).unwrap();
let generation = core
.snapshots()
.acquire()
.unwrap()
.extensions()
.unwrap()
.clone();
let view = runtime.views(&generation).await.unwrap().remove(0);
for (groups, enabled) in [(vec![], true), (vec![group.clone()], false), (vec![], true)] {
environment
.set_account_group_enabled(group.clone(), enabled)
.await;
environment
.store
.admin_ports()
.client_keys()
.update_client_key(
UpdateClientKey {
request_profile_override_updates: Default::default(),
id: ClientApiKeyId::new(id).unwrap(),
name: "facts".into(),
label: None,
group_ids: groups.clone(),
limits: RateLimits::unlimited(),
daily_limit_usd: None,
weekly_limit_usd: None,
},
&crate::support::environment::mutation(),
)
.await
.unwrap();
environment
.store
.admin_ports()
.client_keys()
.set_client_key_enabled(
SetClientKeyEnabled {
id: ClientApiKeyId::new(id).unwrap(),
enabled,
},
&crate::support::environment::mutation(),
)
.await
.unwrap();
let response = runtime
.handle(
&generation,
&view.target,
PluginManagementRequest {
method: "GET".into(),
path: "facts".into(),
query: String::new(),
content_type: None,
body: vec![],
request_id: "key-facts".into(),
},
)
.await
.unwrap();
let result: Value = serde_json::from_slice(&response.body).unwrap();
assert_eq!(
result[0],
json!({"schema_version":1,"client_key_id":id,"enabled":enabled,
"group_ids": groups.iter().map(|id| id.as_str()).collect::<Vec<_>>() })
);
assert_eq!(result[1]["error"], "rejected");
assert_eq!(result[2]["error"], "invalid_input");
}
drop(generation);
environment.release_plugin_accounts(&runtime);
runtime.shutdown().await;
drop(core);
drop(access);
drop(runtime);
environment.close().await;
}
@@ -0,0 +1,460 @@
use std::{sync::Arc, time::Duration};
use gateway_admin::{
PluginManagementService,
model::{
AdminError, MutationContext,
client_keys::ResetClientKeyBudget,
plugin_client_keys::{PluginClientKeyFacts, PluginClientKeyListQuery, PluginClientKeyPage},
plugin_resources::PluginResourceOwner,
plugins::management::PluginManagementRequest,
},
ports::plugin_client_keys::PluginClientKeyAccess,
};
use gateway_core::{
lifecycle::CancellationToken,
policy::ClientApiKeyId,
task::{WorkerContribution, WorkerRunnable},
};
use serde_json::{Value, json};
use tokio::{sync::Notify, time::timeout};
use crate::support::{
environment::{Environment, account_grant},
native,
};
#[tokio::test]
async fn plugin_process_manages_native_budget_through_client_key_service() {
for granted in [false, true] {
let Some(environment) = Environment::create().await else {
eprintln!("SKIP: plugin integration environment absent");
return;
};
let key = ClientApiKeyId::new("key_budget").unwrap();
let plaintext = format!("sk_{}", "a".repeat(43));
environment.client_key(key.as_str(), &plaintext).await;
environment.seed_client_key_budget(key.as_str()).await;
let store = environment.store.admin_ports().client_keys();
let before = store.get_client_key(&key).await.unwrap().unwrap();
environment.install_plugin(json!({
"management_registration":{"routes":[{"method":"POST","path":"reset","request_content_types":[],"response_content_types":["application/json"]}]},
"data_queries":[
{"method":"host.keys.list","query":{"limit":10}},
{"method":"host.keys.reset_budget","query":{"client_key_id":"key_budget","period":"weekly"}},
{"method":"host.keys.reset_budget","query":{"client_key_id":"key_budget"}},
{"method":"host.keys.reset_budget","query":{"client_key_id":"key_budget","period":"monthly"}},
{"method":"host.keys.reset_budget","query":{"client_key_id":"key_budget","period":"all","instance_id":"forged"}},
{"method":"host.keys.reset_budget","query":{"client_key_id":"missing","period":"all"}},
{"method":"host.keys.get_budget","query":{"client_key_id":"key_budget"}},
{"method":"host.keys.update_budget_limits","query":{"client_key_id":"key_budget","weekly_limit_usd":"25.5"}},
{"method":"host.keys.get_budget","query":{"client_key_id":"key_budget"}},
{"method":"host.keys.update_budget_limits","query":{"client_key_id":"key_budget"}},
{"method":"host.keys.update_budget_limits","query":{"client_key_id":"key_budget","daily_limit_usd":"-1"}},
{"method":"host.keys.get_budget","query":{"client_key_id":"missing"}},
{"method":"host.keys.update_budget_limits","query":{"client_key_id":"missing","weekly_limit_usd":"2"}},
{"method":"host.keys.update_budget_limits","query":{"client_key_id":"key_budget","max_concurrency":1}},
{"method":"host.keys.get_budget","query":{"client_key_id":"key_budget","instance_id":"forged"}}
]
}), if granted { vec![account_grant("key_budgets")] } else { vec![] }).await;
let (runtime, core) = environment.runtime().await;
let access = gateway_admin::initialize_plugin_client_keys(
native::admin_registry(),
store.clone(),
core.snapshot_control(),
);
runtime.bind_client_key_ports(&access).unwrap();
let service = PluginManagementService::new(
runtime.clone(),
environment.store.admin_ports().plugins(),
core.snapshots(),
);
let view = service.views().await.unwrap().remove(0);
let reply = service
.handle(
&view.target,
PluginManagementRequest {
method: "POST".into(),
path: "reset".into(),
query: String::new(),
content_type: None,
body: vec![],
request_id: "budget-fixture".into(),
},
)
.await
.unwrap();
let results: Vec<Value> = serde_json::from_slice(&reply.body).unwrap();
let after = store.get_client_key(&key).await.unwrap().unwrap();
let audits = environment.audit_requests("reset_budget").await;
if granted {
assert_eq!(
results[0],
json!({"keys":[{"id":"key_budget","name":"fixture key_budget","enabled":true}],"next_cursor":null})
);
assert_eq!(results[1], json!({"client_key_id":"key_budget"}));
for result in &results[2..5] {
assert_eq!(result, &json!({"error":"invalid_input"}));
}
assert_eq!(results[5], json!({"error":"rejected"}));
assert_eq!(after.budget.weekly_used_usd.canonical(), "0");
assert_eq!(after.budget.daily_used_usd, before.budget.daily_used_usd);
assert_eq!(
after.budget.limits.daily_usd,
before.budget.limits.daily_usd
);
assert_eq!(after.budget.limits.weekly_usd.canonical(), "25.5");
let mut expected = json!({
"client_key_id":"key_budget", "daily_limit_usd":"10", "weekly_limit_usd":"20",
"daily_used_usd":"3", "weekly_used_usd":"0",
"daily_resets_at_ms":before.budget.daily_resets_at.map(|time| chrono::DateTime::<chrono::Utc>::from(time).timestamp_millis()),
"weekly_resets_at_ms":before.budget.weekly_resets_at.map(|time| chrono::DateTime::<chrono::Utc>::from(time).timestamp_millis()),
});
assert_eq!(results[6], expected);
assert_eq!(results[7], json!({"client_key_id":"key_budget"}));
expected["weekly_limit_usd"] = json!("25.5");
assert_eq!(results[8], expected);
for index in [9, 10, 13, 14] {
assert_eq!(results[index], json!({"error":"invalid_input"}));
}
for index in [11, 12] {
assert_eq!(results[index], json!({"error":"rejected"}));
}
assert_eq!(
environment
.audit_requests("update_budget_limits")
.await
.len(),
1
);
assert_eq!(after.budget.daily_resets_at, before.budget.daily_resets_at);
assert_eq!(
after.budget.weekly_resets_at,
before.budget.weekly_resets_at
);
assert_eq!(audits.len(), 1);
assert!(audits[0].starts_with(&format!("plugin:{}:scope:", view.target.instance_id)));
} else {
assert!(
results
.iter()
.all(|value| value == &json!({"error":"permission_denied"}))
);
assert_eq!(after.budget, before.budget);
assert!(audits.is_empty());
}
assert_eq!(
store
.reveal_client_key(&key)
.await
.unwrap()
.unwrap()
.expose_for_response(),
plaintext
);
assert!(
!std::str::from_utf8(&reply.body)
.unwrap()
.contains(&plaintext)
);
drop(service);
drop(access);
environment.release_plugin_accounts(&runtime);
runtime.shutdown().await;
drop(core);
drop(runtime);
drop(store);
environment.close().await;
}
}
async fn seed_budget(environment: &Environment) -> ClientApiKeyId {
let key = ClientApiKeyId::new("key_budget").unwrap();
environment
.client_key(key.as_str(), "sk-budget-fixture")
.await;
environment.seed_client_key_budget(key.as_str()).await;
key
}
fn reset_queries() -> Value {
// 重置无需先查询目录;两个基础接口由测试插件自行选择调用顺序。
json!([
{"method":"host.keys.reset_budget","query":{"client_key_id":"key_budget","period":"weekly"}},
{"method":"host.keys.list","query":{"limit":10}},
{"method":"host.keys.get_budget","query":{"client_key_id":"key_budget"}},
{"method":"host.keys.update_budget_limits","query":{"client_key_id":"key_budget","daily_limit_usd":"10"}}
])
}
fn assert_reset_results(results: &Value) {
assert_eq!(results[0], json!({"client_key_id":"key_budget"}));
assert_eq!(results[1]["keys"][0]["id"], "key_budget");
assert_eq!(results[2]["daily_used_usd"], "3");
assert_eq!(results[2]["weekly_used_usd"], "0");
assert_eq!(results[3], json!({"client_key_id":"key_budget"}));
}
#[tokio::test]
async fn command_plane_resets_budget_with_only_budget_permission() {
let Some(mut environment) = Environment::create_command().await else {
eprintln!("SKIP: plugin integration environment absent");
return;
};
let key = seed_budget(&environment).await;
environment
.install_plugin(
json!({
"command_registration":{"commands":[{"name":"reset","description":"预算接口测试"}]},
"data_queries":reset_queries()
}),
vec![account_grant("key_budgets")],
)
.await;
let (runtime, core) = environment.command_plane().await;
let store = environment.store.admin_ports().client_keys();
let before = store.get_client_key(&key).await.unwrap().unwrap();
let access = gateway_admin::initialize_plugin_client_keys(
native::admin_registry(),
store.clone(),
core.snapshot_control(),
);
runtime.bind_client_key_ports(&access).unwrap();
let instance = environment
.store
.admin_ports()
.plugins()
.load_instances()
.await
.unwrap()
.instances
.remove(0);
let commands = runtime.prepare_command_line().await.unwrap();
environment.store.start_command_line_writes().unwrap();
let reply = commands.execute(&instance.id, "reset", &[]).await.unwrap();
assert_eq!(reply.exit_code, 0);
assert_reset_results(&serde_json::from_str(&reply.stdout).unwrap());
let after = store.get_client_key(&key).await.unwrap().unwrap();
assert_eq!(after.budget.weekly_used_usd.canonical(), "0");
assert_eq!(after.budget.daily_used_usd, before.budget.daily_used_usd);
assert_eq!(environment.audit_requests("reset_budget").await.len(), 1);
commands.shutdown().await;
runtime.shutdown().await;
environment
.store
.shutdown_command_line_writes()
.await
.unwrap();
drop(access);
drop(core);
drop(runtime);
drop(store);
environment.close().await;
}
#[tokio::test]
async fn published_maintenance_resets_budget_with_only_budget_permission() {
let Some(environment) = Environment::create().await else {
eprintln!("SKIP: plugin integration environment absent");
return;
};
let key = seed_budget(&environment).await;
let marker = environment
.directory
.path()
.join("budget-maintenance.jsonl");
environment
.install_plugin(
json!({
"maintenance_fixture":true,
"maintenance_marker":marker,
"data_queries":reset_queries()
}),
vec![account_grant("key_budgets")],
)
.await;
let (runtime, core) = environment.runtime().await;
let store = environment.store.admin_ports().client_keys();
let before = store.get_client_key(&key).await.unwrap().unwrap();
let access = gateway_admin::initialize_plugin_client_keys(
native::admin_registry(),
store.clone(),
core.snapshot_control(),
);
runtime.bind_client_key_ports(&access).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 worker = tokio::spawn(async move { task.run(cancellation).await.unwrap() });
let result: Value = timeout(Duration::from_secs(15), async {
loop {
if let Ok(content) = std::fs::read_to_string(&marker)
&& let Some(result) = content
.lines()
.find_map(|line| serde_json::from_str(line).ok())
{
break result;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await
.expect("maintenance callback completed");
stop.cancel();
timeout(Duration::from_secs(5), worker)
.await
.unwrap()
.unwrap();
assert_reset_results(&result["results"]);
let after = store.get_client_key(&key).await.unwrap().unwrap();
assert_eq!(after.budget.weekly_used_usd.canonical(), "0");
assert_eq!(after.budget.daily_used_usd, before.budget.daily_used_usd);
assert_eq!(environment.audit_requests("reset_budget").await.len(), 1);
environment.release_plugin_accounts(&runtime);
runtime.shutdown().await;
drop(access);
drop(core);
drop(runtime);
drop(store);
environment.close().await;
}
struct HoldCommittedReply {
inner: Arc<dyn PluginClientKeyAccess>,
committed: Arc<Notify>,
}
#[async_trait::async_trait]
impl PluginClientKeyAccess for HoldCommittedReply {
async fn facts(&self, id: &ClientApiKeyId) -> Result<PluginClientKeyFacts, AdminError> {
self.inner.facts(id).await
}
async fn budget(
&self,
id: &ClientApiKeyId,
) -> Result<gateway_core::engine::budget::ClientBudgetStatus, AdminError> {
self.inner.budget(id).await
}
async fn update_budget_limits(
&self,
owner: &PluginResourceOwner,
command: gateway_admin::model::client_keys::UpdateClientKeyBudgetLimits,
context: &MutationContext,
) -> Result<ClientApiKeyId, AdminError> {
self.inner
.update_budget_limits(owner, command, context)
.await
}
async fn reset_budget(
&self,
owner: &PluginResourceOwner,
command: ResetClientKeyBudget,
context: &MutationContext,
) -> Result<ClientApiKeyId, AdminError> {
self.inner.reset_budget(owner, command, context).await?;
// 在真实事务提交后挡住宿主回包,稳定复现插件未收到提交结果的窗口。
self.committed.notify_one();
std::future::pending().await
}
async fn list(
&self,
query: PluginClientKeyListQuery,
) -> Result<PluginClientKeyPage, AdminError> {
self.inner.list(query).await
}
}
#[tokio::test]
async fn losing_callback_reply_after_commit_keeps_budget_and_audit_without_replay() {
let Some(environment) = Environment::create().await else {
eprintln!("SKIP: plugin integration environment absent");
return;
};
let key = seed_budget(&environment).await;
let disconnect = environment.directory.path().join("disconnect-after-commit");
let startups = environment.directory.path().join("budget-startups.jsonl");
environment
.install_plugin(
json!({
"management_registration":{"routes":[{"method":"POST","path":"reset","request_content_types":[],"response_content_types":["application/json"]}]},
"data_queries":reset_queries(),
"startup_marker":startups,
"exit_after_ready_signals":[disconnect]
}),
vec![account_grant("key_budgets")],
)
.await;
let (runtime, core) = environment.runtime().await;
let store = environment.store.admin_ports().client_keys();
let committed = Arc::new(Notify::new());
let access: Arc<dyn PluginClientKeyAccess> = Arc::new(HoldCommittedReply {
inner: gateway_admin::initialize_plugin_client_keys(
native::admin_registry(),
store.clone(),
core.snapshot_control(),
),
committed: committed.clone(),
});
runtime.bind_client_key_ports(&access).unwrap();
let service = PluginManagementService::new(
runtime.clone(),
environment.store.admin_ports().plugins(),
core.snapshots(),
);
let view = service.views().await.unwrap().remove(0);
let request = tokio::spawn(async move {
service
.handle(
&view.target,
PluginManagementRequest {
method: "POST".into(),
path: "reset".into(),
query: String::new(),
content_type: None,
body: vec![],
request_id: "budget-lost-reply".into(),
},
)
.await
});
timeout(Duration::from_secs(15), committed.notified())
.await
.expect("native budget transaction committed");
assert!(!request.is_finished());
// 只断开插件,宿主保持运行,确保连接丢失不会触发透明重放。
std::fs::write(&disconnect, b"").unwrap();
assert!(
timeout(Duration::from_secs(5), request)
.await
.unwrap()
.unwrap()
.is_err()
);
let after = store.get_client_key(&key).await.unwrap().unwrap();
assert_eq!(after.budget.weekly_used_usd.canonical(), "0");
assert_eq!(environment.audit_requests("reset_budget").await.len(), 1);
assert_eq!(
std::fs::read_to_string(startups).unwrap().lines().count(),
1
);
environment.release_plugin_accounts(&runtime);
runtime.shutdown().await;
drop(access);
drop(core);
drop(runtime);
drop(store);
environment.close().await;
}
@@ -1 +1,3 @@
mod key_facts;
mod keys;
mod log;
@@ -22,7 +22,7 @@ use gateway_plugin_sdk::{Capability, Contributions, Permission, Stage};
use serde_json::{Value, json};
#[derive(Default)]
struct Facts(AtomicUsize);
struct Facts(AtomicUsize, AtomicUsize);
fn account() -> AccountRecord {
let at = "2026-01-01T00:00:00Z".parse().unwrap();
@@ -65,6 +65,26 @@ fn account() -> AccountRecord {
#[async_trait]
impl PluginAccountAccess for Facts {
async fn refresh_quota(
&self,
account: &ProviderAccountId,
) -> Result<ProviderQuota, AdminError> {
self.1.fetch_add(1, Ordering::SeqCst);
if account.as_str() == "acct_missing" {
return Err(AdminError::not_found("missing account"));
}
if account.as_str() == "acct_unsupported" {
return Err(AdminError::invalid("unsupported quota refresh"));
}
if account.as_str() == "acct_unavailable" {
return Err(AdminError::unavailable("quota refresh unavailable"));
}
let mut quota = self.get_quota(account).await?;
quota.observed_at = Some("2026-01-01T00:00:00Z".parse().unwrap());
quota.windows[0].used_percent = Some(17.5);
Ok(quota)
}
async fn list(&self, query: PluginAccountListQuery) -> Result<PluginAccountPage, AdminError> {
self.0.fetch_add(1, Ordering::SeqCst);
assert_eq!(query.limit.get(), 1);
@@ -211,3 +231,104 @@ async fn management_facts_are_minimal_bounded_and_separately_authorized() {
super::wait_until_empty(cache.path()).await;
}
}
#[tokio::test]
async fn quota_refresh_requires_observation_permission_and_returns_only_facts() {
for permission in [
None,
Some(Permission::Data),
Some(Permission::QuotaObservations),
] {
let (cache, store, runtime) = super::setup_with_permissions(
Contributions::from([crate::support::contribution(
Capability::Management,
vec![Stage::Management],
vec![],
vec![],
)]),
Default::default(),
permission.into_iter().collect(),
)
.await;
let facts = Arc::new(Facts::default());
let port: Arc<dyn PluginAccountAccess> = facts.clone();
runtime.bind_account_ports(&port).unwrap();
{
let mut snapshot = store.snapshot.lock().unwrap();
let instance = &mut snapshot.instances[0];
instance.grants = permission
.into_iter()
.map(|permission| PluginPermissionGrant {
permission: permission.as_str().into(),
})
.collect();
instance.configuration = json!({
"management_registration":{"routes":[{"method":"GET","path":"facts","request_content_types":[],"response_content_types":["application/json"]}]},
"data_queries":[
{"method":"host.data.quota.get","query":{"account_id":"acct_facts"}},
{"method":"host.quota_observations.refresh","query":{"account_id":"acct_facts"}},
{"method":"host.quota_observations.refresh","query":{"account_id":"acct_facts","provider_id":"forged"}},
{"method":"host.quota_observations.refresh","query":{"account_id":"acct_missing"}},
{"method":"host.quota_observations.refresh","query":{"account_id":"acct_unsupported"}},
{"method":"host.quota_observations.refresh","query":{"account_id":"acct_unavailable"}},
{"method":"host.auth.get","query":{"account_id":"acct_facts"}}
]
});
}
let generation =
ExtensionPreparationPort::prepare(&runtime, ConfigRevision::new(1).unwrap())
.await
.unwrap();
assert!(runtime.middleware_registry().resolve(&generation).is_none());
assert!(runtime.policy_registry().resolve(&generation).is_none());
let view = runtime.views(&generation).await.unwrap().remove(0);
let response = runtime
.handle(
&generation,
&view.target,
PluginManagementRequest {
method: "GET".into(),
path: "facts".into(),
query: String::new(),
content_type: None,
body: vec![],
request_id: "data-fixture".into(),
},
)
.await
.unwrap();
let result: Value = serde_json::from_slice(&response.body).unwrap();
if permission == Some(Permission::QuotaObservations) {
assert_eq!(facts.1.load(Ordering::SeqCst), 4);
assert_eq!(result[0]["observed_at_ms"], Value::Null);
assert_eq!(
result[1],
json!({"schema_version":1,"account_id":"acct_facts","observed_at_ms":1767225600000i64,
"windows":[{"key":"weekly","window_seconds":604800,"used_percent":17.5,"reset_at_ms":null}]})
);
for (index, code) in [
(2, "invalid_input"),
(3, "rejected"),
(4, "invalid_input"),
(5, "fault"),
(6, "permission_denied"),
] {
assert_eq!(result[index]["error"], code);
}
} else {
assert_eq!(facts.1.load(Ordering::SeqCst), 0);
for denied in &result.as_array().unwrap()[1..] {
assert_eq!(denied["error"], "permission_denied");
}
if permission == Some(Permission::Data) {
assert_eq!(result[0]["account_id"], "acct_facts");
} else {
assert_eq!(result[0]["error"], "permission_denied");
}
}
drop(generation);
runtime.shutdown().await;
super::wait_until_empty(cache.path()).await;
}
}
@@ -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;

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