mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-10 00:06:00 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3915aa259b | ||
|
|
d310286e33 | ||
|
|
ea4e3059db | ||
|
|
3bf7030fba | ||
|
|
5d5965c6ca | ||
|
|
0710ed1f7c | ||
|
|
4634d8dc6d | ||
|
|
8a0759905e | ||
|
|
e830b2ef48 | ||
|
|
ae8da00459 | ||
|
|
7bbea42fba | ||
|
|
2e0669d462 | ||
|
|
2df9494e99 | ||
|
|
0534dd8f26 | ||
|
|
525f3d51c1 | ||
|
|
33cc427b86 | ||
|
|
d1bf90d5fd | ||
|
|
eb23a0fb3b | ||
|
|
ca90f2b586 | ||
|
|
cfcf7c486a | ||
|
|
b981e42342 | ||
|
|
02177af921 | ||
|
|
2d20f03bcb | ||
|
|
c12c600b42 | ||
|
|
1c6b5a8f10 |
@@ -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;仍影响升级、恢复的说明保留适用条件与操作。没有文档差异时,只核对行为变化是否使现有说明失真。
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -6,7 +6,6 @@
|
||||
# Frontend
|
||||
/frontend/node_modules/
|
||||
/frontend/dist/
|
||||
/frontend/.vite/
|
||||
/frontend/.pnpm-store/
|
||||
/frontend/.pnpm-cache/
|
||||
*.tsbuildinfo
|
||||
|
||||
+1
-1
@@ -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 包的构建脚本;升级依赖时核对源码并更新锁文件,无需逐次修改构建许可,不开放全部依赖脚本。
|
||||
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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 ¬ifications {
|
||||
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
Reference in New Issue
Block a user