mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-12 00:06:23 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3915aa259b | ||
|
|
d310286e33 | ||
|
|
ea4e3059db | ||
|
|
3bf7030fba | ||
|
|
5d5965c6ca | ||
|
|
0710ed1f7c | ||
|
|
4634d8dc6d | ||
|
|
8a0759905e | ||
|
|
e830b2ef48 | ||
|
|
ae8da00459 | ||
|
|
7bbea42fba | ||
|
|
2e0669d462 | ||
|
|
2df9494e99 |
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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>,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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<()>;
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()),
|
||||
}
|
||||
@@ -3405,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);
|
||||
|
||||
@@ -644,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"))
|
||||
|
||||
@@ -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)?;
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -863,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;
|
||||
@@ -1026,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"))
|
||||
}
|
||||
|
||||
@@ -1411,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(
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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::{
|
||||
|
||||
@@ -479,6 +479,7 @@ pub enum PreferredAccountSelection {
|
||||
pub struct AccountSelection<'a> {
|
||||
candidate: &'a AccountCandidate,
|
||||
preferred: PreferredAccountSelection,
|
||||
policy_choice: bool,
|
||||
}
|
||||
|
||||
impl<'a> AccountSelection<'a> {
|
||||
@@ -491,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 内唯一的账号排序器。
|
||||
@@ -536,6 +543,7 @@ impl AccountSelector {
|
||||
Some(AccountSelection {
|
||||
candidate,
|
||||
preferred,
|
||||
policy_choice: true,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -609,6 +617,7 @@ impl AccountSelector {
|
||||
return Some(AccountSelection {
|
||||
candidate,
|
||||
preferred: PreferredAccountSelection::Hit,
|
||||
policy_choice: false,
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -664,6 +673,7 @@ impl AccountSelector {
|
||||
Some(AccountSelection {
|
||||
candidate,
|
||||
preferred,
|
||||
policy_choice: false,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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 }
|
||||
}
|
||||
}
|
||||
@@ -8,6 +8,7 @@ mod middleware;
|
||||
mod policy;
|
||||
mod probe;
|
||||
mod provider;
|
||||
mod response_control;
|
||||
|
||||
use gateway_core::engine::AttemptTrigger;
|
||||
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
@@ -300,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());
|
||||
}
|
||||
|
||||
@@ -24,6 +24,8 @@
|
||||
"requests",
|
||||
"public_endpoints",
|
||||
"groups",
|
||||
"keys"
|
||||
"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(
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) 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()
|
||||
}
|
||||
|
||||
pub(super) fn encode(value: &impl serde::Serialize) -> Result<RpcReply, PluginFault> {
|
||||
Ok(RpcReply {
|
||||
result: serde_json::json!({}),
|
||||
payload: serde_json::to_vec(value).map_err(|_| invalid())?,
|
||||
})
|
||||
}
|
||||
|
||||
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,8 +43,14 @@ impl PluginData {
|
||||
params: serde_json::Value,
|
||||
payload: &[u8],
|
||||
) -> Result<RpcReply, PluginFault> {
|
||||
// 管理范围的只读授权不继承到客户端请求链,避免借数据查询扩大当前 Key 的范围。
|
||||
if !self.authorized
|
||||
// 基础事实读取不隐含上游访问权;管理范围也不继承到客户端请求链。
|
||||
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
|
||||
@@ -48,9 +61,30 @@ impl PluginData {
|
||||
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
|
||||
@@ -89,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;
|
||||
@@ -70,8 +71,7 @@ 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>,
|
||||
@@ -133,10 +133,13 @@ impl PluginCallbacks {
|
||||
let grants = &instance.grants;
|
||||
Ok(Self {
|
||||
resources: Arc::new(resources::PluginResources::new(instance, ports.resources)),
|
||||
data: Arc::new(data::PluginData::new(ports.accounts.clone(), grants)),
|
||||
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)
|
||||
@@ -385,7 +388,6 @@ impl CallbackHandler for PluginCallbacks {
|
||||
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 {
|
||||
@@ -434,11 +436,14 @@ impl CallbackHandler for PluginCallbacks {
|
||||
) {
|
||||
return resources.call(&context, &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(),
|
||||
"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;
|
||||
@@ -461,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;
|
||||
}
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use super::{
|
||||
accounts::{encode, map_admin_error, mutation_context},
|
||||
admin::{encode, map_admin_error, mutation_context},
|
||||
denied, invalid,
|
||||
};
|
||||
use crate::RpcReply;
|
||||
|
||||
@@ -265,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>,
|
||||
|
||||
@@ -224,14 +224,34 @@ fn callback_allowed(method: &str, stage: Stage, permissions: &[Permission]) -> b
|
||||
&& !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"
|
||||
@@ -241,13 +261,12 @@ 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"
|
||||
"host.data.accounts.list" | "host.data.keys.get"
|
||||
if matches!(
|
||||
stage,
|
||||
Stage::Management | Stage::CommandLine | Stage::Maintenance
|
||||
@@ -255,6 +274,14 @@ fn callback_allowed(method: &str, stage: Stage, permissions: &[Permission]) -> b
|
||||
{
|
||||
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,
|
||||
@@ -271,6 +298,14 @@ fn callback_allowed(method: &str, stage: Stage, permissions: &[Permission]) -> b
|
||||
{
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -714,6 +714,11 @@ fn attempt_middleware_context(effects: Arc<ExecutionEffects>) -> MiddlewareConte
|
||||
fn middleware_headers() -> Vec<MiddlewareHeader> {
|
||||
vec![
|
||||
MiddlewareHeader::new("authorization", Bytes::from_static(b"Bearer private")),
|
||||
MiddlewareHeader::new("X-OpenAI-Fedramp", Bytes::from_static(b"true")),
|
||||
MiddlewareHeader::new(
|
||||
"x-openai-account-routing-override",
|
||||
Bytes::from_static(b"us_cr"),
|
||||
),
|
||||
MiddlewareHeader::new("connection", Bytes::from_static(b"x-hop")),
|
||||
MiddlewareHeader::new("x-hop", Bytes::from_static(b"hidden")),
|
||||
MiddlewareHeader::new("x_team", Bytes::from_static(b"team")),
|
||||
|
||||
@@ -96,7 +96,9 @@ fn assert_exhaustive_permission(permission: Permission) {
|
||||
| Permission::Data
|
||||
| Permission::PublicEndpoints
|
||||
| Permission::Groups
|
||||
| Permission::Keys => {}
|
||||
| Permission::Keys
|
||||
| Permission::KeyBudgets
|
||||
| Permission::QuotaObservations => {}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -236,9 +236,18 @@ async fn execution_allows_each_declared_callback_permission() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn data_callbacks_require_independent_grants_and_management_or_command_stage() {
|
||||
for permissions in [vec![Permission::Accounts], vec![Permission::Data]] {
|
||||
async fn fact_callbacks_require_independent_grants_and_control_plane_stages() {
|
||||
for permissions in [
|
||||
vec![],
|
||||
vec![Permission::Accounts],
|
||||
vec![Permission::Keys],
|
||||
vec![Permission::Models],
|
||||
vec![Permission::KeyBudgets],
|
||||
vec![Permission::QuotaObservations],
|
||||
vec![Permission::Data],
|
||||
] {
|
||||
let authorized = permissions.contains(&Permission::Data);
|
||||
let quota_authorized = permissions.contains(&Permission::QuotaObservations);
|
||||
let callbacks = Arc::new(Callbacks::default());
|
||||
let (_cache, session) = session_with_permissions(Arc::clone(&callbacks), permissions).await;
|
||||
for stage in [
|
||||
@@ -254,10 +263,20 @@ async fn data_callbacks_require_independent_grants_and_management_or_command_sta
|
||||
Stage::PublicManagement,
|
||||
Stage::Management,
|
||||
Stage::CommandLine,
|
||||
Stage::Maintenance,
|
||||
] {
|
||||
for method in ["host.data.accounts.list", "host.data.quota.get"] {
|
||||
for method in [
|
||||
"host.data.accounts.list",
|
||||
"host.data.keys.get",
|
||||
"host.data.quota.get",
|
||||
] {
|
||||
let reply = invoke_callback(&session, stage, method).await;
|
||||
if authorized && matches!(stage, Stage::Management | Stage::CommandLine) {
|
||||
if (authorized || (quota_authorized && method == "host.data.quota.get"))
|
||||
&& matches!(
|
||||
stage,
|
||||
Stage::Management | Stage::CommandLine | Stage::Maintenance
|
||||
)
|
||||
{
|
||||
assert!(reply.is_ok());
|
||||
} else {
|
||||
assert_permission_denied(reply);
|
||||
@@ -266,7 +285,13 @@ async fn data_callbacks_require_independent_grants_and_management_or_command_sta
|
||||
}
|
||||
assert_eq!(
|
||||
callbacks.called.load(Ordering::Relaxed),
|
||||
if authorized { 4 } else { 0 }
|
||||
if authorized {
|
||||
9
|
||||
} else if quota_authorized {
|
||||
3
|
||||
} else {
|
||||
0
|
||||
}
|
||||
);
|
||||
session.shutdown(Duration::from_secs(1)).await;
|
||||
}
|
||||
@@ -296,6 +321,13 @@ async fn managed_resources_require_their_domains_and_control_plane_stages() {
|
||||
("host.groups.ensure", Permission::Groups),
|
||||
("host.groups.change_members", Permission::Groups),
|
||||
("host.keys.ensure", Permission::Keys),
|
||||
("host.keys.reset_budget", Permission::KeyBudgets),
|
||||
("host.keys.get_budget", Permission::KeyBudgets),
|
||||
("host.keys.update_budget_limits", Permission::KeyBudgets),
|
||||
(
|
||||
"host.quota_observations.refresh",
|
||||
Permission::QuotaObservations,
|
||||
),
|
||||
] {
|
||||
for granted in [false, true] {
|
||||
let callbacks = Arc::new(Callbacks::default());
|
||||
@@ -337,7 +369,10 @@ async fn maintenance_has_data_and_state_access_but_cannot_execute_models_or_read
|
||||
let (_cache, session) = session(callbacks).await;
|
||||
for method in [
|
||||
"host.data.accounts.list",
|
||||
"host.data.keys.get",
|
||||
"host.data.quota.get",
|
||||
"host.keys.list",
|
||||
"host.keys.reset_budget",
|
||||
"host.state.put",
|
||||
"host.log",
|
||||
] {
|
||||
@@ -352,9 +387,74 @@ async fn maintenance_has_data_and_state_access_but_cannot_execute_models_or_read
|
||||
"host.auth.save",
|
||||
"host.http.do",
|
||||
"host.model.execute",
|
||||
"host.keys.list",
|
||||
] {
|
||||
assert_permission_denied(invoke_callback(&session, Stage::Maintenance, method).await);
|
||||
}
|
||||
session.shutdown(Duration::from_secs(1)).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn budget_access_does_not_grant_models_or_key_provisioning_and_models_do_not_grant_reset() {
|
||||
for permission in [Permission::KeyBudgets, Permission::Models, Permission::Keys] {
|
||||
let callbacks = Arc::new(Callbacks::default());
|
||||
let (_cache, session) = session_with_permissions(callbacks, vec![permission]).await;
|
||||
for stage in [
|
||||
Stage::Management,
|
||||
Stage::CommandLine,
|
||||
Stage::Maintenance,
|
||||
Stage::Request,
|
||||
Stage::Attempt,
|
||||
Stage::Observation,
|
||||
Stage::Routing,
|
||||
Stage::Scheduling,
|
||||
Stage::Authentication,
|
||||
Stage::Retry,
|
||||
Stage::Registration,
|
||||
Stage::Configuration,
|
||||
Stage::PublicManagement,
|
||||
] {
|
||||
let reply = invoke_callback(&session, stage, "host.keys.list").await;
|
||||
let allowed = match permission {
|
||||
Permission::KeyBudgets => matches!(
|
||||
stage,
|
||||
Stage::Management | Stage::CommandLine | Stage::Maintenance
|
||||
),
|
||||
Permission::Models => matches!(
|
||||
stage,
|
||||
Stage::Management
|
||||
| Stage::CommandLine
|
||||
| Stage::Request
|
||||
| Stage::Attempt
|
||||
| Stage::Observation
|
||||
| Stage::Routing
|
||||
| Stage::Scheduling
|
||||
| Stage::Authentication
|
||||
),
|
||||
_ => false,
|
||||
};
|
||||
if allowed {
|
||||
assert!(reply.is_ok());
|
||||
} else {
|
||||
assert_permission_denied(reply);
|
||||
}
|
||||
if permission != Permission::KeyBudgets {
|
||||
assert_permission_denied(
|
||||
invoke_callback(&session, stage, "host.keys.reset_budget").await,
|
||||
);
|
||||
}
|
||||
}
|
||||
if permission == Permission::KeyBudgets {
|
||||
for method in [
|
||||
"host.model.execute",
|
||||
"host.models.list",
|
||||
"host.keys.ensure",
|
||||
"host.auth.get",
|
||||
] {
|
||||
assert_permission_denied(
|
||||
invoke_callback(&session, Stage::Management, method).await,
|
||||
);
|
||||
}
|
||||
}
|
||||
session.shutdown(Duration::from_secs(1)).await;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -52,6 +52,25 @@ impl Environment {
|
||||
.await;
|
||||
}
|
||||
|
||||
pub async fn seed_client_key_budget(&self, id: &str) {
|
||||
sqlx::query(sqlx::AssertSqlSafe(format!(
|
||||
"update {}.client_api_keys set daily_limit_usd=10, weekly_limit_usd=20 where id=$1",
|
||||
self.schema
|
||||
)))
|
||||
.bind(id)
|
||||
.execute(&self.admin)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(sqlx::AssertSqlSafe(format!(
|
||||
"insert into {}.client_key_budget_windows
|
||||
(client_api_key_id, daily_start, daily_end, daily_used_usd, weekly_start, weekly_end, weekly_used_usd)
|
||||
values ($1, now(), now()+interval '1 day', 3, now(), now()+interval '7 days', 4)",
|
||||
self.schema
|
||||
)))
|
||||
.bind(id)
|
||||
.execute(&self.admin).await.unwrap();
|
||||
}
|
||||
|
||||
pub async fn client_key_with_limits(
|
||||
&self,
|
||||
id: &str,
|
||||
|
||||
@@ -81,6 +81,43 @@ struct Peer {
|
||||
}
|
||||
|
||||
impl Peer {
|
||||
async fn data_queries(&self, id: u64) -> Option<Vec<Value>> {
|
||||
let queries = self.configuration["data_queries"].as_array()?;
|
||||
let mut results = Vec::new();
|
||||
for query in queries {
|
||||
let metadata_query = query["method"] == "host.keys.list";
|
||||
let result = self
|
||||
.callback_payload(
|
||||
id,
|
||||
query["method"].as_str().unwrap(),
|
||||
if metadata_query {
|
||||
query["query"].clone()
|
||||
} else {
|
||||
json!({})
|
||||
},
|
||||
if metadata_query {
|
||||
vec![]
|
||||
} else {
|
||||
serde_json::to_vec(&query["query"]).unwrap()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
results.push(match result {
|
||||
Ok((metadata, payload)) => {
|
||||
if metadata_query {
|
||||
assert!(payload.is_empty());
|
||||
metadata
|
||||
} else {
|
||||
assert_eq!(metadata, json!({}));
|
||||
serde_json::from_slice::<Value>(&payload).unwrap()
|
||||
}
|
||||
}
|
||||
Err(error) => json!({"error":error.code}),
|
||||
});
|
||||
}
|
||||
Some(results)
|
||||
}
|
||||
|
||||
async fn resource_fixture(
|
||||
&self,
|
||||
id: u64,
|
||||
@@ -360,7 +397,15 @@ impl Peer {
|
||||
}
|
||||
match method.as_str() {
|
||||
"plugin.reconcile" => {
|
||||
let result = self.reconcile_fixture(id).await;
|
||||
let result = if let Some(results) = self.data_queries(id).await {
|
||||
self.append_observation_marker(
|
||||
"maintenance_marker",
|
||||
&json!({"phase":"done", "results":results}),
|
||||
);
|
||||
Ok(())
|
||||
} else {
|
||||
self.reconcile_fixture(id).await
|
||||
};
|
||||
match result {
|
||||
Ok(()) => {
|
||||
self.send(
|
||||
@@ -470,25 +515,7 @@ impl Peer {
|
||||
"management_marker",
|
||||
&json!({"method":request.method,"path":request.path}),
|
||||
);
|
||||
if let Some(queries) = self.configuration["data_queries"].as_array() {
|
||||
let mut results = vec![];
|
||||
for query in queries {
|
||||
let result = self
|
||||
.callback_payload(
|
||||
id,
|
||||
query["method"].as_str().unwrap(),
|
||||
json!({}),
|
||||
serde_json::to_vec(&query["query"]).unwrap(),
|
||||
)
|
||||
.await;
|
||||
results.push(match result {
|
||||
Ok((metadata, payload)) => {
|
||||
assert_eq!(metadata, json!({}));
|
||||
serde_json::from_slice::<Value>(&payload).unwrap()
|
||||
}
|
||||
Err(error) => json!({"error":error.code}),
|
||||
});
|
||||
}
|
||||
if let Some(results) = self.data_queries(id).await {
|
||||
self.send(
|
||||
Message::Result {
|
||||
id,
|
||||
@@ -608,7 +635,9 @@ impl Peer {
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
let result = if self.configuration["command_echo"] == true {
|
||||
let result = if let Some(results) = self.data_queries(id).await {
|
||||
json!({"stdout":serde_json::to_string(&results).unwrap(),"stderr":"","exit_code":0})
|
||||
} else if self.configuration["command_echo"] == true {
|
||||
json!({"stdout":serde_json::to_string(&invocation).unwrap(),"stderr":"typed command\n","exit_code":0})
|
||||
} else {
|
||||
self.configuration["command_result"].clone()
|
||||
|
||||
@@ -66,11 +66,13 @@ Provider 固定为宿主内置的 OpenAI 与 xAI。插件提供以下扩展能
|
||||
| `network` | 通过 `host.http.*` 使用宿主受管出站网络;仍受统一代理、超时、大小和流控规则约束 |
|
||||
| `models` | 列出非秘密 Key、按所选 Key 查询模型,以及通过 `host.model.*` 调用模型;调用可能产生消耗 |
|
||||
| `accounts` | 查询账号、读取原始凭据及创建或替换账号;写入仍经过 revision CAS、审计和发布事务 |
|
||||
| `data` | 在管理/命令/维护阶段只读全部账号的最小基础信息及已有额度观测,不包含凭据、写入或预测 |
|
||||
| `data` | 在管理/命令/维护阶段只读账号、Key 的最小基础信息及已有额度观测,不包含凭据、写入或预测 |
|
||||
| `requests` | 参与请求/响应处理、路由、调度、观察及亲和查询 |
|
||||
| `public_endpoints` | 提供无需登录即可访问的已声明静态资源或一次性票据回调 |
|
||||
| `groups` | 创建本实例分组,允许将所有现有及未来新增账号加入或移出这些分组,保留其他分组关系 |
|
||||
| `keys` | 创建仅绑定本实例分组的 Key,返回非秘密身份;不授予其他 Key 的修改或明文读取权限 |
|
||||
| `key_budgets` | 在管理/命令/维护阶段查询全部 Client Key 的预算、修改日/周金额上限及重置用量,不读取密钥或修改其他 Key 配置 |
|
||||
| `quota_observations` | 在管理/命令/维护阶段读取及刷新全部账号的额度观测,不暴露凭据或执行上游额度重置 |
|
||||
|
||||
`host.log` 和本插件声明的 `host.state.*` 是基础设施,不需要额外 permission。公开调用阶段不开放任何
|
||||
宿主回调;权限也不能把一个父调用的句柄、流或上下文转移到另一个调用。
|
||||
@@ -83,14 +85,19 @@ Provider 固定为宿主内置的 OpenAI 与 xAI。插件提供以下扩展能
|
||||
|
||||
纯展示页面声明 `management` 和 `data` 即可,不需要请求处理 binding 或 `accounts` 权限。
|
||||
`data` 是独立的管理员授权域,只允许 `management`、`command_line`、`maintenance` 阶段使用;请求链、注册及公开回调均拒绝。
|
||||
它可读取全部账号的下列最小事实,不继承或授予客户端 Key 的模型执行权。
|
||||
它可读取全部账号与 Key 的下列最小事实,不继承或授予客户端 Key 的模型执行权。
|
||||
|
||||
| SDK 方法 | 回调 | 查询与结果 |
|
||||
| --- | --- | --- |
|
||||
| `call.host.account_facts(query)` | `host.data.accounts.list` | `AccountFactsQuery`:可选 `provider_id`、`cursor`,必填 `limit`(1~200);按账号 ID 升序,`next_cursor=null` 表示本页已结束 |
|
||||
| `call.host.quota_facts(query)` | `host.data.quota.get` | `QuotaFactsQuery { account_id }`:读取 Provider 现有观测,不访问上游刷新 |
|
||||
| `call.host.key_facts(query)` | `host.data.keys.get` | `ClientKeyFactsQuery { client_key_id }`:读取当前 Key 的启用状态及显式分组 ID,不返回密钥 |
|
||||
|
||||
类型在 `call::data`。控制参数为 `{}`,查询和结果使用二进制 JSON;结果固定 `schema_version=1`。
|
||||
|
||||
`key_facts` 每次读取当前管理数据,返回 `client_key_id`、`enabled`、`group_ids`;不存在的 Key 沿用事实接口的 `rejected` 错误。它需要 `data` 权限,`keys`、`key_budgets` 和 `quota_observations` 权限不能替代。原有 Key 目录响应保持不变。旧宿主不支持新增方法,插件应将包含该方法的宿主版本声明为最低兼容版本。
|
||||
|
||||
`group_ids` 表示显式绑定,包含停用分组,不是最终可路由账号集合;空绑定也不代表单账号范围。通过 `account_facts` 关联账号时需完整遍历分页,结果包含停用账号。跨查询关联及同步策略由插件负责,这些调用不构成跨查询事务,管理员修改绑定后应重新检查。
|
||||
账号仅返回 `account_id`、`provider_id`、`group_ids`、`enabled` 和 `updated_at_ms`,不附带姓名、邮箱、令牌或代理信息。
|
||||
额度仅返回观测时间与窗口的 `key`、`window_seconds`、`used_percent`、`reset_at_ms`。
|
||||
时间均为 UTC Unix 毫秒,比例为百分数;未知值保留 `null`,不能解释为 0。`observed_at_ms=null` 表示没有可用观测时间,
|
||||
@@ -110,6 +117,25 @@ for account in page.accounts {
|
||||
已有 `usage` 观察提供请求最终用量、时间和结算事实;该投递有界、不是历史补偿接口。
|
||||
本接口不提供 SQL、池汇总、健康分、额度预测或历史使用记录查询;插件结果保存在自身 `host.state.*` 中。
|
||||
|
||||
### 额度观测刷新
|
||||
|
||||
`quota_observations` 授权访问所有当前及未来账号的额度观测。它与只读 `data`、含原始凭据的 `accounts`
|
||||
分别授权;不会授予账号目录、凭据读取、账号修改、上游额度重置或 Key 预算操作权限。
|
||||
需要枚举账号时另行声明 `data`,已知账号 ID 时可直接查询或刷新。
|
||||
|
||||
- `call.host.quota_facts(query)` / `host.data.quota.get`:`data` 或 `quota_observations` 任一域可读取已有快照,不访问上游。
|
||||
- `call.host.refresh_account_quota(query)` / `host.quota_observations.refresh`:仅 `quota_observations` 可请求 Provider 刷新观测。
|
||||
|
||||
两个方法仅在 `management`、`command_line`、`maintenance` 阶段开放;均使用 `call::data::QuotaFactsQuery`
|
||||
和 `QuotaFacts`,控制参数为 `{}`,输入输出在二进制 JSON 中。刷新复用原生 Provider 的凭据、代理、超时及
|
||||
观测落库规则,返回同一最小事实投影,不返回凭据、预测、个人资料或 Provider 原始响应。
|
||||
刷新观测不会消费额度重置券,也不会自动重置任何 Key;Provider 仍可能依据真实上游响应更新账号可用性观测。
|
||||
|
||||
接口不承诺每次得到不同的周期或样本,也不保证结果对应一个账号重置事件。不存在账号返回 `rejected`;
|
||||
Provider 不支持刷新(如 OpenAI URL + Key 账号)返回 `invalid_input`,其他失败沿用原生管理错误映射。
|
||||
刷新受父调用期限和取消约束;结果未知不表示观测未落库。SDK 不自动重复调用,插件应控制刷新频率并处理失败,
|
||||
不能把失败、空值或旧快照当作额度恢复。维护阶段只允许这条受控原生刷新,不因此获得任意网络访问权。
|
||||
|
||||
### 账号与凭据
|
||||
|
||||
`accounts` 域开放 `host.auth.list/get_runtime/get/save`:
|
||||
@@ -128,7 +154,7 @@ for account in page.accounts {
|
||||
每个实例串行执行,不同实例相互独立;调用限时 30 秒,失败后等待 5 秒重试。维护失败不回滚已经提交的资源,
|
||||
下一次对账继续补齐。只读校验、准备候选与 CLI 帮助不会启动维护;停用、替换和宿主关闭时取消旧任务。
|
||||
|
||||
维护阶段允许日志、私有状态,以及已授权的 `data`、`groups`、`keys` 回调;不开放网络、凭据或模型执行。
|
||||
维护阶段允许日志、私有状态,以及已授权的 `data`、`groups`、`keys`、`key_budgets`、`quota_observations` 回调;不开放任意网络、凭据或模型执行。
|
||||
其他管理/命令入口也可使用下列资源方法:
|
||||
|
||||
| SDK 方法 | 参数与行为 |
|
||||
@@ -149,6 +175,41 @@ Key 明文仍通过宿主管理面查看,插件模型调用使用返回的 Key
|
||||
典型处理器先确保分组存在,再通过 `data` 分页查询账号并增量补齐成员,最后确保 Key 存在。
|
||||
安装 `groups` 域即授权纳入全部当前及未来账号;插件可按自己的配置筛选账号,但该筛选不构成宿主的权限边界。
|
||||
|
||||
### Client Key 预算
|
||||
|
||||
`key_budgets` 是原生 Key 预算访问域,仅在 `management`、`command_line`、`maintenance` 阶段使用。
|
||||
接受该域即允许查询、设置金额上限及清零全部当前及未来 Client Key,包括管理员和其他插件创建的 Key;
|
||||
插件配置中的 Key 筛选不构成宿主权限边界。`keys` 只负责插件自有 Key 的创建,不能替代这一预算管理授权。
|
||||
`key_budgets` 不授予模型执行、密钥读取、Key 创建、名称/分组/画像或并发/RPM 配置修改权限。
|
||||
|
||||
类型位于 `call::key_budgets`;下列预算方法控制参数为 `{}`,输入输出为二进制 JSON:
|
||||
|
||||
| SDK 方法 / 回调 | 输入与结果 |
|
||||
| --- | --- |
|
||||
| `get_key_budget` / `host.keys.get_budget` | `{client_key_id}`;返回 `KeyBudget`,包含 ID、日/周上限及已用金额、日/周重置时间 |
|
||||
| `update_key_budget_limits` / `host.keys.update_budget_limits` | `{client_key_id, daily_limit_usd?, weekly_limit_usd?}`;至少指定一项;返回 `{client_key_id}` |
|
||||
| `reset_key_budget` / `host.keys.reset_budget` | `{client_key_id, period}`;`period` 必填,取 `daily`、`weekly` 或 `all`;返回 `{client_key_id}` |
|
||||
|
||||
`list_keys` / `host.keys.list` 仍可查询非秘密目录,仅返回 `{id,name,enabled}` 和游标,沿用目录的控制参数编码。
|
||||
上述操作均接受 Key ID,不接受 Key 明文,也不要求先查询目录。
|
||||
|
||||
金额为非负十进制字符串,沿用原生金额精度;上限 `"0"` 表示不限。`KeyBudget` 的字段为
|
||||
`daily_limit_usd`、`weekly_limit_usd`、`daily_used_usd`、`weekly_used_usd`、`daily_resets_at_ms`、`weekly_resets_at_ms`;
|
||||
时间为 UTC Unix 毫秒,`null` 表示尚未使用或窗口已过期。读取不触发准入、开启窗口或清零,停用的 Key 仍可管理。
|
||||
|
||||
上限更新仅写入提供的日/周金额;省略或 `null` 的项保持不变。它保留已用金额、窗口到期时间、费用历史及其他 Key 配置。
|
||||
写入复用 Key 行锁,与结算串行;实际变化时授权、修改、配置 revision 和审计在同一事务提交,并通知原生配置发布。
|
||||
相同值再次赋值不产生新 revision 或审计;并发更新按事务顺序生效,同一字段由后提交的值覆盖,接口不提供调用去重或比较交换。
|
||||
将非零上限降低到已用金额及以下会拒绝后续准入,提高限额可恢复准入,在途请求仍按原生规则完成和结算。
|
||||
|
||||
重置只清零所选周期,保留额度上限、密钥、窗口到期时间和历史费用事件;未使用的 Key 不因此开启窗口。
|
||||
费用按完成时间归属,重置前完成但延迟落盘的费用不会重新扣入已重置周期,重置后完成的在途请求仍正常扣额。
|
||||
授权复验、清零和审计在同一事务提交,不推进配置 revision。每次调用均执行新重置,可能清掉两次调用之间的新消费。
|
||||
|
||||
SDK 不自动重试。超时或断连不能证明写入未提交;重试上限赋值也可能覆盖期间其他调用的修改。
|
||||
不存在的 Key 返回 `rejected`,无权限或阶段不符返回 `permission_denied`;写入事务复验发现实例停用、版本或授权变化时返回 `conflict`,
|
||||
非法输入返回 `invalid_input`。账号关联、预算分配和重置触发由插件决定,宿主不自动串联上述接口。
|
||||
|
||||
### Key、模型与模型调用
|
||||
|
||||
`models` 域开放:
|
||||
|
||||
@@ -67,6 +67,8 @@ capability 固定并由工具生成:
|
||||
| `requests` | 查看和处理请求、响应、路由、调度与观察事实 |
|
||||
| `groups` | 创建本实例分组并管理所有当前及未来账号在这些分组中的成员关系 |
|
||||
| `keys` | 创建绑定本实例分组的 Key,不读取密钥明文 |
|
||||
| `key_budgets` | 查询全部 Client Key 的预算、修改日/周金额上限及重置用量,不读取密钥或修改其他配置 |
|
||||
| `quota_observations` | 查询和刷新全部账号的额度观测,不暴露凭据或执行上游额度重置 |
|
||||
| `public_endpoints` | 提供无需登录即可访问的资源或回调 |
|
||||
|
||||
安装时统一接受清单声明的域,不再填写逐方法、用途、Key、账号或 Provider 白名单。日志与清单声明的
|
||||
|
||||
@@ -1,9 +1,27 @@
|
||||
//! 只读基础事实;版本 1 不提供凭据、预测、聚合或主动上游刷新。
|
||||
//! 基础事实投影;主动刷新由独立的 quota_observations 访问域授权。
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
pub const ACCOUNTS_LIST: &str = "host.data.accounts.list";
|
||||
pub const QUOTA_REFRESH: &str = "host.quota_observations.refresh";
|
||||
pub const QUOTA_GET: &str = "host.data.quota.get";
|
||||
pub const KEYS_GET: &str = "host.data.keys.get";
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct ClientKeyFactsQuery {
|
||||
pub client_key_id: String,
|
||||
}
|
||||
|
||||
/// 当前显式分组绑定;空列表不是单账号范围,不包含密钥或凭据。
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct ClientKeyFacts {
|
||||
pub schema_version: u32,
|
||||
pub client_key_id: String,
|
||||
pub enabled: bool,
|
||||
pub group_ids: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
//! Client Key 预算查询、上限更新与用量重置。
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
pub const GET: &str = "host.keys.get_budget";
|
||||
pub const UPDATE_LIMITS: &str = "host.keys.update_budget_limits";
|
||||
pub const RESET: &str = "host.keys.reset_budget";
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum BudgetPeriod {
|
||||
Daily,
|
||||
Weekly,
|
||||
All,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct ResetKeyBudgetRequest {
|
||||
pub client_key_id: String,
|
||||
pub period: BudgetPeriod,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct ResetKeyBudgetResult {
|
||||
pub client_key_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct GetKeyBudgetRequest {
|
||||
pub client_key_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct KeyBudget {
|
||||
pub client_key_id: String,
|
||||
pub daily_limit_usd: String,
|
||||
pub weekly_limit_usd: String,
|
||||
pub daily_used_usd: String,
|
||||
pub weekly_used_usd: String,
|
||||
pub daily_resets_at_ms: Option<i64>,
|
||||
pub weekly_resets_at_ms: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct UpdateKeyBudgetLimitsRequest {
|
||||
pub client_key_id: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub daily_limit_usd: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub weekly_limit_usd: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct UpdateKeyBudgetLimitsResult {
|
||||
pub client_key_id: String,
|
||||
}
|
||||
@@ -4,6 +4,7 @@ pub mod catalog;
|
||||
pub mod data;
|
||||
pub mod frontend_authentication;
|
||||
pub mod host;
|
||||
pub mod key_budgets;
|
||||
pub mod management;
|
||||
pub mod middleware;
|
||||
pub mod model;
|
||||
|
||||
@@ -151,11 +151,13 @@ pub enum Permission {
|
||||
PublicEndpoints,
|
||||
Groups,
|
||||
Keys,
|
||||
KeyBudgets,
|
||||
QuotaObservations,
|
||||
}
|
||||
|
||||
impl Permission {
|
||||
/// 当前公开访问域,安装摘要与授权校验复用同一集合。
|
||||
pub const ALL: [Self; 8] = [
|
||||
pub const ALL: [Self; 10] = [
|
||||
Self::Network,
|
||||
Self::Models,
|
||||
Self::Accounts,
|
||||
@@ -164,6 +166,8 @@ impl Permission {
|
||||
Self::PublicEndpoints,
|
||||
Self::Groups,
|
||||
Self::Keys,
|
||||
Self::KeyBudgets,
|
||||
Self::QuotaObservations,
|
||||
];
|
||||
|
||||
#[must_use]
|
||||
@@ -177,6 +181,8 @@ impl Permission {
|
||||
Self::PublicEndpoints => "public_endpoints",
|
||||
Self::Groups => "groups",
|
||||
Self::Keys => "keys",
|
||||
Self::KeyBudgets => "key_budgets",
|
||||
Self::QuotaObservations => "quota_observations",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -192,6 +198,8 @@ impl Permission {
|
||||
Self::PublicEndpoints => "公开入口",
|
||||
Self::Groups => "专用账号分组",
|
||||
Self::Keys => "专用 API Key",
|
||||
Self::KeyBudgets => "API Key 预算",
|
||||
Self::QuotaObservations => "账号额度观测",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -203,7 +211,7 @@ impl Permission {
|
||||
Self::Models => "查询模型与 API Key 信息并调用模型,可能产生消耗",
|
||||
Self::Accounts => "读取和修改账号,包括访问原始凭据",
|
||||
Self::Data => {
|
||||
"在管理、命令或维护入口只读查询所有账号的基础信息与已有额度观测,不包含凭据或修改权限"
|
||||
"在管理、命令或维护入口只读查询账号、Key 的基础信息与已有额度观测,不包含凭据或修改权限"
|
||||
}
|
||||
Self::Requests => "查看和处理请求、响应、路由及账号选择",
|
||||
Self::PublicEndpoints => "提供无需登录即可访问的资源与回调入口",
|
||||
@@ -211,6 +219,12 @@ impl Permission {
|
||||
"创建本插件的分组,可将所有现有及未来新增账号加入或移出这些分组,不修改其他分组"
|
||||
}
|
||||
Self::Keys => "创建绑定本插件分组的 API Key,不读取密钥明文或修改管理员创建的 Key",
|
||||
Self::QuotaObservations => {
|
||||
"查询所有现有及未来账号的额度观测,并通过宿主刷新上游观测;不读取凭据、不执行上游额度重置"
|
||||
}
|
||||
Self::KeyBudgets => {
|
||||
"查询所有现有及未来新增 Key 的预算、修改日/周金额上限和重置用量,包括管理员和其他插件创建的 Key;不读取密钥或修改其他配置"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,6 +3,30 @@ use crate::{PluginFault, call::data};
|
||||
use super::{HostClient, payload_call};
|
||||
|
||||
impl HostClient {
|
||||
/// 读取 Key 当前启用状态与显式分组绑定,不返回密钥。
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// 未获得 data 权限、调用阶段不符、Key 不存在或宿主读取失败时返回错误。
|
||||
pub async fn key_facts(
|
||||
&self,
|
||||
query: data::ClientKeyFactsQuery,
|
||||
) -> Result<data::ClientKeyFacts, PluginFault> {
|
||||
payload_call(self, data::KEYS_GET, query).await
|
||||
}
|
||||
|
||||
/// 通过宿主刷新账号额度观测,返回与 quota_facts 相同的非秘密投影。
|
||||
/// 需要 quota_observations 权限;不修改上游额度,也不自动重置任何 Key。
|
||||
///
|
||||
/// # Errors
|
||||
/// 未授权、阶段不符、账号不支持刷新或 Provider 查询失败时返回错误。
|
||||
pub async fn refresh_account_quota(
|
||||
&self,
|
||||
query: data::QuotaFactsQuery,
|
||||
) -> Result<data::QuotaFacts, PluginFault> {
|
||||
payload_call(self, data::QUOTA_REFRESH, query).await
|
||||
}
|
||||
|
||||
/// 在已授权的 management/command_line/maintenance 调用中分页读取账号基础事实。
|
||||
///
|
||||
/// # Errors
|
||||
@@ -19,7 +43,7 @@ impl HostClient {
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// 未获得 data 权限、调用阶段不符、账号不存在或宿主读取失败时返回错误。
|
||||
/// 未获得 data 或 quota_observations 权限、调用阶段不符、账号不存在或宿主读取失败时返回错误。
|
||||
pub async fn quota_facts(
|
||||
&self,
|
||||
query: data::QuotaFactsQuery,
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
//! 非秘密 Key 目录与预算管理的类型化调用。
|
||||
|
||||
use crate::{
|
||||
ErrorCode, PluginFault,
|
||||
call::{
|
||||
host::{KeyListRequest, KeyListResult},
|
||||
key_budgets,
|
||||
},
|
||||
};
|
||||
|
||||
use super::{HostClient, SessionError, payload_call};
|
||||
|
||||
impl HostClient {
|
||||
/// 只读查询单个 Key 的预算,不开启或重置窗口。
|
||||
///
|
||||
/// # Errors
|
||||
/// 未授权、阶段不符、Key 不存在或宿主读取失败时返回错误。
|
||||
pub async fn get_key_budget(
|
||||
&self,
|
||||
request: key_budgets::GetKeyBudgetRequest,
|
||||
) -> Result<key_budgets::KeyBudget, PluginFault> {
|
||||
payload_call(self, key_budgets::GET, request).await
|
||||
}
|
||||
|
||||
/// 更新指定日/周上限;省略项不变,零表示不限,不清零用量。
|
||||
///
|
||||
/// # Errors
|
||||
/// 未授权、阶段不符、参数无效、Key 不存在或宿主写入失败时返回错误。
|
||||
pub async fn update_key_budget_limits(
|
||||
&self,
|
||||
request: key_budgets::UpdateKeyBudgetLimitsRequest,
|
||||
) -> Result<key_budgets::UpdateKeyBudgetLimitsResult, PluginFault> {
|
||||
payload_call(self, key_budgets::UPDATE_LIMITS, request).await
|
||||
}
|
||||
|
||||
/// 查询 Key 的非秘密身份;模型或预算访问域决定可用阶段。
|
||||
///
|
||||
/// # Errors
|
||||
/// 权限、阶段或分页参数不合法,或者宿主读取失败时返回错误。
|
||||
pub async fn list_keys(&self, query: KeyListRequest) -> Result<KeyListResult, PluginFault> {
|
||||
let invalid = || PluginFault::new(ErrorCode::InvalidInput, "invalid key list payload");
|
||||
let reply = self
|
||||
.call(
|
||||
"host.keys.list",
|
||||
serde_json::to_value(query).map_err(|_| invalid())?,
|
||||
Vec::new(),
|
||||
)
|
||||
.await
|
||||
.map_err(SessionError::into_plugin_fault)?;
|
||||
if !reply.payload.is_empty() {
|
||||
return Err(invalid());
|
||||
}
|
||||
serde_json::from_value(reply.result).map_err(|_| invalid())
|
||||
}
|
||||
|
||||
/// 清零指定周期,保留限额和到期时间;结果未知时不能盲目重试。
|
||||
///
|
||||
/// # Errors
|
||||
/// 未授权、实例过期、Key 不存在、宿主写入失败时返回错误。
|
||||
pub async fn reset_key_budget(
|
||||
&self,
|
||||
request: key_budgets::ResetKeyBudgetRequest,
|
||||
) -> Result<key_budgets::ResetKeyBudgetResult, PluginFault> {
|
||||
payload_call(self, key_budgets::RESET, request).await
|
||||
}
|
||||
}
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
mod data;
|
||||
mod frame;
|
||||
mod keys;
|
||||
mod middleware;
|
||||
mod plugin;
|
||||
mod resources;
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
use gateway_plugin_sdk::call::key_budgets::{BudgetPeriod, ResetKeyBudgetRequest};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn reset_requires_an_explicit_supported_period_and_rejects_extra_authority() {
|
||||
for period in [BudgetPeriod::Daily, BudgetPeriod::Weekly, BudgetPeriod::All] {
|
||||
let request = ResetKeyBudgetRequest {
|
||||
client_key_id: "key_1".into(),
|
||||
period,
|
||||
};
|
||||
assert_eq!(
|
||||
serde_json::from_value::<ResetKeyBudgetRequest>(
|
||||
serde_json::to_value(&request).unwrap()
|
||||
)
|
||||
.unwrap(),
|
||||
request
|
||||
);
|
||||
}
|
||||
for invalid in [
|
||||
json!({"client_key_id":"key_1"}),
|
||||
json!({"client_key_id":"key_1","period":"monthly"}),
|
||||
json!({"client_key_id":"key_1","period":"weekly","instance_id":"forged"}),
|
||||
] {
|
||||
assert!(serde_json::from_value::<ResetKeyBudgetRequest>(invalid).is_err());
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,6 @@
|
||||
mod frontend_authentication;
|
||||
mod host;
|
||||
mod key_budgets;
|
||||
mod management;
|
||||
mod middleware;
|
||||
mod observation;
|
||||
|
||||
@@ -130,6 +130,84 @@ impl PluginHandler for TestHandler {
|
||||
Box::pin(async move {
|
||||
match call.method.as_str() {
|
||||
"echo" => Ok(CallReply::unary(call.params, call.payload)),
|
||||
"key_facts" => {
|
||||
let result = call
|
||||
.host
|
||||
.key_facts(gateway_plugin_sdk::call::data::ClientKeyFactsQuery {
|
||||
client_key_id: "key_1".into(),
|
||||
})
|
||||
.await?;
|
||||
Ok(CallReply::unary(
|
||||
serde_json::to_value(result).unwrap(),
|
||||
vec![],
|
||||
))
|
||||
}
|
||||
"get_budget" => {
|
||||
let result = call
|
||||
.host
|
||||
.get_key_budget(
|
||||
gateway_plugin_sdk::call::key_budgets::GetKeyBudgetRequest {
|
||||
client_key_id: "key_1".into(),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
Ok(CallReply::unary(
|
||||
serde_json::to_value(result).unwrap(),
|
||||
vec![],
|
||||
))
|
||||
}
|
||||
"update_budget_limits" => {
|
||||
let result = call
|
||||
.host
|
||||
.update_key_budget_limits(
|
||||
gateway_plugin_sdk::call::key_budgets::UpdateKeyBudgetLimitsRequest {
|
||||
client_key_id: "key_1".into(),
|
||||
daily_limit_usd: None,
|
||||
weekly_limit_usd: Some("12.5".into()),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
Ok(CallReply::unary(
|
||||
serde_json::to_value(result).unwrap(),
|
||||
vec![],
|
||||
))
|
||||
}
|
||||
"refresh_quota" => {
|
||||
let result = call
|
||||
.host
|
||||
.refresh_account_quota(gateway_plugin_sdk::call::data::QuotaFactsQuery {
|
||||
account_id: "acct_1".into(),
|
||||
})
|
||||
.await?;
|
||||
Ok(CallReply::unary(
|
||||
serde_json::to_value(result).unwrap(),
|
||||
vec![],
|
||||
))
|
||||
}
|
||||
"reset_budget" => {
|
||||
use gateway_plugin_sdk::call::{
|
||||
host::KeyListRequest,
|
||||
key_budgets::{BudgetPeriod, ResetKeyBudgetRequest},
|
||||
};
|
||||
let keys = call
|
||||
.host
|
||||
.list_keys(KeyListRequest {
|
||||
cursor: None,
|
||||
limit: 10,
|
||||
})
|
||||
.await?;
|
||||
let result = call
|
||||
.host
|
||||
.reset_key_budget(ResetKeyBudgetRequest {
|
||||
client_key_id: keys.keys[0].id.clone(),
|
||||
period: BudgetPeriod::Weekly,
|
||||
})
|
||||
.await?;
|
||||
Ok(CallReply::unary(
|
||||
serde_json::to_value(result).unwrap(),
|
||||
vec![],
|
||||
))
|
||||
}
|
||||
"slow" => {
|
||||
tokio::time::sleep(Duration::from_millis(80)).await;
|
||||
Ok(CallReply::unary(call.params, call.payload))
|
||||
@@ -575,6 +653,76 @@ async fn shutdown(host: &mut HostPeer, task: JoinHandle<Result<(), SessionError>
|
||||
.expect("plugin session shutdown failed");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn typed_budget_helpers_preserve_wire_contract_and_propagate_failure_without_retry() {
|
||||
for failed in [false, true] {
|
||||
let (mut host, task) = start_session(TestHandler::default()).await;
|
||||
send_call(&mut host, 1, "reset_budget", json!({}), vec![]).await;
|
||||
let list = receive(&mut host).await;
|
||||
let Message::Callback {
|
||||
id,
|
||||
parent_id: 1,
|
||||
method,
|
||||
params,
|
||||
} = list.message
|
||||
else {
|
||||
panic!("expected Key list callback")
|
||||
};
|
||||
assert_eq!(method, "host.keys.list");
|
||||
assert_eq!(params, json!({"cursor":null,"limit":10}));
|
||||
assert!(list.payload.is_empty());
|
||||
send_control(&mut host, Message::Result { id, result: json!({"keys":[{"id":"key_1","name":"test","enabled":true}],"next_cursor":null}) }).await;
|
||||
let reset = receive(&mut host).await;
|
||||
let Message::Callback {
|
||||
id,
|
||||
parent_id: 1,
|
||||
method,
|
||||
params,
|
||||
} = reset.message
|
||||
else {
|
||||
panic!("expected budget reset callback")
|
||||
};
|
||||
assert_eq!(method, "host.keys.reset_budget");
|
||||
assert_eq!(params, json!({}));
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<Value>(&reset.payload).unwrap(),
|
||||
json!({"client_key_id":"key_1","period":"weekly"})
|
||||
);
|
||||
if failed {
|
||||
send_control(
|
||||
&mut host,
|
||||
Message::Error {
|
||||
id,
|
||||
error: PluginFault::new(ErrorCode::Conflict, "stale instance"),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let Message::Error { id: 1, error } = receive(&mut host).await.message else {
|
||||
panic!("expected error propagated without another callback")
|
||||
};
|
||||
assert_eq!(error.code, ErrorCode::Conflict);
|
||||
} else {
|
||||
write_frame(
|
||||
&mut host.writer,
|
||||
&Frame {
|
||||
message: Message::Result {
|
||||
id,
|
||||
result: json!({}),
|
||||
},
|
||||
payload: serde_json::to_vec(&json!({"client_key_id":"key_1"})).unwrap(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let Message::Result { id: 1, result } = receive(&mut host).await.message else {
|
||||
panic!("expected successful reset")
|
||||
};
|
||||
assert_eq!(result, json!({"client_key_id":"key_1"}));
|
||||
}
|
||||
shutdown(&mut host, task).await;
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn independent_calls_complete_concurrently() {
|
||||
let (mut host, task) = start_session(TestHandler::default()).await;
|
||||
@@ -1432,3 +1580,89 @@ async fn quiesce_rejects_new_calls_and_shutdown_closes_the_session() {
|
||||
|
||||
shutdown(&mut host, task).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn typed_key_and_quota_calls_keep_payloads_and_do_not_retry_failures() {
|
||||
for (entry, method, request, response) in [
|
||||
(
|
||||
"key_facts",
|
||||
"host.data.keys.get",
|
||||
json!({"client_key_id":"key_1"}),
|
||||
json!({"schema_version":1,"client_key_id":"key_1","enabled":false,"group_ids":["grp_1"]}),
|
||||
),
|
||||
(
|
||||
"get_budget",
|
||||
"host.keys.get_budget",
|
||||
json!({"client_key_id":"key_1"}),
|
||||
json!({
|
||||
"client_key_id":"key_1", "daily_limit_usd":"10", "weekly_limit_usd":"20",
|
||||
"daily_used_usd":"1.25", "weekly_used_usd":"2.5", "daily_resets_at_ms":null, "weekly_resets_at_ms":null,
|
||||
}),
|
||||
),
|
||||
(
|
||||
"update_budget_limits",
|
||||
"host.keys.update_budget_limits",
|
||||
json!({"client_key_id":"key_1","weekly_limit_usd":"12.5"}),
|
||||
json!({"client_key_id":"key_1"}),
|
||||
),
|
||||
(
|
||||
"refresh_quota",
|
||||
"host.quota_observations.refresh",
|
||||
json!({"account_id":"acct_1"}),
|
||||
json!({"schema_version":1,"account_id":"acct_1","observed_at_ms":null,"windows":[]}),
|
||||
),
|
||||
] {
|
||||
for failed in [false, true] {
|
||||
let (mut host, task) = start_session(TestHandler::default()).await;
|
||||
send_call(&mut host, 1, entry, json!({}), vec![]).await;
|
||||
let callback = receive(&mut host).await;
|
||||
let Message::Callback {
|
||||
id,
|
||||
parent_id: 1,
|
||||
method: actual,
|
||||
params,
|
||||
} = callback.message
|
||||
else {
|
||||
panic!("expected callback");
|
||||
};
|
||||
assert_eq!(actual, method);
|
||||
assert_eq!(params, json!({}));
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<Value>(&callback.payload).unwrap(),
|
||||
request
|
||||
);
|
||||
if failed {
|
||||
send_control(
|
||||
&mut host,
|
||||
Message::Error {
|
||||
id,
|
||||
error: PluginFault::new(ErrorCode::Conflict, "changed"),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let Message::Error { id: 1, error } = receive(&mut host).await.message else {
|
||||
panic!("expected propagated error without retry");
|
||||
};
|
||||
assert_eq!(error.code, ErrorCode::Conflict);
|
||||
} else {
|
||||
write_frame(
|
||||
&mut host.writer,
|
||||
&Frame {
|
||||
message: Message::Result {
|
||||
id,
|
||||
result: json!({}),
|
||||
},
|
||||
payload: serde_json::to_vec(&response).unwrap(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let Message::Result { id: 1, result } = receive(&mut host).await.message else {
|
||||
panic!("expected projected response");
|
||||
};
|
||||
assert_eq!(result, response);
|
||||
}
|
||||
shutdown(&mut host, task).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,8 +6,9 @@ use chrono::{DateTime, Utc};
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_admin::model::{
|
||||
MutationContext,
|
||||
client_keys::{ClientKeyBudgetPeriod, ResetClientKeyBudget},
|
||||
client_keys::{ClientKeyBudgetMutationOrigin, ClientKeyBudgetPeriod, ResetClientKeyBudget},
|
||||
};
|
||||
use gateway_admin::ports::store::AdminStoreResult;
|
||||
use gateway_core::{
|
||||
engine::budget::{
|
||||
ClientBudgetCharge, ClientBudgetError, ClientBudgetLimits, ClientBudgetPort,
|
||||
@@ -24,17 +25,41 @@ use crate::{StoreError, StoreResult, mutation_audit, postgres_unavailable};
|
||||
pub(super) async fn reset_client_key_budget(
|
||||
pool: &PgPool,
|
||||
command: ResetClientKeyBudget,
|
||||
origin: ClientKeyBudgetMutationOrigin,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<()> {
|
||||
let mut tx = match &origin {
|
||||
ClientKeyBudgetMutationOrigin::Admin => pool.begin().await.map_err(|_| {
|
||||
crate::admin_store_error(
|
||||
"client API key budget",
|
||||
postgres_unavailable("begin budget reset"),
|
||||
)
|
||||
})?,
|
||||
ClientKeyBudgetMutationOrigin::Plugin(owner) => {
|
||||
super::plugins::begin_authorized_mutation(pool, owner, "key_budgets").await?
|
||||
}
|
||||
};
|
||||
reset_client_key_budget_in_transaction(&mut tx, &command, context)
|
||||
.await
|
||||
.map_err(|error| crate::admin_store_error("client API key budget", error))?;
|
||||
tx.commit().await.map_err(|_| {
|
||||
crate::admin_store_error(
|
||||
"client API key budget",
|
||||
postgres_unavailable("commit budget reset"),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
async fn reset_client_key_budget_in_transaction(
|
||||
tx: &mut Transaction<'_, Postgres>,
|
||||
command: &ResetClientKeyBudget,
|
||||
context: &MutationContext,
|
||||
) -> StoreResult<()> {
|
||||
let mut tx = pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("begin budget reset"))?;
|
||||
// 与准入、结算共用 Key 行锁,重置边界必须在取得锁之后确定。
|
||||
let exists =
|
||||
sqlx::query_scalar::<_, String>("select id from client_api_keys where id = $1 for update")
|
||||
.bind(command.id.as_str())
|
||||
.fetch_optional(&mut *tx)
|
||||
.fetch_optional(&mut **tx)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("lock budget reset key"))?;
|
||||
if exists.is_none() {
|
||||
@@ -65,7 +90,7 @@ pub(super) async fn reset_client_key_budget(
|
||||
.bind(daily)
|
||||
.bind(weekly)
|
||||
.bind(reset_at)
|
||||
.execute(&mut *tx)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("reset client budget"))?;
|
||||
let mut fields = Vec::new();
|
||||
@@ -76,7 +101,7 @@ pub(super) async fn reset_client_key_budget(
|
||||
fields.extend(["weekly_used_usd".to_owned(), "weekly_start".to_owned()]);
|
||||
}
|
||||
super::append_admin_audit_event_in_transaction(
|
||||
&mut tx,
|
||||
tx,
|
||||
mutation_audit(
|
||||
context,
|
||||
"reset_budget",
|
||||
@@ -87,9 +112,7 @@ pub(super) async fn reset_client_key_budget(
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("commit budget reset"))
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub struct PgClientBudgetStore {
|
||||
|
||||
@@ -708,14 +708,105 @@ impl PgAdminClientKeyStore {
|
||||
|
||||
#[async_trait]
|
||||
impl ClientKeyStore for PgAdminClientKeyStore {
|
||||
async fn update_client_key_budget_limits(
|
||||
&self,
|
||||
command: gateway_admin::model::client_keys::UpdateClientKeyBudgetLimits,
|
||||
origin: gateway_admin::model::client_keys::ClientKeyBudgetMutationOrigin,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<Option<gateway_admin::model::Revision>> {
|
||||
use gateway_admin::model::client_keys::ClientKeyBudgetMutationOrigin;
|
||||
let map_error = |error| admin_store_error(ENTITY, error);
|
||||
let mut tx = match &origin {
|
||||
ClientKeyBudgetMutationOrigin::Admin => self
|
||||
.keys
|
||||
.pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|_| map_error(postgres_unavailable("begin budget limits update")))?,
|
||||
ClientKeyBudgetMutationOrigin::Plugin(owner) => {
|
||||
super::plugins::begin_authorized_mutation(&self.keys.pool, owner, "key_budgets")
|
||||
.await?
|
||||
}
|
||||
};
|
||||
// 保持与完整 Key 编辑相同的锁顺序;绝不读出整份配置再覆盖写回。
|
||||
sqlx::query("select config_revision from runtime_settings where id=1 for update")
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|_| map_error(postgres_unavailable("lock budget limits configuration")))?;
|
||||
let current: Option<(String, String)> = sqlx::query_as(
|
||||
"select daily_limit_usd::text, weekly_limit_usd::text from client_api_keys where id=$1 for update"
|
||||
).bind(command.id.as_str()).fetch_optional(&mut *tx).await
|
||||
.map_err(|_| map_error(postgres_unavailable("lock budget limits key")))?;
|
||||
let (daily, weekly) = current.ok_or_else(|| {
|
||||
map_error(StoreError::NotFound {
|
||||
entity: ENTITY,
|
||||
id: command.id.as_str().to_owned(),
|
||||
})
|
||||
})?;
|
||||
let daily: Decimal = daily
|
||||
.parse()
|
||||
.map_err(|_| map_error(postgres_unavailable("decode daily budget limit")))?;
|
||||
let weekly: Decimal = weekly
|
||||
.parse()
|
||||
.map_err(|_| map_error(postgres_unavailable("decode weekly budget limit")))?;
|
||||
let mut fields = Vec::new();
|
||||
if command.daily_limit_usd.is_some_and(|value| value != daily) {
|
||||
fields.push("daily_limit_usd".to_owned());
|
||||
}
|
||||
if command
|
||||
.weekly_limit_usd
|
||||
.is_some_and(|value| value != weekly)
|
||||
{
|
||||
fields.push("weekly_limit_usd".to_owned());
|
||||
}
|
||||
if fields.is_empty() {
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|_| map_error(postgres_unavailable("finish unchanged budget limits")))?;
|
||||
return Ok(None);
|
||||
}
|
||||
sqlx::query(
|
||||
"update client_api_keys set
|
||||
daily_limit_usd=coalesce($2::text::numeric, daily_limit_usd),
|
||||
weekly_limit_usd=coalesce($3::text::numeric, weekly_limit_usd), updated_at=now()
|
||||
where id=$1",
|
||||
)
|
||||
.bind(command.id.as_str())
|
||||
.bind(command.daily_limit_usd.map(|value| value.canonical()))
|
||||
.bind(command.weekly_limit_usd.map(|value| value.canonical()))
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|_| map_error(postgres_unavailable("update budget limits")))?;
|
||||
let revision = super::bump_config_revision_in_transaction(&mut tx)
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
super::append_admin_audit_event_in_transaction(
|
||||
&mut tx,
|
||||
mutation_audit(
|
||||
context,
|
||||
"update_budget_limits",
|
||||
"client_api_key",
|
||||
command.id.as_str(),
|
||||
fields,
|
||||
),
|
||||
revision,
|
||||
)
|
||||
.await
|
||||
.map_err(map_error)?;
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|_| map_error(postgres_unavailable("commit budget limits update")))?;
|
||||
Ok(Some(admin_revision(revision)?))
|
||||
}
|
||||
|
||||
async fn reset_client_key_budget(
|
||||
&self,
|
||||
command: ResetClientKeyBudget,
|
||||
origin: gateway_admin::model::client_keys::ClientKeyBudgetMutationOrigin,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<()> {
|
||||
super::client_budgets::reset_client_key_budget(&self.keys.pool, command, context)
|
||||
super::client_budgets::reset_client_key_budget(&self.keys.pool, command, origin, context)
|
||||
.await
|
||||
.map_err(|error| admin_store_error(ENTITY, error))
|
||||
}
|
||||
|
||||
async fn get_client_key(
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
//! 插件管理写入的事务授权;与实例修改共用控制面锁,避免检查后权限被撤销。
|
||||
|
||||
use gateway_admin::{
|
||||
model::plugin_resources::PluginResourceOwner,
|
||||
ports::store::{AdminStoreError, AdminStoreErrorKind, AdminStoreResult},
|
||||
};
|
||||
use sqlx::{PgPool, Postgres, Transaction};
|
||||
|
||||
fn mutation_unavailable() -> AdminStoreError {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Unavailable,
|
||||
"plugin mutation",
|
||||
"plugin mutation store is unavailable",
|
||||
)
|
||||
}
|
||||
|
||||
fn mutation_denied() -> AdminStoreError {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Conflict,
|
||||
"plugin mutation",
|
||||
"plugin instance authorization changed",
|
||||
)
|
||||
}
|
||||
|
||||
pub(in crate::postgres) async fn begin_authorized_mutation<'a>(
|
||||
pool: &'a PgPool,
|
||||
owner: &PluginResourceOwner,
|
||||
permission: &str,
|
||||
) -> AdminStoreResult<Transaction<'a, Postgres>> {
|
||||
let mut tx = pool.begin().await.map_err(|_| mutation_unavailable())?;
|
||||
// 与全部管理写入保持相同锁顺序;无变化的对账只锁定,不递增 revision。
|
||||
sqlx::query("select config_revision from runtime_settings where id=1 for update")
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|_| mutation_unavailable())?;
|
||||
let id = uuid::Uuid::parse_str(&owner.instance_id).map_err(|_| mutation_denied())?;
|
||||
let revision = i64::try_from(owner.revision.get()).map_err(|_| mutation_denied())?;
|
||||
let allowed: bool = sqlx::query_scalar(
|
||||
"select exists(
|
||||
select 1 from plugin_instances i
|
||||
join plugin_artifacts a on a.sha256=i.artifact_sha256
|
||||
where i.id=$1 and i.enabled and i.revision=$2 and i.artifact_sha256=$3
|
||||
and a.accepted_at is not null and a.metadata_json->'requestedPermissions' ? $4
|
||||
)",
|
||||
)
|
||||
.bind(id)
|
||||
.bind(revision)
|
||||
.bind(&owner.artifact_sha256)
|
||||
.bind(permission)
|
||||
.fetch_one(&mut *tx)
|
||||
.await
|
||||
.map_err(|_| mutation_unavailable())?;
|
||||
if !allowed {
|
||||
return Err(mutation_denied());
|
||||
}
|
||||
Ok(tx)
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
mod artifacts;
|
||||
mod authorization;
|
||||
mod credentials;
|
||||
mod instances;
|
||||
mod resources;
|
||||
@@ -6,3 +7,5 @@ mod sources;
|
||||
mod state;
|
||||
|
||||
pub use artifacts::PgPluginStore;
|
||||
|
||||
pub(super) use authorization::begin_authorized_mutation;
|
||||
|
||||
@@ -62,42 +62,6 @@ fn validate_resource_key(key: &str) -> AdminStoreResult<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
impl PgPluginStore {
|
||||
async fn resource_transaction(
|
||||
&self,
|
||||
owner: &PluginResourceOwner,
|
||||
permission: &str,
|
||||
) -> AdminStoreResult<Transaction<'_, Postgres>> {
|
||||
let mut tx = self.pool.begin().await.map_err(|_| unavailable())?;
|
||||
// 与全部管理写入保持相同锁顺序;无变化的对账只锁定,不递增 revision。
|
||||
sqlx::query("select config_revision from runtime_settings where id=1 for update")
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|_| unavailable())?;
|
||||
let id = uuid::Uuid::parse_str(&owner.instance_id).map_err(|_| denied())?;
|
||||
let revision = i64::try_from(owner.revision.get()).map_err(|_| denied())?;
|
||||
let allowed: bool = sqlx::query_scalar(
|
||||
"select exists(
|
||||
select 1 from plugin_instances i
|
||||
join plugin_artifacts a on a.sha256=i.artifact_sha256
|
||||
where i.id=$1 and i.enabled and i.revision=$2 and i.artifact_sha256=$3
|
||||
and a.accepted_at is not null and a.metadata_json->'requestedPermissions' ? $4
|
||||
)",
|
||||
)
|
||||
.bind(id)
|
||||
.bind(revision)
|
||||
.bind(&owner.artifact_sha256)
|
||||
.bind(permission)
|
||||
.fetch_one(&mut *tx)
|
||||
.await
|
||||
.map_err(|_| unavailable())?;
|
||||
if !allowed {
|
||||
return Err(denied());
|
||||
}
|
||||
Ok(tx)
|
||||
}
|
||||
}
|
||||
|
||||
async fn finish<T>(
|
||||
mut tx: Transaction<'_, Postgres>,
|
||||
context: &MutationContext,
|
||||
@@ -146,7 +110,7 @@ impl PluginResourceStore for PgPluginStore {
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<ResourceMutation<ManagedResource>> {
|
||||
validate_resource_key(&resource_key)?;
|
||||
let mut tx = self.resource_transaction(owner, "groups").await?;
|
||||
let mut tx = super::begin_authorized_mutation(&self.pool, owner, "groups").await?;
|
||||
let instance = uuid::Uuid::parse_str(&owner.instance_id).map_err(|_| denied())?;
|
||||
if let Some(row) = sqlx::query(
|
||||
"select g.id,g.name,g.enabled from plugin_group_resources r
|
||||
@@ -204,7 +168,7 @@ impl PluginResourceStore for PgPluginStore {
|
||||
for group in &groups {
|
||||
validate_resource_key(group)?;
|
||||
}
|
||||
let mut tx = self.resource_transaction(owner, "keys").await?;
|
||||
let mut tx = super::begin_authorized_mutation(&self.pool, owner, "keys").await?;
|
||||
let instance = uuid::Uuid::parse_str(&owner.instance_id).map_err(|_| denied())?;
|
||||
if let Some(row) = sqlx::query(
|
||||
"select k.id,k.name,k.enabled from plugin_key_resources r
|
||||
@@ -274,7 +238,7 @@ impl PluginResourceStore for PgPluginStore {
|
||||
{
|
||||
return Err(invalid());
|
||||
}
|
||||
let mut tx = self.resource_transaction(owner, "groups").await?;
|
||||
let mut tx = super::begin_authorized_mutation(&self.pool, owner, "groups").await?;
|
||||
let instance = uuid::Uuid::parse_str(&owner.instance_id).map_err(|_| denied())?;
|
||||
let group: String = sqlx::query_scalar(
|
||||
"select group_id from plugin_group_resources where instance_id=$1 and resource_key=$2",
|
||||
|
||||
@@ -8,6 +8,20 @@ use gateway_core::provider_ports::ProviderCooldownPort;
|
||||
use super::*;
|
||||
use crate::postgres::ObservabilityQueryBudget;
|
||||
|
||||
fn account_capacity(
|
||||
account: &AccountRecord,
|
||||
default_concurrency: u64,
|
||||
in_flight: Option<&BTreeMap<String, u64>>,
|
||||
) -> gateway_admin::model::accounts::AccountCapacity {
|
||||
gateway_admin::model::accounts::AccountCapacity {
|
||||
used_slots: in_flight.map(|counts| counts.get(&account.id).copied().unwrap_or(0)),
|
||||
total_slots: account.concurrency_limit.map_or_else(
|
||||
|| (default_concurrency > 0).then_some(default_concurrency),
|
||||
|limit| Some(u64::from(limit.get())),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// Admin 账号用例所需的公共账号、留存观测与 revision 事务能力。
|
||||
///
|
||||
/// 三个 PostgreSQL adapter 都保持私有,调用方只能取得 [`AccountStore`] 暴露的领域能力。
|
||||
@@ -376,6 +390,11 @@ impl AccountStore for PgAdminAccountStore {
|
||||
let mut account = admin_account_record(summary)?;
|
||||
account.groups = groups_by_account.remove(&account_id).unwrap_or_default();
|
||||
Ok(AccountPageItem {
|
||||
capacity: account_capacity(
|
||||
&account,
|
||||
page.default_concurrency,
|
||||
runtime.in_flight.as_ref(),
|
||||
),
|
||||
account,
|
||||
projection,
|
||||
})
|
||||
@@ -411,7 +430,21 @@ impl AccountStore for PgAdminAccountStore {
|
||||
.await?;
|
||||
let mut account = admin_account_record(record.summary)?;
|
||||
account.groups = groups.remove(&account_id).unwrap_or_default();
|
||||
let default_concurrency: i64 = sqlx::query_scalar(
|
||||
"select max_concurrent_per_account from runtime_settings where id = 1",
|
||||
)
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_err(|_| admin_store_error(ENTITY, postgres_unavailable("load account concurrency")))?;
|
||||
let default_concurrency = u64::try_from(default_concurrency).map_err(|_| {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Invalid,
|
||||
ENTITY,
|
||||
"invalid default account concurrency",
|
||||
)
|
||||
})?;
|
||||
Ok(Some(AccountPageItem {
|
||||
capacity: account_capacity(&account, default_concurrency, runtime.in_flight.as_ref()),
|
||||
account,
|
||||
projection,
|
||||
}))
|
||||
|
||||
@@ -4,6 +4,7 @@ use super::*;
|
||||
|
||||
pub(crate) struct AdminAccountPageRows {
|
||||
pub(crate) config_revision: AdminRevision,
|
||||
pub(crate) default_concurrency: u64,
|
||||
pub(crate) accounts: Vec<ProviderAccountSummary>,
|
||||
pub(crate) total: u64,
|
||||
pub(crate) summary: AccountSummary,
|
||||
@@ -112,7 +113,7 @@ pub(crate) async fn load_admin_account_page(
|
||||
select count(*)::bigint as filtered_total from filtered
|
||||
),
|
||||
settings as (
|
||||
select config_revision, usage_retention_days
|
||||
select config_revision, usage_retention_days, max_concurrent_per_account
|
||||
from runtime_settings
|
||||
where id = 1
|
||||
)
|
||||
@@ -127,7 +128,7 @@ pub(crate) async fn load_admin_account_page(
|
||||
global_summary.summary_total, global_summary.summary_normal,
|
||||
global_summary.summary_quota_exhausted, global_summary.summary_rate_limited,
|
||||
global_summary.summary_disabled, global_summary.summary_error,
|
||||
settings.config_revision
|
||||
settings.config_revision, settings.max_concurrent_per_account
|
||||
from filtered_total
|
||||
cross join global_summary
|
||||
cross join settings
|
||||
@@ -168,6 +169,7 @@ pub(crate) async fn load_admin_account_page(
|
||||
)
|
||||
})?;
|
||||
let config_revision = revision_from_row(metadata)?;
|
||||
let default_concurrency = unsigned_metadata(metadata, "max_concurrent_per_account")?;
|
||||
let total = unsigned_metadata(metadata, "filtered_total")?;
|
||||
let summary = AccountSummary {
|
||||
total: unsigned_metadata(metadata, "summary_total")?,
|
||||
@@ -190,6 +192,7 @@ pub(crate) async fn load_admin_account_page(
|
||||
.map_err(|error| admin_store_error(ENTITY, error))?;
|
||||
Ok(AdminAccountPageRows {
|
||||
config_revision,
|
||||
default_concurrency,
|
||||
accounts,
|
||||
total,
|
||||
summary,
|
||||
|
||||
@@ -1,12 +1,23 @@
|
||||
use std::time::SystemTime;
|
||||
use std::{collections::BTreeMap, time::SystemTime};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use gateway_admin::{
|
||||
model::{
|
||||
MutationActor, MutationContext,
|
||||
client_keys::{ClientKeyBudgetPeriod, ResetClientKeyBudget, UpdateClientKey},
|
||||
MutationActor, MutationContext, Revision,
|
||||
client_keys::{
|
||||
ClientKeyBudgetMutationOrigin, ClientKeyBudgetPeriod, ResetClientKeyBudget,
|
||||
UpdateClientKey, UpdateClientKeyBudgetLimits,
|
||||
},
|
||||
plugin_resources::PluginResourceOwner,
|
||||
plugins::{
|
||||
PluginSource,
|
||||
instances::{PluginInstance, PluginPermissionGrant},
|
||||
},
|
||||
},
|
||||
ports::{
|
||||
plugins::PluginStore as _,
|
||||
store::{AdminStoreErrorKind, ClientKeyStore as _},
|
||||
},
|
||||
ports::store::ClientKeyStore as _,
|
||||
};
|
||||
use gateway_core::{
|
||||
engine::{
|
||||
@@ -18,7 +29,7 @@ use gateway_core::{
|
||||
};
|
||||
use gateway_store::postgres::{
|
||||
ClientApiKeyRepository as _, PgAdminClientKeyStore, PgClientApiKeyRepository,
|
||||
PgClientBudgetStore,
|
||||
PgClientBudgetStore, PgPluginStore,
|
||||
};
|
||||
|
||||
use super::TestDatabase;
|
||||
@@ -95,6 +106,7 @@ async fn manual_reset_clears_only_selected_budget_and_preserves_policy_history_a
|
||||
id: key_id(key),
|
||||
period,
|
||||
},
|
||||
ClientKeyBudgetMutationOrigin::Admin,
|
||||
&context(),
|
||||
)
|
||||
.await
|
||||
@@ -143,6 +155,7 @@ async fn manual_reset_excludes_old_completions_but_counts_inflight_requests_afte
|
||||
id: key_id("key"),
|
||||
period: ClientKeyBudgetPeriod::Daily
|
||||
},
|
||||
ClientKeyBudgetMutationOrigin::Admin,
|
||||
&reset_context
|
||||
),
|
||||
store.settle(delayed),
|
||||
@@ -183,6 +196,7 @@ async fn manual_reset_leaves_unused_and_expired_windows_inactive_and_reports_mis
|
||||
id: key_id(key),
|
||||
period: ClientKeyBudgetPeriod::All,
|
||||
},
|
||||
ClientKeyBudgetMutationOrigin::Admin,
|
||||
&context(),
|
||||
)
|
||||
.await
|
||||
@@ -203,6 +217,7 @@ async fn manual_reset_leaves_unused_and_expired_windows_inactive_and_reports_mis
|
||||
id: key_id("missing"),
|
||||
period: ClientKeyBudgetPeriod::All,
|
||||
},
|
||||
ClientKeyBudgetMutationOrigin::Admin,
|
||||
&context(),
|
||||
)
|
||||
.await
|
||||
@@ -240,6 +255,7 @@ async fn manual_reset_rolls_back_when_audit_cannot_be_written() {
|
||||
id: key_id("key"),
|
||||
period: ClientKeyBudgetPeriod::All
|
||||
},
|
||||
ClientKeyBudgetMutationOrigin::Admin,
|
||||
&context()
|
||||
)
|
||||
.await
|
||||
@@ -562,3 +578,451 @@ async fn budget_database_outage_fails_closed() {
|
||||
assert!(store.settle(charge("key", "offline", "1")).await.is_err());
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
async fn plugin_reset_owner(database: &TestDatabase) -> PluginResourceOwner {
|
||||
super::plugins::artifacts::initialize_revision(database).await;
|
||||
let store = PgPluginStore::new(database.pool.clone());
|
||||
let mut package = super::plugins::artifacts::artifact('b', &["linux-x86_64"]);
|
||||
package.metadata.requested_permissions = vec!["key_budgets".into()];
|
||||
let installed = store
|
||||
.install_artifact(package, PluginSource::Upload, &context())
|
||||
.await
|
||||
.unwrap();
|
||||
let accepted = store
|
||||
.accept_artifact(&installed.artifact.metadata.sha256, &context())
|
||||
.await
|
||||
.unwrap();
|
||||
let instance = store
|
||||
.save_instance(
|
||||
PluginInstance {
|
||||
id: uuid::Uuid::now_v7().to_string(),
|
||||
name: "budget reset".into(),
|
||||
artifact_sha256: installed.artifact.metadata.sha256,
|
||||
enabled: true,
|
||||
trusted_process: true,
|
||||
configuration: serde_json::json!({}),
|
||||
secrets: BTreeMap::new(),
|
||||
grants: vec![PluginPermissionGrant {
|
||||
permission: "key_budgets".into(),
|
||||
}],
|
||||
bindings: vec![],
|
||||
revision: Revision::new(1).unwrap(),
|
||||
},
|
||||
accepted.config_revision,
|
||||
&context(),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.instance;
|
||||
PluginResourceOwner {
|
||||
instance_id: instance.id,
|
||||
artifact_sha256: instance.artifact_sha256,
|
||||
revision: instance.revision,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn plugin_resets_share_the_native_ledger_and_each_call_resets_again() {
|
||||
let Some(database) = TestDatabase::create("plugin_budget_reset").await else {
|
||||
return;
|
||||
};
|
||||
let owner = plugin_reset_owner(&database).await;
|
||||
seed(&database, "key", "10", "20").await;
|
||||
let budgets = PgClientBudgetStore::new(database.pool.clone());
|
||||
budgets.settle(charge("key", "before", "3")).await.unwrap();
|
||||
let before = status(&database, "key").await;
|
||||
let revision: i64 =
|
||||
sqlx::query_scalar("select config_revision from runtime_settings where id=1")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let store = PgAdminClientKeyStore::new(database.pool.clone());
|
||||
let command = ResetClientKeyBudget {
|
||||
id: key_id("key"),
|
||||
period: ClientKeyBudgetPeriod::Weekly,
|
||||
};
|
||||
store
|
||||
.reset_client_key_budget(
|
||||
command.clone(),
|
||||
ClientKeyBudgetMutationOrigin::Plugin(owner.clone()),
|
||||
&context(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let after = status(&database, "key").await;
|
||||
assert_eq!(after.daily_used_usd.canonical(), "3");
|
||||
assert_eq!(after.weekly_used_usd.canonical(), "0");
|
||||
assert_eq!(after.limits, before.limits);
|
||||
assert_eq!(after.daily_resets_at, before.daily_resets_at);
|
||||
assert_eq!(after.weekly_resets_at, before.weekly_resets_at);
|
||||
|
||||
budgets.settle(charge("key", "after", "2")).await.unwrap();
|
||||
assert_eq!(
|
||||
status(&database, "key").await.weekly_used_usd.canonical(),
|
||||
"2"
|
||||
);
|
||||
store
|
||||
.reset_client_key_budget(
|
||||
command,
|
||||
ClientKeyBudgetMutationOrigin::Plugin(owner),
|
||||
&context(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let after = status(&database, "key").await;
|
||||
assert_eq!(after.daily_used_usd.canonical(), "5");
|
||||
assert_eq!(after.weekly_used_usd.canonical(), "0");
|
||||
let (audits, charges, current_revision): (i64, i64, i64) = sqlx::query_as(
|
||||
"select (select count(*) from admin_audit_events where action='reset_budget' and config_revision is null),
|
||||
(select count(*) from client_key_charge_events),
|
||||
(select config_revision from runtime_settings where id=1)"
|
||||
).fetch_one(&database.pool).await.unwrap();
|
||||
assert_eq!((audits, charges, current_revision), (2, 2, revision));
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn plugin_reset_revalidates_current_authorization_without_changing_the_ledger() {
|
||||
let Some(database) = TestDatabase::create("plugin_budget_authorization").await else {
|
||||
return;
|
||||
};
|
||||
let owner = plugin_reset_owner(&database).await;
|
||||
seed(&database, "key", "10", "20").await;
|
||||
PgClientBudgetStore::new(database.pool.clone())
|
||||
.settle(charge("key", "before", "4"))
|
||||
.await
|
||||
.unwrap();
|
||||
let before = status(&database, "key").await;
|
||||
let store = PgAdminClientKeyStore::new(database.pool.clone());
|
||||
for sql in [
|
||||
"update plugin_instances set revision=revision+1",
|
||||
"update plugin_instances set revision=revision-1, enabled=false",
|
||||
"update plugin_instances set enabled=true; update plugin_artifacts set accepted_at=null",
|
||||
"update plugin_artifacts set accepted_at=now(), metadata_json=jsonb_set(metadata_json,'{requestedPermissions}','[]'::jsonb)",
|
||||
] {
|
||||
sqlx::raw_sql(sql).execute(&database.pool).await.unwrap();
|
||||
assert_eq!(
|
||||
store
|
||||
.reset_client_key_budget(
|
||||
ResetClientKeyBudget {
|
||||
id: key_id("key"),
|
||||
period: ClientKeyBudgetPeriod::All
|
||||
},
|
||||
ClientKeyBudgetMutationOrigin::Plugin(owner.clone()),
|
||||
&context()
|
||||
)
|
||||
.await
|
||||
.unwrap_err()
|
||||
.kind(),
|
||||
AdminStoreErrorKind::Conflict
|
||||
);
|
||||
assert_eq!(status(&database, "key").await, before);
|
||||
}
|
||||
sqlx::query("update plugin_artifacts set metadata_json=jsonb_set(metadata_json,'{requestedPermissions}','[\"key_budgets\"]'::jsonb)")
|
||||
.execute(&database.pool).await.unwrap();
|
||||
let wrong_artifact = PluginResourceOwner {
|
||||
artifact_sha256: "c".repeat(64),
|
||||
..owner.clone()
|
||||
};
|
||||
assert_eq!(
|
||||
store
|
||||
.reset_client_key_budget(
|
||||
ResetClientKeyBudget {
|
||||
id: key_id("key"),
|
||||
period: ClientKeyBudgetPeriod::All
|
||||
},
|
||||
ClientKeyBudgetMutationOrigin::Plugin(wrong_artifact),
|
||||
&context()
|
||||
)
|
||||
.await
|
||||
.unwrap_err()
|
||||
.kind(),
|
||||
AdminStoreErrorKind::Conflict
|
||||
);
|
||||
assert_eq!(status(&database, "key").await, before);
|
||||
assert_eq!(
|
||||
store
|
||||
.reset_client_key_budget(
|
||||
ResetClientKeyBudget {
|
||||
id: key_id("missing"),
|
||||
period: ClientKeyBudgetPeriod::All
|
||||
},
|
||||
ClientKeyBudgetMutationOrigin::Plugin(owner.clone()),
|
||||
&context()
|
||||
)
|
||||
.await
|
||||
.unwrap_err()
|
||||
.kind(),
|
||||
AdminStoreErrorKind::NotFound
|
||||
);
|
||||
// 验证拒绝后的事务和锁已释放,当前身份仍可正常执行。
|
||||
store
|
||||
.reset_client_key_budget(
|
||||
ResetClientKeyBudget {
|
||||
id: key_id("key"),
|
||||
period: ClientKeyBudgetPeriod::All,
|
||||
},
|
||||
ClientKeyBudgetMutationOrigin::Plugin(owner),
|
||||
&context(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
status(&database, "key").await.daily_used_usd.canonical(),
|
||||
"0"
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn plugin_reset_rolls_back_with_audit_and_serializes_with_late_settlement() {
|
||||
let Some(database) = TestDatabase::create("plugin_budget_atomicity").await else {
|
||||
return;
|
||||
};
|
||||
let owner = plugin_reset_owner(&database).await;
|
||||
seed(&database, "key", "10", "20").await;
|
||||
let budgets = PgClientBudgetStore::new(database.pool.clone());
|
||||
budgets.settle(charge("key", "before", "4")).await.unwrap();
|
||||
let before = status(&database, "key").await;
|
||||
let store = PgAdminClientKeyStore::new(database.pool.clone());
|
||||
let command = ResetClientKeyBudget {
|
||||
id: key_id("key"),
|
||||
period: ClientKeyBudgetPeriod::Weekly,
|
||||
};
|
||||
sqlx::raw_sql("create function reject_budget_audit() returns trigger language plpgsql as $$ begin raise exception 'test rollback'; end $$;
|
||||
create trigger reject_budget_audit before insert on admin_audit_events for each row execute function reject_budget_audit()")
|
||||
.execute(&database.pool).await.unwrap();
|
||||
assert!(
|
||||
store
|
||||
.reset_client_key_budget(
|
||||
command.clone(),
|
||||
ClientKeyBudgetMutationOrigin::Plugin(owner.clone()),
|
||||
&context()
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert_eq!(status(&database, "key").await, before);
|
||||
sqlx::query("drop trigger reject_budget_audit on admin_audit_events")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let mutation = context();
|
||||
// 完成时间在重置前,无论谁先拿到行锁,周用量都不能被迟到结算重新扣回。
|
||||
let delayed = charge("key", "delayed", "1");
|
||||
let (reset, settlement) = tokio::join!(
|
||||
store.reset_client_key_budget(
|
||||
command,
|
||||
ClientKeyBudgetMutationOrigin::Plugin(owner),
|
||||
&mutation
|
||||
),
|
||||
budgets.settle(delayed),
|
||||
);
|
||||
reset.unwrap();
|
||||
settlement.unwrap();
|
||||
let after = status(&database, "key").await;
|
||||
assert_eq!(after.daily_used_usd.canonical(), "5");
|
||||
assert_eq!(after.weekly_used_usd.canonical(), "0");
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
fn budget_update(
|
||||
key: &str,
|
||||
daily: Option<&str>,
|
||||
weekly: Option<&str>,
|
||||
) -> UpdateClientKeyBudgetLimits {
|
||||
UpdateClientKeyBudgetLimits {
|
||||
id: key_id(key),
|
||||
daily_limit_usd: daily.map(|value| value.parse().unwrap()),
|
||||
weekly_limit_usd: weekly.map(|value| value.parse().unwrap()),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn plugin_budget_limits_preserve_consumption_and_unrelated_configuration_and_control_admission()
|
||||
{
|
||||
let Some(database) = TestDatabase::create("plugin_budget_limits").await else {
|
||||
return;
|
||||
};
|
||||
let owner = plugin_reset_owner(&database).await;
|
||||
seed(&database, "key", "10", "20").await;
|
||||
sqlx::query("update client_api_keys set label='preserved', max_concurrency=2, requests_per_minute=30 where id='key'")
|
||||
.execute(&database.pool).await.unwrap();
|
||||
let budgets = PgClientBudgetStore::new(database.pool.clone());
|
||||
budgets.settle(charge("key", "before", "4")).await.unwrap();
|
||||
let store = PgAdminClientKeyStore::new(database.pool.clone());
|
||||
let before = store.get_client_key(&key_id("key")).await.unwrap().unwrap();
|
||||
let revision: i64 =
|
||||
sqlx::query_scalar("select config_revision from runtime_settings where id=1")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let command = budget_update("key", None, Some("3"));
|
||||
let changed = store
|
||||
.update_client_key_budget_limits(
|
||||
command.clone(),
|
||||
ClientKeyBudgetMutationOrigin::Plugin(owner.clone()),
|
||||
&context(),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(changed.get(), revision as u64 + 1);
|
||||
let after = store.get_client_key(&key_id("key")).await.unwrap().unwrap();
|
||||
assert_eq!(after.budget.daily_used_usd, before.budget.daily_used_usd);
|
||||
assert_eq!(after.budget.weekly_used_usd, before.budget.weekly_used_usd);
|
||||
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!(
|
||||
after.budget.limits.daily_usd,
|
||||
before.budget.limits.daily_usd
|
||||
);
|
||||
assert_eq!(after.budget.limits.weekly_usd.canonical(), "3");
|
||||
assert_eq!(
|
||||
(
|
||||
&after.name,
|
||||
&after.label,
|
||||
&after.groups,
|
||||
after.limits,
|
||||
&after.request_profile_overrides
|
||||
),
|
||||
(
|
||||
&before.name,
|
||||
&before.label,
|
||||
&before.groups,
|
||||
before.limits,
|
||||
&before.request_profile_overrides
|
||||
)
|
||||
);
|
||||
assert_eq!(
|
||||
budgets.admit(key_id("key")).await.unwrap_err().kind(),
|
||||
GatewayErrorKind::RateLimited
|
||||
);
|
||||
assert!(
|
||||
store
|
||||
.update_client_key_budget_limits(
|
||||
command,
|
||||
ClientKeyBudgetMutationOrigin::Plugin(owner.clone()),
|
||||
&context()
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
let (new_revision, audits): (i64, i64) = sqlx::query_as("select config_revision, (select count(*) from admin_audit_events where action='update_budget_limits') from runtime_settings where id=1")
|
||||
.fetch_one(&database.pool).await.unwrap();
|
||||
assert_eq!((new_revision, audits), (revision + 1, 1));
|
||||
store
|
||||
.update_client_key_budget_limits(
|
||||
budget_update("key", None, Some("5")),
|
||||
ClientKeyBudgetMutationOrigin::Plugin(owner.clone()),
|
||||
&context(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
budgets.admit(key_id("key")).await.unwrap();
|
||||
let mutation = context();
|
||||
let (updated, settled) = tokio::join!(
|
||||
store.update_client_key_budget_limits(
|
||||
budget_update("key", Some("0"), Some("0")),
|
||||
ClientKeyBudgetMutationOrigin::Plugin(owner),
|
||||
&mutation
|
||||
),
|
||||
budgets.settle(charge("key", "inflight", "2")),
|
||||
);
|
||||
updated.unwrap();
|
||||
settled.unwrap();
|
||||
let after = status(&database, "key").await;
|
||||
assert_eq!(after.daily_used_usd.canonical(), "6");
|
||||
assert_eq!(after.weekly_used_usd.canonical(), "6");
|
||||
assert!(!after.limits.is_limited());
|
||||
budgets.admit(key_id("key")).await.unwrap();
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn plugin_budget_limits_revalidate_authority_and_rollback_with_audit() {
|
||||
let Some(database) = TestDatabase::create("plugin_budget_limits_rollback").await else {
|
||||
return;
|
||||
};
|
||||
let owner = plugin_reset_owner(&database).await;
|
||||
seed(&database, "key", "10", "20").await;
|
||||
let store = PgAdminClientKeyStore::new(database.pool.clone());
|
||||
let before = status(&database, "key").await;
|
||||
// 读取和配置赋值都不能开启未使用 Key 的窗口。
|
||||
assert_eq!(before.daily_resets_at, None);
|
||||
let mut stale_owner = owner.clone();
|
||||
stale_owner.revision = Revision::new(owner.revision.get() + 1).unwrap();
|
||||
assert_eq!(
|
||||
store
|
||||
.update_client_key_budget_limits(
|
||||
budget_update("key", Some("2"), None),
|
||||
ClientKeyBudgetMutationOrigin::Plugin(stale_owner),
|
||||
&context()
|
||||
)
|
||||
.await
|
||||
.unwrap_err()
|
||||
.kind(),
|
||||
AdminStoreErrorKind::Conflict
|
||||
);
|
||||
assert_eq!(
|
||||
store
|
||||
.update_client_key_budget_limits(
|
||||
budget_update("missing", Some("2"), None),
|
||||
ClientKeyBudgetMutationOrigin::Plugin(owner.clone()),
|
||||
&context()
|
||||
)
|
||||
.await
|
||||
.unwrap_err()
|
||||
.kind(),
|
||||
AdminStoreErrorKind::NotFound
|
||||
);
|
||||
let revision: i64 =
|
||||
sqlx::query_scalar("select config_revision from runtime_settings where id=1")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::raw_sql("create function reject_limit_audit() returns trigger language plpgsql as $$ begin raise exception 'test rollback'; end $$; create trigger reject_limit_audit before insert on admin_audit_events for each row execute function reject_limit_audit()")
|
||||
.execute(&database.pool).await.unwrap();
|
||||
assert!(
|
||||
store
|
||||
.update_client_key_budget_limits(
|
||||
budget_update("key", Some("2"), None),
|
||||
ClientKeyBudgetMutationOrigin::Plugin(owner.clone()),
|
||||
&context()
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert_eq!(status(&database, "key").await, before);
|
||||
let after_revision: i64 =
|
||||
sqlx::query_scalar("select config_revision from runtime_settings where id=1")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(revision, after_revision);
|
||||
sqlx::query("drop trigger reject_limit_audit on admin_audit_events")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query("update client_api_keys set enabled=false where id='key'")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.update_client_key_budget_limits(
|
||||
budget_update("key", Some("2"), None),
|
||||
ClientKeyBudgetMutationOrigin::Plugin(owner),
|
||||
&context(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let after = status(&database, "key").await;
|
||||
assert_eq!(after.limits.daily_usd.canonical(), "2");
|
||||
assert_eq!(after.daily_resets_at, None);
|
||||
assert_eq!(after.weekly_resets_at, None);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
@@ -14,7 +14,7 @@ use gateway_store::postgres::PgPluginStore;
|
||||
|
||||
use super::super::TestDatabase;
|
||||
|
||||
pub(super) fn artifact(digest: char, platforms: &[&str]) -> InspectedPluginArtifact {
|
||||
pub(in crate::postgres) fn artifact(digest: char, platforms: &[&str]) -> InspectedPluginArtifact {
|
||||
InspectedPluginArtifact {
|
||||
metadata: PluginArtifactMetadata {
|
||||
plugin_id: "test.example".into(),
|
||||
@@ -54,7 +54,7 @@ pub(super) fn context() -> MutationContext {
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn initialize_revision(database: &TestDatabase) {
|
||||
pub(in crate::postgres) async fn initialize_revision(database: &TestDatabase) {
|
||||
sqlx::query(
|
||||
"insert into runtime_settings(id, config_revision, updated_at) values (1,1,now()) on conflict do nothing",
|
||||
)
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
mod artifacts;
|
||||
pub(super) mod artifacts;
|
||||
mod credentials;
|
||||
mod instances;
|
||||
mod sources;
|
||||
|
||||
@@ -1,5 +1,74 @@
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_capacity_should_resolve_limits_for_list_and_detail() {
|
||||
let Some(database) = TestDatabase::create("account_capacity").await else {
|
||||
return;
|
||||
};
|
||||
let repository = PgProviderAccountRepository::new(database.pool.clone());
|
||||
let store = admin_account_store(&database.pool);
|
||||
let mut limited = account("acct_override", "user-override");
|
||||
limited.concurrency_limit = gateway_core::account::AccountConcurrencyLimit::new(3);
|
||||
for record in [account("acct_inherited", "user-inherited"), limited] {
|
||||
repository.insert_provider_account(record).await.unwrap();
|
||||
}
|
||||
for default_limit in [10_i64, 0] {
|
||||
sqlx::query("update runtime_settings set max_concurrent_per_account = $1 where id = 1")
|
||||
.bind(default_limit)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
for counts in [
|
||||
None,
|
||||
Some(BTreeMap::from([("acct_override".to_owned(), 2)])),
|
||||
] {
|
||||
let runtime = AccountRuntimeSnapshot {
|
||||
in_flight: counts,
|
||||
..Default::default()
|
||||
};
|
||||
let page = store
|
||||
.list_accounts(
|
||||
AccountListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(20).unwrap(),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: None,
|
||||
},
|
||||
runtime.clone(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(page.items.len(), 2);
|
||||
for item in page.items {
|
||||
let inherited = item.account.id == "acct_inherited";
|
||||
let expected_limit = if inherited {
|
||||
(default_limit > 0).then_some(10)
|
||||
} else {
|
||||
Some(3)
|
||||
};
|
||||
assert_eq!(item.capacity.total_slots, expected_limit);
|
||||
assert_eq!(
|
||||
item.capacity.used_slots,
|
||||
runtime
|
||||
.in_flight
|
||||
.as_ref()
|
||||
.map(|_| if inherited { 0 } else { 2 })
|
||||
);
|
||||
let detail = store
|
||||
.load_account(&item.account.id, runtime.clone())
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(detail.capacity, item.capacity);
|
||||
}
|
||||
}
|
||||
}
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_filter_applies_before_pagination_even_on_an_empty_page() {
|
||||
let Some(database) = TestDatabase::create("account_provider_filters").await else {
|
||||
|
||||
@@ -1791,23 +1791,63 @@ const fn oauth_error_code(error: &CodexOAuthAdminError) -> &'static str {
|
||||
fn map_quota_error(error: CodexCredentialQuotaError) -> ProviderAdminError {
|
||||
use CodexCredentialQuotaError as Error;
|
||||
use ProviderAdminErrorKind as Kind;
|
||||
let (kind, message) = match error {
|
||||
Error::InvalidCredentialData => (Kind::Invalid, "OpenAI 额度查询凭据无效,请检查账号授权"),
|
||||
Error::NotFound => (Kind::NotFound, "OpenAI 额度查询账号不存在"),
|
||||
Error::RevisionConflict => (Kind::Conflict, "账号凭据已被更新,请刷新账号列表后重试"),
|
||||
let (kind, public_message, upstream_message) = match error {
|
||||
Error::InvalidCredentialData => (
|
||||
Kind::Invalid,
|
||||
"OpenAI 额度查询凭据无效,请检查账号授权",
|
||||
None,
|
||||
),
|
||||
Error::NotFound => (Kind::NotFound, "OpenAI 额度查询账号不存在", None),
|
||||
Error::RevisionConflict => (
|
||||
Kind::Conflict,
|
||||
"账号凭据已被更新,请刷新账号列表后重试",
|
||||
None,
|
||||
),
|
||||
Error::CredentialRefreshRequired => (
|
||||
Kind::Unavailable,
|
||||
"OpenAI 额度查询需要有效凭据,请刷新令牌后重试",
|
||||
None,
|
||||
),
|
||||
Error::Repository(_) | Error::Store { .. } => {
|
||||
(Kind::Unavailable, "OpenAI 额度查询的依赖服务暂不可用")
|
||||
(Kind::Unavailable, "OpenAI 额度查询的依赖服务暂不可用", None)
|
||||
}
|
||||
Error::Upstream {
|
||||
status,
|
||||
code,
|
||||
detail,
|
||||
} => {
|
||||
// 公开文案只解释已知状态和错误码;原始上游材料留在内部诊断,
|
||||
// 额度查询拒绝不作为凭据失效证据,也不写入账号的凭据错误字段。
|
||||
let (kind, public_message) = match (status, code.as_deref()) {
|
||||
(Some(401), Some("token_revoked")) => (
|
||||
Kind::BadGateway,
|
||||
"OpenAI 拒绝了额度查询(HTTP 401,token_revoked):访问令牌已被撤销,请刷新令牌或重新授权",
|
||||
),
|
||||
(Some(401), _) => (
|
||||
Kind::BadGateway,
|
||||
"OpenAI 拒绝了额度查询(HTTP 401),请检查账号授权状态;若令牌已在服务端失效,请重新授权",
|
||||
),
|
||||
(Some(403), _) => (
|
||||
Kind::BadGateway,
|
||||
"OpenAI 拒绝了额度查询(HTTP 403),请检查账号授权状态",
|
||||
),
|
||||
(Some(429), _) => (Kind::BadGateway, "OpenAI 额度查询被限流,请稍后重试"),
|
||||
(Some(500..=599), _) => (Kind::BadGateway, "OpenAI 额度查询服务异常,请稍后重试"),
|
||||
_ => (
|
||||
Kind::Unavailable,
|
||||
"OpenAI 额度查询失败,请检查出站连接与上游服务",
|
||||
),
|
||||
};
|
||||
let upstream_message =
|
||||
format!("OpenAI 额度查询失败:上游 HTTP {status:?},code = {code:?}:{detail}");
|
||||
(kind, public_message, Some(upstream_message))
|
||||
}
|
||||
Error::Upstream { .. } => (
|
||||
Kind::Unavailable,
|
||||
"OpenAI 额度查询失败,请检查出站连接与上游服务",
|
||||
),
|
||||
};
|
||||
provider_admin_error(kind).with_public_message(message)
|
||||
let error = provider_admin_error(kind).with_public_message(public_message);
|
||||
match upstream_message {
|
||||
Some(message) => error.with_message(message),
|
||||
None => error,
|
||||
}
|
||||
}
|
||||
|
||||
fn map_profile_statistics_error(error: CodexProfileStatisticsError) -> ProviderAdminError {
|
||||
|
||||
@@ -202,7 +202,13 @@ pub enum CodexCredentialQuotaError {
|
||||
#[error("Codex quota credential revision is stale")]
|
||||
RevisionConflict,
|
||||
#[error("Codex quota upstream query failed: {detail}")]
|
||||
Upstream { detail: String },
|
||||
Upstream {
|
||||
detail: String,
|
||||
/// 上游 HTTP 状态码;传输失败等无响应场景为 `None`。
|
||||
status: Option<u16>,
|
||||
/// 有界错误码,供内部诊断和已知码映射使用,不直接拼入公开提示。
|
||||
code: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
/// 主动额度重置卡查询/消费失败。
|
||||
@@ -261,6 +267,38 @@ impl From<gateway_core::error::StoreError> for CodexCredentialQuotaError {
|
||||
}
|
||||
}
|
||||
|
||||
/// 客户端错误携带的 HTTP 状态码;传输失败等为 `None`。
|
||||
fn upstream_error_status(error: &CodexClientError) -> Option<u16> {
|
||||
match error {
|
||||
CodexClientError::Upstream { status, .. } => Some(status.as_u16()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 从上游错误体提取稳定错误码:优先 `/error/code`,其次 `/code`。
|
||||
///
|
||||
/// 只接受有界的 ASCII 标识,避免把自由文本作为错误码写入诊断。
|
||||
fn upstream_error_code(error: &CodexClientError) -> Option<String> {
|
||||
let CodexClientError::Upstream { body, .. } = error else {
|
||||
return None;
|
||||
};
|
||||
let value = serde_json::from_str::<Value>(body).ok()?;
|
||||
let code = value
|
||||
.pointer("/error/code")
|
||||
.or_else(|| value.pointer("/code"))
|
||||
.and_then(Value::as_str)?
|
||||
.trim();
|
||||
if code.is_empty()
|
||||
|| code.len() > 64
|
||||
|| !code
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'.'))
|
||||
{
|
||||
return None;
|
||||
}
|
||||
Some(code.to_owned())
|
||||
}
|
||||
|
||||
pub struct CodexCredentialQuotaService {
|
||||
repository: CodexCredentialRepository,
|
||||
store: Arc<dyn ProviderAccountStore>,
|
||||
@@ -969,6 +1007,8 @@ impl CodexCredentialQuotaService {
|
||||
tracing::warn!(
|
||||
account_id = %account.id(),
|
||||
error = %error,
|
||||
upstream_status = upstream_error_status(&error),
|
||||
upstream_code = ?upstream_error_code(&error),
|
||||
"OpenAI quota upstream rejection; refresh cycle will retry later"
|
||||
);
|
||||
}
|
||||
@@ -1499,6 +1539,8 @@ impl CodexCredentialQuotaService {
|
||||
}
|
||||
return Err(CodexCredentialQuotaError::Upstream {
|
||||
detail: error.to_string(),
|
||||
status: upstream_error_status(&error),
|
||||
code: upstream_error_code(&error),
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
@@ -609,7 +609,18 @@ impl CodexCredentialSelector {
|
||||
excluded_accounts: base_excluded.clone(),
|
||||
..context.clone()
|
||||
};
|
||||
let wait_candidates = AccountSelector.wait_candidates(&candidates, &wait_context);
|
||||
let mut wait_candidates =
|
||||
AccountSelector.wait_candidates(&candidates, &wait_context);
|
||||
let preferred_wait = preferred.as_ref().filter(|account_id| {
|
||||
!diagnostic
|
||||
&& queue_policy.max_waiting > 0
|
||||
&& (pinned_account.is_some() || !affinity.inherited)
|
||||
&& wait_candidates.contains(account_id)
|
||||
});
|
||||
if let Some(preferred) = preferred_wait {
|
||||
// 原绑定仍合格时只等待它;并发、请求间隔与已有队列均不应造成临时换号。
|
||||
wait_candidates.retain(|account_id| account_id == preferred);
|
||||
}
|
||||
let capacity = AccountSelector.capacity_snapshot(&candidates, &context);
|
||||
for candidate in &candidates {
|
||||
if !waiting.can_try(candidate.account.id()) {
|
||||
@@ -632,6 +643,19 @@ impl CodexCredentialSelector {
|
||||
return Err(CredentialSelectionError::PolicyUnavailable);
|
||||
}
|
||||
};
|
||||
// 可等待集合使用原始排除集;这里的 Excluded 只可能来自本轮租约争用或队列让位。
|
||||
let selection = selection.filter(|selection| {
|
||||
preferred_wait.is_none()
|
||||
|| selection.is_policy_choice()
|
||||
|| !matches!(
|
||||
selection.preferred(),
|
||||
PreferredAccountSelection::Blocked(
|
||||
AccountSchedulingBlocker::ConcurrencyLimit
|
||||
| AccountSchedulingBlocker::RequestInterval
|
||||
| AccountSchedulingBlocker::Excluded
|
||||
)
|
||||
)
|
||||
});
|
||||
request.attempt.trace().account_selection(
|
||||
&candidates,
|
||||
&context,
|
||||
@@ -673,12 +697,14 @@ impl CodexCredentialSelector {
|
||||
let quota_exhausted = !diagnostic
|
||||
&& statuses.next() == Some(AccountStatus::QuotaExhausted)
|
||||
&& statuses.all(|status| status == AccountStatus::QuotaExhausted);
|
||||
return match shortest_retry {
|
||||
Some(retry_after) => Err(CredentialSelectionError::CapacityUnavailable {
|
||||
retry_after: Some(retry_after),
|
||||
}),
|
||||
None if quota_exhausted => Err(CredentialSelectionError::QuotaExhausted),
|
||||
None => Err(CredentialSelectionError::NoEligibleCredential),
|
||||
return if !wait_candidates.is_empty() || shortest_retry.is_some() {
|
||||
Err(CredentialSelectionError::CapacityUnavailable {
|
||||
retry_after: shortest_retry,
|
||||
})
|
||||
} else if quota_exhausted {
|
||||
Err(CredentialSelectionError::QuotaExhausted)
|
||||
} else {
|
||||
Err(CredentialSelectionError::NoEligibleCredential)
|
||||
};
|
||||
};
|
||||
affinity.observe_preferred_selection(selection.preferred());
|
||||
@@ -714,7 +740,9 @@ impl CodexCredentialSelector {
|
||||
.await?
|
||||
{
|
||||
ProviderLeaseAcquisition::Busy { retry_after } => {
|
||||
affinity.observe_lease_busy(account.id());
|
||||
if preferred_wait != Some(account.id()) {
|
||||
affinity.observe_lease_busy(account.id());
|
||||
}
|
||||
shortest_retry = minimum_duration(shortest_retry, retry_after);
|
||||
excluded.insert(account.id().clone());
|
||||
}
|
||||
|
||||
@@ -837,7 +837,7 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
.with_requested_service_tier(request.service_tier())
|
||||
.with_request_tool_pricing(upstream_model.as_str(), request.tools())
|
||||
.with_raw_sse_passthrough();
|
||||
let mut pre_commit_events = PreCommitClientEvents::new();
|
||||
let mut pre_commit_events = PreCommitClientEvents::new(trace);
|
||||
loop {
|
||||
let Some(stream_deadline) = remaining(context.deadline()) else {
|
||||
if allows_account_state_mutation {
|
||||
@@ -882,7 +882,7 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
let next = match next {
|
||||
Ok(PreCommitPoll::Upstream(next)) => next,
|
||||
Ok(PreCommitPoll::GraceElapsed) => {
|
||||
for event in pre_commit_events.commit_pending() {
|
||||
for event in pre_commit_events.commit_pending(PreCommitReleaseReason::GraceTimeout) {
|
||||
yield event;
|
||||
}
|
||||
continue;
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
//! OpenAI 上游失败分类、恢复决策与稳定错误投影。
|
||||
|
||||
use super::*;
|
||||
use gateway_core::diagnostics::TraceContext;
|
||||
|
||||
/// OpenAI 失败对 Smart 账号分数的结构化 reason 闭集。
|
||||
///
|
||||
@@ -83,23 +84,35 @@ pub(super) async fn wait_for_replay_grace(deadline: Option<Instant>) {
|
||||
|
||||
/// 提交边界前的上游事件预取。
|
||||
///
|
||||
/// 原始 chunk 计数而不是重编码后的 event 大小。时间与 64 KiB 共同限定无感换号
|
||||
/// 原始 chunk 计数而不是重编码后的 event 大小。时间与字节阈值共同限定无感换号
|
||||
/// 窗口;任一边界到达都会提交已缓存 wire,不能因网关私有资源规则伪造上游协议
|
||||
/// 失败。一旦提交,后续事件不再具备无痕重放资格。
|
||||
pub(super) struct PreCommitClientEvents {
|
||||
pending: Vec<ProviderEvent>,
|
||||
prefetched_bytes: usize,
|
||||
replay_grace_deadline: Option<Instant>,
|
||||
replay_grace_started_at: Option<Instant>,
|
||||
committed: bool,
|
||||
trace: TraceContext,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub(super) enum PreCommitReleaseReason {
|
||||
SemanticOutput,
|
||||
Terminal,
|
||||
ByteLimit,
|
||||
GraceTimeout,
|
||||
Eof,
|
||||
}
|
||||
|
||||
impl PreCommitClientEvents {
|
||||
pub(super) const fn new() -> Self {
|
||||
pub(super) const fn new(trace: TraceContext) -> Self {
|
||||
Self {
|
||||
pending: Vec::new(),
|
||||
prefetched_bytes: 0,
|
||||
replay_grace_deadline: None,
|
||||
replay_grace_started_at: None,
|
||||
committed: false,
|
||||
trace,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -120,14 +133,17 @@ impl PreCommitClientEvents {
|
||||
}
|
||||
let starts_replay_grace = incoming.iter().any(ProviderEvent::has_client_event);
|
||||
self.pending.extend(incoming);
|
||||
if timing_signals.semantic_output
|
||||
|| completed
|
||||
|| self.prefetched_bytes > MAX_STREAM_PREFETCH_BYTES
|
||||
{
|
||||
return self.commit_pending();
|
||||
if timing_signals.semantic_output {
|
||||
return self.commit_pending(PreCommitReleaseReason::SemanticOutput);
|
||||
}
|
||||
if starts_replay_grace && self.replay_grace_deadline.is_none() {
|
||||
self.replay_grace_deadline = Instant::now().checked_add(STREAM_REPLAY_GRACE);
|
||||
if completed {
|
||||
return self.commit_pending(PreCommitReleaseReason::Terminal);
|
||||
}
|
||||
if self.prefetched_bytes > MAX_STREAM_PREFETCH_BYTES {
|
||||
return self.commit_pending(PreCommitReleaseReason::ByteLimit);
|
||||
}
|
||||
if starts_replay_grace && self.replay_grace_started_at.is_none() {
|
||||
self.replay_grace_started_at = Some(Instant::now());
|
||||
}
|
||||
Vec::new()
|
||||
}
|
||||
@@ -142,7 +158,7 @@ impl PreCommitClientEvents {
|
||||
if self.committed {
|
||||
events
|
||||
} else {
|
||||
self.commit_pending()
|
||||
self.commit_pending(PreCommitReleaseReason::Eof)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -151,28 +167,38 @@ impl PreCommitClientEvents {
|
||||
return incoming;
|
||||
}
|
||||
self.pending.extend(incoming);
|
||||
self.commit_pending()
|
||||
// 此入口只在同批次已出现语义输出、随后失败时释放原始事件。
|
||||
self.commit_pending(PreCommitReleaseReason::SemanticOutput)
|
||||
}
|
||||
|
||||
pub(super) fn take_for_failure(&mut self, incoming: Vec<ProviderEvent>) -> Vec<ProviderEvent> {
|
||||
self.pending.extend(incoming);
|
||||
self.prefetched_bytes = 0;
|
||||
self.replay_grace_deadline = None;
|
||||
self.replay_grace_started_at = None;
|
||||
std::mem::take(&mut self.pending)
|
||||
}
|
||||
|
||||
pub(super) const fn replay_grace_deadline(&self) -> Option<Instant> {
|
||||
self.replay_grace_deadline
|
||||
pub(super) fn replay_grace_deadline(&self) -> Option<Instant> {
|
||||
self.replay_grace_started_at
|
||||
.and_then(|started| started.checked_add(STREAM_REPLAY_GRACE))
|
||||
}
|
||||
|
||||
pub(super) const fn is_committed(&self) -> bool {
|
||||
self.committed
|
||||
}
|
||||
|
||||
pub(super) fn commit_pending(&mut self) -> Vec<ProviderEvent> {
|
||||
pub(super) fn commit_pending(&mut self, reason: PreCommitReleaseReason) -> Vec<ProviderEvent> {
|
||||
// Provider 只记录释放缓存的原因;实际下游提交仍由 Core 记录和裁决。
|
||||
if self.trace.is_enabled() {
|
||||
self.trace.record("provider.precommit.released", json!({
|
||||
"reason": reason,
|
||||
"prefetchedBytes": self.prefetched_bytes,
|
||||
"waitMs": self.replay_grace_started_at.map_or(0, |started| started.elapsed().as_millis()),
|
||||
}));
|
||||
}
|
||||
self.committed = true;
|
||||
self.prefetched_bytes = 0;
|
||||
self.replay_grace_deadline = None;
|
||||
self.replay_grace_started_at = None;
|
||||
std::mem::take(&mut self.pending)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -120,12 +120,13 @@ const WEBSOCKET_TRANSPORT: &str = "websocket";
|
||||
// 在已观测到的 Codex OAuth 上游 16 MiB 附近消息边界前留出传输 metadata 余量。
|
||||
const WEBSOCKET_HTTP_FALLBACK_THRESHOLD_BYTES: usize = 15 * 1024 * 1024;
|
||||
const MAX_COOKIE_HEADER_BYTES: usize = 16 * 1024;
|
||||
/// 提交边界前最多保留 64 KiB 原始上游 chunk;达到阈值后结束无感换号窗口,
|
||||
/// 提交边界前预取 128 KiB 原始上游 chunk;容纳携带配置回显的前导事件,
|
||||
/// 超过阈值后结束无感换号窗口(最后一个 chunk 可越过阈值),
|
||||
/// 但不会把上游数据改写成协议失败。
|
||||
const MAX_STREAM_PREFETCH_BYTES: usize = 64 * 1024;
|
||||
const MAX_STREAM_PREFETCH_BYTES: usize = 128 * 1024;
|
||||
/// 短暂保留 response.created 等结构事件,让随后到达的明确拒绝可以无感换号;
|
||||
/// 到期即放行,避免模型长时间思考时让客户端一直收不到首事件。
|
||||
const STREAM_REPLAY_GRACE: Duration = Duration::from_millis(1_200);
|
||||
const STREAM_REPLAY_GRACE: Duration = Duration::from_millis(2_500);
|
||||
// 额度拒绝后先给上游额度结算留出时间,再以受限时长同步 usage 快照。
|
||||
const QUOTA_FAILURE_REFRESH_DELAY: Duration = Duration::from_secs(2);
|
||||
const QUOTA_FAILURE_REFRESH_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
@@ -863,6 +864,7 @@ impl CodexProvider {
|
||||
})?
|
||||
.with_authentication(lease.authentication())
|
||||
.with_connection_budget(context.connection_budget().clone())
|
||||
.with_response_control(context.response_control().cloned())
|
||||
.with_middleware_headers(middleware_headers),
|
||||
response_origin: self.responses_url.clone(),
|
||||
request: upstream_request,
|
||||
|
||||
@@ -689,6 +689,7 @@ impl OpenAiUpstreamProtocol {
|
||||
/// Codex HTTP/SSE 上游客户端。
|
||||
#[derive(Clone)]
|
||||
pub struct CodexBackendClient {
|
||||
pub(super) response_control: Option<gateway_core::engine::response_control::ResponseControl>,
|
||||
pub(super) connection_budget: Option<gateway_core::engine::connection::ConnectionBudget>,
|
||||
pub(super) client: Client,
|
||||
pub(super) direct_client: Client,
|
||||
@@ -705,6 +706,14 @@ pub struct CodexBackendClient {
|
||||
}
|
||||
|
||||
impl CodexBackendClient {
|
||||
pub(crate) fn with_response_control(
|
||||
mut self,
|
||||
control: Option<gateway_core::engine::response_control::ResponseControl>,
|
||||
) -> Self {
|
||||
self.response_control = control;
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) fn with_connection_budget(
|
||||
mut self,
|
||||
budget: gateway_core::engine::connection::ConnectionBudget,
|
||||
|
||||
@@ -65,6 +65,7 @@ impl CodexBackendClient {
|
||||
let base_url = base_url.into().trim_end_matches('/').to_string();
|
||||
Self {
|
||||
connection_budget: None,
|
||||
response_control: None,
|
||||
direct_client: client.clone(),
|
||||
client,
|
||||
websocket_origin_key: websocket_origin_key(&base_url),
|
||||
@@ -453,6 +454,7 @@ impl CodexBackendClient {
|
||||
let mut exchange = execute_prepared_response_create_request_stream(
|
||||
&websocket_request,
|
||||
prepared,
|
||||
self.response_control.clone(),
|
||||
context
|
||||
.trace
|
||||
.cloned()
|
||||
|
||||
@@ -30,6 +30,9 @@ pub(super) fn is_managed_identity_header(name: &str) -> bool {
|
||||
| "chatgpt-project-id"
|
||||
| "openai-organization"
|
||||
| "openai-project"
|
||||
// 工作区路由和合规属性属于认证账号,不能继承下游账号的值。
|
||||
| "x-openai-account-routing-override"
|
||||
| "x-openai-fedramp"
|
||||
// 安装身份由当前账号写入 client_metadata,不继承下游安装头。
|
||||
| "x-codex-installation-id"
|
||||
)
|
||||
@@ -217,7 +220,7 @@ impl CodexBackendClient {
|
||||
let provider_headers = headers.keys().cloned().collect::<Vec<_>>();
|
||||
for header in &self.middleware_headers {
|
||||
let name = HeaderName::from_bytes(header.name().as_bytes())?;
|
||||
if provider_headers.contains(&name) {
|
||||
if is_managed_identity_header(name.as_str()) || provider_headers.contains(&name) {
|
||||
return Err(super::client::CodexClientError::MiddlewareHeaderConflict);
|
||||
}
|
||||
headers.append(name, HeaderValue::from_bytes(header.value())?);
|
||||
|
||||
@@ -138,11 +138,13 @@ fn json_value_as_string(value: &Value) -> Option<String> {
|
||||
}
|
||||
}
|
||||
|
||||
/// 从 `response.completed` 旁路提取 response ID,供连接池记录续接能力。
|
||||
/// 从正常完成或官方中断终态旁路提取 response ID,供连接池记录续接能力。
|
||||
///
|
||||
/// 该值不参与客户端 wire 的可交付性判断;无法读取时只是不记录连接内续接状态。
|
||||
pub fn websocket_response_completed_id(value: &Value) -> Option<String> {
|
||||
if value.get("type").and_then(Value::as_str) != Some("response.completed") {
|
||||
if value.get("type").and_then(Value::as_str) != Some("response.completed")
|
||||
&& !websocket_response_is_interrupted(value)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
value
|
||||
@@ -151,6 +153,14 @@ pub fn websocket_response_completed_id(value: &Value) -> Option<String> {
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
pub(crate) fn websocket_response_is_interrupted(value: &Value) -> bool {
|
||||
value.get("type").and_then(Value::as_str) == Some("response.incomplete")
|
||||
&& value
|
||||
.pointer("/response/incomplete_details/reason")
|
||||
.and_then(Value::as_str)
|
||||
== Some("interrupted")
|
||||
}
|
||||
|
||||
/// 生成 Responses WebSocket payload 审计快照。
|
||||
pub fn websocket_payload_audit_snapshot(request: &CodexResponsesRequest) -> PayloadAuditSnapshot {
|
||||
let body = websocket_response_create_payload(request);
|
||||
|
||||
@@ -536,6 +536,9 @@ pub(crate) fn scope_request_to_account(
|
||||
{
|
||||
metadata.remove(*key);
|
||||
}
|
||||
// Guardian 顶层父引用指向原账号的响应;turn metadata 内的同名
|
||||
// 扩展只是客户端关联信息,不能加入各层共用的清理名单。
|
||||
metadata.remove("parent_response_id");
|
||||
}
|
||||
// 官方 Core 在 client_metadata 使用带 x-codex 前缀的键,
|
||||
// turn metadata 内仍使用 installation_id;两处均取当前账号的安装身份。
|
||||
|
||||
@@ -581,6 +581,7 @@ fn continuation_unavailable(
|
||||
pub(crate) async fn execute_prepared_response_create_request_stream(
|
||||
request: &CodexWebSocketRequest,
|
||||
prepared: PreparedWebSocket,
|
||||
response_control: Option<gateway_core::engine::response_control::ResponseControl>,
|
||||
trace: gateway_core::diagnostics::TraceContext,
|
||||
) -> Result<CodexWebSocketStreamingExchange, CodexWebSocketExchangeError> {
|
||||
let PreparedWebSocket {
|
||||
@@ -638,6 +639,7 @@ pub(crate) async fn execute_prepared_response_create_request_stream(
|
||||
reused,
|
||||
stream_idle_timeout,
|
||||
trace,
|
||||
response_control,
|
||||
);
|
||||
exchange.pool_decision = pool_decision;
|
||||
exchange.connection_local_continuation = connection_local_available;
|
||||
|
||||
@@ -6,6 +6,7 @@ use serde_json::Value;
|
||||
use crate::transport::protocol::websocket::{
|
||||
websocket_event_frame, websocket_event_type, websocket_metadata_headers,
|
||||
websocket_metadata_turn_state, websocket_response_completed_id,
|
||||
websocket_response_is_interrupted,
|
||||
};
|
||||
use crate::transport::response_meta;
|
||||
|
||||
@@ -15,6 +16,7 @@ use super::CodexWebSocketExchangeError;
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(in crate::transport::websocket) enum WebSocketTerminalKind {
|
||||
Completed,
|
||||
Interrupted,
|
||||
Incomplete,
|
||||
Failed,
|
||||
}
|
||||
@@ -29,6 +31,7 @@ pub(super) enum ExchangeAction {
|
||||
}
|
||||
|
||||
pub(super) struct ReducedWebSocketEvent {
|
||||
pub(super) created_response_id: Option<String>,
|
||||
pub(super) action: ExchangeAction,
|
||||
pub(super) diagnostic_event_type: Option<String>,
|
||||
pub(super) turn_state_update: Option<String>,
|
||||
@@ -43,6 +46,7 @@ pub(super) fn reduce_websocket_event(
|
||||
// 不可解析的帧不承载可路由的事件类型,忽略。
|
||||
let Ok(value) = serde_json::from_str::<Value>(raw) else {
|
||||
return Ok(ReducedWebSocketEvent {
|
||||
created_response_id: None,
|
||||
action: ExchangeAction::Ignore,
|
||||
diagnostic_event_type: None,
|
||||
turn_state_update: None,
|
||||
@@ -54,6 +58,7 @@ pub(super) fn reduce_websocket_event(
|
||||
metadata.rate_limit_headers.extend(headers);
|
||||
return Ok(ReducedWebSocketEvent {
|
||||
action: ExchangeAction::RateLimits(parsed),
|
||||
created_response_id: None,
|
||||
diagnostic_event_type,
|
||||
turn_state_update: None,
|
||||
});
|
||||
@@ -75,14 +80,15 @@ pub(super) fn reduce_websocket_event(
|
||||
});
|
||||
|
||||
let event = websocket_event_type(&value);
|
||||
if event == Some("response.completed")
|
||||
&& let Some(response_id) = websocket_response_completed_id(&value)
|
||||
{
|
||||
if let Some(response_id) = websocket_response_completed_id(&value) {
|
||||
continuation.record_completed(response_id);
|
||||
}
|
||||
|
||||
let terminal = match event {
|
||||
Some("response.completed") => Some(WebSocketTerminalKind::Completed),
|
||||
Some("response.incomplete") if websocket_response_is_interrupted(&value) => {
|
||||
Some(WebSocketTerminalKind::Interrupted)
|
||||
}
|
||||
Some("response.incomplete") => Some(WebSocketTerminalKind::Incomplete),
|
||||
Some("response.failed" | "error") => Some(WebSocketTerminalKind::Failed),
|
||||
_ => None,
|
||||
@@ -92,6 +98,10 @@ pub(super) fn reduce_websocket_event(
|
||||
None => ExchangeAction::Ignore,
|
||||
};
|
||||
Ok(ReducedWebSocketEvent {
|
||||
created_response_id: (event == Some("response.created"))
|
||||
.then(|| value.pointer("/response/id").and_then(Value::as_str))
|
||||
.flatten()
|
||||
.map(ToOwned::to_owned),
|
||||
action,
|
||||
diagnostic_event_type,
|
||||
turn_state_update,
|
||||
|
||||
@@ -4,6 +4,7 @@ use std::{sync::Arc, time::Duration};
|
||||
|
||||
use bytes::Bytes;
|
||||
use gateway_core::diagnostics::{TraceContext, diagnostic_json};
|
||||
use gateway_core::engine::response_control::{ActiveResponseInterrupt, ResponseControl};
|
||||
use serde_json::json;
|
||||
use tokio::sync::{Mutex, mpsc};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
@@ -65,6 +66,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
|
||||
reused_connection: bool,
|
||||
stream_idle_timeout: Option<Duration>,
|
||||
trace: TraceContext,
|
||||
response_control: Option<ResponseControl>,
|
||||
) -> CodexWebSocketStreamingExchange {
|
||||
let websocket_connection_id = websocket.connection_id();
|
||||
let response_metadata = metadata.clone();
|
||||
@@ -85,6 +87,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
|
||||
);
|
||||
let forward = async move {
|
||||
forward_websocket_response_stream(WebSocketStreamForwardState {
|
||||
response_control,
|
||||
websocket,
|
||||
metadata,
|
||||
pool_return,
|
||||
@@ -124,6 +127,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
|
||||
}
|
||||
|
||||
struct WebSocketStreamForwardState {
|
||||
response_control: Option<ResponseControl>,
|
||||
trace: TraceContext,
|
||||
websocket: PumpedWebSocket,
|
||||
metadata: CodexWebSocketConnectionMetadata,
|
||||
@@ -138,6 +142,7 @@ struct WebSocketStreamForwardState {
|
||||
|
||||
async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
let WebSocketStreamForwardState {
|
||||
response_control,
|
||||
trace,
|
||||
mut websocket,
|
||||
mut metadata,
|
||||
@@ -155,6 +160,10 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
.map(|pool_return| std::mem::take(&mut pool_return.continuation))
|
||||
.unwrap_or_default();
|
||||
let mut last_event_type = None;
|
||||
let mut active_interrupt: Option<ActiveResponseInterrupt> = None;
|
||||
let mut interrupt_requested: futures::future::BoxFuture<'static, ()> =
|
||||
Box::pin(std::future::pending());
|
||||
let mut interrupt_sent = false;
|
||||
// 官方 run_websocket_response_stream 按响应消费 metadata 事件;它们不属于握手头。
|
||||
// 复用连接时恢复握手快照,避免上一轮模型信息及普通响应头随每次请求累积。
|
||||
let opening_response_metadata = metadata.response_metadata.clone();
|
||||
@@ -164,6 +173,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
// 客户端断开下游 SSE 流:立即丢弃连接并释放池 slot,
|
||||
// 不再傻等上游 idle 超时(否则同会话后续请求会一直 bypass/busy)。
|
||||
() = tx.closed() => {
|
||||
drop(active_interrupt.take());
|
||||
trace.record("upstream.cancelled", json!({"reason": "receiver_dropped"}));
|
||||
discard_stream_websocket(
|
||||
websocket,
|
||||
@@ -173,6 +183,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
return;
|
||||
}
|
||||
() = shutdown.cancelled() => {
|
||||
drop(active_interrupt.take());
|
||||
trace.record("upstream.cancelled", json!({"reason": "pool_shutdown"}));
|
||||
discard_stream_websocket(
|
||||
websocket,
|
||||
@@ -181,6 +192,24 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
).await;
|
||||
return;
|
||||
}
|
||||
() = &mut interrupt_requested, if !interrupt_sent => {
|
||||
let Some(active) = active_interrupt.as_ref() else { continue };
|
||||
let payload = json!({
|
||||
"type": "response.interrupt",
|
||||
"response_id": active.response_id(),
|
||||
"mode": "discard_partial_items",
|
||||
}).to_string();
|
||||
// 复用正在消费此响应的原 socket;信号不经过选号或新请求恢复。
|
||||
if let Err(error) = super::super::handshake::send_websocket_request(&websocket, &payload).await {
|
||||
drop(active_interrupt.take());
|
||||
discard_stream_websocket(websocket, pool_return, StreamWebSocketDiscardReason::UpstreamReceiveFailed).await;
|
||||
let _ = tx.send(Err(error)).await;
|
||||
return;
|
||||
}
|
||||
interrupt_sent = true;
|
||||
trace.record("upstream.interrupt.sent", json!({"mode": "discard_partial_items"}));
|
||||
continue;
|
||||
}
|
||||
message = next_websocket_message(
|
||||
&mut websocket,
|
||||
stream_idle_timeout.filter(|timeout| !timeout.is_zero()).unwrap_or(super::super::pool::DEFAULT_STREAM_IDLE_TIMEOUT),
|
||||
@@ -189,6 +218,8 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
let message = match message {
|
||||
Ok(message) => message,
|
||||
Err(error) => {
|
||||
// 先撤销 owner 再发布失败;下一次 attempt 不等待旧 socket 异步清理。
|
||||
drop(active_interrupt.take());
|
||||
let observation = connection_observation(&websocket, &error);
|
||||
trace.record(
|
||||
"upstream.read.failed",
|
||||
@@ -224,6 +255,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
let raw = match message {
|
||||
tungstenite::Message::Text(text) => text.to_string(),
|
||||
tungstenite::Message::Binary(bytes) => {
|
||||
drop(active_interrupt.take());
|
||||
trace.capture("upstream.binary", &bytes);
|
||||
let error = CodexWebSocketExchangeError::UnexpectedBinaryEvent;
|
||||
let observation = connection_observation(&websocket, &error);
|
||||
@@ -239,6 +271,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
return;
|
||||
}
|
||||
tungstenite::Message::Close(frame) => {
|
||||
drop(active_interrupt.take());
|
||||
if let Some(frame) = &frame {
|
||||
trace.dump("upstream.close.reason", frame.reason.as_bytes());
|
||||
}
|
||||
@@ -277,6 +310,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
let reduced = match reduce_websocket_event(&raw, &mut metadata, &mut continuation) {
|
||||
Ok(reduced) => reduced,
|
||||
Err(error) => {
|
||||
drop(active_interrupt.take());
|
||||
let observation =
|
||||
(!matches!(error.classified(), CodexWebSocketExchangeError::Upstream(_)))
|
||||
.then(|| connection_observation(&websocket, &error));
|
||||
@@ -297,6 +331,31 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
if let Some(event_type) = reduced.diagnostic_event_type {
|
||||
last_event_type = Some(event_type);
|
||||
}
|
||||
if active_interrupt.is_none()
|
||||
&& let (Some(control), Some(response_id)) =
|
||||
(&response_control, reduced.created_response_id)
|
||||
{
|
||||
active_interrupt = control.activate(response_id);
|
||||
if let Some(active) = &active_interrupt {
|
||||
// 跨消息复用同一个等待 future,避免每帧留下新的取消等待者。
|
||||
interrupt_requested = Box::pin(active.requested());
|
||||
}
|
||||
if active_interrupt.is_none() {
|
||||
discard_stream_websocket(
|
||||
websocket,
|
||||
pool_return,
|
||||
StreamWebSocketDiscardReason::FailedResponse,
|
||||
)
|
||||
.await;
|
||||
let _ = tx
|
||||
.send(Err(CodexWebSocketExchangeError::PostSendAmbiguous {
|
||||
message: "response control already has an active owner".to_owned(),
|
||||
source: None,
|
||||
}))
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
if let Some(turn_state) = reduced.turn_state_update {
|
||||
let mut pending = response_metadata_updates.lock().await;
|
||||
if pending.turn_state.is_none() {
|
||||
@@ -314,7 +373,12 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
ExchangeAction::Forward { frame, terminal } => (frame, terminal),
|
||||
ExchangeAction::Ignore => continue,
|
||||
};
|
||||
if terminal.is_some() {
|
||||
// 先撤销控制权,再交付终态和归还连接,避免迟到的中断影响下一轮。
|
||||
drop(active_interrupt.take());
|
||||
}
|
||||
if tx.send(Ok(Bytes::from(frame))).await.is_err() {
|
||||
drop(active_interrupt.take());
|
||||
trace.record(
|
||||
"upstream.forward.failed",
|
||||
json!({"reason": "receiver_dropped"}),
|
||||
@@ -333,7 +397,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
json!({"kind": format!("{terminal:?}")}),
|
||||
);
|
||||
match terminal {
|
||||
WebSocketTerminalKind::Completed => {
|
||||
WebSocketTerminalKind::Completed | WebSocketTerminalKind::Interrupted => {
|
||||
metadata.response_metadata = opening_response_metadata;
|
||||
finish_stream_websocket(websocket, metadata, continuation, pool_return.take())
|
||||
.await;
|
||||
@@ -392,6 +456,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
last_event_type,
|
||||
},
|
||||
};
|
||||
drop(active_interrupt.take());
|
||||
let observation = connection_observation(&websocket, &error);
|
||||
discard_stream_websocket(
|
||||
websocket,
|
||||
|
||||
@@ -1976,6 +1976,116 @@ mod errors {
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_refresh_exposes_safe_failure_reasons_and_recovers_without_account_errors() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let account_id = "acct_quota_error";
|
||||
store
|
||||
.seed_oauth_credential(ImportCodexOAuthCredential {
|
||||
account_id: account_id.to_owned(),
|
||||
name: "quota error".to_owned(),
|
||||
secret: secret("quota-error-test-token"),
|
||||
verified_account: profile("chatgpt-quota-error"),
|
||||
next_refresh_at: None,
|
||||
enabled: true,
|
||||
})
|
||||
.await;
|
||||
let before = store.account(account_id).expect("account");
|
||||
let server = MockServer::start().await;
|
||||
let mut config = valid_config();
|
||||
config.config.api.base_url = server.uri();
|
||||
let bundle = provider_openai::initialize(
|
||||
config.config,
|
||||
provider_ports_with(store.clone(), Arc::new(TestOAuthPending::default())),
|
||||
)
|
||||
.await
|
||||
.expect("OpenAI bundle");
|
||||
for (status, code, expected_kind, expected_message) in [
|
||||
(
|
||||
401,
|
||||
"token_revoked",
|
||||
Kind::BadGateway,
|
||||
"OpenAI 拒绝了额度查询(HTTP 401,token_revoked):访问令牌已被撤销,请刷新令牌或重新授权",
|
||||
),
|
||||
(
|
||||
401,
|
||||
"unknown_code",
|
||||
Kind::BadGateway,
|
||||
"OpenAI 拒绝了额度查询(HTTP 401),请检查账号授权状态;若令牌已在服务端失效,请重新授权",
|
||||
),
|
||||
(
|
||||
403,
|
||||
"forbidden",
|
||||
Kind::BadGateway,
|
||||
"OpenAI 拒绝了额度查询(HTTP 403),请检查账号授权状态",
|
||||
),
|
||||
(
|
||||
429,
|
||||
"rate_limited",
|
||||
Kind::BadGateway,
|
||||
"OpenAI 额度查询被限流,请稍后重试",
|
||||
),
|
||||
(
|
||||
503,
|
||||
"unavailable",
|
||||
Kind::BadGateway,
|
||||
"OpenAI 额度查询服务异常,请稍后重试",
|
||||
),
|
||||
(
|
||||
400,
|
||||
"invalid_request",
|
||||
Kind::Unavailable,
|
||||
"OpenAI 额度查询失败,请检查出站连接与上游服务",
|
||||
),
|
||||
] {
|
||||
server.reset().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/codex/usage"))
|
||||
.respond_with(ResponseTemplate::new(status).set_body_json(json!({
|
||||
"error": {"code": code, "message": "raw-secret-marker"}
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let error = bundle
|
||||
.admin_provider()
|
||||
.quota(ProviderQuotaRequest {
|
||||
account_id: before.id().clone(),
|
||||
refresh: true,
|
||||
rolling_usage: None,
|
||||
})
|
||||
.await
|
||||
.expect_err("quota rejection");
|
||||
assert_eq!(error.kind(), expected_kind);
|
||||
assert_eq!(error.public_message(), Some(expected_message));
|
||||
assert!(!format!("{error:?} {error}").contains("raw-secret-marker"));
|
||||
assert_eq!(
|
||||
store.account(account_id).expect("account after rejection"),
|
||||
before
|
||||
);
|
||||
}
|
||||
server.reset().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/codex/usage"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"rate_limit": {"allowed": true, "primary_window": {"used_percent": 8}}
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let quota = bundle
|
||||
.admin_provider()
|
||||
.quota(ProviderQuotaRequest {
|
||||
account_id: before.id().clone(),
|
||||
refresh: true,
|
||||
rolling_usage: None,
|
||||
})
|
||||
.await
|
||||
.expect("quota recovered");
|
||||
assert_eq!(quota.representative_used_percent(), Some(8.0));
|
||||
let current = store.account(account_id).expect("account after recovery");
|
||||
assert_eq!(current.credential_state(), CredentialState::Ready);
|
||||
assert!(current.last_error_message().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn manual_refresh_preserves_banned_evidence_without_promoting_401_to_terminal() {
|
||||
for (status, kind, message) in [
|
||||
|
||||
@@ -0,0 +1,383 @@
|
||||
use futures::{FutureExt, future::BoxFuture};
|
||||
use gateway_core::account::AccountRuntimeSignals;
|
||||
use gateway_core::concurrency::{ConcurrencyQueuePolicy, QueueRejection};
|
||||
use gateway_core::engine::policy::{
|
||||
AccountScheduleDecision, AccountScheduleInput, ModelRouteDecision, ModelRouteInput,
|
||||
RequestPolicyContext, RequestPolicyFault, RequestPolicyPlan,
|
||||
};
|
||||
use gateway_core::runtime::extensions::{ExtensionSetId, ExtensionSetLease, ExtensionSetReference};
|
||||
|
||||
use super::*;
|
||||
|
||||
fn busy_signal(in_flight: u32, last_started_at: Option<SystemTime>) -> AccountRuntimeSignals {
|
||||
AccountRuntimeSignals {
|
||||
in_flight,
|
||||
last_started_at,
|
||||
quota_reset_at: None,
|
||||
quota_remaining_rank: None,
|
||||
cooldown: None,
|
||||
failure_rate_basis_points: None,
|
||||
first_output_latency_ms: None,
|
||||
}
|
||||
}
|
||||
|
||||
async fn bind_first(affinity: &MemorySessionAffinity) -> ProviderSessionAffinityKey {
|
||||
let key = ProviderSessionAffinityKey::try_new("capacity-session").unwrap();
|
||||
affinity
|
||||
.bind(
|
||||
&ProviderKind::new("openai").unwrap(),
|
||||
&key,
|
||||
&ProviderAccountId::new("acct_first").unwrap(),
|
||||
Duration::from_secs(60),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
key
|
||||
}
|
||||
|
||||
fn capacity_attempt(
|
||||
policy: AccountSelectionPolicy,
|
||||
request_policy: Option<RequestPolicyContext>,
|
||||
) -> AttemptContext {
|
||||
AttemptContext::new(
|
||||
RequestAttemptContext::new(
|
||||
ModelRequestId::new("req_capacity_contract").unwrap(),
|
||||
ClientApiKeyId::new("key_codex_contract").unwrap(),
|
||||
)
|
||||
.with_request_policy(request_policy),
|
||||
NonZeroU32::new(1).unwrap(),
|
||||
SystemTime::now() + Duration::from_secs(5),
|
||||
policy,
|
||||
AccountAttemptContext::new(BTreeSet::new(), None, None)
|
||||
.with_account_scope(contract_account_scope()),
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn saturated_snapshot_is_capacity_unavailable_even_with_exhausted_accounts() {
|
||||
for interval in [false, true] {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_first", "test-first");
|
||||
create_account(&store, "acct_second", "test-second");
|
||||
let exhausted =
|
||||
block_on(store.get_account(&ProviderAccountId::new("acct_second").unwrap()))
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
persist_quota_exhaustion(&store, &exhausted, None);
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
leases.signals.lock().unwrap().insert(
|
||||
ProviderAccountId::new("acct_first").unwrap(),
|
||||
busy_signal(
|
||||
if interval { 0 } else { u32::MAX },
|
||||
interval.then(SystemTime::now),
|
||||
),
|
||||
);
|
||||
let selector = selector(&store, leases.clone());
|
||||
let request_attempt = capacity_attempt(
|
||||
AccountSelectionPolicy::new(
|
||||
RotationStrategy::Smart,
|
||||
NonZeroU32::new(2).unwrap(),
|
||||
Duration::from_secs(2),
|
||||
),
|
||||
None,
|
||||
);
|
||||
let url = Url::parse(OFFICIAL_CODEX_BASE_URL).unwrap();
|
||||
let error = block_on(selector.select(&SelectCodexCredential {
|
||||
upstream_model: "gpt-5.4",
|
||||
request_url: &url,
|
||||
attempt: &request_attempt,
|
||||
session_affinity_key: None,
|
||||
}))
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
matches!(error, CredentialSelectionError::CapacityUnavailable { .. }),
|
||||
"{error:?}"
|
||||
);
|
||||
assert!(leases.requests.lock().unwrap().is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn affinity_queue_waits_for_the_original_account_for_local_capacity_blockers() {
|
||||
for blocker in ["lease", "concurrency", "interval"] {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_first", "test-first");
|
||||
create_account(&store, "acct_second", "test-second");
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
let first = ProviderAccountId::new("acct_first").unwrap();
|
||||
if blocker == "lease" {
|
||||
leases.busy_accounts.lock().unwrap().insert(first.clone());
|
||||
} else {
|
||||
leases.signals.lock().unwrap().insert(
|
||||
first.clone(),
|
||||
busy_signal(
|
||||
if blocker == "concurrency" { 2 } else { 0 },
|
||||
(blocker == "interval").then(SystemTime::now),
|
||||
),
|
||||
);
|
||||
}
|
||||
let affinity = Arc::new(MemorySessionAffinity::default());
|
||||
let key = bind_first(&affinity).await;
|
||||
let selector = selector_with_affinity(&store, leases.clone(), affinity.clone());
|
||||
let request_attempt = capacity_attempt(
|
||||
AccountSelectionPolicy::new(
|
||||
RotationStrategy::Smart,
|
||||
NonZeroU32::new(2).unwrap(),
|
||||
Duration::from_secs(2),
|
||||
)
|
||||
.with_queue(ConcurrencyQueuePolicy {
|
||||
max_waiting: 1,
|
||||
timeout: Duration::from_secs(2),
|
||||
}),
|
||||
None,
|
||||
);
|
||||
let url = Url::parse(OFFICIAL_CODEX_BASE_URL).unwrap();
|
||||
let request = SelectCodexCredential {
|
||||
upstream_model: "gpt-5.4",
|
||||
request_url: &url,
|
||||
attempt: &request_attempt,
|
||||
session_affinity_key: Some(&key),
|
||||
};
|
||||
let mut pending = Box::pin(selector.select(&request));
|
||||
assert!(
|
||||
pending.as_mut().now_or_never().is_none(),
|
||||
"blocker={blocker}"
|
||||
);
|
||||
leases.signals.lock().unwrap().clear();
|
||||
leases.busy_accounts.lock().unwrap().clear();
|
||||
let selected = pending.await.unwrap();
|
||||
assert_eq!(selected.account_id(), &first);
|
||||
assert!(selected.affinity_hit());
|
||||
assert_eq!(selected.escape_reason(), None);
|
||||
assert!(!selected.account_switch());
|
||||
assert!(
|
||||
leases
|
||||
.requests
|
||||
.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.all(|request| request.account_id() == &first)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn affinity_queue_full_and_timeout_preserve_binding_despite_free_fallback() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_first", "test-first");
|
||||
create_account(&store, "acct_second", "test-second");
|
||||
let first = ProviderAccountId::new("acct_first").unwrap();
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
leases.busy_accounts.lock().unwrap().insert(first.clone());
|
||||
let affinity = Arc::new(MemorySessionAffinity::default());
|
||||
let key = bind_first(&affinity).await;
|
||||
let selector = selector_with_affinity(&store, leases.clone(), affinity.clone());
|
||||
let head_attempt = queued_attempt(Duration::from_secs(2));
|
||||
let next_attempt = queued_attempt(Duration::from_secs(2));
|
||||
let url = Url::parse(OFFICIAL_CODEX_BASE_URL).unwrap();
|
||||
let request = SelectCodexCredential {
|
||||
upstream_model: "gpt-5.4",
|
||||
request_url: &url,
|
||||
attempt: &head_attempt,
|
||||
session_affinity_key: Some(&key),
|
||||
};
|
||||
let mut head = Box::pin(selector.select(&request));
|
||||
assert!(head.as_mut().now_or_never().is_none());
|
||||
// 槽位已空,但新请求不能越过旧等待者,也不能转而抢占另一账号。
|
||||
leases.busy_accounts.lock().unwrap().clear();
|
||||
let rejected = selector
|
||||
.select(&SelectCodexCredential {
|
||||
attempt: &next_attempt,
|
||||
..request
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(
|
||||
rejected,
|
||||
CredentialSelectionError::QueueRejected(QueueRejection::Full)
|
||||
));
|
||||
drop(head);
|
||||
|
||||
leases.busy_accounts.lock().unwrap().insert(first.clone());
|
||||
let timeout_attempt = queued_attempt(Duration::from_millis(20));
|
||||
let timeout = selector
|
||||
.select(&SelectCodexCredential {
|
||||
attempt: &timeout_attempt,
|
||||
..request
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(
|
||||
timeout,
|
||||
CredentialSelectionError::QueueRejected(QueueRejection::Timeout)
|
||||
));
|
||||
assert_eq!(
|
||||
affinity
|
||||
.load(&ProviderKind::new("openai").unwrap(), &key)
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(first.clone())
|
||||
);
|
||||
leases.busy_accounts.lock().unwrap().clear();
|
||||
let selected = selector
|
||||
.select(&SelectCodexCredential {
|
||||
attempt: &next_attempt,
|
||||
..request
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(selected.account_id(), &first);
|
||||
assert!(
|
||||
leases
|
||||
.requests
|
||||
.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.all(|request| request.account_id() == &first)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn affinity_queue_rechecks_hard_unavailability_before_switching_accounts() {
|
||||
for quota_exhausted in [false, true] {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_first", "test-first");
|
||||
create_account(&store, "acct_second", "test-second");
|
||||
let first = ProviderAccountId::new("acct_first").unwrap();
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
leases.busy_accounts.lock().unwrap().insert(first.clone());
|
||||
let affinity = Arc::new(MemorySessionAffinity::default());
|
||||
let key = bind_first(&affinity).await;
|
||||
let selector = selector_with_affinity(&store, leases, affinity);
|
||||
let request_attempt = queued_attempt(Duration::from_secs(2));
|
||||
let url = Url::parse(OFFICIAL_CODEX_BASE_URL).unwrap();
|
||||
let request = SelectCodexCredential {
|
||||
upstream_model: "gpt-5.4",
|
||||
request_url: &url,
|
||||
attempt: &request_attempt,
|
||||
session_affinity_key: Some(&key),
|
||||
};
|
||||
let mut pending = Box::pin(selector.select(&request));
|
||||
assert!(pending.as_mut().now_or_never().is_none());
|
||||
let account = store.get_account(&first).await.unwrap().unwrap();
|
||||
if quota_exhausted {
|
||||
persist_quota_exhaustion(&store, &account, None);
|
||||
} else {
|
||||
persist_credential_state(&store, &account, CredentialState::Expired);
|
||||
}
|
||||
let selected = pending.await.unwrap();
|
||||
assert_eq!(selected.account_id().as_str(), "acct_second");
|
||||
assert!(selected.account_switch());
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn queue_enabled_without_affinity_still_uses_a_free_account() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_first", "test-first");
|
||||
create_account(&store, "acct_second", "test-second");
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
leases.signals.lock().unwrap().insert(
|
||||
ProviderAccountId::new("acct_first").unwrap(),
|
||||
busy_signal(2, None),
|
||||
);
|
||||
let selector = selector(&store, leases);
|
||||
let request_attempt = queued_attempt(Duration::from_secs(2));
|
||||
let url = Url::parse(OFFICIAL_CODEX_BASE_URL).unwrap();
|
||||
let selected = selector
|
||||
.select(&SelectCodexCredential {
|
||||
upstream_model: "gpt-5.4",
|
||||
request_url: &url,
|
||||
attempt: &request_attempt,
|
||||
session_affinity_key: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(selected.account_id().as_str(), "acct_second");
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Scheduler {
|
||||
explicit: bool,
|
||||
}
|
||||
|
||||
impl RequestPolicyPlan for Scheduler {
|
||||
fn route_model(
|
||||
&self,
|
||||
_: ModelRouteInput,
|
||||
) -> BoxFuture<'static, Result<ModelRouteDecision, RequestPolicyFault>> {
|
||||
Box::pin(async { Ok(ModelRouteDecision::Unhandled) })
|
||||
}
|
||||
|
||||
fn schedule_account(
|
||||
&self,
|
||||
_: AccountScheduleInput,
|
||||
) -> BoxFuture<'static, Result<AccountScheduleDecision, RequestPolicyFault>> {
|
||||
let decision = if self.explicit {
|
||||
AccountScheduleDecision::Pick(ProviderAccountId::new("acct_second").unwrap())
|
||||
} else {
|
||||
AccountScheduleDecision::Delegate
|
||||
};
|
||||
Box::pin(async move { Ok(decision) })
|
||||
}
|
||||
}
|
||||
|
||||
struct ExtensionLease;
|
||||
|
||||
impl ExtensionSetLease for ExtensionLease {
|
||||
fn is_ready(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn affinity_wait_respects_explicit_policy_choices_but_not_delegated_fallbacks() {
|
||||
for explicit in [false, true] {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_first", "test-first");
|
||||
create_account(&store, "acct_second", "test-second");
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
leases.signals.lock().unwrap().insert(
|
||||
ProviderAccountId::new("acct_first").unwrap(),
|
||||
busy_signal(2, None),
|
||||
);
|
||||
let affinity = Arc::new(MemorySessionAffinity::default());
|
||||
let key = bind_first(&affinity).await;
|
||||
let selector = selector_with_affinity(&store, leases.clone(), affinity);
|
||||
let request_attempt = capacity_attempt(
|
||||
account_policy().with_queue(ConcurrencyQueuePolicy {
|
||||
max_waiting: 1,
|
||||
timeout: Duration::from_secs(2),
|
||||
}),
|
||||
Some(RequestPolicyContext::new(
|
||||
Arc::new(Scheduler { explicit }),
|
||||
ExtensionSetReference::new(
|
||||
ExtensionSetId::new("capacity-test".into()).unwrap(),
|
||||
Arc::new(ExtensionLease),
|
||||
),
|
||||
ModelRequestId::new("req_capacity_contract").unwrap(),
|
||||
ClientApiKeyId::new("key_codex_contract").unwrap(),
|
||||
vec![],
|
||||
)),
|
||||
);
|
||||
let url = Url::parse(OFFICIAL_CODEX_BASE_URL).unwrap();
|
||||
let request = SelectCodexCredential {
|
||||
upstream_model: "gpt-5.4",
|
||||
request_url: &url,
|
||||
attempt: &request_attempt,
|
||||
session_affinity_key: Some(&key),
|
||||
};
|
||||
let mut pending = Box::pin(selector.select(&request));
|
||||
if explicit {
|
||||
let selected = pending.as_mut().now_or_never().unwrap().unwrap();
|
||||
assert_eq!(selected.account_id().as_str(), "acct_second");
|
||||
} else {
|
||||
assert!(pending.as_mut().now_or_never().is_none());
|
||||
leases.signals.lock().unwrap().clear();
|
||||
let selected = pending.await.unwrap();
|
||||
assert_eq!(selected.account_id().as_str(), "acct_first");
|
||||
}
|
||||
}
|
||||
}
|
||||
+2
@@ -1,3 +1,5 @@
|
||||
mod capacity;
|
||||
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::num::NonZeroU32;
|
||||
use std::sync::Arc;
|
||||
@@ -60,7 +60,7 @@ fn quota_service(store: &Arc<MemoryAccountStore>) -> CodexCredentialQuotaService
|
||||
)
|
||||
}
|
||||
|
||||
fn quota_service_with_base_url(
|
||||
pub(super) fn quota_service_with_base_url(
|
||||
store: &Arc<MemoryAccountStore>,
|
||||
http: reqwest::Client,
|
||||
base_url: String,
|
||||
@@ -638,6 +638,9 @@ async fn quota_endpoint_auth_rejections_do_not_reclassify_credentials_or_quota()
|
||||
create_account(&store, &account_id).await;
|
||||
let account = store.account(&account_id).expect("account");
|
||||
persist_quota_state(&store, &account, exhausted_quota(None)).await;
|
||||
let before = store
|
||||
.account(&account_id)
|
||||
.expect("account before rejection");
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/codex/usage"))
|
||||
@@ -658,11 +661,65 @@ async fn quota_endpoint_auth_rejections_do_not_reclassify_credentials_or_quota()
|
||||
Err(CodexCredentialQuotaError::Upstream { .. })
|
||||
));
|
||||
let current = store.account(&account_id).expect("preserved account");
|
||||
assert_eq!(current.credential_state(), CredentialState::Ready);
|
||||
assert_eq!(current.quota().access(), QuotaAccessState::Exhausted);
|
||||
assert_eq!(current, before);
|
||||
assert_eq!(
|
||||
service.synchronize().await.expect("quota cycle").transient,
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
store
|
||||
.account(&account_id)
|
||||
.expect("account after quota cycle"),
|
||||
before
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_endpoint_rejection_preserves_details_without_changing_account_facts() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let account_id = "acct_quota_rejection_detail";
|
||||
create_account(&store, account_id).await;
|
||||
let account = store.account(account_id).expect("account");
|
||||
persist_quota_state(&store, &account, exhausted_quota(None)).await;
|
||||
let before = store.account(account_id).expect("account before rejection");
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/codex/usage"))
|
||||
.respond_with(ResponseTemplate::new(401).set_body_json(json!({
|
||||
"error": {
|
||||
"code": "token_revoked",
|
||||
"message": "Encountered invalidated oauth token for user, failing request"
|
||||
}
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let service = quota_service_with_base_url(
|
||||
&store,
|
||||
reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.build()
|
||||
.expect("client"),
|
||||
server.uri(),
|
||||
);
|
||||
|
||||
match service.refresh_account(account.id()).await {
|
||||
Err(CodexCredentialQuotaError::Upstream { status, code, .. }) => {
|
||||
assert_eq!(status, Some(401));
|
||||
assert_eq!(code.as_deref(), Some("token_revoked"));
|
||||
}
|
||||
other => panic!("expected upstream rejection, got {other:?}"),
|
||||
}
|
||||
// 额度查询诊断不能进入凭据快照,否则会干扰在途 OAuth 刷新的终态提交。
|
||||
let current = store.account(account_id).expect("account after refresh");
|
||||
assert_eq!(current, before);
|
||||
store
|
||||
.repository()
|
||||
.load_runtime_credential(&before)
|
||||
.await
|
||||
.expect("unchanged credential");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn payment_required_is_authoritative_quota_exhaustion_without_fabricated_usage() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user